Let the force field resolve the partial charge assignment

news/charge_from_ff.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* The ``assign_offmol_partial_charges`` and ``bulk_assign_partial_charges`` functions can assign charges from a list of OpenFF SMIRNOFF style force fields. Set method=``forcefield`` and provide a list of force field files via the new keyword argument ``forcefields``. This is also supported in the  ``charge-molecules`` CLI command and is set by using a yaml settings file.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/openfe/protocols/openmm_utils/charge_generation.py

@@ -11,6 +11,7 @@
 
 import numpy as np
 from gufe import SmallMoleculeComponent
+from openff.toolkit import ForceField
 from openff.toolkit import Molecule as OFFMol
 from openff.toolkit.utils.base_wrapper import ToolkitWrapper
 from openff.toolkit.utils.toolkit_registry import ToolkitRegistry
@@ -286,10 +287,11 @@ def _generate_offmol_conformers(
 def assign_offmol_partial_charges(
     offmol: OFFMol,
     overwrite: bool,
-    method: Literal["am1bcc", "am1bccelf10", "nagl", "espaloma"],
+    method: Literal["am1bcc", "am1bccelf10", "nagl", "espaloma", "forcefield"],
     toolkit_backend: Literal["ambertools", "openeye", "rdkit"],
     generate_n_conformers: int | None,
     nagl_model: str | None,
+    forcefields: list[str] | None = None,
 ) -> OFFMol:
     """
     Assign partial charges to an OpenFF Molecule based on a selected method.
@@ -299,11 +301,11 @@ def assign_offmol_partial_charges(
     offmol : openff.toolkit.Molecule
       The Molecule to assign partial charges to.
     overwrite : bool
-      Whether or not to overwrite any existing non-zero partial charges.
+      Whether to overwrite any existing non-zero partial charges.
       Note that zeroed charges will always be overwritten.
-    method : Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma']
+    method : Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma', 'forcefield']
       Partial charge assignment method.
-      Supported methods include; am1bcc, am1bccelf10, nagl, and espaloma.
+      Supported methods include; am1bcc, am1bccelf10, nagl, espaloma and forcefield.
     toolkit_backend : Literal['ambertools', 'openeye', 'rdkit']
       OpenFF toolkit backend employed for charge generation.
       Supported options:
@@ -319,6 +321,15 @@ def assign_offmol_partial_charges(
     nagl_model : str | None
       The NAGL model to use for charge assignment if method is ``nagl``.
       If ``None``, the latest am1bcc NAGL charge model is used.
+    forcefields : list[str] | None, default None
+        An optional list of SMIRNOFF style force field offxml paths or strings which should be used to assign partial charges.
+
+    Notes
+    -----
+    Charges are applied based on the following source preferences:
+     - Charges already present on the ligand are retained if overwrite is ``False``
+     - Charges are applied using the input method and settings
+     - the forcefield option will apply the default charges as intended by the force field.
 
     Raises
     ------
@@ -339,6 +350,25 @@ def assign_offmol_partial_charges(
         if not overwrite:
             return offmol
 
+    if method.lower() == "forcefield":
+        if forcefields is None:
+            errmsg = (
+                "The forcefield method requires a force field or list of force fields' to be provided "
+                "via `forcefields`."
+            )
+            raise ValueError(errmsg)
+
+        if isinstance(forcefields, str):
+            forcefields = [forcefields]
+
+        # this expects the full file name of the force field offxml file, e.g. "openff-2.0.0.offxml"
+        # which is different to how settings work which can leave off the .offxml extension
+        ff = ForceField(*forcefields)
+        # let the force field resolve the partial charge assignment
+        charges = ff.get_partial_charges(offmol)
+        offmol.partial_charges = charges
+        return offmol
+
     # Dictionary for each available charge method
     # The idea of this pattern is to allow for maximum flexibility by
     # allowing for swapping out method calls as necessary.
@@ -441,11 +471,12 @@ def assign_offmol_partial_charges(
 def bulk_assign_partial_charges(
     molecules: list[SmallMoleculeComponent],
     overwrite: bool,
-    method: Literal["am1bcc", "am1bccelf10", "nagl", "espaloma"],
+    method: Literal["am1bcc", "am1bccelf10", "nagl", "espaloma", "forcefield"],
     toolkit_backend: Literal["ambertools", "openeye", "rdkit"],
     generate_n_conformers: int | None,
     nagl_model: str | None,
     processors: int = 1,
+    forcefields: list[str] | None = None,
 ) -> list[SmallMoleculeComponent]:
     """
     Assign partial charges to a list of SmallMoleculeComponents using multiprocessing.
@@ -457,7 +488,7 @@ def bulk_assign_partial_charges(
     overwrite : bool
       Whether or not to overwrite any existing non-zero partial charges.
       Note that zeroed charges will always be overwritten.
-    method : Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma']
+    method : Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma', 'forcefield]
       Partial charge assignment method.
       Supported methods include; am1bcc, am1bccelf10, nagl, and espaloma.
     toolkit_backend : Literal['ambertools', 'openeye', 'rdkit']
@@ -477,6 +508,8 @@ def bulk_assign_partial_charges(
       If ``None``, the latest am1bcc NAGL charge model is used.
     processors: int, default 1
         The number of processors which should be used to generate the charges.
+    forcefields : list[str] | None, default None
+        An optional list of SMIRNOFF style force field offxml paths or strings which should be used to assign partial charges.
 
     Raises
     ------
@@ -499,6 +532,7 @@ def bulk_assign_partial_charges(
         "toolkit_backend": toolkit_backend,
         "generate_n_conformers": generate_n_conformers,
         "nagl_model": nagl_model,
+        "forcefields": forcefields,
     }
 
     if processors > 1:

src/openfe/tests/protocols/test_openmmutils.py

@@ -14,6 +14,7 @@
 from gufe.components.errors import ComponentValidationError
 from gufe.settings import OpenMMSystemGeneratorFFSettings, ThermoSettings
 from numpy.testing import assert_allclose, assert_equal
+from openff.toolkit import ForceField
 from openff.toolkit import Molecule as OFFMol
 from openff.toolkit.utils.toolkit_registry import ToolkitRegistry
 from openff.toolkit.utils.toolkits import RDKitToolkitWrapper
@@ -1291,6 +1292,64 @@ def test_openeye_import_error(self, monkeypatch, uncharged_mol):
                 nagl_model=None,
             )
 
+    def test_forcefield_missing_ff(self, uncharged_mol):
+        # Make sure an error is raised if we forget to pass a force field to charge with
+        with pytest.raises(
+            ValueError,
+            match="The forcefield method requires a force field or list of force fields' to be provided via `forcefields`.",
+        ):
+            charge_generation.assign_offmol_partial_charges(
+                uncharged_mol,
+                overwrite=False,
+                method="forcefield",
+                toolkit_backend="rdkit",
+                generate_n_conformers=None,
+                nagl_model=None,
+            )
+
+    def test_forcefield_charges_library(self, uncharged_mol):
+        # Make sure that the forcefield method can assign charges from a library
+        # Create a force field with a charge library for the molecule using a force field with an AM1BCC handler as well
+        ff = ForceField("openff-2.0.0.offxml")
+        lib_handler = ff.get_parameter_handler("LibraryCharges")
+        # add the new parameter
+        charged_mol = copy.deepcopy(uncharged_mol)
+        dummy_charges = np.zeros(charged_mol.n_atoms) * unit.e
+        # no other method should assign all zero charges
+        charged_mol.partial_charges = dummy_charges
+        lib_param = lib_handler._INFOTYPE.from_molecule(charged_mol)
+        lib_handler.add_parameter(parameter=lib_param)
+        del charged_mol
+        charge_generation.assign_offmol_partial_charges(
+            uncharged_mol,
+            overwrite=False,
+            method="forcefield",
+            toolkit_backend="rdkit",
+            generate_n_conformers=None,
+            nagl_model=None,
+            forcefields=ff.to_string(),
+        )
+
+        assert_allclose(uncharged_mol.partial_charges.m, dummy_charges.m)
+
+    @pytest.mark.skipif(not HAS_NAGL, reason="NAGL is not available")
+    def test_forcefield_nagl_charges(self, uncharged_mol):
+        # Make sure that the forcefield method can assign charges from a NAGL model
+        charge_generation.assign_offmol_partial_charges(
+            uncharged_mol,
+            overwrite=False,
+            method="forcefield",
+            toolkit_backend="rdkit",
+            generate_n_conformers=None,
+            # set the model to none this should use the model define in the force field.
+            nagl_model=None,
+            # use a force field that has a NAGL handler and another redundant force field file
+            forcefields=["openff-2.3.0.offxml", "opc-1.0.0.offxml"],
+        )
+
+        assert uncharged_mol.partial_charges is not None
+        assert np.any(uncharged_mol.partial_charges)
+
 
 @pytest.mark.slow
 @pytest.mark.skipif(

src/openfecli/parameters/plan_network_options.py

@@ -223,6 +223,7 @@ def load_yaml_planner_options(path: Optional[str], context) -> PlanNetworkOption
             off_toolkit_backend: ambertools
             number_of_conformers: None
             nagl_model: None
+            forcefields: None
 """
 
 _yaml_help = """
@@ -245,6 +246,7 @@ def load_yaml_planner_options(path: Optional[str], context) -> PlanNetworkOption
     - ``am1bccelf10`` (only possible if ``off_toolkit_backend`` is ``openeye``)
     - ``nagl`` (must have openff-nagl installed)
     - ``espaloma`` (must have espaloma_charge installed)
+    - ``forcefield`` (must supply the chosen force field files via the ``forcefields`` keyword argument. This is useful to get the correct AshGC model or LibraryCharges for a OpenFF force field.)
 
 ``settings:`` allows for passing in any keyword arguments of the method's corresponding Python API.