molstar/lib/mol-math/geometry/gaussian-density/gpu.js

A comprehensive macromolecular library.

/**
 * Copyright (c) 2017-2021 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>
 */
import { __assign } from "tslib";
import { Box3D } from '../../geometry';
import { OrderedSet } from '../../../mol-data/int';
import { Vec3, Tensor, Mat4, Vec2 } from '../../linear-algebra';
import { ValueCell } from '../../../mol-util';
import { createComputeRenderable } from '../../../mol-gl/renderable';
import { decodeFloatRGB } from '../../../mol-util/float-packing';
import { ShaderCode } from '../../../mol-gl/shader-code';
import { createComputeRenderItem } from '../../../mol-gl/webgl/render-item';
import { ValueSpec, AttributeSpec, UniformSpec, TextureSpec, DefineSpec } from '../../../mol-gl/renderable/schema';
import { gaussianDensity_vert } from '../../../mol-gl/shader/gaussian-density.vert';
import { gaussianDensity_frag } from '../../../mol-gl/shader/gaussian-density.frag';
var GaussianDensitySchema = {
    drawCount: ValueSpec('number'),
    instanceCount: ValueSpec('number'),
    aRadius: AttributeSpec('float32', 1, 0),
    aPosition: AttributeSpec('float32', 3, 0),
    aGroup: AttributeSpec('float32', 1, 0),
    uCurrentSlice: UniformSpec('f'),
    uCurrentX: UniformSpec('f'),
    uCurrentY: UniformSpec('f'),
    uBboxMin: UniformSpec('v3', 'material'),
    uBboxSize: UniformSpec('v3', 'material'),
    uGridDim: UniformSpec('v3', 'material'),
    uGridTexDim: UniformSpec('v3', 'material'),
    uGridTexScale: UniformSpec('v2', 'material'),
    uAlpha: UniformSpec('f', 'material'),
    uResolution: UniformSpec('f', 'material'),
    uRadiusFactorInv: UniformSpec('f', 'material'),
    tMinDistanceTex: TextureSpec('texture', 'rgba', 'float', 'nearest'),
    dGridTexType: DefineSpec('string', ['2d', '3d']),
    dCalcType: DefineSpec('string', ['density', 'minDistance', 'groupId']),
};
var GaussianDensityName = 'gaussian-density';
function getFramebuffer(webgl) {
    if (!webgl.namedFramebuffers[GaussianDensityName]) {
        webgl.namedFramebuffers[GaussianDensityName] = webgl.resources.framebuffer();
    }
    return webgl.namedFramebuffers[GaussianDensityName];
}
function getTexture(name, webgl, kind, format, type, filter) {
    var _name = GaussianDensityName + "-" + name;
    if (!webgl.namedTextures[_name]) {
        webgl.namedTextures[_name] = webgl.resources.texture(kind, format, type, filter);
    }
    return webgl.namedTextures[_name];
}
export function GaussianDensityGPU(position, box, radius, props, webgl) {
    // always use texture2d when the gaussian density needs to be downloaded from the GPU,
    // it's faster than texture3d
    // console.time('GaussianDensityTexture2d')
    var tmpTexture = getTexture('tmp', webgl, 'image-uint8', 'rgba', 'ubyte', 'linear');
    var _a = calcGaussianDensityTexture2d(webgl, position, box, radius, false, props, tmpTexture), scale = _a.scale, bbox = _a.bbox, texture = _a.texture, gridDim = _a.gridDim, gridTexDim = _a.gridTexDim, radiusFactor = _a.radiusFactor, resolution = _a.resolution;
    // webgl.waitForGpuCommandsCompleteSync()
    // console.timeEnd('GaussianDensityTexture2d')
    var _b = fieldFromTexture2d(webgl, texture, gridDim, gridTexDim), field = _b.field, idField = _b.idField;
    return { field: field, idField: idField, transform: getTransform(scale, bbox), radiusFactor: radiusFactor, resolution: resolution };
}
export function GaussianDensityTexture(webgl, position, box, radius, props, oldTexture) {
    return webgl.isWebGL2 ?
        GaussianDensityTexture3d(webgl, position, box, radius, props, oldTexture) :
        GaussianDensityTexture2d(webgl, position, box, radius, false, props, oldTexture);
}
export function GaussianDensityTexture2d(webgl, position, box, radius, powerOfTwo, props, oldTexture) {
    return finalizeGaussianDensityTexture(calcGaussianDensityTexture2d(webgl, position, box, radius, powerOfTwo, props, oldTexture));
}
export function GaussianDensityTexture3d(webgl, position, box, radius, props, oldTexture) {
    return finalizeGaussianDensityTexture(calcGaussianDensityTexture3d(webgl, position, box, radius, props, oldTexture));
}
function finalizeGaussianDensityTexture(_a) {
    var texture = _a.texture, scale = _a.scale, bbox = _a.bbox, gridDim = _a.gridDim, gridTexDim = _a.gridTexDim, gridTexScale = _a.gridTexScale, radiusFactor = _a.radiusFactor, resolution = _a.resolution;
    return { transform: getTransform(scale, bbox), texture: texture, bbox: bbox, gridDim: gridDim, gridTexDim: gridTexDim, gridTexScale: gridTexScale, radiusFactor: radiusFactor, resolution: resolution };
}
function getTransform(scale, bbox) {
    var transform = Mat4.identity();
    Mat4.fromScaling(transform, scale);
    Mat4.setTranslation(transform, bbox.min);
    return transform;
}
function calcGaussianDensityTexture2d(webgl, position, box, radius, powerOfTwo, props, texture) {
    // console.log('2d');
    var gl = webgl.gl, resources = webgl.resources, state = webgl.state, _a = webgl.extensions, colorBufferFloat = _a.colorBufferFloat, textureFloat = _a.textureFloat, colorBufferHalfFloat = _a.colorBufferHalfFloat, textureHalfFloat = _a.textureHalfFloat, blendMinMax = _a.blendMinMax;
    var smoothness = props.smoothness, resolution = props.resolution;
    var _b = prepareGaussianDensityData(position, box, radius, props), drawCount = _b.drawCount, positions = _b.positions, radii = _b.radii, groups = _b.groups, scale = _b.scale, expandedBox = _b.expandedBox, dim = _b.dim, maxRadius = _b.maxRadius;
    var dx = dim[0], dy = dim[1], dz = dim[2];
    var _c = getTexture2dSize(dim), texDimX = _c.texDimX, texDimY = _c.texDimY, texCols = _c.texCols, powerOfTwoSize = _c.powerOfTwoSize;
    // console.log({ texDimX, texDimY, texCols, powerOfTwoSize, dim });
    var gridTexDim = Vec3.create(texDimX, texDimY, 0);
    var gridTexScale = Vec2.create(texDimX / powerOfTwoSize, texDimY / powerOfTwoSize);
    var radiusFactor = maxRadius * 2;
    var width = powerOfTwo ? powerOfTwoSize : texDimX;
    var height = powerOfTwo ? powerOfTwoSize : texDimY;
    var minDistTex = getTexture('min-dist-2d', webgl, 'image-uint8', 'rgba', 'ubyte', 'nearest');
    minDistTex.define(width, height);
    var renderable = getGaussianDensityRenderable(webgl, drawCount, positions, radii, groups, minDistTex, expandedBox, dim, gridTexDim, gridTexScale, smoothness, resolution, radiusFactor);
    //
    var _d = renderable.values, uCurrentSlice = _d.uCurrentSlice, uCurrentX = _d.uCurrentX, uCurrentY = _d.uCurrentY;
    var framebuffer = getFramebuffer(webgl);
    framebuffer.bind();
    setRenderingDefaults(webgl);
    if (!texture)
        texture = colorBufferHalfFloat && textureHalfFloat
            ? resources.texture('image-float16', 'rgba', 'fp16', 'linear')
            : colorBufferFloat && textureFloat
                ? resources.texture('image-float32', 'rgba', 'float', 'linear')
                : resources.texture('image-uint8', 'rgba', 'ubyte', 'linear');
    texture.define(width, height);
    // console.log(renderable)
    function render(fbTex, clear) {
        state.currentRenderItemId = -1;
        fbTex.attachFramebuffer(framebuffer, 0);
        if (clear) {
            gl.viewport(0, 0, width, height);
            gl.scissor(0, 0, width, height);
            gl.clear(gl.COLOR_BUFFER_BIT);
        }
        ValueCell.update(uCurrentY, 0);
        var currCol = 0;
        var currY = 0;
        var currX = 0;
        for (var i = 0; i < dz; ++i) {
            if (currCol >= texCols) {
                currCol -= texCols;
                currY += dy;
                currX = 0;
                ValueCell.update(uCurrentY, currY);
            }
            // console.log({ i, currX, currY });
            ValueCell.update(uCurrentX, currX);
            ValueCell.update(uCurrentSlice, i);
            gl.viewport(currX, currY, dx, dy);
            gl.scissor(currX, currY, dx, dy);
            renderable.render();
            ++currCol;
            currX += dx;
        }
        gl.flush();
    }
    setupDensityRendering(webgl, renderable);
    render(texture, true);
    if (blendMinMax) {
        setupMinDistanceRendering(webgl, renderable);
        render(minDistTex, true);
        setupGroupIdRendering(webgl, renderable);
        render(texture, false);
    }
    // printTexture(webgl, minDistTex, 0.75);
    return { texture: texture, scale: scale, bbox: expandedBox, gridDim: dim, gridTexDim: gridTexDim, gridTexScale: gridTexScale, radiusFactor: radiusFactor, resolution: resolution };
}
function calcGaussianDensityTexture3d(webgl, position, box, radius, props, texture) {
    // console.log('3d');
    var gl = webgl.gl, resources = webgl.resources, state = webgl.state, _a = webgl.extensions, colorBufferFloat = _a.colorBufferFloat, textureFloat = _a.textureFloat, colorBufferHalfFloat = _a.colorBufferHalfFloat, textureHalfFloat = _a.textureHalfFloat;
    var smoothness = props.smoothness, resolution = props.resolution;
    var _b = prepareGaussianDensityData(position, box, radius, props), drawCount = _b.drawCount, positions = _b.positions, radii = _b.radii, groups = _b.groups, scale = _b.scale, expandedBox = _b.expandedBox, dim = _b.dim, maxRadius = _b.maxRadius;
    var dx = dim[0], dy = dim[1], dz = dim[2];
    var minDistTex = getTexture('min-dist-3d', webgl, 'volume-uint8', 'rgba', 'ubyte', 'nearest');
    minDistTex.define(dx, dy, dz);
    var gridTexScale = Vec2.create(1, 1);
    var radiusFactor = maxRadius * 2;
    var renderable = getGaussianDensityRenderable(webgl, drawCount, positions, radii, groups, minDistTex, expandedBox, dim, dim, gridTexScale, smoothness, resolution, radiusFactor);
    //
    var uCurrentSlice = renderable.values.uCurrentSlice;
    var framebuffer = getFramebuffer(webgl);
    framebuffer.bind();
    setRenderingDefaults(webgl);
    gl.viewport(0, 0, dx, dy);
    gl.scissor(0, 0, dx, dy);
    if (!texture)
        texture = colorBufferHalfFloat && textureHalfFloat
            ? resources.texture('volume-float16', 'rgba', 'fp16', 'linear')
            : colorBufferFloat && textureFloat
                ? resources.texture('volume-float32', 'rgba', 'float', 'linear')
                : resources.texture('volume-uint8', 'rgba', 'ubyte', 'linear');
    texture.define(dx, dy, dz);
    function render(fbTex, clear) {
        state.currentRenderItemId = -1;
        for (var i = 0; i < dz; ++i) {
            ValueCell.update(uCurrentSlice, i);
            fbTex.attachFramebuffer(framebuffer, 0, i);
            if (clear)
                gl.clear(gl.COLOR_BUFFER_BIT);
            renderable.render();
        }
        gl.flush();
    }
    setupDensityRendering(webgl, renderable);
    render(texture, true);
    setupMinDistanceRendering(webgl, renderable);
    render(minDistTex, true);
    setupGroupIdRendering(webgl, renderable);
    render(texture, false);
    return { texture: texture, scale: scale, bbox: expandedBox, gridDim: dim, gridTexDim: dim, gridTexScale: gridTexScale, radiusFactor: radiusFactor, resolution: resolution };
}
//
function prepareGaussianDensityData(position, box, radius, props) {
    var resolution = props.resolution, radiusOffset = props.radiusOffset;
    var scaleFactor = 1 / resolution;
    var indices = position.indices, x = position.x, y = position.y, z = position.z, id = position.id;
    var n = OrderedSet.size(indices);
    var positions = new Float32Array(n * 3);
    var radii = new Float32Array(n);
    var groups = new Float32Array(n);
    var maxRadius = 0;
    for (var i = 0; i < n; ++i) {
        var j = OrderedSet.getAt(indices, i);
        positions[i * 3] = x[j];
        positions[i * 3 + 1] = y[j];
        positions[i * 3 + 2] = z[j];
        var r = radius(j) + radiusOffset;
        if (maxRadius < r)
            maxRadius = r;
        radii[i] = r;
        groups[i] = id ? id[i] : i;
    }
    var pad = maxRadius * 2 + resolution * 4;
    var expandedBox = Box3D.expand(Box3D(), box, Vec3.create(pad, pad, pad));
    var scaledBox = Box3D.scale(Box3D(), expandedBox, scaleFactor);
    var dim = Box3D.size(Vec3(), scaledBox);
    Vec3.ceil(dim, dim);
    var scale = Vec3.create(resolution, resolution, resolution);
    return { drawCount: n, positions: positions, radii: radii, groups: groups, scale: scale, expandedBox: expandedBox, dim: dim, maxRadius: maxRadius };
}
function getGaussianDensityRenderable(webgl, drawCount, positions, radii, groups, minDistanceTexture, box, gridDim, gridTexDim, gridTexScale, smoothness, resolution, radiusFactor) {
    // console.log('radiusFactor', radiusFactor);
    if (webgl.namedComputeRenderables[GaussianDensityName]) {
        var extent = Vec3.sub(Vec3(), box.max, box.min);
        var v = webgl.namedComputeRenderables[GaussianDensityName].values;
        ValueCell.updateIfChanged(v.drawCount, drawCount);
        ValueCell.updateIfChanged(v.instanceCount, 1);
        ValueCell.update(v.aRadius, radii);
        ValueCell.update(v.aPosition, positions);
        ValueCell.update(v.aGroup, groups);
        ValueCell.updateIfChanged(v.uCurrentSlice, 0);
        ValueCell.updateIfChanged(v.uCurrentX, 0);
        ValueCell.updateIfChanged(v.uCurrentY, 0);
        ValueCell.update(v.uBboxMin, box.min);
        ValueCell.update(v.uBboxSize, extent);
        ValueCell.update(v.uGridDim, gridDim);
        ValueCell.update(v.uGridTexDim, gridTexDim);
        ValueCell.update(v.uGridTexScale, gridTexScale);
        ValueCell.updateIfChanged(v.uAlpha, smoothness);
        ValueCell.updateIfChanged(v.uResolution, resolution);
        ValueCell.updateIfChanged(v.uRadiusFactorInv, 1 / radiusFactor);
        ValueCell.update(v.tMinDistanceTex, minDistanceTexture);
        ValueCell.updateIfChanged(v.dGridTexType, minDistanceTexture.getDepth() > 0 ? '3d' : '2d');
        ValueCell.updateIfChanged(v.dCalcType, 'density');
        webgl.namedComputeRenderables[GaussianDensityName].update();
    }
    else {
        webgl.namedComputeRenderables[GaussianDensityName] = createGaussianDensityRenderable(webgl, drawCount, positions, radii, groups, minDistanceTexture, box, gridDim, gridTexDim, gridTexScale, smoothness, resolution, radiusFactor);
    }
    return webgl.namedComputeRenderables[GaussianDensityName];
}
function createGaussianDensityRenderable(webgl, drawCount, positions, radii, groups, minDistanceTexture, box, gridDim, gridTexDim, gridTexScale, smoothness, resolution, radiusFactor) {
    var extent = Vec3.sub(Vec3(), box.max, box.min);
    var values = {
        drawCount: ValueCell.create(drawCount),
        instanceCount: ValueCell.create(1),
        aRadius: ValueCell.create(radii),
        aPosition: ValueCell.create(positions),
        aGroup: ValueCell.create(groups),
        uCurrentSlice: ValueCell.create(0),
        uCurrentX: ValueCell.create(0),
        uCurrentY: ValueCell.create(0),
        uBboxMin: ValueCell.create(box.min),
        uBboxSize: ValueCell.create(extent),
        uGridDim: ValueCell.create(gridDim),
        uGridTexDim: ValueCell.create(gridTexDim),
        uGridTexScale: ValueCell.create(gridTexScale),
        uAlpha: ValueCell.create(smoothness),
        uResolution: ValueCell.create(resolution),
        uRadiusFactorInv: ValueCell.create(1 / radiusFactor),
        tMinDistanceTex: ValueCell.create(minDistanceTexture),
        dGridTexType: ValueCell.create(minDistanceTexture.getDepth() > 0 ? '3d' : '2d'),
        dCalcType: ValueCell.create('density'),
    };
    var schema = __assign({}, GaussianDensitySchema);
    var shaderCode = ShaderCode(GaussianDensityName, gaussianDensity_vert, gaussianDensity_frag);
    var renderItem = createComputeRenderItem(webgl, 'points', shaderCode, schema, values);
    return createComputeRenderable(renderItem, values);
}
function setRenderingDefaults(ctx) {
    var gl = ctx.gl, state = ctx.state;
    state.disable(gl.CULL_FACE);
    state.enable(gl.BLEND);
    state.disable(gl.DEPTH_TEST);
    state.enable(gl.SCISSOR_TEST);
    state.depthMask(false);
    state.clearColor(0, 0, 0, 0);
}
function setupMinDistanceRendering(webgl, renderable) {
    var gl = webgl.gl, state = webgl.state;
    ValueCell.update(renderable.values.dCalcType, 'minDistance');
    renderable.update();
    state.colorMask(false, false, false, true);
    state.blendFunc(gl.ONE, gl.ONE);
    // the shader writes 1 - dist so we set blending to MAX
    if (!webgl.extensions.blendMinMax) {
        throw new Error('GPU gaussian surface calculation requires EXT_blend_minmax');
    }
    state.blendEquation(webgl.extensions.blendMinMax.MAX);
}
function setupDensityRendering(webgl, renderable) {
    var gl = webgl.gl, state = webgl.state;
    ValueCell.update(renderable.values.dCalcType, 'density');
    renderable.update();
    state.colorMask(false, false, false, true);
    state.blendFunc(gl.ONE, gl.ONE);
    state.blendEquation(gl.FUNC_ADD);
}
function setupGroupIdRendering(webgl, renderable) {
    var gl = webgl.gl, state = webgl.state;
    ValueCell.update(renderable.values.dCalcType, 'groupId');
    renderable.update();
    // overwrite color, don't change alpha
    state.colorMask(true, true, true, false);
    state.blendFunc(gl.ONE, gl.ZERO);
    state.blendEquation(gl.FUNC_ADD);
}
function getTexture2dSize(gridDim) {
    var area = gridDim[0] * gridDim[1] * gridDim[2];
    var squareDim = Math.sqrt(area);
    var powerOfTwoSize = Math.pow(2, Math.ceil(Math.log(squareDim) / Math.log(2)));
    var texDimX = 0;
    var texDimY = gridDim[1];
    var texRows = 1;
    var texCols = gridDim[2];
    if (powerOfTwoSize < gridDim[0] * gridDim[2]) {
        texCols = Math.floor(powerOfTwoSize / gridDim[0]);
        texRows = Math.ceil(gridDim[2] / texCols);
        texDimX = texCols * gridDim[0];
        texDimY *= texRows;
    }
    else {
        texDimX = gridDim[0] * gridDim[2];
    }
    // console.log(texDimX, texDimY, texDimY < powerOfTwoSize ? powerOfTwoSize : powerOfTwoSize * 2);
    return { texDimX: texDimX, texDimY: texDimY, texRows: texRows, texCols: texCols, powerOfTwoSize: texDimY < powerOfTwoSize ? powerOfTwoSize : powerOfTwoSize * 2 };
}
function fieldFromTexture2d(ctx, texture, dim, texDim) {
    // console.time('fieldFromTexture2d')
    var dx = dim[0], dy = dim[1], dz = dim[2];
    var width = texDim[0], height = texDim[1];
    var fboTexCols = Math.floor(width / dx);
    var space = Tensor.Space(dim, [2, 1, 0], Float32Array);
    var data = space.create();
    var field = Tensor.create(space, data);
    var idData = space.create();
    var idField = Tensor.create(space, idData);
    var image = new Uint8Array(width * height * 4);
    var framebuffer = getFramebuffer(ctx);
    framebuffer.bind();
    texture.attachFramebuffer(framebuffer, 0);
    ctx.readPixels(0, 0, width, height, image);
    // printImageData(createImageData(image, width, height), 1/3)
    var j = 0;
    var tmpCol = 0;
    var tmpRow = 0;
    for (var iz = 0; iz < dz; ++iz) {
        if (tmpCol >= fboTexCols) {
            tmpCol = 0;
            tmpRow += dy;
        }
        for (var iy = 0; iy < dy; ++iy) {
            for (var ix = 0; ix < dx; ++ix) {
                var idx = 4 * (tmpCol * dx + (iy + tmpRow) * width + ix);
                data[j] = image[idx + 3] / 255;
                idData[j] = decodeFloatRGB(image[idx], image[idx + 1], image[idx + 2]);
                j++;
            }
        }
        tmpCol++;
    }
    // console.timeEnd('fieldFromTexture2d')
    return { field: field, idField: idField };
}
//# sourceMappingURL=gpu.js.map