extend dims to nD (#61)

Author: kevinyamauchi

scripts/v2/integration_2d_3d/example_2d.py

@@ -81,7 +81,7 @@ def __init__(
             z_slice = z_depth // 2
 
         # Set initial slice index on the scene's dims.
-        self._scene.dims.slice_indices = (z_slice,)
+        self._scene.dims.selection.slice_indices = {0: z_slice}
 
         # Add the image through the standard controller API.
         appearance = ImageAppearance(
@@ -161,7 +161,7 @@ def _setup_ui(self, canvas_widget: QWidget) -> None:
         panel_layout.addWidget(QLabel("Z slice:"))
         self._z_slice_sb = QSpinBox()
         self._z_slice_sb.setRange(0, self._z_max)
-        self._z_slice_sb.setValue(self._scene.dims.slice_indices[0])
+        self._z_slice_sb.setValue(self._scene.dims.selection.slice_indices[0])
         self._z_slice_sb.valueChanged.connect(self._on_z_slice_changed)
         panel_layout.addWidget(self._z_slice_sb)
 
@@ -188,7 +188,7 @@ def _on_level_radio_clicked(self, button: QRadioButton) -> None:
         self._force_level = button.property("force_level")
 
     def _on_z_slice_changed(self, value: int) -> None:
-        self._scene.dims.slice_indices = (value,)
+        self._scene.dims.selection.slice_indices = {0: value}
         # All cached tiles are from the old z-plane — clear and re-fetch.
         gfx_visual = self._get_gfx_visual()
         gfx_visual._block_cache_2d.tile_manager.clear()

scripts/v2/integration_2d_3d/example_combined.py

@@ -253,7 +253,7 @@ def __init__(self, data_store: MultiscaleZarrDataStore) -> None:
         self._scene_2d = self._controller.add_scene(
             dim="2d", coordinate_system=cs, name="scene_2d"
         )
-        self._scene_2d.dims.slice_indices = (z_depth // 2,)
+        self._scene_2d.dims.selection.slice_indices = {0: z_depth // 2}
 
         self._visual_2d = self._controller.add_image(
             data=data_store,
@@ -364,7 +364,7 @@ def _setup_ui(self) -> None:
 
         self._z_slice_sb = QSpinBox()
         self._z_slice_sb.setRange(0, self._z_max)
-        self._z_slice_sb.setValue(self._scene_2d.dims.slice_indices[0])
+        self._z_slice_sb.setValue(self._scene_2d.dims.selection.slice_indices[0])
         self._z_slice_sb.valueChanged.connect(self._on_z_slice_changed)
         pl.addWidget(self._z_slice_sb)
         self._widget_2d.append(self._z_slice_sb)
@@ -493,7 +493,7 @@ def _on_level_radio_clicked(self, button: QRadioButton) -> None:
         self._force_level = button.property("force_level")
 
     def _on_z_slice_changed(self, value: int) -> None:
-        self._scene_2d.dims.slice_indices = (value,)
+        self._scene_2d.dims.selection.slice_indices = {0: value}
         # All cached tiles are from the old z-plane — clear and re-fetch.
         gfx_visual = self._get_gfx_visual_2d()
         gfx_visual._block_cache_2d.tile_manager.clear()

scripts/v2/integration_2d_3d/example_combined_camera_redraw.py

@@ -227,7 +227,7 @@ def __init__(self, data_store: MultiscaleZarrDataStore) -> None:
         self._scene_2d = self._controller.add_scene(
             dim="2d", coordinate_system=cs, name="scene_2d"
         )
-        self._scene_2d.dims.slice_indices = (z_depth // 2,)
+        self._scene_2d.dims.selection.slice_indices = {0: z_depth // 2}
 
         self._visual_2d = self._controller.add_image(
             data=data_store,
@@ -348,7 +348,7 @@ def _setup_ui(self) -> None:
 
         self._z_slice_sb = QSpinBox()
         self._z_slice_sb.setRange(0, self._z_max)
-        self._z_slice_sb.setValue(self._scene_2d.dims.slice_indices[0])
+        self._z_slice_sb.setValue(self._scene_2d.dims.selection.slice_indices[0])
         self._z_slice_sb.valueChanged.connect(self._on_z_slice_changed)
         pl.addWidget(self._z_slice_sb)
         self._widget_2d.append(self._z_slice_sb)
@@ -486,7 +486,7 @@ def _on_level_radio_clicked(self, button: QRadioButton) -> None:
         self._force_level = button.property("force_level")
 
     def _on_z_slice_changed(self, value: int) -> None:
-        self._scene_2d.dims.slice_indices = (value,)
+        self._scene_2d.dims.selection.slice_indices = {0: value}
         # All cached tiles are from the old z-plane — clear and re-fetch.
         gfx_visual = self._get_gfx_visual_2d()
         gfx_visual._block_cache_2d.tile_manager.clear()

scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_cellier_2d.py

@@ -94,7 +94,7 @@ def __init__(
         self._z_slice = z_slice
 
         # Set initial slice index on the scene's dims.
-        self._scene.dims.slice_indices = (self._z_slice,)
+        self._scene.dims.selection.slice_indices = {0: self._z_slice}
 
         self._visual = self._controller.add_image(
             data=data_store,
@@ -201,7 +201,7 @@ def _on_z_slice_changed(self, value: int) -> None:
         self._z_slice = value
         # Update the scene's dims — this is what SceneManager reads
         # when building slice requests.
-        self._scene.dims.slice_indices = (value,)
+        self._scene.dims.selection.slice_indices = {0: value}
         # Clear the 2D tile cache so all tiles are re-fetched at the
         # new Z slice.
         gfx_visual = self._get_gfx_visual()

scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_v2.py

@@ -52,6 +52,7 @@
     QWidget,
 )
 
+from cellier.v2._state import AxisAlignedSelectionState
 from cellier.v2.data.image import MultiscaleZarrDataStore
 from cellier.v2.render import DimsState, RenderManager
 from cellier.v2.render._frustum import frustum_planes_from_corners
@@ -277,7 +278,12 @@ async def _on_update_clicked(self) -> None:
         config.force_level = self._force_level
         config.frustum_cull = self._frustum_cull_cb.isChecked()
 
-        dims_state = DimsState(displayed_axes=(0, 1, 2), slice_indices=())
+        dims_state = DimsState(
+            axis_labels=("z", "y", "x"),
+            selection=AxisAlignedSelectionState(
+                displayed_axes=(0, 1, 2), slice_indices={}
+            ),
+        )
         self._render_manager.trigger_update(dims_state)
 
         visual = self._visual

src/cellier/v2/_state.py

@@ -5,11 +5,40 @@
 from typing import Literal, NamedTuple
 
 
+class AxisAlignedSelectionState(NamedTuple):
+    """Immutable snapshot of an axis-aligned selection."""
+
+    displayed_axes: tuple[int, ...]
+    slice_indices: dict[int, int]
+
+    def to_index_selection(self, ndim: int) -> tuple[int | slice, ...]:
+        """Return a per-axis numpy indexer in axis order.
+
+        Displayed axes → ``slice(None)``, sliced axes → their int value.
+        """
+        result: list[int | slice] = []
+        for axis in range(ndim):
+            if axis in self.slice_indices:
+                result.append(self.slice_indices[axis])
+            else:
+                result.append(slice(None))
+        return tuple(result)
+
+
+class PlaneSelectionState(NamedTuple):
+    """Stub — not yet implemented."""
+
+    pass
+
+
+SelectionState = AxisAlignedSelectionState | PlaneSelectionState
+
+
 class DimsState(NamedTuple):
     """Current dimension display state for a scene."""
 
-    displayed_axes: tuple[int, ...]
-    slice_indices: tuple[int, ...]
+    axis_labels: tuple[str, ...]
+    selection: SelectionState
 
 
 class CameraState(NamedTuple):

src/cellier/v2/controller.py

@@ -8,7 +8,6 @@
 
 import numpy as np
 
-from cellier.v2._state import CameraState, DimsState
 from cellier.v2.events import (
     AppearanceChangedEvent,
     CameraChangedEvent,
@@ -29,7 +28,11 @@
     PerspectiveCamera,
 )
 from cellier.v2.scene.canvas import Canvas
-from cellier.v2.scene.dims import CoordinateSystem, DimsManager
+from cellier.v2.scene.dims import (
+    AxisAlignedSelection,
+    CoordinateSystem,
+    DimsManager,
+)
 from cellier.v2.scene.scene import Scene
 from cellier.v2.viewer_model import DataManager, ViewerModel
 from cellier.v2.visuals._image import MultiscaleImageVisual
@@ -40,6 +43,7 @@
     from psygnal import EmissionInfo
     from PySide6.QtWidgets import QWidget
 
+    from cellier.v2._state import CameraState, DimsState
     from cellier.v2.data._base_data_store import BaseDataStore
     from cellier.v2.visuals._image import ImageAppearance
 
@@ -174,19 +178,23 @@ def add_scene(
         Scene
             The live model object.  Use ``scene.id`` in subsequent calls.
         """
+        ndim = len(coordinate_system.axis_labels)
         if dim == "3d":
-            displayed_axes = (0, 1, 2)
-            slice_indices: tuple[int, ...] = ()
+            displayed_axes = tuple(range(ndim))
+            slice_indices_dict: dict[int, int] = {}
         elif dim == "2d":
-            displayed_axes = (0, 1)
-            slice_indices = (0,)
+            # Display the last two axes; slice all others at 0.
+            displayed_axes = tuple(range(ndim - 2, ndim))
+            slice_indices_dict = {i: 0 for i in range(ndim - 2)}
         else:
             raise ValueError(f"dim must be '2d' or '3d', got {dim!r}")
 
         dims = DimsManager(
             coordinate_system=coordinate_system,
-            displayed_axes=displayed_axes,
-            slice_indices=slice_indices,
+            selection=AxisAlignedSelection(
+                displayed_axes=displayed_axes,
+                slice_indices=slice_indices_dict,
+            ),
         )
         scene = Scene(name=name, dims=dims)
         self._model.scenes[scene.id] = scene
@@ -267,7 +275,7 @@ def add_image(
             self._model.data.stores[data.id] = data
 
         scene = self._model.scenes[scene_id]
-        dim = scene.dims.displayed_axes  # (0,1,2) → 3d, (0,1) → 2d
+        displayed_axes = scene.dims.selection.displayed_axes
 
         downscale_factors = [2**i for i in range(data.n_levels)]
 
@@ -280,12 +288,13 @@ def add_image(
 
         scene.visuals.append(visual_model)
 
-        render_modes = {"3d"} if len(dim) == 3 else {"2d"}
+        render_modes = {"3d"} if len(displayed_axes) == 3 else {"2d"}
 
         level_shapes = list(data.level_shapes)
         gfx_visual = GFXMultiscaleImageVisual.from_cellier_model(
             model=visual_model,
             level_shapes=level_shapes,
+            displayed_axes=displayed_axes,
             render_modes=render_modes,
             block_size=block_size,
             gpu_budget_bytes=gpu_budget_bytes,
@@ -358,11 +367,13 @@ def add_visual(
         scene = self._model.scenes[scene_id]
         scene.visuals.append(visual_model)
 
-        render_modes = {"3d"} if len(scene.dims.displayed_axes) == 3 else {"2d"}
+        displayed_axes = scene.dims.selection.displayed_axes
+        render_modes = {"3d"} if len(displayed_axes) == 3 else {"2d"}
         level_shapes = list(data_store.level_shapes)
         gfx_visual = GFXMultiscaleImageVisual.from_cellier_model(
             model=visual_model,
             level_shapes=level_shapes,
+            displayed_axes=displayed_axes,
             render_modes=render_modes,
             block_size=block_size,
             gpu_budget_bytes=gpu_budget_bytes,
@@ -408,7 +419,9 @@ def add_canvas(
         # Infer dimensionality from the scene if not specified.
         if available_dims is None:
             scene = self._model.scenes[scene_id]
-            available_dims = "3d" if len(scene.dims.displayed_axes) == 3 else "2d"
+            available_dims = (
+                "3d" if len(scene.dims.selection.displayed_axes) == 3 else "2d"
+            )
 
         camera_model = PerspectiveCamera(
             fov=fov,
@@ -487,33 +500,29 @@ def _build_visual_configs_for_scene(
 
     def _dims_state_for_scene(self, scene_id: UUID) -> DimsState:
         """Derive a DimsState from the scene's DimsManager."""
-        dims = self._model.scenes[scene_id].dims
-        return DimsState(
-            displayed_axes=dims.displayed_axes,
-            slice_indices=dims.slice_indices,
-        )
+        return self._model.scenes[scene_id].dims.to_state()
 
     # ------------------------------------------------------------------
     # psygnal bridges
     # ------------------------------------------------------------------
 
     def _wire_dims_model(self, scene: Scene) -> None:
         """Subscribe to all field changes on a scene's DimsManager."""
-        self._dims_cache[scene.id] = scene.dims.displayed_axes
+        self._dims_cache[scene.id] = scene.dims.selection.displayed_axes
         scene.dims.events.connect(self._make_dims_handler(scene.id))
 
     def _make_dims_handler(self, scene_id: UUID) -> Callable:
         """Return a psygnal catch-all handler for a scene's DimsManager."""
 
         def _on_dims_psygnal(info: EmissionInfo) -> None:
-            dims = self._model.scenes[scene_id].dims
-            new_state = DimsState(
-                displayed_axes=dims.displayed_axes,
-                slice_indices=dims.slice_indices,
-            )
+            new_state = self._model.scenes[scene_id].dims.to_state()
             prev_axes = self._dims_cache[scene_id]
-            displayed_axes_changed = prev_axes != new_state.displayed_axes
-            self._dims_cache[scene_id] = new_state.displayed_axes
+            displayed_axes_changed = prev_axes != new_state.selection.displayed_axes
+            self._dims_cache[scene_id] = new_state.selection.displayed_axes
+            if displayed_axes_changed:
+                self._rebuild_visuals_geometry(
+                    scene_id, new_state.selection.displayed_axes
+                )
             self._event_bus.emit(
                 DimsChangedEvent(
                     source_id=self._id,
@@ -525,6 +534,28 @@ def _on_dims_psygnal(info: EmissionInfo) -> None:
 
         return _on_dims_psygnal
 
+    def _rebuild_visuals_geometry(
+        self, scene_id: UUID, displayed_axes: tuple[int, ...]
+    ) -> None:
+        """Rebuild geometry for all visuals in a scene after displayed_axes change."""
+        scene = self._model.scenes[scene_id]
+        scene_manager = self._render_manager._scenes[scene_id]
+
+        for visual_model in scene.visuals:
+            data_store = self._model.data.stores[UUID(visual_model.data_store_id)]
+            level_shapes = list(data_store.level_shapes)
+            gfx_visual = scene_manager.get_visual(visual_model.id)
+
+            old_node, new_node = gfx_visual.rebuild_geometry(
+                level_shapes, displayed_axes
+            )
+
+            # Swap node in the pygfx scene graph.
+            if old_node is not None:
+                scene_manager.scene.remove(old_node)
+            if new_node is not None:
+                scene_manager.scene.add(new_node)
+
     def _wire_appearance(self, visual: MultiscaleImageVisual) -> None:
         """Subscribe to all field changes on a visual's ImageAppearance."""
         visual.appearance.events.connect(self._make_appearance_handler(visual.id))

src/cellier/v2/data/image/_image_requests.py

@@ -19,44 +19,24 @@ class ChunkRequest(NamedTuple):
     may be negative or exceed the store bounds; ``get_data()`` on the
     ``MultiscaleZarrDataStore`` is responsible for clamping and
     zero-padding so the returned array always has the full requested
-    shape ``(z_stop-z_start, y_stop-y_start, x_stop-x_start)`` for 3D
-    requests, or ``(y_stop-y_start, x_stop-x_start)`` for 2D requests
-    (when ``z_slice`` is set).
+    shape.
 
     Parameters
     ----------
     chunk_request_id :
-        Unique ID for this individual chunk.  Used by the callback to look
-        up the corresponding ``TileSlot`` in the ``slot_map`` closure.
+        Unique ID for this individual chunk.
     slice_request_id :
-        Shared ID for all chunks that belong to the same planning event
-        (one Update press → one ``slice_request_id``).  Used by
-        ``AsyncSlicer.cancel()`` to cancel all in-flight work for a
-        single update cycle.
+        Shared ID for all chunks that belong to the same planning event.
     scale_index :
-        0-based index into ``MultiscaleZarrDataStore`` levels
-        (0 = finest).  Note: this is ``scale_index``, not
-        ``BrickKey.level`` (which is 1-based).
-    z_start, y_start, x_start :
-        Padded region origin in voxels at this scale level.  May be
-        negative for boundary bricks.
-    z_stop, y_stop, x_stop :
-        Padded region end (exclusive) in voxels.  May exceed store
-        bounds for boundary bricks.
-    z_slice :
-        If not ``None``, this is a 2D request: the data store reads a
-        single z-plane at this index (at the current scale level) and
-        returns a 2D ``(H, W)`` array.  ``z_start`` and ``z_stop`` are
-        ignored when ``z_slice`` is set.
+        0-based index into ``MultiscaleZarrDataStore`` levels (0 = finest).
+    axis_selections :
+        Per-axis selection in data axis order.  Each element is either:
+        - ``int`` → axis is sliced (single plane, already scaled to this level)
+        - ``(start, stop)`` → axis is displayed (windowed range; may extend
+          outside bounds)
     """
 
     chunk_request_id: UUID
     slice_request_id: UUID
     scale_index: int
-    z_start: int
-    y_start: int
-    x_start: int
-    z_stop: int
-    y_stop: int
-    x_stop: int
-    z_slice: int | None = None
+    axis_selections: tuple[int | tuple[int, int], ...]