{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code your own polarizability tensor workflow, part 1\n",
    "\n",
    "This main goal of this notebook is to show how the `Workflow` class can be used to quickly implement of a workflow. It allows to define a queue of BigDFT jobs to be run in order to compute a quantity afterwards, as a post-processing procedure.\n",
    "\n",
    "We will stick to the polarizability tensor problem, in order to build on what was achieved in the [previous exercise](https://mmoriniere.gitlab.io/MyBigDFT/notebooks/Solution_01.html). The structure of this notebook is very similar to the previous one. You can use the same resources to help you, while the solution to this exercise can be found [here](https://mmoriniere.gitlab.io/MyBigDFT/notebooks/Solution_02.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The [Workflow](https://mmoriniere.gitlab.io/MyBigDFT/workflow.html#mybigdft.workflows.workflow.Workflow) class\n",
    "\n",
    "Using the `Workflow` class allows you to start automatizing the generalization of a workflow. You will generally want to use it in order to create toy workflows, allowing you to easily experiment whether such workflows can be successfully implemented, without having to actually code a workflow class deriving from the `AbstractWorkflow` class (to be presented in the next exercise).\n",
    "\n",
    "One advantage of using this `Workflow` class is that it allows you to properly divide the workflow in a three-step process:\n",
    "\n",
    "- the definition of a list of jobs so as to initialize the workflow,\n",
    "- the sequential running of all the calculations defined in the first step,\n",
    "- the post-processing of of the calculations outputs, so as to extract the quantity of interest, here the polarizability tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define your worflow\n",
    "\n",
    "As already mentioned, a workflow is mainly made of a queue of BigDFT jobs. This means you need to define a list of `Job` instances, to be run sequentially afterwards. \n",
    "\n",
    "We want to compute the polarizability tensor of the N$_2$ molecule using the `Workflow` class, just like in the previous exercise. The main input parameters being given, you only have to add `Job` instances to the `pt_queue` list. You can do so by re-using the same jobs as in the previous notebook (especially the same electric fields and run directories): first, a ground state job (without electric field) and then three calculations with an electric field, along the $x$, $y$ and $z$ axis respectively. Try to use a for loop to add the last three jobs to the queue.\n",
    "\n",
    "Only the base classes of the `MyBigDFT` package are required to do that, and might also require the help of extra packages such as `numpy` (you should not need it yet, though)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mybigdft import Workflow, InputParams, Posinp, Atom, Job\n",
    "import numpy as np\n",
    "\n",
    "# Main input variables:\n",
    "# - use the optimized structure for the N2 molecule (found using the\n",
    "#   LDA exchange-correlation functional)\n",
    "atoms = [Atom('N', [0, 0, 0]), Atom('N', [0, 0, 1.0935])]\n",
    "pos = Posinp(atoms, units=\"angstroem\", boundary_conditions=\"free\")\n",
    "# - use some non-default input parameters to get high quality results\n",
    "inp = InputParams({\"dft\": {\"rmult\": [7, 9], \"hgrids\": 0.35}})\n",
    "# Electric field amplitude to be applied in each direction\n",
    "ef_amplitude = 1.e-4\n",
    "\n",
    "# Initialize the queue of jobs in order to be able to compute the \n",
    "# polarizability tensor\n",
    "pt_queue = []\n",
    "\n",
    "pt_wf = Workflow(queue=pt_queue)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us ensure that you followed the instructions:\n",
    "# - there are four jobs in the queue\n",
    "assert len(pt_wf.queue) == 4\n",
    "# - correct electric fields are applied\n",
    "assert \"elecfield\" not in pt_wf.queue[0].inputparams[\"dft\"]\n",
    "assert pt_wf.queue[1].inputparams[\"dft\"][\"elecfield\"] == [1.e-4, 0.0, 0.0]\n",
    "assert pt_wf.queue[2].inputparams[\"dft\"][\"elecfield\"] == [0.0, 1.e-4, 0.0]\n",
    "assert pt_wf.queue[3].inputparams[\"dft\"][\"elecfield\"] == [0.0, 0.0, 1.e-4]\n",
    "# - correct run directories\n",
    "assert pt_wf.queue[0].run_dir.endswith(\"N2/pol_tensor\")\n",
    "assert pt_wf.queue[1].run_dir.endswith(\"N2/pol_tensor/EF_along_x+\")\n",
    "assert pt_wf.queue[2].run_dir.endswith(\"N2/pol_tensor/EF_along_y+\")\n",
    "assert pt_wf.queue[3].run_dir.endswith(\"N2/pol_tensor/EF_along_z+\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the jobs of your workflow\n",
    "\n",
    "Nothing important to code here, as you might only want to change the value of the attributes to suit your needs. \n",
    "\n",
    "Note that you do not have to use the context manager: no need for a `with` statement before using the `run` method of a `Workflow` instance as it is actually taken care of internally. This is one of the advantage of using a workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_wf.run(nmpi=6, nomp=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that, if you succeeded in performing the previous notebook, the calculations were not run again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us ensure that you followed the instructions:\n",
    "assert pt_wf.is_completed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post-process your results\n",
    "\n",
    "Now that you ran the workflow, you can use the results to extract meaningful quantities. Let us define a `post_proc` function to do just that. Its only argument must be a workflow instance. You should be able to do the post-processing without having to add another argument to the `post_proc` function. Try to loop over the three calculations with an electric field to get their dipoles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def post_proc(workflow):\n",
    "    \"\"\"\n",
    "    Use the logfile of the calculations run by the workflow to\n",
    "    extract the polarizability tensor.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    workflow: Workflow\n",
    "        Workflow of run calculations allowing to compute the\n",
    "        polarizability tensor of a given system.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    numpy array\n",
    "        Polarizability tensor of the system.\n",
    "    \"\"\"\n",
    "    pol_tensor = np.zeros((3, 3))\n",
    "    \n",
    "    return pol_tensor\n",
    "\n",
    "# Run the post-processing of the workflow to compute\n",
    "pol_tensor = post_proc(pt_wf)\n",
    "print(pol_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us ensure that you obtained the expected result:\n",
    "expected_pol_tensor = [[ 1.075753e+01,  5.400000e-04, -9.900000e-04],\n",
    "                       [ 5.400000e-04,  1.075753e+01, -9.900000e-04],\n",
    "                       [-1.740000e-03, -1.740000e-03,  1.525630e+01]]\n",
    "assert np.allclose(pol_tensor, expected_pol_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations, you obtained the same results as in the previous notebook! Note how easier it is to re-use the code implemented here compared to the previous one. This is one of the main advantage of using the `Workflow` class. Another one is that it enables you to create a toy workflow which can be easily converted to a proper class, to make your code even more general and easier to re-use. This latter class must derive from the `AbstractWorkflow` class: you will implement one in the [next exercise](https://mmoriniere.gitlab.io/MyBigDFT/notebooks/Exercise_03_Polarizability_Tensor_Workflow_3.html)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.3"
  },
  "nbsphinx": {
   "execute": "never"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}