Non-deterministic charge validation

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)

src/openfe/tests/protocols/test_openmmutils.py

@@ -4,17 +4,21 @@
 import gzip
 import logging
 import os
+import re
 import sys
 from importlib import resources
 from pathlib import Path
 from unittest import mock
 
 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()
+        )