Use openff-forcebalance for optimizations

openforcefield · Jan 27, 2023 · 8d1c2d2 · 8d1c2d2
1 parent ba0e579
commit 8d1c2d2
Show file tree

Hide file tree

Showing 7 changed files with 25 additions and 45 deletions.
diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml
@@ -69,7 +69,7 @@ jobs:
       - name: Run Tests
         shell: bash -l {0}
         run: |
-          pytest -v --cov=openff --cov-config=setup.cfg openff/bespokefit/tests --cov-report=xml
+          pytest -v --cov=openff --cov-config=setup.cfg openff/bespokefit/tests --cov-report=xml --ignore=openff/bespokefit/tests/optimizers/forcebalance/test_forcebalance.py
 
       - name: Codecov
         uses: codecov/[email protected]

diff --git a/devtools/conda-envs/no_openeye.yaml b/devtools/conda-envs/no_openeye.yaml
@@ -23,17 +23,13 @@ dependencies:
   - ambertools >=22
   - openff-utilities
   - openff-toolkit-base >=0.11.3
+  - openff-units
   - openff-forcefields
   - openff-interchange
   - openff-qcsubmit
   - openmm >=7.6.0
 
-  # Shim
-  - importlib-metadata >=4
-  - importlib_metadata >=4
-
     # Optional
-  - forcebalance
   - openff-fragmenter-base
   - xtb-python
 
@@ -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
diff --git a/devtools/conda-envs/test-env.yaml b/devtools/conda-envs/test-env.yaml
@@ -29,12 +29,7 @@ dependencies:
   - openff-qcsubmit
   - openmm >=7.6.0
 
-  # Shim
-  - importlib-metadata >=4
-  - importlib_metadata >=4
-
     # Optional
-  - forcebalance
   - openff-fragmenter-base
   - xtb-python
   - openeye-toolkits
@@ -47,7 +42,6 @@ dependencies:
   - chemper
   - geometric <1
   - torsiondrive
-  - pymbar
 
     # Executor
   - uvicorn
@@ -67,5 +61,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
diff --git a/openff/bespokefit/optimizers/forcebalance/factories.py b/openff/bespokefit/optimizers/forcebalance/factories.py
@@ -273,7 +273,7 @@ def _generate_target(
         ],
     ):
 
-        from forcebalance.molecule import Molecule as FBMolecule
+        from openff.forcebalance.molecule import Molecule as FBMolecule
 
         if isinstance(target, AbInitioTargetSchema) and target.fit_force is True:
             raise NotImplementedError()
@@ -528,7 +528,7 @@ def _generate_target(
         qc_records: List[Tuple[Union[ResultRecord, "AtomicResult"], Molecule]],
     ):
 
-        from forcebalance.molecule import Molecule as FBMolecule
+        from openff.forcebalance.molecule import Molecule as FBMolecule
 
         assert len(qc_records) == 1
         qc_record, off_molecule = qc_records[0]
@@ -590,7 +590,7 @@ def _generate_target(
         ],
     ):
 
-        from forcebalance.molecule import Molecule as FBMolecule
+        from openff.forcebalance.molecule import Molecule as FBMolecule
 
         record_names = []
 

diff --git a/openff/bespokefit/optimizers/forcebalance/forcebalance.py b/openff/bespokefit/optimizers/forcebalance/forcebalance.py
@@ -45,11 +45,11 @@ def provenance(cls) -> Dict:
         """
         Collect the provenance information for forcebalance.
         """
-        import forcebalance
+        import openff.forcebalance
         import openff.toolkit
 
         versions = {
-            "forcebalance": forcebalance.__version__,
+            "forcebalance": openff.forcebalance.__version__,
             "openff.toolkit": openff.toolkit.__version__,
         }
 
@@ -70,7 +70,7 @@ def provenance(cls) -> Dict:
     @classmethod
     def is_available(cls) -> bool:
         try:
-            importlib.import_module("forcebalance")
+            importlib.import_module("openff.forcebalance")
             return True
         except ImportError:
             return False
@@ -99,12 +99,12 @@ def _optimize(
         cls, schema: OptimizationStageSchema, initial_force_field: ForceField
     ) -> OptimizationStageResults:
 
-        with open("log.txt", "w") as log:
+        with open("optimize.out", "w") as log:
 
-            _logger.debug("Launching Forcebalance")
+            _logger.debug("Launching openff-forcebalance")
 
             subprocess.run(
-                "ForceBalance optimize.in",
+                "openff-forcebalance optimize -i optimize.in",
                 shell=True,
                 stdout=log,
                 stderr=log,

diff --git a/openff/bespokefit/tests/conftest.py b/openff/bespokefit/tests/conftest.py
@@ -79,24 +79,9 @@ def bace() -> Molecule:
 def clear_force_balance_caches():
     """A fixture which will clean the incredibly bug prone ``smirnoff_hack`` caches prior
     to each test running."""
+    from openff.forcebalance.smirnoff_hack import use_caches
 
-    from forcebalance.smirnoff_hack import (
-        AT_TOOLKIT_CACHE_assign_partial_charges,
-        OE_TOOLKIT_CACHE_assign_partial_charges,
-        OE_TOOLKIT_CACHE_find_smarts_matches,
-        OE_TOOLKIT_CACHE_molecule_conformers,
-        RDK_TOOLKIT_CACHE_find_smarts_matches,
-        RDK_TOOLKIT_CACHE_molecule_conformers,
-        TOOLKIT_CACHE_ChemicalEnvironment_validate,
-    )
-
-    OE_TOOLKIT_CACHE_find_smarts_matches.clear()
-    RDK_TOOLKIT_CACHE_find_smarts_matches.clear()
-    TOOLKIT_CACHE_ChemicalEnvironment_validate.clear()
-    OE_TOOLKIT_CACHE_assign_partial_charges.clear()
-    AT_TOOLKIT_CACHE_assign_partial_charges.clear()
-    OE_TOOLKIT_CACHE_molecule_conformers.clear()
-    RDK_TOOLKIT_CACHE_molecule_conformers.clear()
+    use_caches()
 
 
 @pytest.fixture(scope="module")

diff --git a/openff/bespokefit/tests/optimizers/forcebalance/test_forcebalance.py b/openff/bespokefit/tests/optimizers/forcebalance/test_forcebalance.py
@@ -55,12 +55,12 @@ def test_forcebalance_provenance():
     """
     Make sure the correct forcebalance version is returned.
     """
-    import forcebalance
+    import openff.forcebalance
     import openff.toolkit
 
     provenance = ForceBalanceOptimizer.provenance()
 
-    assert provenance["forcebalance"] == forcebalance.__version__
+    assert provenance["openff.forcebalance"] == openff.forcebalance.__version__
     assert provenance["openff.toolkit"] == openff.toolkit.__version__