alpha-orbitals: refactoring

src/examples/alpha-orbitals/controls.tsx

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  */

src/examples/alpha-orbitals/index.ts

@@ -1,27 +1,25 @@
-import { BehaviorSubject } from 'rxjs';
-import { debounceTime, skip } from 'rxjs/operators';
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
- * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Basis, computeIsocontourValues } from '../../extensions/alpha-orbitals/cubes';
+import { Basis } from '../../extensions/alpha-orbitals/cubes';
 import { SphericalBasisOrder } from '../../extensions/alpha-orbitals/orbitals';
+import { CreateOrbitalRepresentation3D, CreateOrbitalVolume, StaticBasisAndOrbitals } from '../../extensions/alpha-orbitals/transforms';
 import { createPluginAsync, DefaultPluginSpec } from '../../mol-plugin';
-import { createVolumeRepresentationParams } from '../../mol-plugin-state/helpers/volume-representation-params';
 import { PluginStateObject } from '../../mol-plugin-state/objects';
-import { StateTransforms } from '../../mol-plugin-state/transforms';
 import { PluginConfig } from '../../mol-plugin/config';
 import { PluginContext } from '../../mol-plugin/context';
-import { StateObjectSelector } from '../../mol-state';
+import { StateObjectSelector, StateTransformer } from '../../mol-state';
 import { Color } from '../../mol-util/color';
 import { ColorNames } from '../../mol-util/color/names';
 import { ParamDefinition } from '../../mol-util/param-definition';
 import { mountControls } from './controls';
 import { DemoMoleculeSDF, DemoOrbitals } from './example-data';
+import { BehaviorSubject } from 'rxjs';
+import { debounceTime, skip } from 'rxjs/operators';
 import './index.html';
