wip, color smoothing experiments

src/mol-geo/geometry/color-data.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -7,18 +7,20 @@
 import { ValueCell } from '../../mol-util';
 import { TextureImage, createTextureImage } from '../../mol-gl/renderable/util';
 import { Color } from '../../mol-util/color';
-import { Vec2, Vec3 } from '../../mol-math/linear-algebra';
+import { Vec2, Vec3, Vec4 } from '../../mol-math/linear-algebra';
 import { LocationIterator } from '../util/location-iterator';
 import { NullLocation } from '../../mol-model/location';
 import { LocationColor, ColorTheme } from '../../mol-theme/color';
 import { Geometry } from './geometry';
 
-export type ColorType = 'uniform' | 'instance' | 'group' | 'groupInstance' | 'vertex' | 'vertexInstance'
+export type ColorType = 'uniform' | 'instance' | 'group' | 'groupInstance' | 'vertex' | 'vertexInstance' | 'volume'
 
 export type ColorData = {
     uColor: ValueCell<Vec3>,
     tColor: ValueCell<TextureImage<Uint8Array>>,
     uColorTexDim: ValueCell<Vec2>,
+    uColorGridDim: ValueCell<Vec3>,
+    uColorGridTransform: ValueCell<Vec4>,
     dColorType: ValueCell<string>,
 }
 
@@ -30,9 +32,12 @@ export function createColors(locationIt: LocationIterator, positionIt: LocationI
         case 'groupInstance': return createGroupInstanceColor(locationIt, colorTheme.color, colorData);
         case 'vertex': return createVertexColor(positionIt, colorTheme.color, colorData);
         case 'vertexInstance': return createVertexInstanceColor(positionIt, colorTheme.color, colorData);
+        case 'volume': return createGridColor((colorTheme as any).grid, 'volume', colorData);
     }
 }
 
+//
+
 export function createValueColor(value: Color, colorData?: ColorData): ColorData {
     if (colorData) {
         ValueCell.update(colorData.uColor, Color.toVec3Normalized(colorData.uColor.ref.value, value));
@@ -43,6 +48,8 @@ export function createValueColor(value: Color, colorData?: ColorData): ColorData
             uColor: ValueCell.create(Color.toVec3Normalized(Vec3(), value)),
             tColor: ValueCell.create({ array: new Uint8Array(3), width: 1, height: 1 }),
             uColorTexDim: ValueCell.create(Vec2.create(1, 1)),
+            uColorGridDim: ValueCell.create(Vec3.create(1, 1, 1)),
+            uColorGridTransform: ValueCell.create(Vec4.create(0, 0, 0, 1)),
             dColorType: ValueCell.create('uniform'),
         };
     }
@@ -53,6 +60,8 @@ export function createUniformColor(locationIt: LocationIterator, color: Location
     return createValueColor(color(NullLocation, false), colorData);
 }
 
+//
+
 export function createTextureColor(colors: TextureImage<Uint8Array>, type: ColorType, colorData?: ColorData): ColorData {
     if (colorData) {
         ValueCell.update(colorData.tColor, colors);
@@ -64,6 +73,8 @@ export function createTextureColor(colors: TextureImage<Uint8Array>, type: Color
             uColor: ValueCell.create(Vec3()),
             tColor: ValueCell.create(colors),
             uColorTexDim: ValueCell.create(Vec2.create(colors.width, colors.height)),
+            uColorGridDim: ValueCell.create(Vec3.create(1, 1, 1)),
+            uColorGridTransform: ValueCell.create(Vec4.create(0, 0, 0, 1)),
             dColorType: ValueCell.create(type),
         };
     }
@@ -138,3 +149,36 @@ export function createVertexInstanceColor(locationIt: LocationIterator, color: L
     }
     return createTextureColor(colors, 'vertexInstance', colorData);
 }
