Check all wildcard combinations for dihedrals and impropers, and add test

wrappers/python/openmm/app/charmmpsffile.py

@@ -34,6 +34,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE.
 from __future__ import division, absolute_import, print_function
 
 from functools import wraps
+from itertools import combinations
 from math import pi, cos, sin, sqrt
 import os
 import re
@@ -603,9 +604,9 @@ class CharmmPsfFile(object):
         - If any parameters that are necessary cannot be found, a
           MissingParameter exception is raised.
         - If any dihedral or improper parameters cannot be found, I will try
-          inserting wildcards (at either end for dihedrals and as the two
-          central atoms in impropers) and see if that matches.  Wild-cards
-          will apply ONLY if specific parameters cannot be found.
+          inserting wildcards (in up to three positions) and see if that
+          matches.  Wild-cards will apply ONLY if specific parameters cannot be
+          found.
         - This method will expand the dihedral_parameter_list attribute by
           adding a separate Dihedral object for each term for types that
           have a multi-term expansion
@@ -655,16 +656,11 @@ class CharmmPsfFile(object):
         for dih in self.dihedral_list:
             # Store the atoms
             a1, a2, a3, a4 = dih.atom1, dih.atom2, dih.atom3, dih.atom4
-            at1, at2, at3, at4 = a1.attype, a2.attype, a3.attype, a4.attype
-            # First see if the exact dihedral is specified
-            key = min((at1,at2,at3,at4), (at4,at3,at2,at1))
-            if not key in parmset.dihedral_types:
-                # Check for wild-cards
-                key = min(('X',at2,at3,'X'), ('X',at3,at2,'X'))
-                if not key in parmset.dihedral_types:
-                    raise MissingParameter('No dihedral parameters found for '
-                                           '%r' % dih)
-            dtlist = parmset.dihedral_types[key]
+            try:
+                at_list = [a1.attype, a2.attype, a3.attype, a4.attype]
+                dtlist = _match_with_wildcards(at_list, parmset.dihedral_types, 3)
+            except KeyError:
+                raise MissingParameter('No dihedral parameters found for %r' % dih)
             for i, dt in enumerate(dtlist):
                 self.dihedral_parameter_list.append(Dihedral(a1,a2,a3,a4,dt))
                 # See if we include the end-group interactions for this
@@ -678,18 +674,11 @@ class CharmmPsfFile(object):
         for imp in self.improper_list:
             # Store the atoms
             a1, a2, a3, a4 = imp.atom1, imp.atom2, imp.atom3, imp.atom4
-            at1, at2, at3, at4 = a1.attype, a2.attype, a3.attype, a4.attype
-            key = min((at1,at2,at3,at4), (at4,at3,at2,at1))
-            if not key in parmset.improper_types:
-                key = min((at1,'X', 'X',at4),(at4,'X','X',at1))
-                if not key in parmset.improper_types:
-                    raise MissingParameter('No improper dihedral parameters found for '
-                                           '%r' % imp)
             try:
-                imp.improper_type = parmset.improper_types[key]
+                at_list = [a1.attype, a2.attype, a3.attype, a4.attype]
+                imp.improper_type = _match_with_wildcards(at_list, parmset.improper_types, 3)
             except KeyError:
-                raise MissingParameter('No improper parameters found for %r' %
-                                       imp)
+                raise MissingParameter('No improper parameters found for %r' % imp)
         # Now do the cmaps. These will not have wild-cards
         for cmap in self.cmap_list:
             # Store the atoms for easy reference
@@ -1854,6 +1843,26 @@ def _mbondi3_radii(atom_list):
 
 # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
+def _match_with_wildcards(at_list, type_dict, max_wc_count):
+    """
+    Tries to find a match for at_list in type_dict.  If one can't be found,
+    tries to insert up to max_wc_count wildcards in different combinations.
+    Assumes that for a given key, min(key, key[::-1]) is the one in type_dict.
+    """
+
+    for wc_count in range(max_wc_count + 1):
+        for wc_indices in combinations(range(len(at_list)), wc_count):
+            key = list(at_list)
+            for wc_index in wc_indices:
+                key[wc_index] = 'X'
+            key = tuple(key)
+            key = min(key, key[::-1])
+            if key in type_dict:
+                return type_dict[key]
+    raise KeyError(key)
+
+# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+
 if __name__ == '__main__':
     import doctest
     doctest.testmod()

wrappers/python/tests/TestCharmmFiles.py

@@ -4,6 +4,7 @@ from openmm.app import *
 from openmm import *
 from openmm.unit import *
 import openmm.app.element as elem
+import itertools
 import math
 import os
 import tempfile
