MolecularDiffusion.modules.models.oareactdiff.graph_tools¶

Graph + batched-tensor helpers for OA-ReactDiff’s multi-object diffusion.

Ported verbatim from oa_reactdiff/utils/_graph_tools.py and oa_reactdiff/diffusion/_utils.py (commit 543aaa8, MIT).

The one idea worth knowing before reading anything else in this package: a “reaction” is three separate node sets (reactant, transition state, product) flat-concatenated into one tensor, with two index vectors over the rows –

combined_mask which sample in the batch a node belongs to n_frag_switch which object (0=R, 1=TS, 2=P) a node belongs to

Edges are fully connected within a sample (get_edges_index), and get_subgraph_mask then marks which of those edges stay inside one object. That mask is what makes the network “object aware”: the centre of gravity is removed per object, not per sample, so each of R/TS/P keeps its own SE(3) frame.

Functions¶

`assert_mean_zero_with_mask`(x, node_mask[, eps])
`cdf_standard_gaussian`(x)
`get_edges_index`(→ torch.Tensor)
`get_mask_for_frag`(→ torch.Tensor)	Get fragment index for each node
`get_n_frag_switch`(→ torch.Tensor)	Get the type of fragments to which each node belongs
`get_subgraph_mask`(→ torch.Tensor)	Filter out edges that have inter-fragment connections.
`remove_mean_batch`(x, indices)
`sample_center_gravity_zero_gaussian_batch`(→ torch.Tensor)
`sample_gaussian`(size, device)
`sum_except_batch`(x, indices, dim_size)

Module Contents¶

MolecularDiffusion.modules.models.oareactdiff.graph_tools.assert_mean_zero_with_mask(x, node_mask, eps=1e-10)¶

MolecularDiffusion.modules.models.oareactdiff.graph_tools.cdf_standard_gaussian(x)¶

MolecularDiffusion.modules.models.oareactdiff.graph_tools.get_edges_index(combined_mask: torch.Tensor, pos: torch.Tensor | None = None, edge_cutoff: float | None = None, remove_self_edge: bool = False) → torch.Tensor¶

Parameters:

combined_mask (Tensor) – Combined mask for all fragments. Edges are built for nodes with the same indexes in the mask.
pos (Optional[Tensor]) – 3D coordinations of nodes. Defaults to None.
edge_cutoff (Optional[float]) – cutoff for building edges within a fragment. Defaults to None.
remove_self_edge (bool) – whether to remove self-connecting edge (i.e., ii). Defaults to False.

Returns:

[2, n_edges], i for node index.

Return type:

Tensor

MolecularDiffusion.modules.models.oareactdiff.graph_tools.get_mask_for_frag(natm: torch.Tensor) → torch.Tensor¶

Get fragment index for each node Example: Tensor([2, 0, 3]) -> [0, 0, 2, 2, 2]

Parameters:: natm (Tensor) – number of nodes per small fragment
Returns:: [n_node], the natural index of fragment a node belongs to
Return type:: Tensor

MolecularDiffusion.modules.models.oareactdiff.graph_tools.get_n_frag_switch(natm_list: List[torch.Tensor]) → torch.Tensor¶

Get the type of fragments to which each node belongs Example: [Tensor(1, 1), Tensor(2, 1)] -> [0, 0, 1, 1 ,1]

Parameters:: natm_list (List[Tensor]) – [Tensor([number of atoms per small fragment])]
Returns:: [n_nodes], type of fragment each node belongs to
Return type:: Tensor

MolecularDiffusion.modules.models.oareactdiff.graph_tools.get_subgraph_mask(edge_index: torch.Tensor, n_frag_switch: torch.Tensor) → torch.Tensor¶

Filter out edges that have inter-fragment connections. Example: edge_index: [

[0, 0, 1, 1, 2, 2], [1, 2, 0, 2, 0, 1], ]

n_frag_switch: [0, 0, 1] -> [1, 0, 1, 0, 0, 0]

Parameters:

edge_index (Tensor) – e_ij
n_frag_switch (Tensor) – fragment that a node belongs to

Returns:

[n_edge], 1 for inner- and 0 for inter-fragment edge

Return type:

Tensor

MolecularDiffusion.modules.models.oareactdiff.graph_tools.remove_mean_batch(x, indices)¶

MolecularDiffusion.modules.models.oareactdiff.graph_tools.sample_center_gravity_zero_gaussian_batch(size: List[int], indices: List[torch.Tensor]) → torch.Tensor¶

MolecularDiffusion.modules.models.oareactdiff.graph_tools.sample_gaussian(size, device)¶

MolecularDiffusion.modules.models.oareactdiff.graph_tools.sum_except_batch(x, indices, dim_size)¶