Merge pull request #1390 from silx-kit/buffer-geo

Refactor and optimise creation of `ScatterPoints` geometry

packages/lib/src/vis/scatter/ScatterPoints.tsx

@@ -1,14 +1,15 @@
 import type { Domain, NumArray, ScaleType } from '@h5web/shared';
 import type { ThreeEvent } from '@react-three/fiber';
 import { useThree } from '@react-three/fiber';
-import { useCallback, useLayoutEffect, useState } from 'react';
-import { BufferAttribute, BufferGeometry } from 'three';
+import { useCallback, useLayoutEffect, useMemo } from 'react';
+import { BufferGeometry } from 'three';
 
 import type { ColorMap } from '../heatmap/models';
 import GlyphMaterial from '../line/GlyphMaterial';
 import { GlyphType } from '../line/models';
-import { useDataToColorScale } from './hooks';
-import { useBufferAttributes } from './hooks';
+import { useVisCanvasContext } from '../shared/VisCanvasProvider';
+import { createBufferAttr } from '../utils';
+import { useValueToColor, useIndexToPosition } from './hooks';
 
 interface Props {
   abscissas: NumArray;
@@ -71,28 +72,46 @@ function ScatterPoints(props: Props) {
     [onPointerOut]
   );
 
-  const [dataGeometry] = useState(() => new BufferGeometry());
+  const { length } = data;
+  const { abscissaScale, ordinateScale } = useVisCanvasContext();
   const invalidate = useThree((state) => state.invalidate);
 
-  const dataToColorScale = useDataToColorScale(
+  const indexToPosition = useIndexToPosition(
+    abscissas,
+    abscissaScale,
+    ordinates,
+    ordinateScale
+  );
+
+  const valueToColor = useValueToColor(
     scaleType,
     domain,
     colorMap,
     invertColorMap
   );
 
-  const { position, color } = useBufferAttributes(
-    abscissas,
-    ordinates,
-    data,
-    dataToColorScale
-  );
+  const dataGeometry = useMemo(() => {
+    const geometry = new BufferGeometry();
+    geometry.setAttribute('position', createBufferAttr(length, 3));
+    geometry.setAttribute('color', createBufferAttr(length, 3));
+    return geometry;
+  }, [length]);
 
   useLayoutEffect(() => {
-    dataGeometry.setAttribute('position', new BufferAttribute(position, 3));
-    dataGeometry.setAttribute('color', new BufferAttribute(color, 3, true));
+    const { color, position } = dataGeometry.attributes;
+
+    data.forEach((val, index) => {
+      const { x, y, z } = indexToPosition(index);
+      position.setXYZ(index, x, y, z);
+
+      const { r, g, b } = valueToColor(val);
+      color.setXYZ(index, r / 255, g / 255, b / 255); // normalize RGB channels
+    });
+
+    color.needsUpdate = true;
+    position.needsUpdate = true;
     invalidate();
-  }, [color, dataGeometry, invalidate, position]);
+  }, [data, dataGeometry, indexToPosition, valueToColor, invalidate]);
 
   return (
     <points

packages/lib/src/vis/scatter/hooks.ts

@@ -1,66 +1,6 @@
-import type { Domain, NumArray, ScaleType } from '@h5web/shared';
-import { rgb } from 'd3-color';
-import { useMemo } from 'react';
+import { createMemo } from '@h5web/shared';
 
-import type { ColorMap } from '../heatmap/models';
-import { getInterpolator } from '../heatmap/utils';
-import { useVisCanvasContext } from '../shared/VisCanvasProvider';
-import { createAxisScale } from '../utils';
+import { getValueToColor, getIndexToPosition } from './utils';
 
-const CAMERA_FAR = 1000; // R3F's default
-
-export function useBufferAttributes(
-  abscissas: NumArray,
-  ordinates: NumArray,
-  data: NumArray,
-  dataToColorScale: (val: number) => [number, number, number]
-) {
-  const { abscissaScale, ordinateScale } = useVisCanvasContext();
-
-  return useMemo(() => {
-    const position = new Float32Array(3 * data.length);
-    const color = new Uint8Array(3 * data.length);
-
-    data.forEach((val, index) => {
-      color.set(dataToColorScale(val), 3 * index);
-
-      const x = abscissaScale(abscissas[index]);
-      const y = ordinateScale(ordinates[index]);
-
-      const hasFiniteCoords = Number.isFinite(x) && Number.isFinite(y);
-
-      /* Render points with NaN/Infinity coordinates (i.e. values <= 0 in log)
-       * at origin to avoid Three warning, and outside of camera's field of view
-       * to hide them and any segments connecting them. */
-      position.set([x, y, hasFiniteCoords ? 0 : CAMERA_FAR], 3 * index);
-    });
-
-    return { position, color };
-  }, [
-    abscissaScale,
-    abscissas,
-    dataToColorScale,
-    data,
-    ordinateScale,
-    ordinates,
-  ]);
-}
-
-export function useDataToColorScale(
-  scaleType: ScaleType,
-  domain: Domain,
-  colorMap: ColorMap,
-  invertColorMap: boolean
-): (v: number) => [number, number, number] {
-  return useMemo(() => {
-    const numScale = createAxisScale(scaleType, {
-      domain,
-      range: [0, 1],
-    });
-    const interpolator = getInterpolator(colorMap, invertColorMap);
-    return (value: number) => {
-      const color = rgb(interpolator(numScale(value)));
-      return [color.r, color.g, color.b];
-    };
-  }, [colorMap, domain, invertColorMap, scaleType]);
-}
+export const useIndexToPosition = createMemo(getIndexToPosition);
+export const useValueToColor = createMemo(getValueToColor);

packages/lib/src/vis/scatter/utils.ts

@@ -0,0 +1,39 @@
+import type { Domain, NumArray, ScaleType } from '@h5web/shared';
+import type { RGBColor } from 'd3-color';
+import { rgb } from 'd3-color';
+
+import type { ColorMap } from '../heatmap/models';
+import { getInterpolator } from '../heatmap/utils';
+import type { AxisScale } from '../models';
+import { createAxisScale } from '../utils';
+
+const CAMERA_FAR = 1000; // R3F's default
+
+export function getIndexToPosition(
+  abscissas: NumArray,
+  abscissaScale: AxisScale,
+  ordinates: NumArray,
+  ordinateScale: AxisScale
+): (index: number) => { x: number; y: number; z: number } {
+  return (index: number) => {
+    const x = abscissaScale(abscissas[index]);
+    const y = ordinateScale(ordinates[index]);
+
+    /* Render points with NaN/Infinity coordinates (i.e. values <= 0 in log)
+     * at origin to avoid Three warning, and outside of camera's field of view
+     * to hide them and any segments connecting them. */
+    const hasFiniteCoords = Number.isFinite(x) && Number.isFinite(y);
+    return hasFiniteCoords ? { x, y, z: 0 } : { x: 0, y: 0, z: CAMERA_FAR };
+  };
+}
+
+export function getValueToColor(
+  scaleType: ScaleType,
+  domain: Domain,
+  colorMap: ColorMap,
+  invertColorMap: boolean
+): (v: number) => RGBColor {
+  const interpolator = getInterpolator(colorMap, invertColorMap);
+  const numScale = createAxisScale(scaleType, { domain, range: [0, 1] });
+  return (value: number) => rgb(interpolator(numScale(value)));
+}

packages/lib/src/vis/utils.ts

@@ -14,6 +14,7 @@ import { tickStep, range } from 'd3-array';
 import type { ScaleLinear, ScaleThreshold } from 'd3-scale';
 import { clamp } from 'lodash';
 import type { IUniform } from 'three';
+import { BufferAttribute } from 'three';
 
 import type {
   Size,
@@ -356,3 +357,10 @@ export function getAxisValues(
 
   return rawValues.slice(0, axisLength);
 }
+
+export function createBufferAttr(
+  dataLength: number,
+  itemSize: number
+): BufferAttribute {
+  return new BufferAttribute(new Float32Array(dataLength * itemSize), itemSize);
+}