MolecularDiffusion.runmodes.train.tasks_painn

Factory wiring the ported PaiNN/OM-Diff EquivNet backbone into the existing, unmodified diffusion task – structured exactly like runmodes/train/tasks_gfmdiff.py::ModelTaskFactory (a new dynamics network dropped into EnVariationalDiffusion + GeomMolecularGenerative, zero edits to either).

Per the approved integration plan (docs/model_integrations/omdiff/INTEGRATION_PLAN.md) this is a §1 backbone swap: OM-Diff’s plain-L2 loss, metal-centre masking, conditional size prior and regressor guidance are out of scope – only the network is ported, so it is named after its architecture family (painn) rather than after the paper.

Attributes

logger

Classes

ModelTaskFactory

Build the PaiNN-backbone diffusion model + task.

Module Contents

class MolecularDiffusion.runmodes.train.tasks_painn.ModelTaskFactory(task_type: str, train_set=None, atom_vocab=None, task_names: list = [], condition_names: list = [], num_interactions: int = 5, hidden_dims: dict = {}, chkpt_path: str = None, **kwargs)

Build the PaiNN-backbone diffusion model + task.

Parameters:

  • task_type (str) – must be "diffusion".

  • train_set – unused; kept for interface parity with the other factories (the atom-count histogram is computed lazily by GeomMolecularGenerative.preprocess).

  • atom_vocab (list) – atom vocabulary used for encoding.

  • task_names (list) – conditional labels (context columns).

  • condition_names (list) – condition names for conditional generation.

  • num_interactions (int) – PaiNN interaction/update blocks.

  • hidden_dims (dict) – keys node_size, edge_size, embedding_dim for the backbone widths.

  • chkpt_path (str) – optional path to a model checkpoint.

  • **kwargs – diffusion + backbone keyword arguments, see configs/tasks/diffusion_painn.yaml.

build()

Build and return the GeomMolecularGenerative task.

atom_vocab = None
chkpt_path = None
condition_names = []
context_node_nf
hidden_dims
in_node_nf
kwargs
num_interactions = 5
task_names = []
task_type
train_set = None
MolecularDiffusion.runmodes.train.tasks_painn.logger