diff --git a/src/mol-math/geometry/gaussian-density/cpu.ts b/src/mol-math/geometry/gaussian-density/cpu.ts index b2d82b4910d3b161e31bdf418902417bf19697e8..3dedac18f884b9a1acc9891f4739bff7428201da 100644 --- a/src/mol-math/geometry/gaussian-density/cpu.ts +++ b/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))