alpha orbitals: density proof of concept

src/examples/alpha-orbitals/index.ts

@@ -98,13 +98,13 @@ export class AlphaOrbitalsExample {
 
     private volumeParams(kind: 'positive' | 'negative', color: Color): StateTransformer.Params<typeof CreateOrbitalRepresentation3D> {
         return {
-            alpha: 1.0,
+            alpha: 0.85,
             color,
             directVolume: this.state.value.gpuSurface,
             kind,
             relativeIsovalue: this.state.value.isoValue,
             pickable: false,
-            xrayShaded: false
+            xrayShaded: true
         };
     }
 

src/extensions/alpha-orbitals/cubes.ts

@@ -14,7 +14,7 @@ import { Grid } from '../../mol-model/volume';
 import { Task } from '../../mol-task';
 import { arrayMax, arrayMin, arrayRms } from '../../mol-util/array';
 import { CollocationParams, sphericalCollocation } from './collocation';
-import { canComputeAlphaOrbitalsOnGPU, gpuComputeAlphaOrbitalsGridValues } from './gpu/compute';
+import { canComputeAlphaOrbitalsOnGPU, gpuComputeAlphaOrbitalsDensityGridValues, gpuComputeAlphaOrbitalsGridValues } from './gpu/compute';
 import { SphericalBasisOrder } from './orbitals';
 
 export interface CubeGridInfo {
@@ -63,7 +63,7 @@ export interface SphericalCollocationParams {
 }
 
 export function createSphericalCollocationGrid(
-    params: SphericalCollocationParams, webgl?: WebGLContext
+    params: SphericalCollocationParams, webgl?: WebGLContext, orbitals?: number[][]
 ): Task<CubeGrid> {
     return Task.create('Spherical Collocation Grid', async (ctx) => {
         const centers = params.basis.atoms.map(a => a.center);
@@ -78,9 +78,11 @@ export function createSphericalCollocationGrid(
 
         let matrix: Float32Array;
         if (canComputeAlphaOrbitalsOnGPU(webgl)) {
-            // console.time('gpu');
-            matrix = gpuComputeAlphaOrbitalsGridValues(webgl!, cParams);
-            // console.timeEnd('gpu');
+            console.time('gpu');
+            // matrix = gpuComputeAlphaOrbitalsGridValues(webgl!, cParams);
+            matrix = gpuComputeAlphaOrbitalsDensityGridValues(webgl!, cParams, orbitals!);
+            console.log(matrix);
+            console.timeEnd('gpu');
         } else {
             // console.time('cpu');
             matrix = await sphericalCollocation(cParams, ctx);

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

@@ -15,7 +15,8 @@ import { ValueCell } from '../../../mol-util';
 import { arrayMin } from '../../../mol-util/array';
 import { isLittleEndian } from '../../../mol-util/is-little-endian';
 import { CollocationParams } from '../collocation';
-import { normalizeBasicOrder } from '../orbitals';
+import { Basis } from '../cubes';
+import { normalizeBasicOrder, SphericalBasisOrder } from '../orbitals';
 import shader_frag from './shader.frag';
 
 const AlphaOrbitalsSchema = {
@@ -32,11 +33,32 @@ const AlphaOrbitalsSchema = {
     uNAlpha: UniformSpec('i'),
     uNCoeff: UniformSpec('i'),
     uMaxCoeffs: UniformSpec('i'),
-    uLittleEndian: UniformSpec('b')
+    uLittleEndian: UniformSpec('b'),
+
+    uOccupancy: UniformSpec('f'),
+    tCumulativeSum: TextureSpec('texture', 'rgba', 'ubyte', 'nearest')
 };
+type AlphaOrbitalsSchema = Values<typeof AlphaOrbitalsSchema>
 const AlphaOrbitalsName = 'alpha-orbitals';
 const AlphaOrbitalsShaderCode = ShaderCode(AlphaOrbitalsName, quad_vert, shader_frag);
-type AlphaOrbitalsRenderable = ComputeRenderable<Values<typeof AlphaOrbitalsSchema>>
+type AlphaOrbitalsRenderable = ComputeRenderable<AlphaOrbitalsSchema>
+
+function getNormalizedAlpha(basis: Basis, alphaOrbitals: number[], sphericalOrder: SphericalBasisOrder) {
+    const alpha = new Float32Array(alphaOrbitals.length);
+
+    let aO = 0;
+    for (const atom of basis.atoms) {
+        for (const shell of atom.shells) {
+            for (const L of shell.angularMomentum) {
+                const a0 = normalizeBasicOrder(L, alphaOrbitals.slice(aO, aO + 2 * L + 1), sphericalOrder);
+                for (let i = 0; i < a0.length; i++) alpha[aO + i] = a0[i];
+                aO += 2 * L + 1;
+            }
+        }
+    }
+
+    return alpha;
+}
 
 function createTextureData({
     basis,
@@ -119,7 +141,11 @@ function createAlphaOrbitalsRenderable(ctx: WebGLContext, params: CollocationPar
     const [nx, ny, nz] = params.grid.dimensions;
     const width = Math.ceil(Math.sqrt(nx * ny * nz));
 
-    const values: Values<typeof AlphaOrbitalsSchema> = {
+    if (!ctx.namedTextures[AlphaOrbitalsName]) {
+        ctx.namedTextures[AlphaOrbitalsName] = ctx.resources.texture('image-uint8', 'rgba', 'ubyte', 'nearest');
+    }
+
+    const values: AlphaOrbitalsSchema = {
         ...QuadValues,
         uDimensions: ValueCell.create(params.grid.dimensions),
         uMin: ValueCell.create(params.grid.box.min),
@@ -134,6 +160,9 @@ function createAlphaOrbitalsRenderable(ctx: WebGLContext, params: CollocationPar
         tCoeff: ValueCell.create({ width: data.nCoeff, height: 1, array: data.coeff }),
         tAlpha: ValueCell.create({ width: data.nAlpha, height: 1, array: data.alpha }),
         uLittleEndian: ValueCell.create(isLittleEndian()),
+
+        uOccupancy: ValueCell.create(0),
+        tCumulativeSum: ValueCell.create(ctx.namedTextures[AlphaOrbitalsName])
     };
 
     const schema = { ...AlphaOrbitalsSchema };
@@ -150,7 +179,7 @@ function createAlphaOrbitalsRenderable(ctx: WebGLContext, params: CollocationPar
 
 function getAlphaOrbitalsRenderable(ctx: WebGLContext, params: CollocationParams): AlphaOrbitalsRenderable {
     if (ctx.namedComputeRenderables[AlphaOrbitalsName]) {
-        const v = ctx.namedComputeRenderables[AlphaOrbitalsName].values;
+        const v = ctx.namedComputeRenderables[AlphaOrbitalsName].values as AlphaOrbitalsSchema;
 
         const data = createTextureData(params);
 
@@ -170,6 +199,7 @@ function getAlphaOrbitalsRenderable(ctx: WebGLContext, params: CollocationParams
         ValueCell.update(v.tCoeff, { width: data.nCoeff, height: 1, array: data.coeff });
         ValueCell.update(v.tAlpha, { width: data.nAlpha, height: 1, array: data.alpha });
         ValueCell.updateIfChanged(v.uLittleEndian, isLittleEndian());
+        ValueCell.updateIfChanged(v.uOccupancy, 0);
 
         ctx.namedComputeRenderables[AlphaOrbitalsName].update();
     } else {
@@ -211,4 +241,68 @@ export function gpuComputeAlphaOrbitalsGridValues(webgl: WebGLContext, params: C
 
 export function canComputeAlphaOrbitalsOnGPU(webgl?: WebGLContext) {
     return !!webgl?.extensions.textureFloat;
+}
+
+const AlphaOrbitalsDensity0 = AlphaOrbitalsName;
+const AlphaOrbitalsDensity1 = AlphaOrbitalsName + '1';
+
+export function gpuComputeAlphaOrbitalsDensityGridValues(webgl: WebGLContext, params: CollocationParams, orbitals: number[][]) {
+    const [nx, ny, nz] = params.grid.dimensions;
+    const renderable = getAlphaOrbitalsRenderable(webgl, params);
+    const width = renderable.values.uWidth.ref.value;
+
+    if (!webgl.namedFramebuffers[AlphaOrbitalsDensity0]) {
+        webgl.namedFramebuffers[AlphaOrbitalsDensity0] = webgl.resources.framebuffer();
+    }
+    const framebuffer = webgl.namedFramebuffers[AlphaOrbitalsDensity0];
+
+    if (!webgl.namedTextures[AlphaOrbitalsDensity0]) {
+        webgl.namedTextures[AlphaOrbitalsDensity0] = webgl.resources.texture('image-uint8', 'rgba', 'ubyte', 'nearest');
+    }
+    if (!webgl.namedTextures[AlphaOrbitalsDensity1]) {
+        webgl.namedTextures[AlphaOrbitalsDensity1] = webgl.resources.texture('image-uint8', 'rgba', 'ubyte', 'nearest');
+    }
+
+
+    const tex = [webgl.namedTextures[AlphaOrbitalsDensity0], webgl.namedTextures[AlphaOrbitalsDensity1]];
+
+    tex[0].define(width, width);
+    tex[1].define(width, width);
+
+    const values = renderable.values as AlphaOrbitalsSchema;
+    const { gl, state } = webgl;
+
+    gl.viewport(0, 0, width, width);
+    gl.scissor(0, 0, width, width);
+    state.disable(gl.SCISSOR_TEST);
+    state.disable(gl.BLEND);
+    state.disable(gl.DEPTH_TEST);
+    state.depthMask(false);
+
+    gl.clearColor(0, 0, 0, 0);
+
+    tex[0].attachFramebuffer(framebuffer, 'color0');
+    gl.clear(gl.COLOR_BUFFER_BIT);
+
+    tex[1].attachFramebuffer(framebuffer, 'color0');
+    gl.clear(gl.COLOR_BUFFER_BIT);
+
+    for (let i = 0; i < 108; i++) {
+        const alpha = getNormalizedAlpha(params.basis, orbitals[i], params.sphericalOrder);
+
+        ValueCell.update(values.uOccupancy, 2);
+        ValueCell.update(values.tCumulativeSum, tex[(i + 1) % 2]);
+        ValueCell.update(values.tAlpha, { width: alpha.length, height: 1, array: alpha });
+        renderable.update();
+        tex[i % 2].attachFramebuffer(framebuffer, 'color0');
+        renderable.render();
+    }
+
+    const array = new Uint8Array(width * width * 4);
+    webgl.readPixels(0, 0, width, width, array);
+
+    tex[0].define(1, 1);
+    tex[1].define(1, 1);
+
+    return new Float32Array(array.buffer, array.byteOffset, nx * ny * nz);
 }

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

@@ -28,6 +28,9 @@ uniform float uWidth;
 uniform int uNCoeff;
 uniform int uNAlpha;
 
+uniform float uOccupancy;
+uniform sampler2D tCumulativeSum;
+
 uniform bool uLittleEndian;
 
 float shiftRight (float v, float amt) {
@@ -67,6 +70,78 @@ vec4 floatToRgba(float texelFloat) {
     );
 }
 
+///////////////////////////////////////////////////////
+
+ivec4 floatsToBytes(vec4 inputFloats) {
+  ivec4 bytes = ivec4(inputFloats * 255.0);
+  return (
+    uLittleEndian
+    ? bytes.abgr
+    : bytes
+  );
+}
+
+// Break the four bytes down into an array of 32 bits.
+void bytesToBits(const in ivec4 bytes, out bool bits[32]) {
+  for (int channelIndex = 0; channelIndex < 4; ++channelIndex) {
+    float acc = float(bytes[channelIndex]);
+    for (int indexInByte = 7; indexInByte >= 0; --indexInByte) {
+      float powerOfTwo = exp2(float(indexInByte));
+      bool bit = acc >= powerOfTwo;
+      bits[channelIndex * 8 + (7 - indexInByte)] = bit;
+      acc = mod(acc, powerOfTwo);
+    }
+  }
+}
+
+// Compute the exponent of the 32-bit float.
+float getExponent(bool bits[32]) {
+  const int startIndex = 1;
+  const int bitStringLength = 8;
+  const int endBeforeIndex = startIndex + bitStringLength;
+  float acc = 0.0;
+  int pow2 = bitStringLength - 1;
+  for (int bitIndex = startIndex; bitIndex < endBeforeIndex; ++bitIndex) {
+    acc += float(bits[bitIndex]) * exp2(float(pow2--));
+  }
+  return acc;
+}
+
+// Compute the mantissa of the 32-bit float.
+float getMantissa(bool bits[32], bool subnormal) {
+  const int startIndex = 9;
+  const int bitStringLength = 23;
+  const int endBeforeIndex = startIndex + bitStringLength;
+  // Leading/implicit/hidden bit convention:
+  // If the number is not subnormal (with exponent 0), we add a leading 1 digit.
+  float acc = float(!subnormal) * exp2(float(bitStringLength));
+  int pow2 = bitStringLength - 1;
+  for (int bitIndex = startIndex; bitIndex < endBeforeIndex; ++bitIndex) {
+    acc += float(bits[bitIndex]) * exp2(float(pow2--));
+  }
+  return acc;
+}
+
+// Parse the float from its 32 bits.
+float bitsToFloat(bool bits[32]) {
+  float signBit = float(bits[0]) * -2.0 + 1.0;
+  float exponent = getExponent(bits);
+  bool subnormal = abs(exponent - 0.0) < 0.01;
+  float mantissa = getMantissa(bits, subnormal);
+  float exponentBias = 127.0;
+  return signBit * mantissa * exp2(exponent - exponentBias - 23.0);
+}
+
+// Decode a 32-bit float from the RGBA color channels of a texel.
+float rgbaToFloat(vec4 texelRGBA) {
+  ivec4 rgbaBytes = floatsToBytes(texelRGBA);
+  bool bits[32];
+  bytesToBits(rgbaBytes, bits);
+  return bitsToFloat(bits);
+}
+
+///////////////////////////////////////////////////////
+
 float L1(vec3 p, float a0, float a1, float a2) {
     return a0 * p.z + a1 * p.x + a2 * p.y;
 }
@@ -224,7 +299,12 @@ void main(void) {
         v += R(R2, coeffStart, coeffEnd, fCoeff) * Y(L, X, aO, fA);
     }
 
-    // TODO: render to single component float32 texture in WebGL2
-    gl_FragColor = floatToRgba(v);
+    
+    if (uOccupancy > 0.0) {
+        float current = rgbaToFloat(texture2D(tCumulativeSum, gl_FragCoord.xy / vec2(uWidth, uWidth)));
+        gl_FragColor = floatToRgba(current + uOccupancy * v * v);
+    } else {
+        gl_FragColor = floatToRgba(v);
+    }
 }
 `;

src/extensions/alpha-orbitals/transforms.ts

@@ -76,7 +76,7 @@ export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
                 sphericalOrder: a.data.order,
                 boxExpand: params.boxExpand,
                 gridSpacing: params.gridSpacing.map(e => [e.atomCount, e.spacing] as [number, number])
-            }, params.forceCpuCompute ? void 0 : plugin.canvas3d?.webgl).runInContext(ctx);
+            }, params.forceCpuCompute ? void 0 : plugin.canvas3d?.webgl, a.data.orbitals.map(o => o.alpha)).runInContext(ctx);
             const volume: Volume = {
                 grid: data.grid,
                 sourceData: { name: 'custom grid', kind: 'alpha-orbitals', data },