<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Check ensemble of OpenMM temperature replica exchange simulations
Note
Step1:  Ensemble validation is particularly useful for validating enhanced sampling methods such
as temperature replica exchange (parallel tempering) molecular dynamics simulations. In
replica exchange MD, configurational swaps are attempted periodically between simulations
running concurrently in separate simulation cells at different temperatures. Since the
acceptance criteria for exchange moves depends on the energies of each replica involved,
it is critical that the energy distributions are correct at all thermodynamic states used
the parallel tempering scheme.
In this example, parallel tempering simulations on a simple coarse-grained oligomer system
were run using the openmmtools framework. 6 temperature states spaced logarithmically over 
the range of 300K to 500K, over which the oligomer undergoes a transition from a helix to 
unfolded state.
Step2:  For convenience, we will use native physical_validation units
Step3:  Create the UnitData object which we will use to inform physical_validation of our choices. Note that the conversion factors will all be 1.0, since we're using native units. The notation here is more general though, and would allow to change any unit set in the previous cell.
Step4:  We first read the replica exchange output file and extract the thermodynamics states used in the simulation. We will use these states to create the EnsembleData objects which inform physical_validation of the sampled ensembles. Note that the example simulations were performed at constant temperature in a non-periodic box. We will therefore set the (undefined) volume of the NVT ensemble to -1.
Step5:  We will then read the replica energies and the state indices. Note that
replica_state_indices[replica, step] denotes the thermodynamics state index sampled by replica replica during step step
replica_energies[replica, state, step] is the reduced potential of replica replica at state state during step step
Step6:  For our analysis, we are only interested in the energies at the states the replicas were sampling at, rather than at all states. It will also make our remaining analysis easier if these energies are organized as time series per state rather than time series per replica.
Also, the replica energies are stored in reduced potential. For our analysis, we are interested in the non-reduced form, which we can obtain by multiplying the result by kT of the respective thermodynamic state.
Step7:  We now have all the required information to create a SimulationData object for each separate state
Step8:  We can now run the ensemble validation on all adjacent temperature pairs
<ASSISTANT_TASK:>
Python Code:
# enable plotting in notebook
%matplotlib notebook
Explanation: Check ensemble of OpenMM temperature replica exchange simulations
Note: This notebook can be run locally by cloning the
Github repository.
The notebook is located in doc/examples/openmm_replica_exchange.ipynb. The input and output files of the simulation are located in doc/examples/simulation_results/openMMTemperatureReplicaExchange/.
Be aware that probabilistic quantities such as error estimates based on bootstrapping
will differ when repeating the analysis.
End of explanation
import numpy as np
import physical_validation
import openmmtools.multistate
import simtk.unit
Explanation: Ensemble validation is particularly useful for validating enhanced sampling methods such
as temperature replica exchange (parallel tempering) molecular dynamics simulations. In
replica exchange MD, configurational swaps are attempted periodically between simulations
running concurrently in separate simulation cells at different temperatures. Since the
acceptance criteria for exchange moves depends on the energies of each replica involved,
it is critical that the energy distributions are correct at all thermodynamic states used
the parallel tempering scheme.
In this example, parallel tempering simulations on a simple coarse-grained oligomer system
were run using the openmmtools framework. 6 temperature states spaced logarithmically over 
the range of 300K to 500K, over which the oligomer undergoes a transition from a helix to 
unfolded state.
End of explanation
energy_unit = simtk.unit.kilojoule_per_mole
length_unit = simtk.unit.nanometer
volume_unit = length_unit ** 3
temperature_unit = simtk.unit.kelvin
pressure_unit = simtk.unit.bar
time_unit = simtk.unit.picosecond
kb = simtk.unit.MOLAR_GAS_CONSTANT_R.value_in_unit(energy_unit / temperature_unit)
Explanation: For convenience, we will use native physical_validation units:
End of explanation
unit_data = physical_validation.data.UnitData(
    kb=kb,
    energy_conversion=energy_unit.conversion_factor_to(simtk.unit.kilojoule_per_mole),
    length_conversion=length_unit.conversion_factor_to(simtk.unit.nanometer),
    volume_conversion=volume_unit.conversion_factor_to(simtk.unit.nanometer ** 3),
    temperature_conversion=temperature_unit.conversion_factor_to(simtk.unit.kelvin),
    pressure_conversion=pressure_unit.conversion_factor_to(simtk.unit.bar),
    time_conversion=time_unit.conversion_factor_to(simtk.unit.picosecond),
    energy_str=energy_unit.get_symbol(),
    length_str=length_unit.get_symbol(),
    volume_str=volume_unit.get_symbol(),
    temperature_str=temperature_unit.get_symbol(),
    pressure_str=pressure_unit.get_symbol(),
    time_str=time_unit.get_symbol(),
)
Explanation: Create the UnitData object which we will use to inform physical_validation of our choices. Note that the conversion factors will all be 1.0, since we're using native units. The notation here is more general though, and would allow to change any unit set in the previous cell.
End of explanation
output_data = "simulation_results/openMMTemperatureReplicaExchange/output/output.nc"
reporter = openmmtools.multistate.MultiStateReporter(output_data, open_mode="r")
states = reporter.read_thermodynamic_states()[0]
num_states = len(states)
ensemble_data = [
    physical_validation.data.EnsembleData(
        ensemble="NVT",
        natoms=state.n_particles,
        volume=-1,
        temperature=state.temperature.value_in_unit(temperature_unit),
    )
    for state in states
]
Explanation: We first read the replica exchange output file and extract the thermodynamics states used in the simulation. We will use these states to create the EnsembleData objects which inform physical_validation of the sampled ensembles. Note that the example simulations were performed at constant temperature in a non-periodic box. We will therefore set the (undefined) volume of the NVT ensemble to -1.
End of explanation
analyzer = openmmtools.multistate.ReplicaExchangeAnalyzer(reporter)
replica_energies, _, _, replica_state_indices = analyzer.read_energies()
Explanation: We will then read the replica energies and the state indices. Note that
replica_state_indices[replica, step] denotes the thermodynamics state index sampled by replica replica during step step
replica_energies[replica, state, step] is the reduced potential of replica replica at state state during step step
End of explanation
# Prepare array of kT values
kT = np.array(
    [kb * state.temperature.value_in_unit(temperature_unit) for state in states]
)
total_steps = replica_energies.shape[2]
potential_energies = []
for state in range(num_states):
    state_energy = np.zeros(total_steps)
    for step in range(total_steps):
        # Find the replica which sampled at state `state` during step `step`
        state_energy[step] = replica_energies[
            np.nonzero(replica_state_indices[:, step] == state), state, step
        ]
    # Append non-reduced potential energy time series
    potential_energies.append(state_energy * kT[state])
Explanation: For our analysis, we are only interested in the energies at the states the replicas were sampling at, rather than at all states. It will also make our remaining analysis easier if these energies are organized as time series per state rather than time series per replica.
Also, the replica energies are stored in reduced potential. For our analysis, we are interested in the non-reduced form, which we can obtain by multiplying the result by kT of the respective thermodynamic state.
End of explanation
simulation_data = []
for ensemble, potential_energy in zip(ensemble_data, potential_energies):
    simulation_data.append(
        physical_validation.data.SimulationData(
            units=unit_data,
            ensemble=ensemble,
            observables=physical_validation.data.ObservableData(
                potential_energy=potential_energy
            ),
        )
    )
Explanation: We now have all the required information to create a SimulationData object for each separate state:
End of explanation
for simulation_lower, simulation_upper in zip(
    simulation_data[:-1], simulation_data[1:]
):
    physical_validation.ensemble.check(
        data_sim_one=simulation_lower, data_sim_two=simulation_upper, screen=True
    )
Explanation: We can now run the ensemble validation on all adjacent temperature pairs:
End of explanation
<END_TASK>