diff --git a/CHANGELOG.md b/CHANGELOG.md index 02bcde8..2c9b1fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ Changelog * Fixed the ABFE lambda schedules ignoring the `restraint_lever` of a user-supplied Boresch restraint, which could leave the restraint uncoupled from the schedule [#229](https://github.com/OpenBioSim/somd2/pull/229). * Fixed replica exchange applying the inverse of the accepted permutation when mixing, which sent configurations to the wrong lambda windows whenever accepted swaps formed a cycle of three or more replicas. The replica exchange transition matrix now also records each configuration's move [#232](https://github.com/OpenBioSim/somd2/pull/232). * Add an `auto` option for `softcore_form`, which is now the default. It uses the Beutler form for the `annihilate` and `decouple` lambda schedules and the Zacharias form otherwise, including for custom schedules. The resolved form is recorded in the config file, so ABFE runs started with the previous default must set `softcore_form="zacharias"` to restart [#234](https://github.com/OpenBioSim/somd2/pull/234). +* Add `pme_alpha`, `pme_grid` and `pme_spacing` options to set the PME parameters explicitly, a `pme_tolerance` option for the Ewald error tolerance, and a `tune_pme` option, on by default, which picks the smallest PME grid that is at least as accurate as the parameters chosen from `pme_tolerance` on the CUDA and OpenCL platforms. The tuned parameters and their relative force error are logged, saved to `pme_parameters.yaml`, and reused on restart. Restarts of runs without this file use the default parameters [#235](https://github.com/OpenBioSim/somd2/pull/235). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 039c913..18c9a1d 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -87,6 +87,7 @@ class Config: _nargs = { "lambda_values": "+", "lambda_energy": "+", + "pme_grid": "+", "rest2_scale": "+", "restraints": "+", } @@ -104,6 +105,11 @@ def __init__( integrator="langevin_middle", cutoff_type="pme", cutoff="9 A", + pme_tolerance=0.0001, + pme_alpha=None, + pme_grid=None, + pme_spacing=None, + tune_pme=True, h_mass_factor=1.5, hmr=True, num_lambda=11, @@ -224,6 +230,27 @@ def __init__( cutoff: str Non-bonded cutoff distance. Use "infinite" for no cutoff. + pme_tolerance: float + The Ewald error tolerance, from which OpenMM chooses the PME parameters. + When 'tune_pme' is set, the tuned parameters are at least as accurate + as those chosen from this tolerance. + + pme_alpha: float + The PME splitting parameter in inverse nanometers. Requires 'pme_grid' + or 'pme_spacing'. If not set, it is derived from the Ewald error tolerance. + + pme_grid: [int] + The PME grid size, either a single integer or one per box vector. + + pme_spacing: str + The maximum PME grid spacing, e.g. "0.12 nm", from which the grid size + is calculated. + + tune_pme: bool + Whether to tune the PME parameters for the fastest settings that are at + least as accurate as the defaults. This only applies on the CUDA and + OpenCL platforms, and is ignored if any of the PME options are set. + h_mass_factor: float Factor by which to scale hydrogen masses. @@ -667,6 +694,11 @@ def __init__( self.integrator = integrator self.cutoff_type = cutoff_type self.cutoff = cutoff + self.pme_tolerance = pme_tolerance + self.pme_alpha = pme_alpha + self.pme_grid = pme_grid + self.pme_spacing = pme_spacing + self.tune_pme = tune_pme self.h_mass_factor = h_mass_factor self.hmr = hmr self.timestep = timestep @@ -1054,6 +1086,89 @@ def cutoff(self, cutoff): else: self._cutoff = cutoff + @property + def pme_tolerance(self): + return self._pme_tolerance + + @pme_tolerance.setter + def pme_tolerance(self, pme_tolerance): + try: + pme_tolerance = float(pme_tolerance) + except Exception: + raise ValueError("'pme_tolerance' must be a float") + if not 0 < pme_tolerance < 1: + raise ValueError("'pme_tolerance' must be between 0 and 1") + self._pme_tolerance = pme_tolerance + + @property + def pme_alpha(self): + return self._pme_alpha + + @pme_alpha.setter + def pme_alpha(self, pme_alpha): + if pme_alpha is not None: + try: + pme_alpha = float(pme_alpha) + except Exception: + raise ValueError("'pme_alpha' must be a float") + if pme_alpha <= 0: + raise ValueError("'pme_alpha' must be positive") + self._pme_alpha = pme_alpha + + @property + def pme_grid(self): + return self._pme_grid + + @pme_grid.setter + def pme_grid(self, pme_grid): + if pme_grid is not None: + if not isinstance(pme_grid, _Iterable) or isinstance(pme_grid, str): + pme_grid = [pme_grid] + try: + pme_grid = [int(x) for x in pme_grid] + except Exception: + raise ValueError("'pme_grid' must be an integer or a list of integers") + if len(pme_grid) not in [1, 3]: + raise ValueError("'pme_grid' must have one or three values") + if any(x < 6 for x in pme_grid): + raise ValueError("Each 'pme_grid' value must be at least 6") + self._pme_grid = pme_grid + + @property + def pme_spacing(self): + return self._pme_spacing + + @pme_spacing.setter + def pme_spacing(self, pme_spacing): + if pme_spacing is not None: + if not isinstance(pme_spacing, str): + raise TypeError("'pme_spacing' must be of type 'str'") + + from sire.units import angstrom + + try: + s = _sr.u(pme_spacing) + except: + raise ValueError( + f"Unable to parse 'pme_spacing' as a Sire GeneralUnit: {pme_spacing}" + ) + if not s.has_same_units(angstrom): + raise ValueError("'pme_spacing' units are invalid.") + + pme_spacing = s + + self._pme_spacing = pme_spacing + + @property + def tune_pme(self): + return self._tune_pme + + @tune_pme.setter + def tune_pme(self, tune_pme): + if not isinstance(tune_pme, bool): + raise ValueError("'tune_pme' must be of type 'bool'") + self._tune_pme = tune_pme + @property def h_mass_factor(self): return self._h_mass_factor diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 1bfeaf5..0723df3 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -268,6 +268,14 @@ def __init__(self, system, config): self._config.use_dispersion_correction ) + # PME parameters. + self._config._extra_args["tolerance"] = self._config.pme_tolerance + + for option in ["pme_alpha", "pme_grid", "pme_spacing"]: + value = getattr(self._config, option) + if value is not None: + self._config._extra_args[option] = value + # GCMC LRC map options. if self._config.gcmc and self._config.use_dispersion_correction: self._config._extra_args["use_gcmc_lrc"] = True @@ -1033,6 +1041,86 @@ def __init__(self, system, config): # Update the maximum number of threads. _sr.legacy.Base.set_max_num_threads(sire_threads) + self._set_tuned_pme_parameters() + + def _set_tuned_pme_parameters(self): + """ + Tune the PME parameters for a new run, saving them so that restarts + reuse the same values, or load the values saved by a previous run. + """ + import yaml as _yaml + + pme_file = _Path(self._filenames["pme"]) + + has_pme_options = any( + getattr(self._config, option) is not None + for option in ["pme_alpha", "pme_grid", "pme_spacing"] + ) + + if self._is_restart: + if pme_file.exists(): + with open(pme_file) as f: + params = _yaml.safe_load(f) + _logger.info(f"Using PME parameters from {pme_file}: {params}") + self._config._extra_args["pme_alpha"] = params["pme_alpha"] + self._config._extra_args["pme_grid"] = params["pme_grid"] + elif self._config.tune_pme and not has_pme_options: + _logger.info( + "No saved PME parameters for this restart, so using the " + "default PME parameters" + ) + return + + if pme_file.exists(): + pme_file.unlink() + + if ( + not self._config.tune_pme + or has_pme_options + or not self._has_space + or self._config.cutoff_type != "pme" + or self._config.platform not in ["cuda", "opencl"] + ): + return + + from sire.convert.openmm import tune_pme as _tune_pme + + system = self._system[0] if isinstance(self._system, list) else self._system + + _logger.info("Tuning PME parameters") + + try: + params = _tune_pme( + system, device=0, return_errors=True, **self._dynamics_kwargs + ) + except Exception as e: + _logger.warning( + f"PME tuning failed, so using the default PME parameters: {str(e)}" + ) + return + + error = f"relative force error {params['pme_error']:.2e} (target {params['pme_target_error']:.2e})" + + if params["pme_error"] == float("inf"): + _logger.warning( + "PME tuning wasn't possible for this system, so using the default " + "PME parameters" + ) + return + + if "pme_alpha" not in params: + _logger.info(f"The default PME parameters are already the fastest, {error}") + return + + _logger.info( + f"Using tuned PME parameters: alpha={params['pme_alpha']:.4f} nm^-1, " + f"grid={params['pme_grid']}, {error}" + ) + + _dict_to_yaml(params, str(pme_file)) + self._config._extra_args["pme_alpha"] = params["pme_alpha"] + self._config._extra_args["pme_grid"] = params["pme_grid"] + @property def _is_abfe_bound(self): """ @@ -1810,6 +1898,9 @@ def _prepare_output(self): filenames["topology0"] = str(self._config.output_directory / "system0.prm7") filenames["topology1"] = str(self._config.output_directory / "system1.prm7") + # File for the tuned PME parameters, so that restarts can reuse them. + filenames["pme"] = str(self._config.output_directory / "pme_parameters.yaml") + return filenames def _check_restart(self): diff --git a/tests/runner/test_pme.py b/tests/runner/test_pme.py new file mode 100644 index 0000000..da053ab --- /dev/null +++ b/tests/runner/test_pme.py @@ -0,0 +1,152 @@ +import os +import tempfile +from pathlib import Path + +import pytest +import yaml + +from somd2.config import Config +from somd2.runner import Runner + + +def _nonbonded_force(runner): + from openmm import NonbondedForce + + d = runner._system.dynamics(**runner._dynamics_kwargs) + + for force in d._d._omm_mols.getSystem().getForces(): + if isinstance(force, NonbondedForce): + return force + + +def _pme_parameters(runner): + from openmm import unit + + alpha, *grid = _nonbonded_force(runner).getPMEParameters() + return alpha.value_in_unit(unit.nanometer**-1), grid + + +def _short_run(tmpdir, **options): + config = { + "runtime": "12fs", + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "CPU", + "max_threads": 1, + "num_lambda": 2, + } + config.update(options) + return Config(**config) + + +def test_pme_config_options(): + """Validate the parsing of the PME options.""" + assert Config().tune_pme + assert Config(pme_tolerance="5e-4").pme_tolerance == pytest.approx(5e-4) + + assert Config(pme_grid=64).pme_grid == [64] + assert Config(pme_grid=["64", "64", "72"]).pme_grid == [64, 64, 72] + assert Config(pme_alpha="3.47").pme_alpha == pytest.approx(3.47) + assert Config(pme_spacing="0.12 nm").pme_spacing.value() == pytest.approx(1.2) + + for options in [ + {"pme_grid": [64, 64]}, + {"pme_grid": 4}, + {"pme_alpha": -1.0}, + {"pme_spacing": "1 ps"}, + {"pme_tolerance": 0.0}, + ]: + with pytest.raises(ValueError): + Config(**options) + + +def test_pme_options_passed(ethane_methanol): + """Validate that explicit PME options reach the OpenMM context.""" + with tempfile.TemporaryDirectory() as tmpdir: + config = Config( + platform="cpu", output_directory=tmpdir, pme_alpha=3.4, pme_grid=32 + ) + + alpha, grid = _pme_parameters(Runner(ethane_methanol, config)) + + assert alpha == pytest.approx(3.4) + assert grid == [32, 32, 32] + + +def test_pme_tolerance_passed(ethane_methanol): + """Validate that the Ewald error tolerance reaches the OpenMM context.""" + with tempfile.TemporaryDirectory() as tmpdir: + config = Config(platform="cpu", output_directory=tmpdir, pme_tolerance=5e-4) + + nbff = _nonbonded_force(Runner(ethane_methanol, config)) + + assert nbff.getEwaldErrorTolerance() == pytest.approx(5e-4) + + +def test_pme_restart(ethane_methanol): + """ + Validate that a restart reuses saved PME parameters, and that a new run + removes stale ones. + """ + with tempfile.TemporaryDirectory() as tmpdir: + runner = Runner(ethane_methanol, _short_run(tmpdir)) + runner.run() + del runner + + params = {"pme_alpha": 3.4, "pme_grid": [32, 32, 32]} + pme_file = Path(tmpdir) / "pme_parameters.yaml" + + with open(pme_file, "w") as f: + yaml.safe_dump(params, f) + + runner = Runner( + ethane_methanol, + _short_run(tmpdir, runtime="24fs", restart=True, overwrite=True), + ) + + alpha, grid = _pme_parameters(runner) + + assert alpha == pytest.approx(3.4) + assert grid == [32, 32, 32] + + del runner + + Runner(ethane_methanol, _short_run(tmpdir, overwrite=True)) + + assert not pme_file.exists() + + +def _has_cuda(): + try: + import openmm + + openmm.Platform.getPlatformByName("CUDA") + return True + except Exception: + return False + + +@pytest.mark.skipif(not _has_cuda(), reason="CUDA platform is not available") +def test_pme_tuning(ethane_methanol, monkeypatch): + """Validate that a new CUDA run tunes and saves the PME parameters.""" + if os.environ.get("CUDA_VISIBLE_DEVICES") is None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + + with tempfile.TemporaryDirectory() as tmpdir: + runner = Runner( + ethane_methanol, Config(platform="cuda", output_directory=tmpdir) + ) + + pme_file = Path(tmpdir) / "pme_parameters.yaml" + + if pme_file.exists(): + with open(pme_file) as f: + params = yaml.safe_load(f) + + alpha, grid = _pme_parameters(runner) + + assert alpha == pytest.approx(params["pme_alpha"]) + assert grid == params["pme_grid"] + assert params["pme_error"] <= params["pme_target_error"]