wip, calculate gaussian density group ids on gpu

src/mol-geo/representation/structure/visual/gaussian-surface-mesh.ts

@@ -29,44 +29,44 @@ async function createGaussianSurfaceMesh(ctx: RuntimeContext, unit: Unit, struct
 
     Mesh.transformImmediate(surface, transform)
 
-    if (props.useGpu) {
-        console.time('find max element radius')
-        const { elements } = unit
-        const n = OrderedSet.size(elements)
-        const l = StructureElement.create(unit)
-        const sizeTheme = SizeTheme({ name: 'physical' })
-        const radius = (index: number) => {
-            l.element = index as ElementIndex
-            return sizeTheme.size(l)
-        }
-        let maxRadius = 0
-        for (let i = 0; i < n; ++i) {
-            const r = radius(OrderedSet.getAt(elements, i)) + radiusOffset
-            if (maxRadius < r) maxRadius = r
-        }
-        console.timeEnd('find max element radius')
+    // if (props.useGpu) {
+    //     console.time('find max element radius')
+    //     const { elements } = unit
+    //     const n = OrderedSet.size(elements)
+    //     const l = StructureElement.create(unit)
+    //     const sizeTheme = SizeTheme({ name: 'physical' })
+    //     const radius = (index: number) => {
+    //         l.element = index as ElementIndex
+    //         return sizeTheme.size(l)
+    //     }
+    //     let maxRadius = 0
+    //     for (let i = 0; i < n; ++i) {
+    //         const r = radius(OrderedSet.getAt(elements, i)) + radiusOffset
+    //         if (maxRadius < r) maxRadius = r
+    //     }
+    //     console.timeEnd('find max element radius')
 
-        console.time('find closest element for vertices')
-        const { lookup3d } = unit
+    //     console.time('find closest element for vertices')
+    //     const { lookup3d } = unit
 
-        const { vertexCount, vertexBuffer, groupBuffer } = surface
-        const vertices = vertexBuffer.ref.value
-        const groups = groupBuffer.ref.value
-        for (let i = 0; i < vertexCount; ++i) {
-            const r = lookup3d.find(vertices[i * 3], vertices[i * 3 + 1], vertices[i * 3 + 2], maxRadius * 2)
-            let minDsq = Infinity
-            let group = 0
-            for (let j = 0, jl = r.count; j < jl; ++j) {
-                const dSq = r.squaredDistances[j]
-                if (dSq < minDsq) {
-                    minDsq = dSq
-                    group = r.indices[j]
-                }
-            }
-            groups[i] = group
-        }
-        console.timeEnd('find closest element for vertices')
-    }
+    //     const { vertexCount, vertexBuffer, groupBuffer } = surface
+    //     const vertices = vertexBuffer.ref.value
+    //     const groups = groupBuffer.ref.value
+    //     for (let i = 0; i < vertexCount; ++i) {
+    //         const r = lookup3d.find(vertices[i * 3], vertices[i * 3 + 1], vertices[i * 3 + 2], maxRadius * 2)
+    //         let minDsq = Infinity
+    //         let group = 0
+    //         for (let j = 0, jl = r.count; j < jl; ++j) {
+    //             const dSq = r.squaredDistances[j]
+    //             if (dSq < minDsq) {
+    //                 minDsq = dSq
+    //                 group = r.indices[j]
+    //             }
+    //         }
+    //         groups[i] = group
+    //     }
+    //     console.timeEnd('find closest element for vertices')
+    // }
 
     Mesh.computeNormalsImmediate(surface)
     Mesh.uniformTriangleGroup(surface)

src/mol-gl/renderable/gaussian-density.ts

@@ -7,7 +7,7 @@
 import { Renderable, RenderableState, createRenderable } from '../renderable'
 import { Context } from '../webgl/context';
 import { createRenderItem } from '../webgl/render-item';
-import { AttributeSpec, Values, UniformSpec, ValueSpec, DefineSpec } from './schema';
+import { AttributeSpec, Values, UniformSpec, ValueSpec, DefineSpec, TextureSpec } from './schema';
 import { GaussianDensityShaderCode } from '../shader-code';
 
 export const GaussianDensitySchema = {
@@ -16,6 +16,7 @@ export const GaussianDensitySchema = {
 
     aRadius: AttributeSpec('float32', 1, 0),
     aPosition: AttributeSpec('float32', 3, 0),
+    aGroup: AttributeSpec('float32', 1, 0),
 
     uCurrentSlice: UniformSpec('f'),
     uCurrentX: UniformSpec('f'),
@@ -25,6 +26,9 @@ export const GaussianDensitySchema = {
     uBboxSize: UniformSpec('v3'),
     uGridDim: UniformSpec('v3'),
     uAlpha: UniformSpec('f'),
+    tMinDistanceTex: TextureSpec('texture2d', 'rgba', 'ubyte', 'nearest'),
+
+    dCalcType: DefineSpec('string', ['density', 'minDistance', 'groupId']),
 }
 export type GaussianDensitySchema = typeof GaussianDensitySchema
 export type GaussianDensityValues = Values<GaussianDensitySchema>

src/mol-gl/shader/gaussian-density.frag

@@ -9,6 +9,13 @@ precision highp float;
 
 varying vec3 vPosition;
 varying float vRadius;
+#if defined(dCalcType_groupId)
+    precision highp sampler3D;
+    uniform sampler3D tMinDistanceTex;
+    varying float vGroup;
+#endif
+
+#pragma glslify: encodeIdRGBA = require(./utils/encode-id-rgba.glsl)
 
 uniform vec3 uBboxSize;
 uniform vec3 uBboxMin;
@@ -19,15 +26,27 @@ uniform float uCurrentX;
 uniform float uCurrentY;
 uniform float uAlpha;
 
-vec4 calc(float x, float y, float z, float radiusSq) {
-    vec3 fragPos = vec3(x, y, z) / uGridDim;
-    float dist = distance(fragPos * uBboxSize, vPosition * uBboxSize);
-    float density = exp(-uAlpha * ((dist * dist) / radiusSq));
-    return vec4(density);
-}
+// encode distance logarithmically with given maxDistance
+const float maxDistance = 10000.0;
+const float distLogFactor = log(maxDistance + 1.0);
+float encodeDistLog(float dist) { return log(dist + 1.0) / distLogFactor; }
+float decodeDistLog(float logDist) { return exp(logDist * distLogFactor) - 1.0; }
 
 void main() {
-    float radiusSq = vRadius * vRadius;
     vec2 v = gl_FragCoord.xy - vec2(uCurrentX, uCurrentY) - 0.5;
-    gl_FragColor = calc(v.x, v.y, uCurrentSlice, radiusSq);
+    vec3 fragPos = vec3(v.x, v.y, uCurrentSlice) / uGridDim;
+    float dist = distance(fragPos * uBboxSize, vPosition * uBboxSize);
+
+    #if defined(dCalcType_density)
+        float radiusSq = vRadius * vRadius;
+        float density = exp(-uAlpha * ((dist * dist) / radiusSq));
+        gl_FragColor = vec4(density);
+    #elif defined(dCalcType_minDistance)
+        gl_FragColor.r = 1.0 - encodeDistLog(dist);
+    #elif defined(dCalcType_groupId)
+        float minDistance = decodeDistLog(1.0 - texture(tMinDistanceTex, fragPos).r);
+        if (dist > minDistance + log(minDistance) / 2.0)
+            discard;
+        gl_FragColor = encodeIdRGBA(vGroup);
+    #endif
 }

src/mol-gl/shader/gaussian-density.vert

@@ -13,6 +13,11 @@ attribute float aRadius;
 varying vec3 vPosition;
 varying float vRadius;
 
+#if defined(dCalcType_groupId)
+    attribute float aGroup;
+    varying float vGroup;
+#endif
+
 uniform vec3 uBboxSize;
 uniform vec3 uBboxMin;
 uniform vec3 uBboxMax;
@@ -21,6 +26,9 @@ uniform float uCurrentSlice;
 
 void main() {
     vRadius = aRadius;
+    #if defined(dCalcType_groupId)
+        vGroup = aGroup;
+    #endif
     float scale = max(uBboxSize.z, max(uBboxSize.x, uBboxSize.y));
     gl_PointSize = (vRadius / scale) * max(uGridDim.x, uGridDim.y) * 6.0;
     vPosition = (aPosition - uBboxMin) / uBboxSize;

src/mol-gl/shader/utils/decode-float-rgba.glsl

@@ -1,5 +1,5 @@
 float decodeFloatRGBA(const in vec4 rgba) {
-    return dot(rgba, vec4(1.0, 1/255.0, 1/65025.0, 1/16581375.0));
+    return dot(rgba, vec4(1.0, 1.0/255.0, 1.0/65025.0, 1.0/16581375.0));
 }
 
 #pragma glslify: export(decodeFloatRGBA)

src/mol-math/geometry/gaussian-density/gpu.ts

@@ -19,18 +19,18 @@ import { Context, createContext, getGLContext } from 'mol-gl/webgl/context';
 import { createFramebuffer } from 'mol-gl/webgl/framebuffer';
 import { createTexture, Texture } from 'mol-gl/webgl/texture';
 import { GLRenderingContext } from 'mol-gl/webgl/compat';
+import { decodeIdRGBA } from 'mol-geo/geometry/picking';
 
 export async function GaussianDensityGPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number, props: GaussianDensityProps): Promise<DensityData> {
     const webgl = defaults(props.webgl, getWebGLContext())
 
     const { transform, texture, gridDimension } = await GaussianDensityTexture(ctx, webgl, position, box, radius, props)
 
-    const field = webgl.maxDrawBuffers > 0 ?
+    const { field, idField } = webgl.isWebGL2 ?
         fieldFromTexture3d(webgl, texture, gridDimension) :
         fieldFromTexture2d(webgl, texture, gridDimension)
 
-    const idData = field.space.create()
-    const idField = Tensor.create(field.space, idData)
+    console.log(idField)
 
     return { field, idField, transform }
 }
@@ -54,9 +54,11 @@ export async function GaussianDensityTexture(ctx: RuntimeContext, webgl: Context
 async function GaussianDensityTexture2d(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number, props: GaussianDensityProps, texture?: Texture) {
     const { smoothness } = props
 
-    const { drawCount, positions, radii, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
+    const { drawCount, positions, radii, groups, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
     const [ dx, dy, dz ] = dim
-    const renderObject = getGaussianDensityRenderObject(webgl, drawCount, positions, radii, expandedBox, dim, smoothness)
+    const minDistanceTexture = getMinDistanceTexture(webgl, dx, dy)
+
+    const renderObject = getGaussianDensityRenderObject(webgl, drawCount, positions, radii, groups, minDistanceTexture, expandedBox, dim, smoothness)
     const renderable = createRenderable(webgl, renderObject)
 
     //
@@ -122,9 +124,12 @@ async function GaussianDensityTexture2d(ctx: RuntimeContext, webgl: Context, pos
 async function GaussianDensityTexture3d(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number, props: GaussianDensityProps, texture?: Texture) {
     const { smoothness } = props
 
-    const { drawCount, positions, radii, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
+    const { drawCount, positions, radii, groups, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
     const [ dx, dy, dz ] = dim
-    const renderObject = getGaussianDensityRenderObject(webgl, drawCount, positions, radii, expandedBox, dim, smoothness)
+    const minDistanceTexture = createTexture(webgl, 'volume-uint8', 'rgba', 'ubyte', 'nearest')
+    minDistanceTexture.define(dx, dy, dz)
+
+    const renderObject = getGaussianDensityRenderObject(webgl, drawCount, positions, radii, groups, minDistanceTexture, expandedBox, dim, smoothness)
     const renderable = createRenderable(webgl, renderObject)
 
     //
@@ -143,8 +148,36 @@ async function GaussianDensityTexture3d(ctx: RuntimeContext, webgl: Context, pos
     }
     texture.define(dx, dy, dz)
 
+    ValueCell.update(renderable.values.dCalcType, 'minDistance')
+    renderable.update()
+    const programMinDistance = renderable.getProgram('draw')
+    programMinDistance.use()
+    gl.blendFunc(gl.ONE, gl.ONE)
+    gl.blendEquation(gl.MAX)
+    for (let i = 0; i < dz; ++i) {
+        ValueCell.update(uCurrentSlice, i)
+        minDistanceTexture.attachFramebuffer(framebuffer, 0, i)
+        renderable.render('draw')
+    }
+
+    ValueCell.update(renderable.values.dCalcType, 'density')
+    renderable.update()
     const programDensity = renderable.getProgram('draw')
     programDensity.use()
+    gl.blendFunc(gl.ONE, gl.ONE)
+    gl.blendEquation(gl.FUNC_ADD)
+    for (let i = 0; i < dz; ++i) {
+        ValueCell.update(uCurrentSlice, i)
+        texture.attachFramebuffer(framebuffer, 0, i)
+        renderable.render('draw')
+    }
+
+    ValueCell.update(renderable.values.dCalcType, 'groupId')
+    renderable.update()
+    const programGroupId = renderable.getProgram('draw')
+    programGroupId.use()
+    gl.blendFuncSeparate(gl.ONE, gl.ZERO, gl.ZERO, gl.ONE)
+    gl.blendEquation(gl.FUNC_ADD)
     for (let i = 0; i < dz; ++i) {
         ValueCell.update(uCurrentSlice, i)
         texture.attachFramebuffer(framebuffer, 0, i)
@@ -184,6 +217,7 @@ async function prepareGaussianDensityData(ctx: RuntimeContext, position: Positio
 
     const positions = new Float32Array(n * 3)
     const radii = new Float32Array(n)
+    const groups = new Float32Array(n)
 
     let maxRadius = 0
 
@@ -196,6 +230,7 @@ async function prepareGaussianDensityData(ctx: RuntimeContext, position: Positio
         const r = radius(j) + radiusOffset
         if (maxRadius < r) maxRadius = r
         radii[i] = r
+        groups[i] = i
 
         if (i % 10000 === 0 && ctx.shouldUpdate) {
             await ctx.update({ message: 'preparing density data', current: i, max: n })
@@ -211,10 +246,16 @@ async function prepareGaussianDensityData(ctx: RuntimeContext, position: Positio
     Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
     console.log('grid dim gpu', dim)
 
-    return { drawCount: n, positions, radii, delta, expandedBox, dim }
+    return { drawCount: n, positions, radii, groups, delta, expandedBox, dim }
 }
 
-function getGaussianDensityRenderObject(webgl: Context, drawCount: number, positions: Float32Array, radii: Float32Array, box: Box3D, dimensions: Vec3, smoothness: number) {
+function getMinDistanceTexture(webgl: Context, width: number, height: number) {
+    const minDistanceTexture = createTexture(webgl, 'image-uint8', 'rgba', 'ubyte', 'nearest')
+    minDistanceTexture.define(width, height)
+    return minDistanceTexture
+}
+
+function getGaussianDensityRenderObject(webgl: Context, drawCount: number, positions: Float32Array, radii: Float32Array, groups: Float32Array, minDistanceTexture: Texture, box: Box3D, dimensions: Vec3, smoothness: number) {
     const extent = Vec3.sub(Vec3.zero(), box.max, box.min)
 
     const values: GaussianDensityValues = {
@@ -223,6 +264,7 @@ function getGaussianDensityRenderObject(webgl: Context, drawCount: number, posit
 
         aRadius: ValueCell.create(radii),
         aPosition: ValueCell.create(positions),
+        aGroup: ValueCell.create(groups),
 
         uCurrentSlice: ValueCell.create(0),
         uCurrentX: ValueCell.create(0),
@@ -232,6 +274,9 @@ function getGaussianDensityRenderObject(webgl: Context, drawCount: number, posit
         uBboxSize: ValueCell.create(extent),
         uGridDim: ValueCell.create(dimensions),
         uAlpha: ValueCell.create(smoothness),
+        tMinDistanceTex: ValueCell.create(minDistanceTexture),
+
+        dCalcType: ValueCell.create('density'),
     }
     const state: RenderableState = {
         visible: true,
@@ -263,6 +308,8 @@ function fieldFromTexture2d(ctx: Context, texture: Texture, dim: Vec3) {
     const space = Tensor.Space(dim, [2, 1, 0], Float32Array)
     const data = space.create()
     const field = Tensor.create(space, data)
+    const idData = space.create()
+    const idField = Tensor.create(space, idData)
 
     const image = new Uint8Array(width * height * 4)
 
@@ -292,7 +339,7 @@ function fieldFromTexture2d(ctx: Context, texture: Texture, dim: Vec3) {
     framebuffer.destroy()
     console.timeEnd('fieldFromTexture2d')
 
-    return field
+    return { field, idField }
 }
 
 function fieldFromTexture3d(ctx: Context, texture: Texture, dim: Vec3) {
@@ -303,6 +350,8 @@ function fieldFromTexture3d(ctx: Context, texture: Texture, dim: Vec3) {
     const space = Tensor.Space(dim, [2, 1, 0], Float32Array)
     const data = space.create()
     const field = Tensor.create(space, data)
+    const idData = space.create()
+    const idField = Tensor.create(space, idData)
 
     const slice = new Uint8Array(width * height * 4)
 
@@ -315,7 +364,9 @@ function fieldFromTexture3d(ctx: Context, texture: Texture, dim: Vec3) {
         gl.readPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, slice)
         for (let iy = 0; iy < height; ++iy) {
             for (let ix = 0; ix < width; ++ix) {
-                data[j] = slice[4 * (iy * width + ix) + 3] / 255
+                const idx = 4 * (iy * width + ix)
+                data[j] = slice[idx + 3] / 255
+                idData[j] = decodeIdRGBA(slice[idx], slice[idx + 1], slice[idx + 2])
                 ++j
             }
         }
@@ -324,5 +375,5 @@ function fieldFromTexture3d(ctx: Context, texture: Texture, dim: Vec3) {
     framebuffer.destroy()
     console.timeEnd('fieldFromTexture3d')
 
-    return field
+    return { field, idField }
 }