|
| 1 | +"""Time what napari does when you zoom, pan or change channel. |
| 2 | +
|
| 3 | +Usage: |
| 4 | + cd workflow/viewer |
| 5 | + pixi run bench /path/to/image.zarr # or image.zarr.zip |
| 6 | + pixi run bench U:\\...\\image.zarr.zip --render |
| 7 | +
|
| 8 | +Reads one typical screen (a plane through the middle, 2048 x 2048 pixels) |
| 9 | +at full resolution and two levels down, from the image and every label |
| 10 | +image, and says where the time goes: |
| 11 | +
|
| 12 | +- how the arrays are chunked, and how many planes a chunk makes zarr |
| 13 | + decompress to show one (a chunk 126 planes deep: 126 times the data); |
| 14 | +- the time for one small read (the latency of the store: network, zip); |
| 15 | +- each view read straight with zarr, read again (the operating system's |
| 16 | + cache), and the way the viewer reads it (dask); |
| 17 | +- with --render, napari drawing a label plane (needs napari). |
| 18 | +
|
| 19 | +Nothing is written. Send the whole output. |
| 20 | +""" |
| 21 | + |
| 22 | +from __future__ import annotations |
| 23 | + |
| 24 | +import argparse |
| 25 | +import os |
| 26 | +import sys |
| 27 | +import time |
| 28 | + |
| 29 | +import numpy as np |
| 30 | + |
| 31 | +from patchworks._io import open_group_any, split_zip_path |
| 32 | +from patchworks.plugins.ome_zarr import read_ngff_attr |
| 33 | + |
| 34 | +SCREEN = 2048 # pixels across one view |
| 35 | + |
| 36 | + |
| 37 | +def _drive(path: str) -> str: |
| 38 | + """Where the store lives: local disk, network, or unknown.""" |
| 39 | + if os.name == "nt": |
| 40 | + import ctypes |
| 41 | + |
| 42 | + root = os.path.splitdrive(os.path.abspath(path))[0] + "\\" |
| 43 | + kind = ctypes.windll.kernel32.GetDriveTypeW(root) |
| 44 | + return {2: "removable", 3: "local disk", 4: "network drive"}.get( |
| 45 | + kind, f"drive type {kind}" |
| 46 | + ) |
| 47 | + try: |
| 48 | + import psutil |
| 49 | + |
| 50 | + best = max( |
| 51 | + ( |
| 52 | + p |
| 53 | + for p in psutil.disk_partitions(all=True) |
| 54 | + if os.path.abspath(path).startswith(p.mountpoint) |
| 55 | + ), |
| 56 | + key=lambda p: len(p.mountpoint), |
| 57 | + ) |
| 58 | + net = best.fstype in ("nfs", "nfs4", "cifs", "smbfs", "fuse.sshfs") |
| 59 | + return f"{'network' if net else 'local'} ({best.fstype})" |
| 60 | + except Exception: # noqa: BLE001 - only informative |
| 61 | + return "unknown" |
| 62 | + |
| 63 | + |
| 64 | +def _levels(group_path: str) -> tuple[list, list[str]]: |
| 65 | + """The pyramid's zarr arrays, finest first, and the axis names.""" |
| 66 | + root = open_group_any(group_path) |
| 67 | + ms = read_ngff_attr(root.attrs, "multiscales")[0] |
| 68 | + axes = [a["name"] if isinstance(a, dict) else a for a in ms["axes"]] |
| 69 | + return [root[d["path"]] for d in ms["datasets"]], axes |
| 70 | + |
| 71 | + |
| 72 | +def _describe(name: str, arrays: list, axes: list[str]) -> None: |
| 73 | + print(f"\n{name}: axes {''.join(axes)}, {len(arrays)} level(s)") |
| 74 | + for i, arr in enumerate(arrays): |
| 75 | + shards = getattr(arr, "shards", None) |
| 76 | + z = axes.index("z") if "z" in axes else None |
| 77 | + depth = arr.chunks[z] if z is not None else 1 |
| 78 | + print( |
| 79 | + f" level {i}: shape {tuple(arr.shape)} {arr.dtype}, chunks " |
| 80 | + f"{tuple(arr.chunks)}" |
| 81 | + + (f" in shards {tuple(shards)}" if shards else "") |
| 82 | + + ( |
| 83 | + f" -> one plane decompresses {depth} plane(s)" |
| 84 | + if z is not None |
| 85 | + else "" |
| 86 | + ) |
| 87 | + ) |
| 88 | + |
| 89 | + |
| 90 | +def _view(arr, axes: list[str], level: int) -> tuple: |
| 91 | + """One screen: the middle plane (channel 0), SCREEN pixels across.""" |
| 92 | + sl = [] |
| 93 | + for ax, n in zip(axes, arr.shape): |
| 94 | + if ax in ("c", "t"): |
| 95 | + sl.append(0) |
| 96 | + elif ax == "z": |
| 97 | + sl.append(n // 2) |
| 98 | + else: |
| 99 | + w = min(SCREEN, n) |
| 100 | + lo = (n - w) // 2 |
| 101 | + sl.append(slice(lo, lo + w)) |
| 102 | + return tuple(sl) |
| 103 | + |
| 104 | + |
| 105 | +def _chunks_touched(arr, sl: tuple) -> int: |
| 106 | + n = 1 |
| 107 | + for s, c, size in zip(sl, arr.chunks, arr.shape): |
| 108 | + if isinstance(s, int): |
| 109 | + continue |
| 110 | + n *= -(-s.stop // c) - s.start // c |
| 111 | + return n |
| 112 | + |
| 113 | + |
| 114 | +def _timed(fn) -> tuple[float, object]: |
| 115 | + t = time.perf_counter() |
| 116 | + out = fn() |
| 117 | + return time.perf_counter() - t, out |
| 118 | + |
| 119 | + |
| 120 | +def _bench(name: str, arrays: list, axes: list[str], dask_level) -> None: |
| 121 | + for level in sorted({0, min(2, len(arrays) - 1)}): |
| 122 | + arr = arrays[level] |
| 123 | + sl = _view(arr, axes, level) |
| 124 | + touched = _chunks_touched(arr, sl) |
| 125 | + chunk_mb = np.prod(arr.chunks) * arr.dtype.itemsize / 1e6 |
| 126 | + cold, data = _timed(lambda: np.asarray(arr[sl])) |
| 127 | + warm, _ = _timed(lambda: np.asarray(arr[sl])) |
| 128 | + via_dask, _ = _timed(lambda: np.asarray(dask_level(level)[sl])) |
| 129 | + print( |
| 130 | + f" {name} level {level}: view {data.shape}, {touched} chunk(s) " |
| 131 | + f"= {touched * chunk_mb:,.0f} MB decompressed for " |
| 132 | + f"{data.nbytes / 1e6:,.0f} MB shown | zarr {cold:.2f} s, again " |
| 133 | + f"{warm:.2f} s | as the viewer reads it (dask) {via_dask:.2f} s" |
| 134 | + ) |
| 135 | + |
| 136 | + |
| 137 | +def _latency(arr) -> float: |
| 138 | + """The best of 5 reads of one voxel in different chunks.""" |
| 139 | + best = float("inf") |
| 140 | + for k in range(5): |
| 141 | + idx = tuple( |
| 142 | + min(n - 1, (k * 7 + 3) * c % max(n, 1)) |
| 143 | + for n, c in zip(arr.shape, arr.chunks) |
| 144 | + ) |
| 145 | + t, _ = _timed(lambda: arr[idx]) |
| 146 | + best = min(best, t) |
| 147 | + return best |
| 148 | + |
| 149 | + |
| 150 | +def _render(arrays: list, axes: list[str]) -> None: |
| 151 | + try: |
| 152 | + import napari |
| 153 | + except ImportError: |
| 154 | + print("\n--render: napari is not installed here") |
| 155 | + return |
| 156 | + arr = arrays[0] |
| 157 | + plane = np.asarray(arr[_view(arr, axes, 0)]).astype("int32") |
| 158 | + print( |
| 159 | + f"\nnapari drawing one {plane.shape} label plane " |
| 160 | + f"(ids up to {int(plane.max()):,})" |
| 161 | + ) |
| 162 | + viewer = napari.Viewer(show=False) |
| 163 | + t, layer = _timed(lambda: viewer.add_labels(plane)) |
| 164 | + print(f" add the layer: {t:.2f} s") |
| 165 | + for k in range(3): |
| 166 | + layer.data = np.roll(plane, k + 1, axis=0) |
| 167 | + t, _ = _timed(layer.refresh) |
| 168 | + print(f" redraw {k + 1}: {t:.2f} s") |
| 169 | + viewer.close() |
| 170 | + |
| 171 | + |
| 172 | +def main() -> None: |
| 173 | + parser = argparse.ArgumentParser(description=__doc__) |
| 174 | + parser.add_argument("store", help="image.zarr, or an image.zarr.zip bundle") |
| 175 | + parser.add_argument( |
| 176 | + "--render", action="store_true", help="also time napari's drawing" |
| 177 | + ) |
| 178 | + args = parser.parse_args() |
| 179 | + store = str(args.store).rstrip("/\\") |
| 180 | + |
| 181 | + zipped = split_zip_path(store) is not None |
| 182 | + print(f"store: {store}") |
| 183 | + print( |
| 184 | + f" {'a .zip bundle' if zipped else 'a directory'} on " |
| 185 | + f"{_drive(split_zip_path(store)[0] if zipped else store)}" |
| 186 | + ) |
| 187 | + |
| 188 | + image, axes = _levels(store) |
| 189 | + _describe("image", image, axes) |
| 190 | + labels = [] |
| 191 | + try: |
| 192 | + names = sorted( |
| 193 | + open_group_any(f"{store}/labels").attrs.get("labels") |
| 194 | + or (open_group_any(f"{store}/labels").attrs.get("ome") or {}).get( |
| 195 | + "labels", [] |
| 196 | + ) |
| 197 | + ) |
| 198 | + except Exception: # noqa: BLE001 - no labels |
| 199 | + names = [] |
| 200 | + for name in names: |
| 201 | + arrays, lab_axes = _levels(f"{store}/labels/{name}") |
| 202 | + _describe(f"labels/{name}", arrays, lab_axes) |
| 203 | + labels.append((name, arrays, lab_axes)) |
| 204 | + |
| 205 | + print( |
| 206 | + f"\none voxel read (a whole chunk fetched and decompressed): " |
| 207 | + f"{1000 * _latency(image[0]):.0f} ms" |
| 208 | + ) |
| 209 | + |
| 210 | + from patchworks._io import load_ome_zarr |
| 211 | + from patchworks.plugins.napari import _multiscale_levels |
| 212 | + |
| 213 | + print(f"\none view ({SCREEN} x {SCREEN} pixels, middle plane):") |
| 214 | + c = 0 if "c" in axes else None |
| 215 | + _bench( |
| 216 | + "image", |
| 217 | + image, |
| 218 | + axes, |
| 219 | + lambda lvl: ( |
| 220 | + load_ome_zarr(store, channel=c, level=lvl)[None] |
| 221 | + if c is not None |
| 222 | + else load_ome_zarr(store, channel=None, level=lvl) |
| 223 | + ), |
| 224 | + ) |
| 225 | + for name, arrays, lab_axes in labels: |
| 226 | + group = f"{store}/labels/{name}" |
| 227 | + _bench( |
| 228 | + f"labels/{name}", |
| 229 | + arrays, |
| 230 | + lab_axes, |
| 231 | + lambda lvl, g=group: _multiscale_levels(g, None)[lvl].astype( |
| 232 | + "int32" |
| 233 | + ), |
| 234 | + ) |
| 235 | + |
| 236 | + if args.render and labels: |
| 237 | + _render(labels[0][1], labels[0][2]) |
| 238 | + print( |
| 239 | + "\nTimes over ~0.5 s per view are felt. 'again' much faster than the " |
| 240 | + "first read: the store's latency or the network is the cost; both " |
| 241 | + "slow: decompressing (chunks much deeper than one plane)." |
| 242 | + ) |
| 243 | + |
| 244 | + |
| 245 | +if __name__ == "__main__": |
| 246 | + sys.exit(main()) |
0 commit comments