torchdrug.transforms

class Compose(transforms)

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

Graph Transformations

TargetNormalize

class TargetNormalize(mean, std)

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

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

Shuffle

class Shuffle(shuffle_node=True, shuffle_edge=True)

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)

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)

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)

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