Accessing context during ReplicaExchange

[clarktemple03]

I am attempting to run Replica Exchange for a system and would ideally like to access the context during the course of the simulation run in order to extract the forces. However I cannot seem to grab the appropriate context to do so.

Below is a code snippet for a simple 2-bead harmonic oscillator I wrote as a test for replica exchange. The primary issue is that . To my understanding of the documentation when one calls the local_context_cache it requires both the thermodynamic state and the integrator that is being used otherwise it spits out the first context associated to the thermodynamic state. At the end of the script I try to obtain the integrators for the two thermodynamic states separately, but it still gives only the context of the first thermodynamic state in both cases. I think the issue comes in defining context_cache from move but it would seem if I had the context cache I would not need to find the context.

import numpy as np
import openmm as mm
import openmm.app as app
import os
from openmmtools import states, mcmc
from openmmtools.multistate import ReplicaExchangeSampler, ReplicaExchangeAnalyzer
from openmmtools import multistate
import openmmtools
from openmm.unit import dalton, kelvin, second, nanometer, kilojoule_per_mole, picosecond
import os
import os.path as osp

if osp.exists('output-1.nc'): os.system('rm output-1.nc')
# System creation
n_atoms = 2
mass = 1.0*dalton
temperature = 2.405*kelvin
timestep = 1e-15*second # (fs)
gamma =  0.5 # (1/ps)
spacing = 0.5
start_pos = np.zeros((n_atoms,3))
start_pos[0,0] = -0.5*spacing
start_pos[1,0] = 0.5*spacing
# 
system = mm.System()
for ix in range(0,n_atoms):
    system.addParticle(mass)
# Specify bonding connectivity
big_k = 500*kilojoule_per_mole/nanometer/nanometer
big_r0 = 0*spacing*nanometer
harmonic_force = mm.HarmonicBondForce()
harmonic_force.addBond(0,1,big_r0,big_k)
system.addForce(harmonic_force)
#

# Replica setup
integrator = mm.LangevinIntegrator(temperature,gamma,timestep)
# platform = mm.Platform.getPlatformByName("CUDA")
# platform_properties = {"DeviceIndex": "0", "Precision": "mixed"}

n_replicas = 2
T_min = 2.405 * kelvin
T_max = 5.440 * kelvin
dT = (T_max-T_min) / n_replicas
temperatures = [T_min + dT * i for i in range(n_replicas)]
value_temperatures = [T_min + (T_max-T_min)/n_replicas * i for i in range(n_replicas)]

sampler_states = list()
thermodynamic_states = list()

for temperature in temperatures:
    thermodynamic_state = openmmtools.states.ThermodynamicState(
        system=system, temperature=temperature
    )
    thermodynamic_states.append(thermodynamic_state)
    sampler_states.append(
        openmmtools.states.SamplerState(start_pos) 
    )

move = mcmc.GHMCMove(
    timestep = timestep,
    collision_rate = gamma/picosecond,
    n_steps = 5,
)

simulation = ReplicaExchangeSampler(mcmc_moves=move, number_of_iterations=1)
reporter = openmmtools.multistate.MultiStateReporter("output-1.nc", checkpoint_interval=1)
simulation.create(thermodynamic_states, sampler_states, reporter)
simulation.minimize()
simulation.run()

def check_compute_forces(coords):
    F = np.zeros(3)
    for ix in range(3):
        F[ix] = -big_k.value_in_unit(kilojoule_per_mole/nanometer/nanometer) * (coords[1][ix]-coords[0][ix]).value_in_unit(nanometer)
    return F

for step in range(5):
    print(step)
    state = reporter.read_last_iteration(last_checkpoint=False)
    for ix in range(n_replicas):
        print("thermodynamic state {}".format(ix))
        r0 = reporter.read_sampler_states(state)[ix].positions
        # r1 = reporter.read_sampler_states(state)[ix].positions[1]
        s0 = simulation.sampler_states[ix].positions
        # s1 = simulation.sampler_states[ix].positions[1]
        print('coords from simulation sampler')
        print(s0)
        print('coords from reporter')
        print(r0)
        #
        local_context_cache = move._get_context_cache(None)
        integrator = move._get_integrator(thermodynamic_states[ix])
        context,integrator = local_context_cache.get_context(thermodynamic_states[ix],integrator)
        coords = context.getState(getPositions=True).getPositions(asNumpy=True)
        forces = context.getState(getForces=True).getForces(asNumpy=True)
        print('coords from context')
        print(coords)
        print('force from context')
        print(forces)
        print(check_compute_forces(coords))
        print('')


    simulation.extend(1)

