fix plugin vsite tests

openforcefield · Jul 8, 2024 · c90a696 · c90a696
1 parent 25807a7
commit c90a696
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 20 deletions.
diff --git a/smirnoff_plugins/_tests/handlers/test_nonbonded.py b/smirnoff_plugins/_tests/handlers/test_nonbonded.py
@@ -996,12 +996,12 @@ def test_non_lj_on_virtual_site(ideal_water_force_field):
             "epsilon": epsilon * unit.kilocalorie_per_mole,
             "type": "DivalentLonePair",
             "match": "once",
-            "distance": -0.010743 * unit.nanometer,
+            "distance": 0.0 * unit.nanometer,
             "outOfPlaneAngle": 0.0 * unit.degree,
             "inPlaneAngle": "None",
             "charge_increment1": 0.0 * unit.elementary_charge,
-            "charge_increment2": 0.53254 * unit.elementary_charge,
-            "charge_increment3": 0.53254 * unit.elementary_charge,
+            "charge_increment2": 0.0 * unit.elementary_charge,
+            "charge_increment3": 0.0 * unit.elementary_charge,
             "name": "EP",
         }
     )

diff --git a/smirnoff_plugins/utilities/openmm.py b/smirnoff_plugins/utilities/openmm.py
@@ -306,24 +306,28 @@ def evaluate_water_energy_at_distances(
 
     energies = []
     for distance in distances:
-        new_positions = openmm.unit.Quantity(
-            numpy.vstack(
-                [
-                    openmm_positions[:n_positions_per_water, :].value_in_unit(
-                        openmm.unit.angstrom
-                    ),
-                    openmm_positions[n_positions_per_water:, :].value_in_unit(
-                        openmm.unit.angstrom
-                    )
-                    # only translate the second water in x
-                    + numpy.array([distance, 0, 0]),
-                ]
-            ),
-            openmm.unit.angstrom,
-        )
 
-        if n_positions_per_water == 4:
-            new_positions = new_positions[[0, 1, 2, 4, 5, 6, 3, 7]]
+        translated_positons = numpy.vstack(
+            [
+                openmm_positions[:3, :].value_in_unit(
+                    openmm.unit.angstrom
+                ),
+                openmm_positions[:3, :].value_in_unit(
+                    openmm.unit.angstrom
+                )
+                # only translate the second water in x
+                + numpy.array([distance, 0, 0]),
+            ]
+        )
+        if n_positions_per_water > 3:
+            # add zeros to pad the positions
+            vsites = numpy.zeros((2 * (n_positions_per_water - 3), 3))
+            new_positions = openmm.unit.Quantity(
+                numpy.vstack([translated_positons, vsites]),
+                openmm.unit.angstrom
+            )
+        else:
+            new_positions = translated_positons * openmm.unit.angstrom
 
         simulation.context.setPositions(
             new_positions.value_in_unit(openmm.unit.nanometer)