Add generate_thermo_ensemble: charge-state-aware conformers with auto-retry

ThermoScreening/thermo/__init__.py

@@ -7,7 +7,7 @@
 from .system import System
 from .thermo import Thermo
 from .screening import screen, rank_by_gibbs
-from .conformers import generate as generate_conformers, write_conformers
+from .conformers import generate as generate_conformers, write_conformers, generate_thermo_ensemble
 from .reactions import reaction_free_energy, reduction_potential
 from .ensemble import boltzmann_weights, ensemble_free_energy, lowest_gibbs, EnsembleThermo
 from .kinetics import eyring_rate_constant, wigner_tunneling_correction

ThermoScreening/thermo/conformers.py

@@ -8,6 +8,8 @@
 
 import numpy as np
 
+from ..exceptions import TSValueError
+
 
 def _import_rdkit():
     """Import RDKit lazily, with a clear message if it is not installed."""
@@ -101,6 +103,88 @@ def generate(
     return [_conformer_to_atoms(molecule, cid) for cid in conformer_ids]
 
 
+def generate_thermo_ensemble(
+    smiles,
+    thermo_fn,
+    charge=0.0,
+    max_conformers=10,
+    max_attempts=5,
+    prune_rms_thresh=0.5,
+    energy_window=None,
+    random_seed=42,
+    **thermo_kwargs,
+):
+    """
+    Generate a conformer ensemble and compute ``Thermo`` for each, retrying
+    past saddle points.
+
+    A conformer that fails to converge to a true minimum (``TSValueError``,
+    e.g. an imaginary frequency) or whose engine process fails outright
+    (``RuntimeError``, e.g. an xtb CLI non-zero exit) is skipped; if fewer
+    than ``max_conformers`` succeed, :func:`generate` is re-run with a new
+    random seed, up to ``max_attempts`` seeds total. This is the loop a
+    charged-species ensemble (radical anion, dianion, ...) otherwise
+    requires by hand before it can be passed to :class:`EnsembleThermo`.
+
+    Parameters
+    ----------
+    smiles : str
+        The molecule as a SMILES string.
+    thermo_fn : callable
+        A ``Thermo``-computing engine, e.g. ``xtb_cli_thermo``, called as
+        ``thermo_fn(atoms, charge=charge, **thermo_kwargs)``.
+    charge : float
+        System charge, forwarded to ``thermo_fn``. Default 0.0.
+    max_conformers : int
+        Target number of successfully-computed conformers. Default 10.
+    max_attempts : int
+        Number of distinct random seeds (each embedding up to
+        ``max_conformers`` conformers) to try before giving up. Default 5.
+    prune_rms_thresh, energy_window : see :func:`generate`.
+    random_seed : int
+        Random seed for the first attempt; later attempts use
+        ``random_seed + 1``, ``random_seed + 2``, etc. Default 42.
+    **thermo_kwargs
+        Forwarded to ``thermo_fn`` (e.g. ``solvent="water"``).
+
+    Returns
+    -------
+    list of Thermo
+        The successfully computed conformers (fewer than ``max_conformers``
+        if ``max_attempts`` is exhausted first).
+
+    Raises
+    ------
+    ValueError
+        If no conformer converges to a true minimum within ``max_attempts``.
+    """
+    thermos = []
+    for attempt in range(max_attempts):
+        atoms_list = generate(
+            smiles,
+            max_conformers=max_conformers,
+            prune_rms_thresh=prune_rms_thresh,
+            energy_window=energy_window,
+            random_seed=random_seed + attempt,
+        )
+        for atoms in atoms_list:
+            if len(thermos) >= max_conformers:
+                break
+            try:
+                thermos.append(thermo_fn(atoms, charge=charge, **thermo_kwargs))
+            except (TSValueError, RuntimeError):
+                continue
+        if len(thermos) >= max_conformers:
+            break
+
+    if not thermos:
+        raise ValueError(
+            f"No conformer of {smiles!r} converged to a true minimum after "
+            f"{max_attempts} attempts."
+        )
+    return thermos
+
+
 def write_conformers(conformers, directory, prefix="conformer"):
     """
     Write conformers to ``directory`` as ``<prefix>_<i>.xyz`` files.

docs/api.rst

@@ -25,6 +25,7 @@ Conformer generation
 --------------------
 
 .. autofunction:: ThermoScreening.thermo.conformers.generate
+.. autofunction:: ThermoScreening.thermo.conformers.generate_thermo_ensemble
 .. autofunction:: ThermoScreening.thermo.conformers.write_conformers
 
 Reactions and redox

docs/usage.rst

@@ -125,6 +125,27 @@ one (possibly non-representative) conformer:
 
    p_ka = pKa(acid, base)
 
+In practice, some conformers -- especially of a charged species -- optimise
+onto a saddle point rather than a true minimum instead of failing outright.
+``generate_thermo_ensemble`` wraps the loop above with automatic retries: it
+generates conformers, computes ``Thermo`` for each, skips any that raise
+(imaginary frequencies, or the engine's process failing outright), and
+reseeds and retries until enough succeed:
+
+.. code-block:: python
+
+   from ThermoScreening.thermo.api import xtb_cli_thermo
+   from ThermoScreening.thermo import generate_thermo_ensemble, EnsembleThermo, pKa
+
+   acid_thermos = generate_thermo_ensemble(
+       "OCCCC(=O)O", xtb_cli_thermo, charge=0, max_conformers=10, solvent="water",
+   )
+   base_thermos = generate_thermo_ensemble(
+       "OCCCC(=O)[O-]", xtb_cli_thermo, charge=-1, max_conformers=10, solvent="water",
+   )
+
+   p_ka = pKa(EnsembleThermo(acid_thermos), EnsembleThermo(base_thermos))
+
 DFT-quality thermochemistry (ORCA)
 ----------------------------------
 

tests/thermo/test_conformers.py

@@ -1,9 +1,25 @@
+import os
+import shutil
+
 import numpy as np
 import pytest
 from ase import Atoms
 
+from ThermoScreening.exceptions import TSValueError
 from ThermoScreening.thermo import conformers
 
+xtb_available = shutil.which("xtb") is not None or "XTB_COMMAND" in os.environ
+
+
+class _FakeThermo:
+    """A stand-in exposing only what a Thermo consumer needs to see it succeeded."""
+
+    def __init__(self, eegtot=-1.0):
+        self._eegtot = eegtot
+
+    def total_EeGtot(self):
+        return self._eegtot
+
 
 def test_generate_returns_ase_conformers():
     # n-butane is flexible enough to embed at least one conformer
@@ -48,7 +64,134 @@ def test_write_conformers_writes_readable_xyz(tmp_path):
 
 
 def test_public_api_is_exported():
-    from ThermoScreening.thermo import generate_conformers, write_conformers
+    from ThermoScreening.thermo import generate_conformers, write_conformers, generate_thermo_ensemble
 
     assert generate_conformers is conformers.generate
     assert write_conformers is conformers.write_conformers
+    assert generate_thermo_ensemble is conformers.generate_thermo_ensemble
+
+
+def test_generate_thermo_ensemble_returns_up_to_max_conformers():
+    def thermo_fn(atoms, charge, **kwargs):
+        return _FakeThermo()
+
+    result = conformers.generate_thermo_ensemble("CCCCCC", thermo_fn, max_conformers=3)
+
+    assert 1 <= len(result) <= 3
+    assert all(isinstance(t, _FakeThermo) for t in result)
+
+
+def test_generate_thermo_ensemble_forwards_charge_and_kwargs():
+    seen = []
+
+    def thermo_fn(atoms, charge, **kwargs):
+        seen.append((charge, kwargs))
+        return _FakeThermo()
+
+    conformers.generate_thermo_ensemble("CCO", thermo_fn, charge=-1.0, max_conformers=1, solvent="water")
+
+    assert seen == [(-1.0, {"solvent": "water"})]
+
+
+def test_generate_thermo_ensemble_default_charge_is_zero():
+    seen_charges = []
+
+    def thermo_fn(atoms, charge, **kwargs):
+        seen_charges.append(charge)
+        return _FakeThermo()
+
+    conformers.generate_thermo_ensemble("CCO", thermo_fn, max_conformers=1)
+
+    assert seen_charges == [0.0]
+
+
+def test_generate_thermo_ensemble_skips_ts_value_error_and_retries():
+    call_count = [0]
+
+    def thermo_fn(atoms, charge, **kwargs):
+        call_count[0] += 1
+        if call_count[0] <= 2:
+            raise TSValueError("Imaginary (non-positive) vibrational frequencies are present")
+        return _FakeThermo()
+
+    result = conformers.generate_thermo_ensemble("CCCCCC", thermo_fn, max_conformers=2, max_attempts=5)
+
+    assert len(result) == 2
+    assert call_count[0] > 2  # the first two failures were skipped and retried past
+
+
+def test_generate_thermo_ensemble_skips_runtime_error():
+    call_count = [0]
+
+    def thermo_fn(atoms, charge, **kwargs):
+        call_count[0] += 1
+        if call_count[0] == 1:
+            raise RuntimeError("xtb failed (exit 1)")
+        return _FakeThermo()
+
+    result = conformers.generate_thermo_ensemble("CCCCCC", thermo_fn, max_conformers=1, max_attempts=3)
+
+    assert len(result) == 1
+
+
+def test_generate_thermo_ensemble_stops_once_target_reached_mid_attempt(monkeypatch):
+    # a single attempt's conformer list can be longer than what's still
+    # needed (e.g. after an earlier attempt partially succeeded) -- the
+    # inner loop must stop as soon as the target is hit, not exhaust the list
+    dummy_atoms = [Atoms("H") for _ in range(3)]
+
+    def fake_generate(smiles, max_conformers, prune_rms_thresh, energy_window, random_seed):
+        return dummy_atoms
+
+    monkeypatch.setattr(conformers, "generate", fake_generate)
+
+    call_count = [0]
+
+    def thermo_fn(atoms, charge, **kwargs):
+        call_count[0] += 1
+        return _FakeThermo()
+
+    result = conformers.generate_thermo_ensemble("CCCC", thermo_fn, max_conformers=2, max_attempts=1)
+
+    assert len(result) == 2
+    assert call_count[0] == 2  # the 3rd dummy conformer was never processed
+
+
+def test_generate_thermo_ensemble_raises_when_nothing_succeeds():
+    def thermo_fn(atoms, charge, **kwargs):
+        raise TSValueError("Imaginary (non-positive) vibrational frequencies are present")
+
+    with pytest.raises(ValueError, match="No conformer"):
+        conformers.generate_thermo_ensemble("CCCC", thermo_fn, max_conformers=2, max_attempts=2)
+
+
+def test_generate_thermo_ensemble_propagates_other_exceptions():
+    def thermo_fn(atoms, charge, **kwargs):
+        raise KeyError("not a recognised failure mode")
+
+    with pytest.raises(KeyError):
+        conformers.generate_thermo_ensemble("CCCC", thermo_fn, max_conformers=1)
+
+
+@pytest.mark.skipif(not xtb_available, reason="the native xtb binary is not available.")
+def test_generate_thermo_ensemble_end_to_end_with_real_xtb(tmp_path):
+    # 4-hydroxybutanoic acid's flexible backbone means naive conformer
+    # generation occasionally lands a starting geometry on a saddle point
+    # (TSValueError) at the anion charge state -- this is exactly the
+    # manual retry-by-hand loop this helper replaces.
+    from ThermoScreening.thermo.api import xtb_cli_thermo
+    from ThermoScreening.thermo.ensemble import EnsembleThermo
+
+    call_count = [0]
+
+    def thermo_fn(atoms, charge, **kwargs):
+        call_count[0] += 1
+        return xtb_cli_thermo(atoms, charge=charge, directory=str(tmp_path / str(call_count[0])), **kwargs)
+
+    thermos = conformers.generate_thermo_ensemble(
+        "OCCCC(=O)[O-]", thermo_fn, charge=-1, max_conformers=3, solvent="water"
+    )
+
+    assert 1 <= len(thermos) <= 3
+    ensemble = EnsembleThermo(thermos)
+    assert ensemble.total_EeGtot() <= min(t.total_EeGtot() for t in thermos)