Add PME parameter options and automatic PME tuning

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
 --------------------------------------------------------------------------------------

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

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):