GCMC simulation explode with "Particle coordinate is nan" or "Energy is NaN"

[Hong-Rui]

Hi developers,

Recently we have tried the grand gcmc for sampling waters, and I installed the latest grand from the master branch and the OpenMM installed from conda with version 7.7.0.
I followed the example scripts in https://grand.readthedocs.io/en/latest/examples.scytalone.html but it does not work for my system.
The system is a urokinase with docked ligand. I've followed the tutorials to prepare the input files as following:

10302_inputs.tar.gz

The ligand pdb, prepi and prmtop files were generated by antechamber using GAFF2 from AmberTools 22, and the whole sovated pdb were generated by tleap.

The simulation scripts were like:

os.chdir('/root/Projects/GCMC_Test/gcmc_sampling_1')

ligand_prmtop_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302.prmtop'
ligand_prepi_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302.prepi'
ligand_pdb_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302.pdb'
ligand_xml_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302.xml'
solvated_pdb_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/system_solvated_10302.pdb'
solvated_pdb_conect_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/system_solvated_10302_conect.pdb'

gcmc_ghost_pdb_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302_gcmc_extra_waters.pdb'
gcmc_ghost_txt_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302_gcmc_extra_waters.txt'
gcmc_log_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302_gcmc.log'
gcmc_dcd_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302_gcmc.dcd'
gcmc_rst_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302_gcmc.rst7'
gcmc_out_file_name = '/root/Projects/GCMC_Test/gcmc_sampling_1/10302_gcmc.out'

# Write CONECT lines to PDB
grand.utils.write_conect(solvated_pdb_file_name, 'MOL', ligand_prepi_file_name, solvated_pdb_conect_file_name)

# Write ligand XML
grand.utils.create_ligand_xml(ligand_prmtop_file_name, ligand_prepi_file_name, 'MOL', ligand_xml_file_name)

system_pdb = PDBFile(solvated_pdb_conect_file_name)
system_pdb.topology, system_pdb.positions, ghosts = grand.utils.add_ghosts(system_pdb.topology, system_pdb.positions, n=10, pdb=gcmc_ghost_pdb_file_name)
system_ff = ForceField('amber14-all.xml', 'amber14/tip3p.xml', ligand_xml_file_name)
system = system_ff.createSystem(system_pdb.topology,
                                nonbondedMethod=PME,
                                nonbondedCutoff=12.0 * angstrom,
                                switchDistance=10.0 * angstrom,
                                constraints=HBonds,
                                rigidWater=True,
                                removeCMMotion=True,
                                hydrogenMass=None,
                                ewaldErrorTolerance=0.0005)

gcmc_reference_atom_list = []
for atom in system_pdb.topology.atoms():
    if atom.residue.name == 'MOL':
        atom_info_dict = {'name': atom.name, 'resname': atom.residue.name, 'resid': atom.residue.id}
        gcmc_reference_atom_list.append(atom_info_dict)

gcmc_mover = grand.samplers.StandardGCMCSphereSampler(system=system,
                                                      topology=system_pdb.topology,
                                                      temperature=298.15 * kelvin,
                                                      referenceAtoms=gcmc_reference_atom_list,
                                                      sphereRadius=4.2 * angstroms,
                                                      ghostFile=gcmc_ghost_txt_file_name,
                                                      log=gcmc_log_file_name,
                                                      dcd=gcmc_dcd_file_name,
                                                      rst=gcmc_rst_file_name,
                                                      overwrite=False)

integrator = LangevinIntegrator(298.15 * kelvin, 1.0 / picosecond, 0.002 * picoseconds)
platform = Platform.getPlatformByName('CUDA')
platform_property = {'Precision': 'mixed', 'DisablePmeStream': 'true', 'DeviceIndex': '0'}
simulation = Simulation(system_pdb.topology, system, integrator, platform, platform_property)
simulation.context.setPositions(system_pdb.positions)

simulation.context.setPeriodicBoxVectors(*system_pdb.topology.getPeriodicBoxVectors())
simulation.context.setVelocitiesToTemperature(298.15)

simulation. Step(250000)

# Prepare the GCMC sphere
gcmc_mover.initialise(simulation.context, ghosts)
gcmc_mover.deleteWatersInGCMCSphere()

print("Equilibration...")
for i in range(50):
    # Carry out 200 moves every 100 fs
    gcmc_mover.move(simulation.context, 200)
    simulation.step(50)
