Merge pull request #52 from EarthyScience/Computation

HUGE Update - Added GPU Compute pipeline for single variable data

Author: TheJeran

src/components/CanvasGeometry.tsx

@@ -7,8 +7,15 @@ import { useEffect, useState, useMemo } from 'react';
 import { useControls } from 'leva'
 import { PointCloud, UVCube, PlotArea, DataCube } from './PlotObjects';
 import { GetColorMapTexture, ArrayToTexture, DefaultCube, colormaps } from './Textures';
-import { Metadata } from './UI';
+import { Metadata, AnalysisWindow } from './UI';
 import { plotContext, DimCoords } from './Contexts';
+import ComputeModule from '@/components/Computation/ComputeModule'
+
+interface Array{
+  data:number[],
+  shape:number[],
+  stride:number[]
+}
 
 const storeURL = "https://s3.bgc-jena.mpg.de:9000/esdl-esdc-v3.0.2/esdc-16d-2.5deg-46x72x1440-3.0.2.zarr"
 
@@ -53,6 +60,7 @@ export function CanvasGeometry() {
   const [timeSeries, setTimeSeries] = useState<number[]>([0]);
   const [showLoading, setShowLoading] = useState<boolean>(false);
   const [metadata,setMetadata] = useState<object[] | null>(null)
+  const [dataArray, setDataArray] = useState<Array | null>(null)
 
   //Timeseries Plotting Information
   const [dimArrays,setDimArrays] = useState([[0],[0],[0]])
@@ -62,6 +70,12 @@ export function CanvasGeometry() {
   const [plotDim,setPlotDim] = useState<number>(0)
   const ZarrDS = useMemo(()=>new ZarrDataset(storeURL),[])
 
+  //Analysis variables
+  const [reduceAxis, setReduceAxis] = useState<number>(0);
+  const [reduceOperation, setReduceOperation] = useState<string>("Mean")
+  const [executeReduction,setExecuteReduction] = useState<boolean>(false)
+  const [showAnalysis, setShowAnalysis] = useState<boolean>(false)
+
   useEffect(()=>{
     setColormap(GetColorMapTexture(colormap,cmap,1,"#000000",0,flipCmap));
   },[cmap, colormap,flipCmap])
@@ -77,6 +91,7 @@ export function CanvasGeometry() {
           data: result.data,
           shape: result.shape
         })
+        setDataArray(result)
         console.log(`logging the shape since we will want to use it in the future for 2D vs 3D actions ${shape}`)
         if (texture instanceof THREE.DataTexture || texture instanceof THREE.Data3DTexture) {
           setTexture(texture)
@@ -157,15 +172,28 @@ export function CanvasGeometry() {
     scaling:{...valueScales,colormap}
   } 
 
+  const analysisSetters = {
+    setAxis:setReduceAxis,
+    setOperation:setReduceOperation,
+    setExecute:setExecuteReduction
+  }
+
+  const analysisVars = {
+    axis:reduceAxis,
+    operation:reduceOperation,
+    execute:executeReduction
+  }
   return (
     <>
+    
     <Loading showLoading={showLoading} />
     <div className='canvas'>
       <Canvas shadows camera={{ position: [-4.5, 3, 4.5], fov: 50 }}
       frameloop="demand"
       >
+        
         <Center top position={[-1, 0, 1]}/>
-
+        {dataArray && showAnalysis && <ComputeModule array={dataArray} cmap={colormap} shape={shape.toArray()} stateVars={analysisVars}/>}
         {/* Volume Plots */}
         {plotter == "volume" && <>
           <DataCube volTexture={texture} shape={shape} colormap={colormap}/>
@@ -180,14 +208,25 @@ export function CanvasGeometry() {
 
       </Canvas>
     </div>
-
+    {showAnalysis && <AnalysisWindow setters={analysisSetters}/>}
     {metadata && <Metadata data={metadata} /> }
 
     <plotContext.Provider value={plotObj} >
+      
       {variable !== "Default" && <PlotArea />}
     </plotContext.Provider>
 
     {/* <Leva theme={lightTheme} /> */}
+    <button
+      style={{
+        position:'fixed',
+        left:'5%',
+        top:'50%'
+      }}
+      onClick={()=>setShowAnalysis(x=>!x)}
+    >
+      {showAnalysis ? 'Hide' : 'Show' } Analysis Stuff
+    </button>
     </>
   )
 }

src/components/Computation/ComputeModule.tsx

@@ -0,0 +1,73 @@
+import { useEffect, useState } from 'react'
+import { OneArrayCompute } from './ComputeShaders'
+import * as THREE from 'three'
+import fragmentShader from '@/components/Computation/shaders/frag.glsl'
+import vertexShader from '@/components/Computation/shaders/vert.glsl'
+
+interface Array{
+    data:number[],
+    shape:number[],
+    stride:number[]
+}
+interface StateVars{
+  axis:number,
+  operation:string,
+  execute:boolean
+}
+
+
+const ComputeModule = ({array,cmap,shape,stateVars}:{array: Array,cmap:THREE.DataTexture,shape:number[],stateVars:StateVars}) => {
+    const {axis,operation,execute} = stateVars;
+    const [planeShape,setPlaneShape] = useState<number[]>(shape.filter((_val,idx)=> idx !== axis))
+
+
+    const GPUCompute = new OneArrayCompute(array)
+    const [texture,setTexture] = useState<THREE.Texture | undefined>(undefined)
+
+    const shaderMaterial = new THREE.ShaderMaterial({
+        glslVersion: THREE.GLSL3,
+        uniforms:{
+          data : {value: texture},
+          cmap : { value : cmap},
+        },
+        vertexShader,
+        fragmentShader,
+        side: THREE.DoubleSide
+    })
+
+    useEffect(()=>{
+      if (array){
+        let newText;
+        switch(operation){
+          case "Max":
+            newText =  GPUCompute.Max(axis);
+            break
+          case "Min":
+            newText =  GPUCompute.Min(axis);
+            break
+          case "Mean":
+            newText =  GPUCompute.Mean(axis);
+            break
+          case "StDev":
+            newText =  GPUCompute.StDev(axis);
+            break;
+          default:
+            newText = texture;
+            break
+        }
+          setTexture(newText)
+          setPlaneShape(shape.filter((_val,idx)=> idx !== axis))
+      }
+    },[execute])
+
+  return (
+    <>
+    <mesh material={shaderMaterial}>
+      <planeGeometry args={[planeShape[1],planeShape[0]]} />
+    </mesh>
+
+    </>
+  )
+}
+
+export default ComputeModule

src/components/Computation/ComputeShaders.ts

@@ -0,0 +1,95 @@
+import * as THREE from 'three'
+import { GPUComputationRenderer } from 'three/examples/jsm/Addons.js'
+import { useThree } from '@react-three/fiber'
+import { MeanFrag, MaxFrag, MinFrag, StDevFrag } from './shaders'
+
+
+interface Array{
+    data:number[],
+    shape:number[],
+    stride:number[]
+}
+
+export class OneArrayCompute{
+    private data: number[];
+    private shape: number[];
+    private stride: number[];
+    private renderer: THREE.WebGLRenderer;
+    private GPUCompute: GPUComputationRenderer;
+    private texture: THREE.Data3DTexture | null;
+    private targetAxis: number;
+    private renderTarget: THREE.WebGLRenderTarget<THREE.Texture>;
+    private initTexture: THREE.DataTexture
+
+    constructor(array: Array){
+        this.data = array.data;
+        this.shape = array.shape;
+        this.stride = array.stride;
+        this.renderer = useThree(state => state.gl) 
+        this.GPUCompute = new GPUComputationRenderer(10,10,this.renderer)
+        this.initTexture = this.GPUCompute.createTexture()
+        this.targetAxis = 5;
+        this.renderTarget = this.GPUCompute.createRenderTarget(10,10,THREE.ClampToEdgeWrapping,THREE.ClampToEdgeWrapping,1006,1006);
+        const size = array.shape[0]*array.shape[1]*array.shape[2]
+        const newArray = new Float32Array(size)
+        for (let i = 0; i<size; i++){
+            newArray[i] = this.data[i]
+        }
+        this.texture = new THREE.Data3DTexture(newArray,this.shape[2],this.shape[1],this.shape[0])
+    }
+    
+    private initAxis(axis:number){
+        const resolution = this.shape.filter((_val,idx)=> idx !== axis)
+        this.GPUCompute = new GPUComputationRenderer(resolution[0],resolution[1],this.renderer)
+        this.targetAxis = axis;
+        this.renderTarget = this.GPUCompute.createRenderTarget(resolution[0],resolution[1],THREE.ClampToEdgeWrapping,THREE.ClampToEdgeWrapping,1006,1006)
+    }
+
+    Mean(axis:number){
+        if (axis !== this.targetAxis){
+            this.initAxis(axis)
+        }
+        const reducer = this.GPUCompute.addVariable("reduction", MeanFrag, this.initTexture);
+        reducer.material.uniforms[`dataArray]`] = { value: this.texture };
+        reducer.material.uniforms['axisSize'] = { value: this.shape[this.targetAxis]}
+        reducer.material.uniforms['axis'] = { value: this.targetAxis}
+        this.GPUCompute.doRenderTarget(reducer.material,this.renderTarget)
+        return this.renderTarget.texture
+    }
+
+    Max(axis:number){
+        if (axis !== this.targetAxis){
+            this.initAxis(axis)
+        }
+        const reducer = this.GPUCompute.addVariable("reduction", MaxFrag, this.initTexture);
+        reducer.material.uniforms[`dataArray]`] = { value: this.texture };
+        reducer.material.uniforms['axisSize'] = { value: this.shape[this.targetAxis]}
+        reducer.material.uniforms['axis'] = { value: this.targetAxis}
+        this.GPUCompute.doRenderTarget(reducer.material,this.renderTarget)
+        return this.renderTarget.texture
+    }
+
+    Min(axis:number){
+        if (axis !== this.targetAxis){
+            this.initAxis(axis)
+        }
+        const reducer = this.GPUCompute.addVariable("reduction", MinFrag, this.initTexture);
+        reducer.material.uniforms[`dataArray]`] = { value: this.texture };
+        reducer.material.uniforms['axisSize'] = { value: this.shape[this.targetAxis]}
+        reducer.material.uniforms['axis'] = { value: this.targetAxis}
+        this.GPUCompute.doRenderTarget(reducer.material,this.renderTarget)
+        return this.renderTarget.texture
+    }
+    
+    StDev(axis:number){
+        if (axis !== this.targetAxis){
+            this.initAxis(axis)
+        }
+        const reducer = this.GPUCompute.addVariable("reduction", StDevFrag, this.initTexture);
+        reducer.material.uniforms[`dataArray]`] = { value: this.texture };
+        reducer.material.uniforms['axisSize'] = { value: this.shape[this.targetAxis]}
+        reducer.material.uniforms['axis'] = { value: this.targetAxis}
+        this.GPUCompute.doRenderTarget(reducer.material,this.renderTarget)
+        return this.renderTarget.texture
+    }
+}

src/components/Computation/shaders/Max.glsl

@@ -0,0 +1,39 @@
+uniform sampler3D dataArray;
+uniform int axis;
+uniform int axisSize;
+
+void main() {
+    vec2 uv = gl_FragCoord.xy / resolution.xy;
+
+    vec3 sampleCoord;
+    
+    // Initialize sampling coordinates based on reduction axis
+    if (axis == 0) {
+        // Reduce along depth (D0), uv maps to (s,t) = (D2,D1)
+        sampleCoord = vec3(uv.x, uv.y, 0.0);
+    } else if (axis == 1) {
+        // Reduce along height (D1), uv maps to (s,r) = (D2,D0)
+        sampleCoord = vec3(uv.x, 0.0, uv.y);
+    } else if (axis == 2) {
+        // Reduce along width (D2), uv maps to (t,r) = (D1,D0)
+        sampleCoord = vec3(0.0, uv.x, uv.y);
+    }
+
+    float maxVal = 0.0;
+    for (int m = 0; m < axisSize; m++) { 
+        float coord = (float(m) + 0.5) / float(axisSize); //0.5 for center of pixel
+        
+        // Vary the appropriate coordinate
+        if (axis == 0) {
+            sampleCoord.z = coord; // Vary r (depth)
+        } else if (axis == 1) {
+            sampleCoord.y = coord; // Vary t (height)
+        } else if (axis == 2) {
+            sampleCoord.x = coord; // Vary s (width)
+        }
+        float sampled = texture(dataArray, sampleCoord).r; // Assuming data in red channel
+        maxVal = sampled > maxVal ? sampled : maxVal; 
+    }
+    
+    gl_FragColor = vec4(maxVal, 0.0, 0.0, 1.0); // Output mean in red channel
+}

src/components/Computation/shaders/Mean.glsl

@@ -0,0 +1,40 @@
+uniform sampler3D dataArray;
+uniform int axis;
+uniform int axisSize;
+
+void main() {
+    vec2 uv = gl_FragCoord.xy / resolution.xy;
+
+    vec3 sampleCoord;
+    
+    // Initialize sampling coordinates based on reduction axis
+    if (axis == 0) {
+        // Reduce along depth (D0), uv maps to (s,t) = (D2,D1)
+        sampleCoord = vec3(uv.x, uv.y, 0.0);
+    } else if (axis == 1) {
+        // Reduce along height (D1), uv maps to (s,r) = (D2,D0)
+        sampleCoord = vec3(uv.x, 0.0, uv.y);
+    } else if (axis == 2) {
+        // Reduce along width (D2), uv maps to (t,r) = (D1,D0)
+        sampleCoord = vec3(0.0, uv.x, uv.y);
+    }
+
+    float sum = 0.0;
+    for (int m = 0; m < axisSize; m++) { 
+        float coord = (float(m) + 0.5) / float(axisSize); //0.5 for center of pixel
+        
+        // Vary the appropriate coordinate
+        if (axis == 0) {
+            sampleCoord.z = coord; // Vary r (depth)
+        } else if (axis == 1) {
+            sampleCoord.y = coord; // Vary t (height)
+        } else if (axis == 2) {
+            sampleCoord.x = coord; // Vary s (width)
+        }
+        
+        sum += texture(dataArray, sampleCoord).r; // Assuming data in red channel
+    }
+    
+    float mean = sum / float(axisSize);
+    gl_FragColor = vec4(mean, 0.0, 0.0, 1.0); // Output mean in red channel
+}

src/components/Computation/shaders/Min.glsl

@@ -0,0 +1,39 @@
+uniform sampler3D dataArray;
+uniform int axis;
+uniform int axisSize;
+
+void main() {
+    vec2 uv = gl_FragCoord.xy / resolution.xy;
+
+    vec3 sampleCoord;
+    
+    // Initialize sampling coordinates based on reduction axis
+    if (axis == 0) {
+        // Reduce along depth (D0), uv maps to (s,t) = (D2,D1)
+        sampleCoord = vec3(uv.x, uv.y, 0.0);
+    } else if (axis == 1) {
+        // Reduce along height (D1), uv maps to (s,r) = (D2,D0)
+        sampleCoord = vec3(uv.x, 0.0, uv.y);
+    } else if (axis == 2) {
+        // Reduce along width (D2), uv maps to (t,r) = (D1,D0)
+        sampleCoord = vec3(0.0, uv.x, uv.y);
+    }
+
+    float minVal = 1.0e15; //BigNum
+    for (int m = 0; m < axisSize; m++) { 
+        float coord = (float(m) + 0.5) / float(axisSize); //0.5 for center of pixel
+        
+        // Vary the appropriate coordinate
+        if (axis == 0) {
+            sampleCoord.z = coord; // Vary r (depth)
+        } else if (axis == 1) {
+            sampleCoord.y = coord; // Vary t (height)
+        } else if (axis == 2) {
+            sampleCoord.x = coord; // Vary s (width)
+        }
+        float sampled = texture(dataArray, sampleCoord).r; // Assuming data in red channel
+        minVal = sampled < minVal ? sampled : minVal; 
+    }
+    
+    gl_FragColor = vec4(minVal, 0.0, 0.0, 1.0); // Output mean in red channel
+}