added atom-id and entity-id color theme

src/mol-theme/color.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -33,6 +33,8 @@ import { OccupancyColorThemeProvider } from './color/occupancy';
 import { OperatorNameColorThemeProvider } from './color/operator-name';
 import { OperatorHklColorThemeProvider } from './color/operator-hkl';
 import { PartialChargeColorThemeProvider } from './color/partial-charge';
+import { AtomIdColorThemeProvider } from './color/atom-id';
+import { EntityIdColorThemeProvider } from './color/entity-id';
 
 export type LocationColor = (location: Location, isSecondary: boolean) => Color
 
@@ -80,10 +82,12 @@ namespace ColorTheme {
     }
 
     export const BuiltIn = {
+        'atom-id': AtomIdColorThemeProvider,
         'carbohydrate-symbol': CarbohydrateSymbolColorThemeProvider,
         'chain-id': ChainIdColorThemeProvider,
         'element-index': ElementIndexColorThemeProvider,
         'element-symbol': ElementSymbolColorThemeProvider,
+        'entity-id': EntityIdColorThemeProvider,
         'entity-source': EntitySourceColorThemeProvider,
         'hydrophobicity': HydrophobicityColorThemeProvider,
         'illustrative': IllustrativeColorThemeProvider,

src/mol-theme/color/atom-id.ts

@@ -0,0 +1,90 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { StructureProperties, StructureElement, Bond, Structure } from '../../mol-model/structure';
+import { Color } from '../../mol-util/color';
+import { Location } from '../../mol-model/location';
+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 DefaultList = 'many-distinct';
+const DefaultColor = Color(0xFAFAFA);
+const Description = 'Gives every atom a color based on its `label_atom_id` value.';
+
+export const AtomIdColorThemeParams = {
+    ...getPaletteParams({ type: 'colors', colorList: DefaultList }),
+};
+export type AtomIdColorThemeParams = typeof AtomIdColorThemeParams
+export function getAtomIdColorThemeParams(ctx: ThemeDataContext) {
+    const params = PD.clone(AtomIdColorThemeParams);
+    return params;
+}
+
+function getAtomIdSerialMap(structure: Structure) {
+    const map = new Map<string, number>();
+    for (const m of structure.models) {
+        const { label_atom_id } = m.atomicHierarchy.atoms;
+        for (let i = 0, il = label_atom_id.rowCount; i < il; ++i) {
+            const id = label_atom_id.value(i);
+            if (!map.has(id)) map.set(id, map.size);
+        }
+    }
+    return map;
+}
+
+export function AtomIdColorTheme(ctx: ThemeDataContext, props: PD.Values<AtomIdColorThemeParams>): ColorTheme<AtomIdColorThemeParams> {
+    let color: LocationColor;
+    let legend: ScaleLegend | TableLegend | undefined;
+
+    if (ctx.structure) {
+        const l = StructureElement.Location.create(ctx.structure.root);
+        const atomIdSerialMap = getAtomIdSerialMap(ctx.structure.root);
+
+        const labelTable = Array.from(atomIdSerialMap.keys());
+        const valueLabel = (i: number) => labelTable[i];
+
+        const palette = getPalette(atomIdSerialMap.size, props, { valueLabel });
+        legend = palette.legend;
+
+        color = (location: Location): Color => {
+            let serial: number | undefined = undefined;
+            if (StructureElement.Location.is(location)) {
+                const id = StructureProperties.atom.label_atom_id(location);
+                serial = atomIdSerialMap.get(id);
+            } else if (Bond.isLocation(location)) {
+                l.unit = location.aUnit;
+                l.element = location.aUnit.elements[location.aIndex];
+                const id = StructureProperties.atom.label_atom_id(l);
+                serial = atomIdSerialMap.get(id);
+            }
+            return serial === undefined ? DefaultColor : palette.color(serial);
+        };
+    } else {
+        color = () => DefaultColor;
+    }
+
+    return {
+        factory: AtomIdColorTheme,
+        granularity: 'group',
+        color,
+        props,
+        description: Description,
+        legend
+    };
+}
+
+export const AtomIdColorThemeProvider: ColorTheme.Provider<AtomIdColorThemeParams, 'atom-id'> = {
+    name: 'atom-id',
+    label: 'Atom Id',
+    category: ColorTheme.Category.Atom,
+    factory: AtomIdColorTheme,
+    getParams: getAtomIdColorThemeParams,
+    defaultValues: PD.getDefaultValues(AtomIdColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure
+};

src/mol-theme/color/entity-id.ts

@@ -0,0 +1,98 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { StructureProperties, StructureElement, Bond, Structure } from '../../mol-model/structure';
+import { Color } from '../../mol-util/color';
+import { Location } from '../../mol-model/location';
+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 DefaultList = 'many-distinct';
+const DefaultColor = Color(0xFAFAFA);
+const Description = 'Gives every chain a color based on its `label_entity_id` value.';
+
+export const EntityIdColorThemeParams = {
+    ...getPaletteParams({ type: 'colors', colorList: DefaultList }),
+};
+export type EntityIdColorThemeParams = typeof EntityIdColorThemeParams
+export function getEntityIdColorThemeParams(ctx: ThemeDataContext) {
+    const params = PD.clone(EntityIdColorThemeParams);
+    return params;
+}
+
+function key(entityId: string, modelIndex: number) {
+    return `${entityId}|${modelIndex}`;
+}
+
+function getEntityIdSerialMap(structure: Structure) {
+    const map = new Map<string, number>();
+    for (let i = 0, il = structure.models.length; i < il; ++i) {
+        const { label_entity_id } = structure.models[i].atomicHierarchy.chains;
+        for (let j = 0, jl = label_entity_id.rowCount; j < jl; ++j) {
+            const k = key(label_entity_id.value(j), i);
+            if (!map.has(k)) map.set(k, map.size);
+        }
+    }
+    return map;
+}
+
+export function EntityIdColorTheme(ctx: ThemeDataContext, props: PD.Values<EntityIdColorThemeParams>): ColorTheme<EntityIdColorThemeParams> {
+    let color: LocationColor;
+    let legend: ScaleLegend | TableLegend | undefined;
+
+    if (ctx.structure) {
+        const l = StructureElement.Location.create(ctx.structure.root);
+        const entityIdSerialMap = getEntityIdSerialMap(ctx.structure.root);
+
+        const labelTable = Array.from(entityIdSerialMap.keys());
+        const valueLabel = (i: number) => labelTable[i];
+
+        const palette = getPalette(entityIdSerialMap.size, props, { valueLabel });
+        legend = palette.legend;
+
+        color = (location: Location): Color => {
+            let serial: number | undefined = undefined;
+            if (StructureElement.Location.is(location)) {
+                const atomId = StructureProperties.chain.label_entity_id(location);
+                const modelIndex = location.structure.models.indexOf(location.unit.model);
+                const k = key(atomId, modelIndex);
+                serial = entityIdSerialMap.get(k);
+            } else if (Bond.isLocation(location)) {
+                l.unit = location.aUnit;
+                l.element = location.aUnit.elements[location.aIndex];
+                const atomId = StructureProperties.chain.label_entity_id(l);
+                const modelIndex = l.structure.models.indexOf(l.unit.model);
+                const k = key(atomId, modelIndex);
+                serial = entityIdSerialMap.get(k);
+            }
+            return serial === undefined ? DefaultColor : palette.color(serial);
+        };
+    } else {
+        color = () => DefaultColor;
+    }
+
+    return {
+        factory: EntityIdColorTheme,
+        granularity: 'group',
+        color,
+        props,
+        description: Description,
+        legend
+    };
+}
+
+export const EntityIdColorThemeProvider: ColorTheme.Provider<EntityIdColorThemeParams, 'entity-id'> = {
+    name: 'entity-id',
+    label: 'Entity Id',
+    category: ColorTheme.Category.Chain,
+    factory: EntityIdColorTheme,
+    getParams: getEntityIdColorThemeParams,
+    defaultValues: PD.getDefaultValues(EntityIdColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure
+};