Repex¶

Replica-exchange simulation algorithms and specific variants.

This module provides a general facility for running replica-exchange simulations, as well as derived classes for special cases such as parallel tempering (in which the states differ only in temperature).

Provided classes include:

ReplicaExchange

Base class for general replica-exchange simulations.
ParallelTempering

Convenience subclass of ReplicaExchange for parallel tempering simulations (one System object, many temperatures).
Reporter

Replica Exchange reporter class to store all variables and data

COPYRIGHT

Current version by Andrea Rizzi <andrea.rizzi@choderalab.org>, Levi N. Naden <levi.naden@choderalab.org> and John D. Chodera<john.chodera@choderalab.org> while at Memorial Sloan-Kettering Cancer Center.

Original version by John D. Chodera <jchodera@gmail.com> while at the University of California Berkeley.

LICENSE

This code is licensed under the latest available version of the MIT License.

class yank.repex.Reporter(storage, open_mode=None, checkpoint_interval=10, checkpoint_storage=None, analysis_particle_indices=())[source]¶

Handle storage write/read operations and different format conventions.

You can use this object to programmatically inspect the data generated by ReplicaExchange.

Parameters:

storage : str

The path to the storage file for analysis.

A second checkpoint file will be determined from either checkpoint_storage or automatically based on the storage option

In the future this will be able to take Storage classes as well.

open_mode : str or None

The mode of the file between ‘r’, ‘w’, and ‘a’ (or equivalently ‘r+’).

If None, the storage file won’t be open on construction, and a call to Reporter.open() will be needed before attempting read/write operations.

checkpoint_interval : int >= 1, Default: 10

The frequency at which checkpointing information is written relative to analysis information.

This is a multiple of the iteration at which energies is written, hence why it must be greater than or equal to 1. Checkpoint information cannot be written on iterations which where iteration % checkpoint_interval != 0.

Attempting to read checkpointing information results in a masked array where only entries which were written are unmasked

checkpoint_storage : str or None, optional

Optional name of the checkpoint point file. This file is used to save trajectory information and other less frequently accessed data.

This should NOT be a full path, and instead just a filename

If None: the derived checkpoint name is the same as storage, less any extension, then “_checkpoint.nc” is added.

The reporter internally tracks what data goes into which file, so its transparent to all other classes In the future, this will be able to take Storage classes as well

analysis_particle_indices : tuple of ints, Optional. Default: () (empty tuple)

Indices of particles which should be treated as special when manipulating read and write functions. If this is an empty tuple, then no particles are treated as special

Attributes

filepath Returns the string file name of the primary storage file

filepath¶

Returns the string file name of the primary storage file

Classes outside the Reporter can access the file string for error messages and such.

analysis_particle_indices¶: Return the tuple of indices of the particles which additional information is stored on for analysis

storage_exists(skip_size=False)[source]¶

Check if the storage files exist on disk.

Reads information on the primary file to see existence of others

Parameters:

skip_size : bool, Optional, Default: False

Skip the check of the file size. Helpful if you have just initialized a storage file but written nothing to it yet and/or its still entirely in memory (e.g. just opened NetCDF files)

Returns:

files_exist : bool

If the primary storage file and its related subfiles exist, returns True. If the primary file or any subfiles do not exist, returns False

is_open()[source]¶: Return True if the Reporter is ready to read/write.

open(mode='r', convention='ReplicaExchange', netcdf_format='NETCDF4')[source]¶

Open the storage file for reading/writing.

Creates and pre-formats the required files if they don’t exist.

This is not necessary if you have indicated in the constructor to open.

Parameters:

mode : str, Optional, Default: ‘r’

The mode of the file between ‘r’, ‘w’, and ‘a’ (or equivalently ‘r+’).

convention : str, Optional, Default: ‘ReplicaExchange’

NetCDF convention to write

netcdf_format : str, Optional, Default: ‘NETCDF4’

The NetCDF file format to use

close()[source]¶: Close the storage files

