SearchNEB.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": [
    "%aiida\n",
    "from apps.reactions.replicawork import ReplicaWorkchain\n",
    "from aiida.work.process import WorkCalculation\n",
    "from aiida_cp2k.calculations import Cp2kCalculation\n",
    "from aiida.orm.data.structure import StructureData\n",
    "from aiida.orm.data.parameter import ParameterData\n",
    "from aiida.orm.data.folder import FolderData\n",
    "\n",
    "from aiida.common.exceptions import NotExistent\n",
    "from aiida.orm.data.base import Str\n",
    "\n",
    "import ipywidgets as ipw\n",
    "from IPython.display import display, clear_output, HTML\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from pprint import pprint\n",
    "\n",
    "import ase.io\n",
    "import tempfile\n",
    "\n",
    "import re\n",
    "import numpy as np\n",
    "import scipy.constants\n",
    "\n",
    "from collections import OrderedDict\n",
    "\n",
    "from tempfile import NamedTemporaryFile\n",
    "from base64 import b64encode\n",
    "def render_thumbnail(atoms):\n",
    "    tmp = NamedTemporaryFile()\n",
    "    ase.io.write(tmp.name, atoms, format='png') # does not accept StringIO\n",
    "    raw = open(tmp.name).read()\n",
    "    tmp.close()\n",
    "    return b64encode(raw)\n",
    "\n",
    "def display_thumbnail(th):\n",
    "    return '<img width=\"400px\" src=\"data:image/png;base64,{}\" title=\"\">'.format(th)\n",
    "def html_thumbnail(th):\n",
    "    return ipw.HTML('<img width=\"400px\" src=\"data:image/png;base64,{}\" title=\"\">'.format(th))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def replicas_from_restart(restart_file_path):\n",
    "    content = open(restart_file_path).read()\n",
    "    m = re.search(r'\\n\\s*&CELL\\n(.*?)\\n\\s*&END CELL\\n', content, re.DOTALL)\n",
    "    cell_lines = [line.strip().split() for line in m.group(1).split(\"\\n\")]\n",
    "    cell_str = [line[1:] for line in cell_lines if line[0] in 'ABC']\n",
    "    cell = np.array(cell_str, np.float64)\n",
    "    \n",
    "    structure_data_list = []\n",
    "    \n",
    "    matches = re.findall(r'\\n\\s*&COORD\\n(.*?)\\n\\s*&END COORD\\n', content, re.DOTALL)\n",
    "\n",
    "    coord_line_sets = [\n",
    "        [line.strip().split() for line in m.split(\"\\n\")] for m in matches\n",
    "    ]\n",
    "    coord_set_with_elements = coord_line_sets[-1]\n",
    "    replica_coord_line_sets = coord_line_sets[:-1]\n",
    "    element_list = [line[0] for line in coord_set_with_elements]\n",
    "    for i_rep, rep_coord_lines in enumerate(replica_coord_line_sets):\n",
    "        positions = np.array(rep_coord_lines, np.float64)\n",
    "        ase_atoms = ase.Atoms(symbols=element_list, positions=positions, cell=cell)\n",
    "        structure_data_list.append(StructureData(ase=ase_atoms))\n",
    "        \n",
    "    return structure_data_list\n",
    "\n",
    "def neb_energies_and_distances_from_output(output_file_path):\n",
    "    content = open(output_file_path).read()\n",
    "    final_energy_str = re.findall(r'ENERGIES \\[au\\] = (.*?)BAND', content, re.DOTALL)[-1]\n",
    "    energies = np.array(final_energy_str.split(),dtype=float) * scipy.constants.physical_constants['Hartree energy in eV'][0]\n",
    "    energies -= energies[0]\n",
    "    \n",
    "    final_dist_str = re.findall(r'DISTANCES REP = (.*?)ENERGIES', content, re.DOTALL)[-1]\n",
    "    distances = np.array(['0.0'] + final_dist_str.split(),dtype=float) * 0.529177249\n",
    "    \n",
    "    return distances, energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_replica_html(structure_data_list, energies):\n",
    "    html = '<table>'\n",
    "    \n",
    "    for i, (rep, en) in enumerate(zip(structure_data_list, energies)):\n",
    "        thumbnail = rep.get_extra('thumbnail')\n",
    "        # The table cell\n",
    "        if i%2 == 0:\n",
    "            html += '<tr>'\n",
    "        html += '<td><img width=\"400px\" src=\"data:image/png;base64,{}\" title=\"\">'.format(thumbnail)\n",
    "        # Output some information about the replica...\n",
    "        html += '<p><b>Nr: </b>{} <br> <b>Energy:</b> {}</p>'\\\n",
    "                .format(i, en)\n",
    "        html += '<p>pk: {}</p>'.format(rep.pk)\n",
    "        # ... and the download link.\n",
    "        html += '<p><a target=\"_blank\" href=\"./export_structure.ipynb?uuid={}\">Export Structure</a></p><td>'\\\n",
    "                .format(rep.uuid)\n",
    "        if i%2 == 1:\n",
    "            html += '</tr>'\n",
    "            \n",
    "    html += '</tr>'\n",
    "    html += '</table>'\n",
    "    return html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sorted_opt_rep_keys(keys):\n",
    "    return sorted([ (int(key.split('_')[2]), key) for key in keys if 'opt_replica' in key])\n",
    "\n",
    "def process_and_show_neb(c):\n",
    "    \n",
    "    btn_show.disabled = True\n",
    "    with main_out:\n",
    "        clear_output()\n",
    "        \n",
    "    \n",
    "    wc = drop_nebs.value\n",
    "    cp2k_calc = wc.get_outputs()[0]\n",
    "    cc_output_dict = cp2k_calc.get_outputs_dict()\n",
    "\n",
    "    #### --------------------------------------------------------------\n",
    "    ## Create a list of structure_data object corresponding to replicas\n",
    "    if 'opt_replica_0' in cc_output_dict:\n",
    "        # NEW: get replicas from output nodes\n",
    "        structure_data_list = []\n",
    "        for rep_key in sorted_opt_rep_keys(cc_output_dict.keys()):\n",
    "            if rep_key[0] >= len(structure_data_list):\n",
    "                structure_data_list.append(cc_output_dict[rep_key[1]])\n",
    "    else:\n",
    "        # OLD: store replicas in nodes and link through extra\n",
    "        structure_data_list = []\n",
    "        for rep_key in sorted_opt_rep_keys(cp2k_calc.get_extras().keys()):\n",
    "            struct_pk = cp2k_calc.get_extras()[rep_key[1]]\n",
    "            structure_data_list.append(load_node(struct_pk))\n",
    "\n",
    "        if len(structure_data_list) == 0:\n",
    "            restart_file_path = cc_output_dict['retrieved'].get_abs_path('aiida-1.restart')\n",
    "            structure_data_list = replicas_from_restart(restart_file_path)\n",
    "            for i_rep, rep in enumerate(structure_data_list):\n",
    "                rep.store()\n",
    "                cp2k_calc.set_extra(\"opt_replica_%d\" % i_rep, rep.pk)\n",
    "                with main_out:\n",
    "                    print(\"Added replica %d with pk %d to database\" % (i_rep, rep.pk))\n",
    "        \n",
    "    #### --------------------------------------------------------------\n",
    "    ## Add thumbnails to replicas if they are not already added\n",
    "    for rep in structure_data_list:\n",
    "        if not \"thumbnail\" in rep.get_extras():\n",
    "            rep.set_extra(\"thumbnail\", render_thumbnail(rep.get_ase()))\n",
    "    \n",
    "    #### --------------------------------------------------------------\n",
    "    output_file_path = cc_output_dict['retrieved'].get_abs_path('aiida.out')\n",
    "    distances, energies = neb_energies_and_distances_from_output(output_file_path)\n",
    "    \n",
    "    replica_html = make_replica_html(structure_data_list, energies)\n",
    "    \n",
    "    with main_out:\n",
    "        plt.plot(energies, 'o-')\n",
    "        plt.ylabel(\"Energy (eV)\")\n",
    "        plt.xlabel(\"Replica nr\")\n",
    "        plt.show()\n",
    "        \n",
    "        display(ipw.HTML(replica_html))\n",
    "    \n",
    "    btn_show.disabled = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qb = QueryBuilder()\n",
    "qb.append(WorkCalculation,\n",
    "          filters={\n",
    "              'attributes._process_label': {'==': 'NEBWorkchain'}\n",
    "          },\n",
    "          tag='wc',\n",
    "          project='*'\n",
    ")\n",
    "qb.append(Cp2kCalculation,\n",
    "          output_of='wc',\n",
    "          filters={\n",
    "              'attributes.state': {'==': 'FINISHED'}}\n",
    ")\n",
    "qb.order_by({WorkCalculation: {'id': 'desc'}})\n",
    "neb_wcs = qb.all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sel_options = OrderedDict([(\"PK %d: \" % (neb_wc[0].pk) + neb_wc[0].description, neb_wc[0]) for neb_wc in neb_wcs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style = {'description_width': '120px'}\n",
    "layout = {'width': '70%'}\n",
    "\n",
    "drop_nebs = ipw.Dropdown(options = sel_options,\n",
    "                            description = 'NEB: ', layout=layout, style=style)\n",
    "\n",
    "btn_show = ipw.Button(description=\"Show\")\n",
    "btn_show.on_click(process_and_show_neb)\n",
    "\n",
    "main_out = ipw.Output()\n",
    "\n",
    "display(drop_nebs, btn_show, main_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}