Merge pull request #1068 from OpenFreeEnergy/cli-partial-charges

Testing out what partial charges could look over the CLI
OpenFreeEnergy · Jan 30, 2025 · 18e6469 · 18e6469
2 parents 0141744 + 1c2d40d
commit 18e6469
Show file tree

Hide file tree

Showing 23 changed files with 1,872 additions and 818 deletions.
diff --git a/docs/environment.yaml b/docs/environment.yaml
@@ -21,6 +21,7 @@ dependencies:
 - nbsphinx
 - nbsphinx-link
 - myst-parser
+- threadpoolctl
 - pip:
   - sphinx-design
   - sphinx-toolbox

diff --git a/docs/guide/cli/cli_basics.rst b/docs/guide/cli/cli_basics.rst
@@ -35,6 +35,7 @@ To get a list of the subcommands and their descriptions, call ``openfe`` (or
     
     Miscellaneous Commands:
       fetch  Fetch tutorial or other resource.
+      charge-molecules  Generate partial charges for a set of molecules.
       test   Run the OpenFE test suite
 
 The ``--log`` option takes a logging configuration file and sets that

diff --git a/docs/guide/cli/cli_yaml.rst b/docs/guide/cli/cli_yaml.rst
@@ -16,10 +16,14 @@ For example, the settings file which re-specifies the default behaviour would lo
       threed: True
       max3d: 0.95
       element_change: True
+  partial_charge:
+    method: am1bcc
+    settings:
+      off_toolkit_backend: ambertools
 
 The name of the algorithm is given behind the ``method:`` key and the arguments to the
 algorithm are then optionally given behind the ``settings:`` key.
-Both the ``network:`` and ``mapper:`` sections are optional.
+All sections of the file ``network:``, ``mapper:``  and ``partial_charge:`` are optional.
 
 The settings YAML file is then provided to the ``-s`` option of ``openfe plan-rbfe-network``: ::
 
@@ -77,3 +81,28 @@ settings file ::
     method: generate_radial_network
     settings:
       central_ligand: '0'
+
+Customising the partial charge generation
+-----------------------------------------
+
+There are a range of partial charge generation schemes available, including
+
+    - ``am1bcc``
+    - ``am1bccelf10`` (only possible if ``off_toolkit_backend`` in settings is set to ``openeye``)
+    - ``nagl`` (must have ``openff-nagl`` installed)
+    - ``espaloma`` (must have ``espaloma_charge`` installed)
+
+The following settings can also be set
+
+    - ``off_toolkit_backend`` The backend to use for partial charge generation. Choose from  ``ambertools`` (default), ``openeye`` or ``rdkit``.
+    - ``number_of_conformers`` The number of conformers to use for partial charge generation. If unset (default), the input conformer will be used.
+    - ``nagl_model``: The NAGL model to use. If unset (default), the latest available production charge model will be used.
+
+For example, to generate the partial charges using the ``am1bccelf10`` method from ``openeye`` the following should be added to the YAML settings file ::
+
+ partial_charge:
+   method: am1bccelf10
+   settings:
+     off_toolkit_backend: openeye
+
+For more information on the different options, please refer to the :class:`.OpenFFPartialChargeSettings`.
diff --git a/docs/reference/cli/charge_molecules.rst b/docs/reference/cli/charge_molecules.rst
@@ -0,0 +1,7 @@
+.. _cli_charge_molecules:
+
+``charge-molecules`` command
+============================
+
+.. click:: openfecli.commands.generate_partial_charges:charge_molecules
+   :prog: openfe charge-molecules
diff --git a/docs/reference/cli/index.rst b/docs/reference/cli/index.rst
@@ -10,3 +10,4 @@ CLI Reference
     plan_rbfe_network
     quickrun
     gather
+    charge_molecules
diff --git a/environment.yml b/environment.yml
@@ -18,7 +18,7 @@ dependencies:
   - pyyaml
   - coverage
   - cinnabar ~=0.4.0
-  - openff-toolkit>=0.13.0
+  - openff-toolkit>=0.16.2
   - openff-nagl-base >=0.3.3
   - openff-units==0.2.0
   - pint<0.22
@@ -42,6 +42,8 @@ dependencies:
   - autodoc-pydantic<2.0
   - pydata-sphinx-theme
   - sphinx-click
+  # Control blas/openmp threads
+  - threadpoolctl
   - pip:
     - sphinx-toolbox
     - openff-nagl-models>=0.1.2

diff --git a/news/cli_charge_molecules.rst b/news/cli_charge_molecules.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+* Added a new CLI command (``charge-molecules``) to bulk assign partial charges to molecules `PR#1068 <https://github.com/OpenFreeEnergy/openfe/pull/1068>`_
+
+**Changed:**
+
+* The ``plan-rhfe-network`` and ``plan-rbfe-network`` CLI commands will now assign partial charges before planning the network if charges are not present, the charge assignment method can be controlled via the yaml settings file `PR#1068 <https://github.com/OpenFreeEnergy/openfe/pull/1068>`_
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>
diff --git a/openfe/protocols/openmm_utils/charge_generation.py b/openfe/protocols/openmm_utils/charge_generation.py
@@ -8,6 +8,7 @@
 import sys
 import warnings
 import numpy as np
+from gufe import SmallMoleculeComponent
 from openff.units import unit
 from openff.toolkit import Molecule as OFFMol
 from openff.toolkit.utils.base_wrapper import ToolkitWrapper
@@ -17,6 +18,7 @@
     RDKitToolkitWrapper
 )
 from openff.toolkit.utils.toolkit_registry import ToolkitRegistry
+from threadpoolctl import threadpool_limits
 
 try:
     import openeye
@@ -285,7 +287,7 @@ def assign_offmol_partial_charges(
     toolkit_backend: Literal['ambertools', 'openeye', 'rdkit'],
     generate_n_conformers: Optional[int],
     nagl_model: Optional[str],
-) -> None:
+) -> OFFMol:
     """
     Assign partial charges to an OpenFF Molecule based on a selected method.
 
@@ -297,7 +299,7 @@ def assign_offmol_partial_charges(
       Whether or not to overwrite any existing non-zero partial charges.
       Note that zeroed charges will always be overwritten.
     method : Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma']
-      Partial charge assignement method.
+      Partial charge assignment method.
       Supported methods include; am1bcc, am1bccelf10, nagl, and espaloma.
     toolkit_backend : Literal['ambertools', 'openeye', 'rdkit']
       OpenFF toolkit backend employed for charge generation.
@@ -318,17 +320,21 @@ def assign_offmol_partial_charges(
     Raises
     ------
     ValueError
-      If the ``toolkit_backend`` is not suported by the selected ``method``.
+      If the ``toolkit_backend`` is not supported by the selected ``method``.
       If ``generate_n_conformers`` is ``None``, but the input ``offmol``
       has no associated conformers.
       If the number of conformers passed or generated exceeds the number
       of conformers selected by the partial charge ``method``.
+
+    Returns
+    -------
+     The Molecule with partial charges assigned.
     """
 
     # If you have non-zero charges and not overwriting, just return
     if (offmol.partial_charges is not None and np.any(offmol.partial_charges)):
         if not overwrite:
-            return
+            return offmol
 
     # Dictionary for each available charge method
     # The idea of this pattern is to allow for maximum flexibility by
@@ -408,12 +414,105 @@ def assign_offmol_partial_charges(
         generate_n_conformers=generate_n_conformers,
     )  # type: ignore
 
-    # Call selected method to assign partial charges
-    CHARGE_METHODS[method.lower()]['charge_func'](
-        offmol=offmol_copy,
-        toolkit_registry=toolkits,
-        **CHARGE_METHODS[method.lower()]['charge_extra_kwargs'],
-    )  # type: ignore
+    # limit the number of threads used by SQM
+    # <https://github.com/openforcefield/openff-toolkit/issues/1831>
+    with threadpool_limits(limits=1):
+        # Call selected method to assign partial charges
+        CHARGE_METHODS[method.lower()]['charge_func'](
+            offmol=offmol_copy,
+            toolkit_registry=toolkits,
+            **CHARGE_METHODS[method.lower()]['charge_extra_kwargs'],
+        )  # type: ignore
 
     # Copy partial charges back
     offmol.partial_charges = offmol_copy.partial_charges
+    return offmol
+
+
+def bulk_assign_partial_charges(
+    molecules: list[SmallMoleculeComponent],
+    overwrite: bool,
+    method: Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma'],
+    toolkit_backend: Literal['ambertools', 'openeye', 'rdkit'],
+    generate_n_conformers: Optional[int],
+    nagl_model: Optional[str],
+    processors: int = 1,
+) -> list[SmallMoleculeComponent]:
+    """
+    Assign partial charges to a list of SmallMoleculeComponents using multiprocessing.
+
+    Parameters
+    ----------
+    molecules : list[gufe.SmallMoleculeComponent]
+      The list of molecules who should have partial charges assigned.
+    overwrite : bool
+      Whether or not to overwrite any existing non-zero partial charges.
+      Note that zeroed charges will always be overwritten.
+    method : Literal['am1bcc', 'am1bccelf10', 'nagl', 'espaloma']
+      Partial charge assignment method.
+      Supported methods include; am1bcc, am1bccelf10, nagl, and espaloma.
+    toolkit_backend : Literal['ambertools', 'openeye', 'rdkit']
+      OpenFF toolkit backend employed for charge generation.
+      Supported options:
+        * ``ambertools``: selects both the AmberTools and RDKit Toolkit Wrapper
+        * ``openeye``: selects the OpenEye toolkit Wrapper
+        * ``rdkit``: selects the RDKit toolkit Wrapper
+      Note that the ``rdkit`` backend cannot be used for `am1bcc` or
+      ``am1bccelf10`` partial charge methods.
+    generate_n_conformers : Optional[int]
+      Number of conformers to generate for partial charge generation.
+      If ``None`` (default), the input conformer will be used.
+      Values greater than 1 can only be used alongside ``am1bccelf10``.
+    nagl_model : Optional[str]
+      The NAGL model to use for charge assignment if method is ``nagl``.
+      If ``None``, the latest am1bcc NAGL charge model is used.
+    processors: int, default 1
+        The number of processors which should be used to generate the charges.
+
+    Raises
+    ------
+    ValueError
+      If the ``toolkit_backend`` is not supported by the selected ``method``.
+      If ``generate_n_conformers`` is ``None``, but the input ``offmol``
+      has no associated conformers.
+      If the number of conformers passed or generated exceeds the number
+      of conformers selected by the partial charge ``method``.
+
+    Returns
+    -------
+        A list of SmallMoleculeComponents with the charges assigned.
+    """
+    import tqdm
+
+    charge_keywords = {
+        "overwrite": overwrite,
+        "method": method,
+        "toolkit_backend": toolkit_backend,
+        "generate_n_conformers": generate_n_conformers,
+        "nagl_model": nagl_model
+    }
+    charged_ligands = []
+
+    if processors > 1:
+        from concurrent.futures import ProcessPoolExecutor, as_completed
+
+        with ProcessPoolExecutor(max_workers=processors) as pool:
+
+            work_list = [
+                pool.submit(
+                    assign_offmol_partial_charges,
+                    m.to_openff(),
+                    **charge_keywords, # type: ignore
+                )
+                for m in molecules
+            ]
+
+            for work in tqdm.tqdm(as_completed(work_list), desc="Generating charges", ncols=80, total=len(molecules)):
+                charged_ligands.append(SmallMoleculeComponent.from_openff(work.result()))
+
+    else:
+        for m in tqdm.tqdm(molecules, desc="Generating charges", ncols=80, total=len(molecules)):
+            mol_with_charge = assign_offmol_partial_charges(m.to_openff(), **charge_keywords) # type: ignore
+            charged_ligands.append(SmallMoleculeComponent.from_openff(mol_with_charge))
+
+    return charged_ligands
diff --git a/openfe/tests/data/cdk8.zip b/openfe/tests/data/cdk8.zip
diff --git a/openfe/tests/data/openmm_rfe/dummy_charge_ligand_23.sdf b/openfe/tests/data/openmm_rfe/dummy_charge_ligand_23.sdf
@@ -0,0 +1,86 @@
+ligand_23
+     RDKit          3D
+
+ 36 38  0  0  0  0  0  0  0  0999 V2000
+   -1.9600   21.5500  -27.3300 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.2900   22.4100  -28.2000 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.9600   22.9500  -29.2900 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -3.3000   22.6200  -29.5400 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -3.9700   21.7600  -28.6500 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -3.2900   21.2300  -27.5600 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.6400   21.3500  -28.9200 Cl  0  0  0  0  0  0  0  0  0  0  0  0
+   -4.0400   23.1500  -30.7100 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -4.2300   22.4200  -31.6700 O   0  0  0  0  0  0  0  0  0  0  0  0
+   -4.5100   24.4100  -30.6500 N   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.2900   25.1100  -31.5800 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.9600   24.5500  -32.6800 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -6.7800   25.3600  -33.4600 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -6.9200   26.6500  -33.1800 N   0  0  0  0  0  0  0  0  0  0  0  0
+   -6.3000   27.2300  -32.1600 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.5100   26.4600  -31.3100 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -6.5400   28.5800  -31.8800 N   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.7100   29.4500  -31.2300 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -4.6200   29.1200  -30.8100 O   0  0  0  0  0  0  0  0  0  0  0  0
+   -6.2300   30.8500  -31.0300 C   0  0  1  0  0  0  0  0  0  0  0  0
+   -5.5800   31.7500  -29.9800 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.3600   32.0200  -31.4700 C   0  0  1  0  0  0  0  0  0  0  0  0
+   -4.1100   31.7200  -32.0100 F   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.1100   24.0500  -30.3300 Cl  0  0  0  0  0  0  0  0  0  0  0  0
+   -7.3100   30.9400  -31.1600 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.8600   32.9000  -31.8800 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.4400   21.1300  -26.4800 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.2600   22.6600  -28.0200 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -3.8100   20.5600  -26.8900 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -4.3200   24.9800  -29.8400 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.8400   23.5000  -32.9100 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -7.3100   24.9400  -34.3000 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -5.0600   26.9100  -30.4300 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -7.4000   29.0100  -32.1800 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -6.3500   32.2000  -29.3600 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -4.9200   31.1500  -29.3500 H   0  0  0  0  0  0  0  0  0  0  0  0
+ 13 32  1  0
+ 13 14  1  0
+ 12 13  2  0
+ 14 15  2  0
+ 12 31  1  0
+ 11 12  1  0
+ 17 34  1  0
+ 15 17  1  0
+ 15 16  1  0
+ 22 23  1  0
+ 17 18  1  0
+ 22 26  1  6
+  8  9  2  0
+ 11 16  2  0
+ 10 11  1  0
+ 20 22  1  0
+ 21 22  1  0
+ 16 33  1  0
+ 18 20  1  0
+ 18 19  2  0
+ 20 25  1  6
+ 20 21  1  0
+  8 10  1  0
+  4  8  1  0
+ 10 30  1  0
+  3 24  1  0
+ 21 35  1  0
+ 21 36  1  0
+  3  4  2  0
+  4  5  1  0
+  2  3  1  0
+  5  7  1  0
+  5  6  2  0
+  2 28  1  0
+  1  2  2  0
+  1  6  1  0
+  6 29  1  0
+  1 27  1  0
+M  END
+>  <atom.dprop.PartialCharge>  (1) 
+-0.097000000000000003 -0.1295 0.050899999999999994 -0.1426 0.050899999999999994 -0.1295 -0.058400000000000007 0.68969999999999998 -0.53810000000000002 -0.46510000000000001 0.1236
+-0.32829999999999998 0.44219999999999998 -0.72699999999999998 0.54220000000000002 -0.27029999999999998 -0.5464 0.69610000000000005 -0.56910000000000005 -0.23069999999999999
+-0.13339999999999999 0.12759999999999999 -0.20530000000000001 -0.058400000000000007 0.1237 0.1027 0.14799999999999999 0.1565 0.1565 0.32550000000000001 0.14199999999999999
+0.022099999999999995 0.20300000000000001 0.33750000000000002 0.094200000000000006 0.094200000000000006
+
+$$$$