molecular surface improvements

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

@@ -10,7 +10,6 @@
 import { fillUniform } from 'mol-util/array';
 import { Vec3, Tensor } from 'mol-math/linear-algebra';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
-import { Lookup3D, Result } from './lookup3d/common';
 import { RuntimeContext } from 'mol-task';
 import { OrderedSet } from 'mol-data/int';
 import { PositionData } from './common';
@@ -44,320 +43,293 @@ function getAngleTables (probePositions: number): AnglesTables {
     return { cosTable, sinTable}
 }
 
-/**
- * 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)
- *
- * Cache the last clipped atom (as very often the same one in subsequent calls)
- * 
- * `a` and `b` must be resolved indices
- */
-function obscured (state: MolSurfCalcState, x: number, y: number, z: number, a: number, b: number) {
-    if (state.lastClip !== -1) {
-        const ai = state.lastClip
-        if (ai !== a && ai !== b && singleAtomObscures(state, ai, x, y, z)) {
-            return ai
-        } else {
-            state.lastClip = -1
-        }
-    }
-
-    for (let j = 0, jl = state.neighbours.count; j < jl; ++j) {
-        const ai = OrderedSet.getAt(state.position.indices, state.neighbours.indices[j])
-        if (ai !== a && ai !== b && singleAtomObscures(state, ai, x, y, z)) {
-            state.lastClip = ai
-            return ai
-        }
-    }
-
-    return -1
-}
+//
 
