MolecularDiffusion.modules.tasks.diffusion_ligandiff

LigandDiff integration with the MolecularDiffusion data pipeline.

LigandDiff (https://github.com/lgd6/LigandDiff, JCTC 2024) regenerates one ligand of a 3D transition-metal complex while keeping the metal and the remaining ligands frozen. It is an EDM-family continuous 3D DDPM with two complementary per-atom masks (context / ligand_diff) plus a categorical ligand-slot embedding (ligand_group), denoised by a GVP network over a runtime fully-connected graph.

See docs/model_integrations/ligandiff/INTEGRATION_PLAN.md for the approved plan this module implements. Ported model code lives under MolecularDiffusion.modules.models.ligandiff.

Two vocabularies, deliberately different widths (INTEGRATION_PLAN.md, “Two vocabularies, one model width”):

• in_node_nf = 8 – the model’s / dataset’s one-hot width, upstream’s heavy-atom vocab {C,N,O,S,Br,Cl,P,F} in src/const.py:11 order. The metal carries an all-zero row and its element travels out of band. The released checkpoint’s edm.dynamics.h_embedding.weight is (192, 8), so this width must not be widened.

• task.atom_vocab – 18 entries (those 8, then the 10 metals). A decode-side symbol table only; it sizes no tensor. It is mandatory because runmodes/generate/tasks_generate.py:1463 resolves reference-structure elements with onehot(..., allow_unknown=False), which raises on every transition-metal complex against an 8-entry vocab.

Attributes

LIGANDIFF_ATOM_VOCAB
LIGANDIFF_FULL_VOCAB
LIGANDIFF_METAL_VOCAB
LIGANDIFF_ROW_DATA_COLUMNS

Classes

LigandDiffTask	Plain nn.Module task (docs/adding_new_models.md §2.6) wrapping
LigandDiffTaskFactory	Factory matching cli/train.py's build() instantiation pattern.
PointCloudToLigandDiffBatch	Dense (B, N, ·) PointCloud batch -> LigandDiff's flat layout.

Functions

LigandDiffToPointCloud(coords_flat, charges_flat, ...)

Flat (N_total, ·) LigandDiff output -> dense PointCloud tensors.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_ligandiff.LigandDiffTask(in_node_nf: int, n_dims: int, ligand_group_node_nf: int, hidden_nf: int, n_layers: int, activation: str, attention: bool, tanh: bool, norm_constant: float, inv_sublayers: int, sin_embedding: bool, normalization_factor: float, aggregation_method: str, model: str, normalization: str | None, condition_time: bool, drop_rate: float, diffusion_steps: int, diffusion_noise_schedule: str, diffusion_noise_precision: float, diffusion_loss_type: str, normalize_factors: tuple, center_of_mass: str, atom_vocab: list | None = None)

Bases: torch.nn.Module

Plain nn.Module task (docs/adding_new_models.md §2.6) wrapping LigandDiff’s EDM.

The EDM is held at self.edm because the released checkpoint’s 238 tensors are all named edm.*; any other attribute name makes every key “unexpected” and cli/generate.py:322’s strict=False load drops them all silently.

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

Validation metric.

forward(batch: dict)

Compute the diffusion loss (src/lightning.py:124-176).

get_extra_state() → dict

Persist the atom-count histogram into the checkpoint.

predict_and_target(batch: dict)

Validation hook: the loss doubles as the prediction.

sample(nodesxsample: torch.Tensor | None = None, batch_size: int | None = None, num_steps: int | None = None, batch: dict | None = None, condition_tensor: torch.Tensor | None = None, condition_mode: str | None = None, outpaint_cfgs: dict | None = None, use_noised_conditioning: bool = False, n_frames: int = 0, n_retrys: int = 0, t_retry: int | None = None, context: torch.Tensor | None = None, **kwargs)

Regenerate one ligand around a fixed metal-complex scaffold.

Dispatched through the bundled outpaint GenerativeFactory path (configs/interference/gen_outpaint.yaml -> tasks_generate.py::structural_guidance), whose mechanism – freeze a reference scaffold, draw a strictly larger total size, grow the difference – is exactly upstream’s sampling.py::reform_data.

condition_tensor is (B, ref_natoms, 3 + len(atom_vocab) + 1) as built by tasks_generate.py::preprocess_ref_structure; only the coordinate block and the trailing charge column are used, because the model’s own 8-wide one-hot has to be rebuilt from the true elements (the metal’s row must be all-zero, which no 18-wide one-hot encodes).

Out of scope this pass, accepted and ignored: batch, num_steps, condition_mode, outpaint_cfgs, use_noised_conditioning, n_frames (no trajectory export), n_retrys/t_retry, context (LigandDiff has no prop_dist_model).

set_extra_state(state: dict) → None

Restore the atom-count histogram from a checkpoint.

T

property atom_count_histogram: dict | None

Total-atom-count histogram accumulated from real batches.

atom_vocab

center_of_mass

property device: torch.device

Device the parameters live on.

edm

in_node_nf

ligand_group_node_nf

loss_type

property model

tasks_generate.py reads T/norm_values off this.

n_dims

property n_node_dist: dict | None

Raw histogram, as core/engine_lightning.py checkpoints it.

ndim_extra = 0

property node_dist_model: MolecularDiffusion.modules.models.difflinker.linker_size.DistributionNodes | None

Sampler over total molecule sizes.

norm_values

prop_dist_model = None

class MolecularDiffusion.modules.tasks.diffusion_ligandiff.LigandDiffTaskFactory(task_type: str, in_node_nf: int = 8, n_dims: int = 3, ligand_group_node_nf: int = 6, hidden_nf: int = 192, n_layers: int = 5, activation: str = 'silu', attention: bool = True, tanh: bool = True, norm_constant: float = 1, inv_sublayers: int = 1, sin_embedding: bool = False, normalization_factor: float = 100, aggregation_method: str = 'sum', model: str = 'gvp_dynamics', normalization: str | None = 'batch_norm', condition_time: bool = True, drop_rate: float = 0.2, diffusion_steps: int = 500, diffusion_noise_schedule: str = 'polynomial_2', diffusion_noise_precision: float = 1e-05, diffusion_loss_type: str = 'vlb', normalize_factors: tuple = (1, 4, 10), center_of_mass: str = 'context', atom_vocab: list | None = None, **kwargs)

Factory matching cli/train.py’s build() instantiation pattern.

The trailing **kwargs is load-bearing, not decoration: this model runs with data.use_ohe_feature: false, so cli/train.py:649-651 computes inferred_extra_dim = 15 and passes node_feature_dim / extra_norm_values into every factory. Both are meaningless here (the 15 columns are consumed by the adapter, not normalised as extra features) and are deliberately absorbed and ignored, exactly as DiffLinkerTaskFactory does.

build() → LigandDiffTask

Instantiate the task module.

activation = 'silu'

aggregation_method = 'sum'

atom_vocab

attention = True

center_of_mass = 'context'

condition_time = True

diffusion_loss_type = 'vlb'

diffusion_noise_precision = 1e-05

diffusion_noise_schedule = 'polynomial_2'

diffusion_steps = 500

drop_rate = 0.2

hidden_nf = 192

in_node_nf = 8

inv_sublayers = 1

kwargs

ligand_group_node_nf = 6

model = 'gvp_dynamics'

n_dims = 3

n_layers = 5

norm_constant = 1

normalization = 'batch_norm'

normalization_factor = 100

normalize_factors = (1, 4, 10)

sin_embedding = False

tanh = True

task_type

class MolecularDiffusion.modules.tasks.diffusion_ligandiff.PointCloudToLigandDiffBatch

Dense (B, N, ·) PointCloud batch -> LigandDiff’s flat layout.

LigandDiff works on a ragged batch: every atom of every molecule concatenated on dim 0, membership via batch_seg. This adapter selects the valid rows and slices the conditioning columns out of node_feature, which the dataset supplies as [one_hot(8) | ligand_diff(1) | ligand_group(6)] through the existing data.use_row_data_features seam (data/component/dataset.py:1328).

The 15 columns are taken from the right-hand end of node_feature, not from index 0. That is the row-data seam’s own guarantee – data/component/dataset.py:1328 appends row.data["node_features"] with torch.cat((node_features, row_feat), dim=1) – and it is load-bearing here: DataModule.load()’s data_type: pointcloud branch (runmodes/train/data.py:294-312) does not forward use_ohe_feature to pointcloud_dataset the way its pyg sibling at :284 does, so the atom OHE is always prepended regardless of the config flag. Slicing from the end is correct either way and needs no change to shared code. The prepended OHE is simply unused: the model’s 8-wide h comes from these columns, so in_node_nf stays 8 and the released checkpoint still loads.

There is no synthesis fallback: the Zenodo splits ship ligand_diff and ligand_group precomputed, so a missing column means a misconfigured dataset and must raise rather than silently train on a fabricated split.

MolecularDiffusion.modules.tasks.diffusion_ligandiff.LigandDiffToPointCloud(coords_flat: torch.Tensor, charges_flat: torch.Tensor, batch_seg: torch.Tensor, batch_size: int, atom_vocab: list)

Flat (N_total, ·) LigandDiff output -> dense PointCloud tensors.

charges_flat must already be the final per-atom atomic number: taken from the known context elements for context rows (which is what stops the metal’s all-zero one-hot decoding to index 0, i.e. carbon) and from the 8-way one_hot.argmax for generated ligand_diff rows. The returned one-hot is re-expanded from those atomic numbers onto the atom_vocab (18 entries), so argmax decoding and charge-based decoding agree and tasks_generate.py:1049’s width check takes its normal branch.

Returns (one_hot, charges, coords, node_mask), all dense.

MolecularDiffusion.modules.tasks.diffusion_ligandiff.LIGANDIFF_ATOM_VOCAB = ['C', 'N', 'O', 'S', 'Br', 'Cl', 'P', 'F']

MolecularDiffusion.modules.tasks.diffusion_ligandiff.LIGANDIFF_FULL_VOCAB = ['C', 'N', 'O', 'S', 'Br', 'Cl', 'P', 'F', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ru', 'Pd', 'Pt']

MolecularDiffusion.modules.tasks.diffusion_ligandiff.LIGANDIFF_METAL_VOCAB = ['Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ru', 'Pd', 'Pt']

MolecularDiffusion.modules.tasks.diffusion_ligandiff.LIGANDIFF_ROW_DATA_COLUMNS