Merge pull request #672 from daico007/box_length_bug

Rework box class including NumPy subclass

Author: mattwthompson

mbuild/box.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 
+__all__ = ['Box']
 
 class Box(object):
     """A box representing the bounds of the system.
@@ -24,23 +25,34 @@ def __init__(self, lengths=None, mins=None, maxs=None, angles=None):
                     "You provided: "
                     "lengths={} mins={} maxs={}".format(lengths, mins, maxs)
                 )
-            self._mins = np.array(mins, dtype=np.float)
-            self._maxs = np.array(maxs, dtype=np.float)
-            self._lengths = self.maxs - self.mins
+            mins = np.array(mins, dtype=np.float)
+            maxs = np.array(maxs, dtype=np.float)
+            assert mins.shape == (3, ), "Given mins have wrong dimensions"
+            assert maxs.shape == (3, ), "Given maxs have wrong dimensions"
+            assert all(mins <= maxs), "Given mins are greater than maxs"
+            self._mins = _BoxArray(array=mins, var="mins", box=self)
+            self._maxs = _BoxArray(array=maxs, var="maxs", box=self)
+            self._lengths = _BoxArray(array=(self.maxs - self.mins), var="lengths", box=self)
         else:
             if mins is not None or maxs is not None:
                 warn(
                     "Provided `lengths` and `mins` and/or `maxs`. Only `lengths` "
                     "is being used. You provided: "
                     "lengths={} mins={} maxs={}".format(lengths, mins, maxs)
                 )
-            self._mins = np.array([0.0, 0.0, 0.0])
-            self._maxs = np.array(lengths, dtype=np.float)
-            self._lengths = np.array(lengths, dtype=np.float)
+            if isinstance(lengths, int) or isinstance(lengths, float):
+                lengths = np.array(lengths*np.ones(3), dtype=np.float)
+            else:
+                lengths = np.array(lengths, dtype=np.float)
+            assert lengths.shape == (3, )
+            assert all(lengths >= 0), "Given lengths are negative"
+            self._mins = _BoxArray(array=(0,0,0), var="mins", box=self)
+            self._maxs = _BoxArray(array=lengths, var="maxs", box=self)
+            self._lengths = _BoxArray(array=lengths, var="lengths", box=self)
         if angles is None:
-            angles = np.array([90.0, 90.0, 90.0])
+            angles = _BoxArray(array=(90.0, 90.0, 90.0), var="angles", box=self)
         elif isinstance(angles, (list, np.ndarray)):
-            angles = np.array(angles, dtype=np.float)
+            angles = _BoxArray(array=angles, var="angles", box=self)
         self._angles = angles
 
     @property
@@ -61,35 +73,78 @@ def angles(self):
 
     @mins.setter
     def mins(self, mins):
-        if isinstance(mins, list):
-            mins = np.array(mins, dtype=np.float)
+        mins = np.array(mins, dtype=np.float)
         assert mins.shape == (3, )
-        self._mins = mins
-        self._lengths = self.maxs - self.mins
+        assert all(mins <= self.maxs), "Given mins is greater than maxs"
+        self._mins = _BoxArray(array=mins, var="mins", box=self)
+        self._lengths = _BoxArray(array=(self.maxs - self.mins), var="lengths", box=self)
 
     @maxs.setter
-    def maxs(self, maxes):
-        if isinstance(maxes, list):
-            maxes = np.array(maxes, dtype=np.float)
-        assert maxes.shape == (3, )
-        self._maxs = maxes
-        self._lengths = self.maxs - self.mins
+    def maxs(self, maxs):
+        maxs = np.array(maxs, dtype=np.float)
+        assert maxs.shape == (3, )
+        assert all(maxs >= self.mins), "Given maxs is less than mins"
+        self._maxs = _BoxArray(array=maxs, var="maxs", box=self)
+        self._lengths = _BoxArray(array= (self.maxs - self.mins), var="lengths", box=self)
 
     @lengths.setter
     def lengths(self, lengths):
-        if isinstance(lengths, list):
+        if isinstance(lengths, int) or isinstance(lengths, float):
+            lengths = np.array(lengths*np.ones(3), dtype=np.float)
+        else:
             lengths = np.array(lengths, dtype=np.float)
         assert lengths.shape == (3, )
-        self._maxs += 0.5*lengths - 0.5*self.lengths
-        self._mins -= 0.5*lengths - 0.5*self.lengths
-        self._lengths = lengths
+        assert all(lengths >= 0), "Given lengths are negative" 
+        self._maxs = _BoxArray(array=(self.maxs + (0.5*lengths - 0.5*self.lengths)), var="maxs", box=self)
+        self._mins = _BoxArray(array=(self.mins - (0.5*lengths - 0.5*self.lengths)), var="mins", box=self)
+        self._lengths = _BoxArray(array=lengths, var="lengths", box=self, dtype=np.float)
 
     @angles.setter
     def angles(self, angles):
-        if isinstance(angles, list):
-            angles = np.array(angles, dtype=np.float)
+        angles = np.array(angles, dtype=np.float)
         assert angles.shape == (3, )
-        self._angles = angles
+        self._angles = _BoxArray(array=angles, var="angles", box=self, dtype=np.float)
 
     def __repr__(self):
         return "Box(mins={}, maxs={}, angles={})".format(self.mins, self.maxs, self.angles)
+
+class _BoxArray(np.ndarray):
+    """Subclass of np.ndarry specifically for mb.Box
+
+    This subclass is meant to be used internally to store Box attribute array.
+    This subclass is modified so that its __setitem__ method is reroute to the
+    corresponding setter method.
+
+    Parameters
+    ----------
+    array : array-like object
+        This can be tuple, list, or any array-like object that can be usually
+        passed to np.array
+    var : str
+        Corresponding Box's attributes like "maxs", "mins", "lengths", "angles"
+    box : mb.Box
+        This is the Box contains this attribute (one level up of this array)
+    """
+    def __new__(cls, array, var=None, box=None, dtype=np.float):
+        _array = np.asarray(array, dtype).view(cls)
+        _array.var = var
+        _array.box = box
+        return _array
+
+    def __setitem__(self, key, val):
+        array = list(self)
+        array[key] = val
+        if self.var == "maxs":
+            msg = "Given max value is less than box's min value" 
+            assert val >= self.box.mins[key], msg
+            self.box.maxs = array
+        elif self.var == "mins":
+            msg = "Given min value is more than box's max value"
+            assert val <= self.box.maxs[key], msg
+            self.box.mins = array
+        elif self.var == "lengths":
+            msg = "Given length value is negative"
+            assert val >= 0, msg
+            self.box.lengths = array
+        else:
+            self.box.angles = array

mbuild/formats/lammpsdata.py

@@ -142,7 +142,6 @@ def write_lammpsdata(structure, filename, atom_style='full',
     box = Box(lengths=np.array([0.1 * val for val in structure.box[0:3]]),
               angles=structure.box[3:6])
     # Divide by conversion factor
-    box.lengths /= sigma_conversion_factor
     box.maxs /= sigma_conversion_factor
 
     # Lammps syntax depends on the functional form

mbuild/tests/test_box.py

@@ -72,3 +72,44 @@ def test_setters_with_lists(self):
         assert (box.maxs - box.mins == 4 * np.ones(3)).all()
         box.angles = [90, 90, 120]
         assert (box.angles == np.array([90, 90, 120])).all()
+
+    def test_single_dimension_setter(self):
+        box = mb.Box(mins=np.zeros(3), maxs=4*np.ones(3))
+        assert (box.lengths == 4*np.ones(3)).all()
+        
+        box.maxs[0] = 5
+        box.mins[2] = 1
+        assert np.allclose(box.mins, np.array([0, 0, 1], dtype=np.float))
+        assert np.allclose(box.maxs, np.array([5, 4, 4], dtype=np.float))
+        assert np.allclose(box.lengths, np.array([5, 4, 3], dtype=np.float))
+
+        box.lengths[1] = 6
+        assert np.allclose(box.mins, np.array([0, -1, 1], dtype=np.float))
+        assert np.allclose(box.maxs, np.array([5, 5, 4], dtype=np.float))
+        assert np.allclose(box.lengths, np.array([5, 6, 3], dtype=np.float))
+
+        new_box = mb.Box(5)
+        assert np.allclose(new_box.lengths, np.array([5, 5, 5], dtype=np.float))
+        assert np.allclose(new_box.mins, np.array([0, 0, 0], dtype=np.float))
+        assert np.allclose(new_box.maxs, np.array([5, 5, 5], dtype=np.float))
+
+        new_box.lengths = 4
+        assert np.allclose(new_box.lengths, np.array([4, 4, 4], dtype=np.float))
+        assert np.allclose(new_box.mins, np.array([0.5, 0.5, 0.5], dtype=np.float))
+        assert np.allclose(new_box.maxs, np.array([4.5, 4.5, 4.5], dtype=np.float))
+
+    def test_sanity_checks(self):
+        # Initialization step
+        with pytest.raises(AssertionError):
+            box = mb.Box(mins=[3,3,3], maxs=[1,1,1])
+        with pytest.raises(AssertionError):
+            box = mb.Box(lengths=-1)
+        
+        # Modifying step
+        box = mb.Box(mins=[2,2,2], maxs=[4,4,4])
+        with pytest.raises(AssertionError):
+            box.mins[1] = box.maxs[1] + 1
+        with pytest.raises(AssertionError): 
+            box.maxs[1] = box.mins[1] - 1
+        with pytest.raises(AssertionError):
+            box.lengths = -1