gaussian density tweaks

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

@@ -22,7 +22,7 @@ uniform float uCurrentSlice;
 void main() {
     radius = aRadius;
     float scale = max(uBboxSize.z, max(uBboxSize.x, uBboxSize.y));
-    gl_PointSize = (radius / scale) * max(uGridDim.x, uGridDim.y) * 3.0;
+    gl_PointSize = (radius / scale) * max(uGridDim.x, uGridDim.y) * 6.0;
     position = (aPosition - uBboxMin) / uBboxSize;
     gl_Position = vec4(position * 2.0 - 1.0, 1.0);
 }

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

@@ -20,7 +20,17 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
     const v = Vec3.zero()
     const p = Vec3.zero()
 
-    const pad = (radiusOffset + 3) * 3 // TODO calculate max radius
+    let maxRadius = 0
+    for (let i = 0; i < n; ++i) {
+        const r = radius(OrderedSet.getAt(indices, i)) + radiusOffset
+        if (maxRadius < r) maxRadius = r
+
+        if (i % 10000 === 0 && ctx.shouldUpdate) {
+            await ctx.update({ message: 'calculating max radius', current: i, max: n })
+        }
+    }
+
+    const pad = maxRadius * 2 + resolution
     const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad))
     const extent = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min)
     const min = expandedBox.min
@@ -43,7 +53,7 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
 
     const alpha = smoothness
 
-    const _r2 = (radiusOffset + 1.4 * 2)
+    const _r2 = maxRadius * 2
     const _radius2 = Vec3.create(_r2, _r2, _r2)
     Vec3.mul(_radius2, _radius2, delta)
     const updateChunk = Math.ceil(10000 / (_radius2[0] * _radius2[1] * _radius2[2]))

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

@@ -6,7 +6,7 @@
  */
 
 import { RuntimeContext } from 'mol-task'
-import { PositionData } from '../common'
+import { PositionData, DensityData } from '../common'
 import { Box3D } from '../../geometry'
 import { GaussianDensityProps, getDelta } from '../gaussian-density'
 import { OrderedSet } from 'mol-data/int'
@@ -18,7 +18,7 @@ import { createRenderable, createGaussianDensityRenderObject } from 'mol-gl/rend
 import { createRenderTarget } from 'mol-gl/webgl/render-target'
 import { Context, createContext } from 'mol-gl/webgl/context';
 
-export async function GaussianDensityGPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps) { // }: Promise<DensityData> {
+export async function GaussianDensityGPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
     const { resolution, radiusOffset, smoothness, readSlices } = props
 
     const { indices, x, y, z } = position
@@ -27,19 +27,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     const positions = new Float32Array(n * 3)
     const radii = new Float32Array(n)
 
-    const pad = (radiusOffset + 3) * 3 // TODO calculate max radius
-    const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad));
-    const extent = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min)
-
-    const delta = getDelta(Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad)), resolution)
-    const dim = Vec3.zero()
-    Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
-    console.log('grid dim', dim)
-
-    const _r2 = (radiusOffset + 1.4 * 2)
-    const _radius2 = Vec3.create(_r2, _r2, _r2)
-    Vec3.mul(_radius2, _radius2, delta)
-    const updateChunk = Math.ceil(10000 / (_radius2[0] * _radius2[1] * _radius2[2]))
+    let maxRadius = 0
 
     for (let i = 0; i < n; ++i) {
         const j = OrderedSet.getAt(indices, i);
@@ -47,13 +35,29 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         positions[i * 3] = x[j]
         positions[i * 3 + 1] = y[j]
         positions[i * 3 + 2] = z[j]
-        radii[i] = radius(j) + radiusOffset
+        const r = radius(j) + radiusOffset
+        if (maxRadius < r) maxRadius = r
+        radii[i] = r
 
         if (i % 10000 === 0 && ctx.shouldUpdate) {
             await ctx.update({ message: 'preparing density data', current: i, max: n })
         }
     }
 
+    const pad = maxRadius * 2 + resolution
+    const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad));
+    const extent = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min)
+
+    const delta = getDelta(Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad)), resolution)
+    const dim = Vec3.zero()
+    Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
+    console.log('grid dim', dim)
+
+    const _r2 = maxRadius * 2
+    const _radius2 = Vec3.create(_r2, _r2, _r2)
+    Vec3.mul(_radius2, _radius2, delta)
+    const updateChunk = Math.ceil(10000 / (_radius2[0] * _radius2[1] * _radius2[2]))
+
     //
 
     const values: GaussianDensityValues = {
@@ -86,7 +90,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     //
 
     // TODO fallback to lower resolution when texture size is not large enough
-    const maxTexSize = webgl.maxTextureSize
+    const maxTexSize = 1024 // webgl.maxTextureSize
     let fboTexDimX = 0
     let fboTexDimY = dim[1]
     let fboTexRows = 1
@@ -100,6 +104,8 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         fboTexDimX = dim[0] * dim[2]
     }
 
+    console.log('dim', dim, 'cols', fboTexCols, 'rows', fboTexRows)
+
     //
 
     const space = Tensor.Space(dim, [2, 1, 0], Float32Array)
@@ -124,7 +130,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     gl.cullFace(gl.BACK)
 
     gl.depthMask(true)
-    gl.clearColor(0, 0, 0, 1)
+    gl.clearColor(0, 0, 0, 0)
     gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT)
     gl.depthMask(false)
 
@@ -173,6 +179,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     if (!readSlices) {
         console.time('gpu gaussian density full')
         renderTarget.getBuffer()
+        const { array } = renderTarget.image
         let idx = 0
         let tmpCol = 0
         let tmpRow = 0
@@ -183,7 +190,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
             }
             for (let iy = 0; iy < dim[1]; ++iy) {
                 for (let ix = 0; ix < dim[0]; ++ix) {
-                    data[idx] = renderTarget.image.array[4 * (tmpCol * dim[0] + (iy + tmpRow) * fboTexDimX + ix)] / 255
+                    data[idx] = array[4 * (tmpCol * dim[0] + (iy + tmpRow) * fboTexDimX + ix)] / 255
                     idx++
                 }
             }
@@ -198,7 +205,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))
     Mat4.setTranslation(transform, expandedBox.min)
 
-    return { field, idField, transform }
+    return { field, idField, transform, renderTarget, bbox: expandedBox, gridDimension: dim }
 }
 
 let webglContext: Context
@@ -206,9 +213,9 @@ function getWebGLContext() {
     if (webglContext) return webglContext
     const canvas = document.createElement('canvas')
     const gl = canvas.getContext('webgl', {
-        alpha: false,
-        antialias: true,
-        depth: true,
+        alpha: true,
+        antialias: false,
+        depth: false,
         preserveDrawingBuffer: true
     })
     if (!gl) throw new Error('Could not create a WebGL rendering context')

src/mol-model/structure/structure/unit/gaussian-density.ts

@@ -14,7 +14,7 @@ import { NumberParam, paramDefaultValues, BooleanParam } from 'mol-view/paramete
 export const GaussianDensityParams = {
     resolution: NumberParam('Resolution', '', 1, 0.1, 10, 0.1),
     radiusOffset: NumberParam('Radius Offset', '', 0, 0, 10, 0.1),
-    smoothness: NumberParam('Smoothness', '', 1.5, 0, 4, 0.1),
+    smoothness: NumberParam('Smoothness', '', 1.5, 0.5, 2.5, 0.1),
     useGpu: BooleanParam('Use GPU', '', true),
     readSlices: BooleanParam('Read Slices', '', false),
     ignoreCache: BooleanParam('Ignore Cache', '', false),