fix: updating bbox for 2d case as well as 3d.

Also updating tests for 2d

Author: lachlangrose

LoopStructural/datatypes/_bounding_box.py

@@ -12,7 +12,7 @@ def __init__(
         maximum: Optional[np.ndarray] = None,
         nsteps: Optional[np.ndarray] = None,
         step_vector: Optional[np.ndarray] = None,
-        dimensions: int = 3,
+        dimensions: Optional[int] = None,
     ):
         """A bounding box for a model, defined by the
         origin, maximum and number of steps in each direction
@@ -32,7 +32,13 @@ def __init__(
             maximum = origin + nsteps * step_vector
         self._origin = np.array(origin)
         self._maximum = np.array(maximum)
-        self.dimensions = dimensions
+        if dimensions is None:
+            if self.origin is None:
+                raise LoopValueError("Origin is not set")
+            self.dimensions = len(self.origin)
+            print(self.dimensions)
+        else:
+            self.dimensions = dimensions
         if nsteps is None:
             self.nsteps = np.array([50, 50, 25])
         self.name_map = {
@@ -182,13 +188,13 @@ def is_inside(self, xyz):
     def regular_grid(self, nsteps=None, shuffle=False, order="C"):
         if nsteps is None:
             nsteps = self.nsteps
-        x = np.linspace(self.origin[0], self.maximum[0], nsteps[0])
-        y = np.linspace(self.origin[1], self.maximum[1], nsteps[1])
-        z = np.linspace(self.origin[2], self.maximum[2], nsteps[2])
-        xx, yy, zz = np.meshgrid(x, y, z, indexing="ij")
-        locs = np.array(
-            [xx.flatten(order=order), yy.flatten(order=order), zz.flatten(order=order)]
-        ).T
+        coordinates = [
+            np.linspace(self.origin[i], self.maximum[i], nsteps[i]) for i in range(self.dimensions)
+        ]
+
+        coordinate_grid = np.meshgrid(*coordinates, indexing="ij")
+
+        locs = np.array([c.flatten(order=order) for c in coordinate_grid]).T
         if shuffle:
             # logger.info("Shuffling points")
             rng.shuffle(locs)

tests/unit/datatypes/test_bounding_box.py

@@ -0,0 +1,107 @@
+from LoopStructural.datatypes import BoundingBox
+import numpy as np
+
+
+def test_create_bounding_box():
+    bbox = BoundingBox(origin=[0, 0, 0], maximum=[1, 1, 1])
+    assert bbox.origin[0] == 0
+    assert bbox.origin[1] == 0
+    assert bbox.origin[2] == 0
+    assert bbox.maximum[0] == 1
+    assert bbox.maximum[1] == 1
+    assert bbox.maximum[2] == 1
+    assert bbox.dimensions == 3
+    assert np.all(bbox.bb == np.array([[0, 0, 0], [1, 1, 1]]))
+    assert bbox.valid is True
+
+
+def test_create_bounding_box_2d():
+    bbox = BoundingBox(origin=[0, 0], maximum=[1, 1], dimensions=2)
+    assert bbox.origin[0] == 0
+    assert bbox.origin[1] == 0
+    assert bbox.maximum[0] == 1
+    assert bbox.maximum[1] == 1
+    assert bbox.dimensions == 2
+    assert np.all(bbox.bb == np.array([[0, 0], [1, 1]]))
+    assert bbox.valid is True
+
+
+def test_create_bounding_box_from_points():
+    bbox = BoundingBox(dimensions=3)
+    bbox.fit(np.array([[0, 0, 0], [1, 1, 1]]))
+    assert bbox.origin[0] == 0
+    assert bbox.origin[1] == 0
+    assert bbox.origin[2] == 0
+    assert bbox.maximum[0] == 1
+    assert bbox.maximum[1] == 1
+    assert bbox.maximum[2] == 1
+    assert np.all(bbox.bb == np.array([[0, 0, 0], [1, 1, 1]]))
+    assert bbox.valid is True
+
+
+def test_create_bounding_box_from_points_2d():
+    bbox = BoundingBox(dimensions=2)
+    bbox.fit(np.array([[0, 0], [1, 1]]))
+    assert bbox.origin[0] == 0
+    assert bbox.origin[1] == 0
+    assert bbox.maximum[0] == 1
+    assert bbox.maximum[1] == 1
+    assert np.all(bbox.bb == np.array([[0, 0], [1, 1]]))
+    assert bbox.valid is True
+
+
+def test_create_3d_bounding_box_from_2d_points():
+    bbox = BoundingBox(dimensions=3)
+    try:
+        bbox.fit(np.array([[0, 0], [1, 1]]))
+    except Exception as e:
+        assert str(e) == "locations array is 2D but bounding box is 3"
+    else:
+        assert False
+
+
+def test_create_with_buffer():
+    bbox = BoundingBox(origin=[0, 0, 0], maximum=[1, 1, 1])
+    bbox = bbox.with_buffer(0.2)
+    assert bbox.origin[0] == -0.2
+    assert bbox.origin[1] == -0.2
+    assert bbox.origin[2] == -0.2
+    assert bbox.maximum[0] == 1.2
+    assert bbox.maximum[1] == 1.2
+    assert bbox.maximum[2] == 1.2
+    assert np.all(bbox.bb == np.array([[-0.2, -0.2, -0.2], [1.2, 1.2, 1.2]]))
+    assert bbox.valid is True
+
+
+def test_is_inside():
+    bbox = BoundingBox(origin=[0, 0, 0], maximum=[1, 1, 1])
+    assert np.all(bbox.is_inside(np.array([0.5, 0.5, 0.5])))
+    assert not np.all(bbox.is_inside(np.array([0.5, 0.5, 1.5])))
+    assert not np.all(bbox.is_inside(np.array([0.5, 0.5, -0.5])))
+    assert not np.all(bbox.is_inside(np.array([0.5, 1.5, 0.5])))
+    assert not np.all(bbox.is_inside(np.array([0.5, -0.5, 0.5])))
+    assert not np.all(bbox.is_inside(np.array([1.5, 0.5, 0.5])))
+    assert not np.all(bbox.is_inside(np.array([-0.5, 0.5, 0.5])))
+
+
+def test_regular_grid_3d():
+    bbox = BoundingBox(origin=[0, 0, 0], maximum=[1, 1, 1])
+    print(bbox.dimensions)
+    grid = bbox.regular_grid((10, 10, 10))
+    assert grid.shape == (10 * 10 * 10, 3)
+
+
+def test_regular_grid_2d():
+    bbox = BoundingBox(origin=[0, 0], maximum=[1, 1], dimensions=2)
+    print(bbox.dimensions)
+    grid = bbox.regular_grid((10, 10))
+    assert grid.shape == (10 * 10, 2)
+
+
+if __name__ == "__main__":
+    test_create_bounding_box()
+    test_create_bounding_box_from_points()
+    test_create_with_buffer()
+    test_is_inside()
+    test_regular_grid_3d()
+    test_regular_grid_2d()