wip, gaussian surface

src/mol-math/geometry/common.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
@@ -31,4 +31,12 @@ export type DensityTextureData = {
     texture: Texture,
     bbox: Box3D,
     gridDimension: Vec3
+}
+
+export function fillGridDim(length: number, start: number, step: number) {
+    const a = new Float32Array(length)
+    for (let i = 0; i < a.length; i++) {
+        a[i] = start + (step * i)
+    }
+    return a
 }

src/mol-math/geometry/gaussian-density/cpu.ts

@@ -1,10 +1,10 @@
 /**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Box3D } from '../../geometry';
+import { Box3D, fillGridDim } from '../../geometry';
 import { Vec3, Mat4, Tensor } from '../../linear-algebra';
 import { RuntimeContext } from 'mol-task';
 import { PositionData, DensityData } from '../common';
@@ -13,6 +13,7 @@ import { GaussianDensityProps, getDelta } from '../gaussian-density';
 
 export async function GaussianDensityCPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
     const { resolution, radiusOffset, smoothness } = props
+    const scaleFactor = 1 / resolution
 
     const { indices, x, y, z } = position
     const n = OrderedSet.size(indices)
@@ -31,11 +32,11 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
 
     const pad = maxRadius * 2 + resolution
     const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad))
-    const extent = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min)
+    const extent = Vec3.sub(Vec3(), expandedBox.max, expandedBox.min)
     const min = expandedBox.min
 
     const delta = getDelta(Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad)), resolution)
-    const dim = Vec3.zero()
+    const dim = Vec3()
     Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
     // console.log('grid dim cpu', dim)
 
@@ -46,60 +47,59 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
     const idData = space.create()
     const idField = Tensor.create(space, idData)
 
-    const iu = dim[2], iv = dim[1], iuv = iu * iv
+    const [ dimX, dimY, dimZ ] = dim
+    const iu = dimZ, iv = dimY, iuv = iu * iv
 
-    const densData = space.create()
+    const gridx = fillGridDim(dim[0], min[0], resolution)
+    const gridy = fillGridDim(dim[1], min[1], resolution)
+    const gridz = fillGridDim(dim[2], min[2], resolution)
 
-    const v = Vec3()
-    const c = Vec3()
+    const densData = space.create()
 
     const alpha = smoothness
+    const updateChunk = Math.ceil(1000000 / (Math.pow(maxRadius * 2, 3) * resolution))
 
-    const _r2 = maxRadius * 2
-    const _radius2 = Vec3.create(_r2, _r2, _r2)
-    Vec3.mul(_radius2, _radius2, delta)
-    const updateChunk = Math.ceil(1000000 / (_radius2[0] * _radius2[1] * _radius2[2]))
-
-    const beg = Vec3()
-    const end = Vec3()
-    const rad2 = Vec3()
-
-    const gridPad = 1 / Math.max(...delta)
-
-    const invDelta = Vec3.inverse(Vec3(), delta)
-    const [ invDeltaX, invDeltaY, invDeltaZ ] = invDelta
-
-    console.time('gaussian density cpu')
+    // console.time('gaussian density cpu')
     for (let i = 0; i < n; ++i) {
         const j = OrderedSet.getAt(indices, i)
-
-        Vec3.set(v, x[j], y[j], z[j])
-        Vec3.sub(v, v, min)
-        const [ vx, vy, vz ] = v
-
-        Vec3.mul(c, v, delta)
+        const vx = x[j], vy = y[j], vz = z[j]
 
         const rad = radii[i]
         const rSq = rad * rad
         const rSqInv = 1 / rSq
 
-        const r2 = radiusOffset + rad * 2 + gridPad
+        const r2 = rad * 2
         const r2sq = r2 * r2
-        Vec3.set(rad2, r2, r2, r2)
-        Vec3.mul(rad2, rad2, delta)
 
-        const [ begX, begY, begZ ] = Vec3.floor(beg, Vec3.sub(beg, c, rad2))
-        const [ endX, endY, endZ ] = Vec3.ceil(end, Vec3.add(end, c, rad2))
+        // Number of grid points, round this up...
+        const ng = Math.ceil(r2 * scaleFactor)
+
+        // Center of the atom, mapped to grid points (take floor)
+        const iax = Math.floor(scaleFactor * (vx - min[0]))
+        const iay = Math.floor(scaleFactor * (vy - min[1]))
+        const iaz = Math.floor(scaleFactor * (vz - min[2]))
+
+        // Extents of grid to consider for this atom
+        const begX = Math.max(0, iax - ng)
+        const begY = Math.max(0, iay - ng)
+        const begZ = Math.max(0, iaz - ng)
+
+        // Add two to these points:
+        // - iax are floor'd values so this ensures coverage
+        // - these are loop limits (exclusive)
+        const endX = Math.min(dimX, iax + ng + 2)
+        const endY = Math.min(dimY, iay + ng + 2)
+        const endZ = Math.min(dimZ, iaz + ng + 2)
 
         for (let xi = begX; xi < endX; ++xi) {
-            const dx = xi * invDeltaX - vx
+            const dx = gridx[xi] - vx
             const xIdx = xi * iuv
             for (let yi = begY; yi < endY; ++yi) {
-                const dy = yi * invDeltaY - vy
+                const dy = gridy[yi] - vy
                 const dxySq = dx * dx + dy * dy
                 const xyIdx = yi * iu + xIdx
                 for (let zi = begZ; zi < endZ; ++zi) {
-                    const dz = zi * invDeltaZ - vz
+                    const dz = gridz[zi] - vz
                     const dSq = dxySq + dz * dz
                     if (dSq <= r2sq) {
                         const dens = Math.exp(-alpha * (dSq * rSqInv))
@@ -118,7 +118,7 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
             await ctx.update({ message: 'filling density grid', current: i, max: n })
         }
     }
-    console.timeEnd('gaussian density cpu')
+    // console.timeEnd('gaussian density cpu')
 
     const transform = Mat4.identity()
     Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))

src/mol-math/geometry/molecular-surface.ts

@@ -15,7 +15,7 @@ import { RuntimeContext } from 'mol-task';
 import { OrderedSet } from 'mol-data/int';
 import { PositionData } from './common';
 import { Mat4 } from 'mol-math/linear-algebra/3d';
-import { Box3D, GridLookup3D } from 'mol-math/geometry';
+import { Box3D, GridLookup3D, fillGridDim } from 'mol-math/geometry';
 import { getDelta } from './gaussian-density';
 
 function normalToLine (out: Vec3, p: Vec3) {
@@ -45,12 +45,6 @@ function getAngleTables (probePositions: number): AnglesTables {
     return { cosTable, sinTable}
 }
 
-function fillGridDim(a: Float32Array, start: number, step: number) {
-    for (let i = 0; i < a.length; i++) {
-        a[i] = start + (step * i)
-    }
-}
-
 /**
  * Is the point at x,y,z obscured by any of the atoms specifeid by indices in neighbours.
  * Ignore indices a and b (these are the relevant atoms in projectPoints/Torii)
@@ -371,13 +365,11 @@ async function createState(ctx: RuntimeContext, position: Required<PositionData>
     fillUniform(data, -1001.0)
     fillUniform(idData, -1)
 
-    const gridx = new Float32Array(dim[0])
-    const gridy = new Float32Array(dim[1])
-    const gridz = new Float32Array(dim[2])
+    const gridx = fillGridDim(dim[0], min[0], resolution)
+    const gridy = fillGridDim(dim[1], min[1], resolution)
+    const gridz = fillGridDim(dim[2], min[2], resolution)
 
-    fillGridDim(gridx, min[0], resolution)
-    fillGridDim(gridy, min[1], resolution)
-    fillGridDim(gridz, min[2], resolution)
+    const updateChunk = Math.ceil(1000000 / (Math.pow(maxRadius, 3) * resolution))
 
     return {
         lastClip: -1,