From 04715342120b97c3c84d8fe6d82ababc1020834f Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Fri, 21 Aug 2026 00:38:48 +0200 Subject: [PATCH 1/2] fix(ts): order generated axes time, then channel, then space `toMultiscales` emitted the caller's dims verbatim, so a channel-last input such as `(z, y, c, x)` produced metadata whose axis order the OME-Zarr specification forbids: axes are ordered by type, time then channel then space. ITK component images and the 4-D/5-D default dims all yield channel-last layouts, so the pipeline reaches that state without the caller asking for it. `canonicalAxisOrder` reorders the axes and moves the data with them, matching the Python port's `_canonical_axis_order`. An image whose dims fall outside `(t, c, z, y, x)` is returned untouched: an axis model that was not expressible before RFC-3 carries no spec ordering to normalize to. A positional `chunks` array indexes the caller's dims, so it follows them through the reordering. Where Python transposes lazily through dask, this reads the array and writes the permuted buffer into a new in-memory zarr array; the downsampling path materializes the image anyway. The transposition helpers move to `utils/transpose.ts`, which carries no toolkit dependency, so `to_multiscales-shared.ts` can reach them without importing `itk-wasm`. The two channel-last downsampling tests now assert what their Python twins in `test_to_ngff_zarr_itkwasm.py` already assert. --- docs/typescript.md | 6 + ts/src/methods/itkwasm-shared.ts | 160 ++---------------------- ts/src/process/to_multiscales-shared.ts | 16 ++- ts/src/utils/axis_order.ts | 94 ++++++++++++++ ts/src/utils/transpose.ts | 125 ++++++++++++++++++ ts/test/canonical_axis_order_test.ts | 106 ++++++++++++++++ ts/test/to_multiscales_itkwasm_test.ts | 20 +-- 7 files changed, 367 insertions(+), 160 deletions(-) create mode 100644 ts/src/utils/axis_order.ts create mode 100644 ts/src/utils/transpose.ts create mode 100644 ts/test/canonical_axis_order_test.ts diff --git a/docs/typescript.md b/docs/typescript.md index aca1b519..50dbdf29 100644 --- a/docs/typescript.md +++ b/docs/typescript.md @@ -464,6 +464,12 @@ async function toMultiscales( **Returns:** NgffMultiscales with generated pyramid levels +Axes are normalized to the OME-Zarr order, time then channel then space, so a +channel-last input such as `(z, y, c, x)` yields `(c, z, y, x)` and the data is +reordered with it. An axis model outside the `(t, c, z, y, x)` vocabulary is +left alone: it carries no spec ordering to normalize to. A positional `chunks` +array indexes the dims you passed and follows them through the reordering. + **Example:** ```typescript // Basic pyramid generation diff --git a/ts/src/methods/itkwasm-shared.ts b/ts/src/methods/itkwasm-shared.ts index ce941c97..9cddafde 100644 --- a/ts/src/methods/itkwasm-shared.ts +++ b/ts/src/methods/itkwasm-shared.ts @@ -14,6 +14,13 @@ import type { NgffImage } from "../types/ngff_image.ts"; import type { ZarrCodec } from "../utils/codecs.ts"; import { defaultCodecs } from "../utils/codecs.ts"; import { zarrGet, zarrSet } from "../utils/worker_pool.ts"; +import { + calculateStride, + componentTypeOf, + transposeArray, +} from "../utils/transpose.ts"; + +export { calculateStride, transposeArray }; export const SPATIAL_DIMS = ["x", "y", "z"]; @@ -231,30 +238,6 @@ function copyTypedArray( } } -/** - * Get ITK component type from typed array - */ -export function getItkComponentType( - data: unknown, -): - | "uint8" - | "int8" - | "uint16" - | "int16" - | "uint32" - | "int32" - | "float32" - | "float64" { - if (data instanceof Uint8Array) return "uint8"; - if (data instanceof Int8Array) return "int8"; - if (data instanceof Uint16Array) return "uint16"; - if (data instanceof Int16Array) return "int16"; - if (data instanceof Uint32Array) return "uint32"; - if (data instanceof Int32Array) return "int32"; - if (data instanceof Float64Array) return "float64"; - return "float32"; -} - /** * Integer component types eligible for the Gaussian float32 workaround */ @@ -362,129 +345,6 @@ export function createIdentityMatrix(dimension: number): Float64Array { return matrix; } -/** - * Calculate stride for array - */ -function calculateStride(shape: number[]): number[] { - const stride = new Array(shape.length); - stride[shape.length - 1] = 1; - for (let i = shape.length - 2; i >= 0; i--) { - stride[i] = stride[i + 1] * shape[i + 1]; - } - return stride; -} - -/** - * Transpose array data according to permutation - */ -export function transposeArray( - data: unknown, - shape: number[], - permutation: number[], - componentType: - | "uint8" - | "int8" - | "uint16" - | "int16" - | "uint32" - | "int32" - | "float32" - | "float64", -): - | Float32Array - | Float64Array - | Uint8Array - | Int8Array - | Uint16Array - | Int16Array - | Uint32Array - | Int32Array { - const typedData = data as - | Float32Array - | Float64Array - | Uint8Array - | Int8Array - | Uint16Array - | Int16Array - | Uint32Array - | Int32Array; - - // Create output array of same type - let output: - | Float32Array - | Float64Array - | Uint8Array - | Int8Array - | Uint16Array - | Int16Array - | Uint32Array - | Int32Array; - const totalSize = typedData.length; - - switch (componentType) { - case "uint8": - output = new Uint8Array(totalSize); - break; - case "int8": - output = new Int8Array(totalSize); - break; - case "uint16": - output = new Uint16Array(totalSize); - break; - case "int16": - output = new Int16Array(totalSize); - break; - case "uint32": - output = new Uint32Array(totalSize); - break; - case "int32": - output = new Int32Array(totalSize); - break; - case "float64": - output = new Float64Array(totalSize); - break; - case "float32": - default: - output = new Float32Array(totalSize); - break; - } - - // Calculate strides for source - const sourceStride = calculateStride(shape); - - // Calculate new shape after permutation - const newShape = permutation.map((i) => shape[i]); - const targetStride = calculateStride(newShape); - - // Perform transpose - const indices = new Array(shape.length).fill(0); - - for (let i = 0; i < totalSize; i++) { - // Calculate source index from multi-dimensional indices - let sourceIdx = 0; - for (let j = 0; j < shape.length; j++) { - sourceIdx += indices[j] * sourceStride[j]; - } - - // Calculate target index with permuted dimensions - let targetIdx = 0; - for (let j = 0; j < permutation.length; j++) { - targetIdx += indices[permutation[j]] * targetStride[j]; - } - - output[targetIdx] = typedData[sourceIdx]; - - // Increment indices - for (let j = shape.length - 1; j >= 0; j--) { - indices[j]++; - if (indices[j] < shape[j]) break; - indices[j] = 0; - } - } - - return output; -} - /** * Convert zarr array to ITK-Wasm Image format * If isVector is true, ensures "c" dimension is last by transposing if needed @@ -533,7 +393,7 @@ export async function zarrToItkImage( result.data, result.shape, permutation, - getItkComponentType(result.data), + componentTypeOf(result.data), ); } else { // "c" already at end or not present, just copy data @@ -556,7 +416,7 @@ export async function zarrToItkImage( const itkImage: Image = { imageType: { dimension: spatialShape.length, - componentType: getItkComponentType(data), + componentType: componentTypeOf(data), pixelType: isVector ? "VariableLengthVector" : "Scalar", components, }, @@ -705,7 +565,7 @@ export async function itkImageToZarr( itkImage.data, currentShape, permutation, - getItkComponentType(itkImage.data), + componentTypeOf(itkImage.data), ); } diff --git a/ts/src/process/to_multiscales-shared.ts b/ts/src/process/to_multiscales-shared.ts index f716d825..557a198e 100644 --- a/ts/src/process/to_multiscales-shared.ts +++ b/ts/src/process/to_multiscales-shared.ts @@ -10,6 +10,7 @@ import { Methods } from "../types/methods.ts"; import type { NgffMultiscales } from "../types/multiscales.ts"; import type { NgffImage } from "../types/ngff_image.ts"; +import { canonicalAxisOrder } from "../utils/axis_order.ts"; import type { ZarrCodec } from "../utils/codecs.ts"; // deno-lint-ignore no-unused-vars import { bytesOnlyCodecs, defaultCodecs } from "../utils/codecs.ts"; @@ -78,18 +79,29 @@ export type DownsampleFunction = ( * @returns NgffMultiscales object */ export async function toMultiscalesCore( - image: NgffImage, + inputImage: NgffImage, options: ToMultiscalesOptions, downsampleItkWasm: DownsampleFunction, ): Promise { const { scaleFactors = [2, 4], method = Methods.ITKWASM_GAUSSIAN, - chunks: _chunks, + chunks: requestedChunks, codecs, orientation, } = options; + // OME-Zarr orders axes time, then channel, then space. Channel-last input + // (ITK component images, the 4-D/5-D default dims) is normalized here so the + // generated metadata and every scale are spec-ordered, and so a model the + // writer would refuse below 0.9.dev1 never reaches it. + const image = await canonicalAxisOrder(inputImage, codecs); + // A positional `chunks` array indexes the caller's dims, so it follows them + // through the reordering. The dim-keyed and scalar forms need no change. + const _chunks = Array.isArray(requestedChunks) && image !== inputImage + ? image.dims.map((dim) => requestedChunks[inputImage.dims.indexOf(dim)]) + : requestedChunks; + // The vector-component axis type (RFC-5 displacement/coordinate fields) is // carried on the input image, mirroring axesUnits / axesOrientations. const axesTypes = image.axesTypes; diff --git a/ts/src/utils/axis_order.ts b/ts/src/utils/axis_order.ts new file mode 100644 index 00000000..a180fc1f --- /dev/null +++ b/ts/src/utils/axis_order.ts @@ -0,0 +1,94 @@ +// SPDX-FileCopyrightText: Copyright (c) Fideus Labs LLC +// SPDX-License-Identifier: MIT +/** + * Normalization of an image's axes to the OME-Zarr specification order. + * + * Mirrors `py/ngff_zarr/methods/_support.py`. + */ + +import * as zarr from "zarrita"; + +import { NgffImage } from "../types/ngff_image.ts"; +import type { ZarrCodec } from "./codecs.ts"; +import { defaultCodecs } from "./codecs.ts"; +import { + calculateStride, + componentTypeOf, + transposeArray, +} from "./transpose.ts"; +import { zarrGet, zarrSet } from "./worker_pool.ts"; + +/** The OME-Zarr specification axis order: time, then channel, then space. */ +export const CANONICAL_AXIS_ORDER = ["t", "c", "z", "y", "x"]; + +/** + * Return `image` with its dims in the spec axis order `(t, c, z, y, x)`. + * + * OME-Zarr requires axes ordered by type: time, then channel, then space. + * Conversion sources produce channel-last layouts, so the multiscale pipeline + * normalizes them here and the generated metadata is spec-ordered. + * + * An image whose dims fall outside the canonical set is returned unchanged: + * an axis model that was not expressible before RFC-3 carries no spec ordering + * to normalize to. + * + * Where the Python port transposes lazily through dask, this reads the array + * and writes the permuted buffer into a new in-memory zarr array. The + * downsampling path materializes the image anyway. + * + * @param image - The image to normalize. + * @param codecs - Codec pipeline for the reordered array; defaults to + * {@link defaultCodecs}. + * @returns The image itself when already ordered, otherwise a reordered copy. + */ +export async function canonicalAxisOrder( + image: NgffImage, + codecs?: ZarrCodec[], +): Promise { + const dims = image.dims; + const newDims = CANONICAL_AXIS_ORDER.filter((dim) => dims.includes(dim)); + if ( + newDims.length !== dims.length || + newDims.every((dim, index) => dim === dims[index]) + ) { + return image; + } + + const permutation = newDims.map((dim) => dims.indexOf(dim)); + const result = await zarrGet(image.data); + const componentType = componentTypeOf(result.data); + const transposed = transposeArray( + result.data, + [...result.shape], + permutation, + componentType, + ); + + const shape = permutation.map((index) => result.shape[index]); + const chunkShape = permutation.map((index) => image.data.chunks[index]); + const store: Map = new Map(); + const array = await zarr.create(zarr.root(store).resolve("/0"), { + shape, + chunk_shape: chunkShape, + data_type: image.data.dtype, + fill_value: 0, + codecs: codecs ?? defaultCodecs(image.data.dtype), + }); + await zarrSet(array, shape.map(() => null), { + data: transposed, + shape, + stride: calculateStride(shape), + }); + + return new NgffImage({ + data: array as zarr.Array, + dims: newDims, + scale: image.scale, + translation: image.translation, + name: image.name, + axesUnits: image.axesUnits, + axesOrientations: image.axesOrientations, + axesTypes: image.axesTypes, + computedCallbacks: image.computedCallbacks, + }); +} diff --git a/ts/src/utils/transpose.ts b/ts/src/utils/transpose.ts new file mode 100644 index 00000000..0d0a4abb --- /dev/null +++ b/ts/src/utils/transpose.ts @@ -0,0 +1,125 @@ +// SPDX-FileCopyrightText: Copyright (c) Fideus Labs LLC +// SPDX-License-Identifier: MIT +/** + * Typed-array transposition, independent of any image toolkit. + * + * Lives outside `methods/` so modules that must not pull in `itk-wasm` can + * reorder array data. + */ + +/** Every numeric typed array the zarr data types map onto. */ +export type NumericTypedArray = + | Float32Array + | Float64Array + | Uint8Array + | Int8Array + | Uint16Array + | Int16Array + | Uint32Array + | Int32Array; + +/** The element type of a {@link NumericTypedArray}. */ +export type ComponentType = + | "uint8" + | "int8" + | "uint16" + | "int16" + | "uint32" + | "int32" + | "float32" + | "float64"; + +/** The {@link ComponentType} of a typed array; `float32` when unrecognized. */ +export function componentTypeOf(data: unknown): ComponentType { + if (data instanceof Uint8Array) return "uint8"; + if (data instanceof Int8Array) return "int8"; + if (data instanceof Uint16Array) return "uint16"; + if (data instanceof Int16Array) return "int16"; + if (data instanceof Uint32Array) return "uint32"; + if (data instanceof Int32Array) return "int32"; + if (data instanceof Float64Array) return "float64"; + return "float32"; +} + +/** Row-major (C-order) strides for `shape`. */ +export function calculateStride(shape: number[]): number[] { + const stride = new Array(shape.length); + stride[shape.length - 1] = 1; + for (let i = shape.length - 2; i >= 0; i--) { + stride[i] = stride[i + 1] * shape[i + 1]; + } + return stride; +} + +/** Allocate a typed array of `length` elements of `componentType`. */ +function allocate( + componentType: ComponentType, + length: number, +): NumericTypedArray { + switch (componentType) { + case "uint8": + return new Uint8Array(length); + case "int8": + return new Int8Array(length); + case "uint16": + return new Uint16Array(length); + case "int16": + return new Int16Array(length); + case "uint32": + return new Uint32Array(length); + case "int32": + return new Int32Array(length); + case "float64": + return new Float64Array(length); + case "float32": + default: + return new Float32Array(length); + } +} + +/** + * Reorder `data` so that axis `permutation[i]` of `shape` becomes axis `i`. + * + * @param data - Row-major buffer of `shape`. + * @param shape - The buffer's current shape. + * @param permutation - Source axis index per target axis. + * @param componentType - The buffer's element type. + * @returns A new buffer of the permuted shape. + */ +export function transposeArray( + data: unknown, + shape: number[], + permutation: number[], + componentType: ComponentType, +): NumericTypedArray { + const typedData = data as NumericTypedArray; + const totalSize = typedData.length; + const output = allocate(componentType, totalSize); + + const sourceStride = calculateStride(shape); + const newShape = permutation.map((i) => shape[i]); + const targetStride = calculateStride(newShape); + + const indices = new Array(shape.length).fill(0); + for (let i = 0; i < totalSize; i++) { + let sourceIdx = 0; + for (let j = 0; j < shape.length; j++) { + sourceIdx += indices[j] * sourceStride[j]; + } + + let targetIdx = 0; + for (let j = 0; j < permutation.length; j++) { + targetIdx += indices[permutation[j]] * targetStride[j]; + } + + output[targetIdx] = typedData[sourceIdx]; + + for (let j = shape.length - 1; j >= 0; j--) { + indices[j]++; + if (indices[j] < shape[j]) break; + indices[j] = 0; + } + } + + return output; +} diff --git a/ts/test/canonical_axis_order_test.ts b/ts/test/canonical_axis_order_test.ts new file mode 100644 index 00000000..07d2856a --- /dev/null +++ b/ts/test/canonical_axis_order_test.ts @@ -0,0 +1,106 @@ +#!/usr/bin/env -S deno test --allow-read --allow-write +// SPDX-FileCopyrightText: Copyright (c) Fideus Labs LLC +// SPDX-License-Identifier: MIT +/** + * Normalization of channel-last input to the OME-Zarr axis order. + * + * Mirrors the Python port's `_canonical_axis_order` coverage: the pipeline + * reorders axes to `(t, c, z, y, x)`, moves the data with them, and leaves an + * axis model outside that vocabulary alone. + */ + +import { assertEquals, assertNotEquals } from "@std/assert"; +import * as zarr from "zarrita"; + +import { + toMultiscales, + toNgffImage, +} from "../src/process/to_multiscales-node.ts"; +import { Methods } from "../src/types/methods.ts"; +import { NgffImage } from "../src/types/ngff_image.ts"; +import { canonicalAxisOrder } from "../src/utils/axis_order.ts"; +import { zarrGet } from "../src/utils/worker_pool.ts"; +import { calculateStride } from "../src/utils/transpose.ts"; + +/** A ramp image over `dims`/`shape`, values `0..n-1` in row-major order. */ +async function rampImage( + dims: string[], + shape: number[], +): Promise { + const total = shape.reduce((a, b) => a * b, 1); + const data = new Float32Array(total); + for (let i = 0; i < total; i++) { + data[i] = i; + } + return await toNgffImage(data, { dims, shape }); +} + +/** Every element of `array`, read back in row-major order. */ +async function readAll( + array: zarr.Array, +): Promise { + const result = await zarrGet(array); + return Array.from(result.data as ArrayLike); +} + +Deno.test("canonicalAxisOrder - channel-last input is reordered", async () => { + const image = await rampImage(["z", "y", "x", "c"], [2, 3, 4, 2]); + const normalized = await canonicalAxisOrder(image); + + assertEquals(normalized.dims, ["c", "z", "y", "x"]); + assertEquals([...normalized.data.shape], [2, 2, 3, 4]); +}); + +Deno.test("canonicalAxisOrder - the data follows the axes", async () => { + // A relabelling that left the buffer alone would keep the ramp intact, so + // compare against the transpose computed independently. + const dims = ["y", "x", "c"]; + const shape = [2, 3, 4]; + const image = await rampImage(dims, shape); + const normalized = await canonicalAxisOrder(image); + + assertEquals(normalized.dims, ["c", "y", "x"]); + const permutation = [2, 0, 1]; + const newShape = permutation.map((i) => shape[i]); + const sourceStride = calculateStride(shape); + const expected = new Array(shape.reduce((a, b) => a * b, 1)); + for (let c = 0; c < newShape[0]; c++) { + for (let y = 0; y < newShape[1]; y++) { + for (let x = 0; x < newShape[2]; x++) { + const target = (c * newShape[1] + y) * newShape[2] + x; + expected[target] = y * sourceStride[0] + x * sourceStride[1] + + c * sourceStride[2]; + } + } + } + assertEquals(await readAll(normalized.data), expected); +}); + +Deno.test("canonicalAxisOrder - already ordered input is untouched", async () => { + const image = await rampImage(["c", "y", "x"], [2, 3, 4]); + assertEquals(await canonicalAxisOrder(image), image); +}); + +Deno.test("canonicalAxisOrder - a non-canonical vocabulary is left alone", async () => { + // RFC-3 axis names carry no spec ordering to normalize to. + const image = await rampImage(["b", "a"], [2, 3]); + const normalized = await canonicalAxisOrder(image); + + assertEquals(normalized, image); + assertEquals(normalized.dims, ["b", "a"]); +}); + +Deno.test("toMultiscales - generated axes are spec-ordered", async () => { + const image = await rampImage(["z", "y", "x", "c"], [8, 16, 16, 2]); + const multiscales = await toMultiscales(image, { + scaleFactors: [], + method: Methods.ITKWASM_GAUSSIAN, + }); + + assertEquals( + multiscales.metadata.axes.map((axis) => axis.name), + ["c", "z", "y", "x"], + ); + assertEquals(multiscales.images[0].dims, ["c", "z", "y", "x"]); + assertNotEquals(multiscales.images[0].dims, image.dims); +}); diff --git a/ts/test/to_multiscales_itkwasm_test.ts b/ts/test/to_multiscales_itkwasm_test.ts index ae8237a0..3331c335 100644 --- a/ts/test/to_multiscales_itkwasm_test.ts +++ b/ts/test/to_multiscales_itkwasm_test.ts @@ -59,10 +59,12 @@ Deno.test("downsample zycx", async () => { const store: MemoryStore = new Map(); await toNgffZarr(store, multiscales); - assertEquals(multiscales.images[0].dims[0], "z"); - assertEquals(multiscales.images[0].dims[2], "c"); - assertEquals(multiscales.images[1].data.shape[0], 16); // z downsampled - assertEquals(multiscales.images[1].data.shape[2], 2); // c unchanged + // The channel-last input is normalized to (c, z, y, x). Mirrors + // py/test/test_to_ngff_zarr_itkwasm.py::test_downsample_zycx. + assertEquals(multiscales.images[0].dims[0], "c"); + assertEquals(multiscales.images[0].dims[1], "z"); + assertEquals(multiscales.images[1].data.shape[0], 2); // c unchanged + assertEquals(multiscales.images[1].data.shape[1], 16); // z downsampled }); Deno.test("downsample cxyz", async () => { @@ -134,12 +136,14 @@ Deno.test("downsample tzycx", async () => { const store: MemoryStore = new Map(); await toNgffZarr(store, multiscales); + // The channel-last input is normalized to (t, c, z, y, x). Mirrors + // py/test/test_to_ngff_zarr_itkwasm.py::test_downsample_tzycx. assertEquals(multiscales.images[0].dims[0], "t"); - assertEquals(multiscales.images[0].dims[1], "z"); - assertEquals(multiscales.images[0].dims[3], "c"); + assertEquals(multiscales.images[0].dims[1], "c"); + assertEquals(multiscales.images[0].dims[2], "z"); assertEquals(multiscales.images[1].data.shape[0], 2); // t unchanged - assertEquals(multiscales.images[1].data.shape[1], 16); // z downsampled - assertEquals(multiscales.images[1].data.shape[3], 2); // c unchanged + assertEquals(multiscales.images[1].data.shape[1], 2); // c unchanged + assertEquals(multiscales.images[1].data.shape[2], 16); // z downsampled }); Deno.test("downsample tcxyz", async () => { From 85375690d54da57bb8118192d244703e53ab5e1c Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Fri, 21 Aug 2026 01:09:02 +0200 Subject: [PATCH 2/2] fix(ts): preserve 64-bit dtypes and stream the axis reordering by chunk `componentTypeOf` fell through to `float32` for the `bigint` arrays, so reordering an `int64` or `uint64` image allocated a `Float32Array` and threw on the first element. It now names those two types and allocates their constructors; `transposeArray` returns the caller's array type, so the ITK paths, which have no 64-bit component type, stay narrowed. The reordering read the whole source and held a full transposed buffer beside the compressed copy, which for a metadata-only request over a remote store meant several times the image in memory. It now walks the source chunk grid, so the region read and the transposed buffer are both chunk-sized. Reading the whole image is inherent to the eager copy and is documented. `(z, y, c, x)` does not place `c` last, so the docs and test comments call such a layout non-canonical. --- docs/typescript.md | 2 +- ts/src/utils/axis_order.ts | 74 +++++++++++++++----- ts/src/utils/transpose.ts | 42 ++++++++--- ts/test/canonical_axis_order_test.ts | 96 +++++++++++++++++++++++++- ts/test/to_multiscales_itkwasm_test.ts | 4 +- 5 files changed, 185 insertions(+), 33 deletions(-) diff --git a/docs/typescript.md b/docs/typescript.md index 50dbdf29..b4310b53 100644 --- a/docs/typescript.md +++ b/docs/typescript.md @@ -465,7 +465,7 @@ async function toMultiscales( **Returns:** NgffMultiscales with generated pyramid levels Axes are normalized to the OME-Zarr order, time then channel then space, so a -channel-last input such as `(z, y, c, x)` yields `(c, z, y, x)` and the data is +non-canonical input such as `(z, y, c, x)` yields `(c, z, y, x)` and the data is reordered with it. An axis model outside the `(t, c, z, y, x)` vocabulary is left alone: it carries no spec ordering to normalize to. A positional `chunks` array indexes the dims you passed and follows them through the reordering. diff --git a/ts/src/utils/axis_order.ts b/ts/src/utils/axis_order.ts index a180fc1f..ee8ea222 100644 --- a/ts/src/utils/axis_order.ts +++ b/ts/src/utils/axis_order.ts @@ -21,6 +21,21 @@ import { zarrGet, zarrSet } from "./worker_pool.ts"; /** The OME-Zarr specification axis order: time, then channel, then space. */ export const CANONICAL_AXIS_ORDER = ["t", "c", "z", "y", "x"]; +/** Every chunk origin of a `shape`/`chunks` grid, in row-major order. */ +function* chunkOrigins(shape: number[], chunks: number[]): Generator { + const counts = shape.map((size, axis) => Math.ceil(size / chunks[axis])); + const total = counts.reduce((a, b) => a * b, 1); + for (let flat = 0; flat < total; flat++) { + const origin = new Array(shape.length); + let rest = flat; + for (let axis = shape.length - 1; axis >= 0; axis--) { + origin[axis] = (rest % counts[axis]) * chunks[axis]; + rest = Math.floor(rest / counts[axis]); + } + yield origin; + } +} + /** * Return `image` with its dims in the spec axis order `(t, c, z, y, x)`. * @@ -32,9 +47,11 @@ export const CANONICAL_AXIS_ORDER = ["t", "c", "z", "y", "x"]; * an axis model that was not expressible before RFC-3 carries no spec ordering * to normalize to. * - * Where the Python port transposes lazily through dask, this reads the array - * and writes the permuted buffer into a new in-memory zarr array. The - * downsampling path materializes the image anyway. + * Where the Python port transposes lazily through dask, this writes a reordered + * copy into a new in-memory zarr array. The copy proceeds one source chunk at + * a time, so peak working memory is a chunk rather than the whole image, but + * the whole image is read: a remote store is fetched in full even when no + * downsampling was requested. * * @param image - The image to normalize. * @param codecs - Codec pipeline for the reordered array; defaults to @@ -55,17 +72,11 @@ export async function canonicalAxisOrder( } const permutation = newDims.map((dim) => dims.indexOf(dim)); - const result = await zarrGet(image.data); - const componentType = componentTypeOf(result.data); - const transposed = transposeArray( - result.data, - [...result.shape], - permutation, - componentType, - ); + const sourceShape = [...image.data.shape]; + const sourceChunks = [...image.data.chunks]; + const shape = permutation.map((index) => sourceShape[index]); + const chunkShape = permutation.map((index) => sourceChunks[index]); - const shape = permutation.map((index) => result.shape[index]); - const chunkShape = permutation.map((index) => image.data.chunks[index]); const store: Map = new Map(); const array = await zarr.create(zarr.root(store).resolve("/0"), { shape, @@ -74,11 +85,38 @@ export async function canonicalAxisOrder( fill_value: 0, codecs: codecs ?? defaultCodecs(image.data.dtype), }); - await zarrSet(array, shape.map(() => null), { - data: transposed, - shape, - stride: calculateStride(shape), - }); + + // One source chunk at a time: the region read and the transposed buffer are + // both chunk-sized, so an image far larger than memory still converts. + for (const origin of chunkOrigins(sourceShape, sourceChunks)) { + const region = origin.map((start, axis) => ({ + start, + stop: Math.min(start + sourceChunks[axis], sourceShape[axis]), + })); + const block = await zarrGet( + image.data, + region.map(({ start, stop }) => zarr.slice(start, stop)), + ); + const blockShape = [...block.shape]; + const transposed = transposeArray( + block.data, + blockShape, + permutation, + componentTypeOf(block.data), + ); + const targetShape = permutation.map((index) => blockShape[index]); + await zarrSet( + array, + permutation.map((index) => + zarr.slice(region[index].start, region[index].stop) + ), + { + data: transposed, + shape: targetShape, + stride: calculateStride(targetShape), + }, + ); + } return new NgffImage({ data: array as zarr.Array, diff --git a/ts/src/utils/transpose.ts b/ts/src/utils/transpose.ts index 0d0a4abb..b7ab820e 100644 --- a/ts/src/utils/transpose.ts +++ b/ts/src/utils/transpose.ts @@ -7,7 +7,7 @@ * reorder array data. */ -/** Every numeric typed array the zarr data types map onto. */ +/** Every number-valued typed array the zarr data types map onto. */ export type NumericTypedArray = | Float32Array | Float64Array @@ -18,7 +18,13 @@ export type NumericTypedArray = | Uint32Array | Int32Array; -/** The element type of a {@link NumericTypedArray}. */ +/** The 64-bit integer arrays, whose elements are `bigint` rather than `number`. */ +export type BigTypedArray = BigInt64Array | BigUint64Array; + +/** Any typed array {@link transposeArray} accepts. */ +export type AnyTypedArray = NumericTypedArray | BigTypedArray; + +/** The element type of an {@link AnyTypedArray}. */ export type ComponentType = | "uint8" | "int8" @@ -26,6 +32,8 @@ export type ComponentType = | "int16" | "uint32" | "int32" + | "int64" + | "uint64" | "float32" | "float64"; @@ -37,6 +45,8 @@ export function componentTypeOf(data: unknown): ComponentType { if (data instanceof Int16Array) return "int16"; if (data instanceof Uint32Array) return "uint32"; if (data instanceof Int32Array) return "int32"; + if (data instanceof BigInt64Array) return "int64"; + if (data instanceof BigUint64Array) return "uint64"; if (data instanceof Float64Array) return "float64"; return "float32"; } @@ -55,7 +65,7 @@ export function calculateStride(shape: number[]): number[] { function allocate( componentType: ComponentType, length: number, -): NumericTypedArray { +): AnyTypedArray { switch (componentType) { case "uint8": return new Uint8Array(length); @@ -69,6 +79,10 @@ function allocate( return new Uint32Array(length); case "int32": return new Int32Array(length); + case "int64": + return new BigInt64Array(length); + case "uint64": + return new BigUint64Array(length); case "float64": return new Float64Array(length); case "float32": @@ -77,6 +91,9 @@ function allocate( } } +/** An indexable view, so one loop can copy `number` and `bigint` elements. */ +type ElementView = { [index: number]: number | bigint; length: number }; + /** * Reorder `data` so that axis `permutation[i]` of `shape` becomes axis `i`. * @@ -84,17 +101,24 @@ function allocate( * @param shape - The buffer's current shape. * @param permutation - Source axis index per target axis. * @param componentType - The buffer's element type. - * @returns A new buffer of the permuted shape. + * @returns A new buffer of the permuted shape, of the same array type as + * `data`. Callers that have already narrowed the element type keep it: `T` + * appears only in the return position, so a caller holding, say, an ITK + * component buffer does not widen to the `bigint` arrays. */ -export function transposeArray( +export function transposeArray( data: unknown, shape: number[], permutation: number[], componentType: ComponentType, -): NumericTypedArray { - const typedData = data as NumericTypedArray; +): T { + const typedData = data as AnyTypedArray; const totalSize = typedData.length; const output = allocate(componentType, totalSize); + // `int64`/`uint64` arrays hold `bigint`; the element type is uniform between + // source and output, but TypeScript cannot narrow the union across the pair. + const source = typedData as unknown as ElementView; + const target = output as unknown as ElementView; const sourceStride = calculateStride(shape); const newShape = permutation.map((i) => shape[i]); @@ -112,7 +136,7 @@ export function transposeArray( targetIdx += indices[permutation[j]] * targetStride[j]; } - output[targetIdx] = typedData[sourceIdx]; + target[targetIdx] = source[sourceIdx]; for (let j = shape.length - 1; j >= 0; j--) { indices[j]++; @@ -121,5 +145,5 @@ export function transposeArray( } } - return output; + return output as T; } diff --git a/ts/test/canonical_axis_order_test.ts b/ts/test/canonical_axis_order_test.ts index 07d2856a..7551c296 100644 --- a/ts/test/canonical_axis_order_test.ts +++ b/ts/test/canonical_axis_order_test.ts @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright (c) Fideus Labs LLC // SPDX-License-Identifier: MIT /** - * Normalization of channel-last input to the OME-Zarr axis order. + * Normalization of a non-canonical axis order to the OME-Zarr axis order. * * Mirrors the Python port's `_canonical_axis_order` coverage: the pipeline * reorders axes to `(t, c, z, y, x)`, moves the data with them, and leaves an @@ -19,7 +19,7 @@ import { import { Methods } from "../src/types/methods.ts"; import { NgffImage } from "../src/types/ngff_image.ts"; import { canonicalAxisOrder } from "../src/utils/axis_order.ts"; -import { zarrGet } from "../src/utils/worker_pool.ts"; +import { zarrGet, zarrSet } from "../src/utils/worker_pool.ts"; import { calculateStride } from "../src/utils/transpose.ts"; /** A ramp image over `dims`/`shape`, values `0..n-1` in row-major order. */ @@ -43,7 +43,7 @@ async function readAll( return Array.from(result.data as ArrayLike); } -Deno.test("canonicalAxisOrder - channel-last input is reordered", async () => { +Deno.test("canonicalAxisOrder - a non-canonical order is normalized", async () => { const image = await rampImage(["z", "y", "x", "c"], [2, 3, 4, 2]); const normalized = await canonicalAxisOrder(image); @@ -104,3 +104,93 @@ Deno.test("toMultiscales - generated axes are spec-ordered", async () => { assertEquals(multiscales.images[0].dims, ["c", "z", "y", "x"]); assertNotEquals(multiscales.images[0].dims, image.dims); }); + +/** An in-memory image over `dims`/`shape` with an explicit dtype and chunking. */ +async function chunkedImage( + dims: string[], + shape: number[], + chunkShape: number[], + dataType: "float32" | "int64", +): Promise { + const total = shape.reduce((a, b) => a * b, 1); + const store: Map = new Map(); + const array = await zarr.create(zarr.root(store).resolve("/0"), { + shape, + chunk_shape: chunkShape, + data_type: dataType, + fill_value: 0, + }); + const data = dataType === "int64" + ? BigInt64Array.from({ length: total }, (_, i) => BigInt(i)) + : Float32Array.from({ length: total }, (_, i) => i); + await zarrSet(array as never, shape.map(() => null), { + data, + shape, + stride: calculateStride(shape), + } as never); + const scale: Record = {}; + const translation: Record = {}; + for (const dim of dims) { + scale[dim] = 1.0; + translation[dim] = 0.0; + } + return new NgffImage({ + data: array as zarr.Array, + dims, + scale, + translation, + name: "image", + axesUnits: undefined, + computedCallbacks: undefined, + }); +} + +Deno.test("canonicalAxisOrder - a multi-chunk array transposes correctly", async () => { + // The copy walks the source chunk grid, so a single-chunk fixture would not + // exercise the per-chunk region arithmetic. + const shape = [4, 6, 2]; + const image = await chunkedImage( + ["y", "x", "c"], + shape, + [2, 3, 1], + "float32", + ); + const normalized = await canonicalAxisOrder(image); + + assertEquals(normalized.dims, ["c", "y", "x"]); + assertEquals([...normalized.data.shape], [2, 4, 6]); + + const sourceStride = calculateStride(shape); + const expected: number[] = []; + for (let c = 0; c < 2; c++) { + for (let y = 0; y < 4; y++) { + for (let x = 0; x < 6; x++) { + expected.push( + y * sourceStride[0] + x * sourceStride[1] + c * sourceStride[2], + ); + } + } + } + assertEquals(await readAll(normalized.data), expected); +}); + +Deno.test("canonicalAxisOrder - a 64-bit integer image transposes", async () => { + // int64 reads back as a BigInt64Array; treating it as float32 would allocate + // the wrong buffer and throw on the first bigint assignment. + const image = await chunkedImage( + ["y", "x", "c"], + [2, 3, 2], + [2, 3, 2], + "int64", + ); + const normalized = await canonicalAxisOrder(image); + + assertEquals(normalized.dims, ["c", "y", "x"]); + assertEquals([...normalized.data.shape], [2, 2, 3]); + const values = (await zarrGet(normalized.data)).data; + assertEquals(values instanceof BigInt64Array, true); + assertEquals( + Array.from(values as BigInt64Array).map(Number), + [0, 2, 4, 6, 8, 10, 1, 3, 5, 7, 9, 11], + ); +}); diff --git a/ts/test/to_multiscales_itkwasm_test.ts b/ts/test/to_multiscales_itkwasm_test.ts index 3331c335..2cade1e1 100644 --- a/ts/test/to_multiscales_itkwasm_test.ts +++ b/ts/test/to_multiscales_itkwasm_test.ts @@ -59,7 +59,7 @@ Deno.test("downsample zycx", async () => { const store: MemoryStore = new Map(); await toNgffZarr(store, multiscales); - // The channel-last input is normalized to (c, z, y, x). Mirrors + // The non-canonical input is normalized to (c, z, y, x). Mirrors // py/test/test_to_ngff_zarr_itkwasm.py::test_downsample_zycx. assertEquals(multiscales.images[0].dims[0], "c"); assertEquals(multiscales.images[0].dims[1], "z"); @@ -136,7 +136,7 @@ Deno.test("downsample tzycx", async () => { const store: MemoryStore = new Map(); await toNgffZarr(store, multiscales); - // The channel-last input is normalized to (t, c, z, y, x). Mirrors + // The non-canonical input is normalized to (t, c, z, y, x). Mirrors // py/test/test_to_ngff_zarr_itkwasm.py::test_downsample_tzycx. assertEquals(multiscales.images[0].dims[0], "t"); assertEquals(multiscales.images[0].dims[1], "c");