Merge pull request #583 from jpattle/model-index-carbon-color

Model index updates & carbon color

CHANGELOG.md

@@ -6,6 +6,9 @@ Note that since we don't clearly distinguish between a public and private interf
 
 ## [Unreleased]
 
+- [Breaking] Rename the ``model-index`` color theme to ``trajectory-index``
+- Add a new ``model-index`` color theme that uniquely colors each loaded model
+- Add the new ``model-index`` color theme as an option for the carbon color in the ``element-symbol`` and ``ilustrative`` color themes
 - Add nearest method to lookup3d
 - Add mipmap-based blur for skybox backgrounds
 

src/mol-model-props/common/custom-model-property.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -71,12 +71,12 @@ namespace CustomModelProperty {
             },
             ref: (data: Model, add: boolean) => data.customProperties.reference(builder.descriptor, add),
             get: (data: Model) => get(data)?.data,
-            set: (data: Model, props: Partial<PD.Values<Params>> = {}) => {
+            set: (data: Model, props: Partial<PD.Values<Params>> = {}, value?: Value) => {
                 const property = get(data);
                 const p = PD.merge(builder.defaultParams, property.props, props);
                 if (!PD.areEqual(builder.defaultParams, property.props, p)) {
                     // this invalidates property.value
-                    set(data, p, undefined);
+                    set(data, p, value);
                     // dispose of assets
                     data.customProperties.assets(builder.descriptor);
                 }
@@ -96,7 +96,7 @@ namespace CustomModelProperty {
             getParams: () => ({ value: PD.Value(defaultValue, { isHidden: true }) }),
             isApplicable: () => true,
             obtain: async (ctx: CustomProperty.Context, data: Model, props: Partial<PD.Values<typeof defaultParams>>) => {
-                return { value: props.value ?? defaultValue };
+                return { ...PD.getDefaultValues(defaultParams), ...props };
             }
         });
     }

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

