Add GBSA/ALPB implicit solvation to the DFTB+ pipeline

ThermoScreening/calculator/dftbplus.py

@@ -141,6 +141,62 @@ def _spin_kwargs(atoms, spin, spin_constants=SPIN_CONSTANTS_3OB):
     return kwargs
 
 
+def _solvation_kwargs(solvent=None, param_file=None, install_root=None):
+    """
+    ASE ``Dftb`` keyword arguments enabling GBSA/ALPB implicit solvation.
+
+    Returns an empty dict when neither ``solvent`` nor ``param_file`` is given,
+    so the gas-phase calculation is left exactly as before. Otherwise it enables
+    ``Solvation = GeneralizedBorn`` with the solvent's parameter file. DFTB+
+    includes the solvation term in the energy, gradient, and Hessian, so the
+    geometry, energy, and frequencies are all computed in solution consistently.
+
+    The parameter files (grimme-lab/gbsa-parameters) are fit for GFN-xTB; used
+    with the DFTB (3ob/mio) Hamiltonians they are an approximation. Pass an
+    explicit ``param_file`` to use a method-consistent set instead.
+
+    Parameters
+    ----------
+    solvent : str, optional
+        Solvent name (e.g. ``"water"``). Resolved to a downloaded parameter file.
+    param_file : str, optional
+        Explicit path to a GBSA parameter file (overrides ``solvent``).
+    install_root : str or Path, optional
+        Root used to locate downloaded solvent parameters.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the resolved parameter file does not exist.
+    """
+    if solvent is None and param_file is None:
+        return {}
+
+    if param_file is not None:
+        path = os.path.abspath(os.path.expanduser(str(param_file)))
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"GBSA solvation parameter file does not exist: {path}"
+            )
+    else:
+        # Lazy import: this is an environment/paths lookup, not a core dependency.
+        from ..cli.dftb_setup import gbsa_param_path
+
+        path = os.path.abspath(str(gbsa_param_path(solvent, install_root)))
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"No GBSA parameter file for solvent {solvent!r} at {path}. "
+                f"Download it first, e.g. `thermo setup-dftb --solvent {solvent}`."
+            )
+
+    # DFTB+ resolves ParamFile relative to the run directory, so pass an absolute
+    # path (jobs run in per-molecule working directories).
+    return {
+        "Hamiltonian_Solvation": "GeneralizedBorn {",
+        "Hamiltonian_Solvation_ParamFile": path,
+    }
+
+
 class Geoopt(Dftb):
     """
     Custom DFTB+ calculator to optimize the system with the 'GeometryOptimisation' driver (Rational).

ThermoScreening/cli/dftb_setup.py

@@ -167,6 +167,118 @@ def install_slakos(
     return parameter_dir.resolve()
 
 
+# GBSA/ALPB implicit-solvation parameter files (grimme-lab/gbsa-parameters). The
+# published sets are fit for GFN-xTB; used with DFTB (3ob/mio) they are an
+# approximation (see the note in ``calculator.dftbplus._solvation_kwargs``).
+GBSA_PARAM_METHOD = "gfn2-0-1"
+GBSA_BASE_URL = "https://raw.githubusercontent.com/grimme-lab/gbsa-parameters/main"
+
+# User-facing solvent name -> parameter-file stem in the repository.
+GBSA_SOLVENTS = {
+    "acetone": "acetone",
+    "acetonitrile": "acetonitrile",
+    "benzene": "benzene",
+    "ch2cl2": "ch2cl2",
+    "dichloromethane": "ch2cl2",
+    "chcl3": "chcl3",
+    "chloroform": "chcl3",
+    "cs2": "cs2",
+    "dmf": "dmf",
+    "dmso": "dmso",
+    "ether": "ether",
+    "diethylether": "ether",
+    "water": "h2o",
+    "h2o": "h2o",
+    "methanol": "methanol",
+    "hexane": "nhexane",
+    "nhexane": "nhexane",
+    "thf": "thf",
+    "toluene": "toluene",
+}
+
+
+def _solvent_stem(solvent: str) -> str:
+    """
+    Map a user-facing solvent name to its parameter-file stem.
+
+    Raises
+    ------
+    ValueError
+        If ``solvent`` is not a known solvent.
+    """
+
+    try:
+        return GBSA_SOLVENTS[solvent.lower()]
+    except KeyError:
+        known = ", ".join(sorted(set(GBSA_SOLVENTS)))
+        raise ValueError(
+            f"Unknown solvent {solvent!r}; choose one of: {known}."
+        )
+
+
+def default_gbsa_dir(install_root: str | Path | None = None) -> Path:
+    """
+    Return the directory holding downloaded GBSA solvation parameter files.
+    """
+
+    root = Path(install_root).expanduser() if install_root is not None else default_install_root()
+    return root / "gbsa" / GBSA_PARAM_METHOD
+
+
+def gbsa_param_path(
+    solvent: str,
+    install_root: str | Path | None = None,
+) -> Path:
+    """
+    Return the local path where the GBSA parameter file for ``solvent`` lives
+    (whether or not it has been downloaded yet).
+    """
+
+    return default_gbsa_dir(install_root) / f"param_gbsa_{_solvent_stem(solvent)}.txt"
+
+
+def install_gbsa_param(
+    solvent: str,
+    install_root: str | Path | None = None,
+    url: str | None = None,
+    force: bool = False,
+) -> Path:
+    """
+    Download the GBSA implicit-solvation parameter file for ``solvent``.
+
+    Parameters
+    ----------
+    solvent : str
+        Solvent name (e.g. ``"water"``); see :data:`GBSA_SOLVENTS`.
+    install_root : str or Path, optional
+        Directory root for downloaded parameters. Defaults to the user-local
+        share directory.
+    url : str, optional
+        Explicit file URL override. Defaults to the grimme-lab release file.
+    force : bool
+        Re-download even when the file already exists.
+    """
+
+    stem = _solvent_stem(solvent)
+    destination = gbsa_param_path(solvent, install_root)
+
+    if destination.exists() and not force:
+        return destination.resolve()
+
+    if url is None:
+        url = f"{GBSA_BASE_URL}/{GBSA_PARAM_METHOD}/param_gbsa_{stem}.txt"
+
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    _download_file(url, destination)
+
+    if not destination.exists():
+        raise FileNotFoundError(
+            f"Downloaded GBSA parameter file is missing: {destination}"
+        )
+
+    return destination.resolve()
+
+
 def dftb_prefix_export(parameter_dir: str | Path) -> str:
     """
     Return the shell export line for a Slater-Koster parameter directory.

