torchdrug.transforms

class Compose(transforms)

Compose a list of transforms into one.

Parameters

transforms (list of callable) – list of transforms

Graph Transformations

NormalizeTarget

class NormalizeTarget(mean, std)

Normalize the target values in a sample.

Parameters

• mean (dict of float) – mean of targets

• std (dict of float) – standard deviation of targets

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.

Parameters

• shuffle_node (bool, optional) – shuffle node order or not

• shuffle_edge (bool, optional) – shuffle edge order or not

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.

Parameters

atom_types (array_like) – vocabulary of atom types

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.

Parameters

• atom_type (int, optional) – type of the virtual atom

• bond_type (int, optional) – type of the virtual bonds

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

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

• kwargs – other attributes of the virtual atoms or virtual bonds

Protein Transformations

TruncateProtein

class TruncateProtein(max_length=None, random=False, keys='graph')

Truncate over long protein sequences into a fixed length.

Parameters

• max_length (int, optional) – maximal length of the sequence. Truncate the sequence if it exceeds this limit.

• random (bool, optional) – truncate the sequence at a random position. If not, truncate the suffix of the sequence.

• keys (str or list of str, optional) – keys for the items that require truncation in a sample

ProteinView

class ProteinView(view, keys='graph')

Convert proteins to a specific view.

Parameters

• view (str) – protein view. Can be atom or residue.

• keys (str or list of str, optional) – keys for the items that require view change in a sample