wip, gpu gaussian density computation

src/apps/canvas/component/representation.tsx

@@ -10,6 +10,7 @@ import { App } from '../app';
 import { Params } from 'mol-view/parameter';
 import { Representation } from 'mol-geo/representation';
 import { ParametersComponent } from 'mol-app/component/parameters';
+import { Progress } from 'mol-task';
 
 export interface RepresentationComponentProps {
     app: App
@@ -24,7 +25,9 @@ export interface RepresentationComponentState {
 export class RepresentationComponent extends React.Component<RepresentationComponentProps, RepresentationComponentState> {
 
     async onChange(k: string, v: any) {
-        await this.props.app.runTask(this.props.repr.createOrUpdate({ [k]: v }).run(), 'Representation Update')
+        await this.props.app.runTask(this.props.repr.createOrUpdate({ [k]: v }).run(
+            progress => console.log(Progress.format(progress))
+        ), 'Representation Update')
         this.props.viewer.add(this.props.repr)
         this.props.viewer.requestDraw(true)
     }

src/mol-geo/representation/structure/visual/gaussian-density-point.ts

@@ -68,6 +68,9 @@ export function GaussianDensityPointVisual(): UnitsVisual<GaussianDensityPointPr
             if (newProps.resolution !== currentProps.resolution) state.createGeometry = true
             if (newProps.radiusOffset !== currentProps.radiusOffset) state.createGeometry = true
             if (newProps.smoothness !== currentProps.smoothness) state.createGeometry = true
+            if (newProps.useGpu !== currentProps.useGpu) state.createGeometry = true
+            if (newProps.readSlices !== currentProps.readSlices) state.createGeometry = true
+            if (newProps.ignoreCache !== currentProps.ignoreCache) state.createGeometry = true
         }
     })
 }

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

@@ -50,6 +50,9 @@ export function GaussianSurfaceVisual(): UnitsVisual<GaussianSurfaceProps> {
             if (newProps.resolution !== currentProps.resolution) state.createGeometry = true
             if (newProps.radiusOffset !== currentProps.radiusOffset) state.createGeometry = true
             if (newProps.smoothness !== currentProps.smoothness) state.createGeometry = true
+            if (newProps.useGpu !== currentProps.useGpu) state.createGeometry = true
+            if (newProps.readSlices !== currentProps.readSlices) state.createGeometry = true
+            if (newProps.ignoreCache !== currentProps.ignoreCache) state.createGeometry = true
         }
     })
 }

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

@@ -52,6 +52,9 @@ export function GaussianWireframeVisual(): UnitsVisual<GaussianWireframeProps> {
             if (newProps.resolution !== currentProps.resolution) state.createGeometry = true
             if (newProps.radiusOffset !== currentProps.radiusOffset) state.createGeometry = true
             if (newProps.smoothness !== currentProps.smoothness) state.createGeometry = true
+            if (newProps.useGpu !== currentProps.useGpu) state.createGeometry = true
+            if (newProps.readSlices !== currentProps.readSlices) state.createGeometry = true
+            if (newProps.ignoreCache !== currentProps.ignoreCache) state.createGeometry = true
         }
     })
 }

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

@@ -24,6 +24,7 @@ export const GaussianDensitySchema = {
     uBboxMax: UniformSpec('v3'),
     uBboxSize: UniformSpec('v3'),
     uGridDim: UniformSpec('v3'),
+    uAlpha: UniformSpec('f'),
 }
 export type GaussianDensitySchema = typeof GaussianDensitySchema
 export type GaussianDensityValues = Values<GaussianDensitySchema>

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

@@ -17,6 +17,7 @@ uniform vec3 uGridDim;
 uniform float uCurrentSlice;
 uniform float uCurrentX;
 uniform float uCurrentY;
+uniform float uAlpha;
 
 void main() {
     vec3 tmpVec = gl_FragCoord.xyz;
@@ -28,8 +29,6 @@ void main() {
         (uCurrentSlice) / uGridDim.z
     );
     float dist = length(fragPos * uBboxSize - position * uBboxSize);
-    float density = 1.0 - smoothstep( 0.0, radius * 2.0, dist);
+    float density = exp(-uAlpha * ((dist * dist) / (radius * radius)));
     gl_FragColor = vec4(1, 1, 1, density);
-    // density = 1.0 - clamp((dist - (radius + 1.4)) + 0.5, 0.0, 1.0);				
-    // gl_FragColor = vec4(vec3(density), 1.0);
 }

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) * 10.0; // * 50.0;
+    gl_PointSize = (radius / scale) * max(uGridDim.x, uGridDim.y) * 3.0;
     position = (aPosition - uBboxMin) / uBboxSize;
     gl_Position = vec4(position * 2.0 - 1.0, 1.0);
 }

