diff --git a/news/charge_from_ff.rst b/news/charge_from_ff.rst new file mode 100644 index 000000000..ea8c4fbc1 --- /dev/null +++ b/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:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/src/openfe/protocols/openmm_utils/charge_generation.py b/src/openfe/protocols/openmm_utils/charge_generation.py index 0196eb5f0..8abf0651b 100644 --- a/src/openfe/protocols/openmm_utils/charge_generation.py +++ b/src/openfe/protocols/openmm_utils/charge_generation.py @@ -6,11 +6,13 @@ import copy import sys +import typing import warnings from typing import Callable, Literal 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 @@ -69,6 +71,12 @@ "openeye": [OpenEyeToolkitWrapper], "rdkit": [RDKitToolkitWrapper], } +# If the user wants to use NAGL to assign charges via the force field option +# then we need to add it to the backend options if it is available +if HAS_NAGL: + BACKEND_OPTIONS["ambertools"].append(NAGLToolkitWrapper) + BACKEND_OPTIONS["openeye"].append(NAGLToolkitWrapper) + BACKEND_OPTIONS["rdkit"].append(NAGLToolkitWrapper) def assign_offmol_espaloma_charges(offmol: OFFMol, toolkit_registry: ToolkitRegistry) -> None: @@ -286,10 +294,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 +308,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 +328,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 ------ @@ -333,12 +350,54 @@ def assign_offmol_partial_charges( ------- The Molecule with partial charges assigned. """ + method_name = method.lower() + toolkit_name = toolkit_backend.lower() + + # validate the input combination first + if method_name == "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) + elif forcefields is not None: + errmsg = f"The `forcefields` option is only valid with the `forcefield` charge method, but got {method_name}." + raise ValueError(errmsg) # If you have non-zero charges and not overwriting, just return if offmol.partial_charges is not None and np.any(offmol.partial_charges): if not overwrite: return offmol + if method_name == "forcefield": + # mypy can't tell we have already validated that forcefields is not None + forcefields = typing.cast(list[str], 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(" with threadpool_limits(limits=1): # Call selected method to assign partial charges - CHARGE_METHODS[method.lower()]["charge_func"]( + CHARGE_METHODS[method_name]["charge_func"]( offmol=offmol_copy, toolkit_registry=toolkits, - **CHARGE_METHODS[method.lower()]["charge_extra_kwargs"], + **CHARGE_METHODS[method_name]["charge_extra_kwargs"], ) # type: ignore # Copy partial charges back @@ -441,11 +500,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 +517,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 +537,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 +561,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: diff --git a/src/openfe/protocols/openmm_utils/omm_settings.py b/src/openfe/protocols/openmm_utils/omm_settings.py index 97c7ad8ce..0b8c575b8 100644 --- a/src/openfe/protocols/openmm_utils/omm_settings.py +++ b/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): diff --git a/src/openfe/tests/protocols/test_openmmutils.py b/src/openfe/tests/protocols/test_openmmutils.py index f1672b6ba..b422f213f 100644 --- a/src/openfe/tests/protocols/test_openmmutils.py +++ b/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,90 @@ 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_forcefields_wrong_method(self, uncharged_mol): + # Make sure an error is raised if we pass in some forcefields but the method isn't forcefield + with pytest.raises( + ValueError, + match="The `forcefields` option is only valid with the `forcefield` charge method, but got am1bcc.", + ): + charge_generation.assign_offmol_partial_charges( + uncharged_mol, + overwrite=False, + method="am1bcc", + toolkit_backend="rdkit", + generate_n_conformers=None, + nagl_model=None, + forcefields=["openff-2.0.0.offxml"], + ) + + 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( diff --git a/src/openfecli/commands/generate_partial_charges.py b/src/openfecli/commands/generate_partial_charges.py index d1a83580f..3243263b9 100644 --- a/src/openfecli/commands/generate_partial_charges.py +++ b/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") diff --git a/src/openfecli/parameters/plan_network_options.py b/src/openfecli/parameters/plan_network_options.py index b349d8972..b2b2a6871 100644 --- a/src/openfecli/parameters/plan_network_options.py +++ b/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. diff --git a/src/openfecli/tests/commands/test_charge_generation.py b/src/openfecli/tests/commands/test_charge_generation.py index d518a7ed7..9fcb27322 100644 --- a/src/openfecli/tests/commands/test_charge_generation.py +++ b/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-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