MolecularDiffusion.modules.tasks.diffusion_loqi

LoQI – low-energy conformer generation on a fixed molecular graph.

One module, two task types, because upstream is one model with two configs:

task_type	interpolant	weights
diffusion_loqi	VDM diffusion, 25 discrete steps, cosine-adaptive, self-conditioned	loqi.ckpt
diffusion_loqi_flow	continuous flow matching, velocity prediction, rigid (Kabsch) OT, linear	loqi_flow.ckpt

This is not de-novo generation. Only x (coordinates) is ever noised. Atom types, formal charges, bond orders and stereochemistry are supplied un-noised as network input at every step and come straight back out of sample() unchanged. Every sample therefore needs an input molecule; see sample_input below.

Data path: data_type: graph3d with bond_collate: raw (a PyG Batch) and kekulize: true.

Bond representation mapping

LoQI’s network consumes a fully-connected directed edge list with a 9-class label: 0=none 1=SINGLE 2=DOUBLE 3=TRIPLE 4=AROMATIC 5=E 6=Z 7=chirality(sym) 8=chirality(directed). Classes 5-8 are not bonds – they are a stereochemistry encoding laid over the bond graph, and the class-8 edges are deliberately asymmetric (the direction is the R/S signal), so they cannot be stored in the platform’s symmetric upper-triangular bond_index/bond_type.

So storage keeps the canonical five, and Graph3DToLoQIAdapter reconstructs the rest per batch:

1. mirror the stored upper-triangular bonds into a directed list (reproducing upstream’s adj.nonzero());

2. rebuild the molecule with build_rdkit_mol – which calls AssignStereochemistryFrom3D, the same source of truth upstream’s add_stereo_bonds(from_3D=True) used – and derive classes 5-8 from it;

3. run upstream’s own make_graph_fully_connected to materialize class 0.

Class 4 never appears: the LoQI pipeline kekulizes (process_chembl3d.py:216, sample_conformers.py:84), so the released weights have never seen one.

There is no per-item cache for step 2 in this first pass. If RDKit dominates step time, the fix is a cache keyed on the dataset index, which needs the dataset to carry one – a separate change.

Attributes

CHARGE_OFFSET
FULL_ATOM_ENCODER
N_CHARGE_CLASSES
logger

Classes

Graph3DToLoQIAdapter	graph3d PyG Batch -> LoQI's flat conditioning dict.
LoQIConformerTask	LoQI wrapped in the platform's duck-typed Task contract (§2.1).
LoQINodeDistribution	Histogram sampler over train_set.graph3d_stats.n_atoms_hist.
ModelTaskFactory	Instantiated by cli/train.py / cli/generate.py from Hydra.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_loqi.Graph3DToLoQIAdapter(atom_vocab: list[str])

Bases: torch.nn.Module

graph3d PyG Batch -> LoQI’s flat conditioning dict.

Output keys: batch (N,), x (N,3) COM-free, h (N,17) one-hot in FULL_ATOM_ENCODER order, charges (N,6) one-hot of fc + 2, edge_index (2,E) fully connected directed, edge_attr (E,9) one-hot.

build_edges(pyg: Any, device) → tuple[torch.Tensor, torch.Tensor]

Mirror stored bonds, append derived stereo edges, then densify.

forward(batch: Any) → dict

vocab_map(device) → torch.Tensor

Platform atom_idx -> LoQI FULL_ATOM_ENCODER index.

Built by SYMBOL, never by position: the platform’s atom_vocab order is a data-config choice, LoQI’s is fixed by the pretrained atom_embedder’s input columns. Getting this wrong loads silently and mislabels every atom.

atom_vocab

class MolecularDiffusion.modules.tasks.diffusion_loqi.LoQIConformerTask(task_type: str, atom_vocab: list[str], interpolant_config: dict, dynamics_config: dict, self_cond_config: dict | None, timesteps: int, loss_scale: float, loss_clamp: float | None, return_step_output: bool, sample_input: str | collections.abc.Sequence[str] | None, sample_pool_limit: int, n_atoms_hist: dict, conditioning_pool: list)

Bases: torch.nn.Module

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

conditioning_pool() → list

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

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

predict_and_target(batch: Any)

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, mols: list | None = None, **kwargs: Any)

Generate conformers for a batch of conditioning molecules.

Returns the platform tuple (one_hot (B,N,V), charges (B,N), coords (B,N,3), node_mask (B,N)) with V = len(atom_vocab). Atom types and charges are the conditioning molecule’s – inputs, not predictions.

nodesxsample is honoured as a count, and molecules are size-matched to it where the pool allows. mode is accepted and ignored.

mols names the conditioning molecules EXPLICITLY (a list of the pool’s Data items, repeats allowed) and bypasses the size-matched draw entirely – that is how the conformer mode asks for “k conformers of THIS molecule”, which a size draw cannot express. When it is None (every pre-existing caller) nothing changes.

T

adapter

atom_vocab

property device: torch.device

dynamics

interpolant

loss_clamp

loss_scale

property model: LoQIConformerTask

property n_node_dist: dict

node_dist_model

prop_dist_model = None

return_step_output

sample_input

sample_pool_limit

self_cond

task_type

timesteps

class MolecularDiffusion.modules.tasks.diffusion_loqi.LoQINodeDistribution(histogram: dict)

Histogram sampler over train_set.graph3d_stats.n_atoms_hist.

Only used so mol_size: [0, 0] works in a generate config; the actual molecule sizes always come from the conditioning molecules.

sample(n_samples: int) → torch.Tensor

histogram

n_node_dist

class MolecularDiffusion.modules.tasks.diffusion_loqi.ModelTaskFactory(task_type: str = 'diffusion_loqi', interpolant: dict | None = None, dynamics: dict | None = None, self_conditioning: dict | None = None, timesteps: int = 25, loss_scale: float = 3.0, loss_clamp: float | None = 10.0, return_step_output: bool = False, sample_input: str | collections.abc.Sequence[str] | None = None, sample_pool_limit: int = 2000, dataset_stats: dict | None = None, atom_vocab: list | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs: Any)

Instantiated by cli/train.py / cli/generate.py from Hydra.

train_set is a declared parameter, which is what makes the declarative injection seam (docs/adding_new_models.md §2.5) hand over the dataset – needed for the size histogram and the default conditioning pool.

build() → LoQIConformerTask

atom_vocab

dataset_stats

dynamics

generation_time_keys = ('sample_input', 'sample_pool_limit')

interpolant

kwargs

loss_clamp = 10.0

loss_scale

return_step_output = False

sample_input = None

sample_pool_limit = 2000

self_conditioning = None

task: LoQIConformerTask | None = None

task_type = 'diffusion_loqi'

timesteps = 25

train_set = None

MolecularDiffusion.modules.tasks.diffusion_loqi.CHARGE_OFFSET = 2

MolecularDiffusion.modules.tasks.diffusion_loqi.FULL_ATOM_ENCODER: dict[str, int]

MolecularDiffusion.modules.tasks.diffusion_loqi.N_CHARGE_CLASSES = 6

MolecularDiffusion.modules.tasks.diffusion_loqi.logger