wip, link repr coloring

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

@@ -13,7 +13,7 @@ import { createUniformSize, SizeData } from '../../util/size-data';
 import { physicalSizeData, getPhysicalRadius } 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 { elementIndexColorData, elementSymbolColorData, instanceIndexColorData, chainIdElementColorData } from '../../theme/structure/color';
 import { ValueCell, defaults } from 'mol-util';
 import { Mesh } from '../../shape/mesh';
 import { RuntimeContext } from 'mol-task';
@@ -34,21 +34,33 @@ export function createTransforms({ units }: Unit.SymmetryGroup, transforms?: Val
     return transforms ? ValueCell.update(transforms, array) : ValueCell.create(array)
 }
 
-export function createColors(group: Unit.SymmetryGroup, vertexMap: VertexMap, props: ColorTheme, colorData?: ColorData) {
+export function createColors(group: Unit.SymmetryGroup, elementCount: number, props: ColorTheme, colorData?: ColorData) {
     switch (props.name) {
         case 'atom-index':
-            return elementIndexColorData({ group, vertexMap }, colorData)
+            return elementIndexColorData({ group, elementCount }, colorData)
         case 'chain-id':
-            return chainIdColorData({ group, vertexMap }, colorData)
+            return chainIdElementColorData({ group, elementCount }, colorData)
         case 'element-symbol':
-            return elementSymbolColorData({ group, vertexMap }, colorData)
+            return elementSymbolColorData({ group, elementCount }, colorData)
         case 'instance-index':
-            return instanceIndexColorData({ group, vertexMap }, colorData)
+            return instanceIndexColorData({ group, elementCount }, colorData)
         case 'uniform':
             return createUniformColor(props, colorData)
     }
 }
 
+// export function createLinkColors(group: Unit.SymmetryGroup, props: ColorTheme, colorData?: ColorData): ColorData {
+//     switch (props.name) {
+//         case 'atom-index':
+//         case 'chain-id':
+//         case 'element-symbol':
+//         case 'instance-index':
+//             return chainIdLinkColorData({ group, vertexMap }, colorData)
+//         case 'uniform':
+//             return createUniformColor(props, colorData)
+//     }
+// }
+
 export function createSizes(group: Unit.SymmetryGroup, vertexMap: VertexMap, props: SizeTheme): SizeData {
     switch (props.name) {
         case 'uniform':

src/mol-geo/representation/structure/visual/element-point.ts

@@ -85,7 +85,7 @@ export default function PointVisual(): UnitsVisual<PointProps> {
             const transforms = createTransforms(group)
 
             if (ctx.shouldUpdate) await ctx.update('Computing point colors');
-            const color = createColors(group, vertexMap, colorTheme)
+            const color = createColors(group, elementCount, colorTheme)
 
             if (ctx.shouldUpdate) await ctx.update('Computing point sizes');
             const size = createSizes(group, vertexMap, sizeTheme)

src/mol-geo/representation/structure/visual/element-sphere.ts

@@ -12,7 +12,6 @@ import { Unit, Element } from 'mol-model/structure';
 import { DefaultStructureProps, UnitsVisual } from '../index';
 import { RuntimeContext } from 'mol-task'
 import { createTransforms, createColors, createElementSphereMesh, markElement, getElementRadius } from '../utils';
-import VertexMap from '../../../shape/vertex-map';
 import { deepEqual, defaults } from 'mol-util';
 import { fillSerial } from 'mol-gl/renderable/util';
 import { RenderableState, MeshValues } from 'mol-gl/renderable';
@@ -39,7 +38,6 @@ export function ElementSphereVisual(): UnitsVisual<ElementSphereProps> {
     let currentProps: typeof DefaultElementSphereProps
     let mesh: Mesh
     let currentGroup: Unit.SymmetryGroup
-    let vertexMap: VertexMap
 
     return {
         renderObjects,
@@ -56,14 +54,12 @@ export function ElementSphereVisual(): UnitsVisual<ElementSphereProps> {
 
             const radius = getElementRadius(unit, sizeTheme)
             mesh = await createElementSphereMesh(ctx, unit, radius, detail, mesh)
-            // console.log(mesh)
-            vertexMap = VertexMap.fromMesh(mesh)
 
             if (ctx.shouldUpdate) await ctx.update('Computing spacefill transforms');
             const transforms = createTransforms(group)
 
             if (ctx.shouldUpdate) await ctx.update('Computing spacefill colors');
-            const color = createColors(group, vertexMap, colorTheme)
+            const color = createColors(group, elementCount, colorTheme)
 
             if (ctx.shouldUpdate) await ctx.update('Computing spacefill marks');
             const marker = createMarkers(instanceCount * elementCount)
@@ -109,8 +105,6 @@ export function ElementSphereVisual(): UnitsVisual<ElementSphereProps> {
                 const radius = getElementRadius(unit, newProps.sizeTheme)
                 mesh = await createElementSphereMesh(ctx, unit, radius, newProps.detail, mesh)
                 ValueCell.update(spheres.values.drawCount, mesh.triangleCount * 3)
-                // TODO update in-place
-                vertexMap = VertexMap.fromMesh(mesh)
                 updateColor = true
             }
 
@@ -119,8 +113,9 @@ export function ElementSphereVisual(): UnitsVisual<ElementSphereProps> {
             }
 
             if (updateColor) {
+                const elementCount = currentGroup.elements.length
                 if (ctx.shouldUpdate) await ctx.update('Computing spacefill colors');
-                createColors(currentGroup, vertexMap, newProps.colorTheme, spheres.values)
+                createColors(currentGroup, elementCount, newProps.colorTheme, spheres.values)
             }
 
             ValueCell.updateIfChanged(spheres.values.uAlpha, newProps.alpha)

src/mol-geo/representation/structure/visual/intra-unit-link-cylinder.ts

@@ -21,11 +21,12 @@ import { Mesh } from '../../../shape/mesh';
 import { PickingId } from '../../../util/picking';
 import { MeshBuilder } from '../../../shape/mesh-builder';
 import { Vec3, Mat4 } from 'mol-math/linear-algebra';
-import { createUniformColor } from '../../../util/color-data';
+// import { createUniformColor } from '../../../util/color-data';
 import { defaults } from 'mol-util';
 import { Loci, isEveryLoci, EmptyLoci } from 'mol-model/loci';
 import { MarkerAction, applyMarkerAction, createMarkers, MarkerData } from '../../../util/marker-data';
 import { SizeTheme } from '../../../theme';
+import { chainIdLinkColorData } from '../../../theme/structure/color/chain-id';
 
 async function createLinkCylinderMesh(ctx: RuntimeContext, unit: Unit, mesh?: Mesh) {
     if (!Unit.isAtomic(unit)) return Mesh.createEmpty(mesh)
@@ -102,14 +103,12 @@ export function IntraUnitLinkVisual(): UnitsVisual<IntraUnitLinkProps> {
 
             mesh = await createLinkCylinderMesh(ctx, unit)
 
-            // console.log(mesh)
-            // vertexMap = VertexMap.fromMesh(mesh)
-
             if (ctx.shouldUpdate) await ctx.update('Computing link transforms');
             const transforms = createTransforms(group)
 
             if (ctx.shouldUpdate) await ctx.update('Computing link colors');
-            const color = createUniformColor({ value: 0xFF0000 })
+            // const color = createUniformColor({ value: 0xFF0000 })
+            const color = chainIdLinkColorData({ group, elementCount })
 
             if (ctx.shouldUpdate) await ctx.update('Computing link marks');
             const marker = createMarkers(instanceCount * elementCount)

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

@@ -7,7 +7,7 @@
 import { Unit, Queries, Element } from 'mol-model/structure';
 
 import { StructureColorDataProps } from '.';
-import { createAttributeOrElementColor, ColorData } from '../../../util/color-data';
+import { ColorData, createElementColor } from '../../../util/color-data';
 import { ColorScale } from 'mol-util/color';
 import { Column } from 'mol-data/db';
 
@@ -38,8 +38,8 @@ function createChainIdMap(unit: Unit) {
     return { map, count: index }
 }
 
-export function chainIdColorData(props: StructureColorDataProps, colorData?: ColorData) {
-    const { group: { units, elements }, vertexMap } = props
+export function chainIdColorData(props: StructureColorDataProps, locationFn: (l: Element.Location, renderElementIdx: number) => void, colorData?: ColorData) {
+    const { group: { units }, elementCount } = props
     const unit = units[0]
 
     const { map, count } = createChainIdMap(unit)
@@ -52,16 +52,74 @@ export function chainIdColorData(props: StructureColorDataProps, colorData?: Col
         asym_id = Queries.props.chain.label_asym_id
     } else if (Unit.isCoarse(unit)) {
         asym_id = Queries.props.coarse.asym_id
+    } else {
+        throw new Error('unhandled unit kind')
     }
 
     const l = Element.Location()
     l.unit = unit
 
-    return createAttributeOrElementColor(vertexMap, {
-        colorFn: (elementIdx: number) => {
-            l.element = elements[elementIdx]
+    // return createAttributeOrElementColor(vertexMap, {
+    //     colorFn: (renderElementIdx: number) => {
+    //         locationFn(l, renderElementIdx)
+    //         console.log(l.element, asym_id(l))
+    //         return scale.color(map.get(asym_id(l)) || 0)
+    //     },
+    //     vertexMap
+    // }, colorData)
+    return createElementColor({
+        colorFn: (renderElementIdx: number) => {
+            locationFn(l, renderElementIdx)
             return scale.color(map.get(asym_id(l)) || 0)
         },
-        vertexMap
+        elementCount
     }, colorData)
 }
+
+export function chainIdElementColorData(props: StructureColorDataProps, colorData?: ColorData) {
+    const elements = props.group.units[0].elements
+    function locationFn(l: Element.Location, renderElementIdx: number) {
+        l.element = elements[renderElementIdx]
+    }
+    return chainIdColorData(props, locationFn, colorData)
+    // const { group: { units, elements }, vertexMap } = props
+    // const unit = units[0]
+
+    // const { map, count } = createChainIdMap(unit)
+
+    // const domain = [ 0, count - 1 ]
+    // const scale = ColorScale.create({ domain })
+
+    // let asym_id: Element.Property<string>
+    // if (Unit.isAtomic(unit)) {
+    //     asym_id = Queries.props.chain.label_asym_id
+    // } else if (Unit.isCoarse(unit)) {
+    //     asym_id = Queries.props.coarse.asym_id
+    // }
+
+    // const l = Element.Location()
+    // l.unit = unit
+
+    // return createAttributeOrElementColor(vertexMap, {
+    //     colorFn: (elementIdx: number) => {
+    //         l.element = elements[elementIdx]
+    //         return scale.color(map.get(asym_id(l)) || 0)
+    //     },
+    //     vertexMap
+    // }, colorData)
+}
+
+export function chainIdLinkColorData(props: StructureColorDataProps, colorData?: ColorData) {
+    const { group: { units } } = props
+    const elements = props.group.units[0].elements
+    const links = (units[0] as Unit.Atomic).links
+    // const { edgeCount, a, b } = links
+    const { a } = links
+
+    function locationFn(l: Element.Location, renderElementIdx: number) {
+        const aI = elements[a[renderElementIdx]]
+        // const bI = elements[b[renderElementIdx]];
+        l.element = aI
+    }
+    return chainIdColorData(props, locationFn, colorData)
+}
\ No newline at end of file

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

@@ -9,15 +9,14 @@ import { StructureColorDataProps } from '.';
 import { createElementInstanceColor, ColorData } from '../../../util/color-data';
 
 export function elementIndexColorData(props: StructureColorDataProps, colorData?: ColorData) {
-    const { group: { units, elements }, vertexMap } = props
+    const { group: { units }, elementCount } = props
     const instanceCount = units.length
-    const elementCount = elements.length
 
     const domain = [ 0, instanceCount * elementCount - 1 ]
     const scale = ColorScale.create({ domain })
     return createElementInstanceColor({
         colorFn: (instanceIdx, elementIdx) => scale.color(instanceIdx * elementCount + elementIdx),
         instanceCount,
-        vertexMap
+        elementCount
     }, colorData)
 }
\ No newline at end of file

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

@@ -7,7 +7,7 @@
 import { ElementSymbol } from 'mol-model/structure/model/types';
 import { Color } from 'mol-util/color';
 import { StructureColorDataProps } from '.';
-import { createAttributeOrElementColor, ColorData } from '../../../util/color-data';
+import { createElementColor, ColorData } from '../../../util/color-data';
 
 // from Jmol http://jmol.sourceforge.net/jscolors/ (or 0xFFFFFF)
 export const ElementSymbolColors: { [k: string]: Color } = {
@@ -22,13 +22,13 @@ export function elementSymbolColor(element: ElementSymbol): Color {
 }
 
 export function elementSymbolColorData(props: StructureColorDataProps, colorData?: ColorData) {
-    const { group: { units, elements }, vertexMap } = props
+    const { group: { units, elements }, elementCount } = props
     const { type_symbol } = units[0].model.atomicHierarchy.atoms
-    return createAttributeOrElementColor(vertexMap, {
+    return createElementColor({
         colorFn: (elementIdx: number) => {
             const e = elements[elementIdx]
             return elementSymbolColor(type_symbol.value(e))
         },
-        vertexMap
+        elementCount
     }, colorData)
 }
\ No newline at end of file

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

@@ -5,14 +5,13 @@
  */
 
 import { Unit } from 'mol-model/structure';
-import VertexMap from '../../../shape/vertex-map';
 
 export interface StructureColorDataProps {
     group: Unit.SymmetryGroup,
-    vertexMap: VertexMap
+    elementCount: number
 }
 
 export { elementIndexColorData } from './element-index'
-export { chainIdColorData } from './chain-id'
+export { chainIdElementColorData } from './chain-id'
 export { elementSymbolColorData } from './element-symbol'
 export { instanceIndexColorData } from './instance-index'
\ No newline at end of file

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

@@ -120,13 +120,12 @@ export function createInstanceColor(props: InstanceColorProps, colorData?: Color
 
 export interface ElementColorProps {
     colorFn: (elementIdx: number) => Color
-    vertexMap: VertexMap
+    elementCount: number
 }
 
 /** Creates color texture with color for each element (i.e. shared across instances/units) */
 export function createElementColor(props: ElementColorProps, colorData?: ColorData): ColorData {
-    const { colorFn, vertexMap } = props
-    const elementCount = vertexMap.offsetCount - 1
+    const { colorFn, elementCount } = props
     const colors = colorData && colorData.tColor.ref.value.array.length >= elementCount * 3 ? colorData.tColor.ref.value : createTextureImage(elementCount, 3)
     for (let i = 0, il = elementCount; i < il; ++i) {
         Color.toArray(colorFn(i), colors.array, i * 3)
@@ -137,13 +136,12 @@ export function createElementColor(props: ElementColorProps, colorData?: ColorDa
 export interface ElementInstanceColorProps {
     colorFn: (instanceIdx: number, elementIdx: number) => Color
     instanceCount: number,
-    vertexMap: VertexMap
+    elementCount: number
 }
 
 /** Creates color texture with color for each element instance (i.e. for each unit) */
 export function createElementInstanceColor(props: ElementInstanceColorProps, colorData?: ColorData): ColorData {
-    const { colorFn, instanceCount, vertexMap } = props
-    const elementCount = vertexMap.offsetCount - 1
+    const { colorFn, instanceCount, elementCount } = props
     const count = instanceCount * elementCount
     const colors = colorData && colorData.tColor.ref.value.array.length >= count * 3 ? colorData.tColor.ref.value : createTextureImage(count, 3)
     let colorOffset = 0
@@ -155,8 +153,3 @@ export function createElementInstanceColor(props: ElementInstanceColorProps, col
     }
     return createTextureColor(colors, 'elementInstance', colorData)
 }
\ No newline at end of file
-
-/** Create color attribute or texture, depending on the vertexMap */
-export function createAttributeOrElementColor(vertexMap: VertexMap, props: AttributeColorProps, colorData?: ColorData) {
-    return vertexMap.idCount < 4 * vertexMap.offsetCount ? createAttributeColor(props, colorData) : createElementColor(props, colorData)
-}
\ No newline at end of file