MolecularDiffusion.modules.tasks.diffusion_jodo

JODO task: joint continuous diffusion over coordinates, atom types, formal charges and an explicit bond tensor.

JODO (Huang et al., NeurIPS 2023, arXiv:2305.12347) diffuses the 2D molecular graph and the 3D geometry together in one continuous VP process: a sample arrives with real bond orders and formal charges, no geometry-based perception needed. What separates it from MiDi – the platform’s other bond-generating model – is that bonds are continuous here, not categorical: a bond is two (or three) float channels that get thresholded back to a class at decode time.

Data path: data_type: graph3d with bond_collate: dense. graph3d_dense_collate already emits the padded (B,N,N) integer bond matrix, so the adapter below is a handful of vectorized lines – upstream’s collate_edge / EdgeComCondTransform are not needed at all.

Bond mapping (identity on the class ids, then JODO’s compressed encoding):

canonical   JODO class   channels (exist, order[, aromatic])
0 = none    0            (0, 0   [, 0])
1 = SINGLE  1            (1, 1/3 [, 0])
2 = DOUBLE  2            (1, 2/3 [, 0])
3 = TRIPLE  3            (1, 1   [, 0])
4 = AROMATIC 4           (1, 0   ,  1)     -- only when edge_ch == 3

edge_ch: 2 (QM9) has no aromatic channel, so the data config MUST set kekulize: true; an aromatic bond would otherwise silently become “no bond” (upstream datasets/build_dataset.py:151). edge_ch: 3 (GEOM) keeps aromatic and takes kekulize: false.

Formal charges are the one place JODO differs from MiDi: there is no categorical charge head. The raw signed charge goes into a single float channel divided by fc_charge_norm and is rounded back at decode.

Conditional generation rides the platform’s existing GenerativeFactory.conditional_generation seam – condition: [gap] selects Cond_DGT_concat and this class implements sample_conditonal.

Out of scope this pass (see the integration plan): multi-property conditioning, the EGNN property classifiers (an evaluation tool, not the model), 2D-only mode, DPM-solver fast sampling, and trajectories/inpainting.

Attributes

Chem
logger

Classes

JodoDiffusionTask	JODO wrapped in the platform's duck-typed Task contract (§2.1).
ModelTaskFactory	Hydra entry point for JODO (configs/tasks/diffusion_jodo.yaml).

Functions

expand_dims(→ torch.Tensor)	v of shape [N] -> [N, 1, ..., 1] with dims dimensions.
get_align_position(→ torch.Tensor)	Rotate the clean positions onto the noisy ones (losses.py:403).
kabsch_batch(→ torch.Tensor)	Batched Kabsch rotation, port of upstream losses.py:424.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_jodo.JodoDiffusionTask(atom_vocab: list, nf: int, n_layers: int, n_heads: int, n_extra_heads: int, dropout: float, mlp_ratio: int, spatial_cut_off: float, edge_ch: int, include_fc_charge: bool, cond_time: bool, dist_gbf: bool, gbf_name: str, trans_name: str, softmax_inf: bool, edge_quan_th: float, com: bool, pred_data: bool, self_cond: bool, noise_align: bool, centered: bool, normalize_factors: list, loss_weights: list, reduce_mean: bool, noise_schedule: str, beta_0: float, beta_1: float, sampling_steps: int, condition: list, normalize_condition: str, sdf_output_path: str | None, n_atoms_hist: dict, prop_dist_model: MolecularDiffusion.modules.models.en_diffusion.DistributionProperty | None = None, property_norms: dict | None = None)

Bases: torch.nn.Module

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

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

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

One training step: port of losses.py:285-382.

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

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

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, context: torch.Tensor | None = None, **kwargs: Any) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Ancestral sampling; port of sampling.py:518-597 + post_process.

Returns the platform’s (one_hot, charges, coords, node_mask). charges carries signed formal charges (FlowMol/MiDi precedent). The generated bond matrix has no slot 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_conditonal(nodesxsample: torch.Tensor = None, target_value: list | None = None, n_frames: int = 0, mode: str | None = None, num_steps: int | None = None, **kwargs: Any)

