openmm / openmm-cookbook

The OpenMM Cookbook and Tutorials
39 stars 10 forks source link

Cookbook entry on computing interaction energies #26

Closed peastman closed 1 year ago

peastman commented 1 year ago

This comes up regularly. It seemed like it would make a good cookbook entry.

sef43 commented 1 year ago

This is a very useful example, it looks good!

I think we should also provide information on how you would compute the Coulomb and LJ interactions for an Amber style forcefield. Am I correct in thinking the Amber forcefields contain both terms in the NonBondedForce so the method described in the example would not work the same way?

peastman commented 1 year ago

It would work equally well for Amber. The NonbondedForce alone would give you the total interaction energy.

peastman commented 1 year ago

I added another paragraph explaining that.

spadavec commented 1 year ago

This is a great cookbook! I'm curious how you might breakout the vdw and Coulomb energies in a similar fashion using AMBER style forcefields?

peastman commented 1 year ago

That might be worth adding. The simplest way is probably to create two different parameters. Something like this.

force.addParticleParameterOffset('solute_coulomb_scale', i, charge, 0, 0)
force.addParticleParameterOffset('solute_lj_scale', i, 0, sigma, epsilon)
spadavec commented 1 year ago

Hi @peastman, thanks for the advice! I am getting some strange behaviors, though:

or force in system.getForces():
    print(i, force, force.getForceGroup())
    if isinstance(force, NonbondedForce):
        force.setForceGroup(0)

        force.addGlobalParameter("solute_coulomb_scale", 1)
        force.addGlobalParameter("solvent_coulomb_scale", 1)

        for i in range(force.getNumParticles()):
            charge, sigma, epsilon = force.getParticleParameters(i)
            # Set the parameters to be 0 when the corresponding parameter is 0,
            # and to have their normal values when it is 1.

            coulomb_param = "solute_coulomb_scale" if i in ligand else "solvent_coulomb_scale"
            force.setParticleParameters(i, 0, 0, 0)
            force.addParticleParameterOffset(
                coulomb_param, i, charge, sigma, epsilon)
        for i in range(force.getNumExceptions()):
            p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(
                i)
            force.setExceptionParameters(i, p1, p2, 0, 0, 0)
    else:
        force.setForceGroup(2)

Results in a nonbonded energy of apprx. -233 kJ/mol in a toy example, while the following:

for force in system.getForces():
    print(i, force, force.getForceGroup())
    if isinstance(force, NonbondedForce):
        force.setForceGroup(0)

        force.addGlobalParameter("solute_coulomb_scale", 1)
        force.addGlobalParameter("solvent_coulomb_scale", 1)

        force.addGlobalParameter("solute_lj_scale", 1)
        force.addGlobalParameter("solvent_lj_scale", 1)

        for i in range(force.getNumParticles()):
            charge, sigma, epsilon = force.getParticleParameters(i)
            # Set the parameters to be 0 when the corresponding parameter is 0,
            # and to have their normal values when it is 1.

            coulomb_param = "solute_coulomb_scale" if i in ligand else "solvent_coulomb_scale"
            lj_param = 'solute_lj_scale' if i in ligand else 'solvent_lj_scale'
            force.setParticleParameters(i, 0, 0, 0)
            force.addParticleParameterOffset(coulomb_param, i, charge, 0, 0)
            force.addParticleParameterOffset(lj_param, i, 0, sigma, epsilon)
        for i in range(force.getNumExceptions()):
            p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(
                i)
            force.setExceptionParameters(i, p1, p2, 0, 0, 0)
    else:
        force.setForceGroup(2)

Has the coulomb energy at apprx. -2135 kJ/mol and the LJ at apprx. 17915 kJ/mol.

Any ideas for the discrepancy?

peastman commented 1 year ago

Can you show your complete code?

spadavec commented 1 year ago

Sure thing:

