torchdrug.transforms

class Compose(transforms)[source]

Compose a list of transforms into one. :param transforms: list of transforms :type transforms: list of callable

Graph Transformations

TargetNormalize

class TargetNormalize(mean, std)[source]

Normalize the target values in a sample. :param mean: mean of targets :type mean: dict of float :param std: standard deviation of targets :type std: dict of float

RandomBFSOrder

class RandomBFSOrder[source]

Order the nodes in a graph according to a random BFS order.

Shuffle

class Shuffle(shuffle_node=True, shuffle_edge=True)[source]

Shuffle the order of nodes and edges in a graph. :param shuffle_node: shuffle node order or not :type shuffle_node: bool, optional :param shuffle_edge: shuffle edge order or not :type shuffle_edge: bool, optional

VirtualNode

class VirtualNode(relation=None, weight=1, node_feature=None, edge_feature=None, **kwargs)[source]

Add a virtual node and connect it with every node in the graph.

Parameters
  • relation (int, optional) – relation of virtual edges. By default, use the maximal relation in the graph plus 1.

  • weight (int, optional) – weight of virtual edges

  • node_feature (array_like, optional) – feature of the virtual node

  • edge_feature (array_like, optional) – feature of virtual edges

  • kwargs – other attributes of the virtual node or virtual edges

Molecule Transformations

RemapAtomType

class RemapAtomType(atom_types)[source]

Map atom types to their index in a vocabulary. Atom types that don’t present in the vocabulary are mapped to -1. :param atom_types: vocabulary of atom types :type atom_types: array_like

VirtualAtom

class VirtualAtom(atom_type=None, bond_type=None, node_feature=None, edge_feature=None, **kwargs)[source]

Add a virtual atom and connect it with every atom in the molecule. :param atom_type: type of the virtual atom :type atom_type: int, optional :param bond_type: type of the virtual bonds :type bond_type: int, optional :param node_feature: feature of the virtual atom :type node_feature: array_like, optional :param edge_feature: feature of virtual bonds :type edge_feature: array_like, optional :param kwargs: other attributes of the virtual atoms or virtual bonds