This produces the following output which is only extracting the coordinates from the second thermodynamic state

0
thermodynamic state 0
coords from simulation sampler
[[ 2.99691269e-03 -1.71080137e-05 -2.60979377e-05]
 [-3.01499991e-03 -2.66069910e-05  7.42300617e-05]] nm
coords from reporter
[[ 2.99691269e-03 -1.71080137e-05 -2.60979377e-05]
 [-3.01499991e-03 -2.66069910e-05  7.42300617e-05]] nm
coords from context
[[-6.26668334e-04  2.88372285e-05 -8.46185067e-05]
 [ 5.67112525e-04  1.84507953e-05 -2.63861348e-05]] nm
force from context
[[ 0.59689045 -0.00519322  0.02911618]
 [-0.59689045  0.00519322 -0.02911618]] kJ/(nm mol)
[-0.59689043  0.00519322 -0.02911619]

thermodynamic state 1
coords from simulation sampler
[[-6.26668334e-04  2.88372285e-05 -8.46185067e-05]
 [ 5.67112525e-04  1.84507953e-05 -2.63861348e-05]] nm
coords from reporter
[[-6.26668334e-04  2.88372285e-05 -8.46185067e-05]
 [ 5.67112525e-04  1.84507953e-05 -2.63861348e-05]] nm
coords from context
[[-6.26668334e-04  2.88372285e-05 -8.46185067e-05]
 [ 5.67112525e-04  1.84507953e-05 -2.63861348e-05]] nm
force from context
[[ 0.59689045 -0.00519322  0.02911618]
 [-0.59689045  0.00519322 -0.02911618]] kJ/(nm mol)
[-0.59689043  0.00519322 -0.02911619]

1
thermodynamic state 0
coords from simulation sampler
[[ 1.86583027e-03 -4.17024785e-05 -8.15357707e-05]
 [-1.85576838e-03 -1.86300713e-05  1.55550500e-04]] nm
coords from reporter
[[ 1.86583027e-03 -4.17024785e-05 -8.15357707e-05]
 [-1.85576838e-03 -1.86300713e-05  1.55550500e-04]] nm
coords from context
[[-3.19459150e-03  7.81324561e-05 -2.17608453e-04]
 [ 3.09490645e-03  5.72364625e-05 -1.90487735e-05]] nm
force from context
[[ 3.14474893 -0.010448    0.09927984]
 [-3.14474893  0.010448   -0.09927984]] kJ/(nm mol)
[-3.14474897  0.010448   -0.09927984]

thermodynamic state 1
coords from simulation sampler
[[-3.19459150e-03  7.81324561e-05 -2.17608453e-04]
 [ 3.09490645e-03  5.72364625e-05 -1.90487735e-05]] nm
coords from reporter
[[-3.19459150e-03  7.81324561e-05 -2.17608453e-04]
 [ 3.09490645e-03  5.72364625e-05 -1.90487735e-05]] nm
coords from context
[[-3.19459150e-03  7.81324561e-05 -2.17608453e-04]
 [ 3.09490645e-03  5.72364625e-05 -1.90487735e-05]] nm
force from context
[[ 3.14474893 -0.010448    0.09927984]
 [-3.14474893  0.010448   -0.09927984]] kJ/(nm mol)
[-3.14474897  0.010448   -0.09927984]

Any help would be appreciated!

[ijpulidos]

Thanks for the feedback. We would need to dig into this, but as far as I know there is no easy way to access the replica positions directly from the python objects. I know you should be able to do this extracting the information from the checkpoint file, which is not ideal.

[clarktemple03]

Shouldn't the context be available though at some point since the simulations need to be updated in time? Couldn't the information be extracted if individual steps are taken and then querying