wip proteopedia wrapper

src/examples/proteopedia-wrapper/coloring.ts

+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+
+import { Unit, StructureProperties, StructureElement, Link } from 'mol-model/structure';
+
+import { Color } from 'mol-util/color';
+import { Location } from 'mol-model/location';
+import { ColorTheme, LocationColor } from 'mol-theme/color';
+import { ParamDefinition as PD } from 'mol-util/param-definition'
+import { ThemeDataContext } from 'mol-theme/theme';
+import { Column } from 'mol-data/db';
+
+const Description = 'Gives every chain a color from a list based on its `asym_id` value.'
+
+export function createProteopediaCustomTheme(colors: number[], defaultColor: number) {
+    const ProteopediaCustomColorThemeParams = {
+        colors: PD.ObjectList({ color: PD.Color(Color(0xffffff)) }, ({ color }) => Color.toHexString(color),
+            { defaultValue: colors.map(c => ({ color: Color(c) })) }),
+        defaultColor: PD.Color(Color(defaultColor))
+    }
+    type ProteopediaCustomColorThemeParams = typeof ProteopediaCustomColorThemeParams
+    function getChainIdColorThemeParams(ctx: ThemeDataContext) {
+        return ProteopediaCustomColorThemeParams // TODO return copy
+    }
+
+    function getAsymId(unit: Unit): StructureElement.Property<string> {
+        switch (unit.kind) {
+            case Unit.Kind.Atomic:
+                return StructureProperties.chain.label_asym_id
+            case Unit.Kind.Spheres:
+            case Unit.Kind.Gaussians:
+                return StructureProperties.coarse.asym_id
+        }
+    }
+
+    function addAsymIds(map: Map<string, number>, data: Column<string>) {
+        let j = map.size
+        for (let o = 0, ol = data.rowCount; o < ol; ++o) {
+            const k = data.value(o)
+            if (!map.has(k)) {
+                map.set(k, j)
+                j += 1
+            }
+        }
+    }
+
+    function ProteopediaCustomColorTheme(ctx: ThemeDataContext, props: PD.Values<ProteopediaCustomColorThemeParams>): ColorTheme<ProteopediaCustomColorThemeParams> {
+        let color: LocationColor
+
+        const colors = props.colors, colorCount = colors.length, defaultColor = props.defaultColor;
+
+        if (ctx.structure) {
+            const l = StructureElement.create()
+            const { models } = ctx.structure
+            const asymIdSerialMap = new Map<string, number>()
+            for (let i = 0, il = models.length; i < il; ++i) {
+                const m = models[i]
+                addAsymIds(asymIdSerialMap, m.atomicHierarchy.chains.label_asym_id)
+                if (m.coarseHierarchy.isDefined) {
+                    addAsymIds(asymIdSerialMap, m.coarseHierarchy.spheres.asym_id)
+                    addAsymIds(asymIdSerialMap, m.coarseHierarchy.gaussians.asym_id)
+                }
+            }
+
+            color = (location: Location): Color => {
+                if (StructureElement.isLocation(location)) {
+                    const asym_id = getAsymId(location.unit);
+                    const o = asymIdSerialMap.get(asym_id(location)) || 0;
+                    return o < colorCount ? colors[o].color : defaultColor;
+                } else if (Link.isLocation(location)) {
+                    const asym_id = getAsymId(location.aUnit)
+                    l.unit = location.aUnit
+                    l.element = location.aUnit.elements[location.aIndex]
+                    const o = asymIdSerialMap.get(asym_id(l)) || 0;
+                    return o < colorCount ? colors[o].color : defaultColor;
+                }
+                return defaultColor
+            }
+        } else {
+            color = () => defaultColor
+        }
+
+        return {
+            factory: ProteopediaCustomColorTheme,
+            granularity: 'group',
+            color,
+            props,
+            description: Description,
+            legend: undefined
+        }
+    }
+
+    const ProteopediaCustomColorThemeProvider: ColorTheme.Provider<ProteopediaCustomColorThemeParams> = {
+        label: 'Proteopedia Custom',
+        factory: ProteopediaCustomColorTheme,
+        getParams: getChainIdColorThemeParams,
+        defaultValues: PD.getDefaultValues(ProteopediaCustomColorThemeParams),
+        isApplicable: (ctx: ThemeDataContext) => !!ctx.structure
+    }
+
+    return ProteopediaCustomColorThemeProvider;
+}
\ No newline at end of file

