MolecularDiffusion.modules.models.painn_dynamics

Dense <-> flat adapter binding the ported PaiNN/OM-Diff EquivNet backbone to the dynamics._forward(t, xh, node_mask, edge_mask, context) contract that EnVariationalDiffusion.phi (modules/models/en_diffusion.py:163) calls.

Mirrors gfmdiff/dynamics.py::GFMDiffDynamics in role, but has real work to do: EnVariationalDiffusion speaks dense padded batches (B, N, ...) while EquivNet speaks a flat concatenation of graphs (n, ...) with an (E, 2) edge list, so this wrapper packs and unpacks around it.

It also owns the pieces OM-Diff kept outside the backbone in their AtomisticModel input/output modules: the linear one-hot embedding, the Fourier time features, and the scalar readout MLP.

Classes

FourierTimeFeatures

Random Fourier expansion of the diffusion time (om-diff

PaiNNDynamics

Denoising network: PaiNN scalar+vector backbone, EDM interface.

Module Contents

class MolecularDiffusion.modules.models.painn_dynamics.FourierTimeFeatures(n_features: int = 16, std: float = 1.0, trainable: bool = False)

Bases: torch.nn.Module

Random Fourier expansion of the diffusion time (om-diff layers/features.py::FourierFeatures). Output is 2 * n_features wide; the projection is a fixed buffer unless trainable.

forward(t: torch.Tensor) → torch.Tensor
class MolecularDiffusion.modules.models.painn_dynamics.PaiNNDynamics(in_node_nf: int, context_node_nf: int = 0, n_dims: int = 3, num_interactions: int = 5, node_size: int = 256, edge_size: int = 64, embedding_dim: int = 256, rbf_features: int = 64, rbf_max_distance: float = 5.0, time_features: int = 16, cutoff: float | None = None, envelop_p: int | None = 6)

Bases: torch.nn.Module

Denoising network: PaiNN scalar+vector backbone, EDM interface.

Parameters:

  • in_node_nf – Node feature channels the diffusion model expects back (atom-type one-hot + atomic number [+ extra values]), excluding time and context, which are added internally.

  • context_node_nf – Conditioning channels, concatenated to the node features before embedding.

  • n_dims – Spatial dimensions (3).

  • num_interactions – Interaction/update blocks.

  • node_size – backbone widths. OM-Diff’s defaults are 256 / 64 / 256.

  • edge_size – backbone widths. OM-Diff’s defaults are 256 / 64 / 256.

  • embedding_dim – backbone widths. OM-Diff’s defaults are 256 / 64 / 256.

  • rbf_features – Gaussian radial basis size and range, in angstrom.

  • rbf_max_distance – Gaussian radial basis size and range, in angstrom.

  • time_features – Fourier time features; contributes 2 x this to the backbone input width.

  • cutoff – Optional edge cutoff in angstrom. None (default) keeps the platform’s dense fully-connected edge_mask; set e.g. 7.5 to reproduce OM-Diff’s radius graph, which they rebuild from the noisy coordinates at every step — this does the same, since it is applied per forward pass.

  • envelop_p – If cutoff is set, order of the polynomial cutoff envelope applied to the radial basis. None disables it.

abstractmethod forward(t, xh, node_mask, edge_mask, context=None)
unwrap_forward()

Return the unbound forward.

wrap_forward(node_mask, edge_mask, context)

Bind the masks so an ODE/SDE solver can call fwd(t, x).

context_node_nf = 0
cutoff = None
embedding
equivnet
in_node_nf
n_dims = 3
readout
time_embedding