From eae5c8939bda1aa46470aae03a0cea8c3e2395a2 Mon Sep 17 00:00:00 2001 From: Jeran Date: Thu, 1 Oct 2026 16:37:10 +0200 Subject: [PATCH 1/2] fix serious bug where time-slicing was always at first chunk spot --- src/GlobalStates/ZarrStore.ts | 4 +- src/components/zarr/GetArray.ts | 85 ++++++++++---------- src/components/zarr/dataFetchers.ts | 115 ---------------------------- 3 files changed, 41 insertions(+), 163 deletions(-) diff --git a/src/GlobalStates/ZarrStore.ts b/src/GlobalStates/ZarrStore.ts index 50e8bcd24..5e5b61486 100644 --- a/src/GlobalStates/ZarrStore.ts +++ b/src/GlobalStates/ZarrStore.ts @@ -23,13 +23,13 @@ type ZarrState = { fetchKey: number; blobKey: string | undefined; // The key for the stored File blob for a local NC - ndSlices: (number | [number, number | null])[]; + ndSlices: [number, number][]; axisMapping: { x: number, y: number, z: number }; setZSlice: (zSlice: [number , number | null]) => void; setYSlice: (ySlice: [number , number | null]) => void; setXSlice: (xSlice: [number , number | null]) => void; - setNdSlices: (ndSlices: (number | [number, number | null])[]) => void; + setNdSlices: (ndSlices: [number, number][]) => void; setAxisMapping: (mapping: { x: number, y: number, z: number }) => void; setCompress: (compress: boolean) => void; setCurrentStore: (currentStore: any) => void; diff --git a/src/components/zarr/GetArray.ts b/src/components/zarr/GetArray.ts index 2bdd5d7de..e1f4e21bb 100644 --- a/src/components/zarr/GetArray.ts +++ b/src/components/zarr/GetArray.ts @@ -18,51 +18,40 @@ export async function GetArray(varOveride?: string) { const meta = await fetcher.getMetadata(targetVariable as string); const { shape, chunkShape, fillValue, dtype } = meta; const rank = shape.length; - // Identify which dimensions are already explicitly mapped - const parseMapped = (val: any) => typeof val === 'number' && !isNaN(val) && val >= 0 ? val : -1; - const mappedX = parseMapped(axisMapping.x); - const mappedY = parseMapped(axisMapping.y); - const mappedZ = parseMapped(axisMapping.z); - - const mappedDims = new Set([mappedX, mappedY, mappedZ].filter(v => v >= 0)); - const unmappedDims: number[] = []; - for (let i = rank - 1; i >= 0; i--) { - if (!mappedDims.has(i)) unmappedDims.push(i); - } - const xDimIndex = mappedX >= 0 ? mappedX : (unmappedDims.length > 0 ? unmappedDims.shift()! : rank - 1); - const yDimIndex = mappedY >= 0 ? mappedY : (unmappedDims.length > 0 ? unmappedDims.shift()! : rank - 2); - const zDimIndex = mappedZ >= 0 ? mappedZ : (unmappedDims.length > 0 ? unmappedDims.shift()! : -1); + // Explicit mappings first (-1 if unmapped/NaN), then fill gaps with + // the remaining dims from last to first, in x, y, z order. + const mapped = [axisMapping.x, axisMapping.y, axisMapping.z] + .map(v => (Number.isInteger(v) && v >= 0 ? v : -1)); + const free = Array.from({ length: rank }, (_, i) => rank - 1 - i) + .filter(i => !mapped.includes(i)); + const [xDimIndex, yDimIndex, zDimIndex] = mapped.map(m => (m >= 0 ? m : free.shift() ?? -1)); const hasZ = zDimIndex >= 0; - const calcDim = (slice: [number, number | null], dimIdx: number) => { - if (dimIdx < 0) return { start: 0, end: 1, size: 0, chunkDim: 1 }; - const dimSize = shape[dimIdx]; - const chunkDim = chunkShape[dimIdx]; - const start = Math.floor(slice[0] / chunkDim); - const sliceEnd = slice[1] ?? dimSize; - return { start, end: Math.ceil(sliceEnd / chunkDim), size: sliceEnd - slice[0], chunkDim, offset: slice[0] % chunkDim }; - }; + // ndSlices is in datashape order, so a dim's slice is just ndSlices[dim] + const sliceOf = (dimIdx: number): [number, number] => (dimIdx >= 0 ? ndSlices[dimIdx] : [0, 1]); + const xSlice = sliceOf(xDimIndex); + const ySlice = sliceOf(yDimIndex); + const zSlice = sliceOf(zDimIndex); - // If an axis is unmapped in the UI (NaN), we MUST fetch it as a scalar (size 1) - // using the collapsed value from ndSlices. Otherwise, it defaults to [0, null] - // and fetches the entire dimension, leading to massive memory requests (OOM). - const getEffectiveSlice = (mappingIdx: number, dimIdx: number): [number, number | null] => { - let sel; - if (mappingIdx >= 0) sel = ndSlices[mappingIdx]; - if (dimIdx >= 0 && ndSlices && ndSlices[dimIdx] !== undefined) sel = ndSlices[dimIdx]; - if (typeof sel === 'number') return [sel, sel + 1]; // If scalar convert to slice - if (Array.isArray(sel)) return sel as [number, number | null]; - return [0, 1]; // Safe fallback to scalar - }; - const xSlice = getEffectiveSlice(axisMapping.x, xDimIndex); - const ySlice = getEffectiveSlice(axisMapping.y, yDimIndex); - const zSlice = getEffectiveSlice(axisMapping.z, zDimIndex); // Carry slices to stores so we know what to slice from the axisDimarrays - usePlotStore.setState({zSlice, ySlice, xSlice}) - useZarrStore.setState({zSlice, ySlice, xSlice}) - const xDim = calcDim(xSlice, xDimIndex); - const yDim = calcDim(ySlice, yDimIndex); - const zDim = calcDim(zSlice, zDimIndex); + usePlotStore.setState({ xSlice, ySlice, zSlice }); + useZarrStore.setState({ xSlice, ySlice, zSlice }); + + const calcDim = (dimIdx: number) => { + if (dimIdx < 0) return { start: 0, end: 1, size: 0, chunkDim: 1, offset: 0 }; + const [lo, hi] = ndSlices[dimIdx]; + const chunkDim = chunkShape[dimIdx]; + return { + start: Math.floor(lo / chunkDim), + end: Math.ceil(hi / chunkDim), + size: hi - lo, + chunkDim, + offset: lo % chunkDim, + }; + }; + const xDim = calcDim(xDimIndex); + const yDim = calcDim(yDimIndex); + const zDim = calcDim(zDimIndex); let outputShape = hasZ ? [zDim.size, yDim.size, xDim.size] : [yDim.size, xDim.size]; if (coarsen) { @@ -108,7 +97,7 @@ export async function GetArray(varOveride?: string) { const zIndexInRaw = activeDims.indexOf(zDimIndex); const yIndexInRaw = activeDims.indexOf(yDimIndex); const xIndexInRaw = activeDims.indexOf(xDimIndex); - + const sliceStart = [ndSlices[zIndexInRaw][0], ndSlices[yIndexInRaw][0], ndSlices[xIndexInRaw][0]] setStatus("Downloading..."); setProgress(0); for (let z = zDim.start; z < zDim.end; z++) { @@ -119,13 +108,17 @@ export async function GetArray(varOveride?: string) { const cachedChunk = cache.get(cacheName); const isCacheValid = cachedChunk && cachedChunk.kernel.kernelSize === (coarsen ? kernelSize : undefined) && - cachedChunk.kernel.kernelDepth === (coarsen ? kernelDepth : undefined); - + cachedChunk.kernel.kernelDepth === (coarsen ? kernelDepth : undefined) + if (isCacheValid) { + if (cachedChunk.sliceStart != sliceStart){ + cache.set(cacheName, {...cachedChunk, + sliceStart + }); + } continue; } else { const raw = await fetcher.fetchChunk({ variable:targetVariable, rank, shape, chunkShape, x, y, z, xDimIndex, yDimIndex, zDimIndex, idx4D, ndSlices, axisMapping }); - const rawData = Number.isFinite(fillValue) ? raw.data.map((v: number) => v === fillValue ? NaN : v) : raw.data; // Don't map if no fillvalue let [chunkF16, newScalingFactor] = ToFloat16(rawData, scalingFactor); @@ -165,7 +158,7 @@ export async function GetArray(varOveride?: string) { scaling: scalingFactor, compressed: compress, coarsened: coarsen, kernel: { kernelDepth: coarsen ? kernelDepth : undefined, kernelSize: coarsen ? kernelSize : undefined }, fullChunkDim: [zDim.chunkDim, yDim.chunkDim, xDim.chunkDim], - sliceStart: [zSlice[0], ySlice[0], xSlice[0]] + sliceStart }); rescaleIDs.push(chunkID); } diff --git a/src/components/zarr/dataFetchers.ts b/src/components/zarr/dataFetchers.ts index 5e3f755c7..00b094965 100644 --- a/src/components/zarr/dataFetchers.ts +++ b/src/components/zarr/dataFetchers.ts @@ -5,7 +5,6 @@ import { GetSize } from './utils'; import * as zarr from 'zarrita'; import { calculateStrides } from '@/utils/HelperFuncs'; - interface FetchOutput{ data: Float32Array shape: number[] @@ -174,118 +173,4 @@ export function NCFetcher() { return { data, shape: collapsedShape, stride: calculateStrides(collapsedShape) }; }, }; -} - - -//---- Class possiblity - -export class ZarrFetcher { - private outVar: zarr.Array | undefined; - private variable: string; - - constructor(variable: string) { - this.variable = variable; - } - - async init(){ - const { currentStore } = useZarrStore.getState(); - const group = await currentStore; - const tempOutVar = await zarr.open(group.resolve(this.variable), { kind: "array" }); - this.outVar = tempOutVar; - } - - async getMetadata(): Promise { - if (!this.outVar) { - await this.init(); - if (!this.outVar){ - throw new Error(`Init failed`); - } - } - if (!this.outVar.is("number") && !this.outVar.is("bigint")) { - throw new Error(`Unsupported data type: ${this.outVar.dtype}`); - } - - const symbols = Object.getOwnPropertySymbols(this.outVar); - const contextSymbol = symbols.find(s => s.toString().includes("zarrita.context")); - const fillValue = contextSymbol && !Number.isNaN((this.outVar as any)[contextSymbol]?.fill_value) - ? (this.outVar as any)[contextSymbol].fill_value - : NaN; - return { - shape: this.outVar.shape, - chunkShape: GetSize(this.outVar)[2], - fillValue, - dtype: this.outVar.dtype, - }; - } - - async fetchChunk({ - rank, - shape, - chunkShape, - x, - y, - z, - xDimIndex, - yDimIndex, - zDimIndex, - idx4D, - ndSlices - }: any): Promise { - if (!this.outVar) { - await this.init(); - if (!this.outVar){ - throw new Error(`Init failed`); - } - } - - const chunkSlice = new Array(rank).fill(0); - if (ndSlices && ndSlices.length === rank) { - for (let i = 0; i < rank; i++) { - if (i === xDimIndex) { - chunkSlice[i] = zarr.slice(x * chunkShape[i], Math.min((x + 1) * chunkShape[i], shape[i])); - } else if (i === yDimIndex) { - chunkSlice[i] = zarr.slice(y * chunkShape[i], Math.min((y + 1) * chunkShape[i], shape[i])); - } else if (i === zDimIndex) { - chunkSlice[i] = zarr.slice(z * chunkShape[i], Math.min((z + 1) * chunkShape[i], shape[i])); - } else { - const sel = ndSlices[i]; - if (Array.isArray(sel)) { - chunkSlice[i] = zarr.slice(sel[0], sel[1]); - } else { - chunkSlice[i] = sel; - } - } - } - } else { - chunkSlice[xDimIndex] = zarr.slice(x * chunkShape[xDimIndex], (x + 1) * chunkShape[xDimIndex]); - chunkSlice[yDimIndex] = zarr.slice(y * chunkShape[yDimIndex], (y + 1) * chunkShape[yDimIndex]); - if (zDimIndex >= 0) { - chunkSlice[zDimIndex] = zarr.slice(z * chunkShape[zDimIndex], (z + 1) * chunkShape[zDimIndex]); - } - if (rank >= 4) { - chunkSlice[0] = idx4D; - } - } - - const chunk = await fetchWithRetry( - () => zarr.get(this.outVar!, chunkSlice), - `variable ${this.variable}`, - useGlobalStore.getState().setStatus - ); - - if (!chunk || chunk.data instanceof BigInt64Array || chunk.data instanceof BigUint64Array) { - throw new Error("BigInt arrays not supported."); - } - - let outShape = chunk.shape; - if (!outShape || outShape.length === 0) { - outShape = chunkShape; // fallback if Zarrita doesn't collapse - } - - return { - data: chunk.data as Float32Array, - shape: outShape as number[], - stride: chunk.stride as number[], - }; - } } \ No newline at end of file From d0bd68ba589933c8436ba79e1a92212cc9d32863 Mon Sep 17 00:00:00 2001 From: Jeran Date: Thu, 1 Oct 2026 16:47:51 +0200 Subject: [PATCH 2/2] update to work with 2D variables --- src/components/zarr/GetArray.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/components/zarr/GetArray.ts b/src/components/zarr/GetArray.ts index e1f4e21bb..5c72faae5 100644 --- a/src/components/zarr/GetArray.ts +++ b/src/components/zarr/GetArray.ts @@ -97,7 +97,8 @@ export async function GetArray(varOveride?: string) { const zIndexInRaw = activeDims.indexOf(zDimIndex); const yIndexInRaw = activeDims.indexOf(yDimIndex); const xIndexInRaw = activeDims.indexOf(xDimIndex); - const sliceStart = [ndSlices[zIndexInRaw][0], ndSlices[yIndexInRaw][0], ndSlices[xIndexInRaw][0]] + const sliceStart = rank > 2 ? [ndSlices[zIndexInRaw][0], ndSlices[yIndexInRaw][0], ndSlices[xIndexInRaw][0]] : [ndSlices[yIndexInRaw][0], ndSlices[xIndexInRaw][0]] + setStatus("Downloading..."); setProgress(0); for (let z = zDim.start; z < zDim.end; z++) {