Property-conditional sampling.

Called by GenerativeFactory.conditional_generation (runmodes/generate/tasks_generate.py:536) – the misspelling is the platform’s, not a typo here. Turns raw target values into the normalized context the conditional backbone was trained on.

T

atom_vocab

backbone

centered

condition

property device: torch.device

edge_ch

include_fc_charge

last_bond_types: torch.Tensor | None = None

loss_weights

property model: JodoDiffusionTask

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

n_atom_types

property n_node_dist: dict

node_dist_model

noise_align

noise_scheduler

normalize_condition

pred_data

prop_dist_model = None

property_norms = None

reduce_mean

sdf_output_path

self_cond

class MolecularDiffusion.modules.tasks.diffusion_jodo.ModelTaskFactory(task_type: str = 'diffusion_jodo', nf: int = 256, n_layers: int = 8, n_heads: int = 16, n_extra_heads: int = 2, dropout: float = 0.1, mlp_ratio: int = 2, spatial_cut_off: float = 2.0, edge_ch: int = 2, include_fc_charge: bool = True, cond_time: bool = True, dist_gbf: bool = True, gbf_name: str = 'CondGaussianLayer', trans_name: str = 'TransMixLayer', softmax_inf: bool = True, edge_quan_th: float = 0.0, com: bool = True, pred_data: bool = True, self_cond: bool = True, noise_align: bool = True, centered: bool = True, normalize_factors: list | None = None, loss_weights: list | None = None, reduce_mean: bool = False, noise_schedule: str = 'cosine', beta_0: float = 0.1, beta_1: float = 20.0, sampling_steps: int = 1000, condition: list | None = None, normalize_condition: str = 'mad', 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 JODO (configs/tasks/diffusion_jodo.yaml).

Declares train_set so cli/train.py’s declarative seam injects the training dataset: the molecule-size histogram (and, for the conditional variant, the property distribution + mean/MAD normalizer) is needed at construction time and is not an nn.Module buffer.

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 – JODO’s whole 2D half – is silently dropped.

build() → JodoDiffusionTask

Construct the task, deriving size/property priors from the data.

atom_vocab

beta_0 = 0.1

beta_1 = 20.0

centered = True

com = True

cond_time = True

condition = []

dist_gbf = True

dropout = 0.1

edge_ch = 2

edge_quan_th = 0.0

gbf_name = 'CondGaussianLayer'

generation_time_keys = ('sdf_output_path',)

include_fc_charge = True

kwargs

loss_weights = [1.0, 0.25, 0.1]

mlp_ratio = 2

n_extra_heads = 2

n_heads = 16

n_layers = 8

nf = 256

noise_align = True

noise_schedule = 'cosine'

normalize_condition = 'mad'

normalize_factors = [1, 4, 4, 1]

pred_data = True

reduce_mean = False

sampling_steps = 1000

sdf_output_path = None

self_cond = True

softmax_inf = True

spatial_cut_off = 2.0

task: JodoDiffusionTask | None = None

task_type = 'diffusion_jodo'

train_set = None

trans_name = 'TransMixLayer'

MolecularDiffusion.modules.tasks.diffusion_jodo.expand_dims(v: torch.Tensor, dims: int) → torch.Tensor

v of shape [N] -> [N, 1, ..., 1] with dims dimensions.

MolecularDiffusion.modules.tasks.diffusion_jodo.get_align_position(z_t: torch.Tensor, xh: torch.Tensor) → torch.Tensor

Rotate the clean positions onto the noisy ones (losses.py:403).

MolecularDiffusion.modules.tasks.diffusion_jodo.kabsch_batch(coords_pred: torch.Tensor, coords_tar: torch.Tensor) → torch.Tensor

Batched Kabsch rotation, port of upstream losses.py:424.

MolecularDiffusion.modules.tasks.diffusion_jodo.Chem = None

MolecularDiffusion.modules.tasks.diffusion_jodo.logger