HeatmapLayer UBO (#9071)

modules/aggregation-layers/src/heatmap-layer/heatmap-layer-uniforms.ts

@@ -0,0 +1,46 @@
+import {Texture} from '@luma.gl/core';
+import type {ShaderModule} from '@luma.gl/shadertools';
+
+const uniformBlock = `\
+uniform weightUniforms {
+  vec4 commonBounds;
+  float radiusPixels;
+  float textureWidth;
+  float weightsScale;
+} weight;
+`;
+export type WeightProps = {
+  commonBounds: [number, number, number, number];
+  radiusPixels: number;
+  textureWidth: number;
+  weightsScale: number;
+  weightsTexture: Texture;
+};
+
+export const weightUniforms = {
+  name: 'weight',
+  vs: uniformBlock,
+  uniformTypes: {
+    commonBounds: 'vec4<f32>',
+    radiusPixels: 'f32',
+    textureWidth: 'f32',
+    weightsScale: 'f32'
+  }
+} as const satisfies ShaderModule<WeightProps>;
+
+export type MaxWeightProps = {
+  inTexture: Texture;
+  textureSize: number;
+};
+
+export const maxWeightUniforms = {
+  name: 'maxWeight',
+  vs: `\
+uniform maxWeightUniforms {
+  float textureSize;
+} maxWeight;
+`,
+  uniformTypes: {
+    textureSize: 'f32'
+  }
+} as const satisfies ShaderModule<MaxWeightProps>;

modules/aggregation-layers/src/heatmap-layer/heatmap-layer-utils.ts

@@ -38,7 +38,11 @@ export function packVertices(points: number[][], dimensions: number = 2): Float3
 }
 
 // Expands boundingBox:[xMin, yMin, xMax, yMax] to match aspect ratio of given width and height
-export function scaleToAspectRatio(boundingBox: number[], width: number, height: number): number[] {
+export function scaleToAspectRatio(
+  boundingBox: [number, number, number, number],
+  width: number,
+  height: number
+): [number, number, number, number] {
   const [xMin, yMin, xMax, yMax] = boundingBox;
 
   const currentWidth = xMax - xMin;

modules/aggregation-layers/src/heatmap-layer/heatmap-layer.ts

@@ -51,6 +51,12 @@ import weightsVs from './weights-vs.glsl';
 import weightsFs from './weights-fs.glsl';
 import maxVs from './max-vs.glsl';
 import maxFs from './max-fs.glsl';
+import {
+  MaxWeightProps,
+  maxWeightUniforms,
+  WeightProps,
+  weightUniforms
+} from './heatmap-layer-uniforms';
 
 const RESOLUTION = 2; // (number of common space pixels) / (number texels)
 const TEXTURE_PROPS: TextureProps = {
@@ -203,7 +209,12 @@ export default class HeatmapLayer<
   };
 
   getShaders(shaders: any) {
-    return super.getShaders({...shaders, modules: [project32]});
+    let modules = [project32];
+    if (shaders.modules) {
+      modules = [...modules, ...shaders.modules];
+    }
+
+    return super.getShaders({...shaders, modules});
   }
 
   initializeState() {
@@ -396,7 +407,7 @@ export default class HeatmapLayer<
     }
   }
 
-  _createWeightsTransform(shaders: {vs: string; fs?: string}) {
+  _createWeightsTransform(shaders: {vs: string; fs?: string; modules: any[]}) {
     let {weightsTransform} = this.state;
     const {weightsTexture} = this.state;
     const attributeManager = this.getAttributeManager()!;
@@ -416,7 +427,8 @@ export default class HeatmapLayer<
         blendAlphaDstFactor: 'one'
       },
       topology: 'point-list',
-      ...shaders
+      ...shaders,
+      modules: [...shaders.modules, weightUniforms]
     } as TextureTransformProps);
 
     this.setState({weightsTransform});
@@ -433,12 +445,14 @@ export default class HeatmapLayer<
     });
     this._createWeightsTransform(weightsTransformShaders);
 
