view_stm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General imports\n",
    "import urllib.parse\n",
    "\n",
    "import ipywidgets as ipw\n",
    "import numpy as np\n",
    "from IPython.display import clear_output, display\n",
    "\n",
    "# AiiDA imports.\n",
    "%load_ext aiida\n",
    "%aiida\n",
    "from aiida import common\n",
    "\n",
    "from surfaces_tools.helpers import HART_2_EV\n",
    "from surfaces_tools.utils import spm\n",
    "from surfaces_tools.widgets import series_plotter\n",
    "\n",
    "# Local imports."
   ]
  },
  {
   "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": [
    "colormaps = [\"gist_heat\", \"seismic\"]\n",
    "\n",
    "e_arr = None\n",
    "\n",
    "\n",
    "def load_pk(b):\n",
    "    global e_arr\n",
    "\n",
    "    new_version = False\n",
    "    workcalc = load_node(pk=pk_select.value)\n",
    "    stm_calc = spm.get_calc_by_label(workcalc, \"stm\")\n",
    "    try:\n",
    "        cp2k_calc = spm.get_calc_by_label(workcalc, \"scf_diag\")\n",
    "    except AssertionError:\n",
    "        try:\n",
    "            dft_out_params = workcalc.outputs.dft_output_parameters.get_dict()\n",
    "            new_version = True\n",
    "        except Exception as exc:\n",
    "            print(\"Incorrect pk.\")\n",
    "            print(exc)\n",
    "            return\n",
    "\n",
    "    geom_info.value = spm.get_slab_calc_info(workcalc.inputs.structure)\n",
    "\n",
    "    # Information about the calculation.\n",
    "    with misc_info:\n",
    "        clear_output()\n",
    "\n",
    "    dft_inp_params = dict(workcalc.inputs[\"dft_params\"])\n",
    "    if not new_version:\n",
    "        dft_out_params = dict(cp2k_calc.outputs.output_parameters)\n",
    "\n",
    "    with misc_info:\n",
    "        if dft_inp_params[\"uks\"]:\n",
    "            print(f\"UKS multiplicity {dft_inp_params['multiplicity']}\")\n",
    "        else:\n",
    "            print(\"RKS\")\n",
    "\n",
    "        print(\n",
    "            f\"Energy [au]: {dft_out_params['energy']:.6f}, [eV]: {dft_out_params['energy'] * HART_2_EV:.6f}\"\n",
    "        )\n",
    "\n",
    "        try:\n",
    "            spm_params = workcalc.inputs.stm_params\n",
    "        except common.NotExistentAttributeError:\n",
    "            spm_params = workcalc.inputs.spm_params\n",
    "\n",
    "        extrap_plane = float(spm_params[\"--eval_region\"][-1][1:])\n",
    "        print(f\"Extrap. plane [ang]: {extrap_plane:.1f}\")\n",
    "\n",
    "        if \"--p_tip_ratios\" in dict(spm_params):\n",
    "            p_tip_ratio = spm_params[\"--p_tip_ratios\"]\n",
    "\n",
    "    ### Load data.\n",
    "    with stm_calc.outputs.retrieved.open(\"stm.npz\", mode=\"rb\") as handle:\n",
    "        loaded_data = np.load(handle.name, allow_pickle=True)\n",
    "    stm_general_info = loaded_data[\"stm_general_info\"][()]\n",
    "    stm_series_info = loaded_data[\"stm_series_info\"]\n",
    "    stm_series_data = loaded_data[\"stm_series_data\"]\n",
    "\n",
    "    e_arr = stm_general_info[\"energies\"]\n",
    "\n",
    "    series_plotter_inst.add_series_collection(\n",
    "        stm_general_info, stm_series_info, stm_series_data\n",
    "    )\n",
    "\n",
    "    series_plotter_inst.setup_added_collections(workcalc.pk)\n",
    "\n",
    "    setup_selection_elements()\n",
    "\n",
    "\n",
    "style = {\"description_width\": \"50px\"}\n",
    "layout = {\"width\": \"70%\"}\n",
    "\n",
    "pk_select = ipw.IntText(value=0, description=\"pk\", style=style, layout=layout)\n",
    "\n",
    "load_pk_btn = ipw.Button(description=\"Load pk\", style=style, layout=layout)\n",
    "load_pk_btn.on_click(load_pk)\n",
    "\n",
    "geom_info = ipw.HTML()\n",
    "\n",
    "display(ipw.HBox([ipw.VBox([pk_select, load_pk_btn]), geom_info]))\n",
    "\n",
    "misc_info = ipw.Output()\n",
    "display(misc_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scanning tunneling microscopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def selected_orbital_indexes():\n",
    "    if tab.selected_index == 0:\n",
    "        # Continuous selection.\n",
    "        min_e, max_e = energy_range_slider.value\n",
    "        ie_1 = np.abs(e_arr - min_e).argmin()\n",
    "        ie_2 = np.abs(e_arr - max_e).argmin() + 1\n",
    "        indexes = np.arange(ie_1, ie_2)\n",
    "\n",
    "    else:\n",
    "        # Discrete selection.\n",
    "        voltages = np.array(voltages_text.value.split(), dtype=float)\n",
    "        filtered_voltages = []\n",
    "        for v in voltages:\n",
    "            if v >= np.min(e_arr) and v <= np.max(e_arr):\n",
    "                filtered_voltages.append(v)\n",
    "            else:\n",
    "                print(f\"Voltage {v:.2f} out of range, skipping.\")\n",
    "\n",
    "        indexes = []\n",
    "        for i_bias, bias in enumerate(filtered_voltages):\n",
    "            indexes.append(np.abs(e_arr - bias).argmin())\n",
    "\n",
    "    return indexes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style = {\"description_width\": \"120px\"}\n",
    "layout = {\"width\": \"40%\"}\n",
    "\n",
    "series_plotter_inst = series_plotter.SeriesPlotter(\n",
    "    select_indexes_function=selected_orbital_indexes, zip_prepend=\"stm\"\n",
    ")\n",
    "\n",
    "# Select energies to plot.\n",
    "\n",
    "energy_range_slider = ipw.FloatRangeSlider(\n",
    "    value=[0.0, 0.0],\n",
    "    min=0.0,\n",
    "    max=0.0,\n",
    "    step=0.1,\n",
    "    description=\"energy range (eV)\",\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation=\"horizontal\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\",\n",
    "    style=style,\n",
    "    layout={\"width\": \"90%\"},\n",
    ")\n",
    "\n",
    "voltages_text = ipw.Text(\n",
    "    description=\"energies (eV)\", value=\"\", style=style, layout={\"width\": \"90%\"}\n",
    ")\n",
    "\n",
    "tab = ipw.Tab(layout={\"width\": \"60%\"})\n",
    "\n",
    "tab.children = [energy_range_slider, voltages_text]\n",
    "tab.set_title(0, \"Continuous selection\")\n",
    "tab.set_title(1, \"Discrete selection\")\n",
    "\n",
    "display(\n",
    "    series_plotter_inst.selector_widget,\n",
    "    tab,\n",
    "    series_plotter_inst.plot_btn,\n",
    "    series_plotter_inst.clear_btn,\n",
    "    series_plotter_inst.plot_output,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_selection_elements():\n",
    "    default_voltages = [-2.0, -1.5, -1.0, -0.5, -0.1, 0.1, 0.5, 1.0, 1.5, 2.0]\n",
    "\n",
    "    # Filter based on energy limits.\n",
    "    default_voltages = [\n",
    "        v for v in default_voltages if v >= np.min(e_arr) and v <= np.max(e_arr)\n",
    "    ]\n",
    "    voltages_text.value = \" \".join([str(v) for v in default_voltages])\n",
    "\n",
    "    energy_range_slider.min = np.min(e_arr)\n",
    "    energy_range_slider.max = np.max(e_arr)\n",
    "    energy_range_slider.step = e_arr[1] - e_arr[0]\n",
    "    energy_range_slider.value = (np.min(e_arr), np.max(e_arr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export\n",
    "Export the currently selected series into a zip file. The raw data in plain txt and IGOR formats are included."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(\n",
    "    ipw.HBox([series_plotter_inst.zip_btn, series_plotter_inst.zip_progress]),\n",
    "    series_plotter_inst.link_out,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clear_tmp(b):\n",
    "    ! rm -rf tmp && mkdir tmp\n",
    "    with series_plotter_inst.link_out:\n",
    "        clear_output()\n",
    "    series_plotter_inst.zip_progress.value = 0.0\n",
    "\n",
    "    if series_plotter_inst.series is not None:\n",
    "        series_plotter_inst.zip_btn.disabled = False\n",
    "\n",
    "\n",
    "clear_tmp_btn = ipw.Button(description=\"clear tmp\")\n",
    "clear_tmp_btn.on_click(clear_tmp)\n",
    "display(clear_tmp_btn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the URL after everything is set up.\n",
    "try:\n",
    "    url = urllib.parse.urlsplit(jupyter_notebook_url)\n",
    "    pk_select.value = urllib.parse.parse_qs(url.query)[\"pk\"][0]\n",
    "    load_pk(0)\n",
    "except Exception as exc:\n",
    "    print(exc)"
   ]
  }
 ],
 "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}