{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "e75b321b-d46b-4509-b2b6-b62e45295f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from astropy.coordinates import SkyCoord\n",
    "from ligo.skymap.io import read_sky_map\n",
    "from ligo.skymap.postprocess import find_greedy_credible_levels\n",
    "import ligo.skymap.plot as skplot\n",
    "from healpy import npix2nside\n",
    "from astropy_healpix import HEALPix\n",
    "import astropy.units as u\n",
    "from gwpy.table import EventTable\n",
    "from gwpy.time import to_gps\n",
    "from io import StringIO, BytesIO\n",
    "from astropy.table import vstack\n",
    "from functools import lru_cache\n",
    "import numpy as np\n",
    "from astropy.io import fits \n",
    "import base64\n",
    "from matplotlib.colors import rgb2hex\n",
    "from astropy.utils.misc import JsonCustomEncoder\n",
    "import json\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "5fbc920b-6fd2-4918-830b-abb564202a85",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "t1 = \"2017-08-16T12:00:00\" \n",
    "t2 = \"2017-08-18T12:00:00\" \n",
    "do_cone_search = False\n",
    "ra = 250\n",
    "dec = -30.\n",
    "radius = 10 #degree\n",
    "level_threshold = 10 #percent\n",
    "contour_levels = \"50,90\" #percent, comma-separated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "43e119cd-7e6e-4cfb-a3bf-ba9637244136",
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = to_gps(t1)\n",
    "t2 = to_gps(t2)\n",
    "#do_cone_search = True if do_cone_search in ('true', 'True', 'yes', True) else False\n",
    "contour_levels = [int(x) for x in contour_levels.split(',')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "be14fb05-06e3-4d96-aa2f-8df40435d077",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: MergeConflictWarning: In merged column 'mass_1_source' the 'unit' attribute does not match (solMass != ).  Using  for merged output [astropy.utils.metadata]\n"
     ]
    }
   ],
   "source": [
    "tabs = []\n",
    "columns_to_add = ['commonName', \n",
    "                 'GPS',\n",
    "                 'luminosity_distance',\n",
    "                 'luminosity_distance_lower',\n",
    "                 'luminosity_distance_upper',\n",
    "                 'mass_1_source', \n",
    "                 'mass_1_source_lower', \n",
    "                 'mass_1_source_upper', \n",
    "                 'mass_2_source', \n",
    "                 'mass_2_source_lower', \n",
    "                 'mass_2_source_upper']\n",
    "for catalog in (\"GWTC-1-confident\", \"GWTC-2\", \"GWTC-3-confident\"):\n",
    "    tabs.append(EventTable.fetch_open_data(catalog, selection=f'GPS > {t1} && GPS < {t2}')[columns_to_add])\n",
    "tab = vstack(tabs)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "6df0bb1c-c769-413c-916d-18a6bd050da3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@lru_cache\n",
    "def reg_idxs(nside, ra, dec, radius):\n",
    "    point = SkyCoord(ra=ra, dec=dec, frame='icrs', unit=u.deg)\n",
    "    hp = HEALPix(nside,order='nested', frame='icrs')\n",
    "    return hp.cone_search_skycoord(point, radius*u.deg)\n",
    "\n",
    "skymap_files = {}\n",
    "mask = []\n",
    "good_levels = {}\n",
    "for i in range(len(tab)):\n",
    "    event = tab['commonName'][i]\n",
    "    if True: #event != 'GW170817':\n",
    "        fits_file = f'{os.getenv(\"ODA_DATA_ROOT\", \".\").rstrip(\"/\")}/skymaps/{event}.fits.gz'\n",
    "        if do_cone_search:\n",
    "            skymap, metadata = read_sky_map(fits_file, moc=False, nest=True)\n",
    "            levels = 100 * find_greedy_credible_levels(skymap)\n",
    "            nside = npix2nside(len(levels))\n",
    "            idxs = reg_idxs(nside, ra, dec, radius)\n",
    "            if levels[idxs].min() < level_threshold:\n",
    "                mask.append(i)\n",
    "                with fits.open(fits_file, 'readonly') as fd:\n",
    "                    skymap_files[event] = {\n",
    "                        'header': fd[1].header.tostring(),\n",
    "                        'data': json.dumps(fd[1].data, cls=JsonCustomEncoder),\n",
    "                    }\n",
    "                good_levels[event] = levels\n",
    "        else:\n",
    "            mask.append(i)\n",
    "            with fits.open(fits_file, 'readonly') as fd:\n",
    "                skymap_files[event] = {\n",
    "                    'header': fd[1].header.tostring(),\n",
    "                    'data': json.dumps(fd[1].data, cls=JsonCustomEncoder),\n",
    "                }\n",
    "    else:\n",
    "        if do_cone_search:\n",
    "            gw170817 = SkyCoord(197.45035417, -23.38148417, frame='icrs', unit=u.deg)\n",
    "            if SkyCoord(ra=ra, dec=dec, frame='icrs', unit=u.deg).separation(gw170817).deg < radius:\n",
    "                mask.append(i)\n",
    "        else:\n",
    "            mask.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "6cca196e-2cc4-465e-915d-fb6dbcc8a3a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "odatab = tab[mask]\n",
    "if do_cone_search:\n",
    "    odatab.meta = {'data': 'GW cone search in time interval',\n",
    "                   't1': t1.gpsSeconds,\n",
    "                   't2': t2.gpsSeconds,\n",
    "                   'ra': ra,\n",
    "                   'dec': dec,\n",
    "                   'radius': radius,\n",
    "                   'level': level_threshold\n",
    "                   }\n",
    "else:\n",
    "    odatab.meta = {'data': 'GW in time interval',\n",
    "                   't1': t1.gpsSeconds,\n",
    "                   't2': t2.gpsSeconds,\n",
    "                   }\n",
    "\n",
    "with StringIO() as fd:\n",
    "    odatab.write(fd, format='ascii.ecsv')\n",
    "    asciicat = fd.getvalue()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "3df01119-e813-40f1-92a0-973610089b4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "contourdict = {}\n",
    "plt.figure(figsize=(8, 4))\n",
    "ax = plt.axes(projection=('astro hours mollweide'))\n",
    "wcs = ax.wcs\n",
    "lpr = []\n",
    "names = []\n",
    "numevents = len(odatab)\n",
    "maxevents = 15\n",
    "ltitle=''\n",
    "ii = np.arange(numevents)\n",
    "if numevents > maxevents:\n",
    "    ii = np.random.randint(0, numevents, maxevents)\n",
    "    ltitle = 'Too much events!\\nShowing random subset'\n",
    "a = np.linspace(0, 1, numevents)\n",
    "np.random.shuffle(a)\n",
    "colors = [rgb2hex(x) for x in plt.cm.gist_rainbow(a)]\n",
    "\n",
    "for i in range(len(odatab)):\n",
    "    event = odatab[i]['commonName']\n",
    "    if False: #event == 'GW170817':\n",
    "        pr, = ax.plot(197.45035417, -23.38148417,\n",
    "                '+',\n",
    "                transform=ax.get_transform('world'),\n",
    "                markersize=10,\n",
    "                markeredgewidth=2, \n",
    "                color=colors[i])    \n",
    "        lpr.append(pr)\n",
    "        names.append('GW170817')\n",
    "    else:\n",
    "        levels = good_levels.get(event, None)\n",
    "        if levels is None:\n",
    "            fits_file = f'{os.getenv(\"ODA_DATA_ROOT\", \".\").rstrip(\"/\")}/skymaps/{event}.fits.gz'\n",
    "            skymap, metadata = read_sky_map(fits_file, nest=True)\n",
    "            levels = 100 * find_greedy_credible_levels(skymap)\n",
    "        cs = ax.contour_hpx(\n",
    "                    (levels, 'ICRS'), nested=True,\n",
    "                     colors=colors[i], linewidths=1, \n",
    "                     levels=contour_levels, \n",
    "                     alpha=1 if i in ii else 0)\n",
    "        if i in ii:\n",
    "            lpr.append(plt.Rectangle((0,0),1,1,fc = colors[i]))\n",
    "            names.append(event)\n",
    "        \n",
    "        contourdict[event] = {}\n",
    "        contourdict[event]['levels'] = cs.levels\n",
    "        contourdict[event]['contours'] = []\n",
    "        for level_idx in range(len(cs.levels)):\n",
    "            contourdict[event]['contours'].append([])\n",
    "            segs = cs.allsegs[level_idx]\n",
    "            for j in range(len(segs)):\n",
    "                world_coords = wcs.all_pix2world(segs[j], 0)\n",
    "                contourdict[event]['contours'][level_idx].append(world_coords)\n",
    "    i+=1\n",
    "\n",
    "if do_cone_search:\n",
    "    pr, = ax.plot(ra, dec,\n",
    "            'rx',\n",
    "            transform=ax.get_transform('world'),\n",
    "            markersize=10,\n",
    "            markeredgewidth=2,\n",
    "            )   \n",
    "    lpr.append(pr)\n",
    "    names.append('Cone search aim point')\n",
    "\n",
    "ax.grid(True)\n",
    "ax.legend(lpr, names, numpoints=1, bbox_to_anchor=(1.05, 1), loc='upper left', title=ltitle)\n",
    "\n",
    "bio = BytesIO()\n",
    "plt.savefig(bio, format='svg', bbox_inches=\"tight\")\n",
    "bio.seek(0)\n",
    "image = base64.b64encode(bio.read())\n",
    "bio.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "67608337-c440-4687-9dd9-0c7f6d773096",
   "metadata": {
    "tags": [
     "outputs"
    ]
   },
   "outputs": [],
   "source": [
    "asciicat = asciicat\n",
    "skymap_files = skymap_files\n",
    "image = image.decode()\n",
    "contours = contourdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d1880f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "7a40ab12092ba9fa9bb6584fd9d0f6172be3f94b893280e26cea05cc6ec15168"
  },
  "kernelspec": {
   "display_name": "Python 3.9.6 ('gw')",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}