fix: Address PR Cephla-Lab#19 review comments

- Add docstring entries for registration_z and registration_t parameters
- Fix type hints: z_level and time_idx now Optional[int] instead of int
- Update read_zarr_tile and read_zarr_region to accept z_level and time_idx
- Zarr reads now respect z-level and timepoint selection instead of hardcoding

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: Alpaca233

src/tilefusion/core.py

@@ -84,6 +84,10 @@ class TileFusion:
         Channel index for registration.
     multiscale_downsample : str
         Either "stride" (default) or "block_mean" to control multiscale reduction.
+    registration_z : int, optional
+        Z-level to use for registration. If None, uses middle z-level.
+    registration_t : int
+        Timepoint to use for registration. Defaults to 0.
     """
 
     def __init__(
@@ -461,7 +465,9 @@ def _update_profiles(self) -> None:
     # I/O methods (delegate to format-specific loaders)
     # -------------------------------------------------------------------------
 
-    def _read_tile(self, tile_idx: int, z_level: int = None, time_idx: int = None) -> np.ndarray:
+    def _read_tile(
+        self, tile_idx: int, z_level: Optional[int] = None, time_idx: Optional[int] = None
+    ) -> np.ndarray:
         """Read a single tile from the input data (all channels)."""
         if z_level is None:
             z_level = self._registration_z  # Default to registration z-level
@@ -471,7 +477,7 @@ def _read_tile(self, tile_idx: int, z_level: int = None, time_idx: int = None) -
         if self._is_zarr_format:
             zarr_ts = self._metadata["tensorstore"]
             is_3d = self._metadata.get("is_3d", False)
-            tile = read_zarr_tile(zarr_ts, tile_idx, is_3d)
+            tile = read_zarr_tile(zarr_ts, tile_idx, is_3d, z_level=z_level, time_idx=time_idx)
         elif self._is_individual_tiffs_format:
             tile = read_individual_tiffs_tile(
                 self._metadata["image_folder"],
@@ -508,8 +514,8 @@ def _read_tile_region(
         tile_idx: int,
         y_slice: slice,
         x_slice: slice,
-        z_level: int = None,
-        time_idx: int = None,
+        z_level: Optional[int] = None,
+        time_idx: Optional[int] = None,
     ) -> np.ndarray:
         """Read a region of a tile from the input data."""
         if z_level is None:
@@ -521,7 +527,14 @@ def _read_tile_region(
             zarr_ts = self._metadata["tensorstore"]
             is_3d = self._metadata.get("is_3d", False)
             region = read_zarr_region(
-                zarr_ts, tile_idx, y_slice, x_slice, self.channel_to_use, is_3d
+                zarr_ts,
+                tile_idx,
+                y_slice,
+                x_slice,
+                self.channel_to_use,
+                is_3d,
+                z_level=z_level,
+                time_idx=time_idx,
             )
         elif self._is_individual_tiffs_format:
             region = read_individual_tiffs_region(

src/tilefusion/io/zarr.py

@@ -100,6 +100,8 @@ def read_zarr_tile(
     zarr_ts: ts.TensorStore,
     tile_idx: int,
     is_3d: bool = False,
+    z_level: int = None,
+    time_idx: int = 0,
 ) -> np.ndarray:
     """
     Read all channels of a tile from Zarr format.
@@ -111,18 +113,27 @@ def read_zarr_tile(
     tile_idx : int
         Index of the tile.
     is_3d : bool
-        If True, data is 3D and max projection is applied.
+        If True, data is 3D.
+    z_level : int, optional
+        Z-level to read. If None and is_3d, uses max projection.
+    time_idx : int
+        Timepoint to read. Defaults to 0.
 
     Returns
     -------
     arr : ndarray of shape (C, Y, X)
         Tile data as float32.
     """
     if is_3d:
-        arr = zarr_ts[0, tile_idx, :, :, :, :].read().result()
-        arr = np.max(arr, axis=1)  # Max projection along Z
+        if z_level is not None:
+            # Read specific z-level
+            arr = zarr_ts[time_idx, tile_idx, :, z_level, :, :].read().result()
+        else:
+            # Max projection along Z (legacy behavior)
+            arr = zarr_ts[time_idx, tile_idx, :, :, :, :].read().result()
+            arr = np.max(arr, axis=1)
     else:
-        arr = zarr_ts[0, tile_idx, :, :, :].read().result()
+        arr = zarr_ts[time_idx, tile_idx, :, :, :].read().result()
     return arr.astype(np.float32)
 
 
@@ -133,6 +144,8 @@ def read_zarr_region(
     x_slice: slice,
     channel_idx: int = 0,
     is_3d: bool = False,
+    z_level: int = None,
+    time_idx: int = 0,
 ) -> np.ndarray:
     """
     Read a region of a single channel from Zarr format.
@@ -149,18 +162,30 @@ def read_zarr_region(
         Channel index.
     is_3d : bool
         If True, data is 3D.
+    z_level : int, optional
+        Z-level to read. If None and is_3d, uses max projection.
+    time_idx : int
+        Timepoint to read. Defaults to 0.
 
     Returns
     -------
     arr : ndarray of shape (1, h, w)
         Tile region as float32.
     """
     if is_3d:
-        arr = zarr_ts[0, tile_idx, channel_idx, :, y_slice, x_slice].read().result()
-        arr = np.max(arr, axis=0)
-        arr = arr[np.newaxis, :, :]
+        if z_level is not None:
+            # Read specific z-level
+            arr = (
+                zarr_ts[time_idx, tile_idx, channel_idx, z_level, y_slice, x_slice].read().result()
+            )
+            arr = arr[np.newaxis, :, :]
+        else:
+            # Max projection along Z (legacy behavior)
+            arr = zarr_ts[time_idx, tile_idx, channel_idx, :, y_slice, x_slice].read().result()
+            arr = np.max(arr, axis=0)
+            arr = arr[np.newaxis, :, :]
     else:
-        arr = zarr_ts[0, tile_idx, channel_idx, y_slice, x_slice].read().result()
+        arr = zarr_ts[time_idx, tile_idx, channel_idx, y_slice, x_slice].read().result()
         arr = arr[np.newaxis, :, :]
     return arr.astype(np.float32)