{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import re\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.constants import hbar, e as eV, c\n",
    "eh = eV/hbar\n",
    "import glob\n",
    "def ri(z): return (z.real, z.imag)\n",
    "#m = re.compile(r\"([^_]+)_r([0-9.]+)nm_\")\n",
    "#removek = re.compile(r\"(k\\([^)]+\\)um-1_)\")\n",
    "remover = re.compile(r\"r[0-9.]+nm_\")\n",
    "\n",
    "\n",
    "markerdict = {\n",
    "    4: \"3\",\n",
    "    -4: \"4\",\n",
    "    3: \"^\",\n",
    "    -3: \"v\",\n",
    "    -2: 'x',\n",
    "    2: '+',\n",
    "    1: 's',\n",
    "    -1: 'd',\n",
    "}\n",
    "\n",
    "\n",
    "D4h_irlabels = {\"B2''\":\"$B_2''$\",\n",
    "                \"B2'\":\"$B_2'$\",\n",
    "                \"A1''\":\"$A_1''$\",\n",
    "                \"A1'\":\"$A_1'$\",\n",
    "                \"A2''\":\"$A_2''$\",\n",
    "                \"B1''\":\"$B_1''$\",\n",
    "                \"A2'\":\"$A_2'$\", \n",
    "                \"B1'\":\"$B_1'$\",\n",
    "                \"E'\":\"$E'$\",\n",
    "                \"E''\":\"$E''$\",}\n",
    "\n",
    "\n",
    "prop_cycle = plt.rcParams['axes.prop_cycle']\n",
    "colors = prop_cycle.by_key()['color']\n",
    "colordict = {i: colors[(i+1)] for i in range(-4,8)}\n",
    "\n",
    "def markerfun(b):\n",
    "    if b in markerdict.keys():\n",
    "        return markerdict[b]\n",
    "    else: return 'X'\n",
    "\n",
    "def colorfun(b):\n",
    "    if (b+1) in colordict.keys():\n",
    "        return colordict[b+1]\n",
    "    else: return colordict[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sph_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L1_cn410.npz 203\n",
      "sph_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L2_cn410.npz 251\n",
      "sph_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L3_cn410.npz 251\n",
      "sph_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L4_cn410.npz 251\n",
      "sph_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L5_cn410.npz 251\n"
     ]
    }
   ],
   "source": [
    "allfiles=glob.glob('*sph*k(0_0)*.npz')\n",
    "allgraphs=dict()\n",
    "for f in allfiles:\n",
    "    base = remover.sub('', f)\n",
    "    if base in allgraphs.keys():\n",
    "        allgraphs[base] += 1\n",
    "    else:\n",
    "        allgraphs[base] = 1\n",
    "for k in sorted(allgraphs.keys()):\n",
    "    print(k, allgraphs[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['sph_r281nm_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L4_cn410.npz', 'sph_r121nm_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L3_cn410.npz', 'sph_r67nm_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L3_cn410.npz', 'sph_r208nm_p580nmx580nm_mAu_n1.52_b-2_k(0_0)um-1_L3_cn410.npz'] collected_XXXsph_rXXX_p580nmx580nm_mAu_n1.52_bX2_k(0_0)um-1_LX_cnXXX.pdf\n"
     ]
    }
   ],
   "source": [
    "projectors = dict()\n",
    "projectors_list = dict()\n",
    "lMaxes = [lMax for lMax in range(1,6)]\n",
    "for lMax in lMaxes:\n",
    "    proj = np.load('projectors_D4h_lMax%d.npz' % lMax)\n",
    "    irlabels = sorted(proj.keys())\n",
    "    proj =  {f: proj[f] for f in irlabels}\n",
    "    proj_list = [proj[irlabels[i]] for i in range(len(proj))]\n",
    "    projectors[lMax] = proj\n",
    "    projectors_list[lMax] = proj_list\n",
    "globpattern = '*sph_r*_p580nmx580nm_mAu_n1.52_b?2_k(0_0)um-1_L?_cn???.npz'\n",
    "filenames=glob.glob(globpattern)\n",
    "plotfilename = 'collected_' + globpattern.replace('*', 'XXX').replace('?', 'X').replace('npz','pdf')\n",
    "print(filenames[:4], plotfilename)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#projectors\n",
    "#glob.glob('cyl_r100nm*L3*3100.npz')\n",
    "#glob.glob('sph_r100*m5*.npz')\n",
    "#dat['meta'][()],list(dat.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inpure result detected [1.         1.         0.92521572 1.         0.99999627 0.99990293\n",
      " 0.99946049] [1.59712906e-05 3.60193407e-05 2.48341492e-01 1.21848930e-03\n",
      " 3.81805601e-03 2.42649228e-02 2.99534246e-02]\n",
      "inpure result detected [1.         1.         0.99999998 0.99999999 0.99999996 0.96608887\n",
      " 0.99999852 0.99999397 0.99998951 0.99999912 0.99982435] [2.66223026e-04 2.12357147e-05 3.54211968e-05 1.06651057e-04\n",
      " 2.79595790e-04 2.41939163e-01 2.17645058e-03 3.41541473e-03\n",
      " 1.14507609e-02 1.49639498e-02 2.33483138e-02]\n",
      "inpure result detected [1.         0.99999999 1.         0.97991334 0.99999996 0.9999989\n",
      " 0.99999983 0.99999966 0.99999322 0.99999721 0.99999653] [3.28735741e-04 2.66532534e-05 2.47011478e-05 1.45012420e-01\n",
      " 2.44785416e-04 7.05405359e-04 1.60203586e-03 1.71245137e-03\n",
      " 1.03244480e-02 9.18732728e-03 1.18651583e-02]\n",
      "inpure result detected [1.         1.         0.99999998 0.99999961 0.93267685 0.99999964\n",
      " 0.99999822 0.99921774 0.99995547 0.99997301] [5.22490396e-04 3.01556792e-05 4.88795563e-05 6.29703960e-04\n",
      " 2.34414238e-01 3.72766210e-03 4.72444059e-03 7.62106094e-02\n",
      " 6.32796684e-02 5.63231562e-02]\n"
     ]
    }
   ],
   "source": [
    "plotdata = {}\n",
    "for file in filenames:\n",
    "    dat = np.load(file, allow_pickle=True)\n",
    "    kx = dat['meta'][()]['k'][0]\n",
    "    radius = dat['meta'][()]['radius']\n",
    "    b = dat['meta'][()]['band_index']\n",
    "    eigvals = dat['eigval']\n",
    "    lMax = dat['meta'][()]['lMax']\n",
    "    residuals = dat['residuals']\n",
    "    ef =dat['empty_freqs']\n",
    "    eigvecs = dat['eigvec']\n",
    "    irweights = []\n",
    "    #for proj in projectors_list[lMax]:\n",
    "    #    try:\n",
    "    #        irweights.append(np.linalg.norm(np.tensordot(proj, eigvecs, axes=(-1, -1)), axis=0,ord=2) if len(proj) != 0 else np.zeros((len(eigvecs),)))\n",
    "    #    except ValueError as err:\n",
    "    #        print(proj, len(proj))\n",
    "    #        raise err\n",
    "    irweights = np.array(irweights)\n",
    "    #print(irweights)\n",
    "    irweights = np.array([np.linalg.norm(np.tensordot(proj, eigvecs, axes=(-1, -1)), axis=0,ord=2) if len(proj) != 0 else np.zeros((len(eigvecs),)) for proj in projectors_list[lMax]]).T\n",
    "    irclass = np.argmax(irweights, axis=-1)\n",
    "    purities = np.amax(irweights, axis=-1)\n",
    "    if (np.any(purities < 0.98)):\n",
    "        print(\"inpure result detected\", purities, residuals)\n",
    "    #print(purities)\n",
    "    \n",
    "    #for i in range(len(residuals)): \n",
    "    #    if residuals[i] < 0.01:\n",
    "    #        vec = eigvecs[i]\n",
    "    #        for irlabel, proj in projectors.items():\n",
    "    #            print(irlabel, np.linalg.norm(np.dot(proj, vec))) #maybe some conj() here?\n",
    "    #        print('--->', irlabels[irclass[i]])\n",
    "\n",
    "    \n",
    "    plotdata[(lMax,radius)] = (eigvals, residuals, b, ef, irclass,)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "collected_XXXsph_rXXX_p580nmx580nm_mAu_n1.52_bX2_k(0_0)um-1_LX_cnXXX.pdf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41488 characters]\n",
      "text/plain": [
       "<Figure size 432x216 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(6,3))\n",
    "axesR = {}\n",
    "axesI = {}\n",
    "for i, lMax in enumerate(lMaxes):\n",
    "    axesR[lMax] = fig.add_subplot(2,len(lMaxes),i+1)\n",
    "    axesR[lMax].set_xlim([50,300])\n",
    "    axesR[lMax].set_ylim([1.25, ef[1]/eh])\n",
    "    axesI[lMax] = fig.add_subplot(2,len(lMaxes),len(lMaxes)+i+1)\n",
    "    axesI[lMax].set_xlim([50,300])\n",
    "    axesI[lMax].set_ylim([-60, 30])\n",
    "    axesR[lMax].set_title('$l_\\max = %d $' % lMax)    \n",
    "    axesR[lMax].tick_params(labelbottom=False) \n",
    "    if i == len(lMaxes)//2:\n",
    "        axesI[lMax].set_xlabel(\"Particle radius / nm\")\n",
    "    if i == 0:\n",
    "        axesR[lMax].set_ylabel('$\\hbar \\Re \\omega / \\mathrm{eV}$')\n",
    "        axesI[lMax].set_ylabel('$\\hbar \\Im \\omega / \\mathrm{meV}$')\n",
    "    else:\n",
    "        axesR[lMax].tick_params(labelleft=False) \n",
    "        axesI[lMax].tick_params(labelleft=False) \n",
    "    axesR[lMax].set_xticks([100,300])\n",
    "    axesI[lMax].set_xticks([100,300])\n",
    "\n",
    "res_thr = 0.00015\n",
    "\n",
    "ir_labeled=set()\n",
    "if True:\n",
    "  for (lMax, radius), (eigvals, residuals, b, ef, irclass) in plotdata.items():\n",
    "    for i, (e, res, iri) in enumerate(zip(eigvals, residuals, irclass)):\n",
    "        #if i == 0:\n",
    "        if res < res_thr:# and e.real < 2.14e15:\n",
    "            if iri in ir_labeled: \n",
    "                label=None\n",
    "            else:\n",
    "                ir_labeled.add(iri)\n",
    "                label=D4h_irlabels[irlabels[iri]]\n",
    "            axesR[lMax].plot(radius*1e9, e.real/eh, \n",
    "                             marker='.',\n",
    "                             #marker=markerfun(b),\n",
    "                    ms=1, #c=colorfun(b)\n",
    "                    c=matplotlib.cm.hsv(iri/9),\n",
    "                    #c = colorfun(iri),\n",
    "                    label=label,\n",
    "                   )\n",
    "            axesI[lMax].plot(radius*1e9, e.imag/eh*1000, \n",
    "                    #marker='x', \n",
    "                    #c=colorfun(b), \n",
    "                    c=matplotlib.cm.hsv(iri/9),#colorfun(iri),\n",
    "                    marker='.', #markerfun(b),\n",
    "                             ms=1,\n",
    "                    #label=label\n",
    "                            )\n",
    "fig.legend(title=\"Irrep\", loc=\"upper right\")\n",
    "plt.subplots_adjust(bottom=0.3, right=0.82,left=0.10)\n",
    "#fig.suptitle('$l_\\mathrm{max}=%d$, residual threshold = %g' % (lMax, res_thr) )\n",
    "fig.savefig(plotfilename)\n",
    "fig.savefig(plotfilename.replace('pdf', 'png'))\n",
    "print(plotfilename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.        , 1.40635433, 1.98888536, 2.81270865, 3.14470387,\n",
       "       3.97777072, 4.21906298, 4.44728287])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ef / eh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1207"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(plotdata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([5, 3, 5, 0, 7, 8, 8, 9, 9])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "irclass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
