improve docs, variable names, test

brainglobe_template_builder/preproc/transform_utils.py

@@ -76,57 +76,60 @@ def apply_transform(
 
 
 def downsample_anisotropic_image_stack(
-    stack: da.Array, xy_downsampling: int, z_downsampling: int
+    stack: da.Array, in_plane_factor: int, axial_factor: int
 ) -> np.ndarray:
     """
 
-    Lazily downsamples a dask array first along axis 1,2 and then along axis 0,
-    using a local mean of the pixels. The (smaller) array is returned
-    in memory (numpy) form at the end.
+    Lazily downsamples a dask array first along axes 1,2 (in-plane) and then
+    along axis 0 (axial), using a local mean of the pixels. The image is
+    zero-padded to allow for the correct dimensions of the averaging
+    neighbourhood, since it uses `skimage.transform.downscale_local_mean`
+    under the hood.
 
     This setup is typical for certain types of microscopy,
-    where z-resolution is lower than x-y-resolution.
+    where axial resolution is lower than in-plane resolution.
 
-    The input dask array must be chunked by x-y slice,
+    The input dask array must be chunked by plane. The (smaller) array
+    is returned in memory (numpy) form at the end.
 
     Parameters:
     ----------
     stack : da.Array
         The input dask array representing the image stack.
-    xy_downsampling : int
-        The downsampling factor for the x and y axes.
-    z_downsampling : int
-        The downsampling factor for the z axis.
+    in_plane_factor : int
+        The in-plane downsampling factor (axes=1,2).
+    axial_factor : int
+        The downsampling factor in axial direction (axis=0).
     Returns:
     -------
     np.ndarray
         The computed downsampled (numpy) array.
     Raises:
     ------
     AssertionError
-        If the array is not chunked slice-wise on axis 0.
+        If the array is not chunked by plane along axis 0.
     """
     # check we have expected slice chunks
     assert np.all(
         np.array(stack.chunks[0]) == 1
-    ), f"Array not chunked slice-wise! Chunks on axis 0 are {stack.chunks[0]}"
+    ), f"Array not chunked by plane! Chunks on axis 0 are {stack.chunks[0]}"
 
     # we have xy slices as chunks, so apply downscaling in xy first
-    downsampled_xy = stack.map_blocks(
+    downsampled_inplane = stack.map_blocks(
         transform.downscale_local_mean,
-        (1, xy_downsampling, xy_downsampling),
+        (1, in_plane_factor, in_plane_factor),
         dtype=np.float64,
     )
 
     # rechunk so we can map_blocks along z
-    downsampled_xy = downsampled_xy.rechunk(
-        {0: downsampled_xy.shape[0], 1: -1, 2: -1}
+    downsampled_inplane = downsampled_inplane.rechunk(
+        {0: downsampled_inplane.shape[0], 1: -1, 2: -1}
     )
 
     # downsample in z
-    downsampled_z = downsampled_xy.map_blocks(
+    downsampled_axial = downsampled_inplane.map_blocks(
         transform.downscale_local_mean,
-        (z_downsampling, 1, 1),
+        (axial_factor, 1, 1),
         dtype=np.float64,
     )
-    return downsampled_z.compute()
+    return downsampled_axial.compute()

tests/test_unit/test_transform_utils.py

@@ -11,7 +11,8 @@
 @pytest.fixture()
 def stack():
     """Create a dask array representing an image stack"""
-    data = np.random.rand(10, 100, 100)  # Random image stack
+    rng = np.random.default_rng()
+    data = rng.random(size=(10, 100, 100))
     return da.from_array(data, chunks=(1, 100, 100))
 
 
@@ -43,8 +44,7 @@ def test_downsample_anisotropic_image_stack(stack):
 
 
 def test_downsample_anisotropic_image_stack_raises(not_slicewise_stack):
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError, match="not chunked by plane!"):
         downsample_anisotropic_image_stack(
-            not_slicewise_stack, xy_downsampling=20, z_downsampling=2
+            not_slicewise_stack, in_plane_factor=20, axial_factor=2
         )
-    assert e.match("not chunked slice-wise!")