Procházet zdrojové kódy

mol-plugin: modify component manager/ui

David Sehnal před 5 roky
rodič
revize
a3ebc4df45

+ 1 - 1
src/mol-canvas3d/canvas3d.ts

@@ -37,7 +37,7 @@ import { Sphere3D } from '../mol-math/geometry';
 import { isDebugMode } from '../mol-util/debug';
 
 export const Canvas3DParams = {
-    cameraMode: PD.Select('perspective', [['perspective', 'Perspective'], ['orthographic', 'Orthographic']]),
+    cameraMode: PD.Select('perspective', [['perspective', 'Perspective'], ['orthographic', 'Orthographic']] as const),
     cameraFog: PD.Numeric(50, { min: 0, max: 100, step: 1 }),
     cameraClipFar: PD.Boolean(true),
     cameraResetDurationMs: PD.Numeric(250, { min: 0, max: 1000, step: 1 }, { description: 'The time it takes to reset the camera.' }),

+ 47 - 0
src/mol-plugin-state/builder/structure.ts

@@ -12,6 +12,8 @@ import { RootStructureDefinition } from '../helpers/root-structure';
 import { StructureComponentParams } from '../helpers/structure-component';
 import { BuildInTrajectoryFormat, TrajectoryFormatProvider } from '../formats/trajectory';
 import { StructureRepresentationBuilder } from './structure/representation';
+import { StructureSelectionQuery } from '../helpers/structure-selection-query';
+import { Task } from '../../mol-task';
 
 export type TrajectoryFormat = 'pdb' | 'cif' | 'gro' | '3dg'
 
@@ -151,6 +153,51 @@ export class StructureBuilder {
         return selector;
     }
 
+    tryCreateQueryComponent(params: { structure: StateObjectRef<SO.Molecule.Structure>, query: StructureSelectionQuery, key: string, label?: string, tags?: string[] }): Promise<StateObjectRef<SO.Molecule.Structure> | undefined> {
+        return this.plugin.runTask(Task.create('Query Component', async taskCtx => {
+            let { structure, query, key, label, tags } = params;        
+            label = (label || '').trim();
+
+            const structureData = StateObjectRef.resolveAndCheck(this.dataState, structure)?.obj?.data;
+
+            if (!structureData) return;
+    
+            const transformParams: StructureComponentParams = query.referencesCurrent
+                ? {
+                    type: { name: 'bundle', params: await StructureSelectionQuery.getBundle(this.plugin, taskCtx, query, structureData) },
+                    nullIfEmpty: true,
+                    label: label || query.label
+                } : {
+                    type: { name: 'expression', params: query.expression },
+                    nullIfEmpty: true,
+                    label: label || query.label
+                };
+
+            if (query.ensureCustomProperties) {
+                await query.ensureCustomProperties({ fetch: this.plugin.fetch, runtime: taskCtx }, structureData);
+            }
+            
+            const state = this.dataState;
+            const root = state.build().to(structure);
+            const keyTag = `structure-component-${key}`;
+            const component = root.applyOrUpdateTagged(keyTag, StateTransforms.Model.StructureComponent, transformParams, { 
+                tags: tags ? [...tags, StructureBuilderTags.Component, keyTag] : [StructureBuilderTags.Component, keyTag]
+            });
+    
+            await this.dataState.updateTree(component).runInContext(taskCtx);
+    
+            const selector = component.selector;
+    
+            if (!selector.isOk || selector.cell?.obj?.data.elementCount === 0) {
+                const del = state.build().delete(selector.ref);
+                await this.plugin.runTask(this.dataState.updateTree(del));
+                return;
+            }
+    
+            return selector; 
+        }))
+    }
+
     constructor(public plugin: PluginContext) {
     }
 }

+ 1 - 1
src/mol-plugin-state/builder/structure/provider.ts

@@ -27,7 +27,7 @@ export namespace StructureRepresentationProvider {
 }
 
 export const enum RepresentationProviderTags {
-    Representation = 'preset-structure-representation',
+    Representation = 'structure-representation',
     Component = 'preset-structure-component'
 }
 

+ 21 - 3
src/mol-plugin-state/builder/structure/representation.ts

@@ -15,6 +15,11 @@ import { PluginContext } from '../../../mol-plugin/context';
 import { PresetStructureReprentations } from './preset';
 import { StructureRepresentationProvider, RepresentationProviderTags } from './provider';
 import { UniqueArray } from '../../../mol-data/generic';
+import { PluginStateObject } from '../../objects';
+import { StructureRepresentation3D, StructureRepresentation3DHelpers } from '../../transforms/representation';
+import { RepresentationProvider } from '../../../mol-repr/representation';
+import { SizeTheme } from '../../../mol-theme/size';
+import { ColorTheme } from '../../../mol-theme/color';
 
 // TODO: support quality
 // TODO: support ignore hydrogens
