feat(classical): implement nonbonded exclusions as pre-factor instead of counting them twice

dmff/api.py

@@ -6,6 +6,7 @@
 import numpy as np
 import jax.numpy as jnp
 from collections import defaultdict
+import xml.etree.ElementTree as ET
 from .admp.disp_pme import ADMPDispPmeForce
 from .admp.multipole import convert_cart2harm, rot_local2global
 from .admp.pairwise import TT_damping_qq_c6_kernel, generate_pairwise_interaction
@@ -31,6 +32,39 @@
 import sys
 
 
+class XMLNodeInfo:
+
+    class XMLElementInfo:
+
+        def __init__(self, name):
+            self.name = name
+            self.attributes = {}
+
+        def addAttribute(self, key, value):
+            self.attributes[key] = value
+
+
+    def __init__(self, name):
+        self.name = name
+        self.attributes = {}
+        self.elements = []
+
+
+    def addAttribute(self, key, value):
+        self.attributes[key] = value
+
+
+    def addElement(self, name, info):
+        element = self.XMLElementInfo(name)
+        for k, v in info.items():
+            element.addAttribute(k, v)
+            self.elements.append(element)
+
+
+    def modResidue(self, residue, atom, key, value):
+        pass
+
+
 def get_line_context(file_path, line_number):
     return linecache.getline(file_path, line_number).strip()
 
@@ -56,6 +90,15 @@ def build_covalent_map(data, max_neighbor):
     return covalent_map
 
 
+def findAtomTypeTexts(attribs, num):
+    typetxt = []
+    for n in range(1, num+1):
+        for key in ["type%i"%n, "class%i"%n]:
+            if key in attribs:
+                typetxt.append((key, attribs[key]))
+                break
+    return typetxt
+
 class ADMPDispGenerator:
     def __init__(self, hamiltonian):
         self.ff = hamiltonian
@@ -630,10 +673,13 @@ def __init__(self, hamiltonian):
         self.params = {"k": [], "length": []}
         self._jaxPotential = None
         self.types = []
+        self.typetexts = []
 
     def registerBondType(self, bond):
+        typetxt = findAtomTypeTexts(bond, 2)
         types = self.ff._findAtomTypes(bond, 2)
         self.types.append(types)
+        self.typetexts.append(typetxt)
         self.params["k"].append(float(bond["k"]))
         self.params["length"].append(float(bond["length"]))  # length := r0
 
@@ -647,10 +693,9 @@ def parseElement(element, hamiltonian):
               <HarmonicBondForce>
                 <Bond type1="ow" type2="hw" length="0.09572000000000001" k="462750.3999999999"/>
                 <Bond type1="hw" type2="hw" length="0.15136000000000002" k="462750.3999999999"/>
