fix(ts): order generated axes time, then channel, then space

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
+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.
+
 **Example:**
 ```typescript
 // Basic pyramid generation

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),
       );
     }
 

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<NgffMultiscales> {
   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;

ts/src/utils/axis_order.ts

@@ -0,0 +1,132 @@
+// 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"];
+
+/** Every chunk origin of a `shape`/`chunks` grid, in row-major order. */
+function* chunkOrigins(shape: number[], chunks: number[]): Generator<number[]> {
+  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<number>(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)`.
+ *
+ * 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 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
+ * {@link defaultCodecs}.
+ * @returns The image itself when already ordered, otherwise a reordered copy.
+ */
+export async function canonicalAxisOrder(
+  image: NgffImage,
+  codecs?: ZarrCodec[],
+): Promise<NgffImage> {
+  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 sourceShape = [...image.data.shape];
+  const sourceChunks = [...image.data.chunks];
+  const shape = permutation.map((index) => sourceShape[index]);
+  const chunkShape = permutation.map((index) => sourceChunks[index]);
+
+  const store: Map<string, Uint8Array> = 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),
+  });
+
+  // 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<zarr.DataType, zarr.Readable>,
+    dims: newDims,
+    scale: image.scale,
+    translation: image.translation,
+    name: image.name,
+    axesUnits: image.axesUnits,
+    axesOrientations: image.axesOrientations,
+    axesTypes: image.axesTypes,
+    computedCallbacks: image.computedCallbacks,
+  });
+}