def minimize(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState, max_iterations: int = 100) -> states.SamplerState: """ Minimize the given system and state, up to a maximum number of steps. This does not return a copy of the samplerstate; it is an update-in-place. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The state at which the system could be minimized sampler_state : openmmtools.states.SamplerState The starting state at which to minimize the system. max_iterations : int, optional, default 100 The maximum number of minimization steps. Default is 100. Returns ------- sampler_state : openmmtools.states.SamplerState The posititions and accompanying state following minimization """ if type(cache.global_context_cache) == cache.DummyContextCache: integrator = openmm.VerletIntegrator( 1.0) #we won't take any steps, so use a simple integrator context, integrator = cache.global_context_cache.get_context( thermodynamic_state, integrator) _logger.debug(f"using dummy context cache") else: _logger.debug(f"using global context cache") context, integrator = cache.global_context_cache.get_context( thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) openmm.LocalEnergyMinimizer.minimize(context, maxIterations=max_iterations) sampler_state.update_from_context(context)
def _update_particle_state_substate(self, particle_state, new_state_subset=False): """ update the particle state from the context, create a particle substate and update from context """ #update the particle state and the particle state subset particle_state.update_from_context( self.context, ignore_velocities=True ) #update the particle state from the context if new_state_subset: self.particle_state_subset = SamplerState( positions=particle_state.positions[list( self._subset_indices_map.keys( ))]) #create a particle state from the subset context else: self.particle_state_subset.positions = particle_state.positions[ list(self._subset_indices_map.keys() )] #update the particle subset positions appropriately self.particle_state_subset.apply_to_context( self.context_subset, ignore_velocities=True ) #apply the subset particle state to its context self.particle_state_subset.update_from_context( self.context_subset, ignore_velocities=True ) #update the subset particle state from its context to updated the potential energy
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float: """ Compute the reduced potential of the given SamplerState under the given ThermodynamicState. Arguments ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state under which to compute the reduced potential sampler_state : openmmtools.states.SamplerState The sampler state for which to compute the reduced potential Returns ------- reduced_potential : float unitless reduced potential (kT) """ if type(cache.global_context_cache) == cache.DummyContextCache: integrator = openmm.VerletIntegrator( 1.0) #we won't take any steps, so use a simple integrator context, integrator = cache.global_context_cache.get_context( thermodynamic_state, integrator) else: context, integrator = cache.global_context_cache.get_context( thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) return thermodynamic_state.reduced_potential(context)
def get_harmonic_oscillator(cls): """Create a harmonic oscillator thermodynamic state to test the trailblaze algorithm.""" from openmmtools.states import (GlobalParameterState, ThermodynamicState, CompoundThermodynamicState, SamplerState) # Create composable state that control offset of harmonic oscillator. class X0State(GlobalParameterState): testsystems_HarmonicOscillator_x0 = GlobalParameterState.GlobalParameter( cls.PAR_NAME_X0, 1.0) testsystems_HarmonicOscillator_K = GlobalParameterState.GlobalParameter( cls.PAR_NAME_K, (1.0 * unit.kilocalories_per_mole / unit.nanometer**2).value_in_unit_system(unit.md_unit_system)) # Create a harmonic oscillator thermo state. k = 1.0 * unit.kilocalories_per_mole / unit.nanometer**2 oscillator = mmtools.testsystems.HarmonicOscillator(K=k) sampler_state = SamplerState(positions=oscillator.positions) thermo_state = ThermodynamicState(oscillator.system, temperature=300 * unit.kelvin) x0_state = X0State( testsystems_HarmonicOscillator_x0=0.0, testsystems_HarmonicOscillator_K=k.value_in_unit_system( unit.md_unit_system)) compound_state = CompoundThermodynamicState( thermo_state, composable_states=[x0_state]) return compound_state, sampler_state
def __init__(self, config_: Config, old_sampler_state=None): """ :param systemLoader: :param config: """ self._times = None self.config = config_ self.logger = make_message_writer(self.config.verbose, self.__class__.__name__) with self.logger("__init__") as logger: self.explicit = self.config.systemloader.explicit self.amber = bool( self.config.systemloader.config.method == 'amber') self._trajs = np.zeros((1, 1)) self._id_number = int(self.config.systemloader.params_written) if self.config.systemloader.system is None: self.system = self.config.systemloader.get_system( self.config.parameters.createSystem) self.topology = self.config.systemloader.topology cache.global_context_cache.set_platform( self.config.parameters.platform, self.config.parameters.platform_config) positions, velocities = self.config.systemloader.get_positions( ), None else: self.system = self.config.systemloader.system self.topology = self.config.systemloader.topology positions, velocities = self.config.systemloader.get_positions( ), None # if self.config.systemloader.config.relax_ligand: # positions, velocities = self.relax_ligand(positions, velocities) thermo_state = ThermodynamicState( system=self.system, temperature=self.config.parameters. integrator_params['temperature'], pressure=1.0 * unit.atmosphere if self.config.systemloader.explicit else None) sampler_state = SamplerState( positions=positions, velocities=velocities, box_vectors=self.config.systemloader.boxvec) self.sampler = MCMCSampler(thermo_state, sampler_state, move=mmWrapperUtils.prepare_mcmc( self.topology, self.config)) self.sampler.minimize( max_iterations=self.config.parameters.minMaxIters) ctx = cache.global_context_cache.get_context( self.sampler.thermodynamic_state)[0] ctx.setVelocitiesToTemperature( self.config.parameters.integrator_params['temperature']) self.sampler.sampler_state.velocities = ctx.getState( getVelocities=True).getVelocities()
def test_metropolized_moves(): """Test Displacement and Rotation moves.""" testsystem = testsystems.AlanineDipeptideVacuum() original_sampler_state = SamplerState(testsystem.positions) thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin) all_metropolized_moves = MetropolizedMove.__subclasses__() for move_class in all_metropolized_moves: move = move_class(atom_subset=range(thermodynamic_state.n_particles)) sampler_state = copy.deepcopy(original_sampler_state) # Start applying the move and remove one at each iteration tyring # to generate both an accepted and rejected move. old_n_accepted, old_n_proposed = 0, 0 while len(move.atom_subset) > 0: initial_positions = copy.deepcopy(sampler_state.positions) move.apply(thermodynamic_state, sampler_state) final_positions = copy.deepcopy(sampler_state.positions) # If the move was accepted the positions should be different. if move.n_accepted > old_n_accepted: assert not np.allclose(initial_positions, final_positions) # If we have generated a rejection and an acceptance, test next move. if move.n_accepted > 0 and move.n_accepted != move.n_proposed: break # Try with a smaller subset. move.atom_subset = move.atom_subset[:-1] old_n_accepted, old_n_proposed = move.n_accepted, move.n_proposed # Check that we were able to generate both an accepted and a rejected move. assert len(move.atom_subset) != 0, ('Could not generate an accepted and rejected ' 'move for class {}'.format(move_class.__name__))
def _geometry_forward(self, topology_proposal, old_sampler_state): """ Run geometry engine to propose new positions and compute logP Parameters ---------- topology_proposal : TopologyProposal Contains old/new Topology and System objects and atom mappings. old_sampler_state : openmmtools.states.SamplerState Configurational properties of the old system atoms. Returns ------- new_sampler_state : openmmtools.states.SamplerState Configurational properties of new atoms proposed by geometry engine calculation. geometry_logp_propose : float The log probability of the forward-only proposal """ if self.verbose: print("Geometry engine proposal...") # Generate coordinates for new atoms and compute probability ratio of old and new probabilities. initial_time = time.time() new_positions, geometry_logp_propose = self.geometry_engine.propose(topology_proposal, old_sampler_state.positions, self.sampler.thermodynamic_state.beta) if self.verbose: print('proposal took %.3f s' % (time.time() - initial_time)) if self.geometry_pdbfile is not None: print("Writing proposed geometry...") from simtk.openmm.app import PDBFile PDBFile.writeFile(topology_proposal.new_topology, new_positions, file=self.geometry_pdbfile) self.geometry_pdbfile.flush() new_sampler_state = SamplerState(new_positions, box_vectors=old_sampler_state.box_vectors) return new_sampler_state, geometry_logp_propose
def relax_ligand(self, positions, velocities): with self.logger("relax_ligand") as logger: cache.global_context_cache.empty() system = copy.deepcopy( self.config.systemloader._unconstrained_system) thermo_state_ = ThermodynamicState( system=system, temperature=self.config.parameters. integrator_params['temperature'], pressure=1.0 * unit.atmosphere if self.config.systemloader.explicit else None) pforce = mmWrapperUtils.get_protein_restraint_force(self.topology, positions, self.explicit, K=5.0) lforce = mmWrapperUtils.get_ligand_restraint_force( self.topology, positions, self.explicit, K=thermo_state_.kT / 3.0**2) del thermo_state_ system.addForce(pforce) system.addForce(lforce) thermo_state = ThermodynamicState( system=system, temperature=self.config.parameters. integrator_params['temperature'], pressure=1.0 * unit.atmosphere if self.config.systemloader.explicit else None) sampler_state = SamplerState( positions=positions, velocities=velocities, box_vectors=self.config.systemloader.boxvec) sampler = MCMCSampler( thermo_state, sampler_state, move=LangevinSplittingDynamicsMove( timestep=0.1 * unit.femtosecond, n_steps=1000, collision_rate=self.config.parameters. integrator_params['collision_rate'], reassign_velocities=True, n_restart_attempts=6, constraint_tolerance=self.config.parameters. integrator_setConstraintTolerance)) sampler.minimize(max_iterations=self.config.parameters.minMaxIters) logger.log("Build unconstrained system. Relaxing for 1 ps") sampler.run(10) positions, velocities = sampler.sampler_state.positions, sampler.sampler_state.velocities cache.global_context_cache.empty() return positions, velocities
def test_move_restart(): """Test optional restart move if NaN is detected.""" n_restart_attempts = 5 # We define a Move that counts the times it is attempted. class MyMove(BaseIntegratorMove): def __init__(self, **kwargs): super(MyMove, self).__init__(n_steps=1, n_restart_attempts=n_restart_attempts, **kwargs) self.attempted_count = 0 def _get_integrator(self, thermodynamic_state): return integrators.GHMCIntegrator(temperature=300 * unit.kelvin) def _before_integration(self, context, thermodynamic_state): self.attempted_count += 1 # Create a system with an extra NaN particle. testsystem = testsystems.AlanineDipeptideVacuum() system = testsystem.system for force in system.getForces(): if isinstance(force, openmm.NonbondedForce): break # Add a non-interacting particle to the system at NaN position. system.addParticle(39.9 * unit.amu) force.addParticle(0.0, 1.0, 0.0) particle_position = np.array([np.nan, 0.2, 0.2]) positions = unit.Quantity(np.vstack( (testsystem.positions, particle_position)), unit=testsystem.positions.unit) # Create and run move. An IntegratoMoveError is raised. sampler_state = SamplerState(positions) thermodynamic_state = ThermodynamicState(system, 300 * unit.kelvin) # We use a local context cache with Reference platform since on the # CPU platform CustomIntegrators raises an error with NaN particles. reference_platform = openmm.Platform.getPlatformByName('Reference') move = MyMove(context_cache=cache.ContextCache( platform=reference_platform)) with nose.tools.assert_raises(IntegratorMoveError) as cm: move.apply(thermodynamic_state, sampler_state) # We have counted the correct number of restart attempts. assert move.attempted_count == n_restart_attempts + 1 # Test serialization of the error. with utils.temporary_directory() as tmp_dir: prefix = os.path.join(tmp_dir, 'prefix') cm.exception.serialize_error(prefix) assert os.path.exists(prefix + '-move.json') assert os.path.exists(prefix + '-system.xml') assert os.path.exists(prefix + '-integrator.xml') assert os.path.exists(prefix + '-state.xml')
def test_minimizer_all_testsystems(): # testsystem_classes = testsystems.TestSystem.__subclasses__() testsystem_classes = [testsystems.AlanineDipeptideVacuum] for testsystem_class in testsystem_classes: class_name = testsystem_class.__name__ logging.info("Testing minimization with testsystem %s" % class_name) testsystem = testsystem_class() sampler_state = SamplerState(testsystem.positions) thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin) # Create sampler for minimization. sampler = MCMCSampler(thermodynamic_state, sampler_state, move=None) sampler.minimize(max_iterations=0) # Check if NaN. err_msg = 'Minimization of system {} yielded NaN'.format(class_name) assert not sampler_state.has_nan(), err_msg
def runEthyleneTest(dir, N): filename = dir.join('ethylene-test_%s' % N) print('Running %s...' % filename) # Set Simulation parameters temperature = 200 * unit.kelvin collision_rate = 1 / unit.picoseconds timestep = 1.0 * unit.femtoseconds n_steps = 20 nIter = 100 reportInterval = 5 alchemical_atoms = [2, 3, 4, 5, 6, 7] platform = openmm.Platform.getPlatformByName('CPU') context_cache = cache.ContextCache(platform) # Load a Parmed Structure for the Topology and create our openmm.Simulation structure_pdb = utils.get_data_filename( 'blues', 'tests/data/ethylene_structure.pdb') structure = parmed.load_file(structure_pdb) nc_reporter = NetCDF4Storage(filename + '_MD.nc', reportInterval) # Iniitialize our Move set rot_move = RandomLigandRotationMove(timestep=timestep, n_steps=n_steps, atom_subset=alchemical_atoms, context_cache=context_cache, reporters=[nc_reporter]) langevin_move = ReportLangevinDynamicsMove(timestep=timestep, collision_rate=collision_rate, n_steps=n_steps, reassign_velocities=True, context_cache=context_cache) # Load our OpenMM System and create Integrator system_xml = utils.get_data_filename('blues', 'tests/data/ethylene_system.xml') with open(system_xml, 'r') as infile: xml = infile.read() system = openmm.XmlSerializer.deserialize(xml) thermodynamic_state = ThermodynamicState(system=system, temperature=temperature) sampler_state = SamplerState( positions=structure.positions.in_units_of(unit.nanometers)) sampler = BLUESSampler(thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, ncmc_move=rot_move, dynamics_move=langevin_move, topology=structure.topology) sampler.run(nIter) return filename
def test_langevin_splitting_move(): """Test that the langevin splitting mcmc move works with different splittings""" splittings = ["V R O R V", "V R R R O R R R V", "O { V R V } O"] testsystem = testsystems.AlanineDipeptideVacuum() sampler_state = SamplerState(testsystem.positions) thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin) for splitting in splittings: move = LangevinSplittingDynamicsMove(splitting=splitting) # Create MCMC sampler sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move) sampler.run(1)
def create_langevin_integrator(htf, constraint_tol): """ create lambda alchemical states, thermodynamic states, sampler states, integrator, and return context, thermostate, sampler_state, integrator """ fast_lambda_alchemical_state = RelativeAlchemicalState.from_system( htf.hybrid_system) fast_lambda_alchemical_state.set_alchemical_parameters( 0.0, LambdaProtocol(functions='default')) fast_thermodynamic_state = CompoundThermodynamicState( ThermodynamicState(htf.hybrid_system, temperature=temperature), composable_states=[fast_lambda_alchemical_state]) fast_sampler_state = SamplerState( positions=htf._hybrid_positions, box_vectors=htf.hybrid_system.getDefaultPeriodicBoxVectors()) integrator_1 = integrators.LangevinIntegrator( temperature=temperature, timestep=4.0 * unit.femtoseconds, splitting='V R O R V', measure_shadow_work=False, measure_heat=False, constraint_tolerance=constraint_tol, collision_rate=5.0 / unit.picoseconds) # mcmc_moves=mcmc.LangevinSplittingDynamicsMove(timestep = 4.0 * unit.femtoseconds, # collision_rate=5.0 / unit.picosecond, # n_steps=1, # reassign_velocities=False, # n_restart_attempts=20, # splitting="V R R R O R R R V", # constraint_tolerance=constraint_tol) #print(integrator_1.getConstraintTolerance()) fast_context, fast_integrator = cache.global_context_cache.get_context( fast_thermodynamic_state, integrator_1) fast_sampler_state.apply_to_context(fast_context) return fast_context, fast_thermodynamic_state, fast_sampler_state, fast_integrator
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float: """ Compute the reduced potential of the given SamplerState under the given ThermodynamicState. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state under which to compute the reduced potential sampler_state : openmmtools.states.SamplerState The sampler state for which to compute the reduced potential Returns ------- reduced_potential : float unitless reduced potential (kT) """ context, integrator = cache.global_context_cache.get_context( thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) return thermodynamic_state.reduced_potential(context)
def test_context_cache(): """Test configuration of the context cache.""" testsystem = testsystems.AlanineDipeptideImplicit() sampler_state = SamplerState(testsystem.positions) thermodynamic_state = ThermodynamicState(testsystem.system, 300 * unit.kelvin) # By default the global context cache is used. cache.global_context_cache.empty() # Clear cache from previous tests. move = SequenceMove([LangevinDynamicsMove(n_steps=5), GHMCMove(n_steps=5)]) move.apply(thermodynamic_state, sampler_state) assert len(cache.global_context_cache) == 2 # Configuring the global cache works correctly. cache.global_context_cache = cache.ContextCache(time_to_live=1) move.apply(thermodynamic_state, sampler_state) assert len(cache.global_context_cache) == 1 # The ContextCache creates only one context with compatible moves. cache.global_context_cache = cache.ContextCache(capacity=10, time_to_live=None) move = SequenceMove([ LangevinDynamicsMove(n_steps=1), LangevinDynamicsMove(n_steps=1), LangevinDynamicsMove(n_steps=1), LangevinDynamicsMove(n_steps=1) ]) move.apply(thermodynamic_state, sampler_state) assert len(cache.global_context_cache) == 1 # We can configure a local context cache instead of global. local_cache = cache.ContextCache() move = SequenceMove([LangevinDynamicsMove(n_steps=5), GHMCMove(n_steps=5)], context_cache=local_cache) for m in move: assert m.context_cache == local_cache # Running with the local cache doesn't affect the global one. cache.global_context_cache = cache.ContextCache() # empty global move.apply(thermodynamic_state, sampler_state) assert len(cache.global_context_cache) == 0 assert len(local_cache) == 2 # DummyContextCache works for all platforms. platforms = utils.get_available_platforms() dummy_cache = cache.DummyContextCache() for platform in platforms: dummy_cache.platform = platform move = LangevinDynamicsMove(n_steps=5, context_cache=dummy_cache) move.apply(thermodynamic_state, sampler_state) assert len(cache.global_context_cache) == 0
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float: """ Compute the reduced potential of the given SamplerState under the given ThermodynamicState. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state under which to compute the reduced potential sampler_state : openmmtools.states.SamplerState The sampler state for which to compute the reduced potential Returns ------- reduced_potential : float unitless reduced potential (kT) """ if type(cache.global_context_cache) == cache.DummyContextCache: integrator = openmm.VerletIntegrator(1.0) #we won't take any steps, so use a simple integrator context, integrator = cache.global_context_cache.get_context(thermodynamic_state, integrator) else: context, integrator = cache.global_context_cache.get_context(thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) return thermodynamic_state.reduced_potential(context)
def from_amber(cls, prmtop, inpcrd, temperature=50 * u.kelvin): prmtop = app.AmberPrmtopFile(prmtop) inpcrd = app.AmberInpcrdFile(inpcrd) system = prmtop.createSystem(nonbondedMethod=app.PME, constraints=app.HBonds, nonbondedCutoff=10 * u.angstroms, switchDistance=8 * u.angstroms) thermodynamic_state = ThermodynamicState(system, temperature=temperature) sampler_state = SamplerState( positions=inpcrd.getPositions(asNumpy=True), box_vectors=inpcrd.boxVectors) return Esmacs(thermodynamic_state, sampler_state, prmtop.topology)
def test_barostat_move_frequency(): """MonteCarloBarostatMove restore barostat's frequency afterwards.""" # Get periodic test case. for test_case in analytical_testsystems: testsystem = test_case[1] if testsystem.system.usesPeriodicBoundaryConditions(): break assert testsystem.system.usesPeriodicBoundaryConditions(), "Can't find periodic test case!" sampler_state = SamplerState(testsystem.positions) thermodynamic_state = ThermodynamicState(testsystem.system, 298*unit.kelvin, 1*unit.atmosphere) move = MonteCarloBarostatMove(n_attempts=5) # Test-precondition: the frequency must be different than 1 or it # will never change during the application of the MCMC move. old_frequency = thermodynamic_state.barostat.getFrequency() assert old_frequency != 1 move.apply(thermodynamic_state, sampler_state) assert thermodynamic_state.barostat.getFrequency() == old_frequency
def get_states(): # Set Simulation parameters temperature = 200 * unit.kelvin collision_rate = 1 / unit.picoseconds timestep = 1.0 * unit.femtoseconds n_steps = 20 nIter = 100 alchemical_atoms = [2, 3, 4, 5, 6, 7] platform = openmm.Platform.getPlatformByName('CPU') context_cache = cache.ContextCache(platform) # Load a Parmed Structure for the Topology and create our openmm.Simulation structure_pdb = utils.get_data_filename( 'blues', 'tests/data/ethylene_structure.pdb') structure = parmed.load_file(structure_pdb) # Load our OpenMM System and create Integrator system_xml = utils.get_data_filename('blues', 'tests/data/ethylene_system.xml') with open(system_xml, 'r') as infile: xml = infile.read() system = openmm.XmlSerializer.deserialize(xml) thermodynamic_state = ThermodynamicState(system=system, temperature=temperature) sampler_state = SamplerState( positions=structure.positions.in_units_of(unit.nanometers)) alch_system = generateAlchSystem(thermodynamic_state.get_system(), alchemical_atoms) alch_state = alchemy.AlchemicalState.from_system(alch_system) alch_thermodynamic_state = ThermodynamicState( alch_system, thermodynamic_state.temperature) alch_thermodynamic_state = CompoundThermodynamicState( alch_thermodynamic_state, composable_states=[alch_state]) return structure, thermodynamic_state, alch_thermodynamic_state
def from_testsystem( # pylint: disable=too-many-arguments self, test, reference_state, thermodynamic_states, pressure=None, storage=None, target_state=0, metadata=None, **kwargs): """Initialize sampler from TestSystem object.""" if not isinstance(thermodynamic_states, Sequence): # a scalar thermodynamic_states = [thermodynamic_states] # check if temp or thermodynamic states or something else if not isinstance(thermodynamic_states[0], ThermodynamicState): # as temperatures temperatures = thermodynamic_states if not isinstance(temperatures[0], unit.Quantity): # no units temperatures = [t * unit.kelvin for t in temperatures] thermodynamic_states = [ ThermodynamicState( system=test.system, temperature=t, pressure=pressure, ) for t in temperatures ] sampler_states = SamplerState(positions=test.positions, box_vectors=test.default_box_vectors) self.create(reference_state, thermodynamic_states, sampler_states, test.topology, target_state=target_state, storage=storage, metadata=metadata, **kwargs)
def subtest_mcmc_expectation(testsystem, move): if debug: print(testsystem.__class__.__name__) print(str(move)) # Retrieve system and positions. [system, positions] = [testsystem.system, testsystem.positions] # Test settings. temperature = 298.0 * unit.kelvin niterations = 500 # number of production iterations if system.usesPeriodicBoundaryConditions(): pressure = 1.0 * unit.atmosphere else: pressure = None # Compute properties. kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA kT = kB * temperature ndof = 3 * system.getNumParticles() - system.getNumConstraints() # Create sampler and thermodynamic state. sampler_state = SamplerState(positions=positions) thermodynamic_state = ThermodynamicState(system=system, temperature=temperature, pressure=pressure) # Create MCMC sampler sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move) # Accumulate statistics. x_n = np.zeros( [niterations], np.float64 ) # x_n[i] is the x position of atom 1 after iteration i, in angstroms potential_n = np.zeros( [niterations], np.float64 ) # potential_n[i] is the potential energy after iteration i, in kT kinetic_n = np.zeros( [niterations], np.float64 ) # kinetic_n[i] is the kinetic energy after iteration i, in kT temperature_n = np.zeros( [niterations], np.float64 ) # temperature_n[i] is the instantaneous kinetic temperature from iteration i, in K volume_n = np.zeros( [niterations], np.float64) # volume_n[i] is the volume from iteration i, in K for iteration in range(niterations): # Update sampler state. sampler.run(1) # Get statistics. potential_energy = sampler.sampler_state.potential_energy kinetic_energy = sampler.sampler_state.kinetic_energy instantaneous_temperature = kinetic_energy * 2.0 / ndof / kB volume = sampler.sampler_state.volume # Accumulate statistics. x_n[iteration] = sampler_state.positions[0, 0] / unit.angstroms potential_n[iteration] = potential_energy / kT kinetic_n[iteration] = kinetic_energy / kT temperature_n[iteration] = instantaneous_temperature / unit.kelvin volume_n[iteration] = volume / (unit.nanometers**3) # Compute expected statistics. if (hasattr(testsystem, 'get_potential_expectation') and testsystem.get_potential_standard_deviation(thermodynamic_state) / kT.unit != 0.0): assert potential_n.std( ) != 0.0, 'Test {} shows no potential fluctuations'.format( testsystem.__class__.__name__) potential_expectation = testsystem.get_potential_expectation( thermodynamic_state) / kT [t0, g, Neff_max] = timeseries.detectEquilibration(potential_n) potential_mean = potential_n[t0:].mean() dpotential_mean = potential_n[t0:].std() / np.sqrt(Neff_max) potential_error = potential_mean - potential_expectation nsigma = abs(potential_error) / dpotential_mean err_msg = ( 'Potential energy expectation\n' 'observed {:10.5f} +- {:10.5f}kT | expected {:10.5f} | ' 'error {:10.5f} +- {:10.5f} ({:.1f} sigma) | t0 {:5d} | g {:5.1f} | Neff {:8.1f}\n' '----------------------------------------------------------------------------' ).format(potential_mean, dpotential_mean, potential_expectation, potential_error, dpotential_mean, nsigma, t0, g, Neff_max) assert nsigma <= NSIGMA_CUTOFF, err_msg.format() if debug: print(err_msg) elif debug: print('Skipping potential expectation test.') if (hasattr(testsystem, 'get_volume_expectation') and testsystem.get_volume_standard_deviation(thermodynamic_state) / (unit.nanometers**3) != 0.0): assert volume_n.std( ) != 0.0, 'Test {} shows no volume fluctuations'.format( testsystem.__class__.__name__) volume_expectation = testsystem.get_volume_expectation( thermodynamic_state) / (unit.nanometers**3) [t0, g, Neff_max] = timeseries.detectEquilibration(volume_n) volume_mean = volume_n[t0:].mean() dvolume_mean = volume_n[t0:].std() / np.sqrt(Neff_max) volume_error = volume_mean - volume_expectation nsigma = abs(volume_error) / dvolume_mean err_msg = ( 'Volume expectation\n' 'observed {:10.5f} +- {:10.5f}kT | expected {:10.5f} | ' 'error {:10.5f} +- {:10.5f} ({:.1f} sigma) | t0 {:5d} | g {:5.1f} | Neff {:8.1f}\n' '----------------------------------------------------------------------------' ).format(volume_mean, dvolume_mean, volume_expectation, volume_error, dvolume_mean, nsigma, t0, g, Neff_max) assert nsigma <= NSIGMA_CUTOFF, err_msg.format() if debug: print(err_msg) elif debug: print('Skipping volume expectation test.')
def create( # pylint: disable=too-many-arguments self, reference_state, thermodynamic_states, sampler_states, top, target_state=0, initial_thermodynamic_states=None, unsampled_thermodynamic_states=None, storage=None, metadata=None, stride=1): """Create new multistate sampler simulation. Override MultistateSampler create. Parameters ---------- thermodynamic_states : list of 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 : states.SamplerState or list One or more sets of initial sampler states. The number of replicas is taken to be the number of sampler states provided. If the sampler states do not have box_vectors attached and the system is periodic, an exception will be thrown. top : Topology or Topography object, optional target_state : int, optional The indef of the reference thermodynamic state. Defaults to 0. initial_thermodynamic_states : None or list or array-like of int of length len(sampler_states), optional, Default: None. Initial thermodynamic_state index for each sampler_state. If no initial distribution is chosen, ``sampler_states`` are distributed between the ``thermodynamic_states`` following these rules: * If ``len(thermodynamic_states) == len(sampler_states)``: 1-to-1 distribution * If ``len(thermodynamic_states) > len(sampler_states)``: First and last state distributed first Remaining ``sampler_states`` spaced evenly by index until ``sampler_states`` are depleted. If there is only one ``sampler_state``, then the only first ``thermodynamic_state`` will be chosen * If ``len(thermodynamic_states) < len(sampler_states)``: each ``thermodynamic_state`` receives an equal number of ``sampler_states`` until there are insufficient number of ``sampler_states`` remaining to give each ``thermodynamic_state`` an equal number. Then the rules from the previous point are followed. unsampled_thermodynamic_states : list of states.ThermodynamicState, optional, default=None 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). 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. metadata : dict, optional, default=None Simulation metadata to be stored in the file. """ if isinstance(thermodynamic_states, ThermodynamicState): # a single state thermodynamic_states = [thermodynamic_states] if not isinstance(sampler_states, (SamplerState, Sequence)): # TODO: check this # as positions sampler_states = SamplerState(sampler_states) # Do not modify passed ref thermodynamic state. self._reference_thermodynamic_state = copy.deepcopy(reference_state) thermodynamic_state = copy.deepcopy( self._reference_thermodynamic_state) self._reference_system = thermodynamic_state.system if storage is None: storage = tempfile.NamedTemporaryFile(delete=False).name + '.nc' self.reporter = self._define_reporter(storage, stride) # set self.topography if isinstance(top, Topography): self.topography = top else: self.topography = Topography(top) self.ref_state = target_state if metadata is None: metadata = {} sampler_full_name = mmtools.utils.typename(self.__class__) metadata['title'] = 'Created using %s on %s' % ( sampler_full_name, time.asctime(time.localtime())) metadata['sampler_full_name'] = sampler_full_name metadata['topography'] = mmtools.utils.serialize(self.topography) metadata['reference_state'] = mmtools.utils.serialize( thermodynamic_state) # the ref thermodynamic state super().create( thermodynamic_states=thermodynamic_states, sampler_states=sampler_states, storage=self.reporter, initial_thermodynamic_states=initial_thermodynamic_states, unsampled_thermodynamic_states=unsampled_thermodynamic_states, metadata=metadata)
def compare_energies(REST_system, other_system, positions, rest_atoms, T_min, T): # Create thermodynamic state lambda_zero_alchemical_state = RESTState.from_system(REST_system) thermostate = ThermodynamicState(REST_system, temperature=T_min) compound_thermodynamic_state = CompoundThermodynamicState(thermostate, composable_states=[lambda_zero_alchemical_state]) # Set alchemical parameters beta_0 = 1 / (kB * T_min) beta_m = 1 / (kB * T) compound_thermodynamic_state.set_alchemical_parameters(beta_0, beta_m) # Minimize and save energy integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond) context = compound_thermodynamic_state.create_context(integrator) context.setPositions(positions) sampler_state = SamplerState.from_context(context) REST_energy = compound_thermodynamic_state.reduced_potential(sampler_state) # Compute energy for non-RESTified system # Determine regions and scaling factors solute = rest_atoms solvent = [i for i in range(other_system.getNumParticles()) if i not in solute] solute_scaling = beta_m / beta_0 inter_scaling = np.sqrt(beta_m / beta_0) # Scale the terms in the bond force appropriately bond_force = other_system.getForce(0) for bond in range(bond_force.getNumBonds()): p1, p2, length, k = bond_force.getBondParameters(bond) if p1 in solute and p2 in solute: bond_force.setBondParameters(bond, p1, p2, length, k * solute_scaling) elif (p1 in solute and p2 in solvent) or (p1 in solvent and p2 in solute): bond_force.setBondParameters(bond, p1, p2, length, k * inter_scaling) # Scale the terms in the angle force appropriately angle_force = other_system.getForce(1) for angle_index in range(angle_force.getNumAngles()): p1, p2, p3, angle, k = angle_force.getAngleParameters(angle_index) if p1 in solute and p2 in solute and p3 in solute: angle_force.setAngleParameters(angle_index, p1, p2, p3, angle, k * solute_scaling) elif set([p1, p2, p3]).intersection(set(solute)) != set() and set([p1, p2, p3]).intersection( set(solvent)) != set(): angle_force.setAngleParameters(angle_index, p1, p2, p3, angle, k * inter_scaling) # Scale the terms in the torsion force appropriately torsion_force = other_system.getForce(2) for torsion_index in range(torsion_force.getNumTorsions()): p1, p2, p3, p4, periodicity, phase, k = torsion_force.getTorsionParameters(torsion_index) if p1 in solute and p2 in solute and p3 in solute and p4 in solute: torsion_force.setTorsionParameters(torsion_index, p1, p2, p3, p4, periodicity, phase, k * solute_scaling) elif set([p1, p2, p3, p4]).intersection(set(solute)) != set() and set([p1, p2, p3, p4]).intersection( set(solvent)) != set(): torsion_force.setTorsionParameters(torsion_index, p1, p2, p3, p4, periodicity, phase, k * inter_scaling) # Scale the exceptions in the nonbonded force appropriately nb_force = other_system.getForce(3) for nb_index in range(nb_force.getNumExceptions()): p1, p2, chargeProd, sigma, epsilon = nb_force.getExceptionParameters(nb_index) if p1 in solute and p2 in solute: nb_force.setExceptionParameters(nb_index, p1, p2, solute_scaling * chargeProd, sigma, solute_scaling * epsilon) elif (p1 in solute and p2 in solvent) or (p1 in solvent and p2 in solute): nb_force.setExceptionParameters(nb_index, p1, p2, inter_scaling * chargeProd, sigma, inter_scaling * epsilon) # Scale nonbonded interactions for solute-solute region by adding exceptions for all pairs of atoms exception_pairs = [tuple(sorted([nb_force.getExceptionParameters(nb_index)[0], nb_force.getExceptionParameters(nb_index)[1]])) for nb_index in range(nb_force.getNumExceptions())] solute_pairs = set([tuple(sorted(pair)) for pair in list(itertools.product(solute, solute))]) for pair in list(solute_pairs): p1 = pair[0] p2 = pair[1] p1_charge, p1_sigma, p1_epsilon = nb_force.getParticleParameters(p1) p2_charge, p2_sigma, p2_epsilon = nb_force.getParticleParameters(p2) if p1 != p2: if pair not in exception_pairs: nb_force.addException(p1, p2, p1_charge * p2_charge * solute_scaling, 0.5 * (p1_sigma + p2_sigma), np.sqrt(p1_epsilon * p2_epsilon) * solute_scaling) # Scale nonbonded interactions for inter region by adding exceptions for all pairs of atoms for pair in list(itertools.product(solute, solvent)): p1 = pair[0] p2 = int(pair[1]) # otherwise, will be a numpy int p1_charge, p1_sigma, p1_epsilon = nb_force.getParticleParameters(p1) p2_charge, p2_sigma, p2_epsilon = nb_force.getParticleParameters(p2) if tuple(sorted(pair)) not in exception_pairs: nb_force.addException(p1, p2, p1_charge * p2_charge * inter_scaling, 0.5 * (p1_sigma + p2_sigma), np.sqrt(p1_epsilon * p2_epsilon) * inter_scaling) # Get energy thermostate = ThermodynamicState(other_system, temperature=T_min) integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond) context = thermostate.create_context(integrator) context.setPositions(positions) sampler_state = SamplerState.from_context(context) nonREST_energy = thermostate.reduced_potential(sampler_state) assert REST_energy - nonREST_energy < 1, f"The energy of the REST system ({REST_energy}) does not match " \ f"that of the non-REST system with terms manually scaled according to REST2({nonREST_energy})."
def HybridTopologyFactory_energies( current_mol='toluene', proposed_mol='1,2-bis(trifluoromethyl) benzene'): """ Test whether the difference in the nonalchemical zero and alchemical zero states is the forward valence energy. Also test for the one states. """ from perses.tests.utils import generate_solvated_hybrid_test_topology, generate_endpoint_thermodynamic_states import openmmtools.cache as cache #Just test the solvated system top_proposal, old_positions, _ = generate_solvated_hybrid_test_topology( current_mol_name=current_mol, proposed_mol_name=proposed_mol) #remove the dispersion correction top_proposal._old_system.getForce(3).setUseDispersionCorrection(False) top_proposal._new_system.getForce(3).setUseDispersionCorrection(False) # run geometry engine to generate old and new positions _geometry_engine = FFAllAngleGeometryEngine(metadata=None, use_sterics=False, n_bond_divisions=100, n_angle_divisions=180, n_torsion_divisions=360, verbose=True, storage=None, bond_softening_constant=1.0, angle_softening_constant=1.0, neglect_angles=False) _new_positions, _lp = _geometry_engine.propose(top_proposal, old_positions, beta) _lp_rev = _geometry_engine.logp_reverse(top_proposal, _new_positions, old_positions, beta) # make the hybrid system, reset the CustomNonbondedForce cutoff HTF = HybridTopologyFactory(top_proposal, old_positions, _new_positions) hybrid_system = HTF.hybrid_system nonalch_zero, nonalch_one, alch_zero, alch_one = generate_endpoint_thermodynamic_states( hybrid_system, top_proposal) # compute reduced energies #for the nonalchemical systems... attrib_list = [(nonalch_zero, old_positions, top_proposal._old_system.getDefaultPeriodicBoxVectors()), (alch_zero, HTF._hybrid_positions, hybrid_system.getDefaultPeriodicBoxVectors()), (alch_one, HTF._hybrid_positions, hybrid_system.getDefaultPeriodicBoxVectors()), (nonalch_one, _new_positions, top_proposal._new_system.getDefaultPeriodicBoxVectors())] rp_list = [] for (state, pos, box_vectors) in attrib_list: context, integrator = cache.global_context_cache.get_context(state) samplerstate = SamplerState(positions=pos, box_vectors=box_vectors) samplerstate.apply_to_context(context) rp = state.reduced_potential(context) rp_list.append(rp) #valence energy definitions forward_added_valence_energy = _geometry_engine.forward_final_context_reduced_potential - _geometry_engine.forward_atoms_with_positions_reduced_potential reverse_subtracted_valence_energy = _geometry_engine.reverse_final_context_reduced_potential - _geometry_engine.reverse_atoms_with_positions_reduced_potential nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[ 0], rp_list[1], rp_list[2], rp_list[3] # print(f"Difference between zeros: {nonalch_zero_rp - alch_zero_rp}; forward added: {forward_added_valence_energy}") # print(f"Difference between ones: {nonalch_zero_rp - alch_zero_rp}; forward added: {forward_added_valence_energy}") assert abs( nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy ) < ENERGY_THRESHOLD, f"The zero state alchemical and nonalchemical energy absolute difference {abs(nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy)} is greater than the threshold of {ENERGY_THRESHOLD}." assert abs( nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy ) < ENERGY_THRESHOLD, f"The one state alchemical and nonalchemical energy absolute difference {abs(nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy)} is greater than the threshold of {ENERGY_THRESHOLD}." print( f"Abs difference in zero alchemical vs nonalchemical systems: {abs(nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy)}" ) print( f"Abs difference in one alchemical vs nonalchemical systems: {abs(nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy)}" )
def run_hybrid_endpoint_overlap(topology_proposal, current_positions, new_positions): """ Test that the variance of the perturbation from lambda={0,1} to the corresponding nonalchemical endpoint is not too large. Parameters ---------- topology_proposal : perses.rjmc.TopologyProposal TopologyProposal object describing the transformation current_positions : np.array, unit-bearing Positions of the initial system new_positions : np.array, unit-bearing Positions of the new system Returns ------- hybrid_endpoint_results : list list of [df, ddf, N_eff] for 1 and 0 """ #create the hybrid system: #hybrid_factory = HybridTopologyFactory(topology_proposal, current_positions, new_positions, use_dispersion_correction=True) hybrid_factory = HybridTopologyFactory( topology_proposal, current_positions, new_positions, use_dispersion_correction=False) # DEBUG #get the relevant thermodynamic states: nonalchemical_zero_thermodynamic_state, nonalchemical_one_thermodynamic_state, lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state = utils.generate_endpoint_thermodynamic_states( hybrid_factory.hybrid_system, topology_proposal) nonalchemical_thermodynamic_states = [ nonalchemical_zero_thermodynamic_state, nonalchemical_one_thermodynamic_state ] alchemical_thermodynamic_states = [ lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state ] #create an MCMCMove, BAOAB with default parameters (but don't restart if we encounter a NaN) mc_move = mcmc.LangevinDynamicsMove(n_restart_attempts=0, n_steps=100) initial_sampler_state = SamplerState( hybrid_factory.hybrid_positions, box_vectors=hybrid_factory.hybrid_system.getDefaultPeriodicBoxVectors( )) hybrid_endpoint_results = [] all_results = [] for lambda_state in (0, 1): result, non, hybrid = run_endpoint_perturbation( alchemical_thermodynamic_states[lambda_state], nonalchemical_thermodynamic_states[lambda_state], initial_sampler_state, mc_move, 100, hybrid_factory, lambda_index=lambda_state) all_results.append(non) all_results.append(hybrid) print('lambda {} : {}'.format(lambda_state, result)) hybrid_endpoint_results.append(result) calculate_cross_variance(all_results) return hybrid_endpoint_results
def run_endpoint_perturbation(lambda_thermodynamic_state, nonalchemical_thermodynamic_state, initial_hybrid_sampler_state, mc_move, n_iterations, factory, lambda_index=0, print_work=False, write_system=False, write_state=False, write_trajectories=False): """ Parameters ---------- lambda_thermodynamic_state : ThermodynamicState The thermodynamic state corresponding to the hybrid system at a lambda endpoint nonalchemical_thermodynamic_state : ThermodynamicState The nonalchemical thermodynamic state for the relevant endpoint initial_hybrid_sampler_state : SamplerState Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state mc_move : MCMCMove The MCMove that will be used for sampling at the lambda endpoint n_iterations : int The number of iterations factory : HybridTopologyFactory The hybrid topology factory lambda_index : int, optional, default=0 The index, 0 or 1, at which to retrieve nonalchemical positions print_work : bool, optional, default=False If True, will print work values write_system : bool, optional, default=False If True, will write alchemical and nonalchemical System XML files write_state : bool, optional, default=False If True, write alchemical (hybrid) State XML files each iteration write_trajectories : bool, optional, default=False If True, will write trajectories Returns ------- df : float Free energy difference between alchemical and nonalchemical systems, estimated with EXP ddf : float Standard deviation of estimate, corrected for correlation, from EXP estimator. """ import mdtraj as md #run an initial minimization: mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state, initial_hybrid_sampler_state, mc_move) mcmc_sampler.minimize(max_iterations=20) new_sampler_state = mcmc_sampler.sampler_state if write_system: with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile: outfile.write( openmm.XmlSerializer.serialize( lambda_thermodynamic_state.system)) with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile: outfile.write( openmm.XmlSerializer.serialize( nonalchemical_thermodynamic_state.system)) #initialize work array w = np.zeros([n_iterations]) non_potential = np.zeros([n_iterations]) hybrid_potential = np.zeros([n_iterations]) #run n_iterations of the endpoint perturbation: hybrid_trajectory = unit.Quantity( np.zeros([ n_iterations, lambda_thermodynamic_state.system.getNumParticles(), 3 ]), unit.nanometers) # DEBUG nonalchemical_trajectory = unit.Quantity( np.zeros([ n_iterations, nonalchemical_thermodynamic_state.system.getNumParticles(), 3 ]), unit.nanometers) # DEBUG for iteration in range(n_iterations): # Generate a new sampler state for the hybrid system mc_move.apply(lambda_thermodynamic_state, new_sampler_state) # Compute the hybrid reduced potential at the new sampler state hybrid_context, integrator = cache.global_context_cache.get_context( lambda_thermodynamic_state) new_sampler_state.apply_to_context(hybrid_context, ignore_velocities=True) hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential( hybrid_context) if write_state: state = hybrid_context.getState(getPositions=True, getParameters=True) state_xml = openmm.XmlSerializer.serialize(state) with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile: outfile.write(state_xml) # Construct a sampler state for the nonalchemical system if lambda_index == 0: nonalchemical_positions = factory.old_positions( new_sampler_state.positions) elif lambda_index == 1: nonalchemical_positions = factory.new_positions( new_sampler_state.positions) else: raise ValueError( "The lambda index needs to be either one or zero for this to be meaningful" ) nonalchemical_sampler_state = SamplerState( nonalchemical_positions, box_vectors=new_sampler_state.box_vectors) if write_trajectories: state = hybrid_context.getState(getPositions=True) hybrid_trajectory[iteration, :, :] = state.getPositions( asNumpy=True) nonalchemical_trajectory[iteration, :, :] = nonalchemical_positions # Compute the nonalchemical reduced potential nonalchemical_context, integrator = cache.global_context_cache.get_context( nonalchemical_thermodynamic_state) nonalchemical_sampler_state.apply_to_context(nonalchemical_context, ignore_velocities=True) nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential( nonalchemical_context) # Compute and store the work w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential non_potential[iteration] = nonalchemical_reduced_potential hybrid_potential[iteration] = hybrid_reduced_potential if print_work: print( f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}' ) if write_trajectories: if lambda_index == 0: nonalchemical_mdtraj_topology = md.Topology.from_openmm( factory._topology_proposal.old_topology) elif lambda_index == 1: nonalchemical_mdtraj_topology = md.Topology.from_openmm( factory._topology_proposal.new_topology) md.Trajectory( hybrid_trajectory / unit.nanometers, factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb') md.Trajectory(nonalchemical_trajectory / unit.nanometers, nonalchemical_mdtraj_topology).save( f'nonalchemical{lambda_index}.pdb') # Analyze data and return results [t0, g, Neff_max] = timeseries.detectEquilibration(w) w_burned_in = w[t0:] [df, ddf] = pymbar.EXP(w_burned_in) ddf_corrected = ddf * np.sqrt(g) results = [df, ddf_corrected, t0, Neff_max] return results, non_potential, hybrid_potential
prmtop = utils.get_data_filename('blues', 'tests/data/eqToluene.prmtop') #TOL-parm inpcrd = utils.get_data_filename('blues', 'tests/data/eqToluene.inpcrd') tol = parmed.load_file(prmtop, xyz=inpcrd) tol.system = tol.createSystem(nonbondedMethod=openmm.app.PME, nonbondedCutoff=10 * unit.angstrom, constraints=openmm.app.HBonds, hydrogenMass=3.024 * unit.dalton, rigidWater=True, removeCMMotion=True, flexibleConstraints=True, splitDihedrals=False) # Create our State objects sampler_state = SamplerState(positions=tol.positions) thermodynamic_state = ThermodynamicState(system=tol.system, temperature=temperature) md_reporter = NetCDF4Storage(outfname + '.nc', reportInterval) state_reporter = BLUESStateDataStorage(reportInterval=reportInterval, title='md', step=True, speed=True, progress=True, totalSteps=int(n_steps * nIter)) ncmc_state_reporter = BLUESStateDataStorage(reportInterval=reportInterval, title='ncmc', step=True,
class Propagator(OMMBIP): """ Propagator pseudocode: Step 1: initialization-- set iteration = 0, n_iterations = n_iterations, lambda = 0 (i.e. iteration / n_iterations); work_accumulated = 0.0 generate sample x_0 ~ e^(-p(x)) evaluate work_incremental = 0 (i.e. u_mm(x_0) - g(x_0), but we presume that g = u_mm(.)) work_accumulated <- work_accumulated + work_incremental x' = x_0 Step 2: sampling for increment in range(n_iterations): x = x' ante_perturbation_potential = (1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x) set iteration <- iteration + 1.0; lambda <- iteration / n_iterations evaluate work_incremental = [(1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)] - ante_perturbation_potential work_accumulated <- work_accumulated + work_incremental create a modified force: modified_f = (1 - lambda) * f_mm + lambda * f_ani_mm_mix (where f_. = -grad(u_.) ) x' = V R O R (where V deterministic update is according to modified_f defined above) w.r.t x NOTE: in this regime, the last x' is propagated w.r.t. a propagator whose invariant distribution respects u_ani_mm_mix; this should _not_ be the case. There should be an exception in the Step 2 for loop that breaks once the final work_incremental is computed and updated to the work_accumulated. Regardless, the distribution of accumulated works is unaffected by this 'bug'; only expectations (as a function of x) w.r.t. these weights may be affected. See: 3.1.1. of https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf (esp. Remark 1.) """ def __init__(self, openmm_pdf_state, openmm_pdf_state_subset, subset_indices_map, integrator, ani_handler, context_cache=None, reassign_velocities=True, n_restart_attempts=0, reporter=None, write_trajectory_interval = 1, **kwargs): """ arguments openmm_pdf_state : openmmtools.states.ThermodynamicState the pdf state of the propagator openmm_pdf_state_subset : openmmtools.states.ThermodynamicState the pdf state of the atom subset subset_indices_map : dict dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index} integrator : openmm.Integrator integrator of dynamics ani_handler : ANI1_force_and_energy handler for ani forces and potential energy context_cache : openmmtools.cache.ContextCache, optional default:None The ContextCache to use for Context creation. If None, the global cache openmmtools.cache.global_context_cache is used. reassign_velocities : bool, optional default:False If True, the velocities will be reassigned from the Maxwell-Boltzmann distribution at the beginning of the move. n_restart_attempts : int, optional default:0 When greater than 0, if after the integration there are NaNs in energies, the move will restart. When the integrator has a random component, this may help recovering. On the last attempt, the ``Context`` is re-initialized in a slower process, but better than the simulation crashing. An IntegratorMoveError is raised after the given number of attempts if there are still NaNs. reporter : coddiwomple.openmm.reporter.OpenMMReporter, default None a reporter object to write trajectories write_trajectory_interval : int frequency of writing trajectory """ super().__init__(openmm_pdf_state, integrator, context_cache, reassign_velocities, n_restart_attempts) #create a pdf state for the subset indices (usually a vacuum system) self.pdf_state_subset = openmm_pdf_state_subset assert self.pdf_state_subset.temperature == self.pdf_state.temperature, f"the temperatures of the pdf states do not match" #create a dictionary for subset indices self._subset_indices_map = subset_indices_map #create an ani handler attribute that can be referenced self.ani_handler = ani_handler #create a context for the subset atoms that can be referenced self.context_subset, _ = cache.global_context_cache.get_context(self.pdf_state_subset) #create a reporter for the accumulated works self._state_works = {} self._state_works_counter = 0 #create a reporter self._write_trajectory = False if reporter is None else True self.reporter=reporter if self._write_trajectory: from coddiwomple.particles import Particle self.particle = Particle(0) self.write_trajectory_interval=write_trajectory_interval else: self.particle = None self.write_trajectory_interval=None def _initialize_state_works(self): """ initialize an empty list and add 0.0 to it (state works) """ self._current_state_works = [] #define an interim (auxiliary) list that will track the thermodynamic work of the current application self._current_state_works.append(0.0) #the first incremental work is always 0 since the importance function is identical to the first target distribution (i.e. fully interacting MM) def _initialize_iterations(self, n_iterations): """ initialize the iteration counter """ self._iteration = 0.0 #define the first iteration as 0 self._n_iterations = n_iterations #the number of iterations in the protocol is equal to the number of steps in the application def _update_particle_state_substate(self, particle_state, new_state_subset=False): """ update the particle state from the context, create a particle substate and update from context """ #update the particle state and the particle state subset particle_state.update_from_context(self.context, ignore_velocities=True) #update the particle state from the context if new_state_subset: self.particle_state_subset = SamplerState(positions = particle_state.positions[list(self._subset_indices_map.keys())]) #create a particle state from the subset context else: self.particle_state_subset.positions = particle_state.positions[list(self._subset_indices_map.keys())] #update the particle subset positions appropriately self.particle_state_subset.apply_to_context(self.context_subset, ignore_velocities=True) #apply the subset particle state to its context self.particle_state_subset.update_from_context(self.context_subset, ignore_velocities=True) #update the subset particle state from its context to updated the potential energy def _update_current_state_works(self, particle_state): """ update the current state and associated works """ #get the reduced potential reduced_potential = self._compute_hybrid_potential(_lambda = self._iteration / self._n_iterations, particle_state = particle_state) perturbed_reduced_potential = self._compute_hybrid_potential(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state) self._current_state_works.append(self._current_state_works[-1] + (perturbed_reduced_potential - reduced_potential)) def _update_force(self, particle_state): """ update the force """ mm_force_matrix = self._compute_hybrid_forces(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state).value_in_unit_system(unit.md_unit_system) self.integrator.setPerDofVariableByName('modified_force', mm_force_matrix) def _before_integration(self, *args, **kwargs): particle_state = args[0] #define the particle state n_iterations = args[1] #define the number of iterations self._initialize_state_works() self._initialize_iterations(n_iterations) #update the particle state and the particle state subset self._update_particle_state_substate(particle_state, new_state_subset=True) self._update_current_state_works(particle_state) self._update_force(particle_state) #report if self._write_trajectory: # the first state is always saved for processing purposes self.particle.update_state(particle_state) self.reporter.record([self.particle]) def _during_integration(self, *args, **kwargs): particle_state = args[0] self._iteration += 1.0 self._update_particle_state_substate(particle_state) #get the reduced potential if self._iteration < self._n_iterations: self._update_current_state_works(particle_state) self._update_force(particle_state) else: #we are done pass if self._write_trajectory and int(self._iteration) % self.write_trajectory_interval == 0: self.particle.update_state(particle_state) if self._iteration == self._n_iterations: self.reporter.record([self.particle], save_to_disk=True) else: self.reporter.record([self.particle], save_to_disk=False) def _after_integration(self, *args, **kwargs): self._state_works[self._state_works_counter] = deepcopy(self._current_state_works) self._state_works_counter += 1 if self._write_trajectory: self.reporter.reset() #self._log_context_parameters() def _compute_hybrid_potential(self,_lambda, particle_state): """ function to compute the hybrid reduced potential defined as follows: U(x_rec, x_lig) = u_mm,rec(x_rec) - lambda*u_mm,lig(x_lig) + lambda*u_ani,lig(x_lig) """ reduced_potential = (self.pdf_state.reduced_potential(particle_state) - _lambda * self.pdf_state_subset.reduced_potential(self.particle_state_subset) + _lambda * self.ani_handler.calculate_energy(self.particle_state_subset.positions) * self.pdf_state.beta) return reduced_potential def _compute_hybrid_forces(self, _lambda, particle_state): """ function to compute a hybrid force matrix of shape num_particles x 3 in the spirit of the _compute_hybrid_potential, we compute the forces in the following way F(x_rec, x_lig) = F_mm(x_rec, x_lig) - lambda * F_mm(x_lig) + lambda * F_ani(x_lig) """ # get the complex mm forces state = self.context.getState(getForces=True) mm_force_matrix = state.getForces(asNumpy=True) # returns forces in kJ/(nm mol) # get the ligand mm forces subset_state = self.context_subset.getState(getForces=True) mm_force_matrix_subset = subset_state.getForces(asNumpy=True) # get the ligand ani forces coords = self.particle_state_subset.positions subset_ani_force_matrix, energie = self.ani_handler.calculate_force(coords) # returns force in kJ/(A mol) #print(f"ani force matrix head: ",subset_ani_force_matrix[0]) # now combine the ligand forces subset_force_matrix = _lambda * (subset_ani_force_matrix - mm_force_matrix_subset) #we are adding two Quantities with different units, but they are compatible #print(f"mm subset force matrix head", mm_force_matrix_subset[0]) # and append to the complex forces... #print(f"mm force matrix head", mm_force_matrix[0]) mm_force_matrix[list(self._subset_indices_map.keys()), :] += subset_force_matrix #and same, here... #print(f"mm force matrix head (after ani modification)", mm_force_matrix[0]) return mm_force_matrix def _get_context_subset_parameters(self): """ return a dictionary of the self.context_subset's parameters returns context_parameters : dict {parameter name <str> : parameter value value <float>} """ swig_parameters = self.context_subset.getParameters() context_parameters = {q: swig_parameters[q] for q in swig_parameters} return context_parameters def _log_context_parameters(self): """ log the context and context subset parameters """ context_parameters = self._get_context_parameters() context_subset_parameters = self._get_context_subset_parameters() _logger.debug(f"\tcontext_parameters during integration:") for key, val in context_parameters.items(): _logger.debug(f"\t\t{key}: {val}") _logger.debug(f"\tcontext subset parameters during integration:") for key, val in context_subset_parameters: _logger.debug(f"\t\t{key}: {val}") @property def state_works(self): return self._state_works
# alanine = testsystems.AlanineDipeptideVacuum() # system = alanine.system # positions = alanine.positions # b) create & load our own pdb = PDBFile(os.path.dirname(os.path.realpath(__file__)) + '/input/alanine-dipeptide-nowater.pdb') forcefield = ForceField('amber10.xml') system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME, nonbondedCutoff=1*unit.nanometer, constraints=HBonds) thermodynamic_state = ThermodynamicState( system=system, temperature=298.15*unit.kelvin) positions = pdb.getPositions() # finally, start with sampling sampler_state = SamplerState(positions=positions) # TODO: decide on nr. of timesteps to go before recalculation V_calculator = TorsionNeglectingPotentialEnergyCalculator(timesteps=5) move = DeterministicRotateDisplaceMove( displacement_sigma=5.0*unit.nanometer, potential_energy_calculator=V_calculator, atom_subset=5, rng_seed=42 ) sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move) # run the sampler sampler.minimize() sampler.run(n_iterations=400) # TODO: increase iterations for final run
def setup(self, n_states, temperature, storage_file, minimisation_steps=100, n_replicas=None, lambda_schedule=None, lambda_protocol=LambdaProtocol(), endstates=True): from perses.dispersed import feptasks hybrid_system = self._factory.hybrid_system positions = self._factory.hybrid_positions lambda_zero_alchemical_state = RelativeAlchemicalState.from_system( hybrid_system) thermostate = ThermodynamicState(hybrid_system, temperature=temperature) compound_thermodynamic_state = CompoundThermodynamicState( thermostate, composable_states=[lambda_zero_alchemical_state]) thermodynamic_state_list = [] sampler_state_list = [] context_cache = cache.ContextCache() if n_replicas is None: _logger.info( f'n_replicas not defined, setting to match n_states, {n_states}' ) n_replicas = n_states elif n_replicas > n_states: _logger.warning( f'More sampler states: {n_replicas} requested greater than number of states: {n_states}. Setting n_replicas to n_states: {n_states}' ) n_replicas = n_states # TODO this feels like it should be somewhere else... just not sure where. Maybe into lambda_protocol if lambda_schedule is None: lambda_schedule = np.linspace(0., 1., n_states) else: assert ( len(lambda_schedule) == n_states ), 'length of lambda_schedule must match the number of states, n_states' assert ( lambda_schedule[0] == 0.), 'lambda_schedule must start at 0.' assert ( lambda_schedule[-1] == 1.), 'lambda_schedule must end at 1.' difference = np.diff(lambda_schedule) assert (all(i >= 0. for i in difference) ), 'lambda_schedule must be monotonicly increasing' #starting with the initial positions generated py geometry.py sampler_state = SamplerState( positions, box_vectors=hybrid_system.getDefaultPeriodicBoxVectors()) for lambda_val in lambda_schedule: compound_thermodynamic_state_copy = copy.deepcopy( compound_thermodynamic_state) compound_thermodynamic_state_copy.set_alchemical_parameters( lambda_val, lambda_protocol) thermodynamic_state_list.append(compound_thermodynamic_state_copy) # now generating a sampler_state for each thermodyanmic state, with relaxed positions context, context_integrator = context_cache.get_context( compound_thermodynamic_state_copy) feptasks.minimize(compound_thermodynamic_state_copy, sampler_state) sampler_state_list.append(copy.deepcopy(sampler_state)) reporter = storage_file # making sure number of sampler states equals n_replicas if len(sampler_state_list) != n_replicas: # picking roughly evenly spaced sampler states # if n_replicas == 1, then it will pick the first in the list idx = np.round( np.linspace(0, len(sampler_state_list) - 1, n_replicas)).astype(int) sampler_state_list = [ state for i, state in enumerate(sampler_state_list) if i in idx ] assert len(sampler_state_list) == n_replicas if endstates: # generating unsampled endstates _logger.info('Generating unsampled endstates.') unsampled_dispersion_endstates = create_endstates( copy.deepcopy(thermodynamic_state_list[0]), copy.deepcopy(thermodynamic_state_list[-1])) self.create( thermodynamic_states=thermodynamic_state_list, sampler_states=sampler_state_list, storage=reporter, unsampled_thermodynamic_states=unsampled_dispersion_endstates) else: self.create(thermodynamic_states=thermodynamic_state_list, sampler_states=sampler_state_list, storage=reporter)
def integrate(self, topology_proposal, initial_sampler_state, proposed_sampler_state, iteration=None): """ Performs NCMC switching to either delete or insert atoms according to the provided `topology_proposal`. For `delete`, the system is first modified from fully interacting to alchemically modified, and then NCMC switching is used to eliminate atoms. For `insert`, the system begins with eliminated atoms in an alchemically noninteracting form and NCMC switching is used to turn atoms on, followed by making system real. Parameters ---------- topology_proposal : TopologyProposal Contains old/new Topology and System objects and atom mappings. initial_sampler_state : openmmtools.states.SamplerState representing the initial (old) system Configurational properties of the atoms at the beginning of the NCMC switching. proposed_sampler_state : openmmtools.states.SamplerState representing the proposed (post-geometry new) system Configurational properties new system atoms at beginning of NCMC switching iteration : int, optional, default=None Iteration number, for storage purposes. Returns ------- final_old_sampler_state : openmmtools.State.SamplerState The final configurational properties of the old system after hybrid alchemical switching final_sampler_state : openmmtools.states.SamplerState The final configurational properties after `nsteps` steps of alchemical switching, and reversion to the nonalchemical system logP_work : float The NCMC work contribution to the log acceptance probability (Eqs. 62 and 63) logP_initial : float The initial logP of the hybrid configuration logP_final : float The final logP of the hybrid configuration """ assert not initial_sampler_state.has_nan() and not proposed_sampler_state.has_nan() #generate or retrieve the hybrid topology factory: hybrid_factory = self.make_alchemical_system(topology_proposal, initial_sampler_state.positions, proposed_sampler_state.positions) if hybrid_factory is None: _logger.warning("Unable to construct hybrid system for {} -> {}".format(topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key)) return initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0 topology = hybrid_factory.hybrid_topology #generate the corresponding thermodynamic and sampler states so that we can use the NonequilibriumSwitchingMove: #First generate the thermodynamic state: hybrid_system = hybrid_factory.hybrid_system hybrid_thermodynamic_state = ThermodynamicState(hybrid_system, temperature=self._temperature, pressure=self._pressure) #Now create an RelativeAlchemicalState from the hybrid system: alchemical_state = RelativeAlchemicalState.from_system(hybrid_system) alchemical_state.set_alchemical_parameters(0.0) #Now create a compound thermodynamic state that combines the hybrid thermodynamic state with the alchemical state: compound_thermodynamic_state = CompoundThermodynamicState(hybrid_thermodynamic_state, composable_states=[alchemical_state]) #construct a sampler state from the hybrid positions and the box vectors of the initial sampler state: initial_hybrid_positions = hybrid_factory.hybrid_positions initial_hybrid_box_vectors = initial_sampler_state.box_vectors initial_hybrid_sampler_state = SamplerState(initial_hybrid_positions, box_vectors=initial_hybrid_box_vectors) final_hybrid_sampler_state = copy.deepcopy(initial_hybrid_sampler_state) #create the nonequilibrium move: #ne_move = NonequilibriumSwitchingMove(self._functions, self._integrator_splitting, self._temperature, self._nsteps, self._timestep, # work_save_interval=self._write_ncmc_interval, top=topology,subset_atoms=None, # save_configuration=self._save_configuration, measure_shadow_work=self._measure_shadow_work) ne_move = ExternalNonequilibriumSwitchingMove(self._functions, nsteps_neq=self._nsteps, timestep=self._timestep, temperature=self._temperature, work_configuration_save_interval=self._work_save_interval, splitting="V R O R V") #run the NCMC protocol try: ne_move.apply(compound_thermodynamic_state, final_hybrid_sampler_state) except Exception as e: _logger.warn("NCMC failed because {}; rejecting.".format(str(e))) logP_work = -np.inf return [initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0] #get the total work: logP_work = - ne_move.cumulative_work[-1] # Compute contribution of transforming to and from the hybrid system: context, integrator = global_context_cache.get_context(hybrid_thermodynamic_state) #set all alchemical parameters to zero: for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) initial_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) initial_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context) #set all alchemical parameters to one: for parameter in self._functions.keys(): context.setParameter(parameter, 1.0) final_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) final_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context) #reset the parameters back to zero just in case for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) #compute the output SamplerState, which has the atoms only for the new system post-NCMC: new_positions = hybrid_factory.new_positions(final_hybrid_sampler_state.positions) new_box_vectors = final_hybrid_sampler_state.box_vectors final_sampler_state = SamplerState(new_positions, box_vectors=new_box_vectors) #compute the output SamplerState for the atoms only in the old system (required for geometry_logP_reverse) old_positions = hybrid_factory.old_positions(final_hybrid_sampler_state.positions) old_box_vectors = copy.deepcopy(new_box_vectors) #these are the same as the new system final_old_sampler_state = SamplerState(old_positions, box_vectors=old_box_vectors) #extract the trajectory and box vectors from the move: trajectory = ne_move.trajectory[::-self._write_ncmc_interval, :, :][::-1] topology = hybrid_factory.hybrid_topology position_varname = "ncmcpositions" nframes = np.shape(trajectory)[0] #extract box vectors: box_vec_varname = "ncmcboxvectors" box_lengths = ne_move.box_lengths[::-self._write_ncmc_interval, :][::-1] box_angles = ne_move.box_angles[::-self._write_ncmc_interval, :][::-1] box_lengths_and_angles = np.stack([box_lengths, box_angles]) #write out the positions of the topology if self._storage: for frame in range(nframes): self._storage.write_configuration(position_varname, trajectory[frame, :, :], topology, iteration=iteration, frame=frame, nframes=nframes) #write out the periodict box vectors: if self._storage: self._storage.write_array(box_vec_varname, box_lengths_and_angles, iteration=iteration) #retrieve the protocol work and write that out too: protocol_work = ne_move.cumulative_work if self._storage: self._storage.write_array("protocolwork", protocol_work, iteration=iteration) # Return return [final_old_sampler_state, final_sampler_state, logP_work, -initial_reduced_potential, -final_reduced_potential]
def integrate(self, topology_proposal, initial_sampler_state, proposed_sampler_state, iteration=None): """ Performs NCMC switching to either delete or insert atoms according to the provided `topology_proposal`. For `delete`, the system is first modified from fully interacting to alchemically modified, and then NCMC switching is used to eliminate atoms. For `insert`, the system begins with eliminated atoms in an alchemically noninteracting form and NCMC switching is used to turn atoms on, followed by making system real. Parameters ---------- topology_proposal : TopologyProposal Contains old/new Topology and System objects and atom mappings. initial_sampler_state : openmmtools.states.SamplerState representing the initial (old) system Configurational properties of the atoms at the beginning of the NCMC switching. proposed_sampler_state : openmmtools.states.SamplerState representing the proposed (post-geometry new) system Configurational properties new system atoms at beginning of NCMC switching iteration : int, optional, default=None Iteration number, for storage purposes. Returns ------- final_old_sampler_state : openmmtools.State.SamplerState The final configurational properties of the old system after hybrid alchemical switching final_sampler_state : openmmtools.states.SamplerState The final configurational properties after `nsteps` steps of alchemical switching, and reversion to the nonalchemical system logP_work : float The NCMC work contribution to the log acceptance probability (Eqs. 62 and 63) logP_initial : float The initial logP of the hybrid configuration logP_final : float The final logP of the hybrid configuration """ assert not initial_sampler_state.has_nan( ) and not proposed_sampler_state.has_nan() #generate or retrieve the hybrid topology factory: hybrid_factory = self.make_alchemical_system( topology_proposal, initial_sampler_state.positions, proposed_sampler_state.positions) if hybrid_factory is None: _logger.warning( "Unable to construct hybrid system for {} -> {}".format( topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key)) return initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0 topology = hybrid_factory.hybrid_topology #generate the corresponding thermodynamic and sampler states so that we can use the NonequilibriumSwitchingMove: #First generate the thermodynamic state: hybrid_system = hybrid_factory.hybrid_system hybrid_thermodynamic_state = ThermodynamicState( hybrid_system, temperature=self._temperature, pressure=self._pressure) #Now create an RelativeAlchemicalState from the hybrid system: alchemical_state = RelativeAlchemicalState.from_system(hybrid_system) alchemical_state.set_alchemical_parameters(0.0) #Now create a compound thermodynamic state that combines the hybrid thermodynamic state with the alchemical state: compound_thermodynamic_state = CompoundThermodynamicState( hybrid_thermodynamic_state, composable_states=[alchemical_state]) #construct a sampler state from the hybrid positions and the box vectors of the initial sampler state: initial_hybrid_positions = hybrid_factory.hybrid_positions initial_hybrid_box_vectors = initial_sampler_state.box_vectors initial_hybrid_sampler_state = SamplerState( initial_hybrid_positions, box_vectors=initial_hybrid_box_vectors) final_hybrid_sampler_state = copy.deepcopy( initial_hybrid_sampler_state) #create the nonequilibrium move: #ne_move = NonequilibriumSwitchingMove(self._functions, self._integrator_splitting, self._temperature, self._nsteps, self._timestep, # work_save_interval=self._write_ncmc_interval, top=topology,subset_atoms=None, # save_configuration=self._save_configuration, measure_shadow_work=self._measure_shadow_work) ne_move = ExternalNonequilibriumSwitchingMove( self._functions, nsteps_neq=self._nsteps, timestep=self._timestep, temperature=self._temperature, work_configuration_save_interval=self._work_save_interval, splitting="V R O R V") #run the NCMC protocol try: ne_move.apply(compound_thermodynamic_state, final_hybrid_sampler_state) except Exception as e: _logger.warn("NCMC failed because {}; rejecting.".format(str(e))) logP_work = -np.inf return [ initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0 ] #get the total work: logP_work = -ne_move.cumulative_work[-1] # Compute contribution of transforming to and from the hybrid system: context, integrator = global_context_cache.get_context( hybrid_thermodynamic_state) #set all alchemical parameters to zero: for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) initial_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) initial_reduced_potential = hybrid_thermodynamic_state.reduced_potential( context) #set all alchemical parameters to one: for parameter in self._functions.keys(): context.setParameter(parameter, 1.0) final_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) final_reduced_potential = hybrid_thermodynamic_state.reduced_potential( context) #reset the parameters back to zero just in case for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) #compute the output SamplerState, which has the atoms only for the new system post-NCMC: new_positions = hybrid_factory.new_positions( final_hybrid_sampler_state.positions) new_box_vectors = final_hybrid_sampler_state.box_vectors final_sampler_state = SamplerState(new_positions, box_vectors=new_box_vectors) #compute the output SamplerState for the atoms only in the old system (required for geometry_logP_reverse) old_positions = hybrid_factory.old_positions( final_hybrid_sampler_state.positions) old_box_vectors = copy.deepcopy( new_box_vectors) #these are the same as the new system final_old_sampler_state = SamplerState(old_positions, box_vectors=old_box_vectors) #extract the trajectory and box vectors from the move: trajectory = ne_move.trajectory[::-self. _write_ncmc_interval, :, :][::-1] topology = hybrid_factory.hybrid_topology position_varname = "ncmcpositions" nframes = np.shape(trajectory)[0] #extract box vectors: box_vec_varname = "ncmcboxvectors" box_lengths = ne_move.box_lengths[::-self. _write_ncmc_interval, :][::-1] box_angles = ne_move.box_angles[::-self._write_ncmc_interval, :][::-1] box_lengths_and_angles = np.stack([box_lengths, box_angles]) #write out the positions of the topology if self._storage: for frame in range(nframes): self._storage.write_configuration(position_varname, trajectory[frame, :, :], topology, iteration=iteration, frame=frame, nframes=nframes) #write out the periodict box vectors: if self._storage: self._storage.write_array(box_vec_varname, box_lengths_and_angles, iteration=iteration) #retrieve the protocol work and write that out too: protocol_work = ne_move.cumulative_work if self._storage: self._storage.write_array("protocolwork", protocol_work, iteration=iteration) # Return return [ final_old_sampler_state, final_sampler_state, logP_work, -initial_reduced_potential, -final_reduced_potential ]
def run_endpoint_perturbation(lambda_thermodynamic_state, nonalchemical_thermodynamic_state, initial_hybrid_sampler_state, mc_move, n_iterations, factory, lambda_index=0, print_work=False, write_system=False, write_state=False, write_trajectories=False): """ Parameters ---------- lambda_thermodynamic_state : ThermodynamicState The thermodynamic state corresponding to the hybrid system at a lambda endpoint nonalchemical_thermodynamic_state : ThermodynamicState The nonalchemical thermodynamic state for the relevant endpoint initial_hybrid_sampler_state : SamplerState Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state mc_move : MCMCMove The MCMove that will be used for sampling at the lambda endpoint n_iterations : int The number of iterations factory : HybridTopologyFactory The hybrid topology factory lambda_index : int, optional, default=0 The index, 0 or 1, at which to retrieve nonalchemical positions print_work : bool, optional, default=False If True, will print work values write_system : bool, optional, default=False If True, will write alchemical and nonalchemical System XML files write_state : bool, optional, default=False If True, write alchemical (hybrid) State XML files each iteration write_trajectories : bool, optional, default=False If True, will write trajectories Returns ------- df : float Free energy difference between alchemical and nonalchemical systems, estimated with EXP ddf : float Standard deviation of estimate, corrected for correlation, from EXP estimator. """ import mdtraj as md #run an initial minimization: mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state, initial_hybrid_sampler_state, mc_move) mcmc_sampler.minimize(max_iterations=20) new_sampler_state = mcmc_sampler.sampler_state if write_system: with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile: outfile.write(openmm.XmlSerializer.serialize(lambda_thermodynamic_state.system)) with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile: outfile.write(openmm.XmlSerializer.serialize(nonalchemical_thermodynamic_state.system)) #initialize work array w = np.zeros([n_iterations]) non_potential = np.zeros([n_iterations]) hybrid_potential = np.zeros([n_iterations]) #run n_iterations of the endpoint perturbation: hybrid_trajectory = unit.Quantity(np.zeros([n_iterations, lambda_thermodynamic_state.system.getNumParticles(), 3]), unit.nanometers) # DEBUG nonalchemical_trajectory = unit.Quantity(np.zeros([n_iterations, nonalchemical_thermodynamic_state.system.getNumParticles(), 3]), unit.nanometers) # DEBUG for iteration in range(n_iterations): # Generate a new sampler state for the hybrid system mc_move.apply(lambda_thermodynamic_state, new_sampler_state) # Compute the hybrid reduced potential at the new sampler state hybrid_context, integrator = cache.global_context_cache.get_context(lambda_thermodynamic_state) new_sampler_state.apply_to_context(hybrid_context, ignore_velocities=True) hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(hybrid_context) if write_state: state = hybrid_context.getState(getPositions=True, getParameters=True) state_xml = openmm.XmlSerializer.serialize(state) with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile: outfile.write(state_xml) # Construct a sampler state for the nonalchemical system if lambda_index == 0: nonalchemical_positions = factory.old_positions(new_sampler_state.positions) elif lambda_index == 1: nonalchemical_positions = factory.new_positions(new_sampler_state.positions) else: raise ValueError("The lambda index needs to be either one or zero for this to be meaningful") nonalchemical_sampler_state = SamplerState(nonalchemical_positions, box_vectors=new_sampler_state.box_vectors) if write_trajectories: state = hybrid_context.getState(getPositions=True) hybrid_trajectory[iteration,:,:] = state.getPositions(asNumpy=True) nonalchemical_trajectory[iteration,:,:] = nonalchemical_positions # Compute the nonalchemical reduced potential nonalchemical_context, integrator = cache.global_context_cache.get_context(nonalchemical_thermodynamic_state) nonalchemical_sampler_state.apply_to_context(nonalchemical_context, ignore_velocities=True) nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(nonalchemical_context) # Compute and store the work w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential non_potential[iteration] = nonalchemical_reduced_potential hybrid_potential[iteration] = hybrid_reduced_potential if print_work: print(f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}') if write_trajectories: if lambda_index == 0: nonalchemical_mdtraj_topology = md.Topology.from_openmm(factory._topology_proposal.old_topology) elif lambda_index == 1: nonalchemical_mdtraj_topology = md.Topology.from_openmm(factory._topology_proposal.new_topology) md.Trajectory(hybrid_trajectory / unit.nanometers, factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb') md.Trajectory(nonalchemical_trajectory / unit.nanometers, nonalchemical_mdtraj_topology).save(f'nonalchemical{lambda_index}.pdb') # Analyze data and return results [t0, g, Neff_max] = timeseries.detectEquilibration(w) w_burned_in = w[t0:] [df, ddf] = pymbar.EXP(w_burned_in) ddf_corrected = ddf * np.sqrt(g) results = [df, ddf_corrected, t0, Neff_max] return results, non_potential, hybrid_potential