Add ColorTheme.palette support

- add example to basic-wrapper that uses it

src/examples/basic-wrapper/custom-theme.ts

+import { isPositionLocation } from '../../mol-geo/util/location-iterator';
+import { Vec3 } from '../../mol-math/linear-algebra';
+import { ColorTheme } from '../../mol-theme/color';
+import { ThemeDataContext } from '../../mol-theme/theme';
+import { Color } from '../../mol-util/color';
+import { ColorNames } from '../../mol-util/color/names';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+
+export function CustomColorTheme(
+    ctx: ThemeDataContext,
+    props: PD.Values<{}>
+): ColorTheme<{}> {
+    const { radius, center } = ctx.structure?.boundary.sphere!;
+    const radiusSq = Math.max(radius * radius, 0.001);
+    const scale = ColorTheme.PaletteScale;
+
+    return {
+        factory: CustomColorTheme,
+        granularity: 'vertex',
+        color: location => {
+            if (!isPositionLocation(location)) return ColorNames.black;
+            const dist = Vec3.squaredDistance(location.position, center);
+            const t = Math.min(dist / radiusSq, 1);
+            return ((t * scale) | 0) as Color;
+        },
+        palette: {
+            colors: [
+                ColorNames.red,
+                ColorNames.pink,
+                ColorNames.violet,
+                ColorNames.orange,
+                ColorNames.yellow,
+                ColorNames.green,
+                ColorNames.blue
+            ]
+        },
+        props: props,
+        description: '',
+    };
+}
+
+export const CustomColorThemeProvider: ColorTheme.Provider<{}, 'basic-wrapper-custom-color-theme'> = {
+    name: 'basic-wrapper-custom-color-theme',
+    label: 'Custom Color Theme',
+    category: ColorTheme.Category.Misc,
+    factory: CustomColorTheme,
+    getParams: () => ({}),
+    defaultValues: { },
+    isApplicable: (ctx: ThemeDataContext) => true,
+};

src/examples/basic-wrapper/index.html

@@ -97,6 +97,7 @@
             addHeader('Misc');
 
             addControl('Apply Stripes', () => BasicMolStarWrapper.coloring.applyStripes());
+            addControl('Apply Custom Theme', () => BasicMolStarWrapper.coloring.applyCustomTheme());
             addControl('Default Coloring', () => BasicMolStarWrapper.coloring.applyDefault());
 
             addHeader('Interactivity');

src/examples/basic-wrapper/index.ts

@@ -18,6 +18,7 @@ import { Asset } from '../../mol-util/assets';
 import { Color } from '../../mol-util/color';
 import { StripedResidues } from './coloring';
 import { CustomToastMessage } from './controls';
+import { CustomColorThemeProvider } from './custom-theme';
 import './index.html';
 import { buildStaticSuperposition, dynamicSuperpositionTest, StaticSuperpositionTestData } from './superposition';
 require('mol-plugin-ui/skin/light.scss');
@@ -42,6 +43,7 @@ class BasicWrapper {
         });
 
         this.plugin.representation.structure.themes.colorThemeRegistry.add(StripedResidues.colorThemeProvider!);
+        this.plugin.representation.structure.themes.colorThemeRegistry.add(CustomColorThemeProvider);
         this.plugin.managers.lociLabels.addProvider(StripedResidues.labelProvider!);
         this.plugin.customModelProperties.register(StripedResidues.propertyProvider, true);
     }
@@ -103,6 +105,13 @@ class BasicWrapper {
                 }
             });
         },
