Merge pull request #26 from EarthyScience/ZarrCache

Implemented caching and timeslicing from loaded array

Author: TheJeran

.eslintrc.cjs

@@ -14,5 +14,6 @@ module.exports = {
       'warn',
       { allowConstantExport: true },
     ],
+    '@typescript-eslint/no-unused-vars': ['off'],
   },
 }

src/components/CanvasGeometry.tsx

@@ -3,8 +3,8 @@ THREE.Cache.enabled = true;
 import { Canvas } from '@react-three/fiber';
 import { Center, OrbitControls, Environment, Html } from '@react-three/drei'
 // import * as zarr from 'zarrita'
-import { variables, GetArray } from '@/components/ZarrLoaderLRU'
-import { useEffect, useState } from 'react';
+import { variables, ZarrDataset } from '@/components/ZarrLoaderLRU'
+import { useEffect, useState, useMemo } from 'react';
 // import { useEffect, useState } from 'react';
 import { useControls } from 'leva'
 // import { Leva } from 'leva'
@@ -55,15 +55,21 @@ export function CanvasGeometry() {
   const [valueScales,setValueScales] = useState({maxVal:1,minVal:-1})
   const [showTimeSeries,setShowTimeSeries] = useState<boolean>(false)
   const [colormap,setColormap] = useState<THREE.DataTexture>(GetColorMapTexture())
+  const [timeSeries, setTimeSeries] = useState<number[]>([0]);
+
+
+  const ZarrDS = useMemo(()=>new ZarrDataset(storeURL),[])
 
   useEffect(()=>{
     setColormap(GetColorMapTexture(colormap,cmap));
   },[cmap, colormap])
 
+
+  //DATA LOADING
   useEffect(() => {
     if (variable != "Default") {
       //Need to add a check somewhere here to swap to 2D or 3D based on shape. Probably export two variables from GetArray
-      GetArray(storeURL, variable).then((result) => {
+      ZarrDS.GetArray(variable).then((result) => {
         // result now contains: { data: TypedArray, shape: number[], dtype: string }
         const [texture, shape, scaling] = ArrayToTexture({
           data: result.data,
@@ -97,6 +103,13 @@ export function CanvasGeometry() {
       }
   }, [variable])
 
+  //TIMESERIES
+  useEffect(()=>{
+    if(ZarrDS){
+      ZarrDS.GetTimeSeries(timeSeriesLocs).then((e)=> setTimeSeries(e))
+    }
+  },[timeSeriesLocs])
+
   return (
     <>
     <div className='canvas'>

src/components/PlotObjects.tsx

@@ -132,8 +132,6 @@ import { useState } from 'react'
 
 export const UVCube = ({shape,setTimeSeriesLocs} : {shape:THREE.Vector3, setTimeSeriesLocs:React.Dispatch<React.SetStateAction<TimeSeriesLocs>>} )=>{
   const [clickPoint, setClickPoint] = useState<THREE.Vector3 | null>(null);
-  console.log("logging the click point", clickPoint)
-
   function TimeSeriesLocs(event: THREE.Intersection){
     const point = event.point;
     const uv = event.uv!;

src/components/ZarrLoaderLRU.tsx

@@ -47,6 +47,7 @@ export async function GetArray(storePath: string, variable: string): Promise<Zar
     if (outVar.is("number") || outVar.is("bigint")) {
         const chunk = await zarr.get(outVar)
         const typedArray = new Float32Array(chunk.data);
+		console.log(chunk);
        // TypeScript will now infer the correct numeric type
         return {
             data: typedArray,
@@ -83,6 +84,8 @@ function parseUVCoords({normal,uv}:{normal:THREE.Vector3,uv:THREE.Vector2}){
 	}
 }
 
+
+
 export async function GetTimeSeries({TimeSeriesObject}:GetTimeSeries){
 	const {uv,normal,variable, storePath} = TimeSeriesObject
 	const d_store = zarr.tryWithConsolidated(
@@ -99,6 +102,90 @@ export async function GetTimeSeries({TimeSeriesObject}:GetTimeSeries){
 	return arr
 }
 
+interface TimeSeriesInfo{
+	uv:THREE.Vector2,
+	normal:THREE.Vector3
+}
+
+export class ZarrDataset{
+	private storePath: string;
+	private variable: string;
+	private cache: { [key: string]: any };
+
+	constructor(storePath: string){
+		this.storePath = storePath;
+		this.variable = "Default";
+		this.cache = {};
+	}
+
+	async GetArray(variable: string){
+		//This checks if variable is stored in cache
+		this.variable = variable;
+		let outVar = null;
+		if (this.cache.hasOwnProperty(variable)){
+			console.log("Using Cache")
+			return this.cache[variable]
+		}
+		const d_store = zarr.tryWithConsolidated(
+			new zarr.FetchStore(this.storePath)
+		);
+	
+		const group = await d_store.then(store => zarr.open(store, {kind: 'group'}))
+		
+		outVar = await zarr.open(group.resolve(variable), {kind:"array"})
+		// Type check using zarr.Array.is
+		if (outVar.is("number") || outVar.is("bigint")) {
+			const chunk = await zarr.get(outVar)
+			const typedArray = new Float32Array(chunk.data);
+			this.cache[variable] = chunk;
+			// TypeScript will now infer the correct numeric type
+			return {
+				data: typedArray,
+				shape: chunk.shape,
+				dtype: outVar.dtype
+			}
+		} else {
+			throw new Error(`Unsupported data type: Only numeric arrays are supported. Got: ${outVar.dtype}`)
+		}
+	}
+
+	async GetTimeSeries(TimeSeriesInfo:TimeSeriesInfo){
+		const {uv,normal} = TimeSeriesInfo
+		if (!this.cache[this.variable]){
+			return [0]
+		}
+		const {data,shape,stride} = this.cache[this.variable]
+		//This is a complicated logic check but it works bb
+		const sliceSize = parseUVCoords({normal,uv})
+
+		const slice = sliceSize.map((value, index) =>
+			value === null || shape[index] === null ? null : Math.round(value * shape[index]));
+
+		const mapDim = slice.indexOf(null);
+		const dimStride = stride[mapDim];
+		const pz = slice[0] == null ? 0 : stride[0]*slice[0]
+		const py = slice[1] == null ? 0 : stride[1]*slice[1]
+		const px = slice[2] == null ? 0 : stride[2]*slice[2]
+
+		// console.log(`
+		// 	Slice:${slice}
+		// 	mapDim:${mapDim}
+		// 	dimStride:${dimStride}
+		// 	pz:${pz}
+		// 	py:${py}
+		// 	px:${px}
+		// 	`)
+		const ts = [];
+
+		for (let i = 0; i < shape[mapDim] ; i++){
+			const idx = i*dimStride+pz+py+px
+			ts.push(data[idx])
+		}
+		return ts;
+	}
+
+}
+