Merge pull request #524 from swails/nbfix_amber

Add support for Amber prmtop files with NBFIX

wrappers/python/simtk/openmm/app/internal/amber_file_parser.py

@@ -12,9 +12,9 @@ Simbios, the NIH National Center for Physics-Based Simulation of
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2012-2013 Stanford University and the Authors.
+Portions copyright (c) 2012-2014 Stanford University and the Authors.
 Authors: Randall J. Radmer, John D. Chodera, Peter Eastman
-Contributors: Christoph Klein, Michael R. Shirts
+Contributors: Christoph Klein, Michael R. Shirts, Jason Swails
 
 Permission is hereby granted, free of charge, to any person obtaining a
 copy of this software and associated documentation files (the "Software"),
@@ -76,6 +76,10 @@ POINTER_LABEL_LIST = POINTER_LABELS.replace(',', '').split()
 VELSCALE = 20.455 # velocity conversion factor to angstroms/picosecond
 TINY = 1.0e-8
 
+class NbfixPresent(Exception):
+    """ Exception raised when NBFIX is used for the Lennard-Jones terms """
+    pass
+
 class PrmtopLoader(object):
     """Parsed AMBER prmtop file.
 
@@ -297,7 +301,11 @@ class PrmtopLoader(object):
         return self.residuePointerDict[iAtom]
 
     def getNonbondTerms(self):
-        """Return list of all rVdw, epsilon pairs for each atom"""
+        """
+        Return list of all rVdw, epsilon pairs for each atom. If off-diagonal
+        elements of the Lennard-Jones A and B coefficient matrices are found,
+        NbfixPresent exception is raised
+        """
         try:
             return self._nonbondTerms
         except AttributeError:
@@ -305,9 +313,10 @@ class PrmtopLoader(object):
         self._nonbondTerms=[]
         lengthConversionFactor = units.angstrom.conversion_factor_to(units.nanometer)
         energyConversionFactor = units.kilocalorie_per_mole.conversion_factor_to(units.kilojoule_per_mole)
-        for iAtom in range(self.getNumAtoms()):
-            numTypes=self.getNumTypes()
+        numTypes = self.getNumTypes()
         atomTypeIndexes=self._getAtomTypeIndexes()
+        type_parameters = [(0, 0) for i in range(numTypes)]
+        for iAtom in range(self.getNumAtoms()):
             index=(numTypes+1)*(atomTypeIndexes[iAtom]-1)
             nbIndex=int(self._raw_data['NONBONDED_PARM_INDEX'][index])-1
             if nbIndex<0:
@@ -320,9 +329,31 @@ class PrmtopLoader(object):
             except ZeroDivisionError:
                 rMin = 1.0
                 epsilon = 0.0
+            type_parameters[atomTypeIndexes[iAtom]-1] = (rMin/2.0, epsilon)
             rVdw = rMin/2.0*lengthConversionFactor
             epsilon = epsilon*energyConversionFactor
             self._nonbondTerms.append( (rVdw, epsilon) )
+        # Check if we have any off-diagonal modified LJ terms that would require
+        # an NBFIX-like solution
+        for i in range(numTypes):
+            for j in range(numTypes):
+                index = int(self._raw_data['NONBONDED_PARM_INDEX'][numTypes*i+j]) - 1
+                rij = type_parameters[i][0] + type_parameters[j][0]
+                wdij = sqrt(type_parameters[i][1] * type_parameters[j][1])
+                a = float(self._raw_data['LENNARD_JONES_ACOEF'][index])
+                b = float(self._raw_data['LENNARD_JONES_BCOEF'][index])
+                if a == 0 or b == 0:
+                    if a != 0 or b != 0 or (wdij != 0 and rij != 0):
+                        del self._nonbondTerms
+                        raise NbfixPresent('Off-diagonal Lennard-Jones elements'
+                                           ' found. Cannot determine LJ '
+                                           'parameters for individual atoms.')
+                elif (abs((a - (wdij * rij ** 12)) / a) > 1e-6 or
+                      abs((b - (2 * wdij * rij**6)) / b) > 1e-6):
+                    del self._nonbondTerms
+                    raise NbfixPresent('Off-diagonal Lennard-Jones elements '
+                                       'found. Cannot determine LJ parameters '
+                                       'for individual atoms.')
         return self._nonbondTerms
 
     def _getBonds(self, bondPointers):
@@ -429,7 +460,47 @@ class PrmtopLoader(object):
                           +self._raw_data["DIHEDRALS_WITHOUT_HYDROGEN"]
         returnList=[]
         charges=self.getCharges()
+        try:
             nonbondTerms = self.getNonbondTerms()
+        except NbfixPresent:
+            # We need to do the unit conversions here, since getNonbondTerms
+            # never finished and it has unit conversions in there
+            length_conv = units.angstrom.conversion_factor_to(units.nanometers)
+            ene_conv = units.kilocalories_per_mole.conversion_factor_to(
+                                units.kilojoules_per_mole)
+            parm_acoef = [float(x) for x in self._raw_data['LENNARD_JONES_ACOEF']]
+            parm_bcoef = [float(x) for x in self._raw_data['LENNARD_JONES_BCOEF']]
+            nbidx = [int(x) for x in self._raw_data['NONBONDED_PARM_INDEX']]
+            numTypes = self.getNumTypes()
+            atomTypeIndexes=self._getAtomTypeIndexes()
+            for ii in range(0, len(dihedralPointers), 5):
+                if int(dihedralPointers[ii+2])>0 and int(dihedralPointers[ii+3])>0:
+                    iAtom = int(dihedralPointers[ii])//3
+                    lAtom = int(dihedralPointers[ii+3])//3
+                    iidx = int(dihedralPointers[ii+4]) - 1
+                    chargeProd = charges[iAtom]*charges[lAtom]
+                    typ1 = atomTypeIndexes[iAtom] - 1
+                    typ2 = atomTypeIndexes[lAtom] - 1
+                    idx = nbidx[numTypes*typ1+typ2] - 1
+                    a = parm_acoef[idx]
+                    b = parm_bcoef[idx]
+                    try:
+                        epsilon = b * b / (4 * a) * ene_conv
+                        rMin = (2 * a / b) ** (1/6.0) * length_conv
+                    except ZeroDivisionError:
+                        rMin = 1
+                        epsilon = 0
+                    try:
+                        iScee = float(self._raw_data['SCEE_SCALE_FACTOR'][iidx])
+                    except KeyError:
+                        iScee = 1.2
+                    try:
+                        iScnb = float(self._raw_data['SCNB_SCALE_FACTOR'][iidx])
+                    except KeyError:
+                        iScnb = 2.0
+                    returnList.append((iAtom, lAtom, chargeProd, rMin, epsilon, iScee, iScnb))
+        else:
+            # This block gets hit if NbfixPresent is _not_ caught
             for ii in range(0,len(dihedralPointers),5):
                 if int(dihedralPointers[ii+2])>0 and int(dihedralPointers[ii+3])>0:
                     iAtom = int(dihedralPointers[ii])//3
@@ -782,10 +853,55 @@ def readAmberSystem(prmtop_filename=None, prmtop_loader=None, shake=None, gbmode
 
     # Add per-particle nonbonded parameters.
     sigmaScale = 2**(-1./6.) * 2.0
-    for (charge, (rVdw, epsilon)) in zip(prmtop.getCharges(), prmtop.getNonbondTerms()):
+    nbfix = False
+    try:
+        nonbondTerms = prmtop.getNonbondTerms()
+    except NbfixPresent:
+        nbfix = True
+        for charge in prmtop.getCharges():
+            force.addParticle(charge, 1.0, 0.0)
+        numTypes = prmtop.getNumTypes()
+        parm_acoef = [float(x) for x in prmtop._raw_data['LENNARD_JONES_ACOEF']]
+        parm_bcoef = [float(x) for x in prmtop._raw_data['LENNARD_JONES_BCOEF']]
+        nbidx = [int(x) for x in prmtop._raw_data['NONBONDED_PARM_INDEX']]
+        acoef = [0 for i in range(numTypes*numTypes)]
+        bcoef = acoef[:] # copy
+        ene_conv = units.kilocalories_per_mole.conversion_factor_to(units.kilojoules_per_mole)
+        length_conv = units.angstroms.conversion_factor_to(units.nanometers)
+        afac = sqrt(ene_conv) * length_conv**6
+        bfac = ene_conv * length_conv**6
+        for i in range(numTypes):
+            for j in range(numTypes):
+                idx = nbidx[numTypes*i+j] - 1
+                acoef[i+numTypes*j] = sqrt(parm_acoef[idx]) * afac
+                bcoef[i+numTypes*j] = parm_bcoef[idx] * bfac
+        cforce = mm.CustomNonbondedForce('(a/r6)^2-b/r6; r6=r^6;'
+                                         'a=acoef(type1, type2);'
+                                         'b=bcoef(type1, type2);')
+        cforce.addTabulatedFunction('acoef',
+                    mm.Discrete2DFunction(numTypes, numTypes, acoef))
+        cforce.addTabulatedFunction('bcoef',
+                    mm.Discrete2DFunction(numTypes, numTypes, bcoef))
+        cforce.addPerParticleParameter('type')
+        for atom in prmtop._getAtomTypeIndexes():
+            cforce.addParticle((atom-1,))
+        # Now set the various properties based on the NonbondedForce object
+        if nonbondedMethod in ('PME', 'Ewald', 'CutoffPeriodic'):
+            cforce.setNonbondedMethod(cforce.CutoffPeriodic)
+            cforce.setCutoffDistance(nonbondedCutoff)
+        elif nonbondedMethod == 'CutoffNonPeriodic':
+            cforce.setNonbondedMethod(cforce.CutoffNonPeriodic)
+            cforce.setCutoffDistance(nonbondedCutoff)
+        elif nonbondedMethod == 'NoCutoff':
+            cforce.setNonbondedMethod(cforce.NoCutoff)
+        else:
+            raise ValueError('Unrecognized cutoff option %s' % nonbondedMethod)
+    else:
+        for (charge, (rVdw, epsilon)) in zip(prmtop.getCharges(), nonbondTerms):
             sigma = rVdw * sigmaScale
             force.addParticle(charge, sigma, epsilon)
 
+
     # Add 1-4 Interactions
     excludedAtomPairs = set()
     sigmaScale = 2**(-1./6.)
@@ -807,6 +923,13 @@ def readAmberSystem(prmtop_filename=None, prmtop_loader=None, shake=None, gbmode
             if min((iAtom, jAtom), (jAtom, iAtom)) in excludedAtomPairs: continue
             force.addException(iAtom, jAtom, excludeParams[0], excludeParams[1], excludeParams[2])
 
+    # Copy the exceptions as exclusions to the CustomNonbondedForce if we have
+    # NBFIX terms
+    if nbfix:
+        for i in range(force.getNumExceptions()):
+            ii, jj, chg, sig, eps = force.getExceptionParameters(i)
+            cforce.addExclusion(ii, jj)
+        system.addForce(cforce)
     system.addForce(force)
 
     # Add virtual sites for water.

wrappers/python/tests/TestAmberPrmtopFile.py

@@ -5,18 +5,15 @@ from simtk.openmm import *
 from simtk.unit import *
 import simtk.openmm.app.element as elem
 
+prmtop1 = AmberPrmtopFile('systems/alanine-dipeptide-explicit.prmtop')
+prmtop2 = AmberPrmtopFile('systems/alanine-dipeptide-implicit.prmtop')
+prmtop3 = AmberPrmtopFile('systems/ff14ipq.parm7')
+inpcrd3 = AmberInpcrdFile('systems/ff14ipq.rst7')
+
 class TestAmberPrmtopFile(unittest.TestCase):
 
     """Test the AmberPrmtopFile.createSystem() method."""
 
-    def setUp(self):
-        """Set up the tests by loading the input files."""
-
-        # alanine dipeptide with explicit water
-        self.prmtop1 = AmberPrmtopFile('systems/alanine-dipeptide-explicit.prmtop')
-        # alanine dipeptide with implicit water
-        self.prmtop2 = AmberPrmtopFile('systems/alanine-dipeptide-implicit.prmtop')
-
     def test_NonbondedMethod(self):
         """Test all five options for the nonbondedMethod parameter."""
 
@@ -25,7 +22,7 @@ class TestAmberPrmtopFile(unittest.TestCase):
                      CutoffPeriodic:NonbondedForce.CutoffPeriodic, 
                      Ewald:NonbondedForce.Ewald, PME: NonbondedForce.PME}
         for method in methodMap:
-            system = self.prmtop1.createSystem(nonbondedMethod=method)
+            system = prmtop1.createSystem(nonbondedMethod=method)
             forces = system.getForces()
             self.assertTrue(any(isinstance(f, NonbondedForce) and 
                                 f.getNonbondedMethod()==methodMap[method] 
@@ -35,7 +32,7 @@ class TestAmberPrmtopFile(unittest.TestCase):
         """Test to make sure the nonbondedCutoff parameter is passed correctly."""
 
         for method in [CutoffNonPeriodic, CutoffPeriodic, Ewald, PME]:
-            system = self.prmtop1.createSystem(nonbondedMethod=method, 
+            system = prmtop1.createSystem(nonbondedMethod=method, 
                                           nonbondedCutoff=2*nanometer, 
                                           constraints=HBonds)
             cutoff_distance = 0.0*nanometer
@@ -49,7 +46,7 @@ class TestAmberPrmtopFile(unittest.TestCase):
         """Test to make sure the ewaldErrorTolerance parameter is passed correctly."""
 
         for method in [Ewald, PME]:
-            system = self.prmtop1.createSystem(nonbondedMethod=method, 
+            system = prmtop1.createSystem(nonbondedMethod=method, 
                                           ewaldErrorTolerance=1e-6, 
                                           constraints=HBonds)
             tolerance = 0
@@ -63,17 +60,17 @@ class TestAmberPrmtopFile(unittest.TestCase):
         """Test both options (True and False) for the removeCMMotion parameter."""
 
         for b in [True, False]:
-            system = self.prmtop1.createSystem(removeCMMotion=b)
+            system = prmtop1.createSystem(removeCMMotion=b)
             forces = system.getForces()
             self.assertEqual(any(isinstance(f, CMMotionRemover) for f in forces), b)
 
     def test_RigidWaterAndConstraints(self):
         """Test all eight options for the constraints and rigidWater parameters."""
 
-        topology = self.prmtop1.topology
+        topology = prmtop1.topology
         for constraints_value in [None, HBonds, AllBonds, HAngles]:
             for rigidWater_value in [True, False]:
-                system = self.prmtop1.createSystem(constraints=constraints_value, 
+                system = prmtop1.createSystem(constraints=constraints_value, 
                                               rigidWater=rigidWater_value)
                 validateConstraints(self, topology, system, 
                                     constraints_value, rigidWater_value)
@@ -84,7 +81,7 @@ class TestAmberPrmtopFile(unittest.TestCase):
 
         """
         for implicitSolvent_value in [HCT, OBC1, OBC2, GBn]:
-            system = self.prmtop2.createSystem(implicitSolvent=implicitSolvent_value)
+            system = prmtop2.createSystem(implicitSolvent=implicitSolvent_value)
             forces = system.getForces()
             if implicitSolvent_value in set([HCT, OBC1, GBn]):
                 force_type = CustomGBForce
@@ -99,7 +96,7 @@ class TestAmberPrmtopFile(unittest.TestCase):
                      CutoffNonPeriodic:NonbondedForce.CutoffNonPeriodic}
         for implicitSolvent_value in [HCT, OBC1, OBC2, GBn]:
             for method in methodMap:
-                system = self.prmtop2.createSystem(implicitSolvent=implicitSolvent_value, 
+                system = prmtop2.createSystem(implicitSolvent=implicitSolvent_value, 
                                     solventDielectric=50.0, soluteDielectric=0.9, nonbondedMethod=method)
                 found_matching_solvent_dielectric=False
                 found_matching_solute_dielectric=False
@@ -136,10 +133,10 @@ class TestAmberPrmtopFile(unittest.TestCase):
     def test_HydrogenMass(self):
         """Test that altering the mass of hydrogens works correctly."""
         
-        topology = self.prmtop1.topology
+        topology = prmtop1.topology
         hydrogenMass = 4*amu
-        system1 = self.prmtop1.createSystem()
-        system2 = self.prmtop1.createSystem(hydrogenMass=hydrogenMass)
+        system1 = prmtop1.createSystem()
+        system2 = prmtop1.createSystem(hydrogenMass=hydrogenMass)
         for atom in topology.atoms():
             if atom.element == elem.hydrogen:
                 self.assertNotEqual(hydrogenMass, system1.getParticleMass(atom.index))
@@ -148,5 +145,35 @@ class TestAmberPrmtopFile(unittest.TestCase):
         totalMass2 = sum([system2.getParticleMass(i) for i in range(system2.getNumParticles())]).value_in_unit(amu)
         self.assertAlmostEqual(totalMass1, totalMass2)
 
+    def test_NBFIX(self):
+        """Test that prmtop files with modified off-diagonal LJ elements are treated properly"""
+        system = prmtop3.createSystem(nonbondedMethod=PME,
+                                      nonbondedCutoff=8*angstroms)
+        # Check the forces
+        has_nonbond_force = has_custom_nonbond_force = False
+        nonbond_exceptions = custom_nonbond_exclusions = 0
+        for force in system.getForces():
+            if isinstance(force, NonbondedForce):
+                has_nonbond_force = True
+                nonbond_exceptions = force.getNumExceptions()
+            elif isinstance(force, CustomNonbondedForce):
+                has_custom_nonbond_force = True
+                custom_nonbond_exceptions = force.getNumExclusions()
+        self.assertTrue(has_nonbond_force)
+        self.assertTrue(has_custom_nonbond_force)
+        self.assertEqual(nonbond_exceptions, custom_nonbond_exceptions)
+        integrator = VerletIntegrator(1.0*femtoseconds)
+        # Use reference platform, since it should always be present and
+        # 'working', and the system is plenty small so this won't be too slow
+        sim = Simulation(prmtop3.topology, system, integrator, Platform.getPlatformByName('Reference'))
+        # Check that the energy is about what we expect it to be
+        sim.context.setPeriodicBoxVectors(*inpcrd3.boxVectors)
+        sim.context.setPositions(inpcrd3.positions)
+        ene = sim.context.getState(getEnergy=True, enforcePeriodicBox=True).getPotentialEnergy()
+        ene = ene.value_in_unit(kilocalories_per_mole)
+        # Make sure the energy is relatively close to the value we get with
+        # Amber using this force field.
+        self.assertAlmostEqual(-7042.3903307/ene, 1, places=3)
+
 if __name__ == '__main__':
     unittest.main()