+        applyCustomTheme: async () => {
+            this.plugin.dataTransaction(async () => {
+                for (const s of this.plugin.managers.structure.hierarchy.current.structures) {
+                    await this.plugin.managers.structure.component.updateRepresentationsTheme(s.components, { color: CustomColorThemeProvider.name as any });
+                }
+            });
+        },
         applyDefault: async () => {
             this.plugin.dataTransaction(async () => {
                 for (const s of this.plugin.managers.structure.hierarchy.current.structures) {

src/mol-geo/geometry/color-data.ts

@@ -18,11 +18,24 @@ export type ColorType = 'uniform' | 'instance' | 'group' | 'groupInstance' | 've
 export type ColorData = {
     uColor: ValueCell<Vec3>,
     tColor: ValueCell<TextureImage<Uint8Array>>,
+    tPalette: ValueCell<TextureImage<Uint8Array>>,
     uColorTexDim: ValueCell<Vec2>,
     dColorType: ValueCell<string>,
+    dUsePalette: ValueCell<boolean>,
 }
 
 export function createColors(locationIt: LocationIterator, positionIt: LocationIterator, colorTheme: ColorTheme<any>, colorData?: ColorData): ColorData {
+    const data = _createColors(locationIt, positionIt, colorTheme, colorData);
+    if (colorTheme.palette) {
+        ValueCell.updateIfChanged(data.dUsePalette, true);
+        updatePaletteTexture(colorTheme.palette, data.tPalette);
+    } else {
+        ValueCell.updateIfChanged(data.dUsePalette, false);
+    }
+    return data;
+}
+
+function _createColors(locationIt: LocationIterator, positionIt: LocationIterator, colorTheme: ColorTheme<any>, colorData?: ColorData): ColorData {
     switch (Geometry.getGranularity(locationIt, colorTheme.granularity)) {
         case 'uniform': return createUniformColor(locationIt, colorTheme.color, colorData);
         case 'instance': return createInstanceColor(locationIt, colorTheme.color, colorData);
@@ -42,18 +55,20 @@ export function createValueColor(value: Color, colorData?: ColorData): ColorData
         return {
             uColor: ValueCell.create(Color.toVec3Normalized(Vec3(), value)),
             tColor: ValueCell.create({ array: new Uint8Array(3), width: 1, height: 1 }),
+            tPalette: ValueCell.create({ array: new Uint8Array(3), width: 1, height: 1 }),
             uColorTexDim: ValueCell.create(Vec2.create(1, 1)),
             dColorType: ValueCell.create('uniform'),
+            dUsePalette: ValueCell.create(false),
         };
     }
 }
 
 /** Creates color uniform */
-export function createUniformColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
+function createUniformColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     return createValueColor(color(NullLocation, false), colorData);
 }
 
-export function createTextureColor(colors: TextureImage<Uint8Array>, type: ColorType, colorData?: ColorData): ColorData {
+function createTextureColor(colors: TextureImage<Uint8Array>, type: ColorType, colorData?: ColorData): ColorData {
     if (colorData) {
         ValueCell.update(colorData.tColor, colors);
         ValueCell.update(colorData.uColorTexDim, Vec2.create(colors.width, colors.height));
@@ -63,14 +78,16 @@ export function createTextureColor(colors: TextureImage<Uint8Array>, type: Color
         return {
             uColor: ValueCell.create(Vec3()),
             tColor: ValueCell.create(colors),
+            tPalette: ValueCell.create({ array: new Uint8Array(3), width: 1, height: 1 }),
             uColorTexDim: ValueCell.create(Vec2.create(colors.width, colors.height)),
             dColorType: ValueCell.create(type),
+            dUsePalette: ValueCell.create(false),
         };
     }
 }
 
 /** Creates color texture with color for each instance */
-export function createInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
+function createInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     const { instanceCount } = locationIt;
     const colors = createTextureImage(Math.max(1, instanceCount), 3, Uint8Array, colorData && colorData.tColor.ref.value.array);
     locationIt.reset();
@@ -83,7 +100,7 @@ export function createInstanceColor(locationIt: LocationIterator, color: Locatio
 }
 
 /** Creates color texture with color for each group (i.e. shared across instances) */
-export function createGroupColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
+function createGroupColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     const { groupCount } = locationIt;
     const colors = createTextureImage(Math.max(1, groupCount), 3, Uint8Array, colorData && colorData.tColor.ref.value.array);
     locationIt.reset();
@@ -95,7 +112,7 @@ export function createGroupColor(locationIt: LocationIterator, color: LocationCo
 }
 
 /** Creates color texture with color for each group in each instance */
-export function createGroupInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
+function createGroupInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     const { groupCount, instanceCount } = locationIt;
     const count = instanceCount * groupCount;
     const colors = createTextureImage(Math.max(1, count), 3, Uint8Array, colorData && colorData.tColor.ref.value.array);
@@ -108,7 +125,7 @@ export function createGroupInstanceColor(locationIt: LocationIterator, color: Lo
 }
 
 /** Creates color texture with color for each vertex (i.e. shared across instances) */
-export function createVertexColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
+function createVertexColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     const { groupCount, stride } = locationIt;
     const colors = createTextureImage(Math.max(1, groupCount), 3, Uint8Array, colorData && colorData.tColor.ref.value.array);
     locationIt.reset();
@@ -124,7 +141,7 @@ export function createVertexColor(locationIt: LocationIterator, color: LocationC
 }
 
 /** Creates color texture with color for each vertex in each instance */
-export function createVertexInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
+function createVertexInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     const { groupCount, instanceCount, stride } = locationIt;
     const count = instanceCount * groupCount;
     const colors = createTextureImage(Math.max(1, count), 3, Uint8Array, colorData && colorData.tColor.ref.value.array);
@@ -138,3 +155,34 @@ export function createVertexInstanceColor(locationIt: LocationIterator, color: L
     }
     return createTextureColor(colors, 'vertexInstance', colorData);
 }
+
+function updatePaletteTexture(palette: ColorTheme.Palette, cell: ValueCell<TextureImage<Uint8Array>>) {
+    let isSynced = true;
+    const texture = cell.ref.value;
+    if (palette.colors.length !== texture.width) {
+        isSynced = false;
+    } else {
+        const data = texture.array;
+        let o = 0;
+        for (const c of palette.colors) {
+            const [r, g, b] = Color.toRgb(c);
+            if (data[o++] !== r || data[o++] !== g || data[o++] !== b) {
+                isSynced = false;
+                break;
+            }
+        }
+    }
+
+    if (isSynced) return;
+
+    const array = new Uint8Array(palette.colors.length * 3);
+    let o = 0;
+    for (const c of palette.colors) {
+        const [r, g, b] = Color.toRgb(c);
+        array[o++] = r;
+        array[o++] = g;
+        array[o++] = b;
+    }
+
+    ValueCell.update(cell, { array, height: 1, width: palette.colors.length });
+}
\ No newline at end of file

src/mol-gl/renderable/schema.ts

@@ -185,7 +185,9 @@ export const ColorSchema = {
     uColor: UniformSpec('v3', 'material'),
     uColorTexDim: UniformSpec('v2'),
     tColor: TextureSpec('image-uint8', 'rgb', 'ubyte', 'nearest'),
+    tPalette: TextureSpec('image-uint8', 'rgb', 'ubyte', 'linear'),
     dColorType: DefineSpec('string', ['uniform', 'attribute', 'instance', 'group', 'groupInstance', 'vertex', 'vertexInstance']),
+    dUsePalette: DefineSpec('boolean'),
 } as const;
 export type ColorSchema = typeof ColorSchema
 export type ColorValues = Values<ColorSchema>

src/mol-gl/shader/chunks/assign-color-varying.glsl.ts

@@ -14,6 +14,10 @@ export const assign_color_varying = `
         vColor.rgb = readFromTexture(tColor, int(aInstance) * uVertexCount + VertexID, uColorTexDim).rgb;
     #endif
 
+    #ifdef dUsePalette
+        vPaletteV = ((vColor.r * 256.0 * 256.0 * 255.0 + vColor.g * 256.0 * 255.0 + vColor.b * 255.0) - 1.0) / 16777215.0;
+    #endif
+
     #ifdef dOverpaint
         vOverpaint = readFromTexture(tOverpaint, aInstance * float(uGroupCount) + group, uOverpaintTexDim);
     #endif

src/mol-gl/shader/chunks/assign-material-color.glsl.ts

 export const assign_material_color = `
 #if defined(dRenderVariant_color)
-    #if defined(dColorType_uniform)
+    #if defined(dUsePalette)
+        vec4 material = vec4(texture2D(tPalette, vec2(vPaletteV, 0.5)).rgb, uAlpha);
+    #elif defined(dColorType_uniform)
         vec4 material = vec4(uColor, uAlpha);
     #elif defined(dColorType_varying)
         vec4 material = vec4(vColor.rgb, uAlpha);

src/mol-gl/shader/chunks/color-frag-params.glsl.ts

@@ -21,4 +21,9 @@ export const color_frag_params = `
     varying float vGroup;
     varying float vTransparency;
 #endif
+
+#ifdef dUsePalette
+    uniform sampler2D tPalette;
+    varying float vPaletteV;
+#endif
 `;
\ No newline at end of file

src/mol-gl/shader/chunks/color-vert-params.glsl.ts

@@ -30,4 +30,8 @@ export const color_vert_params = `
     uniform vec2 uTransparencyTexDim;
     uniform sampler2D tTransparency;
 #endif
+
+#ifdef dUsePalette
+    varying float vPaletteV;
+#endif
 `;
\ No newline at end of file

src/mol-theme/color.ts

@@ -44,6 +44,9 @@ interface ColorTheme<P extends PD.Params> {
     readonly granularity: ColorType
     readonly color: LocationColor
     readonly props: Readonly<PD.Values<P>>
+    // if palette is defined, 24bit RGB color value normalized to interval [0, 1]
+    // is used as index to the colors
+    readonly palette?: Readonly<ColorTheme.Palette>
     readonly contextHash?: number
     readonly description?: string
     readonly legend?: Readonly<ScaleLegend | TableLegend>
@@ -58,6 +61,12 @@ namespace ColorTheme {
         Misc = 'Miscellaneous',
     }
 
+    export interface Palette {
+        colors: Color[]
+    }
+
+    export const PaletteScale = (1 << 24) - 1;
+
     export type Props = { [k: string]: any }
     export type Factory<P extends PD.Params> = (ctx: ThemeDataContext, props: PD.Values<P>) => ColorTheme<P>
     export const EmptyFactory = () => Empty;