Fix vsite fits

.github/workflows/ci.yml

@@ -22,9 +22,12 @@ jobs:
           - macOS-latest
           - ubuntu-latest
         python-version:
+          - "3.9"
+          - "3.10"
           - "3.11"
           - "3.12"
-          - "3.13"
+          # pymbar 3 does not support Python 3.13
+          # - "3.13"
 
     env:
       CI_OS: ${{ matrix.os }}
@@ -35,12 +38,27 @@ jobs:
     - name: Checkout
       uses: actions/checkout@v4
 
-    - name: Set up conda environment
+    - name: Set up conda environment for Python 3.9 or 3.10
       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:
         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 }}
+
+    - 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"
 
     - name: Additional info about the build
       shell: bash
@@ -119,7 +137,7 @@ jobs:
 
     - name: Run tests
       run: |
-        pytest -v --cov=forcebalance --cov-config=setup.cfg --durations=0 --cov-report=xml src/tests/
+        pytest -v --cov=forcebalance --cov-config=setup.cfg --durations=0 --cov-report=xml -k "not TestEvaluatorBromineStudy" src/tests/
 
     - name: Run water study
       run: |

.gitignore

@@ -123,3 +123,12 @@ studies/*/targets/
 studies/*/*.tmp/
 studies/*/*.bak/
 sutdies/*/*.err
+
+# test files from running CI locally
+src/tests/files/XmlScript_out/TIP3G2w.xml
+src/tests/files/targets/dms-liquid/dms.xml
+src/tests/files/temp.bak/
+src/tests/files/test_liquid/targets/Liquid/TIP3G2w.xml
+src/tests/files/test_liquid/targets/Liquid/dms.xml
+src/tests/test.job
+studies/003d_evaluator_liquid_bromine/smirnoff_parameter_assignments.json

devtools/conda-envs/test_env.yaml

@@ -23,7 +23,28 @@ dependencies:
   - ndcctools
   - geometric
   # - gromacs =2019.1
-  - openff-toolkit >=0.11.3
-  - openff-evaluator >= 0.4.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.

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