{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9fd2308b",
   "metadata": {},
   "source": [
    "# Relaxation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af8d6f96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"html/RelaxWorkgraph.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x74cc561d2f10>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from workgraph_collections.ase.espresso.relax import RelaxWorkgraph\n",
    "from ase.build import bulk\n",
    "from ase import Atoms\n",
    "from aiida import load_profile\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",
    "atoms.set_cell(atoms.cell * 1.05, scale_atoms=True)\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 = RelaxWorkgraph.build_graph(\n",
    "    atoms=atoms,\n",
    "    computer=\"localhost\",\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",
    ")\n",
    "wg.to_html()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c06ac3a",
   "metadata": {},
   "source": [
    "## Run the workflow\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "defb31bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "09/04/2025 07:30:18 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "09/04/2025 07:30:19 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180113\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180113\n",
      "09/04/2025 07:30:41 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "09/04/2025 07:30:42 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: inspect_relax\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: inspect_relax\n",
      "INFO:aiida_pythonjob.decorator:Executing process function, current stack status: 47 frames of 3000\n",
      "09/04/2025 07:30:43 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|update_task_state]: Task: inspect_relax, type: PYFUNCTION, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|update_task_state]: Task: inspect_relax, type: PYFUNCTION, finished.\n",
      "09/04/2025 07:30:43 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: RelaxWorkgraph\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: RelaxWorkgraph\n",
      "09/04/2025 07:30:46 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180139\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180139\n",
      "09/04/2025 07:30:50 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "09/04/2025 07:30:50 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180143\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180143\n",
      "09/04/2025 07:31:02 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "09/04/2025 07:31:04 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: inspect_relax\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: inspect_relax\n",
      "INFO:aiida_pythonjob.decorator:Executing process function, current stack status: 47 frames of 3000\n",
      "09/04/2025 07:31:05 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|update_task_state]: Task: inspect_relax, type: PYFUNCTION, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|update_task_state]: Task: inspect_relax, type: PYFUNCTION, finished.\n",
      "09/04/2025 07:31:05 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: RelaxWorkgraph\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: RelaxWorkgraph\n",
      "09/04/2025 07:31:08 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180168\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180168\n",
      "09/04/2025 07:31:10 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180168|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180168|WorkGraphEngine|continue_workgraph]: tasks ready to run: pw_calculator\n",
      "09/04/2025 07:31:11 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180168|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180172\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180168|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 180172\n",
      "09/04/2025 07:31:20 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180168|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180168|WorkGraphEngine|update_task_state]: Task: pw_calculator, type: PYTHONJOB, finished.\n",
      "09/04/2025 07:31:21 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180168|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180168|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "09/04/2025 07:31:21 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180168|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180168|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "09/04/2025 07:31:22 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|update_task_state]: Task: RelaxWorkgraph, type: GRAPH, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|update_task_state]: Task: RelaxWorkgraph, type: GRAPH, finished.\n",
      "09/04/2025 07:31:24 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "09/04/2025 07:31:24 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180139|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180139|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "09/04/2025 07:31:24 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|update_task_state]: Task: RelaxWorkgraph, type: GRAPH, finished.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|update_task_state]: Task: RelaxWorkgraph, type: GRAPH, finished.\n",
      "09/04/2025 07:31:26 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n",
      "09/04/2025 07:31:26 AM <593221> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [180109|WorkGraphEngine|finalize]: Finalize workgraph.\n",
      "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[180109|WorkGraphEngine|finalize]: Finalize workgraph.\n"
     ]
    }
   ],
   "source": [
    "wg.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ab3ef8",
   "metadata": {},
   "source": [
    "## Print the outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c2dc5c49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relaxed structure:\n",
      "Atoms(symbols='Si2', pbc=True, cell=[[0.0, 2.7391864002289, 2.7391864002289], [2.7391864002289, 0.0, 2.7391864002289], [2.7391864002289, 2.7391864002289, 0.0]])\n"
     ]
    }
   ],
   "source": [
    "print(\"Relaxed structure:\")\n",
    "print(wg.outputs.atoms.value.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
}