exploding solvate when using the openmm-alchemy tools

[leo hall]

Hi all,
hopefully this is the right place to ask about this,
I'm trying to understand how to use the openmm-tools library to do TI and solvate a glycine in water, but as soon as lambda is slightly low (0.8) the glycine coordinate explode all over the place. I seems like all intra-molecular forces are getting turned off rather than just electrostatics and LJ with the rest of the system ?
I've tried gradually turning down lambda during equilibration to see if that help but it doesn't, here's my code bellow, honestly any insight would be helpful ^^
I'm also open to other methods of doing this with base openmm functions...

Code: Select all

#always include at the top of simulation
import openmm.app as app
import openmm as mm
import openmm.unit as unit
import os
os.environ['JAX_ENABLE_X64'] = '1' #set JAX to 64bits by default
import openmmtools as tools
from openmmtools import integrators
from openmmtools import forces #still under construction, check for updates occasionally
from openmmtools import alchemy
from sys import stdout
from time import perf_counter as pf
import MDAnalysis as mda
import numpy as np

start = pf()

pdb = app.PDBFile('ZGY_1p5nmPadWaterCube.pdb')
forcefield = app.ForceField('./amber99sb-ZGY-hack.xml','/home/leo/.conda/envs/openmmEnv/lib/python3.10/site-packages/openmm/app/data/tip3p.xml')
#forcefield = app.ForceField('/home/hall2401/ForceFields/amber99sb-ZGY-hack.xml','/home/hall2401/ForceFields/tip3p.xml')

# System Configuration
nonbondedMethod = app.PME
nonbondedCutoff = 1.2*unit.nanometers
ewaldErrorTolerance = 0.0005
constraints = None
lambda_value = 0.8

# Integration Options
dt = 0.0005*unit.picoseconds # 0.5 femtoseconds
temperature = 298*unit.kelvin 
friction = 1.0/unit.picosecond
pressure = 1.0*unit.atmospheres
barostatInterval = 25

# Simulation Options
steps = 20000#00 #1ns with the 0.5fs timestep
equilibrationSteps = 10000#000 #0.5ns equilibration
CPUplatform = mm.Platform.getPlatformByName('CPU')
GPUplatform = mm.Platform.getPlatformByName('CUDA')
GPUplatformProperties = {'Precision': 'single'}
pdbReporter = app.PDBReporter('trajectory-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.pdb', 1000)
dataReporter = app.StateDataReporter('trajectory-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.csv', 1000, totalSteps=steps,
    step=True, speed=True, progress=True, potentialEnergy=True, temperature=True, separator=',')
checkpointReporter = app.CheckpointReporter('checkpoint-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.chk', 10000)

# Create the system
print('Building system...')
topology = pdb.topology
positions = pdb.positions
system = forcefield.createSystem(topology=topology, nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
    constraints=constraints, ewaldErrorTolerance=ewaldErrorTolerance)

#Get residue of interest
target_residue_index = 0  # Index of the target glycine residue
target_atoms = [atom.index for atom in pdb.topology.atoms() if atom.residue.index == target_residue_index]

#create system where target atoms are alchemically modifiable
alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=target_atoms)
factory = alchemy.AbsoluteAlchemicalFactory()
alchemical_system = factory.create_alchemical_system(system,alchemical_region)

integrator = integrators.LangevinIntegrator(temperature, friction, dt)
#integrator = mm.LangevinMiddleIntegrator(temperature, friction, dt)
simulation = app.Simulation(topology, alchemical_system, integrator, GPUplatform, GPUplatformProperties)
simulation.context.setPositions(positions)


print("performing energy minimisation...")
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(temperature)

print("equilibrating system...")
simulation.reporters.append(dataReporter)
simulation.step(equilibrationSteps)

print("applying alchemical transformation")
#apply the Lambda parametre to the system
alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
alchemical_state.lambda_electrostatics = 0.9
alchemical_state.lambda_sterics = 0.9
#alchemical_state.set_alchemical_parameters(lambda_value)
alchemical_state.apply_to_context(simulation.context)

print("equilibrating system again...")
#simulation.reporters.append(dataReporter)
simulation.currentStep = 0
simulation.step(equilibrationSteps)

print("applying alchemical transformation")
#apply the Lambda parametre to the system
alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
alchemical_state.lambda_electrostatics = lambda_value
alchemical_state.lambda_sterics = lambda_value
#alchemical_state.set_alchemical_parameters(lambda_value)
alchemical_state.apply_to_context(simulation.context)

# Simulate
print('Simulating...')
simulation.reporters.append(pdbReporter)
#simulation.reporters.append(dataReporter) #already added during equilibration
simulation.reporters.append(checkpointReporter)
simulation.currentStep = 0
simulation.step(steps)

end = pf()
print('runtime (mn) = ',(end-start)/60)

[Peter Eastman]

I don't know if any of the openmmtools developers monitor this forum. It might be best to open an issue at https://github.com/choderalab/openmmtools.

