MolecularDiffusion.modules.tasks.diffusion_midi

MiDi task: joint diffusion over coordinates, atom types, bonds and charges.

MiDi (Vignac et al., ECML 2023, arXiv:2302.09048) is the first model in this platform that generates the molecular graph itself – bond orders and formal charges are diffused jointly with the 3D coordinates, so a sample arrives with an explicit bond table instead of needing post-hoc perception.

Data path: data_type: graph3d with bond_collate: dense. graph3d_dense_collate already produces MiDi’s exact dense shapes, so the adapter below is a handful of F.one_hot calls and upstream’s utils.to_dense (the PyG -> dense bridge) is not needed at all.

Bond classes are the platform’s canonical five (0=none, 1=SINGLE, 2=DOUBLE, 3=TRIPLE, 4=AROMATIC) and MiDi uses exactly those, in the same order – the mapping is the identity.

Formal charges are stored raw and signed; the offset and class count are applied here (charge_offset/n_charge_classes: QM9 -> +1/3, GEOM -> +2/6), never baked into the dataset.

Out of scope this pass (see the integration plan): ExtraFeatures (every released config sets extra_features: null), the variational-NLL validation path, MiDi’s own molecular metrics, and the size-aware loader.

Property-conditioning + classifier-free guidance use the same config signature as configs/tasks/diffusion.yaml (en_diffusion’s formulation): set condition_names to opt in. The injection point is PlaceHolder.y, MiDi’s own per-graph global feature, already FiLM-injected into every transformer layer – see modules/models/midi/transformer_model.py.

Attributes

Chem
logger

Classes

MidiDiffusionTask	MiDi wrapped in the platform's duck-typed Task contract (§2.1).
ModelTaskFactory	Hydra entry point for the MiDi task (configs/tasks/diffusion_midi).

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_midi.MidiDiffusionTask(atom_vocab: list, n_layers: int, hidden_mlp_dims: dict, hidden_dims: dict, diffusion_steps: int, diffusion_noise_schedule: str, transition: str, nu: dict, lambda_train: list, charge_offset: int, n_charge_classes: int, sdf_output_path: str | None, x_marginals: torch.Tensor | None, e_marginals: torch.Tensor | None, charges_marginals: torch.Tensor | None, n_atoms_hist: dict, condition_names: list = [], context_mask_rate: float = 0.0, mask_value: float = 0.0, normalize_condition: str | None = None, adapter_conditions: list | None = None, use_adapter_module: bool = False)

Bases: torch.nn.Module

MiDi wrapped in the platform’s duck-typed Task contract (§2.1).

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

Reduce accumulated losses into the logged validation metric.

forward(batch: dict) → tuple[torch.Tensor, dict]

One training step: noise the batch, denoise it, weight the losses.

predict_and_target(batch: dict) → tuple[torch.Tensor, torch.Tensor]

Pure-generative stub: the loss is both prediction and target.

preprocess(train_set=None, valid_set=None, test_set=None)

Build self.property_norms for CFG conditioning (train-side only).

Called generically by cli/train.py if this attribute exists. Does NOT touch node_dist_model/n_node_dist – those come from graph3d_stats at __init__ time via ModelTaskFactory, a separate mechanism. Deliberately skips DistributionProperty/prop_dist_model (out of scope – generation always takes an explicit target_value).

sample(batch_size: int | None = None, nodesxsample: torch.Tensor | None = None, num_steps: int | None = None, batch: dict | None = None, mode: str | None = None, n_frames: int = 0, condition: torch.Tensor | None = None, negative_condition: torch.Tensor | None = None, cfg_scale: float = 0.0, cfg_scale_schedule: str | None = None, **kwargs: Any) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Sampling, unconditional unless condition/cfg_scale are set.

Returns the platform’s (one_hot, charges, coords, node_mask) tuple. charges carries signed formal charges (FlowMol’s precedent for this slot) – the element identity is in one_hot. The generated bond matrix has no channel in that tuple, so it is stashed on self.last_bond_types and, when sdf_output_path is set, written as an .sdf sidecar alongside the platform’s .xyz.

sample_guidance_conitional(target_function: Any = None, target_value: list | None = None, negative_target_value: list | None = None, nodesxsample: torch.Tensor | None = None, cfg_scale: float = 1, cfg_scale_schedule: str | None = None, guidance_ver: str = 'cfg', n_frames: int = 0, num_steps: int | None = None) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Classifier-free-guidance generation.

Matches the call signature GenerativeFactory.conditional_generation() hardcodes for task_type == "cfg", and returns (one_hot, charges, x, node_mask) like sample().

T

atom_vocab

backbone

charge_offset

condition = []

context_mask_rate = 0.0

property device: torch.device

Device of the backbone parameters.

input_dims

lambda_train

last_bond_types: torch.Tensor | None = None

mask_value = 0.0

property model: MidiDiffusionTask

tasks_generate.py reads task.model.T; self is the model.

n_adapter_context

n_atom_types

n_charge_classes

n_concat_context

property n_node_dist: dict

{n_atoms: count} histogram used to clamp mol_size.

node_dist_model

normalize_condition = None

output_dims

prop_dist_model = None

property_norms = None

sdf_output_path

task_type = 'diffusion_midi'

class MolecularDiffusion.modules.tasks.diffusion_midi.ModelTaskFactory(task_type: str = 'diffusion_midi', n_layers: int = 12, hidden_mlp_dims: dict | None = None, hidden_dims: dict | None = None, diffusion_steps: int = 500, diffusion_noise_schedule: str = 'cosine', transition: str = 'marginal', nu: dict | None = None, lambda_train: list | None = None, charge_offset: int = 1, n_charge_classes: int = 3, sdf_output_path: str | None = None, atom_vocab: list | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs: Any)

Hydra entry point for the MiDi task (configs/tasks/diffusion_midi).

Declares train_set so cli/train.py’s declarative seam injects the training dataset: MiDi’s noise model needs the atom/bond/charge marginals and the size histogram at construction time, and neither is an nn.Module buffer, so neither survives in a checkpoint.

sdf_output_path is declared generation-time (docs §2.5b): the task is rebuilt from the checkpoint’s training-time config, where it is null, so without this declaration the generate config’s value never arrives and the bond sidecar – the whole 2D half of the model – is silently dropped.

build() → MidiDiffusionTask

Construct the task, deriving marginals/size histogram from data.

atom_vocab

charge_offset = 1

diffusion_noise_schedule = 'cosine'

diffusion_steps = 500

generation_time_keys = ('sdf_output_path',)

hidden_dims

hidden_mlp_dims

kwargs

lambda_train = [3, 0.4, 1, 2, 0]

n_charge_classes = 3

n_layers = 12

nu

sdf_output_path = None

task: MidiDiffusionTask | None = None

task_type = 'diffusion_midi'

train_set = None

transition = 'marginal'

MolecularDiffusion.modules.tasks.diffusion_midi.Chem = None

MolecularDiffusion.modules.tasks.diffusion_midi.logger