Use fast reduce potentials computation for energy matrix
@@ -2512,24 +2512,23 @@ class ReplicaExchange(object):
# Compute energy for all thermodynamic states.
for energies, states in [(energy_thermodynamic_states, self._thermodynamic_states),
(energy_unsampled_states, self._unsampled_states)]:
- for i, state in enumerate(states):
- # Check if we need to request a new Context.
- if i == 0 or not state.is_state_compatible(context_state):
- context_state = state
+ # Group thermodynamic states by compatibility.
+ compatible_groups, original_indices = mmtools.states.group_by_compatibility(states)
+ # Compute the reduced potentials of all the compatible states.
+ for compatible_group, state_indices in zip(compatible_groups, original_indices):
# Get the context, any Integrator works.
- context, integrator = mmtools.cache.global_context_cache.get_context(state)
+ context, integrator = mmtools.cache.global_context_cache.get_context(compatible_group[0])
# Update positions and box vectors. We don't need
# to set Context velocities for the potential.
sampler_state.apply_to_context(context, ignore_velocities=True)
- else:
- # If this state is compatible with the context, just fix the
- # thermodynamic state as positions/box vectors are the same.
- state.apply_to_context(context)
- # Compute energy.
- energies[i] = state.reduced_potential(context)
+ # Compute and update the reduced potentials.
+ compatible_energies = mmtools.states.ThermodynamicState.reduced_potential_at_states(
+ context, compatible_group)
+ for energy_idx, state_idx in enumerate(state_indices):
+ energies[state_idx] = compatible_energies[energy_idx]
# Return the new energies.
return energy_thermodynamic_states, energy_unsampled_states
