Skip to content

Seed perturbed coords from a system clone, not raw arrays #140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lohedges merged 1 commit into from
Apr 10, 2026
+35 −72
Merged

Seed perturbed coords from a system clone, not raw arrays #140

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 21 additions & 22 deletions src/somd2/runner/_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__,
Expand Down Expand Up @@ -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:
Expand Down
64 changes: 14 additions & 50 deletions src/somd2/runner/_repex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,7 @@ def __init__(
dynamics_kwargs,
gcmc_kwargs=None,
output_directory=None,
perturbed_positions=None,
perturbed_box=None,
perturbed_system=None,
):
"""
Constructor.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
Loading