[leo hall]

Will do, thanks !

If that doesn't work, is there a way to do that directly through base openmm tools ? I'd tried using a custom-force to treat part of the system but got an error where exceptions in the non-bonded force and customnonbonded force had to be the same, even if I just used exclusions
should I post my code in this thread or open a new one for that ?

[Peter Eastman]

Sure, you can post your code here.

[leo hall]

Hi !
I gave it a few days to see if I got a response from the github and no so here goes.
I started off trying to create a scalable custom non bonded force for the residue I want to solvate, and creating an exception in the default nonbonded force for the residue's atoms.
I got an "forces must have the same exceptions" error, which I dealt with by trying to only have exclusions in the customnonbonded force, but still got the error.
I played around with openmmtools for a bit, figured meybe the issue is because I want to use PME in the nonbondedforde but that's not supported in customnonbondedForce so there's an interaction issue ?
Then I changed my code back to not using the alchemy package, ran it, and it worked fine.

Given I don't understand what changed, I'd be very grateful if you could take a brief look at my code to see if it makes sense before I run multiple calculations with it and other people in my lab don't do this type of calculations so they wouldn't necessarily know if it's wrong. I'm thinking maybe the issue was I wasn't renaming the files as I was going and didn't comment out the checkpointReporter, so I was getting outdated info from previous calculations ?

I've tried to comment it as thoroughly as I could

Code: Select all

#always include at the top of simulation
import openmm.app as app
import openmm as mm
import openmm.unit as unit
import os
os.environ['JAX_ENABLE_X64'] = '1' #set JAX to 64bits by default
import openmmtools as tools
from openmmtools import integrators
from openmmtools import forces #still under construction, check for updates occasionally
from openmmtools import alchemy
from sys import stdout
from time import perf_counter as pf
import MDAnalysis as mda
import numpy as np

start = pf()
LAMBDA = 10 #LAMBDA
TEMP = 298#TEMP
pdb = app.PDBFile('ZGY_1p5nmPadWaterCube.pdb')
forcefield = app.ForceField('./amber99sb-ZGY-hack.xml','/home/leo/.conda/envs/openmmEnv/lib/python3.10/site-packages/openmm/app/data/tip3p.xml')
#forcefield = app.ForceField('/home/hall2401/ForceFields/amber99sb-ZGY-hack.xml','/home/hall2401/ForceFields/tip3p.xml')

# System Configuration
nonbondedMethod = app.PME
nonbondedCutoff = 1.2*unit.nanometers
ewaldErrorTolerance = 0.0005
constraints = None
lambda_value = LAMBDA/100 #LAMBDA will be given as a percentage for purpose of filenames

# Integration Options
dt = 0.0005*unit.picoseconds # 0.5 femtoseconds
temperature = TEMP*unit.kelvin 
friction = 1.0/unit.picosecond
pressure = 1.0*unit.atmospheres
barostatInterval = 25

# Simulation Options
steps = 20000#00 #1ns with the 0.5fs timestep
equilibrationSteps = 10000#000 #0.5ns equilibration
CPUplatform = mm.Platform.getPlatformByName('CPU')
GPUplatform = mm.Platform.getPlatformByName('CUDA')
GPUplatformProperties = {'Precision': 'single'}
pdbReporter = app.PDBReporter('trajectory-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.pdb', 1000)
dataReporter = app.StateDataReporter('trajectory-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.csv', 1000, totalSteps=steps,
    step=True, speed=True, progress=True, potentialEnergy=True, temperature=True, separator=',')
checkpointReporter = app.CheckpointReporter('checkpoint-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.chk', 10000)

# Create the system
print('Building system...')
topology = pdb.topology
positions = pdb.positions
system = forcefield.createSystem(topology=topology, nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
    constraints=constraints, ewaldErrorTolerance=ewaldErrorTolerance)

#Get residue of interest
target_residue_index = 0  # Index of the target glycine residue
target_atoms = [atom.index for atom in pdb.topology.atoms() if atom.residue.index == target_residue_index]

integrator = mm.LangevinMiddleIntegrator(temperature, friction, dt)
simulation = app.Simulation(topology, system, integrator, GPUplatform, GPUplatformProperties)

simulation.context.setPositions(positions)


# Get the NonbondedForce object
#we get the first element of a list that presumably contains just one element as there should be only one nonbonded force object
nonbonded_force = [force for force in system.getForces() if isinstance(force, mm.NonbondedForce)][0] 
    
# Lambda parametre already defined
#dielectric constant is not considered as we're doing an explicit solvent simulation
# Create a CustomNonbondedForce that scales both Lennard-Jones and Coulombic interactions
custom_force = mm.CustomNonbondedForce(
    'lambda*(4*epsilon*((sigma/r)^12 - (sigma/r)^6) + 138.935456*charge1*charge2/(1/r)); ' 
    'sigma=0.5*(sigma1+sigma2); epsilon=sqrt(epsilon1*epsilon2)')
