MolecularDiffusion.modules.layers.graph_module¶

Classes¶

`BackboneInterface`
`GraphData`	Class to keep graph attributes nicely packaged.
`GraphModelMixin`	Mixin Model class implementing some general convenience properties and methods.
`HeadInterface`

Functions¶

`compute_neighbors`(data, edge_index)
`get_max_neighbors_mask`(natoms, index, atom_distance, ...)	Give a mask that filters out edges so that each atom has at most
`get_pbc_distances`(pos, edge_index, cell, cell_offsets, ...)
`radius_graph_pbc`(data, radius, max_num_neighbors_threshold)

Module Contents¶

class MolecularDiffusion.modules.layers.graph_module.BackboneInterface¶

abstractmethod forward(data: MolecularDiffusion.modules.layers.atomic_data.AtomicData) → dict[str, torch.Tensor]¶

Backbone forward.

Parameters:: data (AtomicData) – Atomic systems as input
Returns:: embedding – Return backbone embeddings for the given input
Return type:: dict[str->torch.Tensor]

class MolecularDiffusion.modules.layers.graph_module.GraphData¶

Class to keep graph attributes nicely packaged.

atomic_numbers_full: torch.Tensor¶

batch_full: torch.Tensor¶

cell_offsets: torch.Tensor¶

edge_distance: torch.Tensor¶

edge_distance_vec: torch.Tensor¶

edge_index: torch.Tensor¶

neighbors: torch.Tensor¶

node_offset: int = 0¶

offset_distances: torch.Tensor¶

class MolecularDiffusion.modules.layers.graph_module.GraphModelMixin¶

Mixin Model class implementing some general convenience properties and methods.

generate_graph(data, cutoff=None, max_neighbors=None, use_pbc=None, otf_graph=None, enforce_max_neighbors_strictly=None, use_pbc_single=False, fully_connected=False)¶

no_weight_decay() → list¶: Returns a list of parameters with no weight decay.

property num_params: int¶

class MolecularDiffusion.modules.layers.graph_module.HeadInterface¶

abstractmethod forward(data: MolecularDiffusion.modules.layers.atomic_data.AtomicData, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]¶

Head forward.

Parameters:

data (AtomicData) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

property use_amp¶

MolecularDiffusion.modules.layers.graph_module.compute_neighbors(data, edge_index)¶

MolecularDiffusion.modules.layers.graph_module.get_max_neighbors_mask(natoms, index, atom_distance, max_num_neighbors_threshold, degeneracy_tolerance: float = 0.01, enforce_max_strictly: bool = False)¶

Give a mask that filters out edges so that each atom has at most max_num_neighbors_threshold neighbors. Assumes that index is sorted.

Enforcing the max strictly can force the arbitrary choice between degenerate edges. This can lead to undesired behaviors; for example, bulk formation energies which are not invariant to unit cell choice.

A degeneracy tolerance can help prevent sudden changes in edge existence from small changes in atom position, for example, rounding errors, slab relaxation, temperature, etc.

MolecularDiffusion.modules.layers.graph_module.get_pbc_distances(pos, edge_index, cell, cell_offsets, neighbors, return_offsets: bool = False, return_distance_vec: bool = False)¶

MolecularDiffusion.modules.layers.graph_module.radius_graph_pbc(data, radius, max_num_neighbors_threshold, enforce_max_neighbors_strictly: bool = False, pbc=None)¶