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_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 __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 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 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 _topology_proposal_to_thermodynamic_states(self, topology_proposal): """ Convert a topology proposal to thermodynamic states for the end systems. This will be used to compute the "logP_energy" quantity. Parameters ---------- topology_proposal : perses.rjmc.TopologyProposal topology proposal for whose endpoint systems we want ThermodynamicStates Returns ------- old_thermodynamic_state : openmmtools.states.ThermodynamicState The old system (nonalchemical) thermodynamic state new_thermodynamic_state : openmmtools.states.ThermodynamicState The new system (nonalchemical) thermodynamic state """ systems = [topology_proposal.old_system, topology_proposal.new_system] thermostates = [] for system in systems: thermodynamic_state = ThermodynamicState( system, temperature=self._temperature, pressure=self._pressure) thermostates.append(thermodynamic_state) return thermostates[0], thermostates[1]
def generate_thermodynamic_states(system: openmm.System, topology_proposal: TopologyProposal): """ Generate endpoint thermodynamic states for the system Parameters ---------- system : openmm.System System object corresponding to thermodynamic state topology_proposal : perses.rjmc.topology_proposal.TopologyProposal TopologyProposal representing transformation Returns ------- nonalchemical_zero_thermodynamic_state : ThermodynamicState Nonalchemical thermodynamic state for lambda zero endpoint nonalchemical_one_thermodynamic_state : ThermodynamicState Nonalchemical thermodynamic state for lambda one endpoint lambda_zero_thermodynamic_state : ThermodynamicState Alchemical (hybrid) thermodynamic state for lambda zero lambda_one_thermodynamic_State : ThermodynamicState Alchemical (hybrid) thermodynamic state for lambda one """ #create the thermodynamic state lambda_zero_alchemical_state = alchemy.AlchemicalState.from_system(system) lambda_one_alchemical_state = copy.deepcopy(lambda_zero_alchemical_state) #ensure their states are set appropriately lambda_zero_alchemical_state.set_alchemical_parameters(0.0) lambda_one_alchemical_state.set_alchemical_parameters(0.0) #create the base thermodynamic state with the hybrid system thermodynamic_state = ThermodynamicState(system, temperature=temperature) #Create thermodynamic states for the nonalchemical endpoints nonalchemical_zero_thermodynamic_state = ThermodynamicState( topology_proposal.old_system, temperature=temperature) nonalchemical_one_thermodynamic_state = ThermodynamicState( topology_proposal.new_system, temperature=temperature) #Now create the compound states with different alchemical states lambda_zero_thermodynamic_state = CompoundThermodynamicState( thermodynamic_state, composable_states=[lambda_zero_alchemical_state]) lambda_one_thermodynamic_state = CompoundThermodynamicState( thermodynamic_state, composable_states=[lambda_one_alchemical_state]) return nonalchemical_zero_thermodynamic_state, nonalchemical_one_thermodynamic_state, lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state
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 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 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 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_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 _system_to_thermodynamic_state(self, system): """ Given an OpenMM system object, create a corresponding ThermodynamicState that has the same temperature and pressure as the current thermodynamic state. Parameters ---------- system : openmm.System The OpenMM system for which to create the thermodynamic state Returns ------- new_thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state object representing the given system """ return ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)
def compare_energy_components(rest_system, other_system, positions, platform=REFERENCE_PLATFORM): """ Get energy components of a given system """ platform = configure_platform(platform) # Create thermodynamic state and sampler state for non-rest system thermostate_other = ThermodynamicState(system=other_system, temperature=temperature) # Create context for non-rest system integrator_other = openmm.VerletIntegrator(1.0 * unit.femtosecond) context_other = thermostate_other.create_context(integrator_other) context_other.setPositions(positions) # Get energy components for non-rest system components_other = [ component[1] for component in compute_potential_components(context_other, beta=beta) ] # Create thermodynamic state for rest_system thermostate_rest = ThermodynamicState(system=rest_system, temperature=temperature) # Create context forrest system integrator_rest = openmm.VerletIntegrator(1.0 * unit.femtosecond) context_rest = thermostate_rest.create_context(integrator_rest) context_rest.setPositions(positions) # Get energy components for rest system components_rest = [ component[1] for component in compute_potential_components(context_rest, beta=beta) ] # Check that bond, angle, and torsion energies match for other, rest in zip(components_other[:3], components_rest[:3]): assert np.isclose( [other], [rest] ), f"The energies do not match for the {other[0]}: {other[1]} (other system) vs. {rest[1]} (REST system)" # Check that nonbonded energies match print(components_rest) nonbonded_other = np.array(components_other[3:]).sum() nonbonded_rest = np.array(components_rest[3:]).sum() assert np.isclose( [nonbonded_other], [nonbonded_rest] ), f"The energies do not match for the NonbondedForce: {nonbonded_other} (other system) vs. {nonbonded_rest} (REST system)"
def __init__(self, system, alchemical_composability=AlchemicalState, temperature=300 * unit.kelvin, pressure=1.0 * unit.atmosphere, **kwargs): """ Subclass of coddiwomple.states.PDFState that specifically handles openmmtools.states.CompoundThermodynamicStates Init method does the following: 1. create a openmmtools.states.ThermodynamicState with the given system, temperature, and pressure 2. create an openmmtools.alchemy.AlchemicalState from the given system 3. create an openmtools.states.CompoundThermodynamicState 1 and 2 arguments system : openmm.System system object to wrap alchemical_composability : openmmtools.alchemy.AlchemicalState (super), default openmmtools.alchemy.AlchemicalState the class holding the method `from_system` which can compose an alchemical state with exposed parameters from a system temperature : float * unit.kelvin (or temperature units), default 300.0 * unit.kelvin temperature of the system pressure : float * unit.atmosphere (or pressure units), default 1.0 * unit.atmosphere pressure of the system init method is adapted from https://github.com/choderalab/openmmtools/blob/110524bc5079af77d31f5fab464edd7b668ff5ac/openmmtools/states.py#L2766-L2783 """ from openmmtools.alchemy import AlchemicalState openmm_pdf_state = ThermodynamicState(system, temperature, pressure) alchemical_state = alchemical_composability.from_system( system, **kwargs) assert isinstance( alchemical_state, IComposableState ), f"alchemical state is not an instance of IComposableState" self.__dict__ = openmm_pdf_state.__dict__ self._composable_states = [alchemical_state] self.set_system(self._standard_system, fix_state=True) #class create an internal context integrator = get_dummy_integrator() platform = configure_platform(utils.get_fastest_platform().getName()) self._internal_context = self.create_context(integrator, platform=platform) _logger.debug( f"successfully instantiated OpenMMPDFState equipped with the following parameters: {self._parameters}" )
def test_create_endstates(): """ test the creation of unsampled endstates """ from pkg_resources import resource_filename smiles_filename = resource_filename("perses", os.path.join("data", "test.smi")) fe_setup = RelativeFEPSetup(ligand_input=smiles_filename, old_ligand_index=0, new_ligand_index=1, forcefield_files=[], small_molecule_forcefield='gaff-2.11', phases=['vacuum']) hybrid_factory = HybridTopologyFactory( topology_proposal=fe_setup._vacuum_topology_proposal, current_positions=fe_setup._vacuum_positions_old, new_positions=fe_setup._vacuum_positions_new, neglected_new_angle_terms=fe_setup._vacuum_forward_neglected_angles, neglected_old_angle_terms=fe_setup._vacuum_reverse_neglected_angles, softcore_LJ_v2=True, interpolate_old_and_new_14s=False) zero_state_error, one_state_error = validate_endstate_energies( fe_setup._vacuum_topology_proposal, hybrid_factory, added_energy=fe_setup._vacuum_added_valence_energy, subtracted_energy=fe_setup._vacuum_subtracted_valence_energy, beta=beta, platform=openmm.Platform.getPlatformByName('Reference'), ENERGY_THRESHOLD=ENERGY_THRESHOLD) lambda_alchemical_state = RelativeAlchemicalState.from_system( hybrid_factory.hybrid_system) lambda_protocol = LambdaProtocol(functions='default') lambda_alchemical_state.set_alchemical_parameters(0.0, lambda_protocol) thermodynamic_state = CompoundThermodynamicState( ThermodynamicState(hybrid_factory.hybrid_system, temperature=temperature), composable_states=[lambda_alchemical_state]) zero_endstate = copy.deepcopy(thermodynamic_state) one_endstate = copy.deepcopy(thermodynamic_state) one_endstate.set_alchemical_parameters(1.0, lambda_protocol) new_endstates = create_endstates(zero_endstate, one_endstate)
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 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_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 test_create_endstates(): """ test the creation of unsampled endstates """ fe_setup = RelativeFEPSetup(ligand_input=f"{os.getcwd()}/test.smi", old_ligand_index=0, new_ligand_index=1, forcefield_files=['gaff.xml'], phases=['vacuum']) hybrid_factory = HybridTopologyFactory( topology_proposal=fe_setup._vacuum_topology_proposal, current_positions=fe_setup._vacuum_positions_old, new_positions=fe_setup._vacuum_positions_new, neglected_new_angle_terms=fe_setup._vacuum_forward_neglected_angles, neglected_old_angle_terms=fe_setup._vacuum_reverse_neglected_angles, softcore_LJ_v2=True, interpolate_old_and_new_14s=False) zero_state_error, one_state_error = validate_endstate_energies( fe_setup._vacuum_topology_proposal, hybrid_factory, added_energy=fe_setup._vacuum_added_valence_energy, subtracted_energy=fe_setup._vacuum_subtracted_valence_energy, beta=beta, ENERGY_THRESHOLD=ENERGY_THRESHOLD) lambda_alchemical_state = RelativeAlchemicalState.from_system( hybrid_factory.hybrid_system) lambda_protocol = LambdaProtocol(functions='default') lambda_alchemical_state.set_alchemical_parameters(0.0, lambda_protocol) thermodynamic_state = CompoundThermodynamicState( ThermodynamicState(hybrid_factory.hybrid_system, temperature=temperature), composable_states=[lambda_alchemical_state]) zero_endstate = copy.deepcopy(thermodynamic_state) one_endstate = copy.deepcopy(thermodynamic_state) one_endstate.set_alchemical_parameters(1.0, lambda_protocol) new_endstates = create_endstates(zero_endstate, one_endstate)
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_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 create_endstates(first_thermostate, last_thermostate): """ utility function to generate unsampled endstates 1. move all alchemical atom LJ parameters from CustomNonbondedForce to NonbondedForce 2. delete the CustomNonbondedForce 3. set PME tolerance to 1e-5 4. enable LJPME to handle long range dispersion corrections in a physically reasonable manner Arguments --------- first_thermostate : openmmtools.states.CompoundThermodynamicState the first thermodynamic state for which an unsampled endstate will be created last_thermostate : openmmtools.states.CompoundThermodynamicState the last thermodynamic state for which an unsampled endstate will be created Returns ------- unsampled_endstates : list of openmmtools.states.CompoundThermodynamicState the corrected unsampled endstates """ unsampled_endstates = [] for master_lambda, endstate in zip([0., 1.], [first_thermostate, last_thermostate]): dispersion_system = endstate.get_system() energy_unit = unit.kilocalories_per_mole # Find the NonbondedForce (there must be only one) forces = { force.__class__.__name__: force for force in dispersion_system.getForces() } # Set NonbondedForce to use LJPME forces['NonbondedForce'].setNonbondedMethod( openmm.NonbondedForce.LJPME) # Set tight PME tolerance TIGHT_PME_TOLERANCE = 1.0e-5 forces['NonbondedForce'].setEwaldErrorTolerance(TIGHT_PME_TOLERANCE) # Move alchemical LJ sites from CustomNonbondedForce back to NonbondedForce for particle_index in range( forces['NonbondedForce'].getNumParticles()): charge, sigma, epsilon = forces[ 'NonbondedForce'].getParticleParameters(particle_index) sigmaA, epsilonA, sigmaB, epsilonB, unique_old, unique_new = forces[ 'CustomNonbondedForce'].getParticleParameters(particle_index) if (epsilon / energy_unit == 0.0) and ((epsilonA > 0.0) or (epsilonB > 0.0)): sigma = (1 - master_lambda) * sigmaA + master_lambda * sigmaB epsilon = (1 - master_lambda) * epsilonA + master_lambda * epsilonB forces['NonbondedForce'].setParticleParameters( particle_index, charge, sigma, epsilon) # Delete the CustomNonbondedForce since we have moved all alchemical particles out of it for force_index, force in enumerate(list( dispersion_system.getForces())): if force.__class__.__name__ == 'CustomNonbondedForce': custom_nonbonded_force_index = force_index break dispersion_system.removeForce(custom_nonbonded_force_index) # Set all parameters to master lambda for force_index, force in enumerate(list( dispersion_system.getForces())): if hasattr(force, 'getNumGlobalParameters'): for parameter_index in range(force.getNumGlobalParameters()): if force.getGlobalParameterName( parameter_index)[0:7] == 'lambda_': force.setGlobalParameterDefaultValue( parameter_index, master_lambda) # Store the unsampled endstate unsampled_endstates.append( ThermodynamicState(dispersion_system, temperature=endstate.temperature)) return unsampled_endstates
def run_equilibrium(equilibrium_result: EquilibriumResult, thermodynamic_state: states.ThermodynamicState, nsteps_equil: int, topology: md.Topology, n_iterations : int, atom_indices_to_save: List[int] = None, trajectory_filename: str = None, splitting: str="V R O R V", timestep: unit.Quantity=1.0*unit.femtoseconds) -> EquilibriumResult: """ Run nsteps of equilibrium sampling at the specified thermodynamic state and return the final sampler state as well as a trajectory of the positions after each application of an MCMove. This means that if the MCMove is configured to run 1000 steps of dynamics, and n_iterations is 100, there will be 100 frames in the resulting trajectory; these are the result of 100,000 steps (1000*100) of dynamics. Parameters ---------- equilibrium_result : EquilibriumResult EquilibriumResult namedtuple containing the information necessary to resume thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state (including context parameters) that should be used nsteps_equil : int The number of equilibrium steps that a move should make when apply is called topology : mdtraj.Topology an MDTraj topology object used to construct the trajectory n_iterations : int The number of times to apply the move. Note that this is not the number of steps of dynamics; it is n_iterations*n_steps (which is set in the MCMove). splitting: str, default "V R O H R V" The splitting string for the dynamics atom_indices_to_save : list of int, default None list of indices to save (when excluding waters, for instance). If None, all indices are saved. trajectory_filename : str, optional, default None Full filepath of trajectory files. If none, trajectory files are not written. splitting: str, default "V R O H R V" The splitting string for the dynamics Returns ------- equilibrium_result : EquilibriumResult Container namedtuple that has the SamplerState for resuming, an MDTraj trajectory, and the reduced potential of the final frame. """ sampler_state = equilibrium_result.sampler_state #get the atom indices we need to subset the topology and positions if atom_indices_to_save is None: atom_indices = list(range(topology.n_atoms)) subset_topology = topology else: subset_topology = topology.subset(atom_indices_to_save) atom_indices = atom_indices_to_save n_atoms = subset_topology.n_atoms #construct the MCMove: mc_move = mcmc.LangevinSplittingDynamicsMove(n_steps=nsteps_equil, splitting=splitting) mc_move.n_restart_attempts = 10 #create a numpy array for the trajectory trajectory_positions = np.zeros([n_iterations, n_atoms, 3]) trajectory_box_lengths = np.zeros([n_iterations, 3]) trajectory_box_angles = np.zeros([n_iterations, 3]) #loop through iterations and apply MCMove, then collect positions into numpy array for iteration in range(n_iterations): mc_move.apply(thermodynamic_state, sampler_state) trajectory_positions[iteration, :] = sampler_state.positions[atom_indices, :].value_in_unit_system(unit.md_unit_system) #get the box lengths and angles a, b, c, alpha, beta, gamma = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(*sampler_state.box_vectors) trajectory_box_lengths[iteration, :] = [a, b, c] trajectory_box_angles[iteration, :] = [alpha, beta, gamma] #construct trajectory object: trajectory = md.Trajectory(trajectory_positions, subset_topology, unitcell_lengths=trajectory_box_lengths, unitcell_angles=trajectory_box_angles) #get the reduced potential from the final frame for endpoint perturbations reduced_potential_final_frame = thermodynamic_state.reduced_potential(sampler_state) #construct equilibrium result object equilibrium_result = EquilibriumResult(sampler_state, reduced_potential_final_frame) #If there is a trajectory filename passed, write out the results here: if trajectory_filename is not None: write_equilibrium_trajectory(equilibrium_result, trajectory, trajectory_filename) return equilibrium_result
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)
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, speed=True, progress=True,
def ANI_endstate_sampler( system, system_subset, subset_indices_map, positions_cache_filename, md_topology, index_to_run, steps_per_application, integrator_kwargs = {'temperature': 300.0 * unit.kelvin, 'collision_rate': 1.0 / unit.picoseconds, 'timestep': 2.0 * unit.femtoseconds, 'splitting': "V R O R F", 'constraint_tolerance': 1e-6, 'pressure': 1.0 * unit.atmosphere}, position_extractor = None ): """ conduct ani endstate sampling arguments system : openmm.System system system_subset : openmm.System subset system subset_indices_map : dict dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index} positions_cache_filename : str path to the cache positions index_to_run : int index of the positions to anneal number_of_applications : int number of applications of the propagator steps_per_application : int number of integration steps per application integrator_kwargs : dict, see default kwargs to pass to OMMLIAIS integrator """ from coddiwomple.particles import Particle from coddiwomple.openmm.states import OpenMMParticleState, OpenMMPDFState #load the endstate cache traj = np.load(positions_cache_filename) positions = traj['positions'][index_to_run,:,:] * unit.nanometers if position_extractor is not None: positions = position_extractor(_positions) try: box_vectors = traj['box_vectors'][index_to_run,:,:] * unit.nanometers except Exception as e: box_vectors = None species_str = ''.join([atom.element.symbol for atom in md_topology.subset(list(subset_indices_map.keys())).atoms]) _logger.info(f"species string: {species_str}") ani_handler = ANI1_force_and_energy(model = torchani.models.ANI1ccx(), atoms=species_str, platform='cpu', temperature=integrator_kwargs['temperature']) #make thermostates pressure = integrator_kwargs['pressure'] if box_vectors is not None else None pdf_state = ThermodynamicState(system = system, temperature = integrator_kwargs['temperature'], pressure=pressure) pdf_state_subset = ThermodynamicState(system = system_subset, temperature = integrator_kwargs['temperature'], pressure = None) #make an integrator integrator = Integrator(**integrator_kwargs) #make a propagator propagator = ANIPropagator(openmm_pdf_state = pdf_state, openmm_pdf_state_subset = pdf_state_subset, subset_indices_map = subset_indices_map, integrator = integrator, ani_handler = ani_handler, context_cache=None, reassign_velocities=True, n_restart_attempts=0, reporter = None) particle = Particle(0) particle_state = OpenMMParticleState(positions = positions, box_vectors = box_vectors) particle.update_state(particle_state) particle_state, _return_dict = propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True) if box_vectors is None: particle_state.box_vectors=None return particle_state, np.array(propagator.state_works[0])
def annealed_importance_sampling(direction, system, system_subset, subset_indices_map, positions_cache_filename, index_to_run, directory_name, trajectory_prefix, md_topology, steps_per_application, integrator_kwargs = {'temperature': 300.0 * unit.kelvin, 'collision_rate': 1.0 / unit.picoseconds, 'timestep': 1.0 * unit.femtoseconds, 'splitting': "V R O R F", 'constraint_tolerance': 1e-6, 'pressure': 1.0 * unit.atmosphere}, save_indices = None, position_extractor = None, write_trajectory_interval=1 ): """ conduct annealed importance sampling in the openmm regime; will write the accumulated work dictionary after each application arguments direction : str forward or backward system : openmm.System system system_subset : openmm.System subset system subset_indices_map : dict dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index} positions_cache_filename : str path to the cache positions index_to_run : int index of the positions to anneal directory_name : str directory that will be written to trajectory_prefix : str .pdb prefix md_topology : mdtraj.Topology topology that will write the trajectory save_indices : list list of indices that will be saved number_of_applications : int number of applications of the propagator steps_per_application : int number of integration steps per application integrator_kwargs : dict, see default kwargs to pass to OMMLIAIS integrator save_indices : list(int) list of indices of md_topology atoms to save to disk position_extractor : function, default None function to extract appropriate positons from the cache write_trajectory_interval : int frequency with which to write trajectory to disk """ from coddiwomple.particles import Particle from coddiwomple.openmm.states import OpenMMParticleState, OpenMMPDFState #load the endstate cache traj = np.load(positions_cache_filename) positions = traj['positions'][index_to_run,:,:] * unit.nanometers if position_extractor is not None: positions = position_extractor(_positions) try: box_vectors = traj['box_vectors'][index_to_run,:,:] * unit.nanometers except Exception as e: box_vectors = None assert direction in ['forward', 'backward'] #make a handle object for ANI species_str = ''.join([atom.element.symbol for atom in md_topology.subset(list(subset_indices_map.keys())).atoms]) _logger.info(f"species string: {species_str}") ani_handler = ANI1_force_and_energy(model = torchani.models.ANI1ccx(), atoms=species_str, platform='cpu', temperature=integrator_kwargs['temperature']) #make thermostates pressure = integrator_kwargs['pressure'] if box_vectors is not None else None pdf_state = ThermodynamicState(system = system, temperature = integrator_kwargs['temperature'], pressure=pressure) pdf_state_subset = ThermodynamicState(system = system_subset, temperature = integrator_kwargs['temperature'], pressure = None) #make a reporter saveable_topology = md_topology.subset(save_indices) reporter = OpenMMReporter(directory_name, trajectory_prefix, saveable_topology, subset_indices = save_indices) #make an integrator integrator = Integrator(**integrator_kwargs) #make a propagator if direction == 'forward': propagator = Propagator(openmm_pdf_state = pdf_state, openmm_pdf_state_subset = pdf_state_subset, subset_indices_map = subset_indices_map, integrator = integrator, ani_handler = ani_handler, context_cache=None, reassign_velocities=True, n_restart_attempts=0, reporter = reporter, write_trajectory_interval = write_trajectory_interval) else: propagator = BackwardPropagator(openmm_pdf_state = pdf_state, openmm_pdf_state_subset = pdf_state_subset, subset_indices_map = subset_indices_map, integrator = integrator, ani_handler = ani_handler, context_cache=None, reassign_velocities=True, n_restart_attempts=0, reporter = reporter, write_trajectory_interval = write_trajectory_interval) particle = Particle(0) particle_state = OpenMMParticleState(positions = positions, box_vectors = box_vectors) particle.update_state(particle_state) particle_state, _return_dict = propagator.apply(particle_state, n_steps = steps_per_application, reset_integrator=True, apply_pdf_to_context=True) if box_vectors is None: particle_state.box_vectors=None return particle_state, np.array(propagator.state_works[0])
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 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_equilibrium( equilibrium_result: EquilibriumResult, thermodynamic_state: states.ThermodynamicState, nsteps_equil: int, topology: md.Topology, n_iterations: int, atom_indices_to_save: List[int] = None, trajectory_filename: str = None, splitting: str = "V R O R V", timestep: unit.Quantity = 1.0 * unit.femtoseconds ) -> EquilibriumResult: """ Run nsteps of equilibrium sampling at the specified thermodynamic state and return the final sampler state as well as a trajectory of the positions after each application of an MCMove. This means that if the MCMove is configured to run 1000 steps of dynamics, and n_iterations is 100, there will be 100 frames in the resulting trajectory; these are the result of 100,000 steps (1000*100) of dynamics. Parameters ---------- equilibrium_result : EquilibriumResult EquilibriumResult namedtuple containing the information necessary to resume thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state (including context parameters) that should be used nsteps_equil : int The number of equilibrium steps that a move should make when apply is called topology : mdtraj.Topology an MDTraj topology object used to construct the trajectory n_iterations : int The number of times to apply the move. Note that this is not the number of steps of dynamics; it is n_iterations*n_steps (which is set in the MCMove). splitting: str, default "V R O H R V" The splitting string for the dynamics atom_indices_to_save : list of int, default None list of indices to save (when excluding waters, for instance). If None, all indices are saved. trajectory_filename : str, optional, default None Full filepath of trajectory files. If none, trajectory files are not written. splitting: str, default "V R O H R V" The splitting string for the dynamics Returns ------- equilibrium_result : EquilibriumResult Container namedtuple that has the SamplerState for resuming, an MDTraj trajectory, and the reduced potential of the final frame. """ sampler_state = equilibrium_result.sampler_state #get the atom indices we need to subset the topology and positions if atom_indices_to_save is None: atom_indices = list(range(topology.n_atoms)) subset_topology = topology else: subset_topology = topology.subset(atom_indices_to_save) atom_indices = atom_indices_to_save n_atoms = subset_topology.n_atoms #construct the MCMove: mc_move = mcmc.LangevinSplittingDynamicsMove(n_steps=nsteps_equil, splitting=splitting) mc_move.n_restart_attempts = 10 #create a numpy array for the trajectory trajectory_positions = np.zeros([n_iterations, n_atoms, 3]) trajectory_box_lengths = np.zeros([n_iterations, 3]) trajectory_box_angles = np.zeros([n_iterations, 3]) #loop through iterations and apply MCMove, then collect positions into numpy array for iteration in range(n_iterations): mc_move.apply(thermodynamic_state, sampler_state) trajectory_positions[iteration, :] = sampler_state.positions[ atom_indices, :].value_in_unit_system(unit.md_unit_system) #get the box lengths and angles a, b, c, alpha, beta, gamma = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles( *sampler_state.box_vectors) trajectory_box_lengths[iteration, :] = [a, b, c] trajectory_box_angles[iteration, :] = [alpha, beta, gamma] #construct trajectory object: trajectory = md.Trajectory(trajectory_positions, subset_topology, unitcell_lengths=trajectory_box_lengths, unitcell_angles=trajectory_box_angles) #get the reduced potential from the final frame for endpoint perturbations reduced_potential_final_frame = thermodynamic_state.reduced_potential( sampler_state) #construct equilibrium result object equilibrium_result = EquilibriumResult(sampler_state, reduced_potential_final_frame) #If there is a trajectory filename passed, write out the results here: if trajectory_filename is not None: write_equilibrium_trajectory(equilibrium_result, trajectory, trajectory_filename) return equilibrium_result