Improve performance of heatmaps layer

src/layers/heatmaps.js

@@ -24,6 +24,7 @@ import yaOctree from 'yaot';
 import { emptyObject } from '../utils/gc';
 import threeDigest from '../utils/digest';
 import { array2BufferAttr, bufferAttr2Array } from '../utils/three-utils';
+import { applyShaderExtensionToMaterial, setRadiusShaderExtend } from '../utils/shaders.js';
 import { color2ShaderArr } from '../utils/color-utils';
 import { cartesian2Polar, polar2Cartesian } from '../utils/coordTranslate';
 import { getGeoKDE } from '../utils/kde';
@@ -34,6 +35,7 @@ import { GLOBE_RADIUS } from '../constants';
 const RES_BW_FACTOR = 3.5; // divider of bandwidth to use in geometry resolution
 const MIN_RESOLUTION = 0.1; // degrees
 const BW_RADIUS_INFLUENCE = 3.5; // multiplier of bandwidth to use in octree for max radius of point influence
+const NUM_COLORS = 100; // to sample in shader
 
 class PointsOctree {
   constructor(points, neighborhoodAngularDistance) {
@@ -103,10 +105,13 @@ export default Kapsule({
     const topAltitudeAccessor = accessorFn(state.heatmapTopAltitude);
 
     threeDigest(state.heatmapsData, state.scene, {
-      createObj: d => {
+      createObj: () => {
         const obj = new THREE.Mesh(
           new THREE.SphereGeometry(GLOBE_RADIUS),
-          new THREE.MeshLambertMaterial({ vertexColors: true, transparent: true })
+          applyShaderExtensionToMaterial(
+            new THREE.MeshLambertMaterial({ vertexColors: true, transparent: true }),
+            setRadiusShaderExtend
+          )
         );
 
         obj.__globeObjType = 'heatmap'; // Add object type
@@ -158,28 +163,21 @@ export default Kapsule({
         });
 
         // Animations
+        const colors = [...new Array(NUM_COLORS)].map((_, idx) => color2ShaderArr(colorFn(idx / (NUM_COLORS - 1))));
         const applyUpdate = td => {
           const { kdeVals, topAlt, saturation } = obj.__currentTargetD = td;
           const maxVal = max(kdeVals.map(Math.abs)) || 1e-15;
 
           // Set vertex colors
           obj.geometry.setAttribute('color', array2BufferAttr(
             // normalization between [0, saturation]
-            kdeVals.map(val => color2ShaderArr(colorFn(val / maxVal * saturation))),
+            kdeVals.map(val => colors[Math.min(NUM_COLORS - 1, Math.round(val / maxVal * saturation * NUM_COLORS))]),
             4
           ));
 
           // Set altitudes
           const altScale = scaleLinear([0, maxVal], [baseAlt, topAlt || baseAlt]);
-          obj.geometry.setAttribute('position', array2BufferAttr(
-            kdeVals.map((val, idx) => {
-              const [lng, lat] = vertexGeoCoords[idx];
-              const alt = altScale(Math.abs(val));
-              const p = polar2Cartesian(lat, lng, alt);
-              return [p.x, p.y, p.z];
-            }),
-            3
-          ));
+          obj.geometry.setAttribute('r', array2BufferAttr(kdeVals.map(v => GLOBE_RADIUS * (1 + altScale(Math.abs(v))))));
         };
 
         const targetD = { kdeVals, topAlt, saturation };

src/utils/shaders.js

@@ -97,9 +97,62 @@ export const invisibleUndergroundShaderExtend = shader => {
   return shader;
 };
 
+export const setRadiusShaderExtend = shader => {
+  shader.vertexShader = `
+    attribute float r;
+    
+    const float PI = 3.1415926535897932384626433832795;
+    float toRad(in float a) {
+      return a * PI / 180.0;
+    }
+    
+    vec3 Polar2Cartesian(in vec3 c) { // [lat, lng, r]
+      float phi = toRad(90.0 - c.x);
+      float theta = toRad(90.0 - c.y);
+      float r = c.z;
+      return vec3( // x,y,z
+        r * sin(phi) * cos(theta),
+        r * cos(phi),
+        r * sin(phi) * sin(theta)
+      );
+    }
+    
+    vec2 Cartesian2Polar(in vec3 p) {
+      float r = sqrt(p.x * p.x + p.y * p.y + p.z * p.z);
+      float phi = acos(p.y / r);
+      float theta = atan(p.z, p.x);
+      return vec2( // lat,lng
+        90.0 - phi * 180.0 / PI,
+        90.0 - theta * 180.0 / PI - (theta < -PI / 2.0 ? 360.0 : 0.0)
+      );
+    }
+    ${shader.vertexShader.replace('}', `                  
+        vec3 pos = Polar2Cartesian(vec3(Cartesian2Polar(position), r));
+        gl_Position = projectionMatrix * modelViewMatrix * vec4(pos, 1.0);
+      }
+    `)}
+  `;
+
+  return shader;
+}
+
+//
+
 export const applyShaderExtensionToMaterial = (material, extensionFn) => {
   material.onBeforeCompile = shader => {
     material.userData.shader = extensionFn(shader);
   };
   return material;
-};
+};
+
+export const setExtendedMaterialUniforms = (material, uniformsFn = u => u) => {
+  if (material.userData.shader) {
+    uniformsFn(material.userData.shader.uniforms);
+  } else {
+    const curFn = material.onBeforeCompile;
+    material.onBeforeCompile = shader => {
+      curFn(shader);
+      uniformsFn(shader.uniforms);
+    };
+  }
+}

src/utils/three-utils.js

@@ -4,11 +4,15 @@ const THREE = window.THREE
   ? window.THREE // Prefer consumption from global THREE, if exists
   : { Float32BufferAttribute };
 
-function array2BufferAttr(data, itemSize, BufferAttributeClass = THREE.Float32BufferAttribute) {
-  const ba = new BufferAttributeClass(data.length * itemSize, itemSize);
-  itemSize === 1
-    ? data.forEach((val, idx) => ba.setX(idx, val))
-    : data.forEach((val, idx) => ba.set(val, idx * itemSize));
+function array2BufferAttr(data, itemSize = 1, ArrayClass = Float32Array) {
+  if (itemSize === 1) { // edge case handle for improved performance
+    return new THREE.BufferAttribute(new ArrayClass(data), itemSize);
+  }
+
+  const ba = new THREE.BufferAttribute(new ArrayClass(data.length * itemSize), itemSize);
+  for (let idx = 0, l = data.length; idx < l; idx++) {
+    ba.set(data[idx], idx * itemSize);
+  }
   return ba;
 }