gt4sd.algorithms.generation.diffusion.geodiff.model.utils module

Summary

Functions:

assemble_atom_pair_feature
clip_norm
extend_graph_order_radius
get_distance
graph_field_network	Transformation to make the epsilon predicted from the diffusion model roto-translational equivariant.
is_local_edge
repeat_batch	type batch Batch
repeat_data	type data Data
set_rdmol_positions	param rdkit_mol An rdkit.Chem.rdchem.Mol object.
set_rdmol_positions_	param rdkit_mol An rdkit.Chem.rdchem.Mol object.

Reference

assemble_atom_pair_feature(node_attr, edge_index, edge_attr)

extend_graph_order_radius(num_nodes, pos, edge_index, edge_type, batch, order=3, cutoff=10.0, extend_order=True, extend_radius=True, is_sidechain=None)

get_distance(pos, edge_index)

graph_field_network(score_d, pos, edge_index, edge_length)

Transformation to make the epsilon predicted from the diffusion model roto-translational equivariant. See equations 5-7 of the GeoDiff Paper https://arxiv.org/pdf/2203.02923.pdf

clip_norm(vec, limit, p=2)

is_local_edge(edge_type)

repeat_data(data, num_repeat)

Parameters

data (Data) – An torch_geometric.data.Data object.

Returns

A copy of data repetead num_repeat times.

Return type

batch

repeat_batch(batch, num_repeat)

Parameters

batch (Batch) – An torch_geometric.data.Batch object.

Returns

A copy of batch repetead num_repeat times.

Return type

batch

set_rdmol_positions(rdkit_mol, pos)

Parameters

• rdkit_mol – An rdkit.Chem.rdchem.Mol object.

• pos – (N_atoms, 3)

Returns

A copy of rdkit_mol with the positions set to pos.

Return type

mol

set_rdmol_positions_(mol, pos)

Parameters

• rdkit_mol – An rdkit.Chem.rdchem.Mol object.

• pos – (N_atoms, 3)

Returns

A copy of rdkit_mol with the positions set to pos.

Return type

mol