More progress on Drude support

parmed/charmm/psf.py

@@ -241,11 +241,14 @@ def __init__(self, psf_name=None):
                     alpha = self._convert(words[9], float, 'alpha')
                     thole = self._convert(words[10], float, 'thole')
                     drude_alpha_thole.append((alpha, thole))
-                if is_drude and i > 1 and drude_alpha_thole[-2] != (0, 0):
+                if is_drude and i >= 1 and drude_alpha_thole[-2] != (0, 0):
+                    # This assumes that the Drude atom is defined immediately after its parent atom.
+                    # This always appears to be the case, but if it proves untrue at some point then
+                    # this will need to be refactored to identify Drude atoms after identifying bonds.
                     my_alpha, my_thole = drude_alpha_thole[-2]
-                    atom = DrudeAtom(name=name, type=attype, charge=charge, mass=mass,
+                    atom = DrudeAtom(name=name, type=attype, charge=charge, mass=mass, parent=self.atoms[-1],
                                      atomic_number=0, alpha=my_alpha, thole=my_thole, drude_type=attype)
-                elif name.startswith('LP') and mass == 0:
+                elif (name.startswith('LP') or (isinstance(attype, str) and attype.startswith("LP"))) and mass == 0:
                     atom = ExtraPoint(name=name, type=attype, charge=charge, mass=mass, atomic_number=0)
                 else:
                     atom = Atom(name=name, type=attype, charge=charge, mass=mass,
@@ -637,10 +640,16 @@ def load_parameters(self, parmset, copy_parameters=True):
             atom.atomic_number = atype.atomic_number
 
         # Next load all of the bonds
+        skipped_bonds = set()
         for bond in self.bonds:
+            # Skip any bonds with drude atoms or virtual sites. They are not stored.
+            # Depending on how Drude support is implemented in Amber (if that ever happens), we
+            # may have to add dummy values here.
+            if isinstance(bond.atom1, (DrudeAtom, ExtraPoint)) or isinstance(bond.atom2, (DrudeAtom, ExtraPoint)):
+                skipped_bonds.add(bond)
+                continue
             # Construct the key
-            key = (min(bond.atom1.type, bond.atom2.type),
-                   max(bond.atom1.type, bond.atom2.type))
+            key = (min(bond.atom1.type, bond.atom2.type), max(bond.atom1.type, bond.atom2.type))
             try:
                 bond.type = parmset.bond_types[key]
             except KeyError:
@@ -649,6 +658,8 @@ def load_parameters(self, parmset, copy_parameters=True):
         # Build the bond_types list
         del self.bond_types[:]
         for bond in self.bonds:
+            if bond in skipped_bonds:
+                continue
             if bond.type.used:
                 continue
             bond.type.used = True

parmed/structure.py

@@ -17,7 +17,7 @@
 from .geometry import (STANDARD_BOND_LENGTHS_SQUARED, box_lengths_and_angles_to_vectors,
                        box_vectors_to_lengths_and_angles, distance2)
 from .topologyobjects import (AcceptorDonor, Angle, Atom, AtomList, Bond, ChiralFrame, Cmap,
-                              Dihedral, DihedralType, DihedralTypeList, ExtraPoint, Group, Improper,
+                              Dihedral, DihedralType, DihedralTypeList, DrudeAtom, ExtraPoint, Group, Improper,
                               MultipoleFrame, NonbondedException, NoUreyBradley, OutOfPlaneBend,
                               OutOfPlaneExtraPointFrame, PiTorsion, ResidueList, StretchBend,
                               ThreeParticleExtraPointFrame, TorsionTorsion, TrackedList,
@@ -234,6 +234,7 @@ class Structure:
     TORSION_TORSION_FORCE_GROUP = 10
     NONBONDED_FORCE_GROUP = 11
     RB_TORSION_FORCE_GROUP = 12
+    DRUDE_FORCE_GROUP = 13
 
     #===================================================
 
@@ -2066,6 +2067,8 @@ def createSystem(self, nonbondedMethod=None,
         if self.box is not None:
             system.setDefaultPeriodicBoxVectors(*reducePeriodicBoxVectors(self.box_vectors))
         self.omm_set_virtual_sites(system)
+        if any(isinstance(atom, DrudeAtom) for atom in self.atoms):
+            self._add_force_to_system(system, self.omm_drude_force(system))
         return system
 
     #===================================================
@@ -2685,6 +2688,13 @@ def exclude_to(origin, atom, level, end):
 
     #===================================================
 
+    def omm_drude_force(self, system):
+        force = mm.DrudeForce()
+        force.setForceGroup(self.DRUDE_FORCE_GROUP)
+        return force
+
+    #===================================================
+
     def _omm_nbfixed_force(self, nonbfrc, nonbondedMethod):
         """ Private method for creating a CustomNonbondedForce with a lookup
         table. This should not be called by users -- you have been warned.

parmed/topologyobjects.py

@@ -946,7 +946,7 @@ def __le__(self, other):
         return not self > other
 
     def __repr__(self):
-        start = f'<Atom {self.name} [{self.idx}]'
+        start = f'<{self.__class__.__name__} {self.name} [{self.idx}]'
         if self.residue is not None and hasattr(self.residue, 'idx'):
             return start + f'; In {self.residue.name} {self.residue.idx}>'
         elif self.residue is not None:
@@ -4198,12 +4198,13 @@ class DrudeAtom(Atom):
         atoms, this is None.
     """
 
-    def __init__(self, alpha=0.0, thole=1.3, drude_type='DRUD', **kwargs):
+    def __init__(self, alpha=0.0, thole=1.3, drude_type='DRUD', parent=None, **kwargs):
         super().__init__(**kwargs)
         self.alpha = alpha
         self.thole = thole
         self.drude_type = drude_type
         self.anisotropy = None
+        self.parent = parent
 
     @property
     def drude_charge(self):
@@ -4231,11 +4232,9 @@ class DrudeAnisotropy(_FourAtomTerm):
     atom4 : :class:`Atom`
         the fourth atom defining the coordinate frame
     a11 : ``float``
-        the scale factor for the polarizability along the direction defined by
-        atom1 and atom2
+        the scale factor for the polarizability along the direction defined by atom1 and atom2
     a22 : ``float``
-        the scale factor for the polarizability along the direction defined by
-        atom3 and atom4
+        the scale factor for the polarizability along the direction defined by atom3 and atom4
     """
 
     def __init__(self, atom1, atom2, atom3, atom4, a11, a22):

test/test_parmed_charmm.py

@@ -11,7 +11,8 @@
 from parmed.charmm import charmmcrds, parameters, psf
 from parmed.charmm._charmmfile import CharmmFile, CharmmStreamFile
 from parmed import exceptions, topologyobjects as to, load_file, ParameterSet
-from parmed.topologyobjects import BondType, AngleType, DihedralType, DihedralTypeList, DrudeAtom, ExtraPoint
+from parmed.topologyobjects import Atom, BondType, AngleType, DihedralType, DihedralTypeList, DrudeAtom, ExtraPoint
+from parmed.periodic_table import Element, Mass
 import parmed.unit as u
 import random
 import unittest
@@ -1030,12 +1031,10 @@ class TestFileWriting(TestCharmmBase):
 
     def test_charmm_file(self):
         """ Test the CharmmFile API and error handling """
-        self.assertRaises(ValueError, lambda:
-                CharmmFile(get_fn('trx.prmtop'), 'x')
-        )
-        self.assertRaises(IOError, lambda:
-                CharmmFile(get_fn('file_does_not_exist'), 'r')
-        )
+        with self.assertRaises(ValueError):
+            CharmmFile(get_fn('trx.prmtop'), 'x')
+        with self.assertRaises(IOError):
+            CharmmFile(get_fn('file_does_not_exist'), 'r')
         with CharmmFile(self.get_fn('newfile.chm', written=True), 'w') as f:
             f.write('abc123\ndef456\nghi789!comment...\n')
         with CharmmFile(self.get_fn('newfile.chm', written=True), 'r') as f:
@@ -1066,10 +1065,8 @@ def test_charmm_file(self):
                 lines.append(line)
                 comments.append(f.comment)
                 line = f.readline()
-            self.assertEqual(lines, ['First line \n', 'Second line \n',
-                                     'Third line \n', 'Fourth line \n'])
-            self.assertEqual(comments, ['! first comment', '! second comment',
-                                        '! third comment', '! fourth comment'])
+            self.assertEqual(lines, ['First line \n', 'Second line \n', 'Third line \n', 'Fourth line \n'])
+            self.assertEqual(comments, ['! first comment', '! second comment', '! third comment', '! fourth comment'])
 
     def test_charmm_stream_file(self):
         """ Test the CharmmStreamFile API """
@@ -1185,6 +1182,20 @@ def test_write_xplor(self):
 class TestCharmmDrudeSystems(unittest.TestCase):
     """ Tests processing of DRUDE systems """
 
+    def test_drude_mass(self):
+        drude_psf = psf.CharmmPsfFile(get_fn("ala3_solv_drude.psf"))
+        params = parameters.CharmmParameterSet(get_fn("toppar_drude_master_protein_2013e.str"))
+        drude_psf.load_parameters(params)
+        # Check this against a known-correct implementation in OpenM
+        self.assertEqual(sum([isinstance(a, DrudeAtom) for a in drude_psf.atoms]), 957)
+        for atom in drude_psf.atoms:
+            if isinstance(atom, DrudeAtom):
+                self.assertIsInstance(atom.parent, Atom)
+        for atom in drude_psf.atoms:
+            if isinstance(atom, DrudeAtom):
+                self.assertAlmostEqual(atom.mass + atom.parent.mass, Mass[atom.parent.element_name], delta=0.01)
+
+
     def test_drude_parsing_na(self):
         """ Test parsing a trinucleotide with Drude particles and lone pairs """
         system = psf.CharmmPsfFile(get_fn("cyt-gua-cyt.psf"))