From 230e9ee5c9569b70a31a654dbc0386f8d2450af0 Mon Sep 17 00:00:00 2001 From: Jeran Date: Thu, 30 Apr 2026 16:13:30 +0200 Subject: [PATCH 1/2] Added IsNan to shaders w/out boilerplate --- src/components/computation/WGSLShaders.ts | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/components/computation/WGSLShaders.ts b/src/components/computation/WGSLShaders.ts index 42339bf6..aab2c557 100644 --- a/src/components/computation/WGSLShaders.ts +++ b/src/components/computation/WGSLShaders.ts @@ -492,6 +492,7 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(2) var outputData: array<${precision}>; @group(0) @binding(3) var params: Params; + ${isNaNFunc} @compute @workgroup_size(16, 16, 1) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; @@ -615,6 +616,7 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(2) var outputData: array<${precision}>; @group(0) @binding(3) var params: Params; + ${isNaNFunc} @compute @workgroup_size(16, 16, 1) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; @@ -697,6 +699,8 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(2) var outputData: array<${precision}>; @group(0) @binding(3) var params: Params; + ${isNaNFunc} + @compute @workgroup_size(16, 16, 1) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; @@ -950,6 +954,7 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(2) var outputData: array<${precision}>; @group(0) @binding(3) var params: Params; + ${isNaNFunc} @compute @workgroup_size(4, 4, 4) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; @@ -1055,7 +1060,7 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(1) var secondData: array<${precision}>; @group(0) @binding(2) var outputData: array<${precision}>; @group(0) @binding(3) var params: Params; - + ${isNaNFunc} @compute @workgroup_size(4, 4, 4) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; @@ -1169,7 +1174,7 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(1) var secondData: array<${precision}>; @group(0) @binding(2) var outputData: array<${precision}>; @group(0) @binding(3) var params: Params; - + ${isNaNFunc} @compute @workgroup_size(4, 4, 4) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; @@ -1411,7 +1416,7 @@ export const createShaders = (precision: Precision) => { @group(0) @binding(0) var inputData: array<${precision}>; @group(0) @binding(1) var outputData: array; @group(0) @binding(2) var params: Params; - + ${isNaNFunc} @compute @workgroup_size(4, 4, 4) fn main(@builtin(global_invocation_id) global_id: vec3) { let zStride = params.zStride; From 10ea1af14e71fed196a3c141ea93b2ff9999b97d Mon Sep 17 00:00:00 2001 From: Jeran Date: Mon, 22 Jun 2026 09:14:03 +0200 Subject: [PATCH 2/2] Gemini fixes --- src/components/computation/WGSLShaders.ts | 28 ++++++++++++++++++----- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/src/components/computation/WGSLShaders.ts b/src/components/computation/WGSLShaders.ts index aab2c557..c890fcaf 100644 --- a/src/components/computation/WGSLShaders.ts +++ b/src/components/computation/WGSLShaders.ts @@ -513,6 +513,8 @@ export const createShaders = (precision: Precision) => { var ySum: f32 = 0; var xSum: f32 = 0.0; + var nanCount: u32 = 0; + // Iterate along the dimension we're averaging if (reduceDim == 0u) { let cCoord = outX * xStride + outY * yStride; @@ -521,6 +523,7 @@ export const createShaders = (precision: Precision) => { let xi: f32 = f32(firstData[inputIndex]); let yi: f32 = f32(secondData[inputIndex]); if (isNaN(xi) || isNaN(yi)){ + nanCount++; continue; } xSum += f32(firstData[inputIndex]); @@ -533,6 +536,7 @@ export const createShaders = (precision: Precision) => { let xi: f32 = f32(firstData[inputIndex]); let yi: f32 = f32(secondData[inputIndex]); if (isNaN(xi) || isNaN(yi)){ + nanCount++; continue; } xSum += f32(firstData[inputIndex]); @@ -545,6 +549,7 @@ export const createShaders = (precision: Precision) => { let xi: f32 = f32(firstData[inputIndex]); let yi: f32 = f32(secondData[inputIndex]); if (isNaN(xi) || isNaN(yi)){ + nanCount++; continue; } xSum += f32(firstData[inputIndex]); @@ -552,11 +557,12 @@ export const createShaders = (precision: Precision) => { } } - let xMean: f32 = xSum / f32(dimLength); - let yMean: f32 = ySum / f32(dimLength); + let xMean: f32 = xSum / f32(dimLength - nanCount); + let yMean: f32 = ySum / f32(dimLength - nanCount); var numSum: f32 = 0; var denomSum: f32 = 0; + if (reduceDim == 0u) { // Average along Z let cCoord = outX * xStride + outY * yStride; for (var z: u32 = 0u; z < dimLength; z++) { @@ -652,20 +658,25 @@ export const createShaders = (precision: Precision) => { iterStride = xStride; } + var nanCount: u32 = 0; + // Single pass: calculate sums, means, and covariance for (var i: u32 = 0u; i < dimLength; i++) { let inputIndex = baseCoord + (i * iterStride); let xi: f32 = f32(firstData[inputIndex]); let yi: f32 = f32(secondData[inputIndex]); if (isNaN(xi) || isNaN(yi)){ + nanCount++; continue; } xSum += xi; ySum += yi; } - let xMean: f32 = xSum / f32(dimLength); - let yMean: f32 = ySum / f32(dimLength); + var N: f32 = f32(dimLength - nanCount); + + let xMean: f32 = xSum / N; + let yMean: f32 = ySum / N; // Second pass for covariance calculation for (var i: u32 = 0u; i < dimLength; i++) { @@ -679,7 +690,7 @@ export const createShaders = (precision: Precision) => { } let outputIndex = outY * xSize + outX; - outputData[outputIndex] = ${precision}(numSum / (f32(dimLength) - 1)); + outputData[outputIndex] = ${precision}(numSum / (N - 1)); } `, @@ -723,6 +734,8 @@ export const createShaders = (precision: Precision) => { var ySum: f32 = 0.0; var yySum: f32 = 0.0; var xySum: f32 = 0.0; + + var nanCount: u32 = 0; // Iterate along the dimension we're averaging if (reduceDim == 0u) { // Average along Z let cCoord = outX * xStride + outY * yStride; @@ -731,6 +744,7 @@ export const createShaders = (precision: Precision) => { let xI = f32(firstData[inputIndex]); let yI = f32(secondData[inputIndex]); if (isNaN(xI) || isNaN(yI)){ + nanCount++; continue; } xSum += xI; @@ -746,6 +760,7 @@ export const createShaders = (precision: Precision) => { let xI = f32(firstData[inputIndex]); let yI = f32(secondData[inputIndex]); if (isNaN(xI) || isNaN(yI)){ + nanCount++; continue; } xSum += xI; @@ -761,6 +776,7 @@ export const createShaders = (precision: Precision) => { let xI = f32(firstData[inputIndex]); let yI = f32(secondData[inputIndex]); if (isNaN(xI) || isNaN(yI)){ + nanCount++; continue; } xSum += xI; @@ -771,7 +787,7 @@ export const createShaders = (precision: Precision) => { } } - let N: f32 = f32(dimLength); + let N: f32 = f32(dimLength - nanCount); let meanX = xSum / N; let meanY = ySum / N; let varX = (xxSum / N) - (meanX * meanX);