wip: source structure manager

src/mol-model-formats/structure/basic/parser.ts

@@ -24,9 +24,16 @@ import { getModelGroupName } from './util';
 
 export async function createModels(data: BasicData, format: ModelFormat, ctx: RuntimeContext) {
     const properties = getProperties(data)
-    return data.ihm_model_list._rowCount > 0
+    const models = data.ihm_model_list._rowCount > 0
         ? await readIntegrative(ctx, data, properties, format)
         : await readStandard(ctx, data, properties, format);
+
+    for (let i = 0; i < models.length; i++) {
+        models[i].trajectoryInfo.index = i;
+        models[i].trajectoryInfo.size = models.length;
+    }
+
+    return models;
 }
 
 /** Standard atomic model */
@@ -62,6 +69,7 @@ function createStandardModel(data: BasicData, atom_site: AtomSite, sourceIndex:
         entry,
         sourceData: format,
         modelNum,
+        trajectoryInfo: { index: 0, size: 1 },
         entities,
         sequence,
         atomicHierarchy: atomic.hierarchy,
@@ -99,6 +107,7 @@ function createIntegrativeModel(data: BasicData, ihm: CoarseData, properties: Mo
         entry,
         sourceData: format,
         modelNum: ihm.model_id,
+        trajectoryInfo: { index: 0, size: 1 },
         entities: ihm.entities,
         sequence,
         atomicHierarchy: atomic.hierarchy,

src/mol-model/structure/model/model.ts

@@ -43,6 +43,14 @@ export interface Model extends Readonly<{
      */
     modelNum: number,
 
+    /**
+     * This is a hack to allow "model-index coloring"
+     */
+    trajectoryInfo: {
+        index: number,
+        size: number
+    },
+
     sourceData: ModelFormat,
 
     entities: Entities,
@@ -109,8 +117,12 @@ export namespace Model {
             const trajectory = trajectoryFromModelAndCoordinates(model, coordinates)
             const bondData = { pairs: topology.bonds, count: model.atomicHierarchy.atoms._rowCount }
             const indexPairBonds = IndexPairBonds.fromData(bondData)
+
+            let index = 0;
             for (const m of trajectory) {
                 IndexPairBonds.Provider.set(m, indexPairBonds)
+                m.trajectoryInfo.index = index++;
+                m.trajectoryInfo.size = trajectory.length;
             }
             return trajectory
         })

src/mol-plugin-state/manager/structure/component.ts

@@ -39,7 +39,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
     }
 
     get currentStructures() {
-        return this.plugin.managers.structure.hierarchy.state.currentStructures;
+        return this.plugin.managers.structure.hierarchy.state.current.structures;
     }
 
     get pivotStructure(): StructureRef | undefined {
@@ -299,7 +299,7 @@ namespace StructureComponentManager {
     export function getAddParams(plugin: PluginContext) {
         return {
             selection: SelectionParam,
-            representation: getRepresentationTypesSelect(plugin, plugin.managers.structure.hierarchy.state.currentStructures[0], [['none', '< None >']]),
+            representation: getRepresentationTypesSelect(plugin, plugin.managers.structure.component.pivotStructure, [['none', '< None >']]),
             label: PD.Text('')
         };
     }

src/mol-plugin-state/manager/structure/hierarchy.ts

@@ -7,39 +7,56 @@
 import { PluginContext } from '../../../mol-plugin/context';
 import { StructureHierarchy, buildStructureHierarchy, ModelRef, StructureComponentRef, StructureRef, HierarchyRef, TrajectoryRef } from './hierarchy-state';
 import { PluginComponent } from '../../component';
+import { StateTransform } from '../../../mol-state';
 
 interface StructureHierarchyManagerState {
     hierarchy: StructureHierarchy,
-    currentTrajectories: ReadonlyArray<TrajectoryRef>,
-    currentModels: ReadonlyArray<ModelRef>,
-    currentStructures: ReadonlyArray<StructureRef>
+    current: {
+        trajectories: ReadonlyArray<TrajectoryRef>,
+        models: ReadonlyArray<ModelRef>,
+        structures: ReadonlyArray<StructureRef>
+    }
 }
 
 export class StructureHierarchyManager extends PluginComponent<StructureHierarchyManagerState> {
     readonly behaviors = {
         current: this.ev.behavior({
             hierarchy: this.state.hierarchy,
-            trajectories: this.state.currentTrajectories,
-            models: this.state.currentModels,
-            structures: this.state.currentStructures
+            trajectories: this.state.current.trajectories,
+            models: this.state.current.models,
+            structures: this.state.current.structures
         })
     }
 
+    private get dataState() {
+        return this.plugin.state.dataState;
+    }
+
     private _currentComponentGroups: ReturnType<typeof StructureHierarchyManager['getComponentGroups']> | undefined = void 0;
 
     get currentComponentGroups() {
         if (this._currentComponentGroups) return this._currentComponentGroups;
-        this._currentComponentGroups = StructureHierarchyManager.getComponentGroups(this.state.currentStructures);
+        this._currentComponentGroups = StructureHierarchyManager.getComponentGroups(this.state.current.structures);
         return this._currentComponentGroups;
     }
 
-    private syncCurrentTrajectories(hierarchy: StructureHierarchy): TrajectoryRef[] {
-        const current = this.state.currentTrajectories;
+    private _currentSelectionSet: Set<string> | undefined = void 0;
+    get currentSeletionSet() {
+        if (this._currentSelectionSet) return this._currentSelectionSet;
+        this._currentSelectionSet = new Set();
+        for (const r of this.state.current.trajectories) this._currentSelectionSet.add(r.cell.transform.ref);
+        for (const r of this.state.current.models) this._currentSelectionSet.add(r.cell.transform.ref);
+        for (const r of this.state.current.structures) this._currentSelectionSet.add(r.cell.transform.ref);
+        return this._currentSelectionSet;
+    }
+
+    private syncCurrentTrajectories(hierarchy: StructureHierarchy, map: Map<StateTransform.Ref, HierarchyRef>): TrajectoryRef[] {
+        const current = this.state.current.trajectories;
         if (current.length === 0) return hierarchy.trajectories.length > 0 ? [hierarchy.trajectories[0]] : [];
 
         const newCurrent: TrajectoryRef[] = [];
         for (const c of current) {
-            const ref = hierarchy.refs.get(c.cell.transform.ref) as TrajectoryRef;
+            const ref = map.get(c.cell.transform.ref) as TrajectoryRef;
             if (ref) newCurrent.push(ref);
         }
 
@@ -47,13 +64,13 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         return newCurrent;
     }
 
-    private syncCurrentModels(hierarchy: StructureHierarchy, currentTrajectories: TrajectoryRef[]): ModelRef[] {
-        const current = this.state.currentModels;
+    private syncCurrentModels(hierarchy: StructureHierarchy, map: Map<StateTransform.Ref, HierarchyRef>, currentTrajectories: TrajectoryRef[]): ModelRef[] {
+        const current = this.state.current.models;
         if (current.length === 0) return currentTrajectories[0]?.models || [];
 
         const newCurrent: ModelRef[] = [];
         for (const c of current) {
-            const ref = hierarchy.refs.get(c.cell.transform.ref) as ModelRef;
+            const ref = map.get(c.cell.transform.ref) as ModelRef;
             if (ref) newCurrent.push(ref);
         }
 
@@ -61,13 +78,13 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         return newCurrent;
     }
 
-    private syncCurrentStructures(hierarchy: StructureHierarchy, currentModels: ModelRef[]): StructureRef[] {
-        const current = this.state.currentStructures;
+    private syncCurrentStructures(map: Map<StateTransform.Ref, HierarchyRef>, currentModels: ModelRef[]): StructureRef[] {
+        const current = this.state.current.structures;
         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;
+            const ref = map.get(c.cell.transform.ref) as StructureRef;
             if (ref) newCurrent.push(ref);
         }
 
@@ -81,19 +98,40 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
             return;
         }
         this._currentComponentGroups = void 0;
+        this._currentSelectionSet = void 0;
 
-        const currentTrajectories = this.syncCurrentTrajectories(update.hierarchy);
-        const currentModels = this.syncCurrentModels(update.hierarchy, currentTrajectories);
-        const currentStructures = this.syncCurrentStructures(update.hierarchy, currentModels);
-        console.log(currentTrajectories, currentModels, currentStructures);
-        this.updateState({ hierarchy: update.hierarchy, currentModels, currentStructures, currentTrajectories });
-
-        this.behaviors.current.next({
-            hierarchy: update.hierarchy,
-            trajectories: currentTrajectories,
-            models: currentModels,
-            structures: currentStructures
-        });
+        const trajectories = this.syncCurrentTrajectories(update.hierarchy, update.hierarchy.refs);
+        const models = this.syncCurrentModels(update.hierarchy, update.hierarchy.refs, trajectories);
+        const structures = this.syncCurrentStructures(update.hierarchy.refs, models);
+
+        this.updateState({ hierarchy: update.hierarchy, current: { trajectories, models, structures }});
+        this.behaviors.current.next({ hierarchy: update.hierarchy, trajectories, models, structures });
+    }
+
+    updateCurrent(refs: HierarchyRef[], action: 'add' | 'remove') {
+
+        console.log(refs, action);
+
+        const hierarchy = this.state.hierarchy;
+        const map = new Map<StateTransform.Ref, HierarchyRef>();
+        const set = this.currentSeletionSet;
+        if (action === 'add') {
+            set.forEach(r => map.set(r, hierarchy.refs.get(r)!))
+            for (const r of refs) map.set(r.cell.transform.ref, r);
+        } else {
+            set.forEach(r => map.set(r, hierarchy.refs.get(r)!))
+            for (const r of refs) map.delete(r.cell.transform.ref);
+        }
+     
+        const trajectories = this.syncCurrentTrajectories(hierarchy, map);
+        const models = this.syncCurrentModels(hierarchy, map, trajectories);
+        const structures = this.syncCurrentStructures(map, models);
+
+        this.updateState({ current: { trajectories, models, structures }});
+
+        console.log(this.state.current);
+
+        this.behaviors.current.next({ hierarchy, trajectories, models, structures });
     }
 
     remove(refs: HierarchyRef[]) {
@@ -103,12 +141,33 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         return this.plugin.runTask(this.plugin.state.dataState.updateTree(deletes));
     }
 
+    createAllModels(trajectory: TrajectoryRef) {
+        return this.plugin.dataTransaction(async () => {
+            if (trajectory.models.length > 0) {
+                await this.clearTrajectory(trajectory);
+            }
+
+            const tr = trajectory.cell.obj?.data!;
+            for (let i = 0; i < tr.length; i++) {
+                const model = await this.plugin.builders.structure.createModel(trajectory.cell, { modelIndex: i });
+                const structure = await this.plugin.builders.structure.createStructure(model, { name: 'deposited', params: { } });
+                await this.plugin.builders.structure.representation.structurePreset(structure, 'auto');
+            }
+        })
+    }
+
+    private clearTrajectory(trajectory: TrajectoryRef) {
+        const builder = this.dataState.build();
+        for (const m of trajectory.models) {
+            builder.delete(m.cell);
+        }
+        return this.plugin.runTask(this.dataState.updateTree(builder));
+    }
+
     constructor(private plugin: PluginContext) {
         super({
             hierarchy: StructureHierarchy(),
-            currentTrajectories: [],
-            currentModels: [],
-            currentStructures: []
+            current: { trajectories: [], models: [], structures: [] }
         });
 
         plugin.state.dataState.events.changed.subscribe(e => {

src/mol-plugin-state/manager/structure/selection.ts

@@ -339,7 +339,7 @@ export class StructureSelectionManager extends PluginComponent<StructureSelectio
     }
 
     private get applicableStructures() {
-        return this.plugin.managers.structure.hierarchy.state.currentStructures
+        return this.plugin.managers.structure.hierarchy.state.current.structures
             .filter(s => !!s.cell.obj)
             .map(s => s.cell.obj!.data);
     }

src/mol-plugin-ui/controls.tsx

@@ -20,6 +20,7 @@ import { StructureSelectionControls } from './structure/selection';
 import { StructureMeasurementsControls } from './structure/measurements';
 import { Icon } from './controls/icons';
 import { StructureComponentControls } from './structure/components';
+import { StructureSourceControls } from './structure/source';
 
 export class TrajectoryViewportControls extends PluginUIComponent<{}, { show: boolean, label: string }> {
     state = { show: false, label: '' }
@@ -266,6 +267,7 @@ export class StructureToolsWrapper extends PluginUIComponent {
         return <div>
             <div className='msp-section-header'><Icon name='code' /> Structure Tools</div>
 
+            <StructureSourceControls />
             <StructureSelectionControls />
             <StructureComponentControls />
             <StructureMeasurementsControls />

src/mol-plugin-ui/controls/action-menu.tsx

@@ -21,18 +21,19 @@ export class ActionMenu extends React.PureComponent<ActionMenu.Props> {
                     <b>{cmd.header}</b>
                 </button>
             </div>}
-            <Section items={cmd.items} onSelect={cmd.onSelect} current={cmd.current} />
+            <Section items={cmd.items} onSelect={cmd.onSelect} current={cmd.current} multiselect={this.props.multiselect} />
         </div>
     }
 }
 
 export namespace ActionMenu {
-    export type Props = { items: Items, onSelect: OnSelect, header?: string, current?: Item | undefined }
+    export type Props = { items: Items, onSelect: OnSelect | OnSelectMany, header?: string, current?: Item, multiselect?: boolean }
 
     export type OnSelect = (item: Item | undefined) => void
+    export type OnSelectMany = (itemOrItems: Item[] | undefined) => void
 
     export type Items =  string | Item | Items[]
-    export type Item = { label: string, icon?: IconName, disabled?: boolean, value: unknown }
+    export type Item = { label: string, icon?: IconName, disabled?: boolean, selected?: boolean, value: unknown }
 
     export function Item(label: string, value: unknown): Item
     export function Item(label: string, icon: string, value: unknown): Item
@@ -78,6 +79,16 @@ export namespace ActionMenu {
         return createItems(param.options, _selectOptions);
     }
 
+    export function hasSelectedItem(items: Items): boolean {
+        if (typeof items === 'string') return false;
+        if (isItem(items)) return !!items.selected;
+        for (const s of items) {
+            const found = hasSelectedItem(s);
+            if (found) return true;
+        }
+        return false;
+    }
+
     export function findItem(items: Items, value: any): Item | undefined {
         if (typeof items === 'string') return;
         if (isItem(items)) return items.value === value ? items : void 0;
@@ -97,14 +108,20 @@ export namespace ActionMenu {
     }
 }
 
-type SectionProps = { header?: string, items: ActionMenu.Items, onSelect: ActionMenu.OnSelect, current: ActionMenu.Item | undefined }
+type SectionProps = { header?: string, items: ActionMenu.Items, onSelect: ActionMenu.OnSelect | ActionMenu.OnSelectMany, current: ActionMenu.Item | undefined, multiselect: boolean | undefined }
 type SectionState = { items: ActionMenu.Items, current: ActionMenu.Item | undefined, isExpanded: boolean }
 
 class Section extends React.PureComponent<SectionProps, SectionState> {
     state = {
         items: this.props.items,
         current: this.props.current,
-        isExpanded: !!this.props.current && !!ActionMenu.findItem(this.props.items, this.props.current.value)
+        isExpanded: this.hasCurrent
+    }
+
+    get hasCurrent() {
+        return this.props.multiselect
+            ? ActionMenu.hasSelectedItem(this.props.items)
+            : !!this.props.current && !!ActionMenu.findItem(this.props.items, this.props.current.value);
     }
 
     toggleExpanded = (e: React.MouseEvent<HTMLButtonElement>) => {
@@ -114,41 +131,85 @@ class Section extends React.PureComponent<SectionProps, SectionState> {
 
     static getDerivedStateFromProps(props: SectionProps, state: SectionState) {
         if (props.items === state.items && props.current === state.current) return null;
-        return { items: props.items, current: props.current, isExpanded: props.current && !!ActionMenu.findItem(props.items, props.current.value) }
+        return {
+            items: props.items,
+            current: props.current,
+            isExpanded: props.multiselect
+                ? ActionMenu.hasSelectedItem(props.items)
+                : props.current && !!ActionMenu.findItem(props.items, props.current.value) }
+    }
+
+    selectAll = () => {
+        const items = collectItems(this.props.items, []).filter(i => !i.selected);
+        this.props.onSelect(items as any);
+    }
+
+    selectNone = () => {
+        const items = collectItems(this.props.items, []).filter(i => !!i.selected);
+        this.props.onSelect(items as any);
+    }
+
+    get multiselectHeader() {
+        const { header } = this.props;
+        const hasCurrent = this.hasCurrent;
+
+        return <div className='msp-control-group-header msp-flex-row' style={{ marginTop: '1px' }}>
+            <button className='msp-btn msp-form-control msp-flex-item' onClick={this.toggleExpanded}>
+                <Icon name={this.state.isExpanded ? 'collapse' : 'expand'} />
+                {hasCurrent ? <b>{header}</b> : header}
+            </button>
+            <button className='msp-btn msp-form-control msp-flex-item' onClick={this.selectAll} style={{ flex: '0 0 50px', textAlign: 'right' }}>
+                <Icon name='check' />
+                All
+            </button>
+            <button className='msp-btn msp-form-control msp-flex-item' onClick={this.selectNone} style={{ flex: '0 0 50px', textAlign: 'right' }}>
+                <Icon name='cancel' />
+                None
+            </button>
+        </div>;
+    }
+
+    get basicHeader() {
+        const { header } = this.props;
+        const hasCurrent = this.hasCurrent;
+
+        return <div className='msp-control-group-header' style={{ marginTop: '1px' }}>
+            <button className='msp-btn msp-btn-block msp-form-control' onClick={this.toggleExpanded}>
+                <Icon name={this.state.isExpanded ? 'collapse' : 'expand'} />
+                {hasCurrent ? <b>{header}</b> : header}
+            </button>
+        </div>;
     }
 
     render() {
         const { header, items, onSelect, current } = this.props;
 
         if (typeof items === 'string') return null;
-        if (isItem(items)) return <Action item={items} onSelect={onSelect} current={current} />
-
-        const hasCurrent = header && current && !!ActionMenu.findItem(items, current.value)
+        if (isItem(items)) return <Action item={items} onSelect={onSelect} current={current} multiselect={this.props.multiselect} />
 
         return <div>
-            {header && <div className='msp-control-group-header' style={{ marginTop: '1px' }}>
-                <button className='msp-btn msp-btn-block' onClick={this.toggleExpanded}>
-                    <Icon name={this.state.isExpanded ? 'collapse' : 'expand'} />
-                    {hasCurrent ? <b>{header}</b> : header}
-                </button>
-            </div>}
+            {header && (this.props.multiselect && this.state.isExpanded ? this.multiselectHeader : this.basicHeader)}
             <div className='msp-control-offset'>
                 {(!header || this.state.isExpanded) && items.map((x, i) => {
                     if (typeof x === 'string') return null;
-                    if (isItem(x)) return <Action key={i} item={x} onSelect={onSelect} current={current} />
-                    return <Section key={i} header={typeof x[0] === 'string' ? x[0] : void 0} items={x} onSelect={onSelect} current={current} />
+                    if (isItem(x)) return <Action key={i} item={x} onSelect={onSelect} current={current} multiselect={this.props.multiselect} />
+                    return <Section key={i} header={typeof x[0] === 'string' ? x[0] : void 0} items={x} onSelect={onSelect} current={current} multiselect={this.props.multiselect} />
                 })}
             </div>
         </div>;
     }
 }
 
-const Action: React.FC<{ item: ActionMenu.Item, onSelect: ActionMenu.OnSelect, current: ActionMenu.Item | undefined }> = ({ item, onSelect, current }) => {
+const Action: React.FC<{
+    item: ActionMenu.Item, 
+    onSelect: ActionMenu.OnSelect | ActionMenu.OnSelectMany, 
+    multiselect: boolean | undefined, 
+    current: ActionMenu.Item | undefined }> = ({ item, onSelect, current, multiselect }) => {
     const isCurrent = current === item;
     return <div className='msp-control-row'>
-        <button onClick={() => onSelect(item)} disabled={item.disabled}>
+        <button onClick={() => onSelect(multiselect ? [item] : item as any)} disabled={item.disabled}>
             {item.icon && <Icon name={item.icon} style={{ fontSize: '80%', marginRight: '6px' }} />}
-            {isCurrent ? <b>{item.label}</b> : item.label}
+            {isCurrent || item.selected ? <b>{item.label}</b> : item.label}
         </button>
     </div>;
 }
@@ -156,4 +217,16 @@ const Action: React.FC<{ item: ActionMenu.Item, onSelect: ActionMenu.OnSelect, c
 function isItem(x: any): x is ActionMenu.Item {
     const v = x as ActionMenu.Item;
     return v && !!v.label && typeof v.value !== 'undefined';
+}
+
+function collectItems(items: ActionMenu.Items, target: ActionMenu.Item[]) {
+    if (typeof items === 'string') return target;
+    if (isItem(items)) {
+        target.push(items);
+        return target;
+    }
+    for (const i of items) {
+        collectItems(i, target);
+    }
+    return target;
 }

src/mol-plugin-ui/skin/base/components/controls.scss

@@ -242,9 +242,9 @@
     background: $default-background;
     > button {
         padding-left: $control-spacing / 2 !important;
-        text-align: left !important;
-        height: 2 * $row-height / 3 !important;
-        line-height: 2 * $row-height / 3 !important;
+        text-align: left;
+        height: 22px !important; // 2 * $row-height / 3 !important;
+        line-height: 22px !important; // 2 * $row-height / 3 !important;
         font-size: 70% !important;
         background: $default-background !important;
         color: color-lower-contrast($font-color, 15%) !important;

src/mol-plugin-ui/skin/base/components/temp.scss

@@ -82,6 +82,23 @@
     }
 }
 
+.msp-flex-row {
+    display:flex;
+    flex-direction:row;
+    width: inherit;
+
+    > .msp-flex-item {
+        margin: 0;
+        flex: 1 1 auto;
+        margin-right: 1px;
+        overflow: hidden;
+    }
+
+    > .msp-flex-item:last-child {
+        margin-right: 0;
+    }
+}
+
 .msp-state-list {
     list-style: none;
     margin-top: $control-spacing;

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

@@ -97,7 +97,7 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
         if (!item) return;
         const mng = this.plugin.managers.structure;
 
-        const structures = mng.hierarchy.state.currentStructures;
+        const structures = mng.hierarchy.state.current.structures;
         if (item.value === null) mng.component.clear(structures);
         else mng.component.applyPreset(structures, item.value as any);
     }
@@ -257,8 +257,13 @@ class StructureComponentGroup extends PurePluginUIComponent<{ group: StructureCo
 
     highlight = (e: React.MouseEvent<HTMLElement>) => {
         e.preventDefault();
+
+        PluginCommands.State.Highlight(this.plugin, { state: this.props.group[0].cell.parent, ref: this.props.group[0].cell.transform.ref });
+
+        let first = true;
         for (const c of this.props.group) {
-            PluginCommands.State.Highlight(this.plugin, { state: c.cell.parent, ref: c.cell.transform.ref });
+            if (first) { first = false; continue; }
+            PluginCommands.State.Highlight(this.plugin, { state: c.cell.parent, ref: c.cell.transform.ref, extend: true });
         }
     }
 

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

@@ -0,0 +1,60 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import * as React from 'react';
+import { HierarchyRef, ModelRef, TrajectoryRef } from '../../mol-plugin-state/manager/structure/hierarchy-state';
+import { CollapsableControls, CollapsableState } from '../base';
+import { ActionMenu } from '../controls/action-menu';
+
+interface StructureSourceControlState extends CollapsableState {
+    isBusy: boolean
+}
+
+export class StructureSourceControls extends CollapsableControls<{}, StructureSourceControlState> {
+    protected defaultState(): StructureSourceControlState {
+        return { header: 'Source', isCollapsed: false, isBusy: false };
+    }
+
+    componentDidMount() {
+        this.subscribe(this.plugin.managers.structure.hierarchy.behaviors.current, () => this.forceUpdate());
+        this.subscribe(this.plugin.behaviors.state.isBusy, v => {
+            this.setState({ isBusy: v })
+        });
+    }
+
+    private item = (ref: HierarchyRef) => {
+        const selected = this.plugin.managers.structure.hierarchy.currentSeletionSet;
+        return { label: ref.cell.obj?.label, selected: selected.has(ref.cell.transform.ref), value: ref } as ActionMenu.Item;
+    }
+
+    getTrajectoryItems = (t: TrajectoryRef): ActionMenu.Items => {
+        if (t.models.length === 0) return this.item(t);
+        // if (t.models.length === 1) return this.getModelItems(t.models[0]);
+        return [t.cell.obj?.label!, ...t.models.map(this.getModelItems)];
+    }
+
+    private getModelItems = (m: ModelRef): ActionMenu.Items => {
+        if (m.structures.length === 0) return this.item(m);
+        // if (m.structures.length === 1) return this.item(m.structures[0]);
+        return [m.cell.obj?.label!, ...m.structures.map(this.item)];
+    }
+
+    get hierarchyItems() {
+        return this.plugin.managers.structure.hierarchy.state.current.trajectories.map(this.getTrajectoryItems);
+    }
+
+    onSelect: ActionMenu.OnSelectMany = (items) => {
+        if (!items || items.length === 0) return 0;
+        this.plugin.managers.structure.hierarchy.updateCurrent(items.map(i => i.value as HierarchyRef), items[0].selected ? 'remove' : 'add')
+    }
+
+    renderControls() {
+        return <>
+            <ActionMenu items={this.hierarchyItems} onSelect={this.onSelect} multiselect />
+            <button onClick={() => this.plugin.managers.structure.hierarchy.createAllModels(this.plugin.managers.structure.hierarchy.state.current.trajectories[0])}>All Models</button>
+        </>;
+    }
+}

src/mol-plugin/behavior/static/state.ts

@@ -104,14 +104,16 @@ function setVisibilityVisitor(t: StateTransform, tree: StateTree, ctx: { state:
 }
 
 export function Highlight(ctx: PluginContext) {
-    PluginCommands.State.Highlight.subscribe(ctx, ({ state, ref }) => {
+    PluginCommands.State.Highlight.subscribe(ctx, ({ state, ref, extend }) => {
         const cell = state.select(ref)[0];
         if (!cell) return;
         if (SO.Molecule.Structure.is(cell.obj)) {
-            ctx.managers.interactivity.lociHighlights.highlightOnly({ loci: Structure.Loci(cell.obj.data) }, false);
+            if (extend) ctx.managers.interactivity.lociHighlights.highlightOnlyExtend({ loci: Structure.Loci(cell.obj.data) }, false);
+            else ctx.managers.interactivity.lociHighlights.highlightOnly({ loci: Structure.Loci(cell.obj.data) }, false);
         } else if (cell && SO.isRepresentation3D(cell.obj)) {
             const { repr } = cell.obj.data
-            ctx.managers.interactivity.lociHighlights.highlightOnly({ loci: repr.getLoci(), repr }, false);
+            if (extend) ctx.managers.interactivity.lociHighlights.highlightOnlyExtend({ loci: repr.getLoci(), repr }, false);
+            else ctx.managers.interactivity.lociHighlights.highlightOnly({ loci: repr.getLoci(), repr }, false);
         } else if (SO.Molecule.Structure.Selections.is(cell.obj)) {
             ctx.managers.interactivity.lociHighlights.clearHighlights();
             for (const entry of cell.obj.data) {

src/mol-plugin/commands.ts

@@ -24,7 +24,7 @@ export const PluginCommands = {
 
         ToggleExpanded: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
         ToggleVisibility: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
-        Highlight: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
+        Highlight: PluginCommand<{ state: State, ref: StateTransform.Ref, extend?: boolean }>(),
         ClearHighlight: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
 
         Snapshots: {

src/mol-state/state/builder.ts

@@ -93,8 +93,9 @@ namespace StateBuilder {
             return new To<StateObject, StateTransformer>(this.state, ref, this);
         }
         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;
+        delete(obj: StateObjectRef) {
+            const ref = StateObjectRef.resolveRef(obj);
+            if (!ref || !this.state.tree.transforms.has(ref)) return this;
             this.editInfo.count++;
             this.state.tree.remove(ref);
             this.state.actions.push({ kind: 'delete', ref });
@@ -229,7 +230,7 @@ namespace StateBuilder {
         to<S extends StateObjectSelector>(selector: S): To<StateObjectSelector.Obj<S>, StateObjectSelector.Transformer<S>>
         to(ref: StateTransform.Ref | StateObjectCell | StateObjectSelector) { return  this.root.to(ref as any); }
         toRoot<A extends StateObject>() { return this.root.toRoot<A>(); }
-        delete(ref: StateTransform.Ref) { return this.root.delete(ref); }
+        delete(ref: StateObjectRef) { return this.root.delete(ref); }
 
         getTree(): StateTree { return buildTree(this.state); }
 

src/mol-theme/color/model-index.ts

@@ -30,18 +30,23 @@ export function ModelIndexColorTheme(ctx: ThemeDataContext, props: PD.Values<Mod
 
     if (ctx.structure) {
         const { models } = ctx.structure.root
-        const palette = getPalette(models.length, props)
+
+        let size = 0;
+        for (const m of models) size = Math.max(size, m.trajectoryInfo.size);
+
+        const palette = getPalette(size, props)
         legend = palette.legend
-        const modelColor = new Map<string, Color>()
+        const modelColor = new Map<number, Color>()
         for (let i = 0, il = models.length; i <il; ++i) {
-            modelColor.set(models[i].id, palette.color(i))
+            const idx = models[i].trajectoryInfo.index;
+            modelColor.set(models[i].trajectoryInfo.index, palette.color(idx))
         }
 
         color = (location: Location): Color => {
             if (StructureElement.Location.is(location)) {
-                return modelColor.get(location.unit.model.id)!
+                return modelColor.get(location.unit.model.trajectoryInfo.index)!
             } else if (Bond.isLocation(location)) {
-                return modelColor.get(location.aUnit.model.id)!
+                return modelColor.get(location.aUnit.model.trajectoryInfo.index)!
             }
             return DefaultColor
         }
@@ -65,5 +70,5 @@ export const ModelIndexColorThemeProvider: ColorTheme.Provider<ModelIndexColorTh
     factory: ModelIndexColorTheme,
     getParams: getModelIndexColorThemeParams,
     defaultValues: PD.getDefaultValues(ModelIndexColorThemeParams),
-    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.models.length > 1
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.elementCount > 0 && ctx.structure.models[0].trajectoryInfo.size > 1
 }