@cosmos.gl/graph 2.6.2 → 2.7.0-beta.0

package/src/modules/ForceManyBody/index.ts

@@ -0,0 +1,525 @@
+import { Buffer, Framebuffer, RenderPass, Texture, UniformStore } from '@luma.gl/core'
+import { Model } from '@luma.gl/engine'
+import { CoreModule } from '@/graph/modules/core-module'
+import calculateLevelFrag from '@/graph/modules/ForceManyBody/calculate-level.frag?raw'
+import calculateLevelVert from '@/graph/modules/ForceManyBody/calculate-level.vert?raw'
+import forceFrag from '@/graph/modules/ForceManyBody/force-level.frag?raw'
+import forceCenterFrag from '@/graph/modules/ForceManyBody/force-centermass.frag?raw'
+import { createIndexesForBuffer } from '@/graph/modules/Shared/buffer'
+import clearFrag from '@/graph/modules/Shared/clear.frag?raw'
+import updateVert from '@/graph/modules/Shared/quad.vert?raw'
+
+type LevelTarget = {
+  texture: Texture;
+  fbo: Framebuffer;
+}
+
+export class ForceManyBody extends CoreModule {
+  private randomValuesTexture: Texture | undefined
+  private pointIndices: Buffer | undefined
+  private levels = 0
+  private levelTargets = new Map<number, LevelTarget>()
+
+  private clearLevelsCommand: Model | undefined
+  private calculateLevelsCommand: Model | undefined
+  private forceCommand: Model | undefined
+  private forceFromItsOwnCentermassCommand: Model | undefined
+
+  private clearLevelsVertexCoordBuffer: Buffer | undefined
+  private forceVertexCoordBuffer: Buffer | undefined
+
+  private calculateLevelsUniformStore: UniformStore<{
+    calculateLevelsUniforms: {
+      pointsTextureSize: number;
+      levelTextureSize: number;
+      cellSize: number;
+    };
+  }> | undefined
+
+  private forceUniformStore: UniformStore<{
+    forceUniforms: {
+      level: number;
+      levels: number;
+      levelTextureSize: number;
+      alpha: number;
+      repulsion: number;
+      spaceSize: number;
+      theta: number;
+    };
+  }> | undefined
+
+  private forceCenterUniformStore: UniformStore<{
+    forceCenterUniforms: {
+      levelTextureSize: number;
+      alpha: number;
+      repulsion: number;
+    };
+  }> | undefined
+
+  private previousPointsTextureSize: number | undefined
+  private previousSpaceSize: number | undefined
+
+  public create (): void {
+    const { device, store } = this
+    if (!store.pointsTextureSize) return
+
+    this.levels = Math.log2(store.adjustedSpaceSize)
+
+    // Allocate quadtree levels
+    for (let level = 0; level < this.levels; level += 1) {
+      const levelTextureSize = Math.pow(2, level + 1)
+      const existingTarget = this.levelTargets.get(level)
+
+      if (
+        existingTarget &&
+        existingTarget.texture.width === levelTextureSize &&
+        existingTarget.texture.height === levelTextureSize
+      ) {
+        // Clear existing texture data to zero
+        existingTarget.texture.copyImageData({
+          data: new Float32Array(levelTextureSize * levelTextureSize * 4).fill(0),
+          bytesPerRow: levelTextureSize,
+          mipLevel: 0,
+          x: 0,
+          y: 0,
+        })
+        continue
+      }
+
+      // Destroy old resources if size changed
+      if (existingTarget) {
+        if (!existingTarget.texture.destroyed) existingTarget.texture.destroy()
+        if (!existingTarget.fbo.destroyed) existingTarget.fbo.destroy()
+      }
+
+      const texture = device.createTexture({
+        width: levelTextureSize,
+        height: levelTextureSize,
+        format: 'rgba32float',
+        usage: Texture.SAMPLE | Texture.RENDER | Texture.COPY_DST,
+      })
+      texture.copyImageData({
+        data: new Float32Array(levelTextureSize * levelTextureSize * 4).fill(0),
+        bytesPerRow: levelTextureSize,
+        mipLevel: 0,
+        x: 0,
+        y: 0,
+      })
+      const fbo = device.createFramebuffer({
+        width: levelTextureSize,
+        height: levelTextureSize,
+        colorAttachments: [texture],
+      })
+      this.levelTargets.set(level, { texture, fbo })
+    }
+
+    // Drop any stale higher-level buffers if space size shrank
+    for (const [level, target] of Array.from(this.levelTargets.entries())) {
+      if (level >= this.levels) {
+        if (!target.texture.destroyed) target.texture.destroy()
+        if (!target.fbo.destroyed) target.fbo.destroy()
+        this.levelTargets.delete(level)
+      }
+    }
+
+    // Random jitter texture to prevent sticking
+    const totalPixels = store.pointsTextureSize * store.pointsTextureSize
+    const randomValuesState = new Float32Array(totalPixels * 4)
+    for (let i = 0; i < totalPixels; ++i) {
+      randomValuesState[i * 4] = store.getRandomFloat(-1, 1) * 0.00001
+      randomValuesState[i * 4 + 1] = store.getRandomFloat(-1, 1) * 0.00001
+    }
+
+    if (!this.randomValuesTexture || this.randomValuesTexture.destroyed) {
+      this.randomValuesTexture = device.createTexture({
+        width: store.pointsTextureSize,
+        height: store.pointsTextureSize,
+        format: 'rgba32float',
+        usage: Texture.SAMPLE | Texture.COPY_DST,
+      })
+    }
+    this.randomValuesTexture.copyImageData({
+      data: randomValuesState,
+      bytesPerRow: store.pointsTextureSize,
+      mipLevel: 0,
+      x: 0,
+      y: 0,
+    })
+
+    // Point index buffer
+    const indexData = createIndexesForBuffer(store.pointsTextureSize)
+    const requiredByteLength = indexData.byteLength
+    if (!this.pointIndices || this.pointIndices.byteLength !== requiredByteLength) {
+      this.pointIndices?.destroy()
+      this.pointIndices = device.createBuffer({
+        data: indexData,
+        usage: Buffer.VERTEX | Buffer.COPY_DST,
+      })
+    } else {
+      this.pointIndices.write(indexData)
+    }
+
+    this.previousPointsTextureSize = store.pointsTextureSize
+    this.previousSpaceSize = store.adjustedSpaceSize
+  }
+
+  public initPrograms (): void {
+    const { device, store, data, points } = this
+    if (!data.pointsNumber || !points || !store.pointsTextureSize) return
+
+    // Clear levels command (fullscreen quad)
+    if (!this.clearLevelsCommand) {
+      if (!this.clearLevelsVertexCoordBuffer) {
+        this.clearLevelsVertexCoordBuffer = device.createBuffer({
+          data: new Float32Array([-1, -1, 1, -1, -1, 1, 1, 1]),
+        })
+      }
+      this.clearLevelsCommand = new Model(device, {
+        fs: clearFrag,
+        vs: updateVert,
+        topology: 'triangle-strip',
+        vertexCount: 4,
+        attributes: {
+          vertexCoord: this.clearLevelsVertexCoordBuffer,
+        },
+        bufferLayout: [
+          { name: 'vertexCoord', format: 'float32x2' },
+        ],
+      })
+    }
+
+    // Calculate levels command (point list)
+    if (!this.calculateLevelsCommand) {
+      if (!this.calculateLevelsUniformStore) {
+        this.calculateLevelsUniformStore = new UniformStore({
+          calculateLevelsUniforms: {
+            uniformTypes: {
+              pointsTextureSize: 'f32',
+              levelTextureSize: 'f32',
+              cellSize: 'f32',
+            },
+            defaultUniforms: {
+              pointsTextureSize: store.pointsTextureSize,
+              levelTextureSize: 0,
+              cellSize: 0,
+            },
+          },
+        })
+      }
+
+      this.calculateLevelsCommand = new Model(device, {
+        fs: calculateLevelFrag,
+        vs: calculateLevelVert,
+        topology: 'point-list',
+        vertexCount: data.pointsNumber,
+        attributes: {
+          pointIndices: this.pointIndices!,
+        },
+        bufferLayout: [
+          { name: 'pointIndices', format: 'float32x2' },
+        ],
+        defines: {
+          USE_UNIFORM_BUFFERS: true,
+        },
+        bindings: {
+          calculateLevelsUniforms: this.calculateLevelsUniformStore.getManagedUniformBuffer(device, 'calculateLevelsUniforms'),
+          positionsTexture: points.previousPositionTexture!,
+        },
+        parameters: {
+          blend: true,
+          blendColorOperation: 'add',
+          blendColorSrcFactor: 'one',
+          blendColorDstFactor: 'one',
+          blendAlphaOperation: 'add',
+          blendAlphaSrcFactor: 'one',
+          blendAlphaDstFactor: 'one',
+          depthWriteEnabled: false,
+          depthCompare: 'always',
+        },
+      })
+    }
+
+    // Force command (fullscreen quad)
+    if (!this.forceCommand) {
+      if (!this.forceUniformStore) {
+        this.forceUniformStore = new UniformStore({
+          forceUniforms: {
+            uniformTypes: {
+              level: 'f32',
+              levels: 'f32',
+              levelTextureSize: 'f32',
+              alpha: 'f32',
+              repulsion: 'f32',
+              spaceSize: 'f32',
+              theta: 'f32',
+            },
+            defaultUniforms: {
+              level: 0,
+              levels: this.levels,
+              levelTextureSize: 0,
+              alpha: store.alpha,
+              repulsion: this.config.simulationRepulsion ?? 0,
+              spaceSize: store.adjustedSpaceSize ?? 0,
+              theta: this.config.simulationRepulsionTheta ?? 0,
+            },
+          },
+        })
+      }
+
+      if (!this.forceVertexCoordBuffer) {
+        this.forceVertexCoordBuffer = device.createBuffer({
+          data: new Float32Array([-1, -1, 1, -1, -1, 1, 1, 1]),
+        })
+      }
+
+      this.forceCommand = new Model(device, {
+        fs: forceFrag,
+        vs: updateVert,
+        topology: 'triangle-strip',
+        vertexCount: 4,
+        attributes: {
+          vertexCoord: this.forceVertexCoordBuffer,
+        },
+        bufferLayout: [
+          { name: 'vertexCoord', format: 'float32x2' },
+        ],
+        defines: {
+          USE_UNIFORM_BUFFERS: true,
+        },
+        bindings: {
+          forceUniforms: this.forceUniformStore.getManagedUniformBuffer(device, 'forceUniforms'),
+          positionsTexture: points.previousPositionTexture!,
+        },
+        parameters: {
+          blend: true,
+          blendColorOperation: 'add',
+          blendColorSrcFactor: 'one',
+          blendColorDstFactor: 'one',
+          blendAlphaOperation: 'add',
+          blendAlphaSrcFactor: 'one',
+          blendAlphaDstFactor: 'one',
+          depthWriteEnabled: false,
+          depthCompare: 'always',
+        },
+      })
+    }
+
+    // Force-from-centermass command (fullscreen quad)
+    if (!this.forceFromItsOwnCentermassCommand) {
+      if (!this.forceCenterUniformStore) {
+        this.forceCenterUniformStore = new UniformStore({
+          forceCenterUniforms: {
+            uniformTypes: {
+              levelTextureSize: 'f32',
+              alpha: 'f32',
+              repulsion: 'f32',
+            },
+            defaultUniforms: {
+              levelTextureSize: 0,
+              alpha: store.alpha,
+              repulsion: this.config.simulationRepulsion ?? 0,
+            },
+          },
+        })
+      }
+
+      if (!this.forceVertexCoordBuffer) {
+        this.forceVertexCoordBuffer = device.createBuffer({
+          data: new Float32Array([-1, -1, 1, -1, -1, 1, 1, 1]),
+        })
+      }
+
+      this.forceFromItsOwnCentermassCommand = new Model(device, {
+        fs: forceCenterFrag,
+        vs: updateVert,
+        topology: 'triangle-strip',
+        vertexCount: 4,
+        attributes: {
+          vertexCoord: this.forceVertexCoordBuffer,
+        },
+        bufferLayout: [
+          { name: 'vertexCoord', format: 'float32x2' },
+        ],
+        defines: {
+          USE_UNIFORM_BUFFERS: true,
+        },
+        bindings: {
+          forceCenterUniforms: this.forceCenterUniformStore.getManagedUniformBuffer(device, 'forceCenterUniforms'),
+          positionsTexture: points.previousPositionTexture!,
+          randomValues: this.randomValuesTexture!,
+        },
+        parameters: {
+          blend: true,
+          blendColorOperation: 'add',
+          blendColorSrcFactor: 'one',
+          blendColorDstFactor: 'one',
+          blendAlphaOperation: 'add',
+          blendAlphaSrcFactor: 'one',
+          blendAlphaDstFactor: 'one',
+          depthWriteEnabled: false,
+          depthCompare: 'always',
+        },
+      })
+    }
+  }
+
+  public run (renderPass?: RenderPass): void {
+    // Skip if sizes changed and create() wasn't called yet
+    if (this.store.pointsTextureSize !== this.previousPointsTextureSize || this.store.adjustedSpaceSize !== this.previousSpaceSize) {
+      return
+    }
+    this.drawLevels()
+    this.drawForces(renderPass)
+  }
+
+  public destroy (): void {
+    this.calculateLevelsUniformStore?.destroy()
+    this.calculateLevelsUniformStore = undefined
+    this.forceUniformStore?.destroy()
+    this.forceUniformStore = undefined
+    this.forceCenterUniformStore?.destroy()
+    this.forceCenterUniformStore = undefined
+
+    this.clearLevelsCommand?.destroy()
+    this.clearLevelsCommand = undefined
+    this.calculateLevelsCommand?.destroy()
+    this.calculateLevelsCommand = undefined
+    this.forceCommand?.destroy()
+    this.forceCommand = undefined
+    this.forceFromItsOwnCentermassCommand?.destroy()
+    this.forceFromItsOwnCentermassCommand = undefined
+
+    this.pointIndices?.destroy()
+    this.pointIndices = undefined
+
+    if (this.randomValuesTexture && !this.randomValuesTexture.destroyed) {
+      this.randomValuesTexture.destroy()
+    }
+    this.randomValuesTexture = undefined
+
+    for (const target of this.levelTargets.values()) {
+      if (!target.texture.destroyed) target.texture.destroy()
+      if (!target.fbo.destroyed) target.fbo.destroy()
+    }
+    this.levelTargets.clear()
+
+    if (this.clearLevelsVertexCoordBuffer && !this.clearLevelsVertexCoordBuffer.destroyed) {
+      this.clearLevelsVertexCoordBuffer.destroy()
+    }
+    this.clearLevelsVertexCoordBuffer = undefined
+
+    if (this.forceVertexCoordBuffer && !this.forceVertexCoordBuffer.destroyed) {
+      this.forceVertexCoordBuffer.destroy()
+    }
+    this.forceVertexCoordBuffer = undefined
+  }
+
+  private drawLevels (): void {
+    const { device, store, data, points } = this
+    if (!points || !data.pointsNumber || !this.calculateLevelsCommand || !this.calculateLevelsUniformStore || !this.clearLevelsCommand) return
+    if (!points.previousPositionTexture || points.previousPositionTexture.destroyed) return
+
+    for (let level = 0; level < this.levels; level += 1) {
+      const target = this.levelTargets.get(level)
+      if (!target || target.fbo.destroyed || target.texture.destroyed) continue
+
+      const levelTextureSize = Math.pow(2, level + 1)
+      const cellSize = (store.adjustedSpaceSize ?? 0) / levelTextureSize
+
+      this.calculateLevelsUniformStore.setUniforms({
+        calculateLevelsUniforms: {
+          pointsTextureSize: store.pointsTextureSize ?? 0,
+          levelTextureSize,
+          cellSize,
+        },
+      })
+
+      this.calculateLevelsCommand.setVertexCount(data.pointsNumber)
+      this.calculateLevelsCommand.setBindings({
+        calculateLevelsUniforms: this.calculateLevelsUniformStore.getManagedUniformBuffer(device, 'calculateLevelsUniforms'),
+        positionsTexture: points.previousPositionTexture!,
+      })
+
+      const levelPass = device.beginRenderPass({
+        framebuffer: target.fbo,
+        clearColor: [0, 0, 0, 0],
+      })
+
+      this.clearLevelsCommand.draw(levelPass)
+      this.calculateLevelsCommand.draw(levelPass)
+
+      levelPass.end()
+    }
+  }
+
+  private drawForces (renderPass?: RenderPass): void {
+    const { device, store, points } = this
+    if (!points || !this.forceCommand || !this.forceUniformStore || !this.forceFromItsOwnCentermassCommand || !this.forceCenterUniformStore) return
+    if (!points.previousPositionTexture || points.previousPositionTexture.destroyed) return
+    if (!this.randomValuesTexture || this.randomValuesTexture.destroyed) return
+    if (!renderPass && (!points.velocityFbo || points.velocityFbo.destroyed)) return
+
+    // Update bindings shared for both force commands
+    this.forceCommand.setBindings({
+      forceUniforms: this.forceUniformStore.getManagedUniformBuffer(device, 'forceUniforms'),
+      positionsTexture: points.previousPositionTexture!,
+    })
+    this.forceFromItsOwnCentermassCommand.setBindings({
+      forceCenterUniforms: this.forceCenterUniformStore.getManagedUniformBuffer(device, 'forceCenterUniforms'),
+      positionsTexture: points.previousPositionTexture!,
+      randomValues: this.randomValuesTexture!,
+    })
+
+    const drawPass = renderPass ?? device.beginRenderPass({
+      framebuffer: points.velocityFbo,
+    })
+
+    for (let level = 0; level < this.levels; level += 1) {
+      const target = this.levelTargets.get(level)
+      if (!target || target.texture.destroyed) continue
+      const levelTextureSize = Math.pow(2, level + 1)
+
+      this.forceUniformStore.setUniforms({
+        forceUniforms: {
+          level,
+          levels: this.levels,
+          levelTextureSize,
+          alpha: store.alpha,
+          repulsion: this.config.simulationRepulsion ?? 0,
+          spaceSize: store.adjustedSpaceSize ?? 0,
+          theta: this.config.simulationRepulsionTheta ?? 0,
+        },
+      })
+
+      this.forceCommand.setBindings({
+        forceUniforms: this.forceUniformStore.getManagedUniformBuffer(device, 'forceUniforms'),
+        positionsTexture: points.previousPositionTexture!,
+        levelFbo: target.texture,
+      })
+
+      this.forceCommand.draw(drawPass)
+
+      // Only the deepest level uses the centermass fallback
+      if (level === this.levels - 1) {
+        this.forceCenterUniformStore.setUniforms({
+          forceCenterUniforms: {
+            levelTextureSize,
+            alpha: store.alpha,
+            repulsion: this.config.simulationRepulsion ?? 0,
+          },
+        })
+        this.forceFromItsOwnCentermassCommand.setBindings({
+          forceCenterUniforms: this.forceCenterUniformStore.getManagedUniformBuffer(device, 'forceCenterUniforms'),
+          positionsTexture: points.previousPositionTexture!,
+          randomValues: this.randomValuesTexture!,
+          levelFbo: target.texture,
+        })
+        this.forceFromItsOwnCentermassCommand.draw(drawPass)
+      }
+    }
+
+    if (!renderPass) {
+      drawPass.end()
+    }
+  }
+}

