Numpyfication of custom GB forces

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

@@ -34,6 +34,7 @@ from __future__ import division, absolute_import
 
 from collections import defaultdict
 import copy
+import numpy
 from simtk.openmm.app import element as E
 from simtk.openmm import CustomGBForce, Discrete2DFunction
 import simtk.unit as u
@@ -48,7 +49,8 @@ def strip_unit(value, unit):
 
 
 # GBn and GBn2 lookup tables
-d0=[2.26685,2.32548,2.38397,2.44235,2.50057,2.55867,2.61663,2.67444,
+d0= numpy.array(
+   [2.26685,2.32548,2.38397,2.44235,2.50057,2.55867,2.61663,2.67444,
     2.73212,2.78965,2.84705,2.9043,2.96141,3.0184,3.07524,3.13196,
     3.18854,3.24498,3.30132,3.35752,3.4136,
     2.31191,2.37017,2.4283,2.48632,2.5442,2.60197,2.65961,2.71711,
@@ -110,9 +112,11 @@ d0=[2.26685,2.32548,2.38397,2.44235,2.50057,2.55867,2.61663,2.67444,
     4.04598,4.09974,4.15347,4.20715,4.26078,
     3.22855,3.28307,3.33751,3.39188,3.4462,3.50046,3.55466,3.6088,
     3.6629,3.71694,3.77095,3.82489,3.8788,3.93265,3.98646,4.04022,
-    4.09395,4.14762,4.20126,4.25485,4.3084]
+    4.09395,4.14762,4.20126,4.25485,4.3084],
+    numpy.double)
 
-m0=[0.0381511,0.0338587,0.0301776,0.027003,0.0242506,0.0218529,
+m0=numpy.array(
+   [0.0381511,0.0338587,0.0301776,0.027003,0.0242506,0.0218529,
     0.0197547,0.0179109,0.0162844,0.0148442,0.0135647,0.0124243,
     0.0114047,0.0104906,0.00966876,0.008928,0.0082587,0.00765255,
     0.00710237,0.00660196,0.00614589,
@@ -195,10 +199,11 @@ m0=[0.0381511,0.0338587,0.0301776,0.027003,0.0242506,0.0218529,
     0.0609642,0.0546169,0.0491183,0.0443295,0.0401388,0.036455,
     0.0332033,0.030322,0.0277596,0.0254732,0.0234266,0.0215892,
     0.0199351,0.018442,0.0170909,0.0158654,0.0147514,0.0137365,
-    0.0128101,0.0119627,0.0111863]
+    0.0128101,0.0119627,0.0111863],
+    numpy.double)
 # Rescale to nanometers
-d0 = [d / 10. for d in d0]
-m0 = [m * 10. for m in m0]
+d0/=10
+m0*=10
 
 
 def _get_bonded_atom_list(topology):
@@ -232,38 +237,45 @@ def _is_carboxylateO(atom, all_bonds):
 
 def _bondi_radii(topology):
     """ Sets the bondi radii """
-    radii = [0.0 for atom in topology.atoms()]
+    default_radius = 1.5
+    element_to_radius = {
+        E.carbon:       1.7,
+        E.hydrogen:     1.2,
+        E.deuterium:    1.2,
+        E.nitrogen:     1.55,
+        E.oxygen:       1.5,
+        E.fluorine:     1.5,
+        E.silicon:      2.1,
+        E.phosphorus:   1.85,
+        E.sulfur:       1.8,
+        E.chlorine:     1.5,
+    }
+    natoms = topology.getNumAtoms()
+    radii = numpy.empty(topology.getNumAtoms(), numpy.double)
     for i, atom in enumerate(topology.atoms()):
-        if atom.element is E.carbon:
-            radii[i] = 1.7
-        elif atom.element in (E.hydrogen, E.deuterium):
-            radii[i] = 1.2
-        elif atom.element is E.nitrogen:
-            radii[i] = 1.55
-        elif atom.element is E.oxygen:
-            radii[i] = 1.5
-        elif atom.element is E.fluorine:
-            radii[i] = 1.5
-        elif atom.element is E.silicon:
-            radii[i] = 2.1
-        elif atom.element is E.phosphorus:
-            radii[i] = 1.85
-        elif atom.element is E.sulfur:
-            radii[i] = 1.8
-        elif atom.element is E.chlorine:
-            radii[i] = 1.5
-        else:
-            radii[i] = 1.5
+        radii[i] = element_to_radius.get(atom.element, default_radius)
     return radii  # converted to nanometers above
 
 
-def _mbondi_radii(topology):
+def _mbondi_radii(topology, all_bonds = None):
     """ Sets the mbondi radii """
-    radii = [0.0 for atom in topology.atoms()]
-    all_bonds = _get_bonded_atom_list(topology)
+    default_radius = 1.5
+    element_to_const_radius = {
+        E.nitrogen:     1.55,
+        E.oxygen:       1.5,
+        E.fluorine:     1.5,
+        E.silicon:      2.1,
+        E.phosphorus:   1.85,
+        E.sulfur:       1.8,
+        E.chlorine:     1.5,
+    }
+    radii = numpy.empty(topology.getNumAtoms(), numpy.double)
+    if all_bonds is None:
+        all_bonds = _get_bonded_atom_list(topology)
     for i, atom in enumerate(topology.atoms()):
         # Radius of H atom depends on element it is bonded to
-        if atom.element in (E.hydrogen, E.deuterium):
+        element = atom.element
+        if element in (E.hydrogen, E.deuterium):
             bondeds = all_bonds[atom]
             if bondeds[0].element in (E.carbon, E.nitrogen):
                 radii[i] = 1.3
@@ -272,67 +284,51 @@ def _mbondi_radii(topology):
             else:
                 radii[i] = 1.2
         # Radius of C atom depends on what type it is
-        elif atom.element is E.carbon:
+        elif element is E.carbon:
             radii[i] = 1.7
         # All other elements have fixed radii for all types/partners
-        elif atom.element is E.nitrogen:
-            radii[i] = 1.55
-        elif atom.element is E.oxygen:
-            radii[i] = 1.5
-        elif atom.element is E.fluorine:
-            radii[i] = 1.5
-        elif atom.element is E.silicon:
-            radii[i] = 2.1
-        elif atom.element is E.phosphorus:
-            radii[i] = 1.85
-        elif atom.element is E.sulfur:
-            radii[i] = 1.8
-        elif atom.element is E.chlorine:
-            radii[i] = 1.5
-        else:
-            radii[i] = 1.5
+        radii[i] = element_to_const_radius.get(element, default_radius)
     return radii  # converted to nanometers above
 
 
-def _mbondi2_radii(topology):
+def _mbondi2_radii(topology, all_bonds = None):
     """ Sets the mbondi2 radii """
-    radii = [0.0 for atom in topology.atoms()]
-    all_bonds = _get_bonded_atom_list(topology)
+    default_radius = 1.5
+    element_to_const_radius = {
+        E.nitrogen:     1.55,
+        E.oxygen:       1.5,
+        E.fluorine:     1.5,
+        E.silicon:      2.1,
+        E.phosphorus:   1.85,
+        E.sulfur:       1.8,
+        E.chlorine:     1.5,
+    }
+    radii = numpy.empty(topology.getNumAtoms(), numpy.double)
+    if all_bonds is None:
+        all_bonds = _get_bonded_atom_list(topology)
     for i, atom in enumerate(topology.atoms()):
+        element = atom.element
         # Radius of H atom depends on element it is bonded to
-        if atom.element in (E.hydrogen, E.deuterium):
+        if element in (E.hydrogen, E.deuterium):
             bondeds = all_bonds[atom]
             if bondeds[0].element is E.nitrogen:
                 radii[i] = 1.3
             else:
                 radii[i] = 1.2
         # Radius of C atom depends on what type it is
-        elif atom.element is E.carbon:
+        elif element is E.carbon:
             radii[i] = 1.7
         # All other elements have fixed radii for all types/partners
-        elif atom.element is E.nitrogen:
-            radii[i] = 1.55
-        elif atom.element is E.oxygen:
-            radii[i] = 1.5
-        elif atom.element is E.fluorine:
-            radii[i] = 1.5
-        elif atom.element is E.silicon:
-            radii[i] = 2.1
-        elif atom.element == E.phosphorus:
-            radii[i] = 1.85
-        elif atom.element == E.sulfur:
-            radii[i] = 1.8
-        elif atom.element == E.chlorine:
-            radii[i] = 1.5
         else:
-            radii[i] = 1.5
+            radii[i] = element_to_const_radius.get(element, default_radius)
     return radii  # Converted to nanometers above
 
 
-def _mbondi3_radii(topology):
+def _mbondi3_radii(topology, all_bonds = None):
     """ Sets the mbondi3 radii """
-    radii = _mbondi2_radii(topology)
-    all_bonds = _get_bonded_atom_list(topology)
+    if all_bonds is None:
+        all_bonds = _get_bonded_atom_list(topology)
+    radii = _mbondi2_radii(topology, all_bonds=all_bonds)
     for i, atom in enumerate(topology.atoms()):
         # carboxylate and HH/HE (ARG)
         if _is_carboxylateO(atom, all_bonds):
@@ -391,12 +387,12 @@ _SCREEN_PARAMETERS = { # normal, GBn, GBn2
         E.sulfur : (0.96, 0.602256336067, -0.703469),
         None : (0.8, 0.5, 0.5) # default
 }
-
+for key, val in _SCREEN_PARAMETERS.items():
+    _SCREEN_PARAMETERS[key] = numpy.array(val, numpy.double)
 _SCREEN_PARAMETERS[E.deuterium] = _SCREEN_PARAMETERS[E.hydrogen]
+
 def _screen_parameter(atom):
-    if atom.element in _SCREEN_PARAMETERS:
-        return _SCREEN_PARAMETERS[atom.element]
-    return _SCREEN_PARAMETERS[None]
+    return _SCREEN_PARAMETERS.get(atom.element, _SCREEN_PARAMETERS[None])
 
 
 class CustomAmberGBForceBase(CustomGBForce):
@@ -439,6 +435,33 @@ class CustomAmberGBForceBase(CustomGBForce):
         self.parameters.append(params)
         return params
 
+    def addParticles(self, params):
+        """Add a set of particles to the force
+
+        Particles are added in order. The total number of particles added to the
+        force must match the number of particles in the system.
+
+        Parameters
+        ----------
+        params : list or numpy array
+            A (natoms * npar) dimensional array of parameters to add to the
+            force. The meaning of the parameters depends on the model. All
+            parameters should be simple floating-point values in OpenMM's
+            default units.
+
+        Returns
+        -------
+        None
+        """
+        if isinstance(params, list):
+            params = numpy.array(params)
+        else:
+            params = copy.deepcopy(params)
+
+        params[:,self.RADIUS_ARG_POSITION] -= self.OFFSET
+        params[:,self.SCREEN_POSITION] *= params[:,self.RADIUS_ARG_POSITION]
+        self.parameters.extend(params.tolist())
+
     @staticmethod
     def getStandardParameters(topology):
         """ Gets list of standard parameters for this GB model based on an input Topology
@@ -528,9 +551,10 @@ class GBSAHCTForce(CustomAmberGBForceBase):
             at the beginning of the list
 
         """
-        radii = [[x/10] for x in _mbondi_radii(topology)]
-        for i, atom in enumerate(topology.atoms()):
-            radii[i].append(_screen_parameter(atom)[0])
+        radii = numpy.empty((topology.getNumAtoms(), 2), numpy.double)
+        radii[:,0] = _mbondi_radii(topology)/10
+        for rad, atom in zip(radii, topology.atoms()):
+            rad[1] = _screen_parameter(atom)[0]
         return radii
 
 
@@ -590,9 +614,10 @@ class GBSAOBC1Force(CustomAmberGBForceBase):
             at the beginning of the list
 
         """
-        radii = [[x/10] for x in _mbondi2_radii(topology)]
-        for i, atom in enumerate(topology.atoms()):
-            radii[i].append(_screen_parameter(atom)[0])
+        radii = numpy.empty((topology.getNumAtoms(), 2), numpy.double)
+        radii[:,0] = _mbondi2_radii(topology)/10
+        for rad, atom in zip(radii, topology.atoms()):
+            rad[1] = _screen_parameter(atom)[0]
         return radii
 
 
@@ -678,6 +703,14 @@ class GBSAGBnForce(CustomAmberGBForceBase):
         if parameters[1] + self.OFFSET < 0.1 or parameters[1] + self.OFFSET > 0.2:
             raise ValueError('Radii must be between 1 and 2 Angstroms for neck lookup')
 
+    def addParticles(self, parameters):
+        if isinstance(parameters, list):
+            parameters = numpy.array(parameters)
+        radii = parameters[:,1]
+        if numpy.any(numpy.logical_or(radii<0.1, radii>0.2)):
+            raise ValueError('Radii must be between 1 and 2 Angstroms for neck lookup')
+        CustomAmberGBForceBase.addParticles(self, parameters)
+
     def finalize(self):
         self._findUniqueRadii()
         self._createRadiusToIndexMap()
@@ -701,9 +734,10 @@ class GBSAGBnForce(CustomAmberGBForceBase):
             at the beginning of the list
 
         """
-        radii = [[x/10] for x in _bondi_radii(topology)]
-        for i, atom in enumerate(topology.atoms()):
-            radii[i].append(_screen_parameter(atom)[1])
+        radii = numpy.empty((topology.getNumAtoms(), 2), numpy.double)
+        radii[:,0] = _bondi_radii(topology)/10
+        for rad, atom in zip(radii, topology.atoms()):
+            rad[1] = _screen_parameter(atom)[1]
         return radii
 
     def _findUniqueRadii(self):
@@ -799,8 +833,57 @@ class GBSAGBn2Force(GBSAGBnForce):
     def __init__(self, solventDielectric=78.5, soluteDielectric=1, SA=None, cutoff=None, kappa=0.0):
         GBSAGBnForce.__init__(self, solventDielectric, soluteDielectric, SA, cutoff, kappa)
 
-    @staticmethod
-    def getStandardParameters(topology):
+
+    # @staticmethod
+    # def getStandardParameters(topology):
+    #     """ Gets list of standard parameters for this GB model based on an input Topology
+    #
+    #     Parameters
+    #     ----------
+    #     topology : simtk.openmm.app.Topology
+    #         Topology of the system to get parameters for
+    #
+    #     Returns
+    #     -------
+    #     list of float
+    #         List of all parameters needed for this GB model. These can be passed
+    #         to addParticle or setParticleParameters after the charge is inserted
+    #         at the beginning of the list
+    #
+    #     """
+    #     import numpy
+    #     natoms = topology.getNumAtoms()
+    #     radii = numpy.empty((6,natoms), numpy.double)
+    #     radii[:,0] = [[x/10] for x in _mbondi3_radii(topology)]
+    #     for i, atom in enumerate(topology.atoms()):
+    #         radii[i,1] = _screen_parameter(atom)[2]
+    #         e = atom.element
+    #         if e in (E.hydrogen, E.deuterium):
+    #             radii[i, 2:]= [0.788440, 0.798699, 0.437334]
+    #         elif atom.element is E.carbon:
+    #             radii[i].extend([0.733756, 0.506378, 0.205844])
+    #         elif atom.element is E.nitrogen:
+    #             radii[i].extend([0.503364, 0.316828, 0.192915])
+    #         elif atom.element is E.oxygen:
+    #             radii[i].extend([0.867814, 0.876635, 0.387882])
+    #         elif atom.element is E.sulfur:
+    #             radii[i].extend([0.867814, 0.876635, 0.387882])
+    #         else:
+    #             radii[i].extend([0.8, 4.85, 0.5])
+    #     return radii
+
+    _atom_params = {
+        E.hydrogen:     numpy.array([0.788440, 0.798699, 0.437334]),
+        E.deuterium:    numpy.array([0.788440, 0.798699, 0.437334]),
+        E.carbon:       numpy.array([0.733756, 0.506378, 0.205844]),
+        E.nitrogen:     numpy.array([0.503364, 0.316828, 0.192915]),
+        E.oxygen:       numpy.array([0.867814, 0.876635, 0.387882]),
+        E.sulfur:       numpy.array([0.867814, 0.876635, 0.387882]),
+    }
+    _default_atom_params = numpy.array([0.8, 4.85, 0.5]);
+
+    @classmethod
+    def getStandardParameters(cls, topology):
         """ Gets list of standard parameters for this GB model based on an input Topology
 
         Parameters
@@ -816,21 +899,12 @@ class GBSAGBn2Force(GBSAGBnForce):
             at the beginning of the list
 
         """
-        radii = [[x/10] for x in _mbondi3_radii(topology)]
-        for i, atom in enumerate(topology.atoms()):
-            radii[i].append(_screen_parameter(atom)[2])
-            if atom.element in (E.hydrogen, E.deuterium):
-                radii[i].extend([0.788440, 0.798699, 0.437334])
-            elif atom.element is E.carbon:
-                radii[i].extend([0.733756, 0.506378, 0.205844])
-            elif atom.element is E.nitrogen:
-                radii[i].extend([0.503364, 0.316828, 0.192915])
-            elif atom.element is E.oxygen:
-                radii[i].extend([0.867814, 0.876635, 0.387882])
-            elif atom.element is E.sulfur:
-                radii[i].extend([0.867814, 0.876635, 0.387882])
-            else:
-                radii[i].extend([0.8, 4.85, 0.5])
+        natoms = topology.getNumAtoms()
+        radii = numpy.empty([natoms,5], numpy.double)
+        radii[:,0] = _mbondi3_radii(topology)/10
+        for atom, rad in zip(topology.atoms(), radii):
+            rad[1] = _screen_parameter(atom)[2]
+            rad[2:] = cls._atom_params.get(atom.element, cls._default_atom_params)
         return radii
 
     def _addEnergyTerms(self):