Commit 21f6d329 authored by Bharath Ramsundar's avatar Bharath Ramsundar
Browse files

Task Splitter implementation

parent c17663cb

deepchem/splits/__init__.py

+1 −3
Original line number Diff line number Diff line
@@ -41,7 +41,7 @@ class Splitter(object): 
    """Creates splitter object."""

    self.verbosity = verbosity



  def k_fold_split(self, dataset, directories, compute_feature_statistics=True):

  def k_fold_split(self, dataset, directories=None, compute_feature_statistics=True):

    """Does K-fold split of dataset."""

    log("Computing K-fold split", self.verbosity)

    k = len(directories)
@@ -59,8 +59,6 @@ class Splitter(object): 
      fold_dataset = rem_dataset.select( 

          fold_dir, fold_inds,

          compute_feature_statistics=compute_feature_statistics)

      # TODO(rbharath): Is making a tempfile the best way to handle remainders?

      # Would be  nice to be able to do in memory dataset construction...

      rem_dir = tempfile.mkdtemp()

      rem_dataset = rem_dataset.select( 

          rem_dir, rem_inds,

deepchem/splits/task_splitter.py

+49 −16
Original line number Diff line number Diff line
@@ -18,6 +18,28 @@ from deepchem.datasets import NumpyDataset 
from deepchem.featurizers.featurize import load_data

from deepchem.splits import Splitter



def merge_fold_datasets(fold_datasets):

  """Merges fold datasets together.



  Assumes that fold_datasets were outputted from k_fold_split. Specifically,

  assumes that each dataset contains the same datapoints, listed in the same

  ordering.

  """

  if not len(fold_datasets):

    return None



  # All datasets share features and identifiers by assumption.

  X = fold_datasets[0].X

  ids = fold_datasets[0].ids



  ys, ws = [], []

  for fold_dataset in fold_datasets:

    ys.append(fold_dataset.y)

    ws.append(fold_dataset.w)

  y = np.concatenate(ys, axis=1)

  w = np.concatenate(ws, axis=1)

  return NumpyDataset(X, y, w, ids)



class TaskSplitter(Splitter):

  """

  Provides a simple interface for splitting datasets task-wise.
@@ -48,16 +70,11 @@ class TaskSplitter(Splitter): 
      Proportion of tasks to be put into test. Rounded to nearest int.

    """

    n_tasks = len(dataset.get_task_names())

    n_train = np.round(frac_train * n_tasks)

    n_valid = np.round(frac_valid * n_tasks)

    n_test = np.round(frac_test * n_tasks)

    n_train = int(np.round(frac_train * n_tasks))

    n_valid = int(np.round(frac_valid * n_tasks))

    n_test = int(np.round(frac_test * n_tasks))

    if n_train + n_valid + n_test != n_tasks:

      raise ValueError("Train/Valid/Test fractions don't split tasks evenly.")

    ########################################### DEBUG

    print("train_valid_test_split")

    print("n_train, n_valid, n_test")

    print(n_train, n_valid, n_test)

    ########################################### DEBUG



    X, y, w, ids = dataset.X, dataset.y, dataset.w, dataset.ids
@@ -66,12 +83,28 @@ class TaskSplitter(Splitter): 
        X, y[:,n_train:n_train+n_valid], w[:,n_train:n_train+n_valid], ids)

    test_dataset = NumpyDataset(

        X, y[:,n_train+n_valid:], w[:,n_train+n_valid:], ids)

    ########################################### DEBUG

    print("train_dataset.get_task_names()")

    print(train_dataset.get_task_names())

    print("valid_dataset.get_task_names()")

    print(valid_dataset.get_task_names())

    print("test_dataset.get_task_names()")

    print(test_dataset.get_task_names())

    ########################################### DEBUG

    return train_dataset, valid_dataset, test_dataset



  def k_fold_split(self, dataset, K):

    """Performs a K-fold split of the tasks for dataset.



    Parameters

    ----------

    dataset: deepchem.datasets.Dataset

      Dataset to be split

    K: int

      Number of splits to be made

    """

    n_tasks = len(dataset.get_task_names())

    n_per_fold = int(np.round(n_tasks/float(K)))

    if K * n_per_fold != n_tasks:

      raise ValueError("Cannot perform a valid %d-way split" % K)

    

    X, y, w, ids = dataset.X, dataset.y, dataset.w, dataset.ids



    fold_datasets = []

    for fold in range(K):

      fold_tasks = range(fold*n_per_fold, (fold+1)*n_per_fold)

      fold_datasets.append(

          NumpyDataset(X, y[:, fold_tasks], w[:, fold_tasks], ids))

    return fold_datasets

deepchem/splits/tests/test_task_splitter.py

+56 −4
Original line number Diff line number Diff line
@@ -13,6 +13,7 @@ __license__ = "GPL" 
import tempfile

import numpy as np

from deepchem.splits.task_splitter import TaskSplitter

from deepchem.splits.task_splitter import merge_fold_datasets

from deepchem.datasets import NumpyDataset

from deepchem.datasets.tests import TestDatasetAPI
@@ -33,10 +34,6 @@ class TestTaskSplitters(TestDatasetAPI): 
    p = .05 # proportion actives

    y = np.random.binomial(1, p, size=(n_samples, n_tasks))

    dataset = NumpyDataset(X, y)

    ########################################### DEBUG

    print("dataset")

    print(dataset)

    ########################################### DEBUG



    task_splitter = TaskSplitter()

    train, valid, test = task_splitter.train_valid_test_split(
@@ -45,3 +42,58 @@ class TestTaskSplitters(TestDatasetAPI): 
    assert len(train.get_task_names()) == 4

    assert len(valid.get_task_names()) == 3

    assert len(test.get_task_names()) == 3



  def test_multitask_K_fold_split(self):

    """

    Test TaskSplitter K-fold split on multitask dataset.

    """

    n_samples = 100

    n_features = 10

    n_tasks = 10

    X = np.random.rand(n_samples, n_features)

    p = .05 # proportion actives

    y = np.random.binomial(1, p, size=(n_samples, n_tasks))

    dataset = NumpyDataset(X, y)

    K = 5



    task_splitter = TaskSplitter()

    fold_datasets = task_splitter.k_fold_split(dataset, K)



    for fold_dataset in fold_datasets:

      assert len(fold_dataset.get_task_names()) == 2



  def test_merge_fold_datasets(self):

    """

    Test that (K-1) folds can be merged into train dataset.

    """

    n_samples = 100

    n_features = 10

    n_tasks = 10

    X = np.random.rand(n_samples, n_features)

    p = .05 # proportion actives

    y = np.random.binomial(1, p, size=(n_samples, n_tasks))

    w = np.ones((n_samples, n_tasks))

    dataset = NumpyDataset(X, y, w)

    K = 5



    task_splitter = TaskSplitter()

    fold_datasets = task_splitter.k_fold_split(dataset, K)

    # Number tasks per fold

    n_per_fold = 2



    for fold in range(K):

      train_inds = list(set(range(K)) - set([fold]))

      train_fold_datasets = [fold_datasets[ind] for ind in train_inds]

      train_dataset = merge_fold_datasets(train_fold_datasets)



      # Find the tasks that correspond to this test fold

      train_tasks = list(

          set(range(10)) - set(range(fold*n_per_fold, (fold+1)*n_per_fold)))



      # Assert that all arrays look like they should

      np.testing.assert_array_equal(train_dataset.X, X)

      np.testing.assert_array_equal(

          train_dataset.y, y[:, train_tasks])

      np.testing.assert_array_equal(

          train_dataset.w, w[:, train_tasks])

      np.testing.assert_array_equal(train_dataset.X, X)

CRA Git | Maintained and supported by SUSTech CRA and CCSE