mol-plugin-state: StructureHierarchy tweaks

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

@@ -20,11 +20,13 @@ export function buildStructureHierarchy(state: State, previous?: StructureHierar
 
 export interface StructureHierarchy {
     trajectories: TrajectoryRef[],
+    models: ModelRef[],
+    structures: StructureRef[],
     refs: Map<StateTransform.Ref, HierarchyRef>
 }
 
 export function StructureHierarchy(): StructureHierarchy {
-    return { trajectories: [], refs: new Map() }
+    return { trajectories: [], models: [], structures: [], refs: new Map() }
 }
 
 interface RefBase<K extends string = string, T extends StateObject = StateObject> {
@@ -47,12 +49,12 @@ function TrajectoryRef(cell: StateObjectCell<SO.Molecule.Trajectory>): Trajector
 }
 
 export interface ModelRef extends RefBase<'model', SO.Molecule.Model> {
-    trajectory: TrajectoryRef,
+    trajectory?: TrajectoryRef,
     properties?: ModelPropertiesRef,
     structures: StructureRef[]
 }
 
-function ModelRef(cell: StateObjectCell<SO.Molecule.Model>, trajectory: TrajectoryRef): ModelRef {
+function ModelRef(cell: StateObjectCell<SO.Molecule.Model>, trajectory?: TrajectoryRef): ModelRef {
     return { kind: 'model', cell, version: cell.transform.version, trajectory, structures: [] };
 }
 
@@ -65,10 +67,9 @@ function ModelPropertiesRef(cell: StateObjectCell<SO.Molecule.Model>, model: Mod
 }
 
 export interface StructureRef extends RefBase<'structure', SO.Molecule.Structure> {
-    model: ModelRef,
+    model?: ModelRef,
     properties?: StructurePropertiesRef,
     components: StructureComponentRef[],
-    // representations: StructureRepresentationRef[],
     currentFocus?: {
         focus?: StructureComponentRef,
         surroundings?: StructureComponentRef,
@@ -76,7 +77,7 @@ export interface StructureRef extends RefBase<'structure', SO.Molecule.Structure
     // volumeStreaming?: ....
 }
 
-function StructureRef(cell: StateObjectCell<SO.Molecule.Structure>, model: ModelRef): StructureRef {
+function StructureRef(cell: StateObjectCell<SO.Molecule.Structure>, model?: ModelRef): StructureRef {
     return { kind: 'structure', cell, version: cell.transform.version, model, components: [] };
 }
 
@@ -160,16 +161,24 @@ const tagMap: [string, (state: BuildState, cell: StateObjectCell) => boolean | v
         state.currentTrajectory = createOrUpdateRefList(state, cell, state.hierarchy.trajectories, TrajectoryRef, cell);
     }, state => state.currentTrajectory = void 0],
     [StructureBuilderTags.Model, (state, cell) => {
-        if (!state.currentTrajectory) return false;
-        state.currentModel = createOrUpdateRefList(state, cell, state.currentTrajectory.models, ModelRef, cell, state.currentTrajectory);
+        if (state.currentTrajectory) {
+            state.currentModel = createOrUpdateRefList(state, cell, state.currentTrajectory.models, ModelRef, cell, state.currentTrajectory);
+        } else {
+            state.currentModel = ModelRef(cell)
+        }
+        state.hierarchy.models.push(state.currentModel);
     }, state => state.currentModel = void 0],
     [StructureBuilderTags.ModelProperties, (state, cell) => {
         if (!state.currentModel) return false;
         state.currentModel.properties = createOrUpdateRef(state, cell, state.currentModel.properties, ModelPropertiesRef, cell, state.currentModel);
     }, state => { }],
     [StructureBuilderTags.Structure, (state, cell) => {
-        if (!state.currentModel) return false;
-        state.currentStructure = createOrUpdateRefList(state, cell, state.currentModel.structures, StructureRef, cell, state.currentModel);
+        if (state.currentModel) {
+            state.currentStructure = createOrUpdateRefList(state, cell, state.currentModel.structures, StructureRef, cell, state.currentModel);
+        } else {
+            state.currentStructure = StructureRef(cell);
+        }
+        state.hierarchy.structures.push(state.currentStructure);
     }, state => state.currentStructure = void 0],
     [StructureBuilderTags.StructureProperties, (state, cell) => {
         if (!state.currentStructure) return false;

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

@@ -7,7 +7,6 @@
 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';
 import { SetUtils } from '../../../mol-util/set';
 
 interface StructureHierarchyManagerState {
@@ -51,45 +50,16 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         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]] : [];
+    private syncCurrent<T extends HierarchyRef>(hierarchy: StructureHierarchy, current: ReadonlyArray<T>, all: ReadonlyArray<T>): T[] {
+        if (current.length === 0) return all.length > 0 ? [all[0]] : [];
 
-        const newCurrent: TrajectoryRef[] = [];
+        const newCurrent: T[] = [];
         for (const c of current) {
-            const ref = map.get(c.cell.transform.ref) as TrajectoryRef;
+            const ref = hierarchy.refs.get(c.cell.transform.ref) as T;
             if (ref) newCurrent.push(ref);
         }
 
-        if (newCurrent.length === 0) return hierarchy.trajectories.length > 0 ? [hierarchy.trajectories[0]] : [];
-        return newCurrent;
-    }
-
-    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 = map.get(c.cell.transform.ref) as ModelRef;
-            if (ref) newCurrent.push(ref);
-        }
-
-        if (newCurrent.length === 0) return currentTrajectories[0]?.models || [];
-        return newCurrent;
-    }
-
-    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 = map.get(c.cell.transform.ref) as StructureRef;
-            if (ref) newCurrent.push(ref);
-        }
-
-        if (newCurrent.length === 0 && currentModels.length > 0) return Array.prototype.concat.apply([], currentModels.map(m => m.structures));
+        if (newCurrent.length === 0) return all.length > 0 ? [all[0]] : [];
         return newCurrent;
     }
 
@@ -98,15 +68,17 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         if (update.added.length === 0 && update.updated.length === 0 && update.removed.length === 0) {
             return;
         }
+
         this._currentComponentGroups = void 0;
         this._currentSelectionSet = void 0;
 
-        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);
+        const { hierarchy } = update;
+        const trajectories = this.syncCurrent(hierarchy, this.state.current.trajectories, hierarchy.trajectories);
+        const models = this.syncCurrent(hierarchy, this.state.current.models, hierarchy.models);
+        const structures = this.syncCurrent(hierarchy, this.state.current.structures, hierarchy.structures);
 
-        this.updateState({ hierarchy: update.hierarchy, current: { trajectories, models, structures }});
-        this.behaviors.current.next({ hierarchy: update.hierarchy, trajectories, models, structures });
+        this.updateState({ hierarchy, current: { trajectories, models, structures }});
+        this.behaviors.current.next({ hierarchy, trajectories, models, structures });
     }
 
     updateCurrent(refs: HierarchyRef[], action: 'add' | 'remove') {

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

@@ -62,8 +62,9 @@ export class StructureSourceControls extends CollapsableControls<{}, StructureSo
         
         if (structures.length === 1) {
             const s = structures[0];
-            if (s.model.trajectory.models.length === 1) return s.cell.obj?.data.label;
-            return `${s.model.cell.obj?.label} | ${s.cell.obj?.data.label}`;
+            if (s.model?.trajectory?.models && s.model.trajectory.models.length === 1) return s.cell.obj?.data.label;
+            if (s.model) return `${s.model.cell.obj?.label} | ${s.cell.obj?.data.label}`;
+            return s.cell.obj?.data.label;
         }
 
         if (structures.length > 1) {