Let pKa/redox helpers consume conformer ensembles via EnsembleThermo

ThermoScreening/thermo/__init__.py

@@ -9,6 +9,6 @@
 from .screening import screen, rank_by_gibbs
 from .conformers import generate as generate_conformers, write_conformers
 from .reactions import reaction_free_energy, reduction_potential
-from .ensemble import boltzmann_weights, ensemble_free_energy, lowest_gibbs
+from .ensemble import boltzmann_weights, ensemble_free_energy, lowest_gibbs, EnsembleThermo
 from .kinetics import eyring_rate_constant, wigner_tunneling_correction
 from .pka import pKa, calibrate_proton_reference, PROTON_AQUEOUS_FREE_ENERGY_KCAL

ThermoScreening/thermo/ensemble.py

@@ -117,3 +117,35 @@ def lowest_gibbs(thermos):
     if not thermos:
         raise ValueError("thermos must contain at least one conformer.")
     return min(thermos, key=lambda thermo: thermo.total_EeGtot())
+
+
+class EnsembleThermo:
+    """
+    A ``Thermo``-like adapter exposing a conformer ensemble's Boltzmann free
+    energy through ``total_EeGtot()``.
+
+    ``pKa``, ``calibrate_proton_reference``, ``reduction_potential``, and
+    ``reaction_free_energy`` only ever call ``.total_EeGtot()`` on the species
+    they're given, so wrapping a conformer ensemble in this class lets it be
+    passed anywhere a single ``Thermo`` is expected today -- no changes needed
+    to those helpers.
+
+    Parameters
+    ----------
+    thermos : iterable of Thermo
+        Conformers of the same species, computed with matching settings.
+    temperature : float
+        Temperature in kelvin, matching how ``thermos`` were computed and any
+        downstream helper it's passed to. Default 298.15.
+
+    Raises
+    ------
+    ValueError
+        If ``thermos`` is empty or ``temperature`` is not positive.
+    """
+
+    def __init__(self, thermos, temperature=298.15):
+        self._total_eegtot = ensemble_free_energy(thermos, temperature=temperature, unit="H")
+
+    def total_EeGtot(self):
+        return self._total_eegtot

docs/api.rst

@@ -51,6 +51,8 @@ Conformer ensembles
 .. autofunction:: ThermoScreening.thermo.ensemble.boltzmann_weights
 .. autofunction:: ThermoScreening.thermo.ensemble.ensemble_free_energy
 .. autofunction:: ThermoScreening.thermo.ensemble.lowest_gibbs
+.. autoclass:: ThermoScreening.thermo.ensemble.EnsembleThermo
+   :members:
 
 Thermochemistry core
 --------------------

docs/usage.rst

@@ -101,6 +101,30 @@ The ensemble free energy lies at or below the lowest conformer's by the mixing
 (conformational) entropy; use ``lowest_gibbs`` when you instead want the single
 dominant structure to carry forward (e.g. into :func:`reaction_free_energy`).
 
+``pKa``, ``calibrate_proton_reference``, ``reduction_potential``, and
+``reaction_free_energy`` all take whatever they're given and call
+``.total_EeGtot()`` on it -- they don't care whether that's a single ``Thermo``
+or something else with the same method. ``EnsembleThermo`` wraps a conformer
+ensemble's Boltzmann free energy behind that same method, so a flexible acid
+or base can be passed to ``pKa`` as an ensemble directly, instead of picking
+one (possibly non-representative) conformer:
+
+.. code-block:: python
+
+   from ThermoScreening.thermo.api import xtb_cli_thermo
+   from ThermoScreening.thermo.conformers import generate
+   from ThermoScreening.thermo import EnsembleThermo, pKa
+
+   def ensemble_thermo(smiles, charge):
+       conformers = generate(smiles, max_conformers=10)
+       thermos = [xtb_cli_thermo(c, charge=charge, solvent="water") for c in conformers]
+       return EnsembleThermo(thermos)
+
+   acid = ensemble_thermo("OCCCC(=O)O", 0)     # 4-hydroxybutanoic acid
+   base = ensemble_thermo("OCCCC(=O)[O-]", -1)
+
+   p_ka = pKa(acid, base)
+
 DFT-quality thermochemistry (ORCA)
 ----------------------------------
 

tests/thermo/test_ensemble.py

@@ -1,13 +1,18 @@
 import math
+import os
+import shutil
 
 import pytest
 
 from ThermoScreening.thermo.ensemble import (
     boltzmann_weights,
     ensemble_free_energy,
     lowest_gibbs,
+    EnsembleThermo,
 )
 
+xtb_available = shutil.which("xtb") is not None or "XTB_COMMAND" in os.environ
+
 _KB_HARTREE_PER_K = 1.380649e-23 / 4.3597447222071e-18  # k_B in Hartree/K
 _H_TO_KCAL = 627.5094740631
 
@@ -96,12 +101,112 @@ def test_non_positive_temperature_raises(func, temperature):
         func([_FakeThermo(0.0)], temperature=temperature)
 
 
