Merge pull request #598 from OpenFreeEnergy/partial_charges_selection

Add charge backends

.github/workflows/ci.yaml

@@ -25,7 +25,7 @@ defaults:
 jobs:
   tests:
     runs-on: ${{ matrix.os }}-latest
-    name: "💻-${{matrix.os }} 🐍-${{ matrix.python-version }} 🗃️${{ matrix.pydantic-version }}"
+    name: "💻-${{matrix.os }} 🐍-${{ matrix.python-version }} 🗃️${{ matrix.pydantic-version }} oechem: ${{ matrix.openeye }}"
     strategy:
       fail-fast: false
       matrix:
@@ -35,13 +35,21 @@ jobs:
           - "3.9"
           - "3.10"
           - "3.11"
+        openeye: ["no"]
         include:
           - os: "macos"
             python-version: "3.11"
             pydantic-version: ">1"
           - os: "ubuntu"
             python-version: "3.11"
             pydantic-version: "<2"
+          - os: "ubuntu"
+            python-version: "3.11"
+            pydantic-version: ">1"
+            openeye: "yes"
+
+    env:
+      OE_LICENSE: ${{ github.workspace }}/oe_license.txt
 
     steps:
       - uses: actions/checkout@v4
@@ -66,6 +74,16 @@ jobs:
             pydantic=${{ matrix.pydantic-version }}
           init-shell: bash
 
+      - name: "Install OpenEye"
+        if: ${{ !github.event.pull_request.head.repo.fork
+                && matrix.openeye == 'yes' }}
+        env:
+          OE_LICENSE_TEXT: ${{ secrets.OE_LICENSE }}
+        run: |
+          echo "${OE_LICENSE_TEXT}" > ${OE_LICENSE}
+          micromamba install -c openeye openeye-toolkits
+          python -c "import openeye; assert openeye.oechem.OEChemIsLicensed(), 'oechem license check failed!'"
+
       - name: "Install GUFE from main@HEAD"
         run: python -m pip install --no-deps git+https://github.com/OpenFreeEnergy/gufe@main
 

.gitignore

@@ -170,3 +170,6 @@ docs/ExampleNotebooks/
 
 # duecredit
 .duecredit.p
+
+# Some charge stuff
+*.model.pt

environment.yml

@@ -19,6 +19,7 @@ dependencies:
   - coverage
   - cinnabar ==0.3.0
   - openff-toolkit>=0.13.0
+  - openff-nagl-base >=0.3.3
   - openff-units==0.2.0
   - pint<0.22
   - openff-models>=0.0.5
@@ -42,4 +43,6 @@ dependencies:
   - sphinx-click
   - pip:
     - sphinx-toolbox
+    - espaloma_charge
+    - openff-nagl-models>=0.1.2
     - git+https://github.com/OpenFreeEnergy/gufe@main

openfe/protocols/openmm_afe/base.py

@@ -62,8 +62,10 @@
 from openfe.protocols.openmm_rfe._rfe_utils import compute
 from ..openmm_utils import (
     settings_validation, system_creation,
-    multistate_analysis
+    multistate_analysis, charge_generation
 )
+from openfe.utils import without_oechem_backend
+
 
 logger = logging.getLogger(__name__)
 
