feat: Support Zarr v3 (Replace zarr.js with zarrita.js)

package.json

@@ -1,6 +1,7 @@
 {
   "name": "@hms-dbmi/vizarr",
   "version": "0.3.0",
+  "type": "module",
   "dependencies": {
     "@hms-dbmi/viv": "^0.16.0",
     "@material-ui/core": "^4.11.0",
@@ -13,8 +14,7 @@
     "quick-lru": "^6.0.0",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "reference-spec-reader": "^0.1.1",
-    "zarr": "^0.5.2"
+    "zarrita": "^0.4.0-next.14"
   },
   "scripts": {
     "start": "vite",
@@ -36,9 +36,9 @@
     "@types/node": "^14.14.5",
     "@types/react": "^18.2.51",
     "@types/react-dom": "^18.2.18",
-    "@vitejs/plugin-react": "^3.1.0",
+    "@vitejs/plugin-react": "^4.3.1",
     "prettier": "^2.2.0",
-    "typescript": "^4.9.5",
-    "vite": "^4.5.0"
+    "typescript": "^5.5.3",
+    "vite": "^5.3.3"
   }
 }

src/ZarrPixelSource.ts

@@ -0,0 +1,133 @@
+import * as zarr from 'zarrita';
+
+import type * as viv from '@vivjs/types';
+import type { Readable } from '@zarrita/storage';
+
+import { assert } from './utils';
+import { getImageSize } from '@hms-dbmi/viv';
+
+// TODO: Export from top-level zarrita
+type Slice = ReturnType<typeof zarr.slice>;
+
+const X_AXIS_NAME = 'x';
+const Y_AXIS_NAME = 'y';
+const RGBA_CHANNEL_AXIS_NAME = '_c';
+const SUPPORTED_DTYPES = ['Uint8', 'Uint16', 'Uint32', 'Float32', 'Int8', 'Int16', 'Int32', 'Float64'] as const;
+
+export class ZarrPixelSource<S extends Array<string> = Array<string>> implements viv.PixelSource<S> {
+  #arr: zarr.Array<zarr.DataType, Readable>;
+  readonly labels: viv.Labels<S>;
+  readonly tileSize: number;
+  readonly dtype: viv.SupportedDtype;
+
+  constructor(
+    arr: zarr.Array<zarr.DataType, Readable>,
+    options: {
+      labels: viv.Labels<S>;
+      tileSize: number;
+    }
+  ) {
+    this.#arr = arr;
+    this.labels = options.labels;
+    this.tileSize = options.tileSize;
+    const vivDtype = capitalize(arr.dtype);
+    assert(isSupportedDtype(vivDtype), `Unsupported viv dtype: ${vivDtype}`);
+    this.dtype = vivDtype;
+  }
+
+  get shape() {
+    return this.#arr.shape;
+  }
+
+  async getRaster(options: {
+    selection: viv.PixelSourceSelection<S> | Array<number>;
+    signal?: AbortSignal;
+  }): Promise<viv.PixelData> {
+    const { selection, signal } = options;
+    return this.#fetchData(buildZarrQuery(this.labels, selection), { signal });
+  }
+
+  async getTile(options: {
+    x: number;
+    y: number;
+    selection: viv.PixelSourceSelection<S> | Array<number>;
+    signal?: AbortSignal;
+  }): Promise<viv.PixelData> {
+    const { x, y, selection, signal } = options;
+    const sel = buildZarrQuery(this.labels, selection);
+
+    const { height, width } = getImageSize(this);
+    const [xStart, xStop] = [x * this.tileSize, Math.min((x + 1) * this.tileSize, width)];
+    const [yStart, yStop] = [y * this.tileSize, Math.min((y + 1) * this.tileSize, height)];
+
+    // Deck.gl can sometimes request edge tiles that don't exist. We throw
+    // a BoundsCheckError which is picked up in `ZarrPixelSource.onTileError`
+    // and ignored by deck.gl.
+    if (xStart === xStop || yStart === yStop) {
+      throw new BoundsCheckError('Tile slice is zero-sized.');
+    }
+    if (xStart < 0 || yStart < 0 || xStop > width || yStop > height) {
+      throw new BoundsCheckError('Tile slice is out of bounds.');
+    }
+
+    sel[this.labels.indexOf(X_AXIS_NAME)] = zarr.slice(xStart, xStop);
+    sel[this.labels.indexOf(Y_AXIS_NAME)] = zarr.slice(yStart, yStop);
+    return this.#fetchData(sel, { signal });
+  }
+
+  onTileError(err: Error): void {
+    if (err instanceof BoundsCheckError) {
+      return;
+    }
+    throw err;
+  }
+
+  async #fetchData(selection: Array<number | Slice>, options: { signal?: AbortSignal }): Promise<viv.PixelData> {
+    const {
+      data,
+      shape: [height, width],
+    } = await zarr.get(this.#arr, selection, {
+      // @ts-expect-error this is ok for now and should be supported by all backends
+      signal: options.signal,
+    });
+    return { data: data as viv.SupportedTypedArray, width, height };
+  }
+}
+
+function buildZarrQuery(labels: string[], selection: Record<string, number> | Array<number>): Array<Slice | number> {
+  let sel: Array<Slice | number>;
+  if (Array.isArray(selection)) {
+    // shallow copy
+    sel = [...selection];
+  } else {
+    // initialize with zeros
+    sel = Array.from({ length: labels.length }, () => 0);
+    // fill in the selection
+    for (const [key, idx] of Object.entries(selection)) {
+      sel[labels.indexOf(key)] = idx;
+    }
+  }
+  sel[labels.indexOf(X_AXIS_NAME)] = zarr.slice(null);
+  sel[labels.indexOf(Y_AXIS_NAME)] = zarr.slice(null);
+  if (RGBA_CHANNEL_AXIS_NAME in labels) {
+    sel[labels.indexOf(RGBA_CHANNEL_AXIS_NAME)] = zarr.slice(null);
+  }
+  return sel;
+}
+
+function capitalize<T extends string>(s: T): Capitalize<T> {
+  // @ts-expect-error - TypeScript can't verify that the return type is correct
+  return s[0].toUpperCase() + s.slice(1);
+}
+
+function isSupportedDtype(dtype: string): dtype is viv.SupportedDtype {
+  // @ts-expect-error - TypeScript can't verify that the return type is correct
+  return SUPPORTED_DTYPES.includes(dtype);
+}
+
+class BoundsCheckError extends Error {
+  name = 'BoundsCheckError';
+  constructor(message?: string) {
+    super(message);
+  }
+}

