alpha-orbitals: density support

src/examples/alpha-orbitals/index.ts

@@ -5,7 +5,7 @@
  */
 
 import { SphericalBasisOrder } from '../../extensions/alpha-orbitals/spherical-functions';
-import { CreateOrbitalRepresentation3D, CreateOrbitalVolume, StaticBasisAndOrbitals } from '../../extensions/alpha-orbitals/transforms';
+import { BasisAndOrbitals, CreateOrbitalDensityVolume, CreateOrbitalRepresentation3D, CreateOrbitalVolume, StaticBasisAndOrbitals } from '../../extensions/alpha-orbitals/transforms';
 import { createPluginAsync, DefaultPluginSpec } from '../../mol-plugin';
 import { PluginStateObject } from '../../mol-plugin-state/objects';
 import { PluginConfig } from '../../mol-plugin/config';
@@ -32,11 +32,22 @@ interface DemoInput {
 }
 
 interface Params {
-    orbitalIndex: number,
+    show: { name: 'orbital', params: { index: number } } | { name: 'density', params: {} },
     isoValue: number,
     gpuSurface: boolean
 }
 
+type Selectors = {
+    type: 'orbital',
+    volume: StateObjectSelector<PluginStateObject.Volume.Data, typeof CreateOrbitalVolume>,
+    positive: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof CreateOrbitalRepresentation3D>
+    negative: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof CreateOrbitalRepresentation3D>
+} | {
+    type: 'density',
+    volume: StateObjectSelector<PluginStateObject.Volume.Data, typeof CreateOrbitalDensityVolume>,
+    positive: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof CreateOrbitalRepresentation3D>
+}
+
 export class AlphaOrbitalsExample {
     plugin: PluginContext;
 
@@ -77,28 +88,72 @@ export class AlphaOrbitalsExample {
     }
 
     readonly params = new BehaviorSubject<ParamDefinition.For<Params>>({} as any);
-    readonly state = new BehaviorSubject<Params>({ orbitalIndex: 32, isoValue: 1, gpuSurface: false });
+    readonly state = new BehaviorSubject<Params>({ show: { name: 'orbital', params: { index: 32 } }, isoValue: 1, gpuSurface: false });
+
+    private selectors?: Selectors = void 0;
+    private basis?: StateObjectSelector<BasisAndOrbitals> = void 0;
 
-    private volume?: StateObjectSelector<PluginStateObject.Volume.Data, typeof CreateOrbitalVolume>;
-    private positive?: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof CreateOrbitalRepresentation3D>;
-    private negative?: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof CreateOrbitalRepresentation3D>;
     private currentParams: Params = { ...this.state.value };
 
-    private async setIndex() {
-        if (!this.volume?.isOk) return;
+    private clearVolume() {
+        if (!this.selectors) return;
+        const v = this.selectors.volume;
+        this.selectors = void 0;
+        return this.plugin.build().delete(v).commit();
+    }
+
+    private async syncVolume() {
+        if (!this.basis?.isOk) return;
+
         const state = this.state.value;
-        await this.plugin.build().to(this.volume).update(CreateOrbitalVolume, () => ({ index: state.orbitalIndex })).commit();
+
+        if (state.show.name !== this.selectors?.type) {
+            await this.clearVolume();
+        }
+
+        const update = this.plugin.build();
+        if (state.show.name === 'orbital') {
+            if (!this.selectors) {
+                const volume = update
+                    .to(this.basis)
+                    .apply(CreateOrbitalVolume, { index: state.show.params.index });
+
+                const positive = volume.apply(CreateOrbitalRepresentation3D, this.volumeParams('positive', ColorNames.blue)).selector;
+                const negative = volume.apply(CreateOrbitalRepresentation3D, this.volumeParams('negative', ColorNames.red)).selector;
+
+                this.selectors = { type: 'orbital', volume: volume.selector, positive, negative };
+            } else {
+                const index = state.show.params.index;
+                update.to(this.selectors.volume).update(CreateOrbitalVolume, () => ({ index }));
+            }
+        } else {
+            if (!this.selectors) {
+                const volume = update
+                    .to(this.basis)
+                    .apply(CreateOrbitalDensityVolume);
+                const positive = volume.apply(CreateOrbitalRepresentation3D, this.volumeParams('positive', ColorNames.blue)).selector;
+                this.selectors = { type: 'density', volume: volume.selector, positive };
+            }
+        }
+
+        await update.commit();
+
         if (this.currentParams.gpuSurface !== this.state.value.gpuSurface) {
             await this.setIsovalue();
         }
+
         this.currentParams = this.state.value;
     }
 
     private setIsovalue() {
+        if (!this.selectors) return;
+
         this.currentParams = this.state.value;
         const update = this.plugin.build();
-        update.to(this.positive!).update(this.volumeParams('positive', ColorNames.blue));
-        update.to(this.negative!).update(this.volumeParams('negative', ColorNames.red));
+        update.to(this.selectors.positive).update(this.volumeParams('positive', ColorNames.blue));
+        if (this.selectors?.type === 'orbital') {
+            update.to(this.selectors.negative).update(this.volumeParams('negative', ColorNames.red));
+        }
         return update.commit();
     }
 
@@ -125,34 +180,28 @@ export class AlphaOrbitalsExample {
         const all = await this.plugin.builders.structure.tryCreateComponentStatic(structure, 'all');
         if (all) await this.plugin.builders.structure.representation.addRepresentation(all, { type: 'ball-and-stick', color: 'element-symbol', colorParams: { carbonColor: { name: 'element-symbol', params: {} } } });
 
-        const state = this.state.value;
 
-        this.volume = await this.plugin.build().toRoot()
+        this.basis = await this.plugin.build().toRoot()
             .apply(StaticBasisAndOrbitals, { basis: input.basis, order: input.order, orbitals: input.orbitals })
-            .apply(CreateOrbitalVolume, { index: state.orbitalIndex })
             .commit();
 
-        if (!this.volume.isOk) {
-            this.volume = void 0;
-            return;
-        }
-
-        const repr = this.plugin.build().to(this.volume);
-
-        this.positive = repr.apply(CreateOrbitalRepresentation3D, this.volumeParams('positive', ColorNames.blue)).selector;
-        this.negative = repr.apply(CreateOrbitalRepresentation3D, this.volumeParams('negative', ColorNames.red)).selector;
-
-        await repr.commit();
+        await this.syncVolume();
 
         this.params.next({
-            orbitalIndex: ParamDefinition.Numeric(this.currentParams.orbitalIndex, { min: 0, max: input.orbitals.length - 1 }, { immediateUpdate: true, isEssential: true }),
+            show: ParamDefinition.MappedStatic('orbital', {
+                'orbital': ParamDefinition.Group({
+                    index: ParamDefinition.Numeric(0, { min: 0, max: input.orbitals.length - 1 }, { immediateUpdate: true, isEssential: true }),
+                }),
+                'density': ParamDefinition.EmptyGroup()
+            }, { cycle: true }),
             isoValue: ParamDefinition.Numeric(this.currentParams.isoValue, { min: 0.5, max: 3, step: 0.1 }, { immediateUpdate: true, isEssential: false }),
-            gpuSurface: ParamDefinition.Boolean(this.currentParams.gpuSurface)
+            gpuSurface: ParamDefinition.Boolean(this.currentParams.gpuSurface, { isHidden: true })
         });
 
         this.state.pipe(skip(1), debounceTime(1000 / 24)).subscribe(async params => {
-            if (params.orbitalIndex !== this.currentParams.orbitalIndex) {
-                this.setIndex();
+            if (params.show.name !== this.currentParams.show.name
+                || (params.show.name === 'orbital' && this.currentParams.show.name === 'orbital' && params.show.params.index !== this.currentParams.show.params.index)) {
+                this.syncVolume();
             } else if (params.isoValue !== this.currentParams.isoValue || params.gpuSurface !== this.currentParams.gpuSurface) {
                 this.setIsovalue();
             }

src/extensions/alpha-orbitals/density.ts

@@ -0,0 +1,123 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { sortArray } from '../../mol-data/util';
+import { WebGLContext } from '../../mol-gl/webgl/context';
+import { Task } from '../../mol-task';
+import { AlphaOrbital, createGrid, CubeGrid, CubeGridComputationParams, initCubeGrid } from './data-model';
+import { canComputeAlphaOrbitalsOnGPU, gpuComputeAlphaOrbitalsDensityGridValues } from './gpu/compute';
+
+export function createSphericalCollocationDensityGrid(
+    params: CubeGridComputationParams, orbitals: AlphaOrbital[], webgl?: WebGLContext
+): Task<CubeGrid> {
+    return Task.create('Spherical Collocation Grid', async (ctx) => {
+        const cubeGrid = initCubeGrid(params);
+
+        let matrix: Float32Array;
+        if (canComputeAlphaOrbitalsOnGPU(webgl)) {
+            // console.time('gpu');
+            matrix = await gpuComputeAlphaOrbitalsDensityGridValues(webgl!, cubeGrid, orbitals, ctx);
+            // console.timeEnd('gpu');
+        } else {
+            throw new Error('Missing OES_texture_float WebGL extension.');
+        }
+
+        const grid = createGrid(cubeGrid, matrix, [0, 1, 2]);
+        let isovalues: { negative?: number, positive?: number } | undefined;
+
+        if (!params.doNotComputeIsovalues) {
+            isovalues = computeDensityIsocontourValues(matrix, 0.85);
+        }
+
+        return { grid, isovalues };
+    });
+}
+
+export function computeDensityIsocontourValues(input: Float32Array, cumulativeThreshold: number) {
+    let weightSum = 0;
+    for (let i = 0, _i = input.length; i < _i; i++) {
+        const v = input[i];
+        const w = Math.abs(v);
+        weightSum += w;
+    }
+    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;
+        size = 0;
+        for (let i = 0, _i = input.length; i < _i; i++) {
+            const v = input[i];
+            const w = Math.abs(v);
+            if (w >= minWeight) {
+                csum += w;
+                size++;
+            }
+        }
+
+        if (csum / weightSum > cumulativeThreshold) {
+            break;
+        }
+
+        minWeight -= avgWeight;
+    }
+
+    const values = new Float32Array(size);
+    const weights = new Float32Array(size);
+    const indices = new Int32Array(size);
+
+    let o = 0;
+    for (let i = 0, _i = input.length; i < _i; i++) {
+        const v = input[i];
+        const w = Math.abs(v);
+        if (w >= minWeight) {
+            values[o] = v;
+            weights[o] = w;
+            indices[o] = o;
+            o++;
+        }
+    }
+
+    sortArray(
+        indices,
+        (indices, i, j) => weights[indices[j]] - weights[indices[i]]
+    );
+
+    let cweight = 0,
+        cutoffIndex = 0;
+    for (let i = 0; i < size; i++) {
+        cweight += weights[indices[i]];
+
+        if (cweight / weightSum >= cumulativeThreshold) {
+            cutoffIndex = i;
+            break;
+        }
+    }
+
+    let positive = Number.POSITIVE_INFINITY,
+        negative = Number.NEGATIVE_INFINITY;
+
+    for (let i = 0; i < cutoffIndex; i++) {
+        const v = values[indices[i]];
+        if (v > 0) {
+            if (v < positive) positive = v;
+        } else if (v < 0) {
+            if (v > negative) negative = v;
+        }
+    }
+
+    return {
+        negative: negative !== Number.NEGATIVE_INFINITY ? negative : void 0,
+        positive: positive !== Number.POSITIVE_INFINITY ? positive : void 0,
+    };
+}

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

@@ -11,6 +11,7 @@ import { ShaderCode } from '../../../mol-gl/shader-code';
 import quad_vert from '../../../mol-gl/shader/quad.vert';
 import { WebGLContext } from '../../../mol-gl/webgl/context';
 import { createComputeRenderItem } from '../../../mol-gl/webgl/render-item';
+import { RuntimeContext } from '../../../mol-task';
 import { ValueCell } from '../../../mol-util';
 import { arrayMin } from '../../../mol-util/array';
 import { isLittleEndian } from '../../../mol-util/is-little-endian';
@@ -242,11 +243,9 @@ export function canComputeAlphaOrbitalsOnGPU(webgl?: WebGLContext) {
     return !!webgl?.extensions.textureFloat;
 }
 
-export function gpuComputeAlphaOrbitalsDensityGridValues(webgl: WebGLContext, grid: CubeGridInfo, orbitals: AlphaOrbital[]) {
-    return _gpuComputeAlphaOrbitalsDensityGridValues(webgl, grid, orbitals);
-}
+export async function gpuComputeAlphaOrbitalsDensityGridValues(webgl: WebGLContext, grid: CubeGridInfo, orbitals: AlphaOrbital[], ctx: RuntimeContext) {
+    await ctx.update({ message: 'Initializing...', isIndeterminate: true });
 
-function _gpuComputeAlphaOrbitalsDensityGridValues(webgl: WebGLContext, grid: CubeGridInfo, orbitals: AlphaOrbital[]) {
     const [nx, ny, nz] = grid.dimensions;
     const renderable = getAlphaOrbitalsRenderable(webgl, grid, orbitals[0]);
     const width = renderable.values.uWidth.ref.value;
@@ -280,19 +279,25 @@ function _gpuComputeAlphaOrbitalsDensityGridValues(webgl: WebGLContext, grid: Cu
 
     ValueCell.update(values.uDensity, true);
 
-    for (let i = 0; i < orbitals.length; i++) {
-        if (orbitals[i].occupancy === 0) continue;
-
-        const alpha = getNormalizedAlpha(grid.params.basis, orbitals[i].alpha, grid.params.sphericalOrder);
+    const nonZero = orbitals.filter(o => o.occupancy !== 0);
+    await ctx.update({ message: 'Computing...', isIndeterminate: false, current: 0, max: nonZero.length });
+    for (let i = 0; i < nonZero.length; i++) {
+        const alpha = getNormalizedAlpha(grid.params.basis, nonZero[i].alpha, grid.params.sphericalOrder);
 
-        ValueCell.update(values.uOccupancy, orbitals[i].occupancy);
+        ValueCell.update(values.uOccupancy, nonZero[i].occupancy);
         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();
+
+        if (ctx.shouldUpdate) {
+            await ctx.update({ current: i + 1 });
+        }
     }
 
+    await ctx.update({ message: 'Finalizing...', isIndeterminate: true });
+
     const array = new Uint8Array(width * width * 4);
     webgl.readPixels(0, 0, width, width, array);
 

src/extensions/alpha-orbitals/orbitals.ts

@@ -41,10 +41,7 @@ export function createSphericalCollocationGrid(
     });
 }
 
-export function computeOrbitalIsocontourValues(
-    input: Float32Array,
-    cumulativeThreshold: number
-) {
+export function computeOrbitalIsocontourValues(input: Float32Array, cumulativeThreshold: number) {
     let weightSum = 0;
     for (let i = 0, _i = input.length; i < _i; i++) {
         const v = input[i];

src/extensions/alpha-orbitals/transforms.ts

@@ -18,6 +18,7 @@ import { StateTransformer } from '../../mol-state';
 import { Theme } from '../../mol-theme/theme';
 import { VolumeRepresentation3DHelpers } from '../../mol-plugin-state/transforms/representation';
 import { AlphaOrbital, Basis, CubeGrid } from './data-model';
+import { createSphericalCollocationDensityGrid } from './density';
 
 export class BasisAndOrbitals extends PluginStateObject.Create<{ basis: Basis, order: SphericalBasisOrder, orbitals: AlphaOrbital[] }>({ name: 'Basis', typeClass: 'Object' }) { }
 
@@ -39,7 +40,6 @@ export const StaticBasisAndOrbitals = PluginStateTransform.BuiltIn({
 });
 
 const CreateOrbitalVolumeParamBase = {
-    forceCpuCompute: PD.Boolean(false),
     cutoffThreshold: PD.Numeric(0.0015, { min: 0, max: 0.1, step: 0.0001 }),
     boxExpand: PD.Numeric(4.5, { min: 0, max: 7, step: 0.1 }),
     gridSpacing: PD.ObjectList({ atomCount: PD.Numeric(0), spacing: PD.Numeric(0.35, { min: 0.1, max: 2, step: 0.01 }) }, e => `Atoms ${e.atomCount}: ${e.spacing}`, {
@@ -76,7 +76,35 @@ export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
                 sphericalOrder: a.data.order,
                 boxExpand: params.boxExpand,
                 gridSpacing: params.gridSpacing.map(e => [e.atomCount, e.spacing] as [number, number])
-            }, a.data.orbitals[params.index], params.forceCpuCompute ? void 0 : plugin.canvas3d?.webgl).runInContext(ctx);
+            }, a.data.orbitals[params.index], plugin.canvas3d?.webgl).runInContext(ctx);
+            const volume: Volume = {
+                grid: data.grid,
+                sourceData: { name: 'custom grid', kind: 'alpha-orbitals', data },
+                customProperties: new CustomProperties(),
+                _propertyData: Object.create(null),
+            };
+
+            return new PluginStateObject.Volume.Data(volume, { label: 'Orbital Volume' });
+        });
+    }
+});
+
+export const CreateOrbitalDensityVolume = PluginStateTransform.BuiltIn({
+    name: 'create-orbital-density-volume',
+    display: 'Orbital Density Volume',
+    from: BasisAndOrbitals,
+    to: PluginStateObject.Volume.Data,
+    params: CreateOrbitalVolumeParamBase
+})({
+    apply({ a, params }, plugin: PluginContext) {
+        return Task.create('Orbital Volume', async ctx => {
+            const data = await createSphericalCollocationDensityGrid({
+                basis: a.data.basis,
+                cutoffThreshold: params.cutoffThreshold,
+                sphericalOrder: a.data.order,
+                boxExpand: params.boxExpand,
+                gridSpacing: params.gridSpacing.map(e => [e.atomCount, e.spacing] as [number, number])
+            }, a.data.orbitals, plugin.canvas3d?.webgl).runInContext(ctx);
             const volume: Volume = {
                 grid: data.grid,
                 sourceData: { name: 'custom grid', kind: 'alpha-orbitals', data },