+
+//
+
+interface ColorVolume {
+    colors: TextureImage<Uint8Array> // | Texture
+    dimension: Vec3
+    transform: Vec4
+}
+
+export function createGridColor(grid: ColorVolume, type: ColorType, colorData?: ColorData): ColorData {
+    const { colors, dimension, transform } = grid;
+    // const width = 'array' in colors ? colors.width : colors.getWidth();
+    // const height = 'array' in colors ? colors.height : colors.getHeight();
+    const width = colors.width;
+    const height = colors.height;
+    if (colorData) {
+        ValueCell.update(colorData.tColor, colors);
+        ValueCell.update(colorData.uColorTexDim, Vec2.create(width, height));
+        ValueCell.update(colorData.uColorGridDim, Vec3.clone(dimension));
+        ValueCell.update(colorData.uColorGridTransform, Vec4.clone(transform));
+        ValueCell.updateIfChanged(colorData.dColorType, type);
+        return colorData;
+    } else {
+        return {
+            uColor: ValueCell.create(Vec3()),
+            tColor: ValueCell.create(colors),
+            uColorTexDim: ValueCell.create(Vec2.create(width, height)),
+            uColorGridDim: ValueCell.create(Vec3.clone(dimension)),
+            uColorGridTransform: ValueCell.create(Vec4.clone(transform)),
+            dColorType: ValueCell.create(type),
+        };
+    }
+}

src/mol-geo/geometry/mesh/mesh.ts

