Fingerprint splitter

from rdkit import DataStructs

from rdkit.Chem import AllChem

from rdkit.ML.Cluster import Butina

from rdkit import DataStructs

from rdkit.Chem.Fingerprints import FingerprintMols

import deepchem as dc

from deepchem.data import DiskDataset

from deepchem.utils import ScaffoldGenerator

    return train_inds, valid_inds, test_inds

class FingerprintSplitter(Splitter):

  """

    Class for doing data splits based on the fingerprints of small molecules

    O(N**2) algorithm

    """

  def split(self,

            dataset,

            frac_train=.8,

            frac_valid=.1,

            frac_test=.1,

            log_every_n=1000):

    """

        Splits internal compounds into train/validation/test by scaffold.

        """

    np.testing.assert_almost_equal(frac_train + frac_valid + frac_test, 1.)

    data_len = len(dataset)

    mols, fingerprints = [], []

    train_inds, valid_inds, test_inds = [], [], []

    for ind, smiles in enumerate(dataset.ids):

      mol = Chem.MolFromSmiles(smiles, sanitize=False)

      mols.append(mol)

      fp = FingerprintMols.FingerprintMol(mol)

      fingerprints.append(fp)

    distances = np.ones(shape=(data_len, data_len))

    for i in range(data_len):

      for j in range(data_len):

        distances[i][j] = DataStructs.FingerprintSimilarity(fingerprints[i],

                                                            fingerprints[j])

    train_cutoff = int(frac_train * len(dataset))

    valid_cutoff = int(frac_valid * len(dataset))

    # Pick the mol closest to everything as the first element of training

    closest_ligand = np.argmin(np.sum(distances, axis=1))

    train_inds.append(closest_ligand)

    cur_distances = [float('inf')] * data_len

    self.update_distances(closest_ligand, cur_distances, distances, train_inds)

    for i in range(1, train_cutoff):

      closest_ligand = np.argmin(cur_distances)

      train_inds.append(closest_ligand)

      self.update_distances(closest_ligand, cur_distances, distances,

                            train_inds)

    # Pick the closest mol from what is left

    index, best_dist = 0, float('inf')

    for i in range(data_len):

      if i in train_inds:

        continue

      dist = np.sum(distances[i])

      if dist < best_dist:

        index, best_dist = i, dist

    valid_inds.append(index)

    leave_out_indexes = train_inds + valid_inds

    cur_distances = [float('inf')] * data_len

    self.update_distances(index, cur_distances, distances, leave_out_indexes)

    for i in range(1, valid_cutoff):

      closest_ligand = np.argmin(cur_distances)

      valid_inds.append(closest_ligand)

      leave_out_indexes.append(closest_ligand)

      self.update_distances(closest_ligand, cur_distances, distances,

                            leave_out_indexes)

    # Test is everything else

    for i in range(data_len):

      if i in leave_out_indexes:

        continue

      test_inds.append(i)

    print(train_inds, valid_inds, test_inds)

    return train_inds, valid_inds, test_inds

  def update_distances(self, last_selected, cur_distances, distance_matrix,

                       dont_update):

    for i in range(len(cur_distances)):

      if i in dont_update:

        cur_distances[i] = float('inf')

        continue

      new_dist = distance_matrix[i][last_selected]

      if new_dist < cur_distances[i]:

        cur_distances[i] = new_dist

class SpecifiedSplitter(Splitter):

  """

    Class that splits data according to user specification.