{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Migrating to AiiDA\n",
    "\n",
    ":::{admonition} Learning Objectives\n",
    ":class: learning-objectives\n",
    "\n",
    "In this section, we will look at how to migrate from running a quantum code from text-based input files, to running it within AiiDA, and understand how AiiDA automates the computation execution and output parsing.\n",
    "\n",
    "We shall take the example of Quantum ESPRESSO, but the same principles apply to any other code.\n",
    "This would be a typical command line script to run a Quantum ESPRESSO relaxation:\n",
    "\n",
    "```console\n",
    "$ mpirun -np 2 pw.x -in pwx.in > pwx.out\n",
    "```\n",
    "\n",
    ":::\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modularising the inputs\n",
    "\n",
    "The first step is to modularise the inputs within the `input.in` file, and any pseudo-potential files.\n",
    "\n",
    "By splitting them into separate components, we can create **re-usable** building blocks for multiple calculations.\n",
    "We shall also see later how these components can be generated from external data sources, such as databases or web APIs.\n",
    "\n",
    "![pw-to-aiida](_static/pw-to-aiida.svg){height=500px}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the diagram above, we have split the input generation into separate entities, handling the different aspects of the calculation and allowing for component re-use.\n",
    "For a `pw.x` calculation, we need to create the following nodes:\n",
    "\n",
    "- {term}`Computer`, which describes how we interface with a compute resource\n",
    "- {term}`Code`, which contains the information on how to execute a single calculation\n",
    "- `StructureData`, which contains the crystal structure\n",
    "- `UpfData`, which contains the pseudo-potentials per atomic\n",
    "- `KpointsData`, which contains the k-point mesh\n",
    "- `Dict` node, which contains the parameters for the calculation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The AiiDA Profile\n",
    "\n",
    "First we need to create a new AiiDA profile.\n",
    "This is where we store all the nodes generated for a project, and the links between them.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "Here we generate a profile with temporary, in-memory storage, which will be destroyed when the Python is restarted.\n",
    "This is useful for testing, but for a real project, you would create a persistent profile connected to a PostgreSQL database,\n",
    "using the `verdi quicksetup` command.\n",
    "\n",
    ":::\n"
   ]
  },
  {
   "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(name=\"qe-to-aiida\")\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import aiida\n",
    "\n",
    "profile = aiida.load_profile(\"qe-to-aiida\")\n",
    "profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi profile show qe-to-aiida"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check on the status of the profile using the `verdi status` command.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi -p qe-to-aiida status --no-rmq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also check the statistics of the profile's storage.\n",
    "Before running any simulations, we see that only a single {term}`User` node has been created, which is the default creator of data for the profile.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi storage info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a compute resource\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An AiiDA {term}`Computer` represents a compute resource, such as a local or remote machine.\n",
    "It contains information on how to connect to the machine, how to **transport** data to/from the compute resource, and how to **schedule** jobs on it.\n",
    "\n",
    "In the following we will use a simple `local_direct` computer, which connects to the local machine, and runs the calculations directly, without any scheduler.\n",
    "\n",
    "AiiDA also has built-in support for a number of {term}`Scheduler`s, including:\n",
    "\n",
    "- `pbspro`\n",
    "- `slurm`\n",
    "- `sge`\n",
    "- `torque`\n",
    "- `lsf`\n",
    "\n",
    "Connections to remote machines can be made using the `SSH` {term}`Transport`, and [aiida-code-registry](https://github.com/aiidateam/aiida-code-registry) provides a collection of example configurations for Swiss based HPC clusters.\n",
    "\n",
    "We can create the computer using the `verdi computer setup` CLI.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi computer setup \\\n",
    "    --non-interactive \\\n",
    "    --label local_direct \\\n",
    "    --hostname localhost \\\n",
    "    --description \"Local computer with direct scheduler\" \\\n",
    "    --transport core.local \\\n",
    "    --scheduler core.direct \\\n",
    "    --work-dir {data.workdir} \\\n",
    "    --mpiprocs-per-machine {data.cpu_count}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi computer configure core.local local_direct \\\n",
    "    --non-interactive \\\n",
    "    --safe-interval 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or we can use the `Computer` class from the `aiida.orm` API module.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "created, computer = aiida.orm.Computer.collection.get_or_create(\n",
    "    label=\"local_direct\",\n",
    "    description=\"local computer with direct scheduler\",\n",
    "    hostname=\"localhost\",\n",
    "    workdir=str(data.workdir),\n",
    "    transport_type=\"core.local\",\n",
    "    scheduler_type=\"core.direct\",\n",
    ")\n",
    "if created:\n",
    "    computer.store()\n",
    "    computer.set_minimum_job_poll_interval(0.0)\n",
    "    computer.set_default_mpiprocs_per_machine(data.cpu_count)\n",
    "    computer.configure()\n",
    "computer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a computer, ready to run calculations on.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi computer show local_direct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up a code plugin\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An AiiDA {term}`Code` represent a single executable, and contain information on how to execute it.\n",
    "The `Code` node is associated with a specific `Computer`, contains the path to the executable, and is associated with a specific {term}`CalcJob` plugin we shall discuss later.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we can use either the CLI or the API to create a new `Code` node.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi code setup \\\n",
    "    --non-interactive \\\n",
    "    --label pw.x \\\n",
    "    --description \"Quantum ESPRESSO pw.x code\" \\\n",
    "    --computer local_direct \\\n",
    "    --remote-abs-path {data.pwx_path} \\\n",
    "    --input-plugin quantumespresso.pw \\\n",
    "    --prepend-text \"export OMP_NUM_THREADS=1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    code = aiida.orm.load_code(\"pw.x@local_direct\")\n",
    "except aiida.common.NotExistent:\n",
    "    code = aiida.orm.Code(\n",
    "        input_plugin_name=\"quantumespresso.pw\",\n",
    "        remote_computer_exec=[computer, data.pwx_path],\n",
    "    )\n",
    "    code.label = \"pw.x\"\n",
    "    code.description = \"Quantum ESPRESSO pw.x code\"\n",
    "    code.set_prepend_text(\"export OMP_NUM_THREADS=1\")\n",
    "    code.store()\n",
    "code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a code ready to run our computations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi code show pw.x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deconstructing the input file\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now take a look at a typical `pw.x` input file, and how we can convert it to the requisite AiiDA nodes.\n",
    "\n",
    ":::{note}\n",
    "\n",
    "Here we are simply generating the inputs from a pre-written input file.\n",
    "But in practice, you would want to generate the inputs from a Python script, or from a database or web API, as we shall see in the next section.\n",
    "\n",
    ":::\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%cat direct_run/pwx.in"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To decompose this file into the components we need, we can use the [qe_tools](https://pypi.org/project/qe-tools/) package, which provides a Python API to parse Quantum ESPRESSO input files.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import qe_tools\n",
    "\n",
    "pw_input = qe_tools.parsers.PwInputFile(open(\"direct_run/pwx.in\").read())\n",
    "pw_input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then generate our AiiDA input {term}`Data` nodes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "structure = aiida.orm.StructureData(cell=pw_input.structure[\"cell\"])\n",
    "for p, s in zip(pw_input.structure[\"positions\"], pw_input.structure[\"atom_names\"]):\n",
    "    structure.append_atom(position=p, symbols=s)\n",
    "structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kpoints = aiida.orm.KpointsData()\n",
    "kpoints.set_cell_from_structure(structure)\n",
    "kpoints.set_kpoints_mesh(\n",
    "    pw_input.k_points[\"points\"],\n",
    "    offset=pw_input.k_points[\"offset\"],\n",
    ")\n",
    "kpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# AiiDA will handle assigning file names to generated input files,\n",
    "# and computing te system type from the structure.\n",
    "_parameters = pw_input.namelists\n",
    "for disallowed in [\"pseudo_dir\", \"outdir\", \"prefix\"]:\n",
    "    _parameters[\"CONTROL\"].pop(disallowed, None)\n",
    "for disallowed in [\"nat\", \"ntyp\"]:\n",
    "    _parameters[\"SYSTEM\"].pop(disallowed, None)\n",
    "parameters = aiida.orm.Dict(dict=_parameters)\n",
    "parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import abspath\n",
    "\n",
    "pseudo_si, _ = aiida.orm.UpfData.get_or_create(\n",
    "    abspath(\"direct_run/pseudo/Si.pbe-n-rrkjus_psl.1.0.0.UPF\")\n",
    ")\n",
    "pseudo_si"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the inputs for a calculation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `verdi plugin list aiida.calculations` we can inspect the full specification for the inputs of the calculation plugin we wish to use.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi plugin list aiida.calculations quantumespresso.pw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we already assigned the `quantumespresso.pw` plugin to our `Code` node, we can load it and use the `get_builder` to generate a template for the inputs, known as the `Builder`.\n",
    "\n",
    "The `Builder` provides us a structured way to add (and validate) the inputs for the calculation.\n",
    "Below we add the input nodes that we have created for our calculation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "code = aiida.orm.load_code(\"pw.x@local_direct\")\n",
    "builder = code.get_builder()\n",
    "builder.structure = structure\n",
    "builder.parameters = parameters\n",
    "builder.kpoints = kpoints\n",
    "builder.pseudos = {\"Si\": pseudo_si}\n",
    "\n",
    "# we can also add metadata like the maximum walltime\n",
    "builder.metadata.options.max_wallclock_seconds = 30 * 60\n",
    "\n",
    "builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the calculation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AiiDA provides two main ways to run a calculation:\n",
    "\n",
    "1. Using the `engine.run` functions, which runs the computation directly and waits for it to complete.\n",
    "2. Using the `engine.submit` function, which submits the calculation to the AiiDA daemon, which can be started in the background and manages the execution of the calculations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = aiida.engine.run_get_node(builder)\n",
    "output.node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How the calculation is run\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On executing the calculation, AiiDA will:\n",
    "\n",
    "1. Generate the input files necessary for the calculation, and the submission script specific to the computer's scheduler.\n",
    "2. Write the input files to the desired location on the local/remote computer.\n",
    "3. Submit the job to the scheduler.\n",
    "4. Monitor the job until it completes.\n",
    "5. Retrieve the output files from the computer.\n",
    "6. Parse the output files and store the results.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generated input files are stored on the `CalcJobNode`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "calcnode_repo = output.node.base.repository\n",
    "print(\"input files: \", calcnode_repo.list_object_names())\n",
    "print(\"-\" * 10 + \"\\naiida.in\\n\" + \"-\" * 10)\n",
    "print(calcnode_repo.get_object_content(\"aiida.in\"))\n",
    "print(\"-\" * 16 + \"\\n_aiidasubmit.sh\\n\" + \"-\" * 16)\n",
    "print(calcnode_repo.get_object_content(\"_aiidasubmit.sh\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are then \"transported\" to the remote computer, into a unique sub-folder of the the working directory.\n",
    "\n",
    ":::{tip}\n",
    "\n",
    "These folders and their contents are not deleted by default after the calculation is completed, and can be inspected at any time with `verdi calcjob gotocomputer <IDENTIFIER>`.\n",
    "\n",
    "Many workflows though can be configured to clean up these folders after the calculation is (successfully) completed, to save disk space.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output.node.get_remote_workdir()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The retrieved output files are stored in the `retrieved` output node from the `CalcJobNode`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "print(\"output files:\", output.node.get_retrieved_node().list_object_names())\n",
    "print(\"-\" * 10 + \"\\naiida.out\\n\" + \"-\" * 10)\n",
    "print(output.node.get_retrieved_node().get_object_content(\"aiida.out\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and finally, the parsed results are stored on defined output nodes from the `CalcJobNode`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "%verdi process show {output.node.pk}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then access key results from the calculation using the `CalcJobNode`s `outputs` method (or loading the node by its identifier)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "output.node.outputs.output_parameters.get_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AiiDA automatically generates links between the inputs, calculation and outputs, to generate the provenance graph.\n",
    "The provence graph is a directed acyclic graph (DAG) that contains the nodes and links between them, and can be used for visualisation of a calculation or workflow, or with advance querying of the stored results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aiida.tools.visualization import Graph\n",
    "\n",
    "graph = Graph()\n",
    "graph.add_incoming(output.node, annotate_links=\"both\")\n",
    "graph.add_outgoing(output.node, annotate_links=\"both\")\n",
    "graph.graphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = aiida.orm.QueryBuilder()\n",
    "query.append(aiida.orm.StructureData, tag=\"initial\", project=\"*\")\n",
    "query.append(\n",
    "    aiida.orm.CalcJobNode,\n",
    "    filters={\"attributes.process_state\": \"finished\"},\n",
    "    tag=\"calculation\",\n",
    "    with_incoming=\"initial\",\n",
    "    project=\"id\",\n",
    ")\n",
    "query.append(\n",
    "    aiida.orm.StructureData, tag=\"final\", with_incoming=\"calculation\", project=\"*\"\n",
    ")\n",
    "query.dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving compute time with caching\n",
    "\n",
    "Over the course of a project, you may end up re-running the same calculations multiple times - perhaps because two workflows include the same calculation.\n",
    "\n",
    "Since AiiDA stores the full provenance of each calculation, it can detect whether a calculation has been run before and, instead of running it again, simply reuse its outputs, thereby saving valuable computational resources. This is what we mean by **caching** in AiiDA.\n",
    "\n",
    "With caching enabled, AiiDA searches the database for a calculation of the same hash. If found, AiiDA creates a copy of the calculation node and its results, thus ensuring that the resulting provenance graph is independent of whether caching is enabled or not.\n",
    "\n",
    "![caching](_static/caching.png){align=center width=\"150px\"}\n",
    "\n",
    "Caching happens on the calculation level (not the workchain level), and is **not** enabled by default.\n",
    "We can enable it by setting the `verdi config` options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi config set caching.enabled_for 'aiida.calculations:quantumespresso.pw'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%verdi config list caching"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, when we run the same calculation again, AiiDA will detect that it has already been run, and will simply reuse the results from the previous run!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = aiida.engine.run_get_node(builder)\n",
    "output.node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the calculation was created from the cache by checking the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output.node.base.caching.is_created_from_cache"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "The [caching how-to documentation](https://aiida.readthedocs.io/projects/aiida-core/en/latest/howto/run_codes.html#how-to-run-codes-caching).\n",
    "\n",
    ":::"
   ]
  }
 ],
 "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": "20c30adb377910d9d5c8112cf74e9f7ecb37538254a701570d770f074373c53e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}