@@ -382,6 +383,122 @@ class TestCharmmFiles(unittest.TestCase):
                     dtheta = math.pi-angle
                 self.assertAlmostEqual(energy, dtheta**2, delta=1e-5)
 
+    def test_TorsionWildcards(self):
+        """Test matching of dihedrals and impropers with wildcards."""
+        for test_improper, wild_1, wild_2, wild_3, wild_4, reverse, want_mismatch in itertools.product((False, True), repeat=7):
+            # Test with up to 3 wildcards.
+            if wild_1 and wild_2 and wild_3 and wild_4:
+                continue
+
+            # Test both dihedrals and impropers.  If want_mismatch is set, make
+            # sure that the non-wildcard atoms ensure no match to the torsion.
+            prm_header = 'IMPH' if test_improper else 'DIHE'
+            type_1 = 'X' if wild_1 else ('C2' if want_mismatch else 'C1')
+            type_2 = 'X' if wild_2 else ('C1' if want_mismatch else 'C2')
+            type_3 = 'X' if wild_3 else ('C4' if want_mismatch else 'C3')
+            type_4 = 'X' if wild_4 else ('C3' if want_mismatch else 'C4')
+
+            if reverse:
+                type_1, type_2, type_3, type_4 = type_4, type_3, type_2, type_1
+
+            psf_dihedral = '' if test_improper else f'{1:10}{2:10}{3:10}{4:10}'
+            psf_improper = f'{1:10}{2:10}{3:10}{4:10}' if test_improper else ''
+
+            with tempfile.TemporaryDirectory() as temp_path:
+                prm_path = os.path.join(temp_path, 'test.prm')
+                psf_path = os.path.join(temp_path, 'test.psf')
+
+                # Write a sample PRM file.
+                with open(prm_path, 'w') as prm_file:
+                    print(f"""*TEST
+*
+
+ATOMS
+MASS -1 C1 12.0110
+MASS -1 C2 12.0110
+MASS -1 C3 12.0110
+MASS -1 C4 12.0110
+
+BOND
+C1 C2 1 1
+C{1 if test_improper else 2} C3 1 1
+C{1 if test_improper else 3} C4 1 1
+
+{prm_header}
+{type_1} {type_2} {type_3} {type_4} 1 1 0
+
+NBON
+C1 0 0 1
+C2 0 0 1
+C3 0 0 1
+C4 0 0 1
+
+END""", file=prm_file)
+
+                # Write a sample PSF file.
+                with open(psf_path, 'w') as psf_file:
+                    print(f"""PSF EXT CMAP CHEQ XPLOR
+
+         1 !NTITLE
+* TEST
+
+         4 !NATOM
+         1 TEST     1        TEST1    C1       C1       0.000000       12.0110           0
+         2 TEST     1        TEST1    C2       C2       0.000000       12.0110           0
+         3 TEST     1        TEST1    C3       C3       0.000000       12.0110           0
+         4 TEST     1        TEST1    C4       C4       0.000000       12.0110           0
+
+         3 !NBOND: bonds
+{1:10}{2:10}{1 if test_improper else 2:10}{3:10}{1 if test_improper else 3:10}{4:10}
+
+         0 !NTHETA: angles
+
+
+{0 if test_improper else 1:10} !NPHI: dihedrals
+{psf_dihedral}
+
+{1 if test_improper else 0:10} !NIMPHI: impropers
+{psf_improper}
+
+         0 !NDON: donors
+
+
+         0 !NACC: acceptors
+
+
+         0 !NNB
+
+         0         0         0         0
+
+         1         0 !NGRP NST2
+         0         0         0
+
+         1 !MOLNT
+         1         1         1         1
+
+         0         0 !NUMLP NUMLPH
+
+         0 !NCRTERM
+
+
+""", file=psf_file)
+
+                prm = CharmmParameterSet(prm_path)
+                psf = CharmmPsfFile(psf_path)
+
+                if want_mismatch:
+                    # Make sure that the system doesn't get parameterized.
+                    with self.assertRaises(internal.charmm.exceptions.MissingParameter):
+                        system = psf.createSystem(prm)
+
+                else:
+                    # Make sure that one dihedral or improper gets added.
+                    system = psf.createSystem(prm)
+                    force_type = CustomTorsionForce if test_improper else PeriodicTorsionForce
+                    force, = (force for force in system.getForces() if isinstance(force, force_type))
+                    assert force.getNumTorsions() == 1
+                    assert force.getTorsionParameters(0)[:4] == [0, 1, 2, 3]
+
     def test_Residues(self):
         """Test that residues are read correctly, even if they have the same RESID while being in separate segments."""
         m14 = (["C{}".format(i) for i in range(1,14)]