Skip to content

Subselection update #780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TheJeran merged 2 commits into from
Oct 1, 2026
Merged

Subselection update #780

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eae5c89
fix serious bug where time-slicing was always at first chunk spot
TheJeran Oct 1, 2026
d0bd68b
update to work with 2D variables
TheJeran Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/GlobalStates/ZarrStore.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
86 changes: 40 additions & 46 deletions src/components/zarr/GetArray.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -108,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 = 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++) {
Expand All @@ -119,13 +109,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);
Expand Down Expand Up @@ -165,7 +159,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);
}
Expand Down
115 changes: 0 additions & 115 deletions src/components/zarr/dataFetchers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import { GetSize } from './utils';
import * as zarr from 'zarrita';
import { calculateStrides } from '@/utils/HelperFuncs';


interface FetchOutput{
data: Float32Array
shape: number[]
Expand Down Expand Up @@ -174,118 +173,4 @@ export function NCFetcher() {
return { data, shape: collapsedShape, stride: calculateStrides(collapsedShape) };
},
};
}


//---- Class possiblity

export class ZarrFetcher {
private outVar: zarr.Array<zarr.DataType, any> | 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<any> {
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<FetchOutput> {
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[],
};
}
}
Loading