drugforge.ml.dataset.GraphDataset

class drugforge.ml.dataset.GraphDataset(*args: Any, **kwargs: Any)[source]

Bases: Dataset

Class for loading SMILES as graphs.

__init__(compounds={}, structures=[], random_iter=False)[source]

Methods

__init__([compounds, structures, random_iter])

from_exp_compounds(exp_compounds[, ...])

from_ligands(ligands[, exp_dict, ...])

classmethod from_exp_compounds(exp_compounds, exp_dict: dict = {}, node_featurizer=None, edge_featurizer=None, random_iter=False)[source]
Parameters:
  • exp_compounds (List[schema.ExperimentalCompoundData]) – List of compounds

  • exp_dict (dict[str, dict[str, int | float]], optional) – Dict mapping compound_id to an experimental results dict. The dict for a compound will be added to the pose representation of each Complex containing a ligand witht that compound_id

  • node_featurizer (BaseAtomFeaturizer, optional) – Featurizer for node data

  • edge_featurizer (BaseBondFeaturizer, optional) – Featurizer for edges

  • random_iter (bool, default=False) – Iterate through the dataset randomly each time

classmethod from_ligands(ligands: list[Ligand], exp_dict: dict = {}, node_featurizer=None, edge_featurizer=None, random_iter=False)[source]
Parameters:
  • ligands (list[Ligands]) – List of Ligand schema objects to build into a GraphDataset object

  • exp_dict (dict[str, dict[str, int | float]], optional) – Dict mapping compound_id to an experimental results dict. The dict for a compound will be added to the pose representation of each Complex containing a ligand witht that compound_id

  • node_featurizer (BaseAtomFeaturizer, optional) – Featurizer for node data

  • edge_featurizer (BaseBondFeaturizer, optional) – Featurizer for edges

  • random_iter (bool, default=False) – Iterate through the dataset randomly each time