ThermoScreening/cli/thermo.py

@@ -8,6 +8,7 @@
     check_dftb_setup,
     dftb_prefix_export,
     format_diagnostics,
+    install_gbsa_param,
     install_slakos,
 )
 from ThermoScreening.thermo.api import execute
@@ -51,6 +52,12 @@ def _command_parser():
         default=None,
         help="Archive URL override (defaults to the release URL for the set).",
     )
+    setup_parser.add_argument(
+        "--solvent",
+        default=None,
+        help="Instead of a parameter set, download the GBSA implicit-solvation "
+        "parameter file for this solvent (e.g. 'water').",
+    )
     setup_parser.add_argument(
         "--force",
         action="store_true",
@@ -94,6 +101,11 @@ def _command_parser():
         help="Slater-Koster parameter set (default '3ob'). Selects the Hamiltonian "
         "parameters and matching spin constants.",
     )
+    screen_parser.add_argument(
+        "--solvent", default=None,
+        help="GBSA/ALPB implicit-solvation solvent (e.g. 'water') applied to every "
+        "molecule. Default gas phase. Install with 'setup-dftb --solvent <name>'.",
+    )
 
     return parser
 
@@ -123,9 +135,19 @@ def parse_args(argv=None):
 
 def run_setup_dftb(parser_args):
     """
-    Download the default DFTB+ Slater-Koster parameter set.
+    Download a DFTB+ Slater-Koster parameter set, or a GBSA solvent parameter.
     """
 
+    if parser_args.solvent is not None:
+        param_file = install_gbsa_param(
+            parser_args.solvent,
+            install_root=parser_args.install_root,
+            url=parser_args.url,
+            force=parser_args.force,
+        )
+        print("GBSA solvation parameters:", param_file)
+        return 0
+
     parameter_dir = install_slakos(
         install_root=parser_args.install_root,
         url=parser_args.url,
@@ -164,6 +186,7 @@ def run_screen(parser_args):
         pressure=parser_args.pressure,
         directory=parser_args.directory,
         parameter_set=parser_args.parameter_set,
+        solvent=parser_args.solvent,
     )
 
     failed = sum(1 for record in results if record["status"] != "ok")

ThermoScreening/thermo/api.py

@@ -18,7 +18,7 @@
 from .thermo import Thermo
 from .atoms import Atom
 from ..calculator import Geoopt, Hessian, Modes
-from ..calculator.dftbplus import _spin_kwargs, SPIN_CONSTANTS_3OB
+from ..calculator.dftbplus import _spin_kwargs, _solvation_kwargs, SPIN_CONSTANTS_3OB
 
 
 logger = logging.getLogger(__package_name__).getChild("api")
