wip, color themes

src/mol-app/ui/transform/backbone.tsx

@@ -15,7 +15,7 @@ import { Toggle } from '../controls/common';
 import { BackboneEntity } from 'mol-view/state/entity';
 import { BackboneUpdate } from 'mol-view/state/transform'
 import { StateContext } from 'mol-view/state/context';
-import { ColorTheme, SizeTheme, ColorThemeNames, ColorThemeName } from 'mol-geo/theme';
+import { ColorThemeProps, SizeTheme, ColorThemeNames, ColorThemeName } from 'mol-geo/theme';
 import { Color, ColorNames } from 'mol-util/color';
 import { Slider } from '../controls/slider';
 import { VisualQuality } from 'mol-geo/representation/util';
@@ -26,7 +26,7 @@ interface BackboneState {
     flipSided: boolean
     flatShaded: boolean
     detail: number
-    colorTheme: ColorTheme
+    colorTheme: ColorThemeProps
     colorValue: Color
     sizeTheme: SizeTheme
     visible: boolean
@@ -43,7 +43,7 @@ export class Backbone extends View<Controller<any>, BackboneState, { transform:
         flipSided: false,
         flatShaded: false,
         detail: 2,
-        colorTheme: { name: 'element-symbol' } as ColorTheme,
+        colorTheme: { name: 'element-symbol' } as ColorThemeProps,
         colorValue: 0x000000,
         sizeTheme: { name: 'uniform' } as SizeTheme,
         visible: true,

src/mol-app/ui/transform/ball-and-stick.tsx

@@ -15,7 +15,7 @@ import { Toggle } from '../controls/common';
 import { DistanceRestraintEntity } from 'mol-view/state/entity';
 import { DistanceRestraintUpdate } from 'mol-view/state/transform'
 import { StateContext } from 'mol-view/state/context';
-import { ColorTheme, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
+import { ColorThemeProps, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
 import { Color, ColorNames } from 'mol-util/color';
 import { Slider } from '../controls/slider';
 import { VisualQuality } from 'mol-geo/representation/util';
@@ -25,7 +25,7 @@ interface BallAndStickState {
     doubleSided: boolean
     flipSided: boolean
     flatShaded: boolean
-    colorTheme: ColorTheme
+    colorTheme: ColorThemeProps
     colorValue: Color
     sizeTheme: SizeTheme
     visible: boolean
@@ -46,7 +46,7 @@ export class BallAndStick extends View<Controller<any>, BallAndStickState, { tra
         doubleSided: true,
         flipSided: false,
         flatShaded: false,
-        colorTheme: { name: 'element-symbol' } as ColorTheme,
+        colorTheme: { name: 'element-symbol' } as ColorThemeProps,
         colorValue: 0x000000,
         sizeTheme: { name: 'uniform' } as SizeTheme,
         visible: true,

src/mol-app/ui/transform/carbohydrate.tsx

@@ -15,7 +15,7 @@ import { Toggle } from '../controls/common';
 import { CarbohydrateEntity } from 'mol-view/state/entity';
 import { CarbohydrateUpdate } from 'mol-view/state/transform'
 import { StateContext } from 'mol-view/state/context';
-import { ColorTheme, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
+import { ColorThemeProps, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
 import { Color, ColorNames } from 'mol-util/color';
 import { Slider } from '../controls/slider';
 import { VisualQuality } from 'mol-geo/representation/util';
@@ -26,7 +26,7 @@ interface CarbohydrateState {
     flipSided: boolean
     flatShaded: boolean
     detail: number
-    colorTheme: ColorTheme
+    colorTheme: ColorThemeProps
     colorValue: Color
     sizeTheme: SizeTheme
     visible: boolean
@@ -47,7 +47,7 @@ export class Carbohydrate extends View<Controller<any>, CarbohydrateState, { tra
         flipSided: false,
         flatShaded: false,
         detail: 2,
-        colorTheme: { name: 'element-symbol' } as ColorTheme,
+        colorTheme: { name: 'element-symbol' } as ColorThemeProps,
         colorValue: 0x000000,
         sizeTheme: { name: 'uniform' } as SizeTheme,
         visible: true,

src/mol-app/ui/transform/cartoon.tsx

@@ -15,7 +15,7 @@ import { Toggle } from '../controls/common';
 import { CartoonEntity } from 'mol-view/state/entity';
 import { CartoonUpdate } from 'mol-view/state/transform'
 import { StateContext } from 'mol-view/state/context';
-import { ColorTheme, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
+import { ColorThemeProps, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
 import { Color, ColorNames } from 'mol-util/color';
 import { Slider } from '../controls/slider';
 import { VisualQuality } from 'mol-geo/representation/util';
@@ -26,7 +26,7 @@ interface CartoonState {
     flipSided: boolean
     flatShaded: boolean
     detail: number
-    colorTheme: ColorTheme
+    colorTheme: ColorThemeProps
     colorValue: Color
     sizeTheme: SizeTheme
     visible: boolean
@@ -43,7 +43,7 @@ export class Cartoon extends View<Controller<any>, CartoonState, { transform: Ca
         flipSided: false,
         flatShaded: false,
         detail: 2,
-        colorTheme: { name: 'element-symbol' } as ColorTheme,
+        colorTheme: { name: 'element-symbol' } as ColorThemeProps,
         colorValue: 0x000000,
         sizeTheme: { name: 'uniform' } as SizeTheme,
         visible: true,

src/mol-app/ui/transform/distance-restraint.tsx

@@ -15,7 +15,7 @@ import { Toggle } from '../controls/common';
 import { DistanceRestraintEntity } from 'mol-view/state/entity';
 import { DistanceRestraintUpdate } from 'mol-view/state/transform'
 import { StateContext } from 'mol-view/state/context';
-import { ColorTheme, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
+import { ColorThemeProps, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
 import { Color, ColorNames } from 'mol-util/color';
 import { Slider } from '../controls/slider';
 import { VisualQuality } from 'mol-geo/representation/util';
@@ -25,7 +25,7 @@ interface DistanceRestraintState {
     doubleSided: boolean
     flipSided: boolean
     flatShaded: boolean
-    colorTheme: ColorTheme
+    colorTheme: ColorThemeProps
     colorValue: Color
     sizeTheme: SizeTheme
     visible: boolean
@@ -46,7 +46,7 @@ export class DistanceRestraint extends View<Controller<any>, DistanceRestraintSt
         doubleSided: true,
         flipSided: false,
         flatShaded: false,
-        colorTheme: { name: 'element-symbol' } as ColorTheme,
+        colorTheme: { name: 'element-symbol' } as ColorThemeProps,
         colorValue: 0x000000,
         sizeTheme: { name: 'uniform' } as SizeTheme,
         visible: true,

src/mol-app/ui/transform/spacefill.tsx

@@ -15,7 +15,7 @@ import { Toggle } from '../controls/common';
 import { SpacefillEntity } from 'mol-view/state/entity';
 import { SpacefillUpdate } from 'mol-view/state/transform'
 import { StateContext } from 'mol-view/state/context';
-import { ColorTheme, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
+import { ColorThemeProps, SizeTheme, ColorThemeName, ColorThemeNames } from 'mol-geo/theme';
 import { Color, ColorNames } from 'mol-util/color';
 import { Slider } from '../controls/slider';
 import { VisualQuality } from 'mol-geo/representation/util';
@@ -26,7 +26,7 @@ interface SpacefillState {
     flipSided: boolean
     flatShaded: boolean
     detail: number
-    colorTheme: ColorTheme
+    colorTheme: ColorThemeProps
     colorValue: Color
     sizeTheme: SizeTheme
     visible: boolean
@@ -43,7 +43,7 @@ export class Spacefill extends View<Controller<any>, SpacefillState, { transform
         flipSided: false,
         flatShaded: false,
         detail: 2,
-        colorTheme: { name: 'element-symbol' } as ColorTheme,
+        colorTheme: { name: 'element-symbol' } as ColorThemeProps,
         colorValue: 0x000000,
         sizeTheme: { name: 'uniform' } as SizeTheme,
         visible: true,

src/mol-geo/representation/structure/index.ts

@@ -7,14 +7,14 @@
 
 import { Structure } from 'mol-model/structure';
 import { Representation, RepresentationProps } from '..';
-import { ColorTheme, SizeTheme } from '../../theme';
+import { ColorThemeProps, SizeTheme } from '../../theme';
 import { DefaultBaseProps, DefaultMeshProps } from '../util';
 
 export interface StructureRepresentation<P extends RepresentationProps = {}> extends Representation<Structure, P> { }
 
 export const DefaultStructureProps = {
     ...DefaultBaseProps,
-    colorTheme: { name: 'instance-index' } as ColorTheme,
+    colorTheme: { name: 'unit-index' } as ColorThemeProps,
     sizeTheme: { name: 'physical' } as SizeTheme,
 }
 export type StructureProps = typeof DefaultStructureProps

src/mol-geo/representation/structure/visual/util/common.ts

@@ -8,15 +8,14 @@
 import { Unit, Structure } from 'mol-model/structure';
 import { Mat4 } from 'mol-math/linear-algebra'
 
-import { createUniformColor, ColorData } from '../../../../util/color-data';
+import { createUniformColor, ColorData, createElementColor, createElementInstanceColor, createInstanceColor } from '../../../../util/color-data';
 import { createUniformSize, SizeData } from '../../../../util/size-data';
 import { physicalSizeData } from '../../../../theme/structure/size/physical';
 import VertexMap from '../../../../shape/vertex-map';
-import { ColorTheme, SizeTheme } from '../../../../theme';
-import { elementIndexColorData, elementSymbolColorData, instanceIndexColorData, chainIdColorData } from '../../../../theme/structure/color';
+import { ColorThemeProps, SizeTheme } from '../../../../theme';
+import { ColorTheme } from '../../../../theme/structure/color';
 import { ValueCell, defaults } from 'mol-util';
 import { LocationIterator } from './location-iterator';
-import { carbohydrateSymbolColorData } from '../../../../theme/structure/color/carbohydrate-symbol';
 import { Mesh } from '../../../../shape/mesh';
 import { MeshValues } from 'mol-gl/renderable';
 import { getMeshData } from '../../../../util/mesh-data';
@@ -41,20 +40,14 @@ export function createIdentityTransform(transforms?: ValueCell<Float32Array>) {
     return transforms ? ValueCell.update(transforms, identityTransform) : ValueCell.create(identityTransform)
 }
 
-export function createColors(locationIt: LocationIterator, props: ColorTheme, colorData?: ColorData) {
-    switch (props.name) {
-        case 'atom-index':
-            return elementIndexColorData(locationIt, colorData)
-        case 'carbohydrate-symbol':
-            return carbohydrateSymbolColorData(locationIt, props, colorData)
-        case 'chain-id':
-            return chainIdColorData(locationIt, colorData)
-        case 'element-symbol':
-            return elementSymbolColorData(locationIt, colorData)
-        case 'instance-index':
-            return instanceIndexColorData(locationIt, colorData)
-        case 'uniform':
-            return createUniformColor(locationIt, () => props.value || 0x000000, colorData)
+export function createColors(locationIt: LocationIterator, props: ColorThemeProps, colorData?: ColorData) {
+    const colorTheme = ColorTheme(props)
+    switch (colorTheme.kind) {
+        case 'uniform': return createUniformColor(locationIt, colorTheme.color, colorData)
+        case 'element': return createElementColor(locationIt, colorTheme.color, colorData)
+        case 'elementInstance': return createElementInstanceColor(locationIt, colorTheme.color, colorData)
+        case 'instance': return createInstanceColor(locationIt, colorTheme.color, colorData)
+        case 'attribute': throw new Error('not implemented') // TODO
     }
 }
 

src/mol-geo/theme/index.ts

@@ -7,19 +7,19 @@
 import { Color } from 'mol-util/color';
 import { Structure } from 'mol-model/structure';
 
-export interface ColorTheme {
-    name: 'atom-index' | 'chain-id'| 'instance-index' | 'uniform' | 'carbohydrate-symbol' | 'element-symbol'
+export interface ColorThemeProps {
+    name: 'element-index' | 'chain-id'| 'unit-index' | 'uniform' | 'carbohydrate-symbol' | 'element-symbol'
     domain?: [number, number]
     value?: Color
     structure?: Structure
 }
 
 export const ColorThemeInfo = {
-    'atom-index': {},
+    'element-index': {},
     'carbohydrate-symbol': {},
     'chain-id': {},
     'element-symbol': {},
-    'instance-index': {},
+    'unit-index': {},
     'uniform': {}
 }
 export type ColorThemeName = keyof typeof ColorThemeInfo

src/mol-geo/theme/structure/color/carbohydrate-symbol.ts

@@ -6,16 +6,16 @@
 
 import { StructureElement, Link, ElementIndex, Unit } from 'mol-model/structure';
 
-import { ColorData, createElementColor } from '../../../util/color-data';
-import { Color } from 'mol-util/color';
-import { LocationIterator, LocationValue } from '../../../representation/structure/visual/util/location-iterator';
-import { ColorTheme } from '../..';
+import { ColorThemeProps } from '../..';
 import { SaccharideColors } from 'mol-model/structure/structure/carbohydrates/constants';
+import { Location } from 'mol-model/location';
+import { ColorTheme } from '.';
+import { LocationColor } from '../../../util/color-data';
 
 const DefaultColor = 0xCCCCCC;
 
-export function carbohydrateSymbolColorData(locationIt: LocationIterator, props: ColorTheme, colorData?: ColorData) {
-    let colorFn: (locationValue: LocationValue) => Color
+export function CarbohydrateSymbolColorTheme(props: ColorThemeProps): ColorTheme {
+    let colorFn: LocationColor
 
     if (props.structure) {
         const { elements, getElementIndex, getAnomericCarbon } = props.structure.carbohydrates
@@ -30,15 +30,14 @@ export function carbohydrateSymbolColorData(locationIt: LocationIterator, props:
             return DefaultColor
         }
 
-        colorFn = (locationValue: LocationValue) => {
-            const { location: l } = locationValue
-            if (locationValue.isSecondary) {
+        colorFn = (location: Location, isSecondary: boolean) => {
+            if (isSecondary) {
                 return SaccharideColors.Secondary
             } else {
-                if (StructureElement.isLocation(l)) {
-                    return getColor(l.unit, l.element)
-                } else if (Link.isLocation(l)) {
-                    return getColor(l.aUnit, l.aUnit.elements[l.aIndex])
+                if (StructureElement.isLocation(location)) {
+                    return getColor(location.unit, location.element)
+                } else if (Link.isLocation(location)) {
+                    return getColor(location.aUnit, location.aUnit.elements[location.aIndex])
                 }
             }
             return DefaultColor
@@ -47,5 +46,8 @@ export function carbohydrateSymbolColorData(locationIt: LocationIterator, props:
         colorFn = () => DefaultColor
     }
 
-    return createElementColor(locationIt, colorFn, colorData)
+    return {
+        kind: 'element',
+        color: colorFn
+    }
 }

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

@@ -6,9 +6,10 @@
 
 import { Unit, StructureProperties, StructureElement, Link } from 'mol-model/structure';
 
-import { ColorData, createElementColor } from '../../../util/color-data';
 import { ColorScale, Color } from 'mol-util/color';
-import { LocationIterator, LocationValue } from '../../../representation/structure/visual/util/location-iterator';
+import { Location } from 'mol-model/location';
+import { ColorThemeProps } from '../..';
+import { ColorTheme } from '.';
 
 function getAsymId(unit: Unit): StructureElement.Property<string> {
     switch (unit.kind) {
@@ -20,11 +21,10 @@ function getAsymId(unit: Unit): StructureElement.Property<string> {
     }
 }
 
-export function chainIdColorData(locationIt: LocationIterator, colorData?: ColorData) {
+export function ChainIdColorTheme(props: ColorThemeProps): ColorTheme {
     const l = StructureElement.create()
 
-    function colorFn(locationValue: LocationValue): Color {
-        const { location } = locationValue
+    function colorFn(location: Location): Color {
         if (StructureElement.isLocation(location)) {
             const map = location.unit.model.properties.asymIdSerialMap
             const scale = ColorScale.create({ domain: [ 0, map.size - 1 ] })
@@ -41,5 +41,8 @@ export function chainIdColorData(locationIt: LocationIterator, colorData?: Color
         return 0
     }
 
-    return createElementColor(locationIt, colorFn, colorData)
+    return {
+        kind: 'element',
+        color: colorFn
+    }
 }

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

@@ -4,18 +4,48 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { ColorScale } from 'mol-util/color';
-import { createElementInstanceColor, ColorData } from '../../../util/color-data';
-import { LocationIterator, LocationValue } from '../../../representation/structure/visual/util/location-iterator';
+import { ColorScale, Color } from 'mol-util/color';
+import { Location } from 'mol-model/location';
+import { ColorThemeProps } from '../..';
+import { StructureElement, Link, Unit } from 'mol-model/structure';
+import { OrderedSet } from 'mol-data/int';
+import { ColorTheme } from '.';
+import { LocationColor } from '../../../util/color-data';
 
-export function elementIndexColorData(locationIt: LocationIterator, colorData?: ColorData) {
-    const { elementCount, instanceCount } = locationIt
+const DefaultColor = 0xCCCCCC;
 
-    const domain = [ 0, instanceCount * elementCount - 1 ]
-    const scale = ColorScale.create({ domain })
-    return createElementInstanceColor(
-        locationIt,
-        (value: LocationValue) => scale.color(value.instanceIndex * elementCount + value.elementIndex),
-        colorData
-    )
+export function ElementIndexColorTheme(props: ColorThemeProps): ColorTheme {
+    let colorFn: LocationColor
+
+    if (props.structure) {
+        const { units } = props.structure
+        const unitCount = units.length
+        const cummulativeElementCount = new Map<number, number>()
+
+        let elementCount = 0
+        for (let i = 0; i < unitCount; ++i) {
+            cummulativeElementCount.set(i, elementCount)
+            elementCount += units[i].elements.length
+        }
+        const scale = ColorScale.create({ domain: [ 0, elementCount ] })
+
+        colorFn = (location: Location): Color => {
+            if (StructureElement.isLocation(location)) {
+                const unitIndex = Unit.findUnitById(location.unit.id, units)
+                const unitElementIndex = OrderedSet.findPredecessorIndex(location.unit.elements, location.element)
+                return scale.color(cummulativeElementCount.get(unitIndex) || 0 + unitElementIndex)
+            } else if (Link.isLocation(location)) {
+                const unitId = Unit.findUnitById(location.aUnit.id, units)
+                return scale.color(cummulativeElementCount.get(unitId) || 0 + location.aIndex)
+            }
+            return 0
+        }
+    } else {
+        colorFn = () => DefaultColor
+    }
+
+    return {
+        kind: 'elementInstance',
+        color: colorFn
+    }
 }

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

@@ -6,9 +6,10 @@
 
 import { ElementSymbol } from 'mol-model/structure/model/types';
 import { Color } from 'mol-util/color';
-import { createElementColor, ColorData } from '../../../util/color-data';
 import { StructureElement, Unit, Link } from 'mol-model/structure';
-import { LocationIterator, LocationValue } from '../../../representation/structure/visual/util/location-iterator';
+import { Location } from 'mol-model/location';
+import { ColorThemeProps } from '../..';
+import { ColorTheme } from '.';
 
 // from Jmol http://jmol.sourceforge.net/jscolors/ (or 0xFFFFFF)
 export const ElementSymbolColors: { [k: string]: Color } = {
@@ -22,9 +23,8 @@ export function elementSymbolColor(element: ElementSymbol): Color {
     return c === void 0 ? DefaultElementSymbolColor : c
 }
 
-export function elementSymbolColorData(locationIt: LocationIterator, colorData?: ColorData) {
-    function colorFn(locationValue: LocationValue): Color {
-        const { location } = locationValue
+export function ElementSymbolColorTheme(props: ColorThemeProps): ColorTheme {
+    function colorFn(location: Location): Color {
         if (StructureElement.isLocation(location)) {
             if (Unit.isAtomic(location.unit)) {
                 const { type_symbol } = location.unit.model.atomicHierarchy.atoms
@@ -39,5 +39,8 @@ export function elementSymbolColorData(locationIt: LocationIterator, colorData?:
         return DefaultElementSymbolColor
     }
 
-    return createElementColor(locationIt, colorFn, colorData)
+    return {
+        kind: 'element',
+        color: colorFn
+    }
 }

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

@@ -4,14 +4,28 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Unit } from 'mol-model/structure';
+import { ColorThemeProps } from '../..';
 
-export interface StructureColorDataProps {
-    group: Unit.SymmetryGroup,
-    elementCount: number
+import { ElementIndexColorTheme } from './element-index';
+import { CarbohydrateSymbolColorTheme } from './carbohydrate-symbol';
+import { ChainIdColorTheme } from './chain-id';
+import { ElementSymbolColorTheme } from './element-symbol';
+import { UnitIndexColorTheme } from './unit-index';
+import { UniformColorTheme } from './uniform';
+import { ColorType, LocationColor } from '../../../util/color-data';
+
+export interface ColorTheme {
+    kind: ColorType
+    color: LocationColor
 }
 
-export { elementIndexColorData } from './element-index'
-export { chainIdColorData } from './chain-id'
-export { elementSymbolColorData } from './element-symbol'
-export { instanceIndexColorData } from './instance-index'
+export function ColorTheme(props: ColorThemeProps): ColorTheme {
+    switch (props.name) {
+        case 'element-index': return ElementIndexColorTheme(props)
+        case 'carbohydrate-symbol': return CarbohydrateSymbolColorTheme(props)
+        case 'chain-id': return ChainIdColorTheme(props)
+        case 'element-symbol': return ElementSymbolColorTheme(props)
+        case 'unit-index': return UnitIndexColorTheme(props)
+        case 'uniform': return UniformColorTheme(props)
+    }
+}

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

@@ -1,21 +0,0 @@
-/**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author Alexander Rose <alexander.rose@weirdbyte.de>
- */
-
-import { ColorScale } from 'mol-util/color';
-import { createInstanceColor, ColorData } from '../../../util/color-data';
-import { LocationIterator, LocationValue } from '../../../representation/structure/visual/util/location-iterator';
-
-export function instanceIndexColorData(locationIt: LocationIterator, colorData?: ColorData) {
-    const { instanceCount } = locationIt
-    const domain = [ 0, instanceCount - 1 ]
-    const scale = ColorScale.create({ domain })
-
-    function colorFn(locationValue: LocationValue) {
-        return scale.color(locationValue.instanceIndex)
-    }
-
-    return createInstanceColor(locationIt, colorFn, colorData)
-}

src/mol-geo/theme/structure/color/uniform.ts

@@ -0,0 +1,19 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ColorThemeProps } from '../..';
+import { ColorTheme } from '.';
+
+const DefaultColor = 0xCCCCCC;
+
+export function UniformColorTheme(props: ColorThemeProps): ColorTheme {
+    const color = props.value || DefaultColor
+
+    return {
+        kind: 'uniform',
+        color: () => color
+    }
+}

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

@@ -0,0 +1,41 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ColorScale, Color } from 'mol-util/color';
+import { Location } from 'mol-model/location';
+import { ColorThemeProps } from '../..';
+import { Unit, StructureElement, Link } from 'mol-model/structure';
+import { ColorTheme } from '.';
+import { LocationColor } from '../../../util/color-data';
+
+const DefaultColor = 0xCCCCCC;
+
+export function UnitIndexColorTheme(props: ColorThemeProps): ColorTheme {
+    let colorFn: LocationColor
+
+    if (props.structure) {
+        const { units } = props.structure
+        const unitCount = units.length
+
+        const scale = ColorScale.create({ domain: [ 0, unitCount ] })
+
+        colorFn = (location: Location): Color => {
+            if (StructureElement.isLocation(location)) {
+                return scale.color(Unit.findUnitById(location.unit.id, units))
+            } else if (Link.isLocation(location)) {
+                return scale.color(Unit.findUnitById(location.aUnit.id, units))
+            }
+            return 0
+        }
+    } else {
+        colorFn = () => DefaultColor
+    }
+
+    return {
+        kind: 'instance',
+        color: colorFn
+    }
+}

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

@@ -8,7 +8,8 @@ import { ValueCell } from 'mol-util';
 import { TextureImage, createTextureImage } from 'mol-gl/renderable/util';
 import { Color } from 'mol-util/color';
 import { Vec2, Vec3 } from 'mol-math/linear-algebra';
-import { LocationIterator, LocationValue, NullLocationValue } from '../representation/structure/visual/util/location-iterator';
+import { LocationIterator } from '../representation/structure/visual/util/location-iterator';
+import { Location, NullLocation } from 'mol-model/location';
 
 export type ColorType = 'uniform' | 'attribute' | 'instance' | 'element' | 'elementInstance'
 
@@ -20,7 +21,7 @@ export type ColorData = {
     dColorType: ValueCell<string>,
 }
 
-export type LocationColor = (locationValue: LocationValue) => Color
+export type LocationColor = (location: Location, isSecondary: boolean) => Color
 
 const emptyColorTexture = { array: new Uint8Array(3), width: 1, height: 1 }
 function createEmptyColorTexture() {
@@ -49,7 +50,7 @@ export function createValueColor(value: Color, colorData?: ColorData): ColorData
 
 /** Creates color uniform */
 export function createUniformColor(locationIt: LocationIterator, colorFn: LocationColor, colorData?: ColorData): ColorData {
-    return createValueColor(colorFn(NullLocationValue), colorData)
+    return createValueColor(colorFn(NullLocation, false), colorData)
 }
 
 // export interface AttributeColorProps {
@@ -110,8 +111,8 @@ export function createInstanceColor(locationIt: LocationIterator, colorFn: Locat
     const { instanceCount} = locationIt
     const colors = colorData && colorData.tColor.ref.value.array.length >= instanceCount * 3 ? colorData.tColor.ref.value : createTextureImage(instanceCount, 3)
     while (locationIt.hasNext && !locationIt.isNextNewInstance) {
-        const value = locationIt.move()
-        Color.toArray(colorFn(value), colors.array, value.index * 3)
+        const v = locationIt.move()
+        Color.toArray(colorFn(v.location, v.isSecondary), colors.array, v.index * 3)
         locationIt.skipInstance()
     }
     return createTextureColor(colors, 'instance', colorData)
@@ -122,9 +123,9 @@ export function createElementColor(locationIt: LocationIterator, colorFn: Locati
     const { elementCount } = locationIt
     const colors = colorData && colorData.tColor.ref.value.array.length >= elementCount * 3 ? colorData.tColor.ref.value : createTextureImage(elementCount, 3)
     while (locationIt.hasNext && !locationIt.isNextNewInstance) {
-        const value = locationIt.move()
-        // console.log(value)
-        Color.toArray(colorFn(value), colors.array, value.elementIndex * 3)
+        const v = locationIt.move()
+        // console.log(v)
+        Color.toArray(colorFn(v.location, v.isSecondary), colors.array, v.elementIndex * 3)
     }
     return createTextureColor(colors, 'element', colorData)
 }
@@ -135,8 +136,8 @@ export function createElementInstanceColor(locationIt: LocationIterator, colorFn
     const count = instanceCount * elementCount
     const colors = colorData && colorData.tColor.ref.value.array.length >= count * 3 ? colorData.tColor.ref.value : createTextureImage(count, 3)
     while (locationIt.hasNext && !locationIt.isNextNewInstance) {
-        const value = locationIt.move()
-        Color.toArray(colorFn(value), colors.array, value.index * 3)
+        const v = locationIt.move()
+        Color.toArray(colorFn(v.location, v.isSecondary), colors.array, v.index * 3)
     }
     return createTextureColor(colors, 'elementInstance', colorData)
 }