better color theme typing

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

@@ -12,10 +12,12 @@ import { Vec2, Vec3, Vec4 } from '../../mol-math/linear-algebra';
 import { LocationIterator } from '../util/location-iterator';
 import { NullLocation } from '../../mol-model/location';
 import { LocationColor, ColorTheme, ColorVolume } from '../../mol-theme/color';
-import { Geometry } from './geometry';
 import { createNullTexture, Texture } from '../../mol-gl/webgl/texture';
 
-export type ColorType = 'uniform' | 'instance' | 'group' | 'groupInstance' | 'vertex' | 'vertexInstance' | 'volume' | 'volumeInstance' | 'direct'
+export type ColorTypeLocation = 'uniform' | 'instance' | 'group' | 'groupInstance' | 'vertex' | 'vertexInstance';
+export type ColorTypeGrid = 'volume' | 'volumeInstance';
+export type ColorTypeDirect = 'direct';
+export type ColorType = ColorTypeLocation | ColorTypeGrid | ColorTypeDirect;
 
 export type ColorData = {
     uColor: ValueCell<Vec3>,
@@ -41,25 +43,18 @@ export function createColors(locationIt: LocationIterator, positionIt: LocationI
 }
 
 function _createColors(locationIt: LocationIterator, positionIt: LocationIterator, colorTheme: ColorTheme<any>, colorData?: ColorData): ColorData {
-    switch (Geometry.getGranularity(locationIt, colorTheme.granularity)) {
+    switch (colorTheme.granularity) {
         case 'uniform': return createUniformColor(locationIt, colorTheme.color, colorData);
-        case 'instance': return createInstanceColor(locationIt, colorTheme.color, colorData);
+        case 'instance':
+            return locationIt.nonInstanceable
+                ? createInstanceColor(locationIt, colorTheme.color, colorData)
+                : createGroupColor(locationIt, colorTheme.color, colorData);
         case 'group': return createGroupColor(locationIt, colorTheme.color, colorData);
         case 'groupInstance': return createGroupInstanceColor(locationIt, colorTheme.color, colorData);
         case 'vertex': return createVertexColor(positionIt, colorTheme.color, colorData);
         case 'vertexInstance': return createVertexInstanceColor(positionIt, colorTheme.color, colorData);
-        case 'volume':
-            if (colorTheme.grid) {
-                return createGridColor(colorTheme.grid, 'volume', colorData);
-            } else {
-                throw new Error('Grid missing for "volume" color theme');
-            }
-        case 'volumeInstance':
-            if (colorTheme.grid) {
-                return createGridColor(colorTheme.grid, 'volumeInstance', colorData);
-            } else {
-                throw new Error('Grid missing for "volume" color theme');
-            }
+        case 'volume': return createGridColor(colorTheme.grid, 'volume', colorData);
+        case 'volumeInstance': return createGridColor(colorTheme.grid, 'volumeInstance', colorData);
         case 'direct': return createDirectColor(colorData);
     }
 }

src/mol-repr/volume/slice.ts

@@ -32,7 +32,7 @@ export async function createImage(ctx: VisualContext, volume: Volume, theme: The
     const isoVal = Volume.IsoValue.toAbsolute(isoValue, volume.grid.stats).absoluteValue;
 
     // TODO more color themes
-    const color = theme.color.color(NullLocation, false);
+    const color = 'color' in theme.color ? theme.color.color(NullLocation, false) : Color(0xffffff);
     const [r, g, b] = Color.toRgbNormalized(color);
 
     const {

src/mol-theme/color.ts

@@ -6,7 +6,7 @@
 
 import { Color } from '../mol-util/color';
 import { Location } from '../mol-model/location';
-import { ColorType } from '../mol-geo/geometry/color-data';
+import { ColorType, ColorTypeDirect, ColorTypeGrid, ColorTypeLocation } from '../mol-geo/geometry/color-data';
 import { CarbohydrateSymbolColorThemeProvider } from './color/carbohydrate-symbol';
 import { UniformColorThemeProvider } from './color/uniform';
 import { deepEqual } from '../mol-util';
@@ -48,12 +48,10 @@ export interface ColorVolume {
 }
 
 export { ColorTheme };
-interface ColorTheme<P extends PD.Params> {
+
+type ColorThemeShared<P extends PD.Params> = {
     readonly factory: ColorTheme.Factory<P>
-    readonly granularity: ColorType
-    readonly color: LocationColor
     readonly props: Readonly<PD.Values<P>>
-    readonly grid?: ColorVolume
     /**
      * if palette is defined, 24bit RGB color value normalized to interval [0, 1]
      * is used as index to the colors
@@ -64,6 +62,26 @@ interface ColorTheme<P extends PD.Params> {
     readonly description?: string
     readonly legend?: Readonly<ScaleLegend | TableLegend>
 }
+
+type ColorThemeLocation<P extends PD.Params> = {
+    readonly granularity: ColorTypeLocation
+    readonly color: LocationColor
+} & ColorThemeShared<P>
+
+type ColorThemeGrid<P extends PD.Params> = {
+    readonly granularity: ColorTypeGrid
+    readonly grid: ColorVolume
+} & ColorThemeShared<P>
+
+type ColorThemeDirect<P extends PD.Params> = {
+    readonly granularity: ColorTypeDirect
+} & ColorThemeShared<P>
+
+type ColorTheme<P extends PD.Params, G extends ColorType = any> =
+    G extends ColorTypeLocation ? ColorThemeLocation<P> :
+        G extends ColorTypeGrid ? ColorThemeGrid<P> :
+            G extends ColorTypeDirect ? ColorThemeDirect<P> : never
+
 namespace ColorTheme {
     export const enum Category {
         Atom = 'Atom Property',

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

@@ -12,6 +12,7 @@ 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';
+import { ColorTypeLocation } from '../../mol-geo/geometry/color-data';
 
 const DefaultList = 'many-distinct';
 const DefaultColor = Color(0xFAFAFA);
@@ -75,7 +76,7 @@ function getAsymIdSerialMap(structure: Structure, type: AsymIdType) {
     return map;
 }
 
-export function ChainIdColorTheme(ctx: ThemeDataContext, props: PD.Values<ChainIdColorThemeParams>): ColorTheme<ChainIdColorThemeParams> {
+export function ChainIdColorTheme(ctx: ThemeDataContext, props: PD.Values<ChainIdColorThemeParams>): ColorTheme<ChainIdColorThemeParams, ColorTypeLocation> {
     let color: LocationColor;
     let legend: ScaleLegend | TableLegend | undefined;
 

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

@@ -12,6 +12,7 @@ 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';
+import { ColorTypeLocation } from '../../mol-geo/geometry/color-data';
 
 const DefaultList = 'many-distinct';
 const DefaultColor = Color(0xFAFAFA);
@@ -65,7 +66,7 @@ function getEntityId(location: StructureElement.Location): string {
     }
 }
 
-export function EntityIdColorTheme(ctx: ThemeDataContext, props: PD.Values<EntityIdColorThemeParams>): ColorTheme<EntityIdColorThemeParams> {
+export function EntityIdColorTheme(ctx: ThemeDataContext, props: PD.Values<EntityIdColorThemeParams>): ColorTheme<EntityIdColorThemeParams, ColorTypeLocation> {
     let color: LocationColor;
     let legend: ScaleLegend | TableLegend | undefined;
 

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

@@ -17,6 +17,7 @@ import { TableLegend, ScaleLegend } from '../../mol-util/legend';
 import { isInteger } from '../../mol-util/number';
 import { ColorLists, getColorListFromName } from '../../mol-util/color/lists';
 import { MmcifFormat } from '../../mol-model-formats/structure/mmcif';
+import { ColorTypeLocation } from '../../mol-geo/geometry/color-data';
 
 const DefaultList = 'dark-2';
 const DefaultColor = Color(0xFAFAFA);
@@ -121,7 +122,7 @@ function getLabelTable(srcKeySerialMap: Map<string, number>) {
     });
 }
 
-export function EntitySourceColorTheme(ctx: ThemeDataContext, props: PD.Values<EntitySourceColorThemeParams>): ColorTheme<EntitySourceColorThemeParams> {
+export function EntitySourceColorTheme(ctx: ThemeDataContext, props: PD.Values<EntitySourceColorThemeParams>): ColorTheme<EntitySourceColorThemeParams, ColorTypeLocation> {
     let color: LocationColor;
     let legend: ScaleLegend | TableLegend | undefined;
 

src/mol-theme/color/molecule-type.ts

@@ -15,6 +15,7 @@ import { ThemeDataContext } from '../theme';
 import { TableLegend } from '../../mol-util/legend';
 import { getAdjustedColorMap } from '../../mol-util/color/color';
 import { getColorMapParams } from '../../mol-util/color/params';
+import { ColorTypeLocation } from '../../mol-geo/geometry/color-data';
 
 export const MoleculeTypeColors = ColorMap({
     water: 0x386cb0,
@@ -57,7 +58,7 @@ export function moleculeTypeColor(colorMap: MoleculeTypeColors, unit: Unit, elem
     return DefaultMoleculeTypeColor;
 }
 
-export function MoleculeTypeColorTheme(ctx: ThemeDataContext, props: PD.Values<MoleculeTypeColorThemeParams>): ColorTheme<MoleculeTypeColorThemeParams> {
+export function MoleculeTypeColorTheme(ctx: ThemeDataContext, props: PD.Values<MoleculeTypeColorThemeParams>): ColorTheme<MoleculeTypeColorThemeParams, ColorTypeLocation> {
     const colorMap = getAdjustedColorMap(props.colors.name === 'default' ? MoleculeTypeColors : props.colors.params, props.saturation, props.lightness);
 
     function color(location: Location): Color {

src/mol-theme/color/operator-name.ts

@@ -12,6 +12,7 @@ import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { ThemeDataContext } from '../theme';
 import { getPaletteParams, getPalette } from '../../mol-util/color/palette';
 import { ScaleLegend, TableLegend } from '../../mol-util/legend';
+import { ColorTypeLocation } from '../../mol-geo/geometry/color-data';
 
 const DefaultList = 'many-distinct';
 const DefaultColor = Color(0xCCCCCC);
@@ -35,7 +36,7 @@ function getOperatorNameSerialMap(structure: Structure) {
     return map;
 }
 
-export function OperatorNameColorTheme(ctx: ThemeDataContext, props: PD.Values<OperatorNameColorThemeParams>): ColorTheme<OperatorNameColorThemeParams> {
+export function OperatorNameColorTheme(ctx: ThemeDataContext, props: PD.Values<OperatorNameColorThemeParams>): ColorTheme<OperatorNameColorThemeParams, ColorTypeLocation> {
     let color: LocationColor;
     let legend: ScaleLegend | TableLegend | undefined;
 

src/mol-theme/color/uniform.ts

@@ -10,6 +10,7 @@ import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { ThemeDataContext } from '../theme';
 import { TableLegend } from '../../mol-util/legend';
 import { defaults } from '../../mol-util';
+import { ColorTypeLocation } from '../../mol-geo/geometry/color-data';
 
 const DefaultColor = Color(0xCCCCCC);
 const Description = 'Gives everything the same, uniform color.';
@@ -22,7 +23,7 @@ export function getUniformColorThemeParams(ctx: ThemeDataContext) {
     return UniformColorThemeParams; // TODO return copy
 }
 
-export function UniformColorTheme(ctx: ThemeDataContext, props: PD.Values<UniformColorThemeParams>): ColorTheme<UniformColorThemeParams> {
+export function UniformColorTheme(ctx: ThemeDataContext, props: PD.Values<UniformColorThemeParams>): ColorTheme<UniformColorThemeParams, ColorTypeLocation> {
     const color = defaults(props.value, DefaultColor);
 
     return {

src/mol-theme/color/volume-value.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2021-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -10,7 +10,6 @@ import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { ThemeDataContext } from '../theme';
 import { ColorNames } from '../../mol-util/color/names';
 
-const DefaultColor = Color(0xCCCCCC);
 const Description = 'Assign color based on the given value of a volume cell.';
 
 export const VolumeValueColorThemeParams = {
@@ -43,7 +42,6 @@ export function VolumeValueColorTheme(ctx: ThemeDataContext, props: PD.Values<Vo
     return {
         factory: VolumeValueColorTheme,
         granularity: 'direct',
-        color: () => DefaultColor,
         props: props,
         description: Description,
         legend: scale.legend,