MolecularDiffusion.modules.models.painn_backbone

PaiNN-style scalar+vector equivariant backbone, ported from OM-Diff.

Source: om-diff @ 976084eaf407cc84a5f169f4204f72d4a1fbdbbc (https://github.com/Aalto-QuML/om-diff — MIT), files src/models/backbones/equivnet.py, src/models/layers/{rbf,norm, readout,mlp,features}.py and src/models/ops.py.

Only the network is ported. OM-Diff’s diffusion objective, metal-centre masking, conditional size prior and sampler are deliberately not part of this file — the platform’s EnVariationalDiffusion supplies all of that (see docs/model_integrations/omdiff/INTEGRATION_PLAN.md).

Two deliberate deviations from upstream:

• EquivNet.forward takes plain tensors instead of OM-Diff’s Batch dataclass, so no part of their data model has to come along.

• EdgeLayer / with_edge_interactions is not ported — it is off in every shipped OM-Diff config and nothing else here would use it.

Classes

BesselRBFLayer	Bessel radial basis (upstream alternative to the Gaussian one).
EnvelopLayer	Polynomial cutoff envelope, smoothly zero beyond xc.
EquivNet	OM-Diff's scalar+vector message-passing network.
EquivNetHParams	Hyperparameters of EquivNet.
GaussianLinearRBFLayer	Evenly spaced Gaussians — OM-Diff's configured default.
GraphLayerNorm	Layer norm whose statistics are pooled per graph, not per node.

Functions

pairwise_distances(→ tuple[torch.Tensor, torch.Tensor])

Edge distances (E, 1) and unit vectors (E, 3), i -> j.

Module Contents

class MolecularDiffusion.modules.models.painn_backbone.BesselRBFLayer(n_features: int = 20, max_distance: float = 5.0, trainable: bool = True)

Bases: RBFLayer

Bessel radial basis (upstream alternative to the Gaussian one).

forward(distances: torch.Tensor) → torch.Tensor

prefactor

r_max

class MolecularDiffusion.modules.models.painn_backbone.EnvelopLayer(p: int = 6, xc: float = 5.0)

Bases: torch.nn.Module

Polynomial cutoff envelope, smoothly zero beyond xc.

forward(distances: torch.Tensor) → torch.Tensor

class MolecularDiffusion.modules.models.painn_backbone.EquivNet(hparams: EquivNetHParams, rbf_layer: RBFLayer | None = None, envelop_layer: EnvelopLayer | None = None)

Bases: torch.nn.Module

OM-Diff’s scalar+vector message-passing network.

Unlike upstream this consumes/returns tensors rather than their Batch dataclass, so it can be driven straight from the dense diffusion batch by PaiNNDynamics.

forward(node_positions: torch.Tensor, node_states: torch.Tensor, edge_index: torch.Tensor, num_nodes: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Run the network over one flat (concatenated) batch of graphs.

Parameters:

• node_positions – (n, 3) coordinates.

• node_states – (n, input_size) embedded node features.

• edge_index – (E, 2) directed edges, [:, 0] -> [:, 1].

• num_nodes – (B,) nodes per graph; must sum to n.

Returns:

(delta_node_positions (n, 3), node_states (n, node_size)).

edge_featurizers

envelop_layer = None

hp

interactions

project_layer

rbf_layer = None

updates

class MolecularDiffusion.modules.models.painn_backbone.EquivNetHParams

Hyperparameters of EquivNet.

edge_size: int = 64

input_size: int = 288

node_size: int = 256

num_interactions: int = 5

update_node_positions: bool = True

class MolecularDiffusion.modules.models.painn_backbone.GaussianLinearRBFLayer(n_features: int = 64, max_distance: float = 5.0, min_distance: float = 0.0)

Bases: RBFLayer

Evenly spaced Gaussians — OM-Diff’s configured default.

forward(distances: torch.Tensor) → torch.Tensor

class MolecularDiffusion.modules.models.painn_backbone.GraphLayerNorm(in_channels: int, eps: float = 1e-06, affine: bool = True)

Bases: torch.nn.Module

Layer norm whose statistics are pooled per graph, not per node.

forward(x: torch.Tensor, splits: torch.Tensor) → torch.Tensor

eps = 1e-06

MolecularDiffusion.modules.models.painn_backbone.pairwise_distances(positions: torch.Tensor, edges: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Edge distances (E, 1) and unit vectors (E, 3), i -> j.

edges is (E, 2) (OM-Diff’s layout, not PyG’s (2, E)).