support colorscheme changes on surfaces

Issue #729

Author: dkoes

src/GLViewer.ts

@@ -17,6 +17,16 @@ import { decode, toRGBA8, encode } from 'upng-js'
 
 export const CONTEXTS_PER_VIEWPORT = 16;
 
+interface SurfObj {
+    geo: Geometry;
+    mat: Material;
+    done: Boolean;
+    finished: Boolean;
+    lastGL?: any;
+    symmetries?: any[];
+    style?: SurfaceStyleSpec;
+}
+
 /**
  * WebGL-based 3Dmol.js viewer
  * Note: The preferred method of instantiating a GLViewer is through {@link createViewer}
@@ -38,7 +48,7 @@ export class GLViewer {
     private glDOM: HTMLCanvasElement | null = null;
 
     private models: GLModel[] = []; // atomistic molecular models
-    private surfaces: any = {};
+    private surfaces: Record<number,SurfObj[]> = {};
     private shapes = []; // Generic shapes
     private labels: Label[] = [];
     private clickables = []; //things you can click on
@@ -1826,7 +1836,6 @@ export class GLViewer {
                         geo.normalsNeedUpdate = true;
                         geo.colorsNeedUpdate = true;
                         geo.buffersNeedUpdate = true;
-                        geo.boundingSphere = null;
                         surfArr[n].mat.needsUpdate = true;
 
                         if (surfArr[n].done)
@@ -2512,7 +2521,7 @@ export class GLViewer {
               viewer.render();
             });
      */
-    public addPropertyLabels(prop: string, sel: AtomSelectionSpec, style: AtomStyleSpec) {
+    public addPropertyLabels(prop: string, sel: AtomSelectionSpec, style: LabelSpec) {
         this.applyToModels("addPropertyLabels", prop, sel, this, style);
         this.show();
         return this;
@@ -4083,19 +4092,22 @@ export class GLViewer {
 
         //set colorArray of there are per-atom colors
         var colorArray = geoGroup.colorArray;
+        let atomArray = geoGroup.atomArray;
 
         if (mat.voldata && mat.volscheme) {
             //convert volumetric data into colors
             var scheme = mat.volscheme;
             var voldata = mat.voldata;
             var range = scheme.range() || [-1, 1];
             for (let i = 0, il = v.length; i < il; i++) {
+                let A = v[i].atomid;
                 let val = voldata.getVal(v[i].x, v[i].y, v[i].z);
                 let col = CC.color(scheme.valueToHex(val, range));
                 let offset = i * 3;
                 colorArray[offset] = col.r;
                 colorArray[offset + 1] = col.g;
                 colorArray[offset + 2] = col.b;
+                atomArray[i] = atoms[A];
             }
         }
         else if (colors.length > 0) { //have atom colors
@@ -4106,6 +4118,7 @@ export class GLViewer {
                 colorArray[offsetA] = colors[A].r;
                 colorArray[offsetA + 1] = colors[A].g;
                 colorArray[offsetA + 2] = colors[A].b;
+                atomArray[i] = atoms[A];
             }
         }
 
@@ -4237,7 +4250,6 @@ export class GLViewer {
         return mat;
     }
 
-
     /**
      * Adds an explicit mesh as a surface object.
      * @param {Mesh}
@@ -4254,7 +4266,7 @@ export class GLViewer {
             finished: false //the rendered finishes surfaces when they are done
         };
         var surfid = this.nextSurfID();
-        this.surfaces[surfid] = surfobj;
+        this.surfaces[surfid] = [surfobj];
         return surfid;
     };
 
@@ -4597,15 +4609,14 @@ export class GLViewer {
         adjustVolumeStyle(style);
         if (this.surfaces[surf]) {
             var surfArr = this.surfaces[surf];
-            surfArr.style = style;
-            for (var i = 0; i < surfArr.length; i++) {
+            for (let i = 0; i < surfArr.length; i++) {
                 var mat = surfArr[i].mat = GLViewer.getMatWithStyle(style);
                 surfArr[i].mat.side = FrontSide;
                 if (style.color) {
-                    surfArr[i].mat.color = style.color;
+                    surfArr[i].mat.color = CC.color(style.color);
                     surfArr[i].geo.colorsNeedUpdate = true;
                     const c = CC.color(style.color);
-                    surfArr[i].geo.setColors(function () { return c; });
+                    surfArr[i].geo.setColor(c);
                 }
                 else if (mat.voldata && mat.volscheme) {
                     //convert volumetric data into colors
@@ -4618,6 +4629,17 @@ export class GLViewer {
                         let col = cc.color(scheme.valueToHex(val, range));
                         return col;
                     });
+                } else {
+                    surfArr[i].geo.colorsNeedUpdate = true;
+                    for(let geo of  surfArr[i].geo.geometryGroups ) {
+                        for(let j = 0; j < geo.vertices; j++) {
+                            let c = getColorFromStyle(geo.atomArray[j],style);
+                            let off = 3*j;
+                            geo.colorArray[off] = c.r;
+                            geo.colorArray[off+1] = c.g;
+                            geo.colorArray[off+2] = c.b;
+                        }
+                    }
                 }
                 surfArr[i].finished = false; // trigger redraw
             }

src/WebGL/core/Geometry.ts

@@ -3,6 +3,7 @@ import { LineBasicMaterial } from '../materials/LineBasicMaterial';
 import { EventDispatcher } from "./EventDispatcher";
 import { Vector3 } from "../math";
 import { CC, Color } from "../../colors";
+import { AtomSpec } from 'specs';
 const BUFFERSIZE = 65535; //limited to 16bit indices
 export class GeometryGroup {
   id: number;
@@ -12,6 +13,7 @@ export class GeometryGroup {
   radiusArray: Float32Array | null = null;
   faceArray: Uint16Array | null = null;
   lineArray: Uint16Array | null = null;
+  atomArray: Array<AtomSpec> = Array<AtomSpec>();
   vertices: number = 0;
   faceidx: number = 0;
   lineidx: number = 0;
@@ -22,7 +24,22 @@ export class GeometryGroup {
     this.id = id;
   }
 
-  setColors(setcolor: (r: number, g: number, b: number) => Color | number): void {
+  public setColor(color: Color | number): void {
+    //apply constant color
+    var v = this.vertexArray;
+    var c = this.colorArray;
+    if (!v) throw new Error("vertex array not initialized");
+    if (!c) throw new Error("color array not initialized");
+
+    let col = CC.color(color);
+    for (var i = 0; i < v.length; i += 3) {
+      c[i] = col.r;
+      c[i + 1] = col.g;
+      c[i + 2] = col.b;
+    }
+  }
+
+  setColors(setcolor: (x: number, y: number, z: number) => Color | number): void {
     //apply a function that takes the vertex coordinate and returns a color
     var v = this.vertexArray;
     var c = this.colorArray;
@@ -448,6 +465,13 @@ export class Geometry extends EventDispatcher {
     }
   }
 
+  setColor(...setcolor: Parameters<GeometryGroup["setColor"]>): void {
+    var len = this.geometryGroups.length;
+    for (var g = 0; g < len; g++) {
+      var geoGroup = this.geometryGroups[g];
+      geoGroup.setColor(...setcolor);
+    }
+  }
 
   setUpWireframe(...lineIndexArgs: Parameters<GeometryGroup["setLineIndices"]>) {
     for (var g = 0; g < this.groups; g++) {

tests/auto/tests/surfacechange.js

@@ -0,0 +1,9 @@
+
+$3Dmol.download("pdb:4DLN",viewer,{},function(){
+  viewer.setStyle({'cartoon':{colorscheme: 'ssPyMol'}});
+  var sb = viewer.addSurface("SAS", {colorscheme: 'ssJmol'} ,{hetflag:false},null,null, function(sb) {
+    viewer.render( /* no callback */);
+    viewer.setSurfaceMaterialStyle(sb, {colorscheme: 'ssPyMol',opacity:.7});
+    viewer.render();
+  });  
+});