alpha-orbitals: improvements and fixes

src/examples/alpha-orbitals/transforms.ts

@@ -56,10 +56,10 @@ export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
                 sphericalOrder: a.data.order,
                 boxExpand: 4.5,
                 gridSpacing: [
-                    [55, 0.6],
-                    [40, 0.5],
+                    [55, 0.5],
+                    [40, 0.45],
                     [25, 0.4],
-                    [0, 0.33],
+                    [0, 0.35],
                 ],
             }, plugin.canvas3d?.webgl).runInContext(ctx);
             const volume: Volume = {

src/extensions/alpha-orbitals/cubes.ts

@@ -11,6 +11,7 @@ import { WebGLContext } from '../../mol-gl/webgl/context';
 import { Box3D } from '../../mol-math/geometry';
 import { Mat4, Tensor, Vec3 } from '../../mol-math/linear-algebra';
 import { Grid } from '../../mol-model/volume';
+import { getNonStandardResidueQueries } from '../../mol-plugin-state/helpers/structure-selection-query';
 import { Task } from '../../mol-task';
 import { arrayMax, arrayMin, arrayRms } from '../../mol-util/array';
 import { CollocationParams, sphericalCollocation } from './collocation';
@@ -67,18 +68,8 @@ export function createSphericalCollocationGrid(
     return Task.create('Spherical Collocation Grid', async (ctx) => {
         const centers = params.basis.atoms.map(a => a.center);
 
-        const grid = initBox(centers, params.gridSpacing, params.boxExpand, true);
-
-        // const cParams: CollocationParams = {
-        //     grid,
-        //     basis: params.basis,
-        //     alphaOrbitals: params.alphaOrbitals,
-        //     cutoffThreshold: params.cutoffThreshold,
-        //     sphericalOrder: params.sphericalOrder
-        // };
-
         const cParams: CollocationParams = {
-            grid,
+            grid: initBox(centers, params.gridSpacing, params.boxExpand),
             basis: params.basis,
             alphaOrbitals: params.alphaOrbitals,
             cutoffThreshold: params.cutoffThreshold,
@@ -100,15 +91,9 @@ export function createSphericalCollocationGrid(
 
         // if (false && webgl) {
         // } else {
-        console.time('cpu');
-        const matrix = await sphericalCollocation({
-            grid: initBox(centers, params.gridSpacing, params.boxExpand, false),
-            basis: params.basis,
-            alphaOrbitals: params.alphaOrbitals,
-            cutoffThreshold: params.cutoffThreshold,
-            sphericalOrder: params.sphericalOrder
-        }, ctx);
-        console.timeEnd('cpu');
+        // console.time('cpu');
+        // const matrix = await sphericalCollocation(cParams, ctx);
+        // console.timeEnd('cpu');
         // // }
 
         console.log(matrixGL);
@@ -122,7 +107,7 @@ export function createSphericalCollocationGrid(
         //     }
         // }
 
-        return createCubeGrid(grid, matrixGL, [0, 1, 2]);
+        return createCubeGrid(cParams.grid, matrixGL, [0, 1, 2]);
     });
 }
 
@@ -174,8 +159,7 @@ function createCubeGrid(gridInfo: CubeGridInfo, values: Float32Array, axisOrder:
 function initBox(
     geometry: Vec3[],
     spacing: SphericalCollocationParams['gridSpacing'],
-    expand: number,
-    isGpu: boolean
+    expand: number
 ): CubeGridInfo {
     const count = geometry.length;
     const box = Box3D.expand(
@@ -197,14 +181,6 @@ function initBox(
 
     const dimensions = Vec3.ceil(Vec3(), Vec3.scale(Vec3(), size, 1 / s));
 
-    // dimensions need to be powers of 2 otherwise it leads to roudning error
-    // TODO: possible to avoid this having to be power of 2?
-    if (isGpu) {
-        dimensions[0] = reasonablePowerOf2(dimensions[0]);
-        dimensions[1] = reasonablePowerOf2(dimensions[1]);
-        dimensions[2] = reasonablePowerOf2(dimensions[2]);
-    }
-
     return {
         box,
         dimensions,
@@ -214,15 +190,6 @@ function initBox(
     };
 }
 
-function reasonablePowerOf2(x: number) {
-    const high = Math.pow(2, Math.ceil(Math.log2(x)));
-    if (high < 129) return high | 0;
-
-    const low = Math.pow(2, Math.floor(Math.log2(x)));
-    if ((x - low) / (high - low) < 0.33) return low | 0;
-    return high | 0;
-}
-
 export function computeIsocontourValues(
     input: Float32Array,
     cumulativeThreshold: number
@@ -236,6 +203,12 @@ export function computeIsocontourValues(
     const avgWeight = weightSum / input.length;
     let minWeight = 3 * avgWeight;
 
+    // do not try to identify isovalues for degenerate data
+    // e.g. all values are almost same
+    if (Math.abs(avgWeight - input[0] * input[0]) < 1e-5) {
+        return { negative: void 0, positive: void 0 };
+    }
+
     let size = 0;
     while (true) {
         let csum = 0;

src/extensions/alpha-orbitals/gpu/pass.ts

@@ -148,12 +148,8 @@ export class AlphaOrbitalsPass {
         const [nx, ny, nz] = this.params.grid.dimensions;
         const width = nx;
         const height = ny * nz;
-
-        // const { x, y, width, height } = this.camera.viewport;
-
         const { gl, state } = this.webgl;
         this.target.bind();
-        // this.setQuadShift(0, 0);
         gl.viewport(0, 0, width, height);
         gl.scissor(0, 0, width, height);
         state.disable(gl.SCISSOR_TEST);

src/extensions/alpha-orbitals/gpu/shader.frag.ts

@@ -62,11 +62,11 @@ vec4 floatToRgba(float texelFloat) {
     );
 }
 
-float L1(vec3 p, float a0, float a1, float a2) {
+float L1(const in vec3 p, const in float a0, const in float a1, const in float a2) {
     return a0 * p.z + a1 * p.x + a2 * p.y;
 }
 
-float L2(vec3 p, float a0, float a1, float a2, float a3, float a4) {
+float L2(const in vec3 p, const in float a0, const in float a1, const in float a2, const in float a3, const in float a4) {
     float x = p.x, y = p.y, z = p.z;
     float xx = x * x, yy = y * y, zz = z * z;
     return (
@@ -78,7 +78,7 @@ float L2(vec3 p, float a0, float a1, float a2, float a3, float a4) {
     );
 }
 
-float L3(vec3 p, float a0, float a1, float a2, float a3, float a4, float a5, float a6) {
+float L3(const in vec3 p, const in float a0, const in float a1, const in float a2, const in float a3, const in float a4, const in float a5, const in float a6) {
     float x = p.x, y = p.y, z = p.z;
     float xx = x * x, yy = y * y, zz = z * z;
     float xxx = xx * x, yyy = yy * y, zzz = zz * z;
@@ -93,21 +93,46 @@ float L3(vec3 p, float a0, float a1, float a2, float a3, float a4, float a5, flo
     );
 }
 
+float L4(const in vec3 p, const in float a0, const in float a1, const in float a2, const in float a3, const in float a4, const in float a5, const in float a6, const in float a7, const in float a8) {
+    float x = p.x, y = p.y, z = p.z;
+    float xx = x * x, yy = y * y, zz = z * z;
+    float xxx = xx * x, yyy = yy * y, zzz = zz * z;
+    float xxxx = xxx * x, yyyy = yyy * y, zzzz = zzz * z;
+    return (
+        a0 * (0.375 * xxxx + 0.75 * xx * yy + 0.375 * yyyy - 3.0 * xx * zz - 3.0 * yy * zz + zzzz) +
+        a1 * (-2.3717082451262845 * xxx * z - 2.3717082451262845 * x * yy * z + 3.1622776601683795 * x * zzz) +
+        a2 * (-2.3717082451262845 * xx * y * z - 2.3717082451262845 * yyy * z + 3.1622776601683795 * y * zzz) +
+        a3 * (-0.5590169943749475 * xxxx + 0.5590169943749475 * yyyy + 3.3541019662496847 * xx * zz - 3.3541019662496847 * yy * zz) +
+        a4 * (-1.118033988749895 * xxx * y - 1.118033988749895 * x * yyy + 6.708203932499369 * x * y * zz) +
+        a5 * (2.091650066335189 * xxx * z + -6.274950199005566 * x * yy * z) +
+        a6 * (6.274950199005566 * xx * y * z + -2.091650066335189 * yyy * z) +
+        a7 * (0.739509972887452 * xxxx - 4.437059837324712 * xx * yy + 0.739509972887452 * yyyy) +
+        a8 * (2.958039891549808 * xxx * y + -2.958039891549808 * x * yyy)
+    );
+}
+
+float alpha(const in float offset, const in float f) {
+    return texture2D(tAlpha, vec2(offset * f, 0.5)).x;
+}
+
+float intDiv(const in float a, const in float b) { return float(int(a) / int(b)); }
+float intMod(const in float a, const in float b) { return a - b * float(int(a) / int(b)); }
+
 void main(void) {
-    // TODO: use "integer division" to avoid rounding errors for non power of 2 dimensions
-    vec3 idx = vec3(gl_FragCoord.x - 0.5, mod(gl_FragCoord.y - 0.5, uDimensions.z), floor((gl_FragCoord.y - 0.5) / uDimensions.y));
+    vec3 idx = vec3(gl_FragCoord.x - 0.5, intMod(gl_FragCoord.y - 0.5, uDimensions.y), intDiv((gl_FragCoord.y - 0.5), uDimensions.y));
+
     float offset = idx.x + idx.y * uDimensions.x + idx.z * uDimensions.x * uDimensions.y;
 
     // TODO: is there render target ordering that does not require the remapping of the coords?
-    float k = mod(offset, uDimensions.z), kk = floor(offset / uDimensions.z);
-    float j = mod(kk, uDimensions.y);
-    float i = floor(kk / uDimensions.y);
+    float k = intMod(offset, uDimensions.z), kk = intDiv(offset, uDimensions.z);
+    float j = intMod(kk, uDimensions.y);
+    float i = intDiv(kk, uDimensions.y);
     
     vec3 xyz = uMin + uDelta * vec3(i, j, k);
 
     float fCenter = 1.0 / float(uNCenters - 1);
     float fCoeff = 1.0 / float(uNCoeff - 1);
-    float fAlpha = 1.0 / float(uNAlpha - 1);
+    float fA = 1.0 / float(uNAlpha - 1);
 
     float v = 0.0;
 
@@ -126,11 +151,10 @@ void main(void) {
         vec4 info = texture2D(tInfo, cCoord);
 
         int L = int(info.x);
-        float alphaOffset = info.y;
+        float aO = info.y;
         int coeffStart = int(info.z);
         int coeffEnd = int(info.w);
 
-
         float gauss = 0.0;
         for (int j = coeffStart; j < coeffEnd; j++) {
             vec2 c = texture2D(tCoeff, vec2(float(j) * fCoeff, 0.5)).xy;
@@ -139,36 +163,27 @@ void main(void) {
 
         float spherical = 0.0;
         if (L == 0) {
-            spherical = texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x;
+            spherical = alpha(aO, fA);
         } else if (L == 1) {
-            spherical = L1(
-                X,
-                texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 1.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 2.0) * fAlpha, 0.5)).x
+            spherical = L1(X,
+                alpha(aO, fA), alpha(aO + 1.0, fA), alpha(aO + 2.0, fA)
             );
         } else if (L == 2) {
-            spherical = L2(
-                X,
-                texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 1.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 2.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 3.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 4.0) * fAlpha, 0.5)).x
+            spherical = L2(X,
+                alpha(aO, fA), alpha(aO + 1.0, fA), alpha(aO + 2.0, fA), alpha(aO + 3.0, fA), alpha(aO + 4.0, fA)
             );
         } else if (L == 3) {
-            spherical = L3(
-                X,
-                texture2D(tAlpha, vec2(alphaOffset * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 1.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 2.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 3.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 4.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 5.0) * fAlpha, 0.5)).x,
-                texture2D(tAlpha, vec2((alphaOffset + 6.0) * fAlpha, 0.5)).x
+            spherical = L3(X,
+                alpha(aO, fA), alpha(aO + 1.0, fA), alpha(aO + 2.0, fA), alpha(aO + 3.0, fA), alpha(aO + 4.0, fA), 
+                alpha(aO + 5.0, fA), alpha(aO + 6.0, fA)
+            );
+        } else if (L == 4) {
+            spherical = L4(X,
+                alpha(aO, fA), alpha(aO + 1.0, fA), alpha(aO + 2.0, fA), alpha(aO + 3.0, fA), alpha(aO + 4.0, fA), 
+                alpha(aO + 5.0, fA), alpha(aO + 6.0, fA), alpha(aO + 7.0, fA), alpha(aO + 8.0, fA)
             );
         } 
-        // TODO: do we need L > 3?
+        // TODO: do we need L > 4?
 
         v += gauss * spherical;
     }