sync()[source]¶: Force any buffer to be flushed to the file

read_thermodynamic_states()[source]¶

Retrieve the stored thermodynamic states from the checkpoint file.

Returns:

thermodynamic_states : list of ThermodynamicStates

The previously stored thermodynamic states. During the simulation these are swapped among replicas.

unsampled_states : list of ThermodynamicState

The previously stored unsampled thermodynamic states.

See also

write_replica_thermodynamic_states

read_sampler_states(iteration, analysis_particles_only=False)[source]¶

Retrieve the stored sampler states on the checkpoint file

If the iteration is not on the checkpoint interval, None is returned. Exception to this is if``analysis_particles_only``, see the Returns for output behavior.

Parameters:

iteration : int

The iteration at which to read the data.

analysis_particles_only : bool, Optional, Default: False

If set, will return the trajectory of ONLY the analysis particles flagged on original creation of the files

Returns:

sampler_states : list of SamplerStates or None

The previously stored sampler states for each replica.

If the iteration is not on the checkpoint_interval and the analysis_particles_only is not set, None is returned

If analysis_particles_only is set, will return only the subset of particles which were defined by the analysis_particle_indices on reporter creation

write_sampler_states(sampler_states, iteration)[source]¶

Store all sampler states for a given iteration on the checkpoint file

If the iteration is not on the checkpoint interval, only the analysis_particle_indices data is written, if set.

Parameters:

sampler_states : list of SamplerStates

The sampler states to store for each replica.

iteration : int

The iteration at which to store the data.

read_replica_thermodynamic_states(iteration=slice(None, None, None))[source]¶

Retrieve the indices of the ThermodynamicStates for each replica on the analysis file

Parameters:

iteration : int or slice

The iteration(s) at which to read the data. The slice(None) allows fetching all iterations at once.

Returns:

state_indices : list of int

At the given iteration, replica i propagated the system in SamplerState sampler_states[i] and ThermodynamicState thermodynamic_states[states_indices[i]].

write_replica_thermodynamic_states(state_indices, iteration)[source]¶

Store the indices of the ThermodynamicStates for each replica on the analysis file

Parameters:

state_indices : list of int of size n_replicas

At the given iteration, replica i propagated the system in SamplerState sampler_states[i] and ThermodynamicState thermodynamic_states[replica_thermodynamic_states[i]].

iteration : int

The iteration at which to store the data.

read_mcmc_moves()[source]¶

Return the MCMCMoves of the yank.repex.ReplicaExchange simulation on the checkpoint

Returns:

mcmc_moves : list of MCMCMove

The MCMCMoves used to propagate the simulation.

write_mcmc_moves(mcmc_moves)[source]¶

Store the MCMCMoves of the yank.repex.ReplicaExchange simulation on the checkpoint

Parameters:

mcmc_moves : list of MCMCMove

The MCMCMoves used to propagate the simulation.

read_energies(iteration=slice(None, None, None))[source]¶

Retrieve the energy matrix at the given iteration on the analysis file

Parameters:

iteration : int or slice

The iteration(s) at which to read the data. The slice(None) allows fetching all iterations at once.

Returns:

energy_thermodynamic_states : n_replicas x n_states numpy.ndarray

energy_thermodynamic_states[iteration, i, j] is the reduced potential computed at SamplerState sampler_states[iteration, i] and ThermodynamicState thermodynamic_states[iteration, j].

energy_unsampled_states : n_replicas x n_unsampled_states numpy.ndarray

energy_unsampled_states[iteration, i, j] is the reduced potential computed at SamplerState sampler_states[iteration, i] and ThermodynamicState unsampled_thermodynamic_states[iteration, j].

write_energies(energy_thermodynamic_states, energy_unsampled_states, iteration)[source]¶

Store the energy matrix at the given iteration on the analysis file

Parameters:

energy_thermodynamic_states : n_replicas x n_states numpy.ndarray

energy_thermodynamic_states[i][j] is the reduced potential computed at SamplerState sampler_states[i] and ThermodynamicState thermodynamic_states[j].

