Floored version of align_with_mag() (#217)

* Floored version of align_with_mag() * Format * Update wkcuber/api/bounding_box.py Co-authored-by: Philipp Otto <[email protected]> * Update wkcuber/api/bounding_box.py Co-authored-by: Jonathan Striebel <[email protected]> * Fix lambda
scalableminds · Jun 9, 2020 · 175dac6 · 175dac6
1 parent b5dee90
commit 175dac6
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 9 deletions.
diff --git a/tests/test_bounding_box.py b/tests/test_bounding_box.py
@@ -2,19 +2,35 @@
 import pytest
 
 
-def test_align_with_mag():
+def test_align_with_mag_ceiled():
+
+    assert BoundingBox((1, 1, 1), (10, 10, 10)).align_with_mag(
+        Mag(2), ceil=True
+    ) == BoundingBox(topleft=(0, 0, 0), size=(12, 12, 12))
+    assert BoundingBox((1, 1, 1), (9, 9, 9)).align_with_mag(
+        Mag(2), ceil=True
+    ) == BoundingBox(topleft=(0, 0, 0), size=(10, 10, 10))
+    assert BoundingBox((1, 1, 1), (9, 9, 9)).align_with_mag(
+        Mag(4), ceil=True
+    ) == BoundingBox(topleft=(0, 0, 0), size=(12, 12, 12))
+    assert BoundingBox((1, 2, 3), (9, 9, 9)).align_with_mag(
+        Mag(2), ceil=True
+    ) == BoundingBox(topleft=(0, 2, 2), size=(10, 10, 10))
+
+
+def test_align_with_mag_floored():
 
     assert BoundingBox((1, 1, 1), (10, 10, 10)).align_with_mag(Mag(2)) == BoundingBox(
-        topleft=(0, 0, 0), size=(12, 12, 12)
+        topleft=(2, 2, 2), size=(8, 8, 8)
     )
     assert BoundingBox((1, 1, 1), (9, 9, 9)).align_with_mag(Mag(2)) == BoundingBox(
-        topleft=(0, 0, 0), size=(10, 10, 10)
+        topleft=(2, 2, 2), size=(8, 8, 8)
     )
     assert BoundingBox((1, 1, 1), (9, 9, 9)).align_with_mag(Mag(4)) == BoundingBox(
-        topleft=(0, 0, 0), size=(12, 12, 12)
+        topleft=(4, 4, 4), size=(4, 4, 4)
     )
     assert BoundingBox((1, 2, 3), (9, 9, 9)).align_with_mag(Mag(2)) == BoundingBox(
-        topleft=(0, 2, 2), size=(10, 10, 10)
+        topleft=(2, 2, 4), size=(8, 8, 8)
     )
 
 

diff --git a/wkcuber/api/bounding_box.py b/wkcuber/api/bounding_box.py
@@ -201,12 +201,22 @@ def in_mag(self, mag: Mag) -> "BoundingBox":
             size=(self.size // np_mag).astype(np.int),
         )
 
-    def align_with_mag(self, mag: Mag):
-        """Rounds the bounding box up, so that both topleft and bottomright are divisible by mag."""
+    def align_with_mag(self, mag: Mag, ceil=False):
+        """Rounds the bounding box, so that both topleft and bottomright are divisible by mag.
+
+        :argument ceil: If true, the bounding box is enlarged when necessary. If false, it's shrinked when necessary.
+        """
 
         np_mag = np.array(mag.to_array())
-        topleft = (self.topleft // np_mag).astype(np.int) * np_mag
-        bottomright = np.ceil(self.bottomright / np_mag).astype(np.int) * np_mag
+
+        align = lambda point, round_fn: round_fn(point / np_mag).astype(np.int) * np_mag
+
+        if ceil:
+            topleft = align(self.topleft, np.floor)
+            bottomright = align(self.bottomright, np.ceil)
+        else:
+            topleft = align(self.topleft, np.ceil)
+            bottomright = align(self.bottomright, np.floor)
         return BoundingBox(topleft, bottomright - topleft)
 
     def contains(self, coord: Shape3D) -> bool: