Repository: scalableminds/webknossos-libs

Issue:
Don't warn about unaligned/compressed writes if these are at dataset boundary
For example, during downsampling there are warnings about unaligned/compressed writes (see cuber CI). I think, these happen when the dataset boundaries are not block-aligned. In that case, the unaligned/compressed writes should not be problematic (since they cannot be avoided except for intentionally writing outside of the dataset's bounding box).

@rschwanhold Could you have a look into this? Let me know if we should have a call about it. I assume that `_handle_compressed_write` needs to be adapted so that the mag's bounding box is taken into account, too.


Discussion:

Resolving patch:
diff --git a/wkcuber/api/mag_view.py b/wkcuber/api/mag_view.py
--- a/wkcuber/api/mag_view.py
+++ b/wkcuber/api/mag_view.py
@@ -19,7 +19,6 @@
 from wkcuber.api.bounding_box import BoundingBox
 from wkcuber.compress_utils import compress_file_job
 from wkcuber.utils import (
-    convert_mag1_size,
     convert_mag1_offset,
     get_executor_for_args,
     wait_and_ensure_success,
@@ -77,23 +76,23 @@ def __init__(
             file_len=file_len,
             block_type=block_type,
         )
+
+        self.layer = layer
+        self.name = name
+
         super().__init__(
             _find_mag_path_on_disk(layer.dataset.path, layer.name, name),
             header,
             cast(
                 Tuple[int, int, int],
-                tuple(
-                    convert_mag1_size(layer.get_bounding_box().bottomright, Mag(name))
-                ),
+                tuple(self._mag_view_bounding_box_at_creation.bottomright),
             ),
             (0, 0, 0),
             False,
             False,
+            None,
         )
 
-        self.layer = layer
-        self.name = name
-
         if create:
             wkw.Dataset.create(
                 join(layer.dataset.path, layer.name, self.name), self.header
@@ -328,6 +327,14 @@ def _get_file_dimensions(self) -> Tuple[int, int, int]:
             Tuple[int, int, int], (self.header.file_len * self.header.block_len,) * 3
         )
 
+    @property
+    def _mag_view_bounding_box_at_creation(self) -> BoundingBox:
+        return (
+            self.layer.get_bounding_box()
+            .align_with_mag(Mag(self.name), ceil=True)
+            .in_mag(Mag(self.name))
+        )
+
     def __repr__(self) -> str:
         return repr(
             "MagView(name=%s, global_offset=%s, size=%s)"
diff --git a/wkcuber/api/view.py b/wkcuber/api/view.py
--- a/wkcuber/api/view.py
+++ b/wkcuber/api/view.py
@@ -29,6 +29,7 @@ def __init__(
         global_offset: Tuple[int, int, int] = (0, 0, 0),
         is_bounded: bool = True,
         read_only: bool = False,
+        mag_view_bbox_at_creation: Optional[BoundingBox] = None,
     ):
         """
         Do not use this constructor manually. Instead use `wkcuber.api.mag_view.MagView.get_view()` to get a `View`.
@@ -41,6 +42,16 @@ def __init__(
         self._is_bounded = is_bounded
         self.read_only = read_only
         self._is_opened = False
+        # The bounding box of the view is used to prevent warnings when writing compressed but unaligned data
+        # directly at the borders of the bounding box.
+        # A View is unable to get this information from the Dataset because it is detached from it.
+        # Adding the bounding box as parameter is a workaround for this.
+        # However, keep in mind that this bounding box is just a snapshot.
+        # This bounding box is not updated if the bounding box of the dataset is updated.
+        # Even though an outdated bounding box can lead to missed (or unwanted) warnings,
+        # this is sufficient for our use case because such scenarios are unlikely and not critical.
+        # This should not be misused to infer the size of the dataset because this might lead to problems.
+        self._mag_view_bbox_at_creation = mag_view_bbox_at_creation
 
     def open(self) -> "View":
         """
@@ -244,6 +255,7 @@ def get_view(
             global_offset=view_offset,
             is_bounded=True,
             read_only=read_only,
+            mag_view_bbox_at_creation=self._mag_view_bounding_box_at_creation,
         )
 
     def _assert_bounds(
@@ -426,24 +438,42 @@ def _handle_compressed_write(
         aligned_offset = absolute_offset_np - margin_to_top_left
         bottom_right = absolute_offset_np + np.array(data.shape[-3:])
         margin_to_bottom_right = (file_bb - (bottom_right % file_bb)) % file_bb
-        aligned_bottom_right = bottom_right + margin_to_bottom_right
-        aligned_shape = aligned_bottom_right - aligned_offset
+        is_bottom_right_aligned = bottom_right + margin_to_bottom_right
+        aligned_shape = is_bottom_right_aligned - aligned_offset
 
         if (
             tuple(aligned_offset) != absolute_offset
             or tuple(aligned_shape) != data.shape[-3:]
         ):
-            # the data is not aligned
-            # read the aligned bounding box
-
-            # We want to read the data at the absolute offset.
-            # The absolute offset might be outside of the current view.
-            # That is the case if the data is compressed but the view does not include the whole file on disk.
-            # In this case we avoid checking the bounds because the aligned_offset and aligned_shape are calculated internally.
-            warnings.warn(
-                "Warning: write() was called on a compressed mag without block alignment. Performance will be degraded as the data has to be padded first.",
-                RuntimeWarning,
+            mag_view_bbox_at_creation = self._mag_view_bounding_box_at_creation
+
+            # Calculate in which dimensions the data is aligned and in which dimensions it matches the bbox of the mag.
+            is_top_left_aligned = aligned_offset == np.array(absolute_offset)
+            is_bottom_right_aligned = is_bottom_right_aligned == bottom_right
+            is_at_border_top_left = np.array(
+                mag_view_bbox_at_creation.topleft
+            ) == np.array(absolute_offset)
+            is_at_border_bottom_right = (
+                np.array(mag_view_bbox_at_creation.bottomright) == bottom_right
             )
+
+            if not (
+                np.all(np.logical_or(is_top_left_aligned, is_at_border_top_left))
+                and np.all(
+                    np.logical_or(is_bottom_right_aligned, is_at_border_bottom_right)
+                )
+            ):
+                # the data is not aligned
+                # read the aligned bounding box
+
+                # We want to read the data at the absolute offset.
+                # The absolute offset might be outside of the current view.
+                # That is the case if the data is compressed but the view does not include the whole file on disk.
+                # In this case we avoid checking the bounds because the aligned_offset and aligned_shape are calculated internally.
+                warnings.warn(
+                    "Warning: write() was called on a compressed mag without block alignment. Performance will be degraded as the data has to be padded first.",
+                    RuntimeWarning,
+                )
             aligned_data = self._read_without_checks(aligned_offset, aligned_shape)
 
             index_slice = (
@@ -484,6 +514,11 @@ def __repr__(self) -> str:
             % (self.path, self.global_offset, self.size)
         )
 
+    @property
+    def _mag_view_bounding_box_at_creation(self) -> BoundingBox:
+        assert self._mag_view_bbox_at_creation is not None
+        return self._mag_view_bbox_at_creation
+
 
 def _assert_positive_dimensions(
     offset: Tuple[int, int, int], size: Tuple[int, int, int]
diff --git a/wkcuber/utils.py b/wkcuber/utils.py
--- a/wkcuber/utils.py
+++ b/wkcuber/utils.py
@@ -496,12 +496,6 @@ def named_partial(func: F, *args: Any, **kwargs: Any) -> F:
     return partial_func
 
 
-def convert_mag1_size(
-    mag1_size: Union[List, np.ndarray], target_mag: Mag
-) -> np.ndarray:
-    return ceil_div_np(np.array(mag1_size), target_mag.as_np())
-
-
 def convert_mag1_offset(
     mag1_offset: Union[List, np.ndarray], target_mag: Mag
 ) -> np.ndarray:


Test patch:
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1048,20 +1048,20 @@ def test_writing_subset_of_compressed_data() -> None:
     # create uncompressed dataset
     Dataset.create(TESTOUTPUT_DIR / "compressed_data", scale=(1, 1, 1)).add_layer(
         "color", LayerCategories.COLOR_TYPE
-    ).add_mag("1", block_len=8, file_len=8).write(
-        (np.random.rand(20, 40, 60) * 255).astype(np.uint8)
+    ).add_mag("2", block_len=8, file_len=8).write(
+        (np.random.rand(120, 140, 160) * 255).astype(np.uint8)
     )
 
     # compress data
     compress_mag_inplace(
         (TESTOUTPUT_DIR / "compressed_data").resolve(),
         layer_name="color",
-        mag=Mag("1"),
+        mag=Mag("2"),
     )
 
     # open compressed dataset
     compressed_mag = (
-        Dataset(TESTOUTPUT_DIR / "compressed_data").get_layer("color").get_mag("1")
+        Dataset(TESTOUTPUT_DIR / "compressed_data").get_layer("color").get_mag("2")
     )
 
     with warnings.catch_warnings():
@@ -1082,6 +1082,26 @@ def test_writing_subset_of_compressed_data() -> None:
                 data=(np.random.rand(10, 10, 10) * 255).astype(np.uint8),
             )
 
+        assert compressed_mag._mag_view_bounding_box_at_creation == BoundingBox(
+            topleft=(
+                0,
+                0,
+                0,
+            ),
+            size=(120, 140, 160),
+        )
+        # Writing unaligned data to the edge of the bounding box of the MagView does not raise an error.
+        # This write operation writes unaligned data into the bottom-right corner of the MagView.
+        compressed_mag.write(
+            offset=(64, 64, 64),
+            data=(np.random.rand(56, 76, 96) * 255).astype(np.uint8),
+        )
+        # This also works for normal Views but they only use the bounding box at the time of creation as reference.
+        compressed_mag.get_view().write(
+            offset=(64, 64, 64),
+            data=(np.random.rand(56, 76, 96) * 255).astype(np.uint8),
+        )
+
         # Writing aligned data does not raise a warning. Therefore, this does not fail with these strict settings.
         compressed_mag.write(data=(np.random.rand(64, 64, 64) * 255).astype(np.uint8))
 
diff --git a/tests/test_downsampling.py b/tests/test_downsampling.py
--- a/tests/test_downsampling.py
+++ b/tests/test_downsampling.py
@@ -1,4 +1,5 @@
 import logging
+import warnings
 from pathlib import Path
 from typing import Tuple, cast
 
@@ -21,7 +22,6 @@
 from wkcuber.downsampling_utils import _mode, non_linear_filter_3d
 import shutil
 
-WKW_CUBE_SIZE = 1024
 CUBE_EDGE_LEN = 256
 
 TESTOUTPUT_DIR = Path("testoutput")
@@ -93,23 +93,20 @@ def downsample_test_helper(use_compress: bool) -> None:
         pass
 
     source_ds = Dataset(source_path)
-    source_layer = source_ds.get_layer("color")
-    mag1 = source_layer.get_mag("1")
+    target_ds = source_ds.copy_dataset(target_path, block_len=16, file_len=16)
+
+    target_layer = target_ds.get_layer("color")
+    mag1 = target_layer.get_mag("1")
+    target_layer.delete_mag("2-2-1")  # This is not needed for this test
 
-    target_ds = Dataset.create(target_path, scale=(1, 1, 1))
-    target_layer = target_ds.add_layer(
-        "color", LayerCategories.COLOR_TYPE, dtype_per_channel="uint8"
-    )
-    # The bounding box has to be set here explicitly because the downsampled data is written to a different dataset.
-    target_layer.set_bounding_box(
-        source_layer.get_bounding_box().topleft, source_layer.get_bounding_box().size
-    )
     mag2 = target_layer._initialize_mag_from_other_mag("2", mag1, use_compress)
 
-    offset = (WKW_CUBE_SIZE, 2 * WKW_CUBE_SIZE, 0)
-    target_offset = cast(Tuple[int, int, int], tuple([o // 2 for o in offset]))
-    source_size = cast(Tuple[int, int, int], (CUBE_EDGE_LEN * 2,) * 3)
-    target_size = cast(Tuple[int, int, int], (CUBE_EDGE_LEN,) * 3)
+    # The actual size of mag1 is (4600, 4600, 512).
+    # To keep this test case fast, we are only downsampling a small part
+    offset = (4096, 4096, 0)
+    target_offset = (2048, 2048, 0)
+    source_size = (504, 504, 512)
+    target_size = (252, 252, 256)
     source_buffer = mag1.read(
         offset=offset,
         size=source_size,
@@ -127,7 +124,7 @@ def downsample_test_helper(use_compress: bool) -> None:
         ),
         [2, 2, 2],
         InterpolationModes.MAX,
-        CUBE_EDGE_LEN,
+        128,
         100,
     )
 
@@ -147,7 +144,9 @@ def test_downsample_cube_job() -> None:
 
 
 def test_compressed_downsample_cube_job() -> None:
-    downsample_test_helper(True)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("error")  # This escalates the warning to an error
+        downsample_test_helper(True)
 
 
 def test_downsample_multi_channel() -> None: