Validation/image_quality_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSST DM Stack Image Quality Walkthrough\n",
    "\n",
    "<br>Author(s): **Keith Bechtol** ([@bechtol](https://github.com/LSSTScienceCollaborations/StackClub/issues/new?body=@bechtol))\n",
    "<br>Maintainer(s):   Douglas Tucker ([@douglasleetucker](https://github.com/LSSTScienceCollaborations/StackClub/issues/new?body=@douglasleetucker)), Alex Drlica-Wagner ([@kadrlica](https://github.com/LSSTScienceCollaborations/StackClub/issues/new?body=@kadrlica))\n",
    "<br>Last Verified to Run: **2021-09-11**\n",
    "<br>Verified Stack Release: **w_2021_33**\n",
    "\n",
    "This notebook provides a walkthrough of several image quality measurements in the LSST DM Stack. The walkthrough is largely based on the following resources, which should be consulted for further details:\n",
    "\n",
    "**For the DESC DC2 data set:**\n",
    "\n",
    "1. _The LSST DESC DC2 Simulated Sky Survey_: https://arxiv.org/abs/2010.05926\n",
    "2. _DESC DC2 Data Release Note_: https://arxiv.org/abs/2101.04855\n",
    "\n",
    "This notebook uses a recent processing of data (\"Run2.2i\") performed at the CC-IN2P3: \n",
    "1. The DESC DRP DC2 Run2.2i Processing Confluence page:  https://confluence.slac.stanford.edu/pages/viewpage.action?pageId=293309149\n",
    "\n",
    "\n",
    "**An older HSC data set came from:** \n",
    "\n",
    "1. _The Hyper Suprime-Cam Software Pipeline_: https://arxiv.org/abs/1705.06766\n",
    "2. _The first-year shear catalog of the Subaru Hyper-Suprime-Cam SSP Survey_: https://arxiv.org/abs/1705.06745\n",
    "3. Systematics Tests In LEnsing (STILE): https://github.com/msimet/Stile\n",
    "\n",
    "and used a reprocessing of data from the first HSC Strategic Survey Program data release 1 (HSC-SSP DR1): \n",
    "1. _HSC Public Release Webpage_: https://hsc-release.mtk.nao.ac.jp/doc/\n",
    "2. _HSC-SSP DR1 Paper_: http://ads.nao.ac.jp/abs/2018PASJ...70S...8A\n",
    "3. Processing information available here: https://confluence.lsstcorp.org/display/DM/S18+HSC+PDR1+reprocessing\n",
    "\n",
    "\n",
    "### Learning Objectives\n",
    "After working through and studying this notebook you should be able to\n",
    "1. Access PSF model sizes and ellipticities\n",
    "2. Access several different shape measurements with and without PSF correction for single-visit source catalogs and object catalogs derived from the full image stack\n",
    "3. Evaluate and visualize PSF model ellipticity residuals \n",
    "\n",
    "Other techniques that are demonstrated, but not empasized, in this notebook are\n",
    "1. Use `pandas` to compute aggregate statisitics and visualize the output \n",
    "2. Use `treecorr` to compute correlation functions\n",
    "\n",
    "### Logistics\n",
    "This notebook is intended to be run at `lsst-lsp-stable.ncsa.illinois.edu` or `data.lsst.cloud` from a local git clone of the [StackClub](https://github.com/LSSTScienceCollaborations/StackClub) repo.\n",
    "\n",
    "### Setup\n",
    "You can find the Stack version by using `eups list -s` on the terminal command line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Site, host, and stack version\n",
    "! echo $EXTERNAL_INSTANCE_URL\n",
    "! echo $HOSTNAME\n",
    "! eups list -s | grep lsst_distrib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matplotlib in Jupyter notebooks is a little fussy: let's ignore warnings from this and other modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning) # eg warning from mpl.use(“Agg”)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prelude: Data Sample\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lsst.daf.persistence as daf_persistence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the data set\n",
    "URL = os.getenv('EXTERNAL_INSTANCE_URL')\n",
    "if URL.endswith('data.lsst.cloud'): # IDF\n",
    "    repo = \"s3://butler-us-central1-dp01\"\n",
    "elif URL.endswith('ncsa.illinois.edu'): # NCSA\n",
    "    repo = \"/repo/dc2\"\n",
    "else:\n",
    "    raise Exception(f\"Unrecognized URL: {URL}\")\n",
    "    \n",
    "dataset='DC2'\n",
    "collection='2.2i/runs/DP0.1'\n",
    "dataId={'band':'i','visit':260,'detector':43}    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the Butler\n",
    "from lsst.daf.butler import Butler \n",
    "butler = Butler(repo,collections=collection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use that dataid to get the image and associated source catalog\n",
    "calexp = butler.get('calexp', **dataId)\n",
    "src = butler.get('src', **dataId)\n",
    "\n",
    "# Note that the calexp contains a PSF object\n",
    "psf = calexp.getPsf()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: PSF Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = psf.computeShape()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PSF Size\n",
    "\n",
    "This walkthrough focuses on the computation of PSF models and shape measurements using adaptive moments. Two common methods of computing the PSF size from adaptive moments are the trace radius and determinant radius.\n",
    "\n",
    "Adaptive second moments:\n",
    "\n",
    "$Q = \\begin{bmatrix}\n",
    "    I_{xx} & I_{xy} \\\\\n",
    "    I_{xy} & I_{yy}\n",
    "  \\end{bmatrix}$\n",
    "  \n",
    "The trace radius is defined as $\\rm{Tr}(Q) = \\sqrt{\\frac{I_{xx} + I_{yy}}{2}}$ and determinant radius is defined as $\\rm{Det(Q)} = (I_{xx} I_{yy} - I_{xy}I_{xy})^\\frac{1}{4}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i_xx, i_yy, i_xy = shape.getIxx(), shape.getIyy(), shape.getIxy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Put in a few assert statements to show exactly what the Stack is doing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trace radius\n",
    "assert np.isclose(shape.getTraceRadius(), np.sqrt((i_xx + i_yy) / 2.))\n",
    "print('trace radius =', shape.getTraceRadius())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determinant radius\n",
    "assert np.isclose(shape.getDeterminantRadius(), (i_xx * i_yy - i_xy**2)**(1. / 4.))\n",
    "print('determinant radius =', shape.getDeterminantRadius())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comment Regarding Units\n",
    "\n",
    "What are the units of the trace radius and determinant radius? How do we convert the trace radius or determinant radius into FWHM in arcseconds? To do this, we need to quickly preview some functionality available from the calibrated image and source catalog.\n",
    "\n",
    "First, let's make sure that the PSF shape second moments from the PSF model in the calibrated exposure (calexp) match the entry for the source catalog for the PSF at the source position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsst.geom import Point2D\n",
    "\n",
    "# PSF shape at coordinates of first indexed souce\n",
    "shape = psf.computeShape(Point2D(src['slot_Centroid_x'][0], src['slot_Centroid_y'][0])) \n",
    "print(shape.getIxx())\n",
    "# Compare to catalog entry for PSF shape\n",
    "print(src['slot_PsfShape_xx'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pixel scale in arcseconds\n",
    "calexp.getWcs().getPixelScale().asArcseconds()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Units of the second moments\n",
    "#src.schema.find('base_SdssShape_psf_xx').getField().getUnits()\n",
    "src.schema.find('slot_PsfShape_xx').getField().getUnits()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find that the second moments are in units of pixels$^2$, so the trace radius and determinant radius are in units of pixels. We found the pixel scale in arcseconds\n",
    "\n",
    "For a one-dimensional Gaussian, the FWHM = $2 \\sqrt{2 \\ln 2}$ RMS = 2.35 RMS \n",
    "\n",
    "**THIS APPROXIMATION FOR THE FWHM SHOULD BE VERIFIED**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Approximate FWHM in arcseconds\n",
    "fwhm = 2 * np.sqrt(2. * np.log(2)) * shape.getTraceRadius() * calexp.getWcs().getPixelScale().asArcseconds()\n",
    "print('FWHM = %.3f arcsec'%(fwhm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PSF Ellipticity\n",
    "\n",
    "**IMPORTANT:** Two conventions are commonly used in weak gravitational lensing to define ellipticity. The LSST Stack uses the \n",
    "\"distortion\" convention as oppose to the \"shear\" convention. See https://en.wikipedia.org/wiki/Gravitational_lensing_formalism#Measures_of_ellipticity for both definitions.\n",
    "\n",
    "The Science Requirements Document (https://docushare.lsst.org/docushare/dsweb/Get/LPM-17) also uses the distortion convention.\n",
    "\n",
    "Formalism ellipticities in the distortion convention:\n",
    "\n",
    "$e_1 = \\frac{I_{xx} - I_{yy}}{I_{xx} + I_{yy}}$\n",
    "\n",
    "$e_2 = \\frac{2 I_{xy}}{I_{xx} + I_{yy}}$\n",
    "\n",
    "$\\tan(2 \\theta) = \\frac{2 I_{xy}}{I_{xx} - I_{yy}} $\n",
    "\n",
    "$e = \\sqrt{e_1^2 + e_2^2}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Put in a few assert statements to explicitly show what the Stack is doing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ellipticity\n",
    "from lsst.afw.geom.ellipses import Quadrupole, SeparableDistortionTraceRadius\n",
    "\n",
    "q = Quadrupole(i_xx, i_yy, i_xy)\n",
    "s = SeparableDistortionTraceRadius(q)\n",
    "assert np.isclose(s.getE1(), (i_xx - i_yy) / (i_xx + i_yy)) # e1\n",
    "print('e1 =', s.getE1())\n",
    "assert np.isclose(s.getE2(), (2. * i_xy) / (i_xx + i_yy)) # e2\n",
    "print('e2 =', s.getE2())\n",
    "assert np.isclose(s.getEllipticity().getTheta(), np.arctan2(2. * i_xy, i_xx - i_yy) / 2.) # theta\n",
    "print('theta =', s.getEllipticity().getTheta())\n",
    "\n",
    "# An alternative way to compute the angle\n",
    "e1, e2 = s.getE1(), s.getE2() \n",
    "assert np.allclose(np.arctan2(e2, e1) / 2., np.arctan2(2. * i_xy, i_xx - i_yy) / 2.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For visualization purposes, let's evaluate the PSF model at grid of points across the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "x_array = np.arange(0, calexp.getDimensions()[0], 200)\n",
    "y_array = np.arange(0, calexp.getDimensions()[1], 200)\n",
    "xx, yy = np.meshgrid(x_array, y_array)\n",
    "\n",
    "print(calexp.getDimensions())\n",
    "\n",
    "size = []\n",
    "i_xx = []\n",
    "i_yy = []\n",
    "i_xy = []\n",
    "for x, y in zip(xx.flatten(), yy.flatten()):\n",
    "    point = Point2D(x, y)\n",
    "    shape = psf.computeShape(point)\n",
    "    size.append(shape.getTraceRadius())\n",
    "    i_xx.append(shape.getIxx())\n",
    "    i_yy.append(shape.getIyy())\n",
    "    i_xy.append(shape.getIxy())\n",
    "size = np.reshape(size, xx.shape)\n",
    "i_xx = np.reshape(i_xx, xx.shape)\n",
    "i_yy = np.reshape(i_yy, xx.shape)\n",
    "i_xy = np.reshape(i_xy, xx.shape)\n",
    "\n",
    "theta = np.arctan2(2. * i_xy, i_xx - i_yy) / 2.\n",
    "e1 = (i_xx - i_yy) / (i_xx + i_yy)\n",
    "e2 = (2. * i_xy) / (i_xx + i_yy)\n",
    "theta_alternate = np.arctan2(e2, e1) / 2.\n",
    "assert np.allclose(theta, theta_alternate)\n",
    "e = np.sqrt(e1**2 + e2**2)\n",
    "ex = e * np.cos(theta)\n",
    "ey = e * np.sin(theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and then plot the results in a few different ways"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20, 6))\n",
    "plt.subplots_adjust(wspace=0.5)\n",
    "\n",
    "plt.subplot(1, 4, 1)\n",
    "plt.quiver(xx, yy, ex, ey, e, headlength=0., headwidth=1., pivot='mid', width=0.01)\n",
    "#plt.quiver(xx, yy, scale=e, angles=phi, headlength=0., headwidth=1., pivot='mid', width=0.005)\n",
    "#colorbar = plt.colorbar(label='r$\\sqrt(e1^{2} + e2^{2})$')\n",
    "colorbar = plt.colorbar(label='e')\n",
    "plt.clim(0., 0.05)\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Ellipticity Sticks')\n",
    "\n",
    "plt.subplot(1, 4, 2)\n",
    "plt.scatter(xx, yy, c=e1, vmin=-0.05, vmax=0.05, cmap='bwr')\n",
    "colorbar = plt.colorbar(label='e1')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Value of e1')\n",
    "\n",
    "plt.subplot(1, 4, 3)\n",
    "plt.scatter(xx, yy, c=e2, vmin=-0.05, vmax=0.05, cmap='bwr')\n",
    "colorbar = plt.colorbar(label='e2')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Value of e2')\n",
    "\n",
    "plt.subplot(1, 4, 4)\n",
    "plt.scatter(xx, yy, c=size)\n",
    "colorbar = plt.colorbar(label='Trace Radius')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Ellipticity Modulus');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Source Catalog Shape Measurements\n",
    "\n",
    "The source catalogs include many measurements of source shapes and PSF model shapes. Let's look at the available columns relating to object shape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# I'm still learning how to use regular expressions...\n",
    "#import re\n",
    "#list(filter(re.compile(\".*shape*\"), src.schema.getOrderedNames()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Columns for shape measurements\n",
    "for name in src.schema.getOrderedNames():\n",
    "    if 'shape' in name.lower():\n",
    "        print(name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Columns for calibration measurements, including indicators of the stars using for PSF modeling and held in reserve for testing the PSF model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Columns for calibration measurements\n",
    "for name in src.schema.getOrderedNames():\n",
    "    if 'calib' in name.lower():\n",
    "        print(name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are also \"slots\" defined for default measurements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# How are the slots defined?\n",
    "for slot in ['slot_PsfShape', 'slot_Shape']:\n",
    "    print('%s -> %s'%(slot, src.schema.getAliasMap()[slot]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To figure out what these columns mean, look at https://arxiv.org/abs/1705.06745\n",
    "\n",
    "**Non PSF-corrected shapes**\n",
    "\n",
    "SDSS algorithm:\n",
    "\n",
    "`base_SdssShape_xx/xy/yy` = Adaptive moments in arcsec$^2$, SDSS algorithm\n",
    "\n",
    "`base_SdssShape_psf_xx/xy/yy` = Adaptive moments of PSF evaluated at object position in arcsec$^2$, SDSS algorithm\n",
    "\n",
    "HSM algorithm: \n",
    "\n",
    "_Galaxy shapes are estimated on the coadded i-band images using the re-Gaussianization PSF correction method (Hirata & Seljak 2003)... In the course of the re-Gaussianization PSF-correction method, corrections are applied to account for dilution of the observed shape by the PSF, including the non-Gaussianity of both the PSF and the galaxy surface brightness profiles._\n",
    "\n",
    "`ext_shapeHSM_HsmPsfMoments_xx/xy/yy` = Adaptive moments of PSF evaluated at object position in arcsec$^2$, HSM algorithm\n",
    "\n",
    "`ext_shapeHSM_HsmSourceMoments_xx/xy/yy` = Adaptive moments in arcsec$^2$, not PSF-corrected , HSM algorithm\n",
    "\n",
    "`ext_shapeHSM_HsmSourceMomentsRound_xx/xy/yy` = ??\n",
    "\n",
    "**Regaussianization shapes based on data alone, PSF-corrected**\n",
    "\n",
    "`ext_shapeHSM_HsmShapeRegauss_e1/e2` = distortion in sky coordinates estimated by regaussianization method defined in distortion, HSM algorithm (**NEED TO CHECK IF THESE ARE SKY OR INSTRUMENT COORDINATES**)\n",
    "\n",
    "`ext_shapeHSM_HsmShapeRegauss_sigma` = non-calibrated shape measurement noise, HSM algorithm\n",
    "\n",
    "`ext_shapeHSM_HsmShapeRegauss_resolution` = resolution of galaxy image defined in equation (4) of https://arxiv.org/abs/1705.06745, HSM algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With those definitions in mind, now verify that the PSF models from SDSS and HSM algorithms are consistent by comparing adaptive moments of the PSF. The plot below shows the ratio of adaptive moment sizes is very nearly unity. The plot also shows no trend in adaptive moments with respect to how well the source is resolved (i.e., source size), which is as expected since we are comparing PSF models as oppose to measurements for individual sources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(src['ext_shapeHSM_HsmShapeRegauss_resolution'], \n",
    "            src['base_SdssShape_psf_xx'] / src['ext_shapeHSM_HsmPsfMoments_xx'], s=1, label='xx')\n",
    "plt.scatter(src['ext_shapeHSM_HsmShapeRegauss_resolution'], \n",
    "            src['base_SdssShape_psf_yy'] / src['ext_shapeHSM_HsmPsfMoments_yy'], s=1, label='yy')\n",
    "plt.xlim(0., 1.)\n",
    "plt.ylim(1. - 2.e-3, 1. + 2.e-3)\n",
    "plt.legend(loc='upper left', markerscale=4)\n",
    "plt.xlabel('ext_shapeHSM_HsmShapeRegauss_resolution')\n",
    "plt.ylabel('base_SdssShape_psf / ext_shapeHSM_HsmPsfMoments');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To illustrate the difference between the PSF-corrected and non-PSF-corrected quantities within the HSM algorithm, now compare the ellipticity computed from the adaptive moments (not-PSF corrected) with the ellipticity including PSF-corrections. In the plot below, we see a clear trend with the amount of source resolution. For sources that are increasingly resolved, the PSF-corrected and non-PSF-corrected ellipticities converge. For sources that are not well resolved, the PSF-corrected and non-PSF-corrected ellipticities diverge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i_xx, i_yy, i_xy = src['ext_shapeHSM_HsmSourceMoments_xx'], src['ext_shapeHSM_HsmSourceMoments_yy'], src['ext_shapeHSM_HsmSourceMoments_xy']\n",
    "e1_non_psf_corrected = (i_xx - i_yy) / (i_xx + i_yy)\n",
    "ratio = e1_non_psf_corrected / src['ext_shapeHSM_HsmShapeRegauss_e1']\n",
    "#trace_radius = np.sqrt((src['base_SdssShape_xx'] + src['base_SdssShape_yy']) / 2.)\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(src['ext_shapeHSM_HsmShapeRegauss_resolution'], ratio, s=1)\n",
    "plt.xlim(0., 1.)\n",
    "plt.ylim(0., 1.2)\n",
    "plt.xlabel('ext_shapeHSM_HsmShapeRegauss_resolution')\n",
    "plt.ylabel('HsmSourceMoments_e1 / HsmShapeRegauss_e1');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intermission\n",
    "\n",
    "In the next section of the notebook, we'll use a full tract worth of data for some of the examples. It will take about 10 minutes to ingest the object catalogs from each patch in the tract. If you don't mind waiting, set the SHORTCUT variable below to `False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Browse the deepCoadd_meas schema\n",
    "coadd_dataid = {'band':'i', 'tract':4024, 'patch': (3) + 7*(4)}\n",
    "\n",
    "coadd_skymap = butler.get('skyMap')\n",
    "coadd_meas   = butler.get('deepCoadd_meas', dataId=coadd_dataid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set SHORTCUT = True for quick evaluation but lower statistics, \n",
    "# SHORTCUT = False to get all the objects from all the patches in the tract (~10 mins)\n",
    "# If you don't take the shortcut, youll need to do some uncommenting below as well.\n",
    "\n",
    "SHORTCUT = True\n",
    "\n",
    "# Pick a random tract and collect all the patches\n",
    "tract = coadd_dataid['tract']\n",
    "\n",
    "patch_array = []\n",
    "for ii in range(coadd_skymap.generateTract(tract).getNumPatches()[0]):\n",
    "    for jj in range(coadd_skymap.generateTract(tract).getNumPatches()[1]):\n",
    "        patch_array.append((ii) + 7*(jj))\n",
    "            \n",
    "tract_array = np.tile(tract, len(patch_array))\n",
    "\n",
    "if SHORTCUT:\n",
    "    # Only get three patches\n",
    "    #df_tract_patch = pd.DataFrame({'tract': [9936, 9936, 9936],\n",
    "    #                               'patch': ['0,0', '0,1', '0,2']})\n",
    "    df_tract_patch = pd.DataFrame({'tract': tract_array[:3],\n",
    "                                   'patch': patch_array[:3]})\n",
    "else:\n",
    "    # Get all the object catalogs from one tract\n",
    "    df_tract_patch = pd.DataFrame({'tract': tract_array,\n",
    "                                   'patch': patch_array})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_columns = ['id',\n",
    "                    'coord_ra', \n",
    "                    'coord_dec',\n",
    "                    'base_SdssCentroid_x',\n",
    "                    'base_SdssCentroid_y',\n",
    "                    'base_ClassificationExtendedness_value',\n",
    "                    'base_ClassificationExtendedness_flag',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_x',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_y',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_xx',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_yy',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_xy',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_flag',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_flag_no_pixels',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_flag_not_contained',\n",
    "                    'ext_shapeHSM_HsmPsfMoments_flag_parent_source',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_e1',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_e2',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_sigma',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_resolution',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_flag',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_flag_no_pixels',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_flag_not_contained',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_flag_parent_source',\n",
    "                    'ext_shapeHSM_HsmShapeRegauss_flag_galsim',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_x',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_y',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_xx',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_yy',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_xy',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_flag',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_flag_no_pixels',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_flag_not_contained',\n",
    "                    'ext_shapeHSM_HsmSourceMoments_flag_parent_source',\n",
    "                    'slot_Centroid_x',\n",
    "                    'slot_Centroid_y',\n",
    "                    'slot_Shape_xx',\n",
    "                    'slot_Shape_yy',\n",
    "                    'slot_Shape_xy',\n",
    "                    'slot_PsfShape_xx',\n",
    "                    'slot_PsfShape_yy',\n",
    "                    'slot_PsfShape_xy',\n",
    "                   ]    \n",
    "\n",
    "if dataset == 'HSC':\n",
    "    fluxCol         = 'slot_PsfFlux_flux'\n",
    "    fluxErrCol      = 'slot_PsfFlux_fluxSigma'\n",
    "    modelFluxCol    = 'slot_ModelFlux_flux'\n",
    "    modelFluxErrCol = 'slot_ModelFlux_fluxSigma'\n",
    "    calibUsedCol    = 'calib_psfUsed'\n",
    "    calibCandCol    = 'calib_psfCandidate'\n",
    "    calibResCol     = 'calib_psf_reserved'\n",
    "\n",
    "elif dataset == 'DC2':\n",
    "    fluxCol         = 'slot_PsfFlux_instFlux'\n",
    "    fluxErrCol      = 'slot_PsfFlux_instFluxErr'\n",
    "    modelFluxCol    = 'slot_ModelFlux_instFlux'\n",
    "    modelFluxErrCol = 'slot_ModelFlux_instFluxErr'\n",
    "    calibUsedCol    = 'calib_psf_used'\n",
    "    calibCandCol    = 'calib_psf_candidate'\n",
    "    calibResCol     = 'calib_psf_reserved'    \n",
    "    \n",
    "selected_columns += [fluxCol, fluxErrCol, modelFluxCol, modelFluxErrCol, calibUsedCol, calibCandCol, calibResCol]\n",
    "fluxRatioMin = 100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "coadd_array = []\n",
    "for ii in range(0, len(df_tract_patch)):\n",
    "    tract, patch = df_tract_patch['tract'][ii], df_tract_patch['patch'][ii] \n",
    "    print(tract, patch)\n",
    "    dataid = dict(coadd_dataid, tract=tract)\n",
    "    coadd_ref = butler.get('deepCoadd_ref', dataId=dataid)\n",
    "    coadd_meas = butler.get('deepCoadd_meas', dataId=dataid)\n",
    "    coadd_calib = butler.get('deepCoadd_calexp', dataId=dataid).getPhotoCalib()\n",
    "    \n",
    "    selected_rows = (coadd_ref['detect_isPrimary']\n",
    "                     & ~coadd_meas['base_SdssCentroid_flag']\n",
    "                     & ~coadd_meas['base_PixelFlags_flag_interpolated']\n",
    "                     & ~coadd_meas['base_PixelFlags_flag_saturated']\n",
    "                     & ~coadd_meas['base_PsfFlux_flag']\n",
    "                     & ~coadd_meas['modelfit_CModel_flag'])\n",
    "    \n",
    "    coadd_array.append(coadd_meas.asAstropy().to_pandas().loc[selected_rows, selected_columns])\n",
    "    coadd_array[-1]['detect_isPrimary'] = coadd_ref['detect_isPrimary'][selected_rows]\n",
    "    coadd_array[-1]['psf_mag'] = coadd_calib.instFluxToMagnitude(coadd_meas[selected_rows], 'base_PsfFlux')[:,0]\n",
    "    coadd_array[-1]['cm_mag'] = coadd_calib.instFluxToMagnitude(coadd_meas[selected_rows], 'modelfit_CModel')[:,0]\n",
    "    \n",
    "df_coadd = pd.concat(coadd_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you just waited 10 minutes to get all the patches, it could be a good idea to write out your tract-level catalog for re-use later. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not SHORTCUT:\n",
    "    outfile = os.path.expandvars('$HOME/DATA/temp.h5')\n",
    "    print(outfile)\n",
    "    df_coadd.to_hdf(outfile, 'df')\n",
    "    \n",
    "# Then, to read in this file, do:\n",
    "# df_coadd = pd.read_hdf(outfile, 'df')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.hexbin(np.degrees(df_coadd['coord_ra']), np.degrees(df_coadd['coord_dec']))\n",
    "#plt.hexbin(df_coadd['base_SdssCentroid_x'], df_coadd['base_SdssCentroid_y'])\n",
    "plt.xlim(plt.xlim()[::-1])\n",
    "plt.xlabel('RA (deg)')\n",
    "plt.ylabel('Dec (deg)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Characterization\n",
    "\n",
    "Show a few examples of how to characterize shape measurements."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ellipticity residuals for PSF calibration stars and reserve stars\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ellipticity(i_xx, i_yy, i_xy):\n",
    "    e1 = (i_xx - i_yy) / (i_xx + i_yy)\n",
    "    e2 = (2. * i_xy) / (i_xx + i_yy)\n",
    "    return e1, e2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ellipticity residuals\n",
    "e1_star, e2_star = ellipticity(df_coadd['slot_Shape_xx'], df_coadd['slot_Shape_yy'], df_coadd['slot_Shape_xy'])\n",
    "e1_psf, e2_psf = ellipticity(df_coadd['slot_PsfShape_xx'], df_coadd['slot_PsfShape_yy'], df_coadd['slot_PsfShape_xy'])\n",
    "e1_res, e2_res = e1_star - e1_psf, e2_star - e2_psf "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# Plots to investigate the S/N selection\n",
    "from matplotlib.colors import LogNorm\n",
    "plt.figure()\n",
    "plt.scatter(df_coadd['psf_mag'], df_coadd['psf_mag'] - df_coadd['cm_mag'], \n",
    "            c=df_coadd[fluxCol] / df_coadd[fluxErrCol],\n",
    "            vmin=1, vmax=500, s=1, norm=LogNorm())\n",
    "plt.colorbar(label='S/N')\n",
    "plt.xlim(16., 30.)\n",
    "plt.ylim(-0.1, 1.)\n",
    "plt.ylabel('psf - model')\n",
    "plt.xlabel('mag')\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(df_coadd[fluxCol]/df_coadd[fluxErrCol], \\\n",
    "         bins=100, range=(0, 1000), \\\n",
    "         histtype='step', lw=2, density=True, log=True)\n",
    "plt.xlim(0, 1000.)\n",
    "plt.xlabel('S/N')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cut_star = ((df_coadd[fluxCol] / df_coadd[fluxErrCol]) > 3.0) & (df_coadd['base_ClassificationExtendedness_value'] == 0)\n",
    "print(\"Number of stars:\",cut_star.sum())\n",
    "\n",
    "plt.figure()\n",
    "#plt.yscale('log')\n",
    "plt.hist(e1_res[cut_star] / e1_psf[cut_star], bins=np.linspace(-1., 1., 51), range=(-1., 1.), histtype='step', lw=2, density=True, label='e1')\n",
    "plt.hist(e2_res[cut_star] / e2_psf[cut_star], bins=np.linspace(-1., 1., 51), range=(-1., 1.), histtype='step', lw=2, density=True, label='e2')\n",
    "plt.xlim(-1., 1.)\n",
    "plt.legend(loc='upper right')\n",
    "plt.xlabel('e_star - e_psf / e_psf')\n",
    "plt.ylabel('PDF');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Repeat the exercise with the single-sensor, but compare the stars used for PSF modeling and those held in reserve for testing. In this case, there are only ~100 stars total on the single sensor so not enough statistics to see any more interesting trends. In practice, one would want to do this test with many more stars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ellipticity residuals\n",
    "e1_star, e2_star = ellipticity(src['slot_Shape_xx'], src['slot_Shape_yy'], src['slot_Shape_xy'])\n",
    "e1_psf, e2_psf = ellipticity(src['slot_PsfShape_xx'], src['slot_PsfShape_yy'], src['slot_PsfShape_xy'])\n",
    "e1_res, e2_res = e1_star - e1_psf, e2_star - e2_psf "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cut_psf_star = src[calibUsedCol]\n",
    "cut_reserved_star = src[calibResCol]\n",
    "\n",
    "print('Number of stars used for PSF modeling =', np.sum(cut_psf_star))\n",
    "print('Number of reserved stars =', np.sum(cut_reserved_star)) # Should be about 20% of the sample used for PSF modeling\n",
    "\n",
    "plt.figure()\n",
    "#plt.yscale('log')\n",
    "plt.hist(e1_res[cut_psf_star] / e1_psf[cut_psf_star], bins=np.linspace(-1., 1., 21), range=(-1., 1.), histtype='step', lw=2, density=True, label='PSF')\n",
    "plt.hist(e1_res[cut_reserved_star] / e1_psf[cut_reserved_star], bins=np.linspace(-1., 1., 21), range=(-1., 1.), histtype='step', lw=2, density=True, label='Reserved')\n",
    "plt.xlim(-1., 1.)\n",
    "plt.legend(loc='upper right')\n",
    "plt.xlabel('e_star - e_psf / e_psf')\n",
    "plt.ylabel('PDF');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute aggregate statistics across the focal plane / tract\n",
    "\n",
    "This example shows how to use pandas to compute aggregate statistics across the focal plane. In this case, we plot the mean PSF model size (trace radius) computed in spatial bins over the sensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert the src object to a pandas dataframe\n",
    "df_sensor = src.asAstropy().to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a helper function to aggregate statistics in two-dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aggregateStatistics2D(df, x_name, x_bins, y_name, y_bins, z_name, operation):\n",
    "    \"\"\"\n",
    "    This function preserves empty bins, which get filled by nan values\n",
    "    \"\"\"\n",
    "    x_centers = x_bins[:-1] + 0.5 * np.diff(x_bins)\n",
    "    y_centers = y_bins[:-1] + 0.5 * np.diff(y_bins)\n",
    "    grouped = df.groupby([pd.cut(df[y_name], bins=y_bins),\n",
    "                          pd.cut(df[x_name], bins=x_bins)])\n",
    "    xx, yy = np.meshgrid(x_centers, y_centers)\n",
    "    zz = grouped[['id', z_name]].agg(operation)[z_name].values.reshape(xx.shape)\n",
    "    return xx, yy, zz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a new column for the PSF model trace radius\n",
    "df_sensor['psf_trace_radius'] =  np.sqrt((src['slot_PsfShape_xx'] + src['slot_PsfShape_yy']) / 2.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next plot should look the same as the plot above that shows the PSF model size evaluated in a grid across the sensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bin_size = 200 # pixels\n",
    "x_bins = np.arange(0, calexp.getDimensions()[0] + bin_size, bin_size)\n",
    "y_bins = np.arange(0, calexp.getDimensions()[1] + bin_size, bin_size)\n",
    "\n",
    "xx, yy, zz = aggregateStatistics2D(df_sensor, \n",
    "                                   'slot_Centroid_x', x_bins, \n",
    "                                   'slot_Centroid_y', y_bins, \n",
    "                                   'psf_trace_radius', 'mean')\n",
    "\n",
    "plt.figure()\n",
    "plt.pcolormesh(xx, yy, zz)\n",
    "plt.colorbar(label='Trace Radius')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.xlim(0, calexp.getDimensions()[0])\n",
    "plt.ylim(0, calexp.getDimensions()[1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now repeat the aggregaion exercise for the coadd tract."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a new column for the PSF model trace radius\n",
    "df_coadd['psf_trace_radius'] =  np.sqrt((df_coadd['ext_shapeHSM_HsmPsfMoments_xx'] + df_coadd['ext_shapeHSM_HsmPsfMoments_yy']) / 2.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bin_size = 200 # pixels\n",
    "x_bins = np.arange(np.min(df_coadd.base_SdssCentroid_x), np.max(df_coadd.base_SdssCentroid_x) + bin_size, bin_size)\n",
    "y_bins = np.arange(np.min(df_coadd.base_SdssCentroid_y), np.max(df_coadd.base_SdssCentroid_y) + bin_size, bin_size)\n",
    "\n",
    "xx, yy, zz = aggregateStatistics2D(df_coadd, \n",
    "                                   'base_SdssCentroid_x', x_bins, \n",
    "                                   'base_SdssCentroid_y', y_bins, \n",
    "                                   'psf_trace_radius', 'mean')\n",
    "\n",
    "plt.figure()\n",
    "plt.pcolormesh(xx, yy, zz)\n",
    "plt.colorbar(label='Trace Radius')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.xlim(np.min(df_coadd.base_SdssCentroid_x), np.max(df_coadd.base_SdssCentroid_x))\n",
    "plt.ylim(np.min(df_coadd.base_SdssCentroid_y), np.max(df_coadd.base_SdssCentroid_y));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ellipticity residual correlation functions\n",
    "\n",
    "Example of how to evaluate correlation functions, in this case, for ellipticity residuals on the coadd. This section reproduces the analysis of `validate_drp` (https://github.com/lsst/validate_drp) in a condensed form for easier reading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import treecorr\n",
    "import astropy.units as u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ellipticity residuals\n",
    "e1_star, e2_star = ellipticity(df_coadd['slot_Shape_xx'], df_coadd['slot_Shape_yy'], df_coadd['slot_Shape_xy'])\n",
    "e1_psf, e2_psf = ellipticity(df_coadd['slot_PsfShape_xx'], df_coadd['slot_PsfShape_yy'], df_coadd['slot_PsfShape_xy'])\n",
    "e1_res, e2_res = e1_star - e1_psf, e2_star - e2_psf "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cut_star = ((df_coadd[fluxCol] / df_coadd[fluxErrCol]) > fluxRatioMin) & (df_coadd['base_ClassificationExtendedness_value'] == 0)\n",
    "cut_star = cut_star & ~np.isnan(e1_res) & ~np.isnan(e2_res)\n",
    "\n",
    "nbins=20\n",
    "min_sep=0.25\n",
    "max_sep=20\n",
    "sep_units='arcmin'\n",
    "verbose=False\n",
    "\n",
    "catTree = treecorr.Catalog(ra=df_coadd['coord_ra'][cut_star], dec=df_coadd['coord_dec'][cut_star], \n",
    "                           g1=e1_res[cut_star], g2=e2_res[cut_star],\n",
    "                           dec_units='radian', ra_units='radian')\n",
    "gg = treecorr.GGCorrelation(nbins=nbins, min_sep=min_sep, max_sep=max_sep,\n",
    "                            sep_units=sep_units,\n",
    "                            verbose=verbose)\n",
    "gg.process(catTree)\n",
    "r = np.exp(gg.meanlogr) * u.arcmin\n",
    "xip = gg.xip * u.Unit('')\n",
    "xip_err = np.sqrt(gg.varxip) * u.Unit('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.errorbar(r.value, xip, yerr=xip_err)\n",
    "plt.xlabel('Separation (arcmin)')\n",
    "plt.ylabel('Median Residual Ellipticity Correlation');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "LSST",
   "language": "python",
   "name": "lsst"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}