src/mol-gl/util.ts

@@ -0,0 +1,25 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+function debugTexture(imageData: ImageData, scale = 1) {
+    const canvas = document.createElement('canvas')
+    canvas.width = imageData.width
+    canvas.height = imageData.height
+    const ctx = canvas.getContext('2d')
+    if (!ctx) throw new Error('Could not create canvas 2d context')
+    ctx.putImageData(imageData, 0, 0)
+    canvas.toBlob(imgBlob => {
+        const objectURL = window.URL.createObjectURL(imgBlob)
+        const img = document.createElement('img')
+        img.src = objectURL
+        img.style.width = imageData.width * scale + 'px'
+        img.style.height = imageData.height * scale + 'px'
+        img.style.position = 'absolute'
+        img.style.top = '0px'
+        img.style.left = '0px'
+        document.body.appendChild(img)
+    }, 'image/png')
+}

src/mol-gl/webgl/context.ts

@@ -80,6 +80,8 @@ export interface Context {
     instanceCount: number
     instancedDrawCount: number
 
+    readonly maxTextureSize: number
+
     unbindFramebuffer: () => void
     readPixels: (x: number, y: number, width: number, height: number, buffer: Uint8Array) => void
     destroy: () => void
@@ -106,6 +108,10 @@ export function createContext(gl: WebGLRenderingContext): Context {
     const shaderCache = createShaderCache()
     const programCache = createProgramCache()
 
+    const parameters = {
+        maxTextureSize: gl.getParameter(gl.MAX_TEXTURE_SIZE)
+    }
+
     return {
         gl,
         extensions: { angleInstancedArrays, standardDerivatives, oesElementIndexUint, oesVertexArrayObject },
@@ -124,6 +130,8 @@ export function createContext(gl: WebGLRenderingContext): Context {
         instanceCount: 0,
         instancedDrawCount: 0,
 
+        get maxTextureSize () { return parameters.maxTextureSize },
+
         unbindFramebuffer: () => unbindFramebuffer(gl),
         readPixels: (x: number, y: number, width: number, height: number, buffer: Uint8Array) => {
             gl.readPixels(x, y, width, height, gl.RGBA, gl.UNSIGNED_BYTE, buffer)

src/mol-gl/webgl/render-target.ts

@@ -9,14 +9,20 @@ import { idFactory } from 'mol-util/id-factory';
 import { createTexture } from './texture';
 import { createFramebuffer } from './framebuffer';
 import { createRenderbuffer } from './renderbuffer';
+import { TextureImage } from '../renderable/util';
 
 const getNextRenderTargetId = idFactory()
 
 export interface RenderTarget {
     readonly id: number
+    readonly width: number
+    readonly height: number
+    readonly image: Readonly<TextureImage>
 
     bind: () => void
     setSize: (width: number, height: number) => void
+    readBuffer: (x: number, y: number, width: number, height: number, dst: Uint8Array) => void
+    getBuffer: () => Uint8Array
     getImageData: () => ImageData
     destroy: () => void
 }
@@ -24,7 +30,7 @@ export interface RenderTarget {
 export function createRenderTarget (ctx: Context, _width: number, _height: number): RenderTarget {
     const { gl } = ctx
 
-    const image = {
+    const image: TextureImage = {
         array: new Uint8Array(_width * _height * 4),
         width: _width,
         height: _height
@@ -43,8 +49,21 @@ export function createRenderTarget (ctx: Context, _width: number, _height: numbe
 
     let destroyed = false
 
+    function readBuffer(x: number, y: number, width: number, height: number, dst: Uint8Array) {
+        framebuffer.bind()
+        ctx.readPixels(x, y, width, height, dst)
+    }
+
+    function getBuffer() {
+        readBuffer(0, 0, _width, _height, image.array)
+        return image.array
+    }
+
     return {
         id: getNextRenderTargetId(),
+        get width () { return _width },
+        get height () { return _height },
+        image,
 
         bind: () => {
             framebuffer.bind()
@@ -60,11 +79,9 @@ export function createRenderTarget (ctx: Context, _width: number, _height: numbe
 
             depthRenderbuffer.setSize(_width, _height)
         },
-        getImageData: () => {
-            framebuffer.bind()
-            ctx.readPixels(0, 0, _width, _height, image.array)
-            return createImageData(image.array, _width, _height)
-        },
+        readBuffer,
+        getBuffer,
+        getImageData: () => createImageData(getBuffer(), _width, _height),
         destroy: () => {
             if (destroyed) return
             targetTexture.destroy()

src/mol-gl/webgl/texture.ts

@@ -95,6 +95,7 @@ export function createTexture(ctx: Context, _format: TextureFormat, _type: Textu
             gl.bindTexture(gl.TEXTURE_2D, texture)
             // unpack alignment of 1 since we use textures only for data
             gl.pixelStorei(gl.UNPACK_ALIGNMENT, 1);
+            // gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, true);
             gl.texImage2D(gl.TEXTURE_2D, 0, format, width, height, 0, format, type, array)
             _width = width
             _height = height

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

@@ -10,7 +10,7 @@ import { RuntimeContext, Task } from 'mol-task';
 import { PositionData, DensityData } from './common';
 import { OrderedSet } from 'mol-data/int';
 import { createRenderable, createGaussianDensityRenderObject } from 'mol-gl/render-object';
-import { createContext } from 'mol-gl/webgl/context';
+import { createContext, Context } from 'mol-gl/webgl/context';
 import { GaussianDensityValues } from 'mol-gl/renderable/gaussian-density';
 import { RenderableState } from 'mol-gl/renderable';
 import { ValueCell } from 'mol-util';
@@ -19,7 +19,9 @@ import { createRenderTarget } from 'mol-gl/webgl/render-target';
 export const DefaultGaussianDensityProps = {
     resolution: 1,
     radiusOffset: 0,
-    smoothness: 1.5
+    smoothness: 1.5,
+    readSlices: false,
+    useGpu: true,
 }
 export type GaussianDensityProps = typeof DefaultGaussianDensityProps
 
@@ -33,16 +35,19 @@ function getDelta(box: Box3D, resolution: number) {
 
 export function computeGaussianDensity(position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps) {
     return Task.create('Gaussian Density', async ctx => {
-        const foo = await GaussianDensityGPU(ctx, position, box, radius, props)
-        console.log('FOOBAR', foo)
         return await GaussianDensity(ctx, position, box, radius, props)
     });
 }
 
 export async function GaussianDensity(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
-    const foo = await GaussianDensityGPU(ctx, position, box, radius, props)
-    console.log('FOOBAR', foo)
+    if (props.useGpu) {
+        return await GaussianDensityGPU(ctx, position, box, radius, props)
+    } else {
+        return await GaussianDensityCPU(ctx, position, box, radius, props)
+    }
+}
 
+export async function GaussianDensityCPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
     const { resolution, radiusOffset, smoothness } = props
 
     const { indices, x, y, z } = position
@@ -52,7 +57,7 @@ export async function GaussianDensity(ctx: RuntimeContext, position: PositionDat
     const p = Vec3.zero()
 
     const pad = (radiusOffset + 3) * 3 // TODO calculate max radius
-    const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad));
+    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
 
@@ -83,8 +88,9 @@ export async function GaussianDensity(ctx: RuntimeContext, position: PositionDat
 
     const gridPad = 1 / Math.max(...delta)
 
+    console.time('gaussian density cpu')
     for (let i = 0; i < n; ++i) {
-        const j = OrderedSet.getAt(indices, i);
+        const j = OrderedSet.getAt(indices, i)
 
         Vec3.set(v, x[j], y[j], z[j])
 
@@ -121,9 +127,10 @@ export async function GaussianDensity(ctx: RuntimeContext, position: PositionDat
         }
 
         if (i % updateChunk === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'filling density grid', current: i, max: n });
+            await ctx.update({ message: 'filling density grid', current: i, max: n })
         }
     }
+    console.timeEnd('gaussian density cpu')
 
     const transform = Mat4.identity()
     Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))
@@ -133,7 +140,7 @@ export async function GaussianDensity(ctx: RuntimeContext, position: PositionDat
 }
 
 export async function GaussianDensityGPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps) { // }: Promise<DensityData> {
-    const { resolution, radiusOffset } = props
+    const { resolution, radiusOffset, smoothness, readSlices } = props
 
     const { indices, x, y, z } = position
     const n = OrderedSet.size(indices)
@@ -149,6 +156,11 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     const dim = Vec3.zero()
     Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
 
+    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]))
+
     for (let i = 0; i < n; ++i) {
         const j = OrderedSet.getAt(indices, i);
 
@@ -158,7 +170,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         radii[i] = radius(j) + radiusOffset
 
         if (i % 10000 === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'preparing density data', current: i, max: n });
+            await ctx.update({ message: 'preparing density data', current: i, max: n })
         }
     }
 
@@ -178,47 +190,52 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         uBboxMax: ValueCell.create(expandedBox.max),
         uBboxSize: ValueCell.create(extent),
         uGridDim: ValueCell.create(dim),
+        uAlpha: ValueCell.create(smoothness),
     }
     const state: RenderableState = {
         visible: true,
         depthMask: false
     }
 
-    const canvas = document.createElement('canvas')
-    const gl = canvas.getContext('webgl', {
-        alpha: false,
-        antialias: true,
-        depth: true,
-        preserveDrawingBuffer: true
-    })
-    if (!gl) throw new Error('Could not create a WebGL rendering context')
-    const webgl = createContext(gl)
+    // TODO do in OffscreenCanvas (https://www.chromestatus.com/feature/5681560598609920)
+    const webgl = getWebGLContext()
 
     const renderObject = createGaussianDensityRenderObject(values, state)
     const renderable = createRenderable(webgl, renderObject)
 
     //
 
-    // get actual max texture size
-	const maxTexSize = 4096; // gl. .limits.maxTextureSize;
-	let fboTexDimX = 0
-	let fboTexDimY = dim[1]
-	let fboTexRows = 1
-	let fboTexCols = dim[0]
-	if(maxTexSize < dim[0] * dim[2]) {
-		fboTexCols =  Math.floor(maxTexSize / dim[0])
-		fboTexRows = Math.ceil(dim[2] / fboTexCols)
-		fboTexDimX = fboTexCols * dim[0]
-		fboTexDimY *= fboTexRows
-	} else {
-		fboTexDimX = dim[0] * dim[2]
-	}
+    // TODO fallback to lower resolution when texture size is not large enough
+    const maxTexSize = webgl.maxTextureSize
+    let fboTexDimX = 0
+    let fboTexDimY = dim[1]
+    let fboTexRows = 1
+    let fboTexCols = dim[0]
+    if (maxTexSize < dim[0] * dim[2]) {
+        fboTexCols =  Math.floor(maxTexSize / dim[0])
+        fboTexRows = Math.ceil(dim[2] / fboTexCols)
+        fboTexDimX = fboTexCols * dim[0]
+        fboTexDimY *= fboTexRows
+    } else {
+        fboTexDimX = dim[0] * dim[2]
+    }
+
+    //
+
+    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 { gl } = webgl
+
     const program = renderable.getProgram('draw')
     const renderTarget = createRenderTarget(webgl, fboTexDimX, fboTexDimY)
-    
+
     program.use()
     renderTarget.bind()
 
@@ -235,29 +252,65 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     gl.blendEquation(gl.FUNC_ADD)
     gl.enable(gl.BLEND)
 
-	gl.finish();
-	let currCol = 0;
-	let currY = 0;
-	let currX = 0;
-	for(let i = 0; i < dim[2]; ++i) {
-		if (currCol >= fboTexCols) {
-			currCol -= fboTexCols
-			currY += dim[1]
-			currX = 0
+    const slice = new Uint8Array(dim[0] * dim[1] * 4)
+
+    console.time('gpu gaussian density slices')
+    let currCol = 0
+    let currY = 0
+    let currX = 0
+    let j = 0
+    for (let i = 0; i < dim[2]; ++i) {
+        if (currCol >= fboTexCols) {
+            currCol -= fboTexCols
+            currY += dim[1]
+            currX = 0
         }
         gl.viewport(currX, currY, dim[0], dim[1])
         ValueCell.update(values.uCurrentSlice, i)
         ValueCell.update(values.uCurrentX, currX)
         ValueCell.update(values.uCurrentY, currY)
         renderable.render('draw')
-		++currCol
-		currX += dim[0]
-	}
-    gl.finish();
-    
-    const imageData = renderTarget.getImageData()
-    console.log(imageData)
-    debugTexture(imageData, 0.4)
+        if (readSlices) {
+            renderTarget.readBuffer(currX, currY, dim[0], dim[1], slice)
+            for (let iy = 0; iy < dim[1]; ++iy) {
+                for (let ix = 0; ix < dim[0]; ++ix) {
+                    data[j] = slice[4 * (iy * dim[0] + ix)] / 255
+                    ++j
+                }
+            }
+        }
+        ++currCol
+        currX += dim[0]
+
+        if (i % updateChunk === 0 && ctx.shouldUpdate) {
+            await ctx.update({ message: 'filling density grid', current: i, max: n })
+        }
+    }
+    console.timeEnd('gpu gaussian density slices')
+
+    //
+
+    if (!readSlices) {
+        console.time('gpu gaussian density full')
+        renderTarget.getBuffer()
+        let idx = 0
+        let tmpCol = 0
+        let tmpRow = 0
+        for (let iz = 0; iz < dim[2]; ++iz) {
+            if (tmpCol >= fboTexCols ) {
+                tmpCol = 0
+                tmpRow += dim[1]
+            }
+            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
+                    idx++
+                }
+            }
+            tmpCol++
+        }
+        console.timeEnd('gpu gaussian density full')
+    }
 
     //
 
@@ -265,25 +318,20 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))
     Mat4.setTranslation(transform, expandedBox.min)
 
-    return { field: imageData, idField: undefined, transform }
+    return { field, idField, transform }
 }
 
-function debugTexture(imageData: ImageData, scale: number) {
-	const canvas = document.createElement('canvas')
-	canvas.width = imageData.width
-	canvas.height = imageData.height
-    const ctx = canvas.getContext('2d')
-    if (!ctx) throw new Error('Could not create canvas 2d context')
-	ctx.putImageData(imageData, 0, 0)
-	canvas.toBlob(function(imgBlob){
-		var objectURL = window.URL.createObjectURL(imgBlob)
-		var img = document.createElement('img')
-		img.src = objectURL
-		img.style.width = imageData.width * scale + 'px'
-        img.style.height = imageData.height * scale + 'px'
-        img.style.position = 'absolute'
-        img.style.top = '0px'
-        img.style.left = '0px'
-		document.body.appendChild(img)
-	}, 'image/png')
+let webglContext: Context
+function getWebGLContext() {
+    if (webglContext) return webglContext
+    const canvas = document.createElement('canvas')
+    const gl = canvas.getContext('webgl', {
+        alpha: false,
+        antialias: true,
+        depth: true,
+        preserveDrawingBuffer: true
+    })
+    if (!gl) throw new Error('Could not create a WebGL rendering context')
+    webglContext = createContext(gl)
+    return webglContext
 }

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

@@ -9,12 +9,15 @@ import { SizeTheme } from 'mol-view/theme/size';
 import { GaussianDensity } from 'mol-math/geometry/gaussian-density';
 import { Task, RuntimeContext } from 'mol-task';
 import { DensityData } from 'mol-math/geometry';
-import { NumberParam, paramDefaultValues } from 'mol-view/parameter';
+import { NumberParam, paramDefaultValues, BooleanParam } from 'mol-view/parameter';
 
 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),
+    useGpu: BooleanParam('Use GPU', '', true),
+    readSlices: BooleanParam('Read Slices', '', false),
+    ignoreCache: BooleanParam('Ignore Cache', '', false),
 }
 export const DefaultGaussianDensityProps = paramDefaultValues(GaussianDensityParams)
 export type GaussianDensityProps = typeof DefaultGaussianDensityProps
@@ -52,8 +55,8 @@ export function computeUnitGaussianDensity(unit: Unit, props: GaussianDensityPro
 export async function computeUnitGaussianDensityCached(unit: Unit, props: GaussianDensityProps, cache: Map<string, DensityData>, ctx?: RuntimeContext) {
     const key = `${props.radiusOffset}|${props.resolution}|${props.smoothness}`
     let density = cache.get(key)
-    if (density) return density
+    if (density && !props.ignoreCache) return density
     density = ctx ? await computeUnitGaussianDensity(unit, props).runInContext(ctx) : await computeUnitGaussianDensity(unit, props).run()
-    cache.set(key, density)
+    if (!props.ignoreCache) cache.set(key, density)
     return density
 }