-              <\HarminicBondForce>
+              <\HarmonicBondForce>
 
         """
-
         generator = HarmonicBondJaxGenerator(hamiltonian)
         hamiltonian.registerGenerator(generator)
         for bondtype in element.findall("Bond"):
@@ -700,7 +745,18 @@ def getJaxPotential(self):
 
     def renderXML(self):
         # generate xml force field file
-        pass
+        finfo = XMLNodeInfo("HarmonicBondForce")
+        for ntype in range(len(self.types)):
+            binfo = {}
+            k1, v1 = self.typetexts[ntype][0]
+            k2, v2 = self.typetexts[ntype][1]
+            binfo[k1] = v1
+            binfo[k2] = v2
+            for key in self.params.keys():
+                binfo[key] = "%.8f"%self.params[key][ntype]
+            finfo.addElement("Bond", binfo)
+        return finfo
+
 
 
 # register all parsers
@@ -1450,7 +1506,10 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args):
             self.params[k] = jnp.array(self.params[k])
 
         mscales_coul = jnp.array([0.0, 0.0, 0.0, 1.0, 1.0, 1.0])  # mscale for PME
-        mscales_coul = mscales_coul.at[2].set(1.0 - self.params["coulomb14scale"][0])
+        mscales_coul = mscales_coul.at[2].set(self.params["coulomb14scale"][0])
+        mscales_lj = jnp.array([0.0, 0.0, 0.0, 1.0, 1.0, 1.0])  # mscale for LJ
+        mscales_lj = mscales_lj.at[2].set(self.params["lj14scale"][0])
+
 
         # Coulomb: only support PME for now
         # set PBC
@@ -1505,29 +1564,6 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args):
         map_nbfix = np.array(map_nbfix, dtype=int).reshape((-1, 2))
 
         colv_map = build_covalent_map(data, 6)
-        map_exclusion = []
-        scale_14 = []
-        npair = 0
-        for ii in range(colv_map.shape[0]):
-            for jj in range(ii + 1, colv_map.shape[1]):
-                if colv_map[ii, jj] > 0 and colv_map[ii, jj] < 4:
-                    map_exclusion.append((ii, jj))
-                    if colv_map[ii, jj] == 3:
-                        scale_14.append(npair)
-                    npair += 1
-
-        map_exclusion = np.array(map_exclusion, dtype=int)
-        scale_14 = np.array(scale_14, dtype=int)
-        scale_lj_exclusion = np.ones((map_exclusion.shape[0],))
-        scale_coul_exclusion = np.ones((map_exclusion.shape[0],))
-        scale_lj_exclusion = jnp.array(scale_lj_exclusion)
-        scale_coul_exclusion = jnp.array(scale_coul_exclusion)
-        scale_lj_exclusion = scale_lj_exclusion.at[scale_14].set(
-            1.0 - self.params["lj14scale"][0]
-        )
-        scale_coul_exclusion = scale_coul_exclusion.at[scale_14].set(
-            1.0 - self.params["coulomb14scale"][0]
-        )
 
         if unit.is_quantity(nonbondedCutoff):
             r_cut = nonbondedCutoff.value_in_unit(unit.nanometer)
@@ -1549,8 +1585,7 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args):
             r_cut,
             map_lj,
             map_nbfix,
-            map_exclusion,
-            scale_lj_exclusion,
+            colv_map,
             isSwitch=ifSwitch,
             isPBC=ifPBC,
             isNoCut=isNoCut,
@@ -1561,14 +1596,10 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args):
             # do not use PME
             if nonbondedMethod in [app.CutoffPeriodic, app.CutoffNonPeriodic]:
                 # use Reaction Field
-                coulforce = CoulReactionFieldForce(
-                    r_cut, map_charge, map_exclusion, scale_coul_exclusion, isPBC=ifPBC
-                )
+                coulforce = CoulReactionFieldForce(r_cut, map_charge, colv_map, isPBC=ifPBC)
             if nonbondedMethod is app.NoCutoff:
                 # use NoCutoff
-                coulforce = CoulNoCutoffForce(
-                    map_charge, map_exclusion, scale_coul_exclusion
-                )
+                coulforce = CoulNoCutoffForce(map_charge, colv_map)
         else:
             coulforce = CoulombPMEForce(box, r_cut, self.ethresh, colv_map)
 
@@ -1584,6 +1615,7 @@ def potential_fn(positions, box, pairs, params):
                 params["sigma"],
                 params["epsfix"],
                 params["sigfix"],
+                mscales_lj
             )
             coulE = coulenergy(positions, box, pairs, params["charge"], mscales_coul)
 
@@ -1604,6 +1636,9 @@ def renderXML(self):
 class Hamiltonian(app.forcefield.ForceField):
     def __init__(self, *xmlnames):
         super().__init__(*xmlnames)
+        # add a function to parse AtomTypes and Residues information
+        self._atomtypes = None
+        self._residues = None
         self._potentials = []
 
     def createPotential(
@@ -1629,6 +1664,23 @@ def createPotential(
                 pass
         return [p for p in self._potentials]
 
+    def render(self, filename):
+        root = ET.Element("ForceField")
+        forceInfos = [g.renderXML() for g in self._forces]
+        for finfo in forceInfos:
+            # create xml nodes
+            if finfo is not None:
+                node = ET.SubElement(root, finfo.name)
+                for key in finfo.attributes.keys():
+                    node.set(key, finfo.attributes[key])
+                for elem in finfo.elements:
+                    subnode = ET.SubElement(node, elem.name)
+                    for key in elem.attributes.keys():
+                        subnode.set(key, elem.attributes[key])
+
+        tree = ET.ElementTree(root)
+        tree.write(filename)
+
 
 if __name__ == "__main__":
     H = Hamiltonian("forcefield.xml")

dmff/classical/inter.py

@@ -19,8 +19,7 @@ def __init__(
         r_cut,
         map_prm,
         map_nbfix,
-        map_exclusion,
-        scale_exclusion,
+        colvmap,
         isSwitch=False,
         isPBC=True,
         isNoCut=False,
@@ -31,10 +30,9 @@ def __init__(
 
         self.map_prm = map_prm
         self.map_nbfix = map_nbfix
-        self.map_exclusion = map_exclusion
-        self.scale_exclusion = scale_exclusion
         self.ifPBC = isPBC
         self.ifNoCut = isNoCut
+        self.colvmap = colvmap
 
     def generate_get_energy(self):
         def get_LJ_energy(dr_vec, sig, eps, box):
@@ -60,7 +58,7 @@ def get_LJ_energy(dr_vec, sig, eps, box):
 
             return E
 
-        def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix):
+        def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix, mscales):
 
             eps_m1 = jnp.repeat(epsilon.reshape((-1, 1)), epsilon.shape[0], axis=1)
             eps_m2 = eps_m1.T
@@ -74,42 +72,33 @@ def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix):
             sig_mat = sig_mat.at[self.map_nbfix[:, 0], self.map_nbfix[:, 1]].set(sigfix)
             sig_mat = sig_mat.at[self.map_nbfix[:, 1], self.map_nbfix[:, 0]].set(sigfix)
 
+            colv_pair = self.colvmap[pairs[:,0],pairs[:,1]]
+            mscale_pair = mscales[colv_pair-1] # in mscale vector, the 0th item is 1-2 scale, the 1st item is 1-3 scale, etc...
+
+
             dr_vec = positions[pairs[:, 0]] - positions[pairs[:, 1]]
             prm_pair0 = self.map_prm[pairs[:, 0]]
             prm_pair1 = self.map_prm[pairs[:, 1]]
             eps = eps_mat[prm_pair0, prm_pair1]
             sig = sig_mat[prm_pair0, prm_pair1]
 
-            E_inter = get_LJ_energy(dr_vec, sig, eps, box)
-
-            # exclusion
-            dr_excl_vec = (
-                positions[self.map_exclusion[:, 0]]
-                - positions[self.map_exclusion[:, 1]]
-            )
-            excl_map0 = self.map_prm[self.map_exclusion[:, 0]]
-            excl_map1 = self.map_prm[self.map_exclusion[:, 1]]
-            eps_excl = eps_mat[excl_map0, excl_map1]
-            sig_excl = sig_mat[excl_map0, excl_map1]
+            eps_scale = eps * mscale_pair
 
-            E_excl = get_LJ_energy(dr_excl_vec, sig_excl, eps_excl, box)
-            E_excl = self.scale_exclusion * E_excl
+            E_inter = get_LJ_energy(dr_vec, sig, eps_scale, box)
 
-            return jnp.sum(E_inter) - jnp.sum(E_excl)
-            #return jnp.sum(E_inter)
+            return jnp.sum(E_inter)
 
         return get_energy
 
 
 class CoulNoCutoffForce:
     # E=\frac{{q}_{1}{q}_{2}}{4\pi\epsilon_0\epsilon_1 r}
 
-    def __init__(self, map_prm, map_exclusion, scale_exclusion, epsilon_1=1.0) -> None:
+    def __init__(self, map_prm, colvmap, epsilon_1=1.0) -> None:
 
         self.eps_1 = epsilon_1
         self.map_prm = map_prm
-        self.map_exclusion = map_exclusion
-        self.scale_exclusion = scale_exclusion
+        self.colvmap = colvmap
 
     def generate_get_energy(self):
         def get_coul_energy(dr_vec, chrgprod, box):
@@ -121,30 +110,21 @@ def get_coul_energy(dr_vec, chrgprod, box):
             return E
 
         def get_energy(positions, box, pairs, charges, mscales):
+
+            colv_pair = self.colvmap[pairs[:,0],pairs[:,1]]
+            mscale_pair = mscales[colv_pair-1]
+
             chrg_map0 = self.map_prm[pairs[:, 0]]
             chrg_map1 = self.map_prm[pairs[:, 1]]
             charge0 = charges[chrg_map0]
             charge1 = charges[chrg_map1]
             chrgprod = charge0 * charge1
+            chrgprod_scale = chrgprod * mscale_pair
             dr_vec = positions[pairs[:, 0]] - positions[pairs[:, 1]]
 
-            E_inter = get_coul_energy(dr_vec, chrgprod, box)
-
-            # exclusion
-            dr_excl_vec = (
-                positions[self.map_exclusion[:, 0]]
-                - positions[self.map_exclusion[:, 1]]
-            )
-            excl_map0 = self.map_prm[self.map_exclusion[:, 0]]
-            excl_map1 = self.map_prm[self.map_exclusion[:, 1]]
-            chrg0_excl = charges[excl_map0]
-            chrg1_excl = charges[excl_map1]
-            chrgprod_excl = chrg0_excl * chrg1_excl
+            E_inter = get_coul_energy(dr_vec, chrgprod_scale, box)
 
-            E_excl = get_coul_energy(dr_excl_vec, chrgprod_excl, box)
-            E_excl = self.scale_exclusion * E_excl
-
-            return jnp.sum(E_inter) - jnp.sum(E_excl)
+            return jnp.sum(E_inter) 
 
         return get_energy
 
@@ -155,8 +135,7 @@ def __init__(
         self,
         r_cut,
         map_prm,
-        map_exclusion,
-        scale_exclusion,
+        colvmap,
         epsilon_1=1.0,
         epsilon_solv=78.5,
         isPBC=True,
@@ -168,8 +147,7 @@ def __init__(
         self.exp_solv = epsilon_solv
         self.eps_1 = epsilon_1
         self.map_prm = map_prm
-        self.map_exclusion = map_exclusion
-        self.scale_exclusion = scale_exclusion
+        self.colvmap = colvmap
         self.ifPBC = isPBC
 
     def generate_get_energy(self):
@@ -191,30 +169,21 @@ def get_rf_energy(dr_vec, chrgprod, box):
             return E
 
         def get_energy(positions, box, pairs, charges, mscales):
+
+            colv_pair = self.colvmap[pairs[:,0],pairs[:,1]]
+            mscale_pair = mscales[colv_pair-1]
+
             chrg_map0 = self.map_prm[pairs[:, 0]]
             chrg_map1 = self.map_prm[pairs[:, 1]]
             charge0 = charges[chrg_map0]
             charge1 = charges[chrg_map1]
             chrgprod = charge0 * charge1
+            chrgprod_scale = chrgprod * mscale_pair
             dr_vec = positions[pairs[:, 0]] - positions[pairs[:, 1]]
 
-            E_inter = get_rf_energy(dr_vec, chrgprod, box)
-
-            # exclusion
-            dr_excl_vec = (
-                positions[self.map_exclusion[:, 0]]
-                - positions[self.map_exclusion[:, 1]]
-            )
-            excl_map0 = self.map_prm[self.map_exclusion[:, 0]]
-            excl_map1 = self.map_prm[self.map_exclusion[:, 1]]
-            chrg0_excl = charges[excl_map0]
-            chrg1_excl = charges[excl_map1]
-            chrgprod_excl = chrg0_excl * chrg1_excl
-
-            E_excl = get_rf_energy(dr_excl_vec, chrgprod_excl, box)
-            E_excl = self.scale_exclusion * E_excl
+            E_inter = get_rf_energy(dr_vec, chrgprod_scale, box)
 
-            return jnp.sum(E_inter) - jnp.sum(E_excl)
+            return jnp.sum(E_inter)
 
         return get_energy