import sys
import time
import openmm
from openff.toolkit.topology import Molecule
from openmmforcefields.generators import SystemGenerator
from simtk import unit
from openmm import app, Platform, LangevinIntegrator, NonbondedForce
from openmm.app import Simulation, Modeller
from rdkit import Chem
from openff.units.openmm import to_openmm
import copy
import pickle
from openmm import *
from openmm.app import *

temperature = 300 * unit.kelvin
equilibration_steps = 5000
reporting_interval = 100

input_mol = 'toyLigand.sdf'

num_steps = 25000

speed = 0

for i in range(Platform.getNumPlatforms()):
    p = Platform.getPlatform(i)
    if p.getSpeed() > speed:
        platform = p
        speed = p.getSpeed()

if platform.getName() == 'CUDA' or platform.getName() == 'OpenCL':
    platform.setPropertyDefaultValue('Precision', 'mixed')
    print('Set precision for platform', platform.getName(), 'to mixed')

mol = Chem.SDMolSupplier(input_mol)
mol = [x for x in mol][0]
ligand_mol = Molecule(mol)

forcefield_kwargs = {'constraints': app.HBonds, 'rigidWater': True,
                     'removeCMMotion': True, 'hydrogenMass': 4 * unit.amu, 'nonbondedMethod': app.PME, 'nonbondedCutoff': 0.9 * unit.nanometer}

system_generator = SystemGenerator(
    forcefields=['amber/ff14SB.xml', 'amber/tip3p_standard.xml'],
    small_molecule_forcefield='gaff-2.11',
    molecules=[ligand_mol],
    periodic_forcefield_kwargs=forcefield_kwargs)

ligand_positions = ligand_mol.conformers[0]
modeller = Modeller(ligand_mol.to_topology().to_openmm(),
                    to_openmm(ligand_positions))

modeller.addSolvent(system_generator.forcefield,
                    model='tip3p', padding=12 * unit.angstroms)

system = system_generator.create_system(
    modeller.topology)

solvent = set([a.index for a in modeller.topology.atoms()
              if a.residue.name == 'HOH'])
ligand = set([a.index for a in modeller.topology.atoms()
             if a.residue.name == 'UNK'])

for force in system.getForces():
    print(i, force, force.getForceGroup())
    if isinstance(force, NonbondedForce):
        force.setForceGroup(0)

        force.addGlobalParameter("solute_coulomb_scale", 1)
        force.addGlobalParameter("solvent_coulomb_scale", 1)

        force.addGlobalParameter("solute_lj_scale", 1)
        force.addGlobalParameter("solvent_lj_scale", 1)

        for i in range(force.getNumParticles()):
            charge, sigma, epsilon = force.getParticleParameters(i)
            # Set the parameters to be 0 when the corresponding parameter is 0,
            # and to have their normal values when it is 1.

            coulomb_param = "solute_coulomb_scale" if i in ligand else "solvent_coulomb_scale"
            lj_param = 'solute_lj_scale' if i in ligand else 'solvent_lj_scale'
            force.setParticleParameters(i, 0, 0, 0)
            force.addParticleParameterOffset(coulomb_param, i, charge, 0, 0)
            force.addParticleParameterOffset(lj_param, i, 0, sigma, epsilon)
        for i in range(force.getNumExceptions()):
            p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(
                i)
            force.setExceptionParameters(i, p1, p2, 0, 0, 0)
    else:
        force.setForceGroup(2)

integrator = LangevinIntegrator(
    temperature, 1 / unit.picosecond, 0.004 * unit.picoseconds)
system.addForce(openmm.MonteCarloBarostat(
    1 * unit.atmospheres, temperature, 25))
simulation = Simulation(modeller.topology, system,
                        integrator, platform=platform)
context = simulation.context
context.setPositions(modeller.positions)
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(temperature)
simulation.step(equilibration_steps)
simulation.step(num_steps)

