Use openff-forcebalance for optimizations

.github/workflows/CI.yaml

@@ -23,6 +23,7 @@ jobs:
         python-version: [3.9]
         openeye: [true, false]
         integration: [true]
+        optimizer: [forcebalance, openff-forcebalance]
 
     env:
       OE_LICENSE: ${{ github.workspace }}/oe_license.txt
@@ -56,6 +57,7 @@ jobs:
           SECRET_OE_LICENSE: ${{ secrets.OE_LICENSE }}
 
       - name: Install Package
+        id: install_package
         shell: bash -l {0}
         run: |
           python setup.py develop --no-deps
@@ -78,10 +80,13 @@ jobs:
           fail_ci_if_error: false
 
       - name: Run Integration Tests
-        if: ${{ matrix.integration == true }}
+        # Run even if unit tests fail as long as bespokefit was installed successfully and
+        # integration tests are on in the matrix
+        if: ${{ always() && matrix.integration == true && steps.install_package.outcome == 'success' }}
         shell: bash -l {0}
         run: |
           openff-bespoke executor run --smiles                 'CC'               \
                                       --workflow               'default'          \
                                       --default-qc-spec        xtb gfn2xtb none   \
-                                      --target-torsion         '[C:1]-[C:2]'
+                                      --target-torsion         '[C:1]-[C:2]'      \
+                                      --optimizer              ${{ matrix.optimizer }}

devtools/conda-envs/no_openeye.yaml

@@ -22,20 +22,16 @@ dependencies:
   - rdkit >=22
   - ambertools >=22
   - openff-utilities
-  - openff-toolkit-base =0.11
+  - openff-toolkit-base =0.14
   - openff-forcefields
   - openff-interchange
   - openff-qcsubmit
   - openmm >=7.6.0
 
-  # Shim
-  - importlib-metadata >=4
-  - importlib_metadata >=4
-
     # Optional
-  - forcebalance
   - openff-fragmenter-base
   - xtb-python
+  - forcebalance
 
     ### Bespoke dependencies
 
@@ -45,7 +41,6 @@ dependencies:
   - chemper
   - geometric >=1
   - torsiondrive
-  - pymbar
 
     # Executor
   - uvicorn
@@ -65,5 +60,8 @@ dependencies:
   - codecov
   - requests-mock
 
-  - importlib-metadata >=4
-
+    # `openff-forcebalance` and dependencies
+  - pymbar =3 # Implicit constraint, can remove after `openff-forcebalance` package
+  - lxml # Implicit dependency, can remove after `openff-forcebalance` package
+  - pip:
+    - git+https://github.com/openforcefield/openff-forcebalance.git

devtools/conda-envs/test-env.yaml

@@ -29,15 +29,11 @@ dependencies:
   - openff-qcsubmit
   - openmm >=7.6.0
 
-  # Shim
-  - importlib-metadata >=4
-  - importlib_metadata >=4
-
     # Optional
-  - forcebalance
   - openff-fragmenter-base
   - xtb-python
   - openeye-toolkits
+  - forcebalance
 
     ### Bespoke dependencies
 
@@ -47,7 +43,6 @@ dependencies:
   - chemper
   - geometric >=1
   - torsiondrive
-  - pymbar
 
     # Executor
   - uvicorn
@@ -67,5 +62,8 @@ dependencies:
   - codecov
   - requests-mock
 
-  - importlib-metadata >=4
-  - importlib_metadata >=4
+    # `openff-forcebalance` and dependencies
+  - pymbar =3 # Implicit constraint, can remove after `openff-forcebalance` package
+  - lxml # Implicit dependency, can remove after `openff-forcebalance` package
+  - pip:
+    - git+https://github.com/openforcefield/openff-forcebalance.git

openff/bespokefit/cli/executor/run.py

