Add conformer generation workflow (#30)

README.md

@@ -124,7 +124,10 @@ DFTB+ integration tests run only when the executables are available and `DFTB_PR
 
 ## Roadmap
 
-Planned work is tracked in GitHub issues rather than in this README. Current roadmap areas include additional engines, conformer generation, broader test coverage, documentation, and batch screening workflows.
+Planned work is tracked in GitHub issues rather than in this README. The tool
+supports the DFTB+, GFN-xTB (tblite) and native-xtb engines, implicit solvation,
+quasi-RRHO, batch screening with resume, and RDKit conformer generation; see the
+issue tracker for further enhancements.
 
 ## License
 

ThermoScreening/cli/thermo.py

@@ -13,10 +13,11 @@
 )
 from ThermoScreening.thermo.api import execute
 from ThermoScreening.thermo.screening import screen
+from ThermoScreening.thermo.conformers import generate as generate_conformers, write_conformers
 from ThermoScreening.version import __version__
 
 
-SUBCOMMANDS = {"setup-dftb", "doctor", "screen"}
+SUBCOMMANDS = {"setup-dftb", "doctor", "screen", "conformers"}
 
 
 def _run_parser():
@@ -128,6 +129,29 @@ def _command_parser():
         "the missing/failed molecules.",
     )
 
+    conf_parser = subparsers.add_parser(
+        "conformers",
+        help="Generate conformers from a SMILES (RDKit ETKDG) and write them as "
+        ".xyz files ready to screen.",
+    )
+    conf_parser.add_argument("smiles", help="Molecule as a SMILES string.")
+    conf_parser.add_argument(
+        "-o", "--out-dir", default="conformers",
+        help="Directory to write <prefix>_<i>.xyz files (default 'conformers').",
+    )
+    conf_parser.add_argument(
+        "--max-conformers", type=int, default=10,
+        help="Maximum number of conformers to embed (default 10).",
+    )
+    conf_parser.add_argument(
+        "--energy-window", type=float, default=None,
+        help="Keep only conformers within this many kcal/mol of the lowest.",
+    )
+    conf_parser.add_argument(
+        "--no-optimize", action="store_true",
+        help="Skip MMFF force-field optimisation of the embedded conformers.",
+    )
+
     return parser
 
 
@@ -222,6 +246,28 @@ def run_screen(parser_args):
     return 1 if failed else 0
 
 
+def run_conformers(parser_args):
+    """
+    Generate conformers from a SMILES and write them as .xyz files.
+    """
+
+    try:
+        conformers = generate_conformers(
+            parser_args.smiles,
+            max_conformers=parser_args.max_conformers,
+            optimize=not parser_args.no_optimize,
+            energy_window=parser_args.energy_window,
+        )
+    except (ValueError, ImportError) as exc:
+        print(f"Conformer generation failed: {exc}", file=sys.stderr)
+        return 1
+
+    paths = write_conformers(conformers, parser_args.out_dir)
+    print(f"Generated {len(paths)} conformer(s) in {parser_args.out_dir}/")
+    print(f"Screen them with: thermo screen {parser_args.out_dir}")
+    return 0
+
+
 def main():
     """
     Main function to run the thermo cli. It parses the command line arguments
@@ -249,6 +295,9 @@ def main():
     if command == "screen":
         return run_screen(parser_args)
 
+    if command == "conformers":
+        return run_conformers(parser_args)
+
     input_file = parser_args.input_file
     verbose = parser_args.verbose
 

ThermoScreening/thermo/__init__.py

@@ -7,3 +7,4 @@
 from .system import System
 from .thermo import Thermo
 from .screening import screen
+from .conformers import generate as generate_conformers, write_conformers

ThermoScreening/thermo/conformers.py

@@ -0,0 +1,121 @@
+"""Conformer generation for molecule screening.
+
+RDKit is the conformer backend: ETKDG distance-geometry embedding, optional
+MMFF force-field optimisation, and RMSD + energy-window pruning. Conformers are
+returned as ASE ``Atoms`` so they feed straight into the screening pipeline
+(e.g. generate -> ``write_conformers`` -> ``screen`` that directory).
+"""
+
+import numpy as np
+
+
+def _import_rdkit():
+    """Import RDKit lazily, with a clear message if it is not installed."""
+    try:
+        from rdkit import Chem
+        from rdkit.Chem import AllChem
+    except ImportError as exc:  # pragma: no cover - exercised only without rdkit
+        raise ImportError(
+            "Conformer generation requires RDKit. Install it with "
+            "`conda install -c conda-forge rdkit` or `pip install rdkit`."
+        ) from exc
+    return Chem, AllChem
+
+
+def _conformer_to_atoms(mol, conformer_id):
+    from ase import Atoms
+
+    conformer = mol.GetConformer(conformer_id)
+    symbols = [atom.GetSymbol() for atom in mol.GetAtoms()]
+    positions = np.asarray(conformer.GetPositions(), dtype=float)
+    return Atoms(symbols=symbols, positions=positions)
+
+
+def generate(
+    smiles,
+    max_conformers=10,
+    optimize=True,
+    prune_rms_thresh=0.5,
+    energy_window=None,
+    random_seed=42,
+):
+    """
+    Generate conformers for a molecule from its SMILES.
+
+    Parameters
+    ----------
+    smiles : str
+        The molecule as a SMILES string.
+    max_conformers : int
+        Maximum number of conformers to embed (ETKDG). Default 10.
+    optimize : bool
+        If True, optimise each conformer with the MMFF force field and sort the
+        results by MMFF energy (lowest first). Default True.
+    prune_rms_thresh : float
+        RMSD threshold (A) for pruning duplicate embedded conformers. Default 0.5.
+    energy_window : float, optional
+        If given (and ``optimize``), keep only conformers within this many
+        kcal/mol of the lowest-energy conformer.
+    random_seed : int
+        Random seed for reproducible embedding. Default 42.
+
+    Returns
+    -------
+    list of ase.Atoms
+        The generated conformers (energy-sorted when optimised).
+
+    Raises
+    ------
+    ImportError
+        If RDKit is not installed.
+    ValueError
+        If the SMILES cannot be parsed or no conformer could be embedded.
+    """
+    Chem, AllChem = _import_rdkit()
+
+    molecule = Chem.MolFromSmiles(smiles)
+    if molecule is None:
+        raise ValueError(f"Could not parse SMILES: {smiles!r}")
+    molecule = Chem.AddHs(molecule)
+
+    params = AllChem.ETKDGv3()
+    params.randomSeed = random_seed
+    params.pruneRmsThresh = prune_rms_thresh
+    conformer_ids = list(AllChem.EmbedMultipleConfs(molecule, numConfs=max_conformers, params=params))
+    if not conformer_ids:
+        raise ValueError(f"Could not embed any conformer for SMILES: {smiles!r}")
+
+    if optimize:
+        results = AllChem.MMFFOptimizeMoleculeConfs(molecule)
+        energies = [energy for _converged, energy in results]
+        order = sorted(range(len(conformer_ids)), key=lambda i: energies[i])
+        conformer_ids = [conformer_ids[i] for i in order]
+        energies = [energies[i] for i in order]
+        if energy_window is not None:
+            lowest = energies[0]
+            conformer_ids = [
+                cid for cid, energy in zip(conformer_ids, energies)
+                if energy - lowest <= energy_window
+            ]
+
+    return [_conformer_to_atoms(molecule, cid) for cid in conformer_ids]
+
+
+def write_conformers(conformers, directory, prefix="conformer"):
+    """
+    Write conformers to ``directory`` as ``<prefix>_<i>.xyz`` files.
+
+    Returns the list of written paths, so the directory can be passed straight to
+    :func:`ThermoScreening.thermo.screening.screen`.
+    """
+    import ase.io
+    from pathlib import Path
+
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    paths = []
+    for index, atoms in enumerate(conformers):
+        path = directory / f"{prefix}_{index}.xyz"
+        ase.io.write(str(path), atoms)
+        paths.append(path)
+    return paths

tests/thermo/test_conformers.py

@@ -0,0 +1,54 @@
+import numpy as np
+import pytest
+from ase import Atoms
+
+from ThermoScreening.thermo import conformers
+
+
+def test_generate_returns_ase_conformers():
+    # n-butane is flexible enough to embed at least one conformer
+    result = conformers.generate("CCCC", max_conformers=5)
+
+    assert 1 <= len(result) <= 5
+    assert all(isinstance(atoms, Atoms) for atoms in result)
+    assert all(atoms.get_chemical_formula() == "C4H10" for atoms in result)
+
+
+def test_generate_respects_max_conformers():
+    result = conformers.generate("CCCCCCCC", max_conformers=3)
+    assert len(result) <= 3
+
+
+def test_generate_energy_window_filters():
+    wide = conformers.generate("CCCCCC", max_conformers=15, energy_window=100.0)
+    narrow = conformers.generate("CCCCCC", max_conformers=15, energy_window=0.1)
+    assert len(narrow) <= len(wide)
+
+
+def test_generate_without_optimize_runs():
+    result = conformers.generate("CCO", max_conformers=3, optimize=False)
+    assert len(result) >= 1
+
+
+def test_generate_rejects_bad_smiles():
+    with pytest.raises(ValueError, match="Could not parse SMILES"):
+        conformers.generate("this-is-not-smiles")
+
+
+def test_write_conformers_writes_readable_xyz(tmp_path):
+    import ase.io
+
+    result = conformers.generate("CCO", max_conformers=2)
+    paths = conformers.write_conformers(result, tmp_path / "confs", prefix="c")
+
+    assert len(paths) == len(result)
+    assert all(p.name.startswith("c_") and p.suffix == ".xyz" for p in paths)
+    back = ase.io.read(str(paths[0]))
+    assert back.get_chemical_formula() == "C2H6O"
+
+
+def test_public_api_is_exported():
+    from ThermoScreening.thermo import generate_conformers, write_conformers
+
+    assert generate_conformers is conformers.generate
+    assert write_conformers is conformers.write_conformers

tests/thermo/test_main.py

@@ -168,6 +168,22 @@ def test_main_runs_doctor(monkeypatch, capsys):
     assert capsys.readouterr().out == "not ready\n"
 
 
+def test_main_runs_conformers(monkeypatch):
+    monkeypatch.setattr(
+        thermo, "parse_args", lambda: argparse.Namespace(command="conformers")
+    )
+    called = {}
+
+    def fake_run_conformers(args):
+        called["ran"] = True
+        return 0
+
+    monkeypatch.setattr(thermo, "run_conformers", fake_run_conformers)
+
+    assert thermo.main() == 0
+    assert called["ran"] is True
+
+
 def test_main_doctor_ignores_missing_optional_backend(monkeypatch, capsys):
     required_ok = type("D", (), {"ok": True, "optional": False})()
     optional_missing = type("D", (), {"ok": False, "optional": True})()

tests/thermo/test_screening.py

@@ -428,3 +428,61 @@ def test_cli_run_screen_returns_failure_count(monkeypatch):
     )
 
     assert cli.run_screen(args) == 1  # one molecule failed
+
+
+def test_cli_parse_args_routes_conformers():
+    import ThermoScreening.cli.thermo as cli
+
+    args = cli.parse_args(
+        ["conformers", "CCCC", "-o", "confs", "--max-conformers", "5",
+         "--energy-window", "2.0", "--no-optimize"]
+    )
+    assert args.command == "conformers"
+    assert args.smiles == "CCCC"
+    assert args.out_dir == "confs"
+    assert args.max_conformers == 5
+    assert args.energy_window == 2.0
+    assert args.no_optimize is True
+
+
+def test_cli_run_conformers(monkeypatch, tmp_path):
+    import ThermoScreening.cli.thermo as cli
+
+    captured = {}
+
+    def fake_generate(smiles, max_conformers=10, optimize=True, energy_window=None):
+        captured.update(smiles=smiles, max_conformers=max_conformers,
+                        optimize=optimize, energy_window=energy_window)
+        return ["conf_a", "conf_b"]
+
+    def fake_write(conformers, out_dir, prefix="conformer"):
+        captured["written"] = len(conformers)
+        return [tmp_path / "a.xyz", tmp_path / "b.xyz"]
+
+    monkeypatch.setattr(cli, "generate_conformers", fake_generate)
+    monkeypatch.setattr(cli, "write_conformers", fake_write)
+
+    args = Namespace(smiles="CCO", out_dir=str(tmp_path / "out"),
+                     max_conformers=7, energy_window=1.5, no_optimize=True)
+
+    assert cli.run_conformers(args) == 0
+    assert captured["smiles"] == "CCO"
+    assert captured["max_conformers"] == 7
+    assert captured["optimize"] is False  # --no-optimize
+    assert captured["energy_window"] == 1.5
+    assert captured["written"] == 2
+
+
+def test_cli_run_conformers_reports_bad_smiles(monkeypatch, capsys):
+    import ThermoScreening.cli.thermo as cli
+
+    def fail(*args, **kwargs):
+        raise ValueError("Could not parse SMILES: 'oops'")
+
+    monkeypatch.setattr(cli, "generate_conformers", fail)
+    args = Namespace(smiles="oops", out_dir="out", max_conformers=10,
+                     energy_window=None, no_optimize=False)
+
+    # clean exit code + message on stderr, not a traceback
+    assert cli.run_conformers(args) == 1
+    assert "Conformer generation failed" in capsys.readouterr().err