element custom properties helper

src/apps/basic-wrapper/coloring.ts

+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { CustomElementProperty } from 'mol-model-props/common/custom-element-property';
+import { Model, ElementIndex } from 'mol-model/structure';
+import { Color } from 'mol-util/color';
+
+export const StripedResidues = CustomElementProperty.create<number>({
+    isStatic: true,
+    name: 'basic-wrapper-residue-striping',
+    display: 'Residue Stripes',
+    getData(model: Model) {
+        const map = new Map<ElementIndex, number>();
+        const residueIndex = model.atomicHierarchy.residueAtomSegments.index;
+        for (let i = 0, _i = model.atomicHierarchy.atoms._rowCount; i < _i; i++) {
+            map.set(i as ElementIndex, residueIndex[i] % 2);
+        }
+        return map;
+    },
+    coloring: {
+        getColor(e) { return e === 0 ? Color(0xff0000) : Color(0x0000ff) },
+        defaultColor: Color(0x777777)
+    },
+    format(e) {
+        return e === 0 ? 'Odd stripe' : 'Even stripe'
+    }
+})
\ No newline at end of file

src/apps/basic-wrapper/index.html

@@ -100,6 +100,10 @@
             addControl('Play Loop', () => BasicMolStarWrapper.animate.modelIndex.loop());
             addControl('Stop', () => BasicMolStarWrapper.animate.modelIndex.stop());
 
+            addHeader('Misc');
+
+            addControl('Apply Stripes', () => BasicMolStarWrapper.coloring.applyStripes());
+
             ////////////////////////////////////////////////////////
 
             function addControl(label, action) {

src/apps/basic-wrapper/index.ts

@@ -11,9 +11,10 @@ import { PluginCommands } from 'mol-plugin/command';
 import { StateTransforms } from 'mol-plugin/state/transforms';
 import { StructureRepresentation3DHelpers } from 'mol-plugin/state/transforms/representation';
 import { Color } from 'mol-util/color';
-import { PluginStateObject as PSO } from 'mol-plugin/state/objects';
+import { PluginStateObject as PSO, PluginStateObject } from 'mol-plugin/state/objects';
 import { AnimateModelIndex } from 'mol-plugin/state/animation/built-in';
 import { StateBuilder } from 'mol-state';
+import { StripedResidues } from './coloring';
 require('mol-plugin/skin/light.scss')
 
 type SupportedFormats = 'cif' | 'pdb'
@@ -30,6 +31,10 @@ class BasicWrapper {
                 showControls: false
             }
         });
+
+        this.plugin.structureRepresentation.themeCtx.colorThemeRegistry.add(StripedResidues.Descriptor.name, StripedResidues.colorTheme!);
+        this.plugin.lociLabels.addProvider(StripedResidues.labelProvider);
+        this.plugin.customModelProperties.register(StripedResidues.propertyProvider);
     }
 
     private download(b: StateBuilder.To<PSO.Root>, url: string) {
@@ -43,6 +48,7 @@ class BasicWrapper {
 
         return parsed
             .apply(StateTransforms.Model.ModelFromTrajectory, { modelIndex: 0 })
+            .apply(StateTransforms.Model.CustomModelProperties, { properties: [StripedResidues.Descriptor.name] }, { ref: 'props', props: { isGhost: false } })
             .apply(StateTransforms.Model.StructureAssemblyFromModel, { id: assemblyId || 'deposited' }, { ref: 'asm' });
     }
 
@@ -110,6 +116,22 @@ class BasicWrapper {
             stop: () => this.plugin.state.animation.stop()
         }
     }
+
+    coloring = {
+        applyStripes: async () => {
+            const state = this.plugin.state.dataState;
+
+            const visuals = state.selectQ(q => q.ofType(PluginStateObject.Molecule.Representation3D).filter(c => c.transform.transformer === StateTransforms.Representation.StructureRepresentation3D));
+            const tree = state.build();
+            const colorTheme = { name: StripedResidues.Descriptor.name, params: this.plugin.structureRepresentation.themeCtx.colorThemeRegistry.get(StripedResidues.Descriptor.name).defaultValues };
+
+            for (const v of visuals) {
+                tree.to(v.transform.ref).update(StateTransforms.Representation.StructureRepresentation3D, old => ({ ...old, colorTheme }));
+            }
+
+            await PluginCommands.State.Update.dispatch(this.plugin, { state, tree });
+        }
+    }
 }
 
 (window as any).BasicMolStarWrapper = new BasicWrapper();
\ No newline at end of file

src/mol-model-props/common/custom-element-property.ts

 /**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  */
@@ -9,7 +9,7 @@ import { StructureElement } from 'mol-model/structure/structure';
 import { Location } from 'mol-model/location';
 import { CustomPropertyRegistry } from './custom-property-registry';
 import { Task } from 'mol-task';
