From 94f780d8532698f17d3c454757f386c6d3ef88e8 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 1 Oct 2026 12:31:32 +0100 Subject: [PATCH 1/3] Add PME parameter options and automatic PME tuning. --- CHANGELOG.md | 1 + src/somd2/config/_config.py | 94 +++++++++++++++++++++++++ src/somd2/runner/_base.py | 67 ++++++++++++++++++ tests/runner/test_pme.py | 133 ++++++++++++++++++++++++++++++++++++ 4 files changed, 295 insertions(+) create mode 100644 tests/runner/test_pme.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 02bcde8..cfa80b3 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, and a `tune_pme` option, on by default, which picks the smallest PME grid that is at least as accurate as the defaults on the CUDA and OpenCL platforms. The tuned parameters are 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..2833a6d 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,10 @@ def __init__( integrator="langevin_middle", cutoff_type="pme", cutoff="9 A", + pme_alpha=None, + pme_grid=None, + pme_spacing=None, + tune_pme=True, h_mass_factor=1.5, hmr=True, num_lambda=11, @@ -224,6 +229,22 @@ def __init__( cutoff: str Non-bonded cutoff distance. Use "infinite" for no cutoff. + 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 +688,10 @@ def __init__( self.integrator = integrator self.cutoff_type = cutoff_type self.cutoff = cutoff + 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 +1079,75 @@ def cutoff(self, cutoff): else: self._cutoff = cutoff + @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..c7ce228 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -268,6 +268,12 @@ def __init__(self, system, config): self._config.use_dispersion_correction ) + # Explicit PME parameters. + 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 +1039,64 @@ 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.update(params) + 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") + + params = _tune_pme(system, device=0, **self._dynamics_kwargs) + + if not params: + _logger.info("The default PME parameters are already the fastest") + return + + _logger.info(f"Using tuned PME parameters: {params}") + + _dict_to_yaml(params, str(pme_file)) + self._config._extra_args.update(params) + @property def _is_abfe_bound(self): """ @@ -1810,6 +1874,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..da67512 --- /dev/null +++ b/tests/runner/test_pme.py @@ -0,0 +1,133 @@ +import os +import tempfile +from pathlib import Path + +import pytest +import yaml + +from somd2.config import Config +from somd2.runner import Runner + + +def _pme_parameters(runner): + from openmm import NonbondedForce, unit + + d = runner._system.dynamics(**runner._dynamics_kwargs) + + for force in d._d._omm_mols.getSystem().getForces(): + if isinstance(force, NonbondedForce): + alpha, *grid = force.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_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"}, + ]: + 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_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"] From a9d124368b7f0c9602d2acb2abf6ced2eb78fba0 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 1 Oct 2026 12:44:15 +0100 Subject: [PATCH 2/3] Add pme_tolerance option and log the tuned PME force errors. --- CHANGELOG.md | 2 +- src/somd2/config/_config.py | 21 +++++++++++++++++++++ src/somd2/runner/_base.py | 25 ++++++++++++++++++------- tests/runner/test_pme.py | 27 +++++++++++++++++++++++---- 4 files changed, 63 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cfa80b3..2c9b1fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +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, and a `tune_pme` option, on by default, which picks the smallest PME grid that is at least as accurate as the defaults on the CUDA and OpenCL platforms. The tuned parameters are 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). +* 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 2833a6d..18c9a1d 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -105,6 +105,7 @@ def __init__( integrator="langevin_middle", cutoff_type="pme", cutoff="9 A", + pme_tolerance=0.0001, pme_alpha=None, pme_grid=None, pme_spacing=None, @@ -229,6 +230,11 @@ 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. @@ -688,6 +694,7 @@ 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 @@ -1079,6 +1086,20 @@ 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 diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index c7ce228..4a15f91 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -268,7 +268,9 @@ def __init__(self, system, config): self._config.use_dispersion_correction ) - # Explicit PME parameters. + # 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: @@ -1060,7 +1062,8 @@ def _set_tuned_pme_parameters(self): 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.update(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 " @@ -1086,16 +1089,24 @@ def _set_tuned_pme_parameters(self): _logger.info("Tuning PME parameters") - params = _tune_pme(system, device=0, **self._dynamics_kwargs) + params = _tune_pme( + system, device=0, return_errors=True, **self._dynamics_kwargs + ) + + error = f"relative force error {params['pme_error']:.2e} (target {params['pme_target_error']:.2e})" - if not params: - _logger.info("The default PME parameters are already the fastest") + 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: {params}") + _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.update(params) + self._config._extra_args["pme_alpha"] = params["pme_alpha"] + self._config._extra_args["pme_grid"] = params["pme_grid"] @property def _is_abfe_bound(self): diff --git a/tests/runner/test_pme.py b/tests/runner/test_pme.py index da67512..da053ab 100644 --- a/tests/runner/test_pme.py +++ b/tests/runner/test_pme.py @@ -9,15 +9,21 @@ from somd2.runner import Runner -def _pme_parameters(runner): - from openmm import NonbondedForce, unit +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): - alpha, *grid = force.getPMEParameters() - return alpha.value_in_unit(unit.nanometer**-1), grid + 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): @@ -38,6 +44,7 @@ def _short_run(tmpdir, **options): 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] @@ -49,6 +56,7 @@ def test_pme_config_options(): {"pme_grid": 4}, {"pme_alpha": -1.0}, {"pme_spacing": "1 ps"}, + {"pme_tolerance": 0.0}, ]: with pytest.raises(ValueError): Config(**options) @@ -67,6 +75,16 @@ def test_pme_options_passed(ethane_methanol): 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 @@ -131,3 +149,4 @@ def test_pme_tuning(ethane_methanol, monkeypatch): assert alpha == pytest.approx(params["pme_alpha"]) assert grid == params["pme_grid"] + assert params["pme_error"] <= params["pme_target_error"] From 337d74a8c5ea7af53c006d90d9a50c16239db05d Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 1 Oct 2026 16:41:35 +0100 Subject: [PATCH 3/3] Catch PME tuning errors. --- src/somd2/runner/_base.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 4a15f91..0723df3 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1089,12 +1089,25 @@ def _set_tuned_pme_parameters(self): _logger.info("Tuning PME parameters") - params = _tune_pme( - system, device=0, return_errors=True, **self._dynamics_kwargs - ) + 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