use Representation3DState for overpaint transform

src/mol-plugin/behavior/dynamic/representation.ts

@@ -7,21 +7,11 @@
 
 import { MarkerAction } from 'mol-geo/geometry/marker-data';
 import { EmptyLoci } from 'mol-model/loci';
-import { QueryContext, StructureElement, StructureSelection } from 'mol-model/structure';
+import { StructureElement } from 'mol-model/structure';
 import { PluginContext } from 'mol-plugin/context';
-import { PluginStateObject } from 'mol-plugin/state/objects';
 import { Representation } from 'mol-repr/representation';
-import Expression from 'mol-script/language/expression';
-import { parseMolScript } from 'mol-script/language/parser';
-import { compile } from 'mol-script/runtime/query/compiler';
-import { transpileMolScript } from 'mol-script/script/mol-script/symbols';
-import { StateObjectTracker, StateSelection } from 'mol-state';
 import { labelFirst } from 'mol-theme/label';
-import { Overpaint } from 'mol-theme/overpaint';
-import { Color } from 'mol-util/color';
-import { ColorNames } from 'mol-util/color/tables';
 import { ButtonsType } from 'mol-util/input/input-observer';
-import { ParamDefinition as PD } from 'mol-util/param-definition';
 import { PluginBehavior } from '../behavior';
 
 export const HighlightLoci = PluginBehavior.create({
@@ -124,105 +114,3 @@ export const DefaultLociLabelProvider = PluginBehavior.create({
     },
     display: { name: 'Provide Default Loci Label' }
 });
-
-type ColorMappings = { query: Expression, color: Color }[]
-namespace ColorMappings {
-    export function areEqual(colorMappingsA: ColorMappings, colorMappingsB: ColorMappings) {
-        return false
-    }
-}
-
-export namespace ColorRepresentation3D {
-    export const Params = {
-        layers: PD.ObjectList({
-            query: PD.ScriptExpression({ language: 'mol-script', expression: '(sel.atom.atom-groups :residue-test (= atom.resname LYS))' }),
-            color: PD.Color(ColorNames.blueviolet)
-        }, e => `${Color.toRgbString(e.color)}`, {
-            defaultValue: [
-                {
-                    query: {
-                        language: 'mol-script',
-                        expression: '(sel.atom.atom-groups :residue-test (= atom.resname LYS))'
-                    },
-                    color: ColorNames.blueviolet
-                },
-                {
-                    query: {
-                        language: 'mol-script',
-                        expression: '(sel.atom.atom-groups :residue-test (= atom.resname ALA))'
-                    },
-                    color: ColorNames.chartreuse
-                }
-            ]
-        }),
-        alpha: PD.Numeric(1, { min: 0, max: 1, step: 0.01 }, { label: 'Opacity' }),
-    }
-    export type Params = PD.Values<typeof Params>
-
-    export class Behavior implements PluginBehavior<Params> {
-        private currentColorMappings: ColorMappings = [];
-        private repr: StateObjectTracker<PluginStateObject.Molecule.Structure.Representation3D>;
-        private structure: StateObjectTracker<PluginStateObject.Molecule.Structure>;
-
-        private updateData() {
-            const reprUpdated = this.repr.update();
-            const strucUpdated = this.structure.update();
-            return reprUpdated || strucUpdated;
-        }
-
-        register(ref: string): void {
-            this.repr.setQuery(StateSelection.Generators.byRef(ref).ancestorOfType([PluginStateObject.Molecule.Structure.Representation3D]));
-            this.structure.setQuery(StateSelection.Generators.byRef(ref).ancestorOfType([PluginStateObject.Molecule.Structure]));
-            this.update(this.params);
-        }
-
-        update(params: Params): boolean {
-            const { layers, alpha } = params
-            const colorMappings: ColorMappings = []
-            for (let i = 0, il = params.layers.length; i < il; ++i) {
-                const { query, color } = layers[i]
-                const parsed = parseMolScript(query.expression);
-                if (parsed.length === 0) throw new Error('No query');
-                colorMappings.push({ query: transpileMolScript(parsed[0]), color })
-            }
-            return this.applyMappings(colorMappings, alpha)
-        }
-
-        private applyMappings(colorMappings: ColorMappings, alpha: number): boolean {
-            if (!this.updateData() && ColorMappings.areEqual(colorMappings, this.currentColorMappings)) return false;
-            this.currentColorMappings = colorMappings;
-            if (!this.repr.data || !this.structure.data) return true;
-
-            const layers: Overpaint.Layer[] = []
-            for (let i = 0, il = this.currentColorMappings.length; i < il; ++i) {
-                const { query, color } = this.currentColorMappings[i]
-                const compiled = compile<StructureSelection>(query);
-                const result = compiled(new QueryContext(this.structure.data));
-                const loci = StructureSelection.toLoci2(result)
-                layers.push({ loci, color })
-            }
-            return this.applyLayers({ alpha, layers })
-        }
-
-        private applyLayers(overpaint: Overpaint): boolean {
-            if (!this.repr.data) return true;
-            this.repr.data.repr.setState({ overpaint })
-            this.ctx.canvas3d.add(this.repr.data.repr);
-            this.ctx.canvas3d.requestDraw(true);
-            return true;
-        }
-
-        unregister(): void {
-            this.applyLayers(Overpaint.Empty) // clear
-            this.repr.cell = void 0;
-            this.structure.cell = void 0;
-        }
-
-        constructor(private ctx: PluginContext, private params: Params) {
-            this.repr = new StateObjectTracker(ctx.state.dataState);
-            this.structure = new StateObjectTracker(ctx.state.dataState);
-        }
-    }
-
-    export class Obj extends PluginStateObject.CreateBehavior<Behavior>({ name: 'Color Representation3D Behavior' }) { }
-}

src/mol-plugin/state/animation/helpers.ts

@@ -7,8 +7,13 @@
 
 import { SymmetryOperator } from 'mol-math/geometry';
 import { Mat4, Vec3 } from 'mol-math/linear-algebra';
-import { Structure } from 'mol-model/structure';
+import { Structure, StructureSelection, QueryContext } from 'mol-model/structure';
 import { StructureUnitTransforms } from 'mol-model/structure/structure/util/unit-transforms';
+import { Color } from 'mol-util/color';
+import { Overpaint } from 'mol-theme/overpaint';
+import { parseMolScript } from 'mol-script/language/parser';
+import { transpileMolScript } from 'mol-script/script/mol-script/symbols';
+import { compile } from 'mol-script/runtime/query/compiler';
 
 const _unwindMatrix = Mat4.zero();
 export function unwindStructureAssembly(structure: Structure, unitTransforms: StructureUnitTransforms, t: number) {
@@ -33,4 +38,22 @@ export function explodeStructure(structure: Structure, unitTransforms: Structure
 
         unitTransforms.setTransform(_transMat, u);
     }
+}
+
+type ScriptLayers = { script: { language: string, expression: string }, color: Color }[]
+export function getStructureOverpaint(structure: Structure, scriptLayers: ScriptLayers, alpha: number): Overpaint {
+    const layers: Overpaint.Layer[] = []
+    for (let i = 0, il = scriptLayers.length; i < il; ++i) {
+        const { script, color } = scriptLayers[i]
+        const parsed = parseMolScript(script.expression)
+        if (parsed.length === 0) throw new Error('No query')
+        const query = transpileMolScript(parsed[0])
+
+        const compiled = compile<StructureSelection>(query)
+        const result = compiled(new QueryContext(structure))
+        const loci = StructureSelection.toLoci2(result)
+
+        layers.push({ loci, color })
+    }
+    return { layers, alpha }
 }

src/mol-plugin/state/transforms/representation.ts

@@ -7,7 +7,6 @@
 
 import { Structure } from 'mol-model/structure';
 import { VolumeData, VolumeIsoValue } from 'mol-model/volume';
-import { ColorRepresentation3D } from 'mol-plugin/behavior/dynamic/representation';
 import { PluginContext } from 'mol-plugin/context';
 import { RepresentationProvider } from 'mol-repr/representation';
 import { BuiltInStructureRepresentationsName } from 'mol-repr/structure/registry';
@@ -26,14 +25,16 @@ import { ColorNames } from 'mol-util/color/tables';
 import { getLabelRepresentation } from 'mol-plugin/util/structure-labels';
 import { ShapeRepresentation } from 'mol-repr/shape/representation';
 import { StructureUnitTransforms } from 'mol-model/structure/structure/util/unit-transforms';
-import { unwindStructureAssembly, explodeStructure } from '../animation/helpers';
+import { unwindStructureAssembly, explodeStructure, getStructureOverpaint } from '../animation/helpers';
+import { Color } from 'mol-util/color';
+import { Overpaint } from 'mol-theme/overpaint';
 
 export { StructureRepresentation3D }
 export { StructureRepresentation3DHelpers }
 export { StructureLabels3D}
 export { ExplodeStructureRepresentation3D }
 export { UnwindStructureAssemblyRepresentation3D }
-export { ColorStructureRepresentation3D }
+export { OverpaintStructureRepresentation3D as ColorStructureRepresentation3D }
 export { VolumeRepresentation3D }
 
 namespace StructureRepresentation3DHelpers {
@@ -278,21 +279,20 @@ const ExplodeStructureRepresentation3D = PluginStateTransform.BuiltIn({
     canAutoUpdate() {
         return true;
     },
-    apply({ a, params, spine }) {
-        const rootStructure = spine.getRootOfType(SO.Molecule.Structure)!.data;
-        const unitTransforms = new StructureUnitTransforms(rootStructure);
-        explodeStructure(rootStructure, unitTransforms, params.t);
+    apply({ a, params }) {
+        const structure = a.data.source.data;
+        const unitTransforms = new StructureUnitTransforms(structure);
+        explodeStructure(structure, unitTransforms, params.t);
         return new SO.Molecule.Structure.Representation3DState({
             state: { unitTransforms },
-            initialState: { unitTransforms: new StructureUnitTransforms(rootStructure) },
-            info: rootStructure,
+            initialState: { unitTransforms: new StructureUnitTransforms(structure) },
+            info: structure,
             source: a
         }, { label: `Explode T = ${params.t.toFixed(2)}` });
     },
-    update({ a, b, newParams, oldParams, spine }) {
-        const rootStructure = spine.getRootOfType(SO.Molecule.Structure)!.data;
+    update({ a, b, newParams, oldParams }) {
         const structure = b.data.info as Structure;
-        if (rootStructure !== structure) return StateTransformer.UpdateResult.Recreate;
+        if (a.data.source.data !== structure) return StateTransformer.UpdateResult.Recreate;
         if (oldParams.t === newParams.t) return StateTransformer.UpdateResult.Unchanged;
         const unitTransforms = b.data.state.unitTransforms!;
         explodeStructure(structure, unitTransforms, newParams.t);
@@ -302,26 +302,62 @@ const ExplodeStructureRepresentation3D = PluginStateTransform.BuiltIn({
     }
 });
 
-type ColorStructureRepresentation3D = typeof ColorStructureRepresentation3D
-const ColorStructureRepresentation3D = PluginStateTransform.BuiltIn({
-    name: 'color-structure-representation-3d',
-    display: 'Color 3D Representation',
+type OverpaintStructureRepresentation3D = typeof OverpaintStructureRepresentation3D
+const OverpaintStructureRepresentation3D = PluginStateTransform.BuiltIn({
+    name: 'overpaint-structure-representation-3d',
+    display: 'Overpaint 3D Representation',
     from: SO.Molecule.Structure.Representation3D,
-    to: ColorRepresentation3D.Obj,
-    params: ColorRepresentation3D.Params
+    to: SO.Molecule.Structure.Representation3DState,
+    params: {
+        layers: PD.ObjectList({
+            script: PD.ScriptExpression({ language: 'mol-script', expression: '(sel.atom.atom-groups :residue-test (= atom.resname LYS))' }),
+            color: PD.Color(ColorNames.blueviolet)
+        }, e => `${Color.toRgbString(e.color)}`, {
+            defaultValue: [
+                {
+                    script: {
+                        language: 'mol-script',
+                        expression: '(sel.atom.atom-groups :residue-test (= atom.resname LYS))'
+                    },
+                    color: ColorNames.blueviolet
+                },
+                {
+                    script: {
+                        language: 'mol-script',
+                        expression: '(sel.atom.atom-groups :residue-test (= atom.resname ALA))'
+                    },
+                    color: ColorNames.chartreuse
+                }
+            ]
+        }),
+        alpha: PD.Numeric(1, { min: 0, max: 1, step: 0.01 }, { label: 'Opacity' }),
+    }
 })({
     canAutoUpdate() {
         return true;
     },
-    apply({ params }, plugin: PluginContext) {
-        return new ColorRepresentation3D.Obj(new ColorRepresentation3D.Behavior(plugin, params), { label: `Coloring` });
+    apply({ a, params }) {
+        const structure = a.data.source.data
+        const overpaint = getStructureOverpaint(structure, params.layers, params.alpha)
+
+        return new SO.Molecule.Structure.Representation3DState({
+            state: { overpaint },
+            initialState: { overpaint: Overpaint.Empty },
+            info: structure,
+            source: a
+        }, { label: `Overpaint (${overpaint.layers.length} Layers)` })
     },
-    update({ b, newParams }) {
-        return Task.create('Update Coloring', async () => {
-            const updated = await b.data.update(newParams);
-            b.label = `Coloring`;
-            return updated ? StateTransformer.UpdateResult.Updated : StateTransformer.UpdateResult.Unchanged;
-        });
+    update({ a, b, newParams, oldParams }) {
+        const structure = b.data.info as Structure
+        if (a.data.source.data !== structure) return StateTransformer.UpdateResult.Recreate
+        const oldOverpaint = b.data.state.overpaint!
+        const newOverpaint = getStructureOverpaint(structure, newParams.layers, newParams.alpha)
+        if (oldParams.alpha === newParams.alpha && Overpaint.areEqual(oldOverpaint, newOverpaint)) return StateTransformer.UpdateResult.Unchanged
+
+        b.data.state.overpaint = newOverpaint
+        b.data.source = a
+        b.label = `Overpaint (${newOverpaint.layers.length} Layers)`
+        return StateTransformer.UpdateResult.Updated
     }
 });