-import { ThemeDataContext } from 'mol-theme/theme';
+import { ThemeDataContext, ThemeProvider } from 'mol-theme/theme';
 import { ColorTheme, LocationColor } from 'mol-theme/color';
 import { Color } from 'mol-util/color';
 import { TableLegend } from 'mol-util/color/tables';
@@ -19,12 +19,12 @@ import { OrderedSet } from 'mol-data/int';
 export { CustomElementProperty };
 
 namespace CustomElementProperty {
-    export interface CreateParams<S, T> {
+    export interface CreateParams<T> {
         isStatic: boolean,
         name: string,
         autoAttach?: boolean,
         display: string,
-        attachableTo: (model: Model) => boolean,
+        attachableTo?: (model: Model) => boolean,
         getData(model: Model): Map<ElementIndex, T> | Promise<Map<ElementIndex, T>>,
         format?(e: T): string,
         coloring?: {
@@ -33,7 +33,7 @@ namespace CustomElementProperty {
         }
     }
 
-    export function create<S, T>(params: CreateParams<S, T>) {
+    export function create<T>(params: CreateParams<T>) {
         const name = params.name;
 
         const Descriptor = ModelPropertyDescriptor({
@@ -51,6 +51,8 @@ namespace CustomElementProperty {
                     model._dynamicPropertyData[name] = data;
                 }
 
+                model.customProperties.add(Descriptor);
+
                 return true;
             })
         }
@@ -58,7 +60,7 @@ namespace CustomElementProperty {
         function getStatic(e: StructureElement) { return e.unit.model._staticPropertyData[name].get(e.element); }
         function getDynamic(e: StructureElement) { return e.unit.model._staticPropertyData[name].get(e.element); }
 
-        const provider: CustomPropertyRegistry.Provider = {
+        const propertyProvider: CustomPropertyRegistry.Provider = {
             option: [name, params.display],
             descriptor: Descriptor,
             defaultSelected: !!params.autoAttach,
@@ -68,12 +70,14 @@ namespace CustomElementProperty {
 
         const get = params.isStatic ? getStatic : getDynamic;
 
+        function has(model: Model) { return model.customProperties.has(Descriptor); }
+
         function Coloring(ctx: ThemeDataContext, props: {}): ColorTheme<{}> {
             let color: LocationColor;
             const getColor = params.coloring!.getColor;
             const defaultColor = params.coloring!.defaultColor;
 
-            if (ctx.structure && !ctx.structure.isEmpty && ctx.structure.models[0].customProperties.has(Descriptor)) {
+            if (ctx.structure && !ctx.structure.isEmpty && has(ctx.structure.models[0])) {
                 color = (location: Location) => {
                     if (StructureElement.isLocation(location)) {
                         const e = get(location);
@@ -95,9 +99,18 @@ namespace CustomElementProperty {
             };
         }
 
+        const colorTheme: ThemeProvider<ColorTheme<{}>, {}> = {
+            label: params.display,
+            factory: Coloring,
+            getParams: () => ({}),
+            defaultValues: {},
+            isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && !ctx.structure.isEmpty && has(ctx.structure.models[0])
+        }
+
         function LabelProvider(loci: Loci): string | undefined {
             if (loci.kind === 'element-loci') {
                 const e = loci.elements[0];
+                if (!has(e.unit.model)) return void 0;
                 return params.format!(get(StructureElement.create(e.unit, e.unit.elements[OrderedSet.getAt(e.indices, 0)])));
             }
             return void 0;
@@ -107,8 +120,8 @@ namespace CustomElementProperty {
             Descriptor,
             attach,
             get,
-            provider,
-            colorTheme: params.coloring ? Coloring : void 0,
+            propertyProvider,
+            colorTheme: params.coloring ? colorTheme : void 0,
             labelProvider: params.format ? LabelProvider : ((loci: Loci) => void 0)
         };
     }

src/mol-state/state/selection.ts

@@ -97,6 +97,14 @@ namespace StateSelection {
             });
         }
 
+        export function ofType(type: StateObject.Ctor) {
+            return build(() => state => {
+                const ctx = { ret: [] as StateObjectCell[], cells: state.cells, type: type.type };
+                StateTree.doPreOrder(state.tree, state.tree.root, ctx, _findOfType);
+                return ctx.ret;
+            });
+        }
+
         function _findRootsOfType(n: StateTransform, _: any, s: { type: StateObject.Type, roots: StateObjectCell[], cells: State.Cells }) {
             const cell = s.cells.get(n.ref);
             if (cell && cell.obj && cell.obj.type === s.type) {
@@ -105,6 +113,14 @@ namespace StateSelection {
             }
             return true;
         }
+
+        function _findOfType(n: StateTransform, _: any, s: { type: StateObject.Type, ret: StateObjectCell[], cells: State.Cells }) {
+            const cell = s.cells.get(n.ref);
+            if (cell && cell.obj && cell.obj.type === s.type) {
+                s.ret.push(cell);
+            }
+            return true;
+        }
     }
 
     registerModifier('flatMap', flatMap);