wip, gaussian surface picking / coloring

src/apps/canvas/structure-view.ts

@@ -68,7 +68,7 @@ export async function StructureView(viewer: Viewer, models: ReadonlyArray<Model>
     const active: { [k: string]: boolean } = {
         cartoon: true,
         point: false,
-        surface: false,
+        surface: true,
         ballAndStick: false,
         carbohydrate: false,
         spacefill: false,
@@ -209,7 +209,7 @@ export async function StructureView(viewer: Viewer, models: ReadonlyArray<Model>
             console.log('createStructureRepr')
             for (const k in structureRepresentations) {
                 if (active[k]) {
-                    await structureRepresentations[k].createOrUpdate({}, structure).run(
+                    await structureRepresentations[k].createOrUpdate({ colorTheme: { name: 'element-index' } }, structure).run(
                         // progress => console.log(Progress.format(progress))
                     )
                     viewer.add(structureRepresentations[k])

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

@@ -15,6 +15,7 @@ import { createMarkers } from '../marker-data';
 import { TransformData } from '../transform-data';
 import { LocationIterator } from '../../util/location-iterator';
 import { createColors } from '../color-data';
+import { ChunkedArray } from 'mol-data/util';
 
 export interface Mesh {
     readonly kind: 'mesh',
@@ -160,6 +161,71 @@ export namespace Mesh {
         });
     }
 
+    /**
+     * Ensure that each vertices of each triangle have the same group id.
+     * Note that normals are copied over and can't be re-created from the new mesh.
+     */
+    export function uniformTriangleGroup(mesh: Mesh) {
+        const { indexBuffer, vertexBuffer, groupBuffer, normalBuffer, triangleCount, vertexCount } = mesh
+        const ib = indexBuffer.ref.value
+        const vb = vertexBuffer.ref.value
+        const gb = groupBuffer.ref.value
+        const nb = normalBuffer.ref.value
+
+        // new
+        const index = ChunkedArray.create(Uint32Array, 3, 1024, triangleCount)
+
+        // re-use
+        const vertex = ChunkedArray.create(Float32Array, 3, 1024, vb)
+        vertex.currentIndex = vertexCount * 3
+        vertex.elementCount = vertexCount
+        const normal = ChunkedArray.create(Float32Array, 3, 1024, nb)
+        normal.currentIndex = vertexCount * 3
+        normal.elementCount = vertexCount
+        const group = ChunkedArray.create(Float32Array, 1, 1024, gb)
+        group.currentIndex = vertexCount
+        group.elementCount = vertexCount
+
+        const v = Vec3.zero()
+        const n = Vec3.zero()
+
+        function add(i: number) {
+            Vec3.fromArray(v, vb, i * 3)
+            Vec3.fromArray(n, nb, i * 3)
+            ChunkedArray.add3(vertex, v[0], v[1], v[2])
+            ChunkedArray.add3(normal, n[0], n[1], n[2])
+        }
+
+        let newVertexCount = vertexCount
+        for (let i = 0, il = triangleCount; i < il; ++i) {
+            const i0 = ib[i * 3], i1 = ib[i * 3 + 1], i2 = ib[i * 3 + 2]
+            const g0 = gb[i0], g1 = gb[i1], g2 = gb[i2]
+            if (g0 !== g1 || g0 !== g2) {
+                add(i0); add(i1); add(i2)
+                ChunkedArray.add3(index, newVertexCount, newVertexCount + 1, newVertexCount + 2)
+                const g = g1 === g2 ? g1 : g0
+                for (let j = 0; j < 3; ++j) ChunkedArray.add(group, g)
+                newVertexCount += 3
+            } else {
+                ChunkedArray.add3(index, i0, i1, i2)
+            }
+        }
+
+        const newIb = ChunkedArray.compact(index)
+        const newVb = ChunkedArray.compact(vertex)
+        const newNb = ChunkedArray.compact(normal)
+        const newGb = ChunkedArray.compact(group)
+
+        mesh.vertexCount = newVertexCount
+
+        ValueCell.update(vertexBuffer, newVb) as ValueCell<Float32Array>
+        ValueCell.update(groupBuffer, newGb) as ValueCell<Float32Array>
+        ValueCell.update(indexBuffer, newIb) as ValueCell<Uint32Array>
+        ValueCell.update(normalBuffer, newNb) as ValueCell<Float32Array>
+
+        return mesh
+    }
+
     //
 
     export const DefaultProps = {

src/mol-geo/representation/structure/visual/gaussian-surface-mesh.ts

@@ -15,17 +15,24 @@ import { computeGaussianDensity, DefaultGaussianDensityProps } from './util/gaus
 
 async function createGaussianSurfaceMesh(ctx: RuntimeContext, unit: Unit, structure: Structure, props: GaussianSurfaceProps, mesh?: Mesh): Promise<Mesh> {
     const { smoothness } = props
-    const { transform, field } = await computeGaussianDensity(unit, structure, props).runAsChild(ctx)
+    const { transform, field, idField } = await computeGaussianDensity(unit, structure, props).runAsChild(ctx)
 
+    console.time('mc')
     const surface = await computeMarchingCubes({
         isoLevel: Math.exp(-smoothness),
         scalarField: field,
+        idField,
         oldSurface: mesh
     }).runAsChild(ctx)
+    console.timeEnd('mc')
 
     Mesh.transformImmediate(surface, transform)
     Mesh.computeNormalsImmediate(surface)
 
+    console.time('uniformTriangleGroup')
+    Mesh.uniformTriangleGroup(surface)
+    console.timeEnd('uniformTriangleGroup')
+
     return surface;
 }
 

src/mol-geo/representation/structure/visual/util/gaussian.ts

@@ -11,7 +11,7 @@ import { Box3D } from 'mol-math/geometry';
 import { SizeTheme } from 'mol-view/theme/size';
 
 export const DefaultGaussianDensityProps = {
-    resolutionFactor: 7,
+    resolutionFactor: 6,
     radiusOffset: 0,
     smoothness: 1.5,
 }
@@ -28,7 +28,7 @@ function getDelta(box: Box3D, resolutionFactor: number) {
     return delta
 }
 
-type Density = { transform: Mat4, field: Tensor }
+type Density = { transform: Mat4, field: Tensor, idField: Tensor }
 
 export function computeGaussianDensity(unit: Unit, structure: Structure, props: GaussianDensityProps) {
     return Task.create('Gaussian Density', async ctx => {
@@ -74,6 +74,7 @@ export async function GaussianDensity(ctx: RuntimeContext, unit: Unit, structure
     const beg = Vec3.zero()
     const end = Vec3.zero()
 
+    console.time('density grid')
     for (let i = 0; i < elementCount; i++) {
         l.element = elements[i]
         pos(elements[i], v)
@@ -109,13 +110,59 @@ export async function GaussianDensity(ctx: RuntimeContext, unit: Unit, structure
             await ctx.update({ message: 'filling density grid', current: i, max: elementCount });
         }
     }
+    console.timeEnd('density grid')
 
     const t = Mat4.identity()
     Mat4.fromScaling(t, Vec3.inverse(Vec3.zero(), delta))
     Mat4.setTranslation(t, expandedBox.min)
 
+    const { dimensions, get } = space
+    const [ xn, yn, zn ] = dimensions
+
+    const n = xn * yn * zn
+    const idData = space.create()
+    const idField = Tensor.create(space, idData)
+
+    const lookup3d = unit.lookup3d
+
+    let i = 0
+
+    // TODO use max radius of radiusTheme instead of 4
+    const _rSq = 4 + 1 / Math.max(...delta)
+    const _minDsq = _rSq * 4
+
+    console.time('id grid')
+    for (let x = 0; x < xn; ++x) {
+        for (let y = 0; y < yn; ++y) {
+            for (let z = 0; z < zn; ++z) {
+                const dens = get(data, x, y, z)
+                let group = -1
+                if (dens > 0 && dens < 2) {
+                    Vec3.set(p, x, y, z)
+                    Vec3.transformMat4(p, p, t)
+                    const r = lookup3d.find(p[0], p[1], p[2], _rSq)
+                    let minDsq = _minDsq
+                    for (let j = 0, jl = r.count; j < jl; ++j) {
+                        const dSq = r.squaredDistances[j]
+                        if (dSq < minDsq) {
+                            minDsq = dSq
+                            group = r.indices[j]
+                        }
+                    }
+                }
+                space.set(idData, x, y, z, group)
+                if (i % 1000000 === 0 && ctx.shouldUpdate) {
+                    await ctx.update({ message: 'filling id grid', current: i, max: n });
+                }
+                ++i
+            }
+        }
+    }
+    console.timeEnd('id grid')
+
     return {
         field,
+        idField,
         transform: t
     }
 }

src/mol-math/geometry/lookup3d/grid.ts

@@ -212,7 +212,7 @@ function build(data: PositionData, cellSize?: Vec3) {
 }
 
 interface QueryContext {
-    structure: Grid3D,
+    grid: Grid3D,
     x: number,
     y: number,
     z: number,
@@ -221,12 +221,12 @@ interface QueryContext {
     isCheck: boolean
 }
 
-function createContext(structure: Grid3D): QueryContext {
-    return { structure, x: 0.1, y: 0.1, z: 0.1, radius: 0.1, result: Result.create(), isCheck: false }
+function createContext(grid: Grid3D): QueryContext {
+    return { grid, x: 0.1, y: 0.1, z: 0.1, radius: 0.1, result: Result.create(), isCheck: false }
 }
 
 function query(ctx: QueryContext): boolean {
-    const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices, radius }, delta, maxRadius } = ctx.structure;
+    const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices, radius }, delta, maxRadius } = ctx.grid;
     const { radius: inputRadius, isCheck, x, y, z, result } = ctx;
 
     const r = inputRadius + maxRadius;