@@ -505,6 +505,8 @@ def dftbplus_thermo(
         directory=None,
         spin=None,
         spin_constants=None,
+        solvent=None,
+        solvation_param_file=None,
         **kwargs
     ):
     """
@@ -534,6 +536,14 @@ def dftbplus_thermo(
         Element -> spin-constant mapping matching the Slater-Koster set in use.
         Defaults to the 3ob constants; pass ``SPIN_CONSTANTS_MIO`` (or the value
         from :func:`resolve_parameter_set`) when running the mio set.
+    solvent : str, optional
+        Solvent name (e.g. ``"water"``) for GBSA/ALPB implicit solvation. When
+        set, the optimisation, energy, and Hessian all run in solution. Requires
+        the solvent's parameter file (``thermo setup-dftb --solvent <name>``).
+        Defaults to a gas-phase calculation.
+    solvation_param_file : str, optional
+        Explicit path to a GBSA parameter file, overriding ``solvent`` (use a
+        method-consistent set instead of the default GFN-fit one).
 
     Other Parameters
     ----------------
@@ -557,14 +567,19 @@ def dftbplus_thermo(
 
     spin_kwargs = _spin_kwargs(atoms, spin, spin_constants)
 
+    # Implicit solvation (empty for the gas-phase default). Applied to both the
+    # optimisation and the Hessian so the geometry and frequencies are consistent.
+    solvation_kwargs = _solvation_kwargs(solvent, solvation_param_file)
+    engine_kwargs = {**spin_kwargs, **solvation_kwargs, **kwargs}
+
     with _run_in_directory(directory):
         # run geometry optimization
-        geoopt = Geoopt(atoms=atoms, charge=charge, **spin_kwargs, **kwargs)
+        geoopt = Geoopt(atoms=atoms, charge=charge, **engine_kwargs)
         potential_energy = geoopt.potential_energy()
         optimized_atoms = geoopt.read()
 
         # run hessian calculation
-        Hessian(atoms=optimized_atoms, charge=charge, **spin_kwargs, **kwargs)
+        Hessian(atoms=optimized_atoms, charge=charge, **engine_kwargs)
 
         # run normal mode calculation
         modes = Modes()

ThermoScreening/thermo/screening.py

@@ -31,9 +31,11 @@
     "path",
     "charge",
     "status",
+    "Eelec_hartree",
     "E_hartree",
     "H_hartree",
     "G_hartree",
+    "G_total_hartree",
     "S_cal_per_mol_K",
     "Cv_cal_per_mol_K",
     "error",
@@ -98,10 +100,15 @@ def _load_jobs(source, charge: float, spin=None):
 
 
 def _thermo_summary(thermo):
+    # Eelec_hartree is the DFTB+ electronic (SCC) energy, which carries the
+    # implicit-solvation term; E/H/G_hartree are the thermal corrections and
+    # G_total_hartree = Eelec + G correction is the absolute Gibbs free energy.
     return {
+        "Eelec_hartree": thermo.electronic_energy(),
         "E_hartree": thermo.total_energy("H"),
         "H_hartree": thermo.total_enthalpy("H"),
         "G_hartree": thermo.total_gibbs_free_energy("H"),
+        "G_total_hartree": thermo.total_EeGtot(),
         "S_cal_per_mol_K": thermo.total_entropy("cal/(mol*K)"),
         "Cv_cal_per_mol_K": thermo.total_heat_capacity("cal/(mol*K)"),
     }
@@ -135,6 +142,7 @@ def screen(
     parameters=None,
     spin=None,
     parameter_set="3ob",
+    solvent=None,
 ):
     """
     Run a thermochemistry screen over a set of molecules.
@@ -163,6 +171,10 @@ def screen(
     parameter_set : str
         Slater-Koster parameter set to use (``"3ob"`` or ``"mio"``). Selects both
         the default Hamiltonian parameters and the matching spin constants.
+    solvent : str, optional
+        Solvent name for GBSA/ALPB implicit solvation applied to every molecule
+        (e.g. ``"water"``). Defaults to gas phase. The solvent parameter file
+        must be installed (``thermo setup-dftb --solvent <name>``).
 
     Returns
     -------
@@ -194,6 +206,7 @@ def screen(
                 directory=str(root / job.name),
                 spin=job.spin,
                 spin_constants=spin_constants,
+                solvent=solvent,
                 **parameters,
             )
             record.update(_thermo_summary(thermo))

tests/calculator/test_dftbplus.py

@@ -448,6 +448,33 @@ def test_resolve_parameter_set_rejects_unknown_set():
         dftbplus_module.resolve_parameter_set("does-not-exist")
 
 
+def test_solvation_kwargs_gas_phase_is_empty():
+    assert dftbplus_module._solvation_kwargs() == {}
+    assert dftbplus_module._solvation_kwargs(solvent=None, param_file=None) == {}
+
+
+def test_solvation_kwargs_explicit_param_file(tmp_path):
+    param = tmp_path / "param_gbsa_h2o.txt"
+    param.write_text("data", encoding="utf-8")
+
+    kw = dftbplus_module._solvation_kwargs(param_file=str(param))
+
+    assert kw["Hamiltonian_Solvation"] == "GeneralizedBorn {"
+    # DFTB+ resolves ParamFile relative to the run dir, so it must be absolute
+    assert os.path.isabs(kw["Hamiltonian_Solvation_ParamFile"])
+    assert kw["Hamiltonian_Solvation_ParamFile"] == str(param.resolve())
+
+
+def test_solvation_kwargs_missing_param_file_raises(tmp_path):
+    with pytest.raises(FileNotFoundError):
+        dftbplus_module._solvation_kwargs(param_file=str(tmp_path / "nope.txt"))
+
+
+def test_solvation_kwargs_missing_solvent_file_hints_setup(tmp_path):
+    with pytest.raises(FileNotFoundError, match="setup-dftb --solvent water"):
+        dftbplus_module._solvation_kwargs(solvent="water", install_root=tmp_path)
+
+
 @pytest.mark.skipif(
     not dftbplus_ready,
     reason="DFTB+ executables are not installed or cannot start.",