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 compute_potential_components(context, beta=beta, platform=DEFAULT_PLATFORM): """ Compute potential energy, raising an exception if it is not finite. Parameters ---------- context : simtk.openmm.Context The context from which to extract, System, parameters, and positions. """ # Make a deep copy of the system. import copy from perses.dispersed.utils import configure_platform platform = configure_platform(platform.getName(), fallback_platform_name='Reference', precision='double') system = context.getSystem() system = copy.deepcopy(system) # Get positions. positions = context.getState(getPositions=True).getPositions(asNumpy=True) # Get Parameters parameters = context.getParameters() # Segregate forces. for index in range(system.getNumForces()): force = system.getForce(index) force.setForceGroup(index) # Create new Context. integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds) context = openmm.Context(system, integrator, platform) context.setPositions(positions) for (parameter, value) in parameters.items(): context.setParameter(parameter, value) energy_components = list() for index in range(system.getNumForces()): force = system.getForce(index) forcename = force.__class__.__name__ groups = 1 << index potential = beta * context.getState( getEnergy=True, groups=groups).getPotentialEnergy() energy_components.append((forcename, potential)) del context, integrator return energy_components
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 __init__(self, openmm_pdf_state, integrator, context_cache=None, reassign_velocities=False, n_restart_attempts=0): """ call the BaseIntegratorMove init method arguments openmm_pdf_state : coddiwomple.openmm_adapters.OpenMMPDFState the pdf state of the propagator integrator : openmm.Integrator integrator of dynamics context_cache : openmmtools.cache.ContextCache, optional The ContextCache to use for Context creation. If None, the global cache openmmtools.cache.global_context_cache is used (default is None). reassign_velocities : bool, optional If True, the velocities will be reassigned from the Maxwell-Boltzmann distribution at the beginning of the move (default is False). n_restart_attempts : int, optional 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. attributes pdf_state : coddiwomple.openmm_adapters.OpenMMPDFState integrator : openmm.Integrator context_cache : openmmtools.cache.ContextCache reassign_velocities : bool n_restart_attempts : int or None """ super().__init__(n_steps = None, context_cache = context_cache, reassign_velocities = reassign_velocities, n_restart_attempts = n_restart_attempts) _logger.debug(f"successfully executed {BaseIntegratorMove.__class__.__name__} init.") import openmmtools.cache as cache from perses.dispersed.utils import check_platform, configure_platform from openmmtools.utils import get_fastest_platform try: cache.global_context_cache.platform = configure_platform(get_fastest_platform().getName()) except Exception as e: _logger.warning(f"platform configuration error: {e}") self.pdf_state = openmm_pdf_state # Check if we have to use the global cache. if self.context_cache is None: self._context_cache = cache.global_context_cache else: self._context_cache = self.context_cache # Create context and reset integrator for good measure self.context, self.integrator = self._context_cache.get_context(self.pdf_state, integrator) self.integrator.reset() _logger.debug(f"successfully equipped integrator: {self.integrator.__class__.__name__}") _logger.debug(f"integrator printable: {self.integrator.pretty_print()}")
############################################################################# # HYBRID SYSTEM SAMPLERS ############################################################################# from perses.annihilation.lambda_protocol import RelativeAlchemicalState, LambdaProtocol from openmmtools.multistate import sams, replicaexchange from openmmtools import cache, utils from perses.dispersed.utils import configure_platform from openmmtools import cache cache.global_context_cache.platform = configure_platform( utils.get_fastest_platform().getName()) from openmmtools.states import * from perses.dispersed.utils import create_endstates import numpy as np import copy import logging _logger = logging.getLogger() _logger.setLevel(logging.INFO) _logger = logging.getLogger("multistate") class HybridCompatibilityMixin(object): """ Mixin that allows the MultistateSampler to accommodate the situation where unsampled endpoints have a different number of degrees of freedom. """ def __init__(self, *args, hybrid_factory=None, **kwargs): self._hybrid_factory = hybrid_factory
def validate_endstate_energies(topology_proposal, htf, added_energy, subtracted_energy, beta=1.0 / kT, ENERGY_THRESHOLD=1e-6, platform=DEFAULT_PLATFORM, trajectory_directory=None): """ Function to validate that the difference between the nonalchemical versus alchemical state at lambda = 0,1 is equal to the difference in valence energy (forward and reverse). Parameters ---------- topology_proposal : perses.topology_proposal.TopologyProposal object top_proposal for relevant transformation htf : perses.new_relative.HybridTopologyFactory object hybrid top factory for setting alchemical hybrid states added_energy : float reduced added valence energy subtracted_energy: float reduced subtracted valence energy Returns ------- zero_state_energy_difference : float reduced potential difference of the nonalchemical and alchemical lambda = 0 state (corrected for valence energy). one_state_energy_difference : float reduced potential difference of the nonalchemical and alchemical lambda = 1 state (corrected for valence energy). """ import copy #import openmmtools.cache as cache #context_cache = cache.global_context_cache from perses.dispersed.utils import configure_platform from perses.utils import data platform = configure_platform(platform.getName(), fallback_platform_name='Reference', precision='double') #create copies of old/new systems and set the dispersion correction top_proposal = copy.deepcopy(topology_proposal) forces = { top_proposal._old_system.getForce(index).__class__.__name__: top_proposal._old_system.getForce(index) for index in range(top_proposal._old_system.getNumForces()) } forces['NonbondedForce'].setUseDispersionCorrection(False) forces = { top_proposal._new_system.getForce(index).__class__.__name__: top_proposal._new_system.getForce(index) for index in range(top_proposal._new_system.getNumForces()) } forces['NonbondedForce'].setUseDispersionCorrection(False) #create copy of hybrid system, define old and new positions, and turn off dispersion correction hybrid_system = copy.deepcopy(htf.hybrid_system) hybrid_system_n_forces = hybrid_system.getNumForces() for force_index in range(hybrid_system_n_forces): forcename = hybrid_system.getForce(force_index).__class__.__name__ if forcename == 'NonbondedForce': hybrid_system.getForce(force_index).setUseDispersionCorrection( False) old_positions, new_positions = htf._old_positions, htf._new_positions #generate endpoint thermostates 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 = [('real-old', nonalch_zero, old_positions, top_proposal._old_system.getDefaultPeriodicBoxVectors()), ('hybrid-old', alch_zero, htf._hybrid_positions, hybrid_system.getDefaultPeriodicBoxVectors()), ('hybrid-new', alch_one, htf._hybrid_positions, hybrid_system.getDefaultPeriodicBoxVectors()), ('real-new', nonalch_one, new_positions, top_proposal._new_system.getDefaultPeriodicBoxVectors())] rp_list = [] for (state_name, state, pos, box_vectors) in attrib_list: integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds) context = state.create_context(integrator, platform) samplerstate = states.SamplerState(positions=pos, box_vectors=box_vectors) samplerstate.apply_to_context(context) rp = state.reduced_potential(context) rp_list.append(rp) energy_comps = compute_potential_components(context) for name, force in energy_comps: print("\t\t\t{}: {}".format(name, force)) _logger.debug( f'added forces:{sum([energy for name, energy in energy_comps])}') _logger.debug(f'rp: {rp}') if trajectory_directory is not None: _logger.info( f'Saving {state_name} state xml to {trajectory_directory}/{state_name}-state.gz' ) state = context.getState(getPositions=True, getVelocities=True, getForces=True, getEnergy=True, getParameters=True) data.serialize(state, f'{trajectory_directory}-{state_name}-state.gz') del context, integrator nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[ 0], rp_list[1], rp_list[2], rp_list[3] ratio = abs((nonalch_zero_rp - alch_zero_rp + added_energy) / (nonalch_zero_rp + alch_zero_rp + added_energy)) assert ratio < ENERGY_THRESHOLD, f"The ratio in energy difference for the ZERO state is {ratio}.\n This is greater than the threshold of {ENERGY_THRESHOLD}.\n real-zero: {nonalch_zero_rp} \n alc-zero: {alch_zero_rp} \nadded-valence: {added_energy}" ratio = abs((nonalch_one_rp - alch_one_rp + subtracted_energy) / (nonalch_one_rp + alch_one_rp + subtracted_energy)) assert ratio < ENERGY_THRESHOLD, f"The ratio in energy difference for the ONE state is {ratio}.\n This is greater than the threshold of {ENERGY_THRESHOLD}.\n real-one: {nonalch_one_rp} \n alc-one: {alch_one_rp} \nsubtracted-valence: {subtracted_energy}" return abs(nonalch_zero_rp - alch_zero_rp + added_energy), abs(nonalch_one_rp - alch_one_rp + subtracted_energy)