cellpack: split color themes

src/extensions/cellpack/color.ts → src/extensions/cellpack/color/generate.ts

@@ -4,29 +4,28 @@
  * @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';
-
+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.';
+const Description = 'Gives every model in a CellPack packing a unique generated 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 const CellPackGenerateColorThemeParams = {};
+export type CellPackGenerateColorThemeParams = typeof CellPackGenerateColorThemeParams
+export function getCellPackGenerateColorThemeParams(ctx: ThemeDataContext) {
+    return CellPackGenerateColorThemeParams; // TODO return copy
 }
 
-export function CellPackColorTheme(ctx: ThemeDataContext, props: PD.Values<CellPackColorThemeParams>): ColorTheme<CellPackColorThemeParams> {
+export function CellPackGenerateColorTheme(ctx: ThemeDataContext, props: PD.Values<CellPackGenerateColorThemeParams>): ColorTheme<CellPackGenerateColorThemeParams> {
     let color: LocationColor;
     let legend: ScaleLegend | TableLegend | undefined;
 
@@ -55,11 +54,7 @@ export function CellPackColorTheme(ctx: ThemeDataContext, props: PD.Values<CellP
         const modelColor = new Map<number, Color>();
         for (let i = 0, il = models.length; i < il; ++i) {
             const idx = models[i].trajectoryInfo.index;
-            let acolor = palette.color(idx);
-            if (info.color){
-                acolor = info.colors[idx];
-            }
-            modelColor.set(models[i].trajectoryInfo.index, acolor);
+            modelColor.set(models[i].trajectoryInfo.index, palette.color(idx));
         }
 
         color = (location: Location): Color => {
@@ -75,7 +70,7 @@ export function CellPackColorTheme(ctx: ThemeDataContext, props: PD.Values<CellP
     }
 
     return {
-        factory: CellPackColorTheme,
+        factory: CellPackGenerateColorTheme,
         granularity: 'instance',
         color,
         props,
@@ -84,13 +79,13 @@ export function CellPackColorTheme(ctx: ThemeDataContext, props: PD.Values<CellP
     };
 }
 
-export const CellPackColorThemeProvider: ColorTheme.Provider<CellPackColorThemeParams, 'cellpack'> = {
-    name: 'cellpack',
-    label: 'CellPack',
+export const CellPackGenerateColorThemeProvider: ColorTheme.Provider<CellPackGenerateColorThemeParams, 'cellpack-generate'> = {
+    name: 'cellpack-generate',
+    label: 'CellPack Generate',
     category: ColorTheme.Category.Chain,
-    factory: CellPackColorTheme,
-    getParams: getCellPackColorThemeParams,
-    defaultValues: PD.getDefaultValues(CellPackColorThemeParams),
+    factory: CellPackGenerateColorTheme,
+    getParams: getCellPackGenerateColorThemeParams,
+    defaultValues: PD.getDefaultValues(CellPackGenerateColorThemeParams),
     isApplicable: (ctx: ThemeDataContext) => {
         return (
             !!ctx.structure && ctx.structure.elementCount > 0 &&
@@ -98,4 +93,5 @@ export const CellPackColorThemeProvider: ColorTheme.Provider<CellPackColorThemeP
             !!CellPackInfoProvider.get(ctx.structure).value
         );
     }
-};
+};
+

src/extensions/cellpack/color/provided.ts

@@ -0,0 +1,72 @@
+/**
+ * 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 { ColorTheme, LocationColor } from '../../../mol-theme/color';
+import { ScaleLegend, TableLegend } from '../../../mol-util/legend';
+import { StructureElement } from '../../../mol-model/structure';
+import { Location } from '../../../mol-model/location';
+import { CellPackInfoProvider } from '../property';
+
+const DefaultColor = Color(0xCCCCCC);
+const Description = 'Gives every model in a CellPack the color provied in the packing data.';
+
+export const CellPackProvidedColorThemeParams = {};
+export type CellPackProvidedColorThemeParams = typeof CellPackProvidedColorThemeParams
+export function getCellPackProvidedColorThemeParams(ctx: ThemeDataContext) {
+    return CellPackProvidedColorThemeParams; // TODO return copy
+}
+
+export function CellPackProvidedColorTheme(ctx: ThemeDataContext, props: PD.Values<CellPackProvidedColorThemeParams>): ColorTheme<CellPackProvidedColorThemeParams> {
+    let color: LocationColor;
+    let legend: ScaleLegend | TableLegend | undefined;
+
+    const info = ctx.structure && CellPackInfoProvider.get(ctx.structure).value;
+
+    if (ctx.structure && info?.colors) {
+        const { models } = ctx.structure.root;
+        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, info.colors[idx]);
+        }
+
+        color = (location: Location): Color => {
+            return StructureElement.Location.is(location)
+                ? modelColor.get(location.unit.model.trajectoryInfo.index)!
+                : DefaultColor;
+        };
+    } else {
+        color = () => DefaultColor;
+    }
+
+    return {
+        factory: CellPackProvidedColorTheme,
+        granularity: 'instance',
+        color,
+        props,
+        description: Description,
+        legend
+    };
+}
+
+export const CellPackProvidedColorThemeProvider: ColorTheme.Provider<CellPackProvidedColorThemeParams, 'cellpack-provided'> = {
+    name: 'cellpack-provided',
+    label: 'CellPack Provided',
+    category: ColorTheme.Category.Chain,
+    factory: CellPackProvidedColorTheme,
+    getParams: getCellPackProvidedColorThemeParams,
+    defaultValues: PD.getDefaultValues(CellPackProvidedColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => {
+        return (
+            !!ctx.structure && ctx.structure.elementCount > 0 &&
+            ctx.structure.models[0].trajectoryInfo.size > 1 &&
+            !!CellPackInfoProvider.get(ctx.structure).value?.colors
+        );
+    }
+};

src/extensions/cellpack/index.ts

@@ -5,9 +5,9 @@
  */
 
 import { PluginBehavior } from '../../mol-plugin/behavior';
-import { CellPackColorThemeProvider } from './color';
 import { LoadCellPackModel } from './model';
-
+import { CellPackGenerateColorThemeProvider } from './color/generate';
+import { CellPackProvidedColorThemeProvider } from './color/provided';
 
 export const CellPack = PluginBehavior.create<{ autoAttach: boolean, showTooltip: boolean }>({
     name: 'cellpack',
@@ -19,12 +19,14 @@ export const CellPack = PluginBehavior.create<{ autoAttach: boolean, showTooltip
     ctor: class extends PluginBehavior.Handler<{ autoAttach: boolean, showTooltip: boolean }> {
         register(): void {
             this.ctx.state.data.actions.add(LoadCellPackModel);
-            this.ctx.representation.structure.themes.colorThemeRegistry.add(CellPackColorThemeProvider);
+            this.ctx.representation.structure.themes.colorThemeRegistry.add(CellPackGenerateColorThemeProvider);
+            this.ctx.representation.structure.themes.colorThemeRegistry.add(CellPackProvidedColorThemeProvider);
         }
 
         unregister() {
             this.ctx.state.data.actions.remove(LoadCellPackModel);
-            this.ctx.representation.structure.themes.colorThemeRegistry.remove(CellPackColorThemeProvider);
+            this.ctx.representation.structure.themes.colorThemeRegistry.remove(CellPackGenerateColorThemeProvider);
+            this.ctx.representation.structure.themes.colorThemeRegistry.remove(CellPackProvidedColorThemeProvider);
         }
     }
 });

src/extensions/cellpack/model.ts

@@ -378,27 +378,23 @@ export function createStructureFromCellPack(plugin: PluginContext, packing: Cell
         const assets: Asset.Wrapper[] = [];
         const trajCache = new TrajectoryCache();
         const structures: Structure[] = [];
-        let colors: Color[] = [];
-        let skip_color: boolean = false;
+        const colors: Color[] = [];
+        let skipColors: boolean = false;
         for (const iName in ingredients) {
             if (ctx.shouldUpdate) await ctx.update(iName);
             const ingredientStructure = await getIngredientStructure(plugin, ingredients[iName], baseUrl, ingredientFiles, trajCache);
             if (ingredientStructure) {
                 structures.push(ingredientStructure.structure);
                 assets.push(...ingredientStructure.assets);
-                if (ingredients[iName].color){
-                    let c = ingredients[iName].color;
-                    if (!c) c = Vec3.create(1, 0, 0);
+                const c = ingredients[iName].color;
+                if (c){
                     colors.push(Color.fromNormalizedRgb(c[0], c[1], c[2]));
                 } else {
-                    colors.push(Color.fromNormalizedRgb(1, 0, 0));
-                    skip_color = true;
+                    skipColors = true;
                 }
             }
         }
-        if (skip_color){
-            colors = [];
-        }
+
         if (ctx.shouldUpdate) await ctx.update(`${name} - units`);
         const builder = Structure.Builder({ label: name });
         let offsetInvariantId = 0;
@@ -420,7 +416,7 @@ export function createStructureFromCellPack(plugin: PluginContext, packing: Cell
             trajectoryInfo.size = il;
             trajectoryInfo.index = i;
         }
-        return { structure, assets, colors };
+        return { structure, assets, colors: skipColors ? undefined : colors };
     });
 }
 

src/extensions/cellpack/preset.ts

@@ -8,7 +8,9 @@ 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';
+import { CellPackGenerateColorThemeProvider } from './color/generate';
+import { CellPackInfoProvider } from './property';
+import { CellPackProvidedColorThemeProvider } from './color/provided';
 
 export const CellpackPackingPresetParams = {
     traceOnly: PD.Boolean(true),
@@ -40,8 +42,10 @@ export const CellpackPackingPreset = StructureRepresentationPresetProvider({
             Object.assign(reprProps, { sizeFactor: 2 });
         }
 
+        const info = structureCell.obj?.data && CellPackInfoProvider.get(structureCell.obj?.data).value;
+        const color = info?.colors ? CellPackProvidedColorThemeProvider.name : CellPackGenerateColorThemeProvider.name;
+
         const { update, builder, typeParams } = StructureRepresentationPresetProvider.reprBuilder(plugin, {});
-        const color = CellPackColorThemeProvider.name;
         const representations = {
             polymer: builder.buildRepresentation<any>(update, components.polymer, { type: params.representation, typeParams: { ...typeParams, ...reprProps }, color }, { tag: 'polymer' })
         };

src/extensions/cellpack/property.ts

@@ -14,12 +14,11 @@ import { CustomPropertyDescriptor } from '../../mol-model/custom-property';
 export type CellPackInfoValue = {
     packingsCount: number
     packingIndex: number
-    color: boolean
-    colors: Color[]
+    colors?: Color[]
 }
 
 const CellPackInfoParams = {
-    info: PD.Value<CellPackInfoValue>({ packingsCount: 1, packingIndex: 0, color: false, colors:[] }, { isHidden: true })
+    info: PD.Value<CellPackInfoValue>({ packingsCount: 1, packingIndex: 0, colors: undefined }, { isHidden: true })
 };
 type CellPackInfoParams = PD.Values<typeof CellPackInfoParams>
 

src/extensions/cellpack/state.ts

@@ -74,9 +74,8 @@ const StructureFromCellpack = PluginStateTransform.BuiltIn({
                 }
             }
             const { structure, assets, colors } = await createStructureFromCellPack(plugin, packing, params.baseUrl, ingredientFiles).runInContext(ctx);
-            const docolor: boolean = (colors.length !== 0);
             await CellPackInfoProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, structure, {
-                info: { packingsCount: a.data.packings.length, packingIndex: params.packing, colors: colors, color: docolor }
+                info: { packingsCount: a.data.packings.length, packingIndex: params.packing, colors }
             });
 
             (cache as any).assets = assets;