Merge pull request #268 from mattwthompson/nbfix-workaround

Add workaround function for applying nbfixes

foyer/tests/test_utils.py

@@ -0,0 +1,57 @@
+import numpy as np
+import parmed as pmd
+import pytest
+
+from foyer import Forcefield
+from foyer.tests.utils import get_fn
+from foyer.utils.nbfixes import apply_nbfix
+
+
+def test_apply_nbfix():
+    opls = Forcefield(name='oplsaa')
+    ethane = pmd.load_file(get_fn('ethane.mol2'), structure=True)
+    ethane = opls.apply(ethane)
+    ethane_tweaked = apply_nbfix(
+        struct=ethane,
+        atom_type1='opls_135',
+        atom_type2='opls_140',
+        sigma=0.4,
+        epsilon=50.0,
+    )
+
+    assert not ethane.has_NBFIX()
+    assert ethane_tweaked.has_NBFIX()
+
+    # 0.44898.... is rmin, which parmed uses internally in place of sigma
+    for atom in ethane_tweaked:
+        if atom.atom_type.name == 'opls_135':
+            assert np.allclose(
+                atom.atom_type.nbfix['opls_140'][:2],
+                [0.44898481932374923, 50.0]
+            )
+        elif atom.atom_type.name == 'opls_140':
+            assert np.allclose(
+                atom.atom_type.nbfix['opls_135'][:2],
+                [0.44898481932374923, 50.0]
+            )
+
+def test_apply_nbfix_bad_atom_type():
+    opls = Forcefield(name='oplsaa')
+    ethane = pmd.load_file(get_fn('ethane.mol2'), structure=True)
+    ethane = opls.apply(ethane)
+    with pytest.raises(ValueError):
+        apply_nbfix(
+            struct=ethane,
+            atom_type1='opls__typo_135',
+            atom_type2='opls_140',
+            sigma=0.4,
+            epsilon=50.0,
+        )
+    with pytest.raises(ValueError):
+        apply_nbfix(
+            struct=ethane,
+            atom_type1='opls_135',
+            atom_type2='opls_141',
+            sigma=0.4,
+            epsilon=50.0,
+        )

foyer/utils/nbfixes.py

@@ -0,0 +1,45 @@
+from copy import deepcopy
+
+
+def apply_nbfix(struct, atom_type1, atom_type2, sigma, epsilon):
+    """Apply a single nbfix to a particular interaction.
+
+    Parameters
+    ----------
+    struct : parmed.structure.Structure
+        The ParmEd structure to which this nbfix will be applied.
+    atom_type1 : str
+        The name of the first atom type in the nbfix pair
+    atom_type2 : str
+        The name of the second atom type in the nbfix pair
+    sigma : float
+        The sigma of the cross-interaction in kcal/mol
+    epsilon : float
+        The epsilon of the cross-interaction in Angstroms
+
+    Returns
+    -------
+    struct : parmed.structure.Structure
+        The input structure with the nbfix applied.
+    """
+    struct_copy = deepcopy(struct)
+
+    atom_types = list({a.atom_type for a in struct_copy.atoms})
+    for atom_type in sorted(atom_types, key=lambda a: a.name):
+        if atom_type.name == atom_type1:
+            a1 = atom_type
+        if atom_type.name == atom_type2:
+            a2 = atom_type
+    try:
+        a1
+        a2
+    except:
+        raise ValueError('Atom types {} and {} not found '
+                'in structure.'.format(atom_type1, atom_type2))
+
+    # Calculate rmin from sigma because parmed uses it internally
+    rmin = sigma * 2 ** (1 / 6)
+    a1.add_nbfix(a2.name, rmin, epsilon)
+    a2.add_nbfix(a1.name, rmin, epsilon)
+
+    return struct_copy