1+ import { useGlobalStore } from "@/GlobalStates/GlobalStore" ;
2+ import { useZarrStore } from "@/GlobalStates/ZarrStore" ;
3+ import { useCacheStore } from "@/GlobalStates/CacheStore" ;
4+ import { useErrorStore } from "@/GlobalStates/ErrorStore" ;
5+ import { calculateStrides } from "@/utils/HelperFuncs" ;
6+ import { ToFloat16 , CompressArray , DecompressArray , copyChunkToArray , RescaleArray , copyChunkToArray2D } from "./utils" ;
7+ import { NCFetcher , zarrFetcher } from "./dataFetchers" ;
8+ import { Convolve } from "../computation/webGPU" ;
9+ import { coarsen3DArray } from "@/utils/HelperFuncs" ;
10+
11+ export async function GetArray ( varOveride ?: string ) {
12+ const { idx4D, initStore, variable, setProgress, setStrides, setStatus } = useGlobalStore . getState ( ) ;
13+ const { compress, xSlice, ySlice, zSlice, coarsen, kernelSize, kernelDepth, fetchNC, setCurrentChunks, setArraySize } = useZarrStore . getState ( ) ;
14+ const { cache } = useCacheStore . getState ( ) ;
15+ const fetcher = fetchNC ? NCFetcher ( ) : zarrFetcher ( )
16+ const targetVariable = varOveride ?? variable ;
17+ const meta = await fetcher . getMetadata ( targetVariable ) ;
18+ const { shape, chunkShape, fillValue, dtype } = meta ;
19+ const rank = shape . length ;
20+ const hasZ = rank >= 3 ;
21+ const xDimIndex = rank - 1 , yDimIndex = rank - 2 , zDimIndex = rank - 3 ;
22+
23+ const calcDim = ( slice : [ number , number | null ] , dimIdx : number ) => { // This function provides information for extraction from each dimension of datarray
24+ if ( dimIdx < 0 ) return { start : 0 , end : 1 , size : 0 , chunkDim : 1 } ;
25+ const dimSize = shape [ dimIdx ] ;
26+ const chunkDim = chunkShape [ dimIdx ] ;
27+ const start = Math . floor ( slice [ 0 ] / chunkDim ) ;
28+ const sliceEnd = slice [ 1 ] ?? dimSize ;
29+ return { start, end : Math . ceil ( sliceEnd / chunkDim ) , size : sliceEnd - slice [ 0 ] , chunkDim } ;
30+ } ;
31+
32+ const xDim = calcDim ( xSlice , xDimIndex ) ;
33+ const yDim = calcDim ( ySlice , yDimIndex ) ;
34+ const zDim = calcDim ( zSlice , zDimIndex ) ;
35+
36+ let outputShape = hasZ ? [ zDim . size , yDim . size , xDim . size ] : [ yDim . size , xDim . size ] ;
37+ if ( coarsen ) {
38+ outputShape = outputShape . map ( ( dim , idx ) => Math . floor ( dim / ( hasZ && idx === 0 ? kernelDepth : kernelSize ) ) ) ;
39+ }
40+
41+ const totalElements = outputShape . reduce ( ( a , b ) => a * b , 1 ) ;
42+ if ( totalElements > 1e9 ) { useErrorStore . getState ( ) . setError ( "largeArray" ) ; throw new Error ( "Cannot allocate array." ) ; }
43+
44+ const destStride = calculateStrides ( outputShape ) ;
45+ setStrides ( destStride ) ;
46+ setArraySize ( totalElements ) ;
47+ setCurrentChunks ( { x : [ xDim . start , xDim . end ] , y : [ yDim . start , yDim . end ] , z : [ zDim . start , zDim . end ] } ) ; // These are used in GetCurrentArray() function
48+
49+ const typedArray = new Float16Array ( totalElements ) ;
50+
51+ let scalingFactor : number | null = null ;
52+ const totalChunks = ( zDim . end - zDim . start ) * ( yDim . end - yDim . start ) * ( xDim . end - xDim . start ) ;
53+ let iter = 1 ;
54+ const rescaleIDs : string [ ] = [ ] ;
55+
56+ setStatus ( "Downloading..." ) ;
57+ setProgress ( 0 ) ;
58+
59+ for ( let z = zDim . start ; z < zDim . end ; z ++ ) {
60+ for ( let y = yDim . start ; y < yDim . end ; y ++ ) {
61+ for ( let x = xDim . start ; x < xDim . end ; x ++ ) {
62+ const chunkID = `z${ z } ${ y } ${ x } ;
63+ const cacheBase = rank > 3 ? `${ initStore } ${ targetVariable } ${ idx4D } : `${ initStore } ${ targetVariable } ;
64+ const cacheName = `${ cacheBase } ${ chunkID } ;
65+ const cachedChunk = cache . get ( cacheName ) ;
66+ const isCacheValid = cachedChunk &&
67+ cachedChunk . kernel . kernelSize === ( coarsen ? kernelSize : undefined ) &&
68+ cachedChunk . kernel . kernelDepth === ( coarsen ? kernelDepth : undefined ) ;
69+
70+ if ( isCacheValid ) {
71+ const chunkData = cachedChunk . compressed ? DecompressArray ( cachedChunk . data ) : cachedChunk . data . slice ( ) ;
72+ copyChunkToArray (
73+ chunkData ,
74+ cachedChunk . shape ,
75+ cachedChunk . stride ,
76+ typedArray ,
77+ outputShape ,
78+ destStride as any , [ z , y , x ] ,
79+ [ zDim . start , yDim . start , xDim . start ]
80+ ) ;
81+ } else {
82+ const raw = await fetcher . fetchChunk ( { variable :targetVariable , rank, shape, chunkShape, x, y, z, xDimIndex, yDimIndex, zDimIndex, idx4D } ) ;
83+
84+ const rawData = Number . isFinite ( fillValue ) ? raw . data . map ( ( v : number ) => v === fillValue ? NaN : v ) : raw . data ; // Don't map if no fillvalue
85+
86+ let [ chunkF16 , newScalingFactor ] = ToFloat16 ( rawData , scalingFactor ) ;
87+ let thisShape = raw . shape ;
88+ let chunkStride = raw . stride ;
89+
90+ if ( coarsen ) {
91+ chunkF16 = await Convolve ( chunkF16 , { shape : chunkShape , strides : chunkStride } , "Mean3D" , { kernelSize, kernelDepth } ) as Float16Array ;
92+ thisShape = thisShape . map ( ( dim , idx ) => Math . floor ( dim / ( idx === 0 ? kernelDepth : kernelSize ) ) ) ;
93+ chunkF16 = coarsen3DArray ( chunkF16 , chunkShape , chunkStride as any , kernelSize , kernelDepth , thisShape . reduce ( ( a , b ) => a * b , 1 ) ) ;
94+ chunkStride = calculateStrides ( thisShape ) ;
95+ }
96+
97+ if ( newScalingFactor != null && newScalingFactor !== scalingFactor ) {
98+ const delta = scalingFactor ? newScalingFactor - scalingFactor : newScalingFactor ;
99+ RescaleArray ( typedArray , delta ) ;
100+ scalingFactor = newScalingFactor ;
101+ for ( const id of rescaleIDs ) {
102+ const tempChunk = cache . get ( `${ cacheBase } ${ id } ) ;
103+ tempChunk . scaling = scalingFactor ;
104+ RescaleArray ( tempChunk . data , delta ) ;
105+ cache . set ( `${ cacheBase } ${ id } , tempChunk ) ;
106+ }
107+ }
108+
109+ if ( hasZ ) {
110+ copyChunkToArray ( chunkF16 , thisShape . slice ( - 3 ) , chunkStride . slice ( - 3 ) as any , typedArray , outputShape , destStride as any , [ z , y , x ] , [ zDim . start , yDim . start , xDim . start ] ) ;
111+ } else {
112+ copyChunkToArray2D ( chunkF16 , thisShape , chunkStride as any , typedArray , outputShape , destStride as any , [ y , x ] , [ yDim . start , xDim . start ] ) ;
113+ }
114+
115+ cache . set ( cacheName , {
116+ data : compress ? CompressArray ( chunkF16 , 7 ) : chunkF16 ,
117+ shape : chunkShape , stride : chunkStride ,
118+ scaling : scalingFactor , compressed : compress , coarsened : coarsen ,
119+ kernel : { kernelDepth : coarsen ? kernelDepth : undefined , kernelSize : coarsen ? kernelSize : undefined }
120+ } ) ;
121+ rescaleIDs . push ( chunkID ) ;
122+ }
123+ setProgress ( Math . round ( iter ++ / totalChunks * 100 ) ) ;
124+ }
125+ }
126+ }
127+
128+ setProgress ( 0 ) ;
129+ return { data : typedArray , shape : outputShape , dtype, scalingFactor } ;
130+ }
0 commit comments