fix: numerical stability on decimal resolutions

ZettaAI · Dec 6, 2024 · 5d2d0f4 · 5d2d0f4
1 parent 358847e
commit 5d2d0f4
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 0 deletions.
diff --git a/tests/unit/geometry/test_bbox.py b/tests/unit/geometry/test_bbox.py
@@ -138,6 +138,46 @@ def test_constructor(
     assert result.bounds == expected_bounds
 
 
+@pytest.mark.parametrize(
+    "bbox, vec, expected",
+    [
+        [
+            BBox3D(bounds=((0, 1), (0, 2), (0, 3))),
+            Vec3D(1, 2, 3),
+            BBox3D(bounds=((0, 1), (0, 1), (0, 1))),
+        ],
+        [
+            BBox3D(bounds=((0, 5734.4), (0, 8601.6), (0, 1))),
+            Vec3D(2.8, 2.8, 1),
+            BBox3D(bounds=((0, 2048), (0, 3072), (0, 1))),
+        ],
+    ],
+)
+def test_truediv(bbox: BBox3D, vec: Vec3D, expected: BBox3D):
+    result = bbox / vec
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    "bbox, vec, expected",
+    [
+        [
+            BBox3D(bounds=((0, 1), (0, 2), (0, 3))),
+            Vec3D(1, 2, 3),
+            BBox3D(bounds=((0, 1), (0, 4), (0, 9))),
+        ],
+        [
+            BBox3D(bounds=((0, 2048), (0, 3072), (0, 1))),
+            Vec3D(2.8, 2.8, 1),
+            BBox3D(bounds=((0, 5734.4), (0, 8601.6), (0, 1))),
+        ],
+    ],
+)
+def test_mul(bbox: BBox3D, vec: Vec3D, expected: BBox3D):
+    result = bbox * vec
+    assert result == expected
+
+
 @pytest.mark.parametrize(
     "bbox, resolution, allow_slice_rounding, expected",
     [

diff --git a/zetta_utils/geometry/bbox.py b/zetta_utils/geometry/bbox.py
@@ -44,6 +44,19 @@ class BBox3D:  # pylint: disable=too-many-public-methods # fundamental class
     unit: str = DEFAULT_UNIT
     pprint_px_resolution: Sequence[float] = (1, 1, 1)
 
+    def __attrs_post_init__(self):
+        object.__setattr__(
+            self,
+            "bounds",
+            tuple(
+                (
+                    round(start, VEC3D_PRECISION),
+                    round(end, VEC3D_PRECISION),
+                )
+                for start, end in self.bounds
+            ),
+        )
+
     @property
     def ndim(self) -> int:
         return 3
@@ -154,6 +167,12 @@ def from_points(
 
         return cls(bounds=bounds, unit=unit)
 
+    def __truediv__(self, vec: Vec3D) -> BBox3D:
+        return BBox3D.from_coords(self.start / vec, self.end / vec)
+
+    def __mul__(self, vec: Vec3D) -> BBox3D:
+        return BBox3D.from_coords(self.start * vec, self.end * vec)
+
     def get_slice(
         self,
         dim: int,

diff --git a/zetta_utils/layer/volumetric/index.py b/zetta_utils/layer/volumetric/index.py
@@ -76,6 +76,16 @@ def from_coords(
             allow_slice_rounding=allow_slice_rounding,
         )
 
+    def __truediv__(self, vec: Vec3D) -> VolumetricIndex:
+        return VolumetricIndex(
+            self.resolution, self.bbox / vec, self.chunk_id, self.allow_slice_rounding
+        )
+
+    def __mul__(self, vec: Vec3D) -> VolumetricIndex:
+        return VolumetricIndex(
+            self.resolution, self.bbox * vec, self.chunk_id, self.allow_slice_rounding
+        )
+
     def to_slices(self):
         """
         Represent this index as a tuple of slices.