Add a colorData parameter to representations, which gets passed throu…

…gh to a StructuredataColormaker class. It also provides a mechanism for passing arbitrary data through to a Colormaker instance - see nglviewer#892

examples/scripts/color/data.js

@@ -0,0 +1,31 @@
+/** Demonstration of data colormaker - passing arbitrary data through to the
+ * colormaker for plotting.
+ */
+stage.loadFile('data://1blu.mmtf').then(function (c) {
+
+  // We might be colouring by some calculated/predicted property (e.g. ML model which outputs
+  // scores/quantities for individual atoms)
+  // To keep this self-contained, we'll just calculate distance from the origin for each
+  // atom:
+
+  var structure = c.structure
+  var distance = new Float64Array(structure.atomCount)
+
+  structure.eachAtom(ap => {
+    distance[ap.index] = Math.sqrt(ap.x**2 + ap.y**2 + ap.z**2)
+  })
+
+  var repr = c.addRepresentation('ball+stick', { 
+    color: 'structuredata',
+    colorData: {atomData: distance},
+    colorDomain: [30.0, 60.0]
+   })
+  c.autoView()
+
+  // If colorData isn't defined, empty etc, falls back to value, here we wait 2 seconds
+  // then update the paramter to an empty array to demonstrate
+  window.setTimeout(() => {
+    repr.setParameters({colorData: [], colorValue: 'green'})
+  }, 2000)
+
+})

src/color/colormaker-registry.ts

@@ -138,7 +138,7 @@ class ColormakerRegistry {
    * @param {Colormaker} scheme - the colormaker
    * @return {undefined}
    */
-  add (id: string, scheme: Colormaker) {
+  add (id: string, scheme: Colormaker | typeof Colormaker) {
     id = id.toLowerCase()
     this.schemes[ id ] = scheme
   }

src/color/colormaker.ts

@@ -46,10 +46,16 @@ export const ScaleDefaultParameters = {
 }
 export type ScaleParameters = typeof ScaleDefaultParameters
 
+export interface ColorData {
+  atomData?: { [key: number]: number }
+  bondData?: { [key: number]: number }
+}
+
 export interface ColormakerParameters extends ScaleParameters {
   structure?: Structure
   volume?: Volume
   surface?: Surface
+  data?: ColorData
 }
 
 export type StuctureColormakerParams = { structure: Structure } & Partial<ColormakerParameters>
@@ -59,7 +65,7 @@ export type ColormakerScale = (v: number) => number
 const tmpColor = new Color()
 
 /** Decorator for optionally linearizing a numeric color */
-type colorFuncType = (value: any) => number // decorator applies to functions with this shape
+type colorFuncType = (value: any, fromTo?: boolean) => number // decorator applies to functions with this shape
 export function manageColor<T extends {parameters: ColormakerParameters}>
   (_target: Object,
    _name: string | symbol,

src/color/structuredata-colormaker.ts

@@ -0,0 +1,43 @@
+/**
+ * @file Colordata Colormaker
+ * @author Fred Ludlow <[email protected]>
+ * @private
+ */
+
+import { ColormakerRegistry } from '../globals'
+import Colormaker, { ColormakerScale, manageColor, StuctureColormakerParams } from './colormaker'
+import AtomProxy from '../proxy/atom-proxy'
+import BondProxy from '../proxy/bond-proxy'
+
+
+class StructuredataColormaker extends Colormaker {
+  atomData?: { [key: number]: number }
+  bondData?: { [key: number]: number }
+  scale: ColormakerScale
+
+  constructor(params: StuctureColormakerParams) {
+    super(params)
+    if (!params.scale) {
+      this.parameters.scale = 'rwb'
+    }
+    this.atomData = this.parameters.data?.atomData
+    this.bondData = this.parameters.data?.bondData
+    this.scale = this.getScale(this.parameters)
+  }
+
+  @manageColor
+  atomColor(a: AtomProxy) {
+    const val = this.atomData?.[a.index]
+    return  (val !== undefined) ? this.scale(val) : this.parameters.value
+  }
+
+  @manageColor
+  bondColor(bond: BondProxy, fromTo: boolean) {
+      const val = this.bondData?.[bond.index]
+      return (val !== undefined) ?  this.scale(val) : this.parameters.value
+  }
+}
+
+ColormakerRegistry.add('structuredata', StructuredataColormaker)
+
+export default StructuredataColormaker

src/ngl.ts

@@ -73,6 +73,7 @@ import './color/randomcoilindex-colormaker'
 import './color/residueindex-colormaker'
 import './color/resname-colormaker'
 import './color/sstruc-colormaker'
+import './color/structuredata-colormaker'
 import './color/uniform-colormaker'
 import './color/value-colormaker'
 import './color/volume-colormaker'

src/representation/representation.ts

@@ -12,7 +12,7 @@ import Queue from '../utils/queue'
 import Counter from '../utils/counter'
 import Viewer from '../viewer/viewer'
 import { BufferParameters, BufferSide, default as Buffer } from '../buffer/buffer';
-import { ColormakerParameters, ColorMode } from '../color/colormaker';
+import { ColorData, ColormakerParameters, ColorMode } from '../color/colormaker';
 
 export interface RepresentationParameters {
   name: string
@@ -25,6 +25,7 @@ export interface RepresentationParameters {
   depthWrite: boolean,
   side: BufferSide,
   wireframe: boolean,
+  colorData: ColorData,
   colorScheme: string,
   colorScale: string | number[],
   colorReverse: boolean,
@@ -65,6 +66,7 @@ export interface RepresentationParameters {
  * @property {String} [side] - which triangle sides to render, "front" front-side,
  *                            "back" back-side, "double" front- and back-side
  * @property {Boolean} [wireframe] - render as wireframe
+ * @property {ColorData} [colorData] - atom or bond indexed data for coloring
  * @property {String} [colorScheme] - color scheme
  * @property {String} [colorScale] - color scale, either a string for a
  *                                 predefined scale or an array of
@@ -111,6 +113,7 @@ class Representation {
   protected depthWrite: boolean
   protected side: BufferSide
   protected wireframe: boolean
+  protected colorData: ColorData
   protected colorScheme: string
   protected colorScale: string | string[]
   protected colorReverse: boolean
@@ -181,6 +184,11 @@ class Representation {
         type: 'boolean', buffer: true
       },
 
+      colorData: {
+        type: 'hidden',
+        update: 'color',
+      },
+
       colorScheme: {
         type: 'select',
         update: 'color',
@@ -284,6 +292,7 @@ class Representation {
 
     this.setColor(p.color, p)
 
+    this.colorData = defaults(p.colorData, undefined)
     this.colorScheme = defaults(p.colorScheme, 'uniform')
     this.colorScale = defaults(p.colorScale, '')
     this.colorReverse = defaults(p.colorReverse, false)
@@ -370,6 +379,7 @@ class Representation {
   getColorParams (p?: {[k: string]: any}): { scheme: string, [k: string]: any } & ColormakerParameters {
     return Object.assign({
 
+      data: this.colorData,
       scheme: this.colorScheme,
       scale: this.colorScale,
       reverse: this.colorReverse,