MolecularDiffusion.modules.tasks.diffusion_ditmc

DiTMC conformer-generation task (task_type: diffusion_ditmc).

Wraps modules/models/ditmc in the platform’s duck-typed Task contract (docs/adding_new_models.md 2.1) and adapts the graph3d PyG Batch (bond_collate: raw) into the three graphs DiTMC consumes.

This model does not invent molecules. You hand it a molecule you already have – its atoms and its bonds – and it returns 3D conformers of exactly that molecule. Nothing about the composition is generated, so:

• DiTMCTask.sample() raises. There is no unconditional mode: without a molecular graph there is nothing to place in 3D, and the platform’s (one_hot, charges, coords, node_mask) return contract has no meaning here. Generation goes through the shared ConformerFactory, pointed at by configs/interference/gen_conformer.yaml – the same route gen_diffdec_scaffold.yaml and gen_apo2mol_pocket.yaml take.

• node_dist_model / n_node_dist are deliberately absent. They exist to let GenerativeFactory choose a molecule size; DiTMC’s size is dictated by the input molecule.

Bond mapping (canonical -> DiTMC): ditmc_class = canonical_class - 1, so 1=SINGLE -> 0, 2=DOUBLE -> 1, 3=TRIPLE -> 2, 4=AROMATIC -> 3. Canonical class 0 (“no bond”) is simply never an edge of the conditioner graph, which is exactly the platform’s storage rule – a pass-through. Do not set ``kekulize: true``: upstream trains on RDKit’s aromatic perception and aromatic is a real class in its 4-wide vocabulary.

Attributes

ModelTaskFactory
N_BOND_CLASSES
logger

Classes

DiTMCTask	DiTMC in the platform's Task contract.
DiTMCTaskFactory	Factory instantiated by cli/train.py.
Graph3DBatchToDiTMCAdapter	One graph3d PyG Batch -> (LatentGraph, CondGraph, PriorGraph).

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_ditmc.DiTMCTask(variant: str, dataset: str, model_kwargs: dict, flow_kwargs: dict, prior: str = 'harmonic', augmentation_bool: bool = True, atom_vocab: list | None = None, task_type: str = 'diffusion_ditmc')

Bases: torch.nn.Module

DiTMC in the platform’s Task contract.

evaluate(pred: torch.Tensor, target: torch.Tensor)

forward(batch: Any)

generate_conformers(items: list, *, num_steps: int = 50, free_guidance_scale: float = 1.0, logarithmic_time_bool: bool = False, return_trajectory: bool = False, generator: torch.Generator | None = None)

predict_and_target(batch: Any, all_loss=None, metric=None)

sample(*args, **kwargs)

adapter

atom_vocab

augmentation_bool = True

dataset

property device: torch.device

property model: DiTMCTask

net

process

task_type = 'diffusion_ditmc'

variant

class MolecularDiffusion.modules.tasks.diffusion_ditmc.DiTMCTaskFactory(task_type: str = 'diffusion_ditmc', variant: str = 'so3', dataset: str = 'qm9', model: dict | None = None, flow: dict | None = None, prior: str = 'harmonic', augmentation_bool: bool = True, atom_vocab: list | None = None, **kwargs: Any)

Factory instantiated by cli/train.py.

train_set is deliberately not declared: DiTMC needs no dataset statistics at construction (no marginals, no valency table, no size histogram), so the declarative injection seam simply does not fire.

build() → DiTMCTask

atom_vocab

augmentation_bool = True

dataset = 'qm9'

flow_kwargs

kwargs

model_kwargs

prior = 'harmonic'

task: DiTMCTask | None = None

task_type = 'diffusion_ditmc'

variant = 'so3'

class MolecularDiffusion.modules.tasks.diffusion_ditmc.Graph3DBatchToDiTMCAdapter(dataset: str = 'qm9')

Bases: torch.nn.Module

One graph3d PyG Batch -> (LatentGraph, CondGraph, PriorGraph).

Everything derived per molecule (node_attr, shortest hops, the Laplacian eigendecomposition) is cached by SMILES, exactly as upstream caches it at dataset-build time.

The one featurization column that is derived rather than copied is the chiral tag: upstream reads it off the GEOM mol, and here it is recovered by RDKit from the stored coordinates. It is the first thing to check if a converted pretrained checkpoint underperforms.

build(items: list, device=None)

forward(batch: Any)

cache

dataset = 'qm9'

MolecularDiffusion.modules.tasks.diffusion_ditmc.ModelTaskFactory

MolecularDiffusion.modules.tasks.diffusion_ditmc.N_BOND_CLASSES = 4

MolecularDiffusion.modules.tasks.diffusion_ditmc.logger