Merge pull request #895 from openforcefield/more-openmm-box-tests

Improve testing of boxes processed from OpenMM
openforcefield · Feb 9, 2024 · 8ff9987 · 8ff9987
2 parents 2da48d0 + e8290c7
commit 8ff9987
Show file tree

Hide file tree

Showing 6 changed files with 181 additions and 23 deletions.
diff --git a/examples/experimental/openmmforcefields/gaff.ipynb b/examples/experimental/openmmforcefields/gaff.ipynb
@@ -76,10 +76,9 @@
     "from openff.interchange.interop.openmm import from_openmm\n",
     "\n",
     "imported = from_openmm(\n",
-    "    molecule.to_topology().to_openmm(),\n",
-    "    system_gaff,\n",
-    "    box_vectors=None,\n",
-    "    positions=ensure_quantity(molecule.conformers[0], \"openmm\"),\n",
+    "    topology=topology.to_openmm(),\n",
+    "    system=system_gaff,\n",
+    "    positions=molecule.conformers[0].to_openmm(),\n",
     ")"
    ]
   },
@@ -88,7 +87,7 @@
    "id": "6",
    "metadata": {},
    "source": [
-    "We can use Interchange's machinery to evaluate each of these objects - the `openmm.System` created by `openmmforcefields` and the `Interchange` object created from it."
+    "This `Interchange` contains an internal representation of approximately everything contained in the `openmm.System`. Topological information (atoms, their elements, bonds, the graph they compose, etc.), forces (including non-bonded settings), and box vectors."
    ]
   },
   {
@@ -97,6 +96,24 @@
    "id": "7",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "imported.visualize()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "We can use Interchange's machinery to evaluate each of these objects - the `openmm.System` created by `openmmforcefields` and the `Interchange` object created from it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "from openff.interchange.drivers.openmm import (\n",
     "    _get_openmm_energies,\n",
@@ -113,7 +130,7 @@
     "        platform=\"Reference\",\n",
     "    ),\n",
     "    system=system_gaff,\n",
-    "    combine_nonbonded_forces=False,\n",
+    "    combine_nonbonded_forces=True,\n",
     "    detailed=False,\n",
     ")\n",
     "\n",
@@ -123,7 +140,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -134,7 +151,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9",
+   "id": "11",
    "metadata": {},
    "source": [
     "Except for some small differences in non-bonded energies due to PME errors, the energies evaluated from each of these objects should match quite closely."
@@ -143,7 +160,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "10",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -173,7 +190,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.11.7"
   },
   "vscode": {
    "interpreter": {

diff --git a/openff/interchange/_tests/data/README.md b/openff/interchange/_tests/data/README.md
@@ -93,3 +93,24 @@ for arg in ["cube", "dodecahedron", "octahedron"]:
 `de-force-1.0.1.offxml`
 
 - [Source](https://github.com/jthorton/de-forcefields/blob/a6f666fc8a3f48d597bfb4db5c46826b9d5d7ed4/deforcefields/offxml/de-force-1.0.1.offxml) with vdW section removed
+
+`system.xml`
+
+Simple OpenMM System XML, written from toolkit objects.
+
+```python
+from openff.toolkit import ForceField, Molecule, unit, Quantity
+import openmm
+
+molecule = Molecule.from_smiles("C")
+
+topology = molecule.to_topology()
+topology.box_vectors = Quantity([2.5, 2.5, 2.5], unit.nanometer)
+
+
+open("system.xml", "w").write(
+    openmm.XmlSerializer.serialize(
+        ForceField("openff-2.1.0.offxml").create_openmm_system(topology)
+    )
+)
+```
diff --git a/openff/interchange/_tests/data/system.xml b/openff/interchange/_tests/data/system.xml
@@ -0,0 +1,63 @@
+<?xml version="1.0" ?>
+<System openmmVersion="8.1.1" type="System" version="1">
+	<PeriodicBoxVectors>
+		<A x="2.5" y="0" z="0"/>
+		<B x="0" y="2.5" z="0"/>
+		<C x="0" y="0" z="2.5"/>
+	</PeriodicBoxVectors>
+	<Particles>
+		<Particle mass="12.01078"/>
+		<Particle mass="1.007947"/>
+		<Particle mass="1.007947"/>
+		<Particle mass="1.007947"/>
+		<Particle mass="1.007947"/>
+	</Particles>
+	<Constraints>
+		<Constraint d=".1090139506109" p1="0" p2="1"/>
+		<Constraint d=".1090139506109" p1="0" p2="2"/>
+		<Constraint d=".1090139506109" p1="0" p2="3"/>
+		<Constraint d=".1090139506109" p1="0" p2="4"/>
+	</Constraints>
+	<Forces>
+		<Force alpha="0" cutoff=".9" dispersionCorrection="1" ewaldTolerance=".0001" exceptionsUsePeriodic="0" forceGroup="0" includeDirectSpace="1" ljAlpha="0" ljnx="0" ljny="0" ljnz="0" method="4" name="Nonbonded force" nx="0" ny="0" nz="0" recipForceGroup="-1" rfDielectric="78.3" switchingDistance=".8" type="NonbondedForce" useSwitchingFunction="1" version="4">
+			<GlobalParameters/>
+			<ParticleOffsets/>
+			<ExceptionOffsets/>
+			<Particles>
+				<Particle eps=".45538911611061844" q="-.10868000239133835" sig=".337953176162662"/>
+				<Particle eps=".06602135607582665" q=".027170000597834587" sig=".26445434132681245"/>
+				<Particle eps=".06602135607582665" q=".027170000597834587" sig=".26445434132681245"/>
+				<Particle eps=".06602135607582665" q=".027170000597834587" sig=".26445434132681245"/>
+				<Particle eps=".06602135607582665" q=".027170000597834587" sig=".26445434132681245"/>
+			</Particles>
+			<Exceptions>
+				<Exception eps="0" p1="0" p2="1" q="0" sig="1"/>
+				<Exception eps="0" p1="0" p2="2" q="0" sig="1"/>
+				<Exception eps="0" p1="1" p2="2" q="0" sig="1"/>
+				<Exception eps="0" p1="0" p2="3" q="0" sig="1"/>
+				<Exception eps="0" p1="1" p2="3" q="0" sig="1"/>
+				<Exception eps="0" p1="2" p2="3" q="0" sig="1"/>
+				<Exception eps="0" p1="0" p2="4" q="0" sig="1"/>
+				<Exception eps="0" p1="1" p2="4" q="0" sig="1"/>
+				<Exception eps="0" p1="2" p2="4" q="0" sig="1"/>
+				<Exception eps="0" p1="3" p2="4" q="0" sig="1"/>
+			</Exceptions>
+		</Force>
+		<Force forceGroup="0" name="PeriodicTorsionForce" type="PeriodicTorsionForce" usesPeriodic="0" version="2">
+			<Torsions/>
+		</Force>
+		<Force forceGroup="0" name="HarmonicAngleForce" type="HarmonicAngleForce" usesPeriodic="0" version="2">
+			<Angles>
+				<Angle a="1.8951470183507508" k="314.14536559165447" p1="1" p2="0" p3="2"/>
+				<Angle a="1.8951470183507508" k="314.14536559165447" p1="1" p2="0" p3="3"/>
+				<Angle a="1.8951470183507508" k="314.14536559165447" p1="1" p2="0" p3="4"/>
+				<Angle a="1.8951470183507508" k="314.14536559165447" p1="2" p2="0" p3="3"/>
+				<Angle a="1.8951470183507508" k="314.14536559165447" p1="2" p2="0" p3="4"/>
+				<Angle a="1.8951470183507508" k="314.14536559165447" p1="3" p2="0" p3="4"/>
+			</Angles>
+		</Force>
+		<Force forceGroup="0" name="HarmonicBondForce" type="HarmonicBondForce" usesPeriodic="0" version="2">
+			<Bonds/>
+		</Force>
+	</Forces>
+</System>
diff --git a/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_import.py b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_import.py
@@ -3,7 +3,7 @@
 
 import numpy
 import pytest
-from openff.toolkit import Molecule, Topology, unit
+from openff.toolkit import Molecule, Quantity, Topology, unit
 from openff.utilities import get_data_file_path, has_package, skip_if_missing
 
 from openff.interchange import Interchange
@@ -29,7 +29,7 @@ def test_simple_roundtrip(self, monkeypatch, sage_unconstrained, ethanol):
         interchange = Interchange.from_smirnoff(
             sage_unconstrained,
             [ethanol],
-            box=[4, 4, 4] * unit.nanometer,
+            box=Quantity([4, 4, 4], unit.nanometer),
         )
 
         system = interchange.to_openmm(combine_nonbonded_forces=True)
@@ -56,6 +56,56 @@ def test_simple_roundtrip(self, monkeypatch, sage_unconstrained, ethanol):
         # OpenMM seems to avoid using the built-in type
         assert converted.box.m.dtype in (float, numpy.float32, numpy.float64)
 
+    @pytest.fixture()
+    def simple_system(self):
+        return openmm.XmlSerializer.deserialize(
+            open(
+                get_data_file_path(
+                    "system.xml",
+                    "openff.interchange._tests.data",
+                ),
+            ).read(),
+        )
+
+    @pytest.mark.parametrize("as_argument", [False, True])
+    def test_different_ways_to_process_box_vectors(
+        self,
+        monkeypatch,
+        as_argument,
+        simple_system,
+    ):
+        monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1")
+
+        if as_argument:
+            box = Interchange.from_openmm(
+                system=simple_system,
+                box_vectors=simple_system.getDefaultPeriodicBoxVectors(),
+            ).box
+        else:
+            box = Interchange.from_openmm(system=simple_system).box
+
+        assert box.shape == (3, 3)
+        assert type(box.m[2][2]) in (float, numpy.float64, numpy.float32)
+        assert type(box.m[1][1]) is not Quantity
+
+    def test_topology_and_system_box_vectors_differ(
+        self,
+        monkeypatch,
+        simple_system,
+    ):
+        """Ensure that, if box vectors specified in the topology and system differ, those in the topology are used."""
+        monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1")
+
+        topology = Molecule.from_smiles("C").to_topology()
+        topology.box_vectors = Quantity([4, 5, 6], unit.nanometer)
+
+        box = Interchange.from_openmm(
+            system=simple_system,
+            topology=topology.to_openmm(),
+        ).box
+
+        assert numpy.diag(box.m_as(unit.nanometer)) == pytest.approx([4, 5, 6])
+
     def test_openmm_roundtrip_metadata(self, monkeypatch):
         monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1")
 

diff --git a/openff/interchange/components/interchange.py b/openff/interchange/components/interchange.py
@@ -156,23 +156,24 @@ class Config:
         arbitrary_types_allowed = True
 
     @validator("box", allow_reuse=True)
-    def validate_box(cls, value):
+    def validate_box(cls, value) -> Optional[Quantity]:
         if value is None:
             return value
-        first_pass = ArrayQuantity.validate_type(value)
-        as_2d = np.atleast_2d(first_pass)
-        if as_2d.shape == (3, 3):
-            box = as_2d
-        elif as_2d.shape == (1, 3):
-            box = as_2d * np.eye(3)
+
+        validated = ArrayQuantity.validate_type(value)
+
+        dimensions = np.atleast_2d(validated).shape
+
+        if dimensions == (3, 3):
+            return validated
+        elif dimensions == (1, 3):
+            return validated * np.eye(3)
         else:
             raise InvalidBoxError(
                 f"Failed to convert value {value} to 3x3 box vectors. Please file an issue if you think this "
                 "input should be supported and the failure is an error.",
             )
 
-        return box
-
     @validator("topology", pre=True)
     def validate_topology(cls, value):
         if value is None:

diff --git a/openff/interchange/interop/openmm/_import/_import.py b/openff/interchange/interop/openmm/_import/_import.py
@@ -77,9 +77,15 @@ def from_openmm(
         interchange.positions = positions
 
     if box_vectors is not None:
-        interchange.box = box_vectors
+        interchange.box = Interchange.validate_box(box_vectors)
+    elif topology is not None:
+        interchange.box = Interchange.validate_box(
+            topology.getPeriodicBoxVectors(),
+        )
     elif system is not None:
-        interchange.box = system.getDefaultPeriodicBoxVectors()
+        interchange.box = Interchange.validate_box(
+            system.getDefaultPeriodicBoxVectors(),
+        )
 
     return interchange