Merge pull request #897 from fredludlow/colordata

Colordata addition

examples/scripts/color/atom-data.js

@@ -0,0 +1,30 @@
+/** Demonstration of structuredata colormaker, coloring atoms by numeric property passed
+ * in as colorData.atomData
+ */
+stage.loadFile('data://1lee.pdb').then(function (c) {
+  // We might be colouring by some calculated/predicted property (e.g. ML model which outputs
+  // scores/quantities for individual atoms)
+
+  // For our example, we'll generate some random "scores" for the ligand atoms in PDB entry 1lee
+  // (Ligand has hetcode R36), leaving all the protein atoms unscored
+
+  var structure = c.structure
+  var scores = []
+
+  var ligandView = structure.getView(new NGL.Selection('R36'))
+
+  ligandView.eachAtom(ap => {
+    // We only populate scores for some of the atoms (those in the R36 residue),
+    // For non-ligand atoms scores[i] will be undefined, and color will fallback to color value
+    scores[ap.index] = Math.random()
+  })
+
+  c.addRepresentation('licorice', {
+    color: 'structuredata',
+    colorData: {atomData: scores},
+    colorDomain: [0.0, 1.0], // This is the default domain
+    colorScale: 'rainbow', // Any of the normal color scales are available
+    colorValue: '#888' // Fallback color for atoms without data
+  })
+  c.autoView('R36')
+})

examples/scripts/color/bond-data.js

@@ -0,0 +1,19 @@
+/**
+ * Use of structuredata-colormaker with bonding info.
+ */
+
+stage.loadFile('data://adrenalin.mol2').then(function (o) {
+  var s = o.structure
+  var bondData = []
+  // Here the data is randomly generated, but could be a bond-specific
+  // property from an external source (e.g. estimated bond enthalpy)
+  s.eachBond(b => {
+    bondData.push(Math.random())
+  })
+
+  o.addRepresentation('ball+stick', {
+    colorScheme: 'structuredata',
+    colorData: { bondData: bondData },
+    colorValue: 'grey' // As colorData.atomData is not defined, the atom colors fall back to colorValue.
+  })
+})

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?: number[],
+  bondData?: number[]
+}
+
 export interface ColormakerParameters extends ScaleParameters {
   structure?: Structure
   volume?: Volume
   surface?: Surface
+  data?: ColorData
 }
 
 export type StuctureColormakerParams = { structure: Structure } & Partial<ColormakerParameters>
@@ -59,14 +65,14 @@ 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,
    descriptor: TypedPropertyDescriptor<colorFuncType>): PropertyDescriptor {
     const originalMethod = descriptor.value
-    const linearize: colorFuncType = function (this: T, value: any) {
-      let result = originalMethod!.bind(this, value)()
+    const linearize: colorFuncType = function (this: T, value: any, fromTo?: boolean) {
+      let result = originalMethod!.bind(this, value, fromTo)()
       if (colorSpace == 'linear') {
         tmpColor.set(result)
         tmpColor.convertSRGBToLinear()

src/color/structuredata-colormaker.ts

@@ -0,0 +1,54 @@
+/**
+ * @file Colordata Colormaker
+ * @author Fred Ludlow <[email protected]>
+ * @private
+ */
+
+import { ColormakerRegistry } from '../globals'
+import Colormaker, { ColorData, ColormakerScale, manageColor, StuctureColormakerParams } from './colormaker'
+import AtomProxy from '../proxy/atom-proxy'
+import BondProxy from '../proxy/bond-proxy'
+
+
+class StructuredataColormaker extends Colormaker {
+  atomData?: ColorData['atomData']
+  bondData?: ColorData['bondData']
+  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]
+
+      // Explicit bond data?
+      if (val !== undefined) return this.scale(val)
+
+
+      if (this.atomProxy) {
+        this.atomProxy.index = fromTo ? bond.atomIndex1 : bond.atomIndex2
+        return this.atomColor(this.atomProxy)
+      } 
+
+      // Fallback
+      return 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,