wip, symmetry-cluster color-theme

src/apps/canvas/assembly-symmetry.ts

@@ -6,12 +6,15 @@
 
 import { AssemblySymmetry } from 'mol-model-props/rcsb/symmetry';
 import { Table } from 'mol-data/db';
-import { Color } from 'mol-util/color';
+import { Color, ColorScale } from 'mol-util/color';
 import { MeshBuilder } from 'mol-geo/mesh/mesh-builder';
 import { Tensor } from 'mol-math/linear-algebra';
 import { addSphere } from 'mol-geo/mesh/builder/sphere';
 import { addCylinder } from 'mol-geo/mesh/builder/cylinder';
 import { Shape } from 'mol-model/shape';
+import { ColorTheme } from 'mol-view/theme/color';
+import { Location } from 'mol-model/location';
+import { StructureElement, Unit, StructureProperties } from 'mol-model/structure';
 
 export function getAxesShape(featureId: number, assemblySymmetry: AssemblySymmetry) {
     const f = assemblySymmetry.db.rcsb_assembly_symmetry_feature
@@ -45,15 +48,52 @@ export function getAxesShape(featureId: number, assemblySymmetry: AssemblySymmet
     return shape
 }
 
-export function getClusterColorTheme(featureId: number, assemblySymmetry: AssemblySymmetry) {
+function getAsymId(unit: Unit): StructureElement.Property<string> {
+    switch (unit.kind) {
+        case Unit.Kind.Atomic:
+            return StructureProperties.chain.auth_asym_id // TODO
+        case Unit.Kind.Spheres:
+        case Unit.Kind.Gaussians:
+            return StructureProperties.coarse.asym_id
+    }
+}
+
+function memberKey (asym_id: string, oper_list_id?: number) {
+    return `${asym_id}|${oper_list_id}`
+}
+
+export function getClusterColorTheme(featureId: number, assemblySymmetry: AssemblySymmetry): ColorTheme {
+    const DefaultColor = Color(0xCCCCCC)
     const f = assemblySymmetry.db.rcsb_assembly_symmetry_feature
     const feature = Table.pickRow(f, i => f.id.value(i) === featureId)
-    if (!feature) return
+    if (!feature) return { kind: 'uniform', color: () => DefaultColor }
 
     const clusters = assemblySymmetry.getClusters(featureId)
-    if (!clusters._rowCount) return
+    if (!clusters._rowCount) return { kind: 'uniform', color: () => DefaultColor }
 
+    const clusterByMember = new Map<string, number>()
     for (let i = 0, il = clusters._rowCount; i < il; ++i) {
-        console.log(clusters.members.value(i), clusters.avg_rmsd.value(i), feature.stoichiometry_value, feature.stoichiometry_description)
+        clusters.members.value(i).forEach(m => {
+            const ms = m.split('_')
+            const asym_id = ms[0]
+            const oper_list_id = ms.length === 2 ? parseInt(ms[1]) : undefined
+            clusterByMember.set(memberKey(asym_id, oper_list_id), i)
+        })
+    }
+
+    const scale = ColorScale.create({ domain: [ 0, clusters._rowCount - 1 ] })
+
+    return {
+        kind: 'instance',
+        color: (location: Location): Color => {
+            if (StructureElement.isLocation(location)) {
+                const ns = location.unit.conformation.operator.name.split('-')
+                const asym_id = getAsymId(location.unit)
+                const oper_list_id = ns.length === 2 ? parseInt(ns[1]) : undefined
+                const cluster = clusterByMember.get(memberKey(asym_id(location), oper_list_id))
+                return cluster !== undefined ? scale.color(cluster) : DefaultColor
+            }
+            return DefaultColor
+        }
     }
 }