Skip to content
Snippets Groups Projects
Commit 2f6b2989 authored by Alexander Rose's avatar Alexander Rose
Browse files

wip, calculate gaussian density group ids on gpu

parent 5996ae6b
No related branches found
No related tags found
No related merge requests found
......@@ -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)
......
......@@ -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>
......
......@@ -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
}
\ No newline at end of file
......@@ -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;
......
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)
\ No newline at end of file
......@@ -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 }
}
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment