Demultiplexing RepEx trajectories

[dlukauskis]

Hi, I'd like to use openmmtools's replica exchange code, but I can't see a way of actually demultiplexing the trajectories and having a look at them. Am I missing something? The docs don't address this point.

Thanks,
Dom

[andrrizzi]

Hi @dlukauskis . The Repex facilities were ported from YANK, but unfortunately we forgot to carry over the utilities to handle the trajectories in the netcdf file. You can take a look at this function https://github.com/choderalab/yank/blob/master/Yank/analyze.py#L1036-L1202 .

[dlukauskis]

I gave the YANK code a try but I'm getting a metadata-related error:

  File "sim_extract_trj.py", line 32, in <module>
    replica_index=0)
  File "/home/dom/Projects/Playground/SAMS/demux_test/extract_trajectory.py", line 54, in extract_trajectory
    reference_system = mmtools.utils.deserialize(metadata['reference_state']).system
KeyError: 'reference_state'

Attaching a directory with the exact code I used - demux_test.zip

[andrrizzi]

Yes, that's because YANK stores some metadata with the reference system and the topology. You will have to modify those bits and provide your own OpenMM Topology object.

[simomarsili]

This should solve the problem (until extract_trajectory will be ported to opemmtools):

https://github.com/simomarsili/mmdemux

[andrrizzi]

Thanks for sharing @simomarsili !

[tanmoy7989]

@andrrizzi, Is demultiplexing now included in MultistateAnalyzer? I'm still fairly new to the openmm environment (and openmmtools as well), so didn't go through the MultiStateAnalyzer methods in detail before asking.

[schuhmc]

Correct me if I'm wrong, but I also could not find a way to extract the trajectories. The solution by @simomarsili should work, but needs yank which does not support Python 3.10 yet (see choderalab/yank#1286 ). The fix to get the script working on python 3.10 is relatively straight-forward: change 'collections.MutableMapping' to 'collections.abc.MutableMapping' here: https://github.com/choderalab/yank/blob/c06059045bcf86d610f2e39c6db3944994b9f392/Yank/utils.py#L279. But it seems to me that yank has been abandoned so I'm not sure if it makes sense to rely on it.

Anyways, below is a a short script to get out all of the stored frames from a replica exchange simulation. It is not very fast but arbitrary reporters can be added to it as well.

import numpy as np
from openmm import unit, XmlSerializer, app
import openmm as mm
from tqdm import tqdm
import copy

from openmmtools import states, multistate

with open('system.xml', 'r') as f:
    system = XmlSerializer.deserialize(f.read())

pdb = app.PDBFile('eq.pdb')
integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 2*unit.femtosecond)

reporter = multistate.MultiStateReporter('out.nc', 'r')
analyzer = multistate.ReplicaExchangeAnalyzer(reporter)

platform = mm.Platform.getPlatformByName('Reference')
simulation = mm.app.Simulation(topology=pdb.topology,
                               system=system,
                               integrator=integrator,
                               platform=platform)

td_states = reporter.read_replica_thermodynamic_states()

simulations = [copy.copy(simulation) for i in range(reporter.n_states)]

for i in range(reporter.n_states):
    simulations[i].reporters = [mm.app.DCDReporter(f'test_traj_{i}.dcd', reportInterval=1, append=False)]

print('Getting coordinates...')
for nstep in tqdm(reporter.read_checkpoint_iterations()[:10]):
    sampler_state = reporter.read_sampler_states(nstep)

    for nstate in range(reporter.n_states):
        state_index = np.where(td_states[nstep] == nstate)[0][0] # get the index of the replica that is currently in the desired state
        positions = sampler_state[state_index].positions
        simulations[nstate].context.setPositions(positions)
        state = simulations[nstate].context.getState(getPositions=True)
        [r.report(simulations[nstate],state) for r in simulations[nstate].reporters]

I would be very interested in a direct openmmtools implementation though.