Fix vsite fits

.github/workflows/ci.yml

@@ -28,6 +28,8 @@ jobs:
           - "3.10"
           - "3.11"
           - "3.12"
+          # pymbar 3 does not support Python 3.13
+          # - "3.13"
         gmx-version:
           - "5"
           - "2019"
@@ -67,6 +69,16 @@ jobs:
     - name: Set up conda environment with gromacs
       if: matrix.gmx-version != '5' && matrix.gmx-version != '4' && matrix.gmx-version != '2022' && matrix.gmx-version != '2023'
       uses: mamba-org/setup-micromamba@v1
+      if: matrix.python-version == '3.9' || matrix.python-version == '3.10'
+      with:
+        # OpenFF Evaluator requires Python >= 3.10, so we use a separate environment file for Python 3.9 and 3.10 builds.
+        environment-file: devtools/conda-envs/test_env_py310.yaml
+        create-args: >- # beware the >- instead of |, we don't split on newlines but on spaces
+          python=${{ matrix.python-version }}
+
+    - name: Set up conda environment for Python 3.11 or higher
+      uses: mamba-org/setup-micromamba@v1
+      if: matrix.python-version != '3.9' && matrix.python-version != '3.10'
       with:
         # ambertools has no Python 3.9 builds; use a separate env file for 3.9
         environment-file: ${{ matrix.python-version == '3.9' && 'devtools/conda-envs/test_env_py39.yaml' || 'devtools/conda-envs/test_env.yaml' }}

devtools/conda-envs/test_env.yaml

@@ -23,7 +23,29 @@ dependencies:
   - ambertools
   - ndcctools
   - geometric
-  - openff-toolkit >=0.11.3
-  - openff-evaluator-base >= 0.4.5
+  # - gromacs =2019.1
+  - openff-toolkit >=0.18.0
+  - openff-interchange >=0.5.1
+  # - openff-evaluator-base >= 0.4.5
   # - openff-recharge
   # - openeye-toolkits (Don't have a license file to use with GH Actions.)
+  - pint <0.25
+
+  # evaluator dependencies
+  - dask ==2024.2.1
+  - distributed >=2.7.0
+  - dask-jobqueue >=0.8.0
+  - dask-kubernetes
+  - uncertainties
+  - openmmtools
+  - pyyaml
+  - requests
+  - python-dateutil
+  - pydantic =2,<2.12
+  - taproom
+  - dataclasses
+  - pandas =2
+
+  - pip:
+    # install branch of evaluator
+    - git+https://github.com/openforcefield/openff-evaluator.git@fix-vsite-fits

devtools/conda-envs/test_env_py310.yaml

@@ -0,0 +1,30 @@
+name: forcebalance-test
+channels:
+  - conda-forge
+  - openeye
+dependencies:
+    # Base depends
+  - python
+  - pip
+    # Testing
+  - pytest
+  - pytest-cov
+  - codecov
+  - numpy
+  - scipy
+  - lxml
+  - networkx
+  - zlib
+  - swig
+  - future
+  - pymbar =3
+  - openmm >= 8
+  # ambertools has no Python 3.9 builds on conda-forge
+  - ndcctools
+  - geometric
+  # - gromacs =2019.1
+  # openff packages require Python >= 3.11; tests are skipped on 3.9 and 3.10
+  # - openff-toolkit-base
+  # - openff-evaluator-base
+  # - openff-recharge
+  # - openeye-toolkits (Don't have a license file to use with GH Actions.)

src/evaluator_io.py

@@ -16,6 +16,7 @@
 import numpy as np
 from forcebalance.nifty import warn_once, printcool, printcool_dictionary
 from forcebalance.output import getLogger
+from forcebalance.smirnoffio import select_virtual_site_parameter
 from forcebalance.target import Target
 
 try:
