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,14 @@ 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.
 
     Raises
     ------
@@ -339,6 +349,42 @@ 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]
+
+        try:
+            # try to parse what the user has provided, due to supporting dropping the offxml extension,
+            # we need to try and catch the OSError and add the extension if needed
+            ff = ForceField(*forcefields)
+        except OSError:
+            # try adding the offxml extension if not present and it's a possible file path
+            forcefields_with_ext = []
+            for _ff in forcefields:
+                # if the string of the force field is passed it should start with the xml header
+                if not _ff.endswith(".offxml") and not _ff.startswith("<?xml"):
+                    ff_with_ext = f"{_ff}.offxml"
+                    forcefields_with_ext.append(ff_with_ext)
+                else:
+                    forcefields_with_ext.append(_ff)
+
+            # try again to load the force field with the added extension, if we fail let it raise the error
+            ff = ForceField(*forcefields_with_ext)
+
+        # make the toolkit registry based on the selected backend
+        toolkits = ToolkitRegistry([i() for i in BACKEND_OPTIONS[toolkit_backend.lower()]])
+        # let the force field resolve the partial charge assignment method
+        charges = ff.get_partial_charges(offmol, toolkit_registry=toolkits)
+        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 +487,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 +504,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 +524,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 +548,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/protocols/openmm_utils/omm_settings.py

@@ -238,7 +238,9 @@ class OpenFFPartialChargeSettings(BasePartialChargeSettings):
     Settings for controlling partial charge assignment using the OpenFF tooling
     """
 
-    partial_charge_method: Literal["am1bcc", "am1bccelf10", "nagl", "espaloma"] = "am1bcc"
+    partial_charge_method: Literal["am1bcc", "am1bccelf10", "nagl", "espaloma", "forcefield"] = (
+        "am1bcc"
+    )
     """
     Selection of method for partial charge generation.
 
@@ -271,11 +273,17 @@ class OpenFFPartialChargeSettings(BasePartialChargeSettings):
       Only ``ambertools`` and ``rdkit`` `off_toolkit_backend`` options
       are supported. A maximum of one conformer is allowed.
 
+    ``forcefield``:
+        Assign partial charges using the OpenFF force field's defined charge model, this is useful to get the correct
+        NAGL(AshGC) model for a specific force field or to use LibraryCharges.
+
     """
     off_toolkit_backend: Literal["ambertools", "openeye", "rdkit"] = "ambertools"
     """
     The OpenFF toolkit registry backend to use for partial charge generation.
 
+    This is always respected regardless of the ``partial_charge_method``.
+
 
     OpenFF backend selection options
     --------------------------------
@@ -314,6 +322,11 @@ class OpenFFPartialChargeSettings(BasePartialChargeSettings):
     If ``None`` (default) and ``partial_charge_method`` is set to ``nagl``,
     the latest available production am1bcc charge model will be used.
     """
+    forcefields: list[str] | None = None
+    """
+    An optional list of SMIRNOFF style force field offxml paths or raw force field contents strings which should be used to assign partial
+    charges if the ``partial_charge_method`` is set to ``forcefield``.
+    """
 
 
 class OpenMMEngineSettings(SettingsBaseModel):

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,74 @@ 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
+        opc = ForceField("opc-1.0.0.offxml")
+        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
+            # this is also testing that missing the file extension doesn't break the code
+            forcefields=["openff-2.3.0", opc.to_string()],
+        )
+
+        assert uncharged_mol.partial_charges is not None
+
+        # get the reference charges to compare with
+        ff = ForceField("openff-2.3.0.offxml")
+        nagl_model = ff.get_parameter_handler("NAGLCharges").model_file
+        copy_mol = copy.deepcopy(uncharged_mol)
+        copy_mol.partial_charges = None
+        copy_mol.assign_partial_charges(partial_charge_method=nagl_model)
+
+        assert_allclose(uncharged_mol.partial_charges.m, copy_mol.partial_charges.m, rtol=1e-4)
+
 
 @pytest.mark.slow
 @pytest.mark.skipif(

src/openfecli/commands/generate_partial_charges.py

@@ -14,6 +14,7 @@
     - ``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.)
 
 ``settings`` allows for passing in any keyword arguments of the method's corresponding Python API.
 
@@ -80,6 +81,7 @@ def charge_molecules(molecules, yaml_settings, output, n_cores, overwrite_charge
         generate_n_conformers=partial_charge.number_of_conformers,
         nagl_model=partial_charge.nagl_model,
         processors=n_cores,
+        forcefields=partial_charge.forcefields,
     )
 
     write("\tDone")

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.
 

src/openfecli/tests/commands/test_charge_generation.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+import yaml
 from click import ClickException
 from click.testing import CliRunner
 from gufe import SmallMoleculeComponent
@@ -188,3 +189,71 @@ def test_charge_settings(
             output_order.append(smc.name)
 
         assert input_order == output_order
+
+
+def test_charge_molecules_missing_force_fields(methane, tmp_path):
+    # make sure an error is raised if we try to use forcefield charges without specifying a forcefield
+    runner = CliRunner()
+    mol_path = tmp_path / "methane.sdf"
+    methane.to_file(str(mol_path), "sdf")
+    out_path = str(tmp_path / "charged_methane.sdf")
+
+    settings = {
+        "partial_charge": {
+            "method": "forcefield",
+        }
+    }
+
+    settings_path = tmp_path / "settings.yaml"
+    yaml.safe_dump(settings, open(settings_path, "w"))
+
+    with runner.isolated_filesystem():
+        # # check an error is raised if we try to overwrite the input
+        with pytest.raises(
+            ValueError,
+            match="The forcefield method requires a force field or list of force fields' to be provided via `forcefields`.",
+        ):
+            _ = runner.invoke(
+                charge_molecules,
+                ["-M", mol_path, "-o", out_path, "-s", settings_path],
+                catch_exceptions=False,
+            )
+
+
+@pytest.mark.skipif(
+    not HAS_NAGL,
+    reason="needs NAGL",
+)
+@pytest.mark.skipif(
+    HAS_OPENEYE, reason="cannot use NAGL with rdkit backend when OpenEye is installed"
+)
+def test_charge_molecules_from_forcefield(methane, tmp_path):
+    # make sure we can use forcefield charges if we specify a forcefield
+    runner = CliRunner()
+    mol_path = tmp_path / "methane.sdf"
+    methane.to_file(str(mol_path), "sdf")
+    out_path = str(tmp_path / "charged_methane.sdf")
+
+    settings = {
+        "partial_charge": {
+            "method": "forcefield",
+            "settings": {"forcefields": ["openff_unconstrained-2.3.0"]},
+        }
+    }
+
+    settings_path = tmp_path / "settings.yaml"
+    yaml.safe_dump(settings, open(settings_path, "w"))
+
+    with runner.isolated_filesystem():
+        result = runner.invoke(
+            charge_molecules,
+            ["-M", mol_path, "-o", out_path, "-s", settings_path],
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert "Partial Charge Generation: forcefield" in result.output
+
+        # make sure the charges have been saved
+        methane_out = SmallMoleculeComponent.from_sdf_file(filename=out_path)
+        off_methane_out = methane_out.to_openff()
+        assert off_methane_out.partial_charges is not None