wip, StructureRepresentationHelper tweaks

src/mol-plugin-ui/structure/representation.tsx

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -13,8 +13,6 @@ import { Color } from '../../mol-util/color';
 import { ButtonSelect, Options } from '../controls/common'
 import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { VisualQuality, VisualQualityOptions } from '../../mol-geo/geometry/base';
-import { StructureRepresentationPresets as P } from '../../mol-plugin/util/structure-representation-helper';
-import { camelCaseToWords } from '../../mol-util/string';
 import { CollapsableControls } from '../base';
 import { StateSelection, StateObject } from '../../mol-state';
 import { PluginStateObject } from '../../mol-plugin/state/objects';
@@ -133,13 +131,6 @@ export class StructureRepresentationControls extends CollapsableControls<Collaps
         this.subscribe(this.plugin.state.dataState.events.isUpdating, v => this.setState({ isDisabled: v }))
     }
 
-    preset = async (value: string) => {
-        const presetFn = P[value as keyof typeof P]
-        if (presetFn) {
-            await presetFn(this.plugin.helpers.structureRepresentation)
-        }
-    }
-
     onChange = async (p: { param: PD.Base<any>, name: string, value: any }) => {
         if (p.name === 'options') {
             await this.plugin.helpers.structureRepresentation.setIgnoreHydrogens(!p.value.showHydrogens)
@@ -178,15 +169,7 @@ export class StructureRepresentationControls extends CollapsableControls<Collaps
     }
 
     renderControls() {
-        const presets = PD.objectToOptions(P, camelCaseToWords);
         return <div>
-            <div className='msp-control-row'>
-                <div className='msp-select-row'>
-                    <ButtonSelect label='Preset' onChange={this.preset}>
-                        <optgroup label='Preset'>{Options(presets)}</optgroup>
-                    </ButtonSelect>
-                </div>
-            </div>
             <EverythingStructureRepresentationControls />
             <SelectionStructureRepresentationControls />
 

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

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -12,8 +12,6 @@ import { PluginContext } from '../context';
 import { StructureRepresentation3DHelpers } from '../state/transforms/representation';
 import Expression from '../../mol-script/language/expression';
 import { compile } from '../../mol-script/runtime/query/compiler';
-import { StructureSelectionQueries as Q } from '../util/structure-selection-helper';
-import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
 import { VisualQuality } from '../../mol-geo/geometry/base';
 
 type StructureTransform = StateObjectCell<PSO.Molecule.Structure, StateTransform<StateTransformer<any, PSO.Molecule.Structure, any>>>
@@ -34,6 +32,12 @@ function getCombinedLoci(mode: SelectionModifier, loci: StructureElement.Loci, c
 
 type SelectionModifier = 'add' | 'remove' | 'only'
 
+type ReprProps = {
+    repr?: {},
+    color?: string | [string, {}],
+    size?: string | [string, {}],
+}
+
 export class StructureRepresentationHelper {
     getRepresentationStructure(rootRef: string, type: string) {
         const state = this.plugin.state.dataState
@@ -49,12 +53,52 @@ export class StructureRepresentationHelper {
         return selections.length > 0 ? selections[0] : undefined
     }
 
-    private async _set(modifier: SelectionModifier, type: string, loci: StructureElement.Loci, structure: StructureTransform, props = {}) {
+    private getRepresentationParams(structure: Structure, type: string, repr: RepresentationTransform | undefined, props: ReprProps = {}) {
+        const reprProps = {
+            ...(repr?.params && repr.params.values.type.params),
+            ignoreHydrogens: this._ignoreHydrogens,
+            quality: this._quality,
+            ...props.repr
+        }
+        const { themeCtx } =  this.plugin.structureRepresentation
+
+        const p: StructureRepresentation3DHelpers.Props = {
+            repr: [
+                this.plugin.structureRepresentation.registry.get(type),
+                () => reprProps
+            ]
+        }
+        if (props.color) {
+            const colorType = props.color instanceof Array ? props.color[0] : props.color
+            const colorTheme = themeCtx.colorThemeRegistry.get(colorType)
+            const colorProps = {
+                ...(repr?.params && repr.params.values.colorTheme.params),
+                ...(props.color instanceof Array ? props.color[1] : {})
+            }
+            p.color = [colorTheme, () => colorProps]
+        }
+        if (props.size) {
+            const sizeType = props.size instanceof Array ? props.size[0] : props.size
+            const sizeTheme = themeCtx.sizeThemeRegistry.get(sizeType)
+            const sizeProps = {
+                ...(repr?.params && repr.params.values.sizeTheme.params),
+                ...(props.size instanceof Array ? props.size[1] : {})
+            }
+            p.size = [sizeTheme, () => sizeProps]
+        }
+        if (props.size) p.size = props.size
+
+        return StructureRepresentation3DHelpers.createParams(this.plugin, structure, p)
+    }
+
+    private async _set(modifier: SelectionModifier, type: string, loci: StructureElement.Loci, structure: StructureTransform, props: ReprProps = {}) {
         const state = this.plugin.state.dataState
         const update = state.build()
         const s = structure.obj!.data
 
+        const repr = this.getRepresentation(structure.transform.ref, type)
         const reprStructure = this.getRepresentationStructure(structure.transform.ref, type)
+        const reprParams = this.getRepresentationParams(s, type, repr, props)
 
         if (reprStructure) {
             const currentLoci = StructureElement.Bundle.toLoci(reprStructure.params!.values.bundle, s)
@@ -64,14 +108,9 @@ export class StructureRepresentationHelper {
                 ...reprStructure.params!.values,
                 bundle: StructureElement.Bundle.fromLoci(combinedLoci)
             })
+            if (repr) update.to(repr).update(reprParams)
         } else {
             const combinedLoci = getCombinedLoci(modifier, loci, StructureElement.Loci.none(s))
-            const params = StructureRepresentation3DHelpers.getDefaultParams(this.plugin, type as any, s)
-
-            const p = params.type.params
-            Object.assign(p, props)
-            if (p.ignoreHydrogens !== undefined) p.ignoreHydrogens = this._ignoreHydrogens
-            if (p.quality !== undefined) p.quality = this._quality
 
             update.to(structure.transform.ref)
                 .apply(
@@ -79,13 +118,13 @@ export class StructureRepresentationHelper {
                     { bundle: StructureElement.Bundle.fromLoci(combinedLoci), label: type },
                     { tags: [ RepresentationManagerTag, getRepresentationManagerTag(type) ] }
                 )
-                .apply( StateTransforms.Representation.StructureRepresentation3D, params)
+                .apply(StateTransforms.Representation.StructureRepresentation3D, reprParams)
         }
 
         await this.plugin.runTask(state.updateTree(update, { doNotUpdateCurrent: true }))
     }
 
-    async set(modifier: SelectionModifier, type: string, lociGetter: (structure: Structure) => StructureElement.Loci, props = {}) {
+    async set(modifier: SelectionModifier, type: string, lociGetter: (structure: Structure) => StructureElement.Loci, props: ReprProps = {}) {
         const state = this.plugin.state.dataState;
         const structures = state.select(StateSelection.Generators.rootsOfType(PSO.Molecule.Structure))
 
@@ -96,7 +135,7 @@ export class StructureRepresentationHelper {
         }
     }
 
-    async setFromExpression(modifier: SelectionModifier, type: string, expression: Expression, props = {}) {
+    async setFromExpression(modifier: SelectionModifier, type: string, expression: Expression, props: ReprProps = {}) {
         return this.set(modifier, type, (structure) => {
             const compiled = compile<StructureSelection>(expression)
             const result = compiled(new QueryContext(structure))
@@ -104,44 +143,76 @@ export class StructureRepresentationHelper {
         }, props)
     }
 
-    async clear() {
+    async eachStructure(callback: (structure: StructureTransform, type: string, update: StateBuilder.Root) => void) {
         const { registry } = this.plugin.structureRepresentation
         const state = this.plugin.state.dataState;
         const update = state.build()
         const structures = state.select(StateSelection.Generators.rootsOfType(PSO.Molecule.Structure))
-        const bundle = StructureElement.Bundle.Empty
 
         for (const structure of structures) {
             for (let i = 0, il = registry.types.length; i < il; ++i) {
                 const type = registry.types[i][0]
                 const reprStructure = this.getRepresentationStructure(structure.transform.ref, type)
-                if (reprStructure) {
-                    update.to(reprStructure).update({ ...reprStructure.params!.values, bundle })
-                }
+                if (reprStructure) callback(reprStructure, type, update)
             }
         }
         await this.plugin.runTask(state.updateTree(update, { doNotUpdateCurrent: true }))
     }
 
-    async eachRepresentation(callback: (repr: RepresentationTransform, update: StateBuilder.Root) => void) {
+    async clear() {
+        const bundle = StructureElement.Bundle.Empty
+        await this.eachStructure((structure, type, update) => {
+            update.to(structure).update({ ...structure.params!.values, bundle })
+        })
+    }
+
+    async clearExcept(exceptTypes: string[]) {
+        const bundle = StructureElement.Bundle.Empty
+        await this.eachStructure((structure, type, update) => {
+            if (!exceptTypes.includes(type)) {
+                update.to(structure).update({ ...structure.params!.values, bundle })
+            }
+        })
+    }
+
+    async eachRepresentation(callback: (repr: RepresentationTransform, type: string, update: StateBuilder.Root) => void) {
         const { registry } = this.plugin.structureRepresentation
         const state = this.plugin.state.dataState;
         const update = state.build()
         const structures = state.select(StateSelection.Generators.rootsOfType(PSO.Molecule.Structure))
         for (const structure of structures) {
             for (let i = 0, il = registry.types.length; i < il; ++i) {
-                const repr = this.getRepresentation(structure.transform.ref, registry.types[i][0])
-                if (repr) callback(repr, update)
+                const type = registry.types[i][0]
+                const repr = this.getRepresentation(structure.transform.ref, type)
+                if (repr) callback(repr, type, update)
             }
         }
         await this.plugin.runTask(state.updateTree(update, { doNotUpdateCurrent: true }))
     }
 
+    setRepresentationParams(repr: RepresentationTransform, type: string, update: StateBuilder.Root, props: ReprProps) {
+        const state = this.plugin.state.dataState;
+        const structures = state.select(StateSelection.Generators.rootsOfType(PSO.Molecule.Structure))
+
+        for (const structure of structures) {
+            const s = structure.obj!.data
+            const reprParams = this.getRepresentationParams(s, type, repr, props)
+            update.to(repr).update(reprParams)
+        }
+    }
+
+    async updateRepresentation(repr: RepresentationTransform, type: string, props: ReprProps) {
+        const state = this.plugin.state.dataState;
+        const update = state.build()
+        this.setRepresentationParams(repr, type, update, props)
+        await this.plugin.runTask(state.updateTree(update, { doNotUpdateCurrent: true }))
+    }
+
     private _ignoreHydrogens = false
     get ignoreHydrogens () { return this._ignoreHydrogens }
     async setIgnoreHydrogens(ignoreHydrogens: boolean) {
         if (ignoreHydrogens === this._ignoreHydrogens) return
-        await this.eachRepresentation((repr, update) => {
+        await this.eachRepresentation((repr, type, update) => {
             if (repr.params && repr.params.values.type.params.ignoreHydrogens !== undefined) {
                 const { name, params } = repr.params.values.type
                 update.to(repr.transform.ref).update(
@@ -157,7 +228,7 @@ export class StructureRepresentationHelper {
     get quality () { return this._quality }
     async setQuality(quality: VisualQuality) {
         if (quality === this._quality) return
-        await this.eachRepresentation((repr, update) => {
+        await this.eachRepresentation((repr, type, update) => {
             if (repr.params && repr.params.values.type.params.quality !== undefined) {
                 const { name, params } = repr.params.values.type
                 update.to(repr.transform.ref).update(
@@ -169,74 +240,7 @@ export class StructureRepresentationHelper {
         this._quality = quality
     }
 
-    async preset() {
-        // TODO option to limit to specific structure
-        const state = this.plugin.state.dataState;
-        const structures = state.select(StateSelection.Generators.rootsOfType(PSO.Molecule.Structure))
-
-        if (structures.length === 0) return
-        const s = structures[0].obj!.data
-
-        if (s.elementCount < 50000) {
-            await polymerAndLigand(this)
-        } else if (s.elementCount < 200000) {
-            await proteinAndNucleic(this)
-        } else {
-            if (s.unitSymmetryGroups[0].units.length > 10) {
-                await capsid(this)
-            } else {
-                await coarseCapsid(this)
-            }
-        }
-    }
-
     constructor(private plugin: PluginContext) {
 
     }
-}
-
-//
-
-async function polymerAndLigand(r: StructureRepresentationHelper) {
-    await r.clear()
-    await r.setFromExpression('add', 'cartoon', Q.polymer.expression)
-    await r.setFromExpression('add', 'carbohydrate', Q.branchedPlusConnected.expression)
-    await r.setFromExpression('add', 'ball-and-stick', MS.struct.modifier.union([
-        MS.struct.combinator.merge([
-            Q.ligandPlusConnected.expression,
-            Q.branchedConnectedOnly.expression,
-            Q.disulfideBridges.expression,
-            Q.nonStandardPolymer.expression,
-            Q.water.expression
-        ])
-    ]))
-}
-
-async function proteinAndNucleic(r: StructureRepresentationHelper) {
-    await r.clear()
-    await r.setFromExpression('add', 'cartoon', Q.protein.expression)
-    await r.setFromExpression('add', 'gaussian-surface', Q.nucleic.expression)
-}
-
-async function capsid(r: StructureRepresentationHelper) {
-    await r.clear()
-    await r.setFromExpression('add', 'gaussian-surface', Q.polymer.expression, {
-        smoothness: 0.5,
-    })
-}
-
-async function coarseCapsid(r: StructureRepresentationHelper) {
-    await r.clear()
-    await r.setFromExpression('add', 'gaussian-surface', Q.trace.expression, {
-        radiusOffset: 1,
-        smoothness: 0.5,
-        visuals: ['structure-gaussian-surface-mesh']
-    })
-}
-
-export const StructureRepresentationPresets = {
-    polymerAndLigand,
-    proteinAndNucleic,
-    capsid,
-    coarseCapsid
 }