wip, overpaint plugin

src/mol-geo/geometry/overpaint-data.ts

@@ -15,10 +15,10 @@ export type OverpaintData = {
     dOverpaint: ValueCell<boolean>,
 }
 
-export function applyOverpaintColor(array: Uint8Array, start: number, end: number, color: Color) {
+export function applyOverpaintColor(array: Uint8Array, start: number, end: number, color: Color, alpha: number) {
     for (let i = start; i < end; ++i) {
         Color.toArray(color, array, i * 4)
-        array[i * 4 + 3] = 255
+        array[i * 4 + 3] = alpha * 255
     }
     return true
 }

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

@@ -25,7 +25,6 @@ type StructureAnimationProps = PD.Values<typeof StructureAnimationParams>
 /**
  * TODO
  * - animation class is just for testing purposes, needs better API
- * - allow per-unit transform `unitTransform: { [unitId: number]: Mat4 }`
  */
 export const StructureAnimation = PluginBehavior.create<StructureAnimationProps>({
     name: 'structure-animation',

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

@@ -1,12 +1,13 @@
 /**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
 import { MarkerAction } from 'mol-geo/geometry/marker-data';
 import { Mat4, Vec3 } from 'mol-math/linear-algebra';
-import { EmptyLoci } from 'mol-model/loci';
+import { EmptyLoci, EveryLoci } from 'mol-model/loci';
 import { StructureUnitTransforms } from 'mol-model/structure/structure/util/unit-transforms';
 import { PluginContext } from 'mol-plugin/context';
 import { PluginStateObject } from 'mol-plugin/state/objects';
@@ -16,7 +17,15 @@ import { ParamDefinition as PD } from 'mol-util/param-definition';
 import { PluginBehavior } from '../behavior';
 import { Representation } from 'mol-repr/representation';
 import { ButtonsType } from 'mol-util/input/input-observer';
-import { StructureElement } from 'mol-model/structure';
+import { StructureElement, StructureSelection, QueryContext } from 'mol-model/structure';
+import { ColorNames } from 'mol-util/color/tables';
+// import { MolScriptBuilder as MS } from 'mol-script/language/builder';
+import Expression from 'mol-script/language/expression';
+import { Color } from 'mol-util/color';
+import { compile } from 'mol-script/runtime/query/compiler';
+import { Overpaint } from 'mol-theme/overpaint';
+import { parseMolScript } from 'mol-script/language/parser';
+import { transpileMolScript } from 'mol-script/script/mol-script/symbols';
 
 export const HighlightLoci = PluginBehavior.create({
     name: 'representation-highlight-loci',
@@ -119,7 +128,6 @@ export const DefaultLociLabelProvider = PluginBehavior.create({
     display: { name: 'Provide Default Loci Label' }
 });
 
-
 export namespace ExplodeRepresentation3D {
     export const Params = {
         t: PD.Numeric(0, { min: 0, max: 1, step: 0.01 })
@@ -196,4 +204,87 @@ export namespace ExplodeRepresentation3D {
     }
 
     export class Obj extends PluginStateObject.CreateBehavior<Behavior>({ name: 'Explode Representation3D Behavior' }) { }
+}
+
+type ColorMappings = { query: Expression, color: Color }[]
+namespace ColorMappings {
+    export function areEqual(colorMappingsA: ColorMappings, colorMappingsB: ColorMappings) {
+        return false
+    }
+}
+
+export namespace ColorRepresentation3D {
+    export const Params = {
+        query: PD.ScriptExpression({ language: 'mol-script', expression: '(sel.atom.atom-groups :residue-test (= atom.resname LYS))' }),
+        color: PD.Color(ColorNames.blueviolet)
+        // colorMappings: PD.Value<ColorMappings>([{ query: MS.struct.generator.atomGroups({
+        //     'residue-test': MS.core.rel.eq([MS.ammp('auth_comp_id'), 'ALA'])
+        // }), color: ColorNames.greenyellow }], { isHidden: true }),
+    }
+    export type Params = PD.Values<typeof Params>
+
+    export class Behavior implements PluginBehavior<Params> {
+        private currentColorMappings: ColorMappings = [];
+        private repr: StateObjectTracker<PluginStateObject.Molecule.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.Representation3D]));
+            this.structure.setQuery(StateSelection.Generators.byRef(ref).ancestorOfType([PluginStateObject.Molecule.Structure]));
+            this.update(this.params);
+        }
+
+        update(params: Params): boolean {
+            const parsed = parseMolScript(params.query.expression);
+            if (parsed.length === 0) throw new Error('No query');
+            const query = transpileMolScript(parsed[0]);
+
+            return this.apply([{ query, color: params.color }])
+        }
+
+        private apply(colorMappings: ColorMappings): 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.Layers = [
+                // unsets the overpaint
+                // TODO do smarter by looking at the current color mappings
+                { loci: EveryLoci, color: ColorNames.black, alpha: 0 }
+            ]
+            // console.log('currentColorMappings', this.currentColorMappings)
+            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, alpha: 1 })
+            }
+            this.repr.data.setOverpaint(layers)
+
+            this.ctx.canvas3d.add(this.repr.data);
+            this.ctx.canvas3d.requestDraw(true);
+
+            return true;
+        }
+
+        unregister(): void {
+            this.apply([])
+            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/index.ts

@@ -45,6 +45,7 @@ export const DefaultPluginSpec: PluginSpec = {
         PluginSpec.Action(StateTransforms.Representation.StructureRepresentation3D),
         PluginSpec.Action(StateTransforms.Representation.StructureLabels3D),
         PluginSpec.Action(StateTransforms.Representation.ExplodeStructureRepresentation3D),
+        PluginSpec.Action(StateTransforms.Representation.ColorStructureRepresentation3D),
         PluginSpec.Action(StateTransforms.Representation.VolumeRepresentation3D),
 
         PluginSpec.Action(StateActions.Structure.StructureFromSelection),

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

@@ -7,7 +7,7 @@
 
 import { Structure } from 'mol-model/structure';
 import { VolumeData, VolumeIsoValue } from 'mol-model/volume';
-import { ExplodeRepresentation3D } from 'mol-plugin/behavior/dynamic/representation';
+import { ExplodeRepresentation3D, 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';
@@ -30,6 +30,7 @@ export { StructureRepresentation3D }
 export { StructureRepresentation3DHelpers }
 export { StructureLabels3D}
 export { ExplodeStructureRepresentation3D }
+export { ColorStructureRepresentation3D }
 export { VolumeRepresentation3D }
 
 namespace StructureRepresentation3DHelpers {
@@ -228,7 +229,6 @@ const StructureLabels3D = PluginStateTransform.BuiltIn({
     }
 });
 
-
 type ExplodeStructureRepresentation3D = typeof ExplodeStructureRepresentation3D
 const ExplodeStructureRepresentation3D = PluginStateTransform.BuiltIn({
     name: 'explode-structure-representation-3d',
@@ -252,6 +252,29 @@ const ExplodeStructureRepresentation3D = PluginStateTransform.BuiltIn({
     }
 });
 
+type ColorStructureRepresentation3D = typeof ColorStructureRepresentation3D
+const ColorStructureRepresentation3D = PluginStateTransform.BuiltIn({
+    name: 'color-structure-representation-3d',
+    display: 'Color 3D Representation',
+    from: SO.Molecule.Representation3D,
+    to: ColorRepresentation3D.Obj,
+    params: ColorRepresentation3D.Params
+})({
+    canAutoUpdate() {
+        return true;
+    },
+    apply({ params }, plugin: PluginContext) {
+        return new ColorRepresentation3D.Obj(new ColorRepresentation3D.Behavior(plugin, params), { label: `Coloring` });
+    },
+    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;
+        });
+    }
+});
+
 //
 
 export namespace VolumeRepresentation3DHelpers {

src/mol-repr/shape/representation.ts

@@ -208,11 +208,11 @@ export function ShapeRepresentation<D, G extends Geometry, P extends Geometry.Pa
             const { groupCount, instanceCount } = locationIt
 
             for (let i = 0, il = layers.length; i < il; ++i) {
-                const { loci, color } = layers[i]
+                const { loci, color, alpha } = layers[i]
                 const apply = (interval: Interval) => {
                     const start = Interval.start(interval)
                     const end = Interval.end(interval)
-                    return applyOverpaintColor(tOverpaint.ref.value.array, start, end, color)
+                    return applyOverpaintColor(tOverpaint.ref.value.array, start, end, color, alpha)
                 }
 
                 if (isEveryLoci(loci) || (Shape.isLoci(loci) && loci.shape === _shape)) {

src/mol-repr/structure/complex-visual.ts

@@ -214,11 +214,11 @@ export function ComplexVisual<G extends Geometry, P extends ComplexParams & Geom
             createOverpaint(layers.length ? groupCount * instanceCount : 0, renderObject.values)
 
             for (let i = 0, il = layers.length; i < il; ++i) {
-                const { loci, color } = layers[i]
+                const { loci, color, alpha } = layers[i]
                 const apply = (interval: Interval) => {
                     const start = Interval.start(interval)
                     const end = Interval.end(interval)
-                    return applyOverpaintColor(tOverpaint.ref.value.array, start, end, color)
+                    return applyOverpaintColor(tOverpaint.ref.value.array, start, end, color, alpha)
                 }
 
                 if (isEveryLoci(loci) || (Structure.isLoci(loci) && loci.structure === currentStructure)) {

src/mol-repr/structure/units-visual.ts

@@ -260,11 +260,11 @@ export function UnitsVisual<G extends Geometry, P extends UnitsParams & Geometry
             createOverpaint(layers.length ? groupCount * instanceCount : 0, renderObject.values)
 
             for (let i = 0, il = layers.length; i < il; ++i) {
-                const { loci, color } = layers[i]
+                const { loci, color, alpha } = layers[i]
                 const apply = (interval: Interval) => {
                     const start = Interval.start(interval)
                     const end = Interval.end(interval)
-                    return applyOverpaintColor(tOverpaint.ref.value.array, start, end, color)
+                    return applyOverpaintColor(tOverpaint.ref.value.array, start, end, color, alpha)
                 }
 
                 if (isEveryLoci(loci) || (Structure.isLoci(loci) && loci.structure === currentStructureGroup.structure)) {

src/mol-theme/overpaint.ts

@@ -0,0 +1,29 @@
+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Loci } from 'mol-model/loci';
+import { Color } from 'mol-util/color';
+
+export { Overpaint }
+namespace Overpaint {
+    export interface Layer {
+        loci: Loci
+        color: Color
+        alpha: number
+    }
+    export type Layers = Layer[]
+
+    export function areEqual(layersA: Layers, layersB: Layers) {
+        if (layersA.length === 0 && layersB.length === 0) return true
+        if (layersA.length !== layersB.length) return false
+        for (let i = 0, il = layersA.length; i < il; ++i) {
+            if (layersA[i].alpha !== layersB[i].alpha) return false
+            if (layersA[i].color !== layersB[i].color) return false
+            if (!Loci.areEqual(layersA[i].loci, layersB[i].loci)) return false
+        }
+        return true
+    }
+}