Initial untested implementation of LibraryChargeHandler

Author: j-wags

openforcefield/tests/test_forcefield.py

@@ -156,8 +156,17 @@
 xml_toolkitam1bcc_ff = '''
 <SMIRNOFF version="0.3" aromaticity_model="OEAroModel_MDL">
   <ToolkitAM1BCC version="0.3"/>
-</SMIRNOFF>'''
+</SMIRNOFF>
+'''
 
+xml_tip3p_library_charges_ff = '''
+<SMIRNOFF version="0.3" aromaticity_model="OEAroModel_MDL">
+    <LibraryCharges version="0.3">
+       <!-- TIP3P water oxygen with charge override -->
+       <LibraryCharge name="TIP3P" smirks="[#1:1]-[#8X2H2+0:2]-[#1:3]" charge1="+0.417*elementary_charge" charge2="-0.834*elementary_charge" charge3="+0.417*elementary_charge"/>
+    </LibraryCharges>
+</SMIRNOFF>
+'''
 #======================================================================
 # TEST UTILITY FUNCTIONS
 #======================================================================
@@ -806,6 +815,70 @@ def test_parameterize_ethanol_to_parmed(self):
 
         parmed_system = forcefield.create_parmed_structure(topology, positions=pdbfile.getPositions())
 
+
+
+
+    @pytest.mark.parametrize("toolkit_registry,registry_description", toolkit_registries)
+    def test_pass_invalid_kwarg_to_create_openmm_system(self, toolkit_registry, registry_description):
+        """Test to ensure an exception is raised when an unrecognized kwarg is passed """
+        from simtk.openmm import app
+
+        file_path = get_data_file_path('test_forcefields/smirnoff99Frosst.offxml')
+        forcefield = ForceField(file_path)
+        pdbfile = app.PDBFile(get_data_file_path('systems/test_systems/1_ethanol.pdb'))
+        molecules = []
+        molecules.append(Molecule.from_smiles('CCO'))
+        topology = Topology.from_openmm(pdbfile.topology, unique_molecules=molecules)
+
+        with pytest.raises(ValueError, match=".* not used by any registered force Handler: {'invalid_kwarg'}.*") as e:
+            omm_system = forcefield.create_openmm_system(topology, invalid_kwarg='aaa', toolkit_registry=toolkit_registry)
+
+
+    @pytest.mark.parametrize("inputs", nonbonded_resolution_matrix)
+    def test_nonbonded_method_resolution(self,
+                                         inputs
+                                         ):
+        """Test predefined permutations of input options to ensure nonbonded handling is correctly resolved"""
+        from simtk.openmm import app
+
+        vdw_method = inputs['vdw_method']
+        electrostatics_method = inputs['electrostatics_method']
+        has_periodic_box = inputs['has_periodic_box']
+        omm_force = inputs['omm_force']
+        exception = inputs['exception']
+        exception_match= inputs['exception_match']
+
+        molecules = [create_ethanol()]
+        forcefield = ForceField('test_forcefields/smirnoff99Frosst.offxml')
+
+        pdbfile = app.PDBFile(get_data_file_path('systems/test_systems/1_ethanol.pdb'))
+        topology = Topology.from_openmm(pdbfile.topology, unique_molecules=molecules)
+
+        if not(has_periodic_box):
+            topology.box_vectors = None
+
+        if exception is None:
+            # The method is validated and may raise an exception if it's not supported.
+            forcefield.get_parameter_handler('vdW', {}).method = vdw_method
+            forcefield.get_parameter_handler('Electrostatics', {}).method = electrostatics_method
+            omm_system = forcefield.create_openmm_system(topology)
+            nonbond_method_matched = False
+            for f_idx in range(omm_system.getNumForces()):
+                force = omm_system.getForce(f_idx)
+                if isinstance(force, openmm.NonbondedForce):
+                    if force.getNonbondedMethod() == omm_force:
+                        nonbond_method_matched = True
+            assert nonbond_method_matched
+        else:
+            with pytest.raises(exception, match=exception_match) as excinfo:
+                # The method is validated and may raise an exception if it's not supported.
+                forcefield.get_parameter_handler('vdW', {}).method = vdw_method
+                forcefield.get_parameter_handler('Electrostatics', {}).method = electrostatics_method
+                omm_system = forcefield.create_openmm_system(topology)
+
+
+
+class TestForceFieldChargeAssignment:
     @pytest.mark.parametrize("toolkit_registry,registry_description", toolkit_registries)
     def test_charges_from_molecule(self, toolkit_registry, registry_description):
         """Test skipping charge generation and instead getting charges from the original Molecule"""
@@ -878,65 +951,49 @@ def test_some_charges_from_molecule(self, toolkit_registry, registry_description
             q, sigma, epsilon = nonbondedForce.getParticleParameters(particle_index)
             assert q != (0. * unit.elementary_charge)
 
+    def test_library_charges_to_single_water(self):
+        """Test assigning charges to one water molecule using library charges"""
+        from simtk.openmm import NonbondedForce
 
+        ff = ForceField(simple_xml_ff, xml_tip3p_library_charges_ff)
+        mol = Molecule.from_smiles('O')
+        omm_system = ff.create_openmm_system(mol.to_topology())
+        nonbondedForce = [f for f in omm_system.getForces() if type(f) == NonbondedForce][0]
+        expected_charges = [-0.834, 0.417, 0.417] * unit.elementary_charge
+        for particle_index, expected_charge in expected_charges:
+            q, sigma, epsilon = nonbondedForce.getParticleParameters(particle_index)
+            assert q == expected_charge
 
-    @pytest.mark.parametrize("toolkit_registry,registry_description", toolkit_registries)
-    def test_pass_invalid_kwarg_to_create_openmm_system(self, toolkit_registry, registry_description):
-        """Test to ensure an exception is raised when an unrecognized kwarg is passed """
-        from simtk.openmm import app
+    def test_library_charges_to_two_waters(self):
+        """Test assigning charges to two water molecules using library charges"""
+        pass
 
-        file_path = get_data_file_path('test_forcefields/smirnoff99Frosst.offxml')
-        forcefield = ForceField(file_path)
-        pdbfile = app.PDBFile(get_data_file_path('systems/test_systems/1_ethanol.pdb'))
-        molecules = []
-        molecules.append(Molecule.from_smiles('CCO'))
-        topology = Topology.from_openmm(pdbfile.topology, unique_molecules=molecules)
+    def test_library_charges_to_two_ethanols_different_atom_ordering(self):
+        """Test assigning charges to two ethanols with different atom orderings"""
+        pass
 
-        with pytest.raises(ValueError, match=".* not used by any registered force Handler: {'invalid_kwarg'}.*") as e:
-            omm_system = forcefield.create_openmm_system(topology, invalid_kwarg='aaa', toolkit_registry=toolkit_registry)
+    def test_library_charges_to_only_some_molecules(self):
+        """Test assigning charges to a topology in which some molecules are covered by library charges
+         but other aren't"""
+        pass
 
+    def test_assign_charges_to_molecule_in_parts_using_multiple_library_charges(self):
+        """Test assigning charges to parts of a molecule using two library charge lines"""
+        pass
 
-    @pytest.mark.parametrize("inputs", nonbonded_resolution_matrix)
-    def test_nonbonded_method_resolution(self,
-                                         inputs
-                                         ):
-        """Test predefined permutations of input options to ensure nonbonded handling is correctly resolved"""
-        from simtk.openmm import app
+    def test_assign_charges_using_library_charges_by_single_atoms(self):
+        """Test assigning charges to parts of a molecule using per-atom library charges"""
+        pass
 
-        vdw_method = inputs['vdw_method']
-        electrostatics_method = inputs['electrostatics_method']
-        has_periodic_box = inputs['has_periodic_box']
-        omm_force = inputs['omm_force']
-        exception = inputs['exception']
-        exception_match= inputs['exception_match']
+    def test_library_charges_dont_parameterize_molecule_because_of_incomplete_coverage(self):
+        """Fail to assign charges to a molecule becau\se not all atoms can be assigned"""
+        pass
 
-        molecules = [create_ethanol()]
-        forcefield = ForceField('test_forcefields/smirnoff99Frosst.offxml')
+    def library_charges_overridden_by_charge_from_molecules(self):
+        """Skip assigning charges to a molecule because it has already had charges assigned
+        using charge_from_molecules"""
+        pass
 
-        pdbfile = app.PDBFile(get_data_file_path('systems/test_systems/1_ethanol.pdb'))
-        topology = Topology.from_openmm(pdbfile.topology, unique_molecules=molecules)
-
-        if not(has_periodic_box):
-            topology.box_vectors = None
-
-        if exception is None:
-            # The method is validated and may raise an exception if it's not supported.
-            forcefield.get_parameter_handler('vdW', {}).method = vdw_method
-            forcefield.get_parameter_handler('Electrostatics', {}).method = electrostatics_method
-            omm_system = forcefield.create_openmm_system(topology)
-            nonbond_method_matched = False
-            for f_idx in range(omm_system.getNumForces()):
-                force = omm_system.getForce(f_idx)
-                if isinstance(force, openmm.NonbondedForce):
-                    if force.getNonbondedMethod() == omm_force:
-                        nonbond_method_matched = True
-            assert nonbond_method_matched
-        else:
-            with pytest.raises(exception, match=exception_match) as excinfo:
-                # The method is validated and may raise an exception if it's not supported.
-                forcefield.get_parameter_handler('vdW', {}).method = vdw_method
-                forcefield.get_parameter_handler('Electrostatics', {}).method = electrostatics_method
-                omm_system = forcefield.create_openmm_system(topology)
 
 #======================================================================
 # TEST CONSTRAINTS

openforcefield/tests/test_parameters.py

@@ -21,7 +21,7 @@
 from openforcefield.typing.engines.smirnoff.parameters import (
     ParameterAttribute, IndexedParameterAttribute, ParameterList,
     ParameterType, BondHandler, ParameterHandler, ProperTorsionHandler,
-    ImproperTorsionHandler, GBSAHandler, SMIRNOFFSpecError,
+    ImproperTorsionHandler, LibraryChargeHandler, GBSAHandler, SMIRNOFFSpecError,
     _ParameterAttributeHandler
     )
 from openforcefield.utils import detach_units, IncompatibleUnitError
@@ -931,6 +931,13 @@ def test_torsion_handler_charmm_potential(self):
             ph1 = ImproperTorsionHandler(potential='charmm', skip_version_check=True)
         ph1 = ImproperTorsionHandler(potential='k*(1+cos(periodicity*theta-phase))', skip_version_check=True)
 
+class TestLibraryChargeHandler:
+    def test_create_library_charge_handler(self):
+        """Test creation of an empty LibraryChargeHandler"""
+        handler = LibraryChargeHandler()
+
+
+
 class TestGBSAHandler:
     def test_create_default_gbsahandler(self):
         """Test creation of an empty GBSAHandler, with all default attributes"""

openforcefield/typing/engines/smirnoff/parameters.py

@@ -2568,7 +2568,7 @@ class ToolkitAM1BCCHandler(ParameterHandler):
 
     _TAGNAME = 'ToolkitAM1BCC'  # SMIRNOFF tag name to process
     _OPENMMTYPE = openmm.NonbondedForce  # OpenMM force class to create or utilize
-    _DEPENDENCIES = [vdWHandler] # vdWHandler must first run NonBondedForce.addParticle for each particle in the topology
+    _DEPENDENCIES = [vdWHandler, LibraryChargeHandler] # vdWHandler must first run NonBondedForce.addParticle for each particle in the topology
     _KWARGS = ['charge_from_molecules', 'toolkit_registry'] # Kwargs to catch when create_force is called
 
     def check_handler_compatibility(self,
@@ -2737,6 +2737,92 @@ def postprocess_system(self, system, topology, **kwargs):
                     # TODO: Calculate exceptions
 
 
+
+class LibraryChargeHandler(ParameterHandler):
+    """Handle SMIRNOFF ``<LibraryCharges>`` tags
+
+    .. warning :: This API is experimental and subject to change.
+    """
+
+    class LibraryChargeType(ParameterType):
+        """A SMIRNOFF Library Charge type.
+
+        .. warning :: This API is experimental and subject to change.
+        """
+        _VALENCE_TYPE = 'Atom'  # ChemicalEnvironment valence type expected for SMARTS
+        _ELEMENT_NAME = 'LibraryCharge'
+
+        charge = IndexedParameterAttribute(unit=unit.elementary_charge)
+
+
+    _TAGNAME = 'LibraryCharges'  # SMIRNOFF tag name to process
+    _INFOTYPE = LibraryChargeType  # info type to store
+    _OPENMMTYPE = openmm.NonbondedForce  # OpenMM force class to create
+    _DEPENDENCIES = [vdWHandler] # vdWHandler must first run NonBondedForce.addParticle for each particle in the topology
+
+
+    def create_force(self, system, topology, **kwargs):
+        existing = [system.getForce(i) for i in range(system.getNumForces())]
+        existing = [f for f in existing if type(f) == self._OPENMMTYPE]
+        if len(existing) == 0:
+            force = self._OPENMMTYPE()
+            system.addForce(force)
+        else:
+            force = existing[0]
+
+        # Iterate over all defined library charge parameters, allowing later matches to override earlier ones.
+        atom_matches = self.find_matches(topology)
+
+        # Create a set of all the topology atom indices for which library charges can be applied
+        assignable_atoms = set()
+        atom_assignments = dict()
+
+        # TODO: This assumes that later matches should always override earlier ones. This may require more
+        #       thought, since matches can be partially overlapping
+        for topology_indices, library_charge in atom_matches.items():
+            for charge_idx, top_idx in enumerate(topology_indices):
+                if top_idx in assignable_atoms:
+                    logger.debug(f'Multiple library charge assignments found for atom {top_idx}')
+                assignable_atoms.add(top_idx)
+                atom_assignments[top_idx] = library_charge.charge[charge_idx]
+        #assignable_atoms = {top_idx for top_idx in top_idx_tuple for top_idx_tuple in atom_matches.keys()}
+        #atom_assignments =
+
+        # TODO: Should header include a residue separator delimiter? Maybe not, since it's not clear how having
+        #       multiple LibraryChargeHandlers could return a single set of matches, while respecting different
+        #       separators.
+
+        # Check to see whether the set contains any complete molecules, and remove the matches if not.
+        for top_mol in topology.topology_molecules:
+            # Check whether any atom in the molecule already has a charge. If so, skip it
+            for top_atom in top_mol.topology_atoms:
+                q, _, _2 = force.getParticleParameters(top_atom.topology_atom_index)
+                if q != 0 * unit.elementary_charge:
+                    logger.debug('Original molecule has at least one atom with existing charge. Skipping library '
+                                 'charge assignment')
+                    continue
+
+            # Ensure all of the atoms in this mol are covered, otherwise skip it
+            if set(top_mol.topology_atoms).intersection(assignable_atoms) != top_mol.n_atoms:
+                logger.debug('Entire molecule is not covered. Skipping library charge assignment.')
+                continue
+
+            # If we pass both tests above, go ahead and assign charges
+            # TODO: We could probably save a little time by looking up this TopologyMolecule's _reference molecule_
+            #       and assigning charges to all other instances of it in this topoligy
+            for top_atom in top_mol.topology_atoms:
+                _, sigma, epsilon = force.getParticleParameters(top_atom.topology_atom_index)
+                force.setParticleParameters(top_atom.topology_atom_index,
+                                            atom_assignments[top_atom.topology_atom_index],
+                                            sigma,
+                                            epsilon)
+
+
+    # TODO: Can we express separate constraints for postprocessing and normal processing?
+    def postprocess_system(self, system, topology, **kwargs):
+        pass
+
+
 class ChargeIncrementModelHandler(ParameterHandler):
     """Handle SMIRNOFF ``<ChargeIncrementModel>`` tags