Browse Source

added TrajectoryHierarchyBuilder

Alexander Rose 5 years ago
parent
commit
a9533b666c

+ 10 - 11
src/mol-plugin-state/actions/structure.ts

@@ -17,8 +17,7 @@ import { StateTransforms } from '../transforms';
 import { Download, ParsePsf } from '../transforms/data';
 import { CoordinatesFromDcd, CustomModelProperties, CustomStructureProperties, TopologyFromPsf, TrajectoryFromModelAndCoordinates } from '../transforms/model';
 import { DataFormatProvider, guessCifVariant } from './data-format';
-import { applyTrajectoryHierarchyPreset } from '../builder/structure/hierarchy-preset';
-import { PresetStructureReprentations } from '../builder/structure/representation-preset';
+import { PresetStructureRepresentations } from '../builder/structure/representation-preset';
 
 // TODO make unitcell creation part of preset
 
@@ -37,7 +36,7 @@ export const MmcifProvider: DataFormatProvider<PluginStateObject.Data.String | P
         return Task.create('mmCIF default builder', async () => {
             const trajectory = await ctx.builders.structure.parseTrajectory(data, 'mmcif');
             const representationPreset = options.visuals ? 'auto' : 'empty';
-            await applyTrajectoryHierarchyPreset(ctx, trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
+            await ctx.builders.structure.hierarchy.applyPreset(trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
         })
     }
 }
