{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AiiDA by example: Computing a band structure\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{admonition} Learning Objectives\n",
    ":class: learning-objectives\n",
    "\n",
    "In this section we will present a complete example of an AiiDA workflow, which defines the sequence of calculations needed to compute the band structure of silicon.\n",
    "\n",
    "How to setup the input data and the details of the workflow execution will be discussed in subsequent sections.\n",
    "Here we simply give an initial overview of what it means to run an AiiDA workflow.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interacting with AiiDA\n",
    "\n",
    "AiiDA can be controlled in two ways:\n",
    "\n",
    "1. Using the `verdi` command line interface (CLI), or `%verdi` magic in Jupyter notebooks.\n",
    "2. Using the `aiida` Python API\n",
    "\n",
    "For each project in AiiDA, we set up a *profile*, which defines the connection to the data storage, and other settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from local_module import load_temp_profile\n",
    "\n",
    "data = load_temp_profile(\n",
    "    name=\"bands_workflow\",\n",
    "    add_computer=True,\n",
    "    add_pw_code=True,\n",
    "    add_sssp=True,\n",
    "    add_structure_si=True,\n",
    ")\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi status --no-rmq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Within this profile, we have stored the initial input components for our workflow, including the pseudo-potentials, and the silicon structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi storage info --detailed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have also set up the compute resource that we will use to run the calculations, and the code (`pw.x`) installed on that computer, which we will use to perform the electronic structure calculations.\n",
    "\n",
    "Here, we will use our \"local\" machine to run the computations, but AiiDA can also be used to submit calculations to remote supercomputer schedulers, transporting data between the local machine and the remote computer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi computer show local_direct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi code show pw.x@local_direct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utilising a pre-defined workflow\n",
    "\n",
    "AiiDA plugins can declare workflow plugins, for use within AiiDA.\n",
    "These are workflows that are pre-defined, and can be used as-is, or as a starting point for your own workflows.\n",
    "\n",
    "Here we utilise the `quantumespresso.pw.bands` workflow defined by the [`aiida-quantumespresso`](https://aiida-quantumespresso.readthedocs.io) plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi plugin list aiida.workflows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi plugin list aiida.workflows quantumespresso.pw.bands"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![quantumespresso.pw.bands](_static/bands-workflow.svg){height=250px align=center}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `quantumespresso.pw.bands` workflow provides a helpful method for setting up the default inputs for a given \"protocol\", as to how fast/precise the calculation should be.\n",
    "This provides a \"builder\" object, which stores all the inputs for the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "from aiida_quantumespresso.workflows.pw.bands import PwBandsWorkChain\n",
    "\n",
    "builder = PwBandsWorkChain.get_builder_from_protocol(\n",
    "    code=data.code, \n",
    "    structure=data.structure,\n",
    "    protocol=\"fast\",\n",
    ")\n",
    "builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the workflow\n",
    "\n",
    "Workflows can be run in the interpreter using the `run` method, in a blocking manner, which we shall do here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "from aiida import engine\n",
    "\n",
    "result = engine.run_get_node(builder)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Typically however, long running workflows are executed by using the `submit` method.\n",
    "This will store the initial state of the workflow in the profile storage, and notify the AiiDA daemon to run the workflow in the background.\n",
    "\n",
    "The AiiDA daemon can be launched using the `verdi daemon start n` command, with `n` being the number of worker processes to launch.\n",
    "Each worker can asynchronously handle 1000s of individual calculations, allowing for a high-throughput of workflow submissions.\n",
    "\n",
    "![daemon illustration](_static/aiida/rabbitmq-sysml.svg){height=200px align=center}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each workflow and node stored in the AiiDA profile is assigned a unique identifier (a.k.a Primary Key), which can be used to reference them.\n",
    "The execution of the workflows can be monitored using the `verdi process list` command, which will show the status of all running processes in the profile (or also finished ones with `-a`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi process list -a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also monitor the progress of individual workflows using the `verdi process status` command, which will show the status of the individual steps of the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi process status {result.node.pk}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This work-chain demonstrates how we can build up a complex workflow from a series of individual calculations.\n",
    "In this case, the workflow is made up of the following steps:\n",
    "\n",
    "1. The `PwRelaxWorkChain` will run multiple Quantum ESPRESSO `vc-relax` calculations, to make sure that there are no Pulay stresses present in the material and that the requested k-points density is respected in case there is a significant volume change in the material.\n",
    "2. Once the geometry has been optimized, SeeK-path will be used to primitivize and standardize the structure, as well as find the standard path along which to calculate the band structure.\n",
    "3. A static calculation (`scf`) is run to calculate the charge density for the structure obtained from SeeK-path.\n",
    "4. Finally, an NSCF is run to calculate the band structure along the path determined by Seek-path.\n",
    "\n",
    "We shall also discuss in subsequent sections, how the `PwBaseWorkChain` can identify and recover from known failure modes, such as reaching the wall-time limit of the scheduler, or convergence failures."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspecting the results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we the workflow has finished, we can inspect the results using the `verdi process show` command, which will show the results of the workflow, and its \"attached\" outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi process show {result.node.pk}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The provenance graph\n",
    "\n",
    "As well as storing the inputs and outputs of the workflow, and its composite calculations, AiiDA also stores the links between them, which can be used to reconstruct the provenance graph of the workflow.\n",
    "\n",
    "This can be visualised using the `verdi node graph generate` command, or using the `Graph` Python API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aiida.tools.visualization import Graph\n",
    "\n",
    "graph = Graph(graph_attr={\"rankdir\": \"TB\", \"size\": \"8!,8!\"})\n",
    "graph.recurse_ancestors(result.node, annotate_links=\"both\")\n",
    "graph.recurse_descendants(result.node, annotate_links=\"both\")\n",
    "graph.graphviz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The output structure\n",
    "\n",
    "AiiDA's `StructureData` class provides integration with both [ASE](https://wiki.fysik.dtu.dk/ase/), and [Pymatgen](https://pymatgen.org/), which can be used to inspect and visualise the structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pym_structure = result.node.outputs.primitive_structure.get_pymatgen()\n",
    "pym_structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.visualize.plot import plot_atoms\n",
    "\n",
    "ase_atoms = result.node.outputs.primitive_structure.get_ase()\n",
    "ax = plot_atoms(ase_atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from ase.visualize.ngl import NGLDisplay\n",
    "\n",
    "NGLDisplay(ase_atoms).gui"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The output band structure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we get to our desired result, the band structure of silicon computed using Quantum ESPRESSO 🎉"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from local_module.bandstructure import plot_bandstructure\n",
    "\n",
    "bands = result.node.outputs.band_structure\n",
    "fig, ax = plot_bandstructure(bands)"
   ]
  }
 ],
 "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"
  },
  "toc-autonumbering": true,
  "toc-showcode": false,
  "toc-showtags": false,
  "vscode": {
   "interpreter": {
    "hash": "20c30adb377910d9d5c8112cf74e9f7ecb37538254a701570d770f074373c53e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}