by-volume-value theme

CHANGELOG.md

@@ -7,6 +7,7 @@ Note that since we don't clearly distinguish between a public and private interf
 ## [Unreleased]
 
 - Make `PluginContext.initContainer` checkered canvas background optional
+- `by-volume-value` theme (coloring or arbitrary geometries by user-selected volume)
 
 ## [v3.23.0] - 2022-10-19
 

src/mol-plugin-ui/controls/parameters.tsx

@@ -18,7 +18,7 @@ import { getPrecision } from '../../mol-util/number';
 import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { ParamMapping } from '../../mol-util/param-mapping';
 import { camelCaseToWords } from '../../mol-util/string';
-import { PluginUIComponent } from '../base';
+import { PluginReactContext, PluginUIComponent } from '../base';
 import { PluginUIContext } from '../context';
 import { ActionMenu } from './action-menu';
 import { ColorOptions, ColorValueOption, CombinedColorControl } from './color';
@@ -505,10 +505,12 @@ export class ValueRefControl extends React.PureComponent<ParamProps<PD.ValueRef<
 
     toggle = () => this.setState({ showOptions: !this.state.showOptions });
 
-    items = memoizeLatest((param: PD.ValueRef) => ActionMenu.createItemsFromSelectOptions(param.getOptions()));
+    private get items() {
+        return ActionMenu.createItemsFromSelectOptions(this.props.param.getOptions(this.context));
+    }
 
     renderControl() {
-        const items = this.items(this.props.param);
+        const items = this.items;
         const current = this.props.value.ref ? ActionMenu.findItem(items, this.props.value.ref) : void 0;
         const label = current
             ? current.label
@@ -521,7 +523,7 @@ export class ValueRefControl extends React.PureComponent<ParamProps<PD.ValueRef<
     renderAddOn() {
         if (!this.state.showOptions) return null;
 
-        const items = this.items(this.props.param);
+        const items = this.items;
         const current = ActionMenu.findItem(items, this.props.value.ref);
 
         return <ActionMenu items={items} current={current} onSelect={this.onSelect} />;
@@ -539,6 +541,7 @@ export class ValueRefControl extends React.PureComponent<ParamProps<PD.ValueRef<
         });
     }
 }
