wip, gpu gaussian density

src/mol-gl/render-object.ts

@@ -4,10 +4,11 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { PointsRenderable, MeshRenderable, RenderableState, MeshValues, PointsValues, LinesValues, LinesRenderable } from './renderable'
+import { PointsRenderable, MeshRenderable, RenderableState, MeshValues, PointsValues, LinesValues, LinesRenderable, Renderable } from './renderable'
 import { RenderableValues } from './renderable/schema';
 import { idFactory } from 'mol-util/id-factory';
 import { Context } from './webgl/context';
+import { GaussianDensityValues, GaussianDensityRenderable } from './renderable/gaussian-density';
 
 const getNextId = idFactory(0, 0x7FFFFFFF)
 
@@ -15,7 +16,8 @@ export interface BaseRenderObject { id: number, type: string, values: Renderable
 export interface MeshRenderObject extends BaseRenderObject { type: 'mesh', values: MeshValues }
 export interface PointsRenderObject extends BaseRenderObject { type: 'points', values: PointsValues }
 export interface LinesRenderObject extends BaseRenderObject { type: 'lines', values: LinesValues }
-export type RenderObject = MeshRenderObject | PointsRenderObject | LinesRenderObject
+export interface GaussianDensityRenderObject extends BaseRenderObject { type: 'gaussian-density', values: GaussianDensityValues }
+export type RenderObject = MeshRenderObject | PointsRenderObject | LinesRenderObject | GaussianDensityRenderObject
 
 export function createMeshRenderObject(values: MeshValues, state: RenderableState): MeshRenderObject {
     return { id: getNextId(), type: 'mesh', values, state }
@@ -26,11 +28,15 @@ export function createPointsRenderObject(values: PointsValues, state: Renderable
 export function createLinesRenderObject(values: LinesValues, state: RenderableState): LinesRenderObject {
     return { id: getNextId(), type: 'lines', values, state }
 }
+export function createGaussianDensityRenderObject(values: GaussianDensityValues, state: RenderableState): GaussianDensityRenderObject {
+    return { id: getNextId(), type: 'gaussian-density', values, state }
+}
 
-export function createRenderable(ctx: Context, o: RenderObject) {
+export function createRenderable(ctx: Context, o: RenderObject): Renderable<any> {
     switch (o.type) {
         case 'mesh': return MeshRenderable(ctx, o.id, o.values, o.state)
         case 'points': return PointsRenderable(ctx, o.id, o.values, o.state)
         case 'lines': return LinesRenderable(ctx, o.id, o.values, o.state)
+        case 'gaussian-density': return GaussianDensityRenderable(ctx, o.id, o.values, o.state)
     }
 }
\ No newline at end of file

src/mol-gl/renderable.ts

@@ -5,7 +5,7 @@
  */
 
 import { Program } from './webgl/program';
-import { RenderableValues, Values, RenderableSchema, BaseValues } from './renderable/schema';
+import { RenderableValues, Values, RenderableSchema } from './renderable/schema';
 import { RenderVariant, RenderItem } from './webgl/render-item';
 import { Sphere3D } from 'mol-math/geometry';
 // import { calculateBoundingSphereFromValues } from './renderable/util';
@@ -17,7 +17,7 @@ export type RenderableState = {
     depthMask: boolean
 }
 
-export interface Renderable<T extends RenderableValues & BaseValues> {
+export interface Renderable<T extends RenderableValues> {
     readonly values: T
     readonly state: RenderableState
     readonly boundingSphere: Sphere3D
@@ -29,7 +29,7 @@ export interface Renderable<T extends RenderableValues & BaseValues> {
     dispose: () => void
 }
 
-export function createRenderable<T extends Values<RenderableSchema> & BaseValues>(renderItem: RenderItem, values: T, state: RenderableState): Renderable<T> {
+export function createRenderable<T extends Values<RenderableSchema>>(renderItem: RenderItem, values: T, state: RenderableState): Renderable<T> {
     // TODO
     let boundingSphere: Sphere3D = Sphere3D.create(Vec3.zero(), 50)
 
@@ -43,7 +43,7 @@ export function createRenderable<T extends Values<RenderableSchema> & BaseValues
             // boundingSphere = calculateBoundingSphereFromValues(values)
             // return boundingSphere
         },
-        get opaque () { return values.uAlpha.ref.value === 1 },
+        get opaque () { return values.uAlpha && values.uAlpha.ref.value === 1 },
 
         render: (variant: RenderVariant) => renderItem.render(variant),
         getProgram: (variant: RenderVariant) => renderItem.getProgram(variant),

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

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Renderable, RenderableState, createRenderable } from '../renderable'
+import { Context } from '../webgl/context';
+import { createRenderItem } from '../webgl/render-item';
+import { AttributeSpec, Values, UniformSpec, ValueSpec } from './schema';
+import { GaussianDensityShaderCode } from '../shader-code';
+
+export const GaussianDensitySchema = {
+    drawCount: ValueSpec('number'),
+    instanceCount: ValueSpec('number'),
+
+    aRadius: AttributeSpec('float32', 1, 0),
+    aPosition: AttributeSpec('float32', 3, 0),
+
+    uCurrentSlice: UniformSpec('f'),
+    uCurrentX: UniformSpec('f'),
+    uCurrentY: UniformSpec('f'),
+    uBboxMin: UniformSpec('v3'),
+    uBboxMax: UniformSpec('v3'),
+    uBboxSize: UniformSpec('v3'),
+    uGridDim: UniformSpec('v3'),
+}
+export type GaussianDensitySchema = typeof GaussianDensitySchema
+export type GaussianDensityValues = Values<GaussianDensitySchema>
+
+export function GaussianDensityRenderable(ctx: Context, id: number, values: GaussianDensityValues, state: RenderableState): Renderable<GaussianDensityValues> {
+    const schema = { ...GaussianDensitySchema }
+    const shaderCode = GaussianDensityShaderCode
+    const renderItem = createRenderItem(ctx, 'points', shaderCode, schema, { ...values })
+
+    return createRenderable(renderItem, values, state);
+}
\ No newline at end of file

src/mol-gl/shader-code.ts

@@ -39,6 +39,11 @@ export const MeshShaderCode = ShaderCode(
     require('mol-gl/shader/mesh.frag')
 )
 
+export const GaussianDensityShaderCode = ShaderCode(
+    require('mol-gl/shader/gaussian-density.vert'),
+    require('mol-gl/shader/gaussian-density.frag')
+)
+
 export type ShaderDefines = {
     [k: string]: ValueCell<DefineType>
 }

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

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Michael Krone <michael.krone@uni-tuebingen.de>
+ */
+
+precision mediump float;
+
+varying vec3 position;
+varying float radius;
+
+uniform vec3 uBboxSize;
+uniform vec3 uBboxMin;
+uniform vec3 uBboxMax;
+uniform vec3 uGridDim;
+uniform float uCurrentSlice;
+uniform float uCurrentX;
+uniform float uCurrentY;
+
+void main() {
+    vec3 tmpVec = gl_FragCoord.xyz;
+    tmpVec.x = tmpVec.x - uCurrentX;
+    tmpVec.y = tmpVec.y - uCurrentY;
+    vec3 fragPos = vec3(
+        (tmpVec.x - 0.5) / uGridDim.x,
+        (tmpVec.y - 0.5) / uGridDim.y,
+        (uCurrentSlice) / uGridDim.z
+    );
+    float dist = length(fragPos * uBboxSize - position * uBboxSize);
+    float density = 1.0 - smoothstep( 0.0, radius * 2.0, dist);
+    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);
+}
\ No newline at end of file

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

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Michael Krone <michael.krone@uni-tuebingen.de>
+ */
+
+precision mediump float;
+
+attribute vec3 aPosition;
+attribute float aRadius;
+
+varying vec3 position;
+varying float radius;
+
+uniform vec3 uBboxSize;
+uniform vec3 uBboxMin;
+uniform vec3 uBboxMax;
+uniform vec3 uGridDim;
+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;
+    position = (aPosition - uBboxMin) / uBboxSize;
+    gl_Position = vec4(position * 2.0 - 1.0, 1.0);
+}
\ No newline at end of file

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

@@ -60,6 +60,13 @@ interface ValueChanges {
     uniforms: boolean
 }
 
+// TODO make `RenderVariantDefines` a parameter for `createRenderItem`
+
+/**
+ * Creates a render item
+ * 
+ * - assumes that `values.drawCount` and `values.instanceCount` exist
+ */
 export function createRenderItem(ctx: Context, drawMode: DrawMode, shaderCode: ShaderCode, schema: RenderableSchema, values: RenderableValues): RenderItem {
     const id = getNextRenderItemId()
     const { programCache } = ctx

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

@@ -9,6 +9,12 @@ import { Vec3, Mat4, Tensor } from '../linear-algebra';
 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 { GaussianDensityValues } from 'mol-gl/renderable/gaussian-density';
+import { RenderableState } from 'mol-gl/renderable';
+import { ValueCell } from 'mol-util';
+import { createRenderTarget } from 'mol-gl/webgl/render-target';
 
 export const DefaultGaussianDensityProps = {
     resolution: 1,
@@ -26,10 +32,17 @@ 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 => await GaussianDensity(ctx, position, box, radius, props));
+    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)
+
     const { resolution, radiusOffset, smoothness } = props
 
     const { indices, x, y, z } = position
@@ -118,3 +131,159 @@ export async function GaussianDensity(ctx: RuntimeContext, position: PositionDat
 
     return { field, idField, transform }
 }
+
+export async function GaussianDensityGPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps) { // }: Promise<DensityData> {
+    const { resolution, radiusOffset } = props
+
+    const { indices, x, y, z } = position
+    const n = OrderedSet.size(indices)
+
+    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))
+
+    for (let i = 0; i < n; ++i) {
+        const j = OrderedSet.getAt(indices, i);
+
+        positions[i * 3] = x[j]
+        positions[i * 3 + 1] = y[j]
+        positions[i * 3 + 2] = z[j]
+        radii[i] = radius(j) + radiusOffset
+
+        if (i % 10000 === 0 && ctx.shouldUpdate) {
+            await ctx.update({ message: 'preparing density data', current: i, max: n });
+        }
+    }
+
+    //
+
+    const values: GaussianDensityValues = {
+        drawCount: ValueCell.create(n),
+        instanceCount: ValueCell.create(1),
+
+        aRadius: ValueCell.create(radii),
+        aPosition: ValueCell.create(positions),
+
+        uCurrentSlice: ValueCell.create(0),
+        uCurrentX: ValueCell.create(0),
+        uCurrentY: ValueCell.create(0),
+        uBboxMin: ValueCell.create(expandedBox.min),
+        uBboxMax: ValueCell.create(expandedBox.max),
+        uBboxSize: ValueCell.create(extent),
+        uGridDim: ValueCell.create(dim),
+    }
+    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)
+
+    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]
+	}
+
+    //
+
+    const program = renderable.getProgram('draw')
+    const renderTarget = createRenderTarget(webgl, fboTexDimX, fboTexDimY)
+    
+    program.use()
+    renderTarget.bind()
+
+    gl.disable(gl.CULL_FACE)
+    gl.frontFace(gl.CCW)
+    gl.cullFace(gl.BACK)
+
+    gl.depthMask(true)
+    gl.clearColor(0, 0, 0, 1)
+    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT)
+    gl.depthMask(false)
+
+    gl.blendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA)
+    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
+        }
+        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)
+
+    //
+
+    const transform = Mat4.identity()
+    Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))
+    Mat4.setTranslation(transform, expandedBox.min)
+
+    return { field: imageData, idField: undefined, 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')
+}
\ No newline at end of file