def coulomb_energy(solute_scale, solvent_scale):
    context.setParameter("solute_coulomb_scale", solute_scale)
    context.setParameter("solvent_coulomb_scale", solvent_scale)
    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

def vdw_energy(solute_scale, solvent_scale):
    context.setParameter("solute_lj_scale", solute_scale)
    context.setParameter("solvent_lj_scale", solvent_scale)
    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

total_coulomb = coulomb_energy(1, 1)
solute_coulomb = coulomb_energy(1, 0)
solvent_coulomb = coulomb_energy(0, 1)
print(f'Coulomb component: {total_coulomb - solute_coulomb - solvent_coulomb}')

total_vdw = vdw_energy(1, 1)
solute_vdw = vdw_energy(1, 0)
solvent_vdw = vdw_energy(0, 1)
print(f'vdw energy component: {total_vdw - solute_vdw - solvent_vdw}')
peastman commented 1 year ago
def coulomb_energy(solute_scale, solvent_scale):
    context.setParameter("solute_coulomb_scale", solute_scale)
    context.setParameter("solvent_coulomb_scale", solvent_scale)
    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

You need to also set

    context.setParameter("solute_lj_scale", 0)
    context.setParameter("solvent_lj_scale", 0)

Otherwise it might or might not also include one or both LJ components, depending on how it was last called. Same for vdw_energy().

spadavec commented 1 year ago

@peastman thank you, that worked perfectly! The only problem I'm seeing now is when you try to set the forces for protein&water in a protein/water/ligand system, it produces a NaN for particle coordinate during equilibration. Relevant code:

import sys
import time
import openmm
from openff.toolkit.topology import Molecule
from openmmforcefields.generators import SystemGenerator
from simtk import unit
from openmm import app, Platform, LangevinIntegrator, NonbondedForce, XmlSerializer
from openmm.app import Simulation, Modeller
from rdkit import Chem
from openff.units.openmm import to_openmm
import copy
import pickle
from openmm import *
from openmm.app import *
from openmm.app import PDBFile, StateDataReporter
from mdtraj.reporters import HDF5Reporter

temperature = 300 * unit.kelvin
equilibration_steps = 5000
reporting_interval = 100

pdb_in = 'input/test_protein.pdb'
protein_pdb = PDBFile(pdb_in)
modeller = Modeller(protein_pdb.topology, protein_pdb.positions)

input_mol = 'ToyLigand.sdf'
output_traj = 'protein_ligand.h5'
num_steps = 2500

speed = 0

for i in range(Platform.getNumPlatforms()):
    p = Platform.getPlatform(i)
    if p.getSpeed() > speed:
        platform = p
        speed = p.getSpeed()

if platform.getName() == 'CUDA' or platform.getName() == 'OpenCL':
    platform.setPropertyDefaultValue('Precision', 'mixed')
    print('Set precision for platform', platform.getName(), 'to mixed')

mol = Chem.SDMolSupplier(input_mol)
mol = [x for x in mol][0]
ligand_mol = Molecule(mol)

forcefield_kwargs = {'constraints': app.HBonds, 'rigidWater': True,
                     'removeCMMotion': True, 'hydrogenMass': 4 * unit.amu, 'nonbondedMethod': app.PME, 'nonbondedCutoff': 0.9 * unit.nanometer}

system_generator = SystemGenerator(
    forcefields=['amber/ff14SB.xml', 'amber/tip3p_standard.xml'],
    small_molecule_forcefield='gaff-2.11',
    molecules=[ligand_mol],
    periodic_forcefield_kwargs=forcefield_kwargs)

ligand_positions = ligand_mol.conformers[0]
modeller.add(ligand_mol.to_topology().to_openmm(), to_openmm(ligand_positions))

modeller.addSolvent(system_generator.forcefield,
                    model='tip3p', padding=12 * unit.angstroms)

system = system_generator.create_system(
    modeller.topology)

ligand = set([a.index for a in modeller.topology.atoms()
             if a.residue.name == 'UNK'])

protein_water = set([a.index for a in modeller.topology.atoms()
                         if a.index not in ligand])

for force in system.getForces():
    print(i, force, force.getForceGroup())
    if isinstance(force, NonbondedForce):
        force.setForceGroup(0)

        force.addGlobalParameter("solute_coulomb_scale", 1)
        force.addGlobalParameter("solvent_coulomb_scale", 1)

        force.addGlobalParameter("solute_lj_scale", 1)
        force.addGlobalParameter("solvent_lj_scale", 1)

        for i in range(force.getNumParticles()):
            charge, sigma, epsilon = force.getParticleParameters(i)
            # Set the parameters to be 0 when the corresponding parameter is 0,
            # and to have their normal values when it is 1.

            coulomb_param = "solute_coulomb_scale" if i in ligand else "solvent_coulomb_scale"
            lj_param = 'solute_lj_scale' if i in ligand else 'solvent_lj_scale'
            force.setParticleParameters(i, 0, 0, 0)
            force.addParticleParameterOffset(coulomb_param, i, charge, 0, 0)
            force.addParticleParameterOffset(lj_param, i, 0, sigma, epsilon)

        for i in range(force.getNumExceptions()):
            p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(
                i)
            force.setExceptionParameters(i, p1, p2, 0, 0, 0)
    else:
        force.setForceGroup(2)

integrator = LangevinIntegrator(
    temperature, 1 / unit.picosecond, 0.004 * unit.picoseconds)
system.addForce(openmm.MonteCarloBarostat(
    1 * unit.atmospheres, temperature, 25))
simulation = Simulation(modeller.topology, system,
                        integrator, platform=platform)
context = simulation.context

context.setPositions(modeller.positions)
simulation.minimizeEnergy()

simulation.context.setVelocitiesToTemperature(temperature)
simulation.step(equilibration_steps)
simulation.step(num_steps)

def coulomb_energy(solute_scale, solvent_scale):
    context.setParameter("solute_coulomb_scale", solute_scale)
    context.setParameter("solvent_coulomb_scale", solvent_scale)
    context.setParameter("solute_lj_scale", 0)
    context.setParameter("solvent_lj_scale", 0)
    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

def vdw_energy(solute_scale, solvent_scale):
    context.setParameter("solute_lj_scale", solute_scale)
    context.setParameter("solvent_lj_scale", solvent_scale)
    context.setParameter("solute_coulomb_scale", 0)
    context.setParameter("solvent_coulomb_scale", 0)

    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

total_coulomb = coulomb_energy(1, 1)
solute_coulomb = coulomb_energy(1, 0)
solvent_coulomb = coulomb_energy(0, 1)
print(
    f'Coulomb component: {(total_coulomb - solute_coulomb - solvent_coulomb) / unit.kilocalorie_per_mole}')

total_vdw = vdw_energy(1, 1)
solute_vdw = vdw_energy(1, 0)
solvent_vdw = vdw_energy(0, 1)
print(
    f'vdw energy component: {(total_vdw - solute_vdw - solvent_vdw) / unit.kilocalorie_per_mole}')

As attempting to add the forces to the protein & water together too large, and making the system blow up?

spadavec commented 1 year ago

