notebook-tutorials/notebook6_interactive_visualization_on_the_gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab72eb27-6c09-4a05-a8da-7dd6e42fe5b1",
   "metadata": {},
   "source": [
    "# Notebook 6: Ray tracing and interactive data visualization on the GPU\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38064b49",
   "metadata": {},
   "source": [
    "Here we illustrate how one can build custom widgets for interactive data visualization. \n",
    "We use the Paicos CUDA GPU implementation of ray tracing to achieve the necessary speed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66a11310-5311-439f-a61c-09dcf43461ae",
   "metadata": {},
   "source": [
    "## Required python packages\n",
    "\n",
    "This notebook requires that you have the GPU requirements installed and available on your system\n",
    "and that you have modified your Paicos user settings to load GPU functionality on startup.\n",
    "Please see the details here: https://paicos.readthedocs.io/en/latest/installation.html#gpu-cuda-requirements \n",
    "\n",
    "This notebook also requires that you have a working version of ipywidgets,\n",
    "which might sometimes be a bit cumbersome to get working (I have this working in a Jupyter notebook\n",
    "but there is a risk that you will have [trouble](https://discourse.jupyter.org/t/jupyter-lab-ipywidgets-no-longer-work-with-ipympl/14539/10) if you are using Jupyter Lab). You can simply try \n",
    "```\n",
    "pip install ipywidgets\n",
    "```\n",
    "and if that does not work then you can try with the detailed instructions found here: https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
    "\n",
    "You can check that ipywidgets is working by uncommenting and executing the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "46f1c3e8-efcd-4e1d-a0e7-d8c8bbcbc864",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from ipywidgets import interact\n",
    "\n",
    "# @interact\n",
    "# def greet(name=\"World\", count=5):\n",
    "#     for _ in range(count):\n",
    "#         print(f\"Hello, {name}!\")\n",
    "\n",
    "# greet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7d4f0f43-66ee-45b8-ad9a-8b2384a86fbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2e838b96-f41d-4749-9016-b55c9ffabc75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "from numba import cuda"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15492619-4f3b-4e4c-98c7-e71b324af219",
   "metadata": {},
   "source": [
    "## Do a manual pre-selection\n",
    "We start by loading a snapshot and selecting just a part of it.\n",
    "We limit ourselves because the GPU has limited memory and the GPU ray\n",
    "tracer class builds a binary tree spanning the entire snapshot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8efef817-b631-423e-902b-e9bc0f5ba69b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import paicos as pa\n",
    "import numpy as np\n",
    "\n",
    "pa.use_units(True)\n",
    "\n",
    "snapnum = 247\n",
    "try:\n",
    "    snap = pa.Snapshot(pa.data_dir + 'highres', snapnum)\n",
    "except FileNotFoundError as e:\n",
    "    print(e)\n",
    "    err_msg = ('This example is much more fun with a large data set.\\nPlease see: '\n",
    "               + 'https://github.com/tberlok/paicos/tree/main/data/highres/README.md'\n",
    "               + ' for download instructions.\\n'\n",
    "              + 'For now we simply load the low resolution data set.')\n",
    "    print(err_msg)\n",
    "    snap = pa.Snapshot(pa.data_dir, snapnum)\n",
    "center = snap.Cat.Group['GroupPos'][0]\n",
    "R200c = snap.Cat.Group['Group_R_Crit200'][0]\n",
    "widths = np.array([10000, 10000, 10000]) * R200c.uq\n",
    "\n",
    "# Create subset of snapshot\n",
    "index = pa.util.get_index_of_radial_range(snap['0_Coordinates'], center, 0, np.max(widths)*np.cbrt(3))\n",
    "snap = snap.select(index, parttype=0)\n",
    "\n",
    "# Pixel dimensions of image\n",
    "nx = ny = 1024"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64068cad-9ea2-4f0c-87d0-37415552e2af",
   "metadata": {},
   "source": [
    "## Initialize the GPU projector\n",
    "\n",
    "Here we use a Paicos orientation class to initialize the view such that the \n",
    "width of the image is along the $x$-coordinate of the simulation and the height\n",
    "of the image is along the $y$-coordinate. The depth of the image is in the $z$-direction.\n",
    "\n",
    "The orientation class has methods for rotating the view around $x$, $y$, and $z$\n",
    "or around the axes of its local coordinate system. When an orientation\n",
    "instance has been passed to an ImageCreator (such as the projector below),\n",
    "then calling these methods will result in a rotation around the center of\n",
    "the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ff0210bf-069a-438c-90b0-37f847bfa8dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attempting to get derived variable: 0_Volume...\t[DONE]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "orientation = pa.Orientation(normal_vector=[0, 0, 1], perp_vector1=[1, 0, 0])\n",
    "projector = pa.GpuRayProjector(snap, center, widths, orientation, npix=nx, threadsperblock=8, do_pre_selection=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7a60f39d-3e84-495c-97de-fcf5fd5937b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5b2c752-dd6c-4000-af9f-79349fe93ade",
   "metadata": {},
   "source": [
    "## Some code for sorting FoF and subfind catalogues\n",
    "\n",
    "This is just for plotting the 20 most massive FoF or subhalos that are inside the projection cube.\n",
    "Mainly for the widget so no need to read this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f2e4f6ac-123f-4209-8982-002244925b13",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_group_and_sub_indices():\n",
    "    info = {}\n",
    "    if hasattr(projector.snap.Cat, 'Sub'):\n",
    "        sub_in_region_bool = pa.util.get_index_of_rotated_cubic_region_plus_thin_layer(projector.snap.Cat.Sub['SubhaloPos'],\n",
    "                                                        projector.center, projector.widths,\n",
    "                                                        projector.snap.Cat.Sub['SubhaloHalfmassRad'],\n",
    "                                                        projector.snap.box, projector.orientation)\n",
    "\n",
    "        info['Subhalo_ids'] = np.arange(sub_in_region_bool.shape[0])[sub_in_region_bool]\n",
    "        info['SubhaloPos'] = projector.snap.Cat.Sub['SubhaloPos'][sub_in_region_bool]\n",
    "        info['SubhaloHalfmassRad'] = projector.snap.Cat.Sub['SubhaloHalfmassRad'][sub_in_region_bool]\n",
    "        info['SubhaloMass'] = projector.snap.Cat.Sub['SubhaloMass'][sub_in_region_bool]\n",
    "        # Sort according to mass\n",
    "        sort_index = np.argsort(info['SubhaloMass'])[::-1]\n",
    "        info['Subhalo_ids'] =  info['Subhalo_ids'][sort_index]\n",
    "        info['SubhaloPos'] =  info['SubhaloPos'][sort_index]\n",
    "        info['SubhaloHalfmassRad'] = info['SubhaloHalfmassRad'][sort_index]\n",
    "        info['SubhaloMass'] =  info['SubhaloMass'][sort_index]\n",
    "        \n",
    "    if hasattr(projector.snap.Cat, 'Group'):\n",
    "        group_in_region_bool = pa.util.get_index_of_rotated_cubic_region_plus_thin_layer(projector.snap.Cat.Group['GroupPos'],\n",
    "                                                projector.center, projector.widths,\n",
    "                                                projector.snap.Cat.Group['Group_R_Crit200'],\n",
    "                                                projector.snap.box, projector.orientation)\n",
    "        info['Group_ids'] = np.arange(group_in_region_bool.shape[0])[group_in_region_bool]\n",
    "        info['GroupPos'] = projector.snap.Cat.Group['GroupPos'][group_in_region_bool]\n",
    "        info['Group_R_Crit200'] = projector.snap.Cat.Group['Group_R_Crit200'][group_in_region_bool]\n",
    "        info['Group_M_Crit200'] = projector.snap.Cat.Group['Group_M_Crit200'][group_in_region_bool]\n",
    "        # Sort according to mass\n",
    "        sort_index = np.argsort(info['Group_M_Crit200'])[::-1]\n",
    "        info['Group_ids'] = info['Group_ids'][sort_index]\n",
    "        info['GroupPos']  = info['GroupPos'] [sort_index]\n",
    "        info['Group_R_Crit200'] = info['Group_R_Crit200'][sort_index]\n",
    "        info['Group_M_Crit200'] = info['Group_M_Crit200'][sort_index]\n",
    "    return info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b1b4708-751b-4435-a9b0-e940e934a083",
   "metadata": {},
   "source": [
    "## The interactive widget\n",
    "The rather long code below defines an interactive ipython widget with a\n",
    "number of hopefully mostly self-explanatory buttons.\n",
    "\n",
    "These can zoom in/out, rotate the image, change width, height and depth etc.\n",
    "\n",
    "Pressing the 'Recording' tick mark will output a png/hdf5 every time the\n",
    "wiev changes (if those boxes are ticked). These will be saved in the\n",
    "directory entered in the box just to the right of png tick box.\n",
    "The recording also saves a .log file with a series of commands\n",
    "that can be used to reproduce an interactive session.\n",
    "This allows for using an interactive session as a starting point for\n",
    "creating an animation of a simulation snapshot.\n",
    "\n",
    "We have left all the code for the widget visible instead of saving it somewhere\n",
    "else and importing it. We hope that it will in this way be easier for someone\n",
    "to extend/modify the code to their own needs.\n",
    "\n",
    "The data included with Paicos is very low resolution and this notebook does\n",
    "not really showcase how well the GPU code works at higher resolution.\n",
    "We have tried with a $12^3$ times better mass resolution simulation (equivalent\n",
    "to the resolution in the TNG300 simulation) and find that an A100 GPU is fast enough\n",
    "to give a smooth user experience. Download instructions for this data set can be found [here](https://github.com/tberlok/paicos/tree/main/data/highres/README.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7d7c3da1-aff2-4db5-87b7-ded73cc4b433",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b649423f37846d3b7b00cd3cbcbeba9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0172126eb1a64217ab511cb47e8ca06e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Button(description='Update', style=ButtonStyle()), Dropdown(description='Field:', index=1, opti…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d40af1ded5854651a3bdd991ed80816c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Button(description='Pan left', style=ButtonStyle()), Button(description='Pan right', style=Butt…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc6c1a8e0db94e5fac874ddd69adf8bf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatSlider(value=10000.0, continuous_update=False, description='width', max=20000.0, min=1.0, …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5370211f754f435987a45508f4205868",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Button(description='Move center', style=ButtonStyle()), BoundedFloatText(value=0.0, description…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ff6ff63b414a4a148d74c5eb6c58d9db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Checkbox(value=False, description='Fix climits', indent=False), FloatText(value=5.8146158647168…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "938e8337800d4b19a6861a24177416a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatSlider(value=15.0, description='Step (Degrees)', max=90.0, step=0.5), Checkbox(value=False…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "61337adb8f654d90933d668102723431",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Dropdown(description='Cmap:', index=3, options=('magma', 'inferno', 'plasma', 'viridis', 'civid…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eb772c6f32df494899b85e4f6f21a3f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Checkbox(value=False, description='Recording', indent=False), Checkbox(value=False, description…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c05ae14909540abb243ff6614ecba67",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Checkbox(value=False, description='Show FoF groups', indent=False), Checkbox(value=False, descr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "def update():\n",
    "    proj = projector.project_variable(var_str.value)\n",
    "    extent = projector.centered_extent\n",
    "\n",
    "    if to_physical.value:\n",
    "        proj = proj.to_physical\n",
    "        extent = extent.to_physical\n",
    "\n",
    "    if to_cgs.value:\n",
    "        proj = proj.cgs\n",
    "        extent = extent.cgs\n",
    "    \n",
    "    if to_astro_units.value:\n",
    "        proj = proj.astro\n",
    "        extent = extent.astro\n",
    "\n",
    "    fig = plt.figure(1)\n",
    "    plt.clf()\n",
    "\n",
    "    # Deal with color limits and do plot\n",
    "    if fix_climits.value:\n",
    "        vmin.disabled = False\n",
    "        vmax.disabled = False\n",
    "        if vmin.value > 0 and vmax.value > 0:\n",
    "            pass\n",
    "        else:\n",
    "            vmin.value = proj.value.min()\n",
    "            vmax.value = proj.value.max()\n",
    "        im = plt.imshow(proj.value, extent=extent.value, origin='lower',\n",
    "                        norm=LogNorm(vmin.value, vmax.value), cmap=cmap_str.value)\n",
    "    else:\n",
    "        vmin.value = proj.value.min()\n",
    "        vmax.value = proj.value.max()\n",
    "        vmin.disabled = True\n",
    "        vmax.disabled = True\n",
    "        im = plt.imshow(proj.value, extent=extent.value, origin='lower',\n",
    "                norm=LogNorm(), cmap=cmap_str.value)\n",
    "    # Labels\n",
    "    plt.xlabel(extent.label())\n",
    "    plt.ylabel(extent.label())\n",
    "    \n",
    "    # Colorbar\n",
    "    cb = plt.colorbar()\n",
    "    cb.set_label(proj.label('\\\\mathrm{' + var_str.value.replace('_', '\\_') + '}\\,'))\n",
    "\n",
    "    # Title\n",
    "    title_str = f'Snapnum: {snapnum}, Age: {snap.age:1.2f}, Redshift: {snap.z:1.2f}'\n",
    "    plt.title(title_str)\n",
    "\n",
    "    # Add subs/groups\n",
    "    # TODO: Get rid of mostly duplicate code for groups/subhalos\n",
    "    select_center.disabled = True\n",
    "    if hasattr(projector.snap, 'Cat'):\n",
    "        info = get_group_and_sub_indices()\n",
    "\n",
    "        if 'Group_ids' in info and show_groups.value:\n",
    "            select_center.disabled = False\n",
    "            orientation = projector.orientation\n",
    "            points = info['GroupPos'] - projector.center\n",
    "            points = np.matmul(orientation.inverse_rotation_matrix, points.T).T\n",
    "            if to_physical.value:\n",
    "                points = points.to_physical\n",
    "                info['Group_R_Crit200'] = info['Group_R_Crit200'].to_physical\n",
    "            if to_cgs.value:\n",
    "                points = points.cgs\n",
    "                info['Group_R_Crit200'] = info['Group_R_Crit200'].cgs\n",
    "            if to_astro_units.value:\n",
    "                points = points.astro\n",
    "                info['Group_R_Crit200'] = info['Group_R_Crit200'].astro\n",
    "                \n",
    "            ax = plt.gca()\n",
    "            options = []\n",
    "            for ii in range(points.shape[0]):\n",
    "                if ii >= 20 or info['Group_M_Crit200'][ii].value == 0:\n",
    "                    break\n",
    "                circ = plt.Circle((points[ii, 0].value, points[ii, 1].value),\n",
    "                           info['Group_R_Crit200'][ii].value, color='k', fill=False)\n",
    "                ax.add_patch(circ)\n",
    "                plt.text(points[ii, 0].value, points[ii, 1].value, f\"G{info['Group_ids'][ii]}\", fontsize=6)\n",
    "                options.append(f\"G{info['Group_ids'][ii]}\")\n",
    "\n",
    "            select_center.options = list(options)\n",
    "\n",
    "        if 'Subhalo_ids' in info and show_subs.value:\n",
    "        \n",
    "            select_center.disabled = False\n",
    "            orientation = projector.orientation\n",
    "            points = info['SubhaloPos'] - projector.center\n",
    "            points = np.matmul(orientation.inverse_rotation_matrix, points.T).T\n",
    "            if to_physical.value:\n",
    "                points = points.to_physical\n",
    "                info['SubhaloHalfmassRad'] = info['SubhaloHalfmassRad'].to_physical\n",
    "            if to_cgs.value:\n",
    "                points = points.cgs\n",
    "                info['SubhaloHalfmassRad'] = info['SubhaloHalfmassRad'].cgs\n",
    "            if to_astro_units.value:\n",
    "                points = points.astro\n",
    "                info['SubhaloHalfmassRad'] = info['SubhaloHalfmassRad'].astro\n",
    "                \n",
    "            ax = plt.gca()\n",
    "            options = []\n",
    "            for ii in range(points.shape[0]):\n",
    "                if ii >= 20 or info['SubhaloMass'][ii].value == 0:\n",
    "                    break\n",
    "                circ = plt.Circle((points[ii, 0].value, points[ii, 1].value),\n",
    "                           info['SubhaloHalfmassRad'][ii].value, color='k', fill=False)\n",
    "                ax.add_patch(circ)\n",
    "                plt.text(points[ii, 0].value, points[ii, 1].value, f\"S{info['Subhalo_ids'][ii]}\", fontsize=6)\n",
    "                options.append(f\"S{info['Subhalo_ids'][ii]}\")\n",
    "\n",
    "            select_center.options = list(options)\n",
    "\n",
    "    if recording.value:\n",
    "        if hdf5.value:\n",
    "            image_file = pa.ArepoImage(projector, basedir=outfolder.value,\n",
    "                               basename=f'{basename.value}_{var_str.value}_frame_{frame_counter.value}')\n",
    "\n",
    "            image_file.save_image(var_str.value, proj)\n",
    "\n",
    "            # Move from temporary filename to final filename\n",
    "            image_file.finalize()\n",
    "            if frame_counter.value == 0:\n",
    "                mylogger(image_file.filename)\n",
    "                mylogger(outfolder.value)\n",
    "                mylogger(basename.value)\n",
    "                mylogger(var_str.value)\n",
    "                org_center = projector.center - projector._diff_center\n",
    "                mylogger(f\"org_center,{org_center[0].value},{org_center[1].value},{org_center[2].value}\")\n",
    "\n",
    "        if png.value:\n",
    "            plt.savefig(f'{outfolder.value}/{basename.value}_{var_str.value}_frame_{frame_counter.value}_{projector.snap.snapnum}.png', dpi=700)\n",
    "        frame_counter.value += 1\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def mylogger(string, line_ending='\\n', mode='a'):\n",
    "    with open(f'{outfolder.value}/{basename.value}.log', mode) as f:\n",
    "        f.write(string + line_ending)\n",
    "\n",
    "width = widgets.FloatSlider(value=projector.width.value, min=1, max=2*widths[0].value, step=1, description=\"width\", continuous_update=False)\n",
    "height = widgets.FloatSlider(value=projector.height.value, min=1, max=2*widths[1].value, step=1, description=\"height\", continuous_update=False)\n",
    "depth = widgets.FloatSlider(value=projector.depth.value, min=1, max=2*widths[2].value, step=1, description=\"depth\", continuous_update=False)\n",
    "zoom_slider = widgets.FloatSlider(value=1, min=0.1, max=10, step=0.1, description=\"Zoom factor\")\n",
    "\n",
    "center_horizontal = widgets.BoundedFloatText(\n",
    "    value=0,\n",
    "    min=-widths[0].value,\n",
    "    max=widths[0].value,\n",
    "    step=1,\n",
    "    description='Horizontal',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "center_vertical = widgets.BoundedFloatText(\n",
    "    value=0,\n",
    "    min=-widths[0].value,\n",
    "    max=widths[0].value,\n",
    "    step=1,\n",
    "    description='Vertical',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "center_depth = widgets.BoundedFloatText(\n",
    "    value=0,\n",
    "    min=-widths[0].value,\n",
    "    max=widths[0].value,\n",
    "    step=1,\n",
    "    description='Depth',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "def click_step_button(b):\n",
    "    if center_horizontal.value != 0:\n",
    "        projector.move_center_along_perp_vector1(center_horizontal.value * projector.width.uq)\n",
    "            \n",
    "    if center_vertical.value != 0:\n",
    "        projector.move_center_along_perp_vector2(center_vertical.value * projector.width.uq)\n",
    "            \n",
    "    if center_depth.value != 0:\n",
    "        projector.move_center_along_normal_vector(center_depth.value * projector.width.uq)\n",
    "\n",
    "    if recording.value:\n",
    "        mylogger(f'move_center,{center_horizontal.value},{center_vertical.value},{center_depth.value}')\n",
    "    update()\n",
    "\n",
    "def press_reset_center_button(b):\n",
    "    diff = projector._diff_center\n",
    "    org_center = projector.center - diff\n",
    "    if recording.value:\n",
    "        mylogger(f\"move_center_sim_coordinates,{diff[0].value},{diff[1].value},{diff[2].value}\", line_ending=',#,')\n",
    "        mylogger(f\"reset_center_to_org_center,{org_center[0].value},{org_center[1].value},{org_center[2].value}\")\n",
    "\n",
    "    projector.center = org_center\n",
    "    projector._diff_center[:] = 0\n",
    "    update()\n",
    "\n",
    "step_button = widgets.Button(description=\"Move center\")\n",
    "recenter_button = widgets.Button(description=\"Reset center\")\n",
    "step_button.on_click(click_step_button)\n",
    "recenter_button.on_click(press_reset_center_button)\n",
    "\n",
    "\n",
    "## Create dropdown for parttype 0 only (TODO: Remove vectors and tensors from list, or add box for selecting component)\n",
    "avail_list = []\n",
    "for key in snap._auto_list:\n",
    "    if key[0] == '0':\n",
    "        avail_list.append(key)\n",
    "var_str =  widgets.Dropdown(options=avail_list,\n",
    "    value='0_Density',\n",
    "    description='Field:',\n",
    ")\n",
    "\n",
    "select_center =  widgets.Dropdown(options=[],\n",
    "    value=None,\n",
    "    description='Center on:',\n",
    "    disabled=True\n",
    ")\n",
    "\n",
    "# Cmap dropdown\n",
    "cmap_str =  widgets.Dropdown(options=plt.colormaps(),\n",
    "    value='viridis',\n",
    "    description='Cmap:',\n",
    ")\n",
    "\n",
    "# vmin and vmax\n",
    "vmin = widgets.FloatText(\n",
    "    description='vmin:',\n",
    "    disabled=True\n",
    ")\n",
    "\n",
    "vmax = widgets.FloatText(\n",
    "    description='vmax:',\n",
    "    disabled=True\n",
    ")\n",
    "\n",
    "fix_climits = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Fix climits',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "# \n",
    "button_left = widgets.Button(description=\"Pan left\")\n",
    "button_right = widgets.Button(description=\"Pan right\")\n",
    "button_up = widgets.Button(description=\"Pan up\")\n",
    "button_down = widgets.Button(description=\"Pan down\")\n",
    "button_clock_wise = widgets.Button(description=\"Clockwise\")\n",
    "button_anti_clock_wise = widgets.Button(description=\"Anti-clockwise\")\n",
    "step_size_in_degrees = widgets.FloatSlider(value=15, min=0, max=90, step=0.5, description=\"Step (Degrees)\")\n",
    "\n",
    "button_update = widgets.Button(description=\"Update\")\n",
    "\n",
    "def call_update(b):\n",
    "    update()\n",
    "\n",
    "button_update.on_click(call_update)\n",
    "\n",
    "def call_double_resolution(b):\n",
    "    projector.double_resolution\n",
    "    if recording.value:\n",
    "        mylogger('double_resolution')\n",
    "    update()\n",
    "\n",
    "def call_half_resolution(b):\n",
    "    projector.half_resolution\n",
    "    if recording.value:\n",
    "        mylogger('half_resolution')\n",
    "    update()\n",
    "\n",
    "button_double = widgets.Button(description=\"Double res\")\n",
    "button_double.on_click(call_double_resolution)\n",
    "\n",
    "button_half = widgets.Button(description=\"Half res\")\n",
    "button_half.on_click(call_half_resolution)\n",
    "\n",
    "def call_zoom(b):\n",
    "    zoom_button_was_pressed.value = True\n",
    "    projector.zoom(zoom_slider.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'zoom,{zoom_slider.value}')\n",
    "        \n",
    "    # Check that new widths are not completely unreasonable!\n",
    "    # Do check\n",
    "\n",
    "    width.value = projector.width.value\n",
    "    height.value = projector.height.value\n",
    "    update()\n",
    "    zoom_button_was_pressed.value = False\n",
    "\n",
    "button_zoom = widgets.Button(description=\"Zoom in/out\")\n",
    "button_zoom.on_click(call_zoom)\n",
    "zoom_button_was_pressed = widgets.Checkbox(\n",
    "    value=False, description='internal boolean for avoiding calling call_update twice when zooming',\n",
    ")\n",
    "\n",
    "to_physical = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='physical units',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "to_cgs = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='cgs units',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "to_astro_units = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description=\"'astro' units\",\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "to_physical.observe(call_update, names=['value'])\n",
    "to_cgs.observe(call_update, names=['value'])\n",
    "to_astro_units.observe(call_update, names=['value'])\n",
    "\n",
    "def change_width(change):\n",
    "    projector.width = width.value * projector.width.uq\n",
    "    if not zoom_button_was_pressed.value:\n",
    "        if recording.value:\n",
    "            mylogger(f'width,{width.value}')\n",
    "        update()\n",
    "\n",
    "def change_height(change):\n",
    "    projector.height = height.value * projector.height.uq\n",
    "    if not zoom_button_was_pressed.value:\n",
    "        if recording.value:\n",
    "            mylogger(f'height,{height.value}')\n",
    "        update()\n",
    "\n",
    "def change_depth(change):\n",
    "    projector.depth = depth.value * projector.depth.uq\n",
    "    if recording.value:\n",
    "        mylogger(f'depth,{depth.value}')\n",
    "    update()\n",
    "\n",
    "def change_var_str(change):\n",
    "    if recording.value:\n",
    "        mylogger(var_str.value)\n",
    "    update()\n",
    "\n",
    "def pan_left(b):\n",
    "    projector.orientation.rotate_around_perp_vector2(degrees=step_size_in_degrees.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'rotate_around_perp_vector2,{step_size_in_degrees.value}')\n",
    "    update()\n",
    "\n",
    "def pan_right(b):\n",
    "    projector.orientation.rotate_around_perp_vector2(degrees=-step_size_in_degrees.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'rotate_around_perp_vector2,{-step_size_in_degrees.value}')\n",
    "    update()\n",
    "\n",
    "def pan_up(b):\n",
    "    projector.orientation.rotate_around_perp_vector1(degrees=step_size_in_degrees.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'rotate_around_perp_vector1,{step_size_in_degrees.value}')\n",
    "    update()\n",
    "\n",
    "def pan_down(b):\n",
    "    projector.orientation.rotate_around_perp_vector1(degrees=-step_size_in_degrees.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'rotate_around_perp_vector1,{-step_size_in_degrees.value}')\n",
    "    update()\n",
    "\n",
    "def clock_wise(b):\n",
    "    projector.orientation.rotate_around_normal_vector(degrees=step_size_in_degrees.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'rotate_around_normal_vector,{step_size_in_degrees.value}')\n",
    "    update()\n",
    "\n",
    "def anti_clock_wise(b):\n",
    "    projector.orientation.rotate_around_normal_vector(degrees=-step_size_in_degrees.value)\n",
    "    if recording.value:\n",
    "        mylogger(f'rotate_around_normal_vector,{-step_size_in_degrees.value}')\n",
    "    update()\n",
    "\n",
    "button_left.on_click(pan_left)\n",
    "button_right.on_click(pan_right)\n",
    "\n",
    "button_up.on_click(pan_up)\n",
    "button_down.on_click(pan_down)\n",
    "\n",
    "button_clock_wise.on_click(clock_wise)\n",
    "button_anti_clock_wise.on_click(anti_clock_wise)\n",
    "\n",
    "var_str.observe(call_update)\n",
    "\n",
    "\n",
    "# Subhalo, groups\n",
    "show_groups = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Show FoF groups',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "show_subs = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Show subhalos',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "show_groups.observe(call_update, names=['value'])\n",
    "show_subs.observe(call_update, names=['value'])\n",
    "\n",
    "# Save hdf5/save\n",
    "\n",
    "frame_counter = widgets.IntSlider(value=0)\n",
    "recording = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Recording',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "def reset_counter(b):\n",
    "    \"\"\"\n",
    "    Reset counter if recording is stopped,\n",
    "    calculate first image if recording is started\n",
    "    \"\"\"\n",
    "    if not recording.value:\n",
    "        frame_counter.value = 0\n",
    "        basename.disabled = False\n",
    "        outfolder.disabled = False\n",
    "    else:\n",
    "        basename.disabled = True\n",
    "        outfolder.disabled = True\n",
    "        mylogger('', mode='w')\n",
    "        update()\n",
    "\n",
    "recording.observe(reset_counter, names=['value'])  \n",
    "\n",
    "hdf5 = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='hdf5',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "png = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='png',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "out = widgets.interactive_output(update, {},)\n",
    "\n",
    "def change_center_using_cat(change):\n",
    "    info = get_group_and_sub_indices()\n",
    "    if select_center.value[0] == 'G':\n",
    "        gr_id = int(select_center.value[1:])\n",
    "\n",
    "        new_center = projector.snap.Cat.Group['GroupPos'][gr_id].T\n",
    "        if recording.value:\n",
    "            diff = new_center - projector._center\n",
    "            mylogger(f\"move_center_sim_coordinates,{diff[0].value},{diff[1].value},{diff[2].value}\", line_ending=',#,')\n",
    "            mylogger(f\"center_on_group,{gr_id},{new_center[0].value},{new_center[1].value},{new_center[2].value}\")\n",
    "        projector._diff_center += new_center - projector._center\n",
    "        projector.center = new_center.copy\n",
    "        show_groups.value = False\n",
    "        select_center.options = []\n",
    "        select_center.value = None\n",
    "    if select_center.value[0] == 'S':\n",
    "        sub_id = int(select_center.value[1:])\n",
    "        new_center = projector.snap.Cat.Sub['SubhaloPos'][sub_id].T\n",
    "        if recording.value:\n",
    "            diff = new_center - projector._center\n",
    "            mylogger(f\"move_center_sim_coordinates,{diff[0].value},{diff[1].value},{diff[2].value}\", line_ending=',#,')\n",
    "            mylogger(f\"center_on_sub,{sub_id},{new_center[0].value},{new_center[1].value},{new_center[2].value}\")\n",
    "\n",
    "        projector._diff_center += new_center - projector._center\n",
    "        projector.center = new_center.copy\n",
    "        show_subs.value = False\n",
    "        select_center.options = []\n",
    "        select_center.value = None\n",
    "    # update()\n",
    "\n",
    "width.observe(change_width, names=['value'])\n",
    "height.observe(change_height, names=['value'])\n",
    "depth.observe(change_depth, names=['value'])\n",
    "var_str.observe(change_var_str, names=['value'])\n",
    "select_center.observe(change_center_using_cat, names=['value'])\n",
    "\n",
    "\n",
    "\n",
    "basename = widgets.Text(value='image')\n",
    "outfolder = widgets.Text(value='./')\n",
    "\n",
    "display(out,\n",
    "       widgets.HBox([button_update, var_str, button_zoom, zoom_slider]),\n",
    "       widgets.HBox([button_left, button_right, button_up, button_down, button_clock_wise, button_anti_clock_wise]),\n",
    "       widgets.HBox([width, height, depth]),\n",
    "       widgets.HBox([step_button, center_horizontal, center_vertical, center_depth, recenter_button]),\n",
    "       widgets.HBox([fix_climits, vmin, vmax]),\n",
    "       widgets.HBox([step_size_in_degrees, to_physical, to_cgs, to_astro_units]),\n",
    "       widgets.HBox([cmap_str, button_double, button_half]),\n",
    "       widgets.HBox([recording, hdf5, png, outfolder, basename]),\n",
    "       widgets.HBox([show_groups, show_subs, select_center]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "324b1b86-c514-42ee-b475-68b6eb25e61e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}