package/src/modules/ForceManyBody/quadtree-frag-shader.ts

@@ -0,0 +1,89 @@
+export function forceFrag (startLevel: number, maxLevels: number): string {
+  startLevel = Math.min(startLevel, maxLevels)
+  const delta = maxLevels - startLevel
+  const calcAdd = `
+    float dist = sqrt(l);
+    if (dist > 0.0) {
+      float c = alpha * repulsion * centermass.b;
+      addVelocity += calcAdd(vec2(x, y), l, c);
+      addVelocity += addVelocity * random.rg;
+    }
+  `
+  function quad (level: number): string {
+    if (level >= maxLevels) {
+      return calcAdd
+    } else {
+      const groupSize = Math.pow(2, level + 1)
+
+      const iEnding = new Array(level + 1 - delta).fill(0).map((_, l) => `pow(2.0, ${level - (l + delta)}.0) * i${l + delta}`).join('+')
+      const jEnding = new Array(level + 1 - delta).fill(0).map((_, l) => `pow(2.0, ${level - (l + delta)}.0) * j${l + delta}`).join('+')
+
+      return `
+      for (float ij${level} = 0.0; ij${level} < 4.0; ij${level} += 1.0) {
+        float i${level} = 0.0;
+        float j${level} = 0.0;
+        if (ij${level} == 1.0 || ij${level} == 3.0) i${level} = 1.0;
+        if (ij${level} == 2.0 || ij${level} == 3.0) j${level} = 1.0;
+        float i = pow(2.0, ${startLevel}.0) * n / width${level + 1} + ${iEnding};
+        float j = pow(2.0, ${startLevel}.0) * m / width${level + 1} + ${jEnding};
+        float groupPosX = (i + 0.5) / ${groupSize}.0;
+        float groupPosY = (j + 0.5) / ${groupSize}.0;
+        
+        vec4 centermass = texture(level[${level}], vec2(groupPosX, groupPosY));
+        if (centermass.r > 0.0 && centermass.g > 0.0 && centermass.b > 0.0) {
+          float x = centermass.r / centermass.b - pointPosition.r;
+          float y = centermass.g / centermass.b - pointPosition.g;
+          float l = x * x + y * y;
+          if ((width${level + 1} * width${level + 1}) / theta < l) {
+            ${calcAdd}
+          } else {
+            ${quad(level + 1)}
+          }
+        }
+      }
+      `
+    }
+  }
+  return `#version 300 es
+precision highp float;
+
+uniform sampler2D positionsTexture;
+uniform sampler2D randomValues;
+uniform float spaceSize;
+uniform float repulsion;
+uniform float theta;
+uniform float alpha;
+uniform sampler2D level[${maxLevels}];
+in vec2 textureCoords;
+out vec4 fragColor;
+
+vec2 calcAdd(vec2 xy, float l, float c) {
+  float distanceMin2 = 1.0;
+  if (l < distanceMin2) l = sqrt(distanceMin2 * l);
+  float add = c / l;
+  return add * xy;
+}
+
+void main() {
+  vec4 pointPosition = texture(positionsTexture, textureCoords);
+  vec4 random = texture(randomValues, textureCoords);
+
+  float width0 = spaceSize;
+
+  vec2 velocity = vec2(0.0);
+  vec2 addVelocity = vec2(0.0);
+
+  ${new Array(maxLevels).fill(0).map((_, i) => `float width${i + 1} = width${i} / 2.0;`).join('\n')}
+
+  for (float n = 0.0; n < pow(2.0, ${delta}.0); n += 1.0) {
+    for (float m = 0.0; m < pow(2.0, ${delta}.0); m += 1.0) {
+      ${quad(delta)}
+    }
+  }
+
+  velocity -= addVelocity;
+
+  fragColor = vec4(velocity, 0.0, 0.0);
+}
+`
+}

package/src/modules/ForceManyBodyQuadtree/calculate-level.frag

@@ -0,0 +1,9 @@
+#ifdef GL_ES
+precision highp float;
+#endif
+
+varying vec4 rgba;
+
+void main() {
+  gl_FragColor = rgba;
+}

package/src/modules/ForceManyBodyQuadtree/calculate-level.vert

@@ -0,0 +1,25 @@
+#ifdef GL_ES
+precision highp float;
+#endif
+
+uniform sampler2D positionsTexture;
+uniform float pointsTextureSize;
+uniform float levelTextureSize;
+uniform float cellSize;
+
+attribute vec2 pointIndices;
+
+varying vec4 rgba;
+
+void main() {
+  vec4 pointPosition = texture2D(positionsTexture, pointIndices / pointsTextureSize);
+  rgba = vec4(pointPosition.rg, 1.0, 0.0);
+
+  float n = floor(pointPosition.x / cellSize);
+  float m = floor(pointPosition.y / cellSize);
+  
+  vec2 levelPosition = 2.0 * (vec2(n, m) + 0.5) / levelTextureSize - 1.0;
+
+  gl_Position = vec4(levelPosition, 0.0, 1.0);
+  gl_PointSize = 1.0;
+}