rfe protocol: partial charges fix

calculate partial charges as separate step and use these as "user charges" throughout

this avoids a case where conformer generation can fail and parametrisation fails

this is technically not fully SMIRNOFF spec, but it probably isn't all that bad either

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -22,7 +22,7 @@
 from collections import defaultdict
 import uuid
 import warnings
-
+from itertools import chain
 import numpy as np
 import numpy.typing as npt
 from openff.units import unit
@@ -538,26 +538,43 @@ def run(self, *, dry=False, verbose=True,
             has_solvent=solvent_comp is not None,
         )
 
+        # workaround for conformer generation failures
+        # see openfe issue #576
+        # calculate partial charges manually if not already given
+        # convert to OpenFF here,
+        # and keep the molecule around to maintain the partial charges
+        off_small_mols = {
+            'stateA': mapping.componentA.to_openff(),
+            'stateB': mapping.componentB.to_openff(),
+            'both': [m.to_openff() for m in small_mols
+                     if (m != mapping.componentA and m != mapping.componentB)]
+        }
+
         # b. force the creation of parameters
         # This is necessary because we need to have the FF generated ahead of
         # solvating the system.
         # Note: by default this is cached to ctx.shared/db.json so shouldn't
         # incur too large a cost
         self.logger.info("Parameterizing molecules")
-        for comp in small_mols:
-            offmol = comp.to_openff()
-            system_generator.create_system(offmol.to_topology().to_openmm(),
-                                           molecules=[offmol])
-            if comp == mapping.componentA:
-                molB = mapping.componentB.to_openff()
-                system_generator.create_system(molB.to_topology().to_openmm(),
-                                               molecules=[molB])
+        for mol in chain(off_small_mols['stateA'], off_small_mols['stateB'],
+                         *off_small_mols['both']):
+            # robustly calculate partial charges;
+            try:
+                # try and follow official spec method
+                mol.assign_partial_charges('am1bcc')
+            except ValueError:  # this is what a confgen failure yields
+                # but fallback to using existing conformer
+                mol.assign_partial_charges('am1bcc',
+                                           use_conformers=mol.conformers)
+
+            system_generator.create_system(mol.to_topology().to_openmm(),
+                                           molecules=[mol])
 
         # c. get OpenMM Modeller + a dictionary of resids for each component
         stateA_modeller, comp_resids = system_creation.get_omm_modeller(
             protein_comp=protein_comp,
             solvent_comp=solvent_comp,
-            small_mols=small_mols,
+            small_mols=[off_small_mols['stateA']] + off_small_mols['both'],
             omm_forcefield=system_generator.forcefield,
             solvent_settings=solvation_settings,
         )
@@ -572,7 +589,7 @@ def run(self, *, dry=False, verbose=True,
         # e. create the stateA System
         stateA_system = system_generator.create_system(
             stateA_modeller.topology,
-            molecules=[s.to_openff() for s in small_mols],
+            molecules=[off_small_mols['stateA']] + off_small_mols['both'],
         )
 
 
@@ -585,14 +602,9 @@ def run(self, *, dry=False, verbose=True,
         )
 
         # b. get a list of small molecules for stateB
-        off_mols_stateB = [mapping.componentB.to_openff(),]
-        for comp in small_mols:
-            if comp != mapping.componentA:
-                off_mols_stateB.append(comp.to_openff())
-
         stateB_system = system_generator.create_system(
             stateB_topology,
-            molecules=off_mols_stateB,
+            molecules=[off_small_mols['stateB']] + off_small_mols['both']
         )
 
         #  c. Define correspondence mappings between the two systems

openfe/protocols/openmm_utils/system_creation.py

@@ -132,7 +132,7 @@ def get_system_generator(
 
 def get_omm_modeller(protein_comp: Optional[ProteinComponent],
                      solvent_comp: Optional[SolventComponent],
-                     small_mols: list[SmallMoleculeComponent],
+                     small_mols: list,
                      omm_forcefield : app.ForceField,
                      solvent_settings : SolvationSettings) -> ModellerReturn:
     """
@@ -145,8 +145,8 @@ def get_omm_modeller(protein_comp: Optional[ProteinComponent],
       Protein Component, if it exists.
     solvent_comp : Optional[ProteinCompoinent]
       Solvent Component, if it exists.
-    small_mols : list[SmallMoleculeComponents]
-      List of SmallMoleculeComponents to add.
+    small_mols : list
+      List of OpenFF Molecule to add.
     omm_forcefield : app.ForceField
       ForceField object for system.
     solvent_settings : SolvationSettings
@@ -162,18 +162,18 @@ def get_omm_modeller(protein_comp: Optional[ProteinComponent],
     """
     component_resids = {}
 
-    def _add_small_mol(comp: SmallMoleculeComponent,
+    def _add_small_mol(comp,
                        system_modeller: app.Modeller,
                        comp_resids: dict[Component, npt.NDArray]):
         """
         Helper method to add OFFMol to an existing Modeller object and
         update a dictionary tracking residue indices for each component.
         """
-        mol = comp.to_openff()
-        omm_top = mol.to_topology().to_openmm()
+        # mol = comp.to_openff()
+        omm_top = comp.to_topology().to_openmm()
         system_modeller.add(
             omm_top,
-            ensure_quantity(mol.conformers[0], 'openmm')
+            ensure_quantity(comp.conformers[0], 'openmm')
         )
 
         nres = omm_top.getNumResidues()