-/**
- * `ai` must be a resolved index
- */
-function singleAtomObscures (state: MolSurfCalcState, ai: number, x: number, y: number, z: number) {
-    const r = state.position.radius[ai]
-    const dx = state.position.x[ai] - x
-    const dy = state.position.y[ai] - y
-    const dz = state.position.z[ai] - z
-    const dSq = dx * dx + dy * dy + dz * dz
-    return dSq < (r * r)
+export const MolecularSurfaceCalculationParams = {
+    resolution: PD.Numeric(0.5, { min: 0.01, max: 10, step: 0.01 }),
+    probeRadius: PD.Numeric(1.4, { min: 0, max: 10, step: 0.1 }),
+    probePositions: PD.Numeric(30, { min: 12, max: 90, step: 1 }),
 }
+export const DefaultMolecularSurfaceCalculationProps = PD.getDefaultValues(MolecularSurfaceCalculationParams)
+export type MolecularSurfaceCalculationProps = typeof DefaultMolecularSurfaceCalculationProps
 
-/**
- * For each atom:
- *     Iterate over a subsection of the grid, for each point:
- *         If current value < 0.0, unvisited, set positive
- *
- *         In any case: Project this point onto surface of the atomic sphere
- *         If this projected point is not obscured by any other atom
- *             Calculate delta distance and set grid value to minimum of
- *             itself and delta
- */
-async function projectPoints (ctx:  RuntimeContext, state: MolSurfCalcState) {
-    const { position, lookup3d, min, space, data, idData, scaleFactor, updateChunk } = state
-    const { gridx, gridy, gridz } = state
-
-    const { indices, x, y, z, radius } = position
-    const n = OrderedSet.size(indices)
 
-    const [ dimX, dimY, dimZ ] = space.dimensions
-    const iu = dimZ, iv = dimY, iuv = iu * iv
+export async function calcMolecularSurface(ctx: RuntimeContext, position: Required<PositionData>, maxRadius: number,  props: MolecularSurfaceCalculationProps) {
+    // Field generation method adapted from AstexViewer (Mike Hartshorn) by Fred Ludlow.
+    // Other parts based heavily on NGL (Alexander Rose) EDT Surface class
 
-    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 = radius[j]
-        const rSq = rad * rad
-
-        state.neighbours = lookup3d.find(vx, vy, vz, rad)
-
-        // Number of grid points, round this up...
-        const ng = Math.ceil(rad * 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 < rSq) {
-                        const idx = zi + xyIdx
-
-                        // if unvisited, make positive
-                        if (data[idx] < 0.0) data[idx] *= -1
-
-                        // Project on to the surface of the sphere
-                        // sp is the projected point ( dx, dy, dz ) * ( ra / d )
-                        const d = Math.sqrt(dSq)
-                        const ap = rad / d
-                        const spx = dx * ap + vx
-                        const spy = dy * ap + vy
-                        const spz = dz * ap + vz
-
-                        if (obscured(state, spx, spy, spz, j, -1) === -1) {
-                            const dd = rad - d
-                            if (dd < data[idx]) {
-                                data[idx] = dd
-                                idData[idx] = i
-                            }
-                        }
-                    }
-                }
+    let lastClip = -1
+
+    /**
+     * 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)
+     *
+     * Cache the last clipped atom (as very often the same one in subsequent calls)
+     *
+     * `a` and `b` must be resolved indices
+     */
+    function obscured(x: number, y: number, z: number, a: number, b: number) {
+        if (lastClip !== -1) {
+            const ai = lastClip
+            if (ai !== a && ai !== b && singleAtomObscures(ai, x, y, z)) {
+                return ai
+            } else {
+                lastClip = -1
             }
         }
 
-        if (i % updateChunk === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'projecting points', current: i, max: n })
+        for (let j = 0, jl = neighbours.count; j < jl; ++j) {
+            const ai = OrderedSet.getAt(position.indices, neighbours.indices[j])
+            if (ai !== a && ai !== b && singleAtomObscures(ai, x, y, z)) {
+                lastClip = ai
+                return ai
+            }
         }
-    }
-}
-
-// Vectors for Torus Projection
-const atob = Vec3()
-const mid = Vec3()
-const n1 = Vec3()
-const n2 = Vec3()
-/**
- * `a` and `b` must be resolved indices
- */
-function projectTorus (state: MolSurfCalcState, a: number, b: number) {
-    const { position, min, space, data, idData } = state
-    const { cosTable, sinTable, probePositions, probeRadius, scaleFactor } = state
-    const { gridx, gridy, gridz } = state
-
-    const [ dimX, dimY, dimZ ] = space.dimensions
-    const iu = dimZ, iv = dimY, iuv = iu * iv
-
-    const ng = Math.max(5, 2 + Math.floor(probeRadius * scaleFactor))
-
-    const rA = position.radius[a]
-    const rB = position.radius[b]
-    const dx = atob[0] = position.x[b] - position.x[a]
-    const dy = atob[1] = position.y[b] - position.y[a]
-    const dz = atob[2] = position.z[b] - position.z[a]
-    const dSq = dx * dx + dy * dy + dz * dz
-
-    // This check now redundant as already done in AVHash.withinRadii
-    // if (dSq > ((rA + rB) * (rA + rB))) { return }
-
-    const d = Math.sqrt(dSq)
-
-    // Find angle between a->b vector and the circle
-    // of their intersection by cosine rule
-    const cosA = (rA * rA + d * d - rB * rB) / (2.0 * rA * d)
-
-    // distance along a->b at intersection
-    const dmp = rA * cosA
-
-    Vec3.normalize(atob, atob)
-
-    // Create normal to line
-    normalToLine(n1, atob)
-    Vec3.normalize(n1, n1)
-
-    // Cross together for second normal vector
-    Vec3.cross(n2, atob, n1)
-    Vec3.normalize(n2, n2)
-
-    // r is radius of circle of intersection
-    const rInt = Math.sqrt(rA * rA - dmp * dmp)
-
-    Vec3.scale(n1, n1, rInt)
-    Vec3.scale(n2, n2, rInt)
-    Vec3.scale(atob, atob, dmp)
-
-    mid[0] = atob[0] + position.x[a]
-    mid[1] = atob[1] + position.y[a]
-    mid[2] = atob[2] + position.z[a]
-
-    state.lastClip = -1
 
-    for (let i = 0; i < probePositions; ++i) {
-        const cost = cosTable[i]
-        const sint = sinTable[i]
-
-        const px = mid[0] + cost * n1[0] + sint * n2[0]
-        const py = mid[1] + cost * n1[1] + sint * n2[1]
-        const pz = mid[2] + cost * n1[2] + sint * n2[2]
+        return -1
+    }
 
-        if (obscured(state, px, py, pz, a, b) === -1) {
-            const iax = Math.floor(scaleFactor * (px - min[0]))
-            const iay = Math.floor(scaleFactor * (py - min[1]))
-            const iaz = Math.floor(scaleFactor * (pz - min[2]))
+    /**
+     * `ai` must be a resolved index
+     */
+    function singleAtomObscures(ai: number, x: number, y: number, z: number) {
+        const r = radius[ai]
+        const dx = px[ai] - x
+        const dy = py[ai] - y
+        const dz = pz[ai] - z
+        const dSq = dx * dx + dy * dy + dz * dz
+        return dSq < (r * r)
+    }
 
+    /**
+     * For each atom:
+     *     Iterate over a subsection of the grid, for each point:
+     *         If current value < 0.0, unvisited, set positive
+     *
+     *         In any case: Project this point onto surface of the atomic sphere
+     *         If this projected point is not obscured by any other atom
+     *             Calculate delta distance and set grid value to minimum of
+     *             itself and delta
+     */
+    function projectPointsRange (begI: number, endI: number) {
+        for (let i = begI; i < endI; ++i) {
+            const j = OrderedSet.getAt(indices, i)
+            const vx = px[j], vy = py[j], vz = pz[j]
+            const rad = radius[j]
+            const rSq = rad * rad
+
+            lookup3d.find(vx, vy, vz, rad)
+
+            // Number of grid points, round this up...
+            const ng = Math.ceil(rad * scaleFactor)
+
+            // Center of the atom, mapped to grid points (take floor)
+            const iax = Math.floor(scaleFactor * (vx - minX))
+            const iay = Math.floor(scaleFactor * (vy - minY))
+            const iaz = Math.floor(scaleFactor * (vz - minZ))
+
+            // 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 = px - gridx[xi]
+                const dx = gridx[xi] - vx
                 const xIdx = xi * iuv
-
                 for (let yi = begY; yi < endY; ++yi) {
-                    const dy = py - gridy[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 = pz - gridz[zi]
+                        const dz = gridz[zi] - vz
                         const dSq = dxySq + dz * dz
 
-                        const idx = zi + xyIdx
-                        const current = data[idx]
-
-                        if (current > 0.0 && dSq < (current * current)) {
-                            data[idx] = Math.sqrt(dSq)
-                            // Is this grid point closer to a or b?
-                            // Take dot product of atob and gridpoint->p (dx, dy, dz)
-                            const dp = dx * atob[0] + dy * atob[1] + dz * atob[2]
-                            idData[idx] = OrderedSet.indexOf(position.indices, dp < 0.0 ? b : a)
+                        if (dSq < rSq) {
+                            const idx = zi + xyIdx
+
+                            // if unvisited, make positive
+                            if (data[idx] < 0.0) data[idx] *= -1
+
+                            // Project on to the surface of the sphere
+                            // sp is the projected point ( dx, dy, dz ) * ( ra / d )
+                            const d = Math.sqrt(dSq)
+                            const ap = rad / d
+                            const spx = dx * ap + vx
+                            const spy = dy * ap + vy
+                            const spz = dz * ap + vz
+
+                            if (obscured(spx, spy, spz, j, -1) === -1) {
+                                const dd = rad - d
+                                if (dd < data[idx]) {
+                                    data[idx] = dd
+                                    idData[idx] = i
+                                }
+                            }
                         }
                     }
                 }
             }
         }
     }
-}
 
-async function projectTorii (ctx: RuntimeContext, state: MolSurfCalcState) {
-    const { n, lookup3d, position, updateChunk } = state
-    const { x, y, z, indices, radius } = position
-    for (let i = 0; i < n; ++i) {
-        const k = OrderedSet.getAt(indices, i)
-        state.neighbours = lookup3d.find(x[k], y[k], z[k], radius[k])
-        for (let j = 0, jl = state.neighbours.count; j < jl; ++j) {
-            const l = OrderedSet.getAt(indices, state.neighbours.indices[j])
-            if (k < l) projectTorus(state, k, l)
+    async function projectPoints() {
+        for (let i = 0; i < n; i += updateChunk) {
+            projectPointsRange(i, Math.min(i + updateChunk, n))
+
+            if (ctx.shouldUpdate) {
+                await ctx.update({ message: 'projecting points', current: i, max: n })
+            }
         }
+    }
 
-        if (i % updateChunk === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'projecting torii', current: i, max: n })
+    // Vectors for Torus Projection
+    const atob = Vec3()
+    const mid = Vec3()
+    const n1 = Vec3()
+    const n2 = Vec3()
+    /**
+     * `a` and `b` must be resolved indices
+     */
+    function projectTorus(a: number, b: number) {
+        const rA = radius[a]
+        const rB = radius[b]
+        const dx = atob[0] = px[b] - px[a]
+        const dy = atob[1] = py[b] - py[a]
+        const dz = atob[2] = pz[b] - pz[a]
+        const dSq = dx * dx + dy * dy + dz * dz
+
+        // This check now redundant as already done in AVHash.withinRadii
+        // if (dSq > ((rA + rB) * (rA + rB))) { return }
+
+        const d = Math.sqrt(dSq)
+
+        // Find angle between a->b vector and the circle
+        // of their intersection by cosine rule
+        const cosA = (rA * rA + d * d - rB * rB) / (2.0 * rA * d)
+
+        // distance along a->b at intersection
+        const dmp = rA * cosA
+
+        Vec3.normalize(atob, atob)
+
+        // Create normal to line
+        normalToLine(n1, atob)
+        Vec3.normalize(n1, n1)
+
+        // Cross together for second normal vector
+        Vec3.cross(n2, atob, n1)
+        Vec3.normalize(n2, n2)
+
+        // r is radius of circle of intersection
+        const rInt = Math.sqrt(rA * rA - dmp * dmp)
+
+        Vec3.scale(n1, n1, rInt)
+        Vec3.scale(n2, n2, rInt)
+        Vec3.scale(atob, atob, dmp)
+
+        mid[0] = atob[0] + px[a]
+        mid[1] = atob[1] + py[a]
+        mid[2] = atob[2] + pz[a]
+
+        lastClip = -1
+
+        for (let i = 0; i < probePositions; ++i) {
+            const cost = cosTable[i]
+            const sint = sinTable[i]
+
+            const px = mid[0] + cost * n1[0] + sint * n2[0]
+            const py = mid[1] + cost * n1[1] + sint * n2[1]
+            const pz = mid[2] + cost * n1[2] + sint * n2[2]
+
+            if (obscured(px, py, pz, a, b) === -1) {
+                const iax = Math.floor(scaleFactor * (px - minX))
+                const iay = Math.floor(scaleFactor * (py - minY))
+                const iaz = Math.floor(scaleFactor * (pz - minZ))
+
+                const begX = Math.max(0, iax - ngTorus)
+                const begY = Math.max(0, iay - ngTorus)
+                const begZ = Math.max(0, iaz - ngTorus)
+
+                const endX = Math.min(dimX, iax + ngTorus + 2)
+                const endY = Math.min(dimY, iay + ngTorus + 2)
+                const endZ = Math.min(dimZ, iaz + ngTorus + 2)
+
+                for (let xi = begX; xi < endX; ++xi) {
+                    const dx = px - gridx[xi]
+                    const xIdx = xi * iuv
+
+                    for (let yi = begY; yi < endY; ++yi) {
+                        const dy = py - gridy[yi]
+                        const dxySq = dx * dx + dy * dy
+                        const xyIdx = yi * iu + xIdx
+
+                        for (let zi = begZ; zi < endZ; ++zi) {
+                            const dz = pz - gridz[zi]
+                            const dSq = dxySq + dz * dz
+
+                            const idx = zi + xyIdx
+                            const current = data[idx]
+
+                            if (current > 0.0 && dSq < (current * current)) {
+                                data[idx] = Math.sqrt(dSq)
+                                // Is this grid point closer to a or b?
+                                // Take dot product of atob and gridpoint->p (dx, dy, dz)
+                                const dp = dx * atob[0] + dy * atob[1] + dz * atob[2]
+                                idData[idx] = OrderedSet.indexOf(position.indices, dp < 0.0 ? b : a)
+                            }
+                        }
+                    }
+                }
+            }
         }
     }
-}
 
-//
+    async function projectTorii () {
+        for (let i = 0; i < n; ++i) {
+            const k = OrderedSet.getAt(indices, i)
+            lookup3d.find(px[k], py[k], pz[k], radius[k])
+            for (let j = 0, jl = neighbours.count; j < jl; ++j) {
+                const l = OrderedSet.getAt(indices, neighbours.indices[j])
+                if (k < l) projectTorus(k, l)
+            }
 
-interface MolSurfCalcState {
-    /** Cached last value for obscured test */
-    lastClip: number
-    /** Neighbours as transient result array from lookup3d */
-    neighbours: Result<number>
-
-    lookup3d: Lookup3D
-    position: Required<PositionData>
-    min: Vec3
-    updateChunk: number
-
-    maxRadius: number
-
-    n: number
-    resolution: number
-    scaleFactor: number
-    probeRadius: number
-
-    /** Angle lookup tables */
-    cosTable: Float32Array
-    sinTable: Float32Array
-    probePositions: number
-
-    /** grid lookup tables */
-    gridx: Float32Array
-    gridy: Float32Array
-    gridz: Float32Array
-
-    expandedBox: Box3D
-    space: Tensor.Space
-    data: Tensor.Data
-    idData: Tensor.Data
-}
+            if (i % updateChunk === 0 && ctx.shouldUpdate) {
+                await ctx.update({ message: 'projecting torii', current: i, max: n })
+            }
+        }
+    }
 
-async function createState(ctx: RuntimeContext, position: Required<PositionData>, maxRadius: number, props: MolecularSurfaceCalculationProps): Promise<MolSurfCalcState> {
+    // console.time('MolecularSurface')
+    // console.time('MolecularSurface createState')
     const { resolution, probeRadius, probePositions } = props
     const scaleFactor = 1 / resolution
+    const ngTorus = Math.max(5, 2 + Math.floor(probeRadius * scaleFactor))
 
     const cellSize =  Vec3.create(maxRadius, maxRadius, maxRadius)
     Vec3.scale(cellSize, cellSize, 2)
     const lookup3d = GridLookup3D(position, cellSize)
+    const neighbours = lookup3d.result
     const box = lookup3d.boundary.box
-    const { indices } = position
+    const { indices, x: px, y: py, z: pz, radius } = position
     const n = OrderedSet.size(indices)
 
     const pad = maxRadius * 2 + resolution
     const expandedBox = Box3D.expand(Box3D(), box, Vec3.create(pad, pad, pad));
-    const min = expandedBox.min
+    const [ minX, minY, minZ ] = expandedBox.min
     const scaledBox = Box3D.scale(Box3D(), expandedBox, scaleFactor)
     const dim = Box3D.size(Vec3(), scaledBox)
     Vec3.ceil(dim, dim)
 
+    const [ dimX, dimY, dimZ ] = dim
+    const iu = dimZ, iv = dimY, iuv = iu * iv
+
     const { cosTable, sinTable } = getAngleTables(probePositions)
 
     const space = Tensor.Space(dim, [0, 1, 2], Float32Array)
@@ -367,81 +339,28 @@ async function createState(ctx: RuntimeContext, position: Required<PositionData>
     fillUniform(data, -1001.0)
     fillUniform(idData, -1)
 
-    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 updateChunk = Math.ceil(1000000 / (Math.pow(maxRadius, 3) * resolution))
-
-    return {
-        lastClip: -1,
-        neighbours: lookup3d.find(0, 0, 0, 0),
-
-        lookup3d,
-        position,
-        min,
-        updateChunk,
-
-        maxRadius,
-
-        n,
-        resolution,
-        scaleFactor,
-        probeRadius,
-
-        cosTable,
-        sinTable,
-        probePositions,
-
-        gridx,
-        gridy,
-        gridz,
-
-        expandedBox,
-        space,
-        data,
-        idData,
-    }
-}
-
-//
-
-export const MolecularSurfaceCalculationParams = {
-    resolution: PD.Numeric(0.5, { min: 0.01, max: 10, step: 0.01 }),
-    probeRadius: PD.Numeric(1.4, { min: 0, max: 10, step: 0.1 }),
-    probePositions: PD.Numeric(30, { min: 12, max: 90, step: 1 }),
-}
-export const DefaultMolecularSurfaceCalculationProps = PD.getDefaultValues(MolecularSurfaceCalculationParams)
-export type MolecularSurfaceCalculationProps = typeof DefaultMolecularSurfaceCalculationProps
-
-
-export async function calcMolecularSurface(ctx: RuntimeContext, position: Required<PositionData>, maxRadius: number,  props: MolecularSurfaceCalculationProps) {
-    // Field generation method adapted from AstexViewer (Mike Hartshorn) by Fred Ludlow.
-    // Other parts based heavily on NGL (Alexander Rose) EDT Surface class
-
-    console.time('MolecularSurface')
-
-    console.time('MolecularSurface createState')
-    const state = await createState(ctx, position, maxRadius, props)
-    console.timeEnd('MolecularSurface createState')
+    const gridx = fillGridDim(dimX, minX, resolution)
+    const gridy = fillGridDim(dimY, minY, resolution)
+    const gridz = fillGridDim(dimZ, minZ, resolution)
 
-    console.time('MolecularSurface projectPoints')
-    await projectPoints(ctx, state)
-    console.timeEnd('MolecularSurface projectPoints')
+    const updateChunk = Math.ceil(100000 / ((Math.pow(Math.pow(maxRadius, 3), 3) * scaleFactor)))
+    // console.timeEnd('MolecularSurface createState')
 
-    console.time('MolecularSurface projectTorii')
-    await projectTorii(ctx, state)
-    console.timeEnd('MolecularSurface projectTorii')
+    // console.time('MolecularSurface projectPoints')
+    await projectPoints()
+    // console.timeEnd('MolecularSurface projectPoints')
 
-    console.timeEnd('MolecularSurface')
+    // console.time('MolecularSurface projectTorii')
+    await projectTorii()
+    // console.timeEnd('MolecularSurface projectTorii')
+    // console.timeEnd('MolecularSurface')
 
-    const field = Tensor.create(state.space, state.data)
-    const idField = Tensor.create(state.space, state.idData)
+    const field = Tensor.create(space, data)
+    const idField = Tensor.create(space, idData)
 
-    const { resolution, expandedBox } = state
     const transform = Mat4.identity()
     Mat4.fromScaling(transform, Vec3.create(resolution, resolution, resolution))
     Mat4.setTranslation(transform, expandedBox.min)
-    console.log({ field, idField, transform, state })
+    // console.log({ field, idField, transform, updateChunk })
     return { field, idField, transform }
 }

src/mol-repr/structure/visual/molecular-surface-mesh.ts

@@ -27,7 +27,6 @@ export type MolecularSurfaceMeshParams = typeof MolecularSurfaceMeshParams
 //
 
 async function createMolecularSurfaceMesh(ctx: VisualContext, unit: Unit, structure: Structure, theme: Theme, props: MolecularSurfaceCalculationProps, mesh?: Mesh): Promise<Mesh> {
-    console.log(props)
 
     const { transform, field, idField } = await computeUnitMolecularSurface(unit, props).runInContext(ctx.runtime)
     const params = {

src/mol-repr/structure/visual/util/molecular-surface.ts

@@ -11,27 +11,27 @@ import { PositionData, DensityData } from 'mol-math/geometry';
 import { MolecularSurfaceCalculationProps, calcMolecularSurface } from 'mol-math/geometry/molecular-surface';
 import { OrderedSet } from 'mol-data/int';
 
-export function computeUnitMolecularSurface(unit: Unit, props: MolecularSurfaceCalculationProps) {
+function getPositionDataAndMaxRadius(unit: Unit, props: MolecularSurfaceCalculationProps) {
     const { position, radius } = getUnitConformationAndRadius(unit)
+    const { indices } = position
+    const n = OrderedSet.size(indices)
+    const radii = new Float32Array(OrderedSet.end(indices))
+
+    let maxRadius = 0
+    for (let i = 0; i < n; ++i) {
+        const j = OrderedSet.getAt(indices, i)
+        const r = radius(j)
+        if (maxRadius < r) maxRadius = r
+        radii[j] = r + props.probeRadius
+    }
+
+    return { position: { ...position, radius: radii }, maxRadius }
+}
 
+export function computeUnitMolecularSurface(unit: Unit, props: MolecularSurfaceCalculationProps) {
+    const { position, maxRadius } = getPositionDataAndMaxRadius(unit, props)
     return Task.create('Molecular Surface', async ctx => {
-        const { indices } = position
-        const n = OrderedSet.size(indices)
-        const radii = new Float32Array(OrderedSet.max(indices))
-
-        let maxRadius = 0
-        for (let i = 0; i < n; ++i) {
-            const j = OrderedSet.getAt(indices, i)
-            const r = radius(j)
-            if (maxRadius < r) maxRadius = r
-            radii[j] = r + props.probeRadius
-
-            if (i % 100000 === 0 && ctx.shouldUpdate) {
-                await ctx.update({ message: 'calculating max radius', current: i, max: n })
-            }
-        }
-
-        return await MolecularSurface(ctx, { ...position, radius: radii }, maxRadius, props);
+        return await MolecularSurface(ctx, position, maxRadius, props);
     });
 }
 

src/tests/browser/render-structure.ts

@@ -84,43 +84,59 @@ async function init() {
     console.timeEnd('computeModelDSSP');
     (models[0].properties as any).secondaryStructure = secondaryStructure
     const structure = await getStructure(models[0])
+
+    const show = {
+        cartoon: false,
+        ballAndStick: true,
+        molecularSurface: true,
+        gaussianSurface: false,
+    }
+
     const cartoonRepr = getCartoonRepr()
-    const ballAndStick = getBallAndStickRepr()
+    const ballAndStickRepr = getBallAndStickRepr()
     const molecularSurfaceRepr = getMolecularSurfaceRepr()
     const gaussianSurfaceRepr = getGaussianSurfaceRepr()
 
-    // cartoonRepr.setTheme({
-    //     color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
-    //     size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
-    // })
-    // await cartoonRepr.createOrUpdate({ ...CartoonRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
-
-    // ballAndStick.setTheme({
-    //     color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
-    //     size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
-    // })
-    // await ballAndStick.createOrUpdate({ ...BallAndStickRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
-
-    molecularSurfaceRepr.setTheme({
-        color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
-        size: reprCtx.sizeThemeRegistry.create('physical', { structure })
-    })
-    console.time('molecular surface')
-    await molecularSurfaceRepr.createOrUpdate({ ...MolecularSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 1.0, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
-    console.timeEnd('molecular surface')
-
-    // gaussianSurfaceRepr.setTheme({
-    //     color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
-    //     size: reprCtx.sizeThemeRegistry.create('physical', { structure })
-    // })
-    // console.time('gaussian surface')
-    // await gaussianSurfaceRepr.createOrUpdate({ ...GaussianSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 1.0, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
-    // console.timeEnd('gaussian surface')
-
-    // canvas3d.add(cartoonRepr)
-    // canvas3d.add(ballAndStick)
-    canvas3d.add(molecularSurfaceRepr)
-    // canvas3d.add(gaussianSurfaceRepr)
+    if (show.cartoon) {
+        cartoonRepr.setTheme({
+            color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
+            size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
+        })
+        await cartoonRepr.createOrUpdate({ ...CartoonRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
+    }
+
+    if (show.ballAndStick) {
+        ballAndStickRepr.setTheme({
+            color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
+            size: reprCtx.sizeThemeRegistry.create('uniform', { structure })
+        })
+        await ballAndStickRepr.createOrUpdate({ ...BallAndStickRepresentationProvider.defaultValues, quality: 'auto' }, structure).run()
+    }
+
+    if (show.molecularSurface) {
+        molecularSurfaceRepr.setTheme({
+            color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
+            size: reprCtx.sizeThemeRegistry.create('physical', { structure })
+        })
+        console.time('molecular surface')
+        await molecularSurfaceRepr.createOrUpdate({ ...MolecularSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 0.5, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
+        console.timeEnd('molecular surface')
+    }
+
+    if (show.gaussianSurface) {
+        gaussianSurfaceRepr.setTheme({
+            color: reprCtx.colorThemeRegistry.create('secondary-structure', { structure }),
+            size: reprCtx.sizeThemeRegistry.create('physical', { structure })
+        })
+        console.time('gaussian surface')
+        await gaussianSurfaceRepr.createOrUpdate({ ...GaussianSurfaceRepresentationProvider.defaultValues, quality: 'custom', alpha: 1.0, flatShaded: true, doubleSided: true, resolution: 0.3 }, structure).run()
+        console.timeEnd('gaussian surface')
+    }
+
+    if (show.cartoon) canvas3d.add(cartoonRepr)
+    if (show.ballAndStick) canvas3d.add(ballAndStickRepr)
+    if (show.molecularSurface) canvas3d.add(molecularSurfaceRepr)
+    if (show.gaussianSurface) canvas3d.add(gaussianSurfaceRepr)
     canvas3d.resetCamera()
 }