@peastman I'm having a really hard time debugging this issue. I'm able to get through minimization (pre-minimization energy : -926543.3647115116 kJ/mol, post : -926462.7088816308 kJ/mol

I'm trying to reduce the number of atoms that are considered "solvent" in this case:

import sys
import time
import openmm
from openff.toolkit.topology import Molecule
from openmmforcefields.generators import SystemGenerator
from simtk import unit
from openmm import app, Platform, LangevinIntegrator, LangevinMiddleIntegrator, NonbondedForce, XmlSerializer
from openmm.app import Simulation, Modeller
from rdkit import Chem
from openff.units.openmm import to_openmm
import copy
import pickle
from openmm import *
from openmm.app import *
from openmm.app import PDBFile, StateDataReporter
from mdtraj.reporters import HDF5Reporter
import mdtraj
from PythonPDBStructures.geometry import get_nearest_neighbors_residues_with_biopython
import numpy
from scipy.spatial import distance

temperature = 300 * unit.kelvin
equilibration_steps = 5000
reporting_interval = 100

pdb_in = ToyProtein.pdb'
protein_pdb = PDBFile(pdb_in)
modeller = Modeller(protein_pdb.topology, protein_pdb.positions)

input_mol = 'ToyLigand.sdf'
output_traj = 'protein_ligand.h5'
num_steps = 2500

speed = 0

for i in range(Platform.getNumPlatforms()):
    p = Platform.getPlatform(i)
    if p.getSpeed() > speed:
        platform = p
        speed = p.getSpeed()

if platform.getName() == 'CUDA' or platform.getName() == 'OpenCL':
    platform.setPropertyDefaultValue('Precision', 'mixed')
    print('Set precision for platform', platform.getName(), 'to mixed')

mol = Chem.SDMolSupplier(input_mol)
mol = [x for x in mol][0]
ligand_mol = Molecule(mol)

forcefield_kwargs = {'constraints': app.HBonds, 'rigidWater': True,
                     'removeCMMotion': True, 'hydrogenMass': 4 * unit.amu, 'nonbondedMethod': app.PME, 'nonbondedCutoff': 0.9 * unit.nanometer}

system_generator = SystemGenerator(
    forcefields=['amber/ff14SB.xml', 'amber/tip3p_standard.xml'],
    small_molecule_forcefield='gaff-2.11',
    molecules=[ligand_mol],
    periodic_forcefield_kwargs=forcefield_kwargs)

ligand_positions = ligand_mol.conformers[0]
modeller.add(ligand_mol.to_topology().to_openmm(), to_openmm(ligand_positions))

ligand_indexes = list(set(
    [a.index for a in modeller.topology.atoms() if a.residue.name == 'UNK']))

modeller.addSolvent(system_generator.forcefield,
                    model='tip3p', padding=12 * unit.angstroms)

all_positions = numpy.array(modeller.getPositions() / unit.nanometer)
distances = distance.cdist(all_positions, all_positions, 'euclidean')
distances_subset = distances[:, ligand_indexes]

CUTOFF = 0.8
matches = numpy.any(distances_subset <= CUTOFF, axis=1)
scaling_indexes = set([i for i, x in enumerate(matches) if x])
print(
    f"A total of {len(scaling_indexes)} atoms are in within {CUTOFF} nm of the ligand...")

system = system_generator.create_system(modeller.topology)

for force in system.getForces():
    print(i, force, force.getForceGroup())
    if isinstance(force, NonbondedForce):
        force.setForceGroup(0)

        force.addGlobalParameter("solute_coulomb_scale", 1)
        force.addGlobalParameter("solvent_coulomb_scale", 1)

        force.addGlobalParameter("solute_lj_scale", 1)
        force.addGlobalParameter("solvent_lj_scale", 1)

        for i in range(force.getNumParticles()):
            charge, sigma, epsilon = force.getParticleParameters(i)

            # Set the parameters to be 0 when the corresponding parameter is 0,
            # and to have their normal values when it is 1.

            if i in ligand_indexes or i in scaling_indexes:
                if i in ligand_indexes:
                    coulomb_param = 'solute_coulomb_scale'
                    lj_param = 'solute_lj_scale'
                elif i in scaling_indexes:
                    coulomb_param = 'solvent_coulomb_scale'
                    lj_param = 'solvent_lj_scale'

                force.setParticleParameters(i, 0, 0, 0)
                force.addParticleParameterOffset(
                    coulomb_param, i, charge, 0, 0)
                force.addParticleParameterOffset(
                    lj_param, i, 0, sigma, epsilon)

                for i in range(force.getNumExceptions()):
                    p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(
                        i)
                    force.setExceptionParameters(i, p1, p2, 0, 0, 0)
    else:
        force.setForceGroup(2)

integrator = LangevinMiddleIntegrator(
    temperature, 1 / unit.picosecond, 0.004 * unit.picoseconds)
system.addForce(openmm.MonteCarloBarostat(
    1 * unit.atmospheres, temperature, 25))
simulation = Simulation(modeller.topology, system,
                        integrator, platform=platform)
context = simulation.context

context.setPositions(modeller.positions)

simulation.context.setVelocitiesToTemperature(temperature)
simulation.step(equilibration_steps)
simulation.step(num_steps)

# Create reporters
simulation.reporters.append(HDF5Reporter(
    output_traj, reporting_interval, enforcePeriodicBox=False))
simulation.reporters.append(StateDataReporter(
    sys.stdout, reporting_interval * 5, step=True, potentialEnergy=True, temperature=True))

def coulomb_energy(solute_scale, solvent_scale):
    context.setParameter("solute_coulomb_scale", solute_scale)
    context.setParameter("solvent_coulomb_scale", solvent_scale)
    context.setParameter("solute_lj_scale", 0)
    context.setParameter("solvent_lj_scale", 0)
    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

def vdw_energy(solute_scale, solvent_scale):
    context.setParameter("solute_lj_scale", solute_scale)
    context.setParameter("solvent_lj_scale", solvent_scale)
    context.setParameter("solute_coulomb_scale", 0)
    context.setParameter("solvent_coulomb_scale", 0)

    return context.getState(getEnergy=True, groups={0}).getPotentialEnergy()

# Solute = Ligand

total_coulomb = coulomb_energy(1, 1)
solute_coulomb = coulomb_energy(1, 0)
solvent_coulomb = coulomb_energy(0, 1)
print(
    f'Coulomb component: {(total_coulomb - solute_coulomb - solvent_coulomb) / unit.kilocalorie_per_mole}')

total_vdw = vdw_energy(1, 1)
solute_vdw = vdw_energy(1, 0)
solvent_vdw = vdw_energy(0, 1)
print(
    f'vdw energy component: {(total_vdw - solute_vdw - solvent_vdw) / unit.kilocalorie_per_mole}')

But this doesn't change anything and I still get a NaN particle position during equilibration. If I completely remove the scaling parameters, everything runs fine. Is there a simple fix to this?

peastman commented 1 year ago

I added a section showing how to do it.

spadavec commented 1 year ago

@peastman Thank you for the update!

I hate to be a pest, but if I try to run the above example without minimization/equilibration (as you've done), then I can run everything just fine and get a reasonable answer. However, if I try to actually minimize, equilibrate, and run MD then I get NaN for particle positions at equilibration, and very quickly--suggesting that something is hinky with the force setup. The minimization runs fine, and reports something in the ballpark I'd expect:

The pre-minimization energy is -595322.9981358416 kJ/mol
Running minimization...
The post-minimization energy is -595218.61226944 kJ/mol

This is the case for both just protein+water systems and protein+water+ligand. I've tried creating systems using pre-solvated system PDBs and using the forcefield to generate the system (as opposed to using system_generator) and I get the same result. I've also tried just limiting the forces to just coulomb or just lj individually and still get the same result.

I've also tried minimizing and equilibrating on CPU first and then switching to GPU, but see the same result.

Are you seeing the same thing on your end? I've attached an example script here for reference.

ExampleScript.zip

peastman commented 1 year ago

The system configured in the cookbook is for computing interaction energies, not for running simulations. It won't work for that purpose. For example, all the nonbonded exceptions were set to zero.

sef43 commented 1 year ago

Looks good, is it ready to merge?

peastman commented 1 year ago

Thanks!