Add option to enforce PBC when retrieving coords

femto/md/config.py

@@ -324,6 +324,12 @@ class HREMD(BaseModel):
         description="The number of cycles to run before saving the current replica "
         "states to DCD trajectory files. If ``None``, no trajectories will be saved.",
     )
+    trajectory_enforce_pbc: bool = pydantic.Field(
+        False,
+        description="Whether to apply periodic boundary conditions when retrieving "
+        "coordinates for writing to trajectory files.",
+    )
+
     checkpoint_interval: int | None = pydantic.Field(
         None,
         description="The number of cycles to run before saving the current replica "

femto/md/hremd.py

@@ -193,6 +193,7 @@ def _propagate_replicas(
     replica_idx_offset: int,
     force_groups: set[int] | int,
     max_retries: int,
+    enforce_pbc: bool = True,
 ):
     """Propagate all replica states forward in time.
 
@@ -233,7 +234,9 @@ def _propagate_replicas(
                 simulation.step(n_steps)
 
                 local_replica_coords = simulation.context.getState(
-                    getPositions=True, getVelocities=True
+                    getPositions=True,
+                    getVelocities=True,
+                    enforcePeriodicBox=enforce_pbc,
                 )
                 femto.md.utils.openmm.check_for_nans(local_replica_coords)
             except openmm.OpenMMException:
@@ -622,7 +625,10 @@ def run_hremd(
         )
 
         if initial_coords is None:
-            coords = [simulation.context.getState(getPositions=True)] * n_replicas
+            coords_0 = simulation.context.getState(
+                getPositions=True, enforcePeriodicBox=config.trajectory_enforce_pbc
+            )
+            coords = [coords_0] * n_replicas
         else:
             coords = [initial_coords[i + replica_idx_offset] for i in range(n_replicas)]
 
@@ -641,6 +647,7 @@ def run_hremd(
                 replica_idx_offset,
                 force_groups,
                 config.max_step_retries,
+                config.trajectory_enforce_pbc,
             )
 
             if mpi_comm.rank == 0:
@@ -672,6 +679,7 @@ def run_hremd(
                 replica_idx_offset,
                 force_groups,
                 config.max_step_retries,
+                config.trajectory_enforce_pbc,
             )
             reduced_potentials = mpi_comm.reduce(reduced_potentials, MPI.SUM, 0)
 

femto/md/simulate.py

@@ -80,6 +80,7 @@ def simulate_state(
     stages: list[femto.md.config.SimulationStage],
     platform: femto.md.constants.OpenMMPlatform,
     reporter: femto.md.reporting.openmm.OpenMMStateReporter | None = None,
+    enforce_pbc: bool = False,
 ) -> openmm.State:
     """Simulate a system following the specified ``stages``, at a given 'state' (i.e.
     a set of context parameters, such as free energy lambda values)
@@ -92,6 +93,8 @@ def simulate_state(
         platform: The accelerator to use.
         reporter: The reporter to use to record system statistics such as volume and
             energy.
+        enforce_pbc: Whether to enforce periodic boundary conditions when retrieving
+            the final coordinates.
 
     Returns:
         The final coordinates and box vectors.
@@ -144,7 +147,11 @@ def simulate_state(
             f"{femto.md.utils.openmm.get_simulation_summary(simulation)}"
         )
         coords = simulation.context.getState(
-            getPositions=True, getVelocities=True, getForces=True, getEnergy=True
+            getPositions=True,
+            getVelocities=True,
+            getForces=True,
+            getEnergy=True,
+            enforcePeriodicBox=enforce_pbc,
         )
 
     return coords
@@ -157,6 +164,7 @@ def simulate_states(
     stages: list[femto.md.config.SimulationStage],
     platform: femto.md.constants.OpenMMPlatform,
     reporter: femto.md.reporting.openmm.OpenMMStateReporter | None = None,
+    enforce_pbc: bool = False,
 ) -> list[openmm.State]:
     """Simulate the system at each 'state' using ``simulate_state``.
 
@@ -171,6 +179,8 @@ def simulate_states(
         platform: The accelerator to use.
         reporter: The reporter to use to record system statistics such as volume and
             energy.
+        enforce_pbc: Whether to enforce periodic boundary conditions when retrieving
+            the final coordinates.
 
     Returns:
         The final coordinates at each state.
@@ -194,6 +204,7 @@ def simulate_states(
                 stages,
                 platform,
                 reporter if state_idx == 0 else None,
+                enforce_pbc=enforce_pbc,
             )
 
         final_coords = femto.md.utils.mpi.reduce_dict(final_coords, mpi_comm)