Add Potential Search with Wildcards in a ForceField (#519)

* Initial Commit for Wildcards support in Potential Search in forcefield [skip-ci] * WIP- Add Circular Masking, ImproperType Search * Fix typo and add warns flag * WIP- Fix format strings. consolidate warning/error messages * WIP- Remove redundant typing import * Fix syntax * WIP- Add basic testing for atom_types; minor quirk fixes * WIP- Minor fixes to function signatures; more tests * WIP- Add More tests for angle/dihedrals * WIP- Revert formatting * WIP- Remove unused deepcopy import * WIP- Additional tests; set warn to false (default) Co-authored-by: Co Quach <[email protected]> Co-authored-by: Co Quach <[email protected]>
mosdef-hub · Apr 13, 2021 · f91e2f5 · f91e2f5
1 parent 8264258
commit f91e2f5
Show file tree

Hide file tree

Showing 9 changed files with 597 additions and 19 deletions.
diff --git a/gmso/core/forcefield.py b/gmso/core/forcefield.py
@@ -1,12 +1,18 @@
-import typing
 from collections import ChainMap
+from typing import Iterable
+import warnings
 
 from lxml import etree
 
-from gmso.utils.ff_utils import (validate,
-                                 parse_ff_metadata,
-                                 parse_ff_atomtypes,
-                                 parse_ff_connection_types)
+from gmso.exceptions import MissingPotentialError
+from gmso.utils._constants import FF_TOKENS_SEPARATOR
+from gmso.utils.ff_utils import (
+    parse_ff_atomtypes,
+    parse_ff_connection_types,
+    parse_ff_metadata,
+    validate,
+)
+from gmso.utils.misc import mask_with, validate_type
 
 
 def _group_by_expression(potential_types):
@@ -183,6 +189,225 @@ def group_improper_types_by_expression(self):
         """
         return _group_by_expression(self.improper_types)
 
+    def get_potential(self, group, key, warn=False):
+        """Returns a specific potential by key in this ForceField
+
+        Parameters
+        ----------
+        group:  {'atom_types', 'bond_types', 'angle_types', 'dihedral_types', 'improper_types'}
+            The potential group to perform this search on
+        key: str or list of str
+            The key to lookup for this potential group
+        warn: bool, default=False
+            If true, raise a warning instead of Error if no match found
+
+        Returns
+        -------
+        gmso.ParametricPotential
+            The parametric potential requested
+
+        Raises
+        ------
+        MissingPotentialError
+            When the potential specified by `key` is not found in the ForceField
+            potential group `group`
+        """
+        group = group.lower()
+
+        potential_extractors = {
+            "atom_type": self._get_atom_type,
+            "bond_type": self._get_bond_type,
+            "angle_type": self._get_angle_type,
+            "dihedral_type": self._get_dihedral_type,
+            "improper_type": self._get_improper_type,
+        }
+
+        if group not in potential_extractors:
+            raise ValueError(f"Cannot get potential for {group}")
+
+        validate_type(
+            [key] if isinstance(key, str) or not isinstance(key, Iterable) else key, str
+        )
+
+        return potential_extractors[group](key, warn=warn)
+
+    def get_parameters(self, group, key, warn=False, copy=False):
+        """Returns parameters for a specific potential by key in this ForceField
+
+        This function uses the `get_potential` function to get Parameters
+
+        See Also
+        --------
+        gmso.ForceField.get_potential
+            Get specific potential/parameters from a forcefield potential group by key
+        """
+        potential = self.get_potential(group, key, warn=warn)
+        return potential.get_parameters(copy=copy)
+
+    def _get_atom_type(self, atom_type, warn=False):
+        """Get a particular atom_type with given `atom_type` from this ForceField"""
+        if isinstance(atom_type, list):
+            atom_type = atom_type[0]
+
+        if not self.atom_types.get(atom_type):
+            msg = f"AtomType {atom_type} is not present in the ForceField"
+            if warn:
+                warnings.warn(msg)
+            else:
+                raise MissingPotentialError(msg)
+
+        return self.atom_types.get(atom_type)
+
+    def _get_bond_type(self, atom_types, warn=False):
+        """Get a particular bond_type between `atom_types` from this ForceField"""
+        if len(atom_types) != 2:
+            raise ValueError(
+                f"BondType potential can only "
+                f"be extracted for two atoms. Provided {len(atom_types)}"
+            )
+
+        forward = FF_TOKENS_SEPARATOR.join(atom_types)
+        reverse = FF_TOKENS_SEPARATOR.join(reversed(atom_types))
+        if forward in self.bond_types:
+            return self.bond_types[forward]
+        if reverse in self.bond_types:
+            return self.bond_types[reverse]
+
+        msg = f"BondType between atoms {atom_types[0]} and {atom_types[1]} " \
+              f"is missing from the ForceField"
+        if warn:
+            warnings.warn(msg)
+            return None
+        else:
+            raise MissingPotentialError(msg)
+
+    def _get_angle_type(self, atom_types, warn=False):
+        """Get a particular angle_type between `atom_types` from this ForceField"""
+        if len(atom_types) != 3:
+            raise ValueError(
+                f"AngleType potential can only "
+                f"be extracted for three atoms. Provided {len(atom_types)}"
+            )
+
+        forward = FF_TOKENS_SEPARATOR.join(atom_types)
+        reverse = FF_TOKENS_SEPARATOR.join(reversed(atom_types))
+        match = None
+        if forward in self.angle_types:
+            match = self.angle_types[forward]
+        if reverse in self.angle_types:
+            match = self.angle_types[reverse]
+
+        msg = f"AngleType between atoms {atom_types[0]}, {atom_types[1]} " \
+              f"and {atom_types[2]} is missing from the ForceField"
+
+        if match:
+            return match
+        elif warn:
+            warnings.warn(msg)
+            return None
+        else:
+            raise MissingPotentialError(msg)
+
+    def _get_dihedral_type(self, atom_types, warn=False):
+        """Get a particular dihedral_type between `atom_types` from this ForceField"""
+        if len(atom_types) != 4:
+            raise ValueError(
+                f"DihedralType potential can only "
+                f"be extracted for four atoms. Provided {len(atom_types)}"
+            )
+
+        forward = FF_TOKENS_SEPARATOR.join(atom_types)
+        reverse = FF_TOKENS_SEPARATOR.join(reversed(atom_types))
+
+        if forward is self.dihedral_types:
+            return self.dihedral_types[forward]
+        if reverse in self.dihedral_types:
+            return self.dihedral_types[reverse]
+
+        match = None
+        for i in range(1, 5):
+            forward_patterns = mask_with(atom_types, i)
+            reverse_patterns = mask_with(reversed(atom_types), i)
+
+            for forward_pattern, reverse_pattern in zip(
+                forward_patterns, reverse_patterns
+            ):
+                forward_match_key = FF_TOKENS_SEPARATOR.join(forward_pattern)
+                reverse_match_key = FF_TOKENS_SEPARATOR.join(reverse_pattern)
+
+                if forward_match_key in self.dihedral_types:
+                    match = self.dihedral_types[forward_match_key]
+                    break
+
+                if reverse_match_key in self.dihedral_types:
+                    match = self.dihedral_types[reverse_match_key]
+                    break
+
+            if match:
+                break
+
+        msg = f"DihedralType between atoms {atom_types[0]}, {atom_types[1]}, "\
+              f"{atom_types[2]} and {atom_types[3]} is missing from the ForceField."
+        if match:
+            return match
+        elif warn:
+            warnings.warn(msg)
+            return None
+        else:
+            raise MissingPotentialError(msg)
+
+    def _get_improper_type(self, atom_types, warn=False):
+        """Get a particular improper_type between `atom_types` from this ForceField"""
+        if len(atom_types) != 4:
+            raise ValueError(
+                f"ImproperType potential can only "
+                f"be extracted for four atoms. Provided {len(atom_types)}"
+            )
+
+        forward = FF_TOKENS_SEPARATOR.join(atom_types)
+        reverse = FF_TOKENS_SEPARATOR.join(
+            [atom_types[0], atom_types[2], atom_types[1], atom_types[3]]
+        )
+
+        if forward is self.improper_types:
+            return self.improper_types[forward]
+        if reverse in self.improper_types:
+            return self.improper_types[reverse]
+
+        match = None
+        for i in range(1, 5):
+            forward_patterns = mask_with(atom_types, i)
+            reverse_patterns = mask_with(
+                [atom_types[0], atom_types[2], atom_types[1], atom_types[3]], i
+            )
+
+            for forward_pattern, reverse_pattern in zip(
+                forward_patterns, reverse_patterns
+            ):
+                forward_match_key = FF_TOKENS_SEPARATOR.join(forward_pattern)
+                reverse_match_key = FF_TOKENS_SEPARATOR.join(reverse_pattern)
+
+                if forward_match_key in self.dihedral_types:
+                    match = self.dihedral_types[forward_match_key]
+                    break
+
+                if reverse_match_key in self.dihedral_types:
+                    match = self.dihedral_types[reverse_match_key]
+                    break
+
+            if match:
+                break
+
+        msg = f"ImproperType between atoms {atom_types[0]}, {atom_types[1]}, "\
+              f"{atom_types[2]} and {atom_types[3]} is missing from the ForceField."
+        if match:
+            return match
+        elif warn:
+            warnings.warn(msg)
+            return None
+        else:
+            raise MissingPotentialError(msg)
+
     def __repr__(self):
         return f"<ForceField {self.name},\n " \
                f"{len(self.atom_types)} AtomTypes,\n " \
@@ -219,7 +444,7 @@ def from_xml(cls, xmls_or_etrees, strict=True, greedy=True):
             A gmso.Forcefield object with a collection of Potential objects
             created using the information in the XML file
         """
-        if not isinstance(xmls_or_etrees, typing.Iterable) or isinstance(xmls_or_etrees, str):
+        if not isinstance(xmls_or_etrees, Iterable) or isinstance(xmls_or_etrees, str):
             xmls_or_etrees = [xmls_or_etrees]
 
         should_parse_xml = False

diff --git a/gmso/core/parametric_potential.py b/gmso/core/parametric_potential.py
@@ -124,6 +124,18 @@ def set_expression(self, expression=None, parameters=None, independent_variables
             parameters=parameters
         )
 
+    def get_parameters(self, copy=False):
+        """Returns parameters for this ParametricPotential"""
+        if copy:
+            params = {
+                k: u.unyt_quantity(v.value, v.units)
+                for k, v in self.parameters.items()
+            }
+        else:
+            params = self.parameters
+
+        return params
+
     @classmethod
     def from_template(cls, potential_template, parameters, topology=None):
         """Create a potential object from the potential_template

diff --git a/gmso/exceptions.py b/gmso/exceptions.py
@@ -25,3 +25,7 @@ class EngineIncompatibilityError(GMSOError):
 
 class MissingAtomTypesError(ForceFieldParseError):
     """Error for missing AtomTypes when creating a ForceField from an XML file"""
+
+
+class MissingPotentialError(ForceFieldError):
+    """Error for missing Potential when searching for Potentials in a ForceField"""
diff --git a/gmso/tests/files/oplsaa-ethane_foyer.xml b/gmso/tests/files/oplsaa-ethane_foyer.xml
@@ -0,0 +1,78 @@
+<?xml version='1.0' encoding='UTF-8'?>
+<ForceField name="oplsaa" version="0.0.1">
+  <FFMetaData electrostatics14Scale="0.5" nonBonded14Scale="0.5">
+    <Units energy="kJ/mol" mass="amu" charge="coulomb" distance="nm"/>
+  </FFMetaData>
+  <AtomTypes expression="ep * ((sigma/r)**12 - (sigma/r)**6) + q / (e0 * r)">
+    <ParametersUnitDef parameter="q" unit="coulomb"/>
+    <ParametersUnitDef parameter="e0" unit="A**2*s**4/(kg*m**3)"/>
+    <ParametersUnitDef parameter="sigma" unit="nm"/>
+    <ParametersUnitDef parameter="ep" unit="kJ/mol"/>
+    <AtomType name="opls_135" atomclass="CT" element="C" charge="0.0" mass="12.01078" definition="[C;X4](C)(H)(H)H" description="" doi="10.1021/ja9621760" overrides="">
+      <Parameters>
+        <Parameter name="ep" value="0.276144"/>
+        <Parameter name="sigma" value="0.35"/>
+        <Parameter name="e0" value="8.8542e-12"/>
+        <Parameter name="q" value="-0.18"/>
+      </Parameters>
+    </AtomType>
+    <AtomType name="opls_140" atomclass="HC" element="H" charge="0.0" mass="1.007947" definition="H[C;X4]" description="" doi="10.1021/ja9621760" overrides="">
+      <Parameters>
+        <Parameter name="ep" value="0.12552"/>
+        <Parameter name="sigma" value="0.25"/>
+        <Parameter name="e0" value="8.8542e-12"/>
+        <Parameter name="q" value="0.06"/>
+      </Parameters>
+    </AtomType>
+  </AtomTypes>
+  <BondTypes expression="k * (r-r_eq)**2">
+    <ParametersUnitDef parameter="r_eq" unit="nm"/>
+    <ParametersUnitDef parameter="k" unit="kJ/nm**2"/>
+    <BondType name="BondType-Harmonic-1" type1="opls_135" type2="opls_135">
+      <Parameters>
+        <Parameter name="k" value="224262.4"/>
+        <Parameter name="r_eq" value="0.1529"/>
+      </Parameters>
+    </BondType>
+    <BondType name="BondType-Harmonic-2" type1="opls_135" type2="opls_140">
+      <Parameters>
+        <Parameter name="k" value="284512.0"/>
+        <Parameter name="r_eq" value="0.109"/>
+      </Parameters>
+    </BondType>
+  </BondTypes>
+  <AngleTypes expression="k * (theta - theta_eq)**2">
+    <ParametersUnitDef parameter="theta_eq" unit="radian"/>
+    <ParametersUnitDef parameter="k" unit="kJ/radian**2"/>
+    <AngleType name="AngleType-Harmonic-1" type1="opls_135" type2="opls_135" type3="opls_140">
+      <Parameters>
+        <Parameter name="k" value="313.8"/>
+        <Parameter name="theta_eq" value="1.932079482"/>
+      </Parameters>
+    </AngleType>
+    <AngleType name="AngleType-Harmonic-2" type1="opls_140" type2="opls_135" type3="opls_140">
+      <Parameters>
+        <Parameter name="k" value="276.144"/>
+        <Parameter name="theta_eq" value="1.8814649337"/>
+      </Parameters>
+    </AngleType>
+  </AngleTypes>
+  <DihedralTypes expression="c0 * cos(phi)**0 + c1 * cos(phi)**1 + c2 * cos(phi)**2 + c3 * cos(phi)**3 + c4 * cos(phi)**4 + c5 * cos(phi)**5">
+    <ParametersUnitDef parameter="c5" unit="kJ/mol"/>
+    <ParametersUnitDef parameter="c4" unit="kJ/mol"/>
+    <ParametersUnitDef parameter="c3" unit="kJ/mol"/>
+    <ParametersUnitDef parameter="c2" unit="kJ/mol"/>
+    <ParametersUnitDef parameter="c1" unit="kJ/mol"/>
+    <ParametersUnitDef parameter="c0" unit="kJ/mol"/>
+    <DihedralType name="DihedralType-RB-Proper-1" type1="opls_140" type2="opls_135" type3="opls_135" type4="opls_140">
+      <Parameters>
+        <Parameter name="c0" value="0.6276"/>
+        <Parameter name="c1" value="1.8828"/>
+        <Parameter name="c2" value="0.0"/>
+        <Parameter name="c3" value="-2.5104"/>
+        <Parameter name="c4" value="0.0"/>
+        <Parameter name="c5" value="0.0"/>
+      </Parameters>
+    </DihedralType>
+  </DihedralTypes>
+</ForceField>