@@ -30,6 +30,7 @@ def _run_cli(
     qc_compute_n_cores: Optional[int],
     qc_compute_max_mem: Optional[float],
     n_optimizer_workers: int,
+    optimizer: Optional[str],
     launch_redis_if_unavailable: bool,
 ):
     """Run bespoke optimization using a temporary executor.
@@ -87,6 +88,7 @@ def _run_cli(
                 workflow_file_name=workflow_file_name,
                 allow_multiple_molecules=False,
                 save_submission=False,
+                optimizer=optimizer,
             )
 
             console.print(Padding("3. running the fitting pipeline", (1, 0, 1, 0)))

openff/bespokefit/cli/executor/submit.py

@@ -18,6 +18,7 @@
     print_header,
 )
 from openff.bespokefit.executor.utilities import handle_common_errors
+from openff.bespokefit.optimizers.base import _optimizers
 
 if TYPE_CHECKING:
     from openff.toolkit.topology import Molecule
@@ -88,6 +89,12 @@ def submit_options(allow_multiple_molecules: bool = False):
             "then 'none' should be specified.",
             required=False,
         ),
+        optgroup.option(
+            "--optimizer",
+            type=str,
+            help=f"The name of the optimizer to use, one of {[*_optimizers]}",
+            required=False,
+        ),
     ]
 
 
@@ -99,6 +106,7 @@ def _to_input_schema(
     default_qc_spec: Optional[Tuple[str, str, str]],
     workflow_name: Optional[str],
     workflow_file_name: Optional[str],
+    optimizer: Optional[str],
 ) -> "BespokeOptimizationSchema":
     from openff.qcsubmit.common_structures import QCSpec, QCSpecificationError
 
@@ -145,6 +153,8 @@ def _to_input_schema(
                     spec_description="CLI provided spec",
                 )
             ]
+        if optimizer is not None:
+            workflow_factory.optimizer.type = _optimizers[optimizer].name()
 
     except FileNotFoundError:
         exit_with_messages(
@@ -195,6 +205,7 @@ def _submit(
     default_qc_spec: Optional[Tuple[str, str, str]],
     workflow_name: Optional[str],
     workflow_file_name: Optional[str],
+    optimizer: Optional[str],
     allow_multiple_molecules: bool,
     save_submission: bool,
 ) -> List[str]:
@@ -259,6 +270,7 @@ def _submit(
                 default_qc_spec,
                 workflow_name,
                 workflow_file_name,
+                optimizer,
             )
         )
 

openff/bespokefit/optimizers/__init__.py

@@ -7,7 +7,10 @@
     list_optimizers,
     register_optimizer,
 )
-from openff.bespokefit.optimizers.forcebalance import ForceBalanceOptimizer
+from openff.bespokefit.optimizers.forcebalance import (
+    ForceBalanceOptimizer,
+    OpenFFForceBalanceOptimizer,
+)
 
 __all__ = [
     "BaseOptimizer",
@@ -16,4 +19,5 @@
     "list_optimizers",
     "register_optimizer",
     "ForceBalanceOptimizer",
+    "OpenFFForceBalanceOptimizer",
 ]

openff/bespokefit/optimizers/base.py

@@ -4,8 +4,11 @@
 from typing import Dict, List, Type, Union
 
 from openff.bespokefit.exceptions import OptimizerError
-from openff.bespokefit.optimizers.forcebalance import ForceBalanceOptimizer
-from openff.bespokefit.optimizers.model import BaseOptimizer
+from openff.bespokefit.optimizers.forcebalance import (
+    ForceBalanceOptimizer,
+    OpenFFForceBalanceOptimizer,
+)
+from openff.bespokefit.optimizers.model import BaseOptimizer, tokenize_name
 
 _optimizers: Dict[str, Type[BaseOptimizer]] = {}
 
@@ -31,10 +34,10 @@ def register_optimizer(optimizer: Type[BaseOptimizer], replace: bool = False) ->
     if not issubclass(optimizer, BaseOptimizer):
         raise OptimizerError(
             f"The optimizer {optimizer} could not be registered it must be a subclass "
-            f"of openff.bespokefit.optimzers.BaseOptimizer"
+            f"of openff.bespokefit.optimizers.BaseOptimizer"
         )
 
-    optimizer_name = optimizer.name().lower()
+    optimizer_name = tokenize_name(optimizer.name())
 
     if optimizer_name in _optimizers and not replace:
         raise OptimizerError(
@@ -58,9 +61,9 @@ def deregister_optimizer(
     """
 
     if isinstance(optimizer, str):
-        optimizer_name = optimizer.lower()
+        optimizer_name = tokenize_name(optimizer)
     else:
-        optimizer_name = optimizer.name().lower()
+        optimizer_name = tokenize_name(optimizer.name())
 
     if _optimizers.pop(optimizer_name, None) is None:
         raise OptimizerError(
@@ -82,7 +85,7 @@ def get_optimizer(optimizer_name: str) -> Type[BaseOptimizer]:
         The requested optimizer matching the given optimizer name.
     """
 
-    optimizer = _optimizers.get(optimizer_name.lower(), None)
+    optimizer = _optimizers.get(tokenize_name(optimizer_name), None)
 
     if optimizer is None:
         raise OptimizerError(
@@ -106,3 +109,4 @@ def list_optimizers() -> List[str]:
 
 # register the built in optimizers
 register_optimizer(ForceBalanceOptimizer)
+register_optimizer(OpenFFForceBalanceOptimizer)

openff/bespokefit/optimizers/forcebalance/__init__.py

@@ -1,4 +1,11 @@
 from openff.bespokefit.optimizers.forcebalance.factories import ForceBalanceInputFactory
-from openff.bespokefit.optimizers.forcebalance.forcebalance import ForceBalanceOptimizer
+from openff.bespokefit.optimizers.forcebalance.forcebalance import (
+    ForceBalanceOptimizer,
+    OpenFFForceBalanceOptimizer,
+)
 
-__all__ = ["ForceBalanceInputFactory", "ForceBalanceOptimizer"]
+__all__ = [
+    "ForceBalanceInputFactory",
+    "ForceBalanceOptimizer",
+    "OpenFFForceBalanceOptimizer",
+]

openff/bespokefit/optimizers/forcebalance/factories.py

@@ -262,6 +262,7 @@ def _generate_target(
             Tuple[Union[TorsionDriveRecord, TorsionDriveResult], Molecule]
         ],
     ):
+        # TODO: This might need to depend on the optimizer if these Molecule classes diverge
         from forcebalance.molecule import Molecule as FBMolecule
 
         if isinstance(target, AbInitioTargetSchema) and target.fit_force is True:
@@ -512,6 +513,7 @@ def _generate_target(
         target: VibrationTargetSchema,
         qc_records: List[Tuple[Union[ResultRecord, "AtomicResult"], Molecule]],
     ):
+        # TODO: This might need to depend on the optimizer if these Molecule classes diverge
         from forcebalance.molecule import Molecule as FBMolecule
 
         assert len(qc_records) == 1
@@ -572,6 +574,7 @@ def _generate_target(
             Tuple[Union[OptimizationRecord, OptimizationResult], Molecule]
         ],
     ):
+        # TODO: This might need to depend on the optimizer if these Molecule classes diverge
         from forcebalance.molecule import Molecule as FBMolecule
 
         record_names = []