@@ -213,6 +213,9 @@ export namespace Model {
     export type Index = number;
     export const Index = CustomModelProperty.createSimple<Index>('index', 'static');
 
+    export type MaxIndex = number;
+    export const MaxIndex = CustomModelProperty.createSimple<MaxIndex>('max_index', 'static');
+
     export function getRoot(model: Model) {
         return model.parent || model;
     }

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

@@ -86,7 +86,7 @@ const allModels = TrajectoryHierarchyPresetProvider({
     id: 'preset-trajectory-all-models',
     display: {
         name: 'All Models', group: 'Preset',
-        description: 'Shows all models; colored by model-index.'
+        description: 'Shows all models; colored by trajectory-index.'
     },
     isApplicable: o => {
         return o.data.frameCount > 1;
@@ -115,7 +115,7 @@ const allModels = TrajectoryHierarchyPresetProvider({
 
             const quality = structure.obj ? getStructureQuality(structure.obj.data, { elementCountFactor: tr.frameCount }) : 'medium';
             const representationPreset = params.representationPreset || plugin.config.get(PluginConfig.Structure.DefaultRepresentationPreset) || PresetStructureRepresentations.auto.id;
-            await builder.representation.applyPreset(structureProperties, representationPreset, { theme: { globalName: 'model-index' }, quality });
+            await builder.representation.applyPreset(structureProperties, representationPreset, { theme: { globalName: 'trajectory-index' }, quality });
         }
 
         return { models, structures };

src/mol-plugin/behavior/dynamic/custom-props/structure-info.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2020-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -52,17 +52,30 @@ export const StructureInfo = PluginBehavior.create({
             return { auth, label };
         }
 
+        private setModelMaxIndex() {
+            const value = this.maxModelIndex;
+            const cells = this.ctx.state.data.select(StateSelection.Generators.rootsOfType(PluginStateObject.Molecule.Model));
+            for (const c of cells) {
+                const m = c.obj?.data;
+                if (m) {
+                    if (Model.MaxIndex.get(m).value !== value) {
+                        Model.MaxIndex.set(m, { value }, value);
+                    }
+                }
+            }
+        }
+
         private handleModel(model: Model, oldModel?: Model) {
             if (Model.Index.get(model).value === undefined) {
                 const oldIndex = oldModel && Model.Index.get(oldModel).value;
                 const value = oldIndex ?? (this.maxModelIndex + 1);
-                Model.Index.set(model, { value });
+                Model.Index.set(model, { value }, value);
             }
 
             if (Model.AsymIdOffset.get(model).value === undefined) {
                 const oldOffset = oldModel && Model.AsymIdOffset.get(oldModel).value;
                 const value = oldOffset ?? { ...this.asymIdOffset };
-                Model.AsymIdOffset.set(model, { value });
+                Model.AsymIdOffset.set(model, { value }, value);
             }
         }
 
@@ -72,7 +85,7 @@ export const StructureInfo = PluginBehavior.create({
 
             const oldIndex = oldStructure && Structure.Index.get(oldStructure).value;
             const value = oldIndex ?? (this.maxStructureIndex + 1);
-            Structure.Index.set(structure, { value });
+            Structure.Index.set(structure, { value }, value);
         }
 
         private handle(ref: string, obj: StateObject<any, StateObject.Type<any>>, oldObj?: StateObject<any, StateObject.Type<any>>) {
@@ -92,10 +105,12 @@ export const StructureInfo = PluginBehavior.create({
         register(): void {
             this.ctx.customModelProperties.register(Model.AsymIdOffset, true);
             this.ctx.customModelProperties.register(Model.Index, true);
+            this.ctx.customModelProperties.register(Model.MaxIndex, true);
             this.ctx.customStructureProperties.register(Structure.Index, true);
 
             this.subscribeObservable(this.ctx.state.data.events.object.created, o => {
                 this.handle(o.ref, o.obj);
+                this.setModelMaxIndex();
             });
 
             this.subscribeObservable(this.ctx.state.data.events.object.updated, o => {
@@ -106,6 +121,7 @@ export const StructureInfo = PluginBehavior.create({
         unregister() {
             this.ctx.customModelProperties.unregister(Model.AsymIdOffset.descriptor.name);
             this.ctx.customModelProperties.unregister(Model.Index.descriptor.name);
+            this.ctx.customModelProperties.unregister(Model.MaxIndex.descriptor.name);
             this.ctx.customStructureProperties.unregister(Structure.Index.descriptor.name);
         }
     }

src/mol-theme/color.ts

@@ -28,7 +28,7 @@ import { UncertaintyColorThemeProvider } from './color/uncertainty';
 import { EntitySourceColorThemeProvider } from './color/entity-source';
 import { IllustrativeColorThemeProvider } from './color/illustrative';
 import { HydrophobicityColorThemeProvider } from './color/hydrophobicity';
-import { ModelIndexColorThemeProvider } from './color/model-index';
+import { TrajectoryIndexColorThemeProvider } from './color/trajectory-index';
 import { OccupancyColorThemeProvider } from './color/occupancy';
 import { OperatorNameColorThemeProvider } from './color/operator-name';
 import { OperatorHklColorThemeProvider } from './color/operator-hkl';
@@ -38,6 +38,7 @@ import { EntityIdColorThemeProvider } from './color/entity-id';
 import { Texture, TextureFilter } from '../mol-gl/webgl/texture';
 import { VolumeValueColorThemeProvider } from './color/volume-value';
 import { Vec3, Vec4 } from '../mol-math/linear-algebra';
+import { ModelIndexColorThemeProvider } from './color/model-index';
 
 export type LocationColor = (location: Location, isSecondary: boolean) => Color
 
@@ -144,6 +145,7 @@ namespace ColorTheme {
         'secondary-structure': SecondaryStructureColorThemeProvider,
         'sequence-id': SequenceIdColorThemeProvider,
         'shape-group': ShapeGroupColorThemeProvider,
+        'trajectory-index': TrajectoryIndexColorThemeProvider,
         'uncertainty': UncertaintyColorThemeProvider,
         'unit-index': UnitIndexColorThemeProvider,
         'uniform': UniformColorThemeProvider,

src/mol-theme/color/element-symbol.ts

@@ -19,6 +19,7 @@ import { OperatorNameColorThemeParams, OperatorNameColorTheme } from './operator
 import { EntityIdColorTheme, EntityIdColorThemeParams } from './entity-id';
 import { assertUnreachable } from '../../mol-util/type-helpers';
 import { EntitySourceColorTheme, EntitySourceColorThemeParams } from './entity-source';
+import { ModelIndexColorTheme, ModelIndexColorThemeParams } from './model-index';
 
 // from Jmol http://jmol.sourceforge.net/jscolors/ (or 0xFFFFFF)
 export const ElementSymbolColors = ColorMap({
@@ -35,6 +36,7 @@ export const ElementSymbolColorThemeParams = {
         'entity-id': PD.Group(EntityIdColorThemeParams),
         'entity-source': PD.Group(EntitySourceColorThemeParams),
         'operator-name': PD.Group(OperatorNameColorThemeParams),
+        'model-index': PD.Group(ModelIndexColorThemeParams),
         'element-symbol': PD.EmptyGroup()
     }, { description: 'Use chain-id coloring for carbon atoms.' }),
     saturation: PD.Numeric(0, { min: -6, max: 6, step: 0.1 }),
@@ -46,7 +48,7 @@ export const ElementSymbolColorThemeParams = {
 };
 export type ElementSymbolColorThemeParams = typeof ElementSymbolColorThemeParams
 export function getElementSymbolColorThemeParams(ctx: ThemeDataContext) {
-    return ElementSymbolColorThemeParams; // TODO return copy
+    return PD.clone(ElementSymbolColorThemeParams);
 }
 
 export function elementSymbolColor(colorMap: ElementSymbolColors, element: ElementSymbol): Color {
@@ -63,8 +65,9 @@ export function ElementSymbolColorTheme(ctx: ThemeDataContext, props: PD.Values<
             pcc.name === 'entity-id' ? EntityIdColorTheme(ctx, pcc.params).color :
                 pcc.name === 'entity-source' ? EntitySourceColorTheme(ctx, pcc.params).color :
                     pcc.name === 'operator-name' ? OperatorNameColorTheme(ctx, pcc.params).color :
-                        pcc.name === 'element-symbol' ? undefined :
-                            assertUnreachable(pcc);
+                        pcc.name === 'model-index' ? ModelIndexColorTheme(ctx, pcc.params).color :
+                            pcc.name === 'element-symbol' ? undefined :
+                                assertUnreachable(pcc);
 
     function elementColor(element: ElementSymbol, location: Location) {
         return (carbonColor && element === 'C')

src/mol-theme/color/illustrative.ts

@@ -17,9 +17,10 @@ import { assertUnreachable } from '../../mol-util/type-helpers';
 import { EntityIdColorTheme, EntityIdColorThemeParams } from './entity-id';
 import { MoleculeTypeColorTheme, MoleculeTypeColorThemeParams } from './molecule-type';
 import { EntitySourceColorTheme, EntitySourceColorThemeParams } from './entity-source';
+import { ModelIndexColorTheme, ModelIndexColorThemeParams } from './model-index';
 
 const DefaultIllustrativeColor = Color(0xEEEEEE);
-const Description = `Assigns an illustrative color that gives every chain a color based on the choosen style but with lighter carbons (inspired by David Goodsell's Molecule of the Month style).`;
+const Description = `Assigns an illustrative color that gives every chain a color based on the chosen style but with lighter carbons (inspired by David Goodsell's Molecule of the Month style).`;
 
 export const IllustrativeColorThemeParams = {
     style: PD.MappedStatic('entity-id', {
@@ -28,6 +29,7 @@ export const IllustrativeColorThemeParams = {
         'entity-id': PD.Group(EntityIdColorThemeParams),
         'entity-source': PD.Group(EntitySourceColorThemeParams),
         'molecule-type': PD.Group(MoleculeTypeColorThemeParams),
+        'model-index': PD.Group(ModelIndexColorThemeParams),
     }),
     carbonLightness: PD.Numeric(0.8, { min: -6, max: 6, step: 0.1 })
 };
@@ -44,7 +46,8 @@ export function IllustrativeColorTheme(ctx: ThemeDataContext, props: PD.Values<I
                 props.style.name === 'entity-id' ? EntityIdColorTheme(ctx, props.style.params) :
                     props.style.name === 'entity-source' ? EntitySourceColorTheme(ctx, props.style.params) :
                         props.style.name === 'molecule-type' ? MoleculeTypeColorTheme(ctx, props.style.params) :
-                            assertUnreachable(props.style);
+                            props.style.name === 'model-index' ? ModelIndexColorTheme(ctx, props.style.params) :
+                                assertUnreachable(props.style);
 
     function illustrativeColor(location: Location, typeSymbol: ElementSymbol) {
         const baseColor = styleColor(location, false);

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

@@ -1,6 +1,7 @@
 /**
- * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
+ * @author Jason Pattle <jpattle@exscientia.co.uk>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
@@ -14,14 +15,14 @@ import { getPaletteParams, getPalette } from '../../mol-util/color/palette';
 import { TableLegend, ScaleLegend } from '../../mol-util/legend';
 
 const DefaultColor = Color(0xCCCCCC);
-const Description = 'Gives every model a unique color based on the position (index) of the model in the list of models in the structure.';
+const Description = 'Gives every model a unique color based on its index.';
 
 export const ModelIndexColorThemeParams = {
-    ...getPaletteParams({ type: 'colors', colorList: 'purples' }),
+    ...getPaletteParams({ type: 'colors', colorList: 'many-distinct' }),
 };
 export type ModelIndexColorThemeParams = typeof ModelIndexColorThemeParams
 export function getModelIndexColorThemeParams(ctx: ThemeDataContext) {
-    return ModelIndexColorThemeParams; // TODO return copy
+    return PD.clone(ModelIndexColorThemeParams);
 }
 
 export function ModelIndexColorTheme(ctx: ThemeDataContext, props: PD.Values<ModelIndexColorThemeParams>): ColorTheme<ModelIndexColorThemeParams> {
@@ -29,24 +30,17 @@ export function ModelIndexColorTheme(ctx: ThemeDataContext, props: PD.Values<Mod
     let legend: ScaleLegend | TableLegend | undefined;
 
     if (ctx.structure) {
-        const { models } = ctx.structure.root;
-
-        let size = 0;
-        for (const m of models) size = Math.max(size, Model.TrajectoryInfo.get(m)?.size || 0);
+        // max-index is the same for all models
+        const size = (Model.MaxIndex.get(ctx.structure.models[0]).value ?? -1) + 1;
 
         const palette = getPalette(size, props);
         legend = palette.legend;
-        const modelColor = new Map<number, Color>();
-        for (let i = 0, il = models.length; i < il; ++i) {
-            const idx = Model.TrajectoryInfo.get(models[i])?.index || 0;
-            modelColor.set(idx, palette.color(idx));
-        }
 
         color = (location: Location): Color => {
             if (StructureElement.Location.is(location)) {
-                return modelColor.get(Model.TrajectoryInfo.get(location.unit.model).index)!;
+                return palette.color(Model.Index.get(location.unit.model).value || 0)!;
             } else if (Bond.isLocation(location)) {
-                return modelColor.get(Model.TrajectoryInfo.get(location.aUnit.model).index)!;
+                return palette.color(Model.Index.get(location.aUnit.model).value || 0)!;
             }
             return DefaultColor;
         };
@@ -71,5 +65,5 @@ export const ModelIndexColorThemeProvider: ColorTheme.Provider<ModelIndexColorTh
     factory: ModelIndexColorTheme,
     getParams: getModelIndexColorThemeParams,
     defaultValues: PD.getDefaultValues(ModelIndexColorThemeParams),
-    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.elementCount > 0 && Model.TrajectoryInfo.get(ctx.structure.models[0]).size > 1
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.elementCount > 0
 };

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

@@ -0,0 +1,75 @@
+/**
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Color } from '../../mol-util/color';
+import { Location } from '../../mol-model/location';
+import { StructureElement, Bond, Model } from '../../mol-model/structure';
+import { ColorTheme, LocationColor } from '../color';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { ThemeDataContext } from '../theme';
+import { getPaletteParams, getPalette } from '../../mol-util/color/palette';
+import { TableLegend, ScaleLegend } from '../../mol-util/legend';
+
+const DefaultColor = Color(0xCCCCCC);
+const Description = 'Gives every model (frame) a unique color based on the index in its trajectory.';
+
+export const TrajectoryIndexColorThemeParams = {
+    ...getPaletteParams({ type: 'colors', colorList: 'purples' }),
+};
+export type TrajectoryIndexColorThemeParams = typeof TrajectoryIndexColorThemeParams
+export function getTrajectoryIndexColorThemeParams(ctx: ThemeDataContext) {
+    return PD.clone(TrajectoryIndexColorThemeParams);
+}
+
+export function TrajectoryIndexColorTheme(ctx: ThemeDataContext, props: PD.Values<TrajectoryIndexColorThemeParams>): ColorTheme<TrajectoryIndexColorThemeParams> {
+    let color: LocationColor;
+    let legend: ScaleLegend | TableLegend | undefined;
+
+    if (ctx.structure) {
+        const { models } = ctx.structure.root;
+
+        let size = 0;
+        for (const m of models) size = Math.max(size, Model.TrajectoryInfo.get(m)?.size || 0);
+
+        const palette = getPalette(size, props);
+        legend = palette.legend;
+        const modelColor = new Map<number, Color>();
+        for (let i = 0, il = models.length; i < il; ++i) {
+            const idx = Model.TrajectoryInfo.get(models[i])?.index || 0;
+            modelColor.set(idx, palette.color(idx));
+        }
+
+        color = (location: Location): Color => {
+            if (StructureElement.Location.is(location)) {
+                return modelColor.get(Model.TrajectoryInfo.get(location.unit.model).index)!;
+            } else if (Bond.isLocation(location)) {
+                return modelColor.get(Model.TrajectoryInfo.get(location.aUnit.model).index)!;
+            }
+            return DefaultColor;
+        };
+    } else {
+        color = () => DefaultColor;
+    }
+
+    return {
+        factory: TrajectoryIndexColorTheme,
+        granularity: 'instance',
+        color,
+        props,
+        description: Description,
+        legend
+    };
+}
+
+export const TrajectoryIndexColorThemeProvider: ColorTheme.Provider<TrajectoryIndexColorThemeParams, 'trajectory-index'> = {
+    name: 'trajectory-index',
+    label: 'Trajectory Index',
+    category: ColorTheme.Category.Chain,
+    factory: TrajectoryIndexColorTheme,
+    getParams: getTrajectoryIndexColorThemeParams,
+    defaultValues: PD.getDefaultValues(TrajectoryIndexColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.elementCount > 0 && Model.TrajectoryInfo.get(ctx.structure.models[0]).size > 1
+};

src/mol-util/color/distinct.ts

@@ -13,6 +13,7 @@ import { deepClone } from '../../mol-util/object';
 import { deepEqual } from '../../mol-util';
 import { arraySum } from '../../mol-util/array';
 import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { ColorNames } from './names';
 
 export const DistinctColorsParams = {
     hue: PD.Interval([1, 360], { min: 0, max: 360, step: 1 }),
@@ -105,7 +106,8 @@ export function distinctColors(count: number, props: Partial<DistinctColorsProps
 
     const samples = getSamples(Math.max(p.minSampleCount, count * 5), p);
     if (samples.length < count) {
-        throw new Error('Not enough samples to generate distinct colors, increase sample count.');
+        console.warn('Not enough samples to generate distinct colors, increase sample count.');
+        return (new Array(count)).fill(ColorNames.lightgrey);
     }
 
     const colors: Lab[] = [];