MolecularDiffusion.modules.models.flowmol_graph3d.priors¶

Priors for the FlowMol3 / CTMC port.

Ported from FlowMol (flowmol/data_processing/priors.py). Only the two the released FlowMol3 config reaches are here:

centered_normal_prior_batched_graph – zero-COM Gaussian for x.
ctmc_masked_prior / ctmc_masked_edge_prior – the all-mask categorical prior for a/c/e.

The marginal, c-given-a, biased, uniform-sample and barycenter priors are deliberately not ported: FlowMol3’s own configure_prior (flowmol/models/flowmol.py:190-193) raises NotImplementedError unless all three categorical priors are ctmc.

Per-item OT / rigid prior alignment (prior_config.x.align: true) is also not ported – it is a dataloader-__getitem__ operation upstream and the platform has no per-item prior hook. It has no state-dict impact, and upstream does no alignment at inference either (flowmol.py:533), so sampling from the released weights is unaffected; only training efficiency is traded.

Functions¶

`centered_normal_prior_batched_graph`(→ torch.Tensor)	Per-molecule zero-COM Gaussian prior for atom positions.
`ctmc_masked_edge_prior`(→ torch.Tensor)	All-mask bond prior, mirrored across the two edge directions.
`ctmc_masked_prior`(→ torch.Tensor)	All-mask categorical prior: every one of `n` rows is the mask token.

Module Contents¶

MolecularDiffusion.modules.models.flowmol_graph3d.priors.centered_normal_prior_batched_graph(g: dgl.DGLGraph, node_batch_idx: torch.Tensor, std: float = 1.0) → torch.Tensor¶: Per-molecule zero-COM Gaussian prior for atom positions.

MolecularDiffusion.modules.models.flowmol_graph3d.priors.ctmc_masked_edge_prior(upper_edge_mask: torch.Tensor, n_bond_types: int) → torch.Tensor¶

All-mask bond prior, mirrored across the two edge directions.

The upper-triangle sample is written to both the upper and the mirrored lower triangle so the two directions of each bond always agree – FlowMol enforces bond symmetry structurally, by this mirroring, not by assertion.

MolecularDiffusion.modules.models.flowmol_graph3d.priors.ctmc_masked_prior(n: int, d: int) → torch.Tensor¶

All-mask categorical prior: every one of n rows is the mask token.

Returns a (n, d + 1) one-hot whose last column (index d, the mask index) is set. d is the number of real classes – the mask token is an internal noise state, never a chemistry class, and never appears in an output head.