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"
@@ -64,21 +66,29 @@ jobs:
     - name: Checkout
       uses: actions/checkout@v4
 
-    - name: Set up conda environment with gromacs
-      if: matrix.gmx-version != '5' && matrix.gmx-version != '4' && matrix.gmx-version != '2022' && matrix.gmx-version != '2023'
+    - name: Set up conda environment with gromacs for Python 3.9 or 3.10
+      if: (matrix.gmx-version != '5' && matrix.gmx-version != '4' && matrix.gmx-version != '2022' && matrix.gmx-version != '2023') && (matrix.python-version == '3.9' || matrix.python-version == '3.10')
       uses: mamba-org/setup-micromamba@v1
       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' }}
+        # Use a reduced legacy environment for Python 3.9/3.10 builds.
+        environment-file: devtools/conda-envs/test_env_py39.yaml
+        create-args: >- # beware the >- instead of |, we don't split on newlines but on spaces
+          python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }}.*=nompi_dblprec*
+
+    - name: Set up conda environment with gromacs for Python 3.11 or higher
+      uses: mamba-org/setup-micromamba@v1
+      if: (matrix.python-version != '3.9' && matrix.python-version != '3.10') && (matrix.gmx-version != '5' && matrix.gmx-version != '4' && matrix.gmx-version != '2022' && matrix.gmx-version != '2023')
+      with:
+        environment-file: devtools/conda-envs/test_env.yaml
         create-args: >- # beware the >- instead of |, we don't split on newlines but on spaces
           python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }}.*=nompi_dblprec*
 
     - name: Set up conda environment without gromacs
       if: matrix.gmx-version == '5' || matrix.gmx-version == '4'
       uses: mamba-org/setup-micromamba@v1
       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' }}
+        # Use the legacy env for Python 3.9/3.10, full env otherwise.
+        environment-file: ${{ (matrix.python-version == '3.9' || matrix.python-version == '3.10') && 'devtools/conda-envs/test_env_py39.yaml' || 'devtools/conda-envs/test_env.yaml' }}
         create-args: >- # beware the >- instead of |, we don't split on newlines but on spaces
           python=${{ matrix.python-version }}
 
@@ -88,11 +98,12 @@ jobs:
       with:
         # dblprec conda-forge builds don't exist for 2022/2023; single-precision is sufficient
         # for the version-rejection test.
-        environment-file: ${{ matrix.python-version == '3.9' && 'devtools/conda-envs/test_env_py39.yaml' || 'devtools/conda-envs/test_env.yaml' }}
+        # Use the legacy env for Python 3.9/3.10, full env otherwise.
+        environment-file: ${{ (matrix.python-version == '3.9' || matrix.python-version == '3.10') && 'devtools/conda-envs/test_env_py39.yaml' || 'devtools/conda-envs/test_env.yaml' }}
         create-args: >-
           python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }}
 
-    - name: Pin setuptools for Python < 3.10
+    - name: Pin setuptools for Python <= 3.10
       if: matrix.python-version == '3.9' || matrix.python-version == '3.10'
       shell: bash -l {0}
       run: pip install "setuptools<82"

devtools/conda-envs/test_env.yaml

@@ -23,7 +23,10 @@ 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.5.3
   # - openff-recharge
   # - openeye-toolkits (Don't have a license file to use with GH Actions.)
+  - pint <0.25

devtools/conda-envs/test_env_py39.yaml

@@ -24,7 +24,7 @@ dependencies:
   - ndcctools
   - geometric
   # - gromacs =2019.1
-  # openff packages require Python >= 3.10; tests are skipped on 3.9
+  # OpenFF stack is covered by test_env.yaml in newer-Python jobs.
   # - openff-toolkit-base
   # - openff-evaluator-base
   # - openff-recharge

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):