{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import healpy as hp\n",
    "import h5py\n",
    "import pandas as pd\n",
    "from cycler import cycler\n",
    "from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)\n",
    "from matplotlib.legend import Legend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def projscatter2(idx, nside, **kwargs):\n",
    "    theta, phi = hp.pix2ang(nside, idx)\n",
    "    hp.projscatter(theta, phi, **kwargs)\n",
    "    return\n",
    "\n",
    "def convert_idxinj_to_radec(idx_inj, nside=16):\n",
    "    theta, phi = hp.pix2ang(nside, idx_inj)\n",
    "    ra = np.degrees(phi) / 15  # Convert phi to RA in hours\n",
    "    dec = 90 - np.degrees(theta)  # Convert theta to Dec in degrees\n",
    "    return ra, dec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "# Update global font settings\n",
    "mpl.rcParams.update({\n",
    "    'font.size': 14,        # general text size\n",
    "    'axes.titlesize': 20,   # axes title size\n",
    "    'axes.labelsize': 20,   # axes labels size\n",
    "    'xtick.labelsize': 20,  # x-axis tick labels\n",
    "    'ytick.labelsize': 20   # y-axis tick labels\n",
    "})\n",
    "mpl.rcParams['text.usetex'] = False\n",
    "mpl.rcParams['savefig.bbox'] = 'tight'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Example Maps\n",
    "\n",
    "Example maps available for the direction of injection: (RA, Dec) = ($\\pi/4$, $\\pi/2$)\n",
    "\n",
    "SNRs of injection = [3, 4, 5]\n",
    "Frequncies = [30, 100, 200, 400, 700, 1000] Hz\n",
    "\n",
    "Shape of the maps: (frequencies, SNRinj, nmaps, npix) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[435]\n"
     ]
    }
   ],
   "source": [
    "#Specify frequency and SNR injection\n",
    "freqs = [30,100,200,400,700, 1000]\n",
    "snr_inj_arr = np.array([3.0, 4.0, 5.0]) #SNR of injection, must be numpy array\n",
    "idx_inj = np.array([hp.ang2pix(nside, np.pi/4, np.pi/2)]) #Direction of injection, must be numpy array #Format: 90 - Declination, Right Ascension\n",
    "print(idx_inj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "savedir = 'plots_paper/'\n",
    "\n",
    "#Create the save directory for plots\n",
    "if not os.path.exists(savedir):\n",
    "    os.makedirs(savedir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30\n",
      "100\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 190312 characters]",
      "text/plain": [
       "<Figure size 850x540 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 176344 characters]",
      "text/plain": [
       "<Figure size 850x540 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 208964 characters]",
      "text/plain": [
       "<Figure size 850x540 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 200728 characters]",
      "text/plain": [
       "<Figure size 850x540 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "peak_color = 'gold'\n",
    "peak_color2 = 'violet'\n",
    "inj_color = 'r'\n",
    "grat_color = 'k'\n",
    "\n",
    "out_dir = 'plots_paper/maps'\n",
    "if not os.path.exists(out_dir):\n",
    "    os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "with h5py.File('maps_fig1.h5', 'r') as maps_fi:\n",
    "    noise_maps = maps_fi['noise_maps'][:]\n",
    "    inj_maps   = maps_fi['inj_maps'][:]\n",
    "\n",
    "file = 'Fig1.h5'\n",
    "with h5py.File(file, 'r') as fi:\n",
    "    for i in range(2):\n",
    "        f = freqs[i]\n",
    "        print(f)\n",
    "        noise_map = noise_maps[i, :]\n",
    "        inj_map = inj_maps[i, :]\n",
    "\n",
    "        N_min_SNR, N_max_SNR = min(noise_map), max(noise_map)\n",
    "        D_min_SNR, D_max_SNR = min(inj_map), max(inj_map)\n",
    "\n",
    "        Npeaks = fi['Npeaks'][str(f)][:]\n",
    "        Dpeaks = fi['Dpeaks'][str(f)][:]\n",
    "        hp.mollview(inj_map, flip='astro', rot=(-180,0,0), title=f'', unit='SNR', cmap = 'Blues', min = min([N_min_SNR, D_min_SNR]), max = max([N_max_SNR, D_max_SNR]),format='%.1f')\n",
    "        projscatter2(np.array(idx_inj), nside, marker='*', color=inj_color, label='Injection')\n",
    "        projscatter2(Dpeaks, nside, marker='^', color=peak_color, label='Noise + Injection Peaks', )\n",
    "        projscatter2(Npeaks, nside, marker='v', color=peak_color2, label='Noise Peaks')\n",
    "\n",
    "        hp.graticule(45,90, color=grat_color);\n",
    "        plt.annotate('12h',(0.5,0.5),xycoords='axes fraction', color=grat_color)\n",
    "        plt.annotate('6h',(0.75,0.5),xycoords='axes fraction', color=grat_color)\n",
    "        plt.annotate('18h',(0.25,0.5),xycoords='axes fraction', color=grat_color)\n",
    "        plt.annotate('45°',(0.04,0.78),xycoords='axes fraction')\n",
    "        plt.annotate('0°',(-0.03,0.49),xycoords='axes fraction')\n",
    "        plt.annotate('-45°',(0.03,0.15),xycoords='axes fraction')\n",
    "\n",
    "        plt.legend(loc='lower left', fontsize=16)\n",
    "        plt.savefig(f'{out_dir}/inj_map_{f}Hz_mollview.pdf', dpi=200)\n",
    "\n",
    "\n",
    "        hp.mollview(noise_map, flip='astro', rot=(-180,0,0), title=f'', unit='SNR', cmap='Blues', min = min([N_min_SNR, D_min_SNR]), max = max([N_max_SNR, D_max_SNR]), format='%.1f')\n",
    "        projscatter2(Npeaks, nside, marker='v', color=peak_color2, label='Noise Peaks')\n",
    "\n",
    "        hp.graticule(45,90, color=grat_color);\n",
    "        plt.annotate('12h',(0.5,0.5),xycoords='axes fraction', color=grat_color)\n",
    "        plt.annotate('6h',(0.75,0.5),xycoords='axes fraction', color=grat_color)\n",
    "        plt.annotate('18h',(0.25,0.5),xycoords='axes fraction', color=grat_color)\n",
    "        plt.annotate('45°',(0.04,0.78),xycoords='axes fraction')\n",
    "        plt.annotate('0°',(-0.03,0.49),xycoords='axes fraction')\n",
    "        plt.annotate('-45°',(0.03,0.15),xycoords='axes fraction')\n",
    "\n",
    "        plt.legend(loc='lower left', fontsize=16)\n",
    "        plt.savefig(f'{out_dir}/noise_map{f}Hz_mollview.pdf', dpi=200)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 3: No. of candidates vs SNR threshold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 173932 characters]",
      "text/plain": [
       "<Figure size 1000x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mpl.rcParams.update({\n",
    "    'font.size': 14,        # general text size\n",
    "    'axes.titlesize': 20,   # axes title size\n",
    "    'axes.labelsize': 20,   # axes labels size\n",
    "    'xtick.labelsize': 20,  # x-axis tick labels\n",
    "    'ytick.labelsize': 20   # y-axis tick labels\n",
    "})\n",
    "mpl.rcParams['text.usetex'] = False\n",
    "mpl.rcParams['savefig.bbox'] = 'tight'\n",
    "\n",
    "j=0 #Select SNR inj = 3\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "ax.grid(which='major', color='#DDDDDD', linewidth=0.8)\n",
    "ax.grid(which='minor', color='#EEEEEE', linewidth=0.6)\n",
    "ax.set_prop_cycle(cycler(color=['r', 'g', 'b', 'c', 'm', 'y', 'k']))\n",
    "ax.set_xlabel('SNR Threshold')\n",
    "ax.set_ylabel('Mean No. of Candidates')\n",
    "\n",
    "ax.set_title(fr\"SNR of injection: {snr_inj_arr[j]}, Direction of injection: (RA, Dec) = (6.2 h, 45°)\")\n",
    "ax.set_yscale('log')\n",
    "\n",
    "with h5py.File('Fig3.h5', 'r') as fi:\n",
    "    for i in range(len(freqs) - 1): #Ignore 1000 Hz\n",
    "        # ax.plot(cutoff, Ncandidate_no_arr_as/nmaps, label='Noise All Sky')\n",
    "        # ax.plot(cutoff, Ncandidate_no_arr_peak/nmaps, label='Noise Peaks')\n",
    "        f = freqs[i]\n",
    "        Dcandidate_no_dens_arr_peak = fi['Dcandidate_no_arr_peak'][str(f)][:]\n",
    "        Dcandidate_no_dens_arr_as   = fi['Dcandidate_no_arr_as'][str(f)][:]\n",
    "        cutoff                      = fi['cutoff'][str(f)][:]\n",
    "\n",
    "        line1, = ax.plot(cutoff, Dcandidate_no_dens_arr_peak, label=f'{freqs[i]} Hz')\n",
    "        line2, = ax.plot(cutoff, Dcandidate_no_dens_arr_as, c = line1.get_color(), linestyle='--')\n",
    "        leg = Legend(ax, [line1, line2], ['Peak finder', 'All Sky'], loc='upper center', frameon=True, fontsize=14)\n",
    "        ax.add_artist(leg)\n",
    "\n",
    "        \n",
    "    plt.legend(fontsize=20)\n",
    "plt.savefig(f'plots_paper/n_candidates_vs_snr_SNRinj{snr_inj_arr[j]}.pdf', dpi=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 4: No. of Peaks vs Frequency and Fig 7: Runtime vs Frequency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36468 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38736 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with h5py.File('Fig4.h5', 'r') as fi:\n",
    "    mean_n_peaks = fi['mean_n_peaks'][:]\n",
    "    std_n_peaks = fi['std_n_peaks'][:]\n",
    "\n",
    "with h5py.File('Fig7.h5', 'r') as fi:\n",
    "    mean_runtime = fi['mean_runtime'][:]\n",
    "    std_runtime = fi['std_runtime'][:]\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.errorbar(freqs, mean_runtime, yerr=std_runtime, fmt='o', capsize=5)\n",
    "ax.set_xlabel('Frequency (Hz)')\n",
    "ax.set_ylabel('Runtime (s)')\n",
    "ax.set_title('Runtime vs Frequency')\n",
    "ax.grid(True)\n",
    "plt.savefig('plots_paper/runtime_vs_freq_v2.pdf', dpi=200)\n",
    "\n",
    "fig1, ax1 = plt.subplots()\n",
    "ax1.errorbar(freqs, mean_n_peaks, yerr=std_n_peaks, fmt='o', capsize=5)\n",
    "ax1.set_xlabel('Frequency (Hz)')\n",
    "ax1.set_ylabel('Number of Peaks')\n",
    "ax1.set_title('Number of Peaks vs Frequency')\n",
    "ax1.grid(True)\n",
    "plt.savefig('plots_paper/n_peaks_vs_freq_v2.pdf', dpi=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 5a: SNR distribution of candidates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Frequency = 400 Hz, SNR inj = 3.0\n"
     ]
    }
   ],
   "source": [
    "#Find peaks in maps\n",
    "i=3 #Frequency index\n",
    "j=0 #SNR inj index\n",
    "print(f\"Frequency = {freqs[i]} Hz, SNR inj = {snr_inj_arr[j]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 76880 characters]",
      "text/plain": [
       "<Figure size 1000x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "mpl.rcParams.update({\n",
    "    'font.size': 20,        # general text size\n",
    "    'axes.titlesize': 20,   # axes title size\n",
    "    'axes.labelsize': 20,   # axes labels size\n",
    "    'xtick.labelsize': 20,  # x-axis tick labels\n",
    "    'ytick.labelsize': 20   # y-axis tick labels\n",
    "})\n",
    "\n",
    "\n",
    "with h5py.File('Fig5a.h5', 'r') as fi:\n",
    "    N_peaks_SNR_arr = fi['N_peaks_SNR_arr'][:]\n",
    "    D_peaks_SNR_arr = fi['D_peaks_SNR_arr'][:]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "ax.grid(which='both', linewidth=0.5)\n",
    "ax.hist(N_peaks_SNR_arr, bins='fd', alpha=0.5, label='Noise', density=True, histtype='step', linestyle='dashdot', lw=2.5)\n",
    "ax.hist(D_peaks_SNR_arr, bins='fd', alpha=0.5, label='Injection', density=True, histtype='step', lw=2.5)\n",
    "ax.set_xlabel('Peak SNR')\n",
    "ax.set_ylabel('Probability Density')\n",
    "ax.set_title(f'SNR of injection: {snr_inj_arr[j]}, Frequency: {freqs[i]} Hz,\\n Direction of injection: (RA, Dec) = (6.2 h , 45°)')\n",
    "plt.legend()\n",
    "plt.savefig(f\"plots_paper/PeakSNR_distribution_SNRinj{snr_inj_arr[j]}_{freqs[i]}Hz.pdf\", dpi=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 5b: FAR, FDR vs SNR threshold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RA and Dec for each idx_inj: [(np.float64(3.0), np.float64(87.07581964294997)), (np.float64(0.0), np.float64(2.388015463268772)), (np.float64(6.1875), np.float64(30.000000000000007)), (np.float64(18.1875), np.float64(-30.000000000000014))]\n"
     ]
    }
   ],
   "source": [
    "idx_inj_arr2 = [0, 1440, 752, 2320]\n",
    "snr_inj_arr2 = np.arange(1,4.5,0.5)\n",
    "freq_arr2 = [30, 60, 100, 200, 400, 700, 1000]\n",
    "\n",
    "ra_dec_list = []\n",
    "for idx in idx_inj_arr2:\n",
    "    theta, phi = hp.pix2ang(nside, idx)\n",
    "    ra = np.degrees(phi) / 15  # Convert phi to RA in hours\n",
    "    dec = 90 - np.degrees(theta)  # Convert theta to Dec in degrees\n",
    "    ra_dec_list.append((ra, dec))\n",
    "\n",
    "print(\"RA and Dec for each idx_inj:\", ra_dec_list)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_folder = './' #This contains all the data I have with the definition of FAR, FDR based on SNR threshold \n",
    "df = pd.read_csv(data_folder +'iterable_comb.csv', header=0, index_col=0)\n",
    "FAR_mat = np.loadtxt(data_folder + 'FAR_mat.txt')\n",
    "FDR_mat = np.loadtxt(data_folder + 'FDR_mat.txt')\n",
    "FAR_err_mat = np.loadtxt(data_folder + 'FAR_err_mat.txt')\n",
    "FDR_err_mat = np.loadtxt(data_folder + 'FDR_err_mat.txt')\n",
    "nSNRthres = 30\n",
    "SNRthres_arr = np.linspace(0,6, nSNRthres)\n",
    "\n",
    "freqs_needed = [30, 100, 200, 400, 700]\n",
    "df2 = df[df['Freq'].isin(freqs_needed)]\n",
    "FAR_mat2= FAR_mat[df2.index]\n",
    "FDR_mat2 = FDR_mat[df2.index]\n",
    "FAR_err_mat2 = FAR_err_mat[df2.index]\n",
    "FDR_err_mat2 = FDR_err_mat[df2.index]\n",
    "\n",
    "df2 = df2.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "90, SNR inj = 3.0, (RA, Dec) = ((np.float64(6.1875), np.float64(30.000000000000007)))\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 235652 characters]",
      "text/plain": [
       "<Figure size 1000x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "j = 90\n",
    "\n",
    "idx_inj2 = df2.loc[j, 'Idx_Inj']\n",
    "snr_inj2 = df2.loc[j, 'SNR_inj']  #Find RA, Dec from Idx_inj\n",
    "theta, phi = hp.pix2ang(nside, int(idx_inj2))\n",
    "ra = np.degrees(phi) / 15  # Convert phi to RA in hours\n",
    "dec = 90 - np.degrees(theta)  # Convert theta to Dec in degrees fig3, ax3 = plt.subplots()\n",
    "\n",
    "print(f\"{j}, SNR inj = {snr_inj2}, (RA, Dec) = ({ra, dec})\")\n",
    "    \n",
    "fig3, ax3 = plt.subplots(figsize=(10,7))\n",
    "ax3.set_xlabel('Peak SNR Threshold')\n",
    "ax3.set_ylabel('False Alarm Rate and False Discard Rate')\n",
    "ax3.set_title(f'SNR of injection: {snr_inj2},\\n Direction of injection: (RA, Dec) = ({ra:.1f} h, {dec:.1f}°)')\n",
    "ax3.grid(which='major', color='#DDDDDD', linewidth=0.8)\n",
    "ax3.grid(which='minor', color='#EEEEEE', linewidth=0.6)\n",
    "ax3.xaxis.set_major_locator(MultipleLocator(0.5))#; ax1.xaxis.set_major_formatter('{x:.0f}')\n",
    "ax3.xaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "ax3.yaxis.set_major_locator(MultipleLocator(0.2))#; ax1.yaxis.set_major_formatter('{x:.0f}')\n",
    "ax3.yaxis.set_minor_locator(MultipleLocator(0.05))\n",
    "ax3.set_prop_cycle(cycler(color=['r', 'g', 'b', 'c', 'm', 'y', 'k']))\n",
    "for i in range(j,j+5):\n",
    "    f = df2['Freq'][i] \n",
    "    # line1,*_ = ax3.errorbar(SNRthres_arr, FAR_mat2[i,:], yerr=FAR_err_mat2[i], label = str(f) + \" Hz\", capsize=2)\n",
    "    # line2,*_ = ax3.errorbar(SNRthres_arr, FDR_mat2[i,:], yerr=FDR_err_mat2[i], linestyle='--', c=line1.get_color(), capsize=2)\n",
    "\n",
    "    line1, = ax3.plot(SNRthres_arr, FAR_mat2[i,:], label = str(f) + \" Hz\")\n",
    "    line2, = ax3.plot(SNRthres_arr, FDR_mat2[i,:], linestyle='--', c=line1.get_color())\n",
    "    ax3.fill_between(SNRthres_arr, \n",
    "                        FAR_mat2[i,:] - FAR_err_mat2[i,:], \n",
    "                        FAR_mat2[i,:] + FAR_err_mat2[i,:], \n",
    "                        alpha=0.3, color=line1.get_color())\n",
    "    ax3.fill_between(SNRthres_arr, \n",
    "                        FDR_mat2[i,:] - FDR_err_mat2[i,:], \n",
    "                        FDR_mat2[i,:] + FDR_err_mat2[i,:], \n",
    "                        alpha=0.3, color=line2.get_color())\n",
    "    \n",
    "    leg = Legend(ax3, [line1, line2], ['FAR', 'FDR'], loc='upper center', frameon=True, fontsize=20, bbox_to_anchor=(0.4, 1))\n",
    "    ax3.add_artist(leg)\n",
    "\n",
    "\n",
    "plt.legend(fontsize=20, bbox_to_anchor = (0.7, 0.1))\n",
    "\n",
    "out_dir = 'plots_paper'\n",
    "\n",
    "if not os.path.exists(out_dir):\n",
    "    os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "plt.savefig(f'{out_dir}/FARFDRvsSNRthres10000_difflim_16_SNRinj' + str(snr_inj2) + '_IdxInj' + str(idx_inj2) + '.pdf', bbox_inches='tight', dpi=200)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 5c: FAR, FDR vs SNR threshold, average over all directions and frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import colors\n",
    "cmap = mpl.colormaps['plasma']\n",
    "norm = colors.Normalize(vmin=snr_inj_arr.min(), vmax=snr_inj_arr.max())\n",
    "\n",
    "snr_inj_arr = np.arange(1, 5.5, 0.5)\n",
    "df3 = df.sort_values(by=['SNR_inj', 'Idx_Inj'], ignore_index=False)\n",
    "new_idx = df3.index.values\n",
    "\n",
    "FAR3 = FAR_mat.take(new_idx, axis=0)\n",
    "FDR3 = FDR_mat.take(new_idx, axis=0)\n",
    "FAR_err3 = FAR_err_mat.take(new_idx, axis=0)\n",
    "FDR_err3 = FDR_err_mat.take(new_idx, axis=0)\n",
    "\n",
    "ndir = 4\n",
    "nfreq = 7\n",
    "\n",
    "var = ndir * nfreq\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 233880 characters]",
      "text/plain": [
       "<Figure size 1000x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots()\n",
    "fig1.set_size_inches(10,7)\n",
    "ax1.set_xlabel('False Alarm Rate ')\n",
    "ax1.set_ylabel('False Discard Rate')\n",
    "ax1.grid(which='major', color='#DDDDDD', linewidth=0.8)\n",
    "ax1.grid(which='minor', color='#EEEEEE', linewidth=0.6)\n",
    "ax1.set_title('Average over 4 directions and 7 frequencies')\n",
    "ax1.xaxis.set_major_locator(MultipleLocator(0.5))#; ax1.xaxis.set_major_formatter('{x:.0f}')\n",
    "ax1.xaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "ax1.yaxis.set_major_locator(MultipleLocator(0.2))#; ax1.yaxis.set_major_formatter('{x:.0f}')\n",
    "ax1.yaxis.set_minor_locator(MultipleLocator(0.05))\n",
    "# ax1.set_prop_cycle(cycler(color=['r', 'g', 'b', 'c', 'm', 'y', 'k']))\n",
    "    #                    cycler(linestyle=['-', '--', '-.', ':', '-', '--', '-.']))\n",
    "\n",
    "\n",
    "for i in range(0,FAR3.shape[0], var):\n",
    "\n",
    "    tFAR = FAR3[i:i+var, :]\n",
    "    tFDR = FDR3[i:i+var, :]\n",
    "    \n",
    "    mean_FAR_arr = np.mean(tFAR, axis=0) #Mean over frequencies\n",
    "    mean_FDR_arr = np.mean(tFDR, axis=0)\n",
    "    std_FAR_arr = np.std(tFAR, axis=0)\n",
    "    std_FDR_arr = np.std(tFDR, axis=0)\n",
    "\n",
    "    # mean_FAR_mat[i//var, :] = mean_FAR_arr\n",
    "    # mean_FDR_mat[i//var, :] = mean_FDR_arr\n",
    "    # std_FAR_mat[i//var, :] = std_FAR_arr\n",
    "    # std_FDR_mat[i//var, :] = std_FDR_arr\n",
    "    # ax1.errorbar(mean_FAR_mat[i//var, :], mean_FDR_mat[i//var, :], xerr=std_FAR_mat[i//var, :], yerr=std_FDR_mat[i//var, :], fmt='o', capsize=2, ms=3, label=f'SNR inj: {snr_inj_arr[i//var]}')\n",
    "    # ax1.errorbar(mean_FAR_arr, mean_FDR_arr, xerr=std_FAR_arr, yerr=std_FDR_arr, fmt='o-', capsize=2, ms=3, label=f'SNR inj: {snr_inj_arr[i%var]}')\n",
    "    snr_val = snr_inj_arr[int(i/var)]\n",
    "    color = cmap(norm(snr_val))\n",
    "    line1, = ax1.plot(mean_FAR_arr, mean_FDR_arr, marker='o', markersize=3, color=color, label=f'SNR inj: {snr_val}')\n",
    "\n",
    "    # if i== var*4:\n",
    "    #     for k, SNRthres in enumerate(SNRthres_arr):\n",
    "    #         ax1.annotate(str(round(SNRthres, 2)), (mean_FAR_arr[k], mean_FDR_arr[k]))\n",
    "\n",
    "    ax1.fill_between(mean_FAR_arr, mean_FDR_arr - std_FDR_arr, mean_FDR_arr + std_FDR_arr, alpha=0.3, color=color)\n",
    "\n",
    "\n",
    "ax1.legend()\n",
    "plt.savefig(f'plots_paper/ROC_SNRinj1-5_avg4dir7freq_annot.pdf', bbox_inches='tight', dpi=200)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 6: FDR vs N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 128208 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 141900 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 155396 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 133056 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 138868 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 153572 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot FDR vs N for multiple SNR injections, and direction of injection\n",
    "\n",
    "wd = './'\n",
    "\n",
    "niter = 36\n",
    "N = 3072\n",
    "\n",
    "idx_inj_arr = [435, 2637]\n",
    "\n",
    "with h5py.File('Fig6.h5', 'r') as fi:\n",
    "    FDR_mat = fi['FDR_mat'][:]\n",
    "    FDR_mat_as = fi['FDR_mat_as'][:]\n",
    "    iterable_df = fi['iterable_df'][:]\n",
    "\n",
    "\n",
    "iterable_df = pd.DataFrame(iterable_df, columns=['Idx_inj', 'SNRinj', 'freq'])\n",
    "\n",
    "nfreq = iterable_df.freq.nunique()\n",
    "ndir = iterable_df.Idx_inj.nunique()\n",
    "nsnr_inj = iterable_df.SNRinj.nunique()\n",
    "\n",
    "for j in range(0, niter, nfreq): \n",
    "\n",
    "    snr_inj = iterable_df.SNRinj[j]\n",
    "    idx_inj = iterable_df.Idx_inj[j]\n",
    "\n",
    "    ra, dec = convert_idxinj_to_radec(int(idx_inj))\n",
    "\n",
    "    N = np.arange(0, 3072)\n",
    "    fig1, ax1 = plt.subplots(figsize=(10, 6))\n",
    "    ax1.set_xlabel('No. of sources that can be followed up (N)')\n",
    "    ax1.set_ylabel('False Dismissal Rate (FDR)')\n",
    "    ax1.set_title(f\"SNR injection = {snr_inj}, \\n Direction of injection: (RA, Dec) = ({ra:.1f} h, {dec:.1f}°)\")\n",
    "\n",
    "    xmax = 15\n",
    "    ax1.set_xlim(1, xmax)\n",
    "    ax1.set_axisbelow(True)\n",
    "\n",
    "    scale='log'\n",
    "    # ax1.set_xscale(scale)\n",
    "    ax1.set_yscale(scale)\n",
    "    # ax1.xaxis.set_minor_locator(AutoMinorLocator(2))\n",
    "    # ax1.set_xticks(np.arange(1, xmax+1, 2))\n",
    "    ax1.grid(which='major', color='#DDDDDD', linewidth=0.8)\n",
    "    ax1.grid(which='minor', color='#EEEEEE', linewidth=0.5)\n",
    "    \n",
    "    ax1.set_prop_cycle(cycler(color=['r', 'g', 'b', 'c', 'm', 'y', 'k']))\n",
    "    #                    cycler(linestyle=['-', '--', '-.', ':', '-', '--', '-.']))\n",
    "\n",
    "    if snr_inj == 3.0:\n",
    "        y_cut = 5e-3\n",
    "    elif snr_inj == 4.0:\n",
    "        y_cut = 2e-3\n",
    "    elif snr_inj == 5.0:\n",
    "        y_cut = 3e-4\n",
    "    \n",
    "    min_FDR = np.min(FDR_mat[j:j+6, 5]) #np.min(FDR_mat[j:j+6, 1:xmax])\n",
    "    max_FDR = np.max(FDR_mat[j:j+6, 1:xmax])\n",
    "\n",
    "    # min_FDR_arg = np.argmin(FDR_mat[j:j+6, 1:xmax])\n",
    "    # max_FDR_arg = np.argmax(FDR_mat[j:j+6, 1:xmax])\n",
    "    # print(f\"Min FDR arg: {min_FDR_arg}, Max FDR arg: {max_FDR_arg} for SNR injection {snr_inj}, idx_inj {idx_inj}\")\n",
    "\n",
    "    ax1.set_ylim(min_FDR - y_cut, max_FDR + 0.005)\n",
    "    for i in range(j, j+6):\n",
    "        # ax1.fill_between(N, FDR_mat4[i] - FDR_mat_err4[i], FDR_mat4[i] + FDR_mat_err4[i], alpha=0.3)\n",
    "        line1, = ax1.plot(N, FDR_mat[i], label=f'{int(iterable_df.freq[i])} Hz', lw=1.5)\n",
    "\n",
    "\n",
    "        # ax1.fill_between(N, FDR_mat_as4[i] - FDR_mat_as_err4[i], FDR_mat_as4[i] + FDR_mat_as_err4[i], alpha=0.3, color=line1.get_color())\n",
    "        line2, = ax1.plot(N, FDR_mat_as[i], ls='--', c=line1.get_color(), lw=1.5)\n",
    "\n",
    "        leg = Legend(ax1, [line1, line2], ['Peak finder', 'Full Sky'], loc = 'upper center', frameon=True, fontsize=16)\n",
    "        ax1.add_artist(leg)\n",
    "    ax1.legend(fontsize=14)\n",
    "\n",
    "    out_dir = f'plots_paper/FDR_vs_N'\n",
    "    if not os.path.exists(out_dir):\n",
    "        os.makedirs(out_dir, exist_ok=True)\n",
    "    \n",
    "    plt.savefig(f'{out_dir}/FDR_vs_N_SNRinj{snr_inj}_IdxInj{idx_inj}.pdf', dpi=200)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "asaf2",
   "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.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