@@ -275,13 +277,16 @@ def _get_system_generator(
         if ffcache is not None:
             ffcache = self.shared_basepath / ffcache
 
-        system_generator = system_creation.get_system_generator(
-            forcefield_settings=settings['forcefield_settings'],
-            integrator_settings=settings['integrator_settings'],
-            thermo_settings=settings['thermo_settings'],
-            cache=ffcache,
-            has_solvent=solvent_comp is not None,
-        )
+        # Block out oechem backend to avoid any issues with
+        # smiles roundtripping between rdkit and oechem
+        with without_oechem_backend():
+            system_generator = system_creation.get_system_generator(
+                forcefield_settings=settings['forcefield_settings'],
+                integrator_settings=settings['integrator_settings'],
+                thermo_settings=settings['thermo_settings'],
+                cache=ffcache,
+                has_solvent=solvent_comp is not None,
+            )
         return system_generator
 
     @staticmethod
@@ -291,6 +296,7 @@ def _assign_partial_charges(
     ) -> None:
         """
         Assign partial charges to SMCs.
+
         Parameters
         ----------
         charge_settings : OpenFFPartialChargeSettings
@@ -300,13 +306,14 @@ def _assign_partial_charges(
           SmallMoleculeComponent.
         """
         for mol in smc_components.values():
-            # don't do this if we have user charges
-            if not (mol.partial_charges is not None and np.any(mol.partial_charges)):
-                # due to issues with partial charge generation in ambertools
-                # we default to using the input conformer for charge generation
-                mol.assign_partial_charges(
-                    'am1bcc', use_conformers=mol.conformers
-                )
+            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,
@@ -328,7 +335,7 @@ def _get_modeller(
         solvent_component : Optional[ProteinCompoinent]
           Solvent Component, if it exists.
         smc_components : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
-          Dicationary of OpenFF Molecules to add, keyed by
+          Dictionary of OpenFF Molecules to add, keyed by
           SmallMoleculeComponent.
         system_generator : openmmforcefields.generator.SystemGenerator
           System Generator to parameterise this unit.
@@ -352,24 +359,26 @@ def _get_modeller(
         # Assign partial charges to smcs
         self._assign_partial_charges(partial_charge_settings, smc_components)
 
+        # TODO: guard the following from non-RDKit backends
         # force the creation of parameters for the small molecules
         # 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 which should
-        # reduce some of the costs.
-        for mol in smc_components.values():
-            system_generator.create_system(
-                mol.to_topology().to_openmm(), molecules=[mol]
-            )
+        # Block out oechem backend to avoid any issues with
+        # smiles roundtripping between rdkit and oechem
+        with without_oechem_backend():
+            for mol in smc_components.values():
+                system_generator.create_system(
+                    mol.to_topology().to_openmm(), molecules=[mol]
+                )
 
-        # get OpenMM modeller + dictionary of resids for each component
-        system_modeller, comp_resids = system_creation.get_omm_modeller(
-            protein_comp=protein_component,
-            solvent_comp=solvent_component,
-            small_mols=smc_components,
-            omm_forcefield=system_generator.forcefield,
-            solvent_settings=solvation_settings,
-        )
+            # get OpenMM modeller + dictionary of resids for each component
+            system_modeller, comp_resids = system_creation.get_omm_modeller(
+                protein_comp=protein_component,
+                solvent_comp=solvent_component,
+                small_mols=smc_components,
+                omm_forcefield=system_generator.forcefield,
+                solvent_settings=solvation_settings,
+            )
 
         return system_modeller, comp_resids
 
@@ -390,7 +399,7 @@ def _get_omm_objects(
           parametrized.
         system_generator : SystemGenerator
           SystemGenerator object to create a System with.
-        smc_components : list[openff.toolkit.Molecules]
+        smc_components : list[openff.toolkit.Molecule]
           A list of openff Molecules to add to the system.
 
         Returns
@@ -405,10 +414,14 @@ def _get_omm_objects(
         topology = system_modeller.getTopology()
         # roundtrip positions to remove vec3 issues
         positions = to_openmm(from_openmm(system_modeller.getPositions()))
-        system = system_generator.create_system(
-            system_modeller.topology,
-            molecules=smc_components,
-        )
+
+        # Block out oechem backend to avoid any issues with
+        # smiles roundtripping between rdkit and oechem
+        with without_oechem_backend():
+            system = system_generator.create_system(
+                system_modeller.topology,
+                molecules=smc_components,
+            )
         return topology, system, positions
 
     def _get_lambda_schedule(
@@ -862,10 +875,10 @@ def run(self, dry=False, verbose=True,
         dry : bool
           Do a dry run of the calculation, creating all necessary alchemical
           system components (topology, system, sampler, etc...) but without
-          running the simulation.
+          running the simulation, default False
         verbose : bool
           Verbose output of the simulation progress. Output is provided via
-          INFO level logging.
+          INFO level logging, default True
         scratch_basepath : pathlib.Path
           Path to the scratch (temporary) directory space.
         shared_basepath : pathlib.Path

openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -191,9 +191,9 @@ def must_be_positive(cls, v):
     """
     Simulation output settings for the solvent transformation.
     """
-
     partial_charge_settings: OpenFFPartialChargeSettings
     """
-    Settings for controlling how to assign partial charges.
-    Currently unused.
+    Settings for controlling how to assign partial charges,
+    including the partial charge assignment method, and the
+    number of conformers used to generate the partial charges.
     """

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -58,7 +58,7 @@
 )
 from ..openmm_utils import system_validation, settings_validation
 from .base import BaseAbsoluteUnit
-from openfe.utils import without_oechem_backend, log_system_probe
+from openfe.utils import log_system_probe
 from openfe.due import due, Doi
 
 
@@ -743,7 +743,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
           A dictionary with the following entries:
             * forcefield_settings : OpenMMSystemGeneratorFFSettings
             * thermo_settings : ThermoSettings
-            * charge_settings: OpenFFPartialChargeSettings
+            * charge_settings : OpenFFPartialChargeSettings
             * solvation_settings : OpenMMSolvationSettings
             * alchemical_settings : AlchemicalSettings
             * lambda_settings : LambdaSettings
@@ -778,9 +778,8 @@ def _execute(
     ) -> dict[str, Any]:
         log_system_probe(logging.INFO, paths=[ctx.scratch])
 
-        with without_oechem_backend():
-            outputs = self.run(scratch_basepath=ctx.scratch,
-                               shared_basepath=ctx.shared)
+        outputs = self.run(scratch_basepath=ctx.scratch,
+                           shared_basepath=ctx.shared)
 
         return {
             'repeat_id': self._inputs['repeat_id'],
@@ -864,9 +863,8 @@ def _execute(
     ) -> dict[str, Any]:
         log_system_probe(logging.INFO, paths=[ctx.scratch])
 
-        with without_oechem_backend():
-            outputs = self.run(scratch_basepath=ctx.scratch,
-                               shared_basepath=ctx.shared)
+        outputs = self.run(scratch_basepath=ctx.scratch,
+                           shared_basepath=ctx.shared)
 
         return {
             'repeat_id': self._inputs['repeat_id'],