@@ -24,6 +29,7 @@ export type StructureRepresentationProviderRef = keyof PresetStructureReprentati
 export class StructureRepresentationBuilder {
     private providers: StructureRepresentationProvider[] = [];
     private providerMap: Map<string, StructureRepresentationProvider> = new Map();
+    private get dataState() { return this.plugin.state.dataState; }
 
     readonly defaultProvider = PresetStructureReprentations.auto;
 
@@ -97,7 +103,7 @@ export class StructureRepresentationBuilder {
         if (!id) return;
 
         const state = this.plugin.state.dataState;
-        const root = StateObjectRef.resolveRef(state, structureRoot) || StateTransform.RootRef;
+        const root = StateObjectRef.resolveRef(structureRoot) || StateTransform.RootRef;
         const reprs = StateSelection.findWithAllTags(state.tree, root, new Set([id, RepresentationProviderTags.Representation]));
 
         const builder = state.build();
@@ -139,8 +145,20 @@ export class StructureRepresentationBuilder {
         return this.plugin.runTask(task);
     }
 
-    // TODO
-    // createOrUpdate(component: any, ) { }
+    async addRepresentation<R extends RepresentationProvider<Structure, any, any>, C extends ColorTheme.Provider<any>, S extends SizeTheme.Provider<any>>
+        (structure: StateObjectRef<PluginStateObject.Molecule.Structure>, props: StructureRepresentation3DHelpers.Props<R, C, S>) {
+
+        const data = StateObjectRef.resolveAndCheck(this.dataState, structure)?.obj?.data;
+        if (!data) return;
+
+        const params = StructureRepresentation3DHelpers.createParams(this.plugin, data, props);
+        const repr = this.dataState.build()
+            .to(structure)
+            .apply(StructureRepresentation3D, params, { tags: RepresentationProviderTags.Representation });
+
+        await this.plugin.runTask(this.dataState.updateTree(repr));
+        return  repr.selector;
+    }
 
     constructor(public plugin: PluginContext) {
         objectForEach(PresetStructureReprentations, r => this.registerPreset(r));

+ 66 - 0
src/mol-plugin-state/helpers/structure-overpaint.ts

@@ -0,0 +1,66 @@
+/**
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Structure, StructureElement } from '../../mol-model/structure';
+import { PluginStateObject } from '../../mol-plugin-state/objects';
+import { StateTransforms } from '../../mol-plugin-state/transforms';
+import { PluginContext } from '../../mol-plugin/context';
+import { StateBuilder, StateObjectCell, StateSelection, StateTransform } from '../../mol-state';
+import { Overpaint } from '../../mol-theme/overpaint';
+import { Color } from '../../mol-util/color';
+import { StructureComponentRef } from '../manager/structure/hierarchy-state';
+import { EmptyLoci, Loci } from '../../mol-model/loci';
+
+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 async function setStructureOverpaint(plugin: PluginContext, components: StructureComponentRef[], color: Color | -1, lociGetter: (structure: Structure) => StructureElement.Loci | EmptyLoci, types?: string[], alpha = 1) {
+    await eachRepr(plugin, components, (update, repr, overpaintCell) => {
+        if (types && !types.includes(repr.params!.values.type.name)) return
+
+        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 (Loci.isEmpty(loci)) return
+
+        const layer = {
+            bundle: StructureElement.Bundle.fromLoci(loci),
+            color: color === -1 ? Color(0) : color,
+            clear: color === -1
+        }
+
+        if (overpaintCell) {
+            const bundleLayers = [...overpaintCell.params!.values.layers, layer]
+            const filtered = getFilteredBundle(bundleLayers, structure)
+            update.to(overpaintCell).update(Overpaint.toBundle(filtered, alpha))
+        } else {
+            const filtered = getFilteredBundle([layer], structure)
+            update.to(repr.transform.ref)
+                .apply(StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle, Overpaint.toBundle(filtered, alpha), { tags: OverpaintManagerTag });
+        }
+    })
+}
+
+async function eachRepr(plugin: PluginContext, components: StructureComponentRef[], callback: OverpaintEachReprCallback) {
+    const state = plugin.state.dataState;    
+    const update = state.build();
+    for (const c of components) {
+        for (const r of c.representations) {
+            const overpaint = state.select(StateSelection.Generators.ofTransformer(StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle, r.cell.transform.ref).withTag(OverpaintManagerTag))
+            callback(update, r.cell, overpaint[0])
+        }
+    }
+
+    await plugin.runTask(state.updateTree(update, { doNotUpdateCurrent: true }));
+}
+
+/** filter overpaint layers for given structure */
+function getFilteredBundle(layers: Overpaint.BundleLayer[], structure: Structure) {
+    const overpaint = Overpaint.ofBundle(layers, 1, structure.root)
+    const merged = Overpaint.merge(overpaint)
+    return Overpaint.filter(merged, structure)
+}

+ 20 - 1
src/mol-plugin-state/helpers/structure-selection-query.ts

@@ -8,7 +8,7 @@
 import { CustomProperty } from '../../mol-model-props/common/custom-property';
 import { AccessibleSurfaceAreaProvider, AccessibleSurfaceAreaSymbols } from '../../mol-model-props/computed/accessible-surface-area';
 import { ValidationReport, ValidationReportProvider } from '../../mol-model-props/rcsb/validation-report';
-import { QueryContext, Structure, StructureQuery, StructureSelection } from '../../mol-model/structure';
+import { QueryContext, Structure, StructureQuery, StructureSelection, StructureElement } from '../../mol-model/structure';
 import { BondType, NucleicBackboneAtoms, ProteinBackboneAtoms, SecondaryStructureType } from '../../mol-model/structure/model/types';
 import { PluginStateObject } from '../objects';
 import { StateTransforms } from '../transforms';
@@ -502,6 +502,8 @@ export const StructureSelectionQueryList = [
     ...StandardNucleicBases.map(v => ResidueQuery(v, StructureSelectionCategory.NucleicBase)),
 ]
 
+export const StructureSelectionQueryOptions: [StructureSelectionQuery, string, string][] = StructureSelectionQueryList.map(q => [q, q.label, q.category])
+
 export function applyBuiltInSelection(to: StateBuilder.To<PluginStateObject.Molecule.Structure>, query: keyof typeof StructureSelectionQueries, customTag?: string) {
     return to.apply(StateTransforms.Model.StructureSelectionFromExpression,
         { expression: StructureSelectionQueries[query].expression, label: StructureSelectionQueries[query].label },
@@ -509,6 +511,18 @@ export function applyBuiltInSelection(to: StateBuilder.To<PluginStateObject.Mole
 }
 
 namespace StructureSelectionQuery {
+    export async function getStructure(plugin: PluginContext, runtime: RuntimeContext, selectionQuery: StructureSelectionQuery, structure: Structure) {
+        const current = plugin.managers.structure.selection.getStructure(structure)
+        const currentSelection = current ? StructureSelection.Singletons(structure, current) : StructureSelection.Empty(structure);
+
+        if (selectionQuery.ensureCustomProperties) {
+            await selectionQuery.ensureCustomProperties({ fetch: plugin.fetch, runtime }, structure)
+        }
+
+        const result = selectionQuery.query(new QueryContext(structure, { currentSelection }))
+        return StructureSelection.unionStructure(result)
+    }
+
     export async function getLoci(plugin: PluginContext, runtime: RuntimeContext, selectionQuery: StructureSelectionQuery, structure: Structure) {
         const current = plugin.managers.structure.selection.getStructure(structure)
         const currentSelection = current ? StructureSelection.Singletons(structure, current) : StructureSelection.Empty(structure);
@@ -520,4 +534,9 @@ namespace StructureSelectionQuery {
         const result = selectionQuery.query(new QueryContext(structure, { currentSelection }))
         return StructureSelection.toLociWithSourceUnits(result)
     }
+
+    export async function getBundle(plugin: PluginContext, runtime: RuntimeContext, selectionQuery: StructureSelectionQuery, structure: Structure) {
+        const loci = await getLoci(plugin, runtime, selectionQuery, structure);
+        return StructureElement.Bundle.fromLoci(loci);
+    }
 }

+ 172 - 18
src/mol-plugin-state/manager/structure/component.ts

@@ -4,16 +4,27 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { StructureRef } from './hierarchy-state'
-import { StructureRepresentationProvider } from '../../builder/structure/provider';
+import { Structure, StructureElement, StructureSelection } from '../../../mol-model/structure';
+import { structureAreIntersecting, structureSubtract, structureUnion } from '../../../mol-model/structure/query/utils/structure-set';
 import { PluginContext } from '../../../mol-plugin/context';
+import { StateBuilder } from '../../../mol-state';
+import { Task } from '../../../mol-task';
+import { UUID } from '../../../mol-util';
+import { Color } from '../../../mol-util/color';
+import { ColorNames } from '../../../mol-util/color/names';
+import { ParamDefinition as PD } from '../../../mol-util/param-definition';
+import { StructureRepresentationProvider } from '../../builder/structure/provider';
+import { StructureComponentParams } from '../../helpers/structure-component';
+import { setStructureOverpaint } from '../../helpers/structure-overpaint';
+import { StructureSelectionQuery, StructureSelectionQueryOptions } from '../../helpers/structure-selection-query';
+import { HierarchyRef, StructureComponentRef, StructureRef, StructureRepresentationRef } from './hierarchy-state';
 
-export { StructureComponentManager }
+export { StructureComponentManager };
 
 class StructureComponentManager {
     applyPreset<P = any, S = {}>(structures: StructureRef[], provider: StructureRepresentationProvider<P, S>, params?: P): Promise<any>  {
-        return this.plugin.runTask(this.dateState.transaction(async () => {
-            await this.removeComponents(structures);
+        return this.plugin.runTask(this.dataState.transaction(async () => {
+            await this.clearComponents(structures);
             for (const s of structures) {
                 await this.plugin.builders.structure.representation.structurePreset(s.cell, provider, params);
             }
@@ -21,19 +32,114 @@ class StructureComponentManager {
     }
 
     clear(structures: StructureRef[]) {
-        return this.removeComponents(structures);
+        return this.clearComponents(structures);
+    }
+
+    removeRepresentations(components: StructureComponentRef[], pivot: StructureRepresentationRef) {
+        if (components.length === 0) return;
+        const index = components[0].representations.indexOf(pivot);
+        if (index < 0) return;
+
+        const toRemove: HierarchyRef[] = [];
+        for (const c of components) {
+            if (index >= c.representations.length) continue;
+            toRemove.push(c.representations[index]);
+        }
+        return this.plugin.managers.structure.hierarchy.remove(toRemove);
+    }
+
+    modify(action: StructureComponentManager.ModifyAction, structures?: ReadonlyArray<StructureRef>) {        
+        return this.plugin.runTask(this.dataState.transaction(async () => {
+            if (!structures) structures = this.plugin.managers.structure.hierarchy.state.currentStructures;
+            if (structures.length === 0) return;
+
+            switch (action.kind) {
+                case 'add': await this.modifyAdd(action, structures); break;
+                case 'merge': await this.modifyMerge(action, structures); break;
+                case 'subtract': await this.modifySubtract(action, structures); break;
+                case 'color': await this.modifyColor(action, structures); break;
+            }
+        }))
+    }
+
+    private async modifyAdd(params: StructureComponentManager.ModifyActionAdd, structures: ReadonlyArray<StructureRef>) {
+        const componentKey = UUID.create22();
+        for (const s of structures) {
+            const component = await this.plugin.builders.structure.tryCreateQueryComponent({ 
+                structure: s.cell,
+                query: params.selection,
+                key: componentKey,
+                label: params.label,
+            });
+            if (params.representation === 'none' || !component) continue;
+            await this.plugin.builders.structure.representation.addRepresentation(component, {
+                repr: this.plugin.structureRepresentation.registry.get(params.representation)
+            });
+        }
+    }
+
+    private updateComponent(builder: StateBuilder.Root, component: StructureComponentRef, by: Structure, action: 'union' | 'subtract') {
+        const structure = component.cell.obj?.data;
+        if (!structure) return;
+        if (!structureAreIntersecting(structure, by)) return;
+
+        const parent = component.structure.cell.obj?.data!;
+        const modified = action === 'union' ? structureUnion(parent, [structure, by]) : structureSubtract(structure, by);
+
+        if (modified.elementCount === 0) {
+            builder.delete(component.cell.transform.ref);
+        } else {
+            const bundle = StructureElement.Bundle.fromLoci(StructureSelection.toLociWithSourceUnits(StructureSelection.Singletons(parent, modified)));
+            const params: StructureComponentParams = {
+                type: { name: 'bundle', params: bundle },
+                nullIfEmpty: true,
+                label: component.cell.obj?.label!
+            };
+            builder.to(component.cell).update(params)
+        }
     }
 
-    modify(structures: StructureRef[], action: StructureComponentManager.ModifyAction) {
+    private async modifyMerge(params: StructureComponentManager.ModifyActionMerge, structures: ReadonlyArray<StructureRef>) {
+        return this.plugin.runTask(Task.create('Merge', async taskCtx => {
+            const b = this.dataState.build();        
+            for (const s of structures) {
+                const by = await StructureSelectionQuery.getStructure(this.plugin, taskCtx, params.selection, s.cell.obj?.data!);
+                for (const c of s.components) {
+                    if (params.componentKey !== 'intersecting' && params.componentKey !== c.key) continue;
+                    this.updateComponent(b, c, by, 'union');
+                }
+            }
+            await this.dataState.updateTree(b).runInContext(taskCtx);
+        }));
+    }
+
+    private async modifySubtract(params: StructureComponentManager.ModifyActionSubtract, structures: ReadonlyArray<StructureRef>) {
+        return this.plugin.runTask(Task.create('Subtract', async taskCtx => {
+            const b = this.dataState.build();        
+            for (const s of structures) {
+                const by = await StructureSelectionQuery.getStructure(this.plugin, taskCtx, params.selection, s.cell.obj?.data!);
+                for (const c of s.components) {
+                    if (params.componentKey !== 'intersecting' && params.componentKey !== c.key) continue;
+                    this.updateComponent(b, c, by, 'subtract');
+                }
+            }
+            await this.dataState.updateTree(b).runInContext(taskCtx);
+        }));
+    }
 
+    private async modifyColor(params: StructureComponentManager.ModifyActionColor, structures: ReadonlyArray<StructureRef>) {
+        const getLoci = (s: Structure) => this.plugin.managers.structure.selection.getLoci(s);
+        for (const s of structures) {
+            await setStructureOverpaint(this.plugin, s.components, params.action.name === 'color' ? params.action.params : -1, getLoci);
+        }
     }
 
-    private get dateState() {
+    private get dataState() {
         return this.plugin.state.dataState;
     }
 
-    private removeComponents(structures: StructureRef[]) {
-        const deletes = this.dateState.build();
+    private clearComponents(structures: StructureRef[]) {
+        const deletes = this.dataState.build();
         for (const s of structures) {
             for (const c of s.components) {
                 deletes.delete(c.cell.transform.ref);
@@ -43,7 +149,7 @@ class StructureComponentManager {
                 if (s.currentFocus.surroundings) deletes.delete(s.currentFocus.surroundings.cell.transform.ref);
             }
         }
-        return this.plugin.runTask(this.dateState.updateTree(deletes));
+        return this.plugin.runTask(this.dataState.updateTree(deletes));
     }
 
     constructor(public plugin: PluginContext) {
@@ -52,14 +158,62 @@ class StructureComponentManager {
 }
 
 namespace StructureComponentManager {
+    export type ActionType = 'add' | 'merge' | 'subtract' | 'color'
 
-    export function getModifyParams() {
-        return 0 as any;
+    const SelectionParam = PD.Select(StructureSelectionQueryOptions[1][0], StructureSelectionQueryOptions)
+
+    function getComponentsOptions(plugin: PluginContext, custom: [string, string][], label?: string) {
+        const types = [
+            ...custom,
+            ...plugin.managers.structure.hierarchy.componentGroups.map(g => [g[0].key!, g[0].cell.obj?.label])
+        ] as [string, string][];
+        return PD.Select(types[0][0], types, { label });
+    }
+
+    function getRepresentationTypes(plugin: PluginContext, pivot: StructureRef | undefined, custom: [string, string][], label?: string) {
+        const types = [
+            ...custom,
+            ...(pivot?.cell.obj?.data
+                ? plugin.structureRepresentation.registry.getApplicableTypes(pivot.cell.obj?.data!)
+                : plugin.structureRepresentation.registry.types)
+        ] as [string, string][];
+        return PD.Select(types[0][0], types, { label });
+    }
+
+    export function getActionParams(plugin: PluginContext, action: ActionType) {
+        switch (action) {
+            case 'add': 
+                return {
+                    kind: PD.Value<ActionType>(action, { isHidden: true }),
+                    selection: SelectionParam,
+                    label: PD.Text(''),
+                    representation: getRepresentationTypes(plugin, plugin.managers.structure.hierarchy.state.currentStructures[0], [['none', '< None >']])
+                };
+            case 'merge':
+            case 'subtract':
+                return {
+                    kind: PD.Value<ActionType>(action, { isHidden: true }),
+                    selection: SelectionParam,
+                    componentKey: getComponentsOptions(plugin, [['intersecting', '< Intersecting >']], 'Target')
+                };
+            case 'color':
+                // TODO: ability to reset
+                return {
+                    kind: PD.Value<ActionType>(action, { isHidden: true }),
+                    action: PD.MappedStatic('color', {
+                        color: PD.Color(ColorNames.black),
+                        reset: PD.EmptyGroup()
+                    }),
+                    // TODO: filter by representation type
+                    // representation: getRepresentationTypes(plugin, void 0, [['all', '< All >']])
+                };
+        }
     }
 
-    export type ModifyAction = 
-        | { kind: 'add', label: string, representationType?: string }
-        | { kind: 'merge', type: { kind: 'intersecting', key: string } | { kind: 'component', key: string } }
-        | { kind: 'subtract', type: { kind: 'all' } | { kind: 'component', key: string } }
-        | { kind: 'color', representationType?: string }
+    export type ModifyActionAdd = { kind: 'add', selection: StructureSelectionQuery, label: string, representation: string }
+    export type ModifyActionMerge = { kind: 'merge', selection: StructureSelectionQuery, componentKey: 'intersecting' | string }
+    export type ModifyActionSubtract = { kind: 'subtract', selection: StructureSelectionQuery, componentKey: 'intersecting' | string }
+    export type ModifyActionColor = { kind: 'color', action: { name: 'color', params: Color } | { name: 'reset', params: any } } //, representationType?: string }
+    
+    export type ModifyAction = ModifyActionAdd | ModifyActionMerge | ModifyActionSubtract | ModifyActionColor
 }

+ 70 - 39
src/mol-plugin-state/manager/structure/hierarchy.ts

@@ -5,28 +5,32 @@
  */
 
 import { PluginContext } from '../../../mol-plugin/context';
-import { StructureHierarchy, buildStructureHierarchy, ModelRef, StructureComponentRef } from './hierarchy-state';
+import { StructureHierarchy, buildStructureHierarchy, ModelRef, StructureComponentRef, StructureRef, HierarchyRef } from './hierarchy-state';
 import { PluginComponent } from '../../component';
 
 interface StructureHierarchyManagerState {
     hierarchy: StructureHierarchy,
     currentModels: ReadonlyArray<ModelRef>,
+    currentStructures: ReadonlyArray<StructureRef>
 }
 
 export class StructureHierarchyManager extends PluginComponent<StructureHierarchyManagerState> {
     readonly behaviors = {
-        hierarchy: this.ev.behavior(this.state.hierarchy),
-        currentModels: this.ev.behavior(this.state.currentModels)
+        current: this.ev.behavior({ hierarchy: this.state.hierarchy, models: this.state.currentModels, structures: this.state.currentStructures })
     }
 
-    private syncCurrent(hierarchy: StructureHierarchy) {
-        const current = this.behaviors.currentModels.value;
+    private _componentGroups: ReturnType<typeof StructureHierarchyManager['getComponentGroups']> | undefined = void 0;
+
+    get componentGroups() {
+        if (this._componentGroups) return this._componentGroups;
+        this._componentGroups = StructureHierarchyManager.getComponentGroups(this.state.currentStructures);
+        return this._componentGroups;
+    }
+
+    private syncCurrentModels(hierarchy: StructureHierarchy): ModelRef[] {
+        const current = this.state.currentModels;
         if (current.length === 0) {
-            const models = hierarchy.trajectories[0]?.models;
-            if (models) {
-                return models;
-            }
-            return [];
+            return hierarchy.trajectories[0]?.models || [];
         }
 
         const newCurrent: ModelRef[] = [];
@@ -36,30 +40,59 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
             newCurrent.push(ref);
         }
 
-        if (newCurrent.length === 0 && hierarchy.trajectories[0]?.models) {
-            return hierarchy.trajectories[0]?.models;
+        if (newCurrent.length === 0) {
+            return hierarchy.trajectories[0]?.models || [];
+        }
+
+        return newCurrent;
+    }
+
+    private syncCurrentStructures(hierarchy: StructureHierarchy, currentModels: ModelRef[]): StructureRef[] {
+        const current = this.state.currentStructures;
+        if (current.length === 0) {
+            return Array.prototype.concat.apply([], currentModels.map(m => m.structures));
+        }
+
+        const newCurrent: StructureRef[] = [];
+        for (const c of current) {
+            const ref = hierarchy.refs.get(c.cell.transform.ref) as StructureRef;
+            if (!ref) continue;
+            newCurrent.push(ref);
+        }
+
+        if (newCurrent.length === 0 && currentModels.length > 0) {
+            return Array.prototype.concat.apply([], currentModels.map(m => m.structures));
         }
 
         return newCurrent;
     }
 
     private sync() {
-        const update = buildStructureHierarchy(this.plugin.state.dataState, this.behaviors.hierarchy.value);
+        const update = buildStructureHierarchy(this.plugin.state.dataState, this.state.hierarchy);
         if (update.added.length === 0 && update.updated.length === 0 && update.removed.length === 0) {
             return;
         }
+        this._componentGroups = void 0;
 
-        const currentModels = this.syncCurrent(update.hierarchy);
-        this.updateState({ hierarchy: update.hierarchy, currentModels });
+        const currentModels = this.syncCurrentModels(update.hierarchy);
+        const currentStructures = this.syncCurrentStructures(update.hierarchy, currentModels);
+        this.updateState({ hierarchy: update.hierarchy, currentModels: currentModels, currentStructures: currentStructures });
 
-        this.behaviors.hierarchy.next(this.state.hierarchy);
-        this.behaviors.currentModels.next(this.state.currentModels);
+        this.behaviors.current.next({ hierarchy: update.hierarchy, models: currentModels, structures: currentStructures });
+    }
+
+    remove(refs: HierarchyRef[]) {
+        if (refs.length === 0) return;
+        const deletes = this.plugin.state.dataState.build();
+        for (const r of refs) deletes.delete(r.cell.transform.ref);
+        return this.plugin.runTask(this.plugin.state.dataState.updateTree(deletes));
     }
 
     constructor(private plugin: PluginContext) {
         super({
             hierarchy: StructureHierarchy(),
-            currentModels: []
+            currentModels: [],
+            currentStructures: []
         });
 
         plugin.state.dataState.events.changed.subscribe(e => {
@@ -74,30 +107,28 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
 }
 
 export namespace StructureHierarchyManager {
-    export function getCommonComponentPivots(models: ReadonlyArray<ModelRef>) {
-        if (!models[0]?.structures?.length) return [];
-        if (models[0]?.structures?.length === 1) return models[0]?.structures[0]?.components || [];
-
-        const pivots = new Map<string, StructureComponentRef>();
-
-        for (const c of models[0]?.structures[0]?.components) {
-            const key = c.key;
-            if (!key) continue;
-            pivots.set(key, c);
-        }
-
-        for (const m of models) {
-            for (const s of m.structures) {
-                for (const c of s.components) {
-                    const key = c.key;
-                    if (!key) continue;
-                    if (!pivots.has(key)) pivots.delete(key);
+    export function getComponentGroups(structures: ReadonlyArray<StructureRef>): StructureComponentRef[][] {
+        if (!structures.length) return [];
+        if (structures.length === 1) return structures[0].components.map(c => [c]);
+
+        const groups: StructureComponentRef[][] = [];
+        const map = new Map<string, StructureComponentRef[]>();
+
+        for (const s of structures) {
+            for (const c of s.components) {
+                const key = c.key;
+                if (!key) continue;
+
+                let component = map.get(key);
+                if (!component) {
+                    component = [];
+                    map.set(key, component);
+                    groups.push(component);
                 }
+                component.push(c);
             }
         }
 
-        const ret: StructureComponentRef[] = [];
-        pivots.forEach(function (this: StructureComponentRef[], p) { this.push(p) }, ret);
-        return ret;
+        return groups;
     }
 }

+ 0 - 2
src/mol-plugin-ui/controls.tsx

@@ -16,7 +16,6 @@ import { StateTransforms } from '../mol-plugin-state/transforms';
 import { StateTransformer } from '../mol-state';
 import { ModelFromTrajectory } from '../mol-plugin-state/transforms/model';
 import { AnimationControls } from './state/animation';
-import { StructureRepresentationControls } from './structure/representation';
 import { StructureSelectionControls } from './structure/selection';
 import { StructureMeasurementsControls } from './structure/measurements';
 import { Icon } from './controls/icons';
@@ -268,7 +267,6 @@ export class StructureToolsWrapper extends PluginUIComponent {
             <div className='msp-section-header'><Icon name='code' /> Structure Tools</div>
 
             <StructureSelectionControls />
-            <StructureRepresentationControls />
             <StructureComponentControls />
             <StructureMeasurementsControls />
         </div>;

+ 101 - 28
src/mol-plugin-ui/structure/components.tsx

@@ -6,7 +6,6 @@
 
 import * as React from 'react';
 import { CollapsableControls, CollapsableState, PurePluginUIComponent } from '../base';
-import { StructureHierarchyManager } from '../../mol-plugin-state/manager/structure/hierarchy';
 import { StructureComponentRef, StructureRepresentationRef } from '../../mol-plugin-state/manager/structure/hierarchy-state';
 import { State, StateAction } from '../../mol-state';
 import { PluginCommands } from '../../mol-plugin/commands';
@@ -16,6 +15,9 @@ import { ActionMenu } from '../controls/action-menu';
 import { ApplyActionControl } from '../state/apply-action';
 import { StateTransforms } from '../../mol-plugin-state/transforms';
 import { Icon } from '../controls/icons';
+import { StructureComponentManager } from '../../mol-plugin-state/manager/structure/component';
+import { ParamDefinition } from '../../mol-util/param-definition';
+import { ParameterControls } from '../controls/parameters';
 
 interface StructureComponentControlState extends CollapsableState {
     isDisabled: boolean
@@ -29,7 +31,7 @@ const MeasurementFocusOptions = {
 
 export class StructureComponentControls extends CollapsableControls<{}, StructureComponentControlState> {
     protected defaultState(): StructureComponentControlState {
-        return { header: 'Components', isCollapsed: false, isDisabled: false };
+        return { header: 'Representation', isCollapsed: false, isDisabled: false };
     }
 
     renderControls() {
@@ -50,6 +52,17 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
         isDisabled: false
     };
 
+    get current() {
+        return this.plugin.managers.structure.hierarchy.behaviors.current;
+    }
+
+    componentDidMount() {
+        this.subscribe(this.current, () => this.setState({ action: void 0 }));
+        this.subscribe(this.plugin.behaviors.state.isBusy, v => {
+            this.setState({ isDisabled: v, action: void 0 })
+        });
+    }
+
     private toggleAction(action: ComponentEditorControlsState['action']) {
         return () => this.setState({ action: this.state.action === action ? void 0 : action });
     }
@@ -99,6 +112,8 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
         else mng.component.applyPreset(this.plugin.managers.structure.hierarchy.state.currentModels[0].structures, item.value as any);
     }
 
+    modifyComponentControls = <div className='msp-control-offset'><ModifyComponentControls onApply={this.hideAction} /></div>
+
     render() {
         return <>
             <div className='msp-control-row msp-select-row'>
@@ -107,61 +122,115 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
                 <ToggleButton icon='cog' label='Options' toggle={this.toggleOptions} isSelected={this.state.action === 'options'} disabled={this.state.isDisabled} />
             </div>
             {this.state.action === 'preset' && this.presetControls}
+            {this.state.action === 'modify' && this.modifyComponentControls}
+            {this.state.action === 'options' && 'TODO'}
+        </>;
+    }
+}
+
+interface ModifyComponentControlsState {
+    action?: StructureComponentManager.ActionType,
+    actionParams?: ParamDefinition.Params,
+    actionParamValues?: StructureComponentManager.ModifyAction
+}
+
+class ModifyComponentControls extends PurePluginUIComponent<{ onApply: () => void }, ModifyComponentControlsState> {
+    state: ModifyComponentControlsState = { };
+
+    private toggleAction(action: StructureComponentManager.ActionType) {
+        return () => {
+            if (this.state.action === action) {
+                this.setState({ action: void 0, actionParams: void 0, actionParamValues: void 0 });
+            } else {
+                const actionParams = StructureComponentManager.getActionParams(this.plugin, action) as any;
+                const actionParamValues = ParamDefinition.getDefaultValues(actionParams) as StructureComponentManager.ModifyAction;
+                this.setState({ action, actionParams, actionParamValues });
+            }
+        }
+    }
+
+    toggleAdd = this.toggleAction('add');
+    toggleMerge = this.toggleAction('merge');
+    toggleSubtract = this.toggleAction('subtract');
+    toggleColor = this.toggleAction('color');
+
+    hideAction = () => this.setState({ action: void 0 });
+
+    apply = () => {
+        this.plugin.managers.structure.component.modify(this.state.actionParamValues!);
+        this.props.onApply();
+    }
+
+    paramsChanged = (actionParamValues: any) => this.setState({ actionParamValues })
+    get paramControls() {
+        if (!this.state.action) return null;
+        return <>
+            <ParameterControls params={this.state.actionParams!} values={this.state.actionParamValues!} onChangeObject={this.paramsChanged} />
+            <button className={`msp-btn msp-btn-block msp-btn-commit msp-btn-commit-on`} onClick={this.apply} style={{ marginTop: '1px' }}>
+                <Icon name='ok' /> Apply
+            </button>
+        </>
+    }
+
+    render() {
+        return <>
+            <div className='msp-control-row msp-select-row'>
+                <ToggleButton icon='plus' label='Add' toggle={this.toggleAdd} isSelected={this.state.action === 'add'} />
+                <ToggleButton icon='flow-branch' label='Merge' toggle={this.toggleMerge} isSelected={this.state.action === 'merge'} />
+                <ToggleButton icon='minus' label='Sub' toggle={this.toggleSubtract} isSelected={this.state.action === 'subtract'} />
+                <ToggleButton icon='brush' label='Color' toggle={this.toggleColor} isSelected={this.state.action === 'color'} />
+            </div>
+            {this.paramControls}
         </>;
     }
 }
 
 class ComponentListControls extends PurePluginUIComponent {
-    get currentModels() {
-        return this.plugin.managers.structure.hierarchy.behaviors.currentModels;
+    get current() {
+        return this.plugin.managers.structure.hierarchy.behaviors.current;
     }
 
     componentDidMount() {
-        this.subscribe(this.currentModels, () => this.forceUpdate());
+        this.subscribe(this.current, () => this.forceUpdate());
     }
 
     render() {
-        const components = StructureHierarchyManager.getCommonComponentPivots(this.currentModels.value)
-        return components.map(c => <StructureComponentEntry key={c.cell.transform.ref} component={c} />)
+        const componentGroups = this.plugin.managers.structure.hierarchy.componentGroups;
+        return componentGroups.map(g => <StructureComponentGroup key={g[0].cell.transform.ref} group={g} />)
     }
 }
 
 type StructureComponentEntryActions = 'add-repr' | 'remove' | 'none'
 
 const createRepr = StateAction.fromTransformer(StateTransforms.Representation.StructureRepresentation3D);
-class StructureComponentEntry extends PurePluginUIComponent<{ component: StructureComponentRef }, { action: StructureComponentEntryActions }> {
+class StructureComponentGroup extends PurePluginUIComponent<{ group: StructureComponentRef[] }, { action: StructureComponentEntryActions }> {
     state = { action: 'none' as StructureComponentEntryActions }
 
-    get ref() {
-        return this.props.component.cell.transform.ref;
+    get pivot() {
+        return this.props.group[0];
     }
 
     componentDidMount() {
         this.subscribe(this.plugin.events.state.cell.stateUpdated, e => {
-            if (State.ObjectEvent.isCell(e, this.props.component.cell)) this.forceUpdate();
+            if (State.ObjectEvent.isCell(e, this.pivot.cell)) this.forceUpdate();
         });
     }
 
     toggleVisible = (e: React.MouseEvent<HTMLElement>) => {
         e.preventDefault();
-        PluginCommands.State.ToggleVisibility(this.plugin, { state: this.props.component.cell.parent, ref: this.ref });
-        e.currentTarget.blur();
-    }
-
-    remove(ref: string) {
-        return () => {
-            this.setState({ action: 'none' });
-            PluginCommands.State.RemoveObject(this.plugin, { state: this.props.component.cell.parent, ref, removeParentGhosts: true });
+        // TODO: check visibility beforehand to set correct value if user made individual change
+        for (const c of this.props.group) {
+            PluginCommands.State.ToggleVisibility(this.plugin, { state: c.cell.parent, ref: c.cell.transform.ref });
         }
+        e.currentTarget.blur();
     }
 
-    
     get removeActions(): ActionMenu.Items {
         const ret = [
-            ActionMenu.Item('Remove Selection', 'remove', this.remove(this.ref))
+            ActionMenu.Item('Remove Selection', 'remove', () => this.plugin.managers.structure.hierarchy.remove(this.props.group))
         ];
-        for (const repr of this.props.component.representations) {
-            ret.push(ActionMenu.Item(`Remove ${repr.cell.obj?.label}`, 'remove', this.remove(repr.cell.transform.ref)))
+        for (const repr of this.pivot.representations) {
+            ret.push(ActionMenu.Item(`Remove ${repr.cell.obj?.label}`, 'remove', () => this.plugin.managers.structure.component.removeRepresentations(this.props.group, repr)))
         }
         return ret;
     }
@@ -176,16 +245,20 @@ class StructureComponentEntry extends PurePluginUIComponent<{ component: Structu
 
     highlight = (e: React.MouseEvent<HTMLElement>) => {
         e.preventDefault();
-        PluginCommands.State.Highlight(this.plugin, { state: this.props.component.cell.parent, ref: this.ref });
+        for (const c of this.props.group) {
+            PluginCommands.State.Highlight(this.plugin, { state: c.cell.parent, ref: c.cell.transform.ref });
+        }
     }
 
     clearHighlight = (e: React.MouseEvent<HTMLElement>) => {
         e.preventDefault();
-        PluginCommands.State.ClearHighlight(this.plugin, { state: this.props.component.cell.parent, ref: this.ref });
+        for (const c of this.props.group) {
+            PluginCommands.State.ClearHighlight(this.plugin, { state: c.cell.parent, ref: c.cell.transform.ref });
+        }
     }
 
     focus = () => {
-        const sphere = this.props.component.cell.obj?.data.boundary.sphere;
+        const sphere = this.pivot.cell.obj?.data.boundary.sphere;
         if (sphere) {
             const { extraRadius, minRadius, durationMs } = MeasurementFocusOptions;
             const radius = Math.max(sphere.radius + extraRadius, minRadius);
@@ -194,7 +267,7 @@ class StructureComponentEntry extends PurePluginUIComponent<{ component: Structu
     }
 
     render() {
-        const component = this.props.component;
+        const component = this.pivot;
         const cell = component.cell;
         const label = cell.obj?.label;
         return <>
@@ -212,7 +285,7 @@ class StructureComponentEntry extends PurePluginUIComponent<{ component: Structu
             <div className='msp-control-offset'>
                 {this.state.action === 'add-repr' && 
                 <ControlGroup header='Add Representation' initialExpanded={true} hideExpander={true} hideOffset={true} onHeaderClick={this.toggleAddRepr} topRightIcon='off'>
-                    <ApplyActionControl plugin={this.plugin} state={cell.parent} action={createRepr} nodeRef={this.ref} hideHeader noMargin onApply={this.toggleAddRepr} applyLabel='Add' />
+                    <ApplyActionControl plugin={this.plugin} state={cell.parent} action={createRepr} nodeRef={component.cell.transform.ref} hideHeader noMargin onApply={this.toggleAddRepr} applyLabel='Add' />
                 </ControlGroup>}
                 {component.representations.map(r => <StructureRepresentationEntry key={r.cell.transform.ref} representation={r} />)}
             </div>

+ 1 - 1
src/mol-repr/structure/visual/label-text.ts

@@ -25,7 +25,7 @@ export const LabelTextParams = {
     backgroundMargin: PD.Numeric(0, { min: 0, max: 1, step: 0.01 }),
     backgroundColor: PD.Color(ColorNames.black),
     backgroundOpacity: PD.Numeric(0.5, { min: 0, max: 1, step: 0.01 }),
-    level: PD.Select('residue', [['chain', 'Chain'], ['residue', 'Residue'], ['element', 'Element']]),
+    level: PD.Select('residue', [['chain', 'Chain'], ['residue', 'Residue'], ['element', 'Element']] as const),
     chainScale: PD.Numeric(10, { min: 0, max: 20, step: 0.1 }),
     residueScale: PD.Numeric(1, { min: 0, max: 20, step: 0.1 }),
     elementScale: PD.Numeric(0.5, { min: 0, max: 20, step: 0.1 }),

+ 1 - 1
src/mol-state/object.ts

@@ -194,7 +194,7 @@ export namespace StateObjectSelector {
 export type StateObjectRef<S extends StateObject = StateObject> = StateObjectSelector<S> | StateObjectCell<S> | StateTransform.Ref
 
 export namespace StateObjectRef {
-    export function resolveRef<S extends StateObject>(state: State, ref?: StateObjectRef<S>): StateTransform.Ref | undefined {
+    export function resolveRef<S extends StateObject>(ref?: StateObjectRef<S>): StateTransform.Ref | undefined {
         if (!ref) return;
         if (typeof ref === 'string') return ref;
         if (StateObjectCell.is(ref)) return ref.transform.ref;

+ 1 - 1
src/mol-theme/color/hydrophobicity.ts

@@ -17,7 +17,7 @@ const Description = 'Assigns a color to every amino acid according to the "Exper
 
 export const HydrophobicityColorThemeParams = {
     list: PD.ColorList<ColorListName>('red-yellow-green', ColorListOptionsScale),
-    scale: PD.Select('DGwif', [['DGwif', 'DG water-membrane'], ['DGwoct', 'DG water-octanol'], ['Oct-IF', 'DG difference']])
+    scale: PD.Select('DGwif', [['DGwif', 'DG water-membrane'], ['DGwoct', 'DG water-octanol'], ['Oct-IF', 'DG difference']] as const)
 }
 export type HydrophobicityColorThemeParams = typeof HydrophobicityColorThemeParams
 export function getHydrophobicityColorThemeParams(ctx: ThemeDataContext) {

+ 2 - 2
src/mol-util/param-definition.ts

@@ -70,13 +70,13 @@ export namespace ParamDefinition {
         return setInfo<Value<T>>({ type: 'value', defaultValue }, info);
     }
 
-    export interface Select<T extends string | number> extends Base<T> {
+    export interface Select<T> extends Base<T> {
         type: 'select'
         /** array of (value, label) tuples */
         options: readonly (readonly [T, string] | readonly [T, string, string])[]
         cycle?: boolean
     }
-    export function Select<T extends string | number>(defaultValue: T, options: readonly (readonly [T, string] | readonly [T, string, string])[], info?: Info & { cycle?: boolean }): Select<T> {
+    export function Select<T>(defaultValue: T, options: readonly (readonly [T, string] | readonly [T, string, string])[], info?: Info & { cycle?: boolean }): Select<T> {
         return setInfo<Select<T>>({ type: 'select', defaultValue: checkDefaultKey(defaultValue, options), options, cycle: info?.cycle }, info)
     }