diff --git a/CHANGELOG.md b/CHANGELOG.md index 50de0be..a0b1e4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,15 @@ Changelog ========= +[2026.2.0](https://github.com/openbiosim/loch/compare/2026.1.0...2026.2.0) - Sep 2026 +-------------------------------------------------------------------------------------- + +* Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29). +* Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context [#30](https://github.com/OpenBioSim/loch/pull/30). +* Resolve unit conversions once when extracting the lambda dependent non-bonded parameters, rather than per atom, which dominated sampler setup [#33](https://github.com/OpenBioSim/loch/pull/33). +* Recount the waters in the GCMC region rather than returning a cached count that was not invalidated when the positions changed, and return the whole box count directly when no region is defined, which previously raised [#36](https://github.com/OpenBioSim/loch/pull/36). +* Stop uploading the GCMC region centre to the GPU twice in `delete_waters` and `num_waters`, which raised on the OpenCL platform, and make `delete_waters` a no-op when no region is defined, which previously raised [#39](https://github.com/OpenBioSim/loch/pull/39). + [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026 ------------------------------------------------------------------------------------- diff --git a/pixi.toml b/pixi.toml index 16c6cff..35498d1 100644 --- a/pixi.toml +++ b/pixi.toml @@ -6,9 +6,9 @@ platforms = ["linux-64", "osx-arm64", "win-64"] [dependencies] python = ">=3.10" # main -biosimspace = ">=2026.1.0,<2026.2.0" +biosimspace = ">=2026.2.0,<2026.3.0" # devel -#biosimspace = "==2026.2.0.dev" +#biosimspace = "==2026.3.0.dev" loguru = "*" pyopencl = "*" diff --git a/recipes/loch/recipe.yaml b/recipes/loch/recipe.yaml index 1e854de..f93af36 100644 --- a/recipes/loch/recipe.yaml +++ b/recipes/loch/recipe.yaml @@ -20,9 +20,9 @@ requirements: - versioningit run: # main - - biosimspace >=2026.1.0,<2026.2.0 + - biosimspace >=2026.2.0,<2026.3.0 # devel - #- biosimspace ==2026.2.0.dev + #- biosimspace ==2026.3.0.dev - loguru - pyopencl - python diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index e2f99c9..3c11e24 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -89,6 +89,8 @@ def __init__( lambda_schedule: _Optional[_Any] = None, lambda_value: float = 0.0, rest2_scale: float = 1.0, + lambda_values: _Optional[list[float]] = None, + rest2_scales: _Optional[list[float]] = None, rest2_selection: _Optional[str] = None, shift_coulomb: str = "1 A", shift_delta: str = "1.5 A", @@ -205,6 +207,16 @@ def __init__( (REST2) for alchemical systems. This should specify the temperature of the REST2 system relative to the rest of the system. + lambda_values: [float] + The lambda values that the sampler will be switched between via + set_lambda(). The non-bonded parameters for each are cached at + construction, so that switching later doesn't need to build an + OpenMM context. If None, only 'lambda_value' is cached. + + rest2_scales: [float] + The REST2 scaling factor for each entry of 'lambda_values'. If + None, 'rest2_scale' is used for all of them. + rest2_selection: str A selection string for atoms to include in the REST2 region in addition to any perturbable molecules. For example, "molidx 0 and @@ -506,6 +518,37 @@ def __init__( raise ValueError("'rest2_scale' must be greater than or equal to 1.0") self._rest2_scale = rest2_scale + # Cache of the lambda dependent non-bonded parameters, keyed by + # (lambda_value, rest2_scale). Populated below and consumed by + # set_lambda(). + self._lambda_params = {} + + if lambda_values is None: + self._lambda_values = None + else: + try: + self._lambda_values = [float(x) for x in lambda_values] + except: + raise ValueError("'lambda_values' must be a list of 'float'") + if not all(0.0 <= x <= 1.0 for x in self._lambda_values): + raise ValueError("'lambda_values' must be between 0 and 1") + + if rest2_scales is None: + self._rest2_scales = [rest2_scale] * len(self._lambda_values) + else: + try: + self._rest2_scales = [float(x) for x in rest2_scales] + except: + raise ValueError("'rest2_scales' must be a list of 'float'") + if len(self._rest2_scales) != len(self._lambda_values): + raise ValueError( + "'rest2_scales' must be the same length as 'lambda_values'" + ) + if any(x < 1.0 for x in self._rest2_scales): + raise ValueError( + "'rest2_scales' must be greater than or equal to 1.0" + ) + if rest2_selection is not None: if not isinstance(rest2_selection, str): raise ValueError("'rest2_selection' must be of type 'str'") @@ -694,6 +737,17 @@ def __init__( self._water_blocks = self._num_waters // self._num_threads + 1 # Initialise the GPU memory. + # Cache the non-bonded parameters for every lambda value that the + # sampler will be switched to, so that set_lambda() never has to build + # an OpenMM context mid-simulation. This is done before the GPU memory + # is initialised, so that the current lambda value is part of the same + # scan rather than needing a context of its own. + if self._lambda_values is not None: + self._precompute_lambdas( + [self._lambda_value] + self._lambda_values, + [self._rest2_scale] + self._rest2_scales, + ) + self._initialise_gpu_memory() # Set the box information. @@ -721,16 +775,23 @@ def __init__( # Zero the number of waters in the sampling volume. self._N = 0 - # Zero the statistics. + # Zero the statistics for the current lambda value. self._num_moves = 0 self._num_accepted = 0 self._num_accepted_attempts = 0 self._num_insertions = 0 self._num_deletions = 0 + # Statistics for lambda values other than the current one, keyed by + # formatted lambda value. The current one is held in the counters above + # and archived here by set_lambda(). + self._stats = {} + # Null the nonbonded forces. self._nonbonded_force = None self._custom_nonbonded_force = None + self._integration_groups = None + self._pme_groups = None # Flag for whether the last move was a bulk sampling move. self._is_bulk = False @@ -1211,7 +1272,10 @@ def num_waters(self, context=None) -> int: context: openmm.Context, optional The OpenMM context to count the waters from. If None, then the internal context is used if one is available, otherwise the count - from the last move is returned. + from the last move is returned. Only omit this immediately after a + move. Pass a context if dynamics have been run since, otherwise + waters will have crossed the region boundary and the stored count + will be out of date. Returns ------- @@ -1373,18 +1437,15 @@ def reset(self) -> None: """ Reset the sampler. """ - # Zero the number of accepted moves. - self._num_accepted = 0 - self._num_insertions = 0 - self._num_deletions = 0 - self._num_moves = 0 - self._num_accepted_attempts = 0 - self._num_accepted_insertions = 0 - self._num_accepted_deletions = 0 + # Zero the number of accepted moves, for every lambda value. + self._zero_stats() + self._stats = {} # Clear the forces. self._nonbonded_force = None self._custom_nonbonded_force = None + self._integration_groups = None + self._pme_groups = None # Clear the OpenMM context. self._openmm_context = None @@ -1392,43 +1453,137 @@ def reset(self) -> None: # The stored region count refers to the cleared context. self._N_region = None - def restore_stats(self, stats: dict) -> None: + @staticmethod + def stats_key(lambda_value: float) -> str: """ - Restore sampler statistics from a dictionary. + Return the key used to store statistics for a lambda value. Parameters ---------- - stats : dict - Dictionary of sampler statistics as returned by ``get_stats()``. + lambda_value: float + The lambda value. + + Returns + ------- + + str + The key. + """ + return f"{float(lambda_value):.5f}" + + def _zero_stats(self) -> None: + """ + Zero the statistics for the current lambda value. + """ + self._num_moves = 0 + self._num_accepted = 0 + self._num_insertions = 0 + self._num_deletions = 0 + self._num_accepted_attempts = 0 + + def _stats_keys(self) -> set: + """ + Return the keys of the lambda values that this sampler visits. + """ + keys = {self.stats_key(self._lambda_value)} + + if self._lambda_values is not None: + keys.update(self.stats_key(x) for x in self._lambda_values) + + return keys + + def _switch_stats(self, lambda_value: float) -> None: + """ + Archive the statistics for the current lambda value and load those for + a new one, zeroing them if it hasn't been visited before. + + Parameters + ---------- + + lambda_value: float + The lambda value being switched to. + """ + old_key = self.stats_key(self._lambda_value) + new_key = self.stats_key(lambda_value) + + if old_key == new_key: + return + + self._stats[old_key] = self._get_current_stats() + + if new_key in self._stats: + self._set_current_stats(self._stats.pop(new_key)) + else: + self._zero_stats() + + def _get_current_stats(self) -> dict: + """ + Return the statistics for the current lambda value. + """ + return { + "num_moves": self._num_moves, + "num_accepted": self._num_accepted, + "num_insertions": self._num_insertions, + "num_deletions": self._num_deletions, + "num_accepted_attempts": self._num_accepted_attempts, + } + + def _set_current_stats(self, stats: dict) -> None: + """ + Set the statistics for the current lambda value. Unrecognised entries + are ignored, so that statistics written by an older version can still + be restored. """ self._num_moves = stats["num_moves"] self._num_accepted = stats["num_accepted"] self._num_insertions = stats["num_insertions"] self._num_deletions = stats["num_deletions"] self._num_accepted_attempts = stats["num_accepted_attempts"] - self._num_accepted_insertions = stats["num_accepted_insertions"] - self._num_accepted_deletions = stats["num_accepted_deletions"] + + def restore_stats(self, stats: dict) -> None: + """ + Restore sampler statistics. + + Parameters + ---------- + + stats : dict + Statistics as returned by ``get_stats()``, i.e. keyed by lambda + value. Entries for lambda values that this sampler doesn't visit + are ignored, so it is safe to pass the statistics for a whole + simulation to each of several samplers. Were they retained, a + sampler would report stale statistics for another's lambda values, + which could then overwrite the live ones when merged. + """ + keys = self._stats_keys() + self._stats = {key: dict(value) for key, value in stats.items() if key in keys} + + # Load the statistics for the current lambda value, if present. + key = self.stats_key(self._lambda_value) + if key in self._stats: + self._set_current_stats(self._stats.pop(key)) + else: + self._zero_stats() def get_stats(self) -> dict: """ - Return the current sampler statistics as a dictionary. + Return the sampler statistics, keyed by lambda value. + + A sampler may be switched between lambda values via ``set_lambda()``, + and accumulates statistics for each of them separately. Non-alchemical + systems have a single key, for the lambda value the sampler was created + with. Returns ------- dict - Dictionary of sampler statistics. + Dictionary of sampler statistics for each lambda value. """ - return { - "num_moves": self._num_moves, - "num_accepted": self._num_accepted, - "num_insertions": self._num_insertions, - "num_deletions": self._num_deletions, - "num_accepted_attempts": self._num_accepted_attempts, - "num_accepted_insertions": self._num_accepted_insertions, - "num_accepted_deletions": self._num_accepted_deletions, - } + stats = {key: dict(value) for key, value in self._stats.items()} + stats[self.stats_key(self._lambda_value)] = self._get_current_stats() + return stats def ghost_residues(self) -> _np.ndarray: """ @@ -1517,7 +1672,11 @@ def move(self, context: _openmm.Context) -> list[int]: self._init_gcmc_lrc(context) # Get the OpenMM state. - state = context.getState(getPositions=True, getEnergy=self._is_pme) + state = context.getState( + getPositions=True, + getEnergy=self._is_pme, + groups=self._pme_groups, + ) # Get the current positions in OpenMM format and in Angstrom. positions_openmm = state.getPositions(asNumpy=True) @@ -1810,7 +1969,7 @@ def move(self, context: _openmm.Context) -> list[int]: # Get the new energy. final_energy = context.getState( - getEnergy=True + getEnergy=True, groups=self._pme_groups ).getPotentialEnergy() # Add the analytic LRC delta so the PME correction sees only @@ -1912,7 +2071,7 @@ def move(self, context: _openmm.Context) -> list[int]: # Get the new energy. final_energy = context.getState( - getEnergy=True + getEnergy=True, groups=self._pme_groups ).getPotentialEnergy() # Add the analytic LRC delta. @@ -2330,6 +2489,217 @@ def _prepare_system(system, water_template, rng, num_ghost_waters): _np.array(water_residues), ) + def _precompute_lambdas( + self, lambda_values: list[float], rest2_scales: list[float] + ) -> None: + """ + Cache the lambda dependent non-bonded parameters for a set of lambda + values, so that set_lambda() can switch between them without touching + OpenMM. + + The system only holds the end-state properties, so the parameters are + obtained from an OpenMM context built using the specified lambda + schedule. A single context is created and scanned over the requested + lambda values, re-reading the GhostNonGhostNonbondedForce at each one, + since building the context is by far the expensive part. + + This is a no-op for non-alchemical systems, and for lambda values that + have already been cached. + + Parameters + ---------- + + lambda_values: [float] + The lambda values to cache. + + rest2_scales: [float] + The REST2 scaling factor for each lambda value. + """ + + if not self._is_fep: + return + + # Only build a context if something is actually missing. + wanted = [] + for lambda_value, rest2_scale in zip(lambda_values, rest2_scales): + key = (float(lambda_value), float(rest2_scale)) + if key not in self._lambda_params and key not in wanted: + wanted.append(key) + + if not wanted: + return + + # Link to the reference state. + mols = _sr.morph.link_to_reference(self._system) + + # Build map of extra options for the dynamics object. + _map = {} + if self._softcore_form == _SoftcoreForm.TAYLOR: + _map["use_taylor_softening"] = True + _map["taylor_power"] = self._taylor_power + elif self._softcore_form == _SoftcoreForm.BEUTLER: + _map["use_beutler_softening"] = True + _map["beutler_alpha"] = self._beutler_alpha + + # Create a single dynamics object, which is then scanned over the + # requested lambda values. + d = mols.dynamics( + cutoff_type=self._cutoff, + cutoff=self._cutoff, + lambda_value=wanted[0][0], + schedule=self._lambda_schedule, + pressure=None, + timestep="2fs", + constraint="h_bonds", + perturbable_constraint="h_bonds_not_heavy_perturbed", + rest2_scale=wanted[0][1], + rest2_selection=self._rest2_selection, + swap_end_states=self._swap_end_states, + platform="cpu", + map=_map, + ) + + # Find the required force. set_lambda() updates it in place, so it + # only needs to be located once. + gng_force = None + for force in d.context().getSystem().getForces(): + if force.getName() == "GhostNonGhostNonbondedForce": + gng_force = force + break + + if gng_force is None: + raise ValueError( + "Could not find the GhostNonGhostNonbondedForce in the system" + ) + + num_particles = gng_force.getNumParticles() + + # Resolve the unit conversions once. Doing this per atom, by formatting + # and re-parsing a string, dominates the setup time for large systems. + nm_to_angstrom = _sr.u("1 nm").to("angstrom") + kj_per_mol_to_kcal_per_mol = _sr.u("1 kJ/mol").to("kcal/mol") + + for lambda_value, rest2_scale in wanted: + d.set_lambda(lambda_value, rest2_scale=rest2_scale) + + # Get the parameters for the GhostNonGhostNonbondedForce. + charges = _np.zeros(self._num_atoms, dtype=_np.float32) + sigmas = _np.zeros(self._num_atoms, dtype=_np.float32) + epsilons = _np.zeros(self._num_atoms, dtype=_np.float32) + alphas = _np.zeros(self._num_atoms, dtype=_np.float32) + for i in range(num_particles): + # Custom force parameters are returned as floats. + q, half_sigma, two_sqrt_epsilon, alpha, _ = ( + gng_force.getParticleParameters(i) + ) + # Charge in |e|, sigma in nm, epsilon in kJ/mol. + charges[i] = q + # Rescale and convert units. + sigmas[i] = 2.0 * half_sigma * nm_to_angstrom + epsilons[i] = (0.5 * two_sqrt_epsilon) ** 2 * kj_per_mol_to_kcal_per_mol + # Store the softening parameter. + alphas[i] = alpha + + self._lambda_params[(lambda_value, rest2_scale)] = ( + charges, + sigmas, + epsilons, + alphas, + ) + + def set_lambda( + self, lambda_value: float, rest2_scale: _Optional[float] = None + ) -> None: + """ + Set the lambda value for the sampler, updating the non-bonded + parameters used to evaluate insertion and deletion energies. + + Parameters for lambda values passed to the constructor are already + cached. Any other value is computed here, which requires building an + OpenMM context and is slow, though the result is then cached too. + + This must be kept consistent with the lambda value of the OpenMM + context that the sampler is used with. + + Parameters + ---------- + + lambda_value: float + The lambda value. + + rest2_scale: float + The REST2 scaling factor. If None, the current value is retained. + """ + + try: + lambda_value = float(lambda_value) + except: + raise ValueError("'lambda_value' must be of type 'float'") + if not 0.0 <= lambda_value <= 1.0: + raise ValueError("'lambda_value' must be between 0 and 1") + + if rest2_scale is None: + rest2_scale = self._rest2_scale + else: + try: + rest2_scale = float(rest2_scale) + except: + raise ValueError("'rest2_scale' must be of type 'float'") + if rest2_scale < 1.0: + raise ValueError("'rest2_scale' must be greater than or equal to 1.0") + + # Nothing to do. + if lambda_value == self._lambda_value and rest2_scale == self._rest2_scale: + return + + # Statistics are accumulated per lambda value, so archive those for the + # current one and load those for the new one. + self._switch_stats(lambda_value) + + # There are no lambda dependent parameters for a non-alchemical system. + if not self._is_fep: + self._lambda_value = lambda_value + self._rest2_scale = rest2_scale + return + + # Make sure the parameters are cached. This builds an OpenMM context, + # so is slow; pass 'lambda_values' to the constructor to avoid it. + self._precompute_lambdas([lambda_value], [rest2_scale]) + + charges, sigmas, epsilons, alphas = self._lambda_params[ + (lambda_value, rest2_scale) + ] + + # Upload the new parameters to the GPU. + self._gpu_charge = self._backend.to_gpu(charges) + self._gpu_sigma = self._backend.to_gpu(sigmas) + self._gpu_epsilon = self._backend.to_gpu(epsilons) + self._gpu_alpha = self._backend.to_gpu(alphas) + + self._lambda_value = lambda_value + self._rest2_scale = rest2_scale + + def set_ghost_file(self, ghost_file: _Optional[str]) -> None: + """ + Set the file that write_ghost_residues() appends to. + + Unlike the constructor, this does not create or truncate the file, + since it is intended for switching between files that are already + being written to. + + Parameters + ---------- + + ghost_file: str + The path to the ghost residue file. If None, ghost residues + cannot be written. + """ + + if ghost_file is not None and not isinstance(ghost_file, str): + raise TypeError("'ghost_file' must be of type 'str'") + + self._ghost_file = ghost_file + def _initialise_gpu_memory(self): """ Initialise the GPU memory. @@ -2394,80 +2764,19 @@ def _initialise_gpu_memory(self): # schedule and value, then extract the required properties from the forces # within the context. (The system just contains the end-state properties.) else: - # Link to the reference state. - mols = _sr.morph.link_to_reference(self._system) - - # Build map of extra options for the dynamics object. - _map = {} - if self._softcore_form == _SoftcoreForm.TAYLOR: - _map["use_taylor_softening"] = True - _map["taylor_power"] = self._taylor_power - elif self._softcore_form == _SoftcoreForm.BEUTLER: - _map["use_beutler_softening"] = True - _map["beutler_alpha"] = self._beutler_alpha - - # Create a dynamics object. - d = mols.dynamics( - cutoff_type=self._cutoff, - cutoff=self._cutoff, - lambda_value=self._lambda_value, - schedule=self._lambda_schedule, - pressure=None, - timestep="2fs", - constraint="h_bonds", - perturbable_constraint="h_bonds_not_heavy_perturbed", - rest2_scale=self._rest2_scale, - rest2_selection=self._rest2_selection, - swap_end_states=self._swap_end_states, - platform="cpu", - map=_map, - ) - - # Flag for the required force. - has_gng = False - - # Find the required forces. - for force in d.context().getSystem().getForces(): - if force.getName() == "GhostNonGhostNonbondedForce": - gng_force = force - has_gng = True - break + # Cache the host arrays so that set_lambda() can switch back to + # this lambda value without rebuilding an OpenMM context. + self._precompute_lambdas([self._lambda_value], [self._rest2_scale]) - # Make sure the force was found. - if not has_gng: - raise ValueError( - "Could not find the GhostNonGhostNonbondedForce in the system" - ) - - # Resolve the unit conversions once. Doing this per atom, by - # formatting and re-parsing a string, dominates the setup time for - # large systems. - nm_to_angstrom = _sr.u("1 nm").to("angstrom") - kj_per_mol_to_kcal_per_mol = _sr.u("1 kJ/mol").to("kcal/mol") - - # Get the parameters for the GhostNonGhostNonbondedForce. - charges = _np.zeros(self._num_atoms, dtype=_np.float32) - sigmas = _np.zeros(self._num_atoms, dtype=_np.float32) - epsilons = _np.zeros(self._num_atoms, dtype=_np.float32) - alphas = _np.zeros(self._num_atoms, dtype=_np.float32) - for i in range(gng_force.getNumParticles()): - # Custom force parameters are returned as floats. - q, half_sigma, two_sqrt_epsilon, alpha, _ = ( - gng_force.getParticleParameters(i) - ) - # Charge in |e|, sigma in nm, epsilon in kJ/mol. - charges[i] = q - # Rescale and convert units. - sigmas[i] = 2.0 * half_sigma * nm_to_angstrom - epsilons[i] = (0.5 * two_sqrt_epsilon) ** 2 * kj_per_mol_to_kcal_per_mol - # Store the softening parameter. - alphas[i] = alpha + charges, sigmas, epsilons, alphas = self._lambda_params[ + (self._lambda_value, self._rest2_scale) + ] # Convert to GPU arrays. - charges = self._backend.to_gpu(charges.astype(_np.float32)) - sigmas = self._backend.to_gpu(sigmas.astype(_np.float32)) - epsilons = self._backend.to_gpu(epsilons.astype(_np.float32)) - alphas = self._backend.to_gpu(alphas.astype(_np.float32)) + charges = self._backend.to_gpu(charges) + sigmas = self._backend.to_gpu(sigmas) + epsilons = self._backend.to_gpu(epsilons) + alphas = self._backend.to_gpu(alphas) # Create the ghost atom array. is_ghost_fep = _np.zeros(self._num_atoms, dtype=_np.int32) @@ -3058,12 +3367,19 @@ def _set_nonbonded_forces(self, context): context: openmm.Context The OpenMM context to use. """ + if self._integration_groups is None: + self._integration_groups = ( + context.getIntegrator().getIntegrationForceGroups() + ) + if self._nonbonded_force is None or ( self._is_fep and self._custom_nonbonded_force is None ): for force in context.getSystem().getForces(): if isinstance(force, _openmm.NonbondedForce): - self._nonbonded_force = force + # Only accept a force actually used for integration. + if self._integration_groups & (1 << force.getForceGroup()): + self._nonbonded_force = force elif self._is_fep and force.getName() == "GhostNonGhostNonbondedForce": self._custom_nonbonded_force = force elif self._pressure is None and "Barostat" in force.getName(): @@ -3085,6 +3401,12 @@ def _set_nonbonded_forces(self, context): _logger.error(msg) raise ValueError(msg) + if self._pme_groups is None: + groups = 1 << self._nonbonded_force.getForceGroup() + if self._is_fep: + groups |= 1 << self._custom_nonbonded_force.getForceGroup() + self._pme_groups = groups + def _get_target_position(self, positions): """ Get the current centre of the GCMC sphere. diff --git a/tests/test_energy.py b/tests/test_energy.py index 703a8a2..4a5ea06 100644 --- a/tests/test_energy.py +++ b/tests/test_energy.py @@ -48,6 +48,9 @@ def test_energy(fixture, softcore_form, platform, request): lambda_schedule=schedule, lambda_value=lambda_value, softcore_form=softcore_form, + # Sample within the region when there is one. Without a reference + # every move is a bulk move regardless. + bulk_sampling_probability=0.0 if reference is not None else 0.1, log_level="debug", ghost_file=None, log_file=None, @@ -79,6 +82,10 @@ def test_energy(fixture, softcore_form, platform, request): map=dyn_map, ) + # Empty the region, since at equilibrium it is full and insertions into it + # are rejected. A no-op when there is no region. + sampler.delete_waters(d.context()) + # Loop until we accept an insertion move. is_accepted = False while not is_accepted: @@ -616,3 +623,168 @@ def test_finalise_system(platform, water_box): assert actual_waters == expected_waters, ( f"Water count mismatch: got {actual_waters}, expected {expected_waters}" ) + + +@pytest.mark.skipif( + "CUDA_VISIBLE_DEVICES" not in os.environ, + reason="Requires CUDA enabled GPU.", +) +@pytest.mark.parametrize("platform", ["cuda", "opencl"]) +def test_set_lambda_uploads_parameters(sd12, platform): + """ + Test that set_lambda() uploads the parameters for the new lambda value. + + A sampler may be re-used across lambda values, re-uploading the lambda + dependent non-bonded parameters rather than rebuilding a context. If the + upload were skipped, or the wrong entry taken from the cache, the sampler + would evaluate insertion and deletion energies against the parameters of + the lambda value it was built at, while reporting the new one. + """ + + mols, reference = sd12 + + # The end states, so that the parameters differ as much as they can. + build_lambda = 0.0 + run_lambda = 1.0 + + sampler = GCMCSampler( + mols, + cutoff_type="rf", + cutoff="10 A", + reference=reference, + lambda_schedule=sr.cas.LambdaSchedule.standard_morph(), + lambda_value=build_lambda, + lambda_values=[run_lambda], + log_level="error", + ghost_file=None, + log_file=None, + test=True, + platform=platform, + ) + + assert sampler._is_fep, "the test system must be perturbable" + + names = ("_gpu_charge", "_gpu_sigma", "_gpu_epsilon", "_gpu_alpha") + built = sampler._lambda_params[(build_lambda, sampler._rest2_scale)] + wanted = sampler._lambda_params[(run_lambda, sampler._rest2_scale)] + + # The two end points must differ, otherwise the test cannot tell whether + # the upload happened. + assert any(not np.allclose(a, b) for a, b in zip(built, wanted)), ( + "the parameters are the same at both lambda values" + ) + + sampler.set_lambda(run_lambda) + assert sampler._lambda_value == run_lambda + + sampler.push() + try: + for name, expected in zip(names, wanted): + on_gpu = np.asarray(sampler._backend.from_gpu(getattr(sampler, name))) + assert np.allclose(on_gpu, expected), ( + f"{name} does not hold the parameters for lambda {run_lambda}" + ) + finally: + sampler.pop() + + +@pytest.mark.skipif( + "CUDA_VISIBLE_DEVICES" not in os.environ, + reason="Requires CUDA enabled GPU.", +) +@pytest.mark.parametrize("platform", ["cuda", "opencl"]) +def test_energy_after_set_lambda(sd12, platform): + """ + Test that the RF energy difference agrees with OpenMM after the sampler has + been switched to a different lambda value. + + This checks the uploaded parameters through the physics rather than by + inspecting them, so it also covers the kernel using them correctly. + + The move has to happen where the perturbation is. Bulk sampling would + place the water anywhere in the box, typically tens of Angstrom from the + perturbable molecule, where the lambda dependent parameters contribute + nothing and the comparison holds however wrong they are. The sphere is + emptied first, since at equilibrium it is full and insertions into it are + rejected. + """ + + mols, reference = sd12 + + schedule = sr.cas.LambdaSchedule.standard_morph() + + # The end states, so that the parameters differ as much as they can. + build_lambda = 0.0 + run_lambda = 1.0 + + sampler = GCMCSampler( + mols, + cutoff_type="rf", + cutoff="10 A", + reference=reference, + lambda_schedule=schedule, + lambda_value=build_lambda, + lambda_values=[run_lambda], + bulk_sampling_probability=0.0, + log_level="debug", + ghost_file=None, + log_file=None, + test=True, + platform=platform, + ) + + sampler.set_lambda(run_lambda) + + d = sampler.system().dynamics( + cutoff_type="rf", + cutoff="10 A", + temperature="298 K", + pressure=None, + constraint="h_bonds", + timestep="2 fs", + schedule=schedule, + lambda_value=run_lambda, + shift_coulomb=str(sampler._shift_coulomb), + shift_delta=str(sampler._shift_delta), + platform=platform, + ) + + def potential_energy(): + return ( + d.context() + .getState(getEnergy=True) + .getPotentialEnergy() + .value_in_unit(openmm.unit.kilocalories_per_mole) + ) + + # Empty the sphere so that the next accepted move is an insertion into it. + sampler.delete_waters(d.context()) + + for _ in range(50): + initial_energy = potential_energy() + moves = sampler.move(d.context()) + if moves: + break + else: + pytest.fail("no GCMC move was accepted") + + energy_difference = potential_energy() - initial_energy + sampler_energy = sampler._debug["energy_coul"] + sampler._debug["energy_lj"] + + # The move must be near the atoms whose parameters perturb, otherwise the + # comparison cannot see them. + charges0 = np.asarray(sampler._lambda_params[(build_lambda, 1.0)][0]) + charges1 = np.asarray(sampler._lambda_params[(run_lambda, 1.0)][0]) + changed = np.where(np.abs(charges0 - charges1) > 1e-9)[0] + positions = ( + d.context().getState(getPositions=True).getPositions(asNumpy=True) + / omm_unit.angstrom + ) + oxygen = positions[sampler._water_indices[sampler._debug["idx"]]] + distances = np.linalg.norm(positions[changed] - oxygen, axis=1) + assert (distances <= 10).any(), ( + "no perturbing atom within the cutoff of the move, so the lambda " + "dependent parameters are not being tested" + ) + + assert math.isclose(energy_difference, sampler_energy, abs_tol=1e-2) diff --git a/tests/test_stats.py b/tests/test_stats.py new file mode 100644 index 0000000..41baeb6 --- /dev/null +++ b/tests/test_stats.py @@ -0,0 +1,224 @@ +import pytest + +from loch import GCMCSampler + + +def make_sampler(lambda_value=0.0, lambda_values=None, is_fep=True): + """ + Create a sampler with only the attributes the statistics use, so that the + bookkeeping can be tested without a system or a GPU. + """ + sampler = object.__new__(GCMCSampler) + sampler._lambda_value = lambda_value + sampler._rest2_scale = 1.0 + sampler._is_fep = is_fep + sampler._lambda_values = lambda_values + sampler._stats = {} + sampler._zero_stats() + return sampler + + +def do_moves(sampler, num_moves): + """Pretend that a number of moves were performed and all were accepted.""" + sampler._num_moves += num_moves + sampler._num_accepted += num_moves + + +def switch(sampler, lambda_value): + """Switch lambda, as set_lambda() does.""" + sampler._switch_stats(lambda_value) + sampler._lambda_value = lambda_value + + +class TestStatsKey: + """Tests for the key used to store statistics.""" + + @pytest.mark.parametrize( + "lambda_value, expected", + [ + (0.0, "0.00000"), + (1, "1.00000"), + (0.33333, "0.33333"), + (1.0 / 3.0, "0.33333"), + ], + ) + def test_key_format(self, lambda_value, expected): + """Keys are formatted to five decimal places, as SOMD2 does.""" + assert GCMCSampler.stats_key(lambda_value) == expected + + def test_key_is_stable_across_representations(self): + """Values that agree to five decimal places share a key.""" + assert GCMCSampler.stats_key(0.1 + 0.2) == GCMCSampler.stats_key(0.3) + + +class TestPerLambdaStats: + """Tests for statistics accumulated per lambda value.""" + + def test_isolated_between_lambdas(self): + """Moves at one lambda must not be counted at another.""" + sampler = make_sampler(lambda_values=[0.0, 0.5]) + + do_moves(sampler, 3) + switch(sampler, 0.5) + + # The new lambda starts from zero. + assert sampler._num_moves == 0 + + do_moves(sampler, 7) + switch(sampler, 0.0) + + # Returning restores the original count, not the total. + assert sampler._num_moves == 3 + + stats = sampler.get_stats() + assert stats["0.00000"]["num_moves"] == 3 + assert stats["0.50000"]["num_moves"] == 7 + + def test_accumulates_across_visits(self): + """Revisiting a lambda continues from where it left off.""" + sampler = make_sampler(lambda_values=[0.0, 0.5]) + + do_moves(sampler, 3) + switch(sampler, 0.5) + do_moves(sampler, 7) + switch(sampler, 0.0) + do_moves(sampler, 2) + + stats = sampler.get_stats() + assert stats["0.00000"]["num_moves"] == 5 + assert stats["0.50000"]["num_moves"] == 7 + + def test_current_lambda_is_reported(self): + """The lambda in use is included alongside the archived ones.""" + sampler = make_sampler(lambda_values=[0.0, 0.5]) + do_moves(sampler, 4) + + assert sampler.get_stats() == { + "0.00000": { + "num_moves": 4, + "num_accepted": 4, + "num_insertions": 0, + "num_deletions": 0, + "num_accepted_attempts": 0, + } + } + + def test_non_alchemical_has_a_single_key(self): + """A non-alchemical system reports the same shape, with one key.""" + sampler = make_sampler(is_fep=False) + do_moves(sampler, 6) + + stats = sampler.get_stats() + assert list(stats) == ["0.00000"] + assert stats["0.00000"]["num_moves"] == 6 + + def test_reset_clears_every_lambda(self): + """reset() zeroes the current lambda and discards the others.""" + sampler = make_sampler(lambda_values=[0.0, 0.5]) + do_moves(sampler, 3) + switch(sampler, 0.5) + do_moves(sampler, 7) + + sampler.reset() + + assert sampler.get_stats() == { + "0.50000": { + "num_moves": 0, + "num_accepted": 0, + "num_insertions": 0, + "num_deletions": 0, + "num_accepted_attempts": 0, + } + } + + +class TestRestoreStats: + """Tests for restoring statistics, e.g. from a checkpoint.""" + + def test_round_trip(self): + """Statistics survive a save and restore.""" + sampler = make_sampler(lambda_values=[0.0, 0.5]) + do_moves(sampler, 3) + switch(sampler, 0.5) + do_moves(sampler, 7) + stats = sampler.get_stats() + + restored = make_sampler(lambda_value=0.5, lambda_values=[0.0, 0.5]) + restored.restore_stats(stats) + + assert restored.get_stats() == stats + assert restored._num_moves == 7 + + def test_unvisited_lambdas_are_ignored(self): + """ + A sampler keeps only its own lambda values. + + Each sampler can be handed the statistics for a whole simulation. If it + kept the others, it would report stale values for lambdas it never + samples, which could overwrite the live ones when merged. + """ + stats = { + "0.00000": { + "num_moves": 5, + "num_accepted": 5, + "num_insertions": 0, + "num_deletions": 0, + "num_accepted_attempts": 0, + }, + "1.00000": { + "num_moves": 9, + "num_accepted": 9, + "num_insertions": 0, + "num_deletions": 0, + "num_accepted_attempts": 0, + }, + } + + sampler = make_sampler(lambda_values=[0.0]) + sampler.restore_stats(stats) + + assert list(sampler.get_stats()) == ["0.00000"] + + def test_merge_order_cannot_clobber(self): + """Merging several samplers is safe regardless of order.""" + first = make_sampler(lambda_value=0.0, lambda_values=[0.0]) + second = make_sampler(lambda_value=1.0, lambda_values=[1.0]) + + do_moves(first, 5) + do_moves(second, 9) + + merged = {} + merged.update(first.get_stats()) + merged.update(second.get_stats()) + + # Restart both from the merged statistics, then advance one of them. + first = make_sampler(lambda_value=0.0, lambda_values=[0.0]) + second = make_sampler(lambda_value=1.0, lambda_values=[1.0]) + first.restore_stats(merged) + second.restore_stats(merged) + do_moves(first, 100) + + for order in ([first, second], [second, first]): + remerged = {} + for sampler in order: + remerged.update(sampler.get_stats()) + assert remerged["0.00000"]["num_moves"] == 105 + assert remerged["1.00000"]["num_moves"] == 9 + + def test_missing_lambda_is_zeroed(self): + """A lambda absent from the statistics starts from zero.""" + sampler = make_sampler(lambda_value=0.5, lambda_values=[0.0, 0.5]) + sampler.restore_stats( + { + "0.00000": { + "num_moves": 5, + "num_accepted": 5, + "num_insertions": 0, + "num_deletions": 0, + "num_accepted_attempts": 0, + } + } + ) + + assert sampler._num_moves == 0 + assert sampler.get_stats()["0.00000"]["num_moves"] == 5