energy_unsampled_states : n_replicas x n_unsampled_states numpy.ndarray

energy_unsampled_states[i][j] is the reduced potential computed at SamplerState sampler_states[i] and ThermodynamicState unsampled_thermodynamic_states[j].

iteration : int

The iteration at which to store the data.

read_mixing_statistics(iteration=slice(None, None, None))[source]¶

Retrieve the mixing statistics for the given iteration on the analysis file

Parameters:

iteration : int or slice

The iteration(s) at which to read the data.

Returns:

n_accepted_matrix : numpy.ndarray with shape (n_states, n_states)

n_accepted_matrix[i][j] is the number of accepted moves from state thermodynamic_states[i] to thermodynamic_states[j] going from iteration-1 to iteration (not cumulative).

n_proposed_matrix : numpy.ndarray with shape (n_states, n_states)

n_proposed_matrix[i][j] is the number of proposed moves from state thermodynamic_states[i] to thermodynamic_states[j] going from iteration-1 to iteration (not cumulative).

write_mixing_statistics(n_accepted_matrix, n_proposed_matrix, iteration)[source]¶

Store the mixing statistics for the given iteration on the analysis file

Parameters:

n_accepted_matrix : numpy.ndarray with shape (n_states, n_states)

n_accepted_matrix[i][j] is the number of accepted moves from state thermodynamic_states[i] to thermodynamic_states[j] going from iteration-1 to iteration (not cumulative).

n_proposed_matrix : numpy.ndarray with shape (n_states, n_states)

n_proposed_matrix[i][j] is the number of proposed moves from state thermodynamic_states[i] to thermodynamic_states[j] going from iteration-1 to iteration (not cumulative).

iteration : int

The iteration for which to store the data.

read_timestamp(iteration=slice(None, None, None))[source]¶

Return the timestamp for the given iteration.

Read from the analysis file, although there is a paired timestamp on the checkpoint file as well

Parameters:

iteration : int or slice

The iteration(s) at which to read the data.

Returns:

timestamp : str

The timestamp at which the iteration was stored.

write_timestamp(iteration)[source]¶

Store a timestamp for the given iteration on both analysis and checkpoint file.

If the iteration is not on the checkpoint_interval, no timestamp is written on the checkpoint file

Parameters:

iteration : int

The iteration at which to read the data.

read_dict(path)[source]¶

Restore a dictionary from the storage file.

The method supports reading only specific dictionary keywords in path notation. If the dictionary is large, this can be quicker. However, note that, depending on how the dictionary was saved, this may end up reading the whole dictionary anyway.

Parameters:

path : str

The path to the dictionary or a keyword in the dictionary.

Returns:

data : dict

The restored data as a dict.

Examples

>>> import os
>>> import openmmtools as mmtools
>>> data = {'info': [1, 2, 3]}
>>> with mmtools.utils.temporary_directory() as temp_dir_path:
...     temp_file_path = os.path.join(temp_dir_path, 'temp.nc')
...     reporter = Reporter(temp_file_path, open_mode='w')
...     reporter.write_dict('data', data)
...     reporter.read_dict('data')
{'info': [1, 2, 3]}
...     reporter.read_dict('data/info')
[1, 2, 3]

write_dict(path, data)[source]¶

Store the contents of a dict.

Parameters:

path : str

The path to the dictionary in the storage file.

data : dict

The dict to store.

read_last_iteration(full_iteration=True)[source]¶

Read the last iteration from file which was written in sequential order.

Parameters:

full_iteration : bool, optional

If True, returns the last checkpoint iteration (default is True).

Returns:

last_iteration : int

Last iteration which was sequentially written.

write_last_iteration(iteration)[source]¶

Tell the reporter what the last iteration which was written in sequential order was to allow resuming and analysis only on valid data.

Call this as the last step of any write_iteration-like routine to ensure analysis will not use junk data left over from an interrupted simulation past the last checkpoint.

