{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# useful to autoreload the module without restarting the kernel\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mppi import InputFiles as I, Calculators as C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial for the QeCalculator class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial describes the usage of the QeCalculator class, that manages the run of a calculation with the QuantumESPRESSO package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_dir = 'QeCalculator_test_new'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform a scf computations for silicon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We init the PwInput object using an exsisting input file. Then we define an inputs with the associated name\n",
    "using the energy_cutoff as a parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_cutoffs = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'control': {'verbosity': \"'high'\",\n",
       "  'pseudo_dir': \"'../pseudos'\",\n",
       "  'calculation': \"'scf'\",\n",
       "  'prefix': \"'ecut_50'\"},\n",
       " 'system': {'force_symmorphic': '.true.',\n",
       "  'occupations': \"'fixed'\",\n",
       "  'ibrav': '2',\n",
       "  'celldm(1)': '10.3',\n",
       "  'ntyp': '1',\n",
       "  'nat': '2',\n",
       "  'ecutwfc': 50},\n",
       " 'electrons': {'conv_thr': '1e-08'},\n",
       " 'ions': {},\n",
       " 'cell': {},\n",
       " 'atomic_species': {'Si': ['28.086', 'Si.pbe-mt_fhi.UPF']},\n",
       " 'atomic_positions': {'type': 'crystal',\n",
       "  'values': [['Si', [0.125, 0.125, 0.125]], ['Si', [-0.125, -0.125, -0.125]]]},\n",
       " 'kpoints': {'type': 'automatic', 'values': ([6, 6, 6], [0.0, 0.0, 0.0])},\n",
       " 'cell_parameters': {},\n",
       " 'file': 'IO_files/si_scf.in'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inp = I.PwInput(file='IO_files/si_scf.in')\n",
    "prefix = 'ecut_%s'%energy_cutoffs\n",
    "inp.set_kpoints(points = [6,6,6])\n",
    "inp.set_prefix(prefix)\n",
    "inp.set_energy_cutoff(energy_cutoffs)\n",
    "name = prefix\n",
    "inp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we have chosen the value of the prefix of the input object as the name of the file. In this way the inp, log and xml file created by QuantumESPRESSO\n",
    "have the same name of the prefix folder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define an intance of the QeCalculator. For this example we use a direct scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mInit signature:\u001b[0m\n",
       "\u001b[0mC\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mQeCalculator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0momp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mmpi\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mmpi_run\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'mpirun -np'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mexecutable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'pw.x'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mscheduler\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'direct'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mskip\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mclean_restart\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m     \n",
       "Perform a QuantumESPRESSO calculation. Computations are managed by a scheduler that,\n",
       "in the actual implementation of the class, can be `direct` or `slurm`.\n",
       "\n",
       "Parameters:\n",
       "   omp (:py:class:`int`) : value of the OMP_NUM_THREADS variable\n",
       "   mpi (:py:class:`int`) : number of mpi processes\n",
       "   mpi_run (:py:class:`string`) : command for the execution of mpirun, e.g. 'mpirun -np' or 'mpiexec -np'\n",
       "   executable (:py:class:`string`) : set the executable (pw.x, ph.x, ..) of the QuantumESPRESSO package\n",
       "   scheduler (:py:class:`string`) : choose the scheduler used to submit the job, actually the choices implemented are\n",
       "        'direct' that runs the computation using the python subprocess package and 'slurm' that creates a slurm script\n",
       "   skip (:py:class:`bool`) : if True evaluate if the computations can be skipped. This is done by checking if the file\n",
       "        $prefix.xml is present in the run_dir folder\n",
       "   clean_restart (:py:class:`bool`) : if True delete the folder $prefix.save before running the computation\n",
       "   verbose (:py:class:`bool`) : set the amount of information provided on terminal\n",
       "   kwargs : other parameters that are stored in the _global_options dictionary\n",
       "\n",
       "Example:\n",
       " >>> code = calculator(omp=1,mpi=4,mpi_run='mpirun -np',skip=True,clean_restart=True,verbose=True,scheduler='direct')\n",
       " >>> code.run(input = ..., run_dir = ...,name = ..., source_dir = ..., **kwargs)\n",
       "\n",
       " where the arguments of the run method are:\n",
       "\n",
       "Args:\n",
       "    run_dir (:py:class:`string`) : the folder in which the simulation is performed\n",
       "    input (:py:class:`string`) : instance of the :class:`PwInput` class\n",
       "        that define the input object\n",
       "    name (:py:class:`string`) : string with the name associated to the input file.\n",
       "        Usually it is convenient to set the name equal to the prefix of the input object so\n",
       "        the name of the input file and the prefix folder built by QuantumESPRESSO are the same\n",
       "    source_dir (:py:class:`string`) : location of the scf source folder for a nscf computation.\n",
       "        If present the class copies this folder in the run_dir with the name $prefix.save\n",
       "    kwargs : other parameters that are stored in the run_options dictionary\n",
       "\n",
       "The calculator looks for the following variables in the run_options dictionary. These options\n",
       "may be useful for interacting with the code in particular in _asincronous_ computation managed\n",
       "the slurm scheduler.\n",
       "\n",
       "    `dry_run=True` with this option the calculator setup the calculations and write the script\n",
       "    for submitting the job, but the computations are not run\n",
       "\n",
       "    `wait_end_run=False` with this option the wait of the end of the run is suppressed\n",
       "\n",
       "    `sbatch_options = [option1,option2,....]` allows the user to include further options in the slurm script\n",
       "\u001b[0;31mFile:\u001b[0m           ~/Applications/MPPI/mppi/Calculators/QeCalculator.py\n",
       "\u001b[0;31mType:\u001b[0m           type\n",
       "\u001b[0;31mSubclasses:\u001b[0m     \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "C.QeCalculator?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize a QuantumESPRESSO calculator with scheduler direct\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'omp': 1,\n",
       " 'mpi': 2,\n",
       " 'mpi_run': 'mpirun -np',\n",
       " 'executable': 'pw.x',\n",
       " 'scheduler': 'direct',\n",
       " 'skip': True,\n",
       " 'clean_restart': True,\n",
       " 'verbose': True}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "code = C.QeCalculator(mpi=2)\n",
    "code.global_options()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We run the computation passing the the input object and the associated name to the run method of the \n",
    "calculator.\n",
    "\n",
    "The effects of various options of the calculator like, skip, clean_restart, dry_run, and wait_end_run can be checked by adding\n",
    "these options in the run methods. For instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "delete log file: QeCalculator_test_new/ecut_50.log\n",
      "delete xml file: QeCalculator_test_new/ecut_50.xml\n",
      "delete folder: QeCalculator_test_new/ecut_50.save\n",
      "run command: cd QeCalculator_test_new; mpirun -np 2 pw.x -inp ecut_50.in > ecut_50.log\n",
      "The wait_end_run is False or the dry_run option is active. The calculator proceedes to the postprocessing\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'QeCalculator_test_new/ecut_50.save/data-file-schema.xml'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = code.run(run_dir=run_dir,input=inp,name=name,other_variable = 1,skip=False, wait_end_run = False)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We perform a simple called to run that perform the computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run command: cd QeCalculator_test_new; mpirun -np 2 pw.x -inp ecut_50.in > ecut_50.log\n",
      "computation ecut_50 still running...\n",
      "computation ecut_50 ended\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'QeCalculator_test_new/ecut_50.save/data-file-schema.xml'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = code.run(run_dir=run_dir,input=inp,name=name,other_variable = 1)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We observe that the output of the run method is a string with the the data-file-schema.xml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the run all the parameters passed to the calculator are written in the run_options attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'omp': 1,\n",
       " 'mpi': 2,\n",
       " 'mpi_run': 'mpirun -np',\n",
       " 'executable': 'pw.x',\n",
       " 'scheduler': 'direct',\n",
       " 'skip': True,\n",
       " 'clean_restart': True,\n",
       " 'verbose': True,\n",
       " 'run_dir': 'QeCalculator_test_new',\n",
       " 'input': {'control': {'verbosity': \"'high'\",\n",
       "   'pseudo_dir': \"'../pseudos'\",\n",
       "   'calculation': \"'scf'\",\n",
       "   'prefix': \"'ecut_50'\"},\n",
       "  'system': {'force_symmorphic': '.true.',\n",
       "   'occupations': \"'fixed'\",\n",
       "   'ibrav': '2',\n",
       "   'celldm(1)': '10.3',\n",
       "   'ntyp': '1',\n",
       "   'nat': '2',\n",
       "   'ecutwfc': 50},\n",
       "  'electrons': {'conv_thr': '1e-08'},\n",
       "  'ions': {},\n",
       "  'cell': {},\n",
       "  'atomic_species': {'Si': ['28.086', 'Si.pbe-mt_fhi.UPF']},\n",
       "  'atomic_positions': {'type': 'crystal',\n",
       "   'values': [['Si', [0.125, 0.125, 0.125]],\n",
       "    ['Si', [-0.125, -0.125, -0.125]]]},\n",
       "  'kpoints': {'type': 'automatic', 'values': ([6, 6, 6], [0.0, 0.0, 0.0])},\n",
       "  'cell_parameters': {},\n",
       "  'file': 'IO_files/si_scf.in'},\n",
       " 'name': 'ecut_50',\n",
       " 'other_variable': 1,\n",
       " 'job': <subprocess.Popen at 0x7f7b3fad35b0>}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "code.run_options"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead, let see what happens if the simulation fails. For instance if we provide an empty input to code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp2 = I.PwInput()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'control': {'prefix': \"'si_scf_test2'\"},\n",
       " 'system': {},\n",
       " 'electrons': {},\n",
       " 'ions': {},\n",
       " 'cell': {},\n",
       " 'atomic_species': {},\n",
       " 'atomic_positions': {},\n",
       " 'kpoints': {},\n",
       " 'cell_parameters': {}}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pref2 = 'si_scf_test2'\n",
    "inp2.set_prefix(pref2)\n",
    "inp2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run command: cd QeCalculator_test_new; mpirun -np 2 pw.x -inp si_scf_test2.in > si_scf_test2.log\n",
      "computation si_scf_test2 still running...\n",
      "computation si_scf_test2 ended\n",
      "Expected file QeCalculator_test_new/si_scf_test2.save/data-file-schema.xml not Found\n",
      "Check if wait_end_run is False or the dry_run option is active. Otherwise a possible error has occured during the computation\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'QeCalculator_test_new/si_scf_test2.save/data-file-schema.xml'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result2 = code.run(input = inp2, run_dir = run_dir,name=pref2) \n",
    "result2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case the calculator provide a warning since the data-file-schema has not been found after the computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Usage of the skip parameter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we repeat a calculation that has been already performed and skip = True the class skip its computation, for instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skip the run of ecut_50\n"
     ]
    }
   ],
   "source": [
    "result = code.run(run_dir=run_dir,input=inp,name=name, skip = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test of the slurm scheduler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the `slurm` scheduler is chosen the calculator prepare the slurm script and submit it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "delete log file: QeCalculator_test_new/ecut_50.log\n",
      "delete xml file: QeCalculator_test_new/ecut_50.xml\n",
      "delete folder: QeCalculator_test_new/ecut_50.save\n",
      "run command: mpirun -np 2 pw.x -inp ecut_50.in > ecut_50.log\n",
      "Dry_run option active. Script not submitted\n",
      "The wait_end_run is False or the dry_run option is active. The calculator proceedes to the postprocessing\n",
      "Expected file QeCalculator_test_new/ecut_50.save/data-file-schema.xml not Found\n",
      "Check if wait_end_run is False or the dry_run option is active. Otherwise a possible error has occured during the computation\n"
     ]
    }
   ],
   "source": [
    "result = code.run(run_dir=run_dir,input=inp,name=name, scheduler = 'slurm', skip = False, dry_run = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The slurm script is written in the run_dir. The execution of the run requires that the slurm scheduler is installed. However with the\n",
    "dry_run option we can write the script on disk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform a nscf computation for silicon. Usage of the source_dir option"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We show how to perform a pw nscf calculation using the results of the first scf run as an input.\n",
    "\n",
    "Before running this computation ensure that the scf one has been computed and that the save folder is written in the run_dir."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_bands = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp.set_nscf(num_bands,force_symmorphic=True)\n",
    "prefix = 'bands_%s'%num_bands\n",
    "inp.set_prefix(prefix)\n",
    "inp.set_energy_cutoff(50)\n",
    "name = prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copy source_dir QeCalculator_test_new/ecut_50.save in the QeCalculator_test_new/bands_8.save\n",
      "run command: cd QeCalculator_test_new; mpirun -np 2 pw.x -inp bands_8.in > bands_8.log\n",
      "computation still running...\n",
      "computation ended\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'QeCalculator_test_new/bands_8.save/data-file-schema.xml'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "result = code.run(input=inp,run_dir=run_dir,name=name,source_dir=os.path.join(run_dir,'ecut_50.save'))\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead, if skip = False the class delete the existing output files before running the computation again. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "delete log file: QeCalculator_test_new/bands_8.log\n",
      "delete xml file: QeCalculator_test_new/bands_8.xml\n",
      "delete folder: QeCalculator_test_new/bands_8.save\n",
      "Copy source_dir QeCalculator_test_new/ecut_50.save in the QeCalculator_test_new/bands_8.save\n",
      "run command: cd QeCalculator_test_new; mpirun -np 2 pw.x -inp bands_8.in > bands_8.log\n",
      "computation still running...\n",
      "computation ended\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'QeCalculator_test_new/bands_8.save/data-file-schema.xml'"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = code.run(input=inp,run_dir=run_dir,name=name,source_dir=os.path.join(run_dir,'ecut_50.save'),skip=False)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}