Non-deterministic charge validation

src/openfe/protocols/openmm_afe/base_afe_units.py

@@ -67,7 +67,6 @@
 )
 from openfe.protocols.openmm_md.plain_md_methods import PlainMDSimulationUnit
 from openfe.protocols.openmm_utils import (
-    charge_generation,
     multistate_analysis,
     omm_compute,
     settings_validation,
@@ -349,33 +348,6 @@ def _pre_equilibrate(
 
         return equilibrated_positions, to_openmm(from_openmm(box))
 
-    @staticmethod
-    def _assign_partial_charges(
-        partial_charge_settings: OpenFFPartialChargeSettings,
-        small_mols: dict[SmallMoleculeComponent, OFFMolecule],
-    ) -> None:
-        """
-        Assign partial charges to the OpenFF Molecules associated with
-        all the SmallMoleculeComponents in the transformation.
-
-        Parameters
-        ----------
-        charge_settings : OpenFFPartialChargeSettings
-          Settings for controlling how the partial charges are assigned.
-        small_mols : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
-          Dictionary of OpenFF Molecules to add, keyed by their
-          associated SmallMoleculeComponent.
-        """
-        for mol in small_mols.values():
-            charge_generation.assign_offmol_partial_charges(
-                offmol=mol,
-                overwrite=False,
-                method=partial_charge_settings.partial_charge_method,
-                toolkit_backend=partial_charge_settings.off_toolkit_backend,
-                generate_n_conformers=partial_charge_settings.number_of_conformers,
-                nagl_model=partial_charge_settings.nagl_model,
-            )
-
     @staticmethod
     def _get_system_generator(
         settings: dict[str, SettingsBaseModel],
@@ -391,7 +363,7 @@ def _get_system_generator(
         ----------
         settings : dict[str, SettingsBaseModel]
           A dictionary of settings object for the unit.
-        solvent_comp : BaseSolventComponent | None
+        solvent_component : BaseSolventComponent | None
           The solvent component of this system, if there is one.
         openff_molecules : list[openff.toolkit.Molecule] | None
           A list of OpenFF Molecules to generate templates for, if any.
@@ -704,9 +676,6 @@ def run(
         # Get settings
         settings = self._get_settings()
 
-        # Assign partial charges now to avoid any discrepancies later
-        self._assign_partial_charges(settings["charge_settings"], small_mols)
-
         # Get OpenMM topology, positions, system, and comp_resids
         omm_topology, omm_system, positions, comp_resids = self._get_omm_objects(
             settings=settings,

src/openfe/protocols/openmm_afe/equil_binding_afe_method.py

@@ -403,6 +403,11 @@ def _validate(
         system_validation.validate_chemical_system(stateB)
         self._validate_endstates(stateA, stateB)
 
+        # validate small molecule charges
+        small_ff = self.settings.forcefield_settings.small_molecule_forcefield
+        system_validation.validate_nondeterministic_charges(stateA, small_ff)
+        system_validation.validate_nondeterministic_charges(stateB, small_ff)
+
         # Validate the complex lambda schedule
         self._validate_lambda_schedule(
             self.settings.complex_lambda_settings,

src/openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -385,6 +385,11 @@ def _validate(
         system_validation.validate_chemical_system(stateB)
         self._validate_endstates(stateA, stateB)
 
+        # validate small molecule charges
+        small_ff = self.settings.solvent_forcefield_settings.small_molecule_forcefield
+        system_validation.validate_nondeterministic_charges(stateA, small_ff)
+        system_validation.validate_nondeterministic_charges(stateB, small_ff)
+
         # Validate the lambda schedule
         for solv_sets in (
             self.settings.solvent_simulation_settings,

src/openfe/protocols/openmm_md/plain_md_methods.py

@@ -50,7 +50,6 @@
     PlainMDProtocolSettings,
 )
 from openfe.protocols.openmm_utils import (
-    charge_generation,
     omm_compute,
     serialization,
     settings_validation,
@@ -197,6 +196,10 @@ def _validate(
         # Validate the ChemicalSystem
         system_validation.validate_chemical_system(stateA)
 
+        # Validate small molecule charges
+        small_ff = self.settings.forcefield_settings.small_molecule_forcefield
+        system_validation.validate_nondeterministic_charges(stateA, small_ff)
+
         # Validate solvent component if present
         nonbond = self.settings.forcefield_settings.nonbonded_method
         system_validation.validate_solvent(stateA, nonbond)
@@ -333,32 +336,6 @@ class PlainMDSetupUnit(PlainMDUnitMixin, gufe.ProtocolUnit):
     Protocol setup unit for plain MD simulations which handles charging, system building and solvation.
     """
 
-    @staticmethod
-    def _assign_partial_charges(
-        charge_settings: OpenFFPartialChargeSettings,
-        smc_components: dict[SmallMoleculeComponent, OFFMolecule],
-    ) -> None:
-        """
-        Assign partial charges to SMCs.
-
-        Parameters
-        ----------
-        charge_settings : OpenFFPartialChargeSettings
-          Settings for controlling how the partial charges are assigned.
-        smc_components : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
-          Dictionary of OpenFF Molecules to add, keyed by
-          SmallMoleculeComponent.
-        """
-        for mol in smc_components.values():
-            charge_generation.assign_offmol_partial_charges(
-                offmol=mol,
-                overwrite=False,
-                method=charge_settings.partial_charge_method,
-                toolkit_backend=charge_settings.off_toolkit_backend,
-                generate_n_conformers=charge_settings.number_of_conformers,
-                nagl_model=charge_settings.nagl_model,
-            )
-
     def run(
         self,
         *,
@@ -450,9 +427,6 @@ def run(
             i: i.to_openff() for i in small_mols
         }
 
-        # a. assign partial charges to smcs
-        self._assign_partial_charges(charge_settings, smc_components)
-
         # b. get a system generator
         if output_settings.forcefield_cache is not None:
             ffcache = self.shared_basepath / output_settings.forcefield_cache

src/openfe/protocols/openmm_rfe/hybridtop_protocols.py

@@ -618,6 +618,11 @@ def _validate(
         system_validation.validate_chemical_system(stateB)
         self._validate_endstates(stateA, stateB)
 
+        # Validate small molecule charges
+        small_ff = self.settings.forcefield_settings.small_molecule_forcefield
+        system_validation.validate_nondeterministic_charges(stateA, small_ff)
+        system_validation.validate_nondeterministic_charges(stateB, small_ff)
+
         # Validate the mapping
         alchem_comps = system_validation.get_alchemical_components(stateA, stateB)
         self._validate_mapping(mapping, alchem_comps)

src/openfe/protocols/openmm_rfe/hybridtop_units.py

@@ -54,7 +54,6 @@
 from ...analysis import plotting
 from ...utils import log_system_probe, without_oechem_backend
 from ..openmm_utils import (
-    charge_generation,
     multistate_analysis,
     omm_compute,
     settings_validation,
@@ -205,33 +204,6 @@ def _get_components(
 
         return solvent_comp, protein_comp, small_mols
 
-    @staticmethod
-    def _assign_partial_charges(
-        charge_settings: OpenFFPartialChargeSettings,
-        small_mols: dict[SmallMoleculeComponent, OFFMolecule],
-    ) -> None:
-        """
-        Assign partial charges to the OpenFF Molecules associated with all
-        the SmallMoleculeComponents in the transformation.
-
-        Parameters
-        ----------
-        charge_settings : OpenFFPartialChargeSettings
-          Settings for controlling how the partial charges are assigned.
-        small_mols : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
-          Dictionary of OpenFF Molecules to add, keyed by
-          their associated SmallMoleculeComponent.
-        """
-        for smc, mol in small_mols.items():
-            charge_generation.assign_offmol_partial_charges(
-                offmol=mol,
-                overwrite=False,
-                method=charge_settings.partial_charge_method,
-                toolkit_backend=charge_settings.off_toolkit_backend,
-                generate_n_conformers=charge_settings.number_of_conformers,
-                nagl_model=charge_settings.nagl_model,
-            )
-
     @staticmethod
     def _get_system_generator(
         settings: dict[str, SettingsBaseModel],
@@ -761,9 +733,6 @@ def run(
         # in the topology
         alchem_resnames = [assigned[alchem_comps["stateA"][0]]]
 
-        # Assign partial charges now to avoid any discrepancies later
-        self._assign_partial_charges(settings["charge_settings"], small_mols)
-
         (
             stateA_system,
             stateA_topology,

src/openfe/protocols/openmm_septop/base_units.py

@@ -69,7 +69,6 @@
 from openfe.utils import log_system_probe, without_oechem_backend
 
 from ..openmm_utils import (
-    charge_generation,
     multistate_analysis,
     settings_validation,
     system_creation,
@@ -457,32 +456,6 @@ def _get_system_generator(
             )
         return system_generator
 
-    @staticmethod
-    def _assign_partial_charges(
-        partial_charge_settings: OpenFFPartialChargeSettings,
-        smc_components: dict[SmallMoleculeComponent, OFFMolecule],
-    ) -> None:
-        """
-        Assign partial charges to OFFMolecules inplace.
-
-        Parameters
-        ----------
-        charge_settings : OpenFFPartialChargeSettings
-          Settings for controlling how the partial charges are assigned.
-        smc_components : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
-          Dictionary of OpenFF Molecules to add, keyed by
-          SmallMoleculeComponent.
-        """
-        for mol in smc_components.values():
-            charge_generation.assign_offmol_partial_charges(
-                offmol=mol,
-                overwrite=False,
-                method=partial_charge_settings.partial_charge_method,
-                toolkit_backend=partial_charge_settings.off_toolkit_backend,
-                generate_n_conformers=partial_charge_settings.number_of_conformers,
-                nagl_model=partial_charge_settings.nagl_model,
-            )
-
     def _get_modeller(
         self,
         protein_component: Optional[ProteinComponent],

src/openfe/protocols/openmm_septop/equil_septop_method.py

@@ -485,6 +485,11 @@ def _validate(
         system_validation.validate_chemical_system(stateB)
         self._validate_endstates(stateA, stateB)
 
+        # Validate small molecule charges
+        small_ff = self.settings.forcefield_settings.small_molecule_forcefield
+        system_validation.validate_nondeterministic_charges(stateA, small_ff)
+        system_validation.validate_nondeterministic_charges(stateB, small_ff)
+
         # Validate the lambda schedule
         self._validate_lambda_schedule(
             self.settings.solvent_lambda_settings,

src/openfe/protocols/openmm_septop/septop_units.py

@@ -711,9 +711,6 @@ def run(
         # 3. Get settings
         settings = self._get_settings()
 
-        # 4. Assign partial charges
-        self._assign_partial_charges(settings["charge_settings"], smc_comps_AB)
-
         # 5. Get the OpenMM systems
         omm_system_A, omm_topology_A, positions_A, modeller_A, comp_resids_A = (
             self.get_system(
@@ -1080,9 +1077,6 @@ def run(
         # 2. Get settings
         settings = self._get_settings()
 
-        # 3. Assign partial charges
-        self._assign_partial_charges(settings["charge_settings"], smc_comps_AB)
-
         # 4. Update the positions of ligand B:
         #    - solvent: Offset ligand B with respect to ligand A
         offset = self._get_ligand_offset(

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,58 @@ 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:
+        try:
+            ff = ForceField(small_molecule_forcefield)
+        except OSError:
+            # try again but adding offxml to the end of the force field name if the user passed one of the installed force fields without the extension
+            try:
+                ff = ForceField(small_molecule_forcefield + ".offxml")
+            except OSError as e:
+                errmsg = f"Could not load force field {small_molecule_forcefield} or {small_molecule_forcefield}.offxml: {e}"
+                raise ProtocolValidationError(errmsg)
+
+    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
+        # While users could use a deterministic charge method with increments we don't currently support this
+        if not labels["LibraryCharges"] 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/analysis/test_plotting.py

@@ -159,15 +159,21 @@ def test_mbar_overlap_plot():
     assert isinstance(ax, matplotlib.axes.Axes)
 
 
-@pytest.mark.parametrize("num", [i for i in range(1, 30)])
-def test_plot_2D_rmsd(num):
+def test_plot_2D_rmsd():
     """
     Smoke test:
       Loop through and test plotting fictitious 2D data
     """
+    num = 29
     points = num * (num - 1) // 2
-    data = [[0.5 for x in range(points)] for i in range(num)]
+    data = [[0.5 for _ in range(points)] for _ in range(num)]
+
     fig = plot_2D_rmsd(data)
+    # check the number of axes and their titles
+    assert len(fig.axes) == 32
+    axis_names = [ax.get_title() for ax in fig.axes if ax.get_title()]
+    assert axis_names == [f"State {i}" for i in range(num)]
+    assert fig._suptitle.get_text() == "Protein 2D RMSD"
     plt.close(fig)