Parameters:

iteration : int

Iteration at which the last good data point was written.

read_mbar_free_energies(iteration)[source]¶

Read the MBAR dimensionless free energy at a given iteration from file.

These free energies are relative to the first state. These are computed through self-consistent iterations from an initial guess.

The initial guess is often 0 for all states, so any state not written is returned as zeros for f_k, and infinity for

Used primarily in online analysis, and should be used in tandem with an yank.analyze.YankPhaseAnalyzer object from the yank.analyze module

Parameters:

iteration : int

iteration to read the free energies from

if the iteration was not written at a the given iteration, then the free_energies are all 0

Returns:

f_k : 1-D ndarray of floats

MBAR free_energies from the iteration.

free_energy : tuple of two Floats

Free energy estimate at the iteration between the end states

write_mbar_free_energies(iteration, f_k, free_energy)[source]¶

Write the mbar free energies at the current iteration. See read_mbar_free_energies() for more information about pymbar’s f_k free energies

Used primarily in online analysis, and should be used in tandem with an yank.analyze.YankPhaseAnalyzer object from the analyze module.

Parameters:

iteration : int,

Iteration at which to save the free energy. Reads the current energy up to this value and stores it in the analysis reporter

f_k : 1D array of floats,

Dimensionless free energies you wish to store. This should come from pymbar.MBAR.f_k from an initialized MBAR object.

free_energy : tuple of two floats

Current estimate of the free energy difference of the phase and its error. This should be of the form (free_energy, error_in_free_energy) Typically output of the pymbar.MBAR.getFreeEnergyDifferences()[i,j]

class yank.repex.ReplicaExchange(mcmc_moves=None, number_of_iterations=1, replica_mixing_scheme='swap-all', online_analysis_interval=None, online_analysis_target_error=0.2, online_analysis_minimum_iterations=50)[source]¶

Replica-exchange simulation facility.

This base class provides a general replica-exchange simulation facility, allowing any set of thermodynamic states to be specified, along with a set of initial positions to be assigned to the replicas in a round-robin fashion.

No distinction is made between one-dimensional and multidimensional replica layout. By default, the replica mixing scheme attempts to mix all replicas to minimize slow diffusion normally found in multidimensional replica exchange simulations (Modification of the ‘replica_mixing_scheme’ setting will allow the traditional ‘neighbor swaps only’ scheme to be used.)

Stored configurations, energies, swaps, and restart information are all written to a single output file using the platform portable, robust, and efficient NetCDF4 library.

Parameters:

mcmc_moves : MCMCMove or list of MCMCMove, optional

The MCMCMove used to propagate the states. If a list of MCMCMoves, they will be assigned to the correspondent thermodynamic state on creation. If None is provided, Langevin dynamics with 2fm timestep, 5.0/ps collision rate, and 500 steps per iteration will be used.

number_of_iterations : int or infinity, optional, default: 1

The number of iterations to perform. Both float('inf') and numpy.inf are accepted for infinity. If you set this to infinity, be sure to set also online_analysis_interval.

replica_mixing_scheme : ‘swap-all’, ‘swap-neighbors’ or None, Default: ‘swap-all’

The scheme used to swap thermodynamic states between replicas.

online_analysis_interval : None or Int >= 1, optional, default None

Choose the interval at which to perform online analysis of the free energy.

After every interval, the simulation will be stopped and the free energy estimated.

If the error in the free energy estimate is at or below online_analysis_target_error, then the simulation will be considered completed.

online_analysis_target_error : float >= 0, optional, default 0.2

The target error for the online analysis measured in kT per phase.

Once the free energy is at or below this value, the phase will be considered complete.

If online_analysis_interval is None, this option does nothing.

online_analysis_minimum_iterations : int >= 0, optional, default 50

Set the minimum number of iterations which must pass before online analysis is carried out.

Since the initial samples likely not to yield a good estimate of free energy, save time and just skip them If online_analysis_interval is None, this does nothing

Examples

Parallel tempering simulation of alanine dipeptide in implicit solvent (replica exchange among temperatures). This is just an illustrative example; use ParallelTempering class for actual production parallel tempering simulations.

Create the system.

>>> import math
>>> from simtk import unit
>>> from openmmtools import testsystems, states, mcmc
>>> testsystem = testsystems.AlanineDipeptideImplicit()

Create thermodynamic states for parallel tempering with exponentially-spaced schedule.

>>> n_replicas = 3  # Number of temperature replicas.
>>> T_min = 298.0 * unit.kelvin  # Minimum temperature.
>>> T_max = 600.0 * unit.kelvin  # Maximum temperature.
>>> temperatures = [T_min + (T_max - T_min) * (math.exp(float(i) / float(nreplicas-1)) - 1.0) / (math.e - 1.0)
...                 for i in range(n_replicas)]
>>> thermodynamic_states = [states.ThermodynamicState(system=testsystem.system, temperature=T)
...                         for T in temperatures]

Initialize simulation object with options. Run with a GHMC integrator.

>>> move = mcmc.GHMCMove(timestep=2.0*unit.femtoseconds, n_steps=50)
>>> simulation = ReplicaExchange(mcmc_moves=move, number_of_iterations=2)

Create simulation with its storage file (in a temporary directory) and run.

>>> storage_path = tempfile.NamedTemporaryFile(delete=False).name + '.nc'
>>> reporter = Reporter(storage_path, checkpoint_interval=1)
>>> simulation.create(thermodynamic_states=thermodynamic_states,
>>>                   sampler_states=states.SamplerState(testsystem.positions),
>>>                   storage=reporter)
>>> simulation.run()  # This runs for a maximum of 2 iterations.
>>> simulation.iteration
2
>>> simulation.run(n_iterations=1)
>>> simulation.iteration
2

To resume a simulation from an existing storage file and extend it beyond the original number of iterations.

>>> del simulation
>>> simulation = ReplicaExchange.from_storage(reporter)
>>> simulation.extend(n_iterations=1)
>>> simulation.iteration
3

You can extract several information from the NetCDF file using the Reporter class while the simulation is running. This reads the SamplerStates of every run iteration.

>>> reporter = Reporter(storage=storage_path, open_mode='r', checkpoint_interval=1)
>>> sampler_states = reporter.read_sampler_states(iteration=range(1, 4))
>>> len(sampler_states)
3
>>> sampler_states[-1].positions.shape  # Alanine dipeptide has 22 atoms.
(22, 3)

Clean up.

>>> os.remove(storage_path)

Parameters:

number_of_iterations – Maximum number of integer iterations that will be run
replica_mixing_scheme – Scheme which describes how replicas are exchanged each iteration as string
online_analysis_interval – How frequently to carry out online analysis in number of iterations
online_analysis_target_error – Target free energy difference error float at which simulation will be stopped during online analysis, in dimensionless energy
online_analysis_minimum_iterations – Minimum number of iterations needed before online analysis is run as int

Attributes

`n_replicas`	The integer number of replicas (read-only).
`iteration`	The integer current iteration of the simulation (read-only).
`mcmc_moves`	A copy of the MCMCMoves list used to propagate the simulation.
`sampler_states`	A copy of the sampler states list at the current iteration.
`metadata`	A copy of the metadata dictionary passed on creation (read-only).
`is_completed`	Check if we have reached any of the stop target criteria (read-only)

classmethod from_storage(storage)[source]¶

Constructor from an existing storage file.

Parameters:

storage : str or Reporter

If str: The path to the storage file. If Reporter: uses the Reporter options In the future this will be able to take a Storage class as well.

Returns:

repex : ReplicaExchange

A new instance of ReplicaExchange in the same state of the last stored iteration.

class Status(iteration, target_error, is_completed)¶

count(value) → integer -- return number of occurrences of value¶

index(value[, start[, stop]]) → integer -- return first index of value.¶: Raises ValueError if the value is not present.