-import { CreateOrbitalVolume, StaticBasisAndOrbitals } from './transforms';
 require('mol-plugin-ui/skin/light.scss');
 
 interface DemoInput {
@@ -37,7 +35,6 @@ interface DemoInput {
 interface Params {
     orbitalIndex: number,
     isoValue: number,
-    staticIsovalues: boolean,
     gpuSurface: boolean,
     cpuCompute: boolean
 }
@@ -74,53 +71,39 @@ export class AlphaOrbitalsExample {
     }
 
     readonly params = new BehaviorSubject<ParamDefinition.For<Params>>({} as any);
-    readonly state = new BehaviorSubject<Params>({ orbitalIndex: 32, isoValue: 1, staticIsovalues: false, gpuSurface: false, cpuCompute: false });
+    readonly state = new BehaviorSubject<Params>({ orbitalIndex: 32, isoValue: 1, gpuSurface: false, cpuCompute: false });
 
     private volume?: StateObjectSelector<PluginStateObject.Volume.Data, typeof CreateOrbitalVolume>;
-    private positive?: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof StateTransforms.Representation.VolumeRepresentation3D>;
-    private negative?: StateObjectSelector<PluginStateObject.Volume.Representation3D, typeof StateTransforms.Representation.VolumeRepresentation3D>;
-    private isovalues: { negative?: number, positive?: number } = { negative: void 0, positive: void 0 }
+    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;
         const state = this.state.value;
         await this.plugin.build().to(this.volume).update(CreateOrbitalVolume, () => ({ index: state.orbitalIndex, cpuCompute: state.cpuCompute })).commit();
-        if (!state.staticIsovalues) {
-            this.isovalues = computeIsocontourValues(this.volume.data!.grid.cells.data as any, 0.85);
-        }
-        if (!state.staticIsovalues || this.currentParams.gpuSurface !== this.state.value.gpuSurface) {
+        if (this.currentParams.gpuSurface !== this.state.value.gpuSurface) {
             await this.setIsovalue();
         }
         this.currentParams = this.state.value;
     }
 
     private setIsovalue() {
-        const { positive, negative } = this.isovalues;
         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.positive!).update(this.volumeParams('positive', ColorNames.blue));
+        update.to(this.negative!).update(this.volumeParams('negative', ColorNames.red));
         return update.commit();
     }
 
-    private volumeParams(value: number | undefined, color: Color) {
-        return createVolumeRepresentationParams(this.plugin, this.volume!.data!, this.currentParams.gpuSurface ? {
-            type: 'direct-volume',
-            typeParams: {
-                renderMode: {
-                    name: 'isosurface',
-                    params: { isoValue: { kind: 'absolute', absoluteValue: (value ?? 1000) * this.state.value.isoValue }, singleLayer: true }
-                }
-            },
-            color: 'uniform',
-            colorParams: { value: color }
-        } : {
-            type: 'isosurface',
-            typeParams: { isoValue: { kind: 'absolute', absoluteValue: (value ?? 1000) * this.state.value.isoValue } },
-            color: 'uniform',
-            colorParams: { value: color }
-        });
+    private volumeParams(kind: 'positive' | 'negative', color: Color): StateTransformer.Params<typeof CreateOrbitalRepresentation3D> {
+        return {
+            alpha: 1.0,
+            color,
+            directVolume: this.state.value.gpuSurface,
+            kind,
+            relativeIsovalue: this.state.value.isoValue
+        };
     }
 
     async load(input: DemoInput) {
@@ -138,7 +121,7 @@ export class AlphaOrbitalsExample {
 
         this.volume = await this.plugin.build().toRoot()
             .apply(StaticBasisAndOrbitals, { basis: input.basis, order: input.order, orbitals: input.orbitals })
-            .apply(CreateOrbitalVolume, { index: state.orbitalIndex, cpuCompute: this.currentParams.cpuCompute })
+            .apply(CreateOrbitalVolume, { index: state.orbitalIndex, forceCpuCompute: this.currentParams.cpuCompute })
             .commit();
 
         if (!this.volume.isOk) {
@@ -146,29 +129,22 @@ export class AlphaOrbitalsExample {
             return;
         }
 
-        // TODO: the isovalues are being computed twice. Need to add more flexible support to Volume object
-        //       for controlling them
-        this.isovalues = computeIsocontourValues(this.volume.data!.grid.cells.data as any, 0.85);
-        const { positive, negative } = this.isovalues;
-
         const repr = this.plugin.build().to(this.volume);
 
-        this.positive = repr.apply(StateTransforms.Representation.VolumeRepresentation3D, this.volumeParams(positive, ColorNames.blue)).selector;
-        this.negative = repr.apply(StateTransforms.Representation.VolumeRepresentation3D, this.volumeParams(negative, ColorNames.red)).selector;
+        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();
 
         this.params.next({
             orbitalIndex: ParamDefinition.Numeric(this.currentParams.orbitalIndex, { min: 0, max: input.orbitals.length - 1 }, { immediateUpdate: true, isEssential: true }),
             isoValue: ParamDefinition.Numeric(this.currentParams.isoValue, { min: 0.5, max: 3, step: 0.1 }, { immediateUpdate: true, isEssential: false }),
-            staticIsovalues: ParamDefinition.Boolean(this.currentParams.staticIsovalues),
             gpuSurface: ParamDefinition.Boolean(this.currentParams.gpuSurface),
-            cpuCompute: ParamDefinition.Boolean(this.currentParams.gpuSurface)
+            cpuCompute: ParamDefinition.Boolean(this.currentParams.cpuCompute, { label: 'CPU Compute' })
         });
 
         this.state.pipe(skip(1), debounceTime(1000 / 24)).subscribe(async params => {
             if (params.orbitalIndex !== this.currentParams.orbitalIndex
-                || params.staticIsovalues !== this.currentParams.staticIsovalues
                 || params.cpuCompute !== this.currentParams.cpuCompute) {
                 this.setIndex();
             } else if (params.isoValue !== this.currentParams.isoValue || params.gpuSurface !== this.currentParams.gpuSurface) {

src/examples/alpha-orbitals/transforms.ts

@@ -1,77 +0,0 @@
-/**
- * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- */
-
-import { PluginStateObject, PluginStateTransform } from '../../mol-plugin-state/objects';
-import { Basis, createSphericalCollocationGrid } from '../../extensions/alpha-orbitals/cubes';
-import { ParamDefinition as PD } from '../../mol-util/param-definition';
-import { Task } from '../../mol-task';
-import { CustomProperties } from '../../mol-model/custom-property';
-import { SphericalBasisOrder } from '../../extensions/alpha-orbitals/orbitals';
-import { Volume } from '../../mol-model/volume';
-import { PluginContext } from '../../mol-plugin/context';
-
-export class BasisAndOrbitals extends PluginStateObject.Create<{ basis: Basis, order: SphericalBasisOrder, orbitals: { energy: number, alpha: number[] }[] }>({ name: 'Basis', typeClass: 'Object' }) { }
-
-export const StaticBasisAndOrbitals = PluginStateTransform.BuiltIn({
-    name: 'static-basis-and-orbitals',
-    display: 'Basis and Orbitals',
-    from: PluginStateObject.Root,
-    to: BasisAndOrbitals,
-    params: {
-        label: PD.Text('Orbital Data', { isHidden: true }),
-        basis: PD.Value<Basis>(void 0 as any, { isHidden: true }),
-        order: PD.Text<SphericalBasisOrder>('gaussian' as SphericalBasisOrder, { isHidden: true }),
-        orbitals: PD.Value<{ energy: number, alpha: number[] }[]>([], { isHidden: true })
-    },
-})({
-    apply({ params }) {
-        return new BasisAndOrbitals({ basis: params.basis, order: params.order, orbitals: params.orbitals }, { label: params.label });
-    }
-});
-
-export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
-    name: 'create-orbital-volume',
-    display: 'Orbital Volume',
-    from: BasisAndOrbitals,
-    to: PluginStateObject.Volume.Data,
-    params: (a) => {
-        if (!a) {
-            return { index: PD.Numeric(0), cpuCompute: PD.Boolean(false) };
-        }
-
-        return {
-            index: PD.Select(0, a.data.orbitals.map((o, i) => [i, `[${i + 1}] ${o.energy.toFixed(4)}`])),
-            cpuCompute: PD.Boolean(false)
-        };
-    }
-})({
-    apply({ a, params }, plugin: PluginContext) {
-        return Task.create('Orbital Volume', async ctx => {
-            const data = await createSphericalCollocationGrid({
-                basis: a.data.basis,
-                cutoffThreshold: 0.0015,
-                alphaOrbitals: a.data.orbitals[params.index].alpha,
-                sphericalOrder: a.data.order,
-                boxExpand: 4.5,
-                gridSpacing: [
-                    [55, 0.5],
-                    [40, 0.45],
-                    [25, 0.4],
-                    [0, 0.35],
-                ],
-                doNotComputeIsovalues: true
-            }, params.cpuCompute ? void 0 : plugin.canvas3d?.webgl).runInContext(ctx);
-            const volume: Volume = {
-                grid: data.grid,
-                sourceData: { name: 'custom grid', kind: 'custom', data },
-                customProperties: new CustomProperties(),
-                _propertyData: Object.create(null),
-            };
-
-            return new PluginStateObject.Volume.Data(volume, { label: 'Orbital Volume' });
-        });
-    }
-});

src/extensions/alpha-orbitals/_spec/collocation.spec.ts

@@ -1,3 +1,8 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
 
 import { Box3D } from '../../../mol-math/geometry';
 import { Vec3 } from '../../../mol-math/linear-algebra';

src/extensions/alpha-orbitals/cubes.ts

@@ -78,13 +78,13 @@ export function createSphericalCollocationGrid(
 
         let matrix: Float32Array;
         if (canComputeAlphaOrbitalsOnGPU(webgl)) {
-            console.time('gpu');
+            // console.time('gpu');
             matrix = gpuComputeAlphaOrbitalsGridValues(webgl!, cParams);
-            console.timeEnd('gpu');
+            // console.timeEnd('gpu');
         } else {
-            console.time('cpu');
+            // console.time('cpu');
             matrix = await sphericalCollocation(cParams, ctx);
-            console.timeEnd('cpu');
+            // console.timeEnd('cpu');
         }
 
         return createCubeGrid(cParams.grid, matrix, [0, 1, 2], !params.doNotComputeIsovalues);

src/extensions/alpha-orbitals/transforms.ts

@@ -0,0 +1,160 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { PluginStateObject, PluginStateTransform } from '../../mol-plugin-state/objects';
+import { Basis, createSphericalCollocationGrid, CubeGrid } from './cubes';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { Task } from '../../mol-task';
+import { CustomProperties } from '../../mol-model/custom-property';
+import { SphericalBasisOrder } from './orbitals';
+import { Volume } from '../../mol-model/volume';
+import { PluginContext } from '../../mol-plugin/context';
+import { ColorNames } from '../../mol-util/color/names';
+import { createVolumeRepresentationParams } from '../../mol-plugin-state/helpers/volume-representation-params';
+import { StateTransformer } from '../../mol-state';
+import { Theme } from '../../mol-theme/theme';
+import { VolumeRepresentation3DHelpers } from '../../mol-plugin-state/transforms/representation';
+
+export class BasisAndOrbitals extends PluginStateObject.Create<{ basis: Basis, order: SphericalBasisOrder, orbitals: { energy: number, alpha: number[] }[] }>({ name: 'Basis', typeClass: 'Object' }) { }
+
+export const StaticBasisAndOrbitals = PluginStateTransform.BuiltIn({
+    name: 'static-basis-and-orbitals',
+    display: 'Basis and Orbitals',
+    from: PluginStateObject.Root,
+    to: BasisAndOrbitals,
+    params: {
+        label: PD.Text('Orbital Data', { isHidden: true }),
+        basis: PD.Value<Basis>(void 0 as any, { isHidden: true }),
+        order: PD.Text<SphericalBasisOrder>('gaussian' as SphericalBasisOrder, { isHidden: true }),
+        orbitals: PD.Value<{ energy: number, alpha: number[] }[]>([], { isHidden: true })
+    },
+})({
+    apply({ params }) {
+        return new BasisAndOrbitals({ basis: params.basis, order: params.order, orbitals: params.orbitals }, { label: params.label });
+    }
+});
+
+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}`, {
+        defaultValue: [
+            { atomCount: 55, spacing: 0.5 },
+            { atomCount: 40, spacing: 0.45 },
+            { atomCount: 25, spacing: 0.4 },
+            { atomCount: 0, spacing: 0.35 },
+        ]
+    })
+};
+
+export const CreateOrbitalVolume = PluginStateTransform.BuiltIn({
+    name: 'create-orbital-volume',
+    display: 'Orbital Volume',
+    from: BasisAndOrbitals,
+    to: PluginStateObject.Volume.Data,
+    params: (a) => {
+        if (!a) {
+            return { index: PD.Numeric(0), ...CreateOrbitalVolumeParamBase };
+        }
+
+        return {
+            index: PD.Select(0, a.data.orbitals.map((o, i) => [i, `[${i + 1}] ${o.energy.toFixed(4)}`])),
+            ...CreateOrbitalVolumeParamBase
+        };
+    }
+})({
+    apply({ a, params }, plugin: PluginContext) {
+        return Task.create('Orbital Volume', async ctx => {
+            const data = await createSphericalCollocationGrid({
+                basis: a.data.basis,
+                cutoffThreshold: params.cutoffThreshold,
+                alphaOrbitals: a.data.orbitals[params.index].alpha,
+                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);
+            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 CreateOrbitalRepresentation3D = PluginStateTransform.BuiltIn({
+    name: 'create-orbital-representation-3d',
+    display: 'Orbital Representation 3D',
+    from: PluginStateObject.Volume.Data,
+    to: PluginStateObject.Volume.Representation3D,
+    params: {
+        directVolume: PD.Boolean(false),
+        relativeIsovalue: PD.Numeric(1, { min: 0.01, max: 5, step: 0.01 }),
+        kind: PD.Select<'positive' | 'negative'>('positive', [['positive', 'Positive'], ['negative', 'Negative']]),
+        color: PD.Color(ColorNames.blue),
+        alpha: PD.Numeric(1, { min: 0, max: 1, step: 0.01 })
+    }
+})({
+    canAutoUpdate() {
+        return true;
+    },
+    apply({ a, params: srcParams }, plugin: PluginContext) {
+        return Task.create('Orbitals Representation 3D', async ctx => {
+            const params = volumeParams(plugin, a, srcParams);
+
+            const propertyCtx = { runtime: ctx, assetManager: plugin.managers.asset };
+            const provider = plugin.representation.volume.registry.get(params.type.name);
+            if (provider.ensureCustomProperties) await provider.ensureCustomProperties.attach(propertyCtx, a.data);
+            const props = params.type.params || {};
+            const repr = provider.factory({ webgl: plugin.canvas3d?.webgl, ...plugin.representation.volume.themes }, provider.getParams);
+            repr.setTheme(Theme.create(plugin.representation.volume.themes, { volume: a.data }, params));
+            await repr.createOrUpdate(props, a.data).runInContext(ctx);
+            return new PluginStateObject.Volume.Representation3D({ repr, source: a }, { label: provider.label, description: VolumeRepresentation3DHelpers.getDescription(props) });
+        });
+    },
+    update({ a, b, newParams: srcParams }, plugin: PluginContext) {
+        return Task.create('Orbitals Representation 3D', async ctx => {
+            const newParams = volumeParams(plugin, a, srcParams);
+
+            const props = { ...b.data.repr.props, ...newParams.type.params };
+            b.data.repr.setTheme(Theme.create(plugin.representation.volume.themes, { volume: a.data }, newParams));
+            await b.data.repr.createOrUpdate(props, a.data).runInContext(ctx);
+            b.description = VolumeRepresentation3DHelpers.getDescription(props);
+            return StateTransformer.UpdateResult.Updated;
+        });
+    }
+});
+
+function volumeParams(plugin: PluginContext, volume: PluginStateObject.Volume.Data, params: StateTransformer.Params<typeof CreateOrbitalRepresentation3D>) {
+    if (volume.data.sourceData.kind !== 'alpha-orbitals') throw new Error('Invalid data source kind.');
+
+    const { isovalues } = volume.data.sourceData.data as CubeGrid;
+    if (!isovalues) throw new Error('Isovalues are not computed.');
+
+    const value = isovalues[params.kind];
+
+    return createVolumeRepresentationParams(plugin, volume.data, params.directVolume ? {
+        type: 'direct-volume',
+        typeParams: {
+            alpha: params.alpha,
+            renderMode: {
+                name: 'isosurface',
+                params: { isoValue: { kind: 'absolute', absoluteValue: (value ?? 1000) * params.relativeIsovalue }, singleLayer: false }
+            }
+        },
+        color: 'uniform',
+        colorParams: { value: params.color }
+    } : {
+        type: 'isosurface',
+        typeParams: { isoValue: { kind: 'absolute', absoluteValue: (value ?? 1000) * params.relativeIsovalue }, alpha: params.alpha },
+        color: 'uniform',
+        colorParams: { value: params.color }
+    });
+}