Seed perturbed coords from a system clone, not raw arrays.

Author: lohedges

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:

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 <sire.system.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 <sire.system.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: