add NBFIX/VSTIE2 handling

wrappers/python/simtk/openmm/app/gromacstopfile.py

@@ -43,7 +43,8 @@ import math
 import os
 import re
 import distutils.spawn
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
+from itertools import combinations
 
 HBonds = ff.HBonds
 AllBonds = ff.AllBonds
@@ -115,7 +116,9 @@ class GromacsTopFile(object):
             self.exclusions = []
             self.pairs = []
             self.cmaps = []
+            self.vsites2 = []
             self.has_virtual_sites = False
+            self.has_nbfix_terms = False
 
     def _processFile(self, file):
         append = ''
@@ -249,6 +252,10 @@ class GromacsTopFile(object):
                 self._processPairType(line)
             elif self._currentCategory == 'cmaptypes':
                 self._processCmapType(line)
+            elif self._currentCategory == 'nonbond_params':
+                self._processNonbondType(line)
+            elif self._currentCategory == 'virtual_sites2':
+                self._processVirtualSites2(line)
             elif self._currentCategory.startswith('virtual_sites'):
                 if self._currentMoleculeType is None:
                     raise ValueError('Found %s before [ moleculetype ]' %
@@ -435,6 +442,22 @@ class GromacsTopFile(object):
             raise ValueError('Unsupported function type in [ cmaptypes ] line: '+line)
         self._cmapTypes[tuple(fields[:5])] = fields
 
+    def _processNonbondType(self, line):
+        """Process a line in the [ nonbond_params ] category."""
+        fields = line.split()
+        if len(fields) < 5:
+            raise ValueError('Too few fields in [ nonbond_params ] line: '+line)
+        if fields[2] != '1':
+            raise ValueError('Unsupported function type in [ nonbond_params ] line: '+line)
+        self._nonbondTypes[tuple(sorted(fields[:2]))] = fields
+
+    def _processVirtualSites2(self, line):
+        """Process a line in the [ virtual_sites2 ] category."""
+        fields = line.split()
+        if len(fields) < 5:
+            raise ValueError('Too few fields in [ virtual_sites2 ] line: ' + line)
+        self._currentMoleculeType.vsites2.append(fields[:5])
+
     def __init__(self, file, periodicBoxVectors=None, unitCellDimensions=None, includeDir=None, defines=None):
         """Load a top file.
 
@@ -484,6 +507,7 @@ class GromacsTopFile(object):
         self._implicitTypes = {}
         self._pairTypes = {}
         self._cmapTypes = {}
+        self._nonbondTypes = {}
         self._processFile(file)
 
         # Create the Topology from it.
@@ -643,6 +667,25 @@ class GromacsTopFile(object):
                 for types in dihedralTypeTable.values():
                     types.append(key)
 
+        # NBFIX
+
+        atom_types = []
+        for moleculeName, moleculeCount in self._molecules:
+            moleculeType = self._moleculeTypes[moleculeName]
+            for _ in range(moleculeCount):
+                for atom in moleculeType.atoms:
+                    atom_types.append(atom[1])
+        has_nbfix_terms = any([pair in self._nonbondTypes for pair in combinations(sorted(atom_types), 2)])
+        has_nbfix_terms = True
+
+        if has_nbfix_terms:
+            # Build a lookup table and angle/dihedral indices list to
+            # let us handle exclusion manually.
+            angleIndices = []
+            torsionIndices = []
+            atom_partners = defaultdict(lambda : defaultdict(set))
+            atom_charges = []
+
         # Loop over molecules and create the specified number of each type.
 
         for moleculeName, moleculeCount in self._molecules:
@@ -846,12 +889,19 @@ class GromacsTopFile(object):
                     else:
                         q = float(params[4])
 
+                    if has_nbfix_terms:
+                        if self._defaults[1] != '2':
+                            raise NotImplemented
+                        nb.addParticle(q, 1.0, 0.0)
+                        atom_charges.append(q)
+                    else:
                         if self._defaults[1] == '1':
                             nb.addParticle(q, 1.0, 0.0)
                             lj.addParticle([float(params[6]), float(params[7])])
                         elif self._defaults[1] == '2':
                             nb.addParticle(q, float(params[6]), float(params[7]))
 
+                for fields in moleculeType.atoms:
                     if implicitSolvent is OBC2:
                         if fields[1] not in self._implicitTypes:
                             raise ValueError('No implicit solvent parameters specified for atom type: '+fields[1])
@@ -860,6 +910,23 @@ class GromacsTopFile(object):
                 for fields in moleculeType.bonds:
                     atoms = [int(x)-1 for x in fields[:2]]
                     bondIndices.append((baseAtomIndex+atoms[0], baseAtomIndex+atoms[1]))
+                if has_nbfix_terms:
+                    for fields in moleculeType.bonds:
+                        atoms = [int(x)-1 for x in fields[:2]]
+                        atom_partners[baseAtomIndex+atoms[0]]['bond'].add(baseAtomIndex+atoms[1])
+                        atom_partners[baseAtomIndex+atoms[1]]['bond'].add(baseAtomIndex+atoms[0])
+                    for fields in moleculeType.angles:
+                        atoms = [int(x)-1 for x in fields[:3]]
+                        angleIndices.append((baseAtomIndex+atoms[0], baseAtomIndex+atoms[1], baseAtomIndex+atoms[2]))
+                        for pair in combinations(atoms, 2):
+                            atom_partners[baseAtomIndex+pair[0]]['angle'].add(baseAtomIndex+pair[1])
+                            atom_partners[baseAtomIndex+pair[1]]['angle'].add(baseAtomIndex+pair[0])
+                    for fields in moleculeType.dihedrals:
+                        atoms = [int(x)-1 for x in fields[:4]]
+                        torsionIndices.append((baseAtomIndex+atoms[0], baseAtomIndex+atoms[1], baseAtomIndex+atoms[2], baseAtomIndex+atoms[3]))
+                        for pair in combinations(atoms, 2):
+                            atom_partners[baseAtomIndex+pair[0]]['torsion'].add(baseAtomIndex+pair[1])
+                            atom_partners[baseAtomIndex+pair[1]]['torsion'].add(baseAtomIndex+pair[0])
 
                 # Record nonbonded exceptions.
 
@@ -886,12 +953,66 @@ class GromacsTopFile(object):
                         if atom > atoms[0]:
                             exclusions.append((baseAtomIndex+atoms[0], baseAtomIndex+atom))
 
+                # Record virtual sites
+
+                for fields in moleculeType.vsites2:
+                    atoms = [int(x)-1 for x in fields[:3]]
+                    c1 = float(fields[4])
+                    vsite = mm.TwoParticleAverageSite(baseAtomIndex+atoms[1], baseAtomIndex+atoms[2], (1-c1), c1)
+                    sys.setVirtualSite(baseAtomIndex+atoms[0], vsite)
 
         # Create nonbonded exceptions.
 
+        if not has_nbfix_terms:
             nb.createExceptionsFromBonds(bondIndices, fudgeQQ, fudgeLJ)
+        else:
+            excluded_atom_pairs = set() # save these pairs so we don't zero them out
+            for tor in torsionIndices:
+                # First check to see if atoms 1 and 4 are already excluded because
+                # they are 1-2 or 1-3 pairs (would happen in 6-member rings or
+                # fewer). Then check that they're not already added as exclusions
+                if 'bond' in atom_partners[tor[3]] and tor[0] in atom_partners[tor[3]]['bond']: continue
+                if 'angle' in atom_partners[tor[3]] and tor[0] in atom_partners[tor[3]]['angle']: continue
+                key = min((tor[0], tor[3]),
+                          (tor[3], tor[0]))
+                if key in excluded_atom_pairs: continue # multiterm...
+                params1 = self._atomTypes[atom_types[tor[0]]]
+                params4 = self._atomTypes[atom_types[tor[3]]]
+                q1 = atom_charges[tor[0]]
+                rmin1 = float(params1[6])
+                eps1 = float(params1[7])
+                q4 = atom_charges[tor[3]]
+                rmin4 = float(params4[6])
+                eps4 = float(params4[7])
+
+                charge_prod = (q1*q4)
+                epsilon = (math.sqrt(abs(eps1 * eps4)))
+                rmin14 = (rmin1 + rmin4) / 2
+                nb.addException(tor[0], tor[3], charge_prod, rmin14, epsilon)
+                excluded_atom_pairs.add(key)
+
+            # Add excluded atoms
+            for atom_idx, atom in atom_partners.items():
+                # Exclude all bonds and angles
+                for atom2 in atom['bond']:
+                    if atom2 > atom_idx:
+                        nb.addException(atom_idx, atom2, 0.0, 1.0, 0.0)
+                        excluded_atom_pairs.add((atom_idx, atom2))
+                for atom2 in atom['angle']:
+                    if ((atom_idx, atom2) in excluded_atom_pairs):
+                        continue
+                    if atom2 > atom_idx:
+                        nb.addException(atom_idx, atom2, 0.0, 1.0, 0.0)
+                        excluded_atom_pairs.add((atom_idx, atom2))
+                for atom2 in atom['dihedral']:
+                    if atom2 <= atom_idx: continue
+                    if ((atom_idx, atom2) in excluded_atom_pairs):
+                        continue
+                    nb.addException(atom_idx, atom2, 0.0, 1.0, 0.0)
+
         for exclusion in exclusions:
             nb.addException(exclusion[0], exclusion[1], 0.0, 1.0, 0.0, True)
+
         if self._defaults[1] == '1':
             # We're using a CustomNonbondedForce for LJ interactions, so also create a CustomBondForce
             # to handle the exceptions.
@@ -932,6 +1053,88 @@ class GromacsTopFile(object):
             lj.setNonbondedMethod(methodMap[nonbondedMethod])
             lj.setCutoffDistance(nonbondedCutoff)
 
+        if has_nbfix_terms:
+            if self._defaults[1] != '2':
+                raise NotImplemented
+
+            atom_nbfix_types = set([])
+            for pair in self._nonbondTypes:
+                atom_nbfix_types.add(pair[0])
+                atom_nbfix_types.add(pair[1])
+
+            lj_idx_list = [0 for _ in atom_types]
+            lj_radii, lj_depths = [], []
+            num_lj_types = 0
+            lj_type_list = []
+            for i,atom_type in enumerate(atom_types):
+                atom = self._atomTypes[atom_type]
+                if lj_idx_list[i]: continue # already assigned
+                atom_rmin = atom[6]
+                atom_epsilon = atom[7]
+                num_lj_types += 1
+                lj_idx_list[i] = num_lj_types
+                ljtype = (atom_rmin, atom_epsilon)
+                lj_type_list.append(atom)
+                lj_radii.append(float(atom_rmin))
+                lj_depths.append(float(atom_epsilon))
+                for j in range(i+1, len(atom_types)):
+                    atom_type2 = atom_types[j]
+                    if lj_idx_list[j] > 0: continue # already assigned
+                    atom2 = self._atomTypes[atom_type2]
+                    atom2_rmin = atom2[6]
+                    atom2_epsilon = atom2[7]
+                    if atom2 is atom:
+                        lj_idx_list[j] = num_lj_types
+                    elif atom_type not in atom_nbfix_types:
+                        # Only non-NBFIXed atom types can be compressed
+                        ljtype2 = (atom2_rmin, atom2_epsilon)
+                        if ljtype == ljtype2:
+                            lj_idx_list[j] = num_lj_types
+
+            # Now everything is assigned. Create the A-coefficient and
+            # B-coefficient arrays
+            acoef = [0 for i in range(num_lj_types*num_lj_types)]
+            bcoef = acoef[:]
+            for i in range(num_lj_types):
+                namei = lj_type_list[i][0]
+                for j in range(num_lj_types):
+                    namej = lj_type_list[j][0]
+                    try:
+                        params = self._nonbondTypes[tuple(sorted((namei, namej)))]
+                        rij = float(params[3])
+                        wdij = float(params[4])
+                    except KeyError:
+                        rij = (lj_radii[i] + lj_radii[j]) / 2
+                        wdij = math.sqrt(lj_depths[i] * lj_depths[j])
+                    acoef[i+num_lj_types*j] = 2 * math.sqrt(wdij) * rij**6
+                    bcoef[i+num_lj_types*j] = 4 * wdij * rij**6
+            cforce = mm.CustomNonbondedForce('(a/r6)^2-b/r6; r6=r^6;'
+                                             'a=acoef(type1, type2);'
+                                             'b=bcoef(type1, type2)')
+            cforce.addTabulatedFunction('acoef',
+                    mm.Discrete2DFunction(num_lj_types, num_lj_types, acoef))
+            cforce.addTabulatedFunction('bcoef',
+                    mm.Discrete2DFunction(num_lj_types, num_lj_types, bcoef))
+            cforce.addPerParticleParameter('type')
+            if (nonbondedMethod in (ff.PME, ff.LJPME, ff.Ewald, ff.CutoffPeriodic)):
+                cforce.setNonbondedMethod(cforce.CutoffPeriodic)
+                cforce.setCutoffDistance(nonbondedCutoff)
+                cforce.setUseLongRangeCorrection(True)
+            elif nonbondedMethod is ff.NoCutoff:
+                cforce.setNonbondedMethod(cforce.NoCutoff)
+            elif nonbondedMethod is ff.CutoffNonPeriodic:
+                cforce.setNonbondedMethod(cforce.CutoffNonPeriodic)
+                cforce.setCutoffDistance(nonbondedCutoff)
+            else:
+                raise ValueError('Unrecognized nonbonded method')
+            for i in lj_idx_list:
+                cforce.addParticle((i - 1,)) # adjust for indexing from 0
+
+            for i in range(nb.getNumExceptions()):
+                ii, jj, q, eps, sig = nb.getExceptionParameters(i)
+                cforce.addExclusion(ii, jj)
+            sys.addForce(cforce)
+
         # Adjust masses.
 
         if hydrogenMass is not None: