Move all of the logic for GB parameter assignment into the new subclasses

Previously, parameters were defined in multiple places -- amber_file_parser.py
and charmmpsffile.py.

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

@@ -32,9 +32,7 @@ DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
-from __future__ import division
-from __future__ import absolute_import
-from __future__ import print_function
+from __future__ import division, absolute_import, print_function
 
 from functools import wraps
 from math import pi, cos, sin, sqrt
@@ -666,88 +664,6 @@ class CharmmPsfFile(object):
 
         return topology
 
-    def _get_gb_params(self, gb_model=HCT):
-        """ Gets the GB parameters. Need this method to special-case GB neck """
-        screen = [0 for atom in self.atom_list]
-        if gb_model is GBn:
-            radii = _bondi_radii(self.atom_list)
-            screen = [0.5 for atom in self.atom_list]
-            for i, atom in enumerate(self.atom_list):
-                if atom.type.atomic_number == 6:
-                    screen[i] = 0.48435382330
-                elif atom.type.atomic_number == 1:
-                    screen[i] = 1.09085413633
-                elif atom.type.atomic_number == 7:
-                    screen[i] = 0.700147318409
-                elif atom.type.atomic_number == 8:
-                    screen[i] = 1.06557401132
-                elif atom.type.atomic_number == 16:
-                    screen[i] = 0.602256336067
-        elif gb_model is GBn2:
-            radii = _mbondi3_radii(self.atom_list)
-            # Add non-optimized values as defaults
-            alpha = [1.0 for i in self.atom_list]
-            beta = [0.8 for i in self.atom_list]
-            gamma = [4.85 for i in self.atom_list]
-            screen = [0.5 for i in self.atom_list]
-            for i, atom in enumerate(self.atom_list):
-                if atom.type.atomic_number == 6:
-                    screen[i] = 1.058554
-                    alpha[i] = 0.733756
-                    beta[i] = 0.506378
-                    gamma[i] = 0.205844
-                elif atom.type.atomic_number == 1:
-                    screen[i] = 1.425952
-                    alpha[i] = 0.788440
-                    beta[i] = 0.798699
-                    gamma[i] = 0.437334
-                elif atom.type.atomic_number == 7:
-                    screen[i] = 0.733599
-                    alpha[i] = 0.503364
-                    beta[i] = 0.316828
-                    gamma[i] = 0.192915
-                elif atom.type.atomic_number == 8:
-                    screen[i] = 1.061039
-                    alpha[i] = 0.867814
-                    beta[i] = 0.876635
-                    gamma[i] = 0.387882
-                elif atom.type.atomic_number == 16:
-                    screen[i] = -0.703469
-                    alpha[i] = 0.867814
-                    beta[i] = 0.876635
-                    gamma[i] = 0.387882
-        else:
-            # Set the default screening parameters
-            for i, atom in enumerate(self.atom_list):
-                if atom.type.atomic_number == 1:
-                    screen[i] = 0.85
-                elif atom.type.atomic_number == 6:
-                    screen[i] = 0.72
-                elif atom.type.atomic_number == 7:
-                    screen[i] = 0.79
-                elif atom.type.atomic_number == 8:
-                    screen[i] = 0.85
-                elif atom.type.atomic_number == 9:
-                    screen[i] = 0.88
-                elif atom.type.atomic_number == 15:
-                    screen[i] = 0.86
-                elif atom.type.atomic_number == 16:
-                    screen[i] = 0.96
-                else:
-                    screen[i] = 0.8
-            # Determine which radii set we need
-            if gb_model is OBC1 or gb_model is OBC2:
-                radii = _mbondi2_radii(self.atom_list)
-            elif gb_model is HCT:
-                radii = _mbondi_radii(self.atom_list)
-
-        length_conv = u.angstrom.conversion_factor_to(u.nanometer)
-        radii = [x * length_conv for x in radii]
-
-        if gb_model is GBn2:
-            return zip(radii, screen, alpha, beta, gamma)
-        return zip(radii, screen)
-
     def createSystem(self, params, nonbondedMethod=ff.NoCutoff,
                      nonbondedCutoff=1.0*u.nanometer,
                      switchDistance=0.0*u.nanometer,
@@ -1249,8 +1165,6 @@ class CharmmPsfFile(object):
         # Add GB model if we're doing one
         if implicitSolvent is not None:
             if verbose: print('Adding GB parameters...')
-            gb_parms = self._get_gb_params(implicitSolvent)
-
             # If implicitSolventKappa is None, compute it from salt
             # concentration
             if implicitSolventKappa is None:
@@ -1289,6 +1203,7 @@ class CharmmPsfFile(object):
             elif implicitSolvent is GBn2:
                 gb = GBSAGBn2Force(solventDielectric, soluteDielectric, None,
                                    cutoff, kappa=implicitSolventKappa)
+            gb_parms = gb.getStandardParameters(self.topology)
             for atom, gb_parm in zip(self.atom_list, gb_parms):
                 gb.addParticle([atom.charge] + list(gb_parm))
             # Set cutoff method

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

@@ -34,8 +34,7 @@ DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
-from __future__ import absolute_import
-from __future__ import print_function
+from __future__ import absolute_import, print_function
 
 #=============================================================================================
 # GLOBAL IMPORTS
@@ -553,87 +552,6 @@ class PrmtopLoader(object):
             self._excludedAtoms.append(atomList)
         return self._excludedAtoms
 
-    def getGBParms(self, gbmodel, elements):
-        """Return list giving GB params, Radius and screening factor"""
-        gb_List=[]
-        radii=[float(r) for r in self._raw_data["RADII"]]
-        screen=[float(s) for s in self._raw_data["SCREEN"]]
-        if gbmodel == 'GBn2':
-            alpha = [0 for i in self._raw_data['RADII']]
-            beta = [0 for i in self._raw_data['RADII']]
-            gamma = [0 for i in self._raw_data['RADII']]
-        # Update screening parameters for GBn if specified
-        if gbmodel == 'GBn':
-            if elements is None:
-                raise Exception('GBn model requires element information')
-            for i, element in enumerate(elements):
-                if element is elem.carbon:
-                    screen[i] = 0.48435382330
-                elif element is elem.hydrogen:
-                    screen[i] = 1.09085413633
-                elif element is elem.nitrogen:
-                    screen[i] = 0.700147318409
-                elif element is elem.oxygen:
-                    screen[i] = 1.06557401132
-                elif element is elem.sulfur:
-                    screen[i] = 0.602256336067
-                else:
-                    screen[i] = 0.5
-                # radii is currently in Angstroms right now. GBn lookup tables
-                # only support radii between 1.0 and 2.0
-                if radii[i] < 1.0 or radii[i] > 2.0:
-                    raise ValueError('GBn requires intrinsic radii between 1 and '
-                                     '2 Angstroms (%.3f found for atom %d)' %
-                                     (radii[i], i))
-        if gbmodel == 'GBn2':
-            if elements is None:
-                raise Exception('GBn2 model requires element information')
-            for i, element in enumerate(elements):
-                if element is elem.carbon:
-                    screen[i] = 1.058554
-                    alpha[i] = 0.733756
-                    beta[i] = 0.506378
-                    gamma[i] = 0.205844
-                elif element is elem.hydrogen:
-                    screen[i] = 1.425952
-                    alpha[i] = 0.788440
-                    beta[i] = 0.798699
-                    gamma[i] = 0.437334
-                elif element is elem.nitrogen:
-                    screen[i] = 0.733599
-                    alpha[i] = 0.503364
-                    beta[i] = 0.316828
-                    gamma[i] = 0.192915
-                elif element is elem.oxygen:
-                    screen[i] = 1.061039
-                    alpha[i] = 0.867814
-                    beta[i] = 0.876635
-                    gamma[i] = 0.387882
-                elif element is elem.sulfur:
-                    screen[i] = -0.703469
-                    alpha[i] = 0.867814
-                    beta[i] = 0.876635
-                    gamma[i] = 0.387882
-                else: # not optimized
-                    screen[i] = 0.5
-                    alpha[i] = 1.0
-                    beta[i] = 0.8
-                    gamma[i] = 4.85
-                # radii is currently in Angstroms right now. GBn lookup tables
-                # only support radii between 1.0 and 2.0
-                if radii[i] < 1.0 or radii[i] > 2.0:
-                    raise ValueError('GBn2 requires intrinsic radii between 1 and '
-                                     '2 Angstroms (%.3f found for atom %d)' %
-                                     (radii[i], i))
-        lengthConversionFactor = units.angstrom.conversion_factor_to(units.nanometer)
-        if gbmodel == 'GBn2':
-            for rad, scr, alp, bet, gam in zip(radii, screen, alpha, beta, gamma):
-                gb_List.append((rad*lengthConversionFactor, scr, alp, bet, gam))
-        else:
-            for rad, scr in zip(radii, screen):
-                gb_List.append((rad*lengthConversionFactor, scr))
-        return gb_List
-
     def getBoxBetaAndDimensions(self):
         """Return periodic boundary box beta angle and dimensions"""
         beta=float(self._raw_data["BOX_DIMENSIONS"][0])
@@ -1053,7 +971,6 @@ def readAmberSystem(topology, prmtop_filename=None, prmtop_loader=None, shake=No
             cutoff = nonbondedCutoff
             if units.is_quantity(cutoff):
                 cutoff = cutoff.value_in_unit(units.nanometers)
-        gb_parms = prmtop.getGBParms(gbmodel, elements)
         if gbmodel == 'HCT':
             gb = customgb.GBSAHCTForce(solventDielectric, soluteDielectric, 'ACE', cutoff, implicitSolventKappa)
         elif gbmodel == 'OBC1':
@@ -1070,7 +987,24 @@ def readAmberSystem(topology, prmtop_filename=None, prmtop_loader=None, shake=No
         elif gbmodel == 'GBn2':
             gb = customgb.GBSAGBn2Force(solventDielectric, soluteDielectric, 'ACE', cutoff, implicitSolventKappa)
         else:
-            raise Exception("Illegal value specified for implicit solvent model")
+            raise ValueError("Illegal value specified for implicit solvent model")
+        gb_parms = gb.getStandardParameters(gbmodel, elements)
+        # Replace radii and screen, but screen *only* gets replaced by the
+        # prmtop contents for HCT, OBC1, and OBC2. GBn and GBn2 both override
+        # the prmtop screen factors from LEaP in sander and pmemd
+        if gbmodel in ('HCT', 'OBC1', 'OBC2'):
+            screen = [float(s) for s in self._raw_data['SCREEN']]
+        else:
+            screen = [gb_parm[1] for gb_parm in gb_parms]
+        radii = [float(r)/10 for r in self._raw_data['RADII']]
+        warned = False
+        for i, (r, s) in enumerate(zip(radii, screen)):
+            if abs(r - gb_parms[0]) > 1e-4 or abs(s - gb_parms[1]) > 1e-4:
+                if not warned:
+                    warnings.warn('Non-optimal GB parameters detected for GB '
+                                  'model %s' % gbmodel)
+                    warned = True
+            gb_parms[0], gb_parms[1] = r, s
 
         for charge, gb_parm in zip(charges, gb_parms):
             if gbmodel == 'OBC2' and implicitSolventKappa == 0:

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

@@ -29,10 +29,11 @@ OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
 
-from __future__ import division
-from __future__ import absolute_import
+from __future__ import division, absolute_import
 
+from collections import defaultdict
 import copy
+import simtk.openmm.app.element as E
 from simtk.openmm import CustomGBForce, Continuous2DFunction
 import simtk.unit as u
 
@@ -196,6 +197,141 @@ for i in range (len(d0)):
     d0[i]=d0[i]/10
     m0[i]=m0[i]*10
 
+def _get_bonded_atom_list(topology):
+    """ Returns a list of atoms bonded to each other atom in a dict """
+    bondeds = defaultdict(list)
+    for a1, a2 in topology.bonds():
+        bondeds[a1].append(a2)
+        bondeds[a2].append(a1)
+    return bondeds
+
+def _is_carboxylateO(atom, all_bonds):
+    if atom is not E.oxygen: return False
+    bondeds = all_bonds[atom]
+    if len(bondeds) != 1:
+        return False
+    if bondeds[0].element is not E.carbon:
+        return False
+    bondedsC = all_bonds[bondeds[0]]
+    if len(bondedsC) != 3:
+        return False
+    for a3 in bondedsC:
+        if a3 is atom: continue
+        if a3.element is E.oxygen:
+            break
+    else:
+        return False
+    # If we got here, must be a carboxylate
+    return True
+
+def _bondi_radii(topology):
+    """ Sets the bondi radii """
+    radii = [0.0 for atom in topology.getAtoms()]
+    for i, atom in enumerate(topology.getAtoms()):
+        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
+    return radii  # converted to nanometers above
+
+def _mbondi_radii(topology):
+    """ Sets the mbondi radii """
+    radii = [0.0 for atom in topology.getAtoms()]
+    all_bonds = _get_bonded_atom_list(topology)
+    for i, atom in enumerate(topology.getAtoms()):
+        # Radius of H atom depends on element it is bonded to
+        if atom.element in (E.hydrogen, E.deuterium):
+            bondeds = all_bonds[atom]
+            if bondeds[0].element in (E.carbon, E.nitrogen):
+                radii[i] = 1.3
+            elif bondeds[0].type.atomic_number in (E.oxygen, E.sulfur):
+                radii[i] = 0.8
+            else:
+                radii[i] = 1.2
+        # Radius of C atom depends on what type it is
+        elif atom.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
+    return radii  # converted to nanometers above
+
+def _mbondi2_radii(topology):
+    """ Sets the mbondi2 radii """
+    radii = [0.0 for atom in atom_list]
+    all_bonds = _get_bonded_atom_list(topology)
+    for i, atom in enumerate(atom_list):
+        # Radius of H atom depends on element it is bonded to
+        if atom.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:
+            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
+    return radii  # Converted to nanometers above
+
+def _mbondi3_radii(topology):
+    """ Sets the mbondi3 radii """
+    radii = _mbondi2_radii(topology)
+    all_bonds = _get_bonded_atom_list(topology)
+    for i, atom in enumerate(topology.getAtoms()):
+        # carboxylate and HH/HE (ARG)
+        if _is_carboxylateO(atom, all_bonds):
+            radii[i] = 1.4
+        elif atom.residue.name == 'ARG':
+            if atom.name.startswith('HH') or atom.name.startswith('HE'):
+                radii[i] = 1.17
+    return radii  # Converted to nanometers above
 
 def _createEnergyTerms(force, solventDielectric, soluteDielectric, SA, cutoff, kappa, offset):
     # Add the energy terms to the CustomGBForce.  These are identical for all the GB models.
@@ -231,6 +367,22 @@ def _createEnergyTerms(force, solventDielectric, soluteDielectric, SA, cutoff, k
             force.addEnergyTerm("-138.935485*(1/soluteDielectric-1/solventDielectric)*q1*q2*(1/f-"+str(1/cutoff)+");"
                                 "f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))"+params, CustomGBForce.ParticlePairNoExclusions)
 
+_SCREEN_PARAMETERS = { # normal, GBn, GBn2
+        E.hydrogen : (0.85, 1.09085413633, 1.425952),
+        E.carbon : (0.72, 0.48435382330, 1.058554),
+        E.nitrogen : (0.79, 0.700147318409, 0.733599),
+        E.oxygen : (0.85, 1.06557401132, 1.061039),
+        E.fluorine : (0.88, 0.5, 0.5),
+        E.phosphorus : (0.86, 0.5, 0.5),
+        E.sulfur : (0.96, 0.602256336067, -0.703469),
+        None : (0.8, 0.5, 0.5) # default
+}
+_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]
+
 class CustomAmberGBForce(CustomGBForce):
 
     OFFSET = 0.009
@@ -254,6 +406,24 @@ class CustomAmberGBForce(CustomGBForce):
         CustomGBForce.addParticle(self, params)
         return params
 
+    @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
+        """
+        raise NotImplementedError('Must override getStandardParameters in subclasses')
+
 """
 Amber Equivalent: igb = 1
 """
@@ -275,6 +445,13 @@ class GBSAHCTForce(CustomAmberGBForce):
         self.addComputedValue("B", "1/(1/or-I)", CustomGBForce.SingleParticle)
         _createEnergyTerms(self, solventDielectric, soluteDielectric, SA, cutoff, kappa, 0.009)
 
+    @staticmethod
+    def getStandardParameters(topology):
+        radii = [[x/10] for x in _mbondi_radii(topology)]
+        for i, atom in zip(topology.getAtoms()):
+            radii[i].append(_screen_parameter(atom)[0])
+        return radii
+
 """
 Amber Equivalents: igb = 2
 """
@@ -297,10 +474,17 @@ class GBSAOBC1Force(CustomAmberGBForce):
                                    "psi=I*or; radius=or+offset; offset=0.009", CustomGBForce.SingleParticle)
         _createEnergyTerms(self, solventDielectric, soluteDielectric, SA, cutoff, kappa, 0.009)
 
+    @staticmethod
+    def getStandardParameters(topology):
+        radii = [[x/10] for x in _mbondi2_radii(topology)]
+        for i, atom in zip(topology.getAtoms()):
+            radii[i].append(_screen_parameter(atom)[0])
+        return radii
+
 """
 Amber Equivalents: igb = 5
 """
-class GBSAOBC2Force(CustomAmberGBForce):
+class GBSAOBC2Force(GBSAOBC1Force):
 
     def __init__(self, solventDielectric=78.5, soluteDielectric=1, SA=None,
                  cutoff=None, kappa=0.0):
@@ -359,6 +543,13 @@ class GBSAGBnForce(CustomAmberGBForce):
         if parameters[1] < 0.1 or parameters[1] > 0.2:
             raise ValueError('Radii must be between 1 and 2 Angstroms for neck lookup')
 
+    @staticmethod
+    def getStandardParameters(topology):
+        radii = [[x/10] for x in _bondi_radii(topology)]
+        for i, atom in zip(topology.getAtoms()):
+            radii[i].append(_screen_parameter(atom)[1])
+        return radii
+
 """
 Amber Equivalents: igb = 8
 """
@@ -393,3 +584,22 @@ class GBSAGBn2Force(GBSAGBnForce):
         self.addComputedValue("B", "1/(1/or-tanh(alpha*psi-beta*psi^2+gamma*psi^3)/radius);"
                                      "psi=I*or; radius=or+offset; offset=0.0195141", CustomGBForce.SingleParticle)
         _createEnergyTerms(self, solventDielectric, soluteDielectric, SA, cutoff, kappa, 0.0195141)
+
+    @staticmethod
+    def getStandardParameters(topology):
+        radii = [[x/10] for x in _mbondi3_radii(topology)]
+        for i, atom in zip(topology.getAtoms()):
+            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])
+        return radii

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

@@ -71,20 +71,17 @@ class Topology(object):
     def getNumAtoms(self):
         """Return the number of atoms in the Topology.
         """
-        natom = self._numAtoms
-        return natom
+        return self._numAtoms
 
     def getNumResidues(self):
         """Return the number of residues in the Topology.
         """
-        nres = self._numResidues
-        return nres
+        return self._numResidues
 
     def getNumChains(self):
         """Return the number of chains in the Topology.
         """
-        nchain = len(self._chains)
-        return nchain
+        return len(self._chains)
 
     def addChain(self, id=None):
         """Create a new Chain and add it to the Topology.