Better support for non-aligned axes

Author: manzt

src/io.ts

@@ -230,8 +230,8 @@ export function initLayerStateFromSource(source: SourceData & { id: string }): L
         lut: label.lut,
       })),
       transformSourceSelection: getTransformSourceSelectionFromLabels(
-        source.labels.map((label) => label.loader[0].labels),
-        source.loader[0].labels,
+        source.labels.map((label) => label.loader[0]),
+        source.loader[0],
       ),
     };
   }
@@ -247,25 +247,44 @@ export function initLayerStateFromSource(source: SourceData & { id: string }): L
   };
 }
 
-function getTransformSourceSelectionFromLabels(labelAxesNames: Array<Array<string>>, sourceAxesNames: Array<string>) {
+function getTransformSourceSelectionFromLabels(
+  labelsResolutions: Array<{ shape: Array<number>; labels: Array<string> }>,
+  source: { shape: Array<number>; labels: Array<string> },
+) {
+  // representative source for labels
+  const labelsSource = labelsResolutions[0];
   utils.assert(
-    labelAxesNames[0].every((name) => sourceAxesNames.includes(name)),
-    () =>
-      `Label axes names be a subset of source. Source: ${JSON.stringify(sourceAxesNames)} Labels: ${JSON.stringify(labelAxesNames)}`,
+    source.shape.length === source.labels.length,
+    `Image source axes and shape are not same rank. Got ${JSON.stringify(source)}`,
   );
-  for (const axes of labelAxesNames) {
+  utils.assert(
+    labelsSource.shape.length === labelsSource.labels.length,
+    `Label axes and shape are not same rank. Got ${JSON.stringify(labelsSource)}`,
+  );
+  utils.assert(
+    utils.zip(labelsSource.labels, source.labels).every(([a, b]) => a === b),
+    `Label axes MUST be a subset of source. Source: ${JSON.stringify(source.labels)} Labels: ${JSON.stringify(labelsSource.labels)}`,
+  );
+  for (const { labels, shape } of labelsResolutions.slice(1)) {
     utils.assert(
-      axes.every((name, i) => name === labelAxesNames[0][i]),
-      () => {
-        const mismatchedLabels = labelAxesNames.map((axes) => ({ axes }));
-        return `Error: All labels must share the same axes. Mismatched labels found: ${JSON.stringify(mismatchedLabels)}`;
-      },
+      utils.zip(labels, labelsSource.labels).every(([a, b]) => a === b),
+      `Error: All labels must share the same axes. Mismatched labels found: ${JSON.stringify(labels)}`,
+    );
+    utils.assert(
+      utils.zip(shape, labelsSource.shape).every(([a, b]) => a === b),
+      `Error: All labels must share the same shape. Mismatched labels found: ${JSON.stringify(shape)}`,
     );
   }
+  // Identify labels that should always map to 0, regardless of the source selection.
+  const excludeFromTransformedSelection = new Set(
+    utils
+      .zip(labelsSource.labels, labelsSource.shape)
+      .filter(([_, size]) => size === 1)
+      .map(([name, _]) => name),
+  );
   return (sourceSelection: Array<number>): Array<number> => {
-    return labelAxesNames[0].map((name) => {
-      if (name === "c") return 0; // skip
-      return sourceSelection[sourceAxesNames.indexOf(name)];
-    });
+    return labelsSource.labels.map((name) =>
+      excludeFromTransformedSelection.has(name) ? 0 : sourceSelection[source.labels.indexOf(name)],
+    );
   };
 }

src/utils.ts

@@ -577,3 +577,17 @@ export function coordinateTransformationsToMatrix(multiscales: Array<Ome.Multisc
 
   return mat;
 }
+
+/**
+ * Builds N-tuples of elements from the given N arrays with matching indices,
+ * stopping when the smallest array's end is reached.
+ */
+export function zip<T extends unknown[]>(...arrays: { [K in keyof T]: ReadonlyArray<T[K]> }): T[] {
+  const minLength = arrays.reduce((minLength, arr) => Math.min(arr.length, minLength), Number.POSITIVE_INFINITY);
+  const result: T[] = new Array(minLength);
+  for (let i = 0; i < minLength; i += 1) {
+    const arr = arrays.map((it) => it[i]);
+    result[i] = arr as T;
+  }
+  return result;
+}