MolecularDiffusion.modules.tasks.diffusion_ligandiff_multi

multi-LigandDiff integration with the MolecularDiffusion data pipeline.

multi-LigandDiff (https://github.com/Neon8988/multi_LigandDiff, ChemRxiv 2024) regenerates several ligands at once around a fixed transition-metal centre – “partial to total generation”, up to every ligand from a bare metal. It is a direct extension of LigandDiff, already integrated here as modules/tasks/diffusion_ligandiff.py.

The network is not re-ported. EDM and Dynamics are imported unchanged from MolecularDiffusion.modules.models.ligandiff and constructed with ligand_group_node_nf=7; that reproduces the released checkpoint’s shapes exactly (strict=True load, 238/238 tensors matched). See docs/model_integrations/ligandiff_multi/INTEGRATION_PLAN.md for the experiment. Only the 20-metal geometry helpers are new, in modules/models/ligandiff_multi/utils.py.

What is new relative to ligandiff:

• the conditioning channel is 7 wide, cat([ligand_group(6), coord_site(1)]) (upstream src/lightning.py:126), where coord_site is a per-atom “coordinates the metal” flag. ligand_group itself is still 6 slots.

• the row-data contract is therefore 16 columns, not 15.

• normalize_factors is [10, 4, 1], not [1, 4, 10].

• sample() grows several ligands, picking an octahedral denticity partition from CN_OCT (upstream generate.py::reform_data).

Two vocabularies, deliberately different widths:

• 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 (byte-identical to LigandDiff’s). The metal carries an all-zero row and its element travels out of band.

• task.atom_vocab – 28 entries (those 8, then the 20 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).

Attributes

LIGANDIFF_ATOM_VOCAB
LIGANDIFF_MULTI_FULL_VOCAB
LIGANDIFF_MULTI_METAL_VOCAB
LIGANDIFF_MULTI_ROW_DATA_COLUMNS

Classes

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

Functions

LigandDiffMultiToPointCloud(coords_flat, charges_flat, ...)

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

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_ligandiff_multi.LigandDiffMultiTask(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, denticity_split: list | None = None)

Bases: torch.nn.Module

Plain nn.Module task (docs/adding_new_models.md §2.6) wrapping the imported, unmodified LigandDiff 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)

Grow several ligands 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 generate.py::reform_data. The scaffold may be as small as a bare metal (upstream’s []_[...] total-generation case).

Out of scope this pass, accepted and ignored: batch, num_steps, condition_mode, use_noised_conditioning, n_frames (no trajectory export), n_retrys/t_retry, context (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

denticity_split

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_multi.LigandDiffMultiTaskFactory(task_type: str, in_node_nf: int = 8, n_dims: int = 3, ligand_group_node_nf: int = 7, 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 = 'learned', diffusion_noise_precision: float = 1e-05, diffusion_loss_type: str = 'vlb', normalize_factors: tuple = (10, 4, 1), center_of_mass: str = 'context', atom_vocab: list | None = None, denticity_split: list | None = None, **kwargs)

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

The trailing **kwargs is load-bearing, not decoration: cli/train.py:649-651 passes node_feature_dim / extra_norm_values into every factory. Both are meaningless here (the 16 row-data columns are consumed by the adapter, not normalised as extra features) and are deliberately absorbed and ignored.

build() → LigandDiffMultiTask

Instantiate the task module.

activation = 'silu'

aggregation_method = 'sum'

atom_vocab

attention = True

center_of_mass = 'context'

condition_time = True

denticity_split

diffusion_loss_type = 'vlb'

diffusion_noise_precision = 1e-05

diffusion_noise_schedule = 'learned'

diffusion_steps = 500

drop_rate = 0.2

hidden_nf = 192

in_node_nf = 8

inv_sublayers = 1

kwargs

ligand_group_node_nf = 7

model = 'gvp_dynamics'

n_dims = 3

n_layers = 5

norm_constant = 1

normalization = 'batch_norm'

normalization_factor = 100

normalize_factors = (10, 4, 1)

sin_embedding = False

tanh = True

task_type

class MolecularDiffusion.modules.tasks.diffusion_ligandiff_multi.PointCloudToLigandDiffMultiBatch

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

Identical in shape to ligandiff’s adapter, one column wider. The dataset supplies [one_hot(8) | ligand_diff(1) | ligand_group(6) | coord_site(1)] through the existing data.use_row_data_features seam (data/component/dataset.py:1328).

The 16 columns are taken from the right-hand end of node_feature, not from index 0: DataModule.load()’s data_type: pointcloud branch (runmodes/train/data.py:294-312) does not forward use_ohe_feature the way its pyg sibling at :284 does, so the 28-wide 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.

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

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

Flat (N_total, ·) 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 rows. The returned one-hot is re-expanded from those atomic numbers onto atom_vocab (28 entries).

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

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

MolecularDiffusion.modules.tasks.diffusion_ligandiff_multi.LIGANDIFF_MULTI_FULL_VOCAB = ['C', 'N', 'O', 'S', 'Br', 'Cl', 'P', 'F', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn',...

MolecularDiffusion.modules.tasks.diffusion_ligandiff_multi.LIGANDIFF_MULTI_METAL_VOCAB = ['Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Zr', 'Mo', 'Ru', 'Rh', 'Pd', 'Cd', 'W',...

MolecularDiffusion.modules.tasks.diffusion_ligandiff_multi.LIGANDIFF_MULTI_ROW_DATA_COLUMNS