src/codecs/jpeg2k.ts

@@ -5,15 +5,14 @@
 import { JpxImage } from 'https://cdn.jsdelivr.net/gh/mozilla/pdf.js@30bd5f0/src/core/jpx.js';
 
 export default class Jpeg2k {
-  static codecId: 'jpeg2k';
-  static fromConfig(_: Record<string, any>): Jpeg2k {
+  kind = 'bytes_to_bytes' as const;
+  static codecId = 'jpeg2k';
+  static fromConfig(): Jpeg2k {
     return new Jpeg2k();
   }
-
   encode(_: Uint8Array): never {
     throw new Error('encode not implemented');
   }
-
   async decode(data: Uint8Array): Promise<Uint8Array> {
     const img = new JpxImage();
     img.failOnCorruptedImage = true;

src/codecs/register.ts

@@ -1,31 +1,2 @@
-import { addCodec } from 'zarr';
-import type { CodecConstructor } from 'numcodecs';
-
-type CodecModule = { default: CodecConstructor<Record<string, unknown>> };
-
-const cache: Map<string, Promise<CodecModule['default']>> = new Map();
-
-function add(name: string, load: () => Promise<CodecModule>): void {
-  // Cache import to avoid duplicate loads
-  const loadAndCache = () => {
-    if (!cache.has(name)) {
-      const promise = load()
-        .then((m) => m.default)
-        .catch((err) => {
-          cache.delete(name);
-          throw err;
-        });
-      cache.set(name, promise);
-    }
-    return cache.get(name)!;
-  };
-  addCodec(name, loadAndCache);
-}
-
-// Add dynmaically imported codecs to Zarr.js registry.
-add('lz4', () => import('numcodecs/lz4'));
-add('gzip', () => import('numcodecs/gzip'));
-add('zlib', () => import('numcodecs/zlib'));
-add('zstd', () => import('numcodecs/zstd'));
-add('blosc', () => import('numcodecs/blosc'));
-add('jpeg2k', () => import('./jpeg2k'));
+import { registry } from '@zarrita/core';
+registry.set('jpeg2k', () => import('./jpeg2k').then((m) => m.default));

src/components/Viewer.tsx

@@ -1,5 +1,5 @@
 import * as React from 'react';
-import { useAtom, WritableAtom } from 'jotai';
+import { useAtom, type WritableAtom } from 'jotai';
 import { useAtomValue } from 'jotai';
 import DeckGL from 'deck.gl';
 import { OrthographicView } from '@deck.gl/core';

src/gridLayer.ts

@@ -3,8 +3,10 @@ import { SolidPolygonLayer, TextLayer } from '@deck.gl/layers';
 import type { CompositeLayerProps } from '@deck.gl/core/lib/composite-layer';
 import pMap from 'p-map';
 
-import { XRLayer, ZarrPixelSource, ColorPaletteExtension } from '@hms-dbmi/viv';
+import { XRLayer, ColorPaletteExtension } from '@hms-dbmi/viv';
 import type { BaseLayerProps } from './state';
+import type { ZarrPixelSource } from './ZarrPixelSource';
+import { assert } from './utils';
 
 export interface GridLoader {
   loader: ZarrPixelSource<string[]>;
@@ -51,9 +53,8 @@ function validateWidthHeight(d: { data: { width: number; height: number } }[]) {
   const { width, height } = first.data;
   // Verify that all grid data is same shape (ignoring undefined)
   d.forEach(({ data }) => {
-    if (data?.width !== width || data?.height !== height) {
-      throw new Error('Grid data is not same shape.');
-    }
+    if (!data) return;
+    assert(data.width === width && data.height === height, 'Grid data is not same shape.');
   });
   return { width, height };
 }

src/index.tsx

@@ -7,7 +7,7 @@ import { ThemeProvider } from '@material-ui/styles';
 import Menu from './components/Menu';
 import Viewer from './components/Viewer';
 import './codecs/register';
-import { addImageAtom, ImageLayerConfig, ViewState, atomWithEffect } from './state';
+import { addImageAtom, type ImageLayerConfig, type ViewState, atomWithEffect } from './state';
 import { defer, typedEmitter } from './utils';
 import theme from './theme';