#138.9... is the value of the Coulomb constant in Openmm in kJ/nm
custom_force.addGlobalParameter('lambda', lambda_value)
#charge, sigma and epsilon are named parametres whose value will be input for each relevant particle
custom_force.addPerParticleParameter('charge')
custom_force.addPerParticleParameter('sigma')
custom_force.addPerParticleParameter('epsilon')
#give it the same parametres as other nonbonded forces
custom_force.setCutoffDistance(nonbondedCutoff)
custom_force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)

# Add particles to the custom force using parameters from the existing force field
#the whole system must be included in the custom force because the residue can interact with the whole system
for i in range(system.getNumParticles()):
    charge, sigma, epsilon = nonbonded_force.getParticleParameters(i)
    custom_force.addParticle([charge, sigma, epsilon])
    
    # Enable lambda scaling for interactions involving the target glycine residue
    if i in target_atoms:
        #setParticleParametre means interactions will be calculated for this particle
        custom_force.setParticleParameters(i, [charge, sigma, epsilon])
        nonbonded_force.setParticleParameters(i, 0, 0, 0) #remove target atoms from the nonbonded force

# Add exclusions to avoid double-counting interactions that are already included in bonded interactions
#the residue presumably only interacts with itself through bonded interactions
# Extract bonds from the topology and convert to the required format
bonds = [(bond[0].index, bond[1].index) for bond in pdb.topology.bonds()]

# Add exclusions for 1-2 interactions (atoms directly bonded to each other - this is what the second argument is for)
custom_force.createExclusionsFromBonds(bonds, 1)

# Add the custom force to the system
system.addForce(custom_force)

# Disable interactions in the default NonbondedForce for target atoms
#for atom_index in target_atoms:
 #   for other_index in range(system.getNumParticles()):
        #if other_index not in target_atoms: #interactions between atoms from the target should also be excluded as they're taken care of by the custom force
  #      nonbonded_force.addException(atom_index, other_index, 0.0, 0.0, 0.0,replace=False) #there might be exceptions already for bonded atoms




print("performing energy minimisation...")
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(temperature)

print("equilibrating system...")
simulation.reporters.append(dataReporter)
simulation.step(equilibrationSteps)


# Simulate
print('Simulating...')
simulation.reporters.append(pdbReporter)
#simulation.reporters.append(dataReporter) #already added during equilibration
simulation.reporters.append(checkpointReporter)
simulation.currentStep = 0
simulation.step(steps)

end = pf()
print('runtime (mn) = ',(end-start)/60)

[Peter Eastman]

When your system includes multiple nonbonded forces (NonbondedForce, CustomNonbondedForce, CustomGBForce), they need to exclude exactly the same interactions. That's because they all get computed together in a single GPU kernel based on a single neighbor list. If one of them excludes an interaction, all of them need to. The easiest way to make sure that's true is to copy the exclusions from the nonbonded force.

Code: Select all

for i in range(nonbonded_force.getNumExceptions()):
    params = nonbonded_force.getExceptionParameters(i)
    custom_force.addExclusion(params[0], params[1])

NonbondedForce calls them "exceptions" rather than "exclusions" because it includes the ability to compute them with nonstandard parameters that don't follow the regular combination rules. As far as the nonbonded kernel is concerned, though, they're completely excluded. You can think of NonbondedForce as really being two forces: a nonbonded force that excludes all of the exceptions, plus a bonded force that computes them with different parameters.

NonbondedForce can do that because it's a hardcoded functional form, and it knows how the combination rules work. CustomNonbondedForce can't because it's generic. The combination rules are built into the expression you provide, and it doesn't know how to separate them out. But you can add a CustomBondForce to compute the interactions with nonstandard parameters, if you want.

[leo hall]

ok but here I exclude the residue atoms from the nonBondedForce by doing

Code: Select all

nonbonded_force.setParticleParameters(i, 0, 0, 0) #remove target atoms from the nonbonded force

for i in the target molecule.

If I then do what you suggest, aren't I also going to exclude those from the CustomNonBondedForce that I'm creating specifically to calcuiate them ?

knowing that the code I pasted in this morning's message actually runs without that, I'm just not quite sure why it doesn't give the error and so I don't 100% know I can trust it

[Peter Eastman]

Setting parameters to 0 is not the same as an exception. Exceptions are used for 1-2, 1-3, and 1-4 interactions. They tell the nonbonded kernel to skip calculating an interaction. Setting the parameters for a particle to 0 doesn't cause anything to be skipped. It just guarantees that all interactions involving that particle will end up with an energy of 0.

The NonbondedForce and CustomNonbondedForce need to skip the same interactions. But it's fine for an interaction to end up with an energy of 0 in one of them and not in the other.

[leo hall]

ok, I must've set an exception in the first itteration of the code and then missunderstood the API.

I don't know why this code previously brought up the exceptions error, I guess I didn't comment a line I should've or something stupid...

thanks a lot for the help !