view_replica.ipynb

{
 "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": [
    "%load_ext aiida\n",
    "%aiida\n",
    "import urllib.parse\n",
    "\n",
    "import ipywidgets as ipw\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output, display\n",
    "from scipy.constants import physical_constants\n",
    "\n",
    "from surfaces_tools import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2ev = physical_constants[\"Hartree energy in eV\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Replica chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_checked_pk_list(replica_calc, check_list):\n",
    "    list_of_pks = [\n",
    "        replica_calc.outputs.structures[rep].pk\n",
    "        for rep in replica_calc.outputs.structures\n",
    "    ]\n",
    "    with output_checked_pks:\n",
    "        clear_output()\n",
    "        print(\"List of all replica PKs:\")\n",
    "        rep_pk_str = \"[\"\n",
    "        for pk, cb in zip(list_of_pks, check_list):\n",
    "            if cb.value:\n",
    "                rep_pk_str += \"%d \" % pk\n",
    "        print(rep_pk_str[:-1] + \"]\")\n",
    "\n",
    "\n",
    "def generate_energy_cv_plot(plot_energy_scf, plot_energy_frc, plot_colvar):\n",
    "    ref_en_scf = plot_energy_scf[0]\n",
    "    ref_en_frc = plot_energy_frc[0]\n",
    "\n",
    "    plot_energy_scf = np.array(plot_energy_scf) * h2ev\n",
    "    plot_energy_scf -= ref_en_scf * h2ev\n",
    "\n",
    "    plot_energy_frc = np.array(plot_energy_frc) * h2ev\n",
    "    plot_energy_frc -= ref_en_frc * h2ev\n",
    "\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    plt.ylabel(\"Energy/eV\")\n",
    "    plt.xlabel(\"Distance from replica 0 (Å)\")\n",
    "    plt.plot(plot_colvar, plot_energy_scf, \"go-\", label=\"SCF energy\")\n",
    "    plt.plot(plot_colvar, plot_energy_frc, \"bo-\", label=\"FORCE_EVAL energy\")\n",
    "\n",
    "    #    for i, rep in enumerate(replica_calc[\"replicas\"]):\n",
    "    #        # if the replica has no cv_target and has energy, it's an initial one, paint as red\n",
    "    #        cv_target, energy, struct_pk = rep\n",
    "    #        if cv_target == None and energy is not None:\n",
    "    #            plt.plot(plot_colvar[i], plot_energy_scf[i], \"ro-\")\n",
    "\n",
    "    plt.legend()\n",
    "    plt.grid()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def generate_thumbnail_html(wc):\n",
    "    html_list = []\n",
    "    check_list = []  # check boxes to show in final pk list\n",
    "\n",
    "    plot_energy_scf = []\n",
    "    plot_energy_frc = []\n",
    "    plot_colvar = []\n",
    "\n",
    "    dist = 0\n",
    "    for rep in wc.outputs.structures:\n",
    "        struct_node = wc.outputs.structures[rep]\n",
    "        cvs_target = wc.outputs.details[rep][\"cvs_target\"]\n",
    "        cvs_actual = wc.outputs.details[rep][\"cvs_actual\"]\n",
    "        energy = [\n",
    "            wc.outputs.details[rep][\"output_parameters\"][\"energy_scf\"],\n",
    "            wc.outputs.details[rep][\"output_parameters\"][\"energy\"],\n",
    "        ]\n",
    "        plot_energy_scf.append(energy[0])\n",
    "        plot_energy_frc.append(energy[1])\n",
    "\n",
    "        html = '<table style=\"border-spacing:20px 0px;border-collapse:separate;\">'\n",
    "\n",
    "        thumbnail = utils.ase_to_thumbnail(struct_node.get_ase())\n",
    "\n",
    "        cv_target_str = \" \".join([\"%.2f\" % val for val in cvs_target])\n",
    "\n",
    "        energy_scf_str = \"%.6f\" % energy[0]\n",
    "        energy_frc_str = \"%.6f\" % energy[1]\n",
    "\n",
    "        colvar_actual_str = \" \".join([\"%.2f\" % val for val in cvs_actual])\n",
    "        d2prev = wc.outputs.details[rep][\"d2prev\"]\n",
    "        dist += d2prev\n",
    "        plot_colvar.append(dist)\n",
    "        d2prev_str = \"%.4f\" % float(d2prev)\n",
    "\n",
    "        check_me = ipw.Checkbox(\n",
    "            value=True, description=\"select\", disabled=False, layout=layout\n",
    "        )\n",
    "        check_me.observe(\n",
    "            lambda x, rc=wc, cl=check_list: print_checked_pk_list(rc, cl),\n",
    "            \"value\",\n",
    "        )\n",
    "        check_list.append(check_me)\n",
    "\n",
    "        html = '<img width=\"400px\" src=\"data:image/png;base64,{}\" title=\"\">'.format(\n",
    "            thumbnail\n",
    "        )\n",
    "\n",
    "        html += '<table style=\"border-spacing:6px 0px;border-collapse:separate;\">'\n",
    "\n",
    "        # Output some information about the replica...\n",
    "        html += '<tr><td align=\"left\"><b>Target:</b></td><td align=\"right\">{}</td></tr>'.format(\n",
    "            cv_target_str\n",
    "        )\n",
    "        html += '<tr><td align=\"left\"><b>CV actual:</b></td><td align=\"right\">{}</td></tr>'.format(\n",
    "            colvar_actual_str\n",
    "        )\n",
    "        html += '<tr><td align=\"left\"><b>Energy (scf) /au:</b></td><td align=\"right\">{}</td></tr>'.format(\n",
    "            energy_scf_str\n",
    "        )\n",
    "        html += '<tr><td align=\"left\"><b>Energy (force) /au:</b></td><td align=\"right\">{}</td></tr>'.format(\n",
    "            energy_frc_str\n",
    "        )\n",
    "        html += '<tr><td align=\"left\"><b>dist to prev:</b></td><td align=\"right\">{}</td></tr>'.format(\n",
    "            d2prev_str\n",
    "        )\n",
    "        html += '<tr><td align=\"left\">pk: </td><td align=\"right\">{}</td></tr>'.format(\n",
    "            struct_node.pk\n",
    "        )\n",
    "\n",
    "        # ... and the download link.\n",
    "        html += '<tr><td align=\"left\"><a target=\"_blank\" href=\"export_structure.ipynb?uuid={}\">View & export</a></td></tr>'.format(\n",
    "            struct_node.uuid\n",
    "        )\n",
    "\n",
    "        html += \"</table>\"\n",
    "\n",
    "        html_list.append(html)\n",
    "\n",
    "    return html_list, check_list, plot_energy_scf, plot_energy_frc, plot_colvar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_and_show_replica(_=None):\n",
    "    wc = load_node(node.value)\n",
    "    (\n",
    "        html_list,\n",
    "        check_list,\n",
    "        plot_energy_scf,\n",
    "        plot_energy_frc,\n",
    "        plot_colvar,\n",
    "    ) = generate_thumbnail_html(wc)\n",
    "\n",
    "    with output_plot:\n",
    "        clear_output()\n",
    "        generate_energy_cv_plot(plot_energy_scf, plot_energy_frc, plot_colvar)\n",
    "\n",
    "    n_col = 4\n",
    "    layout = {\"width\": \"%.1f%%\" % (100.0 / n_col)}\n",
    "    with output_thumbnails:\n",
    "        clear_output()\n",
    "        cur_row = []\n",
    "        for i, (html, cb) in enumerate(zip(html_list, check_list)):\n",
    "            cur_row.append(ipw.VBox([ipw.HTML(html), cb], layout=layout))\n",
    "            if (i + 1) % n_col == 0:\n",
    "                display(ipw.HBox(cur_row))\n",
    "                cur_row = []\n",
    "        if len(cur_row) != 0:\n",
    "            display(ipw.HBox(cur_row))\n",
    "\n",
    "    print_checked_pk_list(wc, check_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style = {\"description_width\": \"120px\"}\n",
    "layout = {\"width\": \"70%\"}\n",
    "\n",
    "node = ipw.IntText(description=\"Load node: \", layout=layout, style=style)\n",
    "\n",
    "btn_show = ipw.Button(description=\"Show\")\n",
    "btn_show.on_click(process_and_show_replica)\n",
    "\n",
    "# output_header = ipw.Output()\n",
    "output_plot = ipw.Output()\n",
    "output_thumbnails = ipw.Output()\n",
    "output_checked_pks = ipw.Output()\n",
    "\n",
    "display(node, btn_show, output_plot, output_thumbnails, output_checked_pks)"
   ]
  },
  {
   "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",
    "    node.value = urllib.parse.parse_qs(url.query)[\"pk\"][0]\n",
    "    process_and_show_replica(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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}