diff --git a/js/src/widget_ngl.ts b/js/src/widget_ngl.ts
index 60201dd7..13b257f6 100644
--- a/js/src/widget_ngl.ts
+++ b/js/src/widget_ngl.ts
@@ -21,12 +21,12 @@ NGL.nglview_debug = false
 
 // From NGL
 // http://www.broofa.com/Tools/Math.uuid.htm
-const chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'.split('')
-const uuid = new Array(36)
+const chars: string[] = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'.split('')
+const uuid: string[] = new Array(36)
 
-function generateUUID() {
+function generateUUID(): string {
     let rnd = 0
-    let r
+    let r: number
 
     for (let i = 0; i < 36; i++) {
         if (i === 8 || i === 13 || i === 18 || i === 23) {
@@ -45,7 +45,7 @@ function generateUUID() {
 }
 
 
-async function createView(that, trait_name) {
+async function createView(that: any, trait_name: string): Promise<any> {
     // Create a view for the model with given `trait_name`
     // e.g: in backend, 'view.<trait_name>`
     console.log("Creating view for model " + trait_name);
@@ -57,7 +57,7 @@ async function createView(that, trait_name) {
 
 export
     class NGLModel extends widgets.DOMWidgetModel {
-    defaults() {
+    defaults(): any {
         return _.extend(widgets.DOMWidgetModel.prototype.defaults(), {
             _model_name: 'NGLModel',
             _model_module: 'nglview-js-widgets',
@@ -71,7 +71,21 @@ export
 
 export
     class NGLView extends widgets.DOMWidgetView {
-    render() {
+    stage: any;
+    $container: any;
+    uuid: string;
+    stage_widget: any;
+    comp_uuids: string[];
+    _synced_model_ids: any;
+    _synced_repr_model_ids: any;
+    ngl_view_id: string;
+    _ngl_focused: number;
+    btn_pview_fullscreen: any;
+    player_pview: any;
+    pgui_view: any;
+    image_btn_pview: any;
+
+    render(): void {
         this.beforeDisplay()
         this.displayed.then(function () {
             // move all below code inside 'displayed'
@@ -85,7 +99,7 @@ export
 
     }
 
-    beforeDisplay() {
+    beforeDisplay(): void {
         this.model.on("change:_parameters", this.parametersChanged, this);
         this.model.on("change:gui_style", this.GUIStyleChanged, this);
         this.model.set('_ngl_version', NGL.Version);
@@ -102,7 +116,7 @@ export
         }
     }
 
-    createStage() {
+    createStage(): void {
         // init NGL stage
         var stage_params = {
             // Shallow copy so that _ngl_full_stage_parameters is not updated yet
@@ -145,11 +159,11 @@ export
         }
     }
 
-    isEmbeded() {
+    isEmbeded(): boolean {
         return (this.model.get("_ngl_serialize") || (this.model.comm == undefined))
     }
 
-    handleMessage() {
+    handleMessage(): void {
         this.model.on("msg:custom", function (msg) {
             this.on_msg(msg);
         }, this);
@@ -167,7 +181,7 @@ export
     }
 
 
-    finalizeDisplay() {
+    finalizeDisplay(): void {
         // for callbacks from Python
         // must be after initializing NGL.Stage
         this.send({
@@ -190,7 +204,7 @@ export
         // this.touch()
     }
 
-    handleSignals() {
+    handleSignals(): void {
         var container = this.stage.viewer.container;
         var that = this;
         container.addEventListener('mouseover', function (e) {
@@ -283,7 +297,7 @@ export
 
     }
 
-    handlePicking() {
+    handlePicking(): void {
         this.$pickingInfo = $("<div></div>")
             .css("position", "absolute")
             .css("top", "5%")
@@ -331,7 +345,7 @@ export
         }, this);
     }
 
-    async mouseOverDisplay(type) {
+    async mouseOverDisplay(type: string): Promise<void> {
         var that = this;
         if (this.btn_pview_fullscreen) {
             var btn = await this.btn_pview_fullscreen
@@ -355,7 +369,7 @@ export
         }
     }
 
-    updateNGLTheme(css_content) {
+    updateNGLTheme(css_content: string): void {
         var ele = document.getElementById("nglview_style")
         if (ele != undefined) {
             document.head.removeChild(ele)
@@ -367,24 +381,24 @@ export
         document.head.appendChild(style)
     }
 
-    serialize_camera_orientation() {
+    serialize_camera_orientation(): void {
         var m = this.stage.viewerControls.getOrientation();
         this.model.set('_camera_orientation', m.elements);
         this.touch();
     }
 
-    set_camera_orientation(orientation) {
+    set_camera_orientation(orientation: any): void {
         if (orientation.length > 0) {
             this.stage.viewerControls.orient(orientation);
             this.serialize_camera_orientation();
         }
     }
 
-    executeCode(code) {
+    executeCode(code: string): void {
         eval(code);
     }
 
-    _handleEmbedBeforeStage() {
+    _handleEmbedBeforeStage(): void {
         // Only need to reconstruct colors in embeding mode (outside notebook)
         // FIXME: remove this function
         var that = this
@@ -398,7 +412,7 @@ export
         }
     }
 
-    async handleEmbed() {
+    async handleEmbed(): Promise<void> {
         var that = this;
         var ngl_msg_archive = that.model.get("_ngl_msg_archive");
         var ngl_stage_params = that.model.get('_ngl_full_stage_parameters');
@@ -465,7 +479,7 @@ export
         that.handleResize() // FIXME: really need this?
     }
 
-    updateCoordinatesFromDict(cdict, frame_index) {
+    updateCoordinatesFromDict(cdict: any, frame_index: number): void {
         // update coordinates for given "index"
         // cdict = Dict[int, List[base64]]
         var keys = Object.keys(cdict).filter(k => (k !== 'n_frames'));
@@ -479,20 +493,20 @@ export
         }
     }
 
-    requestFrame() {
+    requestFrame(): void {
         this.send({
             'type': 'request_frame',
             'data': 'frame'
         });
     }
 
-    requestUpdateStageParameters() {
+    requestUpdateStageParameters(): void {
         var updated_params = this.stage.getParameters();
         this.model.set('_ngl_full_stage_parameters', updated_params);
         this.touch();
     }
 
-    requestReprParameters(component_index, repr_index) {
+    requestReprParameters(component_index: number, repr_index: number): void {
         var comp = this.stage.compList[component_index];
         var repr = comp.reprList[repr_index];
         var msg = repr.repr.getParameters();
@@ -506,7 +520,7 @@ export
         }
     }
 
-    request_repr_dict() {
+    request_repr_dict(): void {
         var repr_dict = this.getReprDictFrontEnd()
         this.send({
             // make sure we are using "request_repr_dict" name
@@ -529,7 +543,7 @@ export
         }
     }
 
-    getReprDictFrontEnd() {
+    getReprDictFrontEnd(): any {
         var repr_dict = {};
         var n_components = this.stage.compList.length;
         for (var i = 0; i < n_components; i++) {
@@ -546,7 +560,7 @@ export
         return repr_dict
     }
 
-    syncReprForAllViews() {
+    syncReprForAllViews(): void {
         var repr_dict_backend = this.model.get("_ngl_repr_dict")
         var repr_dict_frontend = this.getReprDictFrontEnd()
         if (JSON.stringify(repr_dict_frontend) !== JSON.stringify(repr_dict_backend)) {
@@ -554,7 +568,7 @@ export
         }
     }
 
-    async syncReprWithMe() {
+    async syncReprWithMe(): Promise<void> {
         // Make sure views of the same model has the same representations
         // Only needed if we use Sidebar that connects to specific view.
         var that = this
@@ -568,26 +582,26 @@ export
         this.request_repr_dict()
     }
 
-    setSyncRepr(model_ids) {
+    setSyncRepr(model_ids: any): void {
         this._synced_repr_model_ids = model_ids
     }
 
-    setSyncCamera(model_ids) {
+    setSyncCamera(model_ids: any): void {
         this._synced_model_ids = model_ids
     }
 
-    viewXZPlane() {
+    viewXZPlane(): void {
         var m = new NGL.Matrix4().makeRotationX(Math.PI / 2);
         var q = new NGL.Quaternion().setFromRotationMatrix(m);
         this.stage.viewerControls.rotate(q);
     }
 
-    set_representation_from_backend() {
+    set_representation_from_backend(): void {
         var repr_dict = this.model.get('_ngl_repr_dict')
         this._set_representation_from_repr_dict(repr_dict)
     }
 
-    _set_representation_from_repr_dict(repr_dict) {
+    _set_representation_from_repr_dict(repr_dict: any): void {
         var compList = this.stage.compList
         if (compList.length > 0) {
             for (var index in repr_dict) {
@@ -604,7 +618,7 @@ export
         }
     }
 
-    async createView(trait_name) {
+    async createView(trait_name: string): Promise<any> {
         // Create a view for the model with given `trait_name`
         // e.g: in backend, 'view.<trait_name>`
         console.log("Creating view for model " + trait_name);
@@ -613,13 +627,13 @@ export
         return await manager.create_view(await manager.get_model(model_id))
     }
 
-    getPlayerModel() {
+    getPlayerModel(): any {
         // return a Promise
         var model_id = this.model.get("_iplayer").replace("IPY_MODEL_", "");
         return this.model.widget_manager.get_model(model_id)
     }
 
-    async createIPlayer() {
+    async createIPlayer(): Promise<void> {
         this.player_pview = this.createView("_iplayer");
         var view = await this.player_pview
         var that = this;
@@ -633,7 +647,7 @@ export
         pe.style.display = 'none'
     }
 
-    async createImageBtn() {
+    async createImageBtn(): Promise<void> {
         this.image_btn_pview = this.createView("_ibtn_image");
         var view = await this.image_btn_pview
         var pe = view.el
@@ -646,7 +660,7 @@ export
         this.stage.viewer.container.append(view.el);
     }
 
-    async createFullscreenBtn() {
+    async createFullscreenBtn(): Promise<void> {
         this.btn_pview_fullscreen = this.createView("_ibtn_fullscreen");
         var view = await this.btn_pview_fullscreen
         var stage = this.stage;
@@ -675,7 +689,7 @@ export
     }
 
 
-    async createGUI() {
+    async createGUI(): Promise<void> {
         this.pgui_view = this.createView("_igui");
         var view = await this.pgui_view
         var pe = view.el
@@ -688,12 +702,12 @@ export
     }
 
 
-    createNglGUI() {
+    createNglGUI(): void {
         this.stage_widget = new StageWidget(this)
     }
 
 
-    setVisibilityForRepr(component_index, repr_index, value) {
+    setVisibilityForRepr(component_index: number, repr_index: number, value: boolean): void {
         // value = True/False
         var component = this.stage.compList[component_index];
         var repr = component.reprList[repr_index];
@@ -703,7 +717,7 @@ export
         }
     }
 
-    removeRepresentation(component_index, repr_index) {
+    removeRepresentation(component_index: number, repr_index: number): void {
         var component = this.stage.compList[component_index];
         var repr = component.reprList[repr_index]
 
@@ -712,7 +726,7 @@ export
         }
     }
 
-    removeRepresentationsByName(repr_name, component_index) {
+    removeRepresentationsByName(repr_name: string, component_index: number): void {
         var component = this.stage.compList[component_index];
 
         if (component) {
@@ -724,7 +738,7 @@ export
         }
     }
 
-    updateRepresentationForComponent(repr_index, component_index, params) {
+    updateRepresentationForComponent(repr_index: number, component_index: number, params: any): void {
         var component = this.stage.compList[component_index];
         var that = this;
         var repr = component.reprList[repr_index];
@@ -733,7 +747,7 @@ export
         }
     }
 
-    updateRepresentationsByName(repr_name, component_index, params) {
+    updateRepresentationsByName(repr_name: string, component_index: number, params: any): void {
         var component = this.stage.compList[component_index];
         var that = this;
 
@@ -747,7 +761,7 @@ export
         }
     }
 
-    setRepresentation(name, params, component_index, repr_index) {
+    setRepresentation(name: string, params: any, component_index: number, repr_index: number): void {
         var component = this.stage.compList[component_index];
         var repr = component.reprList[repr_index];
         var that = this;
@@ -765,7 +779,7 @@ export
         }
     }
 
-    setColorByResidue(colors, component_index, repr_index) {
+    setColorByResidue(colors: any, component_index: number, repr_index: number): void {
         var repr = this.stage.compList[component_index].reprList[repr_index];
         var schemeId = NGL.ColormakerRegistry.addScheme(function (params) {
             this.atomColor = function (atom) {
@@ -777,7 +791,7 @@ export
         repr.setColor(schemeId);
     }
 
-    addShape(name, shapes) {
+    addShape(name: string, shapes: any[]): void {
         // shapes: List[Tuple[str, ...]]
         // e.g: [('sphere', ...), ('cone', ...)]
         var shape = new NGL.Shape(name);
@@ -808,7 +822,7 @@ export
         shapeComp.addRepresentation("buffer");
     }
 
-    addBuffer(name, kwargs) {
+    addBuffer(name: string, kwargs: any): void {
         var class_dict = {
             "arrow": NGL.ArrowBuffer,
             "box": NGL.BoXbuffer,
@@ -834,7 +848,7 @@ export
         shapeComp.addRepresentation("buffer");
     }
 
-    async replaceStructure(structure) {
+    async replaceStructure(structure: any): Promise<void> {
         var blob = new Blob([structure.data], { type: "text/plain" });
         var stage = this.stage
         var params = structure.params || {};
@@ -856,14 +870,14 @@ export
         this._handleLoadFileFinished();
     }
 
-    superpose(cindex0, cindex1, align, sele0, sele1) {
+    superpose(cindex0: number, cindex1: number, align: boolean, sele0: string, sele1: string): void {
         // superpose two components with given params
         var component0 = this.stage.compList[cindex0];
         var component1 = this.stage.compList[cindex1];
         component1.superpose(component0, align, sele0, sele1);
     }
 
-    decode_base64(base64) {
+    decode_base64(base64: string): ArrayBuffer {
         // lightly adapted from Niklas
 
         /*
@@ -904,7 +918,7 @@ export
         return arraybuffer;
     }
 
-    updateCoordinates(coordinates, model) {
+    updateCoordinates(coordinates: ArrayBuffer, model: number): void {
         // coordinates must be ArrayBuffer (use this.decode_base64)
         var component = this.stage.compList[model];
         if (coordinates && component) {
@@ -916,7 +930,7 @@ export
         }
     }
 
-    handleResizable() {
+    handleResizable(): void {
         this.$container.resizable({
             resize: function (event, ui) {
                 this.setSize(ui.size.width + "px", ui.size.height + "px");
@@ -924,7 +938,7 @@ export
         })
     }
 
-    handleResize() {
+    handleResize(): void {
         var width = this.$el.width()
         var height = this.$el.height() + "px"
         if (this.stage_widget) {
@@ -934,13 +948,13 @@ export
         this.setSize(width, height)
     }
 
-    setSize(width, height) {
+    setSize(width: string, height: string): void {
         this.stage.viewer.container.style.width = width;
         this.stage.viewer.container.style.height = height;
         this.stage.handleResize();
     }
 
-    GUIStyleChanged() {
+    GUIStyleChanged(): void {
         var style = this.model.get("gui_style");
         if (style === 'ngl') {
             this.createNglGUI();
@@ -956,12 +970,12 @@ export
         }
     }
 
-    parametersChanged() {
+    parametersChanged(): void {
         var _parameters = this.model.get("_parameters");
         this.setParameters(_parameters);
     }
 
-    setParameters(parameters) {
+    setParameters(parameters: any): void {
         this.stage.setParameters(parameters);
 
         // do not set _ngl_full_stage_parameters here
@@ -974,14 +988,14 @@ export
         })
     }
 
-    async _downloadImage(filename, params) {
+    async _downloadImage(filename: string, params: any): Promise<void> {
         if (this.ngl_view_id == this.get_last_child_id()) {
             var blob = await this.stage.makeImage(params)
             NGL.download(blob, filename);
         }
     }
 
-    async _exportImage(wid, params) {
+    async _exportImage(wid: string, params: any): Promise<void> {
         if (this.ngl_view_id == this.get_last_child_id()) {
             var blob = await this.stage.makeImage(params)
             var reader = new FileReader();
@@ -1001,7 +1015,7 @@ export
         }
     }
 
-    async handleMovieMaking(render_params) {
+    async handleMovieMaking(render_params: any): Promise<void> {
         console.log('handleMovieMaking: render_params', render_params)
         if (this.ngl_view_id == this.get_last_child_id()) {
             var blob = await this.stage.makeImage(render_params)
@@ -1020,11 +1034,11 @@ export
         }
     }
 
-    async _handleLoadFileFinished() {
+    async _handleLoadFileFinished(): Promise<void> {
         this.send({ 'type': 'async_message', 'data': 'ok' });
     }
 
-    _getLoadFilePromise(msg) {
+    _getLoadFilePromise(msg: any): any {
         // args = [{'type': ..., 'data': ...}]
         var args0 = msg.args[0];
         if (args0.type == 'blob') {
@@ -1060,12 +1074,12 @@ export
         }
     }
 
-    get_last_child_id() {
+    get_last_child_id(): string {
         var keys = this.model.get('_ngl_view_id')
         return keys[keys.length - 1]
     }
 
-    async _handleStageLoadFile(msg) {
+    async _handleStageLoadFile(msg: any): Promise<void> {
         // args = [{'type': ..., 'data': ...}]
         if (this.ngl_view_id != this.get_last_child_id() && msg.last_child) {
             return
@@ -1074,7 +1088,7 @@ export
         this._handleLoadFileFinished();
     }
 
-    addColorScheme(args, label) {
+    addColorScheme(args: any, label: string): string {
         var id = NGL.ColormakerRegistry.addSelectionScheme(args, label);
         var scheme = NGL.ColormakerRegistry.userSchemes[id];
         NGL.ColormakerRegistry.removeScheme(id);
@@ -1083,7 +1097,7 @@ export
         return label
     }
 
-    async on_msg(msg) {
+    async on_msg(msg: any): Promise<void> {
         if (msg.type === 'call_method') {
             await this.handleCallMethod(msg);
         } else if (msg.type === 'base64_single') {
@@ -1097,7 +1111,7 @@ export
         }
     }
 
-    async handleCallMethod(msg) {
+    async handleCallMethod(msg: any): Promise<void> {
         var index, component, func, stage;
         var new_args = msg.args.slice();
         new_args.push(msg.kwargs);
@@ -1138,7 +1152,7 @@ export
         }
     }
 
-    async handleStageMethod(msg, new_args) {
+    async handleStageMethod(msg: any, new_args: any[]): Promise<void> {
         var stage_func = this.stage[msg.methodName];
         if (msg.methodName === 'screenshot') {
             NGL.screenshot(this.stage.viewer, msg.kwargs);
@@ -1155,7 +1169,7 @@ export
         }
     }
 
-    handleBase64Single(coordinatesDict) {
+    handleBase64Single(coordinatesDict: any): void {
         var keys = Object.keys(coordinatesDict);
         for (var i = 0; i < keys.length; i++) {
             var traj_index = keys[i];
@@ -1166,7 +1180,7 @@ export
         }
     }
 
-    handleBinarySingle(msg) {
+    handleBinarySingle(msg: any): void {
         var coordinateMeta = msg.data;
         var keys = Object.keys(coordinateMeta);
         for (var i = 0; i < keys.length; i++) {
@@ -1181,7 +1195,7 @@ export
         }
     }
 
-    handleGetRequest(data) {
+    handleGetRequest(data: string): void {
         if (data === 'camera') {
             this.send(JSON.stringify(this.stage.viewer.camera));
         } else if (data === 'parameters') {
diff --git a/notebooks/custom_color.ipynb b/notebooks/custom_color.ipynb
index e10e0f43..1f5f61bc 100644
--- a/notebooks/custom_color.ipynb
+++ b/notebooks/custom_color.ipynb
@@ -9,9 +9,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0b70ddacc2444deba95428ae8859ce96",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": []
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "import nglview as nv\n",
     "from nglview.color import ColormakerRegistry\n",
@@ -23,9 +36,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bcbbe36f03144db19498457546c6a197",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "NGLWidget()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "v0 = nv.demo(default=False)\n",
     "v0.center()\n",
@@ -35,7 +63,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -51,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -70,9 +98,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "91641312855c4cc6b4396b4de21ace78",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "NGLWidget()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "v1 = nv.demo(default=False)\n",
     "v1.center()\n",
@@ -82,7 +125,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -96,18 +139,44 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "76ff3bd245e3478bad19c1bbd4f624b8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "ThemeManager()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "v1.gui_style = 'ngl'"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'2.2.2'"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "v1._ngl_version"
    ]