improved finding of closest element for vertices

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

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Unit, Structure } from 'mol-model/structure';
+import { Unit, Structure, StructureElement, ElementIndex } from 'mol-model/structure';
 import { UnitsVisual, VisualUpdateState } from '..';
 import { RuntimeContext } from 'mol-task'
 import { Mesh } from '../../../geometry/mesh/mesh';
@@ -13,9 +13,11 @@ import { StructureElementIterator, getElementLoci, markElement } from './util/el
 import { computeMarchingCubesMesh } from '../../../util/marching-cubes/algorithm';
 import { GaussianDensityProps, GaussianDensityParams } from 'mol-model/structure/structure/unit/gaussian-density';
 import { paramDefaultValues } from 'mol-view/parameter';
+import { SizeTheme } from 'mol-view/theme/size';
+import { OrderedSet } from 'mol-data/int';
 
 async function createGaussianSurfaceMesh(ctx: RuntimeContext, unit: Unit, structure: Structure, props: GaussianDensityProps, mesh?: Mesh): Promise<Mesh> {
-    const { smoothness } = props
+    const { smoothness, radiusOffset } = props
     const { transform, field, idField } = await unit.computeGaussianDensity(props, ctx)
 
     const params = {
@@ -28,16 +30,30 @@ async function createGaussianSurfaceMesh(ctx: RuntimeContext, unit: Unit, struct
     Mesh.transformImmediate(surface, transform)
 
     if (props.useGpu) {
-        console.time('find closest atom for vertices')
+        console.time('find max element radius')
+        const { elements } = unit
+        const n = OrderedSet.size(elements)
+        const l = StructureElement.create(unit)
+        const sizeTheme = SizeTheme({ name: 'physical' })
+        const radius = (index: number) => {
+            l.element = index as ElementIndex
+            return sizeTheme.size(l)
+        }
+        let maxRadius = 0
+        for (let i = 0; i < n; ++i) {
+            const r = radius(OrderedSet.getAt(elements, i)) + radiusOffset
+            if (maxRadius < r) maxRadius = r
+        }
+        console.timeEnd('find max element radius')
+
+        console.time('find closest element for vertices')
         const { lookup3d } = unit
-        const maxRadius = 2
-        const maxRadiusSq = maxRadius * maxRadius
 
         const { vertexCount, vertexBuffer, groupBuffer } = surface
         const vertices = vertexBuffer.ref.value
         const groups = groupBuffer.ref.value
         for (let i = 0; i < vertexCount; ++i) {
-            const r = lookup3d.find(vertices[i * 3], vertices[i * 3 + 1], vertices[i * 3 + 2], maxRadiusSq)
+            const r = lookup3d.find(vertices[i * 3], vertices[i * 3 + 1], vertices[i * 3 + 2], maxRadius * 2)
             let minDsq = Infinity
             let group = 0
             for (let j = 0, jl = r.count; j < jl; ++j) {
@@ -49,7 +65,7 @@ async function createGaussianSurfaceMesh(ctx: RuntimeContext, unit: Unit, struct
             }
             groups[i] = group
         }
-        console.timeEnd('find closest atom for vertices')
+        console.timeEnd('find closest element for vertices')
     }
 
     Mesh.computeNormalsImmediate(surface)