From f6e0d6a3544b1a6f9197a3a1b9688c1f304b21ff Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 10 Apr 2026 10:55:07 +0100 Subject: [PATCH] Seed perturbed coords from a system clone, not raw arrays. --- src/somd2/runner/_base.py | 43 +++++++++++++------------ src/somd2/runner/_repex.py | 64 +++++++++----------------------------- 2 files changed, 35 insertions(+), 72 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index dcfd78c..f4c2ed7 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -94,28 +94,6 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) - # Make sure the coordinates property is linked. - perturbed_system = _sr.morph.link_to_perturbed( - self._config.perturbed_system - ) - - # Store the positions. - self._perturbed_positions = _sr.io.get_coords_array(perturbed_system) - - # Store the box vectors. - cell = self._config.perturbed_system.space().box_matrix() - c0 = cell.column0() - c1 = cell.column1() - c2 = cell.column2() - self._perturbed_box = ( - (c0.x().value(), c0.y().value(), c0.z().value()), - (c1.x().value(), c1.y().value(), c1.z().value()), - (c2.x().value(), c2.y().value(), c2.z().value()), - ) - else: - self._perturbed_positions = None - self._perturbed_box = None - # Log the versions of somd2 and sire. from somd2 import ( __version__, @@ -523,6 +501,27 @@ def __init__(self, system, config): # Store the current system as a reference. self._reference_system = self._system.clone() + # Create a clone of the fully-prepared reference system with the + # perturbed end-state coordinates and periodic space. This is done + # after all system preparation so that the clone inherits the same + # topology and properties. It is used to seed starting coordinates + # for lambda > 0.5 replicas. + if self._config.replica_exchange and self._config.perturbed_system is not None: + from sire.legacy.IO import setCoordinates as _setCoordinates + + pert_coords = _sr.io.get_coords_array( + _sr.morph.link_to_perturbed(self._config.perturbed_system) + ) + self._perturbed_system = _sr.system.System( + _setCoordinates(self._system._system, pert_coords.tolist()) + ) + self._perturbed_system.set_space(self._config.perturbed_system.space()) + + # Link properties to the lambda = 0 end state. + self._perturbed_system = _sr.morph.link_to_reference(self._perturbed_system) + else: + self._perturbed_system = None + # Check for a valid restart. if self._config.restart: if self._config.use_backup: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 931d26b..b019030 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -52,8 +52,7 @@ def __init__( dynamics_kwargs, gcmc_kwargs=None, output_directory=None, - perturbed_positions=None, - perturbed_box=None, + perturbed_system=None, ): """ Constructor. @@ -82,13 +81,9 @@ def __init__( output_directory: pathlib.Path The directory for simulation output. - perturbed_positions: numpy.ndarray - The positions for the perturbed state. If None, then the perturbed state - is not used. - - perturbed_box: numpy.ndarray - The box vectors for the perturbed state. If None, then the perturbed state - is not used. + perturbed_system: :class: `System ` + The perturbed end-state system used to seed starting coordinates for + lambda > 0.5 replicas. If None, the perturbed state is not used. """ # Warn if the number of replicas is not a multiple of the number of GPUs. @@ -121,8 +116,7 @@ def __init__( dynamics_kwargs, gcmc_kwargs=gcmc_kwargs, output_directory=output_directory, - perturbed_positions=perturbed_positions, - perturbed_box=perturbed_box, + perturbed_system=perturbed_system, ) def __setstate__(self, state): @@ -173,8 +167,7 @@ def _create_dynamics( dynamics_kwargs, gcmc_kwargs=None, output_directory=None, - perturbed_positions=None, - perturbed_box=None, + perturbed_system=None, ): """ Create the dynamics objects. @@ -203,13 +196,9 @@ def _create_dynamics( output_directory: pathlib.Path The directory for simulation output. - perturbed_positions: numpy.ndarray - The positions for the perturbed state. If None, then the perturbed state - is not used. - - perturbed_box: numpy.ndarray - The box vectors for the perturbed state. If None, then the perturbed state - is not used. + perturbed_system: :class: `System ` + The perturbed end-state system used to seed starting coordinates for + lambda > 0.5 replicas. If None, the perturbed state is not used. """ from math import floor @@ -256,7 +245,10 @@ def _create_dynamics( # This is a restart, get the system for this replica. if isinstance(system, list): mols = system[i] - # This is a new simulation. + # This is a new simulation. For lambda > 0.5, use the perturbed + # system to seed the starting coordinates and periodic space. + elif perturbed_system is not None and lam > 0.5: + mols = perturbed_system else: mols = system @@ -314,33 +306,6 @@ def _create_dynamics( _logger.error(msg) raise RuntimeError(msg) from e - # Update the box vectors and positions if the perturbed state is used. - if ( - perturbed_positions is not None - and perturbed_box is not None - and lam > 0.5 - ): - from openmm.unit import angstrom - - # Get the positions from the context. - positions = ( - dynamics.context() - .getState(getPositions=True) - .getPositions(asNumpy=True) - ) / angstrom - - # The positions array also contains the ghost water atoms that - # were added during the GCMC setup. We need to make sure that - # we copy these over to the perturbed positions array. - diff = len(positions) - len(perturbed_positions) - if diff != 0: - perturbed_positions = _np.concatenate( - [perturbed_positions, positions[-diff:]] - ) - - dynamics.context().setPeriodicBoxVectors(*perturbed_box * angstrom) - dynamics.context().setPositions(perturbed_positions * angstrom) - # Bind the GCMC sampler to the dynamics object. This allows the # dynamics object to reset the water state in its internal OpenMM # context following a crash recovery. @@ -782,8 +747,7 @@ def __init__(self, system, config): self._num_gpus, dynamics_kwargs, gcmc_kwargs=self._gcmc_kwargs, - perturbed_positions=self._perturbed_positions, - perturbed_box=self._perturbed_box, + perturbed_system=self._perturbed_system, output_directory=self._config.output_directory, ) else: