move ncores cli flag, fix tests, add tests for charge generation

environment.yml

@@ -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

openfe/protocols/openmm_utils/charge_generation.py

@@ -18,7 +18,7 @@
     RDKitToolkitWrapper
 )
 from openff.toolkit.utils.toolkit_registry import ToolkitRegistry
-from openfe.utils.context_managers import temp_env
+from threadpoolctl import threadpool_limits
 
 try:
     import openeye
@@ -416,7 +416,7 @@ def assign_offmol_partial_charges(
 
     # limit the number of threads used by SQM
     # <https://github.com/openforcefield/openff-toolkit/issues/1831>
-    with temp_env({"OMP_NUM_THREADS": "2"}):
+    with threadpool_limits(limits=2):
         # Call selected method to assign partial charges
         CHARGE_METHODS[method.lower()]['charge_func'](
             offmol=offmol_copy,
@@ -495,9 +495,8 @@ def bulk_assign_partial_charges(
 
     if processors > 1:
         from concurrent.futures import ProcessPoolExecutor, as_completed
-        from multiprocessing import get_context
 
-        with ProcessPoolExecutor(max_workers=processors, mp_context=get_context("spawn")) as pool:
+        with ProcessPoolExecutor(max_workers=processors) as pool:
 
             work_list = [
                 pool.submit(

openfecli/commands/generate_partial_charges.py

@@ -1,6 +1,6 @@
 import click
 from openfecli import OFECommandPlugin
-from openfecli.parameters import MOL_DIR, YAML_OPTIONS, OUTPUT_FILE_AND_EXT, WORKERS
+from openfecli.parameters import MOL_DIR, YAML_OPTIONS, OUTPUT_FILE_AND_EXT, NCORES
 
 
 @click.command(
@@ -17,8 +17,8 @@
 @OUTPUT_FILE_AND_EXT.parameter(
     help="The name of the SDF file the charged ligands should be writen to."
 )
-@WORKERS.parameter(
-    help=WORKERS.kwargs["help"],
+@NCORES.parameter(
+    help=NCORES.kwargs["help"],
     default=1,
 )
 @click.option(
@@ -31,7 +31,7 @@ def charge_molecules(
         molecules,
         yaml_settings,
         output,
-        workers,
+        n_cores,
         overwrite_charges
 ):
     """
@@ -69,7 +69,7 @@ def charge_molecules(
         toolkit_backend=partial_charge.off_toolkit_backend,
         generate_n_conformers=partial_charge.number_of_conformers,
         nagl_model=partial_charge.nagl_model,
-        processors=workers
+        processors=n_cores
     )
 
     write("\tDone")

openfecli/commands/plan_rbfe_network.py

@@ -5,7 +5,7 @@
 from openfecli.utils import write, print_duration
 from openfecli import OFECommandPlugin
 from openfecli.parameters import (
-    MOL_DIR, PROTEIN, OUTPUT_DIR, COFACTORS, YAML_OPTIONS, WORKERS
+    MOL_DIR, PROTEIN, OUTPUT_DIR, COFACTORS, YAML_OPTIONS, NCORES
 )
 
 def plan_rbfe_network_main(
@@ -116,16 +116,16 @@ def plan_rbfe_network_main(
     help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.",
     default="alchemicalNetwork",
 )
-@WORKERS.parameter(
-    help=WORKERS.kwargs["help"],
+@NCORES.parameter(
+    help=NCORES.kwargs["help"],
     default=1,
 )
 @print_duration
 def plan_rbfe_network(
         molecules: list[str], protein: str, cofactors: tuple[str],
         yaml_settings: str,
         output_dir: str,
-        workers: int
+        n_cores: int
 ):
     """
     Plan a relative binding free energy network, saved as JSON files for
@@ -211,7 +211,7 @@ def plan_rbfe_network(
         protein=protein,
         cofactors=cofactors,
         partial_charge_settings=partial_charge,
-        processors=workers
+        processors=n_cores
     )
     write("\tDone")
     write("")

openfecli/commands/plan_rhfe_network.py

@@ -8,7 +8,7 @@
 from openfecli.utils import write, print_duration
 from openfecli import OFECommandPlugin
 from openfecli.parameters import (
-    MOL_DIR, MAPPER, OUTPUT_DIR, YAML_OPTIONS, WORKERS
+    MOL_DIR, MAPPER, OUTPUT_DIR, YAML_OPTIONS, NCORES
 )
 
 def plan_rhfe_network_main(
@@ -88,12 +88,12 @@ def plan_rhfe_network_main(
     help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.",
     default="alchemicalNetwork",
 )
-@WORKERS.parameter(
-    help=WORKERS.kwargs["help"],
+@NCORES.parameter(
+    help=NCORES.kwargs["help"],
     default=1,
 )
 @print_duration
-def plan_rhfe_network(molecules: List[str], yaml_settings: str, output_dir: str, workers: int):
+def plan_rhfe_network(molecules: List[str], yaml_settings: str, output_dir: str, n_cores: int):
     """
     Plan a relative hydration free energy network, saved as JSON files for
     the quickrun command.
@@ -169,7 +169,7 @@ def plan_rhfe_network(molecules: List[str], yaml_settings: str, output_dir: str,
         small_molecules=small_molecules,
         solvent=solvent,
         partial_charge_settings=partial_charge,
-        processors=workers
+        processors=n_cores
     )
     write("\tDone")
     write("")

openfecli/parameters/__init__.py

@@ -8,4 +8,4 @@
 from .protein import PROTEIN
 from .molecules import MOL_DIR, COFACTORS
 from .plan_network_options import YAML_OPTIONS
-from .utils import WORKERS
+from .misc import NCORES

openfecli/parameters/misc.py

@@ -0,0 +1,7 @@
+from plugcli.params import Option
+
+NCORES = Option(
+    "-n",
+    "--n-cores",
+    help="The number of cores which should be used for multiprocessing stages."
+)

openfecli/parameters/utils.py

@@ -1,7 +1,7 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/openfe
 
-from plugcli.params import NOT_PARSED, Option
+from plugcli.params import NOT_PARSED
 from openfecli.utils import import_thing
 
 
@@ -27,11 +27,4 @@ def import_parameter(import_str: str):
         result = import_thing(import_str)
     except (ImportError, AttributeError):
         result = NOT_PARSED
-    return result
-
-
-WORKERS = Option(
-    "-w",
-    "--workers",
-    help="The number of workers which should be used for multiprocessing stages."
-)
+    return result

openfecli/tests/commands/test_charge_generation.py

@@ -0,0 +1,133 @@
+import pytest
+from click.testing import CliRunner
+
+from gufe import SmallMoleculeComponent
+from openfecli.commands.generate_partial_charges import charge_molecules
+from openff.toolkit import Molecule
+from openff.units import unit
+import numpy as np
+
+@pytest.fixture
+def yaml_nagl_settings():
+    return """\
+partial_charge:
+  method: nagl
+  settings:
+    nagl_model: openff-gnn-am1bcc-0.1.0-rc.3.pt
+"""
+
+@pytest.fixture
+def methane() -> Molecule:
+    # ensure consistent atom ordering
+    methane = Molecule.from_mapped_smiles("[C:1]([H:2])([H:3])([H:4])[H:5]")
+    methane.generate_conformers(n_conformers=1)
+    return methane
+
+@pytest.fixture
+def methane_with_charges(methane) -> Molecule:
+    methane._partial_charges = [-1.0, 0.25, 0.25, 0.25, 0.25] * unit.elementary_charge
+    return methane
+
+def test_charge_molecules_default(methane, tmpdir):
+
+    runner = CliRunner()
+    mol_path = tmpdir / "methane.sdf"
+    methane.to_file(str(mol_path), "sdf")
+    output_file = str(tmpdir / "charged_methane.sdf")
+
+    with runner.isolated_filesystem():
+        # make sure the charges are picked up
+        with pytest.warns(match="Partial charges have been provided, these will "):
+            result = runner.invoke(
+                charge_molecules,
+                [
+                    "-M",
+                    mol_path,
+                    "-o",
+                    output_file
+                ]
+            )
+
+        assert result.exit_code == 0
+        assert "Partial Charge Generation: am1bcc"
+        assert "assigning ligand partial charges -- this may be slow"
+
+        # make sure the charges have been saved
+        methane = SmallMoleculeComponent.from_sdf_file(filename=output_file)
+        off_methane = methane.to_openff()
+        assert off_methane.partial_charges is not None
+        assert len(off_methane.partial_charges) == 5
+
+@pytest.mark.parametrize("overwrite, expected_charges", [
+    pytest.param(False, [-1.0, 0.25, 0.25, 0.25, 0.25], id="Don't overwrite"),
+    pytest.param(True, [-0.1084, 0.0271, 0.0271, 0.0271, 0.0271], id="Overwrite")
+])
+def test_charge_molecules_overwrite(overwrite, tmpdir, methane_with_charges, expected_charges):
+    runner = CliRunner()
+    mol_path = tmpdir / "methane.sdf"
+    methane_with_charges.to_file(str(mol_path), "sdf")
+    output_file = str(tmpdir / "charged_methane.sdf")
+
+    args = [
+        "-M",
+        mol_path,
+        "-o",
+        output_file
+    ]
+    if overwrite:
+        args.append("--overwrite-charges")
+
+    with runner.isolated_filesystem():
+        # make sure the charges are picked up
+        with pytest.warns(match="Partial charges have been provided, these will "):
+            result = runner.invoke(
+                charge_molecules,
+                args,
+            )
+
+        assert result.exit_code == 0
+        assert "Partial Charge Generation: am1bcc"
+        assert "assigning ligand partial charges -- this may be slow"
+
+        # make sure the charges have not changed from the inputs
+        methane = SmallMoleculeComponent.from_sdf_file(filename=output_file)
+        off_methane = methane.to_openff()
+        assert np.allclose(off_methane.partial_charges.m, expected_charges)
+
+
+
+def test_charge_settings(methane, tmpdir, yaml_nagl_settings):
+    runner = CliRunner()
+    mol_path = tmpdir / "methane.sdf"
+    methane.to_file(str(mol_path), "sdf")
+    output_file = str(tmpdir / "charged_methane.sdf")
+
+    # use nagl charges for CI speed!
+    settings_path = tmpdir / "settings.yaml"
+    with open(settings_path, "w") as f:
+        f.write(yaml_nagl_settings)
+
+    with runner.isolated_filesystem():
+        # make sure the charges are picked up
+        with pytest.warns(match="Partial charges have been provided, these will "):
+            result = runner.invoke(
+                charge_molecules,
+                [
+                    "-M",
+                    mol_path,
+                    "-o",
+                    output_file,
+                    "-s",
+                    settings_path
+                ]
+            )
+
+        assert result.exit_code == 0
+        assert "Partial Charge Generation: nagl"
+        assert "assigning ligand partial charges -- this may be slow"
+
+        # make sure the charges have been saved
+        methane = SmallMoleculeComponent.from_sdf_file(filename=output_file)
+        off_methane = methane.to_openff()
+        assert off_methane.partial_charges is not None
+        assert len(off_methane.partial_charges) == 5