molstar
Version:
A comprehensive macromolecular library.
261 lines • 13.4 kB
JavaScript
/**
* Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
*
* @author Alexander Rose <alexander.rose@weirdbyte.de>
*/
import { __assign } from "tslib";
import { ValueCell } from '../../../mol-util';
import { createComputeRenderable } from '../../../mol-gl/renderable';
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 { quad_vert } from '../../../mol-gl/shader/quad.vert';
import { normalize_frag } from '../../../mol-gl/shader/compute/color-smoothing/normalize.frag';
import { QuadSchema, QuadValues } from '../../../mol-gl/compute/util';
import { Vec2, Vec3, Vec4 } from '../../../mol-math/linear-algebra';
import { Box3D } from '../../../mol-math/geometry';
import { accumulate_frag } from '../../../mol-gl/shader/compute/color-smoothing/accumulate.frag';
import { accumulate_vert } from '../../../mol-gl/shader/compute/color-smoothing/accumulate.vert';
export var ColorAccumulateSchema = {
drawCount: ValueSpec('number'),
instanceCount: ValueSpec('number'),
stride: ValueSpec('number'),
uTotalCount: UniformSpec('i'),
uInstanceCount: UniformSpec('i'),
uGroupCount: UniformSpec('i'),
aTransform: AttributeSpec('float32', 16, 1),
aInstance: AttributeSpec('float32', 1, 1),
aSample: AttributeSpec('float32', 1, 0),
uGeoTexDim: UniformSpec('v2', 'buffered'),
tPosition: TextureSpec('texture', 'rgba', 'float', 'nearest'),
tGroup: TextureSpec('texture', 'rgba', 'float', 'nearest'),
uColorTexDim: UniformSpec('v2'),
tColor: TextureSpec('image-uint8', 'rgb', 'ubyte', 'nearest'),
dColorType: DefineSpec('string', ['group', 'groupInstance', 'vertex', 'vertexInstance']),
uCurrentSlice: UniformSpec('f'),
uCurrentX: UniformSpec('f'),
uCurrentY: UniformSpec('f'),
uBboxMin: UniformSpec('v3', 'material'),
uBboxSize: UniformSpec('v3', 'material'),
uResolution: UniformSpec('f', 'material'),
};
var ColorAccumulateName = 'color-accumulate';
function getSampleBuffer(sampleCount, stride) {
var sampleBuffer = new Float32Array(sampleCount);
for (var i = 0; i < sampleCount; ++i) {
sampleBuffer[i] = i * stride;
}
return sampleBuffer;
}
function getAccumulateRenderable(ctx, input, box, resolution, stride) {
if (ctx.namedComputeRenderables[ColorAccumulateName]) {
var extent = Vec3.sub(Vec3(), box.max, box.min);
var v = ctx.namedComputeRenderables[ColorAccumulateName].values;
var sampleCount = Math.round(input.vertexCount / stride);
if (sampleCount > v.drawCount.ref.value || stride !== v.stride.ref.value) {
ValueCell.update(v.aSample, getSampleBuffer(sampleCount, stride));
}
ValueCell.updateIfChanged(v.drawCount, sampleCount);
ValueCell.updateIfChanged(v.instanceCount, input.instanceCount);
ValueCell.updateIfChanged(v.stride, stride);
ValueCell.updateIfChanged(v.uTotalCount, input.vertexCount);
ValueCell.updateIfChanged(v.uInstanceCount, input.instanceCount);
ValueCell.updateIfChanged(v.uGroupCount, input.groupCount);
ValueCell.update(v.aTransform, input.transformBuffer);
ValueCell.update(v.aInstance, input.instanceBuffer);
ValueCell.update(v.uGeoTexDim, Vec2.set(v.uGeoTexDim.ref.value, input.positionTexture.getWidth(), input.positionTexture.getHeight()));
ValueCell.update(v.tPosition, input.positionTexture);
ValueCell.update(v.tGroup, input.groupTexture);
ValueCell.update(v.uColorTexDim, Vec2.set(v.uColorTexDim.ref.value, input.colorData.width, input.colorData.height));
ValueCell.update(v.tColor, input.colorData);
ValueCell.updateIfChanged(v.dColorType, input.colorType);
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.updateIfChanged(v.uResolution, resolution);
ctx.namedComputeRenderables[ColorAccumulateName].update();
}
else {
ctx.namedComputeRenderables[ColorAccumulateName] = createAccumulateRenderable(ctx, input, box, resolution, stride);
}
return ctx.namedComputeRenderables[ColorAccumulateName];
}
function createAccumulateRenderable(ctx, input, box, resolution, stride) {
var extent = Vec3.sub(Vec3(), box.max, box.min);
var sampleCount = Math.round(input.vertexCount / stride);
var values = {
drawCount: ValueCell.create(sampleCount),
instanceCount: ValueCell.create(input.instanceCount),
stride: ValueCell.create(stride),
uTotalCount: ValueCell.create(input.vertexCount),
uInstanceCount: ValueCell.create(input.instanceCount),
uGroupCount: ValueCell.create(input.groupCount),
aTransform: ValueCell.create(input.transformBuffer),
aInstance: ValueCell.create(input.instanceBuffer),
aSample: ValueCell.create(getSampleBuffer(sampleCount, stride)),
uGeoTexDim: ValueCell.create(Vec2.create(input.positionTexture.getWidth(), input.positionTexture.getHeight())),
tPosition: ValueCell.create(input.positionTexture),
tGroup: ValueCell.create(input.groupTexture),
uColorTexDim: ValueCell.create(Vec2.create(input.colorData.width, input.colorData.height)),
tColor: ValueCell.create(input.colorData),
dColorType: ValueCell.create(input.colorType),
uCurrentSlice: ValueCell.create(0),
uCurrentX: ValueCell.create(0),
uCurrentY: ValueCell.create(0),
uBboxMin: ValueCell.create(box.min),
uBboxSize: ValueCell.create(extent),
uResolution: ValueCell.create(resolution),
};
var schema = __assign({}, ColorAccumulateSchema);
var shaderCode = ShaderCode('accumulate', accumulate_vert, accumulate_frag);
var renderItem = createComputeRenderItem(ctx, 'points', shaderCode, schema, values);
return createComputeRenderable(renderItem, values);
}
function setAccumulateDefaults(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);
state.blendFunc(gl.ONE, gl.ONE);
state.blendEquation(gl.FUNC_ADD);
}
//
export var ColorNormalizeSchema = __assign(__assign({}, QuadSchema), { tColor: TextureSpec('texture', 'rgba', 'float', 'nearest'), uTexSize: UniformSpec('v2') });
var ColorNormalizeName = 'color-normalize';
function getNormalizeRenderable(ctx, color) {
if (ctx.namedComputeRenderables[ColorNormalizeName]) {
var v = ctx.namedComputeRenderables[ColorNormalizeName].values;
ValueCell.update(v.tColor, color);
ValueCell.update(v.uTexSize, Vec2.set(v.uTexSize.ref.value, color.getWidth(), color.getHeight()));
ctx.namedComputeRenderables[ColorNormalizeName].update();
}
else {
ctx.namedComputeRenderables[ColorNormalizeName] = createColorNormalizeRenderable(ctx, color);
}
return ctx.namedComputeRenderables[ColorNormalizeName];
}
function createColorNormalizeRenderable(ctx, color) {
var values = __assign(__assign({}, QuadValues), { tColor: ValueCell.create(color), uTexSize: ValueCell.create(Vec2.create(color.getWidth(), color.getHeight())) });
var schema = __assign({}, ColorNormalizeSchema);
var shaderCode = ShaderCode('normalize', quad_vert, normalize_frag);
var renderItem = createComputeRenderItem(ctx, 'triangles', shaderCode, schema, values);
return createComputeRenderable(renderItem, values);
}
function setNormalizeDefaults(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);
state.blendFunc(gl.ONE, gl.ONE);
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 };
}
export function calcTextureMeshColorSmoothing(input, resolution, stride, webgl, texture) {
var gl = webgl.gl, resources = webgl.resources, state = webgl.state, _a = webgl.extensions, colorBufferHalfFloat = _a.colorBufferHalfFloat, textureHalfFloat = _a.textureHalfFloat;
var isInstanceType = input.colorType.endsWith('Instance');
var box = Box3D.fromSphere3D(Box3D(), isInstanceType ? input.boundingSphere : input.invariantBoundingSphere);
var scaleFactor = 1 / resolution;
var scaledBox = Box3D.scale(Box3D(), box, scaleFactor);
var gridDim = Box3D.size(Vec3(), scaledBox);
Vec3.ceil(gridDim, gridDim);
Vec3.add(gridDim, gridDim, Vec3.create(2, 2, 2));
var min = box.min;
var dx = gridDim[0], dy = gridDim[1], dz = gridDim[2];
var _b = getTexture2dSize(gridDim), width = _b.texDimX, height = _b.texDimY, texCols = _b.texCols;
// console.log({ width, height, texCols, dim, resolution });
if (!webgl.namedTextures[ColorAccumulateName]) {
webgl.namedTextures[ColorAccumulateName] = colorBufferHalfFloat && textureHalfFloat
? resources.texture('image-float16', 'rgba', 'fp16', 'nearest')
: resources.texture('image-float32', 'rgba', 'float', 'nearest');
}
var accumulateTexture = webgl.namedTextures[ColorAccumulateName];
accumulateTexture.define(width, height);
var accumulateRenderable = getAccumulateRenderable(webgl, input, box, resolution, stride);
//
var _c = accumulateRenderable.values, uCurrentSlice = _c.uCurrentSlice, uCurrentX = _c.uCurrentX, uCurrentY = _c.uCurrentY;
if (!webgl.namedFramebuffers[ColorAccumulateName]) {
webgl.namedFramebuffers[ColorAccumulateName] = webgl.resources.framebuffer();
}
var framebuffer = webgl.namedFramebuffers[ColorAccumulateName];
framebuffer.bind();
setAccumulateDefaults(webgl);
state.currentRenderItemId = -1;
accumulateTexture.attachFramebuffer(framebuffer, 0);
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);
accumulateRenderable.render();
++currCol;
currX += dx;
}
// const accImage = new Float32Array(width * height * 4);
// accumulateTexture.attachFramebuffer(framebuffer, 0);
// webgl.readPixels(0, 0, width, height, accImage);
// console.log(accImage);
// printTextureImage({ array: accImage, width, height }, 1 / 4);
// normalize
if (!texture)
texture = resources.texture('image-uint8', 'rgb', 'ubyte', 'linear');
texture.define(width, height);
var normalizeRenderable = getNormalizeRenderable(webgl, accumulateTexture);
setNormalizeDefaults(webgl);
state.currentRenderItemId = -1;
texture.attachFramebuffer(framebuffer, 0);
gl.viewport(0, 0, width, height);
gl.scissor(0, 0, width, height);
gl.clear(gl.COLOR_BUFFER_BIT);
normalizeRenderable.render();
// const normImage = new Uint8Array(width * height * 4);
// texture.attachFramebuffer(framebuffer, 0);
// webgl.readPixels(0, 0, width, height, normImage);
// console.log(normImage);
// printTextureImage({ array: normImage, width, height }, 1 / 4);
var gridTransform = Vec4.create(min[0], min[1], min[2], scaleFactor);
var type = isInstanceType ? 'volumeInstance' : 'volume';
return { texture: texture, gridDim: gridDim, gridTexDim: Vec2.create(width, height), gridTransform: gridTransform, type: type };
}
//# sourceMappingURL=color-smoothing.js.map