Exemple #1
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.

    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)
Exemple #2
0
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)
Exemple #3
0
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
Exemple #4
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)
    """
    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)
Exemple #5
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)
Exemple #6
0
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
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_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
Exemple #9
0
    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
        ]
Exemple #10
0
    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]
Exemple #11
0
def run_endpoint_perturbation(lambda_thermodynamic_state,
                              nonalchemical_thermodynamic_state,
                              initial_hybrid_sampler_state,
                              mc_move,
                              n_iterations,
                              factory,
                              lambda_index=0):
    """

    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

    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.
    """
    #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

    #initialize work array
    w = np.zeros([n_iterations])

    #run n_iterations of the endpoint perturbation:
    for iteration in range(n_iterations):
        mc_move.apply(lambda_thermodynamic_state, new_sampler_state)

        #compute the reduced potential at the new 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)

        #generate 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)

        #compute the reduced potential at the nonalchemical system as well:
        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)

        w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential

    [t0, g, Neff_max] = timeseries.detectEquilibration(w)
    print(Neff_max)
    w_burned_in = w[t0:]

    [df, ddf] = pymbar.EXP(w_burned_in)
    ddf_corrected = ddf * np.sqrt(g)

    return [df, ddf_corrected, Neff_max]
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