src/examples/proteopedia-wrapper/helpers.ts

@@ -92,9 +92,24 @@ export interface LoadParams {
 export interface RepresentationStyle {
     sequence?: RepresentationStyle.Entry,
     hetGroups?: RepresentationStyle.Entry,
+    snfg3d?: { hide?: boolean },
     water?: RepresentationStyle.Entry
 }
 
 export namespace RepresentationStyle {
-    export type Entry = { kind?: BuiltInStructureRepresentationsName, coloring?: BuiltInColorThemeName }
+    export type Entry = { hide?: boolean, kind?: BuiltInStructureRepresentationsName, coloring?: BuiltInColorThemeName }
+}
+
+export enum StateElements {
+    Model = 'model',
+    ModelProps = 'model-props',
+    Assembly = 'assembly',
+
+    Sequence = 'sequence',
+    SequenceVisual = 'sequence-visual',
+    Het = 'het',
+    HetVisual = 'het-visual',
+    Het3DSNFG = 'het-3dsnfg',
+    Water = 'water',
+    WaterVisual = 'water-visual'
 }
\ No newline at end of file

src/examples/proteopedia-wrapper/index.html

@@ -55,7 +55,11 @@
             </select>
         </div>
         <div id="app"></div>
-        <script>  
+        <script>
+            // it might be a good idea to define these colors in a separate script file 
+            var CustomColors = [0x00ff00, 0x0000ff];
+            var DefaultCustomColor = 0xff0000;
+
             // create an instance of the plugin
             var PluginWrapper = new MolStarProteopediaWrapper();
 
@@ -78,9 +82,19 @@
             // var format = 'pdb';
             // var assemblyId = 'deposited';
 
-            PluginWrapper.init('app' /** or document.getElementById('app') */);
+            var representationStyle = {
+                sequence: { coloring: 'proteopedia-custom' }, // or just { }
+                hetGroups: { kind: 'ball-and-stick' }, // or 'spacefill
+                water: { hide: true },
+                snfg3d: { hide: false }
+            };
+
+            PluginWrapper.init('app' /** or document.getElementById('app') */, {
+                customColorList: CustomColors,
+                customColorDefault: DefaultCustomColor
+            });
             PluginWrapper.setBackground(0xffffff);
-            PluginWrapper.load({ url: url, format: format, assemblyId: assemblyId });
+            PluginWrapper.load({ url: url, format: format, assemblyId: assemblyId, representationStyle: representationStyle });
             PluginWrapper.toggleSpin();
 
             PluginWrapper.events.modelInfo.subscribe(function (info) {
@@ -92,6 +106,22 @@
 
             addSeparator();
 
+            addHeader('Representation');
+
+            addControl('Custom Chain Colors', () => PluginWrapper.updateStyle({ sequence: { coloring: 'proteopedia-custom' } }, true));
+            addControl('Default Chain Colors', () => PluginWrapper.updateStyle({ sequence: { } }, true));
+
+            addControl('HET Spacefill', () => PluginWrapper.updateStyle({ hetGroups: { kind: 'spacefill' } }, true));
+            addControl('HET Ball-and-stick', () => PluginWrapper.updateStyle({ hetGroups: { kind: 'ball-and-stick' } }, true));
+
+            addControl('Hide 3DSNFG', () => PluginWrapper.updateStyle({ snfg3d: { hide: true } }, true));
+            addControl('Show 3DSNFG', () => PluginWrapper.updateStyle({ snfg3d: { hide: false } }, true));
+
+            addControl('Hide Water', () => PluginWrapper.updateStyle({ water: { hide: true } }, true));
+            addControl('Show Water', () => PluginWrapper.updateStyle({ water: { } }, true));
+
+            addSeparator();
+
             addHeader('Camera');
             addControl('Toggle Spin', () => PluginWrapper.toggleSpin());
             

src/examples/proteopedia-wrapper/index.ts

@@ -15,10 +15,12 @@ import { PluginStateObject as PSO, PluginStateObject } from 'mol-plugin/state/ob
 import { AnimateModelIndex } from 'mol-plugin/state/animation/built-in';
 import { StateBuilder, StateObject } from 'mol-state';
 import { EvolutionaryConservation } from './annotation';
-import { LoadParams, SupportedFormats, RepresentationStyle, ModelInfo } from './helpers';
+import { LoadParams, SupportedFormats, RepresentationStyle, ModelInfo, StateElements } from './helpers';
 import { RxEventHelper } from 'mol-util/rx-event-helper';
 import { ControlsWrapper } from './ui/controls';
 import { PluginState } from 'mol-plugin/state';
+import { Scheduler } from 'mol-task';
+import { createProteopediaCustomTheme } from './coloring';
 require('mol-plugin/skin/light.scss')
 
 class MolStarProteopediaWrapper {
@@ -33,9 +35,15 @@ class MolStarProteopediaWrapper {
 
     plugin: PluginContext;
 
-    init(target: string | HTMLElement) {
+    init(target: string | HTMLElement, options?: {
+        customColorList?: number[],
+        customColorDefault?: number
+    }) {
         this.plugin = createPlugin(typeof target === 'string' ? document.getElementById(target)! : target, {
             ...DefaultPluginSpec,
+            animations: [
+                AnimateModelIndex
+            ],
             layout: {
                 initial: {
                     isExpanded: false,
@@ -47,6 +55,9 @@ class MolStarProteopediaWrapper {
             }
         });
 
+        const customColoring = createProteopediaCustomTheme((options && options.customColorList) || [], (options && options.customColorDefault) || 0x777777);
+
+        this.plugin.structureRepresentation.themeCtx.colorThemeRegistry.add('proteopedia-custom', customColoring);
         this.plugin.structureRepresentation.themeCtx.colorThemeRegistry.add(EvolutionaryConservation.Descriptor.name, EvolutionaryConservation.colorTheme!);
         this.plugin.lociLabels.addProvider(EvolutionaryConservation.labelProvider);
         this.plugin.customModelProperties.register(EvolutionaryConservation.propertyProvider);
@@ -66,43 +77,87 @@ class MolStarProteopediaWrapper {
             : b.apply(StateTransforms.Model.TrajectoryFromPDB);
 
         return parsed
-            .apply(StateTransforms.Model.ModelFromTrajectory, { modelIndex: 0 }, { ref: 'model' });
+            .apply(StateTransforms.Model.ModelFromTrajectory, { modelIndex: 0 }, { ref: StateElements.Model });
     }
 
     private structure(assemblyId: string) {
-        const model = this.state.build().to('model');
+        const model = this.state.build().to(StateElements.Model);
+
+        const s = model
+            .apply(StateTransforms.Model.CustomModelProperties, { properties: [EvolutionaryConservation.Descriptor.name] }, { ref: StateElements.ModelProps, state: { isGhost: false } })
+            .apply(StateTransforms.Model.StructureAssemblyFromModel, { id: assemblyId || 'deposited' }, { ref: StateElements.Assembly });
+
+        s.apply(StateTransforms.Model.StructureComplexElement, { type: 'atomic-sequence' }, { ref: StateElements.Sequence });
+        s.apply(StateTransforms.Model.StructureComplexElement, { type: 'atomic-het' }, { ref: StateElements.Het });
+        s.apply(StateTransforms.Model.StructureComplexElement, { type: 'water' }, { ref: StateElements.Water });
 
-        return model
-            .apply(StateTransforms.Model.CustomModelProperties, { properties: [EvolutionaryConservation.Descriptor.name] }, { ref: 'props', state: { isGhost: false } })
-            .apply(StateTransforms.Model.StructureAssemblyFromModel, { id: assemblyId || 'deposited' }, { ref: 'asm' });
+        return s;
     }
 
-    private visual(ref: string, style?: RepresentationStyle) {
-        const structure = this.getObj<PluginStateObject.Molecule.Structure>(ref);
+    private visual(_style?: RepresentationStyle, partial?: boolean) {
+        const structure = this.getObj<PluginStateObject.Molecule.Structure>(StateElements.Assembly);
         if (!structure) return;
 
-        const root = this.state.build().to(ref);
-
-        root.apply(StateTransforms.Model.StructureComplexElement, { type: 'atomic-sequence' }, { ref: 'sequence' })
-            .apply(StateTransforms.Representation.StructureRepresentation3D,
-                StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin,
-                    (style && style.sequence && style.sequence.kind) || 'cartoon',
-                    (style && style.sequence && style.sequence.coloring) || 'unit-index', structure),
-                    { ref: 'sequence-visual' });
-        root.apply(StateTransforms.Model.StructureComplexElement, { type: 'atomic-het' }, { ref: 'het' })
-            .apply(StateTransforms.Representation.StructureRepresentation3D,
-                StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin,
-                    (style && style.hetGroups && style.hetGroups.kind) || 'ball-and-stick',
-                    (style && style.hetGroups && style.hetGroups.coloring), structure),
-                    { ref: 'het-visual' });
-        root.apply(StateTransforms.Model.StructureComplexElement, { type: 'water' }, { ref: 'water' })
-            .apply(StateTransforms.Representation.StructureRepresentation3D,
-                StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin,
-                    (style && style.water && style.water.kind) || 'ball-and-stick',
-                    (style && style.water && style.water.coloring), structure, { alpha: 0.51 }),
-                    { ref: 'water-visual' });
-
-        return root;
+        const style = _style || { };
+
+        const update = this.state.build();
+
+        if (!partial || (partial && style.sequence)) {
+            const root = update.to(StateElements.Sequence);
+            if (style.sequence && style.sequence.hide) {
+                root.delete(StateElements.SequenceVisual);
+            } else {
+                root.applyOrUpdate(StateElements.SequenceVisual, StateTransforms.Representation.StructureRepresentation3D,
+                    StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin,
+                        (style.sequence && style.sequence.kind) || 'cartoon',
+                        (style.sequence && style.sequence.coloring) || 'unit-index', structure));
+            }
+        }
+
+        if (!partial || (partial && style.hetGroups)) {
+            const root = update.to(StateElements.Het);
+            if (style.hetGroups && style.hetGroups.hide) {
+                root.delete(StateElements.HetVisual);
+            } else {
+                if (style.hetGroups && style.hetGroups.hide) {
+                    root.delete(StateElements.HetVisual);
+                } else {
+                    root.applyOrUpdate(StateElements.HetVisual, StateTransforms.Representation.StructureRepresentation3D,
+                        StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin,
+                            (style.hetGroups && style.hetGroups.kind) || 'ball-and-stick',
+                            (style.hetGroups && style.hetGroups.coloring), structure));
+                }
+            }
+        }
+
+        if (!partial || (partial && style.snfg3d)) {
+            const root = update.to(StateElements.Het);
+            if (style.hetGroups && style.hetGroups.hide) {
+                root.delete(StateElements.HetVisual);
+            } else {
+                if (style.snfg3d && style.snfg3d.hide) {
+                    root.delete(StateElements.Het3DSNFG);
+                } else {
+                    root.applyOrUpdate(StateElements.Het3DSNFG, StateTransforms.Representation.StructureRepresentation3D,
+                        StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin, 'carbohydrate', void 0, structure));
+                }
+            }
+        }
+
+        if (!partial || (partial && style.water)) {
+            const root = update.to(StateElements.Het);
+            if (style.water && style.water.hide) {
+                root.delete(StateElements.Water);
+            } else {
+                root.applyOrUpdate(StateElements.Water, StateTransforms.Model.StructureComplexElement, { type: 'water' })
+                    .applyOrUpdate(StateElements.WaterVisual, StateTransforms.Representation.StructureRepresentation3D,
+                        StructureRepresentation3DHelpers.getDefaultParamsWithTheme(this.plugin,
+                            (style.water && style.water.kind) || 'ball-and-stick',
+                            (style.water && style.water.coloring), structure, { alpha: 0.51 }));
+            }
+        }
+
+        return update;
     }
 
     private getObj<T extends StateObject>(ref: string): T['data'] {
@@ -134,7 +189,7 @@ class MolStarProteopediaWrapper {
         if (this.loadedParams.url !== url || this.loadedParams.format !== format) {
             loadType = 'full';
         } else if (this.loadedParams.url === url) {
-            if (state.select('asm').length > 0) loadType = 'update';
+            if (state.select(StateElements.Assembly).length > 0) loadType = 'update';
         }
 
         if (loadType === 'full') {
@@ -146,18 +201,18 @@ class MolStarProteopediaWrapper {
             await this.applyState(structureTree);
         } else {
             const tree = state.build();
-            tree.to('asm').update(StateTransforms.Model.StructureAssemblyFromModel, p => ({ ...p, id: assemblyId || 'deposited' }));
+            tree.to(StateElements.Assembly).update(StateTransforms.Model.StructureAssemblyFromModel, p => ({ ...p, id: assemblyId || 'deposited' }));
             await this.applyState(tree);
         }
 
         await this.updateStyle(representationStyle);
 
         this.loadedParams = { url, format, assemblyId };
-        PluginCommands.Camera.Reset.dispatch(this.plugin, { });
+        Scheduler.setImmediate(() => PluginCommands.Camera.Reset.dispatch(this.plugin, { }));
     }
 
-    async updateStyle(style?: RepresentationStyle) {
-        const tree = this.visual('asm', style);
+    async updateStyle(style?: RepresentationStyle, partial?: boolean) {
+        const tree = this.visual(style, partial);
         if (!tree) return;
         await PluginCommands.State.Update.dispatch(this.plugin, { state: this.plugin.state.dataState, tree });
     }
@@ -186,7 +241,7 @@ class MolStarProteopediaWrapper {
 
     coloring = {
         evolutionaryConservation: async () => {
-            await this.updateStyle({ sequence: { kind: 'spacefill' } });
+            await this.updateStyle({ sequence: { kind: 'spacefill' } }, true);
 
             const state = this.state;
 
@@ -194,7 +249,7 @@ class MolStarProteopediaWrapper {
             const tree = state.build();
             const colorTheme = { name: EvolutionaryConservation.Descriptor.name, params: this.plugin.structureRepresentation.themeCtx.colorThemeRegistry.get(EvolutionaryConservation.Descriptor.name).defaultValues };
 
-            tree.to('sequence-visual').update(StateTransforms.Representation.StructureRepresentation3D, old => ({ ...old, colorTheme }));
+            tree.to(StateElements.SequenceVisual).update(StateTransforms.Representation.StructureRepresentation3D, old => ({ ...old, colorTheme }));
             // for (const v of visuals) {
             // }
 

src/mol-state/state/builder.ts

@@ -84,6 +84,7 @@ namespace StateBuilder {
         }
         toRoot<A extends StateObject>() { return new To<A>(this.state, this.state.tree.root.ref, this); }
         delete(ref: StateTransform.Ref) {
+            if (!this.state.tree.transforms.has(ref)) return this;
             this.editInfo.count++;
             this.state.tree.remove(ref);
             this.state.actions.push({ kind: 'delete', ref });
@@ -113,6 +114,19 @@ namespace StateBuilder {
             return new To(this.state, t.ref, this.root);
         }
 
+        /**
+         * If the ref is present, the transform is applied.
+         * Otherwise a transform with the specifed ref is created.
+         */
+        applyOrUpdate<T extends StateTransformer<A, any, any>>(ref: StateTransform.Ref, tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
+            if (this.state.tree.transforms.has(ref)) {
+                this.to(ref).update(params);
+                return this.to(ref) as To<StateTransformer.To<T>>;
+            } else {
+                return this.apply(tr, params, { ...options, ref });
+            }
+        }
+
         /**
          * A helper to greate a group-like state object and keep the current type.
          */