{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Equation of state (EOS)\n",
    "\n",
    "## Visualizing the WorkGraph Builder\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"html/eos_workgraph0.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x77bce36fba10>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from workgraph_collections.cp2k.eos import eos_workgraph\n",
    "\n",
    "\n",
    "task = eos_workgraph.TaskCls()\n",
    "task.to_html()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the WorkGraph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"html/EOS.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x77bce12e3850>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from workgraph_collections.cp2k.eos import eos_workgraph\n",
    "from aiida import load_profile\n",
    "\n",
    "load_profile()\n",
    "\n",
    "wg = eos_workgraph(\"eos\")\n",
    "wg.to_html()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare inputs and run\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aiida import load_profile, orm\n",
    "from ase.build import bulk\n",
    "from workgraph_collections.cp2k.eos import eos_workgraph\n",
    "\n",
    "\n",
    "load_profile()\n",
    "\n",
    "atoms = bulk(\"Si\")\n",
    "structure = orm.StructureData(ase=atoms)\n",
    "code = orm.load_code(\"cp2k-ssmp-2024.1@localhost\")\n",
    "# Parameters.\n",
    "parameters = {\n",
    "    \"GLOBAL\": {\n",
    "        \"RUN_TYPE\": \"ENERGY_FORCE\",\n",
    "    },\n",
    "    \"FORCE_EVAL\": {\n",
    "        \"METHOD\": \"Quickstep\",\n",
    "        \"DFT\": {\n",
    "            \"BASIS_SET_FILE_NAME\": \"BASIS_MOLOPT\",\n",
    "            \"POTENTIAL_FILE_NAME\": \"POTENTIALS\",\n",
    "            \"SCF\": {\n",
    "                \"ADDED_MOS\": 10,\n",
    "                \"SMEAR\": {\n",
    "                    \"METHOD\": \"FERMI_DIRAC\",\n",
    "                    \"ELECTRONIC_TEMPERATURE\": 500,\n",
    "                },\n",
    "            },\n",
    "            \"KPOINTS\": {\n",
    "                \"SCHEME\": \"MONKHORST-PACK 5 5 5\",\n",
    "            },\n",
    "            \"QS\": {\n",
    "                \"EPS_DEFAULT\": 1.0e-12,\n",
    "                \"METHOD\": \"GPW\",\n",
    "            },\n",
    "            \"MGRID\": {\n",
    "                \"NGRIDS\": 4,\n",
    "                \"CUTOFF\": 500,\n",
    "                \"REL_CUTOFF\": 50,\n",
    "            },\n",
    "            \"XC\": {\n",
    "                \"XC_FUNCTIONAL\": {\n",
    "                    \"_\": \"PBE\",\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "        \"SUBSYS\": {\n",
    "            \"KIND\": [\n",
    "                {\n",
    "                    \"_\": \"Si\",\n",
    "                    \"BASIS_SET\": \"DZVP-MOLOPT-GTH\",\n",
    "                    \"POTENTIAL\": \"GTH-PBE\",\n",
    "                },\n",
    "            ],\n",
    "        },\n",
    "    },\n",
    "}\n",
    "metadata = {\n",
    "    \"options\": {\n",
    "        \"resources\": {\n",
    "            \"num_machines\": 1,\n",
    "            \"num_mpiprocs_per_machine\": 1,\n",
    "        },\n",
    "    }\n",
    "}\n",
    "basis_pseudo_files = {\n",
    "    \"basis\": orm.load_node(9041),\n",
    "    \"pseudo\": orm.load_node(9042),\n",
    "}\n",
    "# ===============================================================================\n",
    "wg = eos_workgraph(\n",
    "    structure=structure,\n",
    "    code=code,\n",
    "    scales=[0.98, 0.99, 1.0, 1.01, 1.02],\n",
    "    parameters=parameters,\n",
    "    basis_pseudo_files=basis_pseudo_files,\n",
    "    metadata=metadata,\n",
    ")\n",
    "wg.name = \"CP2K-EOS-Si\"\n",
    "#------------------------- Submit the calculation -------------------\n",
    "wg.submit(wait=True, timeout=300)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print out the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#------------------------- Print the output -------------------------\n",
    "data = wg.tasks[\"fit_eos\"].outputs[\"result\"].value.get_dict()\n",
    "print('\\nResult: \\nB: {B}\\nv0: {v0}\\ne0: {e0}\\nv0: {v0}'.format(**data))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aiida",
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}