+def test_ensemble_thermo_matches_ensemble_free_energy():
+    thermos = [_FakeThermo(-1.0), _FakeThermo(-1.2), _FakeThermo(-0.8)]
+
+    wrapped = EnsembleThermo(thermos)
+
+    assert wrapped.total_EeGtot() == pytest.approx(ensemble_free_energy(thermos, unit="H"))
+
+
+def test_ensemble_thermo_single_conformer_is_identity():
+    assert EnsembleThermo([_FakeThermo(-2.0)]).total_EeGtot() == pytest.approx(-2.0)
+
+
+def test_ensemble_thermo_forwards_temperature():
+    thermos = [_FakeThermo(-1.0), _FakeThermo(-1.2)]
+
+    wrapped = EnsembleThermo(thermos, temperature=350.0)
+
+    assert wrapped.total_EeGtot() == pytest.approx(
+        ensemble_free_energy(thermos, temperature=350.0, unit="H")
+    )
+
+
+def test_ensemble_thermo_empty_raises():
+    with pytest.raises(ValueError, match="at least one conformer"):
+        EnsembleThermo([])
+
+
+def test_ensemble_thermo_drops_into_pKa():
+    from ThermoScreening.thermo.pka import pKa
+
+    acid_thermos = [_FakeThermo(-10.0), _FakeThermo(-9.999)]
+    base_thermos = [_FakeThermo(-10.0 + 0.01), _FakeThermo(-9.995)]
+
+    ensemble_result = pKa(
+        EnsembleThermo(acid_thermos), EnsembleThermo(base_thermos), reference_free_energy=0.0
+    )
+    manual_result = pKa(
+        _FakeThermo(ensemble_free_energy(acid_thermos, unit="H")),
+        _FakeThermo(ensemble_free_energy(base_thermos, unit="H")),
+        reference_free_energy=0.0,
+    )
+
+    assert ensemble_result == pytest.approx(manual_result)
+
+
+def test_ensemble_thermo_drops_into_reduction_potential():
+    from ThermoScreening.thermo.reactions import reduction_potential
+
+    oxidized_thermos = [_FakeThermo(-10.0), _FakeThermo(-9.999)]
+    reduced_thermos = [_FakeThermo(-10.01), _FakeThermo(-10.005)]
+
+    ensemble_result = reduction_potential(
+        EnsembleThermo(oxidized_thermos), EnsembleThermo(reduced_thermos)
+    )
+    manual_result = reduction_potential(
+        _FakeThermo(ensemble_free_energy(oxidized_thermos, unit="H")),
+        _FakeThermo(ensemble_free_energy(reduced_thermos, unit="H")),
+    )
+
+    assert ensemble_result == pytest.approx(manual_result)
+
+
 def test_public_api_exported():
     from ThermoScreening.thermo import boltzmann_weights as bw
     from ThermoScreening.thermo import ensemble_free_energy as efe
     from ThermoScreening.thermo import lowest_gibbs as lg
+    from ThermoScreening.thermo import EnsembleThermo as et
     from ThermoScreening.thermo import ensemble
 
     assert bw is ensemble.boltzmann_weights
     assert efe is ensemble.ensemble_free_energy
     assert lg is ensemble.lowest_gibbs
+    assert et is ensemble.EnsembleThermo
+
+
+@pytest.mark.skipif(not xtb_available, reason="the native xtb binary is not available.")
+def test_ensemble_thermo_end_to_end_with_pKa(tmp_path):
+    # 4-hydroxybutanoic acid has a flexible C-C-C-C backbone (unlike e.g.
+    # glycolic acid, whose conformers RDKit's RMSD pruning collapses to one)
+    # -- generate() gives it several genuinely distinct conformers, so this
+    # exercises real Boltzmann averaging, not a degenerate single-conformer
+    # ensemble. EnsembleThermo wraps the real xtb-cli energies over them and
+    # duck-types as a Thermo, so it drops straight into pKa() with no other
+    # changes.
+    from ThermoScreening.thermo.api import xtb_cli_thermo
+    from ThermoScreening.thermo.conformers import generate
+    from ThermoScreening.thermo.pka import pKa
+
+    def ensemble_thermo(smiles, charge, tag):
+        conformers = generate(smiles, max_conformers=5)
+        assert len(conformers) > 1  # otherwise this isn't testing an ensemble
+        thermos = [
+            xtb_cli_thermo(
+                conf, charge=charge, solvent="water", directory=str(tmp_path / f"{tag}_{i}")
+            )
+            for i, conf in enumerate(conformers)
+        ]
+        return thermos, EnsembleThermo(thermos)
+
+    acid_thermos, acid_ensemble = ensemble_thermo("OCCCC(=O)O", 0, "acid")
+    base_thermos, base_ensemble = ensemble_thermo("OCCCC(=O)[O-]", -1, "base")
+
+    # the ensemble free energy is at or below every individual conformer's
+    assert acid_ensemble.total_EeGtot() <= min(t.total_EeGtot() for t in acid_thermos)
+    assert base_ensemble.total_EeGtot() <= min(t.total_EeGtot() for t in base_thermos)
+
+    ensemble_pKa = pKa(acid_ensemble, base_ensemble)
+
+    assert math.isfinite(ensemble_pKa)