Commit 1689dba8 authored by Yutong Zhao's avatar Yutong Zhao
Browse files

MaxMin Splitter implementation.

parent f37eb55b
Loading
Loading
Loading
Loading
+66 −1
Original line number Diff line number Diff line
@@ -18,8 +18,8 @@ from rdkit import Chem
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
from rdkit.SimDivFilters.rdSimDivPickers import MaxMinPicker
import deepchem as dc
from deepchem.data import DiskDataset
from deepchem.utils import ScaffoldGenerator
@@ -597,6 +597,71 @@ class MolecularWeightSplitter(Splitter):
            sortidx[valid_cutoff:])


class MaxMinSplitter(Splitter):
  """
  Class for doing splits based on the MaxMin diversity algorithm. Intuitively,
  the test set is comprised of the most diverse compounds of the entire dataset.
  Furthermore, the validation set is comprised of diverse compounds under
  the test set.
  """

  def split(self,
            dataset,
            seed=None,
            frac_train=.8,
            frac_valid=.1,
            frac_test=.1,
            log_every_n=None):
    """
    Splits internal compounds randomly into train/validation/test.
    """
    np.testing.assert_almost_equal(frac_train + frac_valid + frac_test, 1.)
    if not seed is None:
      np.random.seed(seed)

    num_datapoints = len(dataset)

    train_cutoff = int(frac_train * num_datapoints)
    valid_cutoff = int((frac_train + frac_valid) * num_datapoints)

    num_train = train_cutoff
    num_valid = valid_cutoff - train_cutoff
    num_test = num_datapoints - valid_cutoff

    all_mols = []
    for ind, smiles in enumerate(dataset.ids):
      all_mols.append(Chem.MolFromSmiles(smiles))

    fps = [AllChem.GetMorganFingerprintAsBitVect(x, 2, 1024) for x in all_mols]

    def distance(i, j):
      return 1 - DataStructs.DiceSimilarity(fps[i], fps[j])

    picker = MaxMinPicker()
    testIndices = picker.LazyPick(
        distFunc=distance, poolSize=num_datapoints, pickSize=num_test, seed=23)

    validTestIndices = picker.LazyPick(
        distFunc=distance,
        poolSize=num_datapoints,
        pickSize=num_valid + num_test,
        firstPicks=testIndices,
        seed=23)

    allSet = set(range(num_datapoints))
    testSet = set(testIndices)
    validSet = set(validTestIndices) - testSet

    trainSet = allSet - testSet - validSet

    assert len(testSet & validSet) == 0
    assert len(testSet & trainSet) == 0
    assert len(validSet & trainSet) == 0
    assert (validSet | trainSet | testSet) == allSet

    return sorted(list(trainSet)), sorted(list(validSet)), sorted(list(testSet))


class RandomSplitter(Splitter):
  """
    Class for doing random data splits.
+18 −6
Original line number Diff line number Diff line
@@ -136,16 +136,28 @@ class TestSplitters(unittest.TestCase):
        [train_data, valid_data, test_data])
    assert sorted(merged_dataset.ids) == (sorted(solubility_dataset.ids))

  def test_singletask_maxmin_split(self):
    """
    Test singletask MaxMinSplitter class.
    """
    solubility_dataset = dc.data.tests.load_butina_data()
    maxmin_splitter = dc.splits.MaxMinSplitter()
    train_data, valid_data, test_data = \
      maxmin_splitter.train_valid_test_split(
        solubility_dataset)
    assert len(train_data) == 8
    assert len(valid_data) == 1
    assert len(test_data) == 1

  def test_singletask_butina_split(self):
    """
    Test singletask ScaffoldSplitter class.
    Test singletask ButinaSplitter class.
    """
    solubility_dataset = dc.data.tests.load_butina_data()
    scaffold_splitter = dc.splits.ButinaSplitter()
    butina_splitter = dc.splits.ButinaSplitter()
    train_data, valid_data, test_data = \
      scaffold_splitter.train_valid_test_split(
      butina_splitter.train_valid_test_split(
        solubility_dataset)
    print(len(train_data), len(valid_data))
    assert len(train_data) == 7
    assert len(valid_data) == 3
    assert len(test_data) == 0
@@ -292,7 +304,7 @@ class TestSplitters(unittest.TestCase):
    y[:n_positives] = 1
    w = np.ones((n_samples, n_tasks))
    # Set half the positives to have zero weight
    w[:n_positives / 2] = 0
    w[:n_positives // 2] = 0
    ids = np.arange(n_samples)

    stratified_splitter = dc.splits.RandomStratifiedSplitter()
@@ -340,7 +352,7 @@ class TestSplitters(unittest.TestCase):
    y = np.random.binomial(1, p, size=(n_samples, n_tasks))
    w = np.ones((n_samples, n_tasks))
    # Mask half the examples
    w[:n_samples / 2] = 0
    w[:n_samples // 2] = 0

    stratified_splitter = dc.splits.RandomStratifiedSplitter()
    split_indices = stratified_splitter.get_task_split_indices(