gaussian density, avoid large async function

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

@@ -24,10 +24,6 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
         const r = radius(OrderedSet.getAt(indices, i)) + radiusOffset
         if (maxRadius < r) maxRadius = r
         radii[i] = r
-
-        if (i % 100000 === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'calculating max radius', current: i, max: n })
-        }
     }
 
     const pad = maxRadius * 2 + resolution
@@ -54,68 +50,78 @@ export async function GaussianDensityCPU(ctx: RuntimeContext, position: Position
     const densData = space.create()
 
     const alpha = smoothness
-    const updateChunk = Math.ceil(1000000 / (Math.pow(maxRadius * 2, 3) * resolution))
-
-    console.time('gaussian density cpu')
-    for (let i = 0; i < n; ++i) {
-        const j = OrderedSet.getAt(indices, i)
-        const vx = x[j], vy = y[j], vz = z[j]
-
-        const rad = radii[i]
-        const rSq = rad * rad
-        const rSqInv = 1 / rSq
-
-        const r2 = rad * 2
-        const r2sq = r2 * r2
-
-        // 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 = gridx[xi] - vx
-            const xIdx = xi * iuv
-            for (let yi = begY; yi < endY; ++yi) {
-                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 = gridz[zi] - vz
-                    const dSq = dxySq + dz * dz
-                    if (dSq <= r2sq) {
-                        const dens = Math.exp(-alpha * (dSq * rSqInv))
-                        const idx = zi + xyIdx
-                        data[idx] += dens
-                        if (dens > densData[idx]) {
-                            densData[idx] = dens
-                            idData[idx] = i
+    const updateChunk = Math.ceil(100000 / ((Math.pow(Math.pow(maxRadius, 3), 3) * scaleFactor)))
+
+    function accumulateRange(begI: number, endI: number) {
+        for (let i = begI; i < endI; ++i) {
+            const j = OrderedSet.getAt(indices, i)
+            const vx = x[j], vy = y[j], vz = z[j]
+
+            const rad = radii[i]
+            const rSq = rad * rad
+            const rSqInv = 1 / rSq
+
+            const r2 = rad * 2
+            const r2sq = r2 * r2
+
+            // 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 = gridx[xi] - vx
+                const xIdx = xi * iuv
+                for (let yi = begY; yi < endY; ++yi) {
+                    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 = gridz[zi] - vz
+                        const dSq = dxySq + dz * dz
+                        if (dSq <= r2sq) {
+                            const dens = Math.exp(-alpha * (dSq * rSqInv))
+                            const idx = zi + xyIdx
+                            data[idx] += dens
+                            if (dens > densData[idx]) {
+                                densData[idx] = dens
+                                idData[idx] = i
+                            }
                         }
                     }
                 }
             }
         }
+    }
+
+    async function accumulate() {
+        for (let i = 0; i < n; i += updateChunk) {
+            accumulateRange(i, Math.min(i + updateChunk, n))
 
-        if (i % updateChunk === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'filling density grid', current: i, max: n })
+            if (ctx.shouldUpdate) {
+                await ctx.update({ message: 'filling density grid', current: i, max: n })
+            }
         }
     }
-    console.timeEnd('gaussian density cpu')
+
+    // console.time('gaussian density cpu')
+    await accumulate()
+    // console.timeEnd('gaussian density cpu')
 
     const transform = Mat4.identity()
     Mat4.fromScaling(transform, Vec3.create(resolution, resolution, resolution))