Fix restrained atom indices
@@ -1688,11 +1688,19 @@ class MultiStateSamplerAnalyzer(PhaseAnalyzer):
restraint_force = copy.deepcopy(restraint_force)
is_periodic = restraint_force.usesPeriodicBoundaryConditions()
- # Store the original indices of the restrained atoms.
- original_restrained_atom_indices1 = restraint_force.restrained_atom_indices1
- original_restrained_atom_indices2 = restraint_force.restrained_atom_indices2
- original_restrained_atom_indices = (original_restrained_atom_indices1 +
- original_restrained_atom_indices2)
+ # Store the indices of the restrained atoms in the reduced system.
+ analysis_indices = self._reporter.analysis_particle_indices
+
+ mapped_restrained_indices1 = restraint_force.restrained_atom_indices1
+ mapped_restrained_indices2 = restraint_force.restrained_atom_indices2
+
+ mapped_restrained_indices1 = [analysis_indices.index(index)
+ for index in mapped_restrained_indices1]
+ mapped_restrained_indices2 = [analysis_indices.index(index)
+ for index in mapped_restrained_indices2]
+
+ mapped_restrained_indices = (mapped_restrained_indices1 +
+ mapped_restrained_indices2)
# Create new system with only solute and restraint forces.
reduced_system = openmm.System()
@@ -1745,7 +1753,7 @@ class MultiStateSamplerAnalyzer(PhaseAnalyzer):
analysis_particles_only=True)
for replica_idx, sampler_state in enumerate(sampler_states):
- sliced_sampler_state = sampler_state[original_restrained_atom_indices]
+ sliced_sampler_state = sampler_state[mapped_restrained_indices]
sliced_sampler_state.apply_to_context(context)
potential_energy = context.getState(getEnergy=True).getPotentialEnergy()
self._restraint_energies[iteration][replica_idx] = potential_energy / _OPENMM_ENERGY_UNIT
@@ -1762,11 +1770,11 @@ class MultiStateSamplerAnalyzer(PhaseAnalyzer):
dtype=np.float32)
trajectory.image_molecules(inplace=True, anchor_molecules=anchor_molecules,
other_molecules=imaged_molecules)
- positions_group1 = trajectory.xyz[0][original_restrained_atom_indices1]
- positions_group2 = trajectory.xyz[0][original_restrained_atom_indices2]
+ positions_group1 = trajectory.xyz[0][mapped_restrained_indices1]
+ positions_group2 = trajectory.xyz[0][mapped_restrained_indices2]
else:
- positions_group1 = sampler_state.positions[original_restrained_atom_indices1]
- positions_group2 = sampler_state.positions[original_restrained_atom_indices2]
+ positions_group1 = sampler_state.positions[mapped_restrained_indices1]
+ positions_group2 = sampler_state.positions[mapped_restrained_indices2]
positions_group1 /= _MDTRAJ_DISTANCE_UNIT
positions_group2 /= _MDTRAJ_DISTANCE_UNIT
