minimization with ANI mixed system is slow #94
Description
Hi everyone, thanks for these amazing advancements. I recently noticed that minimizing a mixed ML/MM system with ligand atoms treated with ANI and protein atoms + solvent using amber params takes much longer than expected when compared to actual integration steps. Below is a minimal script and files to reproduce (although I have tested on multiple systems to confirm). In short:
-
ttc md min =12s
-
ttc 200 langevin steps =2.6s.
-
ttc mdml min = 12.8mins
-
ttc 200 langevin mdml steps = 40s.
I'm just wondering if this is expected behavior, given the ratio of min/integration time is so much longer than the same ratio for plain MD systems?
I'm using PDB code 4gj2 and the associated ligand 0XH.
import openmm as mm
import openmm.app as app
import openmm.unit as unit
from openmmml import MLPotential
from openmmforcefields.generators import SystemGenerator, SMIRNOFFTemplateGenerator
from rdkit import Chem
from openff.toolkit.topology import Molecule as OFFMolecule
from openmm import app, unit
import openmmtools
import sys
pdb = app.PDBFile('4gj2_fixed.pdb')
modeller = app.Modeller(pdb.topology, pdb.positions)
sdf_file = '4gj2_B_0XH.sdf'
supplier = Chem.SDMolSupplier(sdf_file,removeHs=False)
ligand_mols = []
for mol in supplier:
name = mol.GetProp('_Name')
mol = Chem.AddHs(mol,addCoords=True)
ligand_mols.append(mol)
off_ligand = OFFMolecule.from_rdkit(ligand_mols[0])
smirnoff = SMIRNOFFTemplateGenerator(molecules=off_ligand)
forcefield = app.ForceField('amber/protein.ff14SB.xml','amber/tip3p_standard.xml')
forcefield.registerTemplateGenerator(smirnoff.generator)
lig_top = off_ligand.to_topology()
modeller.add(lig_top.to_openmm(), lig_top.get_positions().to_openmm())
print('Modeller object has %d atoms' % modeller.topology.getNumAtoms())
print('starting to build system')
modeller.addSolvent(forcefield)
system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.PME,
nonbondedCutoff=1*unit.nanometer, constraints=app.HBonds)
ml_atoms = [i.index for i in [i for i in modeller.topology.chains()][1].atoms()]
potential = MLPotential('ani2x')
ml_system = potential.createMixedSystem(modeller.topology, system, ml_atoms)
temperature = 298.15 * unit.kelvin
frictionCoeff = 1 / unit.picosecond
timeStep = 0.5 * unit.femtosecond
integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)
# ML/MD system
simulation = app.Simulation(modeller.topology, ml_system, integrator)
#MD System
# simulation = app.Simulation(modeller.topology, system, integrator)
simulation.context.setPositions(modeller.positions)
simulation.context.setVelocitiesToTemperature(temperature)
# Minimize the system
simulation.minimizeEnergy()
# Configure a reporter to print to the console every 10 steps and write to a PDB file
reporter = app.StateDataReporter(file=sys.stdout, reportInterval=10, step=True, time=True, potentialEnergy=True, temperature=True)
simulation.reporters.append(reporter)
simulation.reporters.append(app.PDBReporter('traj.pdb', 10))
# Simulate 0.5 ps
simulation.step(200)