MolecularDiffusion.modules.models.ditmc.meshgraphnet

MeshGraphNet bond-graph conditioner.

Port of neural_network_layer.MeshGraphNetLayer + encoder.EncoderModel. This is the only place the covalent bonds enter as message passing; the DiT attention itself runs over the fully connected latent graph.

Classes

MeshGraphNetEncoder	make_graph_mesh_net_encoder: embeddings + a stack of layers.
MeshGraphNetLayer	One edge update followed by one node update, both LayerNorm-then-MLP.

Module Contents

class MolecularDiffusion.modules.models.ditmc.meshgraphnet.MeshGraphNetEncoder(node_attr_dim: int, num_layers: int, num_features: int, activation_fn: str = 'silu')

Bases: torch.nn.Module

make_graph_mesh_net_encoder: embeddings + a stack of layers.

forward(graph: MolecularDiffusion.modules.models.ditmc.graphs.CondGraph) → torch.Tensor

Returns node features in e3x form, (N, 1, 1, num_features).

edge_embedding

layers

node_embedding

num_features

class MolecularDiffusion.modules.models.ditmc.meshgraphnet.MeshGraphNetLayer(num_node_features: int, num_edge_features: int, activation_fn: str = 'silu')

Bases: torch.nn.Module

One edge update followed by one node update, both LayerNorm-then-MLP.

forward(graph: MolecularDiffusion.modules.models.ditmc.graphs.CondGraph, node_features: torch.Tensor, edge_features: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

edge_mlp

edge_norm

node_mlp

node_norm