overpaint improvements

src/mol-plugin/index.ts

@@ -51,7 +51,7 @@ export const DefaultPluginSpec: PluginSpec = {
         PluginSpec.Action(StateTransforms.Representation.StructureLabels3D),
         PluginSpec.Action(StateTransforms.Representation.ExplodeStructureRepresentation3D),
         PluginSpec.Action(StateTransforms.Representation.UnwindStructureAssemblyRepresentation3D),
-        PluginSpec.Action(StateTransforms.Representation.OverpaintStructureRepresentation3D),
+        PluginSpec.Action(StateTransforms.Representation.OverpaintStructureRepresentation3DFromScript),
         PluginSpec.Action(StateTransforms.Representation.TransparencyStructureRepresentation3D),
 
         PluginSpec.Action(StateTransforms.Volume.VolumeFromCcp4),

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

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Structure, StructureSelection, QueryContext } from '../../../mol-model/structure';
+import { Structure, StructureSelection, QueryContext, StructureElement } from '../../../mol-model/structure';
 import { Color } from '../../../mol-util/color';
 import { Overpaint } from '../../../mol-theme/overpaint';
 import { parseMolScript } from '../../../mol-script/language/parser';
@@ -25,7 +25,7 @@ function scriptToLoci(structure: Structure, script: Script) {
     return StructureSelection.toLoci2(result)
 }
 
