Fix replica mixing.

src/somd2/runner/_repex.py

@@ -173,8 +173,10 @@ def mix_states(self):
         for i, (old_state, new_state) in enumerate(zip(self._old_states, self._states)):
             # The state has changed.
             if old_state != new_state:
-                positions, velocities = self._omm_states[new_state]
-                self._dynamics[i]._d._omm_mols.setStates(self._omm_states[new_state])
+                _logger.debug(
+                    f"Replica {i} changed state from {old_state} to {new_state}"
+                )
+                self._dynamics[i]._d._omm_mols.setState(self._omm_states[new_state])
 
 
 class RepexRunner(_RunnerBase):
@@ -591,6 +593,9 @@ def _mix_replicas(num_replicas, energy_matrix, states):
             The new states.
         """
 
+        # Copy the states.
+        states = states.copy()
+
         for swap in range(num_replicas**3):
             # Choose two replicas to swap.
             replica_i = _np.random.randint(num_replicas)