corinwagen
commented
3 years ago I've been seeing the exact same error, using a modified version of the ReplicaExchange class from this tutorial. I tried increasing n_restart_attempts to 100 but it didn't help at all!
Here's my (pretty bad) code, almost entirely copied from the tutorial:
import math, yaml, sys, cctk, os
from datetime import datetime
import numpy as np
from random import random, randint
import mdtraj
from simtk import openmm, unit
from simtk.openmm.app import AmberPrmtopFile, AmberInpcrdFile, PME
from openmmtools import testsystems, integrators, forces, alchemy, states, cache, mcmc
import cctk.helper_functions as helper
# usage: python run_rest2.py config.yaml
class ReplicaExchange:
def __init__(self, thermodynamic_states, sampler_states, mcmc_move, save_interval=1000, swap_interval=1000, total_steps=10000, logfile="rest2.log", traj_prefix="rest2"):
self._thermodynamic_states = thermodynamic_states
self._replicas_sampler_states = sampler_states
self._mcmc_move = mcmc_move
self.save_interval = save_interval
self.swap_interval = swap_interval
self.total_steps = total_steps
self.current_step = 0
self.logfile = logfile
self.traj_prefix = traj_prefix
# initial minimization
self.log("beginning energy minimization")
for thermo_state, sampler_state in zip(self._thermodynamic_states, self._replicas_sampler_states):
context, integrator = cache.global_context_cache.get_context(thermo_state)
sampler_state.apply_to_context(context)
openmm.LocalEnergyMinimizer.minimize(context)
sampler_state.update_from_context(context)
def log(self, text):
with open(self.logfile, "a") as f:
time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
f.write(f"{time}\t{text}\n")
def run(self):
self.log("beginning replica exchange simulation")
for iteration in range(int(self.total_steps/self.save_interval)):
self._propagate_replicas()
self.current_step += self.save_interval
with open(self.logfile, "a") as f:
time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
f.write(f"{time}\t{self.current_step} steps completed\n")
self.save_positions()
if self.current_step % self.swap_interval == 0:
self._mix_replicas()
def _propagate_replicas(self):
for thermo_state, sampler_state in zip(self._thermodynamic_states, self._replicas_sampler_states):
self._mcmc_move.apply(thermo_state, sampler_state)
def _mix_replicas(self, n_attempts=5):
for i in range(len(self._thermodynamic_states) - 1):
j = i+1
sampler_state_i, sampler_state_j = (self._replicas_sampler_states[k] for k in [i, j])
thermo_state_i, thermo_state_j = (self._thermodynamic_states[k] for k in [i, j])
energy_ii = self._compute_reduced_potential(sampler_state_i, thermo_state_i)
energy_jj = self._compute_reduced_potential(sampler_state_j, thermo_state_j)
energy_ij = self._compute_reduced_potential(sampler_state_i, thermo_state_j)
energy_ji = self._compute_reduced_potential(sampler_state_j, thermo_state_i)
log_p_accept = - (energy_ij + energy_ji) + energy_ii + energy_jj
if log_p_accept >= 0.0 or random() < math.exp(log_p_accept):
self._thermodynamic_states[i] = thermo_state_j
self._thermodynamic_states[j] = thermo_state_i
with open(self.logfile, "a") as f:
f.write(f"\tswapping {i} and {j} ({log_p_accept:.3f})\n")
def _compute_reduced_potential(self, sampler_state, thermo_state):
context, integrator = cache.global_context_cache.get_context(thermo_state)
sampler_state.apply_to_context(context)
return thermo_state.reduced_potential(context)
def save_positions(self):
... # omitted for space
config_file = sys.argv[1]
with open(config_file, "r") as f:
config = yaml.safe_load(f)
assert config["swap_interval"] % config["save_interval"] == 0, "swap_interval must be a multiple of save_interval"
num_replicas = config["num_replicas"]
prmtop = AmberPrmtopFile(config["prmtop"])
inpcrd = AmberInpcrdFile(config["inpcrd"])
system = prmtop.createSystem(nonbondedMethod=PME, nonbondedCutoff=config["nonbonded_cutoff"] * unit.angstrom, constraints=None)
mdtop = mdtraj.Topology.from_openmm(prmtop.topology)
high_atoms = mdtop.select(f"resname {config['high_resname']}")
min_L = config["min_lambda"]
protocol = {
'temperature' : [config["temperature"]] * num_replicas * unit.kelvin,
'lambda_sterics' : list(np.geomspace(1.0, min_L, num_replicas)),
'lambda_electrostatics' : list(np.geomspace(1.0, min_L, num_replicas)),
'lambda_angles' : list(np.geomspace(1.0, min_L, num_replicas)),
'lambda_torsions' : list(np.geomspace(1.0, min_L, num_replicas)),
}
alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=high_atoms, alchemical_angles=True, alchemical_torsions=True)
factory = alchemy.AbsoluteAlchemicalFactory()
alchemical_system = factory.create_alchemical_system(system, alchemical_region)
alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
compound_states = states.create_thermodynamic_state_protocol(alchemical_system, protocol=protocol, composable_states=[alchemical_state])
sampler_states = [states.SamplerState(positions=inpcrd.getPositions(), box_vectors=inpcrd.getBoxVectors()) for _ in compound_states]
langevin_move = mcmc.LangevinSplittingDynamicsMove(timestep=0.5*unit.femtosecond, n_steps=config["swap_interval"])
langevin_move.n_restart_attempts = 100
parallel_tempering = ReplicaExchange(
compound_states,
sampler_states,
langevin_move,
total_steps=config["total_steps"],
swap_interval=config["swap_interval"],
save_interval=config["save_interval"],
logfile=config["logfile"],
traj_prefix=config["traj_prefix"],
)
parallel_tempering.run()Using the following input file, I only make it through one LangevinSplittingDynamicsMove per replica before dying. This is surprising because the LocalEnergyMinimizer runs without dying, so it doesn't seem like the system is just bogus/contains massive conflicts. Visually, everything looks fine (inputs generated using antechamber/parmchk/PACKMOL/tleap).
config.yaml:
prmtop: solvated_13precomplex.prmtop
inpcrd: solvated_13precomplex.inpcrd
logfile: 13precomplex.log
traj_prefix: replica
nonbonded_cutoff: 8 # Å
temperature: 298
high_resname: MOL
total_steps: 5000000
swap_interval: 2000
save_interval: 500
num_replicas: 24
min_lambda: 0.66Output:
Potential energy is NaN after 0 attempts of integration with move LangevinSplittingDynamicsMove Attempting a restart...
Potential energy is NaN after 1 attempts of integration with move LangevinSplittingDynamicsMove Attempting a restart...
Potential energy is NaN after 2 attempts of integration with move LangevinSplittingDynamicsMove Attempting a restart...
Potential energy is NaN after 3 attempts of integration with move LangevinSplittingDynamicsMove Attempting a restart...
...
Potential energy is NaN after 98 attempts of integration with move LangevinSplittingDynamicsMove Attempting a restart...
Potential energy is NaN after 99 attempts of integration with move LangevinSplittingDynamicsMove Trying to reinitialize Context as a last-resort restart attempt...
Potential energy is NaN after 100 attempts of integration with move LangevinSplittingDynamicsMove
Traceback (most recent call last):
File "run_rest2.py", line 142, in <module>
parallel_tempering.run()
File "run_rest2.py", line 44, in run
self._propagate_replicas()
File "run_rest2.py", line 56, in _propagate_replicas
self._mcmc_move.apply(thermo_state, sampler_state)
File "/n/home03/cwagen/.conda/envs/openmm/lib/python3.7/site-packages/openmmtools/mcmc.py", line 1114, in apply
super(LangevinDynamicsMove, self).apply(thermodynamic_state, sampler_state)
File "/n/home03/cwagen/.conda/envs/openmm/lib/python3.7/site-packages/openmmtools/mcmc.py", line 707, in apply
raise IntegratorMoveError(err_msg, self, context)
openmmtools.mcmc.IntegratorMoveError: Potential energy is NaN after 100 attempts of integration with move LangevinSplittingDynamicsMoveAny advice would be greatly appreciated!
mjw99
Dan-Burns
@dominicrufa and I were trying to run repex on the endstates of a capped amino acid transformation (ALA->THR) (using the new RepartitionedHybridTopologyFactory in perses) with softened electrostatics, sterics, and torsion. We observed that when we use the
LangevinSplittingDynamicsMove, the simulation runs fine at the ALA endstate, but nans immediately for the THR endstate. When we useGHMCMoveinstead, the simulation runs slower, but doesn't nan.In the example code below, I show that the nans occur when I apply the
LangevinSplittingDynamicsMove(i.e. without even intializing the repex sampler). Note that when I ran 1000 steps of MD at each of the 11 thermodynamic states using theLangevinIntegrator, no nans occurred.Here is the code:
Here is the output from the for loop in the last code snippet: