ADD Custom color scheme API (#839)

js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "nglview-js-widgets",
-  "version": "2.6.2",
+  "version": "2.6.5",
   "description": "nglview-js-widgets",
   "author": "Hai Nguyen <[email protected]>, Alexander Rose <[email protected]>",
   "license": "MIT",

js/src/base.js

@@ -0,0 +1,34 @@
+var widgets = require("@jupyter-widgets/base")
+var NGL = require("ngl")
+
+
+var BaseView = widgets.DOMWidgetView.extend({
+
+    render: function(){
+        this.handleMessage();
+        this.displayed.then(() =>{
+            this.model.set("_ready", true)
+            this.touch()
+        })
+    },
+
+    executeCode: function(code){
+        eval(code);
+    },
+
+    handleMessage: function(){
+        this.model.on("msg:custom", function(msg){
+            this.on_msg(msg)
+        }.bind(this))
+    },
+
+    on_msg: function(msg){
+        if (msg.type == 'callMethod'){
+            this[msg.methodName].apply(this, msg.args, msg.kwargs)
+        }
+    }
+})
+
+module.exports = {
+    "BaseView": BaseView
+}

js/src/color.js

@@ -0,0 +1,60 @@
+var _ = require('underscore')
+var NGL = require("ngl")
+var BaseView = require("./base").BaseView
+var widgets = require("@jupyter-widgets/base")
+
+
+var ColormakerRegistryModel = widgets.DOMWidgetModel.extend({
+    defaults: function(){
+        return _.extend(widgets.DOMWidgetModel.prototype.defaults(), {
+            _model_name: 'ColormakerRegistryModel',
+            _model_module: 'nglview-js-widgets',
+            _model_module_version: require("../package.json").version,
+            _view_name: "ColormakerRegistryView",
+            _view_module: "nglview-js-widgets",
+            _view_module_version: require("../package.json").version,
+        });
+    }
+})
+
+
+var ColormakerRegistryView = BaseView.extend({
+    addSelectionScheme: function(label, args){
+        var id = NGL.ColormakerRegistry.addSelectionScheme(args, label)
+        this._updateId(id, label)
+    },
+
+    addSelectionSchemeOriginal: function(label, args){
+        var id = NGL.ColormakerRegistry.addSelectionScheme(args, label);
+        var scheme = NGL.ColormakerRegistry.userSchemes[id];
+        NGL.ColormakerRegistry.removeScheme(id);
+        // hard code the scheme ID
+        NGL.ColormakerRegistry.add(label, scheme);
+    },
+
+    addScheme: function(label, func_str){
+        var func = Function("return " + func_str)()
+        console.log(func)
+        var id = NGL.ColormakerRegistry.addScheme(function(params){
+            this.atomColor = func
+        })
+        this._updateId(id, label)
+    },
+
+    _updateId: function(oldId, newId){
+        var scheme = NGL.ColormakerRegistry.userSchemes[oldId]
+        console.log(oldId, scheme)
+        NGL.ColormakerRegistry.add(newId, scheme)
+        NGL.ColormakerRegistry.removeScheme(oldId)
+    },
+
+    removeScheme: function(schemeId){
+        NGL.ColormakerRegistry.removeScheme(schemeId)
+    },
+})
+
+
+module.exports = {
+    ColormakerRegistryView: ColormakerRegistryView,
+    ColormakerRegistryModel: ColormakerRegistryModel
+}

js/src/index.js

@@ -2,7 +2,7 @@ module.exports = {};
 
 var loadedModules = [
     require("./widget_ngl.js"),
-    //require("./ngl.js"),
+    require("./color.js"),
 ]
 
 for (var i in loadedModules) {

js/src/widget_ngl.js

@@ -1,6 +1,8 @@
 var Jupyter
 var widgets = require("@jupyter-widgets/base")
 var NGL = require('ngl')
+var ColormakerRegistryModel = require('./color').ColormakerRegistryModel
+var BaseView = require('./base').BaseView
 var $ = require('jquery')
 var _ = require('underscore')
 require("./lib/signals.min.js")
@@ -336,10 +338,13 @@ var NGLView = widgets.DOMWidgetView.extend({
         }
     },
 
-    execute_code: function(code){
+    executeCode: function(code){
         eval(code);
     },
 
+    handleCustomColor: function(){
+    },
+
     handle_embed: function(){
         var that = this;
         var ngl_coordinate_resource = that.model.get("_ngl_coordinate_resource");
@@ -350,8 +355,33 @@ var NGLView = widgets.DOMWidgetView.extend({
         var label
 
         // reconstruct colors
-        for (label in ngl_color_dict){
-            that.addColorScheme(ngl_color_dict[label], label);
+        if (this.model.comm === undefined){
+            var model_dict = this.model.widget_manager._models
+            var models = []
+            for (let k in model_dict){
+                models.push(model_dict[k])
+            }
+            Promise.all(models).then(models => {
+                for (var i in models){
+                    var model = models[i]
+                    if (model instanceof ColormakerRegistryModel){
+                        var k = Object.keys(model.views)[0] // singleton
+                        model.views[k].then(view =>{
+                            view.model.get("_msg_ar").forEach(msg =>{
+                                view.on_msg(msg)
+                            })
+                        })
+                    }
+                }
+            })
+
+            // Outside the notebook
+            // Old API (_ColorScheme)
+            for (label in ngl_color_dict){
+                if (!NGL.ColormakerRegistry.hasScheme(label)){
+                    that.addColorScheme(ngl_color_dict[label], label);
+                }
+            }
         }
 
         _.each(ngl_msg_archive, function(msg){
@@ -1125,7 +1155,7 @@ var FullscreenModel = widgets.DOMWidgetModel.extend({
     }
 })
 
-var FullscreenView = widgets.DOMWidgetView.extend({
+var FullscreenView = BaseView.extend({
     render: function() {
         this.stage = new NGL.Stage()
         var that = this
@@ -1153,13 +1183,13 @@ var FullscreenView = widgets.DOMWidgetView.extend({
         })
     },
 
-    execute_code: function(code){
+    executeCode: function(code){
         eval(code);
     },
 
     on_msg: function(msg){
-        if ('execute_code' in msg){
-            this.execute_code(msg.execute_code)
+        if ('executeCode' in msg){
+            this.executeCode(msg.executeCode)
         }
     }
 

nglview/_frontend.py

@@ -1 +1 @@
-__frontend_version__ = '2.6.2'
+__frontend_version__ = '2.6.5'

nglview/base.py

@@ -0,0 +1,24 @@
+from ipywidgets import DOMWidget
+from traitlets import Bool, List
+
+
+class BaseWidget(DOMWidget):
+    _msg_q = []
+    _msg_ar = List().tag(sync=True)
+    _ready = Bool(False).tag(sync=True)
+
+    def _js(self, code):
+        self._call("executeCode", code)
+
+    def _call(self, method_name, *args, **kwargs):
+        msg = {"type": "callMethod",
+               "methodName": method_name,
+                "args": args, "kwargs": kwargs}
+        if not self._ready:
+            # fire later
+            self._msg_q.append(msg)
+        else:
+            self.send(msg)
+        msg_ar = self._msg_ar[:]
+        msg_ar.append(msg)
+        self._msg_ar = msg_ar # trigger sync

nglview/color.py

@@ -1,3 +1,11 @@
+from ipywidgets import DOMWidget
+from traitlets import Unicode, Bool, observe
+from .base import BaseWidget
+from ._frontend import __frontend_version__
+from IPython.display import display
+import time
+
+
 COLOR_SCHEMES = [
     " ", "picking", "random", "uniform", "atomindex", "residueindex",
     "chainindex", "modelindex", "sstruc", "element", "resname", "bfactor",
@@ -7,6 +15,7 @@
 _USER_COLOR_DICT = {}
 
 
+
 class _ColorScheme:
     _color_dict = {}
 
@@ -19,3 +28,94 @@ def __init__(self, args, label):
     @property
     def data(self):
         return {'data': self._color_scheme, 'label': self._label}
+
+
+def _singleton(cls):
+    # https://www.python.org/dev/peps/pep-0318/#examples
+    instances = {}
+    def getinstance():
+        if cls not in instances:
+            instances[cls] = cls()
+        return instances[cls]
+    return getinstance
+
+
+@_singleton
+class _ColormakerRegistry(BaseWidget):
+    _view_name = Unicode("ColormakerRegistryView").tag(sync=True)
+    _view_module = Unicode("nglview-js-widgets").tag(sync=True)
+    _view_module_version = Unicode(__frontend_version__).tag(sync=True)
+    _model_name = Unicode("ColormakerRegistryModel").tag(sync=True)
+    _model_module = Unicode("nglview-js-widgets").tag(sync=True)
+    _model_module_version = Unicode(__frontend_version__).tag(sync=True)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        display(self)
+
+    @observe("_ready")
+    def _on_ready(self, change):
+        if change.new:
+            while self._msg_q:
+                msg = self._msg_q.pop(0)
+                self.send(msg)
+
+    def _ipython_display_(self, **kwargs):
+        if self._ready:
+            return str(self)
+        super()._ipython_display_(**kwargs)
+
+    def add_selection_scheme(self, scheme_id, arg):
+        """
+        Examples
+        --------
+        >>> ColormakerRegistry.add_selection_scheme("my_custom_scheme",
+                [['blue', '1-10']])
+        >>> view.add_cartoon(color="my_custom_scheme")
+        """
+        self._call("addSelectionScheme", scheme_id, arg)
+
+    def add_scheme_func(self, scheme_id, func_str):
+        """
+
+        Examples
+        --------
+        >>> func_str = '''
+             this.atomColor = function (atom) {
+             if (atom.serial < 1000) {
+               return 0x0000FF // blue
+             } else if (atom.serial > 2000) {
+               return 0xFF0000 // red
+             } else {
+               return 0x00FF00 // green
+             }
+             }
+         '''
+        >>> ColormakerRegistry.add_scheme_func('awesome', func_str)
+        >>> view.add_cartoon(color='awesome')
+        """
+
+        code = """
+            var schemeId = NGL.ColormakerRegistry.addScheme(function (params) {
+                %s
+            })
+            
+            this._updateId(schemeId, '%s')
+        """ % (func_str, scheme_id)
+        self._js(code)
+
+    def add_scheme(self, scheme_id, obj):
+        """
+        Parameters
+        ----------
+        obj: List of List or str (of JS function)
+        """
+        if isinstance(obj, list):
+            self.add_selection_scheme(scheme_id, obj)
+        elif isinstance(obj, str):
+            self.add_scheme_func(scheme_id, obj)
+        else:
+            raise ValueError(f"{obj} must be either list of list or string")
+
+
+ColormakerRegistry = _ColormakerRegistry()