{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DIA Analysis: Supernovae from the Run 1.2p Test\n",
    "Michael Wood-Vasey\n",
    "Last Verified to Run: 2019-07-17\n",
    "\n",
    "After completing this Notebook, the user will be able to\n",
    "1. Get a list of simulated SN from the Run 1.2p truth catalog\n",
    "2. Select SNe within the test patch in the Run 1.2p DIA Test\n",
    "3. Plot lightcurves from the DIA analysis\n",
    "4. Plot lightcurves from the truth catalog\n",
    "5. Compare the input and recovered lightcurves.\n",
    "\n",
    "See the Truth GCR Variables for a basic overview of the Truth information and how to access it for variables.\n",
    "https://github.com/LSSTDESC/DC2-analysis/blob/master/tutorials/truth_gcr_variables.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inject gcr-catalogs that supports DIA source into path.\n",
    "import os\n",
    "import math\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from astropy.coordinates import SkyCoord\n",
    "import astropy.units as u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lsst.afw.display as afwDisplay\n",
    "import lsst.afw.geom as afwGeom\n",
    "from lsst.daf.persistence import Butler\n",
    "from lsst.geom import SpherePoint\n",
    "import lsst.geom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import GCRCatalogs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Polygon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "repo = '/global/cscratch1/sd/rearmstr/new_templates/diffim_template'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "butler = Butler(repo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diaSrc = GCRCatalogs.load_catalog('dc2_dia_source_run1.2p_test')\n",
    "diaObject = GCRCatalogs.load_catalog('dc2_dia_object_run1.2p_test')\n",
    "truth_cat = GCRCatalogs.load_catalog('dc2_truth_run1.2_variable_summary')\n",
    "truth_lc = GCRCatalogs.load_catalog('dc2_truth_run1.2_variable_lightcurve')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(We presently will get a warning from the catalog reader in the initalization above because there is no u-band in the subtractions.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select SNe from the truth catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['ra', 'dec', 'redshift', 'uniqueId', 'galaxy_id', 'sn']\n",
    "truth_all_sne = pd.DataFrame(truth_cat.get_quantities(columns, filters=[f'sn == 1']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll get a dataset.value deprecation warning.  Don't worry about this.  The Data Access Team will fix this someday."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'We found {len(truth_all_sne)} simulated SNe.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select SNe from the truth catalog in the DIA test region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tract = 4849\n",
    "patch = (6, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "skymap = butler.get('deepCoadd_skyMap')\n",
    "tract_info = skymap[tract]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "foo = tract_info.getPatchInfo(patch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bar = foo.getOuterSkyPolygon(tract_info.getWcs())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tract_box = afwGeom.Box2D(tract_info.getBBox())\n",
    "tract_pos_list = tract_box.getCorners()\n",
    "wcs = tract_info.getWcs()\n",
    "corners = wcs.pixelToSky(tract_pos_list)\n",
    "corners = np.array([[c.getRa().asDegrees(), c.getDec().asDegrees()] for c in corners])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(corners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ra = corners[:, 0]\n",
    "dec = corners[:, 1]\n",
    "min_ra, max_ra = np.min(ra), np.max(ra)\n",
    "min_dec, max_dec = np.min(dec), np.max(dec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "area_cut = [f'ra > {min_ra}', f'ra < {max_ra}', f'dec > {min_dec}', f'dec < {max_dec}']\n",
    "sn_cut = ['sn == 1']\n",
    "all_cuts = area_cut + sn_cut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['ra', 'dec', 'redshift', 'uniqueId', 'galaxy_id', 'sn']\n",
    "truth_sne = pd.DataFrame(truth_cat.get_quantities(columns, filters=all_cuts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'We found {len(truth_sne)} simulated SNe.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_dec = (min_dec + max_dec)/2\n",
    "size = 8\n",
    "dec_size = size\n",
    "ra_size = dec_size * np.cos(np.deg2rad(avg_dec))\n",
    "aspect_ratio = dec_size / ra_size\n",
    "\n",
    "# fig = plt.figure(figsize=(size, size))\n",
    "ax = plt.gca()\n",
    "ax.set_aspect(aspect_ratio)\n",
    "\n",
    "patch_region = Polygon(corners, color='red', fill=False)\n",
    "\n",
    "ax.scatter(truth_sne['ra'], truth_sne['dec'])\n",
    "ax.set_xlabel('RA')\n",
    "ax.set_ylabel('Dec')\n",
    "ax.set_xlim(plt.xlim()[::-1])\n",
    "ax.add_patch(patch_region);\n",
    "ax.set_title(f'tract: {tract}, patch: {patch}');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Lightcurve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Search for diaObjects that match input simulated SNe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "sn = truth_sne.iloc[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Oops, we need to fix up the `dtype`s somewhere.  Those shouldn't all be floats."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the lightcurve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['obshistid', 'mjd', 'mag', 'filter']\n",
    "\n",
    "sn_lc = pd.DataFrame(truth_lc.get_quantities(columns,\n",
    "                                             native_filters=[f'uniqueId == {sn[\"uniqueId\"]}']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn_lc.rename(columns={'filter': 'filter_code'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Translate filter codes to filter names\n",
    "filter_names = ['u', 'g', 'r', 'i', 'z', 'y']\n",
    "sn_lc['filter'] = [filter_names[f] for f in sn_lc['filter_code']]\n",
    "sn_lc = sn_lc.sort_values('mjd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_lightcurve(df, plot='mag', flux_col_names=None,\n",
    "                    title=None, marker='o', linestyle='none',\n",
    "                    colors=None, label_prefix='',\n",
    "                    **kwargs):\n",
    "    \"\"\"Plot a lightcurve from a DataFrame.\n",
    "    \"\"\"\n",
    "    # At lexigraphical order, if not wavelength order.\n",
    "    # Assume fixed list of filters.\n",
    "    filter_order = ['u', 'g', 'r', 'i', 'z', 'y']\n",
    "\n",
    "    if colors is None:\n",
    "        colors = {'u': 'violet', 'g': 'indigo', 'r': 'blue', 'i': 'green', 'z': 'orange', 'y': 'red'}\n",
    "    \n",
    "    if flux_col_names is not None:\n",
    "        flux_col, flux_err_col = flux_col_names\n",
    "    else:\n",
    "        if plot == 'flux':\n",
    "            flux_col = 'psFlux'\n",
    "            flux_err_col = 'psFluxErr'\n",
    "        else:\n",
    "            flux_col = 'mag'\n",
    "            flux_err_col = 'mag_err'\n",
    "        \n",
    "    for filt in filter_order:\n",
    "        this_filter = df.query(f'filter == \"{filt}\"')\n",
    "        if this_filter.empty:\n",
    "            continue\n",
    "        # This if sequence is a little silly.\n",
    "        plot_kwargs = {'linestyle': linestyle, 'marker': marker, 'color': colors[filt],\n",
    "                       'label': f'{label_prefix} {filt}'}\n",
    "        plot_kwargs.update(kwargs)\n",
    "\n",
    "        if flux_err_col in this_filter.columns:\n",
    "            plt.errorbar(this_filter['mjd'], this_filter[flux_col], this_filter[flux_err_col],\n",
    "                         **plot_kwargs)\n",
    "                        \n",
    "        else:\n",
    "            if marker is None:\n",
    "                plt.plot(this_filter['mjd'], this_filter[flux_col], **plot_kwargs)\n",
    "\n",
    "            else:\n",
    "                plot_kwargs.pop('linestyle')\n",
    "                plt.scatter(this_filter['mjd'], this_filter[flux_col], **plot_kwargs)\n",
    "\n",
    "\n",
    "\n",
    "    plt.xlabel('MJD')\n",
    "\n",
    "    if plot == 'flux':\n",
    "        plt.ylabel('psFlux [nJy]')\n",
    "    else:\n",
    "        # Ensure that y-axis decreases as one goes up\n",
    "        # Because plot_lightcurve could be called several times on the same axis,\n",
    "        # simply inverting is not correct.  We have to reverse a sorted list.\n",
    "        plt.ylim(sorted(plt.ylim(), reverse=True))\n",
    "        plt.ylabel('mag [AB]')\n",
    "\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 8))\n",
    "plot_lightcurve(sn_lc, plot='mag', linestyle='-', marker=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Match to DIAObject Catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ra, dec = sn['ra'], sn['dec']\n",
    "print(ra, dec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Match on RA, Dec\n",
    "sn_position = SkyCoord(sn['ra'], sn['dec'], unit='deg')\n",
    "\n",
    "diaObject_cat = diaObject.get_quantities(['ra', 'dec', 'diaObjectId'])\n",
    "diaObject_positions = SkyCoord(diaObject_cat['ra'], diaObject_cat['dec'], unit='deg')\n",
    "\n",
    "idx, sep2d, _ = sn_position.match_to_catalog_sky(diaObject_positions)\n",
    "print(f'Index: {idx} is {sep2d.to(u.arcsec)[0]:0.6f} away')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diaObjectId = diaObject_cat['diaObjectId'][idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How did we do with the lightcurve?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn_lc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can't use a direct filters = match in the GCR wrapper for the diaSrc table.\n",
    "# So we have to use a lambda function here to match the ID\n",
    "dia_lc = pd.DataFrame(diaSrc.get_quantities(['visit', 'mjd', 'psFlux', 'psFluxErr', 'mag', 'mag_err', 'filter'],\n",
    "                                            filters=[(lambda x: x == diaObjectId, 'diaObjectId')]))\n",
    "dia_lc = dia_lc.sort_values('mjd')                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 8))\n",
    "plot_lightcurve(sn_lc, plot='mag', linestyle='-', marker=None, label_prefix='Sim')\n",
    "plot_lightcurve(dia_lc, plot='mag', label_prefix='DIA')\n",
    "\n",
    "sim_date_range = [min(sn_lc['mjd']), max(sn_lc['mjd'])]\n",
    "sim_date_delta = sim_date_range[1] - sim_date_range[0]\n",
    "buffer_fraction = 0.05\n",
    "plot_date_range = sim_date_range\n",
    "plot_date_range[0] -= sim_date_delta * buffer_fraction\n",
    "plot_date_range[1] += sim_date_delta * buffer_fraction\n",
    "\n",
    "plt.xlim(plot_date_range)\n",
    "plt.ylim(25, 19);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The shape seems good.  But the calibration is magnitudes off.  It does look like it's a constant magnitude offset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn_lc.query('(60567 < mjd) & (mjd < 60568)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dia_lc.query('(60567 < mjd) & (mjd < 60568)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MJD for the same visits (recall `visit` == `obshistid`) are slightly different between the truth catalog and the DIA lightcurve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(60567.219180 - 60567.219782) * 24 * 3600"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Why the 52 second difference?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's match on `obshistid == visit` to get a matched set of dataframes that we can compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_lc = pd.merge(sn_lc, dia_lc, left_on='obshistid', right_on='visit', suffixes=('_sim', '_dia'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_lc['mjd'] = joint_lc['mjd_dia']\n",
    "joint_lc['filter'] = joint_lc['filter_sim']\n",
    "joint_lc['delta_mag'] = joint_lc['mag_dia'] - joint_lc['mag_sim']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_lc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_lightcurve(joint_lc, flux_col_names=('delta_mag', 'mag_err'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hmmm.... so not really a constant offset in magnitude.\n",
    "Nevertheless, I can only suspect there's some calibration or flux interpretation error somewhere."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If one wants to look at the postage stamps of this SN, get started with the examples in\n",
    "[dia_source_object_stamp.ipynb](dia_source_object_stamp.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "desc-stack",
   "language": "python",
   "name": "desc-stack"
  },
  "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": 2
}