print("{}/{} equilibration GCMC moves accepted. N = {}".format(gcmc_mover.n_accepted,
                                                               gcmc_mover.n_moves,
                                                               gcmc_mover.N))



simulation. Reporters = []
simulation.context.setTime(0)
simulation.currentStep = 0

state_data_reporter = StateDataReporter(gcmc_out_file_name,
                                        100,
                                        step=True,
                                        time=True,
                                        potentialEnergy=True,
                                        temperature=True,
                                        kineticEnergy=True,
                                        totalEnergy=True,
                                        speed=True,
                                        separator='\t',
                                        progress=True,
                                        remainingTime=True,
                                        totalSteps=25000000)

simulation.reporters.append(state_data_reporter)

num_gcmc_moves = 100
total_n_steps = 25000
gcmc_n_steps = 1000
num_gcmc_round = int(total_n_steps / gcmc_n_steps)

gcmc_mover.reset()
for i in range(num_gcmc_round):
    # Carry out 100 GCMC moves per 2 ps of MD
    simulation. Step(gcmc_n_steps)
    gcmc_mover.move(simulation. Context, num_gcmc_moves)
    gcmc_mover.report(simulation)

The simulation explodes even when it's just equilibrating the system using GCMC. However, if I remove all gcmc steps, the simulation runs smoothly without errors. I've tried to do energy minimization before the gcmc moves, but it still does not work.

I've inspect the error, and found that the system explode during the InsertonMoves. Normally the simulation would run for like 1 or 2 minutes and then crash.

It's very helpful for be to hear your response and see how to debug this. Thanks!!

[Hong-Rui]

The error message is like

---------------------------------------------------------------------------
OpenMMException                           Traceback (most recent call last)
Cell In[5], line 25
     22 print("Equilibration...")
     23 for i in range(50):
     24     # Carry out 200 moves every 100 fs
---> 25     gcmc_mover.move(simulation.context, 200)
     26     simulation.step(50)
     27 print("{}/{} equilibration GCMC moves accepted. N = {}".format(gcmc_mover.n_accepted,
     28                                                                gcmc_mover.n_moves,
     29                                                                gcmc_mover.N))

