refactored cellpack loader

- properly handle coloring
- support hiv_lipids loading

src/apps/viewer/extensions/cellpack/color.ts

@@ -0,0 +1,97 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ThemeDataContext } from '../../../../mol-theme/theme'
+import { ParamDefinition as PD } from '../../../../mol-util/param-definition'
+import { Color } from '../../../../mol-util/color'
+import { getPalette } from '../../../../mol-util/color/palette'
+import { ColorTheme, LocationColor } from '../../../../mol-theme/color'
+import { ScaleLegend, TableLegend } from '../../../../mol-util/legend'
+import { StructureElement, Bond } from '../../../../mol-model/structure'
+import { Location } from '../../../../mol-model/location';
+import { CellPackInfoProvider } from './property'
+import { distinctColors } from '../../../../mol-util/color/distinct'
+import { Hcl } from '../../../../mol-util/color/spaces/hcl'
+
+
+const DefaultColor = Color(0xCCCCCC)
+const Description = 'Gives every model in a CellPack packing a unique color similar to other models in the packing.'
+
+export const CellPackColorThemeParams = {}
+export type CellPackColorThemeParams = typeof CellPackColorThemeParams
+export function getCellPackColorThemeParams(ctx: ThemeDataContext) {
+    return CellPackColorThemeParams // TODO return copy
+}
+
+export function CellPackColorTheme(ctx: ThemeDataContext, props: PD.Values<CellPackColorThemeParams>): ColorTheme<CellPackColorThemeParams> {
+    let color: LocationColor
+    let legend: ScaleLegend | TableLegend | undefined
+
+    const info = ctx.structure && CellPackInfoProvider.get(ctx.structure).value
+    console.log(info)
+
+    if (ctx.structure && info) {
+        const colors = distinctColors(info.packingsCount)
+        const hcl = Hcl.fromColor(Hcl(), colors[info.packingIndex])
+        const hue = [Math.max(0, hcl[0] - 35), Math.min(360, hcl[0] + 35)] as [number, number]
+
+        const { models } = ctx.structure.root
+
+        let size = 0;
+        for (const m of models) size = Math.max(size, m.trajectoryInfo.size);
+
+        const palette = getPalette(size, { palette: {
+            name: 'generate',
+            params: {
+                hue, chroma: [30, 80], luminance: [15, 85],
+                clusteringStepCount: 50, minSampleCount: 800, maxCount: 75,
+                minLabel: 'Min', maxLabel: 'Max', valueLabel: (i: number) => `${i + 1}`,
+            }
+        }})
+        legend = palette.legend
+        const modelColor = new Map<number, Color>()
+        for (let i = 0, il = models.length; i <il; ++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.trajectoryInfo.index)!
+            } else if (Bond.isLocation(location)) {
+                return modelColor.get(location.aUnit.model.trajectoryInfo.index)!
+            }
+            return DefaultColor
+        }
+    } else {
+        color = () => DefaultColor
+    }
+
+    return {
+        factory: CellPackColorTheme,
+        granularity: 'instance',
+        color,
+        props,
+        description: Description,
+        legend
+    }
+}
+
+export const CellPackColorThemeProvider: ColorTheme.Provider<CellPackColorThemeParams, 'cellpack'> = {
+    name: 'cellpack',
+    label: 'CellPack',
+    category: ColorTheme.Category.Chain,
+    factory: CellPackColorTheme,
+    getParams: getCellPackColorThemeParams,
+    defaultValues: PD.getDefaultValues(CellPackColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => {
+        return (
+            !!ctx.structure && ctx.structure.elementCount > 0 &&
+            ctx.structure.models[0].trajectoryInfo.size > 1 &&
+            !!CellPackInfoProvider.get(ctx.structure).value
+        )
+    }
+}

src/apps/viewer/extensions/cellpack/model.ts

@@ -17,8 +17,6 @@ import { Mat4, Vec3, Quat } from '../../../../mol-math/linear-algebra';
 import { SymmetryOperator } from '../../../../mol-math/geometry';
 import { Task, RuntimeContext } from '../../../../mol-task';
 import { StateTransforms } from '../../../../mol-plugin-state/transforms';
-import { distinctColors } from '../../../../mol-util/color/distinct';
-import { Hcl } from '../../../../mol-util/color/spaces/hcl';
 import { ParseCellPack, StructureFromCellpack, DefaultCellPackBaseUrl } from './state';
 import { MolScriptBuilder as MS } from '../../../../mol-script/language/builder';
 import { getMatFromResamplePoints } from './curve';
@@ -27,8 +25,10 @@ import { CifCategory, CifField } from '../../../../mol-io/reader/cif';
 import { mmCIF_Schema } from '../../../../mol-io/reader/cif/schema/mmcif';
 import { Column } from '../../../../mol-data/db';
 import { createModels } from '../../../../mol-model-formats/structure/basic/parser';
-import { CellpackPackingsPreset } from './preset';
+import { CellpackPackingPreset, CellpackMembranePreset } from './preset';
 import { AjaxTask } from '../../../../mol-util/data-source';
+import { CellPackInfoProvider } from './property';
+import { CellPackColorThemeProvider } from './color';
 
 function getCellPackModelUrl(fileName: string, baseUrl: string) {
     return `${baseUrl}/results/${fileName}`
@@ -313,17 +313,57 @@ async function loadHivMembrane(plugin: PluginContext, runtime: RuntimeContext, s
         .apply(StateTransforms.Model.TrajectoryFromMmCif, undefined, { state: { isGhost: true } })
         .apply(StateTransforms.Model.ModelFromTrajectory, undefined, { state: { isGhost: true } })
         .apply(StateTransforms.Model.StructureFromModel)
-    await state.updateTree(membraneBuilder).runInContext(runtime)
+    const membraneObj = await state.updateTree(membraneBuilder).runInContext(runtime)
+    console.log(membraneObj.data)
 
-    const packingParams = {
-        traceOnly: params.preset.traceOnly,
-        polymerOnly: false,
+    const membraneParams = {
         representation: params.preset.representation,
-        hue: [0, 0] as [number, number]
     }
     const membrane = state.build().to(membraneBuilder.ref)
     await plugin.updateDataState(membrane, { revertOnError: true });
-    await CellpackPackingsPreset.apply(membrane.selector, packingParams, plugin)
+    await CellpackMembranePreset.apply(membrane.selector, membraneParams, plugin)
+}
+
+async function loadPackings(plugin: PluginContext, runtime: RuntimeContext, state: State, params: LoadCellPackModelParams) {
+    let cellPackJson: StateBuilder.To<PSO.Format.Json, StateTransformer<PSO.Data.String, PSO.Format.Json>>
+    if (params.source.name === 'id') {
+        const url = getCellPackModelUrl(params.source.params, params.baseUrl)
+        cellPackJson = state.build().toRoot()
+            .apply(StateTransforms.Data.Download, { url, isBinary: false, label: params.source.params }, { state: { isGhost: true } })
+    } else {
+        const file = params.source.params
+        cellPackJson = state.build().toRoot()
+            .apply(StateTransforms.Data.ReadFile, { file, isBinary: false, label: file.name }, { state: { isGhost: true } })
+    }
+
+    const cellPackBuilder = cellPackJson
+        .apply(StateTransforms.Data.ParseJson, undefined, { state: { isGhost: true } })
+        .apply(ParseCellPack)
+
+    const cellPackObject = await state.updateTree(cellPackBuilder).runInContext(runtime)
+    const { packings } = cellPackObject.data
+
+    await handleHivRna({ runtime, fetch: plugin.fetch }, packings, params.baseUrl)
+
+    for (let i = 0, il = packings.length; i < il; ++i) {
+        const p = { packing: i, baseUrl: params.baseUrl }
+
+        const packing = state.build().to(cellPackBuilder.ref).apply(StructureFromCellpack, p)
+        await plugin.updateDataState(packing, { revertOnError: true });
+
+        const structure = packing.selector.obj?.data
+        if (structure) {
+            await CellPackInfoProvider.attach({ fetch: plugin.fetch, runtime }, structure, {
+                info: { packingsCount: packings.length, packingIndex: i }
+            })
+        }
+
+        const packingParams = {
+            traceOnly: params.preset.traceOnly,
+            representation: params.preset.representation,
+        }
+        await CellpackPackingPreset.apply(packing.selector, packingParams, plugin)
+    }
 }
 
 const LoadCellPackModelParams = {
@@ -341,7 +381,7 @@ const LoadCellPackModelParams = {
     baseUrl: PD.Text(DefaultCellPackBaseUrl),
     preset: PD.Group({
         traceOnly: PD.Boolean(false),
-        representation: PD.Select('spacefill', PD.arrayToOptions(['spacefill', 'gaussian-surface', 'point', 'ellipsoid']))
+        representation: PD.Select('gaussian-surface', PD.arrayToOptions(['spacefill', 'gaussian-surface', 'point', 'ellipsoid']))
     }, { isExpanded: true })
 }
 type LoadCellPackModelParams = PD.Values<typeof LoadCellPackModelParams>
@@ -351,46 +391,13 @@ export const LoadCellPackModel = StateAction.build({
     params: LoadCellPackModelParams,
     from: PSO.Root
 })(({ state, params }, ctx: PluginContext) => Task.create('CellPack Loader', async taskCtx => {
-    let cellPackJson: StateBuilder.To<PSO.Format.Json, StateTransformer<PSO.Data.String, PSO.Format.Json>>
-    if (params.source.name === 'id') {
-        if (params.source.params === 'hiv_lipids.bcif') {
-            await loadHivMembrane(ctx, taskCtx, state, params)
-            return
-        } else {
-            const url = getCellPackModelUrl(params.source.params, params.baseUrl)
-            cellPackJson = state.build().toRoot()
-                .apply(StateTransforms.Data.Download, { url, isBinary: false, label: params.source.params }, { state: { isGhost: true } })
-        }
-    } else {
-        const file = params.source.params
-        cellPackJson = state.build().toRoot()
-            .apply(StateTransforms.Data.ReadFile, { file, isBinary: false, label: file.name }, { state: { isGhost: true } })
+    if (!ctx.representation.structure.themes.colorThemeRegistry.has(CellPackColorThemeProvider)) {
+        ctx.representation.structure.themes.colorThemeRegistry.add(CellPackColorThemeProvider)
     }
 
-    const cellPackBuilder = cellPackJson
-        .apply(StateTransforms.Data.ParseJson, undefined, { state: { isGhost: true } })
-        .apply(ParseCellPack)
-
-    const cellPackObject = await state.updateTree(cellPackBuilder).runInContext(taskCtx)
-    const { packings } = cellPackObject.data
-
-    await handleHivRna({ runtime: taskCtx, fetch: ctx.fetch }, packings, params.baseUrl)
-
-    const colors = distinctColors(packings.length)
-    for (let i = 0, il = packings.length; i < il; ++i) {
-        const hcl = Hcl.fromColor(Hcl(), colors[i])
-        const hue = [Math.max(0, hcl[0] - 35), Math.min(360, hcl[0] + 35)] as [number, number]
-        const p = { packing: i, baseUrl: params.baseUrl }
-
-        const packing = state.build().to(cellPackBuilder.ref).apply(StructureFromCellpack, p)
-        await ctx.updateDataState(packing, { revertOnError: true });
-
-        const packingParams = {
-            traceOnly: params.preset.traceOnly,
-            polymerOnly: true,
-            representation: params.preset.representation,
-            hue
-        }
-        await CellpackPackingsPreset.apply(packing.selector, packingParams, ctx)
+    if (params.source.name === 'id' && params.source.params === 'hiv_lipids.bcif') {
+        await loadHivMembrane(ctx, taskCtx, state, params)
+    } else {
+        await loadPackings(ctx, taskCtx, state, params)
     }
 }));

src/apps/viewer/extensions/cellpack/preset.ts

@@ -7,19 +7,19 @@
 import { StateObjectRef } from '../../../../mol-state';
 import { StructureRepresentationPresetProvider, presetStaticComponent } from '../../../../mol-plugin-state/builder/structure/representation-preset';
 import { ParamDefinition as PD } from '../../../../mol-util/param-definition';
+import { ColorNames } from '../../../../mol-util/color/names';
+import { CellPackColorThemeProvider } from './color';
 
-export const CellpackPackingsPresetParams = {
+export const CellpackPackingPresetParams = {
     traceOnly: PD.Boolean(true),
-    polymerOnly: PD.Boolean(true),
     representation: PD.Select('gaussian-surface', PD.arrayToOptions(['gaussian-surface', 'spacefill', 'point', 'orientation'])),
-    hue: PD.Interval([0, 360])
 }
-export type CellpackPackingsPresetParams = PD.ValuesFor<typeof CellpackPackingsPresetParams>
+export type CellpackPackingPresetParams = PD.ValuesFor<typeof CellpackPackingPresetParams>
 
-export const CellpackPackingsPreset = StructureRepresentationPresetProvider({
-    id: 'preset-structure-representation-cellpack',
-    display: { name: 'CellPack' },
-    params: () => CellpackPackingsPresetParams,
+export const CellpackPackingPreset = StructureRepresentationPresetProvider({
+    id: 'preset-structure-representation-cellpack-packing',
+    display: { name: 'CellPack Packing' },
+    params: () => CellpackPackingPresetParams,
     async apply(ref, params, plugin) {
         const structureCell = StateObjectRef.resolveAndCheck(plugin.state.data, ref);
         if (!structureCell) return {};
@@ -29,11 +29,8 @@ export const CellpackPackingsPreset = StructureRepresentationPresetProvider({
             traceOnly: params.traceOnly
         };
         const components = {
-            structure: params.polymerOnly
-                ? await presetStaticComponent(plugin, structureCell, 'polymer')
-                : await presetStaticComponent(plugin, structureCell, 'all')
+            polymer: await presetStaticComponent(plugin, structureCell, 'polymer')
         };
-        const selectionType = params.polymerOnly ? 'polymer' : 'all'
 
         if (params.representation === 'gaussian-surface') {
             Object.assign(reprProps, {
@@ -44,19 +41,48 @@ export const CellpackPackingsPreset = StructureRepresentationPresetProvider({
         }
 
         const { update, builder, typeParams } = StructureRepresentationPresetProvider.reprBuilder(plugin, {});
-        const color = 'model-index'
-        const colorParams = {
-            palette: {
-                name: 'generate',
-                params: {
-                    hue: params.hue, chroma: [30, 80], luminance: [15, 85],
-                    clusteringStepCount: 50, minSampleCount: 800,
-                    maxCount: 75
-                }
-            }
+        const color = CellPackColorThemeProvider.name
+        const representations = {
+            polymer: builder.buildRepresentation<any>(update, components.polymer, { type: params.representation, typeParams: { ...typeParams, ...reprProps }, color }, { tag: 'polymer' })
+        };
+
+        await plugin.updateDataState(update, { revertOnError: true });
+        return { components, representations };
+    }
+});
+
+//
+
+export const CellpackMembranePresetParams = {
+    representation: PD.Select('gaussian-surface', PD.arrayToOptions(['gaussian-surface', 'spacefill', 'point', 'orientation'])),
+}
+export type CellpackMembranePresetParams = PD.ValuesFor<typeof CellpackMembranePresetParams>
+
+export const CellpackMembranePreset = StructureRepresentationPresetProvider({
+    id: 'preset-structure-representation-cellpack-membrane',
+    display: { name: 'CellPack Membrane' },
+    params: () => CellpackMembranePresetParams,
+    async apply(ref, params, plugin) {
+        const structureCell = StateObjectRef.resolveAndCheck(plugin.state.data, ref);
+        if (!structureCell) return {};
+        console.log(ref, structureCell)
+
+        const reprProps = {
+            ignoreHydrogens: true,
+        };
+        const components = {
+            membrane: await presetStaticComponent(plugin, structureCell, 'all')
+        };
+
+        if (params.representation === 'gaussian-surface') {
+            Object.assign(reprProps, {
+                quality: 'custom', resolution: 10, radiusOffset: 2, doubleSided: false
+            })
         }
+
+        const { update, builder, typeParams } = StructureRepresentationPresetProvider.reprBuilder(plugin, {});
         const representations = {
-            structure: builder.buildRepresentation<any>(update, components.structure, { type: params.representation, typeParams: { ...typeParams, ...reprProps }, color, colorParams }, { tag: selectionType })
+            membrane: builder.buildRepresentation(update, components.membrane, { type: 'gaussian-surface', typeParams: { ...typeParams, ...reprProps }, color: 'uniform', colorParams: { value: ColorNames.lightgrey } }, { tag: 'all' })
         };
 
         await plugin.updateDataState(update, { revertOnError: true });

src/apps/viewer/extensions/cellpack/property.ts

@@ -0,0 +1,32 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { CustomStructureProperty } from '../../../../mol-model-props/common/custom-structure-property'
+import { Structure, CustomPropertyDescriptor } from '../../../../mol-model/structure'
+import { CustomProperty } from '../../../../mol-model-props/common/custom-property'
+import { ParamDefinition as PD } from '../../../../mol-util/param-definition'
+
+export type CellPackInfoValue = {
+    packingsCount: number
+    packingIndex: number
+}
+
+const CellPackInfoParams = {
+    info: PD.Value<CellPackInfoValue>({ packingsCount: 1, packingIndex: 0 }, { isHidden: true })
+}
+type CellPackInfoParams = PD.Values<typeof CellPackInfoParams>
+
+export const CellPackInfoProvider: CustomStructureProperty.Provider<typeof CellPackInfoParams, CellPackInfoValue> = CustomStructureProperty.createProvider({
+    label: 'CellPack Info',
+    descriptor: CustomPropertyDescriptor({ name: 'cellpack-info' }),
+    type: 'root',
+    defaultParams: CellPackInfoParams,
+    getParams: (data: Structure) => CellPackInfoParams,
+    isApplicable: (data: Structure) => true,
+    obtain: async (ctx: CustomProperty.Context, data: Structure, props: CellPackInfoParams) => {
+        return { ...CellPackInfoParams.info.defaultValue, ...props.info }
+    }
+})

src/mol-theme/theme.ts

@@ -121,20 +121,24 @@ export class ThemeRegistry<T extends ColorTheme<any> | SizeTheme<any>> {
     }
 
     remove(provider: ThemeProvider<T, any>) {
-        this._list.splice(this._list.findIndex(e => e.name === name), 1)
-        const p = this._map.get(name);
+        this._list.splice(this._list.findIndex(e => e.name === provider.name), 1)
+        const p = this._map.get(provider.name);
         if (p) {
-            this._map.delete(name);
+            this._map.delete(provider.name);
             this._name.delete(p);
         }
     }
 
+    has(provider: ThemeProvider<T, any>): boolean {
+        return this._map.has(provider.name)
+    }
+
     get<P extends PD.Params>(name: string): ThemeProvider<T, P> {
         return this._map.get(name) || this.emptyProvider
     }
 
     getName(provider: ThemeProvider<T, any>): string {
-        if (!this._name.has(provider)) throw new Error(`'${provider.label}' is not a registered represenatation provider.`);
+        if (!this._name.has(provider)) throw new Error(`'${provider.label}' is not a registered theme provider.`);
         return this._name.get(provider)!;
     }