From 7264e0c3485ace96976c94b3401fe1fc87731e7c Mon Sep 17 00:00:00 2001 From: Finlay Clark Date: Mon, 13 Jul 2026 17:03:41 +0100 Subject: [PATCH 1/2] Allow the user to supply starting conformers from an SDF --- docs/changelog.md | 6 ++ docs/how-to/index.md | 1 + docs/how-to/use-starting-conformers.md | 69 +++++++++++++ presto/load_molecules.py | 81 ++++++++++++++++ presto/msm.py | 14 +-- presto/sample.py | 53 +++++++--- presto/settings.py | 103 +++++++++++++++++++- presto/tests/conftest.py | 23 +++++ presto/tests/unit/test_load_molecules.py | 117 +++++++++++++++++++++++ presto/tests/unit/test_msm.py | 44 +++++++++ presto/tests/unit/test_sample.py | 109 +++++++++++++++++++++ presto/tests/unit/test_settings.py | 83 ++++++++++++++++ 12 files changed, 681 insertions(+), 22 deletions(-) create mode 100644 docs/how-to/use-starting-conformers.md create mode 100644 presto/tests/unit/test_load_molecules.py diff --git a/docs/changelog.md b/docs/changelog.md index cf4d5fd..27ff662 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -4,11 +4,17 @@ ### Features +- Add an optional `starting_conformers` setting to each sampling stage (`training_sampling_settings`, `testing_sampling_settings`) and to `msm_settings`. When set to an SDF path, that stage starts from the supplied conformers (matched to each molecule by graph isomorphism and realigned automatically) instead of generating them with ETKDG; `n_conformers` is ignored for that stage. The default remains ETKDG. - Add `presto.create_types.add_library_charges_to_forcefield` to write custom partial charges from OpenFF `Molecule` objects into a force field as library charges, addressing [#64](https://github.com/cole-group/presto/issues/64). +### Fixes + +- Fix a latent `IndexError` in the MSM step when fewer conformers were available than `n_conformers`; the conformer loop now iterates the conformers actually present. + ### Documentation - Add [Use custom charges](how-to/use-custom-charges.md) how-to guide. +- Add [Use your own starting conformers](how-to/use-starting-conformers.md) how-to guide. - Add a `CITATION.cff` file and cite the presto preprint in the README and docs, and point users to the OpenFF publications to cite, in [#76](https://github.com/cole-group/presto/pull/76). - Document installing `presto` from `conda-forge` as `presto-fit` (note this comes without the MLP dependencies, which must be installed separately) in [#70](https://github.com/cole-group/presto/pull/70). diff --git a/docs/how-to/index.md b/docs/how-to/index.md index ef1bab0..6c0e185 100644 --- a/docs/how-to/index.md +++ b/docs/how-to/index.md @@ -7,6 +7,7 @@ Short, task-oriented recipes. Each page assumes you've finished **[Get started]( - **[Fit a single molecule](fit-single-molecule.md)** — defaults you may want to change. - **[Fit a congeneric series](fit-congeneric-series.md)** — share parameters across related molecules using `max_extend_distance`. - **[Use SDF inputs](use-sdf-inputs.md)** — switch from SMILES to one or more `.sdf` files. +- **[Use your own starting conformers](use-starting-conformers.md)** — seed a sampling stage (or MSM) from an SDF instead of ETKDG. - **[Use custom charges](use-custom-charges.md)** — bake your own partial charges into the force field as library charges. - **[Wipe output and rerun](clean-rerun.md)** diff --git a/docs/how-to/use-starting-conformers.md b/docs/how-to/use-starting-conformers.md new file mode 100644 index 0000000..31f9d3f --- /dev/null +++ b/docs/how-to/use-starting-conformers.md @@ -0,0 +1,69 @@ +# Use your own starting conformers + +By default each sampling stage generates its starting conformers with RDKit ETKDG. You can +instead start from your own conformers (for example a curated ensemble, docked or crystal +poses, or a torsion scan) by pointing a stage at an SDF file with `starting_conformers`. + +This is available **per stage**, so you can mix supplied conformers and ETKDG freely: + +- `training_sampling_settings.starting_conformers` — training-data MD +- `testing_sampling_settings.starting_conformers` — test-data MD +- `param_settings.msm_settings.starting_conformers` — the MSM Hessian step that seeds the + initial bond/angle parameters + +Any stage left unset keeps generating conformers with ETKDG. + +## Behaviour + +- **The SDF takes precedence.** When `starting_conformers` is set for a stage, that stage + starts from **every** conformer in the file, and the stage's `n_conformers` is **ignored** + (a log line makes this explicit). +- **Conformers are matched by connectivity.** Each record is matched to the molecule being + fitted by graph isomorphism, and its atom ordering is **remapped automatically** to match. + You do not need to worry about atom ordering, and a single SDF can hold conformers for + several molecules in a multi-molecule fit — each molecule picks up only its own records. +- **Missing molecules fail fast.** If a configured SDF contains no conformer for a molecule + being fitted, the run stops during settings validation, before any expensive work. + +## CLI form + +```bash +presto train \ + --param-settings.molecules "CCO" \ + --training-sampling-settings.starting-conformers my_conformers.sdf +``` + +## YAML form + +```yaml +param_settings: + molecule_input_type: smiles + molecules: + - CCO + msm_settings: + # optional: also seed the MSM Hessian step from the same (or a different) SDF + starting_conformers: my_conformers.sdf + +training_sampling_settings: + sampling_protocol: mm_md_metadynamics_torsion_minimisation + starting_conformers: my_conformers.sdf + +testing_sampling_settings: + sampling_protocol: ml_md + # left unset -> ETKDG +``` + +## Preparing the SDF + +The SDF should contain one record per conformer, all of the same molecule (multiple +molecules are fine — group them however you like, they are matched by connectivity). For +example, from an existing OpenFF molecule: + +```python +from rdkit import Chem + +rdkit_molecule = molecule.to_rdkit() # molecule with N conformers +with Chem.SDWriter("my_conformers.sdf") as writer: + for conformer_id in range(rdkit_molecule.GetNumConformers()): + writer.write(rdkit_molecule, confId=conformer_id) +``` diff --git a/presto/load_molecules.py b/presto/load_molecules.py index 4b6054d..d30f1cd 100644 --- a/presto/load_molecules.py +++ b/presto/load_molecules.py @@ -5,8 +5,11 @@ from typing import Literal from openff.toolkit import Molecule +from openff.units import Quantity from rdkit import Chem +from .utils.typing import PathLike + MoleculeInputType = Literal["smiles", "sdf"] MoleculeLoader = Callable[[str], list[Molecule]] @@ -72,3 +75,81 @@ def load_sdf_molecules(input_value: str) -> list[Molecule]: "smiles": load_smiles_molecules, "sdf": load_sdf_molecules, } + + +def load_conformers_for_molecule( + molecule: Molecule, sdf_path: PathLike +) -> list[Quantity]: + """Load the conformers of ``molecule`` from an SDF, aligned to its atom ordering. + + Every record in the SDF that is graph-isomorphic to ``molecule`` is treated as a + conformer of it. Because the atom ordering in the SDF need not match ``molecule``, + each matching record is remapped onto ``molecule``'s atom ordering before its + coordinates are extracted, guaranteeing the returned conformers are valid starting + positions for a topology built from ``molecule``. Records that do not match are + ignored (they may belong to another molecule in a multi-molecule SDF). + + Parameters + ---------- + molecule : openff.toolkit.Molecule + The molecule whose conformers should be loaded. Defines the canonical atom + ordering the returned conformers are aligned to. + sdf_path : PathLike + Path to an SDF file containing one or more conformers of ``molecule`` (and, + optionally, of other molecules). + + Returns: + ------- + list[openff.units.Quantity] + The matching conformers, each aligned to ``molecule``'s atom ordering. + + Raises: + ------ + ValueError + If the path does not exist, does not end in ``.sdf``, or contains no record + matching ``molecule``. + """ + path = Path(sdf_path) + + if not path.exists(): + raise ValueError(f"SDF file does not exist: {path}") + + if path.suffix.lower() != ".sdf": + raise ValueError(f"Expected an SDF file path ending in .sdf: {path}") + + try: + supplier = Chem.SDMolSupplier(str(path), removeHs=False) + except Exception as exc: + raise ValueError(f"Failed to read SDF file: {path}") from exc + + conformers: list[Quantity] = [] + + for rdkit_molecule in supplier: + if rdkit_molecule is None: + continue + + try: + record = Molecule.from_rdkit(rdkit_molecule, allow_undefined_stereo=True) + except Exception: + # Records that cannot be interpreted as molecules cannot be a conformer of + # ``molecule``, so skip them rather than failing the whole load. + continue + + matched, atom_map = Molecule.are_isomorphic( + record, molecule, return_atom_map=True + ) + if not matched or atom_map is None: + continue + + # ``atom_map`` maps record atom indices -> ``molecule`` atom indices, so remapping + # the record with ``current_to_new`` yields a molecule in ``molecule``'s ordering. + aligned = record.remap(atom_map, current_to_new=True) + conformers.extend(aligned.conformers) + + if not conformers: + raise ValueError( + f"SDF file {path} contains no conformers matching the molecule " + f"{_molecule_identity(molecule.to_rdkit())}." + ) + + return conformers diff --git a/presto/msm.py b/presto/msm.py index e9c7cf8..958879d 100644 --- a/presto/msm.py +++ b/presto/msm.py @@ -698,9 +698,11 @@ def apply_msm_to_molecule( Tuple of (bond_params_dict, angle_params_dict). If multiple conformers are used, the parameters are averaged over all conformers. """ - # Generate conformers + # Generate conformers (or load the supplied starting conformers) mol_with_conformers = _copy_mol_and_add_conformers( - mol, n_conformers=settings.n_conformers + mol, + n_conformers=settings.n_conformers, + starting_conformers=settings.starting_conformers, ) simulation, integrator = _build_ml_simulation( mol_with_conformers, @@ -717,14 +719,12 @@ def apply_msm_to_molecule( defaultdict(list) ) - for conf_idx in track( - range(settings.n_conformers), + for conformer in track( + mol_with_conformers.conformers, description="Finding MSM parameters for conformers", ): # Set positions for this conformer - simulation.context.setPositions( - mol_with_conformers.conformers[conf_idx].to_openmm() - ) + simulation.context.setPositions(conformer.to_openmm()) # Minimize to local energy minimum simulation.minimizeEnergy(maxIterations=0, tolerance=settings.tolerance) diff --git a/presto/sample.py b/presto/sample.py index 6ca0b93..3cd84f9 100644 --- a/presto/sample.py +++ b/presto/sample.py @@ -33,6 +33,7 @@ DEFAULT_TORSIONS_TO_INCLUDE_SMARTS, get_rot_torsions_by_rot_bond, ) +from .load_molecules import load_conformers_for_molecule from .metadynamics import Metadynamics from .outputs import OutputType, get_mol_path from .utils.gpu import cleanup_simulation @@ -77,16 +78,36 @@ def __call__(self, **kwargs: Unpack[SampleFnArgs]) -> list[datasets.Dataset]: def _copy_mol_and_add_conformers( mol: openff.toolkit.Molecule, n_conformers: int, + starting_conformers: pathlib.Path | None = None, ) -> openff.toolkit.Molecule: - """Copy a molecule and add conformers to it.""" + """Copy a molecule and add starting conformers to it. + + If ``starting_conformers`` is None (default), ``n_conformers`` conformers are + generated with ETKDG. Otherwise the conformers are loaded from the given SDF (matched + to ``mol`` by graph and aligned to its atom ordering) and ``n_conformers`` is ignored. + """ mol = copy.deepcopy(mol) - mol.generate_conformers(n_conformers=n_conformers, rms_cutoff=0.0 * _ANGSTROM) - n_gen_conformers = len(mol.conformers) - if n_gen_conformers < n_conformers: - logger.warning( - f"Only {n_gen_conformers} conformers were generated, which is less than the requested {n_conformers}." - f" As a result, {n_gen_conformers / n_conformers * 100:.1f}% of the requested samples will be generated." - ) + + if starting_conformers is None: + mol.generate_conformers(n_conformers=n_conformers, rms_cutoff=0.0 * _ANGSTROM) + n_gen_conformers = len(mol.conformers) + if n_gen_conformers < n_conformers: + logger.warning( + f"Only {n_gen_conformers} conformers were generated, which is less than the requested {n_conformers}." + f" As a result, {n_gen_conformers / n_conformers * 100:.1f}% of the requested samples will be generated." + ) + return mol + + # Drop any incidental conformer (e.g. one carried by an SDF-loaded molecule) before + # attaching the supplied starting conformers. + mol._conformers = [] + conformers = load_conformers_for_molecule(mol, starting_conformers) + logger.info( + f"Starting from {len(conformers)} supplied conformers in {starting_conformers} " + f"(n_conformers={n_conformers} ignored)." + ) + for conformer in conformers: + mol.add_conformer(conformer) return mol @@ -338,7 +359,9 @@ def sample_mmmd( all_datasets = [] for mol_idx, mol in enumerate(mols): - mol_with_conformers = _copy_mol_and_add_conformers(mol, settings.n_conformers) + mol_with_conformers = _copy_mol_and_add_conformers( + mol, settings.n_conformers, settings.starting_conformers + ) interchange = openff.interchange.Interchange.from_smirnoff( off_ff, openff.toolkit.Topology.from_molecules(mol_with_conformers) ) @@ -433,7 +456,9 @@ def sample_mlmd( all_datasets = [] for mol_idx, mol in enumerate(mols): - mol_with_conformers = _copy_mol_and_add_conformers(mol, settings.n_conformers) + mol_with_conformers = _copy_mol_and_add_conformers( + mol, settings.n_conformers, settings.starting_conformers + ) ml_simulation, integrator = _build_ml_simulation( mol_with_conformers, mol_with_conformers.to_topology().to_openmm(), @@ -559,7 +584,9 @@ def sample_mmmd_metadynamics( all_datasets = [] for mol_idx, mol in enumerate(mols): - mol_with_conformers = _copy_mol_and_add_conformers(mol, settings.n_conformers) + mol_with_conformers = _copy_mol_and_add_conformers( + mol, settings.n_conformers, settings.starting_conformers + ) interchange = openff.interchange.Interchange.from_smirnoff( off_ff, openff.toolkit.Topology.from_molecules(mol_with_conformers) ) @@ -1175,7 +1202,9 @@ def sample_mmmd_metadynamics_with_torsion_minimisation( all_datasets = [] for mol_idx, mol in enumerate(mols): - mol_with_conformers = _copy_mol_and_add_conformers(mol, settings.n_conformers) + mol_with_conformers = _copy_mol_and_add_conformers( + mol, settings.n_conformers, settings.starting_conformers + ) interchange = openff.interchange.Interchange.from_smirnoff( off_ff, openff.toolkit.Topology.from_molecules(mol_with_conformers) ) diff --git a/presto/settings.py b/presto/settings.py index a6ff7bb..a899a8b 100644 --- a/presto/settings.py +++ b/presto/settings.py @@ -29,7 +29,11 @@ DEFAULT_TORSIONS_TO_EXCLUDE_SMARTS, DEFAULT_TORSIONS_TO_INCLUDE_SMARTS, ) -from .load_molecules import MOLECULE_LOADERS, MoleculeInputType +from .load_molecules import ( + MOLECULE_LOADERS, + MoleculeInputType, + load_conformers_for_molecule, +) from .outputs import OutputType, WorkflowPathManager from .utils.dicts import deep_update from .utils.typing import ( @@ -169,6 +173,27 @@ def validate_no_runtime_placeholders(self) -> Self: return self +def _validate_starting_conformers_path(value: Path | None) -> Path | None: + """Validate an optional starting-conformers SDF path. + + Only checks the obvious, molecule-independent problems (missing file, wrong + suffix). Whether the file actually contains conformers for the molecules being + fitted can only be checked once the molecules are known — see + ``WorkflowSettings._check_starting_conformers_match_molecules``. + """ + if value is None: + return value + if value.suffix.lower() != ".sdf": + raise InvalidSettingsError( + f"starting_conformers must be an SDF file ending in .sdf: {value}" + ) + if not value.exists(): + raise InvalidSettingsError( + f"starting_conformers SDF file does not exist: {value}" + ) + return value + + class _SamplingSettingsBase(_DefaultSettings, ABC): """Settings for sampling (usually molecular dynamics).""" @@ -204,9 +229,25 @@ class _SamplingSettingsBase(_DefaultSettings, ABC): n_conformers: int = Field( 10, - description="The number of conformers to generate, from which sampling is started", + description="The number of conformers to generate, from which sampling is started. " + "Ignored when `starting_conformers` is set.", ) + starting_conformers: Path | None = Field( + None, + description="Optional path to an SDF of starting conformers for this sampling " + "stage. If set, sampling starts from every conformer in the file that matches " + "the molecule (matched by graph, atom order aligned automatically) and " + "`n_conformers` is ignored for this stage. If None (default), conformers are " + "generated with ETKDG.", + ) + + @field_validator("starting_conformers") + @classmethod + def _validate_starting_conformers(cls, value: Path | None) -> Path | None: + """Validate that any supplied starting-conformers path is an existing SDF.""" + return _validate_starting_conformers_path(value) + equilibration_sampling_time_per_conformer: OpenMMQuantity[unit.picoseconds] = Field( # type: ignore[type-arg] default=0.0 * unit.picoseconds, description="Equilibration sampling time per conformer. No snapshots are saved during " @@ -712,9 +753,25 @@ class MSMSettings(_DefaultSettings): n_conformers: int = Field( 1, description="Number of conformers to generate and calculate MSM parameters for. " - "The resulting bond and angle parameters will be averaged over all conformers.", + "The resulting bond and angle parameters will be averaged over all conformers. " + "Ignored when `starting_conformers` is set.", ) + starting_conformers: Path | None = Field( + None, + description="Optional path to an SDF of starting conformers for the MSM Hessian " + "calculation. If set, MSM parameters are calculated for every conformer in the " + "file that matches the molecule (matched by graph, atom order aligned " + "automatically) and averaged, and `n_conformers` is ignored. If None (default), " + "conformers are generated with ETKDG.", + ) + + @field_validator("starting_conformers") + @classmethod + def _validate_starting_conformers(cls, value: Path | None) -> Path | None: + """Validate that any supplied starting-conformers path is an existing SDF.""" + return _validate_starting_conformers_path(value) + class ParamSettings(_DefaultSettings): """Settings controlling the initial parameterisation.""" @@ -945,6 +1002,46 @@ def validate_parameterisation_training_consistency(self) -> Self: return self + @model_validator(mode="after") + def validate_starting_conformers_match_molecules(self) -> Self: + """Fail fast if a configured starting-conformers SDF lacks a molecule being fitted. + + This runs before the (slow) parameterisation stage so a mismatch between the + supplied conformers and the fitted molecules surfaces immediately rather than + mid-run. + """ + msm_settings = self.param_settings.msm_settings + stages: list[tuple[str, Path | None]] = [ + ( + "training_sampling_settings", + getattr(self.training_sampling_settings, "starting_conformers", None), + ), + ( + "testing_sampling_settings", + getattr(self.testing_sampling_settings, "starting_conformers", None), + ), + ( + "param_settings.msm_settings", + None if msm_settings is None else msm_settings.starting_conformers, + ), + ] + + configured = [(name, path) for name, path in stages if path is not None] + if not configured: + return self + + molecules = self.param_settings.openff_molecules + for name, path in configured: + for molecule in molecules: + try: + load_conformers_for_molecule(molecule, path) + except ValueError as exc: + raise InvalidSettingsError( + f"{name}.starting_conformers ({path}) is invalid: {exc}" + ) from exc + + return self + @property def device(self) -> torch.device: """Return a torch.device corresponding to the configured device_type.""" diff --git a/presto/tests/conftest.py b/presto/tests/conftest.py index 7a902a5..7cc9151 100644 --- a/presto/tests/conftest.py +++ b/presto/tests/conftest.py @@ -27,6 +27,29 @@ def skip_if_model_unavailable(model_name: str) -> None: pytest.skip(f"{pkg} not installed") +@pytest.fixture +def write_multiconformer_sdf(): + """Return a helper that writes molecules' conformers as separate SDF records. + + ``Molecule.to_file(..., "SDF")`` only writes a single conformer, so tests that need a + genuine multi-record SDF (the shape a user supplies as starting conformers) go via + RDKit. The returned helper accepts a single ``Molecule`` or an iterable of them and + writes every conformer of each as its own record to ``path``. + """ + from rdkit import Chem + + def _write(molecules, path) -> None: + if isinstance(molecules, Molecule): + molecules = [molecules] + with Chem.SDWriter(str(path)) as writer: + for molecule in molecules: + rdkit_molecule = molecule.to_rdkit() + for conformer_id in range(rdkit_molecule.GetNumConformers()): + writer.write(rdkit_molecule, confId=conformer_id) + + return _write + + # From Simon Boothroyd @pytest.fixture def tmp_cwd(tmp_path, monkeypatch) -> Path: diff --git a/presto/tests/unit/test_load_molecules.py b/presto/tests/unit/test_load_molecules.py new file mode 100644 index 0000000..cb34be7 --- /dev/null +++ b/presto/tests/unit/test_load_molecules.py @@ -0,0 +1,117 @@ +"""Unit tests for load_molecules.py, focusing on starting-conformer loading.""" + +import numpy as np +import pytest +from openff.toolkit import Molecule +from openff.units import unit +from scipy.spatial.distance import pdist + +from presto.load_molecules import load_conformers_for_molecule + + +def _sorted_pairwise_distances(conformer: unit.Quantity) -> np.ndarray: + """Atom-relabelling-invariant fingerprint of a conformer's geometry.""" + return np.sort(pdist(conformer.m_as(unit.angstrom))) + + +def _structures_match(loaded, originals, tol: float = 1e-2) -> bool: + """Check each loaded conformer matches some original geometry up to atom relabelling.""" + original_fps = [_sorted_pairwise_distances(c) for c in originals] + for conformer in loaded: + fp = _sorted_pairwise_distances(conformer) + if not any(np.max(np.abs(fp - o)) < tol for o in original_fps): + return False + return True + + +@pytest.fixture +def pentanol_with_conformers(): + """Pentanol with several distinct conformers (rms pruning disabled).""" + molecule = Molecule.from_smiles("CCCCCO") + molecule.generate_conformers(n_conformers=4, rms_cutoff=0.0 * unit.angstrom) + return molecule + + +def test_loads_all_matching_conformers( + pentanol_with_conformers, tmp_path, write_multiconformer_sdf +): + """All records for the molecule are returned as conformers.""" + sdf = tmp_path / "confs.sdf" + write_multiconformer_sdf(pentanol_with_conformers, sdf) + + target = Molecule.from_smiles("CCCCCO") + conformers = load_conformers_for_molecule(target, sdf) + + assert len(conformers) == pentanol_with_conformers.n_conformers + assert all(c.shape == (target.n_atoms, 3) for c in conformers) + assert _structures_match(conformers, pentanol_with_conformers.conformers) + + +def test_atom_order_is_aligned_to_target( + pentanol_with_conformers, tmp_path, write_multiconformer_sdf +): + """Records with a permuted atom order are realigned to the target's ordering.""" + n_atoms = pentanol_with_conformers.n_atoms + permutation = list(range(n_atoms)) + np.random.default_rng(1).shuffle(permutation) + mapping = {i: permutation[i] for i in range(n_atoms)} + remapped = pentanol_with_conformers.remap(mapping, current_to_new=True) + + sdf = tmp_path / "permuted.sdf" + write_multiconformer_sdf(remapped, sdf) + + target = Molecule.from_smiles("CCCCCO") + conformers = load_conformers_for_molecule(target, sdf) + + # Returned conformers use the target's atom count/order, and the physical geometry is + # preserved despite the permuted input ordering. + assert len(conformers) == pentanol_with_conformers.n_conformers + assert all(c.shape == (target.n_atoms, 3) for c in conformers) + assert _structures_match(conformers, pentanol_with_conformers.conformers) + + +def test_multi_molecule_sdf_returns_only_matching(tmp_path, write_multiconformer_sdf): + """An SDF holding conformers of two molecules yields only the matching ones.""" + pentanol = Molecule.from_smiles("CCCCCO") + pentanol.generate_conformers(n_conformers=3, rms_cutoff=0.0 * unit.angstrom) + butane = Molecule.from_smiles("CCCC") + butane.generate_conformers(n_conformers=2, rms_cutoff=0.0 * unit.angstrom) + + sdf = tmp_path / "mixed.sdf" + write_multiconformer_sdf([pentanol, butane], sdf) + + pentanol_target = Molecule.from_smiles("CCCCCO") + butane_target = Molecule.from_smiles("CCCC") + + assert ( + len(load_conformers_for_molecule(pentanol_target, sdf)) == pentanol.n_conformers + ) + assert len(load_conformers_for_molecule(butane_target, sdf)) == butane.n_conformers + + +def test_no_matching_molecule_raises( + pentanol_with_conformers, tmp_path, write_multiconformer_sdf +): + """A molecule absent from the SDF raises a clear error.""" + sdf = tmp_path / "confs.sdf" + write_multiconformer_sdf(pentanol_with_conformers, sdf) + + benzene = Molecule.from_smiles("c1ccccc1") + with pytest.raises(ValueError, match="no conformers matching"): + load_conformers_for_molecule(benzene, sdf) + + +def test_missing_file_raises(tmp_path): + """A missing SDF path raises.""" + target = Molecule.from_smiles("CCO") + with pytest.raises(ValueError, match="does not exist"): + load_conformers_for_molecule(target, tmp_path / "missing.sdf") + + +def test_non_sdf_suffix_raises(tmp_path): + """A non-.sdf path raises.""" + target = Molecule.from_smiles("CCO") + other = tmp_path / "confs.mol2" + other.write_text("") + with pytest.raises(ValueError, match=r"ending in \.sdf"): + load_conformers_for_molecule(target, other) diff --git a/presto/tests/unit/test_msm.py b/presto/tests/unit/test_msm.py index c25e6a0..fd2d627 100644 --- a/presto/tests/unit/test_msm.py +++ b/presto/tests/unit/test_msm.py @@ -7,6 +7,7 @@ import json import math from importlib.resources import files +from unittest.mock import patch import numpy as np import pytest @@ -842,6 +843,49 @@ def test_multiple_conformers(self, base_forcefield): assert ap.force_constant.magnitude > 0 assert 0 < ap.angle.m_as(_ANGLE_UNIT) < np.pi + def test_supplied_starting_conformers( + self, base_forcefield, tmp_path, write_multiconformer_sdf + ): + """MSM runs from a supplied conformer SDF, iterating the actual conformers.""" + import presto.msm as msm_module + + n_supplied = 3 + source = Molecule.from_smiles("CCO") + source.generate_conformers( + n_conformers=n_supplied, rms_cutoff=0.0 * off_unit.angstrom + ) + assert source.n_conformers == n_supplied + sdf = tmp_path / "confs.sdf" + write_multiconformer_sdf(source, sdf) + + mol = Molecule.from_smiles("CCO") + ff = base_forcefield + labels = ff.label_molecules(mol.to_topology())[0] + bond_indices = list(labels["Bonds"].keys()) + angle_indices = list(labels["Angles"].keys()) + + # n_conformers deliberately differs from the file's conformer count. The Hessian + # (computed once per conformer) is spied on to prove the loop iterates every + # supplied conformer rather than n_conformers of them. + settings = MSMSettings( + n_conformers=1, + mlp_settings=MLPSettings(ml_potential="aimnet2"), + starting_conformers=sdf, + ) + + real_calculate_hessian = msm_module.calculate_hessian + with patch.object( + msm_module, "calculate_hessian", side_effect=real_calculate_hessian + ) as spy_hessian: + bond_params, angle_params = apply_msm_to_molecule( + mol, bond_indices, angle_indices, settings, device=torch.device("cpu") + ) + + # One Hessian per supplied conformer, independent of n_conformers=1. + assert spy_hessian.call_count == n_supplied + assert len(bond_params) == len(bond_indices) + assert len(angle_params) == len(angle_indices) + @pytest.mark.slow class TestApplyMSMToMolecules: diff --git a/presto/tests/unit/test_sample.py b/presto/tests/unit/test_sample.py index 3e53e73..a9fa49d 100644 --- a/presto/tests/unit/test_sample.py +++ b/presto/tests/unit/test_sample.py @@ -575,6 +575,68 @@ def test_returns_fewer_conformers_than_requested_for_simple_molecule(self): # Should return a molecule with conformers (may be fewer than requested) assert len(result.conformers) >= 1 + def test_uses_supplied_conformers_and_ignores_n_conformers( + self, tmp_path, write_multiconformer_sdf + ): + """When a starting-conformers SDF is given, use it and ignore n_conformers.""" + from openff.units import unit as off_unit + + from presto.sample import _copy_mol_and_add_conformers + + source = Molecule.from_smiles("CCCCCO") + source.generate_conformers(n_conformers=4, rms_cutoff=0.0 * off_unit.angstrom) + n_supplied = source.n_conformers + sdf = tmp_path / "confs.sdf" + write_multiconformer_sdf(source, sdf) + + target = Molecule.from_smiles("CCCCCO") + # Request a different count to prove n_conformers is ignored. + result = _copy_mol_and_add_conformers( + target, n_conformers=n_supplied + 7, starting_conformers=sdf + ) + + assert len(result.conformers) == n_supplied + + def test_supplied_path_does_not_call_etkdg( + self, tmp_path, write_multiconformer_sdf + ): + """The ETKDG generator is not invoked when conformers are supplied.""" + from unittest.mock import patch + + from presto.sample import _copy_mol_and_add_conformers + + source = Molecule.from_smiles("CCCCCO") + source.generate_conformers(n_conformers=2) + sdf = tmp_path / "confs.sdf" + write_multiconformer_sdf(source, sdf) + + target = Molecule.from_smiles("CCCCCO") + with patch.object( + Molecule, "generate_conformers", autospec=True + ) as mock_generate: + result = _copy_mol_and_add_conformers( + target, n_conformers=10, starting_conformers=sdf + ) + + mock_generate.assert_not_called() + assert len(result.conformers) == source.n_conformers + + def test_default_path_still_uses_etkdg(self): + """With no supplied conformers, ETKDG is still used.""" + from unittest.mock import patch + + from presto.sample import _copy_mol_and_add_conformers + + mol = Molecule.from_smiles("CCCC") + mol.generate_conformers(n_conformers=1) + + with patch.object( + Molecule, "generate_conformers", autospec=True + ) as mock_generate: + _copy_mol_and_add_conformers(mol, n_conformers=3) + + mock_generate.assert_called_once() + class TestGetMoleculeFromDataset: """Tests for _get_molecule_from_dataset function.""" @@ -1192,6 +1254,53 @@ def test_sample_mmmd_single_molecule_minimal(self, tmp_path): assert len(entry["energy_weights"]) == len(entry["energy"]) assert len(entry["forces_weights"]) == len(entry["energy"]) + def test_sample_mmmd_uses_starting_conformers( + self, tmp_path, write_multiconformer_sdf + ): + """sample_mmmd starts from a supplied SDF's conformers, ignoring n_conformers.""" + from openff.units import unit as off_unit + + # Build a 3-conformer SDF for the molecule being sampled. + source = Molecule.from_smiles("CCCCCO") + source.generate_conformers(n_conformers=3, rms_cutoff=0.0 * off_unit.angstrom) + n_supplied = source.n_conformers + assert n_supplied > 1 + sdf = tmp_path / "confs.sdf" + write_multiconformer_sdf(source, sdf) + + mol = Molecule.from_smiles("CCCCCO") + ff = ForceField("openff_unconstrained-2.3.0.offxml") + + settings_obj = MMMDSamplingSettings( + sampling_protocol="mm_md", + timestep=1.0 * omm_unit.femtoseconds, + temperature=300.0 * omm_unit.kelvin, + n_conformers=1, # deliberately != number supplied; must be ignored + starting_conformers=sdf, + equilibration_sampling_time_per_conformer=0.0 * omm_unit.picoseconds, + production_sampling_time_per_conformer=0.002 * omm_unit.picoseconds, + snapshot_interval=0.001 * omm_unit.picoseconds, # 2 snapshots/conformer + ) + + pdb_dir = tmp_path / "pdb" + pdb_dir.mkdir() + output_paths = {OutputType.PDB_TRAJECTORY: pdb_dir} + + with patch("presto.sample.mlp.get_ml_omm_system") as mock_ml_sys: + mock_system = openmm.System() + for _ in range(mol.n_atoms): + mock_system.addParticle(12.0) + mock_system.addForce(openmm.CustomExternalForce("0")) + mock_ml_sys.return_value = mock_system + + result = sample_mmmd( + [mol], ff, torch.device("cpu"), settings_obj, output_paths + ) + + # 2 snapshots per conformer * the number of supplied conformers (not n_conformers). + expected_snapshots = 2 * n_supplied + assert len(result[0][0]["energy"]) == expected_snapshots + def test_sample_mmmd_with_pdb_output(self, tmp_path): """Test sample_mmmd with PDB trajectory output.""" mol = Molecule.from_smiles("C") # Methane diff --git a/presto/tests/unit/test_settings.py b/presto/tests/unit/test_settings.py index 0dc0482..db20272 100644 --- a/presto/tests/unit/test_settings.py +++ b/presto/tests/unit/test_settings.py @@ -894,3 +894,86 @@ def test_outlier_filter_settings_yaml_round_trip(self, tmp_path): assert loaded.outlier_filter_settings.energy_outlier_threshold == 7.5 assert loaded.outlier_filter_settings.force_outlier_threshold == 30.0 assert loaded.outlier_filter_settings.min_conformations == 2 + + +def _write_single_conformer_sdf(smiles: str, path: Path) -> None: + """Write a molecule with one conformer to an SDF file.""" + molecule = Molecule.from_smiles(smiles) + molecule.generate_conformers(n_conformers=1) + molecule.to_file(str(path), "SDF") + + +class TestStartingConformersField: + """Tests for the optional starting_conformers field on sampling and MSM settings.""" + + def test_default_is_none(self): + """starting_conformers defaults to None (ETKDG) on every stage.""" + assert MMMDSamplingSettings().starting_conformers is None + assert MLMDSamplingSettings().starting_conformers is None + assert MSMSettings().starting_conformers is None + + def test_accepts_existing_sdf(self, tmp_path): + """A valid SDF path is accepted.""" + sdf = tmp_path / "confs.sdf" + _write_single_conformer_sdf("CCO", sdf) + + settings = MMMDSamplingSettings(starting_conformers=sdf) + assert settings.starting_conformers == sdf + + msm = MSMSettings(starting_conformers=sdf) + assert msm.starting_conformers == sdf + + def test_missing_file_rejected(self, tmp_path): + """A missing SDF path is rejected at construction.""" + with pytest.raises(ValidationError, match="does not exist"): + MMMDSamplingSettings(starting_conformers=tmp_path / "missing.sdf") + + def test_non_sdf_suffix_rejected(self, tmp_path): + """A non-.sdf path is rejected at construction.""" + other = tmp_path / "confs.mol2" + other.write_text("") + with pytest.raises(ValidationError, match="must be an SDF file"): + MMMDSamplingSettings(starting_conformers=other) + + def test_workflow_accepts_matching_conformers(self, tmp_path): + """WorkflowSettings validates when the SDF contains the fitted molecule.""" + sdf = tmp_path / "confs.sdf" + _write_single_conformer_sdf("CCO", sdf) + + settings = WorkflowSettings( + param_settings=ParamSettings(molecule_input_type="smiles", molecules="CCO"), + device_type="cpu", + training_sampling_settings=MMMDSamplingSettings(starting_conformers=sdf), + ) + assert settings.training_sampling_settings.starting_conformers == sdf + + def test_workflow_rejects_missing_molecule(self, tmp_path): + """WorkflowSettings fails fast when a fitted molecule is absent from the SDF.""" + sdf = tmp_path / "confs.sdf" + _write_single_conformer_sdf("CCO", sdf) + + with pytest.raises(ValidationError, match="starting_conformers"): + WorkflowSettings( + param_settings=ParamSettings( + molecule_input_type="smiles", molecules="c1ccccc1" + ), + device_type="cpu", + training_sampling_settings=MMMDSamplingSettings( + starting_conformers=sdf + ), + ) + + def test_workflow_validates_msm_conformers(self, tmp_path): + """The MSM starting_conformers path is also cross-checked against molecules.""" + sdf = tmp_path / "confs.sdf" + _write_single_conformer_sdf("CCO", sdf) + + with pytest.raises(ValidationError, match=r"msm_settings\.starting_conformers"): + WorkflowSettings( + param_settings=ParamSettings( + molecule_input_type="smiles", + molecules="c1ccccc1", + msm_settings=MSMSettings(starting_conformers=sdf), + ), + device_type="cpu", + ) From 6835b1ea81e9154613eff09d3ff628925e088484 Mon Sep 17 00:00:00 2001 From: Finlay Clark Date: Mon, 13 Jul 2026 17:05:43 +0100 Subject: [PATCH 2/2] Update changelog --- docs/changelog.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/changelog.md b/docs/changelog.md index 27ff662..14c00bc 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -4,7 +4,7 @@ ### Features -- Add an optional `starting_conformers` setting to each sampling stage (`training_sampling_settings`, `testing_sampling_settings`) and to `msm_settings`. When set to an SDF path, that stage starts from the supplied conformers (matched to each molecule by graph isomorphism and realigned automatically) instead of generating them with ETKDG; `n_conformers` is ignored for that stage. The default remains ETKDG. +- Add an optional `starting_conformers` setting to each sampling stage (`training_sampling_settings`, `testing_sampling_settings`) and to `msm_settings`. When set to an SDF path, that stage starts from the supplied conformers (matched to each molecule by graph isomorphism and realigned automatically) instead of generating them with ETKDG; `n_conformers` is ignored for that stage. The default remains ETKDG.In [#78](https://github.com/cole-group/presto/pull/78). - Add `presto.create_types.add_library_charges_to_forcefield` to write custom partial charges from OpenFF `Molecule` objects into a force field as library charges, addressing [#64](https://github.com/cole-group/presto/issues/64). ### Fixes