MolecularDiffusion.runmodes.train.tasks_gfmdiff

Factory wiring the ported GFMDiff Dual-Track Transformer backbone into the existing, unmodified diffusion task – structured like runmodes/train/tasks_egt.py::ModelTaskFactory (a new dynamics network dropped into EnVariationalDiffusion + GeomMolecularGenerative, zero edits to either).

Per the approved integration plan (revision 2), this is a §1 backbone swap: GFMDiff’s own noise schedule/VLB loss/sampler and the Geometric-Facilitated Loss (+ diffused “degree” channel) are out of scope – only the Dual-Track Transformer architecture (GFMDiffDynamics / EquiGNN) is ported, and modules.models.en_diffusion.EnVariationalDiffusion / modules.tasks.diffusion.GeomMolecularGenerative are reused unmodified.

Attributes

logger

Classes

ModelTaskFactory

Factory to build the GFMDiff-backbone diffusion model + task.

Module Contents

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

Factory to build the GFMDiff-backbone diffusion model + task.

Parameters:

  • task_type (str) – must be “diffusion”.

  • train_set – Training dataset, used to infer input node feature dimensions (kept for interface parity with other factories; unused directly here).

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

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

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

  • num_layers (int) – Number of Dual-Track Transformer blocks.

  • hidden_dims (dict) – expects keys “emb_dim”, “hidden_dim”, “num_heads” for the GFMDiff backbone.

  • chkpt_path (str) – Optional path to model checkpoint.

  • **kwargs – diffusion + backbone-specific keyword arguments (see diffusion_gfmdiff.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_layers = 5
task_names = []
task_type
train_set = None
MolecularDiffusion.runmodes.train.tasks_gfmdiff.logger