UNPKG

molstar

Version:

A comprehensive macromolecular library.

318 lines (317 loc) 16 kB
/** * Copyright (c) 2021-2022 mol* contributors, licensed under MIT, See LICENSE file for more info. * * @author Alexander Rose <alexander.rose@weirdbyte.de> */ import { createTextureImage } from '../../../mol-gl/renderable/util'; import { Box3D } from '../../../mol-math/geometry'; import { lerp } from '../../../mol-math/interpolate'; import { Vec2, Vec3, Vec4 } from '../../../mol-math/linear-algebra'; import { getVolumeTexture2dLayout } from '../../../mol-repr/volume/util'; import { ValueCell } from '../../../mol-util'; export function calcMeshColorSmoothing(input, resolution, stride, webgl, texture) { var colorType = input.colorType, vertexCount = input.vertexCount, groupCount = input.groupCount, positionBuffer = input.positionBuffer, transformBuffer = input.transformBuffer, groupBuffer = input.groupBuffer, itemSize = input.itemSize; var isInstanceType = colorType.endsWith('Instance'); var box = Box3D.fromSphere3D(Box3D(), isInstanceType ? input.boundingSphere : input.invariantBoundingSphere); var pad = 1 + resolution; var expandedBox = Box3D.expand(Box3D(), box, Vec3.create(pad, pad, pad)); var scaleFactor = 1 / resolution; var scaledBox = Box3D.scale(Box3D(), expandedBox, scaleFactor); var gridDim = Box3D.size(Vec3(), scaledBox); Vec3.ceil(gridDim, gridDim); Vec3.add(gridDim, gridDim, Vec3.create(2, 2, 2)); var min = expandedBox.min; var xn = gridDim[0], yn = gridDim[1]; var _a = getVolumeTexture2dLayout(gridDim), width = _a.width, height = _a.height; // console.log({ width, height, dim }); var data = new Float32Array(width * height * itemSize); var count = new Float32Array(width * height); var grid = new Uint8Array(width * height * itemSize); var textureImage = { array: grid, width: width, height: height, filter: 'linear' }; var instanceCount = isInstanceType ? input.instanceCount : 1; var colors = input.colorData.array; function getIndex(x, y, z) { var column = Math.floor(((z * xn) % width) / xn); var row = Math.floor((z * xn) / width); var px = column * xn + x; return itemSize * ((row * yn * width) + (y * width) + px); } var p = 2; var dimX = gridDim[0], dimY = gridDim[1], dimZ = gridDim[2]; var v = Vec3(); for (var i = 0; i < instanceCount; ++i) { for (var j = 0; j < vertexCount; j += stride) { Vec3.fromArray(v, positionBuffer, j * 3); if (isInstanceType) Vec3.transformMat4Offset(v, v, transformBuffer, 0, 0, i * 16); Vec3.sub(v, v, min); Vec3.scale(v, v, scaleFactor); var vx = v[0], vy = v[1], vz = v[2]; // vertex mapped to grid var x = Math.floor(vx); var y = Math.floor(vy); var z = Math.floor(vz); // group colors var ci = (i * groupCount + groupBuffer[j]) * itemSize; // Extents of grid to consider for this atom var begX = Math.max(0, x - p); var begY = Math.max(0, y - p); var begZ = Math.max(0, z - p); // Add two to these points: // - x, y, z are floor'd values so this ensures coverage // - these are loop limits (exclusive) var endX = Math.min(dimX, x + p + 2); var endY = Math.min(dimY, y + p + 2); var endZ = Math.min(dimZ, z + p + 2); for (var xi = begX; xi < endX; ++xi) { var dx = xi - vx; for (var yi = begY; yi < endY; ++yi) { var dy = yi - vy; for (var zi = begZ; zi < endZ; ++zi) { var dz = zi - vz; var d = Math.sqrt(dx * dx + dy * dy + dz * dz); if (d > p) continue; var s = p - d; var index = getIndex(xi, yi, zi); for (var k = 0; k < itemSize; ++k) { data[index + k] += colors[ci + k] * s; } count[index / itemSize] += s; } } } } } for (var i = 0, il = count.length; i < il; ++i) { var is = i * itemSize; var c = count[i]; for (var k = 0; k < itemSize; ++k) { grid[is + k] = Math.round(data[is + k] / c); } } var gridTexDim = Vec2.create(width, height); var gridTransform = Vec4.create(min[0], min[1], min[2], scaleFactor); var type = isInstanceType ? 'volumeInstance' : 'volume'; if (webgl) { if (!texture) { var format = itemSize === 4 ? 'rgba' : itemSize === 3 ? 'rgb' : 'alpha'; texture = webgl.resources.texture('image-uint8', format, 'ubyte', 'linear'); } texture.load(textureImage); return { kind: 'volume', texture: texture, gridTexDim: gridTexDim, gridDim: gridDim, gridTransform: gridTransform, type: type }; } else { var interpolated = getTrilinearlyInterpolated({ vertexCount: vertexCount, instanceCount: instanceCount, transformBuffer: transformBuffer, positionBuffer: positionBuffer, colorType: type, grid: grid, gridDim: gridDim, gridTexDim: gridTexDim, gridTransform: gridTransform, vertexStride: 3, colorStride: itemSize, outputStride: itemSize }); return { kind: 'vertex', texture: interpolated, texDim: Vec2.create(interpolated.width, interpolated.height), type: isInstanceType ? 'vertexInstance' : 'vertex' }; } } export function getTrilinearlyInterpolated(input) { var vertexCount = input.vertexCount, positionBuffer = input.positionBuffer, transformBuffer = input.transformBuffer, grid = input.grid, gridDim = input.gridDim, gridTexDim = input.gridTexDim, gridTransform = input.gridTransform, vertexStride = input.vertexStride, colorStride = input.colorStride; var itemOffset = input.itemOffset || 0; var outputStride = input.outputStride; if (outputStride + itemOffset > colorStride) { throw new Error('outputStride + itemOffset must NOT be larger than colorStride'); } var isInstanceType = input.colorType.endsWith('Instance'); var instanceCount = isInstanceType ? input.instanceCount : 1; var image = createTextureImage(Math.max(1, instanceCount * vertexCount), outputStride, Uint8Array); var array = image.array; var xn = gridDim[0], yn = gridDim[1]; var width = gridTexDim[0]; var min = Vec3.fromArray(Vec3(), gridTransform, 0); var scaleFactor = gridTransform[3]; function getIndex(x, y, z) { var column = Math.floor(((z * xn) % width) / xn); var row = Math.floor((z * xn) / width); var px = column * xn + x; return colorStride * ((row * yn * width) + (y * width) + px); } var v = Vec3(); var v0 = Vec3(); var v1 = Vec3(); var vd = Vec3(); for (var i = 0; i < instanceCount; ++i) { for (var j = 0; j < vertexCount; ++j) { Vec3.fromArray(v, positionBuffer, j * vertexStride); if (isInstanceType) Vec3.transformMat4Offset(v, v, transformBuffer, 0, 0, i * 16); Vec3.sub(v, v, min); Vec3.scale(v, v, scaleFactor); Vec3.floor(v0, v); Vec3.ceil(v1, v); Vec3.sub(vd, v, v0); Vec3.sub(v, v1, v0); Vec3.div(vd, vd, v); var x0 = v0[0], y0 = v0[1], z0 = v0[2]; var x1 = v1[0], y1 = v1[1], z1 = v1[2]; var xd = vd[0], yd = vd[1], zd = vd[2]; var i000 = getIndex(x0, y0, z0) + itemOffset; var i100 = getIndex(x1, y0, z0) + itemOffset; var i001 = getIndex(x0, y0, z1) + itemOffset; var i101 = getIndex(x1, y0, z1) + itemOffset; var i010 = getIndex(x0, y1, z0) + itemOffset; var i110 = getIndex(x1, y1, z0) + itemOffset; var i011 = getIndex(x0, y1, z1) + itemOffset; var i111 = getIndex(x1, y1, z1) + itemOffset; var o = (i * vertexCount + j) * outputStride; for (var k = 0; k < outputStride; ++k) { var s000 = grid[i000 + k]; var s100 = grid[i100 + k]; var s001 = grid[i001 + k]; var s101 = grid[i101 + k]; var s010 = grid[i010 + k]; var s110 = grid[i110 + k]; var s011 = grid[i011 + k]; var s111 = grid[i111 + k]; var s00 = lerp(s000, s100, xd); var s01 = lerp(s001, s101, xd); var s10 = lerp(s010, s110, xd); var s11 = lerp(s011, s111, xd); var s0 = lerp(s00, s10, yd); var s1 = lerp(s01, s11, yd); array[o + k] = lerp(s0, s1, zd); } } } return image; } // function isSupportedColorType(x) { return x === 'group' || x === 'groupInstance'; } export function applyMeshColorSmoothing(values, resolution, stride, webgl, colorTexture) { if (!isSupportedColorType(values.dColorType.ref.value)) return; var smoothingData = calcMeshColorSmoothing({ vertexCount: values.uVertexCount.ref.value, instanceCount: values.uInstanceCount.ref.value, groupCount: values.uGroupCount.ref.value, transformBuffer: values.aTransform.ref.value, instanceBuffer: values.aInstance.ref.value, positionBuffer: values.aPosition.ref.value, groupBuffer: values.aGroup.ref.value, colorData: values.tColor.ref.value, colorType: values.dColorType.ref.value, boundingSphere: values.boundingSphere.ref.value, invariantBoundingSphere: values.invariantBoundingSphere.ref.value, itemSize: 3 }, resolution, stride, webgl, colorTexture); if (smoothingData.kind === 'volume') { ValueCell.updateIfChanged(values.dColorType, smoothingData.type); ValueCell.update(values.tColorGrid, smoothingData.texture); ValueCell.update(values.uColorTexDim, smoothingData.gridTexDim); ValueCell.update(values.uColorGridDim, smoothingData.gridDim); ValueCell.update(values.uColorGridTransform, smoothingData.gridTransform); } else if (smoothingData.kind === 'vertex') { ValueCell.updateIfChanged(values.dColorType, smoothingData.type); ValueCell.update(values.tColor, smoothingData.texture); ValueCell.update(values.uColorTexDim, smoothingData.texDim); } } function isSupportedOverpaintType(x) { return x === 'groupInstance'; } export function applyMeshOverpaintSmoothing(values, resolution, stride, webgl, colorTexture) { if (!isSupportedOverpaintType(values.dOverpaintType.ref.value)) return; var smoothingData = calcMeshColorSmoothing({ vertexCount: values.uVertexCount.ref.value, instanceCount: values.uInstanceCount.ref.value, groupCount: values.uGroupCount.ref.value, transformBuffer: values.aTransform.ref.value, instanceBuffer: values.aInstance.ref.value, positionBuffer: values.aPosition.ref.value, groupBuffer: values.aGroup.ref.value, colorData: values.tOverpaint.ref.value, colorType: values.dOverpaintType.ref.value, boundingSphere: values.boundingSphere.ref.value, invariantBoundingSphere: values.invariantBoundingSphere.ref.value, itemSize: 4 }, resolution, stride, webgl, colorTexture); if (smoothingData.kind === 'volume') { ValueCell.updateIfChanged(values.dOverpaintType, smoothingData.type); ValueCell.update(values.tOverpaintGrid, smoothingData.texture); ValueCell.update(values.uOverpaintTexDim, smoothingData.gridTexDim); ValueCell.update(values.uOverpaintGridDim, smoothingData.gridDim); ValueCell.update(values.uOverpaintGridTransform, smoothingData.gridTransform); } else if (smoothingData.kind === 'vertex') { ValueCell.updateIfChanged(values.dOverpaintType, smoothingData.type); ValueCell.update(values.tOverpaint, smoothingData.texture); ValueCell.update(values.uOverpaintTexDim, smoothingData.texDim); } } function isSupportedTransparencyType(x) { return x === 'groupInstance'; } export function applyMeshTransparencySmoothing(values, resolution, stride, webgl, colorTexture) { if (!isSupportedTransparencyType(values.dTransparencyType.ref.value)) return; var smoothingData = calcMeshColorSmoothing({ vertexCount: values.uVertexCount.ref.value, instanceCount: values.uInstanceCount.ref.value, groupCount: values.uGroupCount.ref.value, transformBuffer: values.aTransform.ref.value, instanceBuffer: values.aInstance.ref.value, positionBuffer: values.aPosition.ref.value, groupBuffer: values.aGroup.ref.value, colorData: values.tTransparency.ref.value, colorType: values.dTransparencyType.ref.value, boundingSphere: values.boundingSphere.ref.value, invariantBoundingSphere: values.invariantBoundingSphere.ref.value, itemSize: 1 }, resolution, stride, webgl, colorTexture); if (smoothingData.kind === 'volume') { ValueCell.updateIfChanged(values.dTransparencyType, smoothingData.type); ValueCell.update(values.tTransparencyGrid, smoothingData.texture); ValueCell.update(values.uTransparencyTexDim, smoothingData.gridTexDim); ValueCell.update(values.uTransparencyGridDim, smoothingData.gridDim); ValueCell.update(values.uTransparencyGridTransform, smoothingData.gridTransform); } else if (smoothingData.kind === 'vertex') { ValueCell.updateIfChanged(values.dTransparencyType, smoothingData.type); ValueCell.update(values.tTransparency, smoothingData.texture); ValueCell.update(values.uTransparencyTexDim, smoothingData.texDim); } } function isSupportedSubstanceType(x) { return x === 'groupInstance'; } export function applyMeshSubstanceSmoothing(values, resolution, stride, webgl, colorTexture) { if (!isSupportedSubstanceType(values.dSubstanceType.ref.value)) return; var smoothingData = calcMeshColorSmoothing({ vertexCount: values.uVertexCount.ref.value, instanceCount: values.uInstanceCount.ref.value, groupCount: values.uGroupCount.ref.value, transformBuffer: values.aTransform.ref.value, instanceBuffer: values.aInstance.ref.value, positionBuffer: values.aPosition.ref.value, groupBuffer: values.aGroup.ref.value, colorData: values.tSubstance.ref.value, colorType: values.dSubstanceType.ref.value, boundingSphere: values.boundingSphere.ref.value, invariantBoundingSphere: values.invariantBoundingSphere.ref.value, itemSize: 4 }, resolution, stride, webgl, colorTexture); if (smoothingData.kind === 'volume') { ValueCell.updateIfChanged(values.dSubstanceType, smoothingData.type); ValueCell.update(values.tSubstanceGrid, smoothingData.texture); ValueCell.update(values.uSubstanceTexDim, smoothingData.gridTexDim); ValueCell.update(values.uSubstanceGridDim, smoothingData.gridDim); ValueCell.update(values.uSubstanceGridTransform, smoothingData.gridTransform); } else if (smoothingData.kind === 'vertex') { ValueCell.updateIfChanged(values.dSubstanceType, smoothingData.type); ValueCell.update(values.tSubstance, smoothingData.texture); ValueCell.update(values.uSubstanceTexDim, smoothingData.texDim); } }