Cleanup

      raise ValueError("No data in dataset.")

    return next(self.metadata_df.iterrows())[1]['task_names']

  def reshard(self, shard_size):

    """Reshards data to have specified shard size."""

    # Create temp directory to store resharded version

    reshard_dir = tempfile.mkdtemp()

    new_metadata = []

    # Write data in new shards

    ind = 0

    tasks = self.get_task_names()

    X_next = np.zeros((0,) + self.get_data_shape())

    y_next = np.zeros((0,) + (len(tasks),))

    w_next = np.zeros((0,) + (len(tasks),))

    ids_next = np.zeros((0,), dtype=object)

    for (X, y, w, ids) in self.itershards():

      X_next = np.vstack([X_next, X])

      y_next = np.vstack([y_next, y])

      w_next = np.vstack([w_next, w])

      ids_next = np.concatenate([ids_next, ids])

      while len(X_next) > shard_size:

        X_batch, X_next = X_next[:shard_size], X_next[shard_size:]

        y_batch, y_next = y_next[:shard_size], y_next[shard_size:]

        w_batch, w_next = w_next[:shard_size], w_next[shard_size:]

        ids_batch, ids_next = ids_next[:shard_size], ids_next[shard_size:]

        new_basename = "reshard-%d" % ind

        new_metadata.append(DiskDataset.write_data_to_disk(

            reshard_dir, new_basename, tasks, X_batch, y_batch, w_batch, ids_batch))

        ind += 1

    # Handle spillover from last shard

    new_basename = "reshard-%d" % ind

    new_metadata.append(DiskDataset.write_data_to_disk(

        reshard_dir, new_basename, tasks, X_next, y_next, w_next, ids_next))

    ind += 1

    # Get new metadata rows

    resharded_dataset = DiskDataset(

        data_dir=reshard_dir, tasks=tasks, metadata_rows=new_metadata,

        verbosity=self.verbosity)

    shutil.rmtree(self.data_dir)

    shutil.move(reshard_dir, self.data_dir)

    self.metadata_df = resharded_dataset.metadata_df

    self.save_to_disk()

  def get_data_shape(self):

    """

    Gets array shape of datapoints in this dataset.

    loader = dc.data.DataLoader(

        tasks=emols_tasks, smiles_field="smiles",

        featurizer=featurizer)

    dataset = loader.featurize(dataset_file, debug=True, logging=False)

    dataset = loader.featurize(dataset_file)

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

    print("ids.shape, X.shape, y.shape, w.shape")

    assert y_orig.shape == y_new.shape

    assert w_orig.shape == w_new.shape

  def test_reshard_shuffle(self):

    """Test that datasets can be merged."""

    current_dir = os.path.dirname(os.path.realpath(__file__))

    dataset_file = os.path.join(

        current_dir, "../../models/tests/example.csv")

    featurizer = dc.feat.CircularFingerprint(size=1024)

    tasks = ["log-solubility"]

    loader = dc.data.DataLoader(

        tasks=tasks, smiles_field="smiles", featurizer=featurizer)

    dataset = loader.featurize(

        dataset_file, shard_size=2)

    X_orig, y_orig, w_orig, orig_ids = (dataset.X, dataset.y, dataset.w,

                                        dataset.ids)

    orig_len = len(dataset)

    dataset.reshard_shuffle(reshard_size=1)

    X_new, y_new, w_new, new_ids = (dataset.X, dataset.y, dataset.w,

                                    dataset.ids)

    assert len(dataset) == orig_len

    # The shuffling should have switched up the ordering

    assert not np.array_equal(orig_ids, new_ids)

    # But all the same entries should still be present

    assert sorted(orig_ids) == sorted(new_ids)

    # All the data should have same shape

    assert X_orig.shape == X_new.shape

    assert y_orig.shape == y_new.shape

    assert w_orig.shape == w_new.shape

  def test_shuffle_each_shard(self):

    """Test that shuffle_each_shard works."""

    n_samples = 100

class ComplexFeaturizer(object):

  """"

  Abstract class for calculating features for mol/protein complexes.

  Class Attributes

  ----------------

  name : str or list

      Name (or names) of this featurizer (used for scripting).

  """

  name = None

  def featurize_complexes(self, mol_pdbs, protein_pdbs, verbose=True,

                          log_every_n=1000):

      List of PDBs for proteins. Each PDB should be a list of lines of the

      PDB file.

    """

    features = []

    for i, (mol_pdb, protein_pdb) in enumerate(zip(mol_pdbs, protein_pdbs)):

      if verbose and i % log_every_n == 0:

  Abstract class for calculating a set of features for a molecule.

  Child classes implement the _featurize method for calculating features

  for a single molecule. The feature matrix returned by featurize has a

  shape that is inferred from the shape of the features returned by

  _featurize, and is either indexed by molecules or by molecules and

  conformers depending on the value of the conformers class attribute.

  Class Attributes

  ----------------

  conformers : bool, optional (default False)

      Whether features are calculated for conformers. If True, the first

      two axes of the feature matrix will index molecules and conformers,

      respectively. If False, only molecule-level features are calculated

      and the feature matrix will not have a separate conformer dimension.

      This is a class attribute because some featurizers take 3D

      conformation into account and others do not, and this is not

      typically an instance-level decision.

  name : str or list

      Name (or names) of this featurizer (used for scripting).

  topo_view : bool (default False)

      Whether the calculated features represent a topological view of the

      data.

  for a single molecule.

  """

  conformers = False

  name = None

  topo_view = False

  def featurize(self, mols, verbose=True, log_every_n=1000):

    """

      else:

        features.append(np.array([]))

    if self.conformers:

      features = self.conformer_container(mols, features)

    else:

    features = np.asarray(features)

    return features

    """

    return self.featurize(mols)

  def conformer_container(self, mols, features):

    """

    Put features into a container with an extra dimension for

    conformers. Conformer indices that are not used are masked.

    For example, if mols contains 3 molecules with 1, 2, 5 conformers,

    respectively, then the final container will have 3 entries on its

    first axis and 5 entries on its second axis. The remaining axes

    correspond to feature dimensions.

    Parameters

    ----------

    mols : iterable

        RDKit Mol objects.

    features : list

        Features calculated for molecule conformers. Each element

        corresponds to the features for a molecule and should be an

        ndarray with conformers on the first axis.

    """

    # get the maximum number of conformers

    max_confs = max([mol.GetNumConformers() for mol in mols])

    if not max_confs:

      max_confs = 1

    # construct the new container

    # - first axis = # mols

    # - second axis = max # conformers

    # - remaining axes = determined by feature shape

    features_shape = None

    for i in range(len(features)):

      for j in range(len(features[i])):

        if features[i][j] is not None:

          features_shape = features[i][0].shape

          break

      if features_shape is not None:

        break

    if features_shape is None:

      raise ValueError('Cannot find any features.')

    shape = (len(mols), max_confs) + features_shape

    x = np.ma.masked_all(shape)

    # fill in the container

    for i, (mol, mol_features) in enumerate(zip(mols, features)):

      n_confs = max(mol.GetNumConformers(), 1)

      try:

        x[i, :n_confs] = mol_features

      except ValueError:  # handle None conformer values

        for j in range(n_confs):

          x[i, j] = mol_features[j]

    return x

class UserDefinedFeaturizer(Featurizer):

  """Directs usage of user-computed featurizations."""

  def __init__(self, feature_fields):

    """Creates user-defined-featurizer."""

    self.feature_fields = feature_fields

    input_file = os.path.join(current_dir, "user_specified_example.csv")

    loader = dc.data.DataLoader(

        tasks=tasks, smiles_field="smiles", featurizer=featurizer)

    dataset = loader.featurize(input_file, debug=True)

    dataset = loader.featurize(input_file)

    splitter = dc.splits.SpecifiedSplitter(input_file, "split")

    train_dataset, test_dataset = splitter.train_test_split(dataset)