Skip to content

Commit db7236a

Browse files
TheJeranTheJeran
authored
Merge pull request #542 from EarthyScience/jp/2d-zarr-coarse
Update 2D getter and coarser
2 parents 7269f57 + 0eb7473 commit db7236a

File tree

3 files changed

+76
-49
lines changed

3 files changed

+76
-49
lines changed

‎src/components/ui/MainPanel/MetaDataInfo.tsx‎

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,8 +93,37 @@ const MetaDataInfo = ({ meta, metadata, setShowMeta, setOpenVariables, popoverSi
9393
const zLength = useMemo(() => meta.shape ? meta.shape[shapeLength-3] : 0, [meta])
9494
const yLength = useMemo(() => meta.shape ? meta.shape[shapeLength-2] : 0, [meta])
9595
const xLength = useMemo(() => meta.shape ? meta.shape[shapeLength-1] : 0, [meta])
96-
const dataShape = coarsen ? meta.shape.map((val: number, idx: number) => Math.floor(idx === 0 ? val/kernelDepth : val/kernelSize)): meta.shape
97-
const chunkShape = coarsen ? meta.chunks.map((val: number, idx: number) => Math.floor(idx === 0 ? val/kernelDepth : val/kernelSize)): meta.chunks
96+
const { dataShape, chunkShape } = useMemo(() => {
97+
if (coarsen) {
98+
const n = meta.shape.length;
99+
const Coarsen = (val: number, idx:number ) =>{
100+
if (n <= 2) {
101+
return Math.floor(val / kernelSize);
102+
}
103+
const thirdLast = n - 3;
104+
const secondLast = n - 2;
105+
const last = n - 1;
106+
if (idx === thirdLast) {
107+
return Math.floor(val / kernelDepth);
108+
}
109+
if (idx === secondLast || idx === last) {
110+
return Math.floor(val / kernelSize);
111+
}
112+
return val;
113+
}
114+
const dataShape = meta.shape.map((val: number, idx: number) => {
115+
return Coarsen(val, idx)
116+
});
117+
const chunkShape = meta.chunks.map((val: number, idx: number) => {
118+
return Coarsen(val, idx)
119+
});
120+
return { dataShape, chunkShape };
121+
}
122+
return {
123+
dataShape: meta.shape,
124+
chunkShape: meta.chunks
125+
};
126+
}, [meta, coarsen, kernelDepth, kernelSize]);
98127

99128
// ---- Booleans ---- //
100129
const is3D = useMemo(() => meta.shape ? meta.shape.length == 3 : false, [meta])
@@ -235,7 +264,11 @@ const MetaDataInfo = ({ meta, metadata, setShowMeta, setOpenVariables, popoverSi
235264
</div>
236265
<Hider show={coarsen} className="mt-2">
237266
<div className="grid grid-cols-2 gap-x-1">
238-
<div className="grid grid-cols-[auto_50px]">
267+
<div className="grid grid-cols-[auto_50px]"
268+
style={{
269+
visibility: dataShape.length >= 3 ? 'visible' : 'hidden'
270+
}}
271+
>
239272
<b>Temporal Coarsening</b>
240273
<Input type='number' min='0' step={1} value={displayDepth}
241274
onChange={e=>{

‎src/components/ui/MainPanel/PlayButton.tsx‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,8 @@ const PlayInterFace = ({visible, setKeepOpen}:{visible : boolean, setKeepOpen: R
126126
let timeSlice = timeArray?.slice(zSlice[0], zSlice[1] ?? undefined)
127127
const timeLength = timeArray?.length || 1
128128
let sliceDist = zSlice[1] ? zSlice[1] - zSlice[0] : timeLength - zSlice[0]
129-
if (coarsen) {
129+
130+
if (coarsen && timeSlice) {
130131
timeSlice = coarsenFlatArray(timeSlice, kernel.kernelDepth)
131132
sliceDist = Math.floor(sliceDist/kernel.kernelDepth)
132133
}

‎src/components/zarr/ZarrGetters.ts‎

Lines changed: 38 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { useZarrStore, useCacheStore, useGlobalStore, useErrorStore } from "@/Gl
22
import * as zarr from 'zarrita';
33
import { CompressArray, DecompressArray, ZarrError, RescaleArray, ToFloat16, copyChunkToArray } from "./ZarrLoaderLRU";
44
import { GetSize } from "./GetMetadata";
5-
import { Convolve } from "../computation/webGPU";
5+
import { Convolve, Convolve2D } from "../computation/webGPU";
66
import { coarsen3DArray, calculateStrides } from "@/utils/HelperFuncs";
77

88
export async function GetZarrDims(variable: string){
@@ -149,52 +149,45 @@ export async function GetZarrArray(){
149149
if (is4D){
150150
chunkShape = chunkShape.slice(1);
151151
}
152+
const is2D = outVar.shape.length === 2;
152153

153-
const notChunked = fullShape.every((dim, idx) => dim > chunkShape[idx]);
154-
155-
// ---- I THINK THIS CODE IS REDUNDANT. BUT NEED TO TEST ON AN UNCHUNCKED ZARR DATASET TO BE SURE ---- //
156-
157-
// //---- Strategy 1. Download whole array (No time chunks) ----//
158-
// if (notChunked){
159-
// setStatus("Downloading...")
160-
// const chunk = await fetchWithRetry(
161-
// () => is4D ? zarr.get(outVar, [idx4D, null, null, null]) : zarr.get(outVar),
162-
// `variable ${variable}`,
163-
// setStatus
164-
// );
165-
// if (!chunk) throw new Error('Unexpected: chunk was not assigned'); // This is redundant but satisfies TypeScript
166-
// if (chunk.data instanceof BigInt64Array || chunk.data instanceof BigUint64Array) {
167-
// throw new Error("BigInt arrays not supported.");
168-
// }
169-
// const shape = is4D ? outVar.shape.slice(1) : outVar.shape;
170-
// const strides = chunk.stride;
171-
// setStrides(strides) // Need strides for the point cloud
172-
173-
// let [typedArray, scalingFactor] = ToFloat16(chunk.data.map((v: number) => v === fillValue ? NaN : v) as Float32Array, null)
174-
// if (coarsen){
175-
// typedArray = await Convolve(typedArray, {shape, strides}, "Mean", {kernelSize, kernelDepth}) as Float16Array
176-
// const newShape = shape.map((dim, idx) => Math.ceil(dim / (idx === 0 ? kernelDepth : kernelSize)))
177-
// let newStrides = newShape.slice()
178-
// newStrides = newStrides.map((val, idx) => {
179-
// return newStrides.reduce((a, b, i) => a * (i < idx ? b : 1), 1)
180-
// })
181-
// console.log(newStrides)
182-
// }
183-
// const cacheChunk = {
184-
// data: compress ? CompressArray(typedArray, 7) : typedArray,
185-
// shape: chunk.shape,
186-
// stride: chunk.stride,
187-
// scaling: scalingFactor,
188-
// compressed: compress
189-
// }
190-
// cache.set(is4D ? `${initStore}_${idx4D}_${variable}` : `${initStore}_${variable}`, cacheChunk)
191-
// setStatus(null)
192-
// return { data: typedArray, shape, dtype: outVar.dtype, scalingFactor };
193-
// }
154+
// //---- Strategy 1. Download whole array (Chunking Logic doesn't work on 2D atm) ----//
155+
if (is2D){
156+
setStatus("Downloading...")
157+
const chunk = await fetchWithRetry(
158+
() => zarr.get(outVar),
159+
`variable ${variable}`,
160+
setStatus
161+
);
162+
if (!chunk) throw new Error('Unexpected: chunk was not assigned'); // This is redundant but satisfies TypeScript
163+
if (chunk.data instanceof BigInt64Array || chunk.data instanceof BigUint64Array) {
164+
throw new Error("BigInt arrays not supported.");
165+
}
166+
const shape = outVar.shape;
167+
const strides = chunk.stride;
168+
setStrides(strides) // Need strides for the point cloud
194169

195-
//---- Strategy 2. Download Chunks ----//
196-
const is2D = outVar.shape.length === 2;
197-
170+
let [typedArray, scalingFactor] = ToFloat16(chunk.data.map((v: number) => v === fillValue ? NaN : v) as Float32Array, null)
171+
if (coarsen){
172+
typedArray = await Convolve2D(typedArray, {shape, strides}, "Mean2D", kernelSize) as Float16Array
173+
const newShape = shape.map((dim) => Math.ceil(dim / kernelSize))
174+
let newStrides = newShape.slice()
175+
newStrides = newStrides.map((_val, idx) => {
176+
return newStrides.reduce((a, b, i) => a * (i < idx ? b : 1), 1)
177+
})
178+
}
179+
const cacheChunk = {
180+
data: compress ? CompressArray(typedArray, 7) : typedArray,
181+
shape: chunk.shape,
182+
stride: chunk.stride,
183+
scaling: scalingFactor,
184+
compressed: compress
185+
}
186+
cache.set(`${initStore}_${variable}`, cacheChunk)
187+
setStatus(null)
188+
return { data: typedArray, shape, dtype: outVar.dtype, scalingFactor };
189+
}
190+
198191
// Calculate Indices
199192
const zIndexOffset = is4D ? 1 : 0;
200193

0 commit comments

Comments
 (0)