Cut the main-thread cost of opening and re-rendering a visualisation

packages/app/src/App.tsx

@@ -22,6 +22,7 @@ import { useDataContext } from './providers/DataProvider';
 import Sidebar from './Sidebar';
 import VisConfigProvider from './VisConfigProvider';
 import Visualizer from './visualizer/Visualizer';
+import { useWarmUpWebGL } from './webgl-warm-up';
 
 const SIDEBAR_ID = 'h5w-sidebar';
 const MAIN_AREA_ID = 'h5w-main-area';
@@ -47,6 +48,8 @@ function App(props: Props) {
   const [selectedPath, setSelectedPath] = useState<string>(initialPath);
   const [isInspecting, setInspecting] = useState(false);
 
+  useWarmUpWebGL();
+
   const { valuesStore } = useDataContext();
   function onSelectPath(path: string) {
     setSelectedPath(path);

packages/app/src/vis-packs/core/complex/utils.test.ts

@@ -0,0 +1,154 @@
+import { type H5WebComplex } from '@h5web/shared/hdf5-models';
+import { cplx } from '@h5web/shared/hdf5-utils';
+import { describe, expect, it } from 'vitest';
+
+import { getPhaseAmplitude, getPhaseAmplitudeArrays } from './utils';
+
+const TWO_PI = 2 * Math.PI;
+
+/* Straightforward reference implementations, kept deliberately naive. The real
+   ones are written to avoid per-element allocation on large datasets, so they
+   are checked against these rather than against hand-written expectations. */
+function referencePhaseAmplitude(values: H5WebComplex[]): {
+  phase: number[];
+  amplitude: number[];
+} {
+  return {
+    phase: values.map(([real, imag]) => Math.atan2(imag, real)),
+    amplitude: values.map(([real, imag]) => Math.hypot(real, imag)),
+  };
+}
+
+function referenceUnwrap(values: number[]): number[] {
+  const unwrapped: number[] = [];
+
+  values.forEach((val, i) => {
+    if (i === 0) {
+      unwrapped.push(val);
+      return;
+    }
+
+    const diff = val - unwrapped[i - 1];
+    unwrapped.push(val - TWO_PI * Math.round(diff / TWO_PI));
+  });
+
+  return unwrapped;
+}
+
+/* Deterministic generator (MINSTD), so a failing case is reproducible. Values
+   straddle the branch cut of `atan2` at the negative real axis, where phase
+   jumps by 2π and unwrapping has to do real work. */
+function makeComplex(seed: number, length: number): H5WebComplex[] {
+  let state = seed;
+
+  function next(): number {
+    state = (state * 48_271) % 2_147_483_647;
+    return state / 2_147_483_647;
+  }
+
+  return Array.from({ length }, () => {
+    const angle = (next() - 0.5) * 8 * Math.PI; // several turns, so phase wraps
+    const radius = next() * 10;
+    return cplx(radius * Math.cos(angle), radius * Math.sin(angle));
+  });
+}
+
+function unwrapOf(values: H5WebComplex[]): number[] {
+  const { unwrappedPhaseArrays } = getPhaseAmplitudeArrays([values]);
+  return [...unwrappedPhaseArrays[0]];
+}
+
+describe('getPhaseAmplitude', () => {
+  it('should match the reference implementation', () => {
+    for (let seed = 1; seed <= 50; seed += 1) {
+      const values = makeComplex(seed, 64);
+      const { phase, amplitude } = getPhaseAmplitude(values);
+      const expected = referencePhaseAmplitude(values);
+
+      expect({ seed, phase: [...phase] }).toStrictEqual({
+        seed,
+        phase: expected.phase,
+      });
+      expect({ seed, amplitude: [...amplitude] }).toStrictEqual({
+        seed,
+        amplitude: expected.amplitude,
+      });
+    }
+  });
+
+  it('should handle the values that break polar conversions', () => {
+    // Both zero (phase undefined, conventionally 0) and the negative real axis
+    const values = [cplx(0, 0), cplx(-1, 0), cplx(0, -1), cplx(-0, -0)];
+    const { phase, amplitude } = getPhaseAmplitude(values);
+
+    expect([...amplitude]).toStrictEqual([0, 1, 1, 0]);
+    expect([...phase]).toStrictEqual(referencePhaseAmplitude(values).phase);
+  });
+
+  it('should return empty arrays for no values', () => {
+    const { phase, amplitude } = getPhaseAmplitude([]);
+
+    expect(phase).toHaveLength(0);
+    expect(amplitude).toHaveLength(0);
+  });
+
+  /* Property: amplitude and phase are polar coordinates, so they must
+     reconstruct the original components. This pins down the pairing of the two
+     outputs, which comparing each against a reference separately does not. */
+  it('should be invertible back to the original components', () => {
+    const values = makeComplex(7, 200);
+    const { phase, amplitude } = getPhaseAmplitude(values);
+
+    values.forEach(([real, imag], i) => {
+      expect(amplitude[i] * Math.cos(phase[i])).toBeCloseTo(real, 10);
+      expect(amplitude[i] * Math.sin(phase[i])).toBeCloseTo(imag, 10);
+    });
+  });
+});
+
+describe('getPhaseAmplitudeArrays', () => {
+  it('should match the reference unwrapping', () => {
+    for (let seed = 1; seed <= 50; seed += 1) {
+      const values = makeComplex(seed, 64);
+      const { phase } = getPhaseAmplitude(values);
+
+      expect({ seed, unwrapped: unwrapOf(values) }).toStrictEqual({
+        seed,
+        unwrapped: referenceUnwrap([...phase]),
+      });
+    }
+  });
+
+  /* Properties that define phase unwrapping: it may add whole turns to each
+     value, and must leave no jump bigger than half a turn between neighbours. */
+  it('should remove every 2-pi discontinuity without moving any value off its turn', () => {
+    const values = makeComplex(11, 500);
+    const { phase } = getPhaseAmplitude(values);
+    const unwrapped = unwrapOf(values);
+
+    expect(unwrapped).toHaveLength(values.length);
+    expect(unwrapped[0]).toBe(phase[0]); // first value is never shifted
+
+    unwrapped.forEach((val, i) => {
+      // Differs from the wrapped phase only by a whole number of turns
+      const turns = (val - phase[i]) / TWO_PI;
+      expect(turns).toBeCloseTo(Math.round(turns), 9);
+    });
+
+    // No jump bigger than half a turn survives anywhere in the sequence
+    const jumps = unwrapped
+      .slice(1)
+      .map((val, i) => Math.abs(val - unwrapped[i]));
+
+    expect(Math.max(...jumps)).toBeLessThanOrEqual(Math.PI + 1e-9);
+  });
+
+  it('should treat real values as complex with zero imaginary part', () => {
+    const { phaseArrays, amplitudeArrays } = getPhaseAmplitudeArrays([
+      [-3, 0, 4],
+    ]);
+
+    expect([...phaseArrays[0]]).toStrictEqual([0, 0, 0]);
+    expect([...amplitudeArrays[0]]).toStrictEqual([3, 0, 4]);
+  });
+});

packages/app/src/vis-packs/core/complex/utils.ts

@@ -18,33 +18,42 @@ export const COMPLEX_VIS_TYPE_LABELS = {
   [ComplexVisType.PhaseAmplitude]: 'Phase & Amplitude',
 } satisfies Record<ComplexVisType, string>;
 
+/* The output arrays are `Float64Array`s rather than `number[]`: an array built
+ * with `Array.from({ length })` is filled with `undefined`, which fixes its V8
+ * elements kind as `PACKED_ELEMENTS`, so every phase and amplitude written into
+ * it is boxed as a separate heap number. */
 export function getPhaseAmplitude(values: H5WebComplex[]): {
-  phase: number[];
-  amplitude: number[];
+  phase: NumArray;
+  amplitude: NumArray;
 } {
-  const phase: number[] = Array.from({ length: values.length });
-  const amplitude: number[] = Array.from({ length: values.length });
+  const phase = new Float64Array(values.length);
+  const amplitude = new Float64Array(values.length);
 
-  values.forEach(([real, imag], i) => {
+  /* Iterate by index: destructuring in a `forEach` callback allocates per
+   * element, which dominates the scan once the boxing is gone. */
+  // eslint-disable-next-line unicorn/no-for-loop -- see above
+  for (let i = 0; i < values.length; i += 1) {
+    const [real, imag] = values[i];
     phase[i] = Math.atan2(imag, real);
     amplitude[i] = Math.hypot(real, imag);
-  });
+  }
 
   return { phase, amplitude };
 }
 
 // Unwrap phase values by removing 2π discontinuities
-function unwrapPhase(values: number[]): number[] {
-  const unwrapped: number[] = Array.from({ length: values.length });
-
-  for (const [i, val] of values.entries()) {
-    if (i === 0) {
-      unwrapped[0] = val;
-      continue;
-    }
+function unwrapPhase(values: NumArray): NumArray {
+  const unwrapped = new Float64Array(values.length);
 
-    const diff = val - unwrapped[i - 1];
-    unwrapped[i] = val - TWO_PI * Math.round(diff / TWO_PI);
+  /* Carry the previous unwrapped value rather than reading it back, and index
+   * rather than using `values.entries()`, which allocates a pair per element. */
+  let previous = 0;
+  // eslint-disable-next-line unicorn/no-for-loop -- see above
+  for (let i = 0; i < values.length; i += 1) {
+    const val = values[i];
+    const diff = val - previous;
+    previous = i === 0 ? val : val - TWO_PI * Math.round(diff / TWO_PI);
+    unwrapped[i] = previous;
   }
 
   return unwrapped;

packages/app/src/webgl-warm-up.ts

@@ -0,0 +1,31 @@
+import { useEffect } from 'react';
+
+/* Creating the first WebGL context in a page costs a few hundred milliseconds
+   of blocking main-thread time, most of it GPU-process and graphics-driver
+   initialisation rather than anything h5web does. Pay it in the background,
+   while the user is still browsing the file tree, so that opening the first
+   visualisation doesn't have to. */
+export function useWarmUpWebGL(): void {
+  useEffect(() => {
+    let context: WebGL2RenderingContext | null = null;
+    let handle: number | undefined;
+
+    // Warming up is best-effort, so skip it where `requestIdleCallback` is unsupported
+    if ('requestIdleCallback' in globalThis) {
+      handle = globalThis.requestIdleCallback(() => {
+        /* Keep the context around for the session: releasing it lets the
+           browser tear down the GPU process again, which is the cost being
+           avoided in the first place. */
+        context = document.createElement('canvas').getContext('webgl2');
+      });
+    }
+
+    return () => {
+      if (handle !== undefined) {
+        globalThis.cancelIdleCallback(handle);
+      }
+
+      context?.getExtension('WEBGL_lose_context')?.loseContext();
+    };
+  }, []);
+}

packages/lib/src/vis/heatmap/HeatmapMaterial.tsx

@@ -1,7 +1,8 @@
 import { type Domain, ScaleType } from '@h5web/shared/vis-models';
+import { useThree } from '@react-three/fiber';
 import { rgb, type RGBColor } from 'd3-color';
 import { type NdArray } from 'ndarray';
-import { memo, useMemo } from 'react';
+import { memo, useLayoutEffect, useMemo } from 'react';
 import {
   DataTexture,
   DoubleSide,
@@ -91,32 +92,35 @@ function HeatmapMaterial(props: Props) {
     ? visScaleType
     : [visScaleType as ScaleType, 1];
 
-  const scaledDomain = scaleDomain(domain, scaleType);
+  const [scaledMin, scaledMax] = scaleDomain(domain, scaleType);
+  const [alphaMin, alphaMax] = alphaDomain;
 
   const badColorAsRgb = typeof badColor === 'string' ? rgb(badColor) : badColor;
-
-  const shader = {
-    uniforms: getUniforms({
-      data: dataTexture,
-      mask: maskTexture,
-      colorMap: colorMapTexture,
-      min: scaledDomain[0],
-      oneOverRange: 1 / (scaledDomain[1] - scaledDomain[0]),
-      gammaExponent,
-      normRevertFactor: values.dtype === 'uint8' ? 255 : 1, // revert WebGL's automatic normalization of UNSIGNED_BYTE with RED format
-      alpha: alphaTexture,
-      withAlpha: alphaValues ? 1 : 0,
-      alphaMin: alphaDomain[0],
-      oneOverAlphaRange: 1 / (alphaDomain[1] - alphaDomain[0]),
-      badColor: new Vector4(
-        badColorAsRgb.r / 255,
-        badColorAsRgb.g / 255,
-        badColorAsRgb.b / 255,
-        badColorAsRgb.opacity,
-      ),
-    }),
-    vertexShader: VERTEX_SHADER,
-    fragmentShader: `
+  const badColorVector = useMemo(() => new Vector4(), []); // updated in place, like the other uniform values
+
+  /* The scale type is the only prop that appears in the shader source, so the
+     material only has to be re-created when it changes. Everything else is a
+     uniform, updated in place below: passing a new object in `args` would make
+     R3F re-instantiate the material, which recompiles and relinks the GLSL
+     program — on every slice of a 3D dataset, for instance. */
+  const shader = useMemo(
+    () => ({
+      uniforms: getUniforms({
+        data: null,
+        mask: null,
+        colorMap: null,
+        min: 0,
+        oneOverRange: 1,
+        gammaExponent: 1,
+        normRevertFactor: 1,
+        alpha: null,
+        withAlpha: 0,
+        alphaMin: 0,
+        oneOverAlphaRange: 1,
+        badColor: badColorVector,
+      }),
+      vertexShader: VERTEX_SHADER,
+      fragmentShader: `
       uniform sampler2D data;
       uniform sampler2D colorMap;
 
@@ -163,7 +167,35 @@ function HeatmapMaterial(props: Props) {
         }
       }
     `,
-  };
+    }),
+    [scaleType, badColorVector],
+  );
+
+  const invalidate = useThree((state) => state.invalidate);
+  const { uniforms } = shader; // same object the material holds, so mutating it updates the material
+
+  useLayoutEffect(() => {
+    uniforms.data.value = dataTexture;
+    uniforms.mask.value = maskTexture;
+    uniforms.colorMap.value = colorMapTexture;
+    uniforms.min.value = scaledMin;
+    uniforms.oneOverRange.value = 1 / (scaledMax - scaledMin);
+    uniforms.gammaExponent.value = gammaExponent;
+    uniforms.normRevertFactor.value = values.dtype === 'uint8' ? 255 : 1; // revert WebGL's automatic normalization of UNSIGNED_BYTE with RED format
+    uniforms.alpha.value = alphaTexture;
+    uniforms.withAlpha.value = alphaValues ? 1 : 0;
+    uniforms.alphaMin.value = alphaMin;
+    uniforms.oneOverAlphaRange.value = 1 / (alphaMax - alphaMin);
+
+    badColorVector.set(
+      badColorAsRgb.r / 255,
+      badColorAsRgb.g / 255,
+      badColorAsRgb.b / 255,
+      badColorAsRgb.opacity,
+    );
+
+    invalidate();
+  });
 
   return <shaderMaterial args={[shader]} side={DoubleSide} />;
 }