MolecularDiffusion.modules.tasks.elucidation_generator¶

The one run() loop behind every measurement-conditioned generator.

Structure elucidation is a different problem from the ones the existing seams solve. You are handed a measurement of an unknown compound – a tandem mass spectrum, an NMR pair, an IR trace – and you return a ranked shortlist of candidate structures. The molecule is the unknown; the conditioning is read per item from a corpus of measurements; and the answer is a set, not a sample.

Every existing generator was surveyed before this one was written, and each fails on a named mechanism rather than on taste:

GenerativeFactory conditions on target_values – a flat scalar list broadcast identically to every molecule in the run – or on one reference .xyz tiled to the batch. Elucidation needs a different payload per output. Its size logic is wrong in kind too (nodesxsample from a prior), and every writer in it is coordinate-based.
PocketGenerator has the right candidate-set shape but is single-context by construction (item = self._pocket(), called once), and its output contract is hard-wired to coordinates (_accept(one_hot, coords, node_mask), save_xyz_file). TSGenerator rides it precisely because a transition state is coordinates.
ConformerFactory is closest by loop shape, but its pool items are molecules (_to_mol(item, item.pos...)), it writes a reference.sdf of its input, and its metric is RMSD against that input geometry. Here the input is a measurement and the molecule is the unknown.

Hence a new shared seam rather than a bandaid. What differs per model lives in hooks:

Four constraints keep this seam general rather than shaped around any one model. Each was checked against two structurally different elucidation models – one graph-based and coordinate-free, one coordinate-first – before being written down:

``_priors`` is an explicit named slot, not a kwarg blob. Both models are handed the molecular formula; composition-is-known is the norm in this problem, not an extra.
``_decode`` returns SMILES plus nullable 3D. A graph-based model emits connectivity and never fills coords; a coordinate-first model fills it and reaches SMILES through bond perception. Both slots always exist. _decode is also allowed to be slow and to fail per candidate.
``_rank`` defaults to identity. No ranking framework is built here. A model that votes over repeated draws overrides it (e.g. canonical-InChI frequency); a model with no scorer at all takes the default and is not forced into a no-op override.
``_reference`` may return ``None``. Scoring is a separate, optional pass. The core loop must run on genuinely unknown spectra – that is the actual use case, and an evaluation harness that cannot run without labels is not one.

Padding is deliberately not in the seam: it belongs to _repeat(), because one model pads to a dataset-wide max_n_atoms while another is variable-size PyG.

Outputs, per run:

<output_path>/record_XXXX/candidates.{sdf,smi}   ranked candidates
<output_path>/record_XXXX/ranking.csv            rank, smiles, score
<output_path>/predictions.csv                    run-level, one row per record
<output_path>/metrics.json                       only if _reference gave labels

Attributes¶

logger

Classes¶

`Candidate`	One proposed structure.
`ElucidationGenerator`	Walk a corpus of measurements; emit a ranked candidate set per record.

Functions¶

move_to_device(→ Any)

Recursively place a batch on device, leaving unknown types alone.

Module Contents¶

class MolecularDiffusion.modules.tasks.elucidation_generator.Candidate¶

One proposed structure.

smiles¶: canonical SMILES. Always present – it is the common currency between a graph-generating model and a coordinate-generating one.

mol¶: the RDKit molecule, when the model produced one.

coords¶: (n_atoms, 3) positions, or None. Nullable on purpose: a 2D elucidation model never fills it, and forcing a fake conformer in would be inventing geometry the model did not predict.

score¶: rank score, when the model has a scorer. None means the candidate is ordered by generation order.

coords: numpy.ndarray | None = None¶

meta: dict¶

mol: Any = None¶

score: float | None = None¶

smiles: str¶

class MolecularDiffusion.modules.tasks.elucidation_generator.ElucidationGenerator(task: Any, spectra_source: str | None = None, spectra_index: int = 0, max_records: int | None = None, num_candidates: int = 100, batch_size: int = 10, num_steps: int | None = None, guidance_scale: float | None = None, split: str | None = None, drop_channels: Sequence[str] = (), top_k: Sequence[int] = (1, 10), seed: int = 42, device: str | None = None, output_path: str = 'generated_elucidation', **kwargs: Any)¶

Walk a corpus of measurements; emit a ranked candidate set per record.

Not instantiable from a config on its own – _records(), _condition(), _repeat() and _decode() are what a model must supply.

Configure the run.

Parameters:

task – the loaded elucidation task; its elucidate() is called.
spectra_source – where measurements are read from. A subclass may alias this under its own name via source_key.
spectra_index – index of the first record to process.
max_records – how many records to process; None => all of them.
num_candidates – candidates to draw per record.
batch_size – how many of those to draw at once. k is a batch dimension, never a Python loop.
num_steps – reverse-process steps; None => the model’s default.
guidance_scale – classifier-free guidance strength – the w in (1+w)*cond - w*uncond. None => the model’s own default. Note 0.0 is a meaningful, different value: it means unconditional. Rejected when the model declares supports_guidance = False.
split – which fold of a labelled corpus to elucidate (e.g. test). None => whatever the model’s own reader defaults to. Corpora with no folds ignore it.
drop_channels – measurement channels to blank out before the model is shown the record, by name. Valid names are this model’s maskable_channels; anything else is rejected naming the valid set, and a model declaring none rejects any non-empty list. () => the measurement exactly as recorded.
top_k – which top-k accuracies to report, when references exist.
seed – torch/random/numpy seed.
device – None => cuda if available.
output_path – directory for the per-record candidate sets.
**kwargs – rejected by the base, on purpose.

Raises:

ValueError – on any unrecognised interference key. A silently ignored key makes a run look configured when it is not – the same reason PocketGenerator rejects rather than ignores.

run() → None¶

batch_size = 10¶

device = 'cuda'¶

drop_channels¶

guidance_scale = None¶

maskable_channels: ClassVar[tuple[str, ...]] = ()¶

max_records = None¶

num_candidates = 100¶

num_steps = None¶

output_path = 'generated_elucidation'¶

seed = 42¶

source_key: ClassVar[str] = 'spectra_source'¶

source_required_msg: ClassVar[str | None] = None¶

spectra_index = 0¶

spectra_source = None¶

split = None¶

supports_guidance: ClassVar[bool] = True¶

tag: ClassVar[str] = 'elucidation'¶

task¶

top_k¶

MolecularDiffusion.modules.tasks.elucidation_generator.move_to_device(obj: Any, device: Any) → Any¶

Recursively place a batch on device, leaving unknown types alone.

The conditioning payload is opaque to this seam by design, but “the batch must be on the same device as the model” is true for every model that will ever ride it – and getting it wrong surfaces as a RuntimeError deep inside an encoder, not as anything that names the batch. So the base applies this to whatever ElucidationGenerator._repeat() returns. It handles tensors, dicts, lists/tuples and anything with a .to() (a PyG Batch, for instance); anything else is returned untouched, so a model whose payload needs custom placement can simply do it in _repeat – this then finds nothing left to move.

MolecularDiffusion.modules.tasks.elucidation_generator.logger¶