@@ -54,7 +53,7 @@ export const PdbProvider: DataFormatProvider<any> = {
         return Task.create('PDB default builder', async () => {
             const trajectory = await ctx.builders.structure.parseTrajectory(data, 'pdb');
             const representationPreset = options.visuals ? 'auto' : 'empty';
-            await applyTrajectoryHierarchyPreset(ctx, trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
+            await ctx.builders.structure.hierarchy.applyPreset(trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
         })
     }
 }
@@ -71,7 +70,7 @@ export const GroProvider: DataFormatProvider<any> = {
         return Task.create('GRO default builder', async () => {
             const trajectory = await ctx.builders.structure.parseTrajectory(data, 'gro');
             const representationPreset = options.visuals ? 'auto' : 'empty';
-            await applyTrajectoryHierarchyPreset(ctx, trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
+            await ctx.builders.structure.hierarchy.applyPreset(trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
         })
     }
 }
@@ -88,7 +87,7 @@ export const Provider3dg: DataFormatProvider<any> = {
         return Task.create('3DG default builder', async () => {
             const trajectory = await ctx.builders.structure.parseTrajectory(data, '3dg');
             const representationPreset = options.visuals ? 'auto' : 'empty';
-            await applyTrajectoryHierarchyPreset(ctx, trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
+            await ctx.builders.structure.hierarchy.applyPreset(trajectory, 'first-model', { showUnitcell: options.visuals, representationPreset });
         })
     }
 }
@@ -129,7 +128,7 @@ export const DcdProvider: DataFormatProvider<any> = {
 
 const DownloadModelRepresentationOptions = (plugin: PluginContext) => PD.Group({
     type: RootStructureDefinition.getParams(void 0, 'auto').type,
-    representation: PD.Select(PresetStructureReprentations.auto.id,
+    representation: PD.Select(PresetStructureRepresentations.auto.id,
         plugin.builders.structure.representation.getPresets().map(p => [p.id, p.display.name] as any),
         { description: 'Which representation preset to use.' }),
     asTrajectory: PD.Optional(PD.Boolean(false, { description: 'Load all entries into a single trajectory.' }))
@@ -224,8 +223,8 @@ const DownloadStructure = StateAction.build({
         default: throw new Error(`${(src as any).name} not supported.`);
     }
 
-    const representationPreset: any = params.source.params.options.representation || PresetStructureReprentations.auto.id;
-    const showUnitcell = representationPreset !== PresetStructureReprentations.empty.id;
+    const representationPreset: any = params.source.params.options.representation || PresetStructureRepresentations.auto.id;
+    const showUnitcell = representationPreset !== PresetStructureRepresentations.empty.id;
 
     const structure = src.params.options.type.name === 'auto' ? void 0 : src.params.options.type;
 
@@ -237,7 +236,7 @@ const DownloadStructure = StateAction.build({
             }, { state: { isGhost: true } });
             const trajectory = await plugin.builders.structure.parseTrajectory(blob, { formats: downloadParams.map((_, i) => ({ id: '' + i, format: 'cif' as 'cif' })) });
 
-            await applyTrajectoryHierarchyPreset(plugin, trajectory, 'first-model', {
+            await plugin.builders.structure.hierarchy.applyPreset(trajectory, 'first-model', {
                 structure,
                 showUnitcell,
                 representationPreset
@@ -247,7 +246,7 @@ const DownloadStructure = StateAction.build({
                 const data = await plugin.builders.data.download(download, { state: { isGhost: true } });
                 const trajectory = await plugin.builders.structure.parseTrajectory(data, format);
 
-                await applyTrajectoryHierarchyPreset(plugin, trajectory, 'first-model', {
+                await plugin.builders.structure.hierarchy.applyPreset(trajectory, 'first-model', {
                     structure,
                     showUnitcell,
                     representationPreset

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

@@ -18,6 +18,7 @@ import { StructureElement } from '../../mol-model/structure';
 import { ModelSymmetry } from '../../mol-model-formats/structure/property/symmetry';
 import { SpacegroupCell } from '../../mol-math/geometry';
 import Expression from '../../mol-script/language/expression';
+import { TrajectoryHierarchyBuilder } from './structure/hierarchy';
 
 export class StructureBuilder {
     private get dataState() {
@@ -40,6 +41,7 @@ export class StructureBuilder {
         return trajectory.selector;
     }
 
+    readonly hierarchy = new TrajectoryHierarchyBuilder(this.plugin);
     readonly representation = new StructureRepresentationBuilder(this.plugin);
 
     async parseTrajectory(data: StateObjectRef<SO.Data.Binary | SO.Data.String>, format: BuiltInTrajectoryFormat | TrajectoryFormatProvider): Promise<StateObjectSelector<SO.Molecule.Trajectory>>

+ 69 - 26
src/mol-plugin-state/builder/structure/hierarchy-preset.ts

@@ -10,10 +10,10 @@ import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { StateObjectRef, StateTransformer } from '../../../mol-state';
 import { StateTransforms } from '../../transforms';
 import { RootStructureDefinition } from '../../helpers/root-structure';
-import { PresetStructureReprentations } from './representation-preset';
+import { PresetStructureRepresentations } from './representation-preset';
 import { PluginContext } from '../../../mol-plugin/context';
-import { isProductionMode } from '../../../mol-util/debug';
-import { Task } from '../../../mol-task';
+import { Vec3 } from '../../../mol-math/linear-algebra';
+import { Model } from '../../../mol-model/structure';
 
 export interface TrajectoryHierarchyPresetProvider<P = any, S = {}> extends PresetProvider<PluginStateObject.Molecule.Trajectory, P, S> { }
 export namespace TrajectoryHierarchyPresetProvider {
@@ -25,12 +25,12 @@ export function TrajectoryHierarchyPresetProvider<P, S>(preset: TrajectoryHierar
 const CommonParams = (a: PluginStateObject.Molecule.Trajectory | undefined, plugin: PluginContext) => ({
     modelProperties: PD.Optional(PD.Group(StateTransformer.getParamDefinition(StateTransforms.Model.CustomModelProperties, void 0, plugin))),
     structureProperties: PD.Optional(PD.Group(StateTransformer.getParamDefinition(StateTransforms.Model.CustomStructureProperties, void 0, plugin))),
-    representationPreset: PD.Optional(PD.Text<keyof PresetStructureReprentations>('auto' as const)),
+    representationPreset: PD.Optional(PD.Text<keyof PresetStructureRepresentations>('auto' as const)),
 })
 
 const FirstModelParams = (a: PluginStateObject.Molecule.Trajectory | undefined, plugin: PluginContext) =>  ({
     model: PD.Optional(PD.Group(StateTransformer.getParamDefinition(StateTransforms.Model.ModelFromTrajectory, a, plugin))),
-    showUnitcell: PD.Optional(PD.Boolean(true)),
+    showUnitcell: PD.Optional(PD.Boolean(false)),
     structure: PD.Optional(RootStructureDefinition.getParams(void 0, 'assembly').type),
     ...CommonParams(a, plugin)
 });
@@ -38,6 +38,9 @@ const FirstModelParams = (a: PluginStateObject.Molecule.Trajectory | undefined,
 const firstModel = TrajectoryHierarchyPresetProvider({
     id: 'preset-trajectory-first-model',
     display: { name: 'First Model', group: 'Preset' },
+    isApplicable: o => {
+        return true
+    },
     params: FirstModelParams,
     async apply(trajectory, params, plugin) {
         const builder = plugin.builders.structure;
@@ -65,6 +68,9 @@ const firstModel = TrajectoryHierarchyPresetProvider({
 const allModels = TrajectoryHierarchyPresetProvider({
     id: 'preset-trajectory-all-models',
     display: { name: 'All Models', group: 'Preset' },
+    isApplicable: o => {
+        return o.data.length > 1
+    },
     params: CommonParams,
     async apply(trajectory, params, plugin) {
         const tr = StateObjectRef.resolveAndCheck(plugin.state.data, trajectory)?.obj?.data;
@@ -82,34 +88,71 @@ const allModels = TrajectoryHierarchyPresetProvider({
 
             models.push(model);
             structures.push(structure);
-            await builder.representation.applyPreset(structureProperties, params.representationPreset || 'auto', { globalThemeName: 'model-index' });
+            await builder.representation.applyPreset(structureProperties, params.representationPreset || 'auto', { globalThemeName: 'model-index', quality: 'medium' });
         }
 
         return { models, structures };
     }
 });
 
-export const PresetStructureTrajectoryHierarchy = {
-    'first-model': firstModel,
-    'all-models': allModels
-};
-export type PresetStructureTrajectoryHierarchy = typeof PresetStructureTrajectoryHierarchy;
+const CrystalSymmetryParams = (a: PluginStateObject.Molecule.Trajectory | undefined, plugin: PluginContext) => ({
+    model: PD.Optional(PD.Group(StateTransformer.getParamDefinition(StateTransforms.Model.ModelFromTrajectory, a, plugin))),
+    ...CommonParams(a, plugin)
+});
 
-// TODO: should there be a registry like for representations?
+async function applyCrystalSymmetry(ijk: { ijkMin: Vec3, ijkMax: Vec3 }, trajectory: StateObjectRef<PluginStateObject.Molecule.Trajectory>, params: PD.ValuesFor<ReturnType<typeof CrystalSymmetryParams>>, plugin: PluginContext) {
+    const builder = plugin.builders.structure;
+
+    const model = await builder.createModel(trajectory, params.model);
+    const modelProperties = await builder.insertModelProperties(model, params.modelProperties);
+
+    const structure = await builder.createStructure(modelProperties || model, {
+        name: 'symmetry',
+        params: ijk
+    });
+    const structureProperties = await builder.insertStructureProperties(structure, params.structureProperties);
+
+    const unitcell = await builder.tryCreateUnitcell(modelProperties, undefined, { isHidden: false });
+    const representation =  await plugin.builders.structure.representation.applyPreset(structureProperties, params.representationPreset || 'auto');
+
+    return {
+        model,
+        modelProperties,
+        unitcell,
+        structure,
+        structureProperties,
+        representation
+    };
+}
 
-export function applyTrajectoryHierarchyPreset<K extends keyof PresetStructureTrajectoryHierarchy>(plugin: PluginContext, parent: StateObjectRef<PluginStateObject.Molecule.Trajectory>, preset: K, params?: Partial<TrajectoryHierarchyPresetProvider.Params<PresetStructureTrajectoryHierarchy[K]>>): Promise<TrajectoryHierarchyPresetProvider.State<PresetStructureTrajectoryHierarchy[K]>> | undefined
-export function applyTrajectoryHierarchyPreset<P = any, S = {}>(plugin: PluginContext, parent: StateObjectRef<PluginStateObject.Molecule.Trajectory>, provider: TrajectoryHierarchyPresetProvider<P, S>, params?: P): Promise<S> | undefined
-export function applyTrajectoryHierarchyPreset(plugin: PluginContext, parent: StateObjectRef, providerRef: string | TrajectoryHierarchyPresetProvider, params?: any): Promise<any> | undefined {
-    const provider = typeof providerRef === 'string' ? (PresetStructureTrajectoryHierarchy as any)[providerRef] : providerRef;
-    if (!provider) return;
+const unitcell = TrajectoryHierarchyPresetProvider({
+    id: 'preset-trajectory-unitcell',
+    display: { name: 'Unitcell', group: 'Preset' },
+    isApplicable: o => {
+        return Model.hasCrystalSymmetry(o.data[0])
+    },
+    params: CrystalSymmetryParams,
+    async apply(trajectory, params, plugin) {
+        return await applyCrystalSymmetry({ ijkMin: Vec3.create(0, 0, 0), ijkMax: Vec3.create(0, 0, 0) }, trajectory, params, plugin);
+    }
+});
 
-    const state = plugin.state.data;
-    const cell = StateObjectRef.resolveAndCheck(state, parent);
-    if (!cell) {
-        if (!isProductionMode) console.warn(`Applying hierarchy preset provider to bad cell.`);
-        return;
+const supercell = TrajectoryHierarchyPresetProvider({
+    id: 'preset-trajectory-supercell',
+    display: { name: 'Supercell', group: 'Preset' },
+    isApplicable: o => {
+        return Model.hasCrystalSymmetry(o.data[0])
+    },
+    params: CrystalSymmetryParams,
+    async apply(trajectory, params, plugin) {
+        return await applyCrystalSymmetry({ ijkMin: Vec3.create(-1, -1, -1), ijkMax: Vec3.create(1, 1, 1) }, trajectory, params, plugin);
     }
-    const prms = { ...PD.getDefaultValues(provider.params(cell.obj!, plugin) as PD.Params), ...params };
-    const task = Task.create(`${provider.display.name}`, () => provider.apply(cell, prms, plugin) as Promise<any>);
-    return plugin.runTask(task);
-}
+});
+
+export const PresetTrajectoryHierarchy = {
+    'first-model': firstModel,
+    'all-models': allModels,
+    unitcell,
+    supercell,
+};
+export type PresetTrajectoryHierarchy = typeof PresetTrajectoryHierarchy;

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

@@ -0,0 +1,107 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { arrayFind } from '../../../mol-data/util';
+import { PluginContext } from '../../../mol-plugin/context';
+import { StateObjectRef } from '../../../mol-state';
+import { Task } from '../../../mol-task';
+import { isProductionMode } from '../../../mol-util/debug';
+import { objectForEach } from '../../../mol-util/object';
+import { ParamDefinition as PD } from '../../../mol-util/param-definition';
+import { PluginStateObject } from '../../objects';
+import { PresetTrajectoryHierarchy, TrajectoryHierarchyPresetProvider } from './hierarchy-preset';
+
+// TODO factor out code shared with StructureRepresentationBuilder?
+
+export type TrajectoryHierarchyPresetProviderRef = keyof PresetTrajectoryHierarchy | TrajectoryHierarchyPresetProvider | string
+
+export class TrajectoryHierarchyBuilder {
+    private _providers: TrajectoryHierarchyPresetProvider[] = [];
+    private providerMap: Map<string, TrajectoryHierarchyPresetProvider> = new Map();
+
+    readonly defaultProvider = PresetTrajectoryHierarchy['first-model'];
+
+    private resolveProvider(ref: TrajectoryHierarchyPresetProviderRef) {
+        return typeof ref === 'string'
+            ? PresetTrajectoryHierarchy[ref as keyof PresetTrajectoryHierarchy] ?? arrayFind(this._providers, p => p.id === ref)
+            : ref;
+    }
+
+    hasPreset(t: PluginStateObject.Molecule.Trajectory) {
+        for (const p of this._providers) {
+            if (!p.isApplicable || p.isApplicable(t, this.plugin)) return true;
+        }
+        return false;
+    }
+
+    get providers(): ReadonlyArray<TrajectoryHierarchyPresetProvider> { return this._providers; }
+
+    getPresets(t?: PluginStateObject.Molecule.Trajectory) {
+        if (!t) return this.providers;
+        const ret = [];
+        for (const p of this._providers) {
+            if (p.isApplicable && !p.isApplicable(t, this.plugin)) continue;
+            ret.push(p);
+        }
+        return ret;
+    }
+
+    getPresetSelect(t?: PluginStateObject.Molecule.Trajectory): PD.Select<string> {
+        const options: [string, string][] = [];
+        for (const p of this._providers) {
+            if (t && p.isApplicable && !p.isApplicable(t, this.plugin)) continue;
+            options.push([p.id, p.display.name]);
+        }
+        return PD.Select('auto', options);
+    }
+
+    getPresetsWithOptions(t: PluginStateObject.Molecule.Trajectory) {
+        const options: [string, string][] = [];
+        const map: { [K in string]: PD.Any } = Object.create(null);
+        for (const p of this._providers) {
+            if (p.isApplicable && !p.isApplicable(t, this.plugin)) continue;
+
+            options.push([p.id, p.display.name]);
+            map[p.id] = p.params ? PD.Group(p.params(t, this.plugin)) : PD.EmptyGroup()
+        }
+        if (options.length === 0) return PD.MappedStatic('', { '': PD.EmptyGroup() });
+        return PD.MappedStatic(options[0][0], map, { options });
+    }
+
+    registerPreset(provider: TrajectoryHierarchyPresetProvider) {
+        if (this.providerMap.has(provider.id)) {
+            throw new Error(`Hierarchy provider with id '${provider.id}' already registered.`);
+        }
+        this._providers.push(provider);
+        this.providerMap.set(provider.id, provider);
+    }
+
+    applyPreset<K extends keyof PresetTrajectoryHierarchy>(parent: StateObjectRef<PluginStateObject.Molecule.Trajectory>, preset: K, params?: Partial<TrajectoryHierarchyPresetProvider.Params<PresetTrajectoryHierarchy[K]>>): Promise<TrajectoryHierarchyPresetProvider.State<PresetTrajectoryHierarchy[K]>> | undefined
+    applyPreset<P = any, S = {}>(parent: StateObjectRef<PluginStateObject.Molecule.Trajectory>, provider: TrajectoryHierarchyPresetProvider<P, S>, params?: P): Promise<S> | undefined
+    applyPreset(parent: StateObjectRef, providerRef: string | TrajectoryHierarchyPresetProvider, params?: any): Promise<any> | undefined {
+        const provider = this.resolveProvider(providerRef);
+        if (!provider) return;
+
+        const state = this.plugin.state.data;
+        const cell = StateObjectRef.resolveAndCheck(state, parent);
+        if (!cell) {
+            if (!isProductionMode) console.warn(`Applying hierarchy preset provider to bad cell.`);
+            return;
+        }
+
+        const prms = params || (provider.params
+            ? PD.getDefaultValues(provider.params(cell.obj, this.plugin) as PD.Params)
+            : {})
+
+        const task = Task.create(`${provider.display.name}`, () => provider.apply(cell, prms, this.plugin) as Promise<any>);
+        return this.plugin.runTask(task);
+    }
+
+    constructor(public plugin: PluginContext) {
+        objectForEach(PresetTrajectoryHierarchy, r => this.registerPreset(r));
+    }
+}

+ 2 - 2
src/mol-plugin-state/builder/structure/representation-preset.ts

@@ -227,7 +227,7 @@ export function presetSelectionComponent(plugin: PluginContext, structure: State
     return plugin.builders.structure.tryCreateComponentFromSelection(structure, Q[query], `selection-${query}`);
 }
 
-export const PresetStructureReprentations = {
+export const PresetStructureRepresentations = {
     empty,
     auto,
     'atomic-detail': atomicDetail,
@@ -236,4 +236,4 @@ export const PresetStructureReprentations = {
     'protein-and-nucleic': proteinAndNucleic,
     'coarse-surface': coarseSurface
 };
-export type PresetStructureReprentations = typeof PresetStructureReprentations;
+export type PresetStructureRepresentations = typeof PresetStructureRepresentations;

+ 9 - 8
src/mol-plugin-state/builder/structure/representation.ts

@@ -14,20 +14,22 @@ import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { createStructureRepresentationParams, StructureRepresentationBuiltInProps, StructureRepresentationProps } from '../../helpers/structure-representation-params';
 import { PluginStateObject } from '../../objects';
 import { StructureRepresentation3D } from '../../transforms/representation';
-import { PresetStructureReprentations, StructureRepresentationPresetProvider } from './representation-preset';
+import { PresetStructureRepresentations, StructureRepresentationPresetProvider } from './representation-preset';
 
-export type StructureRepresentationPresetProviderRef = keyof PresetStructureReprentations | StructureRepresentationPresetProvider | string
+// TODO factor out code shared with TrajectoryHierarchyBuilder?
+
+export type StructureRepresentationPresetProviderRef = keyof PresetStructureRepresentations | StructureRepresentationPresetProvider | string
 
 export class StructureRepresentationBuilder {
     private _providers: StructureRepresentationPresetProvider[] = [];
     private providerMap: Map<string, StructureRepresentationPresetProvider> = new Map();
     private get dataState() { return this.plugin.state.data; }
 
-    readonly defaultProvider = PresetStructureReprentations.auto;
+    readonly defaultProvider = PresetStructureRepresentations.auto;
 
     private resolveProvider(ref: StructureRepresentationPresetProviderRef) {
         return typeof ref === 'string'
-            ? PresetStructureReprentations[ref as keyof PresetStructureReprentations] ?? arrayFind(this._providers, p => p.id === ref)
+            ? PresetStructureRepresentations[ref as keyof PresetStructureRepresentations] ?? arrayFind(this._providers, p => p.id === ref)
             : ref;
     }
 
@@ -74,13 +76,13 @@ export class StructureRepresentationBuilder {
 
     registerPreset(provider: StructureRepresentationPresetProvider) {
         if (this.providerMap.has(provider.id)) {
-            throw new Error(`Repr. provider with id '${provider.id}' already registered.`);
+            throw new Error(`Representation provider with id '${provider.id}' already registered.`);
         }
         this._providers.push(provider);
         this.providerMap.set(provider.id, provider);
     }
 
-    applyPreset<K extends keyof PresetStructureReprentations>(parent: StateObjectRef<PluginStateObject.Molecule.Structure>, preset: K, params?: StructureRepresentationPresetProvider.Params<PresetStructureReprentations[K]>): Promise<StructureRepresentationPresetProvider.State<PresetStructureReprentations[K]>> | undefined
+    applyPreset<K extends keyof PresetStructureRepresentations>(parent: StateObjectRef<PluginStateObject.Molecule.Structure>, preset: K, params?: StructureRepresentationPresetProvider.Params<PresetStructureRepresentations[K]>): Promise<StructureRepresentationPresetProvider.State<PresetStructureRepresentations[K]>> | undefined
     applyPreset<P = any, S = {}>(parent: StateObjectRef<PluginStateObject.Molecule.Structure>, provider: StructureRepresentationPresetProvider<P, S>, params?: P): Promise<S> | undefined
     applyPreset(parent: StateObjectRef<PluginStateObject.Molecule.Structure>, providerId: string, params?: any): Promise<any> | undefined
     applyPreset(parent: StateObjectRef, providerRef: string | StructureRepresentationPresetProvider, params?: any): Promise<any> | undefined {
@@ -98,7 +100,6 @@ export class StructureRepresentationBuilder {
             ? PD.getDefaultValues(provider.params(cell.obj, this.plugin) as PD.Params)
             : {})
 
-
         const task = Task.create(`${provider.display.name}`, () => provider.apply(cell, prms, this.plugin) as Promise<any>);
         return this.plugin.runTask(task);
     }
@@ -127,7 +128,7 @@ export class StructureRepresentationBuilder {
     }
 
     constructor(public plugin: PluginContext) {
-        objectForEach(PresetStructureReprentations, r => this.registerPreset(r));
+        objectForEach(PresetStructureRepresentations, r => this.registerPreset(r));
     }
 }
 

+ 8 - 12
src/mol-plugin-state/manager/structure/hierarchy.ts

@@ -2,6 +2,7 @@
  * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
 import { PluginContext } from '../../../mol-plugin/context';
@@ -9,7 +10,7 @@ import { StructureHierarchy, buildStructureHierarchy, ModelRef, StructureCompone
 import { PluginComponent } from '../../component';
 import { SetUtils } from '../../../mol-util/set';
 import { StateTransform } from '../../../mol-state';
-import { applyTrajectoryHierarchyPreset } from '../../builder/structure/hierarchy-preset';
+import { TrajectoryHierarchyPresetProvider } from '../../builder/structure/hierarchy-preset';
 import { setSubtreeVisibility } from '../../../mol-plugin/behavior/static/state';
 
 export class StructureHierarchyManager extends PluginComponent {
@@ -150,21 +151,16 @@ export class StructureHierarchyManager extends PluginComponent {
         }
     }
 
-    createModels(trajectories: ReadonlyArray<TrajectoryRef>, kind: 'single' | 'all' = 'single') {
+    applyPreset<P = any, S = {}>(trajectories: ReadonlyArray<TrajectoryRef>, provider: TrajectoryHierarchyPresetProvider<P, S>, params?: P): Promise<any>  {
         return this.plugin.dataTransaction(async () => {
-            for (const trajectory of trajectories) {
-                if (trajectory.models.length > 0) {
-                    await this.clearTrajectory(trajectory);
+            for (const t of trajectories) {
+                if (t.models.length > 0) {
+                    await this.clearTrajectory(t);
                 }
 
-                if (trajectory.models.length === 0) return;
+                if (t.models.length === 0) return;
 
-                const tr = trajectory.cell.obj?.data!;
-                if (kind === 'all' && tr.length > 1) {
-                    await applyTrajectoryHierarchyPreset(this.plugin, trajectory.cell, 'all-models');
-                } else {
-                    await applyTrajectoryHierarchyPreset(this.plugin, trajectory.cell, 'first-model');
-                }
+                await this.plugin.builders.structure.hierarchy.applyPreset(t.cell, provider, params);
             }
         });
     }

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

@@ -103,7 +103,7 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
         if (!item) return;
         const mng = this.plugin.managers.structure;
 
-        const structures = mng.hierarchy.selection.structures;
+        const { structures } = mng.hierarchy.selection;
         if (item.value === null) mng.component.clear(structures);
         else mng.component.applyPreset(structures, item.value as any);
     }

+ 20 - 23
src/mol-plugin-ui/structure/source.tsx

@@ -13,13 +13,11 @@ import { IconButton } from '../controls/common';
 import { ParameterControls } from '../controls/parameters';
 import { PluginCommands } from '../../mol-plugin/commands';
 import { StateTransforms } from '../../mol-plugin-state/transforms';
-import { memoize1 } from '../../mol-util/memoize';
-import { StructureHierarchyManager } from '../../mol-plugin-state/manager/structure/hierarchy';
 import { GenericEntry } from './generic';
 
 interface StructureSourceControlState extends CollapsableState {
     isBusy: boolean,
-    show?: 'hierarchy' | 'actions'
+    show?: 'hierarchy' | 'presets'
 }
 
 export class StructureSourceControls extends CollapsableControls<{}, StructureSourceControlState> {
@@ -185,29 +183,28 @@ export class StructureSourceControls extends CollapsableControls<{}, StructureSo
     }
 
     toggleHierarchy = () => this.setState({ show: this.state.show !== 'hierarchy' ? 'hierarchy' : void 0 });
-    toggleActions = () => this.setState({ show: this.state.show !== 'actions' ? 'actions' : void 0 });
+    togglePreset = () => this.setState({ show: this.state.show !== 'presets' ? 'presets' : void 0 });
 
-    actions = memoize1((sel: StructureHierarchyManager['selection']) => this._actions);
+    get presetActions() {
+        const actions: ActionMenu.Item[] = [];
+        const { trajectories } = this.plugin.managers.structure.hierarchy.selection;
+        if (trajectories.length !== 1) return actions
 
-    get _actions() {
-        const ret: ActionMenu.Items = [];
-
-        const { selection } = this.plugin.managers.structure.hierarchy;
-        if (selection.trajectories.some(t => t.cell.obj?.data.length! > 1 && t.cell.obj?.data.length! !== t.models.length)) {
-            ret.push(ActionMenu.Item('Load all models', () => this.plugin.managers.structure.hierarchy.createModels(selection.trajectories, 'all')));
+        const providers = this.plugin.builders.structure.hierarchy.getPresets(trajectories[0].cell.obj)
+        for (const p of providers) {
+            actions.push(ActionMenu.Item(p.display.name, p));
         }
-        if (selection.trajectories.some(t => t.models.length > 1)) {
-            ret.push(ActionMenu.Item('Load single model', () => this.plugin.managers.structure.hierarchy.createModels(selection.trajectories, 'single')));
-        }
-
-        // TODO: remove actions?
-        return ret;
+        return actions;
     }
 
-    selectAction: ActionMenu.OnSelect = item => {
-        if (!item) return;
+    applyPreset: ActionMenu.OnSelect = item => {
         this.setState({ show: void 0 });
-        (item?.value as any)();
+
+        if (!item) return;
+        const mng = this.plugin.managers.structure;
+
+        const { trajectories } = mng.hierarchy.selection;
+        mng.hierarchy.applyPreset(trajectories, item.value as any);
     }
 
     updateStructureModel = async (params: any) => {
@@ -276,17 +273,17 @@ export class StructureSourceControls extends CollapsableControls<{}, StructureSo
 
     renderControls() {
         const disabled = this.state.isBusy || this.isEmpty;
-        const actions = this.actions(this.plugin.managers.structure.hierarchy.selection);
+        const presets = this.presetActions;
         const label = this.label;
         return <>
             <div className='msp-btn-row-group' style={{ marginTop: '1px' }}>
                 <button className='msp-btn msp-form-control msp-flex-item msp-no-overflow' onClick={this.toggleHierarchy} style={{ overflow: 'hidden', textOverflow: 'ellipsis' }} disabled={disabled} title={label}>
                     {label}
                 </button>
-                {actions.length > 0 && <IconButton customClass='msp-form-control' style={{ flex: '0 0 32px' }} onClick={this.toggleActions} icon='dot-3' title='Actions' toggleState={this.state.show === 'actions'} disabled={disabled} />}
+                {presets.length > 0 && <IconButton customClass='msp-form-control' style={{ flex: '0 0 32px' }} onClick={this.togglePreset} icon='bookmarks' title='Presets' toggleState={this.state.show === 'presets'} disabled={disabled} />}
             </div>
             {this.state.show === 'hierarchy' && <ActionMenu items={this.hierarchyItems} onSelect={this.selectHierarchy} multiselect />}
-            {this.state.show === 'actions' && <ActionMenu items={actions} onSelect={this.selectAction} />}
+            {this.state.show === 'presets' && <ActionMenu items={presets} onSelect={this.applyPreset} />}
             {this.modelIndex}
             {this.structureType}
             <div style={{ marginTop: '6px' }}>