def process_replica_exchange_data(
output_data="output/output.nc",
output_directory="output",
series_per_page=4,
write_data_file=True,
plot_production_only=False,
print_timing=False,
equil_nskip=1,
frame_begin=0,
frame_end=-1,
):
"""
Read replica exchange simulation data, detect equilibrium and decorrelation time, and plot replica exchange results.
:param output_data: path to output .nc file from replica exchange simulation, (default='output/output.nc')
:type output_data: str
:param output_directory: path to which output files will be written (default='output')
:type output_directory: stry
:param series_per_page: number of replica data series to plot per pdf page (default=4)
:type series_per_page: int
:param write_data_file: Option to write a text data file containing the state_energies array (default=True)
:type write_data_file: Boolean
:param plot_production_only: Option to plot only the production region, as determined from pymbar detectEquilibration (default=False)
:type plot_production_only: Boolean
:param equil_nskip: skip this number of frames to sparsify the energy timeseries for pymbar detectEquilibration (default=1) - this is used only when frame_begin=0 and the trajectory has less than 40000 frames.
:type equil_nskip: Boolean
:param frame_begin: analyze starting from this frame, discarding all prior as equilibration period (default=0)
:type frame_begin: int
:param frame_end: analyze up to this frame only, discarding the rest (default=-1).
:type frame_end: int
:returns:
- replica_energies ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps] ), simtk.unit ) ) - The potential energies for all replicas at all (printed) time steps
- replica_state_indices ( np.int64( [number_replicas,number_simulation_steps] ), simtk.unit ) - The thermodynamic state assignments for all replicas at all (printed) time steps
- production_start ( int - The frame at which the production region begins for all replicas, as determined from pymbar detectEquilibration
- sample_spacing ( int - The number of frames between uncorrelated state energies, estimated using heuristic algorithm )
- n_transit ( np.float( [number_replicas] ) ) - Number of half-transitions between state 0 and n for each replica
- mixing_stats ( tuple ( np.float( [number_replicas x number_replicas] ) , np.float( [ number_replicas ] ) , float( statistical inefficiency ) ) ) - transition matrix, corresponding eigenvalues, and statistical inefficiency
"""
t1 = time.perf_counter()
# Read the simulation coordinates for individual temperature replicas
reporter = MultiStateReporter(output_data, open_mode="r")
t2 = time.perf_counter()
if print_timing:
print(f"open data time: {t2-t1}")
# figure out what the time between output is.
# We assume all use the same time step (which i think is required)
mcmove = reporter.read_mcmc_moves()[0]
time_interval = mcmove.n_steps * mcmove.timestep
t3 = time.perf_counter()
if print_timing:
print(f"read_mcmc_moves time: {t3-t2}")
# figure out what the temperature list is
states = reporter.read_thermodynamic_states()[0]
t4 = time.perf_counter()
if print_timing:
print(f"read_thermodynamics_states time: {t4-t3}")
temperature_list = []
for s in states:
temperature_list.append(s.temperature)
analyzer = ReplicaExchangeAnalyzer(reporter)
t5 = time.perf_counter()
(
replica_energies,
unsampled_state_energies,
neighborhoods,
replica_state_indices,
) = analyzer.read_energies()
# Truncate output of read_energies() to last frame of interest
if frame_end > 0:
# Use frames from frame_begin to frame_end
replica_energies = replica_energies[:, :, :frame_end]
unsampled_state_energies = unsampled_state_energies[:, :, :frame_end]
neighborhoods = neighborhoods[:, :, :frame_end]
replica_state_indices = replica_state_indices[:, :frame_end]
t6 = time.perf_counter()
if print_timing:
print(f"read_energies time: {t6-t5}")
n_particles = np.shape(
reporter.read_sampler_states(iteration=0)[0].positions)[0]
temps = np.array([temp._value for temp in temperature_list])
beta_k = 1 / (kB * temps)
n_replicas = len(temperature_list)
for k in range(n_replicas):
replica_energies[:, k, :] *= beta_k[k]**(-1)
t7 = time.perf_counter()
if print_timing:
print(f"reduce replica energies time: {t7-t6}")
total_steps = len(replica_energies[0][0])
state_energies = np.zeros([n_replicas, total_steps])
t8 = time.perf_counter()
# there must be some better way to do this as list comprehension.
for step in range(total_steps):
for state in range(n_replicas):
state_energies[state, step] = replica_energies[np.where(
replica_state_indices[:, step] == state)[0], 0, step]
t9 = time.perf_counter()
if print_timing:
print(f"assign state energies time: {t9-t8}")
# can run physical-valication on these state_energies
# Use pymbar timeseries module to detect production period
t10 = time.perf_counter()
# Start of equilibrated data:
t0 = np.zeros((n_replicas))
# Statistical inefficiency:
g = np.zeros((n_replicas))
subsample_indices = {}
# If sufficiently large, discard the first 20000 frames as equilibration period and use
# subsampleCorrelatedData to get the energy decorrelation time.
if total_steps >= 40000 or frame_begin > 0:
if frame_begin > 0:
# If specified, use frame_begin as the start of the production region
production_start = frame_begin
else:
# Otherwise, use frame 20000
production_start = 20000
for state in range(n_replicas):
subsample_indices[state] = timeseries.subsampleCorrelatedData(
state_energies[state][production_start:],
conservative=True,
)
g[state] = subsample_indices[state][1] - subsample_indices[state][0]
else:
# For small trajectories, use detectEquilibration
for state in range(n_replicas):
t0[state], g[state], Neff_max = timeseries.detectEquilibration(
state_energies[state], nskip=equil_nskip)
# Choose the latest equil timestep to apply to all states
production_start = int(np.max(t0))
# Assume a normal distribution (very rough approximation), and use mean plus
# the number of standard deviations which leads to (n_replica-1)/n_replica coverage
# For 12 replicas this should be the mean + 1.7317 standard deviations
# x standard deviations is the solution to (n_replica-1)/n_replica = erf(x/sqrt(2))
# This is equivalent to a target of 23/24 CDF value
print(f"g: {g.astype(int)}")
def erf_fun(x):
return np.power((erf(x / np.sqrt(2)) - (n_replicas - 1) / n_replicas),
2)
# x must be larger than zero
opt_g_results = minimize_scalar(erf_fun, bounds=(0, 10))
if not opt_g_results.success:
print("Error solving for correlation time, exiting...")
print(f"erf opt results: {opt_g_results}")
exit()
sample_spacing = int(np.ceil(np.mean(g) + opt_g_results.x * np.std(g)))
t11 = time.perf_counter()
if print_timing:
print(f"detect equil and subsampling time: {t11-t10}")
print("state mean energies variance")
for state in range(n_replicas):
state_mean = np.mean(state_energies[state,
production_start::sample_spacing])
state_std = np.std(state_energies[state,
production_start::sample_spacing])
print(f" {state:4d} {state_mean:10.6f} {state_std:10.6f}")
t12 = time.perf_counter()
if write_data_file == True:
f = open(os.path.join(output_directory, "replica_energies.dat"), "w")
for step in range(total_steps):
f.write(f"{step:10d}")
for replica_index in range(n_replicas):
f.write(
f"{replica_energies[replica_index,replica_index,step]:12.6f}"
)
f.write("\n")
f.close()
t13 = time.perf_counter()
if print_timing:
print(f"Optionally write .dat file: {t13-t12}")
t14 = time.perf_counter()
if plot_production_only == True:
plot_replica_exchange_energies(
state_energies[:, production_start:],
temperature_list,
series_per_page,
time_interval=time_interval,
time_shift=production_start * time_interval,
file_name=f"{output_directory}/rep_ex_ener.pdf",
)
plot_replica_exchange_energy_histograms(
state_energies[:, production_start:],
temperature_list,
file_name=f"{output_directory}/rep_ex_ener_hist.pdf",
)
plot_replica_exchange_summary(
replica_state_indices[:, production_start:],
temperature_list,
series_per_page,
time_interval=time_interval,
time_shift=production_start * time_interval,
file_name=f"{output_directory}/rep_ex_states.pdf",
)
plot_replica_state_matrix(
replica_state_indices[:, production_start:],
file_name=f"{output_directory}/state_probability_matrix.pdf",
)
else:
plot_replica_exchange_energies(
state_energies,
temperature_list,
series_per_page,
time_interval=time_interval,
file_name=f"{output_directory}/rep_ex_ener.pdf",
)
plot_replica_exchange_energy_histograms(
state_energies,
temperature_list,
file_name=f"{output_directory}/rep_ex_ener_hist.pdf",
)
plot_replica_exchange_summary(
replica_state_indices,
temperature_list,
series_per_page,
time_interval=time_interval,
file_name=f"{output_directory}/rep_ex_states.pdf",
)
plot_replica_state_matrix(
replica_state_indices,
file_name=f"{output_directory}/state_probability_matrix.pdf",
)
t15 = time.perf_counter()
if print_timing:
print(f"plotting time: {t15-t14}")
# Analyze replica exchange state transitions
# For each replica, how many times does the thermodynamic state go between state 0 and state n
# For consistency with the other mixing statistics, use only the production region here
replica_state_indices_prod = replica_state_indices[:, production_start:]
# Number of one-way transitions from states 0 to n or states n to 0
n_transit = np.zeros((n_replicas, 1))
# Replica_state_indices is [n_replicas x n_iterations]
for rep in range(n_replicas):
last_bound = None
for i in range(replica_state_indices_prod.shape[1]):
if replica_state_indices_prod[
rep, i] == 0 or replica_state_indices_prod[rep, i] == (
n_replicas - 1):
if last_bound is None:
# This is the first time state 0 or n is visited
pass
else:
if last_bound != replica_state_indices_prod[rep, i]:
# This is a completed transition from 0 to n or n to 0
n_transit[rep] += 1
last_bound = replica_state_indices_prod[rep, i]
t16 = time.perf_counter()
if print_timing:
print(f"replica transition analysis: {t16-t15}")
# Compute transition matrix from the analyzer
mixing_stats = analyzer.generate_mixing_statistics(
number_equilibrated=production_start)
t17 = time.perf_counter()
if print_timing:
print(f"compute transition matrix: {t17-t16}")
print(f"total time elapsed: {t17-t1}")
return (replica_energies, replica_state_indices, production_start,
sample_spacing, n_transit, mixing_stats)
def make_replica_dcd_files(topology,
timestep=5 * unit.femtosecond,
time_interval=200,
output_dir="output",
output_data="output.nc",
checkpoint_data="output_checkpoint.nc",
frame_begin=0,
frame_stride=1):
"""
Make dcd files from replica exchange simulation trajectory data.
:param topology: OpenMM Topology
:type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_
:param timestep: Time step used in the simulation (default=5*unit.femtosecond)
:type timestep: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>` float * simtk.unit
:param time_interval: frequency, in number of time steps, at which positions were recorded (default=200)
:type time_interval: int
:param output_directory: path to which we will write the output (default='output')
:type output_directory: str
:param output_data: name of output .nc data file (default='output.nc')
:type output_data: str
:param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
:type checkpoint_data: str
:param frame_begin: Frame at which to start writing the dcd trajectory (default=0)
:type frame_begin: int
:param frame_stride: advance by this many time intervals when writing dcd trajectories (default=1)
:type frame_stride: int
"""
file_list = []
output_data_path = os.path.join(output_dir, output_data)
# Get number of replicas:
reporter = MultiStateReporter(output_data_path, open_mode="r")
states = reporter.read_thermodynamic_states()[0]
n_replicas = len(states)
sampler_states = reporter.read_sampler_states(iteration=0)
xunit = sampler_states[0].positions[0].unit
for replica_index in range(n_replicas):
replica_trajectory = extract_trajectory(
topology,
replica_index=replica_index,
output_data=output_data_path,
checkpoint_data=checkpoint_data,
frame_begin=frame_begin,
frame_stride=frame_stride)
file_name = f"{output_dir}/replica_{replica_index+1}.dcd"
file = open(file_name, "wb")
dcd_file = DCDFile(file,
topology,
timestep,
firstStep=frame_begin,
interval=time_interval)
for positions in replica_trajectory:
# Add the units consistent with replica_energies
positions *= xunit
DCDFile.writeModel(dcd_file, positions)
file.close()
file_list.append(file_name)
return file_list
def make_state_dcd_files(topology,
timestep=5 * unit.femtosecond,
time_interval=200,
output_dir="output",
output_data="output.nc",
checkpoint_data="output_checkpoint.nc",
frame_begin=0,
frame_stride=1,
center=True):
"""
Make dcd files by state from replica exchange simulation trajectory data.
Note: these are discontinuous trajectories with constant temperature state.
:param topology: OpenMM Topology
:type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_
:param timestep: Time step used in the simulation (default=5*unit.femtosecond)
:type timestep: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>` float * simtk.unit
:param time_interval: frequency, in number of time steps, at which positions were recorded (default=200)
:type time_interval: int
:param output_directory: path to which we will write the output (default='output')
:type output_directory: str
:param output_data: name of output .nc data file (default='output.nc')
:type output_data: str
:param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
:type checkpoint_data: str
:param frame_begin: Frame at which to start writing the dcd trajectory (default=0)
:type frame_begin: int
:param frame_stride: advance by this many time intervals when writing dcd trajectories (default=1)
:type frame_stride: int
:param center: align the center of mass of each structure in the discontinuous state trajectory (default=True)
:type center: Boolean
"""
file_list = []
output_data_path = os.path.join(output_dir, output_data)
# Get number of states:
reporter = MultiStateReporter(output_data_path, open_mode="r")
states = reporter.read_thermodynamic_states()[0]
sampler_states = reporter.read_sampler_states(iteration=0)
xunit = sampler_states[0].positions[0].unit
for state_index in range(len(states)):
state_trajectory = extract_trajectory(topology,
state_index=state_index,
output_data=output_data_path,
checkpoint_data=checkpoint_data,
frame_begin=frame_begin,
frame_stride=frame_stride)
file_name = f"{output_dir}/state_{state_index+1}.dcd"
file = open(file_name, "wb")
dcd_file = DCDFile(file,
topology,
timestep,
firstStep=frame_begin,
interval=time_interval)
# TODO: replace this with MDTraj alignment tool
if center == True:
center_x = np.mean(state_trajectory[0, :, 0])
center_y = np.mean(state_trajectory[0, :, 1])
center_z = np.mean(state_trajectory[0, :, 2])
for positions in state_trajectory:
if center == True:
positions[:, 0] += (center_x - np.mean(positions[:, 0]))
positions[:, 1] += (center_y - np.mean(positions[:, 1]))
positions[:, 2] += (center_z - np.mean(positions[:, 2]))
# Add the units consistent with replica_energies
positions *= xunit
DCDFile.writeModel(dcd_file, positions)
file.close()
file_list.append(file_name)
return file_list
def make_state_pdb_files(topology,
output_dir="output",
output_data="output.nc",
checkpoint_data="output_checkpoint.nc",
frame_begin=0,
frame_stride=1,
center=True):
"""
Make pdb files by state from replica exchange simulation trajectory data.
Note: these are discontinuous trajectories with constant temperature state.
:param topology: OpenMM Topology
:type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_
:param output_directory: path to which we will write the output (default='output')
:type output_directory: str
:param output_data: name of output .nc data file (default='output.nc')
:type output_data: str
:param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
:type checkpoint_data: str
:param frame_begin: Frame at which to start writing the pdb trajectory (default=0)
:type frame_begin: int
:param frame_stride: advance by this many frames when writing pdb trajectories (default=1)
:type frame_stride: int
:param center: align the center of mass of each structure in the discontinuous state trajectory (default=True)
:type center: Boolean
:returns:
- file_list ( List( str ) ) - A list of names for the files that were written
"""
file_list = []
output_data_path = os.path.join(output_dir, output_data)
# Get number of states:
reporter = MultiStateReporter(output_data_path, open_mode="r")
states = reporter.read_thermodynamic_states()[0]
sampler_states = reporter.read_sampler_states(iteration=0)
xunit = sampler_states[0].positions[0].unit
for state_index in range(len(states)):
state_trajectory = extract_trajectory(topology,
state_index=state_index,
output_data=output_data_path,
checkpoint_data=checkpoint_data,
frame_begin=frame_begin,
frame_stride=frame_stride)
file_name = f"{output_dir}/state_{state_index+1}.pdb"
file = open(file_name, "w")
PDBFile.writeHeader(topology, file=file)
modelIndex = 1
# TODO: replace this with MDTraj alignment tool
if center == True:
center_x = np.mean(state_trajectory[0, :, 0])
center_y = np.mean(state_trajectory[0, :, 1])
center_z = np.mean(state_trajectory[0, :, 2])
for positions in state_trajectory:
if center == True:
positions[:, 0] += (center_x - np.mean(positions[:, 0]))
positions[:, 1] += (center_y - np.mean(positions[:, 1]))
positions[:, 2] += (center_z - np.mean(positions[:, 2]))
# Add the units consistent with replica_energies
positions *= xunit
PDBFile.writeModel(topology,
positions,
file=file,
modelIndex=modelIndex)
PDBFile.writeFooter(topology, file=file)
file.close()
file_list.append(file_name)
return file_list
def physical_validation_ensemble(output_data="output.nc",
output_directory="ouput",
plotfile='ensemble_check',
pairs='single',
ref_state_index=0):
"""
Run ensemble physical validation test for 2 states in replica exchange simulation
:param output_data: Path to the output data for a NetCDF-formatted file containing replica exchange simulation data
:type output_data: str
:param plotfile: Filename for outputting ensemble check plot
:type plotfile: str
:param pairs: Option for running ensemble validation on all replica pair combinations ('all'), adjacent pairs ('adjacent'), or single pair with optimal spacing ('single')
:param ref_state_index: Index in temperature_list to use as one of the states in the ensemble check. The other state will be chosen based on the energy standard deviation at the reference state. Ignored if pairs='all'
:type ref_state_index: int
"""
# Get temperature list and read the energies for individual temperature replicas
reporter = MultiStateReporter(output_data, open_mode="r")
analyzer = ReplicaExchangeAnalyzer(reporter)
states = reporter.read_thermodynamic_states()[0]
temperature_list = []
for s in states:
temperature_list.append(s.temperature)
(
replica_energies,
unsampled_state_energies,
neighborhoods,
replica_state_indices,
) = analyzer.read_energies()
n_particles = np.shape(
reporter.read_sampler_states(iteration=0)[0].positions)[0]
T_unit = temperature_list[0].unit
temps = np.array([temp.value_in_unit(T_unit) for temp in temperature_list])
beta_k = 1 / (kB.value_in_unit(unit.kilojoule_per_mole / T_unit) * temps)
n_replicas = len(temperature_list)
for k in range(n_replicas):
replica_energies[:, k, :] *= beta_k[k]**(-1)
total_steps = len(replica_energies[0][0])
state_energies = np.zeros([n_replicas, total_steps])
for step in range(total_steps):
for state in range(n_replicas):
state_energies[state, step] = replica_energies[np.where(
replica_state_indices[:, step] == state)[0], 0, step]
state_energies *= unit.kilojoule_per_mole
T_array = np.zeros(len(temperature_list))
for i in range(len(temperature_list)):
T_array[i] = temperature_list[i].value_in_unit(T_unit)
if pairs.lower() != 'single' and pairs.lower(
) != 'adjacent' and pairs.lower() != 'all':
print(
f"Error: Pair option '{pairs}' not recognized, using default option 'single'"
)
pairs = 'single'
if pairs.lower() == 'single':
# Run ensemble validation on one optimally spaced temperature pair
quantiles = {}
# Find optimal state pair for ensemble check:
# Compute standard deviations of each energy distribution:
state_energies_std = np.std(state_energies, axis=1)
# Select reference state point
T_ref = temperature_list[ref_state_index]
std_ref = state_energies_std[ref_state_index]
# Compute optimal spacing:
deltaT = 2 * kB * T_ref**2 / std_ref
#print("DeltaT: %r" %deltaT)
# Find closest match
T_diff = np.abs(T_ref.value_in_unit(T_unit) - T_array)
T_opt_index = np.argmin(np.abs(deltaT.value_in_unit(T_unit) - T_diff))
T_opt = temperature_list[T_opt_index]
# Set SimulationData for physical validation
state1_index = ref_state_index
state2_index = T_opt_index
sim_data1, sim_data2 = set_simulation_data(state_energies, T_array,
state1_index, state2_index)
# Run physical validation
try:
quantiles_ij = pv.ensemble.check(sim_data1,
sim_data2,
total_energy=False,
filename=plotfile)
quantiles[
f"state{state1_index}_state{state2_index}"] = quantiles_ij[0]
except InputError:
print(
f"Insufficient overlap between trajectories for states {state1_index} and {state2_index}. Skipping..."
)
elif pairs.lower() == 'adjacent':
# Run ensemble validation on all adjacent temperature pairs
quantiles = {}
for i in range(len(temperature_list) - 1):
# Set SimulationData for physical validation
state1_index = i
state2_index = i + 1
sim_data1, sim_data2 = set_simulation_data(state_energies, T_array,
state1_index,
state2_index)
# Run physical validation
try:
quantiles_ij = pv.ensemble.check(
sim_data1,
sim_data2,
total_energy=False,
filename=f"{plotfile}_{state1_index}_{state2_index}")
quantiles[
f"state{state1_index}_state{state2_index}"] = quantiles_ij[
0]
except InputError:
print(
f"Insufficient overlap between trajectories for states {state1_index} and {state2_index}. Skipping..."
)
elif pairs.lower() == 'all':
# Run ensemble validation on all combinations of temperature pairs
quantiles = {}
for i in range(len(temperature_list)):
for j in range(i + 1, len(temperature_list)):
# Set SimulationData for physical validation
state1_index = i
state2_index = j
sim_data1, sim_data2 = set_simulation_data(
state_energies, T_array, state1_index, state2_index)
# Run physical validation
try:
quantiles_ij = pv.ensemble.check(
sim_data1,
sim_data2,
total_energy=False,
filename=f"{plotfile}_{state1_index}_{state2_index}")
quantiles[
f"state{state1_index}_state{state2_index}"] = quantiles_ij[
0]
except InputError:
print(
f"Insufficient overlap between trajectories for states {state1_index} and {state2_index}. Skipping..."
)
return quantiles