Cube files viewer

handle_cubes.ipynb

@@ -0,0 +1,216 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%javascript\n",
+    "IPython.OutputArea.prototype._should_scroll = function(lines) {\n",
+    "    return false;\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# AiiDA imports.\n",
+    "%load_ext aiida\n",
+    "%aiida\n",
+    "\n",
+    "import io\n",
+    "import os\n",
+    "import tempfile\n",
+    "import urllib\n",
+    "\n",
+    "import ase.io.cube\n",
+    "import ipywidgets as ipw\n",
+    "import nglview\n",
+    "import numpy as np\n",
+    "import traitlets as tl\n",
+    "from aiida import orm\n",
+    "from aiida_shell import ShellJob\n",
+    "from cubehandler import Cube\n",
+    "from IPython.display import clear_output, display"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CubeArrayData3dViewerWidget(ipw.VBox):\n",
+    "    \"\"\"Widget to View 3-dimensional AiiDA ArrayData object in 3D.\"\"\"\n",
+    "\n",
+    "    cube = tl.Instance(Cube, allow_none=True)\n",
+    "\n",
+    "    def __init__(self, **kwargs):\n",
+    "        self.data_3d = None\n",
+    "        self.structure = None\n",
+    "        self.viewer = nglview.NGLWidget()\n",
+    "        self.orb_isosurf_slider = ipw.FloatSlider(\n",
+    "            continuous_update=False,\n",
+    "            value=1e-3,\n",
+    "            min=1e-4,\n",
+    "            max=1e-2,\n",
+    "            step=1e-4,\n",
+    "            description=\"Isovalue\",\n",
+    "            readout_format=\".1e\",\n",
+    "        )\n",
+    "        self.orb_isosurf_slider.observe(\n",
+    "            lambda c: self.set_cube_isosurf([c[\"new\"], -c[\"new\"]], [\"red\", \"blue\"]),\n",
+    "            names=\"value\",\n",
+    "        )\n",
+    "        super().__init__([self.viewer, self.orb_isosurf_slider], **kwargs)\n",
+    "\n",
+    "    @tl.observe(\"cube\")\n",
+    "    def on_observe_cube(self, _=None):\n",
+    "        \"\"\"Update object attributes when cube trait is modified.\"\"\"\n",
+    "\n",
+    "        self.data_3d = self.cube.data\n",
+    "        self.structure = self.cube.ase_atoms\n",
+    "        self.update_plot()\n",
+    "\n",
+    "    def update_plot(self):\n",
+    "        \"\"\"Update the 3D plot.\"\"\"\n",
+    "        while hasattr(self.viewer, \"component_0\"):\n",
+    "            self.viewer.component_0.clear_representations()\n",
+    "            self.viewer.remove_component(self.viewer.component_0.id)\n",
+    "        self.setup_cube_plot()\n",
+    "        self.set_cube_isosurf(\n",
+    "            [\n",
+    "                self.orb_isosurf_slider.value,\n",
+    "                -self.orb_isosurf_slider.value,\n",
+    "            ],\n",
+    "            [\"red\", \"blue\"],\n",
+    "        )\n",
+    "\n",
+    "    def setup_cube_plot(self):\n",
+    "        \"\"\"Setup cube plot.\"\"\"\n",
+    "        n_repeat = 2\n",
+    "        atoms_xn = self.structure.repeat((n_repeat, 1, 1))\n",
+    "        data_xn = np.tile(self.data_3d, (n_repeat, 1, 1))\n",
+    "        self.viewer.add_component(nglview.ASEStructure(atoms_xn))\n",
+    "        with tempfile.NamedTemporaryFile(mode=\"w\") as tempf:\n",
+    "            ase.io.cube.write_cube(tempf, atoms_xn, data_xn)\n",
+    "            c_2 = self.viewer.add_component(tempf.name, ext=\"cube\")\n",
+    "            c_2.clear()\n",
+    "\n",
+    "    def set_cube_isosurf(self, isovals, colors):\n",
+    "        \"\"\"Set cube isosurface.\"\"\"\n",
+    "        if hasattr(self.viewer, \"component_1\"):\n",
+    "            c_2 = self.viewer.component_1\n",
+    "            c_2.clear()\n",
+    "            for isov, col in zip(isovals, colors):\n",
+    "                c_2.add_surface(color=col, isolevelType=\"value\", isolevel=isov)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class HandleCubeFiles(ipw.VBox):\n",
+    "\n",
+    "    node_pk = tl.Int(None, allow_none=True)\n",
+    "\n",
+    "    def __init__(self):\n",
+    "        self.node = None\n",
+    "        self.output = ipw.Output()\n",
+    "        self.show_selected = ipw.Button(description=\"Show selected cube\")\n",
+    "        self.show_selected.on_click(self.show_selected_cube)\n",
+    "        self.dict_cube_files = {}\n",
+    "\n",
+    "        self.select_calc_widget = ipw.Dropdown(\n",
+    "            description=\"Calculation:\", options=self.get_calcs()\n",
+    "        )\n",
+    "\n",
+    "        tl.dlink(\n",
+    "            (self, \"node_pk\"),\n",
+    "            (self.select_calc_widget, \"value\"),\n",
+    "            transform=lambda x: orm.load_node(x).uuid if x else None,\n",
+    "        )\n",
+    "        self.select_calc_widget.observe(self.select_calculation)\n",
+    "        self.select_calculation()\n",
+    "        super().__init__([self.select_calc_widget, self.show_selected, self.output])\n",
+    "\n",
+    "    def list_cube_files(self):\n",
+    "        # Keep only the files that end with .cube\n",
+    "        cube_files = self.node.list_object_names()\n",
+    "        with self.output:\n",
+    "            clear_output()\n",
+    "            for f in cube_files:\n",
+    "                cube_selector = ipw.Checkbox(\n",
+    "                    description=f,\n",
+    "                    value=True,\n",
+    "                    style={\"description_width\": \"initial\"},\n",
+    "                    layout={\"width\": \"initial\"},\n",
+    "                )\n",
+    "                display(cube_selector)\n",
+    "                self.dict_cube_files[f] = cube_selector\n",
+    "        return\n",
+    "\n",
+    "    def show_selected_cube(self, _):\n",
+    "        for name, widget in self.dict_cube_files.items():\n",
+    "            if widget.value:\n",
+    "                cube = Cube.from_content(self.node.get_object_content(name))\n",
+    "                viewer = CubeArrayData3dViewerWidget(cube=cube)\n",
+    "                display(viewer)\n",
+    "\n",
+    "    def get_calcs(self):\n",
+    "        query = QueryBuilder()\n",
+    "        query.append(ShellJob, filters={\"label\": \"cube-shrink\"}, project=\"uuid\")\n",
+    "        return query.all(flat=True)\n",
+    "\n",
+    "    def select_calculation(self, _=None):\n",
+    "        if not self.select_calc_widget.value:\n",
+    "            return\n",
+    "        calc = orm.load_node(self.select_calc_widget.value)\n",
+    "        self.node = calc.outputs.out_cubes\n",
+    "        self.list_cube_files()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cube_files = HandleCubeFiles()\n",
+    "url = urllib.parse.urlsplit(jupyter_notebook_url)\n",
+    "try:\n",
+    "    cube_files.node_pk = int(urllib.parse.parse_qs(url.query)[\"pk\"][0])\n",
+    "except KeyError:\n",
+    "    pass\n",
+    "display(cube_files)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

start.py

@@ -28,6 +28,7 @@ def get_start_widget(appbase, jupbase):
     <td valign="top"><ul>
         <li><a href="{appbase}/submit_spm.ipynb" target="_blank">Scanning probe microscopy</a>
         <li><a href="{appbase}/submit_pdos.ipynb" target="_blank">Projected density of states</a>
+        <li><a href="{appbase}/handle_cubes.ipynb" target="_blank">Handle cube files</a>
     </ul></td>
 
     <!--

surfaces_tools/widgets/search_structures_widget.py

@@ -18,6 +18,7 @@
     "CP2K_Phonons": "view_phonons.ipynb",
     "CP2K_NEB": "view_neb.ipynb",
     "CP2K_Replica": "view_replica.ipynb",
+    "cube-shrink": "handle_cubes.ipynb",
     "ReplicaWorkChain": "view_replica.ipynb",
 }