MolecularDiffusion.modules.models.ditmc.graphs¶

The three graphs DiTMC consumes, as plain dataclasses.

Upstream splits every sample into three jraph.GraphsTuple``s (``data_loader/utils.py:504-531). Nothing about jraph is needed here beyond the field layout, so these are dataclasses of flat tensors instead.

Two conventions inherited from upstream that silently produce wrong models if reversed:

``senders = j``, ``receivers = i``. create_graph_latent does receivers=centers, senders=others where centers is the row index. DiTEdgeEmbed and RadialSphericalEdgeEmbedding then compute displacements = positions[senders] - positions[receivers], i.e. \(\vec r_j - \vec r_i\). Reversing this flips the sign of every odd-l spherical harmonic.
The all-pairs edge list is in C-order over ``(i, j)``, which is the only reason shortest_hops (extracted as M[~eye]) lines up with it element-for-element.

Classes¶

`CondGraph`	The covalent-bond graph the MeshGraphNet conditioner runs over.
`LatentGraph`	The fully connected graph the DiT attention runs over.
`PriorGraph`	Bond-graph Laplacian eigendecomposition, for the harmonic prior.

Module Contents¶

class MolecularDiffusion.modules.models.ditmc.graphs.CondGraph¶

The covalent-bond graph the MeshGraphNet conditioner runs over.

edge_attr: torch.Tensor¶

node_attr: torch.Tensor¶

receivers: torch.Tensor¶

senders: torch.Tensor¶

class MolecularDiffusion.modules.models.ditmc.graphs.LatentGraph¶

The fully connected graph the DiT attention runs over.

replace(**kwargs) → LatentGraph¶

atomic_numbers: torch.Tensor¶

batch_segments: torch.Tensor¶

cond_scaling_edges: torch.Tensor¶

cond_scaling_nodes: torch.Tensor¶

node_attr: torch.Tensor¶

num_graphs: int¶

property num_nodes: int¶

positions: torch.Tensor¶

receivers: torch.Tensor¶

self_cond: torch.Tensor | None = None¶

senders: torch.Tensor¶

shortest_hops: torch.Tensor¶

x1: torch.Tensor | None = None¶

class MolecularDiffusion.modules.models.ditmc.graphs.PriorGraph¶

Bond-graph Laplacian eigendecomposition, for the harmonic prior.

Stored in upstream’s sparse form: node_attr is \(1/\sqrt{\lambda}\) per eigen-index (with the zero modes set to 0, which is what removes the centre of mass), and edge_attr is the flattened eigenvector matrix over a complete per-molecule index grid, so a segment sum evaluates \(P\,\mathrm{diag}(1/\sqrt\lambda)\,z\).

edge_attr: torch.Tensor¶

node_attr: torch.Tensor¶

receivers: torch.Tensor¶

senders: torch.Tensor¶