{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "22d177dc-6cfb-4de2-9509-f1eb45e10cf2",
   "metadata": {},
   "source": [
    "# X-ray photoelectron spectroscopy (XPS)\n",
    "For the source code, see [xps](https://workgraph-collections.readthedocs.io/en/latest/ase/espresso/module.html#workgraph_collections.ase.espresso.xps.xps_workgraph).\n",
    "\n",
    "\n",
    "## Example: ETFA molecule\n",
    "\n",
    "### Prepare the inputs and submit the workflow\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ee799d2-0b5b-4609-957f-6b3f2cd451f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WorkGraph process created, PK: 35068\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<WorkChainNode: uuid: 71341155-996d-4cb3-85a8-9d0caa328b1a (pk: 35068) (aiida.workflows:workgraph.engine)>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ase.io import read\n",
    "from aiida import load_profile\n",
    "from workgraph_collections.ase.espresso.xps import XpsWorkgraph\n",
    "\n",
    "load_profile()\n",
    "\n",
    "# create input structure\n",
    "atoms = read(\"../datas/ETFA.xyz\")\n",
    "atoms.center(vacuum=5.0)\n",
    "input_data = {\n",
    "        \"CONTROL\": {\n",
    "            \"calculation\": \"scf\",\n",
    "        },\n",
    "        \"SYSTEM\": {\n",
    "            \"ecutwfc\": 50,\n",
    "            \"ecutrho\": 400,\n",
    "            \"occupations\": \"fixed\",\n",
    "        },\n",
    "    }\n",
    "kpts = (1, 1, 1)\n",
    "# Pseudos from https://github.com/superstar54/xps-data/tree/main/pseudo_demo/pseudo_demo_pbe\n",
    "pseudopotentials = {\n",
    "    \"H\": \"H.pbe-kjpaw_psl.1.0.0.UPF\",\n",
    "    \"F\": \"F.pbe-n-rrkjus_psl.1.0.0.UPF\",\n",
    "}\n",
    "# corrections from https://github.com/superstar54/xps-data/blob/main/pseudo_demo/pseudo_demo_pbe/datas.py\n",
    "core_hole_pseudos = {\n",
    "    \"C\": {\"ground\": \"C.pbe-n-kjgipaw_psl.1.0.0.UPF\",\n",
    "          \"core_hole\": \"C.star1s.pbe-n-kjgipaw_psl.1.0.0.UPF\",\n",
    "          \"correction\": 345.99 - 6.2,\n",
    "    },\n",
    "    \"O\": {\"ground\": \"O.pbe-n-kjpaw_psl.0.1.UPF\",\n",
    "          \"core_hole\": \"O.star1s.pbe-n-kjpaw_psl.0.1.UPF\",\n",
    "          \"correction\": 676.47 - 8.25,\n",
    "    }\n",
    "}\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",
    "scf_inputs = {\n",
    "    \"command\": \"mpirun -np 4 pw.x\",\n",
    "    \"computer\": \"localhost\",\n",
    "    \"metadata\": metadata,\n",
    "    \"input_data\": input_data,\n",
    "    \"kpts\": kpts,\n",
    "    \"pseudopotentials\": pseudopotentials,\n",
    "    \"pseudo_dir\": \"/home/wang_x3/datas/pseudos/xps/pbe\",\n",
    "    \"core_hole_treatment\": \"xch\",\n",
    "}\n",
    "\n",
    "# ===============================================================================\n",
    "wg = XpsWorkgraph.build_graph(\n",
    "    atoms=atoms,\n",
    "    marked_structures_inputs={\"atom_list\": [0, 1, 2, 3],\n",
    "                              \"min_cell_length\": 14,\n",
    "                              \"is_molecule\": True,\n",
    "                              },\n",
    "    scf_inputs=scf_inputs,\n",
    "    core_hole_pseudos=core_hole_pseudos,\n",
    "    metadata=metadata,\n",
    ")\n",
    "wg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b934e777",
   "metadata": {},
   "source": [
    "Run the workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a78771f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "wg.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a533e6c8",
   "metadata": {},
   "source": [
    "## Print the results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5a25d691",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C_0: 298.32 eV\n",
      "C_1: 295.27 eV\n",
      "C_2: 293.01 eV\n",
      "C_3: 291.45 eV\n"
     ]
    }
   ],
   "source": [
    "for key, energy in wg.tasks[\"get_binding_energy\"].outputs[\"result\"].value.value.items():\n",
    "    print(f\"{key}: {energy:.2f} eV\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e65640ce",
   "metadata": {},
   "source": [
    "### Run on HPC\n",
    "Here is an example of how to run the pw.x calculation on the CSCS eiger cluster.\n",
    "\n",
    "\n",
    "```python \n",
    "# run on CSCS eiger cluster\n",
    "metadata_eiger = {\n",
    "    \"options\": {\n",
    "        'prepend_text' : \"\"\"\n",
    "module load cray/22.05  cpeIntel/22.05 QuantumESPRESSO/7.0\n",
    "eval \"$(/users/xingwang/miniconda3/bin/conda shell.posix hook)\"\n",
    "conda activate py3.11\n",
    "export OMP_NUM_THREADS=1\n",
    "        \"\"\",\n",
    "        'custom_scheduler_commands' : '#SBATCH --account=mr32',\n",
    "        'resources': {\n",
    "                    'num_machines' : 1,\n",
    "                    'num_mpiprocs_per_machine' : 128,\n",
    "        }\n",
    "    }\n",
    "}\n",
    "scf_inputs = {\n",
    "    \"command\": \"srun -n 128 pw.x\",\n",
    "    \"computer\": \"eiger\",\n",
    "    \"metadata\": metadata_eiger,\n",
    "    \"input_data\": input_data,\n",
    "    \"kpts\": kpts,\n",
    "    \"pseudopotentials\": pseudopotentials,\n",
    "    # \"pseudo_dir\": \"/home/wang_x3/datas/pseudos/xps/pbe\",\n",
    "    \"pseudo_dir\": \"/users/xingwang/datas/pseudos/xps/pbe\",\n",
    "    \"core_hole_pseudos\": core_hole_pseudos,\n",
    "    \"is_molecule\": True,\n",
    "}\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "2f450c1ff08798c4974437dd057310afef0de414c25d1fd960ad375311c3f6ff"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}