-    const maxWeightsTransformShaders = this.getShaders({vs: maxVs, fs: maxFs});
+    const maxWeightsTransformShaders = this.getShaders({
+      vs: maxVs,
+      fs: maxFs,
+      modules: [maxWeightUniforms]
+    });
     const maxWeightTransform = new TextureTransform(device, {
       id: `${this.id}-max-weights-transform`,
-      bindings: {inTexture: weightsTexture},
-      uniforms: {textureSize},
-      targetTexture: maxWeightsTexture,
+      targetTexture: maxWeightsTexture!,
       ...maxWeightsTransformShaders,
       vertexCount: textureSize * textureSize,
       topology: 'point-list',
@@ -453,6 +467,11 @@ export default class HeatmapLayer<
       }
     });
 
+    const maxWeightProps: MaxWeightProps = {inTexture: weightsTexture!, textureSize};
+    maxWeightTransform.model.shaderInputs.setProps({
+      maxWeight: maxWeightProps
+    });
+
     this.setState({
       weightsTexture,
       maxWeightsTexture,
@@ -570,7 +589,7 @@ export default class HeatmapLayer<
 
   _updateWeightmap() {
     const {radiusPixels, colorDomain, aggregation} = this.props;
-    const {worldBounds, textureSize, weightsScale} = this.state;
+    const {worldBounds, textureSize, weightsScale, weightsTexture} = this.state;
     const weightsTransform = this.state.weightsTransform!;
     this.state.isWeightMapDirty = false;
 
@@ -593,16 +612,22 @@ export default class HeatmapLayer<
     const attributeManager = this.getAttributeManager()!;
     const attributes = attributeManager.getAttributes();
     const moduleSettings = this.getModuleSettings();
-    const uniforms = {radiusPixels, commonBounds, textureWidth: textureSize, weightsScale};
     this._setModelAttributes(weightsTransform.model, attributes);
     weightsTransform.model.setVertexCount(this.getNumInstances());
-    weightsTransform.model.setUniforms(uniforms);
     weightsTransform.model.updateModuleSettings(moduleSettings);
 
+    const weightProps: WeightProps = {
+      radiusPixels,
+      commonBounds,
+      textureWidth: textureSize,
+      weightsScale,
+      weightsTexture: weightsTexture!
+    };
     const {viewport, devicePixelRatio, coordinateSystem, coordinateOrigin} = moduleSettings;
     const {modelMatrix} = this.props;
     weightsTransform.model.shaderInputs.setProps({
-      project: {viewport, devicePixelRatio, modelMatrix, coordinateSystem, coordinateOrigin}
+      project: {viewport, devicePixelRatio, modelMatrix, coordinateSystem, coordinateOrigin},
+      weight: weightProps
     });
     weightsTransform.run({
       parameters: {viewport: [0, 0, textureSize, textureSize]},
@@ -634,7 +659,10 @@ export default class HeatmapLayer<
   // input: worldBounds: [minLong, minLat, maxLong, maxLat]
   // input: opts.useLayerCoordinateSystem : layers coordiante system is used
   // optput: commonBounds: [minX, minY, maxX, maxY] scaled to fit the current texture
-  _worldToCommonBounds(worldBounds, opts: {useLayerCoordinateSystem?: boolean} = {}) {
+  _worldToCommonBounds(
+    worldBounds,
+    opts: {useLayerCoordinateSystem?: boolean} = {}
+  ): [number, number, number, number] {
     const {useLayerCoordinateSystem = false} = opts;
     const [minLong, minLat, maxLong, maxLat] = worldBounds;
     const {viewport} = this.context;

modules/aggregation-layers/src/heatmap-layer/max-vs.glsl.ts

@@ -1,15 +1,14 @@
 export default `\
 #version 300 es
 uniform sampler2D inTexture;
-uniform float textureSize;
 out vec4 outTexture;
 
 void main()
 {
   // Sample every pixel in texture
-  int yIndex = gl_VertexID / int(textureSize);
-  int xIndex = gl_VertexID - (yIndex * int(textureSize));
-  vec2 uv = (0.5 + vec2(float(xIndex), float(yIndex))) / textureSize;
+  int yIndex = gl_VertexID / int(maxWeight.textureSize);
+  int xIndex = gl_VertexID - (yIndex * int(maxWeight.textureSize));
+  vec2 uv = (0.5 + vec2(float(xIndex), float(yIndex))) / maxWeight.textureSize;
   outTexture = texture(inTexture, uv);
 
   gl_Position = vec4(0.0, 0.0, 0.0, 1.0);

modules/aggregation-layers/src/heatmap-layer/triangle-layer-fragment.glsl.ts

@@ -24,10 +24,8 @@ export default `\
 
 precision highp float;
 
-uniform float opacity;
 uniform sampler2D weightsTexture;
 uniform sampler2D colorTexture;
-uniform float aggregationMode;
 
 in vec2 vTexCoords;
 in float vIntensityMin;
@@ -46,7 +44,7 @@ void main(void) {
   vec4 weights = texture(weightsTexture, vTexCoords);
   float weight = weights.r;
 
-  if (aggregationMode > 0.5) {
+  if (triangle.aggregationMode > 0.5) {
     weight /= max(1.0, weights.a);
   }
 
@@ -56,7 +54,7 @@ void main(void) {
   }
 
   vec4 linearColor = getLinearColor(weight);
-  linearColor.a *= opacity;
+  linearColor.a *= layer.opacity;
   fragColor = linearColor;
 }
 `;

modules/aggregation-layers/src/heatmap-layer/triangle-layer-uniforms.ts

@@ -0,0 +1,33 @@
+import {Texture} from '@luma.gl/core';
+import type {ShaderModule} from '@luma.gl/shadertools';
+
+const uniformBlock = `\
+uniform triangleUniforms {
+  float aggregationMode;
+  vec2 colorDomain;
+  float intensity;
+  float threshold;
+} triangle;
+`;
+
+export type TriangleProps = {
+  aggregationMode: number;
+  colorDomain: [number, number];
+  intensity: number;
+  threshold: number;
+  colorTexture: Texture;
+  maxTexture: Texture;
+  weightsTexture: Texture;
+};
+
+export const triangleUniforms = {
+  name: 'triangle',
+  vs: uniformBlock,
+  fs: uniformBlock,
+  uniformTypes: {
+    aggregationMode: 'f32',
+    colorDomain: 'vec2<f32>',
+    intensity: 'f32',
+    threshold: 'f32'
+  }
+} as const satisfies ShaderModule<TriangleProps>;

modules/aggregation-layers/src/heatmap-layer/triangle-layer-vertex.glsl.ts

@@ -25,10 +25,6 @@ export default `\
 #define SHADER_NAME heatp-map-layer-vertex-shader
 
 uniform sampler2D maxTexture;
-uniform float intensity;
-uniform vec2 colorDomain;
-uniform float threshold;
-uniform float aggregationMode;
 
 in vec3 positions;
 in vec2 texCoords;
@@ -41,14 +37,14 @@ void main(void) {
   gl_Position = project_position_to_clipspace(positions, vec3(0.0), vec3(0.0));
   vTexCoords = texCoords;
   vec4 maxTexture = texture(maxTexture, vec2(0.5));
-  float maxValue = aggregationMode < 0.5 ? maxTexture.r : maxTexture.g;
-  float minValue = maxValue * threshold;
-  if (colorDomain[1] > 0.) {
+  float maxValue = triangle.aggregationMode < 0.5 ? maxTexture.r : maxTexture.g;
+  float minValue = maxValue * triangle.threshold;
+  if (triangle.colorDomain[1] > 0.) {
     // if user specified custom domain use it.
-    maxValue = colorDomain[1];
-    minValue = colorDomain[0];
+    maxValue = triangle.colorDomain[1];
+    minValue = triangle.colorDomain[0];
   }
-  vIntensityMax = intensity / maxValue;
-  vIntensityMin = intensity / minValue;
+  vIntensityMax = triangle.intensity / maxValue;
+  vIntensityMin = triangle.intensity / minValue;
 }
 `;

modules/aggregation-layers/src/heatmap-layer/triangle-layer.ts

@@ -23,11 +23,12 @@ import {Model} from '@luma.gl/engine';
 import {Layer, LayerContext, project32} from '@deck.gl/core';
 import vs from './triangle-layer-vertex.glsl';
 import fs from './triangle-layer-fragment.glsl';
+import {TriangleProps, triangleUniforms} from './triangle-layer-uniforms';
 
 type _TriangleLayerProps = {
   data: {attributes: {positions: Buffer; texCoords: Buffer}};
-  colorDomain: number[];
-  aggregationMode: string;
+  colorDomain: [number, number];
+  aggregationMode: number;
   threshold: number;
   intensity: number;
   vertexCount: number;
@@ -46,20 +47,19 @@ export default class TriangleLayer extends Layer<_TriangleLayerProps> {
   };
 
   getShaders() {
-    return {vs, fs, modules: [project32]};
+    return super.getShaders({vs, fs, modules: [project32, triangleUniforms]});
   }
 
   initializeState({device}: LayerContext): void {
     this.setState({model: this._getModel(device)});
   }
 
   _getModel(device: Device): Model {
-    const {vertexCount, data, weightsTexture, maxTexture, colorTexture} = this.props;
+    const {vertexCount, data} = this.props;
 
     return new Model(device, {
       ...this.getShaders(),
       id: this.props.id,
-      bindings: {weightsTexture, maxTexture, colorTexture},
       attributes: data.attributes,
       bufferLayout: [
         {name: 'positions', format: 'float32x3'},
@@ -70,16 +70,27 @@ export default class TriangleLayer extends Layer<_TriangleLayerProps> {
     });
   }
 
-  draw({uniforms}): void {
+  draw(): void {
     const {model} = this.state;
-    const {intensity, threshold, aggregationMode, colorDomain} = this.props;
-    model.setUniforms({
-      ...uniforms,
+    const {
+      aggregationMode,
+      colorDomain,
       intensity,
       threshold,
+      colorTexture,
+      maxTexture,
+      weightsTexture
+    } = this.props;
+    const triangleProps: TriangleProps = {
       aggregationMode,
-      colorDomain
-    });
+      colorDomain,
+      intensity,
+      threshold,
+      colorTexture,
+      maxTexture,
+      weightsTexture
+    };
+    model.shaderInputs.setProps({triangle: triangleProps});
     model.draw(this.context.renderPass);
   }
 }

modules/aggregation-layers/src/heatmap-layer/weights-vs.glsl.ts

@@ -4,22 +4,18 @@ in vec3 positions;
 in vec3 positions64Low;
 in float weights;
 out vec4 weightsTexture;
-uniform float radiusPixels;
-uniform float textureWidth;
-uniform vec4 commonBounds;
-uniform float weightsScale;
 
 void main()
 {
-  weightsTexture = vec4(weights * weightsScale, 0., 0., 1.);
+  weightsTexture = vec4(weights * weight.weightsScale, 0., 0., 1.);
 
-  float radiusTexels = project_pixel_size(radiusPixels) * textureWidth / (commonBounds.z - commonBounds.x);
+  float radiusTexels = project_pixel_size(weight.radiusPixels) * weight.textureWidth / (weight.commonBounds.z - weight.commonBounds.x);
   gl_PointSize = radiusTexels * 2.;
 
   vec3 commonPosition = project_position(positions, positions64Low);
 
   // // map xy from commonBounds to [-1, 1]
-  gl_Position.xy = (commonPosition.xy - commonBounds.xy) / (commonBounds.zw - commonBounds.xy) ;
+  gl_Position.xy = (commonPosition.xy - weight.commonBounds.xy) / (weight.commonBounds.zw - weight.commonBounds.xy) ;
   gl_Position.xy = (gl_Position.xy * 2.) - (1.);
   gl_Position.w = 1.0;
 }