File [~/Modules/miniconda3/envs/gcmc/lib/python3.9/site-packages/grand/samplers.py:1092](https://vscode-remote+ssh-002dremote-002bydcg1030401-002ebohrium-002etech.vscode-resource.vscode-cdn.net/root/Projects/GCMC_Test/~/Modules/miniconda3/envs/gcmc/lib/python3.9/site-packages/grand/samplers.py:1092), in StandardGCMCSphereSampler.move(self, context, n)
   1088 for i in range(n):
   1089     # Insert or delete a water, based on random choice
   1090     if np.random.randint(2) == 1:
   1091         # Attempt to insert a water
-> 1092         self.insertionMove()
   1093     else:
   1094         # Attempt to delete a water
   1095         self.deletionMove()

File [~/Modules/miniconda3/envs/gcmc/lib/python3.9/site-packages/grand/samplers.py:1112](https://vscode-remote+ssh-002dremote-002bydcg1030401-002ebohrium-002etech.vscode-resource.vscode-cdn.net/root/Projects/GCMC_Test/~/Modules/miniconda3/envs/gcmc/lib/python3.9/site-packages/grand/samplers.py:1112), in StandardGCMCSphereSampler.insertionMove(self)
   1109 print(atom_indices)
   1111 # Recouple this water
...
   8392     This method has several limitations. The only information it updates is the parameters of particles and exceptions. All other aspects of the Force (the nonbonded method, the cutoff distance, etc.) are unaffected and can only be changed by reinitializing the Context. Furthermore, only the chargeProd, sigma, and epsilon values of an exception can be changed; the pair of particles involved in the exception cannot change. Finally, this method cannot be used to add new particles or exceptions, only to change the parameters of existing ones.
   8393     """
-> 8394     return _openmm.NonbondedForce_updateParametersInContext(self, context)

OpenMMException: Particle coordinate is nan

[olliemelling]

Hi,

My initial thought was that using OpenMM 7.7 would be causing these problems as we had issues with OpenMM 7.5-7.7 when running on the CUDA platform with mixed precision.

However, when I try and simulate your system with no GCMC it is still crashing almost immediately. I have tried a number of usual fixes (minimising, reducing the time step etc) and have run in both OpenMM 7.7 and 8.0 but unable to simulate without the system exploding.

I would suggest using OpenMM 8.0.0 to see if that works first but if not then could you send me a script that you are using to simulate this system (just plain MD, no GCMC) and a .yml file for the environment you're using.

Thanks,
Ollie

[Hong-Rui]

Hello,
After a long time for debugging this, I discovered the problem that is causing this. The problem is that when we initialize the GCMC sphere, i.e. gcmc_mover.initialise(simulation. Context, ghosts) the context positions starts to explode, not even doing any gcmc trial.
And after deeper digging, I found the problem actually lies in adjustSpecificWater. Every further function that call adjustSpecificWater would cause the context positions vary a lot from the beginning. And the real cause of this is actually when we setParticleParameters for any nonbonded force or custom nb forces, the positions changes, while it should not be, since we are just modifying some forcefield parameters!

This should be an OpenMM bug and I think we should raise an issue on this. This bug is actually severe since lots of code use this such as openmmtools or perses doing FEP on OpenMM.

In the sampler.py , I just add a small workaround to bypass this:

    def adjustSpecificWater(self, atoms, new_lambda):
        """
        Adjust the coupling of a specific water molecule, by adjusting the lambda value

        Parameters
        ----------
        atoms : list
            List of the atom indices of the water to be adjusted
        new_lambda : float
            Value to set lambda to for this particle
        """
        # Get lambda values
        lambda_vdw, lambda_ele = get_lambda_values(new_lambda)

        state = self.context.getState(getPositions=True, enforcePeriodicBox=True, getEnergy=True)
        original_positions = state.getPositions(asNumpy=True)
        original_box_vectors = state.getPeriodicBoxVectors(asNumpy=True)
#        print(original_positions)
#        print(original_box_vectors)

        # Loop over parameters
        for i, atom_idx in enumerate(atoms):
            # Obtain original parameters
            atom_params = self.water_params[i]
            # Update charge in NonbondedForce
            self.nonbonded_force.setParticleParameters(atom_idx,
                                                       charge=(lambda_ele * atom_params["charge"]),
                                                       sigma=atom_params["sigma"],
                                                       epsilon=abs(0.0))
            # Update lambda in CustomNonbondedForce
            self.custom_nb_force.setParticleParameters(atom_idx,
                                                       [atom_params["sigma"], atom_params["epsilon"], lambda_vdw])

        # Update context with new parameters
        self.nonbonded_force.updateParametersInContext(self.context)
        self.custom_nb_force.updateParametersInContext(self.context)

#        self.context.reinitialize()
        self.context.setPositions(original_positions)
        self.context.setPeriodicBoxVectors(*original_box_vectors)

        state = self.context.getState(getPositions=True, enforcePeriodicBox=True, getEnergy=True)
 #       print(state.getPositions(asNumpy=True))
 #       print(state.getPeriodicBoxVectors(asNumpy=True))
 #       print(state.getPotentialEnergy())

        return None

Now the gcmc simulation works normally, and I'll raise a PR for this fix later.
Hope to hear any advice from you!

[olliemelling]

I've had a look into this - it does appear that the positions do change very slightly when the context parameters are updated, however these changes are happening on a level of precision that should not be affecting the simulation (approx. 1x10^-7 nm). These changes have not caused us any problems before while running GCMC and given they are within the limits of machine precision, I wouldn't say that they constitute a bug in OpenMM.

How large are the changes in positions you were seeing? If they are of a similar magnitude to what I saw then it is safe to ignore them.

Your system runs fine for me when using OpenMM 8.0.0 and the BAOABIntegrator (from openmmtools) - without requiring any changes to the positions after the parameters have been updated (which will impact performance).

I would recommend using OpenMM version 8.0.0 and also the BAOABIntegrator. We have previously noticed some instability issues with the LangevinIntegrator when perturbing parameters.

Let me know if this works for you.

[Hong-Rui]

okay...
In my test the position change after updating parameters is about 2 order of magnitude, and larger and larger from time to time.... The error accumulates and finally you get NaNs.

I haven't tried OpenMM 8.0.0 since I cannot install it from conda, I might need to compile it by myself and see what happens.
Thanks!

[Hong-Rui]

Hi,
I've just tried OpenMM 8.0.0, and it works. Certainly, there's something going on for 7.7 and this bug was fixed in 8.0.
In OpenMM 7.7 with BAOABIntegrator, after updating nb force parameters, the positions changes a least 2 order of magnitudes and this error would be accumulated until it gets NaN. In OpenMM 8.0 with BAOABIntegrator, the error reduced to about 1e-6.

The PR may not be needed but I think I'll still add it just to make sure the positions would not changes after updating parameters.

Thanks!