add structure-index color theme

CHANGELOG.md

@@ -8,8 +8,9 @@ Note that since we don't clearly distinguish between a public and private interf
 
 - [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 the new ``model-index`` and ``structure-index`` color themes as an option for the carbon color in the ``element-symbol`` and ``ilustrative`` color themes
+- Add ``structure-index`` color theme that uniquely colors each root structure
+- Add ``nearest`` method to ``Lookup3D``
 - Add mipmap-based blur for skybox backgrounds
 
 ## [v3.19.0] - 2022-10-01

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

@@ -1357,6 +1357,9 @@ namespace Structure {
     export type Index = number;
     export const Index = CustomStructureProperty.createSimple<Index>('index', 'root');
 
+    export type MaxIndex = number;
+    export const MaxIndex = CustomStructureProperty.createSimple<MaxIndex>('max_index', 'root');
+
     const PrincipalAxesProp = '__PrincipalAxes__';
     export function getPrincipalAxes(structure: Structure): PrincipalAxes {
         if (structure.currentPropertyData[PrincipalAxesProp]) return structure.currentPropertyData[PrincipalAxesProp];

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

@@ -65,6 +65,19 @@ export const StructureInfo = PluginBehavior.create({
             }
         }
 
+        private setStructureMaxIndex() {
+            const value = this.maxModelIndex;
+            const cells = this.ctx.state.data.select(StateSelection.Generators.rootsOfType(PluginStateObject.Molecule.Structure));
+            for (const c of cells) {
+                const s = c.obj?.data;
+                if (s) {
+                    if (Structure.MaxIndex.get(s).value !== value) {
+                        Structure.MaxIndex.set(s, { value }, value);
+                    }
+                }
+            }
+        }
+
         private handleModel(model: Model, oldModel?: Model) {
             if (Model.Index.get(model).value === undefined) {
                 const oldIndex = oldModel && Model.Index.get(oldModel).value;
@@ -107,10 +120,12 @@ export const StructureInfo = PluginBehavior.create({
             this.ctx.customModelProperties.register(Model.Index, true);
             this.ctx.customModelProperties.register(Model.MaxIndex, true);
             this.ctx.customStructureProperties.register(Structure.Index, true);
+            this.ctx.customStructureProperties.register(Structure.MaxIndex, true);
 
             this.subscribeObservable(this.ctx.state.data.events.object.created, o => {
                 this.handle(o.ref, o.obj);
                 this.setModelMaxIndex();
+                this.setStructureMaxIndex();
             });
 
             this.subscribeObservable(this.ctx.state.data.events.object.updated, o => {
@@ -123,6 +138,7 @@ export const StructureInfo = PluginBehavior.create({
             this.ctx.customModelProperties.unregister(Model.Index.descriptor.name);
             this.ctx.customModelProperties.unregister(Model.MaxIndex.descriptor.name);
             this.ctx.customStructureProperties.unregister(Structure.Index.descriptor.name);
+            this.ctx.customStructureProperties.unregister(Structure.MaxIndex.descriptor.name);
         }
     }
 });

src/mol-theme/color.ts

@@ -39,6 +39,7 @@ 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';
+import { StructureIndexColorThemeProvider } from './color/structure-index';
 
 export type LocationColor = (location: Location, isSecondary: boolean) => Color
 
@@ -145,6 +146,7 @@ namespace ColorTheme {
         'secondary-structure': SecondaryStructureColorThemeProvider,
         'sequence-id': SequenceIdColorThemeProvider,
         'shape-group': ShapeGroupColorThemeProvider,
+        'structure-index': StructureIndexColorThemeProvider,
         'trajectory-index': TrajectoryIndexColorThemeProvider,
         'uncertainty': UncertaintyColorThemeProvider,
         'unit-index': UnitIndexColorThemeProvider,

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

@@ -20,6 +20,7 @@ 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';
+import { StructureIndexColorTheme, StructureIndexColorThemeParams } from './structure-index';
 
 // from Jmol http://jmol.sourceforge.net/jscolors/ (or 0xFFFFFF)
 export const ElementSymbolColors = ColorMap({
@@ -37,6 +38,7 @@ export const ElementSymbolColorThemeParams = {
         'entity-source': PD.Group(EntitySourceColorThemeParams),
         'operator-name': PD.Group(OperatorNameColorThemeParams),
         'model-index': PD.Group(ModelIndexColorThemeParams),
+        'structure-index': PD.Group(StructureIndexColorThemeParams),
         'element-symbol': PD.EmptyGroup()
     }, { description: 'Use chain-id coloring for carbon atoms.' }),
     saturation: PD.Numeric(0, { min: -6, max: 6, step: 0.1 }),
@@ -66,8 +68,9 @@ export function ElementSymbolColorTheme(ctx: ThemeDataContext, props: PD.Values<
                 pcc.name === 'entity-source' ? EntitySourceColorTheme(ctx, pcc.params).color :
                     pcc.name === 'operator-name' ? OperatorNameColorTheme(ctx, pcc.params).color :
                         pcc.name === 'model-index' ? ModelIndexColorTheme(ctx, pcc.params).color :
-                            pcc.name === 'element-symbol' ? undefined :
-                                assertUnreachable(pcc);
+                            pcc.name === 'structure-index' ? StructureIndexColorTheme(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

@@ -18,6 +18,7 @@ import { EntityIdColorTheme, EntityIdColorThemeParams } from './entity-id';
 import { MoleculeTypeColorTheme, MoleculeTypeColorThemeParams } from './molecule-type';
 import { EntitySourceColorTheme, EntitySourceColorThemeParams } from './entity-source';
 import { ModelIndexColorTheme, ModelIndexColorThemeParams } from './model-index';
+import { StructureIndexColorTheme, StructureIndexColorThemeParams } from './structure-index';
 
 const DefaultIllustrativeColor = Color(0xEEEEEE);
 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).`;
@@ -30,6 +31,7 @@ export const IllustrativeColorThemeParams = {
         'entity-source': PD.Group(EntitySourceColorThemeParams),
         'molecule-type': PD.Group(MoleculeTypeColorThemeParams),
         'model-index': PD.Group(ModelIndexColorThemeParams),
+        'structure-index': PD.Group(StructureIndexColorThemeParams),
     }),
     carbonLightness: PD.Numeric(0.8, { min: -6, max: 6, step: 0.1 })
 };
@@ -47,7 +49,8 @@ export function IllustrativeColorTheme(ctx: ThemeDataContext, props: PD.Values<I
                     props.style.name === 'entity-source' ? EntitySourceColorTheme(ctx, props.style.params) :
                         props.style.name === 'molecule-type' ? MoleculeTypeColorTheme(ctx, props.style.params) :
                             props.style.name === 'model-index' ? ModelIndexColorTheme(ctx, props.style.params) :
-                                assertUnreachable(props.style);
+                                props.style.name === 'structure-index' ? StructureIndexColorTheme(ctx, props.style.params) :
+                                    assertUnreachable(props.style);
 
     function illustrativeColor(location: Location, typeSymbol: ElementSymbol) {
         const baseColor = styleColor(location, false);

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

@@ -0,0 +1,67 @@
+/**
+ * Copyright (c) 2022 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, Structure } from '../../mol-model/structure';
+import { ColorTheme, LocationColor } from '../color';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { ThemeDataContext } from '../../mol-theme/theme';
+import { getPaletteParams, getPalette } from '../../mol-util/color/palette';
+import { TableLegend, ScaleLegend } from '../../mol-util/legend';
+
+const DefaultColor = Color(0xCCCCCC);
+const Description = 'Gives every structure a unique color based on its index.';
+
+export const StructureIndexColorThemeParams = {
+    ...getPaletteParams({ type: 'colors', colorList: 'many-distinct' }),
+};
+export type StructureIndexColorThemeParams = typeof StructureIndexColorThemeParams
+export function getStructureIndexColorThemeParams(ctx: ThemeDataContext) {
+    return PD.clone(StructureIndexColorThemeParams);
+}
+
+export function StructureIndexColorTheme(ctx: ThemeDataContext, props: PD.Values<StructureIndexColorThemeParams>): ColorTheme<StructureIndexColorThemeParams> {
+    let color: LocationColor;
+    let legend: ScaleLegend | TableLegend | undefined;
+
+    if (ctx.structure) {
+        const size = (Structure.MaxIndex.get(ctx.structure).value ?? -1) + 1;
+
+        const palette = getPalette(size, props);
+        legend = palette.legend;
+
+        color = (location: Location): Color => {
+            if (StructureElement.Location.is(location)) {
+                return palette.color(Structure.Index.get(location.structure).value || 0)!;
+            } else if (Bond.isLocation(location)) {
+                return palette.color(Structure.Index.get(location.aStructure).value || 0)!;
+            }
+            return DefaultColor;
+        };
+    } else {
+        color = () => DefaultColor;
+    }
+
+    return {
+        factory: StructureIndexColorTheme,
+        granularity: 'instance',
+        color,
+        props,
+        description: Description,
+        legend
+    };
+}
+
+export const StructureIndexColorThemeProvider: ColorTheme.Provider<StructureIndexColorThemeParams, 'structure-index'> = {
+    name: 'structure-index',
+    label: 'Structure Index',
+    category: ColorTheme.Category.Chain,
+    factory: StructureIndexColorTheme,
+    getParams: getStructureIndexColorThemeParams,
+    defaultValues: PD.getDefaultValues(StructureIndexColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.elementCount > 0
+};