diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..fbd7c05 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,47 @@ +name: release + +on: + release: + types: [published] + +permissions: + contents: read + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 + with: + persist-credentials: false + - uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6 + with: + python-version: "3.13" + enable-cache: false + - name: Verify tag matches package version + shell: bash + run: | + package_version="$(uv version --short)" + test "${GITHUB_REF_NAME}" = "v${package_version}" + - run: uv build + - run: uvx twine check dist/* + - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 + with: + name: python-package-distributions + path: dist/ + if-no-files-found: error + + publish: + needs: build + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/p/lodstone + permissions: + id-token: write + steps: + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4 + with: + name: python-package-distributions + path: dist/ + - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9f60e3b..26d6dc8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -5,16 +5,27 @@ on: branches: [main] pull_request: +permissions: + contents: read + jobs: test: - runs-on: ubuntu-latest + runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: + os: [ubuntu-latest] python-version: ["3.11", "3.12", "3.13", "3.14"] + include: + - os: macos-latest + python-version: "3.13" + - os: windows-latest + python-version: "3.13" steps: - - uses: actions/checkout@v4 - - uses: astral-sh/setup-uv@v6 + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 + with: + persist-credentials: false + - uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6 with: python-version: ${{ matrix.python-version }} - run: uv run --extra test pytest @@ -22,8 +33,10 @@ jobs: quality: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: astral-sh/setup-uv@v6 + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 + with: + persist-credentials: false + - uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6 with: python-version: "3.13" - run: uv run --group dev ruff check . diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..f1a78b1 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,13 @@ +# Changelog + +Lodstone follows semantic versioning. Viewer adapters remain provisional during +the 0.1 alpha series and may change between prereleases. + +## 0.1.0a0 - Unreleased + +- Add renderer-neutral multiscale chunk planning and progressive streaming. +- Add bounded caches, cancellation, request reprioritization, and diagnostics. +- Add dense resident windows with coarse-to-fine composition. +- Add NumPy, Zarr, and OME-Zarr sources. +- Add experimental napari and ndv adapters. +- Support shared runtimes across layers and channels. diff --git a/README.md b/README.md index 958ef81..2748c0d 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,10 @@ into progressive array `Update`s accepted by a viewer-specific `Target`. It is intended to be shared by clients such as ChimeraX, napari, and ndv. Lodstone does not create windows, textures, shaders, layers, or viewer models. +Lodstone is currently an alpha. The renderer-neutral core is the compatibility +boundary for the 0.1 series; viewer adapters are experimental and may change +between prereleases as their host applications establish public streaming APIs. + ```text Source + View + Target │ @@ -99,6 +103,22 @@ stream = Stream( ) ``` +Viewers with several layers or channels should share a `Runtime`. It owns one +asynchronous scheduler and a bounded CPU staging pool, while each stream keeps +its independent request state. Heavy `stage`, `stage_prepare`, and +`stage_phase` work runs in that pool instead of blocking I/O and cancellation: + +```python +from lodstone import Runtime, Stream + +runtime = Runtime(compute_workers=2) +streams = [ + Stream(source, target, runtime=runtime, dispatch=run_on_viewer_thread) + for source, target in channels +] +# Close streams first, then the shared runtime. +``` + ## napari adapter The optional adapter uses the rendering architecture from napari PR #9067: a @@ -129,6 +149,32 @@ This currently requires the `lodstone-integration` napari branch based on PR `examples/napari_ome_zarr.py` for a two-channel remote example. The core package still has no napari or Qt dependency. +## ndv adapter + +The ndv adapter presents immutable dense phase snapshots through ndv's common +`ArrayCanvas` API, so the same target works with its VisPy and pygfx renderers. +Create an empty `ndv.ArrayViewer`, pass it and a source to `NDVController`, then +submit renderer-neutral `View` snapshots: + +```python +import ndv + +from lodstone.adapters.ndv import NDVController + +viewer = ndv.ArrayViewer() +controller = NDVController(viewer, source) +controller.update(view) +viewer.show() +ndv.run_app() +controller.close() +``` + +The initial adapter supports translated dense 2-D and 3-D windows, hidden-axis +selections, camera-driven replanning, progressive phase replacement, independent +per-image world transforms, and shared runtimes. See `examples/ndv_dense.py`. +These capabilities currently require ndv's `lodstone-integration` branch until +its camera, dispatch, and image-transform APIs are released. + `examples/napari_zebrahub.py` opens one lazy timepoint from the public ZSNS001 Zebrahub light-sheet series in 3-D. Its approximately 32 MiB native chunks make it a useful stress test for cancellation, interaction holds, and @@ -148,6 +194,8 @@ the image's contrast or texture values. - **View** — displayed axes, hidden-axis selections, viewport, and camera matrix. - **Target** — desired dense/tiled/bricked layout and update delivery. - **Planner** — deterministic visible-tile and LOD selection. +- **PlanCoverage / PlanDelta** — stable coverage identity plus retained, + requested, reprioritized, and released work across camera changes. - **Stream** — cancellation, priorities, native-chunk reuse, CPU caching, batching, progressive delivery, and stale-generation rejection. - **Composition** — transform-aware nearest-neighbor backdrop sampling and @@ -161,6 +209,11 @@ for 32-cubed bricks while the Zarr source stores 16 by 64 by 64 chunks. Lodstone reads each overlapping native chunk once and assembles the requested display updates from its decoded cache. +Progressive planning starts at the coarsest level by default. Renderer +integrations can set `Planner(max_initial_voxel_footprint=...)` to choose the +coarsest initial level whose projected voxels stay within that many screen +pixels; the normal target level and napari's default behavior are unchanged. + `stream.diagnostics` separates renderer tiles from native storage activity for the current or most recent generation. `stream.cache_events` records recent `queued`, `loading`, `ready`, `failed`, and `evicted` transitions, while @@ -208,6 +261,23 @@ call `stream.pause()` and `stream.resume()` without discarding the active pass. `bytes_per_second` can pace aggregate source reads when decoding or remote I/O would otherwise compete with interaction and rendering. +`prepare` may return a residency lease with dynamic `available_keys` and +`pending_keys` sets plus `release(keys)`. A lease confirms which target storage +survives replanning, allowing the stream to retain delivered overlap while it +keeps loading native chunks shared by the old and new request. Queued work is +rebuilt in the newest priority order and work outside the new coverage is +canceled. Legacy targets that return no lease retain conservative pass +replacement behavior. `stream.delta` exposes the latest `PlanDelta`. + +Viewers may attach `InteractionState` to a `View` to describe camera motion +and angular, translation, and zoom velocity. Existing integrations can omit +it and retain their current policy. + +Targets that need an atomic presentation point between coarse-to-fine stages +may also implement `phase_complete(view, plan, phase)`. The hook is optional; +existing targets continue to receive the same prepare, apply, complete, and +redraw calls. + Dense targets can use `ResidentArrays` to avoid allocating complete pyramid levels. It stages one full-ND bounding window per desired level, preserves overlapping content when the camera moves, translates absolute updates into @@ -215,9 +285,9 @@ window-relative writes, and retires coarse/replaced storage on completion. The viewer still owns the corresponding grid, texture, or volume objects: ```python -from lodstone import Layout, ResidentArrays +from lodstone import Layout, ResidentArrays, ResidentLease -resident = ResidentArrays(source.pyramid) +resident = ResidentArrays(source.pyramid, compose=True) def layout(view, pyramid): @@ -228,11 +298,13 @@ def prepare(view, plan): transition = resident.prepare(plan) # Create renderer resources for transition.prepared and remove # transition.retired resources. + desired = plan.desired or plan.wanted + return ResidentLease(resident, frozenset(tile.key for tile in desired)) def apply(updates): for change in resident.apply(updates): - # Patch or invalidate the renderer resource for change.window. + # Patch or invalidate change.regions in the renderer resource. pass @@ -241,6 +313,11 @@ def complete(view, plan): # Present resident.active[plan.target_level] and retire old resources. ``` +With `compose=True`, coarse updates initialize and repair unloaded native +chunks in finer pending windows using the pyramid transforms. Directly loaded +fine chunks are never overwritten. Leaving composition disabled preserves the +original fill-value and same-level overlap behavior. + The initial expected layouts are: | Client | Typical layout | @@ -252,6 +329,22 @@ The initial expected layouts are: Lodstone deliberately stops before physical GPU allocation. The target owns textures, double buffering, shader indirection, and renderer invalidation. +## Viewer compatibility + +The first alpha is intended for integration development. It does not make the +streaming paths available in unmodified stable releases of every viewer. + +| Client | Initial support | Required host version | Status | +| --- | --- | --- | --- | +| ChimeraX OME-Zarr | 3-D images, channels, one selected timepoint | `chimerax-ome-zarr` PR 22 | Experimental | +| napari | 2-D/3-D Image and Labels layers | napari PR 34 based on PR 9067 | Experimental | +| ndv + VisPy | 2-D/3-D dense clipmaps and camera replanning | ndv `lodstone-integration` branch | Reference ndv backend | +| ndv + PyGFX | Same renderer-neutral data path | ndv `lodstone-integration` branch | Experimental visual parity | + +Integrations should pin an exact Lodstone prerelease. Compatibility is only +claimed for combinations exercised by the integration's native tests and smoke +tests; adapters remain provisional throughout the 0.1 alpha series. + ## Development ```bash @@ -263,3 +356,6 @@ uv run --group dev pyright src The test suite is network-independent. Remote opening and reading has also been checked against the EBI IDR OME-Zarr v0.4 store used by `chimerax-ome-zarr`. + +Release maintainers should follow [`RELEASING.md`](RELEASING.md). Changes are +recorded in [`CHANGELOG.md`](CHANGELOG.md). diff --git a/RELEASING.md b/RELEASING.md new file mode 100644 index 0000000..8948d22 --- /dev/null +++ b/RELEASING.md @@ -0,0 +1,36 @@ +# Releasing Lodstone + +## One-time PyPI setup + +1. Create a pending trusted publisher for the `lodstone` project on PyPI. +2. Set the owner to `kephale`, repository to `lodstone`, workflow to + `release.yml`, and environment to `pypi`. +3. Create a protected `pypi` environment in the GitHub repository. Requiring a + reviewer is recommended. + +No PyPI token is stored in GitHub. The release workflow uses OpenID Connect +trusted publishing and uploads PyPI attestations. + +## Release checklist + +1. Ensure `CHANGELOG.md` describes the release and has a release date. +2. Set the version in `pyproject.toml` and refresh `uv.lock`. +3. Run the complete local gate: + + ```bash + uv run --group dev ruff check . + uv run --group dev ruff format --check . + uv run --extra ome-zarr --group dev pyright src + uv run --extra test pytest + uv build + uvx twine check dist/* + ``` + +4. Merge the release commit to `main` and wait for CI. +5. Create a GitHub release whose tag is exactly `v`. +6. Verify the `release` workflow and the files and attestations on PyPI. +7. Install the release into a clean environment and run the README example. +8. Update the compatible napari, ndv, and ChimeraX integrations to the released + version and run their native smoke tests. + +The workflow refuses to publish when the Git tag and package version differ. diff --git a/examples/_ndv_block_overlay.py b/examples/_ndv_block_overlay.py new file mode 100644 index 0000000..2ea754c --- /dev/null +++ b/examples/_ndv_block_overlay.py @@ -0,0 +1,134 @@ +"""Qt overlay for inspecting Lodstone's projected dense block plans.""" + +from __future__ import annotations + +from itertools import product +from typing import TYPE_CHECKING, Any + +import numpy as np +from qtpy.QtCore import QEvent, QObject, QPointF, Qt +from qtpy.QtGui import QColor, QPainter, QPen +from qtpy.QtWidgets import QWidget + +if TYPE_CHECKING: + from lodstone import Plan, Pyramid, Region + + +class BlockOverlay(QWidget): + """Draw planned blocks and their enclosing dense focus box over a canvas.""" + + _EDGES = ( + (0, 1), + (0, 2), + (0, 4), + (1, 3), + (1, 5), + (2, 3), + (2, 6), + (3, 7), + (4, 5), + (4, 6), + (5, 7), + (6, 7), + ) + + def __init__(self, canvas: Any, pyramid: Pyramid) -> None: + parent = canvas.frontend_widget() + super().__init__(parent) + self._canvas = canvas + self._pyramid = pyramid + self._plan: Plan | None = None + self.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents) + self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground) + self.setGeometry(parent.rect()) + parent.installEventFilter(self) + self.show() + self.raise_() + + def set_plan(self, plan: Plan) -> None: + self._plan = plan + self.update() + + def eventFilter(self, watched: QObject, event: QEvent) -> bool: + if watched is self.parent() and event.type() == QEvent.Type.Resize: + self.setGeometry(self.parent().rect()) + return super().eventFilter(watched, event) + + def paintEvent(self, event: Any) -> None: + del event + if self._plan is None: + return + try: + viewport, world_to_clip = self._canvas.camera_state() + except RuntimeError: + return + + painter = QPainter(self) + painter.setRenderHint(QPainter.RenderHint.Antialiasing) + context_level = len(self._pyramid.levels) - 1 + target_level = self._plan.target_level + by_level: dict[int, list[Any]] = {} + for tile in self._plan.desired: + by_level.setdefault(tile.level, []).append(tile) + + # Context is cool blue; the camera-selected focus is warm orange. + for level, tiles in by_level.items(): + color = QColor(55, 205, 255, 105) + if level == target_level and level != context_level: + color = QColor(255, 165, 45, 155) + painter.setPen(QPen(color, 1.0)) + transform = self._pyramid.levels[level].voxel_to_world + for tile in tiles: + self._draw_region( + painter, tile.region, transform, world_to_clip, viewport + ) + + focus = by_level.get(target_level, []) + if target_level != context_level and focus: + region = _enclosing_region(focus) + pen = QPen(QColor(255, 245, 170, 230), 2.0, Qt.PenStyle.DashLine) + painter.setPen(pen) + self._draw_region( + painter, + region, + self._pyramid.levels[target_level].voxel_to_world, + world_to_clip, + viewport, + ) + painter.end() + + def _draw_region( + self, + painter: QPainter, + region: Region, + voxel_to_world: np.ndarray, + world_to_clip: np.ndarray, + viewport: tuple[int, int], + ) -> None: + corners = np.asarray( + [(*corner, 1.0) for corner in product(*zip(region.start, region.stop))], + dtype=np.float64, + ) + clip = (world_to_clip @ voxel_to_world @ corners.T).T + if np.any(np.abs(clip[:, 3]) < 1e-12): + return + ndc = clip[:, :3] / clip[:, 3, None] + if not np.all(np.isfinite(ndc)) or np.all(clip[:, 3] < 0): + return + width, height = viewport + points = [ + QPointF((point[0] + 1.0) * width / 2, (1.0 - point[1]) * height / 2) + for point in ndc + ] + for start, stop in self._EDGES: + painter.drawLine(points[start], points[stop]) + + +def _enclosing_region(tiles: list[Any]) -> Region: + from lodstone import Region + + ndim = tiles[0].region.ndim + return Region( + tuple(min(tile.region.start[axis] for tile in tiles) for axis in range(ndim)), + tuple(max(tile.region.stop[axis] for tile in tiles) for axis in range(ndim)), + ) diff --git a/examples/napari_zebrahub.py b/examples/napari_zebrahub.py index 55a8f32..37b989c 100644 --- a/examples/napari_zebrahub.py +++ b/examples/napari_zebrahub.py @@ -124,31 +124,31 @@ def main() -> None: vector_style="line", ) viewer.dims.ndisplay = arguments.ndisplay - viewer.reset_view() - if any( - value is not None - for value in ( - arguments.camera_center, - arguments.camera_angles, - arguments.zoom, - ) - ): - from qtpy.QtCore import QTimer - - def restore_camera() -> None: - # Let the initial volume texture and GL objects be created before - # the saved pose triggers a new progressive-loading pass. - if arguments.camera_center is not None: - viewer.camera.center = tuple(arguments.camera_center) - if arguments.camera_angles is not None: - viewer.camera.angles = tuple(arguments.camera_angles) - if arguments.zoom is not None: - viewer.camera.zoom = arguments.zoom - - QTimer.singleShot(1000, restore_camera) + from qtpy.QtCore import QTimer + + def initialize_camera() -> None: + # Progressive loading enables asynchronous slicing. reset_view() + # before the Qt event loop has realized the initial 3-D slice uses + # raw data extents for some axes and transformed world extents for + # others, centering fine clipmap pages far outside the visible data. + # Initialize from the realized layer extent, then apply any requested + # reproducible pose. + viewer.reset_view() + if arguments.camera_center is None: + world_extent = np.asarray(controller.layer.extent.world, dtype=float) + viewer.camera.center = tuple( + np.mean(world_extent, axis=0)[-arguments.ndisplay :] + ) + else: + viewer.camera.center = tuple(arguments.camera_center) + if arguments.camera_angles is not None: + viewer.camera.angles = tuple(arguments.camera_angles) + if arguments.zoom is not None: + viewer.camera.zoom = arguments.zoom + + QTimer.singleShot(250, initialize_camera) if arguments.screenshot: - from qtpy.QtCore import QTimer def save() -> None: viewer.screenshot(arguments.screenshot, canvas_only=True) diff --git a/examples/ndv_dense.py b/examples/ndv_dense.py new file mode 100644 index 0000000..5c841fc --- /dev/null +++ b/examples/ndv_dense.py @@ -0,0 +1,232 @@ +"""Display a real Zebrahub timepoint through ndv's dense 3-D canvas. + +The coarsest source level remains visible as a clipmap-style context while +translated, camera-selected higher-resolution windows stream above it. +""" + +from __future__ import annotations + +import argparse + +import ndv +import numpy as np + +from lodstone import Plan, View +from lodstone.adapters.ndv import NDVController, NDVPublication +from lodstone.sources import OMEZarrSource + +DEFAULT_URL = ( + "https://public.czbiohub.org/royerlab/zebrahub/imaging/" + "single-objective/ZSNS001.ome.zarr" +) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("url", nargs="?", default=DEFAULT_URL) + parser.add_argument("--time", type=int, default=400) + parser.add_argument( + "--backend", + choices=("vispy", "pygfx"), + help="ndv canvas backend (defaults to ndv's automatic selection)", + ) + arguments = parser.parse_args() + + if arguments.backend is not None: + ndv.set_canvas_backend(arguments.backend) + + print(f"Opening Zebrahub timepoint {arguments.time} from {arguments.url}") + source = OMEZarrSource.open( + arguments.url, + fixed_index={"t": arguments.time, "c": 0}, + ) + finest = source.pyramid.levels[0] + linear = finest.voxel_to_world[:-1, :-1] + world_extent = np.abs(linear) @ np.asarray(finest.shape, dtype=np.float64) + # ndv's scene axes are XYZ while the source axes are ZYX. Bootstrap a + # top-down full-volume view; camera snapshots take over after publication. + world_to_clip = np.zeros((4, 4), dtype=np.float64) + world_to_clip[0, 2] = 2.0 / world_extent[2] + world_to_clip[1, 1] = 2.0 / world_extent[1] + world_to_clip[2, 0] = 2.0 / world_extent[0] + world_to_clip[3, 3] = 1.0 + world_to_clip[:-1, -1] = -1.0 + + viewer = ndv.ArrayViewer() + widget = viewer.widget() + from _ndv_block_overlay import BlockOverlay + from qtpy.QtCore import Qt, QTimer + from qtpy.QtWidgets import ( + QCheckBox, + QFormLayout, + QFrame, + QHBoxLayout, + QLabel, + QSlider, + QWidget, + ) + + block_overlay = BlockOverlay(viewer.canvas, source.pyramid) + viewer.canvas.cameraChanged.connect(lambda: block_overlay.update()) + + backend_name = { + "VispyArrayCanvas": "VisPy", + "GfxArrayCanvas": "PyGFX", + }.get(type(viewer.canvas).__name__, type(viewer.canvas).__name__) + indicator = QLabel(f"Lodstone · {backend_name} · waiting", widget) + indicator.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents) + indicator.setStyleSheet( + "background: rgba(0, 0, 0, 190); color: white; " + "padding: 6px 9px; border-radius: 4px; font-weight: bold;" + ) + indicator.move(12, 12) + indicator.adjustSize() + indicator.show() + indicator.raise_() + active_level: int | None = None + target_level: int | None = None + + def update_indicator() -> None: + active = "waiting" if active_level is None else f"L{active_level}" + target = "?" if target_level is None else f"L{target_level}" + indicator.setText( + f"Lodstone · {backend_name} · active {active} · target {target}" + ) + indicator.adjustSize() + indicator.raise_() + + def show_presented_level(publication: NDVPublication) -> None: + nonlocal active_level + active_level = publication.level + spacing = " × ".join(f"{scale:g}" for scale in publication.scales) + title = f"Lodstone Zebrahub — rendered L{publication.level} ({spacing})" + print(title) + widget.window().setWindowTitle(title) + update_indicator() + + def show_target_level(plan: Plan) -> None: + nonlocal target_level + target_level = plan.target_level + block_overlay.set_plan(plan) + tiles = [tile for tile in plan.desired if tile.level == plan.target_level] + if tiles: + ndim = tiles[0].region.ndim + start = tuple( + min(tile.region.start[axis] for tile in tiles) for axis in range(ndim) + ) + stop = tuple( + max(tile.region.stop[axis] for tile in tiles) for axis in range(ndim) + ) + shape = tuple(upper - lower for lower, upper in zip(start, stop)) + level = source.pyramid.levels[plan.target_level] + dense_mib = np.prod(shape) * level.dtype.itemsize / 1024**2 + block_text = "×".join(str(value) for value in shape) + geometry_value.setText( + f"L{plan.target_level} · {len(tiles)} blk · {dense_mib:.1f} MiB\n" + f"{block_text}" + ) + update_indicator() + + controller = NDVController( + viewer, + source, + on_presented=show_presented_level, + on_targeted=show_target_level, + ) + + controls = QFrame(widget) + controls.setStyleSheet( + "QFrame { background: rgba(0, 0, 0, 190); color: white; " + "border-radius: 4px; } QLabel { color: white; }" + ) + form = QFormLayout(controls) + form.setContentsMargins(9, 7, 9, 7) + form.setSpacing(5) + + def slider_row( + minimum: int, maximum: int, value: int + ) -> tuple[QWidget, QSlider, QLabel]: + row = QWidget(controls) + row_layout = QHBoxLayout(row) + row_layout.setContentsMargins(0, 0, 0, 0) + slider = QSlider(Qt.Orientation.Horizontal, row) + slider.setRange(minimum, maximum) + slider.setValue(value) + slider.setMinimumWidth(150) + readout = QLabel(row) + readout.setMinimumWidth(48) + row_layout.addWidget(slider) + row_layout.addWidget(readout) + return row, slider, readout + + budget_row, budget_slider, budget_value = slider_row(32, 512, 128) + depth_row, depth_slider, depth_value = slider_row(0, 100, 60) + lod_row, lod_slider, lod_value = slider_row(50, 300, 200) + depth_slider.setToolTip( + "0% favors high-resolution coverage across the canvas; " + "100% favors a central column reaching front-to-back" + ) + geometry_value = QLabel("waiting", controls) + blocks_toggle = QCheckBox("Show planned block wireframes", controls) + blocks_toggle.setChecked(True) + blocks_toggle.toggled.connect(block_overlay.setVisible) + form.addRow("Focus MiB", budget_row) + form.addRow("Depth reach", depth_row) + form.addRow("LOD bias", lod_row) + form.addRow("Focus box", geometry_value) + form.addRow("Diagnostics", blocks_toggle) + + tune_timer = QTimer(controls) + tune_timer.setSingleShot(True) + tune_timer.setInterval(250) + + def configure_focus() -> None: + memory_mib = budget_slider.value() + # Span a wide perceptual range: the canvas-heavy endpoint strongly + # penalizes far chunks, while the depth-heavy endpoint retains a small + # near-to-far tie breaker instead of reverting to data-axis ordering. + depth_fraction = depth_slider.value() / 100.0 + depth_weight = 8.0 * (0.5 / 8.0) ** depth_fraction + lod_bias = lod_slider.value() / 100.0 + budget_value.setText(str(memory_mib)) + depth_value.setText(f"{depth_slider.value()}% depth") + lod_value.setText(f"{lod_bias:.2f}") + controller.set_focus_policy( + memory_limit=memory_mib * 1024**2, + focus_depth_weight=depth_weight, + lod_bias=lod_bias, + replan=False, + ) + tune_timer.start() + + tune_timer.timeout.connect(lambda: controller.set_focus_policy()) + for slider in (budget_slider, depth_slider, lod_slider): + slider.valueChanged.connect(configure_focus) + configure_focus() + tune_timer.stop() + controls.move(12, 50) + controls.adjustSize() + controls.show() + controls.raise_() + + plan = controller.update( + View( + displayed_axes=(0, 1, 2), + index=(None, None, None), + viewport=(600, 600), + world_to_clip=world_to_clip, + ) + ) + print( + f"Loading pyramid level {plan.target_level} " + f"from {len(plan.wanted)} logical blocks" + ) + viewer.show() + try: + ndv.run_app() + finally: + controller.close() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 9a7c6d3..24bb529 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,8 @@ dev = ["pyright>=1.1", "ruff>=0.12"] [project.urls] Repository = "https://github.com/kephale/lodstone" +Issues = "https://github.com/kephale/lodstone/issues" +Changelog = "https://github.com/kephale/lodstone/blob/main/CHANGELOG.md" [tool.hatch.build.targets.wheel] packages = ["src/lodstone"] diff --git a/src/lodstone/__init__.py b/src/lodstone/__init__.py index 9613d8c..662c495 100644 --- a/src/lodstone/__init__.py +++ b/src/lodstone/__init__.py @@ -1,5 +1,7 @@ """Renderer-neutral multiscale chunk planning and streaming.""" +from importlib.metadata import version + from .chunks import ( ChunkGrid, chunk_boundaries, @@ -23,9 +25,12 @@ from .model import ( ChunkEvent, ChunkState, + InteractionState, Layout, Level, Plan, + PlanCoverage, + PlanDelta, Pyramid, Region, Status, @@ -45,34 +50,46 @@ from .resident import ( ResidentArrays, ResidentChange, + ResidentLease, ResidentTransition, ResidentWindow, ) +from .runtime import Runtime from .source import Source from .stream import Stream -from .target import PassTarget, Target +from .target import PassTarget, PhaseTarget, ResidencyLease, StagingTarget, Target from .virtual import MultiScaleVirtualData, VirtualArrayView, VirtualData +__version__ = version("lodstone") + __all__ = [ "ChunkEvent", "ChunkGrid", "ChunkState", + "InteractionState", "Layout", "Level", "LevelDiagnosticArray", "MultiScaleVirtualData", "PassTarget", + "PhaseTarget", "Plan", "PlanComparison", + "PlanCoverage", + "PlanDelta", "PlanTrace", "Planner", "Pyramid", "Region", + "ResidencyLease", "ResidentArrays", "ResidentChange", + "ResidentLease", "ResidentTransition", "ResidentWindow", + "Runtime", "Source", + "StagingTarget", "Status", "Stream", "StreamDiagnostics", @@ -83,6 +100,7 @@ "View", "VirtualArrayView", "VirtualData", + "__version__", "anisotropic_extent_for_bytes", "available_tile_keys", "chunk_boundaries", diff --git a/src/lodstone/adapters/__init__.py b/src/lodstone/adapters/__init__.py index 5cb95fa..97c2463 100644 --- a/src/lodstone/adapters/__init__.py +++ b/src/lodstone/adapters/__init__.py @@ -7,9 +7,13 @@ add_lodstone_image, add_lodstone_labels, ) +from .ndv import NDVController, NDVPublication, NDVTarget __all__ = [ "ArraySource", + "NDVController", + "NDVPublication", + "NDVTarget", "NapariController", "add_lodstone_diagnostics", "add_lodstone_image", diff --git a/src/lodstone/adapters/ndv.py b/src/lodstone/adapters/ndv.py new file mode 100644 index 0000000..abc5100 --- /dev/null +++ b/src/lodstone/adapters/ndv.py @@ -0,0 +1,522 @@ +"""Dense Lodstone target for ndv's renderer-independent canvas API.""" + +from __future__ import annotations + +from collections.abc import Callable, Collection, Sequence +from dataclasses import dataclass, replace +from itertools import product +from threading import Condition, Thread +from time import monotonic +from typing import Any, Self + +import numpy as np + +from ..model import Layout, Plan, Pyramid, Region, TileKey, Update, View +from ..planner import Planner +from ..resident import ResidentArrays, ResidentLease +from ..runtime import Runtime +from ..source import Source +from ..stream import Stream + + +@dataclass(frozen=True, slots=True) +class NDVPublication: + """Immutable dense phase snapshot ready for an ndv image handle.""" + + level: int + region: Region + data: np.ndarray + data_axes: tuple[int, ...] + scales: tuple[float, ...] + origins: tuple[float, ...] + + +@dataclass(frozen=True, slots=True) +class NDVPreparation: + desired_keys: frozenset[TileKey] + + +class NDVTarget: + """Present complete dense phases through an ndv ``ArrayCanvas``. + + The target only relies on ndv's canvas and image-handle abstraction, so the + same adapter works with its VisPy and pygfx backends. Each pyramid level is + a separate image handle so a coarse context can remain visible around the + camera-focused high-resolution window. + """ + + def __init__( + self, + canvas: Any, + pyramid: Pyramid, + *, + memory_limit: int = 64 * 1024**2, + block_shape: tuple[int, ...] = (64, 128, 128), + focus_depth_weight: float = 0.5, + on_presented: Callable[[NDVPublication], None] | None = None, + ) -> None: + self.canvas = canvas + self.pyramid = pyramid + self.memory_limit = int(memory_limit) + if self.memory_limit <= 0: + raise ValueError("memory_limit must be positive") + if not np.isfinite(focus_depth_weight) or focus_depth_weight < 0: + raise ValueError("focus_depth_weight must be finite and nonnegative") + self.block_shape = tuple(int(value) for value in block_shape) + self.focus_depth_weight = float(focus_depth_weight) + self.resident = ResidentArrays(pyramid, compose=True) + self.handles: dict[int, Any] = {} + self.publications: dict[int, NDVPublication] = {} + self.handle: Any | None = None + self._context_level = len(pyramid.levels) - 1 + self._masked_focus: tuple[int, Region] | None = None + self.on_presented = on_presented + + def layout(self, view: View, pyramid: Pyramid) -> Layout: + return Layout( + kind="dense", + block_shape=self.block_shape[-len(view.displayed_axes) :], + mixed_lod=True, + memory_limit=self.memory_limit, + squeeze_hidden=False, + max_axis_extent=512, + memory_policy="crop", + focus_depth_weight=self.focus_depth_weight, + ) + + def stage_prepare(self, view: View, plan: Plan) -> NDVPreparation: + self.resident.prepare(plan) + desired = plan.desired or plan.wanted + return NDVPreparation(frozenset(tile.key for tile in desired)) + + def prepare( + self, + view: View, + plan: Plan, + prepared: NDVPreparation, + ) -> ResidentLease: + return ResidentLease(self.resident, prepared.desired_keys) + + def stage(self, updates: Sequence[Update]): + return self.resident.apply(updates) + + def apply(self, updates: Sequence[Update]) -> None: + # Resident writes occur in ``stage``. ndv receives one immutable array + # only when the corresponding phase is complete. + return None + + def stage_phase( + self, + view: View, + plan: Plan, + phase: int, + ) -> NDVPublication: + desired = plan.desired or plan.wanted + levels = {tile.level for tile in desired if tile.phase == phase} + if len(levels) != 1: + raise RuntimeError("an ndv phase must describe exactly one pyramid level") + level = levels.pop() + window = self.resident.windows[level] + data_axes = tuple( + axis for axis in range(window.region.ndim) if axis in view.displayed_axes + ) + hidden = tuple( + axis + for axis in range(window.region.ndim) + if axis not in view.displayed_axes + ) + if any(window.data.shape[axis] != 1 for axis in hidden): + raise RuntimeError("hidden axes must be singleton resident selections") + data = np.squeeze(window.data, axis=hidden).copy() + data.setflags(write=False) + transform = self.pyramid.levels[level].voxel_to_world + scales = tuple( + float(np.linalg.norm(transform[:-1, axis])) for axis in data_axes + ) + world_origin = transform[:-1, :-1] @ np.asarray(window.region.start) + world_origin += transform[:-1, -1] + origins = tuple(float(world_origin[axis]) for axis in data_axes) + return NDVPublication(level, window.region, data, data_axes, scales, origins) + + def phase_complete( + self, + view: View, + plan: Plan, + phase: int, + publication: NDVPublication, + ) -> None: + ndim = len(view.displayed_axes) + self.canvas.set_ndim(ndim) + wants_level = any(tile.level == publication.level for tile in plan.wanted) + previous = self.publications.get(publication.level) + if ( + publication.level == self._context_level + and not wants_level + and previous is not None + and previous.region == publication.region + ): + # Stream phases also complete when every context tile was cached. + # Keep the existing visual instead of uploading the same array. + return + handle = self.handles.get(publication.level) + first_publication = not self.handles + if handle is None: + factory = self.canvas.add_image if ndim == 2 else self.canvas.add_volume + handle = factory(publication.data, reset_range=False) + self.handles[publication.level] = handle + else: + handle.set_data(publication.data) + handle.set_clims(_data_clims(publication.data)) + handle.set_world_transform(publication.scales, publication.origins) + handle.set_visible(True) + self.publications[publication.level] = publication + if publication.level == self._context_level: + self._masked_focus = None + self.handle = handle + if first_publication: + self.canvas.set_range() + if publication.level != self._context_level: + self._mask_context(publication) + self.canvas.refresh() + if self.on_presented is not None: + self.on_presented(publication) + + def discard(self, keys: Collection[TileKey]) -> None: + self.resident.discard(keys) + + def complete(self, view: View, plan: Plan) -> None: + self.resident.complete(plan, retain_levels=True) + keep = {self._context_level, plan.target_level} + for level in tuple(self.handles): + if level not in keep: + self.handles.pop(level).remove() + self.publications.pop(level, None) + context = self.publications.get(self._context_level) + if context is not None: + if plan.target_level == self._context_level: + self.handles[self._context_level].set_data(context.data) + self._masked_focus = None + elif (focus := self.publications.get(plan.target_level)) is not None: + self._mask_context(focus) + + def _mask_context(self, focus: NDVPublication) -> None: + """Punch the fine focus footprint out of the additive coarse visual.""" + context = self.publications.get(self._context_level) + handle = self.handles.get(self._context_level) + identity = (focus.level, focus.region) + if context is None or handle is None or self._masked_focus == identity: + return + overlap = _region_in_level( + focus.region, + self.pyramid.levels[focus.level].voxel_to_world, + self.pyramid.levels[self._context_level].voxel_to_world, + ).intersection(context.region) + masked = context.data.copy() + if overlap is not None: + if np.count_nonzero(focus.data) == focus.data.size: + slices = tuple( + slice( + overlap.start[axis] - context.region.start[axis], + overlap.stop[axis] - context.region.start[axis], + ) + for axis in context.data_axes + ) + masked[slices] = 0 + else: + _mask_occupied_context( + masked, + context, + focus, + self.pyramid.levels[focus.level].voxel_to_world, + self.pyramid.levels[self._context_level].voxel_to_world, + ) + handle.set_data(masked) + handle.set_world_transform(context.scales, context.origins) + self._masked_focus = identity + + def redraw(self) -> None: + self.canvas.refresh() + + def close(self) -> None: + for handle in self.handles.values(): + handle.remove() + self.handles.clear() + self.publications.clear() + self.handle = None + self.resident.clear() + + +def _region_in_level( + region: Region, + source_to_world: np.ndarray, + destination_to_world: np.ndarray, +) -> Region: + """Return a conservative destination-level box for ``region``.""" + ndim = region.ndim + destination_from_source = np.linalg.solve(destination_to_world, source_to_world) + corners = np.asarray( + [(*corner, 1.0) for corner in product(*zip(region.start, region.stop))], + dtype=np.float64, + ) + mapped = (destination_from_source @ corners.T).T[:, :ndim] + start = tuple(max(0, int(np.floor(value))) for value in mapped.min(axis=0)) + stop = tuple(max(0, int(np.ceil(value))) for value in mapped.max(axis=0)) + return Region(start, stop) + + +def _data_clims(data: np.ndarray) -> tuple[float, float]: + """Return finite nondegenerate limits independent of renderer defaults.""" + if np.issubdtype(data.dtype, np.inexact): + finite = data[np.isfinite(data)] + if not finite.size: + return (0.0, 1.0) + lower = float(np.min(finite)) + upper = float(np.max(finite)) + else: + lower = float(np.min(data)) + upper = float(np.max(data)) + if lower == upper: + upper = lower + 1.0 + return lower, upper + + +def _mask_occupied_context( + masked: np.ndarray, + context: NDVPublication, + focus: NDVPublication, + focus_to_world: np.ndarray, + context_to_world: np.ndarray, +) -> None: + """Mask context samples covered by nonzero fine data. + + A dense focus publication can contain a great deal of empty space. Masking + its complete box removes useful coarse landmarks in sparse volumes, while + masking only occupied fine samples still prevents additive brightening for + dense data. Work one leading-axis plane at a time to bound temporary memory. + """ + context_from_focus = np.linalg.solve(context_to_world, focus_to_world) + ndim = focus.region.ndim + base = np.asarray(focus.region.start, dtype=np.float64) + 0.5 + data = focus.data + for leading in range(data.shape[0]): + occupied = np.argwhere(data[leading] != 0) + if not len(occupied): + continue + local = np.column_stack( + (np.full(len(occupied), leading, dtype=np.int64), occupied) + ) + points = np.broadcast_to(base, (len(local), ndim)).copy() + for local_axis, data_axis in enumerate(focus.data_axes): + points[:, data_axis] += local[:, local_axis] + homogeneous = np.column_stack((points, np.ones(len(points)))) + mapped = (context_from_focus @ homogeneous.T).T[:, :ndim] + indices = np.floor(mapped).astype(np.int64) + context_indices = np.column_stack( + [ + indices[:, axis] - context.region.start[axis] + for axis in context.data_axes + ] + ) + valid = np.ones(len(context_indices), dtype=bool) + for axis, size in enumerate(masked.shape): + valid &= (context_indices[:, axis] >= 0) & (context_indices[:, axis] < size) + if np.any(valid): + masked[tuple(context_indices[valid].T)] = 0 + + +class NDVController: + """Connect a Lodstone source to an ndv ``ArrayViewer`` or canvas. + + The viewer integration intentionally accepts explicit :class:`View` + snapshots while ndv's public canvas API does not yet expose a camera matrix. + This is enough for full-view 2-D/3-D use and keeps camera-event policy out of + the renderer target. + """ + + def __init__( + self, + viewer_or_canvas: Any, + source: Source, + *, + dispatch: Callable[[Callable[[], None]], None] | None = None, + runtime: Runtime | None = None, + memory_limit: int = 64 * 1024**2, + block_shape: tuple[int, ...] = (64, 128, 128), + focus_depth_weight: float = 0.5, + on_presented: Callable[[NDVPublication], None] | None = None, + on_targeted: Callable[[Plan], None] | None = None, + camera_debounce_ms: int = 180, + **stream_options: Any, + ) -> None: + self.viewer = viewer_or_canvas if hasattr(viewer_or_canvas, "canvas") else None + self.canvas = ( + viewer_or_canvas.canvas if self.viewer is not None else viewer_or_canvas + ) + if dispatch is None: + dispatch = getattr(viewer_or_canvas, "dispatch", None) + self._dispatch = dispatch or (lambda callback: callback()) + self._on_targeted = on_targeted + self._camera_debounce_seconds = camera_debounce_ms / 1000 + if self._camera_debounce_seconds < 0: + raise ValueError("camera_debounce_ms must be nonnegative") + self._camera_condition = Condition() + self._pending_camera_view: View | None = None + self._camera_deadline = 0.0 + self._camera_generation = 0 + self._closed = False + self._presenting = False + self.runtime = runtime or Runtime() + self._owns_runtime = runtime is None + self.target = NDVTarget( + self.canvas, + source.pyramid, + memory_limit=memory_limit, + block_shape=block_shape, + focus_depth_weight=focus_depth_weight, + on_presented=on_presented, + ) + stream_options.setdefault( + "planner", + Planner( + lod_bias=2.0, + progressive=True, + max_intermediate_levels=0, + max_initial_voxel_footprint=4.0, + ), + ) + self.stream = Stream( + source, + self.target, + dispatch=self._dispatch_presentation, + runtime=self.runtime, + **stream_options, + ) + self._last_view: View | None = None + self._camera_signal = getattr(self.canvas, "cameraChanged", None) + if self._camera_signal is not None: + self._camera_signal.connect(self._camera_changed) + self._camera_thread = Thread( + target=self._camera_worker, + name="lodstone-ndv-camera", + daemon=True, + ) + self._camera_thread.start() + + def update(self, view: View) -> Plan: + self._last_view = view + plan = self.stream.update(view) + self._notify_targeted(plan) + return plan + + def set_focus_policy( + self, + *, + memory_limit: int | None = None, + focus_depth_weight: float | None = None, + lod_bias: float | None = None, + replan: bool = True, + ) -> Plan | None: + """Update interactive focus controls and optionally replan the view.""" + if memory_limit is not None: + if memory_limit <= 0: + raise ValueError("memory_limit must be positive") + self.target.memory_limit = int(memory_limit) + if focus_depth_weight is not None: + if not np.isfinite(focus_depth_weight) or focus_depth_weight < 0: + raise ValueError("focus_depth_weight must be finite and nonnegative") + self.target.focus_depth_weight = float(focus_depth_weight) + if lod_bias is not None: + planner = self.stream.planner + self.stream.planner = Planner( + lod_bias=lod_bias, + progressive=planner.progressive, + max_intermediate_levels=planner.max_intermediate_levels, + max_initial_voxel_footprint=planner.max_initial_voxel_footprint, + ) + if replan and self._last_view is not None: + return self.update(self._last_view) + return None + + def _camera_changed(self) -> None: + if self._last_view is None or self._presenting: + return + viewport, world_to_clip = self.canvas.camera_state() + # Scene-bound changes can alter only the depth/near-far transform when + # a replacement volume is published. They are not camera interaction + # and must not replay the same plan. The first two clip rows fully + # describe screen-space pan, zoom, and rotation for LOD selection. + if viewport == self._last_view.viewport and np.allclose( + world_to_clip[:2], + self._last_view.world_to_clip[:2], + rtol=1e-7, + atol=1e-7, + ): + return + view = replace( + self._last_view, + viewport=viewport, + world_to_clip=world_to_clip, + ) + self._last_view = view + with self._camera_condition: + if self._closed: + return + self._camera_generation += 1 + self._pending_camera_view = view + self._camera_deadline = monotonic() + self._camera_debounce_seconds + self._camera_condition.notify() + + def _camera_worker(self) -> None: + while True: + with self._camera_condition: + while self._pending_camera_view is None and not self._closed: + self._camera_condition.wait() + if self._closed: + return + remaining = self._camera_deadline - monotonic() + if remaining > 0: + self._camera_condition.wait(remaining) + continue + view = self._pending_camera_view + generation = self._camera_generation + self._pending_camera_view = None + assert view is not None + plan = self.stream.plan(view) + with self._camera_condition: + if self._closed or generation != self._camera_generation: + continue + self.stream.submit(view, plan) + self._notify_targeted(plan) + + def _dispatch_presentation(self, callback: Callable[[], None]) -> None: + def guarded() -> None: + self._presenting = True + try: + callback() + finally: + self._presenting = False + + self._dispatch(guarded) + + def _notify_targeted(self, plan: Plan) -> None: + callback = self._on_targeted + if callback is not None: + self._dispatch(lambda: callback(plan)) + + def close(self) -> None: + if self._camera_signal is not None: + self._camera_signal.disconnect(self._camera_changed) + with self._camera_condition: + self._closed = True + self._camera_condition.notify() + self._camera_thread.join(timeout=1) + self.stream.close() + self.target.close() + if self._owns_runtime: + self.runtime.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() diff --git a/src/lodstone/composition.py b/src/lodstone/composition.py index f7e4231..1331406 100644 --- a/src/lodstone/composition.py +++ b/src/lodstone/composition.py @@ -86,7 +86,8 @@ def fill_unloaded_chunks( return () boundaries = [np.concatenate(([0], np.cumsum(axis))) for axis in chunk_grid] - loaded_set = set(loaded) + loaded_regions = tuple(loaded) + loaded_set = set(loaded_regions) per_axis = [] for axis, bounds in enumerate(boundaries): starts, stops = bounds[:-1], bounds[1:] @@ -102,7 +103,9 @@ def fill_unloaded_chunks( tuple(start for start, _stop in bounds), tuple(stop for _start, stop in bounds), ) - if chunk in loaded_set: + if chunk in loaded_set or any( + region.intersection(chunk) == chunk for region in loaded_regions + ): continue write = chunk.intersection(overlap) if write is None: diff --git a/src/lodstone/model.py b/src/lodstone/model.py index 057b551..27ea435 100644 --- a/src/lodstone/model.py +++ b/src/lodstone/model.py @@ -153,6 +153,25 @@ def ndim(self) -> int: return len(self.axes) +@dataclass(frozen=True, slots=True) +class InteractionState: + """Optional motion metadata supplied by an interactive viewer.""" + + moving: bool = False + angular_velocity: float = 0.0 + translation_velocity: float = 0.0 + zoom_velocity: float = 0.0 + + def __post_init__(self) -> None: + velocities = ( + self.angular_velocity, + self.translation_velocity, + self.zoom_velocity, + ) + if any(not np.isfinite(value) for value in velocities): + raise ValueError("interaction velocities must be finite") + + @dataclass(frozen=True, slots=True) class View: """A host-neutral snapshot of a 2-D or 3-D viewer camera and selection. @@ -167,6 +186,7 @@ class View: viewport: tuple[int, int] world_to_clip: npt.NDArray[np.float64] eye: tuple[float, float, float] | None = None + interaction: InteractionState | None = None def __post_init__(self) -> None: if len(self.displayed_axes) not in (2, 3): @@ -201,6 +221,8 @@ class Layout: memory_limit: int = 1 << 30 squeeze_hidden: bool = True max_axis_extent: int | None = None + memory_policy: Literal["coarsen", "crop"] = "coarsen" + focus_depth_weight: float | None = None def __post_init__(self) -> None: if self.block_shape is not None and any( @@ -211,6 +233,12 @@ def __post_init__(self) -> None: raise ValueError("memory_limit must be positive") if self.max_axis_extent is not None and self.max_axis_extent <= 0: raise ValueError("max_axis_extent must be positive") + if self.memory_policy not in {"coarsen", "crop"}: + raise ValueError("memory_policy must be 'coarsen' or 'crop'") + if self.focus_depth_weight is not None and ( + not np.isfinite(self.focus_depth_weight) or self.focus_depth_weight < 0 + ): + raise ValueError("focus_depth_weight must be finite and nonnegative") @dataclass(frozen=True, slots=True) @@ -236,6 +264,31 @@ def level(self) -> int: return self.key.level +@dataclass(frozen=True, slots=True) +class PlanCoverage: + """Order-independent identity of the data covered by a plan. + + Tile priority and progressive phase are intentionally excluded. Hidden-axis + selections are included explicitly as well as through :class:`TileKey`, so + integrations can inspect them without decoding tile identities. + """ + + target_level: int + tile_regions: frozenset[tuple[TileKey, Region]] + retained_keys: frozenset[TileKey] + hidden_axis_selections: frozenset[tuple[int, ...]] + + +@dataclass(frozen=True, slots=True) +class PlanDelta: + """Stable coverage changes between two successive plans.""" + + retained: frozenset[TileKey] + requested: tuple[Tile, ...] + reprioritized: tuple[TileKey, ...] + released: frozenset[TileKey] + + @dataclass(frozen=True, slots=True) class Plan: """A complete desired tile set and the reads needed to reach it. @@ -251,6 +304,55 @@ class Plan: target_level: int desired: tuple[Tile, ...] = () + @property + def coverage(self) -> PlanCoverage: + """Return the requested coverage, ignoring delivery order and priority.""" + + tiles = self.desired or self.wanted + keys = (*[tile.key for tile in tiles], *self.retain) + return PlanCoverage( + target_level=self.target_level, + tile_regions=frozenset((tile.key, tile.region) for tile in tiles), + retained_keys=self.retain, + hidden_axis_selections=frozenset(key.selection for key in keys), + ) + + def delta(self, previous: Plan | None) -> PlanDelta: + """Compare stable coverage while preserving current request order.""" + + current_tiles = self.desired or self.wanted + current = {tile.key: tile for tile in current_tiles} + if previous is None: + return PlanDelta( + retained=frozenset(), + requested=tuple(self.wanted), + reprioritized=(), + released=frozenset(), + ) + previous_tiles = previous.desired or previous.wanted + old = {tile.key: tile for tile in previous_tiles} + retained = frozenset( + key + for key, tile in current.items() + if key in old and old[key].region == tile.region + ) + requested = tuple(tile for tile in self.wanted if tile.key not in retained) + reprioritized = tuple( + tile.key + for tile in current_tiles + if tile.key in retained + and ( + old[tile.key].priority != tile.priority + or old[tile.key].phase != tile.phase + ) + ) + released = frozenset( + key + for key, tile in old.items() + if key not in retained or key not in self.retain + ) + return PlanDelta(retained, requested, reprioritized, released) + @dataclass(frozen=True, slots=True) class Update: @@ -316,6 +418,9 @@ class StreamDiagnostics: evictions: int = 0 cache_chunks: int = 0 cache_bytes: int = 0 + prepare_stage_seconds: float = 0.0 + update_stage_seconds: float = 0.0 + phase_stage_seconds: float = 0.0 def identity_transform(ndim: int) -> npt.NDArray[np.float64]: diff --git a/src/lodstone/planner.py b/src/lodstone/planner.py index f59919c..4cb927d 100644 --- a/src/lodstone/planner.py +++ b/src/lodstone/planner.py @@ -20,11 +20,30 @@ def __init__( *, lod_bias: float = 1.0, progressive: bool = True, + max_intermediate_levels: int | None = None, + max_initial_voxel_footprint: float | None = None, ) -> None: if lod_bias <= 0: raise ValueError("lod_bias must be positive") + if max_intermediate_levels is not None and ( + not isinstance(max_intermediate_levels, int) + or isinstance(max_intermediate_levels, bool) + or max_intermediate_levels < 0 + ): + raise ValueError( + "max_intermediate_levels must be a nonnegative integer or None" + ) + if max_initial_voxel_footprint is not None and ( + not math.isfinite(max_initial_voxel_footprint) + or max_initial_voxel_footprint <= 0 + ): + raise ValueError( + "max_initial_voxel_footprint must be positive and finite or None" + ) self.lod_bias = float(lod_bias) self.progressive = bool(progressive) + self.max_intermediate_levels = max_intermediate_levels + self.max_initial_voxel_footprint = max_initial_voxel_footprint def plan( self, @@ -33,29 +52,72 @@ def plan( layout: Layout, *, available: frozenset[TileKey] = frozenset(), + previous_target_level: int | None = None, + lod_hysteresis: float = 0.0, ) -> Plan: """Return an ordered, cache-aware plan for ``view``.""" self._validate(pyramid, view, layout) - target_level = self._select_level(pyramid, view) + if not 0 <= lod_hysteresis < 1: + raise ValueError("lod_hysteresis must be in [0, 1)") + if previous_target_level is not None and not 0 <= previous_target_level < len( + pyramid.levels + ): + raise ValueError("previous_target_level is outside the pyramid") + target_level = self._select_level( + pyramid, + view, + previous_target_level=previous_target_level, + lod_hysteresis=lod_hysteresis, + ) + context_level = len(pyramid.levels) - 1 + target_depth_weight = ( + None + if layout.mixed_lod and target_level == context_level + else layout.focus_depth_weight + ) while target_level < len(pyramid.levels) - 1: target_tiles = self._tiles_for_level( - pyramid, view, layout, target_level, phase=0 + pyramid, + view, + layout, + target_level, + phase=0, + focus_depth_weight=target_depth_weight, ) - if _tiles_nbytes(target_tiles, pyramid) <= layout.memory_limit: + if layout.memory_policy == "crop": + if target_tiles: + break + elif _tiles_nbytes(target_tiles, pyramid) <= layout.memory_limit: break target_level += 1 - levels = [target_level] - if self.progressive: - levels = list(range(len(pyramid.levels) - 1, target_level - 1, -1)) + target_depth_weight = ( + None + if layout.mixed_lod and target_level == context_level + else layout.focus_depth_weight + ) + levels = self._levels(pyramid, view, target_level) + if layout.mixed_lod and context_level not in levels: + levels.insert(0, context_level) wanted: list[Tile] = [] desired: list[Tile] = [] retain: set[TileKey] = set() for phase, level_index in enumerate(levels): - tiles = self._tiles_for_level(pyramid, view, layout, level_index, phase) + tiles = self._tiles_for_level( + pyramid, + view, + layout, + level_index, + phase, + focus_depth_weight=( + target_depth_weight if level_index == target_level else None + ), + ) desired.extend(tiles) - if level_index == target_level: + if level_index == target_level or ( + layout.mixed_lod and level_index == context_level + ): retain.update(tile.key for tile in tiles) wanted.extend(tile for tile in tiles if tile.key not in available) @@ -92,9 +154,10 @@ def plan_region( if target_region.ndim != pyramid.ndim: raise ValueError("target region dimensionality does not match pyramid") - levels = [target_level] - if self.progressive: - levels = list(range(len(pyramid.levels) - 1, target_level - 1, -1)) + levels = self._levels(pyramid, view, target_level) + context_level = len(pyramid.levels) - 1 + if layout.mixed_lod and context_level not in levels: + levels.insert(0, context_level) desired: list[Tile] = [] wanted: list[Tile] = [] retain: set[TileKey] = set() @@ -123,7 +186,9 @@ def plan_region( phase, ) desired.extend(tiles) - if level_index == target_level: + if level_index == target_level or ( + layout.mixed_lod and level_index == context_level + ): retain.update(tile.key for tile in tiles) if fetch_intermediate or level_index == target_level: wanted.extend(tile for tile in tiles if tile.key not in available) @@ -196,8 +261,32 @@ def _validate(self, pyramid: Pyramid, view: View, layout: Layout) -> None: "block_shape must match displayed or complete data dimensionality" ) - def _select_level(self, pyramid: Pyramid, view: View) -> int: - threshold = self.lod_bias + def _select_level( + self, + pyramid: Pyramid, + view: View, + *, + previous_target_level: int | None = None, + lod_hysteresis: float = 0.0, + ) -> int: + selected = self._select_level_at_threshold(pyramid, view, self.lod_bias) + if ( + previous_target_level is None + or selected == previous_target_level + or lod_hysteresis == 0 + ): + return selected + factor = ( + 1 + lod_hysteresis + if selected < previous_target_level + else 1 - lod_hysteresis + ) + return self._select_level_at_threshold(pyramid, view, self.lod_bias * factor) + + @staticmethod + def _select_level_at_threshold( + pyramid: Pyramid, view: View, threshold: float + ) -> int: selected = 0 for index, level in enumerate(pyramid.levels): footprint = _voxel_footprint_px(level.voxel_to_world, level.shape, view) @@ -207,6 +296,27 @@ def _select_level(self, pyramid: Pyramid, view: View) -> int: break return selected + def _levels(self, pyramid: Pyramid, view: View, target_level: int) -> list[int]: + if not self.progressive: + return [target_level] + coarsest = len(pyramid.levels) - 1 + footprint_limit = self.max_initial_voxel_footprint + if footprint_limit is not None: + coarsest = target_level + for index in range(target_level + 1, len(pyramid.levels)): + footprint = _voxel_footprint_px( + pyramid.levels[index].voxel_to_world, + pyramid.levels[index].shape, + view, + ) + if footprint <= footprint_limit: + coarsest = index + levels = list(range(coarsest, target_level - 1, -1)) + limit = self.max_intermediate_levels + if limit is None or len(levels) <= limit + 2: + return levels + return [coarsest, *levels[-(limit + 1) :]] + def _tiles_for_level( self, pyramid: Pyramid, @@ -214,6 +324,8 @@ def _tiles_for_level( layout: Layout, level_index: int, phase: int, + *, + focus_depth_weight: float | None = None, ) -> list[Tile]: level = pyramid.levels[level_index] grids = _display_grid(layout, level, view.displayed_axes) @@ -242,6 +354,14 @@ def _tiles_for_level( key = TileKey(level_index, tuple(grid_index), selection) result.append(Tile(key, region, projected.priority, phase)) + if layout.memory_policy == "crop": + return _crop_dense_tiles( + result, + level, + layout.memory_limit, + view=view, + focus_depth_weight=focus_depth_weight, + ) return result @@ -453,6 +573,57 @@ def _tiles_in_region( return tiles +def _crop_dense_tiles( + tiles: Sequence[Tile], + level: Level, + memory_limit: int, + *, + view: View, + focus_depth_weight: float | None, +) -> list[Tile]: + """Select a priority-ordered focus window whose dense bounds fit memory.""" + + selected: list[Tile] = [] + start: tuple[int, ...] | None = None + stop: tuple[int, ...] | None = None + if focus_depth_weight is None or len(view.displayed_axes) != 3: + priority = lambda item: (item.priority, item.key.grid_index) + else: + + def priority(item: Tile) -> tuple[float, tuple[int, ...]]: + projection = _project_region(level.voxel_to_world, item.region, view) + # Work in perceptual NDC units. Euclidean screen distance grows + # linearly from the visual center, while depth is normalized from + # near=0 to far=1. A large depth weight therefore fills the canvas + # on near planes first; a small nonzero weight builds a central + # column front-to-back without falling back to data-axis ordering. + screen_distance = math.sqrt(max(0.0, projection.center_distance)) + normalized_depth = min(1.0, max(0.0, (projection.depth + 1.0) / 2.0)) + score = screen_distance + focus_depth_weight * normalized_depth + return score, item.key.grid_index + + for tile in sorted(tiles, key=priority): + candidate_start = ( + tile.region.start + if start is None + else tuple(min(a, b) for a, b in zip(start, tile.region.start, strict=True)) + ) + candidate_stop = ( + tile.region.stop + if stop is None + else tuple(max(a, b) for a, b in zip(stop, tile.region.stop, strict=True)) + ) + size = math.prod( + upper - lower + for lower, upper in zip(candidate_start, candidate_stop, strict=True) + ) + if size * level.dtype.itemsize > memory_limit: + continue + selected.append(tile) + start, stop = candidate_start, candidate_stop + return selected + + def _local_world( matrix: np.ndarray, data_points: np.ndarray, displayed_axes: tuple[int, ...] ) -> np.ndarray: @@ -496,11 +667,20 @@ def _voxel_footprint_px( class _Projection: - __slots__ = ("priority", "visible") + __slots__ = ("center_distance", "depth", "priority", "visible") - def __init__(self, visible: bool, priority: float) -> None: + def __init__( + self, + visible: bool, + priority: float, + *, + center_distance: float = math.inf, + depth: float = math.inf, + ) -> None: self.visible = visible self.priority = priority + self.center_distance = center_distance + self.depth = depth def _project_region(matrix: np.ndarray, region: Region, view: View) -> _Projection: @@ -532,5 +712,11 @@ def _project_region(matrix: np.ndarray, region: Region, view: View) -> _Projecti depth = float(np.min(clip[:, 2])) priority = depth * 1_000_000.0 + center_distance else: + depth = 0.0 priority = center_distance - return _Projection(visible, priority) + return _Projection( + visible, + priority, + center_distance=center_distance, + depth=depth, + ) diff --git a/src/lodstone/resident.py b/src/lodstone/resident.py index 5c331c5..68fe5d4 100644 --- a/src/lodstone/resident.py +++ b/src/lodstone/resident.py @@ -3,11 +3,15 @@ from __future__ import annotations from collections import defaultdict -from collections.abc import Collection, Sequence +from collections.abc import Callable, Collection, Sequence from dataclasses import dataclass, field +from threading import RLock +from time import perf_counter +from typing import Any import numpy as np +from .composition import fill_unloaded_chunks, nearest_resample_region from .model import Plan, Pyramid, Region, Tile, TileKey, Update @@ -50,7 +54,14 @@ class ResidentChange: """Updates written to one resident window.""" window: ResidentWindow - updates: tuple[Update, ...] + updates: tuple[Update, ...] = () + repaired: tuple[Region, ...] = () + + @property + def regions(self) -> tuple[Region, ...]: + """Every region changed directly or by coarse-backdrop repair.""" + + return (*tuple(update.region for update in self.updates), *self.repaired) @dataclass(frozen=True, slots=True) @@ -61,12 +72,31 @@ class ResidentTransition: retired: tuple[ResidentWindow, ...] = () +@dataclass(slots=True) +class ResidentLease: + """Dynamic lease over keys held by a :class:`ResidentArrays` instance.""" + + resident: ResidentArrays + desired_keys: frozenset[TileKey] + + @property + def available_keys(self) -> frozenset[TileKey]: + return self.resident.available_keys() + + @property + def pending_keys(self) -> frozenset[TileKey]: + return self.desired_keys - self.available_keys + + def release(self, keys: Collection[TileKey]) -> None: + self.resident.discard(keys) + + class ResidentArrays: """Manage bounded dense CPU windows for a renderer-specific target. ``prepare`` stages the complete desired ladder, ``apply`` writes updates - into those staged windows, and ``complete`` promotes the target-level - window while retiring coarse and obsolete buffers. Window arrays retain + into those staged windows, and ``complete`` promotes the retained-level + windows while retiring obsolete buffers. Window arrays retain all pyramid dimensions; targets using this helper should request ``Layout(squeeze_hidden=False)``. """ @@ -76,6 +106,8 @@ def __init__( pyramid: Pyramid, *, dtypes: Sequence[np.dtype | str | type] | None = None, + compose: bool = False, + on_timing: Callable[[str, float, int, int, Region], Any] | None = None, ) -> None: self.pyramid = pyramid self.dtypes = tuple( @@ -84,6 +116,9 @@ def __init__( ) if len(self.dtypes) != len(pyramid.levels): raise ValueError("dtypes must contain one value per pyramid level") + self.compose = bool(compose) + self.on_timing = on_timing + self.lock = RLock() self.active: dict[int, ResidentWindow] = {} self.pending: dict[int, ResidentWindow] | None = None @@ -95,13 +130,27 @@ def windows(self) -> dict[int, ResidentWindow]: @property def nbytes(self) -> int: """Physical bytes held across active and staged windows.""" - unique = {id(window): window for window in self.active.values()} - if self.pending is not None: - unique.update({id(window): window for window in self.pending.values()}) - return sum(window.nbytes for window in unique.values()) + with self.lock: + unique = {id(window): window for window in self.active.values()} + if self.pending is not None: + unique.update({id(window): window for window in self.pending.values()}) + return sum(window.nbytes for window in unique.values()) + + def available_keys(self) -> frozenset[TileKey]: + """Return a synchronized snapshot of every logically resident key.""" + with self.lock: + return frozenset( + key + for window in _unique_windows(self.active, self.pending or {}) + for key in window.key_regions + ) def prepare(self, plan: Plan) -> ResidentTransition: """Stage bounded windows covering every desired level.""" + with self.lock: + return self._prepare(plan) + + def _prepare(self, plan: Plan) -> ResidentTransition: regions = _level_regions(plan.desired or plan.wanted) previous_pending = self.pending or {} candidates = {**self.active, **previous_pending} @@ -115,10 +164,31 @@ def prepare(self, plan: Plan) -> ResidentTransition: continue window = self._allocate(level, region) if previous is not None: + started = perf_counter() _copy_overlap(previous, window) + overlap = previous.region.intersection(window.region) + if overlap is not None: + self._timing( + "overlap_copy", + started, + int(np.prod(overlap.shape)) * window.data.dtype.itemsize, + level, + overlap, + ) pending[level] = window prepared.append(window) + if self.compose: + started = perf_counter() + _compose_windows(self.pyramid, pending) + self._timing( + "composition", + started, + sum(window.nbytes for window in pending.values()), + min(pending) if pending else -1, + next(iter(pending.values())).region if pending else Region((), ()), + ) + reused = set(pending.values()) retired = tuple( window @@ -130,6 +200,10 @@ def prepare(self, plan: Plan) -> ResidentTransition: def apply(self, updates: Sequence[Update]) -> tuple[ResidentChange, ...]: """Write a batch and group renderer notifications by window.""" + with self.lock: + return self._apply(updates) + + def _apply(self, updates: Sequence[Update]) -> tuple[ResidentChange, ...]: grouped: dict[ResidentWindow, list[Update]] = defaultdict(list) windows = self.windows for update in updates: @@ -141,22 +215,60 @@ def apply(self, updates: Sequence[Update]) -> tuple[ResidentChange, ...]: ) from error window.write(update) grouped[window].append(update) + repaired: dict[ResidentWindow, list[Region]] = defaultdict(list) + if self.compose: + started = perf_counter() + for window, regions in _compose_windows( + self.pyramid, + windows, + source_levels={window.level for window in grouped}, + ).items(): + repaired[window].extend(regions) + self._timing( + "composition", + started, + sum(window.nbytes for window in windows.values()), + min(windows) if windows else -1, + next(iter(windows.values())).region if windows else Region((), ()), + ) + changed = set(grouped) | set(repaired) return tuple( - ResidentChange(window, tuple(batch)) for window, batch in grouped.items() + ResidentChange( + window, + tuple(grouped.get(window, ())), + tuple(repaired.get(window, ())), + ) + for window in sorted(changed, key=lambda item: item.level, reverse=True) ) def discard(self, keys: Collection[TileKey]) -> None: """Forget logical tile ownership without changing window storage.""" - for window in _unique_windows(self.active, self.pending or {}): - for key in keys: - window.key_regions.pop(key, None) - - def complete(self, plan: Plan) -> ResidentTransition: - """Promote the target window and retire coarse or replaced storage.""" + with self.lock: + for window in _unique_windows(self.active, self.pending or {}): + for key in keys: + window.key_regions.pop(key, None) + + def complete( + self, plan: Plan, *, retain_levels: bool = False + ) -> ResidentTransition: + """Promote target storage and optionally every retained LOD window.""" + with self.lock: + return self._complete(plan, retain_levels=retain_levels) + + def _complete( + self, plan: Plan, *, retain_levels: bool = False + ) -> ResidentTransition: if self.pending is None: return ResidentTransition() - target = self.pending.get(plan.target_level) - next_active = {} if target is None else {plan.target_level: target} + active_levels = ( + {key.level for key in plan.retain} if retain_levels else {plan.target_level} + ) + active_levels.add(plan.target_level) + next_active = { + level: window + for level, window in self.pending.items() + if level in active_levels + } retained = set(next_active.values()) retired = tuple( window @@ -169,20 +281,40 @@ def complete(self, plan: Plan) -> ResidentTransition: def clear(self) -> ResidentTransition: """Retire all active and staged windows.""" - retired = tuple(_unique_windows(self.active, self.pending or {})) - self.active = {} - self.pending = None - return ResidentTransition(retired=retired) + with self.lock: + retired = tuple(_unique_windows(self.active, self.pending or {})) + self.active = {} + self.pending = None + return ResidentTransition(retired=retired) def _allocate(self, level: int, region: Region) -> ResidentWindow: info = self.pyramid.levels[level] + started = perf_counter() data = np.full( region.shape, info.fill_value, dtype=self.dtypes[level], ) + self._timing("allocation_fill", started, data.nbytes, level, region) return ResidentWindow(level, region, data, info.voxel_to_world) + def _timing( + self, + operation: str, + started: float, + bytes_processed: int, + level: int, + region: Region, + ) -> None: + if self.on_timing is not None: + self.on_timing( + operation, + perf_counter() - started, + bytes_processed, + level, + region, + ) + def _level_regions(tiles: Sequence[Tile]) -> dict[int, Region]: bounds: dict[int, tuple[list[int], list[int]]] = {} @@ -231,6 +363,47 @@ def _copy_overlap(source: ResidentWindow, destination: ResidentWindow) -> None: ) +def _compose_windows( + pyramid: Pyramid, + windows: dict[int, ResidentWindow], + *, + source_levels: Collection[int] | None = None, +) -> dict[ResidentWindow, tuple[Region, ...]]: + """Fill unloaded fine chunks from the most detailed available backdrop.""" + + changed: dict[ResidentWindow, list[Region]] = defaultdict(list) + sources = windows if source_levels is None else source_levels + for source_level in sorted(sources, reverse=True): + source = windows[source_level] + if not source.key_regions: + continue + for destination_level in sorted( + (level for level in windows if level < source_level), reverse=True + ): + destination = windows[destination_level] + content = nearest_resample_region( + source.data, + source.region, + source.transform, + destination.region, + destination.transform, + ) + repaired = fill_unloaded_chunks( + destination.data, + destination.region, + content, + destination.region, + tuple( + pyramid.levels[destination_level].chunk_sizes(axis) + for axis in range(pyramid.ndim) + ), + destination.key_regions.values(), + ) + if repaired: + changed[destination].extend(repaired) + return {window: tuple(regions) for window, regions in changed.items()} + + def _unique_windows(*collections: dict[int, ResidentWindow]) -> list[ResidentWindow]: unique: dict[int, ResidentWindow] = {} for windows in collections: diff --git a/src/lodstone/runtime.py b/src/lodstone/runtime.py new file mode 100644 index 0000000..0a07d5a --- /dev/null +++ b/src/lodstone/runtime.py @@ -0,0 +1,107 @@ +"""Shared execution resources for one or more Lodstone streams.""" + +from __future__ import annotations + +import asyncio +import threading +from collections.abc import Callable +from concurrent.futures import Executor, ThreadPoolExecutor +from functools import partial +from typing import Any, Self, TypeVar + +T = TypeVar("T") + + +class Runtime: + """Own the asynchronous scheduler and bounded CPU staging pool. + + A runtime may be shared by multiple streams in one viewer. Sharing avoids + creating one scheduler thread and one staging pool per layer or channel and + establishes the coordination boundary for future cross-stream cache and + priority policies. Streams created without an explicit runtime retain the + original self-contained lifecycle by creating and closing their own runtime. + + Callers that supply an executor retain ownership of it. Callers that share + a runtime must close it after all of its streams have been closed. + """ + + def __init__( + self, + *, + compute_workers: int = 1, + executor: Executor | None = None, + ) -> None: + if compute_workers <= 0: + raise ValueError("compute_workers must be positive") + self._lock = threading.RLock() + self._closed = False + self._executor = executor or ThreadPoolExecutor( + max_workers=compute_workers, + thread_name_prefix="lodstone-stage", + ) + self._owns_executor = executor is None + self._loop = asyncio.new_event_loop() + self._started = threading.Event() + self._thread = threading.Thread( + target=self._run_loop, + name="lodstone-runtime", + daemon=True, + ) + self._thread.start() + self._started.wait() + + @property + def loop(self) -> asyncio.AbstractEventLoop: + """Return the runtime event loop while it is open.""" + + with self._lock: + if self._closed: + raise RuntimeError("runtime is closed") + return self._loop + + @property + def closed(self) -> bool: + with self._lock: + return self._closed + + async def run_cpu(self, function: Callable[..., T], /, *args: Any) -> T: + """Run bounded CPU staging without blocking scheduling or I/O.""" + + with self._lock: + if self._closed: + raise RuntimeError("runtime is closed") + executor = self._executor + loop = asyncio.get_running_loop() + return await loop.run_in_executor(executor, partial(function, *args)) + + def close(self) -> None: + """Stop scheduling and release runtime-owned compute threads.""" + + with self._lock: + if self._closed: + return + self._closed = True + self._loop.call_soon_threadsafe(self._loop.stop) + if threading.current_thread() is not self._thread: + self._thread.join(timeout=5) + if self._owns_executor: + self._executor.shutdown(wait=True, cancel_futures=True) + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._started.set() + self._loop.run_forever() + pending = asyncio.all_tasks(self._loop) + for task in pending: + task.cancel() + if pending: + self._loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True) + ) + self._loop.close() diff --git a/src/lodstone/stream.py b/src/lodstone/stream.py index 22ec498..f9294ef 100644 --- a/src/lodstone/stream.py +++ b/src/lodstone/stream.py @@ -4,6 +4,7 @@ import asyncio import threading +import time from collections import OrderedDict, deque from collections.abc import Callable, Collection, Sequence from concurrent.futures import Future @@ -17,6 +18,7 @@ ChunkEvent, ChunkState, Plan, + PlanDelta, Region, Status, StreamDiagnostics, @@ -26,8 +28,9 @@ View, ) from .planner import Planner +from .runtime import Runtime from .source import Source -from .target import Target +from .target import ResidencyLease, Target Dispatch = Callable[[Callable[[], None]], None] StatusCallback = Callable[[Status], None] @@ -57,6 +60,7 @@ def __init__( inflight: int = 256 << 20, batch_size: int = 8, bytes_per_second: float | None = None, + runtime: Runtime | None = None, ) -> None: if ( workers <= 0 @@ -77,20 +81,21 @@ def __init__( self.inflight_limit = inflight self.batch_size = batch_size self.bytes_per_second = bytes_per_second + self.runtime = runtime or Runtime() + self._owns_runtime = runtime is None self._state_lock = threading.RLock() self._generation = 0 self._status = Status() self._available: set[TileKey] = set() + self._lease: Any | None = None + self._plan: Plan | None = None + self._delta = PlanDelta(frozenset(), (), (), frozenset()) self._status_callbacks: list[StatusCallback] = [] self._closed = False self._paused = False - self._loop = asyncio.new_event_loop() - self._thread = threading.Thread( - target=self._run_loop, name="lodstone-stream", daemon=True - ) - self._thread.start() + self._loop = self.runtime.loop self._active: Future[Any] | None = None self._read_semaphore: asyncio.Semaphore | None = None self._resume_event: asyncio.Event | None = None @@ -116,8 +121,17 @@ def status(self) -> Status: @property def available(self) -> frozenset[TileKey]: with self._state_lock: + if self._lease is not None: + return frozenset(self._lease.available_keys) return frozenset(self._available) + @property + def delta(self) -> PlanDelta: + """Coverage changes applied by the current or most recent request.""" + + with self._state_lock: + return self._delta + @property def diagnostics(self) -> StreamDiagnostics: """Native-read counters for the current or most recent pass.""" @@ -147,16 +161,39 @@ def disconnect() -> None: return disconnect - def update(self, view: View) -> Plan: - """Plan and start streaming the newest view, returning its plan.""" - + def plan( + self, + view: View, + *, + previous_target_level: int | None = None, + lod_hysteresis: float = 0.0, + ) -> Plan: + """Plan a view without changing the active generation.""" with self._state_lock: if self._closed: raise RuntimeError("stream is closed") - available = frozenset(self._available) + # Lease-aware targets explicitly confirm storage that survives a + # replan. Legacy targets retain conservative generation behavior. + if self._lease is not None: + available = frozenset(self._lease.available_keys) + elif self._status.state == "complete": + available = frozenset(self._available) + else: + available = frozenset() layout = self.target.layout(view, self.source.pyramid) - plan = self.planner.plan(self.source.pyramid, view, layout, available=available) - return self._start(view, plan, layout) + return self.planner.plan( + self.source.pyramid, + view, + layout, + available=available, + previous_target_level=previous_target_level, + lod_hysteresis=lod_hysteresis, + ) + + def update(self, view: View) -> Plan: + """Plan and start streaming the newest view, returning its plan.""" + + return self.submit(view, self.plan(view)) def submit(self, view: View, plan: Plan) -> Plan: """Execute an adapter-supplied plan for the newest view. @@ -180,6 +217,8 @@ def _start(self, view: View, plan: Plan, layout: Any) -> Plan: raise RuntimeError("stream is closed") self._generation += 1 generation = self._generation + self._delta = plan.delta(self._plan) + self._plan = plan available = frozenset(self._available) cache_chunks = self._diagnostics.cache_chunks cache_bytes = self._diagnostics.cache_bytes @@ -275,8 +314,10 @@ def close(self) -> None: if self._active is not None: self._active.cancel() self._set_status(Status(generation=generation, state="closed")) - self._loop.call_soon_threadsafe(self._loop.stop) - self._thread.join(timeout=5) + if not self._loop.is_closed(): + self._loop.call_soon_threadsafe(self._cancel_stream_tasks) + if self._owns_runtime: + self.runtime.close() def __enter__(self) -> Self: return self @@ -284,34 +325,26 @@ def __enter__(self) -> Self: def __exit__(self, *args: object) -> None: self.close() - def _run_loop(self) -> None: - asyncio.set_event_loop(self._loop) - self._read_semaphore = asyncio.Semaphore(self.workers) - self._resume_event = asyncio.Event() - self._rate_lock = asyncio.Lock() - if not self._paused: - self._resume_event.set() - self._loop.run_forever() - pending = asyncio.all_tasks(self._loop) - for task in pending: - task.cancel() - if pending: - self._loop.run_until_complete( - asyncio.gather(*pending, return_exceptions=True) - ) - self._loop.close() - async def _execute( self, generation: int, view: View, plan: Plan, layout: Any ) -> None: completed = 0 bytes_read = 0 try: + await self._reconcile_chunk_tasks(self._native_chunk_keys(plan.wanted)) await self._wait_until_resumed() - await self._deliver_prepare(generation, view, plan) - phases = sorted({tile.phase for tile in plan.wanted}) + lease = await self._deliver_prepare(generation, view, plan) + available = ( + frozenset(lease.available_keys) if lease is not None else frozenset() + ) + wanted = tuple(tile for tile in plan.wanted if tile.key not in available) + desired = plan.desired or plan.wanted + # Presentation is a property of the logical plan, not only its + # cache misses. Targets may need to publish an already-resident + # phase into newly prepared renderer resources. + phases = sorted({tile.phase for tile in desired}) for phase in phases: - phase_tiles = [tile for tile in plan.wanted if tile.phase == phase] + phase_tiles = [tile for tile in wanted if tile.phase == phase] for window in self._tile_windows(phase_tiles): await self._wait_until_resumed() if not self._is_current(generation): @@ -336,13 +369,14 @@ async def _execute( Status( generation=generation, state="loading", - wanted=len(plan.wanted), + wanted=len(wanted), resident=len(self.available), - inflight=max(0, len(plan.wanted) - completed), + inflight=max(0, len(wanted) - completed), bytes_read=bytes_read, - progress=completed / len(plan.wanted), + progress=completed / len(wanted), ) ) + await self._deliver_phase_complete(generation, view, plan, phase) if not self._is_current(generation): return @@ -355,7 +389,7 @@ async def _execute( Status( generation=generation, state="complete", - wanted=len(plan.wanted), + wanted=len(wanted), resident=len(self.available), bytes_read=bytes_read, progress=1.0, @@ -379,6 +413,21 @@ async def _execute( ) ) + async def _reconcile_chunk_tasks( + self, desired: frozenset[tuple[int, tuple[int, ...]]] + ) -> None: + """Keep loading overlap and rebuild queued work in newest priority order.""" + + for key, task in tuple(self._chunk_tasks.items()): + state = self._chunk_states.get(key, ChunkState.NEW) + if key not in desired or state is ChunkState.QUEUED: + if self._chunk_tasks.get(key) is task: + self._chunk_tasks.pop(key, None) + task.cancel() + # Let cancellation release semaphore waiters before requesting the + # newest ordered tile windows. + await asyncio.sleep(0) + def _tile_windows(self, tiles: Sequence[Tile]) -> list[list[Tile]]: """Batch tiles without exceeding count or decoded-byte backpressure.""" @@ -511,6 +560,13 @@ def _forget_chunk_task( if self._chunk_tasks.get(key) is task: self._chunk_tasks.pop(key, None) + def _cancel_stream_tasks(self) -> None: + """Cancel reads owned by this stream without stopping a shared runtime.""" + + for task in tuple(self._chunk_tasks.values()): + task.cancel() + self._chunk_tasks.clear() + async def _fetch_chunk( self, level_index: int, @@ -523,10 +579,8 @@ async def _fetch_chunk( ) start = tuple(value for value, _stop in bounds) stop = tuple(value for _start, value in bounds) - # ``update()`` may be called immediately after construction, before the - # runtime thread has entered ``run_forever``. Initialise lazily on the - # runtime loop as well as eagerly in ``_run_loop`` to make that race - # harmless. + # Streams may share a runtime, so per-stream concurrency primitives are + # initialized lazily on that runtime's event loop. if self._read_semaphore is None: self._read_semaphore = asyncio.Semaphore(self.workers) key = (level_index, chunk_index) @@ -620,7 +674,7 @@ def _transition_chunk( ChunkEvent(generation, key, previous, current, reason) ) - def _increment_diagnostics(self, generation: int, **changes: int) -> None: + def _increment_diagnostics(self, generation: int, **changes: float) -> None: with self._state_lock: if self._diagnostics.generation != generation: return @@ -648,7 +702,15 @@ async def _deliver_updates( # stream thread so the dispatched host/UI callback only submits the # prepared rendering update. stage = getattr(self.target, "stage", None) - prepared = stage(updates) if stage is not None else updates + stage_started = time.perf_counter() + prepared = ( + await self.runtime.run_cpu(stage, updates) if stage is not None else updates + ) + if stage is not None: + self._increment_diagnostics( + generation, + update_stage_seconds=time.perf_counter() - stage_started, + ) def apply() -> None: if not self._is_current(generation): @@ -660,16 +722,40 @@ def apply() -> None: await self._run_on_target(apply) - async def _deliver_prepare(self, generation: int, view: View, plan: Plan) -> None: + async def _deliver_prepare( + self, generation: int, view: View, plan: Plan + ) -> Any | None: prepare = getattr(self.target, "prepare", None) if prepare is None: - return + return None + + stage_prepare = getattr(self.target, "stage_prepare", None) + prepared = None + if stage_prepare is not None and self._is_current(generation): + stage_started = time.perf_counter() + prepared = await self.runtime.run_cpu(stage_prepare, view, plan) + self._increment_diagnostics( + generation, + prepare_stage_seconds=time.perf_counter() - stage_started, + ) + + result: Any | None = None def run() -> None: + nonlocal result if self._is_current(generation): - prepare(view, plan) + if stage_prepare is None: + result = prepare(view, plan) + else: + result = prepare(view, plan, prepared) await self._run_on_target(run) + lease = result if isinstance(result, ResidencyLease) else None + if lease is not None and self._is_current(generation): + with self._state_lock: + self._lease = lease + self._available = set(lease.available_keys) + return lease async def _deliver_complete(self, generation: int, view: View, plan: Plan) -> None: complete = getattr(self.target, "complete", None) @@ -682,13 +768,42 @@ def run() -> None: await self._run_on_target(run) + async def _deliver_phase_complete( + self, generation: int, view: View, plan: Plan, phase: int + ) -> None: + phase_complete = getattr(self.target, "phase_complete", None) + if phase_complete is None: + return + + stage_phase = getattr(self.target, "stage_phase", None) + prepared = None + if stage_phase is not None and self._is_current(generation): + stage_started = time.perf_counter() + prepared = await self.runtime.run_cpu(stage_phase, view, plan, phase) + self._increment_diagnostics( + generation, + phase_stage_seconds=time.perf_counter() - stage_started, + ) + + def run() -> None: + if self._is_current(generation): + if stage_phase is None: + phase_complete(view, plan, phase) + else: + phase_complete(view, plan, phase, prepared) + + await self._run_on_target(run) + async def _deliver_discard( self, generation: int, keys: Collection[TileKey] ) -> None: def discard() -> None: if not self._is_current(generation): return - self.target.discard(keys) + if self._lease is None: + self.target.discard(keys) + else: + self._lease.release(keys) with self._state_lock: self._available.difference_update(keys) diff --git a/src/lodstone/target.py b/src/lodstone/target.py index b254f69..8e79868 100644 --- a/src/lodstone/target.py +++ b/src/lodstone/target.py @@ -21,11 +21,24 @@ def discard(self, keys: Collection[TileKey]) -> None: ... def redraw(self) -> None: ... +@runtime_checkable +class ResidencyLease(Protocol): + """Target-confirmed storage that survives request replanning.""" + + @property + def available_keys(self) -> frozenset[TileKey]: ... + + @property + def pending_keys(self) -> frozenset[TileKey]: ... + + def release(self, keys: Collection[TileKey]) -> None: ... + + @runtime_checkable class PassTarget(Protocol): """Optional lifecycle implemented by resident-window render targets.""" - def prepare(self, view: View, plan: Plan) -> None: + def prepare(self, view: View, plan: Plan) -> ResidencyLease | None: """Prepare target residency before the pass starts.""" ... @@ -34,6 +47,41 @@ def complete(self, view: View, plan: Plan) -> None: ... +@runtime_checkable +class PassStagingTarget(Protocol): + """Optional CPU staging performed before pass preparation dispatch.""" + + def stage_prepare(self, view: View, plan: Plan) -> Any: + """Prepare CPU residency on the stream thread.""" + ... + + def prepare(self, view: View, plan: Plan, prepared: Any) -> ResidencyLease | None: + """Apply staged preparation on the host thread.""" + ... + + +@runtime_checkable +class PhaseTarget(Protocol): + """Optional target hook invoked after one progressive phase is delivered.""" + + def phase_complete(self, view: View, plan: Plan, phase: int) -> None: + """Present or reconcile a completed coarse-to-fine phase.""" + ... + + +@runtime_checkable +class PhaseStagingTarget(Protocol): + """Optional CPU staging performed before phase publication dispatch.""" + + def stage_phase(self, view: View, plan: Plan, phase: int) -> Any: + """Build an immutable phase publication on the stream thread.""" + ... + + def phase_complete(self, view: View, plan: Plan, phase: int, prepared: Any) -> None: + """Publish a staged phase on the host thread.""" + ... + + @runtime_checkable class StagingTarget(Protocol): """Optional target hook for CPU preparation before host dispatch.""" diff --git a/tests/test_model.py b/tests/test_model.py index 2300283..91d6f28 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -3,9 +3,14 @@ import numpy as np import pytest +import lodstone from lodstone import Level, Pyramid, Region, View, identity_transform +def test_package_exposes_version() -> None: + assert lodstone.__version__ + + def test_region_intersection_and_slices() -> None: region = Region((2, 3), (8, 9)) overlap = region.intersection(Region((0, 7), (5, 12))) diff --git a/tests/test_ndv_adapter.py b/tests/test_ndv_adapter.py new file mode 100644 index 0000000..85ef4e8 --- /dev/null +++ b/tests/test_ndv_adapter.py @@ -0,0 +1,352 @@ +from __future__ import annotations + +import time + +import numpy as np + +from lodstone import Layout, View +from lodstone.adapters.ndv import NDVController +from lodstone.testing import SimulatedSource + + +class Handle: + def __init__(self) -> None: + self.history = [] + self.removed = False + self.visible = True + self.scales = None + self.origins = None + self.clims = None + + def set_data(self, data) -> None: + self.history.append(np.asarray(data)) + + def set_world_transform(self, scales, origins) -> None: + self.scales = scales + self.origins = origins + + def set_clims(self, clims) -> None: + self.clims = clims + + def set_visible(self, visible) -> None: + self.visible = visible + + def remove(self) -> None: + self.removed = True + + +class Canvas: + def __init__(self) -> None: + self.handle = None + self.handles = [] + self.kinds = [] + self.ndim = None + self.scales = None + self.origins = None + self.range_resets = 0 + self.refreshes = 0 + + def set_ndim(self, ndim) -> None: + self.ndim = ndim + + def add_image(self, data=None, *, reset_range=True): + self.kinds.append("image") + self.handle = Handle() + self.handles.append(self.handle) + if data is not None: + self.handle.set_data(data) + return self.handle + + def add_volume(self, data=None, *, reset_range=True): + self.kinds.append("volume") + self.handle = Handle() + self.handles.append(self.handle) + if data is not None: + self.handle.set_data(data) + return self.handle + + def set_scales(self, scales, *, reset_range=True) -> None: + self.scales = scales + + def set_origins(self, origins) -> None: + self.origins = origins + + def set_range(self) -> None: + self.range_resets += 1 + + def refresh(self) -> None: + self.refreshes += 1 + + +class Event: + def __init__(self) -> None: + self.callback = None + + def connect(self, callback) -> None: + self.callback = callback + + def disconnect(self, callback) -> None: + if self.callback == callback: + self.callback = None + + def emit(self) -> None: + if self.callback is not None: + self.callback() + + +class CameraCanvas(Canvas): + def __init__(self) -> None: + super().__init__() + self.cameraChanged = Event() + self.world_to_clip = np.eye(4) + + def camera_state(self): + return (64, 64), self.world_to_clip.copy() + + +class PublishingCameraCanvas(CameraCanvas): + def add_image(self, data=None, *, reset_range=True): + handle = super().add_image(data, reset_range=reset_range) + original = handle.set_world_transform + + def set_world_transform(scales, origins): + original(scales, origins) + self.cameraChanged.emit() + + handle.set_world_transform = set_world_transform + return handle + + +def test_ndv_target_presents_dense_camera_phase(ortho_view, wait) -> None: + fine = np.arange(64, dtype=np.uint16).reshape(8, 8) + coarse = fine[::2, ::2] + source = SimulatedSource( + [fine, coarse], + transforms=[np.eye(3), np.diag([2.0, 2.0, 1.0])], + chunks=[(4, 4), (4, 4)], + ) + canvas = Canvas() + controller = NDVController(canvas, source) + try: + controller.update(ortho_view(fine.shape, viewport=(64, 64))) + wait(lambda: controller.stream.status.state == "complete") + + assert canvas.kinds == ["image", "image"] + assert canvas.ndim == 2 + assert len(canvas.handle.history) == 1 + np.testing.assert_array_equal(canvas.handle.history[0], fine) + assert canvas.handle.scales == (1.0, 1.0) + assert canvas.handle.origins == (0.0, 0.0) + assert canvas.handle.clims == (0.0, 63.0) + assert canvas.range_resets == 1 + finally: + handle = canvas.handle + controller.close() + + assert handle.removed + + +def test_ndv_target_squeezes_hidden_axes_for_volume(ortho_view, wait) -> None: + data = np.arange(2 * 4 * 5 * 6, dtype=np.uint16).reshape(2, 4, 5, 6) + transform = np.diag([1.0, 2.0, 3.0, 4.0, 1.0]) + source = SimulatedSource( + [data], + transforms=[transform], + chunks=[(1, 4, 5, 6)], + ) + canvas = Canvas() + controller = NDVController(canvas, source) + view = ortho_view( + data.shape, + displayed_axes=(1, 2, 3), + index=(1, None, None, None), + viewport=(64, 64), + ) + try: + controller.update(view) + wait(lambda: controller.stream.status.state == "complete") + + assert canvas.kinds == ["volume"] + assert canvas.ndim == 3 + np.testing.assert_array_equal(canvas.handle.history[-1], data[1]) + assert canvas.handle.scales == (2.0, 3.0, 4.0) + assert canvas.handle.origins == (0.0, 0.0, 0.0) + finally: + controller.close() + + +def test_ndv_target_places_translated_dense_window(wait) -> None: + data = np.arange(32 * 32, dtype=np.uint16).reshape(32, 32) + source = SimulatedSource([data], chunks=[(8, 8)]) + canvas = Canvas() + presented = [] + controller = NDVController( + canvas, + source, + memory_limit=512, + block_shape=(8, 8), + on_presented=presented.append, + ) + world_to_clip = np.eye(4) + world_to_clip[:2, -1] = -15 + view = View( + displayed_axes=(0, 1), + index=(None, None), + viewport=(64, 64), + world_to_clip=world_to_clip, + ) + try: + controller.update(view) + wait(lambda: controller.stream.status.state == "complete") + + assert canvas.handle.origins == (8.0, 8.0) + assert canvas.handle.history[-1].shape == (16, 16) + assert presented[-1].region.start == (8, 8) + finally: + controller.close() + + +def test_ndv_target_keeps_masked_coarse_context_around_fine_focus( + ortho_view, wait +) -> None: + fine = np.full((32, 32), 7, dtype=np.uint16) + coarse = np.full((8, 8), 3, dtype=np.uint16) + source = SimulatedSource( + [fine, coarse], + transforms=[np.eye(3), np.diag([4.0, 4.0, 1.0])], + chunks=[(8, 8), (4, 4)], + ) + canvas = Canvas() + controller = NDVController( + canvas, + source, + memory_limit=512, + block_shape=(8, 8), + ) + try: + controller.update(ortho_view(fine.shape, viewport=(128, 128))) + wait(lambda: controller.stream.status.state == "complete") + + assert canvas.kinds == ["image", "image"] + context, focus = canvas.handles + assert context.scales == (4.0, 4.0) + assert focus.scales == (1.0, 1.0) + assert np.count_nonzero(context.history[-1]) < context.history[-1].size + assert np.count_nonzero(context.history[-1]) > 0 + assert np.all(focus.history[-1] == 7) + assert canvas.range_resets == 1 + finally: + controller.close() + + +def test_ndv_coarse_mask_preserves_context_beneath_empty_fine_space( + ortho_view, wait +) -> None: + fine = np.zeros((32, 32), dtype=np.uint16) + fine[16, 16] = 7 + coarse = np.full((8, 8), 3, dtype=np.uint16) + source = SimulatedSource( + [fine, coarse], + transforms=[np.eye(3), np.diag([4.0, 4.0, 1.0])], + chunks=[(8, 8), (4, 4)], + ) + canvas = Canvas() + controller = NDVController( + canvas, + source, + memory_limit=512, + block_shape=(8, 8), + ) + try: + controller.update(ortho_view(fine.shape, viewport=(128, 128))) + wait(lambda: controller.stream.status.state == "complete") + + context, focus = canvas.handles + assert np.count_nonzero(focus.history[-1]) == 1 + assert np.count_nonzero(context.history[-1]) == context.history[-1].size - 1 + finally: + controller.close() + + +def test_ndv_camera_updates_are_debounced_off_interaction_thread(ortho_view, wait): + data = np.arange(32 * 32, dtype=np.uint16).reshape(32, 32) + source = SimulatedSource([data], chunks=[(8, 8)]) + canvas = CameraCanvas() + targets = [] + controller = NDVController( + canvas, + source, + memory_limit=512, + block_shape=(8, 8), + camera_debounce_ms=30, + on_targeted=lambda plan: targets.append(plan.target_level), + ) + try: + controller.update(ortho_view(data.shape, viewport=(64, 64))) + wait(lambda: controller.stream.status.state == "complete") + targets.clear() + + for offset in range(10): + canvas.world_to_clip[:2, -1] = -offset + canvas.cameraChanged.emit() + + # Camera callbacks only capture the newest view; planning happens after + # the trailing debounce on the dedicated camera worker. + assert targets == [] + wait(lambda: len(targets) == 1) + assert targets == [0] + canvas.cameraChanged.emit() + time.sleep(0.06) + assert targets == [0] + finally: + controller.close() + + +def test_ndv_ignores_camera_events_caused_by_its_own_publication(ortho_view, wait): + source = SimulatedSource([np.ones((16, 16), np.uint8)], chunks=[(8, 8)]) + canvas = PublishingCameraCanvas() + targets = [] + controller = NDVController( + canvas, + source, + camera_debounce_ms=10, + on_targeted=lambda plan: targets.append(plan.target_level), + ) + try: + controller.update(ortho_view((16, 16), viewport=(64, 64))) + wait(lambda: controller.stream.status.state == "complete") + time.sleep(0.04) + assert targets == [0] + finally: + controller.close() + + +def test_ndv_focus_policy_can_be_tuned_and_replanned(ortho_view, wait) -> None: + source = SimulatedSource([np.ones((16, 16), np.uint8)], chunks=[(8, 8)]) + controller = NDVController(Canvas(), source) + view = ortho_view((16, 16), viewport=(64, 64)) + try: + controller.update(view) + wait(lambda: controller.stream.status.state == "complete") + plan = controller.set_focus_policy( + memory_limit=32 * 1024**2, + focus_depth_weight=0.25, + lod_bias=1.5, + ) + + assert plan is not None + assert controller.target.memory_limit == 32 * 1024**2 + assert controller.target.focus_depth_weight == 0.25 + assert controller.stream.planner.lod_bias == 1.5 + assert controller.target.layout(view, source.pyramid) == Layout( + kind="dense", + block_shape=(128, 128), + mixed_lod=True, + memory_limit=32 * 1024**2, + squeeze_hidden=False, + max_axis_extent=512, + memory_policy="crop", + focus_depth_weight=0.25, + ) + finally: + controller.close() diff --git a/tests/test_planner.py b/tests/test_planner.py index d8b654c..c924a12 100644 --- a/tests/test_planner.py +++ b/tests/test_planner.py @@ -1,11 +1,17 @@ from __future__ import annotations +from dataclasses import replace + import numpy as np +import pytest from lodstone import ( Layout, + Plan, Planner, Region, + Tile, + TileKey, available_tile_keys, merge_plans, plan_from_slices, @@ -13,6 +19,25 @@ from lodstone.sources import ArrayPyramidSource +def test_plan_delta_retains_coverage_and_reports_priority_changes() -> None: + key0 = TileKey(0, (0, 0), ()) + key1 = TileKey(0, (0, 1), ()) + key2 = TileKey(0, (0, 2), ()) + tile0 = Tile(key0, Region((0, 0), (4, 4)), 1.0) + tile1 = Tile(key1, Region((0, 4), (4, 8)), 2.0) + first = Plan((tile0, tile1), frozenset({key0, key1}), 0, (tile0, tile1)) + moved0 = replace(tile0, priority=-1.0) + tile2 = Tile(key2, Region((0, 8), (4, 12)), 3.0) + second = Plan((moved0, tile2), frozenset({key0, key2}), 0, (moved0, tile2)) + + delta = second.delta(first) + + assert delta.retained == frozenset({key0}) + assert delta.requested == (tile2,) + assert delta.reprioritized == (key0,) + assert delta.released == frozenset({key1}) + + def _pyramid() -> ArrayPyramidSource: return ArrayPyramidSource( [ @@ -64,6 +89,114 @@ def test_progressive_plan_puts_coarse_coverage_first(ortho_view) -> None: assert {tile.level for tile in plan.desired} == {0, 1, 2} +def test_progressive_plan_can_skip_intermediate_levels(ortho_view) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + plan = Planner(progressive=True, max_intermediate_levels=0).plan( + source.pyramid, view, Layout(block_shape=(32, 32)) + ) + + assert [tile.level for tile in plan.wanted] == sorted( + (tile.level for tile in plan.wanted), reverse=True + ) + assert {tile.level for tile in plan.desired} == {0, 2} + + +def test_progressive_plan_can_adapt_initial_level_to_voxel_footprint( + ortho_view, +) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + plan = Planner( + progressive=True, + max_intermediate_levels=0, + max_initial_voxel_footprint=4.1, + ).plan(source.pyramid, view, Layout(block_shape=(32, 32))) + + assert {tile.level for tile in plan.desired} == {0, 1} + assert plan.desired[0].level == 1 + + +def test_progressive_initial_level_remains_coarsest_by_default(ortho_view) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + plan = Planner(progressive=True, max_intermediate_levels=0).plan( + source.pyramid, view, Layout(block_shape=(32, 32)) + ) + + assert {tile.level for tile in plan.desired} == {0, 2} + + +def test_mixed_lod_retains_coarsest_context_with_target(ortho_view) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + plan = Planner( + progressive=True, + max_intermediate_levels=0, + max_initial_voxel_footprint=1.1, + ).plan( + source.pyramid, + view, + Layout(block_shape=(32, 32), mixed_lod=True), + ) + + assert plan.target_level == 0 + assert {tile.level for tile in plan.desired} == {0, 2} + assert {key.level for key in plan.retain} == {0, 2} + + +@pytest.mark.parametrize("value", [0, -1, float("inf"), float("nan")]) +def test_progressive_initial_voxel_footprint_must_be_positive_and_finite(value) -> None: + with pytest.raises(ValueError, match="max_initial_voxel_footprint"): + Planner(max_initial_voxel_footprint=value) + + +def test_lod_hysteresis_resists_small_level_boundary_crossings(ortho_view) -> None: + source = _pyramid() + layout = Layout(block_shape=(32, 32)) + slightly_inside_fine = ortho_view((256, 256), viewport=(512, 512), extent_scale=3.9) + slightly_inside_coarse = ortho_view( + (256, 256), viewport=(512, 512), extent_scale=4.1 + ) + + assert ( + Planner(progressive=False) + .plan(source.pyramid, slightly_inside_fine, layout) + .target_level + == 0 + ) + assert ( + Planner(progressive=False) + .plan( + source.pyramid, + slightly_inside_fine, + layout, + previous_target_level=1, + lod_hysteresis=0.2, + ) + .target_level + == 1 + ) + assert ( + Planner(progressive=False) + .plan(source.pyramid, slightly_inside_coarse, layout) + .target_level + == 1 + ) + assert ( + Planner(progressive=False) + .plan( + source.pyramid, + slightly_inside_coarse, + layout, + previous_target_level=0, + lod_hysteresis=0.2, + ) + .target_level + == 0 + ) + + def test_gpu_budget_can_select_a_coarser_level(ortho_view) -> None: source = _pyramid() view = ortho_view((256, 256), viewport=(512, 512)) @@ -75,6 +208,114 @@ def test_gpu_budget_can_select_a_coarser_level(ortho_view) -> None: assert plan.target_level == 1 +def test_dense_crop_budget_preserves_camera_selected_level(ortho_view) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + memory_limit = 128 * 128 * 2 + plan = Planner(progressive=False).plan( + source.pyramid, + view, + Layout( + block_shape=(32, 32), + memory_limit=memory_limit, + memory_policy="crop", + ), + ) + + assert plan.target_level == 0 + start = tuple( + min(tile.region.start[axis] for tile in plan.desired) for axis in range(2) + ) + stop = tuple( + max(tile.region.stop[axis] for tile in plan.desired) for axis in range(2) + ) + assert np.prod(np.subtract(stop, start)) * 2 <= memory_limit + + +def test_dense_crop_can_balance_focus_depth_against_screen_center(ortho_view) -> None: + shape = (64, 64, 64) + source = ArrayPyramidSource( + [np.zeros(shape, dtype=np.uint8)], + axes=("z", "y", "x"), + chunks=[(16, 16, 16)], + ) + view = ortho_view(shape, viewport=(256, 256)) + planner = Planner(progressive=False) + common = { + "block_shape": (16, 16, 16), + "memory_limit": 8 * 16**3, + "memory_policy": "crop", + } + + slab = planner.plan(source.pyramid, view, Layout(**common)) + volume = planner.plan( + source.pyramid, + view, + Layout(**common, focus_depth_weight=0.5), + ) + + def extent(plan, axis): + return max(tile.region.stop[axis] for tile in plan.desired) - min( + tile.region.start[axis] for tile in plan.desired + ) + + # This camera maps data axis 2 to clip depth. The balanced focus spends + # the same byte budget on two depth layers instead of one front slab. + assert extent(slab, 2) == 16 + assert extent(volume, 2) == 32 + assert sum(tile.region.size for tile in slab.desired) == sum( + tile.region.size for tile in volume.desired + ) + + +def test_dense_crop_tunes_canvas_coverage_against_depth_reach(ortho_view) -> None: + shape = (64, 64, 64) + source = ArrayPyramidSource( + [np.zeros(shape, dtype=np.uint8)], + axes=("z", "y", "x"), + chunks=[(16, 16, 16)], + ) + view = ortho_view(shape, viewport=(256, 256)) + planner = Planner(progressive=False) + common = { + "block_shape": (16, 16, 16), + "memory_limit": 8 * 16**3, + "memory_policy": "crop", + } + + canvas = planner.plan( + source.pyramid, + view, + Layout(**common, focus_depth_weight=8.0), + ) + depth = planner.plan( + source.pyramid, + view, + Layout(**common, focus_depth_weight=0.5), + ) + + def extents(plan): + return tuple( + max(tile.region.stop[axis] for tile in plan.desired) + - min(tile.region.start[axis] for tile in plan.desired) + for axis in range(3) + ) + + canvas_extents = extents(canvas) + depth_extents = extents(depth) + # This camera maps axis 2 to depth. The canvas-heavy policy spends more of + # the same dense budget in axes 0/1; the depth-heavy policy reaches farther + # along axis 2 around the screen center. + assert canvas_extents[0] * canvas_extents[1] > (depth_extents[0] * depth_extents[1]) + assert canvas_extents[2] < depth_extents[2] + + +@pytest.mark.parametrize("value", [-1, float("inf"), float("nan")]) +def test_focus_depth_weight_must_be_finite_and_nonnegative(value) -> None: + with pytest.raises(ValueError, match="focus_depth_weight"): + Layout(focus_depth_weight=value) + + def test_available_tiles_are_retained_but_not_requested(ortho_view) -> None: source = _pyramid() view = ortho_view((256, 256), viewport=(512, 512)) @@ -92,6 +333,55 @@ def test_available_tiles_are_retained_but_not_requested(ortho_view) -> None: assert available <= {tile.key for tile in updated.desired} +def test_plan_coverage_ignores_tile_order_priority_and_phase(ortho_view) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + plan = Planner(progressive=True).plan( + source.pyramid, view, Layout(block_shape=(32, 32)) + ) + reordered = replace( + plan, + desired=tuple( + replace(tile, priority=-tile.priority, phase=tile.phase + 10) + for tile in reversed(plan.desired) + ), + ) + + assert reordered.coverage == plan.coverage + + +def test_plan_coverage_detects_regions_retention_and_hidden_selection( + ortho_view, +) -> None: + source = _pyramid() + view = ortho_view((256, 256), viewport=(512, 512)) + plan = Planner(progressive=False).plan( + source.pyramid, view, Layout(block_shape=(32, 32)) + ) + tile = plan.desired[0] + shifted_region = Region( + (tile.region.start[0] + 1, *tile.region.start[1:]), + tile.region.stop, + ) + changed_region = replace( + plan, + desired=(replace(tile, region=shifted_region), *plan.desired[1:]), + ) + changed_retention = replace(plan, retain=frozenset()) + selected_key = TileKey(tile.level, tile.key.grid_index, (4, -1)) + changed_selection = replace( + plan, + desired=( + Tile(selected_key, tile.region, tile.priority, tile.phase), + *plan.desired[1:], + ), + ) + + assert changed_region.coverage != plan.coverage + assert changed_retention.coverage != plan.coverage + assert changed_selection.coverage != plan.coverage + + def test_3d_tiles_are_prioritized_front_to_back(ortho_view) -> None: source = ArrayPyramidSource( [np.zeros((16, 16, 16), dtype=np.uint8)], diff --git a/tests/test_resident.py b/tests/test_resident.py index 5c6d33e..efb00d3 100644 --- a/tests/test_resident.py +++ b/tests/test_resident.py @@ -6,6 +6,7 @@ Plan, Region, ResidentArrays, + ResidentLease, Tile, TileKey, Update, @@ -103,6 +104,28 @@ def test_shifted_window_preserves_overlap_and_loaded_keys() -> None: assert first.key not in shifted.key_regions +def test_resident_lease_tracks_confirmed_storage_and_release() -> None: + source = ArrayPyramidSource([np.zeros((20, 20), dtype=np.uint8)], chunks=[(10, 10)]) + arrays = ResidentArrays(source.pyramid) + first = _tile(0, (0, 0), (10, 10)) + second = _tile(0, (0, 10), (10, 20)) + plan = _plan(first, second) + arrays.prepare(plan) + lease = ResidentLease(arrays, frozenset({first.key, second.key})) + + assert lease.available_keys == frozenset() + assert lease.pending_keys == frozenset({first.key, second.key}) + + arrays.apply( + [Update(first.key, first.region, np.ones((10, 10), np.uint8), np.eye(3))] + ) + assert lease.available_keys == frozenset({first.key}) + assert lease.pending_keys == frozenset({second.key}) + + lease.release({first.key}) + assert lease.available_keys == frozenset() + + def test_complete_keeps_target_and_retires_coarse_ladder() -> None: source = ArrayPyramidSource( [ @@ -157,3 +180,57 @@ def test_full_nd_hidden_axis_updates_are_required() -> None: ) ] ) + + +def test_composed_residency_repairs_only_unloaded_fine_chunks() -> None: + fine_transform = np.eye(3) + coarse_transform = np.diag([2.0, 2.0, 1.0]) + source = ArrayPyramidSource( + [ + np.zeros((8, 8), dtype=np.uint8), + np.zeros((4, 4), dtype=np.uint8), + ], + chunks=[(4, 4), (2, 2)], + transforms=[fine_transform, coarse_transform], + ) + arrays = ResidentArrays(source.pyramid, compose=True) + coarse = _tile(1, (0, 0), (4, 4), phase=0) + fine_left = _tile(0, (0, 0), (4, 4), phase=1) + fine_right = _tile(0, (0, 4), (4, 8), phase=1) + plan = _plan(coarse, fine_left, fine_right) + arrays.prepare(plan) + + coarse_update = Update( + coarse.key, + coarse.region, + np.full((4, 4), 3, dtype=np.uint8), + coarse_transform, + ) + changes = arrays.apply([coarse_update]) + + fine_window = arrays.windows[0] + fine_change = next(change for change in changes if change.window is fine_window) + assert fine_change.updates == () + assert fine_change.repaired + assert np.all(fine_window.data == 3) + + fine_update = Update( + fine_left.key, + fine_left.region, + np.full((4, 4), 9, dtype=np.uint8), + fine_transform, + ) + arrays.apply([fine_update]) + arrays.apply( + [ + Update( + coarse.key, + coarse.region, + np.full((4, 4), 5, dtype=np.uint8), + coarse_transform, + ) + ] + ) + + assert np.all(fine_window.data[:, :4] == 9) + assert np.all(fine_window.data[:, 4:] == 5) diff --git a/tests/test_stream.py b/tests/test_stream.py index 3f2f79c..fd41a07 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import threading import time from collections.abc import Callable @@ -13,6 +14,9 @@ Plan, Planner, Region, + ResidentArrays, + ResidentLease, + Runtime, Stream, Tile, TileKey, @@ -37,6 +41,113 @@ def test_stream_reuses_native_chunks_for_smaller_display_tiles( stream.close() +def test_plan_does_not_start_or_replace_a_generation(ortho_view) -> None: + source = SimulatedSource([np.zeros((8, 8), dtype=np.uint8)], chunks=[(4, 4)]) + target = RecordingTarget(Layout(block_shape=(4, 4))) + stream = Stream(source, target, planner=Planner(progressive=False)) + try: + plan = stream.plan(ortho_view((8, 8), viewport=(64, 64))) + + assert plan.wanted + assert stream.status.generation == 0 + assert stream.status.state == "idle" + assert target.updates == [] + finally: + stream.close() + + +def test_lease_confirmed_tiles_are_reused_while_request_is_loading( + ortho_view, wait +) -> None: + class LeaseTarget(RecordingTarget): + def __init__(self, source) -> None: + super().__init__(Layout(block_shape=(4, 4), squeeze_hidden=False)) + self.resident = ResidentArrays(source.pyramid) + + def prepare(self, _view, plan): + self.resident.prepare(plan) + desired = plan.desired or plan.wanted + return ResidentLease(self.resident, frozenset(tile.key for tile in desired)) + + def apply(self, updates) -> None: + self.resident.apply(updates) + super().apply(updates) + + def discard(self, keys) -> None: + self.resident.discard(keys) + super().discard(keys) + + def complete(self, _view, plan) -> None: + self.resident.complete(plan) + + source = SimulatedSource( + [np.zeros((8, 8), dtype=np.uint8)], chunks=[(4, 4)], latency=0.04 + ) + target = LeaseTarget(source) + stream = Stream( + source, + target, + planner=Planner(progressive=False), + workers=1, + batch_size=1, + ) + view = ortho_view((8, 8), viewport=(64, 64)) + try: + stream.update(view) + wait(lambda: stream.status.state == "loading" and len(target.updates) == 1) + + replacement = stream.plan(view) + + assert len(replacement.wanted) == 3 + assert target.updates[0].key not in {tile.key for tile in replacement.wanted} + finally: + stream.close() + + +def test_replan_keeps_loading_overlap_and_cancels_obsolete_queue( + ortho_view, wait +) -> None: + source = SimulatedSource( + [np.zeros((4, 12), dtype=np.uint8)], chunks=[(4, 4)], latency=0.08 + ) + target = RecordingTarget(Layout(block_shape=(4, 4))) + stream = Stream( + source, + target, + planner=Planner(progressive=False), + workers=1, + batch_size=3, + ) + view = ortho_view((4, 12), viewport=(64, 64)) + tiles = tuple( + Tile( + TileKey(0, (0, index), ()), + Region((0, index * 4), (4, (index + 1) * 4)), + float(index), + ) + for index in range(3) + ) + first = Plan(tiles, frozenset(tile.key for tile in tiles), 0, tiles) + replacement_tiles = (tiles[0], tiles[2]) + replacement = Plan( + replacement_tiles, + frozenset(tile.key for tile in replacement_tiles), + 0, + replacement_tiles, + ) + try: + stream.submit(view, first) + wait(lambda: stream.chunk_states.get((0, (0, 0))) is ChunkState.LOADING) + stream.submit(view, replacement) + wait(lambda: stream.status.state == "complete") + + assert source.reads.count((0, tiles[0].region)) == 1 + assert (0, tiles[1].region) not in source.reads + assert (0, tiles[2].region) in source.reads + finally: + stream.close() + + def test_stream_diagnostics_separate_tiles_from_native_reads(ortho_view, wait) -> None: data = np.arange(64, dtype=np.uint16).reshape(8, 8) source = SimulatedSource([data], chunks=[(4, 4)]) @@ -198,6 +309,86 @@ def apply(self, prepared) -> None: stream.close() +def test_target_stages_phase_before_host_dispatch(ortho_view, wait) -> None: + calls: Queue[Callable[[], None]] = Queue() + + class PhaseStagingTarget(RecordingTarget): + def __init__(self) -> None: + super().__init__(Layout(kind="tiled", block_shape=(4, 4))) + self.stage_thread: int | None = None + self.publish_thread: int | None = None + self.prepared = None + + def stage_phase(self, view, plan, phase): + self.stage_thread = threading.get_ident() + return (plan.target_level, phase) + + def phase_complete(self, view, plan, phase, prepared) -> None: + self.publish_thread = threading.get_ident() + self.prepared = prepared + + target = PhaseStagingTarget() + source = SimulatedSource([np.zeros((4, 4), dtype=np.uint8)], chunks=[(4, 4)]) + stream = Stream(source, target, dispatch=calls.put) + main_thread = threading.get_ident() + try: + stream.update(ortho_view((4, 4), viewport=(64, 64))) + deadline = time.monotonic() + 5 + while stream.status.state != "complete": + assert time.monotonic() < deadline + try: + callback = calls.get(timeout=0.05) + except Empty: + continue + callback() + assert target.stage_thread is not None + assert target.stage_thread != main_thread + assert target.publish_thread == main_thread + assert target.prepared == (0, 0) + assert stream.diagnostics.phase_stage_seconds > 0 + finally: + stream.close() + + +def test_target_stages_prepare_before_host_dispatch(ortho_view) -> None: + calls: Queue[Callable[[], None]] = Queue() + + class PrepareStagingTarget(RecordingTarget): + def __init__(self) -> None: + super().__init__(Layout(kind="tiled", block_shape=(4, 4))) + self.stage_thread: int | None = None + self.prepare_thread: int | None = None + + def stage_prepare(self, view, plan): + self.stage_thread = threading.get_ident() + return plan.target_level + + def prepare(self, view, plan, prepared) -> None: + self.prepare_thread = threading.get_ident() + assert prepared == plan.target_level + + target = PrepareStagingTarget() + source = SimulatedSource([np.zeros((4, 4), dtype=np.uint8)], chunks=[(4, 4)]) + stream = Stream(source, target, dispatch=calls.put) + main_thread = threading.get_ident() + try: + stream.update(ortho_view((4, 4), viewport=(64, 64))) + deadline = time.monotonic() + 5 + while stream.status.state != "complete": + assert time.monotonic() < deadline + try: + callback = calls.get(timeout=0.05) + except Empty: + continue + callback() + assert target.stage_thread is not None + assert target.stage_thread != main_thread + assert target.prepare_thread == main_thread + assert stream.diagnostics.prepare_stage_seconds > 0 + finally: + stream.close() + + def test_stream_reads_rectilinear_native_chunks(ortho_view, wait) -> None: data = np.arange(7 * 9, dtype=np.uint16).reshape(7, 9) source = SimulatedSource( @@ -296,6 +487,26 @@ def test_new_generation_rejects_stale_delivery(ortho_view, wait) -> None: stream.close() +def test_loading_replan_redelivers_tiles_from_superseded_pass(ortho_view, wait) -> None: + data = np.arange(16 * 16, dtype=np.uint16).reshape(16, 16) + source = SimulatedSource([data], chunks=[(4, 4)], latency=0.05) + target = RecordingTarget(Layout(kind="tiled", block_shape=(4, 4))) + stream = Stream(source, target, planner=Planner(progressive=False), batch_size=1) + view = ortho_view(data.shape, viewport=(128, 128)) + try: + first = stream.update(view) + wait(lambda: bool(stream.available) and stream.status.state == "loading") + + second = stream.update(view) + + assert len(stream.available) < len(first.desired) + assert len(second.wanted) == len(second.desired) + wait(lambda: stream.status.state == "complete") + assert len(target.updates) >= len(second.desired) + finally: + stream.close() + + def test_status_callbacks_observe_completion(ortho_view, wait) -> None: source = SimulatedSource([np.zeros((8, 8), dtype=np.uint8)], chunks=[(4, 4)]) target = RecordingTarget(Layout(block_shape=(4, 4))) @@ -438,6 +649,144 @@ def complete(self, view, plan) -> None: stream.close() +def test_phase_lifecycle_follows_each_progressive_phase(ortho_view, wait) -> None: + class PhaseRecordingTarget(RecordingTarget): + def __init__(self) -> None: + super().__init__(Layout(block_shape=(4, 4))) + self.phases = [] + + def phase_complete(self, view, plan, phase) -> None: + self.phases.append(phase) + + source = SimulatedSource( + [ + np.zeros((8, 8), dtype=np.uint8), + np.zeros((4, 4), dtype=np.uint8), + ], + transforms=[np.eye(3), np.diag([2.0, 2.0, 1.0])], + chunks=[(4, 4), (4, 4)], + ) + target = PhaseRecordingTarget() + stream = Stream(source, target, planner=Planner(progressive=True)) + try: + stream.update(ortho_view((8, 8), viewport=(64, 64))) + wait(lambda: stream.status.state == "complete") + assert target.phases == [0, 1] + finally: + stream.close() + + +def test_phase_lifecycle_presents_already_resident_phases(ortho_view, wait) -> None: + class PhaseRecordingTarget(RecordingTarget): + def __init__(self) -> None: + super().__init__(Layout(block_shape=(4, 4))) + self.phases: list[int] = [] + + def phase_complete(self, view, plan, phase) -> None: + self.phases.append(phase) + + source = SimulatedSource( + [ + np.zeros((8, 8), dtype=np.uint8), + np.zeros((4, 4), dtype=np.uint8), + ], + transforms=[np.eye(3), np.diag([2.0, 2.0, 1.0])], + chunks=[(4, 4), (4, 4)], + ) + target = PhaseRecordingTarget() + stream = Stream(source, target, planner=Planner(progressive=True)) + view = ortho_view((8, 8), viewport=(64, 64)) + try: + stream.update(view) + wait(lambda: stream.status.state == "complete") + first_reads = tuple(source.reads) + + cached = stream.plan(view) + assert {tile.phase for tile in cached.wanted} == {0} + assert any(tile.phase == 1 for tile in cached.desired) + assert not any(tile.phase == 1 for tile in cached.wanted) + stream.submit(view, cached) + wait( + lambda: stream.status.state == "complete" and stream.status.generation == 2 + ) + + assert target.phases == [0, 1, 0, 1] + assert tuple(source.reads) == first_reads + finally: + stream.close() + + +def test_shared_runtime_stages_without_blocking_other_streams(ortho_view, wait) -> None: + stage_started = threading.Event() + release_stage = threading.Event() + + class BlockingTarget(RecordingTarget): + def stage(self, updates): + stage_started.set() + assert release_stage.wait(timeout=5) + return updates + + runtime = Runtime(compute_workers=2) + slow = Stream( + SimulatedSource([np.zeros((4, 4), dtype=np.uint8)], chunks=[(4, 4)]), + BlockingTarget(Layout(block_shape=(4, 4))), + runtime=runtime, + ) + fast = Stream( + SimulatedSource([np.zeros((4, 4), dtype=np.uint8)], chunks=[(4, 4)]), + RecordingTarget(Layout(block_shape=(4, 4))), + runtime=runtime, + ) + view = ortho_view((4, 4), viewport=(64, 64)) + try: + slow.update(view) + assert stage_started.wait(timeout=5) + fast.update(view) + wait(lambda: fast.status.state == "complete") + assert slow.status.state == "loading" + + release_stage.set() + wait(lambda: slow.status.state == "complete") + slow.close() + assert not runtime.closed + finally: + release_stage.set() + slow.close() + fast.close() + runtime.close() + + assert runtime.closed + + +def test_closing_stream_cancels_its_reads_on_shared_runtime(ortho_view, wait) -> None: + started = threading.Event() + cancelled = threading.Event() + + class BlockingSource(SimulatedSource): + async def read(self, level, region): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + runtime = Runtime() + stream = Stream( + BlockingSource([np.zeros((4, 4), dtype=np.uint8)], chunks=[(4, 4)]), + RecordingTarget(Layout(block_shape=(4, 4))), + runtime=runtime, + ) + try: + stream.update(ortho_view((4, 4), viewport=(64, 64))) + assert started.wait(timeout=5) + stream.close() + wait(cancelled.is_set) + assert not runtime.closed + finally: + stream.close() + runtime.close() + + def test_pause_holds_reads_until_resume(ortho_view, wait) -> None: source = SimulatedSource([np.zeros((8, 8), dtype=np.uint8)], chunks=[(4, 4)]) target = RecordingTarget(Layout(block_shape=(4, 4)))