diff --git a/src/openfe/protocols/openmm_utils/system_validation.py b/src/openfe/protocols/openmm_utils/system_validation.py index 906999910..6ba0bc47a 100644 --- a/src/openfe/protocols/openmm_utils/system_validation.py +++ b/src/openfe/protocols/openmm_utils/system_validation.py @@ -22,6 +22,8 @@ SolventComponent, ) from gufe.components.errors import ComponentValidationError +from gufe.protocols.errors import ProtocolValidationError +from openff.toolkit import ForceField from openff.toolkit import Molecule as OFFMol logger = logging.getLogger(__name__) @@ -371,3 +373,50 @@ def validate_chemical_system(system: ChemicalSystem): except ComponentValidationError as e: errmsg = f"Component {entry} from ChemicalSystem {system.name} failed validation: {e}" raise ComponentValidationError(errmsg) + + +def validate_nondeterministic_charges(system: ChemicalSystem, small_molecule_forcefield: str): + """ + Validate that the SmallMoleculeComponents of the system will have deterministic partial charges. + + This is determined by checking for charges on the molecules before checking what would be assigned by the force field. + + Parameters + ---------- + system : ChemicalSystem + The ChemicalSystem to validate with SmallMoleculeComponents. + small_molecule_forcefield : str + The force field to be used for the SmallMoleculeComponents. + + Raises + ------ + ProtocolValidationError + If any SmallMoleculeComponents in the system would have am1bcc charges generated at runtime. + """ + smcs: list[SmallMoleculeComponent] = system.get_components_of_type(SmallMoleculeComponent) + if "espaloma" in small_molecule_forcefield or "gaff" in small_molecule_forcefield: + # this will always generate charges at runtime, so raise an error for missing charges + ff = None + else: + # We do not check for offxml in the name as users can pass the force field contents as a string + ff = ForceField(small_molecule_forcefield) + + for smc in smcs: + offmol = smc.to_openff() + if offmol.partial_charges is not None and np.any(offmol.partial_charges): + continue + + # check the labels assigned for an openff force field + if ff is not None: + labels = ff.label_molecules(offmol.to_topology())[0] + else: + # return a dummy label as the gaff and espaloma should always give am1bcc charges + labels = {"LibraryCharges": {}} + + # We count library and nagl charges as deterministic + if len(labels["LibraryCharges"]) != offmol.n_atoms and "NAGLCharges" not in labels: + errmsg = ( + f"{smc} from system {system.name} would have am1bcc charges generated at runtime which is non-deterministic. " + f"Please provide a molecule with pre-computed charges or use library charges instead." + ) + raise ProtocolValidationError(errmsg) diff --git a/src/openfe/tests/protocols/test_openmmutils.py b/src/openfe/tests/protocols/test_openmmutils.py index d4c50a5da..d76cd7466 100644 --- a/src/openfe/tests/protocols/test_openmmutils.py +++ b/src/openfe/tests/protocols/test_openmmutils.py @@ -4,6 +4,7 @@ import gzip import logging import os +import re import sys from importlib import resources from pathlib import Path @@ -11,10 +12,13 @@ import numpy as np import pytest +from gufe import ChemicalSystem from gufe.components.errors import ComponentValidationError +from gufe.protocols.errors import ProtocolValidationError from gufe.settings import OpenMMSystemGeneratorFFSettings, ThermoSettings from numpy.testing import assert_allclose, assert_equal from openff.toolkit import Molecule as OFFMol +from openff.toolkit import ForceField from openff.toolkit.utils.toolkit_registry import ToolkitRegistry from openff.toolkit.utils.toolkits import RDKitToolkitWrapper from openff.units import unit @@ -1283,3 +1287,76 @@ def test_set_metadata_none_clears(): _set_offmol_metadata(mol, "residue_name", "LIG") _set_offmol_metadata(mol, "residue_name", None) assert all("residue_name" not in a.metadata for a in mol.atoms) + + +class TestChargeValidation: + """Test validation of nondeterministic partial charge assignment.""" + + @pytest.fixture + def benzene_charged_system(self, benzene_modifications): + return ChemicalSystem({"ligand": benzene_modifications["benzene"]}, name="charged") + + @pytest.fixture + def benzene_no_charge_system(self, benzene_modifications_uncharged): + return ChemicalSystem({"ligand": benzene_modifications_uncharged["benzene"]}, name="no charges") + + def test_gaff_with_molecule_charges(self, benzene_charged_system): + system_validation.validate_nondeterministic_charges( + benzene_charged_system, + small_molecule_forcefield="gaff-2.11" + ) + + def test_gaff_no_charges(self, benzene_no_charge_system): + with pytest.raises(ProtocolValidationError, match=re.escape("SmallMoleculeComponent(name=benzene) from system no charges would have am1bcc charges")): + system_validation.validate_nondeterministic_charges( + benzene_no_charge_system, + small_molecule_forcefield="gaff-2.11" + ) + + def test_espaloma_with_molecule_charges(self, benzene_charged_system): + system_validation.validate_nondeterministic_charges( + benzene_charged_system, + small_molecule_forcefield="espaloma-0.3.2" + ) + + def test_espaloma_no_charges(self, benzene_no_charge_system): + with pytest.raises(ProtocolValidationError, match=re.escape("SmallMoleculeComponent(name=benzene) from system no charges would have am1bcc charges")): + system_validation.validate_nondeterministic_charges( + benzene_no_charge_system, + small_molecule_forcefield="espaloma-0.3.2" + ) + + def test_openff_with_molecule_charges(self, benzene_charged_system): + system_validation.validate_nondeterministic_charges( + benzene_charged_system, + small_molecule_forcefield="openff-2.2.0.offxml" + ) + + def test_openff_no_charges(self, benzene_no_charge_system): + with pytest.raises(ProtocolValidationError, match=re.escape("SmallMoleculeComponent(name=benzene) from system no charges would have am1bcc charges")): + system_validation.validate_nondeterministic_charges( + benzene_no_charge_system, + small_molecule_forcefield="openff-2.2.0.offxml" + ) + + def test_openff_nagl_no_charges(self, benzene_no_charge_system): + system_validation.validate_nondeterministic_charges( + benzene_no_charge_system, + small_molecule_forcefield="openff-2.3.0.offxml" + ) + + def test_openff_lib_charges_no_charges(self, benzene_no_charge_system, benzene_modifications): + # add some library charges to the force field and test using a string + benzene = benzene_modifications["benzene"].to_openff() + # use a force field with an am1bcc handler + ff = ForceField("openff-2.2.0.offxml") + lib_charge_handler = ff.get_parameter_handler("LibraryCharges") + # make a new parameter + lib_charge = lib_charge_handler._INFOTYPE.from_molecule(benzene) + # add it to the handler + lib_charge_handler.add_parameter(parameter=lib_charge) + # run the validation + system_validation.validate_nondeterministic_charges( + benzene_no_charge_system, + small_molecule_forcefield=ff.to_string() + )