ENH: Better handling of box data

openff/system/components/system.py

@@ -30,15 +30,15 @@ class Config:
     def validate_box(cls, val):
         if val is None:
             return val
+        if type(val) == omm_unit.Quantity:
+            val = simtk_to_pint(val)
         if val.shape == (3, 3):
-            pass
+            return val
         elif val.shape == (3,):
-            val *= np.eye(3)
+            val = val * np.eye(3)
+            return val
         else:
             raise ValueError  # InvalidBoxError
-        if type(val) == omm_unit.Quantity:
-            val = simtk_to_pint(val)
-            return val
 
 
 System.to_parmed = to_parmed

openff/system/interop/parmed.py

@@ -130,9 +130,12 @@ def to_parmed(off_system: Any) -> pmd.Structure:
 
 def _convert_box(box: unit.Quantity, structure: pmd.Structure) -> None:
     # TODO: Convert box vectors to box lengths + angles
-    lengths = box.to(unit("angstrom")).diagonal().magnitude
-    angles = 3 * [90]
-    structure.box = np.hstack([lengths, angles])
+    if box is None:
+        structure.box = [0, 0, 0, 90, 90, 90]
+    else:
+        lengths = box.to(unit("angstrom")).diagonal().magnitude
+        angles = 3 * [90]
+        structure.box = np.hstack([lengths, angles])
 
 
 def _lj_params_from_potential(potential):

openff/system/stubs.py

@@ -2,6 +2,7 @@
 Monkeypatching external classes with custom functionality
 """
 import functools
+from typing import Optional, Union
 
 from openforcefield.topology.topology import Topology
 from openforcefield.typing.engines.smirnoff import ForceField
@@ -12,6 +13,7 @@
     ProperTorsionHandler,
     vdWHandler,
 )
+from simtk import unit as omm_unit
 
 from openff.system.components.potentials import PotentialHandler
 from openff.system.components.smirnoff import (
@@ -22,9 +24,15 @@
     SMIRNOFFvdWHandler,
 )
 from openff.system.components.system import System
+from openff.system.types import UnitArray
 
 
-def to_openff_system(self, topology: Topology, **kwargs) -> System:
+def to_openff_system(
+    self,
+    topology: Topology,
+    box: Optional[Union[omm_unit.Quantity, UnitArray]] = None,
+    **kwargs,
+) -> System:
     """
     A method, patched onto ForceField, that creates a System object
 
@@ -50,9 +58,13 @@ def to_openff_system(self, topology: Topology, **kwargs) -> System:
         )
         sys_out.handlers.update({"Electrostatics": charges})
 
-    sys_out.box = sys_out.validate_box(topology.box_vectors)
+    if box is None:
+        sys_out.box = sys_out.validate_box(topology.box_vectors)
+    else:
+        sys_out.box = sys_out.validate_box(box)
+
     sys_out.topology = topology
-    assert "topology" in dir(sys_out)
+
     return sys_out
 
 

openff/system/tests/test_interop/test_parmed.py

@@ -10,9 +10,9 @@
 
 
 class TestParmedConversion(BaseTest):
-    @pytest.mark.skip
     def test_box(self, argon_ff, argon_top):
-        sys_out = argon_ff.create_openff_system(argon_top)
+        box = UnitArray(np.array([4, 4, 4]), units="nanometer")
+        sys_out = argon_ff.create_openff_system(topology=argon_top, box=box)
         sys_out.positions = UnitArray(
             np.zeros(
                 shape=(argon_top.n_topology_atoms, 3),
@@ -23,7 +23,7 @@ def test_box(self, argon_ff, argon_top):
 
         assert np.allclose(
             struct.box[:3],
-            argon_top.box_vectors.value_in_unit(omm_unit.angstrom).diagonal(),
+            [40, 40, 40],
         )
 
     def test_basic_conversion_argon(self, argon_ff, argon_top):