multiscale chunked rendering

[kevinyamauchi]

Overview

I think it would be nice to be able to view larger-than-memory images in Python. I have a multiscale rendering prototype working with pygfx and as I saw multi-resolution data on the roadmap, I was wondering if ndv could be a good home for it. I'm sorry, it's a bit of a brain dump. Happy to jump on a call if you'd like a walkthrough.

Current prototype

I have made a prototype of 2d and 3d multiscaled chunk rendering in pygfx (https://github.com/kevinyamauchi/demo-multiscale). The basic concept is that when the view is updated (i.e., camera moves or dims state changes), only the data that is both in view and fits in the GPU is loaded. As we have multiple resolutions of the data, we can select data that has a resolution such that the level of detail matches that of the screen. I've given a brief overview below of how it works.

Data structure and IO
Currently, I am using tensorstore to load ome-zarr v0.5 files. It should be straightforward to make adapters for other formats and IO libraries. In ndv, I believe this would be a DataWrapper

Selecting data in view
We can use the camera model to determine which region in the displayed world space is viewable by the camera (e.g., the camera frustum in 3d). Then we can select all of the chunks that are in view. Currently I take any chunk that has a corner in the frustum, but there are other ways to do it.

Selecting which resolution to load from
In 2d, the level of detail is selected for the full slice to match the screen pixel size to the data pixel size. For 3d selection is done on a per-brick level, selecting the scale to match the brick size to the screen pixel. For both 2d and 3d, there is also an lod_bias parameter that can tune the bias towards finer or coarser levels of detail.

Texture indirection
One of the performance bottlenecks is loading the data onto the GPU. As you man imagine when the camera moves, we can often reuse some of the chunks as they remain in view. We can use a cache on the GPU to re-use data and save both the time to load the chunk from disk and the time to load the chunk to the GPU. Rather than moving the data from the cache to the sampling texture, we can use a technique called "texture indirection" point the shader sampler to the right element in the cache. This writeup on Kiln has a nice explanation of how texture indirection can be used in volume rendering (works similarly in 2d). The implementation in my prototype has a similar design to Kiln.

Current ndv slicing flow

My understanding of the current slicing flow is described below. Is that right?

1. Dims slider changes
2. ArrayViewer._on_view_current_index_changed updates the model current_index
3. The model.current_index change emits a signal that calls Array_viewer._re_resolve(). This gets the current view state (dims indices, displayed dims) and calls ArrayViewer._apply_changes()
4. ArrayViewer._apply_changes() requests new data if the new state requires it (via ArrayViewer._request_data() )
5. ArrayViewer._request_data() builds the slice requests and then submits them to the threadpool. ArrayViewer._on_data_response_ready() is attached to each request future as a callback
6. When the future is completed, ArrayViewer._on_data_response_ready() iterates through the response and sets the new data on ChannelController.update_texture_data()
7. ChannelController.update_texture_data() sends the data to the GPU using ImageHandle.set_data()

Multiscale chunked rendering slicing flow

I think the main differences between the current ndv slicing flow and the multiscale chunked rendering slicing flow are:

• We need to be able to trigger a reslicing when the camera state changes (likely with some debounce and/or throttling)
• We need the camera state to determine which data is in view
• We need the cache state to know what data are already on the GPU (and thus don't need to be loaded)
• we need the size of the canvas to determine how to select the level of detail

The current slicing flow is (see `demo_3d.py`):

1. Each frame, draw_frame() (demo_3d.py) calls _camera_changed(), which compares the camera's current world position and rotation against the last known state. If they differ, reslice_3d() is called. (this should probably be debounced and/or throttled in some way)
2. reslice_3d() (demo_3d.py) calls visual.cancel_pending() to cancel all in-flight requested chunks from the previous frame. It then encapsulates the current view state in a DimsState object, and calls visual.build_slice_request()
3. GFXMultiscaleImageVisual.build_slice_request()(demo_multiscale/render_visual.py) determines which chunks to request based on the camera view (both what is visible and selects level of detail)
   1. LOD selection: select_levels_from_cache()(demo_multiscale/_level_of_detail_3d.py) assigns each brick a resolution level based on its distance from the camera and the camera's field-of-view
   2. frustum cull: bricks_in_frustum_arr() (demo_multiscale/_frustum.py) discards bricks whose AABBs lie entirely outside the camera frustum
   3. gpu budget: bricks beyond the GPU slot budget are dropped (farthest from the camera dropped first)
   4. staging: tile_manager.stage()(demo_multiscale/block_cache/_tile_manager_3d.py) determines which of the bricks need to be requested based on the cache state and returns a list of the requests.
4. reslice_3d() calls slicer.submit(requests, fetch_fn=data_store.get_data, callback=on_batch) (demo_multiscale/slicer.py). AsyncSlicer.submit() signals any already-running thread for the same slice_request_id to cancel, then spawns a new thread running AsyncSlicer._run().
5. The chunk requests are carried out asynchronously and uploaded to the GPU in batches. The batching can be eliminated for simplicity.
   1. AsyncSlicer._run() (demo_multiscale/slicer.py) splits the request list into batches of batch_size and uses the shared ThreadPoolExecutor to call fetch_fn (i.e. data_store.get_data() on each request in a batch concurrently.
   2. OMEZarrImageDataStore.get_data() (demo_multiscale/data_store.py) receives a single ChunkRequest and returns the chunk as a numpy array.
   3. After each batch completes, GFXMultiscaleImageVisual.on_data_ready()(demo_multiscale/render_visual.py) and writes each chunk in the batch to the CPU-side cache_data. It then calls cache_tex.update_range() to mark that region dirty for GPU upload. tile_manager.commit() then records the brick as resident in the tilemap.
   4. After processing all bricks in the batch, the lookup table is rebuilt. on_data_ready() calls self._lut_manager_3d.rebuild(self._block_cache_3d.tile_manager)(demo_multiscale/lut_indirection/_lut_indirection_manager.py). rebuild_lut() to update the indirection look up table and marks the updated regions of the textures as requiring upload to GPU (lut_tex and brick_max_tex )
6. On next animation frame, draw_frame()callsrenderer.render(scene, camera)`. data are uploaded to the GPU and the scene is rendered

Proposed changes

I am still wrapping my mind around it, so these are just my current ideas. I wanted to start the issue to get some early feedback to see if this seems viable at all, if another approach is needed, or if it's not the right time to consider these types of changes to slicing. Of course, no worries if you decide that this isn't viable!

To try and understand what it would take to port my prototype to ndv I made a version of the demo that is as aligned with my understanding of the ndv components/concepts (see demo_2d_3d_ndv.py).

Demo movie (sorry for the jumpy zooms, I was trying to find zoom levels that trigger a different level of detail 😅 ):

multiscale_rendering_demo.mov

DataWrapper: I think we need a MultiscaleDataWrapper. The main difference is that it needs to have data about the scales (e.g., transformations between them, shapes) and the isel() method needs to take the level/scale to select from. I made a new base class here (not sure it should a subclass of DataWrapper) and implemented the ome-zarr wrapper here. I don't think I implemented all fields from the base class correctly (e.g., coords), but hopefully you get the idea 😬

MultiscaleImageHandle/MultiscaleVolumeHandle: I think it makes sense to define a new handle base class for multiscale images. I defined a new base class (here) and made Pygfx implementations based on my prototype (here). I didn't implement all fields on the base class yet. Just wanted to get something to show how the multiscale stuff could work.

CacheQuery2D/3D interface: Many multiscale rendering methods will have a GPU cache for performance. it is necessary to be able to query that cache when constructing the chunk requests to avoid re-requesting data already on the cache. I think it makes sense to define a generic query interface. That way specific handle implementations can have their own cache implementation, but the controller always knows how to check what is in the cache , clear the cache etc. I defined an protocol and then made it a property of the MultiscaleImageHandle/MultiscaleVolumeHandle base classes.

Camera state: We need to know the region of world space in view in order to do the chunk culling in 2d/3d. I couldn't find a camera model in ndv. If that is correct, we could either add a camera model or add a data class that defines the current region of world space and some info about the canvas pixel size (e.g., CameraView2D, CameraView3D). The CameraView2D/3D approach would be more general across different camera types and implementations. We would also need some sort of mechanisms for checking if the camera state has changed to request reslicing (probably want throttling/debouncing or a "settle timer" on there). I defined the data class and functions to build them from Pygfx components here.

Changes to reslicing: the biggest change to reslicing is that we need to determine which chunks to load based on the camera and cache state. We can still use the existing resolve() function to get the displayed dims. However, we need to add functions to determine which chunks are in view and what level of detail should be loaded. As mentioned above, this requires the addition of the camera and cache state. I've linked to a potential way to do this below.

Here is an example from the demo of how 2d slicing could be implemented (script here):

Expand for code

def reslice_2d(
    image_handle: MultiscaleImageHandle,
    wrapper: OMEZarrDataWrapper,
    camera_view: CameraView2D,
    slicer: AsyncSlicer,
    resolved: ResolvedDisplayState,
    upload_queue: collections.deque,
    lod_bias: float = LOD_BIAS,
) -> tuple[tuple[int, int], ...]:
    """Plan and submit asynchronous 2-D tile requests for the current slice.

    Cancels any pending work, selects visible tiles for the current camera
    view and slice position, evicts tiles from finer LOD levels that are no
    longer needed, allocates GPU cache slots for cache misses, and submits the
    resulting fetch requests to the asynchronous slicer. Completed batches are
    appended to ``upload_queue`` for later GPU upload by the draw loop.

    Parameters
    ----------
    image_handle : MultiscaleImageHandle
        Handle that owns the 2-D GPU cache and accepts tile writes.
    wrapper : OMEZarrDataWrapper
        Multiscale data source used to load each requested chunk.
    camera_view : CameraView2D
        Frozen snapshot of the orthographic camera (bounds, viewport size).
    slicer : AsyncSlicer
        Background loader that executes chunk reads and posts completed
        batches back to the Qt thread.
    resolved : ResolvedDisplayState
        Resolved ndv display state describing the visible axes and any fixed
        indices for non-visible axes.
    upload_queue : collections.deque
        Receives completed ``(request, data)`` pairs until the draw loop
        drains and uploads them to the GPU.
    lod_bias : float, optional
        Scale applied to LOD thresholds. Values > 1 favour coarser levels;
        values < 1 favour finer levels.

    Returns
    -------
    tuple of tuple of int
        Sorted ``(axis, index)`` pairs encoding the current non-visible axis
        positions. The caller passes this to ``on_data_ready_2d`` so uploaded
        tiles are committed into the correct 2-D slice cache entry.
    """
    # Discard any in-flight requests and queued uploads from the previous frame.
    image_handle.invalidate_pending()
    upload_queue.clear()

    # Derive a stable cache key for the current slice position: sorted
    # (axis, index) pairs for all non-displayed integer axes.
    visible = set(resolved.visible_axes)
    slice_coord: tuple[tuple[int, int], ...] = tuple(sorted(
        (ax, v) for ax, v in resolved.current_index.items()
        if isinstance(v, int) and ax not in visible
    ))

    # Advance the LRU frame counter so cache residency checks this frame use
    # a consistent timestamp.
    image_handle.advance_frame()

    # Derive ZYX voxel scales from the wrapper (full nD scales, take last 3).
    # this will have to be done properly using the resolved display state.
    voxel_scales = np.asarray(wrapper.voxel_sizes, dtype=np.float64)[-3:]

    # View selection: choose and prioritize visible tiles based on camera
    # bounds, LOD thresholds, and viewport culling. No cache interaction here.
    required_block_keys, target_level = select_visible_bricks_2d(
        camera_view,
        image_handle.brick_layout,
        voxel_scales,
        slice_coord,
        lod_bias,
    )

    # Remove tiles at finer resolution than the current target level; they
    # would be immediately superseded and waste cache slots.
    n_evicted = image_handle.evict_finer_than(target_level)

    # Build inverse level transforms (world → level-k data coordinates) needed
    # to map non-displayed axis indices into each resolution level.
    world_to_level = [AffineTransform(matrix=t.inverse_matrix) for t in wrapper.level_transforms]

    # Cache-aware pass: skip resident tiles, allocate slots for misses, and
    # assemble MultiscaleChunkRequest objects ready for the slicer.
    requests = build_fetch_requests_2d(
        required_block_keys,
        slice_coord,
        image_handle.cache_query(),
        resolved,
        list(wrapper.level_shapes),
        world_to_level,
        image_handle.expand_fetch_index,
        image_handle.brick_layout.block_size,
    )

    # If evictions changed the cache but no new fetches are needed, rebuild
    # the LUT immediately so the display reflects the evictions right away.
    if n_evicted > 0 and not requests:
        image_handle.rebuild_lut(slice_coord)

    # Capture the slice ID so the callback can drop stale batches if a newer
    # reslice supersedes this one before results arrive.
    slice_id = requests[0].slice_request_id if requests else None

    def on_batch(batch):
        if slicer.current_slice_id != slice_id:
            return
        upload_queue.extend(batch)

    slicer.submit(
        requests,
        fetch_fn=lambda req: wrapper.isel(req.index, level=req.level),
        callback=on_batch,
    )

    return slice_coord

Questions

I would love to hear your thoughts! As mentioned, no stress if this isn't in scope for now. It was super interesting for me to go through the exercise of seeing how this could potentially fit into ndv in any case. In addition to general feedback, I have a couple of questions:

1. Am I correct in understanding that currently, an ndv viewer has only one DataWrapper assigned to it (i.e., only one data source)?
2. How would you want to handle the histogram? I think it would be pretty tricky to compute on a large image while keeping performance.
3. What is the RectangularROI for? I wasn't sure if that would need any special support.

[tlambert03]

amazing. thank you @kevinyamauchi! I look forward to digging into this. Still traveling at the moment, but will digest soon :)

[kevinyamauchi]

No stress, take your time! Looking forward to hearing your thoughts!

[gselzer]

Hey, this is really cool - I've bee dreaming of adding something like this into ndv - glad to see you're taking the plunge!

Maybe there are a few questions I could answer (although @tlambert03 will have valuable insight too, I'm sure):

1. Am I correct in understanding that currently, an ndv viewer has only one DataWrapper assigned to it (i.e., only one data source)?

This is correct, at least for now. I'm pretty far along on a scenex-powered ArrayViewer, and I had a different prototype that was capable of visualizing data from many DataWrappers, but that's a little bit in the future. scenex usage could indeed be a large part of the puzzle in making these things work in ndv, as it'd also give you convenient access to camera state and events on changes.

2. How would you want to handle the histogram? I think it would be pretty tricky to compute on a large image while keeping performance.

Without knowing exactly what your concerns are, a couple thoughts:

1. In some cases we probably just don't want to display the histogram. In such settings, we can just hide the histogram button, avoiding the problem. For situations where we only want to visualize a few chunks of a larger image, this may be a good starting point. Of course, we'll still likely need to compute contrast limits though...
2. If we are delivering data in chunks, it shouldn't be a huge problem to compute the histograms on each chunk and sum them up to a histogram. That many-DataWrapper prototype did that and it worked alright.

3. What is the RectangularROI for? I wasn't sure if that would need any special support.

Within ndv there is little usage, and it's a bit of a half-baked feature (for example, no 3D support). I developed it with the dream of being able to display histograms of subsections of the displayed image, but I never actually wrote that code. Nowadays there is some downstream usage (example: pymmcore-plus/pymmcore-gui#150) though.

[kevinyamauchi]

Thanks for the feedback, @gselzer !

This is correct, at least for now. I'm pretty far along on a scenex-powered ArrayViewer, and I had a different prototype that was capable of visualizing data from many DataWrappers, but that's a little bit in the future. scenex usage could indeed be a large part of the puzzle in making these things work in ndv, as it'd also give you convenient access to camera state and events on changes.

I just had a look at #276 and it looks interesting. I will have to checkout scenex!

Without knowing exactly what your concerns are, a couple thoughts:

• In some cases we probably just don't want to display the histogram. In such settings, we can just hide the histogram button, avoiding the problem.
• For situations where we only want to visualize a few chunks of a larger image, this may be a good starting point. Of course, we'll still likely need to compute contrast limits though...
  If we are delivering data in chunks, it shouldn't be a huge problem to compute the histograms on each chunk and sum them up to a histogram. That many-DataWrapper prototype did that and it worked alright.

This is helpful feedback. My concern is that computing a histogram on a large image would take a lot of time. It's good to know that you would be open to just not computing the histogram in some cases. The idea to accumulate the histogram as the chunks stream in is an interesting one. I think there's some nuance for multiscale data (e.g., would we weight based on the world size of the voxels?), but it's an interesting direction to explore. It feels like disabling the histogram first and then exploring how to best implement a histogram as a follow-up would be a pragmatic approach.

[tlambert03]

just want to say, i think about this often (with much appreciation @kevinyamauchi). A bit overwhelmed by the size of the refactor, but also consider it super high value. So, sorry for continued delays... unburying myself at the moment.

A couple very quick responses to topics above:

• Currently, I am using tensorstore to load ome-zarr v0.5 files. It should be straightforward to make adapters for other formats and IO libraries. In ndv, I believe this would be a DataWrapper.

  this is actually slightly problematic. ndv doesn't currently (and nor do I want it to) rely by default on any io libraries. DataWrappers are, strictly speaking, objects to adapt in-memory objects from various libraries. They are not objects that map filepaths via io. Naturally, we can/should add optional io extras to support this ... but this would be the first case of that (outside of unmerged experiments in feat: add io and cli #21).

• Camera state: We need to know the region of world space in view in order to do the chunk culling in 2d/3d. I couldn't find a camera model in ndv.

  yes, the camera model is too implicit at the moment, and as @gselzer, integration of a proper camera abstraction has always been deferred to the more complete scenegraph abstraction which is scenex. However, I see no reason why we can't hack around that in the meantime here just for pygfx. As long as it's the only backend that we intend to use to support multiscale/culled/chunked rendering, there's no urgent need to abstract it. let's just do the pragmatic thing and hook it up directly.

• Am I correct in understanding that currently, an ndv viewer has only one DataWrapper assigned to it (i.e., only one data source)?

  yes, if it helps as a mental model, you can think of ndv as, for the moment, explicitly handling just the case of a single Image layer equivalent in napari. Not trying to coordinate shared coordinate spaces and slicing. You just have a single, square, object.

• How would you want to handle the histogram? I think it would be pretty tricky to compute on a large image while keeping performance.

  I think I would want to use a graceful degradation. That is, if the number of pixels is too large, just use a stride of 2 or something (skip every other pixel). Yes, it's an approximation, but given the PSF of most images, it would likely still be a very useful approximation for most images (and we could have a little "~" icon or something to indicate when the approximation threshold was triggered).

  And yes, if we're moving to a model where the post-sliced data comes back in more than 1 chunk, we will need to move to a progressive histogramming model (i.e. one that has state, reset, add-chunk, signals, etc)

• What is the RectangularROI for? I wasn't sure if that would need any special support.

  no, we don't need to support it for a v1.

[kevinyamauchi]

Thanks for the additional context, @tlambert03 ! No worries about the timeline to figure this out. This isn't blocking me and I don't mind taking a bit of extra time to try and get it "right".

Regarding the tensorstore dependency, I totally understand. I took a closer look at the existing DataWrappers and I think I get what you are going for.

Would it make sense to create a MultiscaleDataWrapper that takes either a dict or list of arrays (one per scale) and a dict or list of scale-to-world transforms (eg affine matrix or scale + translation)? We would then have to extend the 'isel' method to take the resolution/scale level in addition to the slice coordinates. In this case, I think we would be able to use the existing machinery for slicing the arrays with little to no modifications (famous last words 😂).

The upside of this is approach is that I think it would allow us to continue having no default io dependencies. The downside is that the user has to extract the arrays and metadata from the files themselves, but I think that is inline with the current viewer experience.

I am traveling for the next week or so, but I could try updating the demo repo when I get back. It might be helpful to see if it's somewhat along the lines of what you are thinking.