Commit a4f345a6 authored by leswing's avatar leswing
Browse files

Splitter code

parent 7fb141b5
Loading
Loading
Loading
Loading
+8 −2
Original line number Diff line number Diff line
@@ -53,7 +53,7 @@ class Splitter(object):
    """Creates splitter object."""
    self.verbose = verbose

  def k_fold_split(self, dataset, k, directories=None):
  def k_fold_split(self, dataset, k, directories=None, **kwargs):
    """Does K-fold split of dataset."""
    log("Computing K-fold split", self.verbose)
    if directories is None:
@@ -332,7 +332,13 @@ class SingletaskStratifiedSplitter(Splitter):
    self.task_number = task_number
    self.verbose = verbose

  def k_fold_split(self, dataset, k, seed=None, log_every_n=None):
  def k_fold_split(self,
                   dataset,
                   k,
                   directories=None,
                   seed=None,
                   log_every_n=None,
                   **kwargs):
    """
        Splits compounds into k-folds using stratified sampling.
        Overriding base class k_fold_split.