@@ -1,13 +1,13 @@
 /**
- * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
 import { ValueCell } from '../../../mol-util';
-import { Vec3, Mat4, Mat3, Vec4 } from '../../../mol-math/linear-algebra';
-import { Sphere3D } from '../../../mol-math/geometry';
+import { Vec3, Mat4, Mat3, Vec4, Vec2 } from '../../../mol-math/linear-algebra';
+import { Box3D, Sphere3D } from '../../../mol-math/geometry';
 import { transformPositionArray, transformDirectionArray, computeIndexedVertexNormals, GroupMapping, createGroupMapping} from '../../util';
 import { GeometryUtils } from '../geometry';
 import { createMarkers } from '../marker-data';
@@ -16,7 +16,7 @@ import { LocationIterator, PositionLocation } from '../../util/location-iterator
 import { createColors } from '../color-data';
 import { ChunkedArray, hashFnv32a } from '../../../mol-data/util';
 import { ParamDefinition as PD } from '../../../mol-util/param-definition';
-import { calculateInvariantBoundingSphere, calculateTransformBoundingSphere } from '../../../mol-gl/renderable/util';
+import { calculateInvariantBoundingSphere, calculateTransformBoundingSphere, createTextureImage, TextureImage } from '../../../mol-gl/renderable/util';
 import { Theme } from '../../../mol-theme/theme';
 import { MeshValues } from '../../../mol-gl/renderable/mesh';
 import { Color } from '../../../mol-util/color';
@@ -25,6 +25,8 @@ import { createEmptyOverpaint } from '../overpaint-data';
 import { createEmptyTransparency } from '../transparency-data';
 import { createEmptyClipping } from '../clipping-data';
 import { RenderableState } from '../../../mol-gl/renderable';
+import { getVolumeTexture2dLayout } from '../../../mol-repr/volume/util';
+import { arraySetAdd } from '../../../mol-util/array';
 
 export interface Mesh {
     readonly kind: 'mesh',
@@ -334,6 +336,7 @@ export namespace Mesh {
         flatShaded: PD.Boolean(false, BaseGeometry.ShadingCategory),
         ignoreLight: PD.Boolean(false, BaseGeometry.ShadingCategory),
         xrayShaded: PD.Boolean(false, BaseGeometry.ShadingCategory),
+        smoothColors: PD.Select('off', PD.arrayToOptions(['off', 'vertex', 'volume'] as const), { isHidden: true }),
     };
     export type Params = typeof Params
 
@@ -385,7 +388,7 @@ export namespace Mesh {
         const invariantBoundingSphere = Sphere3D.clone(mesh.boundingSphere);
         const boundingSphere = calculateTransformBoundingSphere(invariantBoundingSphere, transform.aTransform.ref.value, instanceCount);
 
-        return {
+        const values = {
             aPosition: mesh.vertexBuffer,
             aNormal: mesh.normalBuffer,
             aGroup: mesh.groupBuffer,
@@ -407,6 +410,197 @@ export namespace Mesh {
             dIgnoreLight: ValueCell.create(props.ignoreLight),
             dXrayShaded: ValueCell.create(props.xrayShaded),
         };
+
+        return processValues(values, props);
+    }
+
+    function processValues(values: MeshValues, props: PD.Values<Params>) {
+        if (props.smoothColors === 'vertex') {
+            values = vertexColors(values, props);
+        } else if (props.smoothColors === 'volume') {
+            values = gridColors(values, props);
+        }
+        return values;
+    }
+
+    function gridColors(values: MeshValues, props: PD.Values<Params>): MeshValues {
+        const type = values.dColorType.ref.value;
+        if (type === 'uniform' || type.startsWith('vertex')) return values;
+        // TODO handle instance and vertex types
+
+        const box = Box3D.fromSphere3D(Box3D(), values.invariantBoundingSphere.ref.value);
+
+        const resolution = (props as any).resolution ?? 0.5; // TODO better default
+        const scaleFactor = 1 / resolution;
+        const scaledBox = Box3D.scale(Box3D(), box, scaleFactor);
+        const dim = Box3D.size(Vec3(), scaledBox);
+        Vec3.ceil(dim, dim);
+        const { min } = box;
+
+        const [ xn, yn ] = dim;
+        const { width, height } = getVolumeTexture2dLayout(dim);
+
+        const itemSize = 3;
+        const data = new Float32Array(width * height * itemSize);
+        const count = new Float32Array(width * height);
+
+        const array = new Uint8Array(width * height * itemSize);
+        const textureImage = { array, width, height };
+
+        const vertexCount = values.uVertexCount.ref.value;
+        const positions = values.aPosition.ref.value;
+        const groupColors = values.tColor.ref.value.array;
+        const groups = values.aGroup.ref.value;
+
+        function getIndex(x: number, y: number, z: number) {
+            const column = Math.floor(((z * xn) % width) / xn);
+            const row = Math.floor((z * xn) / width);
+            const px = column * xn + x;
+            return itemSize * ((row * yn * width) + (y * width) + px);
+        }
+
+        const p = 2;
+        const f = Math.sqrt(2) * 2 * p * resolution;
+
+        for (let i = 0; i < vertexCount; ++i) {
+            // vertex mapped to grid
+            const x = Math.floor(scaleFactor * (positions[i * 3] - min[0]));
+            const y = Math.floor(scaleFactor * (positions[i * 3 + 1] - min[1]));
+            const z = Math.floor(scaleFactor * (positions[i * 3 + 2] - min[2]));
+
+            for (let xi = x - p, xl = x + p; xi < xl; ++ xi) {
+                for (let yi = y - p, yl = y + p; yi < yl; ++ yi) {
+                    for (let zi = z - p, zl = z + p; zi < zl; ++ zi) {
+                        const index = getIndex(xi, yi, zi);
+
+                        const dx = min[0] + (resolution * xi) - positions[i * 3];
+                        const dy = min[1] + (resolution * yi) - positions[i * 3 + 1];
+                        const dz = min[2] + (resolution * zi) - positions[i * 3 + 2];
+                        const d = Math.sqrt(dx * dx + dy * dy + dz * dz);
+
+                        const s = 2 - (d / f);
+
+                        const g = groups[i];
+                        data[index] += groupColors[g * 3] * s;
+                        data[index + 1] += groupColors[g * 3 + 1] * s;
+                        data[index + 2] += groupColors[g * 3 + 2] * s;
+                        count[index / 3] += 1 * s;
+                    }
+                }
+            }
+        }
+
+        for (let i = 0, il = count.length; i < il; ++i) {
+            const i3 = i * 3;
+            const c = count[i];
+            array[i3] = Math.floor(data[i3] / c);
+            array[i3 + 1] = Math.floor(data[i3 + 1] / c);
+            array[i3 + 2] = Math.floor(data[i3 + 2] / c);
+        }
+
+        ValueCell.updateIfChanged(values.dColorType, 'volume');
+        ValueCell.update(values.tColor, textureImage);
+        ValueCell.update(values.uColorTexDim, Vec2.create(width, height));
+        ValueCell.update(values.uColorGridDim, dim);
+        ValueCell.update(values.uColorGridTransform, Vec4.create(min[0], min[1], min[2], scaleFactor));
+
+        return values;
+    }
+
+    function getNeighboursMap(values: MeshValues) {
+        const vertexCount = values.uVertexCount.ref.value;
+        const elements = values.elements.ref.value;
+
+        const neighboursMap: number[][] = [];
+        for (let i = 0; i < vertexCount; ++i) {
+            neighboursMap[i] = [];
+        }
+
+        for (let i = 0; i < values.drawCount.ref.value / 3; ++i) {
+            const v1 = elements[i * 3];
+            const v2 = elements[i * 3 + 1];
+            const v3 = elements[i * 3 + 2];
+            arraySetAdd(neighboursMap[v1], v1);
+            arraySetAdd(neighboursMap[v1], v2);
+            arraySetAdd(neighboursMap[v1], v3);
+            arraySetAdd(neighboursMap[v2], v1);
+            arraySetAdd(neighboursMap[v2], v2);
+            arraySetAdd(neighboursMap[v2], v3);
+            arraySetAdd(neighboursMap[v3], v1);
+            arraySetAdd(neighboursMap[v3], v2);
+            arraySetAdd(neighboursMap[v3], v3);
+        }
+        return neighboursMap;
+    }
+
+    function vertexColors(values: MeshValues, props: PD.Values<Params>): MeshValues {
+        const type = values.dColorType.ref.value;
+        if (type === 'uniform' || type.startsWith('volume')) return values;
+
+        const vertexCount = values.uVertexCount.ref.value;
+        const groupCount = values.uGroupCount.ref.value;
+        const instanceCount = type.endsWith('Instance') ? values.uInstanceCount.ref.value : 1;
+        const iterations = Math.round((vertexCount / groupCount) / 20); // TODO better formula?
+        const neighboursMap = getNeighboursMap(values);
+
+        let vertexColors: Uint8Array;
+        let colors: TextureImage<Uint8Array>;
+
+        if (type.startsWith('vertex')) {
+            vertexColors = values.tColor.ref.value.array;
+            colors = values.tColor.ref.value;
+        } else {
+            colors = createTextureImage(vertexCount * instanceCount, 3, Uint8Array);
+            vertexColors = colors.array;
+
+            const groupColors = values.tColor.ref.value.array;
+            const groups = values.aGroup.ref.value;
+
+            for (let i = 0; i < instanceCount; ++i) {
+                const og = i * groupCount * 3;
+                const ov = i * vertexCount * 3;
+                for (let j = 0; j < vertexCount; ++j) {
+                    const neighbours = neighboursMap[j];
+                    const nc = neighbours.length;
+                    let r = 0, g = 0, b = 0;
+                    for (let k = 0; k < nc; ++k) {
+                        const neighbourGroup = groups[neighbours[k]];
+                        r += groupColors[og + neighbourGroup * 3];
+                        g += groupColors[og + neighbourGroup * 3 + 1];
+                        b += groupColors[og + neighbourGroup * 3 + 2];
+                    }
+                    vertexColors[ov + j * 3] = Math.round(r / nc);
+                    vertexColors[ov + j * 3 + 1] = Math.round(g / nc);
+                    vertexColors[ov + j * 3 + 2] = Math.round(b / nc);
+                }
+            }
+        }
+
+        for (let x = 0; x < iterations; ++x) {
+            for (let i = 0; i < instanceCount; ++i) {
+                const ov = i * vertexCount * 3;
+                for (let j = 0; j < vertexCount; ++j) {
+                    const neighbours = neighboursMap[j];
+                    const nc = neighbours.length;
+                    let r = 0, g = 0, b = 0;
+                    for (let k = 0; k < nc; ++k) {
+                        const neighbour = neighbours[k];
+                        r += vertexColors[ov + neighbour * 3];
+                        g += vertexColors[ov + neighbour * 3 + 1];
+                        b += vertexColors[ov + neighbour * 3 + 2];
+                    }
+                    vertexColors[ov + j * 3] = Math.round(r / nc);
+                    vertexColors[ov + j * 3 + 1] = Math.round(g / nc);
+                    vertexColors[ov + j * 3 + 2] = Math.round(b / nc);
+                }
+            }
+        }
+
+        ValueCell.updateIfChanged(values.dColorType, type.endsWith('Instance') ? 'vertexInstance' : 'vertex');
+        ValueCell.update(values.tColor, colors);
+        ValueCell.update(values.uColorTexDim, Vec2.create(colors.width, colors.height));
+
+        return values;
     }
 
     function createValuesSimple(mesh: Mesh, props: Partial<PD.Values<Params>>, colorValue: Color, sizeValue: number, transform?: TransformData) {
@@ -422,6 +616,8 @@ export namespace Mesh {
         ValueCell.updateIfChanged(values.dFlipSided, props.flipSided);
         ValueCell.updateIfChanged(values.dIgnoreLight, props.ignoreLight);
         ValueCell.updateIfChanged(values.dXrayShaded, props.xrayShaded);
+
+        processValues(values, props);
     }
 
     function updateBoundingSphere(values: MeshValues, mesh: Mesh) {

src/mol-gl/renderable/schema.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -184,6 +184,8 @@ export const ColorSchema = {
     // aColor: AttributeSpec('float32', 3, 0), // TODO
     uColor: UniformSpec('v3', 'material'),
     uColorTexDim: UniformSpec('v2'),
+    uColorGridDim: UniformSpec('v3'),
+    uColorGridTransform: UniformSpec('v4'),
     tColor: TextureSpec('image-uint8', 'rgb', 'ubyte', 'nearest'),
     dColorType: DefineSpec('string', ['uniform', 'attribute', 'instance', 'group', 'groupInstance', 'vertex', 'vertexInstance']),
 } as const;

src/mol-gl/shader/chunks/assign-color-varying.glsl.ts

@@ -12,6 +12,9 @@ export const assign_color_varying = `
         vColor.rgb = readFromTexture(tColor, VertexID, uColorTexDim).rgb;
     #elif defined(dColorType_vertexInstance)
         vColor.rgb = readFromTexture(tColor, int(aInstance) * uVertexCount + VertexID, uColorTexDim).rgb;
+    #elif defined(dColorType_grid)
+        vec3 gridPos = (uColorGridTransform.w * (position - uColorGridTransform.xyz)) / uColorGridDim;
+        vColor.rgb = texture3dFrom2dLinear(tColor, gridPos, uColorGridDim, uColorTexDim).rgb;
     #endif
 
     #ifdef dOverpaint

src/mol-gl/shader/chunks/color-vert-params.glsl.ts

@@ -9,6 +9,12 @@ export const color_vert_params = `
         varying vec4 vColor;
         uniform vec2 uColorTexDim;
         uniform sampler2D tColor;
+    #elif defined(dColorType_grid)
+        varying vec4 vColor;
+        uniform vec2 uColorTexDim;
+        uniform vec3 uColorGridDim;
+        uniform vec4 uColorGridTransform;
+        uniform sampler2D tColor;
     #endif
 
     #ifdef dOverpaint

src/mol-gl/shader/chunks/common.glsl.ts

@@ -13,7 +13,11 @@ export const common = `
     #define dColorType_texture
 #endif
 
-#if defined(dColorType_attribute) || defined(dColorType_texture)
+#if defined(dColorType_volume)
+    #define dColorType_grid
+#endif
+
+#if defined(dColorType_attribute) || defined(dColorType_texture) || defined(dColorType_volume)
     #define dColorType_varying
 #endif
 

src/mol-gl/shader/mesh.vert.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -15,6 +15,10 @@ precision highp sampler2D;
 #include color_vert_params
 #include common_clip
 
+#if defined(dColorType_volume)
+    #include texture3d_from_2d_linear
+#endif
+
 #ifdef dGeoTexture
     uniform vec2 uGeoTexDim;
     uniform sampler2D tPosition;
@@ -32,10 +36,10 @@ varying vec3 vNormal;
 
 void main(){
     #include assign_group
-    #include assign_color_varying
     #include assign_marker_varying
     #include assign_clipping_varying
     #include assign_position
+    #include assign_color_varying
     #include clip_instance
 
     #ifdef dGeoTexture

src/mol-repr/structure/representation/gaussian-surface.ts

@@ -22,6 +22,7 @@ const GaussianSurfaceVisuals = {
 export const GaussianSurfaceParams = {
     ...GaussianSurfaceMeshParams,
     ...GaussianWireframeParams,
+    smoothColors: PD.Select('off', PD.arrayToOptions(['off', 'vertex', 'volume'] as const)),
     visuals: PD.MultiSelect(['gaussian-surface-mesh'], PD.objectToOptions(GaussianSurfaceVisuals)),
 };
 export type GaussianSurfaceParams = typeof GaussianSurfaceParams

src/mol-repr/structure/representation/molecular-surface.ts

@@ -21,6 +21,7 @@ const MolecularSurfaceVisuals = {
 export const MolecularSurfaceParams = {
     ...MolecularSurfaceMeshParams,
     ...MolecularSurfaceWireframeParams,
+    smoothColors: PD.Select('off', PD.arrayToOptions(['off', 'vertex', 'volume'] as const)),
     visuals: PD.MultiSelect(['molecular-surface-mesh'], PD.objectToOptions(MolecularSurfaceVisuals)),
 };
 export type MolecularSurfaceParams = typeof MolecularSurfaceParams

src/mol-repr/structure/units-visual.ts

@@ -157,7 +157,11 @@ export function UnitsVisual<G extends Geometry, P extends StructureParams & Geom
                 updateState.updateColor = true;
                 updateState.updateSize = true;
             }
-            if (newTheme.color.granularity.startsWith('vertex')) {
+            if (newTheme.color.granularity.startsWith('vertex') ||
+                renderObject.values.dColorType.ref.value.startsWith('vertex') ||
+                newTheme.color.granularity.startsWith('volume') ||
+                renderObject.values.dColorType.ref.value.startsWith('volume')
+            ) {
                 updateState.updateColor = true;
             }
         }
@@ -315,6 +319,7 @@ export function UnitsMeshVisual<P extends UnitsMeshParams>(builder: UnitsMeshVis
         setUpdateState: (state: VisualUpdateState, newProps: PD.Values<P>, currentProps: PD.Values<P>, newTheme: Theme, currentTheme: Theme, newStructureGroup: StructureGroup, currentStructureGroup: StructureGroup) => {
             builder.setUpdateState(state, newProps, currentProps, newTheme, currentTheme, newStructureGroup, currentStructureGroup);
             if (!SizeTheme.areEqual(newTheme.size, currentTheme.size)) state.createGeometry = true;
+            if (newProps.smoothColors !== currentProps.smoothColors) state.updateColor = true;
         },
         geometryUtils: Mesh.Utils
     }, materialId);