def read_replica_exchange_data(system=None,topology=None,temperature_list=None,output_data="output.nc",print_frequency=None):
"""
Read replica exchange simulation data.
:param system: OpenMM system object, default = None
:type system: `System() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1openmm_1_1System.html>`_
:param topology: OpenMM topology object, default = None
:type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_
:param temperature_list: List of temperatures that will be used to define different replicas (thermodynamics states), default = None
:type temperature_list: List( `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_ * number_replicas )
:param output_data: Path to the output data for a Yank, NetCDF-formatted file containing replica exchange simulation data, default = None
:type output_data: str
:param print_frequency: Number of simulation steps to skip when writing data, default = None
:type print_frequency: 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_positions ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps,cgmodel.num_beads,3] ), simtk.unit ) ) - The positions 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
:Example:
>>> from foldamers.cg_model.cgmodel import CGModel
>>> from cg_openmm.simulation.rep_exch import *
>>> cgmodel = CGModel()
>>> replica_energies,replica_positions,replica_state_indices = run_replica_exchange(cgmodel.topology,cgmodel.system,cgmodel.positions)
>>> replica_energies,replica_positions,replica_state_indices = read_replica_exchange_data(system=cgmodel.system,topology=cgmodel.topology,temperature_list=,output_data="output.nc",print_frequency=None)
"""
# Read the simulation coordinates for individual temperature replicas
reporter = MultiStateReporter(output_data, open_mode='r')
analyzer = ReplicaExchangeAnalyzer(reporter)
replica_energies,unsampled_state_energies,neighborhoods,replica_state_indices = analyzer.read_energies()
temps = np.array([temp._value for temp in temperature_list])
beta_k = 1 / (kB * temps)
for k in range(len(replica_energies)):
replica_energies[:,k,:]*=beta_k[k]**(-1)
total_steps = len(replica_energies[0][0])
replica_positions = unit.Quantity(np.zeros([len(temperature_list),total_steps,system.getNumParticles(),3]),unit.nanometer)
for step in range(total_steps):
sampler_states = reporter.read_sampler_states(iteration=step)
if type(sampler_states) == None:
print("ERROR: no data found for step "+str(step))
exit()
else:
for replica_index in range(len(temperature_list)):
for particle in range(system.getNumParticles()):
replica_positions[replica_index][step][particle] = sampler_states[replica_index].positions[particle]
return(replica_energies,replica_positions,replica_state_indices)
def run_replica_exchange(topology,system,positions,temperature_list=None,simulation_time_step=None,total_simulation_time=1.0 * unit.picosecond,output_data='output.nc',print_frequency=100,verbose_simulation=False,exchange_attempts=None,test_time_step=False,output_directory=None):
"""
Run a Yank replica exchange simulation using an OpenMM coarse grained model.
: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 system: OpenMM System()
:type system: `System() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1openmm_1_1System.html>`_
:param positions: Positions array for the model we would like to test
:type positions: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.array( [cgmodel.num_beads,3] ), simtk.unit )
:param temperature_list: List of temperatures for which to perform replica exchange simulations, default = None
:type temperature: List( float * simtk.unit.temperature )
:param simulation_time_step: Simulation integration time step
:type simulation_time_step: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_
:param total_simulation_time: Total run time for individual simulations
:type total_simulation_time: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_
:param output_data: Name of NETCDF file where we will write simulation data
:type output_data: string
:param print_frequency: Number of simulation steps to skip when writing to output, Default = 100
:type print_frequence: int
:param verbose_simulation: Determines how much output is printed during a simulation run. Default = False
:type verbose_simulation: Logical
:param exchange_attempts: Number of exchange attempts to make during a replica exchange simulation run, Default = None
:type exchange_attempts: int
:param test_time_step: Logical variable determining if a test of the time step will be performed, Default = False
:type test_time_step: Logical
:param output_directory: Path to which we will write the output from simulation runs.
:type output_directory: str
: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_positions ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps,cgmodel.num_beads,3] ), simtk.unit ) ) - The positions 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
:Example:
>>> from foldamers.cg_model.cgmodel import CGModel
>>> from cg_openmm.simulation.rep_exch import *
>>> cgmodel = CGModel()
>>> replica_energies,replica_positions,replica_state_indices = run_replica_exchange(cgmodel.topology,cgmodel.system,cgmodel.positions)
"""
if simulation_time_step == None:
simulation_time_step,force_threshold = get_simulation_time_step(topology,system,positions,temperature_list[-1],total_simulation_time)
simulation_steps = int(round(total_simulation_time.__div__(simulation_time_step)))
if exchange_attempts == None:
if simulation_steps > 10000:
exchange_attempts = round(simulation_steps/1000)
else:
exchange_attempts = 10
if temperature_list == None:
temperature_list = [(300.0 * unit.kelvin).__add__(i * unit.kelvin) for i in range(-50,50,10)]
num_replicas = len(temperature_list)
sampler_states = list()
thermodynamic_states = list()
# Define thermodynamic states.
box_vectors = system.getDefaultPeriodicBoxVectors()
for temperature in temperature_list:
thermodynamic_state = mmtools.states.ThermodynamicState(system=system, temperature=temperature)
thermodynamic_states.append(thermodynamic_state)
sampler_states.append(mmtools.states.SamplerState(positions,box_vectors=box_vectors))
# Create and configure simulation object.
move = mmtools.mcmc.LangevinDynamicsMove(timestep=simulation_time_step,collision_rate=5.0/unit.picosecond,n_steps=exchange_attempts, reassign_velocities=True)
simulation = ReplicaExchangeSampler(mcmc_moves=move, number_of_iterations=exchange_attempts)
if os.path.exists(output_data): os.remove(output_data)
reporter = MultiStateReporter(output_data, checkpoint_interval=1)
simulation.create(thermodynamic_states, sampler_states, reporter)
config_root_logger(verbose_simulation)
if not test_time_step:
num_attempts = 0
while num_attempts < 5:
try:
simulation.run()
#print("Replica exchange simulations succeeded with a time step of: "+str(simulation_time_step))
break
except:
num_attempts = num_attempts + 1
if num_attempts >= 5:
print("Replica exchange simulation attempts failed, try verifying your model/simulation settings.")
exit()
else:
simulation_time_step,force_threshold = get_simulation_time_step(topology,system,positions,temperature_list[-1],total_simulation_time)
print("The suggested time step for a simulation with this model is: "+str(simulation_time_step))
while simulation_time_step.__div__(2.0) > 0.001 * unit.femtosecond:
try:
print("Running replica exchange simulations with Yank...")
print("Using a time step of "+str(simulation_time_step))
print("Running each trial simulation for 1000 steps, with 10 exchange attempts.")
move = mmtools.mcmc.LangevinDynamicsMove(timestep=simulation_time_step,collision_rate=20.0/unit.picosecond,n_steps=10, reassign_velocities=True)
simulation = ReplicaExchangeSampler(replica_mixing_scheme='swap-neighbors',mcmc_moves=move,number_of_iterations=10)
reporter = MultiStateReporter(output_data, checkpoint_interval=1)
simulation.create(thermodynamic_states, sampler_states, reporter)
simulation.run()
print("Replica exchange simulations succeeded with a time step of: "+str(simulation_time_step))
break
except:
del simulation
os.remove(output_data)
print("Simulation attempt failed with a time step of: "+str(simulation_time_step))
if simulation_time_step.__div__(2.0) > 0.001 * unit.femtosecond:
simulation_time_step = simulation_time_step.__div__(2.0)
else:
print("Error: replica exchange simulation attempt failed with a time step of: "+str(simulation_time_step))
print("Please check the model and simulations settings, and try again.")
exit()
replica_energies,replica_positions,replica_state_indices = read_replica_exchange_data(system=system,topology=topology,temperature_list=temperature_list,output_data=output_data,print_frequency=print_frequency)
steps_per_stage = round(simulation_steps/exchange_attempts)
if output_directory != None:
plot_replica_exchange_energies(replica_energies,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage,output_directory=output_directory)
plot_replica_exchange_summary(replica_state_indices,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage,output_directory=output_directory)
else:
plot_replica_exchange_energies(replica_energies,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage)
plot_replica_exchange_summary(replica_state_indices,temperature_list,simulation_time_step,steps_per_stage=steps_per_stage)
return(replica_energies,replica_positions,replica_state_indices)
def main():
parser = argparse.ArgumentParser(
description=
'Compute a potential of mean force (PMF) for porin permeation.')
parser.add_argument('--index',
dest='index',
action='store',
type=int,
help='Index of ')
parser.add_argument('--output',
dest='output_filename',
action='store',
default='output.nc',
help='output netcdf filename (default: output.nc)')
args = parser.parse_args()
index = args.index
output_filename = args.output_filename
logger = logging.getLogger(__name__)
logging.root.setLevel(logging.DEBUG)
logging.basicConfig(level=logging.DEBUG)
yank.utils.config_root_logger(verbose=True, log_file_path=None)
# Configure ContextCache, platform and precision
from yank.experiment import ExperimentBuilder
platform = ExperimentBuilder._configure_platform('CUDA', 'mixed')
try:
openmmtools.cache.global_context_cache.platform = platform
except RuntimeError:
# The cache has been already used. Empty it before switching platform.
openmmtools.cache.global_context_cache.empty()
openmmtools.cache.global_context_cache.platform = platform
# Topology
pdbx = app.PDBxFile('mem_prot_md_system.pdbx')
# This system contains the CVforce with parameters different than zero
with open('openmm_system.xml', 'r') as infile:
openmm_system = XmlSerializer.deserialize(infile.read())
####### Indexes of configurations in trajectory ############################
configs = [
39, 141, 276, 406, 562, 668, 833, 1109, 1272, 1417, 1456, 1471, 1537,
1645, 1777, 1882
]
####### Indexes of states for series of replica exchange simulations #######
limits = [(0, 9), (10, 19), (20, 29), (30, 39), (40, 49), (50, 59),
(60, 69), (70, 79), (80, 89), (90, 99), (100, 109), (110, 119),
(120, 129), (130, 139), (140, 149), (150, 159)]
####### Reading positions from mdtraj trajectory ###########################
topology = md.Topology.from_openmm(pdbx.topology)
t = md.load('../../steered_md/comp7/forward/seed_0/steered_forward.nc',
top=topology)
positions = t.openmm_positions(configs[index])
thermodynamic_state_deserialized = states.ThermodynamicState(
system=openmm_system,
temperature=310 * unit.kelvin,
pressure=1.0 * unit.atmospheres)
sampler_state = states.SamplerState(
positions=positions,
box_vectors=t.unitcell_vectors[configs[index], :, :] * unit.nanometer)
logger.debug(type(sampler_state))
move = mcmc.LangevinDynamicsMove(timestep=2 * unit.femtosecond,
collision_rate=1.0 / unit.picoseconds,
n_steps=500,
reassign_velocities=False)
simulation = ReplicaExchangeSampler(mcmc_moves=move,
number_of_iterations=1)
analysis_particle_indices = topology.select(
'(protein and mass > 3.0) or (resname MER and mass > 3.0)')
reporter = MultiStateReporter(
output_filename,
checkpoint_interval=2000,
analysis_particle_indices=analysis_particle_indices)
first, last = limits[index]
# Initialize compound thermodynamic states
protocol = {
'lambda_restraints': [i / 159 for i in range(first, last + 1)],
'K_parallel': [
1250 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
],
'Kmax': [
500 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
],
'Kmin': [
500 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
]
}
my_composable_state = MyComposableState.from_system(openmm_system)
compound_states = states.create_thermodynamic_state_protocol(
thermodynamic_state_deserialized,
protocol=protocol,
composable_states=[my_composable_state])
simulation.create(thermodynamic_states=compound_states,
sampler_states=[sampler_state],
storage=reporter)
simulation.equilibrate(50, mcmc_moves=move)
simulation.run()
ts = simulation._thermodynamic_states[0]
context, _ = openmmtools.cache.global_context_cache.get_context(ts)
files_names = [
'state_{}_{}.log'.format(index, i) for i in range(first, last + 1)
]
files = []
for i, file in enumerate(files_names):
files.append(open(file, 'w'))
mpi.distribute(write_cv,
range(simulation.n_replicas),
context,
simulation,
files,
send_results_to=None)
for i in range(10000):
simulation.extend(n_iterations=2)
mpi.distribute(write_cv,
range(simulation.n_replicas),
context,
simulation,
files,
send_results_to=None)
def _setup(self):
"""
Set up calculation.
"""
# Signal that system has now been set up
if self._setup_complete:
raise Exception(
"System has already been set up---cannot run again.")
self._setup_complete = True
# Compute thermal energy and inverse temperature
self.kT = kB * self.temperature
self.beta = 1.0 / self.kT
# Create the system
self.system = self._create_system()
# Add a barostat
# TODO: Is this necessary, sicne ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat.
self._add_barostat()
# Restrain protein atoms in space
# TODO: Allow protein atom selection to be configured
selection = '((residue 342 and resname GLY) or (residue 97 and resname ASP) or (residue 184 and resname SER)) and (name CA)'
protein_atoms_to_restrain = self.mdtraj_topology.select(selection)
#self._restrain_protein(protein_atoms_to_restrain)
# Create SamplerState for initial conditions
self.sampler_state = states.SamplerState(
positions=self.grofile.positions,
box_vectors=self.grofile.getPeriodicBoxVectors())
# Create reference thermodynamic state
self.reference_thermodynamic_state = states.ThermodynamicState(
system=self.system,
temperature=self.temperature,
pressure=self.pressure)
# Anneal ligand into binding site
if self.anneal_ligand:
self._anneal_ligand()
# Create ThermodynamicStates for umbrella sampling along pore
self.thermodynamic_states = self._create_thermodynamic_states(
self.reference_thermodynamic_state)
# Minimize initial thermodynamic state
# TODO: Select initial thermodynamic state based on which state has minimum energy
initial_state_index = 0
self.sampler_state = self._minimize_sampler_state(
self.thermodynamic_states[initial_state_index], self.sampler_state)
# Set up simulation
from yank.multistate import SAMSSampler, MultiStateReporter
# TODO: Change this to LangevinSplittingDynamicsMove
move = mcmc.LangevinDynamicsMove(timestep=self.timestep,
collision_rate=self.collision_rate,
n_steps=self.n_steps_per_iteration,
reassign_velocities=False)
self.simulation = SAMSSampler(
mcmc_moves=move,
number_of_iterations=self.n_iterations,
online_analysis_interval=None,
gamma0=self.gamma0,
flatness_threshold=self.flatness_threshold)
self.reporter = MultiStateReporter(
self.output_filename,
checkpoint_interval=self.checkpoint_interval,
analysis_particle_indices=self.analysis_particle_indices)
self.simulation.create(
thermodynamic_states=self.thermodynamic_states,
unsampled_thermodynamic_states=[
self.reference_thermodynamic_state
],
sampler_states=[self.sampler_state],
initial_thermodynamic_states=[initial_state_index],
storage=self.reporter)
def setup_class(cls):
"""Shared test cases and variables."""
cls.checkpoint_interval = 2
# Make sure we collect the same number of samples for all tests to avoid instabilities in MBAR.
base_steps = 50
cls.n_steps = int(np.ceil(base_steps / cls.N_SAMPLERS))
# Test case with host guest in vacuum at 3 different positions and alchemical parameters.
# -----------------------------------------------------------------------------------------
hostguest_test = mmtools.testsystems.HostGuestVacuum()
factory = mmtools.alchemy.AbsoluteAlchemicalFactory()
alchemical_region = mmtools.alchemy.AlchemicalRegion(
alchemical_atoms=range(126, 156))
hostguest_alchemical = factory.create_alchemical_system(
hostguest_test.system, alchemical_region)
# Add restraint force between host and guest.
restraint_force = mmtools.forces.HarmonicRestraintBondForce(
spring_constant=2.0 * unit.kilojoule_per_mole / unit.angstrom**2,
restrained_atom_index1=10,
restrained_atom_index2=16,
)
hostguest_alchemical.addForce(copy.deepcopy(restraint_force))
# Translate the sampler states to be different one from each other.
positions = hostguest_test.positions
box_vectors = hostguest_test.system.getDefaultPeriodicBoxVectors()
hostguest_sampler_states = [
mmtools.states.SamplerState(
positions=positions + 10 * i * unit.nanometers,
box_vectors=box_vectors) for i in range(cls.N_SAMPLERS)
]
# Create the basic thermodynamic states.
hostguest_thermodynamic_states = [
mmtools.states.ThermodynamicState(hostguest_alchemical,
300 * unit.kelvin)
for _ in range(cls.N_STATES)
]
# Create alchemical states at different parameter values.
alchemical_states = [
mmtools.alchemy.AlchemicalState.from_system(hostguest_alchemical)
for _ in range(cls.N_STATES)
]
for i, alchemical_state in enumerate(alchemical_states):
alchemical_state.set_alchemical_parameters(
float(i) / (cls.N_STATES - 1))
# Create compound states.
hostguest_compound_states = list()
for i in range(cls.N_STATES):
hostguest_compound_states.append(
mmtools.states.CompoundThermodynamicState(
thermodynamic_state=hostguest_thermodynamic_states[i],
composable_states=[alchemical_states[i]]))
# Unsampled states.
nonalchemical_system = copy.deepcopy(hostguest_test.system)
nonalchemical_system.addForce(copy.deepcopy(restraint_force))
nonalchemical_state = mmtools.states.ThermodynamicState(
nonalchemical_system, 300 * unit.kelvin)
nonalchemical_compound_state = mmtools.states.CompoundThermodynamicState(
thermodynamic_state=nonalchemical_state,
composable_states=[RestraintState(lambda_restraints=1.0)])
hostguest_unsampled_states = [
copy.deepcopy(nonalchemical_compound_state) for _ in range(2)
]
cls.hostguest_test = (hostguest_compound_states,
hostguest_sampler_states,
hostguest_unsampled_states)
# Run a quick simulation
thermodynamic_states, sampler_states, unsampled_states = copy.deepcopy(
cls.hostguest_test)
n_states = len(thermodynamic_states)
# Prepare metadata for analysis.
reference_state = mmtools.states.ThermodynamicState(
hostguest_test.system, 300 * unit.kelvin)
topography = Topography(hostguest_test.topology,
ligand_atoms=range(126, 156))
metadata = {
'standard_state_correction': 4.0,
'reference_state': mmtools.utils.serialize(reference_state),
'topography': mmtools.utils.serialize(topography)
}
analysis_atoms = topography.receptor_atoms
# Create simulation and storage file.
cls.tmp_dir = tempfile.mkdtemp()
storage_path = os.path.join(cls.tmp_dir, 'test_analyze.nc')
move = mmtools.mcmc.LangevinDynamicsMove(n_steps=1)
cls.sampler = cls.SAMPLER(mcmc_moves=move,
number_of_iterations=cls.n_steps)
cls.reporter = MultiStateReporter(
storage_path,
checkpoint_interval=cls.checkpoint_interval,
analysis_particle_indices=analysis_atoms)
cls.call_sampler_create(cls.sampler,
cls.reporter,
thermodynamic_states,
sampler_states,
unsampled_states,
metadata=metadata)
# run some iterations
cls.n_replicas = cls.N_SAMPLERS
cls.n_states = n_states
cls.analysis_atoms = analysis_atoms
cls.sampler.run(cls.n_steps - 1) # Initial config
cls.repex_name = "RepexAnalyzer" # kind of an unused test
# Debugging Messages to sent to Nose with --nocapture enabled
output_descr = "Testing Analyzer: {} -- States: {} -- Samplers: {}".format(
cls.SAMPLER.__name__, cls.N_STATES, cls.N_SAMPLERS)
len_output = len(output_descr)
print("#" * len_output)
print(output_descr)
print("#" * len_output)
protocol=protocol,
constants=constants,
composable_states=[composable_state])
# assign sampler_state
sampler_state = states.SamplerState(
positions=pdb.positions, box_vectors=system.getDefaultPeriodicBoxVectors())
# Set up the context for mtd simulation
# at this step the CV and the system are separately passed to Metadynamics
from yank.multistate import SAMSSampler, MultiStateReporter
# TODO: You can use heavy hydrogens and 4 fs timesteps
move = mcmc.LangevinDynamicsMove(timestep=2.0 * unit.femtoseconds,
collision_rate=1.0 / unit.picosecond,
n_steps=1000,
reassign_velocities=False)
print("Done specifying integrator for simulation.")
simulation = SAMSSampler(mcmc_moves=move,
number_of_iterations=iteration,
online_analysis_interval=None,
gamma0=1.0,
flatness_threshold=0.2)
storage_path = os.path.join(work_dir, 'traj.nc')
reporter = MultiStateReporter(storage_path, checkpoint_interval=500)
# We should also add unsampled_states with the fully-interacting system
simulation.create(thermodynamic_states=thermo_states,
sampler_states=[sampler_state],
storage=reporter)
print("Done specifying simulation.")
simulation.run()
print(f"Done with {iteration} iterations.")
def _setup(self, topology, positions, box):
"""
Set up calculation.
"""
# Signal that system has now been set up
if self._setup_complete:
raise Exception(
"System has already been set up---cannot run again.")
self._setup_complete = True
# Compute thermal energy and inverse temperature
self.kT = kB * self.temperature
self.beta = 1.0 / self.kT
# Add a barostat
# TODO: Is this necessary, since ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat.
self._add_barostat()
# Create SamplerState for initial conditions
self.sampler_state = states.SamplerState(positions=positions,
box_vectors=self.box)
# Create reference thermodynamic state
self.reference_thermodynamic_state = states.ThermodynamicState(
system=self.system,
temperature=self.temperature,
pressure=self.pressure)
# Anneal ligand into binding site
if self.anneal_ligand:
self._anneal_ligand()
positions = self.sampler_state.positions
structure = pmd.openmm.load_topology(topology,
system=self.system,
xyz=positions)
# Create ThermodynamicStates for umbrella sampling along pore
self.thermodynamic_states = self._auto_create_thermodynamic_states(
structure, topology, self.reference_thermodynamic_state)
# Minimize initial thermodynamic state
# TODO: Select initial thermodynamic state based on which state has minimum energy
initial_state_index = 0
#self.sampler_state = self._minimize_sampler_state(self.thermodynamic_states[initial_state_index], self.sampler_state)
#pickle.dump(self.sampler_state, open('sampler_state.obj', 'wb'))
#pickle.dump(self.reference_thermodynamic_state, open('reference_thermodynamic_state.obj', 'wb'))
#pickle.dump(self.thermodynamic_states, open('thermodynamic_states.obj', 'wb'))
#pickle.dump(self.system, open('system.obj', 'wb'))
# Set up simulation
from yank.multistate import SAMSSampler, MultiStateReporter
# TODO: Change this to LangevinSplittingDynamicsMove
move = mcmc.LangevinDynamicsMove(timestep=self.timestep,
collision_rate=self.collision_rate,
n_steps=self.n_steps_per_iteration,
reassign_velocities=False)
self.simulation = SAMSSampler(
mcmc_moves=move,
number_of_iterations=self.n_iterations,
online_analysis_interval=None,
gamma0=self.gamma0,
flatness_threshold=self.flatness_threshold)
self.reporter = MultiStateReporter(
self.output_filename,
checkpoint_interval=self.checkpoint_interval,
analysis_particle_indices=self.analysis_particle_indices)
self.simulation.create(
thermodynamic_states=self.thermodynamic_states,
unsampled_thermodynamic_states=[
self.reference_thermodynamic_state
],
sampler_states=[self.sampler_state],
initial_thermodynamic_states=[initial_state_index],
storage=self.reporter)