submit_spm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit scanning probe microscopy simulation"
   ]
  },
  {
   "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",
    "# AiiDAlab imports.\n",
    "import aiidalab_widgets_base as awb\n",
    "\n",
    "# General imports.\n",
    "import ipywidgets as ipw\n",
    "from aiida import orm, plugins\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# Custom imports.\n",
    "from surfaces_tools.utils import spm, wfn\n",
    "from surfaces_tools.widgets import computational_resources, inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cp2kStmWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.stm\")\n",
    "Cp2kHrstmWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.hrstm\")\n",
    "Cp2kAfmWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.afm\")\n",
    "Cp2kOrbitalsWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.orbitals\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Select structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Structure selector.\n",
    "input_details = inputs.InputDetails()\n",
    "\n",
    "structure_selector = awb.StructureManagerWidget(\n",
    "    importers=[\n",
    "        awb.StructureUploadWidget(title=\"Import from computer\"),\n",
    "        awb.StructureBrowserWidget(title=\"AiiDA database\"),\n",
    "    ],\n",
    "    editors=[],\n",
    "    storable=False,\n",
    "    node_class=\"StructureData\",\n",
    ")\n",
    "display(structure_selector)\n",
    "\n",
    "ipw.dlink((structure_selector, \"structure\"), (input_details, \"structure\"))\n",
    "\n",
    "\n",
    "# Code.\n",
    "resources = computational_resources.ProcessResourcesWidget()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DFT parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style = {\"description_width\": \"80px\"}\n",
    "layout = {\"width\": \"70%\"}\n",
    "\n",
    "\n",
    "uks_input = inputs.UksSectionWidget()\n",
    "\n",
    "diagonalisation_smearing = inputs.DiagonalisationSmearingWidget()\n",
    "\n",
    "\n",
    "protocol = ipw.Dropdown(\n",
    "    value=\"standard\",\n",
    "    options=[\n",
    "        (\"Standard\", \"standard\"),\n",
    "        (\"Low accuracy\", \"low_accuracy\"),\n",
    "        (\"Debug\", \"debug\"),\n",
    "    ],\n",
    "    description=\"Protocol:\",\n",
    "    style={\"description_width\": \"120px\"},\n",
    ")\n",
    "\n",
    "elpa_check = ipw.Checkbox(value=True, description=\"use ELPA\")\n",
    "display(uks_input, diagonalisation_smearing, protocol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SPM parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SPM type.\n",
    "drop_spm_type = ipw.Dropdown(\n",
    "    description=\"SPM type\", options=[\"STM\", \"AFM\", \"HRSTM\", \"ORBITALS\"], value=\"STM\"\n",
    ")\n",
    "spm_parameters_box = ipw.VBox([])\n",
    "display(drop_spm_type, spm_parameters_box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Codes dropdown.\n",
    "stm_code = awb.ComputationalResourcesWidget(\n",
    "    description=\"STM code:\", default_calc_job_plugin=\"nanotech_empa.stm\"\n",
    ")\n",
    "hrstm_code = awb.ComputationalResourcesWidget(\n",
    "    description=\"HRSTM code:\", default_calc_job_plugin=\"nanotech_empa.hrstm\"\n",
    ")\n",
    "ppafm_code = awb.ComputationalResourcesWidget(\n",
    "    description=\"PPAFM code:\", default_calc_job_plugin=\"nanotech_empa.afm\"\n",
    ")\n",
    "twopp_code = awb.ComputationalResourcesWidget(\n",
    "    description=\"2PP code:\", default_calc_job_plugin=\"nanotech_empa.afm\"\n",
    ")\n",
    "cp2k_code = awb.ComputationalResourcesWidget(\n",
    "    description=\"CP2K code:\", default_calc_job_plugin=\"cp2k\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style = {\"description_width\": \"initial\"}  #'140px'\n",
    "layout = {\"width\": \"50%\"}\n",
    "layout_small = {\"width\": \"25%\"}\n",
    "\n",
    "# STM\n",
    "elim_float_slider = ipw.FloatRangeSlider(\n",
    "    value=[-2.0, 2.0],\n",
    "    min=-4.0,\n",
    "    max=4.0,\n",
    "    step=0.1,\n",
    "    description=\"Emin, Emax (eV):\",\n",
    "    continuous_update=False,\n",
    "    orientation=\"horizontal\",\n",
    "    readout=True,\n",
    "    readout_format=\".1f\",\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "\n",
    "de_floattext = ipw.BoundedFloatText(\n",
    "    description=\"dE (eV)\",\n",
    "    min=0.01,\n",
    "    max=1.00,\n",
    "    step=0.01,\n",
    "    value=0.04,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "\n",
    "fwhms_text = ipw.Text(\n",
    "    description=\"FWHMs (eV)\", value=\"0.08\", style=style, layout=layout\n",
    ")\n",
    "\n",
    "extrap_plane_floattext = ipw.BoundedFloatText(\n",
    "    description=\"Extrap plane (ang)\",\n",
    "    min=1.0,\n",
    "    max=10.0,\n",
    "    step=0.1,\n",
    "    value=4.0,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "\n",
    "height_text = ipw.Text(\n",
    "    description=\"Const. H (ang)\", value=\"4.0 6.0\", style=style, layout=layout\n",
    ")\n",
    "\n",
    "const_current_text = ipw.Text(\n",
    "    description=\"Const. cur. (isoval)\", value=\"1e-7\", style=style, layout=layout\n",
    ")\n",
    "\n",
    "ptip_floattext = ipw.BoundedFloatText(\n",
    "    description=\"p tip ratio\",\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    value=0.0,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Orbitals.\n",
    "n_homo_inttext = ipw.IntText(\n",
    "    description=\"num HOMO\", min=0, max=100, value=10, style=style, layout=layout_small\n",
    ")\n",
    "n_lumo_inttext = ipw.IntText(\n",
    "    description=\"num LUMO\", min=0, max=100, value=10, style=style, layout=layout_small\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# AFM\n",
    "scanstep_floattxt = ipw.BoundedFloatText(\n",
    "    description=r\"Scan d\\(\\mathbf{x}\\) (Å)\",\n",
    "    min=0.05,\n",
    "    max=0.5,\n",
    "    step=0.05,\n",
    "    value=0.1,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "\n",
    "scanminz_floattxt = ipw.BoundedFloatText(\n",
    "    description=r\"Scan \\(z_{min}\\) (Å)\",\n",
    "    min=0.0,\n",
    "    max=5.0,\n",
    "    step=0.1,\n",
    "    value=3.5,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "\n",
    "scanmaxz_floattxt = ipw.BoundedFloatText(\n",
    "    description=r\"Scan \\(z_{max}\\) (Å)\",\n",
    "    min=5.0,\n",
    "    max=10.0,\n",
    "    step=0.1,\n",
    "    value=8.5,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "\n",
    "amp_floattxt = ipw.FloatText(\n",
    "    description=\"Amplitude (Å)\", step=0.1, value=1.4, style=style, layout=layout_small\n",
    ")\n",
    "\n",
    "f0_cantilever_floattxt = ipw.FloatText(\n",
    "    description=r\"Cantilever (\\(f_0\\)) \",\n",
    "    step=0.1,\n",
    "    value=22352.5,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "\n",
    "drop_2pp_resp = ipw.Dropdown(\n",
    "    description=\"2PP RESP model\",\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    "    options={  # ChargeCuUp, ChargeCuDown, Ccharge, Ocharge\n",
    "        \"pentacene\": [-0.0669933, -0.0627402, 0.212718, -0.11767],\n",
    "        \"ptcda\": [-0.05, -0.07, 0.23, -0.13],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HRSTM\n",
    "voltext_ipw = ipw.Label(value=\"Voltage Range:\", style=style, layout=layout_small)\n",
    "volstep_ipw = ipw.BoundedFloatText(\n",
    "    description=\"Step\",\n",
    "    value=0.1,\n",
    "    min=0.01,\n",
    "    max=0.5,\n",
    "    step=0.01,\n",
    "    style=style,\n",
    "    layout=layout_small,\n",
    ")\n",
    "volmin_ipw = ipw.FloatText(\n",
    "    description=\"Min\", value=-0.3, style=style, layout=layout_small\n",
    ")\n",
    "volmax_ipw = ipw.FloatText(\n",
    "    description=\"Max\", value=0.3, style=style, layout=layout_small\n",
    ")\n",
    "\n",
    "fwhm_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"FWHM for DOS of Sample (\\(e\\)V)\",\n",
    "    value=0.05,\n",
    "    min=0.01,\n",
    "    max=0.5,\n",
    "    step=0.01,\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "workfun_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"Workfunction of Sample (\\(e\\)V)\",\n",
    "    value=5.0,\n",
    "    min=1.0,\n",
    "    max=10.0,\n",
    "    step=0.1,\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "wfnstep_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"Meshwidth for Grid Orbitals d\\(\\mathbf{x}\\) (Å)\",\n",
    "    value=0.2,\n",
    "    min=0.05,\n",
    "    max=1.0,\n",
    "    step=0.05,\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "extrap_ipw = ipw.BoundedFloatText(\n",
    "    description=\"Extrapolation Plane (Å)\",\n",
    "    value=4.0,\n",
    "    min=1.0,\n",
    "    max=10.0,\n",
    "    step=0.1,\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "# Tip stuff.\n",
    "tiptype_ipw = ipw.ToggleButtons(\n",
    "    description=\"Tip Type\",\n",
    "    value=\"blunt\",\n",
    "    options=[\"parametrized\", \"blunt\"],\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "rotate_ipw = ipw.Checkbox(\n",
    "    description=\"Rotate Tip Coefficients\", value=True, style=style, layout=layout\n",
    ")\n",
    "orbstip_ipw = ipw.BoundedIntText(\n",
    "    description=\"Maximal Tip Orbital\",\n",
    "    value=1,\n",
    "    min=0,\n",
    "    max=1,\n",
    "    step=1,\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "fwhmtip_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"FWHM for DOS of Tip (\\(e\\)V)\",\n",
    "    value=0.00,\n",
    "    min=0.00,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    disabled=(tiptype_ipw.value == \"parametrized\"),\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "# Parametrized tip info.\n",
    "stip_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"\\(s\\)-Value\",\n",
    "    value=0.15,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    disabled=(not tiptype_ipw.value == \"parametrized\"),\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "pytip_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"\\(p_y\\)-Value\",\n",
    "    value=0.5,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    disabled=(not tiptype_ipw.value == \"parametrized\"),\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "pztip_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"\\(p_z\\)-Value\",\n",
    "    value=0.0,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    disabled=(not tiptype_ipw.value == \"parametrized\"),\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "pxtip_ipw = ipw.BoundedFloatText(\n",
    "    description=r\"\\(p_x\\)-Value\",\n",
    "    value=0.5,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    disabled=(not tiptype_ipw.value == \"parametrized\"),\n",
    "    style=style,\n",
    "    layout=layout,\n",
    ")\n",
    "para_list = [stip_ipw, pytip_ipw, pztip_ipw, pxtip_ipw, fwhmtip_ipw]\n",
    "\n",
    "\n",
    "def set_para_values(value):\n",
    "    if value[\"new\"] == \"parametrized\":\n",
    "        stip_ipw.disabled = False\n",
    "        pytip_ipw.disabled = False\n",
    "        pztip_ipw.disabled = False\n",
    "        pxtip_ipw.disabled = False\n",
    "        fwhmtip_ipw.disabled = True\n",
    "        # Always using p-orbitals for parametrized tip\n",
    "        orbstip_ipw.value = 1\n",
    "        orbstip_ipw.disabled = True\n",
    "    else:\n",
    "        stip_ipw.disabled = True\n",
    "        pytip_ipw.disabled = True\n",
    "        pztip_ipw.disabled = True\n",
    "        pxtip_ipw.disabled = True\n",
    "        fwhmtip_ipw.disabled = False\n",
    "        orbstip_ipw.disabled = False\n",
    "\n",
    "\n",
    "tiptype_ipw.observe(set_para_values, \"value\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_description = ipw.Text(\n",
    "    description=\"Workflow description:\",\n",
    "    placeholder=\"Provide the description here.\",\n",
    "    style={\"description_width\": \"initial\"},\n",
    "    layout={\"width\": \"70%\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_out = ipw.Output()\n",
    "\n",
    "\n",
    "def get_builder():\n",
    "    with submit_out:\n",
    "        clear_output()\n",
    "        if structure_selector.structure is None:\n",
    "            print(\"Please select a structure.\")\n",
    "            return\n",
    "        if cp2k_code.value is None:\n",
    "            print(\"Please select CP2K code firstß.\")\n",
    "            return\n",
    "\n",
    "        dft_params_dict = {\n",
    "            \"elpa_switch\": elpa_check.value,\n",
    "            \"sc_diag\": diagonalisation_smearing.enable_diagonalisation.value,\n",
    "            \"periodic\": \"XYZ\",\n",
    "        }\n",
    "\n",
    "        dft_params_dict.update(uks_input.return_dict()[\"dft_params\"])\n",
    "\n",
    "        if diagonalisation_smearing.smearing_enabled:\n",
    "            dft_params_dict[\"smear_t\"] = (\n",
    "                diagonalisation_smearing.smearing_temperature.value\n",
    "            )\n",
    "            dft_params_dict[\"force_multiplicity\"] = (\n",
    "                diagonalisation_smearing.force_multiplicity.value\n",
    "            )\n",
    "\n",
    "        dft_params = orm.Dict(dict=dft_params_dict)\n",
    "\n",
    "        struct = structure_selector.structure\n",
    "\n",
    "        if drop_spm_type.value == \"STM\":\n",
    "            builder = Cp2kStmWorkChain.get_builder()\n",
    "            builder.metadata.label = \"CP2K_STM\"\n",
    "            builder.metadata.description = workflow_description.value\n",
    "            builder.protocol = orm.Str(protocol.value)\n",
    "            builder.spm_code = orm.load_node(stm_code.value)\n",
    "            builder.spm_params = orm.Dict(\n",
    "                dict=spm.create_stm_parameterdata(\n",
    "                    extrap_plane_floattext.value,\n",
    "                    height_text.value.split(),\n",
    "                    struct.symbols,\n",
    "                    \"parent_calc_folder/\",\n",
    "                    elim_float_slider.value[0],\n",
    "                    elim_float_slider.value[1],\n",
    "                    de_floattext.value,\n",
    "                    const_current_text.value.split(),\n",
    "                    fwhms_text.value.split(),\n",
    "                    ptip_floattext.value,\n",
    "                )\n",
    "            )\n",
    "        elif drop_spm_type.value == \"AFM\":\n",
    "            builder = Cp2kAfmWorkChain.get_builder()\n",
    "            builder.metadata.label = \"CP2K_AFM\"\n",
    "            builder.metadata.description = workflow_description.value\n",
    "            builder.protocol = orm.Str(protocol.value)\n",
    "            builder.afm_pp_params = orm.Dict(\n",
    "                dict=spm.create_pp_parameterdata(\n",
    "                    struct,\n",
    "                    scanstep_floattxt.value,\n",
    "                    scanminz_floattxt.value,\n",
    "                    scanmaxz_floattxt.value,\n",
    "                    amp_floattxt.value,\n",
    "                    f0_cantilever_floattxt.value,\n",
    "                )\n",
    "            )\n",
    "            builder.afm_2pp_params = orm.Dict(\n",
    "                dict=spm.create_2pp_parameterdata(\n",
    "                    struct,\n",
    "                    scanstep_floattxt.value,\n",
    "                    drop_2pp_resp.value,\n",
    "                    scanminz_floattxt.value,\n",
    "                    scanmaxz_floattxt.value,\n",
    "                    amp_floattxt.value,\n",
    "                    f0_cantilever_floattxt.value,\n",
    "                )\n",
    "            )\n",
    "            builder.afm_pp_code = orm.load_node(ppafm_code.value)\n",
    "            builder.afm_2pp_code = orm.load_node(twopp_code.value)\n",
    "\n",
    "        elif drop_spm_type.value == \"ORBITALS\":\n",
    "            builder = Cp2kOrbitalsWorkChain.get_builder()\n",
    "            builder.metadata.label = \"CP2K_ORBITALS\"\n",
    "            builder.metadata.description = workflow_description.value\n",
    "            builder.protocol = orm.Str(protocol.value)\n",
    "            builder.spm_code = orm.load_node(stm_code.value)\n",
    "            builder.spm_params = orm.Dict(\n",
    "                dict=spm.create_orbitals_parameterdata(\n",
    "                    extrap_plane_floattext.value,\n",
    "                    height_text.value.split(),\n",
    "                    \"parent_calc_folder/\",\n",
    "                    n_homo_inttext.value,\n",
    "                    n_lumo_inttext.value,\n",
    "                    const_current_text.value.split(),\n",
    "                    fwhms_text.value.split(),\n",
    "                    ptip_floattext.value,\n",
    "                )\n",
    "            )\n",
    "        elif drop_spm_type.value == \"HRSTM\":\n",
    "            builder = Cp2kHrstmWorkChain.get_builder()\n",
    "            builder.metadata.label = \"CP2K_HRSTM\"\n",
    "            builder.metadata.description = workflow_description.value\n",
    "            builder.protocol = orm.Str(protocol.value)\n",
    "            builder.hrstm_code = orm.load_node(hrstm_code.value)\n",
    "            builder.ppm_code = orm.load_node(twopp_code.value)\n",
    "            ppm_params = spm.create_2pp_parameterdata(\n",
    "                struct,\n",
    "                scanstep_floattxt.value,\n",
    "                drop_2pp_resp.value,\n",
    "                scanminz_floattxt.value,\n",
    "                scanmaxz_floattxt.value,\n",
    "                amp_floattxt.value,\n",
    "                f0_cantilever_floattxt.value,\n",
    "            )\n",
    "            ppm_params[\"PBC\"] = \"True\"\n",
    "            builder.ppm_params = orm.Dict(dict=ppm_params)\n",
    "            builder.hrstm_params = orm.Dict(\n",
    "                dict=spm.create_hrstm_parameterdata(\n",
    "                    builder.hrstm_code,\n",
    "                    \"parent_calc_folder/\",\n",
    "                    \"ppm_calc_folder/\",\n",
    "                    struct,\n",
    "                    ppm_params,\n",
    "                    tiptype_ipw.value,\n",
    "                    stip_ipw.value,\n",
    "                    pytip_ipw.value,\n",
    "                    pztip_ipw.value,\n",
    "                    pxtip_ipw.value,\n",
    "                    volmin_ipw.value,\n",
    "                    volmax_ipw.value,\n",
    "                    volstep_ipw.value,\n",
    "                    volstep_ipw.min,\n",
    "                    fwhm_ipw.value,\n",
    "                    wfnstep_ipw.value,\n",
    "                    extrap_ipw.value,\n",
    "                    workfun_ipw.value,\n",
    "                    orbstip_ipw.value,\n",
    "                    fwhmtip_ipw.value,\n",
    "                    rotate_ipw.value,\n",
    "                )\n",
    "            )\n",
    "        builder.cp2k_code = orm.load_node(cp2k_code.value)\n",
    "        builder.structure = structure_selector.structure_node\n",
    "        builder.dft_params = dft_params\n",
    "        # Check if a restart wfn is available.\n",
    "        wave_function = None\n",
    "        #         if structure_selector.structure_node.is_stored:\n",
    "        #             wave_function = wfn.structure_available_wfn(\n",
    "        #             node=structure_selector.structure_node,\n",
    "        #             relative_replica_id=None,\n",
    "        #             current_hostname=orm.load_node(cp2k_code.value).computer.hostname,\n",
    "        #             return_path=False,\n",
    "        #             dft_params=dft_params_dict,\n",
    "        #             )\n",
    "        if wave_function is not None:\n",
    "            print(f\"Restarting from wfn in folder: {wave_function.pk}\")\n",
    "            builder.parent_calc_folder = wave_function\n",
    "\n",
    "        # Resources.\n",
    "        builder.options = {\n",
    "            \"max_wallclock_seconds\": resources.walltime_seconds,\n",
    "            \"resources\": {\n",
    "                \"num_machines\": resources.nodes,\n",
    "                \"num_mpiprocs_per_machine\": resources.tasks_per_node,\n",
    "                \"num_cores_per_mpiproc\": resources.threads_per_task,\n",
    "            },\n",
    "        }\n",
    "        return builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code selector\n",
    "\n",
    "\n",
    "widgets = {\n",
    "    \"STM\": [\n",
    "        elim_float_slider,\n",
    "        de_floattext,\n",
    "        fwhms_text,\n",
    "        extrap_plane_floattext,\n",
    "        height_text,\n",
    "        const_current_text,\n",
    "        ptip_floattext,\n",
    "    ],\n",
    "    \"AFM\": [\n",
    "        scanstep_floattxt,\n",
    "        scanminz_floattxt,\n",
    "        scanmaxz_floattxt,\n",
    "        amp_floattxt,\n",
    "        f0_cantilever_floattxt,\n",
    "        drop_2pp_resp,\n",
    "    ],\n",
    "    \"ORBITALS\": [\n",
    "        n_homo_inttext,\n",
    "        n_lumo_inttext,\n",
    "        height_text,\n",
    "        const_current_text,\n",
    "        fwhms_text,\n",
    "        extrap_plane_floattext,\n",
    "        ptip_floattext,\n",
    "    ],\n",
    "    \"HRSTM\": [\n",
    "        scanstep_floattxt,\n",
    "        scanminz_floattxt,\n",
    "        scanmaxz_floattxt,\n",
    "        amp_floattxt,\n",
    "        f0_cantilever_floattxt,\n",
    "        drop_2pp_resp,\n",
    "        ipw.HBox(\n",
    "            [voltext_ipw, volmin_ipw, volmax_ipw, volstep_ipw],\n",
    "            style=style,\n",
    "            layout=layout,\n",
    "        ),\n",
    "        fwhm_ipw,\n",
    "        workfun_ipw,\n",
    "        wfnstep_ipw,\n",
    "        extrap_ipw,\n",
    "        tiptype_ipw,\n",
    "        rotate_ipw,\n",
    "        orbstip_ipw,\n",
    "        fwhmtip_ipw,\n",
    "        stip_ipw,\n",
    "        pytip_ipw,\n",
    "        pztip_ipw,\n",
    "        pxtip_ipw,\n",
    "    ],\n",
    "}\n",
    "\n",
    "spm_code_box = ipw.VBox([])\n",
    "\n",
    "\n",
    "def on_spm_change(c):\n",
    "    spm_parameters_box.children = widgets[drop_spm_type.value]\n",
    "    if drop_spm_type.value == \"STM\":\n",
    "        spm_code_box.children = [stm_code]\n",
    "        submit_button.children = [btn_submit_stm]\n",
    "        fwhms_text.value = \"0.08\"\n",
    "        height_text.description = \"Const. H (ang)\"\n",
    "        height_text.value = \"4.0 6.0\"\n",
    "        const_current_text.description = \"Const. cur. (isoval)\"\n",
    "    elif drop_spm_type.value == \"ORBITALS\":\n",
    "        spm_code_box.children = [stm_code]\n",
    "        submit_button.children = [btn_submit_orb]\n",
    "        fwhms_text.value = \"0.04\"\n",
    "        height_text.description = \"Heights (ang)\"\n",
    "        height_text.value = \"3.0 5.0\"\n",
    "        const_current_text.description = \"Isovalues\"\n",
    "    elif drop_spm_type.value == \"AFM\":\n",
    "        spm_code_box.children = [ppafm_code, twopp_code]\n",
    "        submit_button.children = [btn_submit_afm]\n",
    "        scanminz_floattxt.value = 3.5\n",
    "        scanminz_floattxt.min = 0.0\n",
    "        scanminz_floattxt.max = 5.0\n",
    "        scanminz_floattxt.step = 0.1\n",
    "        scanmaxz_floattxt.min = 5.0\n",
    "        scanmaxz_floattxt.max = 10.0\n",
    "        scanmaxz_floattxt.step = 0.1\n",
    "        scanmaxz_floattxt.value = 8.5\n",
    "        submit_button.children = [btn_submit_afm]\n",
    "    elif drop_spm_type.value == \"HRSTM\":\n",
    "        spm_code_box.children = [hrstm_code, twopp_code]\n",
    "        submit_button.children = [btn_submit_hrstm]\n",
    "        scanminz_floattxt.value = 4.5\n",
    "        scanminz_floattxt.min = 3.0\n",
    "        scanminz_floattxt.max = 10.0\n",
    "        scanminz_floattxt.step = 0.1\n",
    "        scanmaxz_floattxt.min = 3.0\n",
    "        scanmaxz_floattxt.max = 10.0\n",
    "        scanmaxz_floattxt.step = 0.1\n",
    "        scanmaxz_floattxt.value = 7.5\n",
    "        submit_button.children = [btn_submit_hrstm]\n",
    "\n",
    "\n",
    "btn_submit_stm = awb.SubmitButtonWidget(Cp2kStmWorkChain, inputs_generator=get_builder)\n",
    "btn_submit_afm = awb.SubmitButtonWidget(Cp2kAfmWorkChain, inputs_generator=get_builder)\n",
    "btn_submit_orb = awb.SubmitButtonWidget(\n",
    "    Cp2kOrbitalsWorkChain, inputs_generator=get_builder\n",
    ")\n",
    "btn_submit_hrstm = awb.SubmitButtonWidget(\n",
    "    Cp2kHrstmWorkChain, inputs_generator=get_builder\n",
    ")\n",
    "submit_button = ipw.HBox([])\n",
    "\n",
    "drop_spm_type.observe(on_spm_change)\n",
    "\n",
    "on_spm_change(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resources estimation.\n",
    "resources_estimation = computational_resources.ResourcesEstimatorWidget()\n",
    "resources_estimation.link_to_resources_widget(resources)\n",
    "ipw.dlink((input_details, \"details\"), (resources_estimation, \"details\"))\n",
    "ipw.dlink((input_details, \"details\"), (uks_input, \"details\"))\n",
    "ipw.dlink((uks_input, \"uks\"), (resources_estimation, \"uks\"))\n",
    "_ = ipw.dlink((cp2k_code, \"value\"), (resources_estimation, \"selected_code\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(ipw.HBox([ipw.VBox([cp2k_code, elpa_check]), spm_code_box]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(resources, resources_estimation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(workflow_description, submit_button, submit_out)"
   ]
  }
 ],
 "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
}