{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solution\n",
    "\n",
    "Here is the solution to the [second exercise](https://mmoriniere.gitlab.io/MyBigDFT/notebooks/Exercise_02_Polarizability_Tensor_Workflow_1.html) regarding 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": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import deepcopy\n",
    "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",
    "# - Add the ground state job (without any electric field)\n",
    "base_run_dir = os.path.relpath(\"N2/pol_tensor/\")\n",
    "gs = Job(inputparams=inp, posinp=pos, name=\"N2\",\n",
    "         run_dir=base_run_dir)\n",
    "pt_queue = [gs]\n",
    "# - Add the jobs with an electric field\n",
    "for i, coord in enumerate([\"x\", \"y\", \"z\"]):\n",
    "    # Add the electric field to the base input parameters\n",
    "    elecfield = [0.0] * 3\n",
    "    elecfield[i] = ef_amplitude\n",
    "    new_inp = deepcopy(inp)\n",
    "    new_inp[\"dft\"][\"elecfield\"] = elecfield\n",
    "    # Define the correct run directory\n",
    "    new_run_dir = os.path.join(gs.run_dir,\n",
    "                               \"EF_along_{}+/\".format(coord))\n",
    "    # Append the new job to the end of the queue\n",
    "    pt_queue.append(Job(inputparams=new_inp, posinp=pos,\n",
    "                        name=\"N2\", run_dir=new_run_dir))\n",
    "pt_wf = Workflow(queue=pt_queue)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/maximemoriniere/Documents/Python/MyBigDFT/doc/source/notebooks/N2/pol_tensor\n",
      "Logfile log-N2.yaml already exists!\n",
      "\n",
      "/Users/maximemoriniere/Documents/Python/MyBigDFT/doc/source/notebooks/N2/pol_tensor/EF_along_x+\n",
      "Logfile log-N2.yaml already exists!\n",
      "\n",
      "/Users/maximemoriniere/Documents/Python/MyBigDFT/doc/source/notebooks/N2/pol_tensor/EF_along_y+\n",
      "Logfile log-N2.yaml already exists!\n",
      "\n",
      "/Users/maximemoriniere/Documents/Python/MyBigDFT/doc/source/notebooks/N2/pol_tensor/EF_along_z+\n",
      "Logfile log-N2.yaml already exists!\n",
      "\n"
     ]
    }
   ],
   "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": 4,
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 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"
     ]
    }
   ],
   "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",
    "    delta_ef = 1.e-4\n",
    "    d0 = np.array(workflow.queue[0].logfile.dipole)\n",
    "    for i, job in enumerate(workflow.queue[1:4]):\n",
    "        d1 = np.array(job.logfile.dipole)\n",
    "        pol_tensor[:, i] = (d1 - d0) / delta_ef\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": 6,
   "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)"
   ]
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}