+ValueRefControl.contextType = PluginReactContext;
 
 export class IntervalControl extends React.PureComponent<ParamProps<PD.Interval>, { isExpanded: boolean }> {
     state = { isExpanded: false };

src/mol-theme/color.ts

@@ -40,6 +40,7 @@ import { VolumeValueColorThemeProvider } from './color/volume-value';
 import { Vec3, Vec4 } from '../mol-math/linear-algebra';
 import { ModelIndexColorThemeProvider } from './color/model-index';
 import { StructureIndexColorThemeProvider } from './color/structure-index';
+import { ByVolumeValueColorThemeProvider } from './color/by-volume-value';
 
 export type LocationColor = (location: Location, isSecondary: boolean) => Color
 
@@ -152,6 +153,7 @@ namespace ColorTheme {
         'unit-index': UnitIndexColorThemeProvider,
         'uniform': UniformColorThemeProvider,
         'volume-value': VolumeValueColorThemeProvider,
+        'by-volume-value': ByVolumeValueColorThemeProvider,
     };
     type _BuiltIn = typeof BuiltIn
     export type BuiltIn = keyof _BuiltIn

src/mol-theme/color/by-volume-value.ts

@@ -0,0 +1,126 @@
+/**
+ * Copyright (c) 2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Color, ColorScale } from '../../mol-util/color';
+import { Location } from '../../mol-model/location';
+import { ColorTheme } from '../color';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { ThemeDataContext } from '../theme';
+import { Grid, Volume } from '../../mol-model/volume';
+import { type PluginContext } from '../../mol-plugin/context';
+import { isPositionLocation } from '../../mol-geo/util/location-iterator';
+import { Mat4, Vec3 } from '../../mol-math/linear-algebra';
+import { lerp } from '../../mol-math/interpolate';
+
+const Description = `Assigns a color based volume value at a given vertex.`;
+
+export const ByVolumeValueColorThemeParams = {
+    volume: PD.ValueRef<Volume>(
+        (ctx: PluginContext) => {
+            const volumes = ctx.state.data.selectQ(q => q.root.subtree().filter(c => Volume.is(c.obj?.data)));
+            return volumes.map(v => [v.transform.ref, v.obj?.label ?? '<unknown>'] as [string, string]);
+        },
+        (ref, getData) => getData(ref),
+    ),
+    domain: PD.MappedStatic('auto', {
+        custom: PD.Interval([-1, 1]),
+        auto: PD.EmptyGroup()
+    }),
+    list: PD.ColorList('red-white-blue', { presetKind: 'scale' }),
+    defaultColor: PD.Color(Color(0xcccccc))
+};
+export type ByVolumeValueColorThemeParams = typeof ByVolumeValueColorThemeParams
+
+export function ByVolumeValueColorTheme(ctx: ThemeDataContext, props: PD.Values<ByVolumeValueColorThemeParams>): ColorTheme<ByVolumeValueColorThemeParams> {
+    let volume: Volume | undefined;
+    try {
+        volume = props.volume.getValue();
+    } catch {
+        // .getValue() is resolved during state reconciliation => would throw from UI
+    }
+
+    const domain: [number, number] = props.domain.name === 'custom' ? props.domain.params : volume ? [volume.grid.stats.min, volume.grid.stats.max] : [-1, 1];
+    const scale = ColorScale.create({ domain, listOrName: props.list.colors });
+
+    // NOTE: this will currently not work with GPU iso-surfaces since it requires vertex coloring
+    // TODO: create texture to be able to do the sampling on the GPU
+
+    let color;
+    if (volume) {
+        const cartnToGrid = Grid.getGridToCartesianTransform(volume.grid);
+        Mat4.invert(cartnToGrid, cartnToGrid);
+        const gridCoords = Vec3();
+
+        const { dimensions, get } = volume.grid.cells.space;
+        const data = volume.grid.cells.data;
+
+        const [mi, mj, mk] = dimensions;
+
+        color = (location: Location): Color => {
+            if (!isPositionLocation(location)) {
+                return props.defaultColor;
+            }
+
+            Vec3.copy(gridCoords, location.position);
+            Vec3.transformMat4(gridCoords, gridCoords, cartnToGrid);
+
+            const i = Math.floor(gridCoords[0]);
+            const j = Math.floor(gridCoords[1]);
+            const k = Math.floor(gridCoords[2]);
+
+            if (i < 0 || i >= mi || j < 0 || j >= mj || k < 0 || k >= mk) {
+                return props.defaultColor;
+            }
+
+            const u = gridCoords[0] - i;
+            const v = gridCoords[1] - j;
+            const w = gridCoords[2] - k;
+
+            // Tri-linear interpolation for the value
+            const ii = Math.min(i + 1, mi - 1);
+            const jj = Math.min(j + 1, mj - 1);
+            const kk = Math.min(k + 1, mk - 1);
+
+            let a = get(data, i, j, k);
+            let b = get(data, ii, j, k);
+            let c = get(data, i, jj, k);
+            let d = get(data, ii, jj, k);
+            const x = lerp(lerp(a, b, u), lerp(c, d, u), v);
+
+            a = get(data, i, j, kk);
+            b = get(data, ii, j, kk);
+            c = get(data, i, jj, kk);
+            d = get(data, ii, jj, kk);
+            const y = lerp(lerp(a, b, u), lerp(c, d, u), v);
+
+            const value = lerp(x, y, w);
+
+            return scale.color(value);
+        };
+    } else {
+        color = () => props.defaultColor;
+    }
+
+    return {
+        factory: ByVolumeValueColorTheme,
+        granularity: 'vertex',
+        preferSmoothing: true,
+        color,
+        props,
+        description: Description,
+        legend: scale ? scale.legend : undefined
+    };
+}
+
+export const ByVolumeValueColorThemeProvider: ColorTheme.Provider<ByVolumeValueColorThemeParams, 'by-volume-value'> = {
+    name: 'by-volume-value',
+    label: 'By Volume Value',
+    category: ColorTheme.Category.Misc,
+    factory: ByVolumeValueColorTheme,
+    getParams: () => ByVolumeValueColorThemeParams,
+    defaultValues: PD.getDefaultValues(ByVolumeValueColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => true, // TODO
+};

src/mol-util/param-definition.ts

@@ -290,9 +290,9 @@ export namespace ParamDefinition {
     // getValue needs to be assigned by a runtime because it might not be serializable
     export interface ValueRef<T = any> extends Base<{ ref: string, getValue: () => T }> {
         type: 'value-ref',
-        resolveRef: (ref: string) => T,
+        resolveRef: (ref: string, getData: (ref: string) => any) => T,
         // a provider because the list changes over time
-        getOptions: () => Select<string>['options'],
+        getOptions: (ctx: any) => Select<string>['options'],
     }
     export function ValueRef<T>(getOptions: ValueRef['getOptions'], resolveRef: ValueRef<T>['resolveRef'], info?: Info & { defaultRef?: string }) {
         return setInfo<ValueRef<T>>({ type: 'value-ref', defaultValue: { ref: info?.defaultRef ?? '', getValue: unsetGetValue as any }, getOptions, resolveRef }, info);
@@ -365,8 +365,8 @@ export namespace ParamDefinition {
         return d as Values<T>;
     }
 
-    function _resolveRef(resolve: (ref: string) => any, ref: string) {
-        return () => resolve(ref);
+    function _resolveRef(resolve: (ref: string, getData: (ref: string) => any) => any, ref: string, getData: (ref: string) => any) {
+        return () => resolve(ref, getData);
     }
 
     function resolveRefValue(p: Any, value: any, getData: (ref: string) => any) {
@@ -375,11 +375,11 @@ export namespace ParamDefinition {
         if (p.type === 'value-ref') {
             const v = value as ValueRef['defaultValue'];
             if (!v.ref) v.getValue = () => { throw new Error('Unset ref in ValueRef value.'); };
-            else v.getValue = _resolveRef(p.resolveRef, v.ref);
+            else v.getValue = _resolveRef(p.resolveRef, v.ref, getData);
         } else if (p.type === 'data-ref') {
             const v = value as ValueRef['defaultValue'];
             if (!v.ref) v.getValue = () => { throw new Error('Unset ref in ValueRef value.'); };
-            else v.getValue = _resolveRef(getData, v.ref);
+            else v.getValue = _resolveRef(getData, v.ref, getData);
         } else if (p.type === 'group') {
             resolveRefs(p.params, value, getData);
         } else if (p.type === 'mapped') {