is_completed¶: Alias for field number 2

iteration¶: Alias for field number 0

target_error¶: Alias for field number 1

classmethod read_status(storage)[source]¶

Read the status of the calculation from the storage file.

This class method can be used to quickly check the status of the simulation before loading the full ReplicaExchange object from disk.

Parameters:

storage : str or Reporter

The path to the storage file or the reporter object.

Returns:

status : ReplicaExchange.Status

The status of the replica-exchange calculation. It has three fields: iteration, target_error, and is_completed.

n_replicas¶: The integer number of replicas (read-only).

iteration¶

The integer current iteration of the simulation (read-only).

If the simulation has not been created yet, this is None.

mcmc_moves¶

A copy of the MCMCMoves list used to propagate the simulation.

This can be set only before creation.

sampler_states¶

A copy of the sampler states list at the current iteration.

This can be set only before running.

online_analysis_interval¶: interval to carry out online analysis

metadata¶: A copy of the metadata dictionary passed on creation (read-only).

is_completed¶: Check if we have reached any of the stop target criteria (read-only)

create(thermodynamic_states, sampler_states, storage, unsampled_thermodynamic_states=None, metadata=None)[source]¶

Create new replica-exchange simulation.

Parameters:

thermodynamic_states : list of openmmtools.states.ThermodynamicState

Thermodynamic states to simulate, where one replica is allocated per state. Each state must have a system with the same number of atoms.

sampler_states : openmmtools.states.SamplerState or list

One or more sets of initial sampler states. If a list of SamplerStates, they will be assigned to replicas in a round-robin fashion.

storage : str or instanced Reporter

If str: the path to the storage file. Default checkpoint options from Reporter class are used If Reporter: Uses the reporter options and storage path In the future this will be able to take a Storage class as well.

unsampled_thermodynamic_states : list of openmmtools.states.ThermodynamicState, optional

These are ThermodynamicStates that are not propagated, but their reduced potential is computed at each iteration for each replica. These energy can be used as data for reweighting schemes (default is None).

metadata : dict, optional

Simulation metadata to be stored in the file.

minimize(tolerance=Quantity(value=1.0, unit=kilojoule/(nanometer*mole)), max_iterations=0)[source]¶

Minimize all replicas.

Minimized positions are stored at the end.

Parameters:

tolerance : simtk.unit.Quantity, optional

Minimization tolerance (units of energy/mole/length, default is 1.0 * unit.kilojoules_per_mole / unit.nanometers).

max_iterations : int, optional

Maximum number of iterations for minimization. If 0, minimization continues until converged.

equilibrate(n_iterations, mcmc_moves=None)[source]¶

Equilibrate all replicas.

This does not increase the iteration counter. The equilibrated positions are stored at the end.

Parameters:

n_iterations : int

Number of equilibration iterations.

mcmc_moves : MCMCMove or list of MCMCMove, optional

Optionally, the MCMCMoves to use for equilibration can be different from the ones used in production.

run(n_iterations=None)[source]¶

Run the replica-exchange simulation.

This runs at most number_of_iterations iterations. Use extend() to pass the limit.

Parameters:

n_iterations : int, optional

If specified, only at most the specified number of iterations will be run (default is None).

extend(n_iterations)[source]¶

Extend the simulation by the given number of iterations.

Contrarily to run(), this will extend the number of iterations past number_of_iteration if requested.

Parameters:

n_iterations : int

The number of iterations to run.

class yank.repex.ParallelTempering(mcmc_moves=None, number_of_iterations=1, replica_mixing_scheme='swap-all', online_analysis_interval=None, online_analysis_target_error=0.2, online_analysis_minimum_iterations=50)[source]¶

Parallel tempering simulation facility.

This class provides a facility for parallel tempering simulations. It is a subclass of ReplicaExchange, but provides efficiency improvements for parallel tempering simulations, so should be preferred for this type of simulation. In particular, this makes use of the fact that the reduced potentials are linear in inverse temperature.