-export function getStructureOverpaint(structure: Structure, scriptLayers: { script: Script, color: Color, clear: boolean }[], alpha: number): Overpaint {
+export function getStructureOverpaintFromScript(structure: Structure, scriptLayers: { script: Script, color: Color, clear: boolean }[], alpha: number): Overpaint {
     const layers: Overpaint.Layer[] = []
     for (let i = 0, il = scriptLayers.length; i < il; ++i) {
         const { script, color, clear } = scriptLayers[i]
@@ -34,6 +34,16 @@ export function getStructureOverpaint(structure: Structure, scriptLayers: { scri
     return { layers, alpha }
 }
 
+export function getStructureOverpaintFromBundle(structure: Structure, bundleLayers: { bundle: StructureElement.Bundle, color: Color, clear: boolean }[], alpha: number): Overpaint {
+    const layers: Overpaint.Layer[] = []
+    for (let i = 0, il = bundleLayers.length; i < il; ++i) {
+        const { bundle, color, clear } = bundleLayers[i]
+        const loci = StructureElement.Bundle.toLoci(bundle, structure.root)
+        layers.push({ loci, color, clear })
+    }
+    return { layers, alpha }
+}
+
 export function getStructureTransparency(structure: Structure, script: Script, value: number, variant: Transparency.Variant): Transparency {
     return { loci: scriptToLoci(structure, script), value, variant }
 }

src/mol-plugin/util/structure-overpaint-helper.ts

@@ -8,22 +8,12 @@ import { PluginStateObject } from '../../mol-plugin/state/objects';
 import { StateTransforms } from '../../mol-plugin/state/transforms';
 import { StateSelection, StateObjectCell, StateTransform, StateBuilder } from '../../mol-state';
 import { Structure, StructureElement } from '../../mol-model/structure';
-import { isEmptyLoci, EmptyLoci } from '../../mol-model/loci';
 import { PluginContext } from '../context';
 import { Color } from '../../mol-util/color';
-import { MolScriptBuilder } from '../../mol-script/language/builder';
-import { formatMolScript } from '../../mol-script/language/expression-formatter';
 
-type OverpaintEachReprCallback = (update: StateBuilder.Root, repr: StateObjectCell<PluginStateObject.Molecule.Structure.Representation3D, StateTransform<typeof StateTransforms.Representation.StructureRepresentation3D>>, rootStructure: Structure, overpaint?: StateObjectCell<any, StateTransform<typeof StateTransforms.Representation.OverpaintStructureRepresentation3D>>) => void
+type OverpaintEachReprCallback = (update: StateBuilder.Root, repr: StateObjectCell<PluginStateObject.Molecule.Structure.Representation3D, StateTransform<typeof StateTransforms.Representation.StructureRepresentation3D>>, overpaint?: StateObjectCell<any, StateTransform<typeof StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle>>) => void
 const OverpaintManagerTag = 'overpaint-controls'
 
-export function getExpression(loci: StructureElement.Loci | EmptyLoci) {
-    const scriptExpression = isEmptyLoci(loci)
-        ? MolScriptBuilder.struct.generator.empty()
-        : StructureElement.Loci.toScriptExpression(loci)
-    return formatMolScript(scriptExpression)
-}
-
 export class StructureOverpaintHelper {
     private async eachRepr(callback: OverpaintEachReprCallback) {
         const state = this.plugin.state.dataState;
@@ -31,29 +21,29 @@ export class StructureOverpaintHelper {
 
         const update = state.build();
         for (const r of reprs) {
-            const overpaint = state.select(StateSelection.Generators.ofTransformer(StateTransforms.Representation.OverpaintStructureRepresentation3D, r.transform.ref).withTag(OverpaintManagerTag));
-
-            const structure = r.obj!.data.source.data
-
-            callback(update, r, structure.root, overpaint[0])
+            const overpaint = state.select(StateSelection.Generators.ofTransformer(StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle, r.transform.ref).withTag(OverpaintManagerTag))
+            callback(update, r, overpaint[0])
         }
 
         await this.plugin.runTask(state.updateTree(update, { doNotUpdateCurrent: true }));
     }
 
     async set(color: Color | -1, lociGetter: (structure: Structure) => StructureElement.Loci, types?: string[]) {
-        await this.eachRepr((update, repr, rootStructure, overpaint) => {
+        await this.eachRepr((update, repr, overpaint) => {
+            console.log(types, repr.params!.values.type.name)
             if (types && !types.includes(repr.params!.values.type.name)) return
 
-            // TODO cleanup when loci is full structure or empty
-            // TODO add & use QueryOverpaintStructureRepresentation3D
+            // TODO merge overpaint layers, delete shadowed ones
+            // TODO filter overpaint layers for given structure
 
-            const loci = lociGetter(rootStructure)
-            if (loci.elements.length === 0) return
-            const expression = getExpression(loci)
+            const structure = repr.obj!.data.source.data
+            // always use the root structure to get the loci so the overpaint
+            // stays applicable as long as the root structure does not change
+            const loci = lociGetter(structure.root)
+            if (StructureElement.Loci.isEmpty(loci)) return
 
             const layer = {
-                script: { language: 'mol-script', expression },
+                bundle: StructureElement.Bundle.fromLoci(loci),
                 color: color === -1 ? Color(0) : color,
                 clear: color === -1
             }
@@ -62,7 +52,7 @@ export class StructureOverpaintHelper {
                 update.to(overpaint).update({ layers: [ ...overpaint.params!.values.layers, layer ], alpha: 1 })
             } else {
                 update.to(repr.transform.ref)
-                    .apply(StateTransforms.Representation.OverpaintStructureRepresentation3D, { layers: [ layer ], alpha: 1 }, { tags: OverpaintManagerTag });
+                    .apply(StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle, { layers: [ layer ], alpha: 1 }, { tags: OverpaintManagerTag });
             }
         })
     }

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

@@ -16,6 +16,7 @@ import { Task } from '../../mol-task';
 import { PickingId } from '../../mol-geo/geometry/picking';
 import { EmptyLoci, Loci } from '../../mol-model/loci';
 import { MarkerAction } from '../../mol-util/marker-action';
+import { Overpaint } from '../../mol-theme/overpaint';
 
 export function ComplexRepresentation<P extends StructureParams>(label: string, ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, P>, visualCtor: (materialId: number) => ComplexVisual<P>): StructureRepresentation<P> {
     let version = 0
@@ -73,7 +74,11 @@ export function ComplexRepresentation<P extends StructureParams>(label: string,
         }
         if (state.alphaFactor !== undefined && visual) visual.setAlphaFactor(state.alphaFactor)
         if (state.pickable !== undefined && visual) visual.setPickable(state.pickable)
-        if (state.overpaint !== undefined && visual) visual.setOverpaint(state.overpaint)
+        if (state.overpaint !== undefined && visual) {
+            // Remap loci from equivalent structure to the current structure
+            const remappedOverpaint = Overpaint.remap(state.overpaint, _structure)
+            visual.setOverpaint(remappedOverpaint)
+        }
         if (state.transparency !== undefined && visual) visual.setTransparency(state.transparency)
         if (state.transform !== undefined && visual) visual.setTransform(state.transform)
         if (state.unitTransforms !== undefined && visual) {

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

@@ -19,6 +19,7 @@ import { Task } from '../../mol-task';
 import { PickingId } from '../../mol-geo/geometry/picking';
 import { Loci, EmptyLoci, isEmptyLoci } from '../../mol-model/loci';
 import { MarkerAction } from '../../mol-util/marker-action';
+import { Overpaint } from '../../mol-theme/overpaint';
 
 export const UnitsParams = {
     ...StructureParams,
@@ -181,7 +182,11 @@ export function UnitsRepresentation<P extends UnitsParams>(label: string, ctx: R
         if (visible !== undefined) visuals.forEach(({ visual }) => visual.setVisibility(visible))
         if (alphaFactor !== undefined) visuals.forEach(({ visual }) => visual.setAlphaFactor(alphaFactor))
         if (pickable !== undefined) visuals.forEach(({ visual }) => visual.setPickable(pickable))
-        if (overpaint !== undefined) visuals.forEach(({ visual }) => visual.setOverpaint(overpaint))
+        if (overpaint !== undefined) {
+            // Remap loci from equivalent structure to the current `_structure`
+            const remappedOverpaint = Overpaint.remap(overpaint, _structure)
+            visuals.forEach(({ visual }) => visual.setOverpaint(remappedOverpaint))
+        }
         if (transparency !== undefined) visuals.forEach(({ visual }) => visual.setTransparency(transparency))
         if (transform !== undefined) visuals.forEach(({ visual }) => visual.setTransform(transform))
         if (unitTransforms !== undefined) {

src/mol-theme/overpaint.ts

@@ -6,6 +6,7 @@
 
 import { Loci } from '../mol-model/loci';
 import { Color } from '../mol-util/color';
+import { Structure } from '../mol-model/structure';
 
 export { Overpaint }
 
@@ -26,4 +27,13 @@ namespace Overpaint {
         }
         return true
     }
+
+    export function remap(overpaint: Overpaint, structure: Structure) {
+        const layers: Overpaint.Layer[] = []
+        for (const layer of overpaint.layers) {
+            const { loci, color, clear } = layer
+            layers.push({ loci: Loci.remap(loci, structure), color, clear })
+        }
+        return { layers, alpha: overpaint.alpha }
+    }
 }