{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "22d177dc-6cfb-4de2-9509-f1eb45e10cf2",
   "metadata": {},
   "source": [
    "# Equation of state (EOS)\n",
    "\n",
    "For the source code, see [eos](https://workgraph-collections.readthedocs.io/en/latest/ase/espresso/module.html#workgraph_collections.ase.espresso.eos.eos_workgraph).\n",
    "\n",
    "## Example: Silicon\n",
    "\n",
    "### Prepare the inputs and submit the workflow\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8ee799d2-0b5b-4609-957f-6b3f2cd451f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"html/EosWorkGraph.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7e21b28d2750>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ase.build import bulk\n",
    "from ase import Atoms\n",
    "from aiida import load_profile\n",
    "from workgraph_collections.ase.espresso.eos import EosWorkGraph\n",
    "from aiida_pythonjob.config import config\n",
    "\n",
    "load_profile()\n",
    "\n",
    "config[\"allow_pickle\"] = True\n",
    "\n",
    "# create input structure\n",
    "atoms = bulk(\"Si\")\n",
    "\n",
    "metadata = {\n",
    "    \"options\": {\n",
    "        'prepend_text' : \"\"\"eval \"$(conda shell.posix hook)\"\n",
    "        conda activate aiida\n",
    "        export OMP_NUM_THREADS=1\n",
    "        \"\"\",\n",
    "    }\n",
    "}\n",
    "pseudopotentials = {\"Si\": \"Si.pbe-nl-rrkjus_psl.1.0.0.UPF\"}\n",
    "pseudo_dir = \"/home/xing/data/ase/espresso_pseudo\"\n",
    "# pseudo_dir = \"/home/wang_x3/datas/pseudos/psl\"\n",
    "input_data = {\n",
    "    \"control\": {\"calculation\": \"scf\"},\n",
    "    \"system\": {\"ecutwfc\": 30, \"ecutrho\": 240,\n",
    "               \"occupations\": \"smearing\",\n",
    "                \"degauss\": 0.01,\n",
    "                \"smearing\": \"cold\",},\n",
    "}\n",
    "#------------------------- Set the inputs -------------------------\n",
    "wg = EosWorkGraph.build_graph(atoms=atoms,\n",
    "                   computer=\"localhost\",\n",
    "                   scales=[0.95, 0.98, 1.0, 1.02, 1.05],\n",
    "                   command=\"mpirun -np 2 pw.x\",\n",
    "                   pseudopotentials=pseudopotentials,\n",
    "                   pseudo_dir=pseudo_dir,\n",
    "                   input_data=input_data,\n",
    "                   kpts=[4, 4, 4],\n",
    "                   metadata=metadata)\n",
    "wg.to_html()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c213904",
   "metadata": {},
   "source": [
    "Run the calculation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d1e32cd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "09/03/2025 06:13:07 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "09/03/2025 06:13:08 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178210\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178210\n",
      "09/03/2025 06:13:30 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "09/03/2025 06:13:32 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: generate_scaled_atoms\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: generate_scaled_atoms\n",
      "INFO:aiida_pythonjob.decorator:Executing process function, current stack status: 47 frames of 3000\n",
      "09/03/2025 06:13:33 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|update_task_state]: Task: generate_scaled_atoms, type: PYFUNCTION, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|update_task_state]: Task: generate_scaled_atoms, type: PYFUNCTION, finished.\n",
      "09/03/2025 06:13:34 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: calc_all_structures\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: calc_all_structures\n",
      "09/03/2025 06:13:42 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178253\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178253\n",
      "09/03/2025 06:13:49 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator,pw_calculator1,pw_calculator2,pw_calculator3,pw_calculator4\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator,pw_calculator1,pw_calculator2,pw_calculator3,pw_calculator4\n",
      "09/03/2025 06:13:58 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178257, 178261, 178265, 178269, 178273\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178257, 178261, 178265, 178269, 178273\n",
      "09/03/2025 06:14:25 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "09/03/2025 06:14:31 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "09/03/2025 06:14:31 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178261, 178265, 178269, 178273\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 178261, 178265, 178269, 178273\n",
      "09/03/2025 06:14:43 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|update_task_state]: Task: pw_calculator1, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|update_task_state]: Task: pw_calculator1, type: PYTHONJOB, finished.\n",
      "09/03/2025 06:14:44 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|update_task_state]: Task: pw_calculator2, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|update_task_state]: Task: pw_calculator2, type: PYTHONJOB, finished.\n",
      "09/03/2025 06:14:45 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|update_task_state]: Task: pw_calculator3, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|update_task_state]: Task: pw_calculator3, type: PYTHONJOB, finished.\n",
      "09/03/2025 06:14:45 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|update_task_state]: Task: pw_calculator4, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|update_task_state]: Task: pw_calculator4, type: PYTHONJOB, finished.\n",
      "09/03/2025 06:14:51 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "09/03/2025 06:14:52 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178253|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178253|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "09/03/2025 06:14:53 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|update_task_state]: Task: calc_all_structures, type: GRAPH, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|update_task_state]: Task: calc_all_structures, type: GRAPH, finished.\n",
      "09/03/2025 06:14:56 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: fit_eos\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: fit_eos\n",
      "INFO:aiida_pythonjob.decorator:Executing process function, current stack status: 47 frames of 3000\n",
      "09/03/2025 06:14:57 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|update_task_state]: Task: fit_eos, type: PYFUNCTION, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|update_task_state]: Task: fit_eos, type: PYFUNCTION, finished.\n",
      "09/03/2025 06:14:57 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "09/03/2025 06:14:57 PM <77904> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [178206|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[178206|WorkGraphEngine|finalize]: Finalize workgraph.\n"
     ]
    }
   ],
   "source": [
    "wg.run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6d4e61",
   "metadata": {},
   "source": [
    "Print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4aae567b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Result: \n",
      "B: 88.541858293246\n",
      "v0: 41.224378688688\n",
      "e0: -277.89646661704\n",
      "v0: 41.224378688688\n"
     ]
    }
   ],
   "source": [
    "#------------------------- Print the output -------------------------\n",
    "print('\\nResult: \\nB: {B}\\nv0: {v0}\\ne0: {e0}\\nv0: {v0}'.format(**wg.outputs.result.value))\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": 5
}