{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "328349f5",
   "metadata": {},
   "source": [
    "# Horizon Distance Calculation with Mass-Dependent SNR Thresholds\n",
    "\n",
    "## Find the threshold for BNS systems\n",
    "\n",
    "* choose 'mass1_source' $\\in [1,3]$\n",
    "* we will use pre-selected injection catalogue data. \n",
    "* I will take threshold for individual detector same as that of detector network, for simplicity.\n",
    "* default sensitivity is consider with O4 psds. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cc9098e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# original catalogue data: https://zenodo.org/records/16740117/files/samples-rpo4a_v2_20250503133839UTC-1366933504-23846400.hdf?download=1\n",
    "# injection_data_bns.json is the reduced data extracted from the above hdf file for testing purpose\n",
    "! wget https://raw.githubusercontent.com/hemantaph/gwsnr/refs/heads/main/tests/unit/injection_data_bns.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "58dd2996",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████| 10/10 [00:10<00:00,  1.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best SNR threshold: 11.53\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from gwsnr.threshold import SNRThresholdFinder\n",
    "\n",
    "from gwsnr.utils import get_param_from_json\n",
    "params = get_param_from_json('injection_data_bns.json')\n",
    "\n",
    "gstlal_far = params['gstlal_far']\n",
    "observed_snr_net = params['observed_snr_net']\n",
    "z = params['z']\n",
    "mass1_source = params['mass1_source']\n",
    "\n",
    "test = SNRThresholdFinder(\n",
    "    catalog_file = None,\n",
    "    # below are all default values. You can omit them if you want. \n",
    "    npool=4,\n",
    "    original_detection_statistic = dict(\n",
    "                key_name='gstlal_far',\n",
    "                threshold=1,  # 1 per year\n",
    "                parameter=gstlal_far,\n",
    "            ),\n",
    "    projected_detection_statistic = dict(\n",
    "                key_name='observed_snr_net',\n",
    "                threshold=None, # to be determined\n",
    "                threshold_search_bounds=(8, 14),\n",
    "                parameter=observed_snr_net,\n",
    "            ),\n",
    "    parameters_to_fit = dict(\n",
    "        key_name = 'z',\n",
    "        parameter=z,\n",
    "    ),\n",
    "    sample_size=20000,\n",
    "    selection_range = dict(\n",
    "        key_name = 'mass1_source',\n",
    "        range = (1, 3), # in solar masses\n",
    "        parameter = mass1_source,\n",
    "    ),\n",
    ")\n",
    "\n",
    "best_thr, _, _, _, _ = test.find_threshold(iteration=10, print_output=True, no_multiprocessing=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ffb5d989",
   "metadata": {},
   "outputs": [],
   "source": [
    "snr_th, snr_th_net = best_thr, best_thr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "631a5b82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Initializing GWSNR class...\n",
      "\n",
      "psds not given. Choosing bilby's default psds\n",
      "Interpolator will be loaded for L1 detector from ./interpolator_pickle/L1/partialSNR_dict_4.pickle\n",
      "Interpolator will be loaded for H1 detector from ./interpolator_pickle/H1/partialSNR_dict_4.pickle\n",
      "Interpolator will be loaded for V1 detector from ./interpolator_pickle/V1/partialSNR_dict_4.pickle\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from gwsnr import GWSNR \n",
    "\n",
    "gwsnr = GWSNR(gwsnr_verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b54f5930",
   "metadata": {},
   "source": [
    "## Horizon Distance - Analytic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "93ce4996",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'L1': array([284.7999009]),\n",
       " 'H1': array([284.7999009]),\n",
       " 'V1': array([218.75404729])}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hd_dict_ =  gwsnr.horizon_distance_analytical(mass_1=1.4, mass_2=1.4, snr_th=snr_th)\n",
    "hd_dict_ # in Mpc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92cc6a2b",
   "metadata": {},
   "source": [
    "## Horizon Distance - Numerical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f2c659a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'L1': 328.46671006811084,\n",
       " 'H1': 328.46671006811084,\n",
       " 'V1': 252.2944005149766,\n",
       " 'snr_net': 463.3658527478692}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hd_dict_, optimal_sky_dict =  gwsnr.horizon_distance_numerical(mass_1=1.4, mass_2=1.4, snr_th=snr_th)\n",
    "\n",
    "hd_dict_ # in Mpc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "381c238c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'L1': (2.751464995516151, -0.5327666466087082),\n",
       " 'H1': (2.2513115566737323, -0.8113033727061859),\n",
       " 'V1': (1.3776350432199378, 0.7615162877515576),\n",
       " 'snr_net': (2.650441663192078, -0.6782093614739017)}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# optimal sky position\n",
    "optimal_sky_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96609162",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343de068",
   "metadata": {},
   "outputs": [],
   "source": [
    "        pip install jax\n",
    "        pip install scikit-learn==1.7.1\n",
    "        pip install tensorflow==2.17.1\n",
    "        pip install ml-dtypes --upgrade"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gwsnr",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}