{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ABINIT DOS Analysis\n",
    "\n",
    "This notebook parses ABINIT DOS files to compare the $d$-orbital ($l=2$) electronic density of states for two different atoms.\n",
    "\n",
    "## Files\n",
    "- Atom 1: `tspin_3o_DS2_DOS_AT0001`\n",
    "- Atom 2: `tspin_3o_DS3_DOS_AT0002`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_abinit_dos(filepath):\n",
    "    \"\"\"\n",
    "    Parses an ABINIT DOS file to extract Fermi energy and DOS data.\n",
    "    Returns fermi_energy, spin_up_data, spin_dn_data.\n",
    "    Data is returned as a list of dictionaries or a pandas DataFrame-like structure (numpy array here).\n",
    "    \"\"\"\n",
    "    fermi_energy = 0.0\n",
    "    spin_up_lines = []\n",
    "    spin_dn_lines = []\n",
    "    \n",
    "    current_section = None # 'up' or 'dn'\n",
    "    \n",
    "    with open(filepath, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "        \n",
    "    for line in lines:\n",
    "        if \"Fermi energy\" in line:\n",
    "            # Format: # Fermi energy :       0.52895276\n",
    "            parts = line.split()\n",
    "            if ':' in line:\n",
    "                 # Find the element after ':'\n",
    "                idx = parts.index(':')\n",
    "                fermi_energy = float(parts[idx+1])\n",
    "            else:\n",
    "                 # Sometimes it might just be space separated? check format\n",
    "                 # Based on file view: # Fermi energy :       0.52895276\n",
    "                 fermi_energy = float(parts[-1])\n",
    "\n",
    "        if \"Spin-up DOS\" in line:\n",
    "            current_section = 'up'\n",
    "            continue\n",
    "        if \"Spin-dn DOS\" in line:\n",
    "            current_section = 'dn'\n",
    "            continue\n",
    "            \n",
    "        if line.strip().startswith('#'):\n",
    "            continue\n",
    "            \n",
    "        if not line.strip():\n",
    "            continue\n",
    "            \n",
    "        # Data lines\n",
    "        # Format: energy(Ha) l=0 l=1 l=2 l=3 l=4 ...\n",
    "        try:\n",
    "            vals = [float(x) for x in line.split()]\n",
    "            if current_section == 'up':\n",
    "                spin_up_lines.append(vals)\n",
    "            elif current_section == 'dn':\n",
    "                spin_dn_lines.append(vals)\n",
    "        except ValueError:\n",
    "            continue\n",
    "            \n",
    "    return fermi_energy, np.array(spin_up_lines), np.array(spin_dn_lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fermi Energy (File 1): 0.52895276 Ha\n",
      "Fermi Energy (File 2): 0.52895276 Ha\n",
      "Data shape 1 (up/dn): (7201, 11), (7201, 11)\n",
      "Data shape 2 (up/dn): (7201, 11), (7201, 11)\n"
     ]
    }
   ],
   "source": [
    "# File paths\n",
    "file1 = 'tspin_3o_DS2_DOS_AT0001'\n",
    "file2 = 'tspin_3o_DS3_DOS_AT0002'\n",
    "\n",
    "ef1, up1, dn1 = parse_abinit_dos(file1)\n",
    "ef2, up2, dn2 = parse_abinit_dos(file2)\n",
    "\n",
    "print(f\"Fermi Energy (File 1): {ef1} Ha\")\n",
    "print(f\"Fermi Energy (File 2): {ef2} Ha\")\n",
    "\n",
    "print(f\"Data shape 1 (up/dn): {up1.shape}, {dn1.shape}\")\n",
    "print(f\"Data shape 2 (up/dn): {up2.shape}, {dn2.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 66740 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Setup Plotting\n",
    "l_channel = 3 # Column index 3 corresponds to l=2 (0: Energy, 1: l=0, 2: l=1, 3: l=2)\n",
    "\n",
    "energy1 = up1[:, 0] - ef1\n",
    "dos_up1 = up1[:, l_channel]\n",
    "dos_dn1 = dn1[:, l_channel]\n",
    "\n",
    "energy2 = up2[:, 0] - ef2\n",
    "dos_up2 = up2[:, l_channel]\n",
    "dos_dn2 = dn2[:, l_channel]\n",
    "\n",
    "# Plot\n",
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "# Atom 1\n",
    "plt.plot(energy1, dos_up1, label='AT1 Spin-up ($d$)', color='blue')\n",
    "plt.plot(energy1, -dos_dn1, label='AT1 Spin-down ($d$)', color='blue', linestyle='--')\n",
    "\n",
    "# Atom 2\n",
    "plt.plot(energy2, dos_up2, label='AT2 Spin-up ($d$)', color='red')\n",
    "plt.plot(energy2, -dos_dn2, label='AT2 Spin-down ($d$)', color='red', linestyle='--')\n",
    "\n",
    "plt.axvline(0, color='black', linewidth=0.8, linestyle='-.', label='$E_F$')\n",
    "plt.axhline(0, color='black', linewidth=0.5)\n",
    "\n",
    "plt.xlabel('Energy (Ha) - $E_F$')\n",
    "plt.ylabel('DOS (electrons/Ha)')\n",
    "plt.title('Comparison of $d$-orbital DOS: AT1 vs AT2')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.xlim(-0.5, 0.5)  # Adjust range as needed\n",
    "\n",
    "plt.savefig('tspin_3_comparison.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dft",
   "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
