From 052f446e990c3118c05caa5f626f1935885446cb Mon Sep 17 00:00:00 2001 From: Kevin Date: Thu, 26 Mar 2026 22:29:52 +0100 Subject: [PATCH] extend dims to nD --- scripts/v2/integration_2d_3d/example_2d.py | 6 +- .../v2/integration_2d_3d/example_combined.py | 6 +- .../example_combined_camera_redraw.py | 6 +- .../example_cellier_2d.py | 4 +- .../example_v2.py | 8 +- src/cellier/v2/_state.py | 33 +- src/cellier/v2/controller.py | 81 ++-- src/cellier/v2/data/image/_image_requests.py | 40 +- .../v2/data/image/_zarr_multiscale_store.py | 145 +++----- src/cellier/v2/render/scene_manager.py | 13 +- src/cellier/v2/render/visuals/_image.py | 346 +++++++++++++----- src/cellier/v2/scene/dims.py | 99 ++++- tests/v2/events/test_event_bus.py | 9 +- tests/v2/render/test_render_manager.py | 50 +-- tests/v2/scene/test_dims.py | 94 ++++- tests/v2/scene/test_scene.py | 8 +- tests/v2/test_controller.py | 12 +- tests/v2/test_viewer_model.py | 8 +- 18 files changed, 654 insertions(+), 314 deletions(-) diff --git a/scripts/v2/integration_2d_3d/example_2d.py b/scripts/v2/integration_2d_3d/example_2d.py index 6325c7aa..e1566c04 100644 --- a/scripts/v2/integration_2d_3d/example_2d.py +++ b/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() diff --git a/scripts/v2/integration_2d_3d/example_combined.py b/scripts/v2/integration_2d_3d/example_combined.py index 7a4b308f..760cc31b 100644 --- a/scripts/v2/integration_2d_3d/example_combined.py +++ b/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() diff --git a/scripts/v2/integration_2d_3d/example_combined_camera_redraw.py b/scripts/v2/integration_2d_3d/example_combined_camera_redraw.py index aa8de59c..90fa8f1b 100644 --- a/scripts/v2/integration_2d_3d/example_combined_camera_redraw.py +++ b/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() diff --git a/scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_cellier_2d.py b/scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_cellier_2d.py index c7f2f7d7..d0bfc262 100644 --- a/scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_cellier_2d.py +++ b/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() diff --git a/scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_v2.py b/scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_v2.py index e8fdc285..5cf16e21 100644 --- a/scripts/v2/lut_texture_multiscale_frustum_async_cellier/example_v2.py +++ b/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 diff --git a/src/cellier/v2/_state.py b/src/cellier/v2/_state.py index 57949567..60119a17 100644 --- a/src/cellier/v2/_state.py +++ b/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): diff --git a/src/cellier/v2/controller.py b/src/cellier/v2/controller.py index 208f6162..615fa15b 100644 --- a/src/cellier/v2/controller.py +++ b/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,11 +500,7 @@ 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 @@ -499,21 +508,21 @@ def _dims_state_for_scene(self, scene_id: UUID) -> DimsState: 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)) diff --git a/src/cellier/v2/data/image/_image_requests.py b/src/cellier/v2/data/image/_image_requests.py index 668316fc..7b4b7f9f 100644 --- a/src/cellier/v2/data/image/_image_requests.py +++ b/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], ...] diff --git a/src/cellier/v2/data/image/_zarr_multiscale_store.py b/src/cellier/v2/data/image/_zarr_multiscale_store.py index da368e92..9055ed95 100644 --- a/src/cellier/v2/data/image/_zarr_multiscale_store.py +++ b/src/cellier/v2/data/image/_zarr_multiscale_store.py @@ -129,7 +129,7 @@ class MultiscaleZarrDataStore(BaseDataStore): n_levels : Number of scale levels (length of ``scale_names``). level_shapes : - List of ``(depth, height, width)`` tuples, one per level. + List of shape tuples, one per level. """ # ── Public pydantic fields ────────────────────────────────────────── @@ -165,30 +165,18 @@ def n_levels(self) -> int: return len(self._ts_stores) @property - def level_shapes(self) -> list[tuple[int, int, int]]: - """``(depth, height, width)`` for each scale level, finest first.""" - return [ - tuple(int(d) for d in store.domain.shape) # type: ignore[return-value] - for store in self._ts_stores - ] + def level_shapes(self) -> list[tuple[int, ...]]: + """Shape for each scale level, finest first.""" + return [tuple(int(d) for d in store.domain.shape) for store in self._ts_stores] # ── Async data access ─────────────────────────────────────────────── async def get_data(self, request: ChunkRequest) -> np.ndarray: """Read a single padded brick, returning a zero-padded float32 array. - Boundary bricks (where the padded region extends outside the store) - are handled by allocating a zeroed output of the full requested size - and copying the valid clamped region into the correct destination - offset. - - When ``request.z_slice`` is set, a 2D tile is read from a 3D store - by extracting a single z-plane. The returned array has shape - ``(y_stop-y_start, x_stop-x_start)``. - - The ``await`` on ``store[...].read()`` suspends the coroutine and - yields to the Qt event loop while tensorstore fetches the chunk(s) - from disk. + Interprets ``request.axis_selections`` generically: displayed axes + (tuple ranges) become slice dimensions in the output; sliced axes + (int values) become point selections. Parameters ---------- @@ -199,88 +187,53 @@ async def get_data(self, request: ChunkRequest) -> np.ndarray: Returns ------- out : - ``float32`` array. Shape is ``(dz, dy, dx)`` for 3D requests - or ``(dy, dx)`` for 2D requests (when ``z_slice`` is set). - Out-of-bounds regions are filled with zero. + ``float32`` array. Shape has one dimension per displayed axis + (those with tuple selections). Out-of-bounds regions are + filled with zero. """ - if request.z_slice is not None: - return await self._get_data_2d(request) - return await self._get_data_3d(request) + return await self._get_data_and(request) - async def _get_data_3d(self, request: ChunkRequest) -> np.ndarray: - """Read a 3D padded brick.""" + async def _get_data_and(self, request: ChunkRequest) -> np.ndarray: + """Read a padded brick using generic nD axis_selections.""" store = self._ts_stores[request.scale_index] - - dz = request.z_stop - request.z_start - dy = request.y_stop - request.y_start - dx = request.x_stop - request.x_start - out = np.zeros((dz, dy, dx), dtype=np.float32) - - sd, sh, sw = (int(d) for d in store.domain.shape) - sz0 = max(request.z_start, 0) - sz1 = min(request.z_stop, sd) - sy0 = max(request.y_start, 0) - sy1 = min(request.y_stop, sh) - sx0 = max(request.x_start, 0) - sx1 = min(request.x_stop, sw) - - if sz1 > sz0 and sy1 > sy0 and sx1 > sx0: - dest_z0 = sz0 - request.z_start - dest_y0 = sy0 - request.y_start - dest_x0 = sx0 - request.x_start - + store_shape = tuple(int(d) for d in store.domain.shape) + + # Output shape: one dimension per displayed (tuple) axis. + out_shape = tuple( + stop - start + for sel in request.axis_selections + if isinstance(sel, tuple) + for start, stop in [sel] + ) + out = np.zeros(out_shape, dtype=np.float32) + + # Build the clamped store index and track destination offsets. + store_idx: list[int | slice] = [] + dest_starts: list[int] = [] + valid = True + + for axis_i, sel in enumerate(request.axis_selections): + size = store_shape[axis_i] + if isinstance(sel, int): + # Point selection: clamp to valid range. + store_idx.append(max(0, min(sel, size - 1))) + else: + start, stop = sel + c_start = max(start, 0) + c_stop = min(stop, size) + if c_stop <= c_start: + valid = False + break + store_idx.append(slice(c_start, c_stop)) + dest_starts.append(c_start - start) + + if valid: region = np.asarray( - await store[sz0:sz1, sy0:sy1, sx0:sx1].read(), + await store[tuple(store_idx)].read(), dtype=np.float32, ) - out[ - dest_z0 : dest_z0 + (sz1 - sz0), - dest_y0 : dest_y0 + (sy1 - sy0), - dest_x0 : dest_x0 + (sx1 - sx0), - ] = region - - return out - - async def _get_data_2d(self, request: ChunkRequest) -> np.ndarray: - """Read a 2D padded tile (single z-plane from a 3D store).""" - store = self._ts_stores[request.scale_index] - ndim = len(store.domain.shape) - - dy = request.y_stop - request.y_start - dx = request.x_stop - request.x_start - out = np.zeros((dy, dx), dtype=np.float32) - - if ndim == 3: - sd, sh, sw = (int(d) for d in store.domain.shape) - elif ndim == 2: - sh, sw = (int(d) for d in store.domain.shape) - else: - raise ValueError(f"Unexpected store ndim={ndim}") - - sy0 = max(request.y_start, 0) - sy1 = min(request.y_stop, sh) - sx0 = max(request.x_start, 0) - sx1 = min(request.x_stop, sw) - - if sy1 > sy0 and sx1 > sx0: - dest_y0 = sy0 - request.y_start - dest_x0 = sx0 - request.x_start - - if ndim == 3: - z_mid = request.z_slice - region = np.asarray( - await store[z_mid, sy0:sy1, sx0:sx1].read(), - dtype=np.float32, - ) - else: - region = np.asarray( - await store[sy0:sy1, sx0:sx1].read(), - dtype=np.float32, - ) - - out[ - dest_y0 : dest_y0 + (sy1 - sy0), - dest_x0 : dest_x0 + (sx1 - sx0), - ] = region + # Compute the destination slice in out for each displayed axis. + dest_idx = tuple(slice(d, d + s) for d, s in zip(dest_starts, region.shape)) + out[dest_idx] = region return out diff --git a/src/cellier/v2/render/scene_manager.py b/src/cellier/v2/render/scene_manager.py index 1c31090a..15e60c51 100644 --- a/src/cellier/v2/render/scene_manager.py +++ b/src/cellier/v2/render/scene_manager.py @@ -177,13 +177,14 @@ def _build_slice_requests_3d( else None ) - n_levels = visual._volume_geometry.n_levels + n_levels = visual.n_levels thresholds = self._compute_thresholds_3d(request, n_levels, cfg.lod_bias) chunk_requests = visual.build_slice_request( camera_pos=request.camera_pos, frustum_planes=frustum_planes, thresholds=thresholds, + dims_state=request.dims_state, force_level=cfg.force_level, ) if chunk_requests: @@ -199,14 +200,6 @@ def _build_slice_requests_2d( """2D planning path using orthographic camera and viewport culling.""" result: dict[UUID, list[ChunkRequest]] = {} - # Extract z-slice from dims_state. - # For 2D scenes, slice_indices contains the z-slice index. - z_slice = ( - request.dims_state.slice_indices[0] - if request.dims_state.slice_indices - else 0 - ) - world_width, world_height = request.world_extent viewport_width_px, viewport_height_px = request.screen_size_px @@ -233,7 +226,7 @@ def _build_slice_requests_2d( world_width=world_width, view_min=view_min if cfg.frustum_cull else None, view_max=view_max if cfg.frustum_cull else None, - z_slice=z_slice, + dims_state=request.dims_state, lod_bias=cfg.lod_bias, force_level=cfg.force_level, use_culling=cfg.frustum_cull, diff --git a/src/cellier/v2/render/visuals/_image.py b/src/cellier/v2/render/visuals/_image.py index 777a0f46..0256820b 100644 --- a/src/cellier/v2/render/visuals/_image.py +++ b/src/cellier/v2/render/visuals/_image.py @@ -53,6 +53,7 @@ ) if TYPE_CHECKING: + from cellier.v2._state import AxisAlignedSelectionState, DimsState from cellier.v2.events._events import ( AppearanceChangedEvent, DataStoreContentsChangedEvent, @@ -85,8 +86,10 @@ class VolumeGeometry: Parameters ---------- - level_shapes : list[tuple[int, int, int]] - Volume shape ``(D, H, W)`` at each scale level, finest first. + level_shapes : list[tuple[int, ...]] + Displayed-axis shape at each scale level, finest first. + For 3D rendering this is ``(D, H, W)``; for 2D it is ``(H, W)``. + The caller extracts the displayed dimensions before passing. downscale_factors : list[int] Downscale factor for each level relative to the original data, e.g. ``[1, 2, 4]``. Must have the same length as @@ -97,7 +100,7 @@ class VolumeGeometry: def __init__( self, - level_shapes: list[tuple[int, int, int]], + level_shapes: list[tuple[int, ...]], downscale_factors: list[int], block_size: int, ) -> None: @@ -110,7 +113,7 @@ def __init__( def from_cellier_model( cls, model: MultiscaleImageVisual, - level_shapes: list[tuple[int, int, int]], + level_shapes: list[tuple[int, ...]], block_size: int, ) -> VolumeGeometry: """Build a ``VolumeGeometry`` from a ``MultiscaleImageVisual`` model. @@ -119,9 +122,8 @@ def from_cellier_model( ---------- model : MultiscaleImageVisual The visual model. Only ``downscale_factors`` is read. - level_shapes : list[tuple[int, int, int]] - Volume shape ``(D, H, W)`` per level, finest first. Must be - obtained from the data store by the caller. + level_shapes : list[tuple[int, ...]] + Displayed-axis shape per level, finest first. block_size : int Rendering brick side length in voxels. @@ -135,7 +137,7 @@ def from_cellier_model( block_size=block_size, ) - def _rebuild(self, level_shapes: list[tuple[int, int, int]]) -> None: + def _rebuild(self, level_shapes: list[tuple[int, ...]]) -> None: self.level_shapes = list(level_shapes) self.layouts = [ BlockLayout3D(volume_shape=shape, block_size=self.block_size) @@ -144,7 +146,7 @@ def _rebuild(self, level_shapes: list[tuple[int, int, int]]) -> None: self.base_layout = self.layouts[0] self._level_grids = build_level_grids(self.base_layout, self.n_levels) - def update(self, level_shapes: list[tuple[int, int, int]]) -> None: + def update(self, level_shapes: list[tuple[int, ...]]) -> None: """Rebuild from new level shapes after a DataStoreMutated event.""" self._rebuild(level_shapes) @@ -161,9 +163,9 @@ class ImageGeometry2D: Parameters ---------- - level_shapes : list[tuple[int, int, int]] - Volume shape ``(D, H, W)`` at each scale level, finest first. - The 2D image shape is ``(H, W)`` from each level. + level_shapes : list[tuple[int, int]] + Image shape ``(H, W)`` at each scale level, finest first. + The caller extracts the two displayed dimensions before passing. block_size : int Tile side length in pixels. n_levels : int @@ -172,7 +174,7 @@ class ImageGeometry2D: def __init__( self, - level_shapes: list[tuple[int, int, int]], + level_shapes: list[tuple[int, int]], block_size: int, n_levels: int, ) -> None: @@ -182,11 +184,20 @@ def __init__( # Build 2D base layout from finest level (H, W). self.base_layout = BlockLayout2D.from_shape( - shape=(level_shapes[0][1], level_shapes[0][2]), + shape=(level_shapes[0][0], level_shapes[0][1]), block_size=block_size, ) self._level_grids = build_tile_grids_2d(self.base_layout, n_levels) + def update(self, level_shapes: list[tuple[int, int]]) -> None: + """Rebuild from new level shapes after displayed axes change.""" + self.level_shapes = list(level_shapes) + self.base_layout = BlockLayout2D.from_shape( + shape=(level_shapes[0][0], level_shapes[0][1]), + block_size=self.block_size, + ) + self._level_grids = build_tile_grids_2d(self.base_layout, self.n_levels) + # --------------------------------------------------------------------------- # Coordinate helpers @@ -210,6 +221,28 @@ def _brick_key_to_padded_coords( return z0, y0, x0, z0 + padded, y0 + padded, x0 + padded +def _build_axis_selections( + sel: AxisAlignedSelectionState, + ndim: int, + display_coords: list[tuple[int, int]], + scale: int, +) -> tuple[int | tuple[int, int], ...]: + """Map brick/tile window coords and slice indices onto the full nD axis list. + + ``display_coords`` are in the order of ``sel.displayed_axes``. + Non-displayed axes get their slice index divided by ``scale`` + (coarser levels correspond to smaller indices). + """ + display_pos = {ax: i for i, ax in enumerate(sel.displayed_axes)} + result: list[int | tuple[int, int]] = [] + for data_axis in range(ndim): + if data_axis in display_pos: + result.append(display_coords[display_pos[data_axis]]) + else: + result.append(sel.slice_indices[data_axis] // scale) + return tuple(result) + + def _block_key_2d_to_padded_coords( key: BlockKey2D, block_size: int, @@ -262,7 +295,8 @@ class GFXMultiscaleImageVisual: def __init__( self, visual_model_id: UUID, - volume_geometry: VolumeGeometry, + volume_geometry: VolumeGeometry | None, + image_geometry_2d: ImageGeometry2D | None, render_modes: set[str], colormap: gfx.TextureMap | None = None, clim: tuple[float, float] = (0.0, 1.0), @@ -273,52 +307,58 @@ def __init__( self.visual_model_id = visual_model_id self.render_modes = render_modes self._volume_geometry = volume_geometry + self._image_geometry_2d = image_geometry_2d + self._gpu_budget_bytes = gpu_budget_bytes self._frame_number = 0 self._pending_slot_map: dict[UUID, tuple[BlockKey3D, TileSlot]] = {} self._pending_slot_map_2d: dict[UUID, tuple[BlockKey2D, TileSlot2D]] = {} self._last_plan_stats: dict = {} - # ── 3D GPU resources ─────────────────────────────────────────── - cache_parameters_3d = compute_block_cache_parameters_3d( - block_size=volume_geometry.block_size, - gpu_budget_bytes=gpu_budget_bytes, - ) - self._block_cache_3d = BlockCache3D(cache_parameters=cache_parameters_3d) - self._lut_manager_3d = LutIndirectionManager3D( - base_layout=volume_geometry.base_layout, - n_levels=volume_geometry.n_levels, - ) - self._lut_params_buffer_3d = build_lut_params_buffer_3d( - volume_geometry.base_layout, - cache_parameters_3d, - proxy_voxels_per_brick=volume_geometry.block_size, - ) - self._block_scales_buffer_3d = build_block_scales_buffer_3d( - volume_geometry.downscale_factors - ) - - # ── 2D GPU resources ─────────────────────────────────────────── - cache_parameters_2d = compute_block_cache_parameters_2d( - gpu_budget_bytes=gpu_budget_bytes, block_size=volume_geometry.block_size - ) - self._block_cache_2d = BlockCache2D(cache_parameters=cache_parameters_2d) - - self._image_geometry_2d = ImageGeometry2D( - level_shapes=volume_geometry.level_shapes, - block_size=volume_geometry.block_size, - n_levels=volume_geometry.n_levels, - ) + # ── 3D GPU resources (only when volume_geometry is provided) ─── + self._block_cache_3d: BlockCache3D | None = None + self._lut_manager_3d: LutIndirectionManager3D | None = None + self._lut_params_buffer_3d = None + self._block_scales_buffer_3d = None + if volume_geometry is not None: + cache_parameters_3d = compute_block_cache_parameters_3d( + block_size=volume_geometry.block_size, + gpu_budget_bytes=gpu_budget_bytes, + ) + self._block_cache_3d = BlockCache3D(cache_parameters=cache_parameters_3d) + self._lut_manager_3d = LutIndirectionManager3D( + base_layout=volume_geometry.base_layout, + n_levels=volume_geometry.n_levels, + ) + self._lut_params_buffer_3d = build_lut_params_buffer_3d( + volume_geometry.base_layout, + cache_parameters_3d, + proxy_voxels_per_brick=volume_geometry.block_size, + ) + self._block_scales_buffer_3d = build_block_scales_buffer_3d( + volume_geometry.downscale_factors + ) - self._lut_manager_2d = LutIndirectionManager2D( - base_layout=self._image_geometry_2d.base_layout, - n_levels=volume_geometry.n_levels, - ) - self._lut_params_buffer_2d = build_lut_params_buffer_2d( - self._image_geometry_2d.base_layout, cache_parameters_2d - ) - self._block_scales_buffer_2d = build_block_scales_buffer_2d( - volume_geometry.n_levels - ) + # ── 2D GPU resources (only when image_geometry_2d is provided) ─ + self._block_cache_2d: BlockCache2D | None = None + self._lut_manager_2d: LutIndirectionManager2D | None = None + self._lut_params_buffer_2d = None + self._block_scales_buffer_2d = None + if image_geometry_2d is not None: + cache_parameters_2d = compute_block_cache_parameters_2d( + gpu_budget_bytes=gpu_budget_bytes, + block_size=image_geometry_2d.block_size, + ) + self._block_cache_2d = BlockCache2D(cache_parameters=cache_parameters_2d) + self._lut_manager_2d = LutIndirectionManager2D( + base_layout=image_geometry_2d.base_layout, + n_levels=image_geometry_2d.n_levels, + ) + self._lut_params_buffer_2d = build_lut_params_buffer_2d( + image_geometry_2d.base_layout, cache_parameters_2d + ) + self._block_scales_buffer_2d = build_block_scales_buffer_2d( + image_geometry_2d.n_levels + ) if colormap is None: colormap = gfx.cm.viridis @@ -327,7 +367,7 @@ def __init__( self.node_3d: gfx.Volume | None = None self.material_3d: VolumeBlockMaterial | None = None self._proxy_tex_3d: gfx.Texture | None = None - if "3d" in render_modes: + if "3d" in render_modes and volume_geometry is not None: self.node_3d, self.material_3d, self._proxy_tex_3d = self._build_3d_node( colormap=colormap, clim=clim, @@ -339,7 +379,7 @@ def __init__( self.node_2d: gfx.Image | None = None self.material_2d: ImageBlockMaterial | None = None self._proxy_tex_2d: gfx.Texture | None = None - if "2d" in render_modes: + if "2d" in render_modes and image_geometry_2d is not None: self.node_2d, self.material_2d, self._proxy_tex_2d = self._build_2d_node( colormap=colormap, clim=clim, interpolation=interpolation ) @@ -348,7 +388,8 @@ def __init__( def from_cellier_model( cls, model: MultiscaleImageVisual, - level_shapes: list[tuple[int, int, int]], + level_shapes: list[tuple[int, ...]], + displayed_axes: tuple[int, ...], render_modes: set[str], block_size: int = 32, gpu_budget_bytes: int = 1 * 1024**3, @@ -361,8 +402,10 @@ def from_cellier_model( ---------- model : MultiscaleImageVisual Source visual model. - level_shapes : list[tuple[int, int, int]] - Volume shape ``(D, H, W)`` per level, finest first. + level_shapes : list[tuple[int, ...]] + Full nD shape per level, finest first. + displayed_axes : tuple[int, ...] + Which axes are displayed (from dims.selection.displayed_axes). render_modes : set[str] Which nodes to build: ``{"3d"}``, ``{"2d"}``, or ``{"2d", "3d"}``. @@ -379,15 +422,35 @@ def from_cellier_model( ------- GFXMultiscaleImageVisual """ - volume_geometry = VolumeGeometry.from_cellier_model( - model, level_shapes, block_size - ) + # Extract the displayed-axis shapes. + displayed_level_shapes = [ + tuple(shape[ax] for ax in displayed_axes) for shape in level_shapes + ] + + # Build 3D geometry only when 3D rendering is requested. + volume_geometry: VolumeGeometry | None = None + if "3d" in render_modes and len(displayed_axes) == 3: + volume_geometry = VolumeGeometry.from_cellier_model( + model, displayed_level_shapes, block_size + ) + + # Build 2D geometry only when 2D rendering is requested. + image_geometry_2d: ImageGeometry2D | None = None + if "2d" in render_modes and len(displayed_axes) == 2: + level_shapes_2d = [(s[0], s[1]) for s in displayed_level_shapes] + image_geometry_2d = ImageGeometry2D( + level_shapes=level_shapes_2d, + block_size=block_size, + n_levels=len(level_shapes), + ) + colormap = model.appearance.color_map.to_pygfx(N=256) clim = model.appearance.clim return cls( visual_model_id=model.id, volume_geometry=volume_geometry, + image_geometry_2d=image_geometry_2d, render_modes=render_modes, colormap=colormap, clim=clim, @@ -396,6 +459,123 @@ def from_cellier_model( gpu_budget_bytes=gpu_budget_bytes, ) + # ── Properties ───────────────────────────────────────────────────── + + @property + def n_levels(self) -> int: + """Number of LOD levels.""" + if self._volume_geometry is not None: + return self._volume_geometry.n_levels + if self._image_geometry_2d is not None: + return self._image_geometry_2d.n_levels + raise RuntimeError("No geometry available") + + # ── Geometry rebuild ───────────────────────────────────────────── + + def rebuild_geometry( + self, + level_shapes: list[tuple[int, ...]], + displayed_axes: tuple[int, ...], + ) -> tuple[gfx.WorldObject | None, gfx.WorldObject | None]: + """Rebuild geometry and GPU resources after a displayed_axes change. + + Returns ``(old_node, new_node)`` for the active rendering mode so + the caller can swap the node in the scene graph. + + Parameters + ---------- + level_shapes : list[tuple[int, ...]] + Full nD shape per level from the data store. + displayed_axes : tuple[int, ...] + The new set of displayed axes. + + Returns + ------- + tuple[old_node, new_node] + The previous and replacement pygfx nodes (may be ``None``). + """ + displayed_level_shapes = [ + tuple(shape[ax] for ax in displayed_axes) for shape in level_shapes + ] + + old_node: gfx.WorldObject | None = None + new_node: gfx.WorldObject | None = None + + if "3d" in self.render_modes and len(displayed_axes) == 3: + old_node = self.node_3d + if self._volume_geometry is not None: + self._volume_geometry.update(displayed_level_shapes) + self._rebuild_3d_resources() + new_node = self.node_3d + + if "2d" in self.render_modes and len(displayed_axes) == 2: + old_node = self.node_2d + shapes_2d = [(s[0], s[1]) for s in displayed_level_shapes] + if self._image_geometry_2d is not None: + self._image_geometry_2d.update(shapes_2d) + self._rebuild_2d_resources() + new_node = self.node_2d + + return old_node, new_node + + def _rebuild_3d_resources(self) -> None: + """Rebuild 3D GPU resources after geometry update.""" + geo = self._volume_geometry + # Clear cache + self._block_cache_3d.tile_manager.release_all_in_flight() + # Rebuild LUT manager + self._lut_manager_3d = LutIndirectionManager3D( + base_layout=geo.base_layout, + n_levels=geo.n_levels, + ) + # Rebuild param buffers + self._lut_params_buffer_3d = build_lut_params_buffer_3d( + geo.base_layout, + self._block_cache_3d.info, + proxy_voxels_per_brick=geo.block_size, + ) + self._block_scales_buffer_3d = build_block_scales_buffer_3d( + geo.downscale_factors + ) + # Rebuild node preserving current appearance + if self.node_3d is not None: + colormap = self.material_3d.map + clim = self.material_3d.clim + threshold = self.material_3d.threshold + interpolation = self.material_3d.interpolation + self.node_3d, self.material_3d, self._proxy_tex_3d = self._build_3d_node( + colormap=colormap, + clim=clim, + threshold=threshold, + interpolation=interpolation, + ) + self._pending_slot_map = {} + + def _rebuild_2d_resources(self) -> None: + """Rebuild 2D GPU resources after geometry update.""" + geo2d = self._image_geometry_2d + # Clear cache + self._block_cache_2d.tile_manager.release_all_in_flight() + # Rebuild LUT manager + self._lut_manager_2d = LutIndirectionManager2D( + base_layout=geo2d.base_layout, + n_levels=geo2d.n_levels, + ) + # Rebuild param buffers + self._lut_params_buffer_2d = build_lut_params_buffer_2d( + geo2d.base_layout, self._block_cache_2d.info + ) + self._block_scales_buffer_2d = build_block_scales_buffer_2d(geo2d.n_levels) + # Rebuild node preserving current appearance + if self.node_2d is not None: + colormap = self.material_2d.map + clim = self.material_2d.clim + interpolation = self.material_2d.interpolation + self.node_2d, self.material_2d, self._proxy_tex_2d = self._build_2d_node( + colormap=colormap, clim=clim, interpolation=interpolation + ) + self._pending_slot_map_2d = {} + # ── 3D SliceCoordinator interface ────────────────────────────────── def build_slice_request( @@ -403,6 +583,7 @@ def build_slice_request( camera_pos: np.ndarray, frustum_planes: np.ndarray | None, thresholds: list[float] | None, + dims_state: DimsState | None = None, force_level: int | None = None, ) -> list[ChunkRequest]: """Run the synchronous 3D planning phase and return ChunkRequests. @@ -474,16 +655,20 @@ def build_slice_request( z0, y0, x0, z1, y1, x1 = _brick_key_to_padded_coords( brick_key, geo.block_size, self._block_cache_3d.info.overlap ) + scale = 2 ** (brick_key.level - 1) + display_coords = [(z0, z1), (y0, y1), (x0, x1)] + if dims_state is not None: + ndim = len(dims_state.axis_labels) + axis_selections = _build_axis_selections( + dims_state.selection, ndim, display_coords, scale + ) + else: + axis_selections = tuple(display_coords) req = ChunkRequest( chunk_request_id=chunk_id, slice_request_id=slice_id, scale_index=brick_key.level - 1, - z_start=z0, - y_start=y0, - x_start=x0, - z_stop=z1, - y_stop=y1, - x_stop=x1, + axis_selections=axis_selections, ) chunk_requests.append(req) self._pending_slot_map[chunk_id] = (brick_key, slot) @@ -583,7 +768,7 @@ def build_slice_request_2d( world_width: float, view_min: np.ndarray | None, view_max: np.ndarray | None, - z_slice: int, + dims_state: DimsState, lod_bias: float = 1.0, force_level: int | None = None, use_culling: bool = True, @@ -605,8 +790,8 @@ def build_slice_request_2d( Viewport AABB minimum ``(x, y)``. ``None`` disables culling. view_max : ndarray, shape (2,) or None Viewport AABB maximum ``(x, y)``. ``None`` disables culling. - z_slice : int - Z-slice index at the finest scale level. + dims_state : DimsState + Current dimension display state. lod_bias : float Multiplicative LOD bias. force_level : int or None @@ -678,26 +863,23 @@ def build_slice_request_2d( overlap = self._block_cache_2d.info.overlap + ndim = len(dims_state.axis_labels) + sel = dims_state.selection + for tile_key, slot in fill_plan: chunk_id = uuid4() y0, x0, y1, x1 = _block_key_2d_to_padded_coords( tile_key, block_size, overlap ) - # Compute z-slice at this scale level. scale = 2 ** (tile_key.level - 1) - z_slice_at_level = z_slice // scale + display_coords = [(y0, y1), (x0, x1)] + axis_selections = _build_axis_selections(sel, ndim, display_coords, scale) req = ChunkRequest( chunk_request_id=chunk_id, slice_request_id=slice_id, scale_index=tile_key.level - 1, - z_start=0, - y_start=y0, - x_start=x0, - z_stop=0, - y_stop=y1, - x_stop=x1, - z_slice=z_slice_at_level, + axis_selections=axis_selections, ) chunk_requests.append(req) self._pending_slot_map_2d[chunk_id] = (tile_key, slot) @@ -888,7 +1070,7 @@ def _build_2d_node( geometry = gfx.Geometry(grid=proxy_tex) image = gfx.Image(geometry, material) - bs = self._volume_geometry.block_size + bs = self._image_geometry_2d.block_size image.local.scale = (bs, bs, 1.0) image.local.position = (bs * 0.5, bs * 0.5, 0) diff --git a/src/cellier/v2/scene/dims.py b/src/cellier/v2/scene/dims.py index b9d68a24..7fa393e8 100644 --- a/src/cellier/v2/scene/dims.py +++ b/src/cellier/v2/scene/dims.py @@ -3,11 +3,17 @@ from __future__ import annotations import uuid -from typing import Annotated, Any +from typing import Annotated, Any, Literal, Union from uuid import uuid4 from psygnal import EventedModel -from pydantic import UUID4, AfterValidator, Field +from pydantic import UUID4, AfterValidator, Field, model_validator + +from cellier.v2._state import ( + AxisAlignedSelectionState, + DimsState, + PlaneSelectionState, +) class CoordinateSystem(EventedModel): @@ -25,6 +31,54 @@ class CoordinateSystem(EventedModel): axis_labels: tuple[str, ...] +class AxisAlignedSelection(EventedModel): + """Mutable selection model for axis-aligned slicing. + + Parameters + ---------- + selector_type : Literal["axis_aligned"] + Discriminator field. + displayed_axes : tuple[int, ...] + Indices into the coordinate system that are rendered. + Length 2 → 2D; length 3 → 3D. + slice_indices : dict[int, int] + Mapping of axis index → slice value for non-displayed axes. + """ + + selector_type: Literal["axis_aligned"] = "axis_aligned" + displayed_axes: tuple[int, ...] + slice_indices: dict[int, int] = Field(default_factory=dict) + + def to_state(self) -> AxisAlignedSelectionState: + """Return an immutable snapshot of this selection.""" + return AxisAlignedSelectionState( + displayed_axes=self.displayed_axes, + slice_indices=dict(self.slice_indices), + ) + + +class PlaneSelection(EventedModel): + """Stub — not yet implemented. + + Parameters + ---------- + selector_type : Literal["plane"] + Discriminator field. + """ + + selector_type: Literal["plane"] = "plane" + + def to_state(self) -> PlaneSelectionState: + """Return an immutable snapshot of this selection.""" + raise NotImplementedError("PlaneSelection is not yet implemented.") + + +SelectionType = Annotated[ + Union[AxisAlignedSelection, PlaneSelection], + Field(discriminator="selector_type"), +] + + class DimsManager(EventedModel): """Tracks which axes are displayed and the slice index for non-displayed axes. @@ -36,25 +90,46 @@ class DimsManager(EventedModel): Unique identifier. Auto-generated. coordinate_system : CoordinateSystem Defines the world axis names. - displayed_axes : tuple[int, ...] - Indices into ``coordinate_system.axis_labels`` currently rendered. - Length 2 → 2D; length 3 → 3D. - slice_indices : tuple[int, ...] - One integer index per non-displayed axis, in the same axis order as - ``coordinate_system``. Length must equal - ``len(axis_labels) - len(displayed_axes)``. + selection : SelectionType + The current axis selection (axis-aligned or plane). """ id: UUID4 | Annotated[str, AfterValidator(lambda x: uuid.UUID(x, version=4))] = ( Field(frozen=True, default_factory=lambda: uuid4()) ) coordinate_system: CoordinateSystem - displayed_axes: tuple[int, ...] - slice_indices: tuple[int, ...] + selection: SelectionType + + @model_validator(mode="after") + def _validate_axis_coverage(self) -> DimsManager: + """Verify displayed + sliced axes cover all coordinate axes.""" + if isinstance(self.selection, AxisAlignedSelection): + ndim = len(self.coordinate_system.axis_labels) + covered = set(self.selection.displayed_axes) | set( + self.selection.slice_indices.keys() + ) + expected = set(range(ndim)) + if covered != expected: + raise ValueError( + f"Axis coverage mismatch: displayed_axes | slice_indices.keys() " + f"= {covered}, expected {expected} for ndim={ndim}" + ) + return self def model_post_init(self, __context: Any) -> None: - """Wire coordinate system event relay after model initialization.""" + """Wire event relays after model initialization.""" self.coordinate_system.events.all.connect(self._on_coordinate_system_updated) + self.selection.events.all.connect(self._on_selection_updated) def _on_coordinate_system_updated(self, info: Any) -> None: self.events.coordinate_system.emit(self.coordinate_system) + + def _on_selection_updated(self, info: Any) -> None: + self.events.selection.emit(self.selection) + + def to_state(self) -> DimsState: + """Return an immutable snapshot of the current dims state.""" + return DimsState( + axis_labels=self.coordinate_system.axis_labels, + selection=self.selection.to_state(), + ) diff --git a/tests/v2/events/test_event_bus.py b/tests/v2/events/test_event_bus.py index 302a4ee5..aefc34bc 100644 --- a/tests/v2/events/test_event_bus.py +++ b/tests/v2/events/test_event_bus.py @@ -7,7 +7,7 @@ import pytest -from cellier.v2._state import DimsState +from cellier.v2._state import AxisAlignedSelectionState, DimsState from cellier.v2.events import ( DimsChangedEvent, EventBus, @@ -22,7 +22,12 @@ def _make_dims_event(scene_id=None) -> DimsChangedEvent: return DimsChangedEvent( source_id=uuid4(), scene_id=scene_id or uuid4(), - 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={} + ), + ), displayed_axes_changed=False, ) diff --git a/tests/v2/render/test_render_manager.py b/tests/v2/render/test_render_manager.py index a443dead..5d07ba59 100644 --- a/tests/v2/render/test_render_manager.py +++ b/tests/v2/render/test_render_manager.py @@ -12,6 +12,7 @@ import numpy as np import pytest +from cellier.v2._state import AxisAlignedSelectionState from cellier.v2.data.image import ChunkRequest from cellier.v2.render._requests import DimsState, ReslicingRequest from cellier.v2.render._scene_config import VisualRenderConfig @@ -24,7 +25,13 @@ def _make_dims_state() -> DimsState: - return DimsState(displayed_axes=(0, 1, 2), slice_indices=()) + return DimsState( + axis_labels=("z", "y", "x"), + selection=AxisAlignedSelectionState( + displayed_axes=(0, 1, 2), + slice_indices={}, + ), + ) def _make_reslicing_request( @@ -92,15 +99,23 @@ def cancel(self, slice_request_id) -> bool: def test_dims_state_construction() -> None: - ds = DimsState(displayed_axes=(0, 1, 2), slice_indices=()) - assert ds.displayed_axes == (0, 1, 2) - assert ds.slice_indices == () + ds = DimsState( + axis_labels=("z", "y", "x"), + selection=AxisAlignedSelectionState(displayed_axes=(0, 1, 2), slice_indices={}), + ) + assert ds.selection.displayed_axes == (0, 1, 2) + assert ds.selection.slice_indices == {} def test_dims_state_2d() -> None: - ds = DimsState(displayed_axes=(1, 2), slice_indices=(5,)) - assert ds.displayed_axes == (1, 2) - assert ds.slice_indices == (5,) + ds = DimsState( + axis_labels=("z", "y", "x"), + selection=AxisAlignedSelectionState( + displayed_axes=(1, 2), slice_indices={0: 5} + ), + ) + assert ds.selection.displayed_axes == (1, 2) + assert ds.selection.slice_indices == {0: 5} def test_reslicing_request_construction() -> None: @@ -166,12 +181,7 @@ def test_scene_manager_build_slice_requests_all_visuals() -> None: chunk_request_id=uuid4(), slice_request_id=uuid4(), scale_index=0, - z_start=0, - y_start=0, - x_start=0, - z_stop=32, - y_stop=32, - x_stop=32, + axis_selections=((0, 32), (0, 32), (0, 32)), ) v1 = _make_mock_visual() v1.build_slice_request.return_value = [chunk] @@ -202,12 +212,7 @@ def test_scene_manager_build_slice_requests_target_filtering() -> None: chunk_request_id=uuid4(), slice_request_id=uuid4(), scale_index=0, - z_start=0, - y_start=0, - x_start=0, - z_stop=32, - y_stop=32, - x_stop=32, + axis_selections=((0, 32), (0, 32), (0, 32)), ) v1 = _make_mock_visual() v1.build_slice_request.return_value = [chunk] @@ -271,12 +276,7 @@ def _make_coordinator_with_scene( chunk_request_id=uuid4(), slice_request_id=uuid4(), scale_index=0, - z_start=0, - y_start=0, - x_start=0, - z_stop=32, - y_stop=32, - x_stop=32, + axis_selections=((0, 32), (0, 32), (0, 32)), ) visuals = [] diff --git a/tests/v2/scene/test_dims.py b/tests/v2/scene/test_dims.py index 6595059f..96ce86cf 100644 --- a/tests/v2/scene/test_dims.py +++ b/tests/v2/scene/test_dims.py @@ -1,6 +1,13 @@ -"""Tests for CoordinateSystem and DimsManager models.""" +"""Tests for CoordinateSystem, DimsManager, and AxisAlignedSelection models.""" -from cellier.v2.scene.dims import CoordinateSystem, DimsManager +import pytest + +from cellier.v2._state import AxisAlignedSelectionState +from cellier.v2.scene.dims import ( + AxisAlignedSelection, + CoordinateSystem, + DimsManager, +) def test_coordinate_system_roundtrip(tmp_path): @@ -15,8 +22,10 @@ def test_dims_manager_roundtrip(tmp_path): # 3D case — all axes displayed, no slice indices original_3d = DimsManager( coordinate_system=CoordinateSystem(name="world", axis_labels=("z", "y", "x")), - displayed_axes=(0, 1, 2), - slice_indices=(), + selection=AxisAlignedSelection( + displayed_axes=(0, 1, 2), + slice_indices={}, + ), ) path = tmp_path / "dims_3d.json" path.write_text(original_3d.model_dump_json()) @@ -26,10 +35,83 @@ def test_dims_manager_roundtrip(tmp_path): # 2D slice through 3D volume original_2d = DimsManager( coordinate_system=CoordinateSystem(name="world", axis_labels=("z", "y", "x")), - displayed_axes=(1, 2), - slice_indices=(32,), + selection=AxisAlignedSelection( + displayed_axes=(1, 2), + slice_indices={0: 32}, + ), ) path2 = tmp_path / "dims_2d.json" path2.write_text(original_2d.model_dump_json()) deserialized_2d = DimsManager.model_validate_json(path2.read_text()) assert original_2d.model_dump_json() == deserialized_2d.model_dump_json() + + +def test_axis_aligned_selection_to_state(): + sel = AxisAlignedSelection( + displayed_axes=(1, 2), + slice_indices={0: 42}, + ) + state = sel.to_state() + assert isinstance(state, AxisAlignedSelectionState) + assert state.displayed_axes == (1, 2) + assert state.slice_indices == {0: 42} + + +def test_to_index_selection_3d(): + """3D data, all axes displayed, empty slice_indices.""" + state = AxisAlignedSelectionState( + displayed_axes=(0, 1, 2), + slice_indices={}, + ) + result = state.to_index_selection(ndim=3) + assert result == (slice(None), slice(None), slice(None)) + + +def test_to_index_selection_2d_from_3d(): + """3D data, 2 axes displayed, one sliced.""" + state = AxisAlignedSelectionState( + displayed_axes=(1, 2), + slice_indices={0: 42}, + ) + result = state.to_index_selection(ndim=3) + assert result == (42, slice(None), slice(None)) + + +def test_to_index_selection_5d(): + """5D data, 3 displayed axes, 2 sliced.""" + state = AxisAlignedSelectionState( + displayed_axes=(2, 3, 4), + slice_indices={0: 5, 1: 1}, + ) + result = state.to_index_selection(ndim=5) + assert result == (5, 1, slice(None), slice(None), slice(None)) + + +def test_dims_manager_validates_axis_coverage(): + """Mismatched axes should raise ValidationError.""" + with pytest.raises(ValueError, match="Axis coverage mismatch"): + DimsManager( + coordinate_system=CoordinateSystem( + name="world", axis_labels=("z", "y", "x") + ), + selection=AxisAlignedSelection( + displayed_axes=(0, 1), + slice_indices={}, # Missing axis 2 + ), + ) + + +def test_dims_manager_to_state(): + dims = DimsManager( + coordinate_system=CoordinateSystem( + name="world", axis_labels=("t", "c", "z", "y", "x") + ), + selection=AxisAlignedSelection( + displayed_axes=(2, 3, 4), + slice_indices={0: 5, 1: 1}, + ), + ) + state = dims.to_state() + assert state.axis_labels == ("t", "c", "z", "y", "x") + assert state.selection.displayed_axes == (2, 3, 4) + assert state.selection.slice_indices == {0: 5, 1: 1} diff --git a/tests/v2/scene/test_scene.py b/tests/v2/scene/test_scene.py index 8e2e15b8..97355b55 100644 --- a/tests/v2/scene/test_scene.py +++ b/tests/v2/scene/test_scene.py @@ -4,7 +4,7 @@ from cellier.v2.scene.cameras import OrbitCameraController, 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.visuals._image import ImageAppearance, MultiscaleImageVisual @@ -12,8 +12,10 @@ def test_scene_roundtrip(tmp_path): dims = DimsManager( coordinate_system=CoordinateSystem(name="world", axis_labels=("z", "y", "x")), - displayed_axes=(0, 1, 2), - slice_indices=(), + selection=AxisAlignedSelection( + displayed_axes=(0, 1, 2), + slice_indices={}, + ), ) visual = MultiscaleImageVisual( name="volume", diff --git a/tests/v2/test_controller.py b/tests/v2/test_controller.py index 70474397..be159cb4 100644 --- a/tests/v2/test_controller.py +++ b/tests/v2/test_controller.py @@ -171,9 +171,9 @@ def test_dims_bridge_emits_event(): fired = [] controller.on_dims_changed(scene.id, fired.append) - controller.get_scene(scene.id).dims.slice_indices = (5,) + controller.get_scene(scene.id).dims.selection.slice_indices = {0: 5} assert len(fired) == 1 - assert fired[0].slice_indices == (5,) + assert fired[0].selection.slice_indices == {0: 5} def test_dims_bridge_displayed_axes_flag_true(): @@ -188,7 +188,7 @@ def test_dims_bridge_displayed_axes_flag_true(): entity_id=scene.id, ) - controller.get_scene(scene.id).dims.displayed_axes = (0, 1) + controller.get_scene(scene.id).dims.selection.displayed_axes = (0, 1) assert len(events) == 1 assert events[0].displayed_axes_changed is True @@ -205,7 +205,7 @@ def test_dims_bridge_displayed_axes_flag_false(): entity_id=scene.id, ) - controller.get_scene(scene.id).dims.slice_indices = (3,) + controller.get_scene(scene.id).dims.selection.slice_indices = {0: 3} assert len(events) == 1 assert events[0].displayed_axes_changed is False @@ -352,9 +352,9 @@ def test_on_dims_changed_callback(): fired = [] controller.on_dims_changed(scene.id, fired.append) - controller.get_scene(scene.id).dims.slice_indices = (7,) + controller.get_scene(scene.id).dims.selection.slice_indices = {0: 7} assert len(fired) == 1 - assert fired[0].slice_indices == (7,) + assert fired[0].selection.slice_indices == {0: 7} def test_unsubscribe_all_cleans_up(small_zarr_store): diff --git a/tests/v2/test_viewer_model.py b/tests/v2/test_viewer_model.py index 4f2e0339..fa706f14 100644 --- a/tests/v2/test_viewer_model.py +++ b/tests/v2/test_viewer_model.py @@ -3,7 +3,7 @@ from cellier.v2.data.image._zarr_multiscale_store import MultiscaleZarrDataStore from cellier.v2.scene.cameras import OrbitCameraController, 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 ImageAppearance, MultiscaleImageVisual @@ -34,8 +34,10 @@ def _build_viewer(small_zarr_store): cs = CoordinateSystem(name="world", axis_labels=("z", "y", "x")) dims = DimsManager( coordinate_system=cs, - displayed_axes=(0, 1, 2), - slice_indices=(), + selection=AxisAlignedSelection( + displayed_axes=(0, 1, 2), + slice_indices={}, + ), ) camera = PerspectiveCamera( fov=70.0,