@@ -308,20 +309,40 @@ def _parameter_value_from_gradient_key(self, gradient_key):
         bool
             Returns True if the parameter is a cosmetic one.
         """
-        # try:
-        #     import openmm.unit as simtk_unit
-        # except ImportError:
-        #     import simtk.unit as simtk_unit
         from openff.units import unit as openff_unit
 
-
         parameter_handler = self.FF.openff_forcefield.get_parameter_handler(
             gradient_key.tag
         )
-        parameter = (
-            parameter_handler if gradient_key.smirks is None
-            else parameter_handler.parameters[gradient_key.smirks]
-        )
+
+        if gradient_key.smirks is None:
+            parameter = parameter_handler
+        elif gradient_key.tag != "VirtualSites":
+            parameter = parameter_handler.parameters[gradient_key.smirks]
+        else:
+            # VirtualSite parameters are not uniquely identifiable by SMIRKS alone.
+            # Require explicit type/name/match metadata in every VirtualSites key.
+            if gradient_key.virtual_site_type is None:
+                raise KeyError(
+                    f"Gradient key {gradient_key} is missing required virtual_site_type"
+                )
+            if gradient_key.virtual_site_name is None:
+                raise KeyError(
+                    f"Gradient key {gradient_key} is missing required virtual_site_name"
+                )
+            if gradient_key.virtual_site_match is None:
+                raise KeyError(
+                    f"Gradient key {gradient_key} is missing required virtual_site_match"
+                )
+
+            parameter = select_virtual_site_parameter(
+                parameters=parameter_handler.parameters,
+                smirks=gradient_key.smirks,
+                virtual_site_type=gradient_key.virtual_site_type,
+                virtual_site_name=gradient_key.virtual_site_name,
+                virtual_site_match=gradient_key.virtual_site_match,
+                error_context=f"gradient key {gradient_key}",
+            )
 
         attribute_split = re.split(r"(\d+)", gradient_key.attribute)
         attribute_split = list(filter(None, attribute_split))
@@ -474,14 +495,27 @@ def submit_jobs(self, mvals, AGrad=True, AHess=True):
                 string_key = field_list[0]
                 key_split = string_key.split("/")
 
+                virtual_site_kwargs = {}
+
                 if len(key_split) == 3 and key_split[0] == "":
                     parameter_tag = key_split[1].strip()
                     parameter_smirks = None
                     parameter_attribute = key_split[2].strip()
-                elif len(key_split) == 4:
+                elif len(key_split) >= 4:
                     parameter_tag = key_split[0].strip()
                     parameter_smirks = key_split[3].strip()
                     parameter_attribute = key_split[2].strip()
+
+                    if parameter_tag == "VirtualSites":
+                        # VirtualSites keys must include positional identity metadata:
+                        # VirtualSites/<tag>/<attribute>/<smirks>/<type>/<name>/<match>
+                        if len(key_split) != 7:
+                            raise KeyError(
+                                f"VirtualSites parameter key must include type/name/match: {string_key}"
+                            )
+                        virtual_site_kwargs["virtual_site_type"] = key_split[4].strip()
+                        virtual_site_kwargs["virtual_site_name"] = key_split[5].strip()
+                        virtual_site_kwargs["virtual_site_match"] = key_split[6].strip()
                 else:
                     raise NotImplementedError()
 
@@ -490,6 +524,7 @@ def submit_jobs(self, mvals, AGrad=True, AHess=True):
                     tag=parameter_tag,
                     smirks=parameter_smirks,
                     attribute=parameter_attribute,
+                    **virtual_site_kwargs,
                 )
 
                 # Find the unit of the gradient parameter.
@@ -525,9 +560,7 @@ def submit_jobs(self, mvals, AGrad=True, AHess=True):
         )
 
         if (
-            self._pending_estimate_request.results(
-                True, polling_interval=self._options.polling_interval
-            )[0] is None
+            self._pending_estimate_request.results(False)[0] is None
         ):
 
             raise RuntimeError(
@@ -738,6 +771,19 @@ def get(self, mvals, AGrad=True, AHess=True):
         AGrad = bool(AGrad)
         AHess = bool(AHess)
 
+        # Block until the Evaluator computation is finished.  This is intentionally
+        # placed here (not in submit_jobs) so that Work Queue tasks submitted by other
+        # targets can run concurrently while we wait.
+        estimation_results, _ = self._pending_estimate_request.results(
+            True, polling_interval=self._options.polling_interval
+        )
+        if estimation_results is None:
+            raise RuntimeError(
+                "No `EvaluatorServer` could be found to retrieve results from. "
+                "Please double check that a server is running, and that the connection "
+                "settings specified in the input script are correct."
+            )
+
         # Extract the properties estimated using the unperturbed parameters.
         estimated_data_set, estimated_gradients = self._extract_property_data(
             self._pending_estimate_request, mvals, AGrad

src/forcefield.py

@@ -747,7 +747,7 @@ def addff_xml(self, ffname):
                 res = re.search(r'^[-+]?[0-9]*\.?[0-9]*([eEdD][-+]?[0-9]+)?', quantity_str)
                 value_str, unit_str = quantity_str[:res.end()], quantity_str[res.end():]
                 # LPW 2023-01-23: Behavior of parameter unit string for "evaluated" parameter is undefined.
-                unit_str = ""
+                # unit_str = ""
                 quantity_str = e.get(parameter_name)
                 self.offxml_unit_strs[dest] = unit_str
 

src/openmmio.py

@@ -270,6 +270,11 @@ def GetVirtualSiteParameters(system):
                 vsprm.append(_openmm.OutOfPlaneSite_getWeight12(vs))
                 vsprm.append(_openmm.OutOfPlaneSite_getWeight13(vs))
                 vsprm.append(_openmm.OutOfPlaneSite_getWeightCross(vs))
+            elif vstype == 'LocalCoordinatesSite':
+                vsprm.extend(_openmm.LocalCoordinatesSite_getOriginWeights(vs))
+                vsprm.extend(_openmm.LocalCoordinatesSite_getXWeights(vs))
+                vsprm.extend(_openmm.LocalCoordinatesSite_getYWeights(vs))
+                vsprm.extend(_openmm.LocalCoordinatesSite_getLocalPosition(vs))
     return np.array(vsprm)
 
 def GetDrudeParameters(system):