eelr_modeling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from functools import partial\n",
    "import pickle\n",
    "\n",
    "# for parallel computing\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "# use a better colormap and don't interpolate the pixels\n",
    "matplotlib.rc('image', cmap='inferno', interpolation='none', origin='lower')\n",
    "\n",
    "from scipy.signal import find_peaks\n",
    "from scipy.ndimage import gaussian_filter\n",
    "\n",
    "# our code for source separation\n",
    "import scarlet\n",
    "import scarlet.display\n",
    "# code to detect sources\n",
    "import sep\n",
    "\n",
    "# to open fits files\n",
    "from astropy.io import fits\n",
    "from astropy.table import Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set data path on local machine\n",
    "data_dir = '/home/czhao/Synced/Documents/PrincetonStuff/2019-20/spring/IW/hsc_images'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# source_list = Table.read('parent_sample/source_list.fits')\n",
    "# source_list = Table.read('good_ir_merged.fits')\n",
    "source_list = Table.read('parent_sample/EELR_HSCmag_from_SDSSspec.fits')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find directory for source at RA/DEC\n",
    "\n",
    "def get_source_i(ra, dec):\n",
    "    d = (source_list['RA'] - ra)**2 + (source_list['DEC'] - dec)**2\n",
    "    i = np.argmin(d)\n",
    "    if d[i] * 3600 > 1:\n",
    "        print(\"Closest match more than 1 arcsec away: Proceed with caution!\")\n",
    "    return i\n",
    "\n",
    "# most prominent candidate\n",
    "# ra, dec=37.77, -3.75\n",
    "# source_i = get_source_i(ra, dec)\n",
    "\n",
    "source_i = 2\n",
    "\n",
    "source_id = int(source_list[source_i][\"OBJID\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "\n",
    "def get_images_and_psfs(source_id):\n",
    "    # open files of the source\n",
    "    bands = ['G', 'R', 'I', 'Z', 'Y']\n",
    "    images = []\n",
    "    psfs = []\n",
    "    for b in bands:\n",
    "        file = glob.glob(\"{}/{:05d}/cutout_HSC-{}_*_src_*.fits\".format(data_dir, source_id, b))[0]\n",
    "        hdulist = fits.open(file)\n",
    "        images.append(hdulist[1].data)\n",
    "        hdulist.close()\n",
    "\n",
    "        file = glob.glob(\"{}/{:05d}/psf_HSC-{}_*_src_*.fits\".format(data_dir, source_id, b))[0]\n",
    "        hdulist = fits.open(file)\n",
    "        psfs.append(hdulist[0].data)\n",
    "        hdulist.close()\n",
    "    images = (np.array(images)[:,40:-40,40:-40]).copy()\n",
    "\n",
    "    # pad PSFs to the same shape\n",
    "    psf_height = max(psf.shape[0] for psf in psfs)\n",
    "    psf_width = max(psf.shape[1] for psf in psfs)\n",
    "    psfs = np.stack([np.pad(psf, (((psf_height - psf.shape[0]) // 2,), ((psf_width - psf.shape[1]) // 2,)))\n",
    "                     for psf in psfs])\n",
    "    \n",
    "    return images, psfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images, psfs = get_images_and_psfs(source_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get magnitude of line emitters\n",
    "def mag2amplitude(mags):\n",
    "    dlambda = np.array([.14, .14, .16, .13, .11])\n",
    "    fnu_Jy = 10**((48.6+mags)/-2.5)\n",
    "    photons_1Jy = 1.51e7 / dlambda\n",
    "    return photons_1Jy * fnu_Jy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from Ai-Lei spectral decomposition tables\n",
    "# this will require a read-out method to work with other sources\n",
    "# mags = np.array([26.3652646789765, 23.73486734207137, 21.898052328748093, 28, 21.05645664757267])\n",
    "\n",
    "# mags = np.array([source_list[source_i][f\"speclineMag_{band.lower()}\"] for band in bands])\n",
    "\n",
    "mags = np.nan_to_num(source_list[source_i]['MAG_AB_LINEONLY'], nan=30.0)\n",
    "\n",
    "print(mags)\n",
    "\n",
    "# last element (Y band) appears untrustworthy\n",
    "band_mask = [0,0,0,0,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# detect all sources in the image\n",
    "def makeCatalog(img):\n",
    "    detect = img.mean(axis=0)\n",
    "    bkg = sep.Background(detect)\n",
    "    #catalog = sep.extract(detect, 1.5, err=bkg.globalrms,deblend_nthresh=64,deblend_cont=3e-4)\n",
    "    catalog, segmap = sep.extract(detect, 1.2, err=bkg.globalrms,deblend_nthresh=64,deblend_cont=3e-4, segmentation_map=True)\n",
    "    bg_rms = np.array([sep.Background(band).globalrms for band in img])\n",
    "    return catalog, segmap, bg_rms\n",
    "\n",
    "def display_img(images, norm, catalog, ax=None):\n",
    "    # make and image and label all sources\n",
    "    img_rgb = scarlet.display.img_to_rgb(images, norm=norm)\n",
    "\n",
    "    if ax is None:\n",
    "        plt.figure(figsize=(6,6))\n",
    "        ax = plt.gca()\n",
    "    ax.imshow(img_rgb)\n",
    "    # Mark all of the sources from the detection cataog\n",
    "    for k, src in enumerate(catalog):\n",
    "        ax.text(src[\"x\"], src[\"y\"], str(k), color=\"w\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catalog, segmap, bg_rms = makeCatalog(images)\n",
    "# first define color stretch and convert 5 bands to RGB channels\n",
    "stretch = 1\n",
    "Q = 5\n",
    "norm = scarlet.display.AsinhMapping(minimum=0, stretch=stretch, Q=Q)\n",
    "display_img(images, norm, catalog)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display psfs\n",
    "psfs_min = min(psf.min() for psf in psfs)\n",
    "pnorm = scarlet.display.AsinhMapping(minimum=psfs_min, stretch=1e-2, Q=1)\n",
    "prgb = scarlet.display.img_to_rgb(psfs, norm=pnorm)\n",
    "plt.figure()\n",
    "plt.imshow(prgb)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define Frame and Observation:\n",
    "model_psf = scarlet.PSF(partial(scarlet.psf.gaussian, sigma=.8), shape=(None, 8, 8))\n",
    "bands = ['g', 'r', 'i', 'z', 'y']\n",
    "\n",
    "frame = scarlet.Frame(images.shape, psfs=model_psf, channels=bands)\n",
    "\n",
    "# no weight maps, use flat background noise variance instead\n",
    "# weights = np.ones_like(images) / (bg_rms[:,None,None]**2)\n",
    "observation = scarlet.Observation(images, psfs=psfs, channels=bands).match(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SEDConstraint(scarlet.Constraint):\n",
    "    def __init__(self, sed):\n",
    "        self._sed = sed\n",
    "\n",
    "    def __call__(self, X, step):\n",
    "        S = self._sed\n",
    "        # closest X that is in the direction of S\n",
    "        # allows for flux rescaling: only direction is constrained\n",
    "        if not np.ma.is_masked(S):\n",
    "            X[:] = np.maximum(np.dot(X, S) / np.dot(S, S) * S, 0)\n",
    "        else:\n",
    "            X_ = X[~S.mask]\n",
    "            S_ = S[~S.mask]\n",
    "            X[:][~S.mask] =  np.maximum(np.dot(X_, S_) / np.dot(S_, S_) * S_, 0)\n",
    "        return X\n",
    "    \n",
    "class RadialMaskConstraint(scarlet.Constraint):\n",
    "    def __init__(self, shape, pixel_center, R):\n",
    "        c, ny, nx = shape\n",
    "        dy = np.arange(ny) - pixel_center[0]\n",
    "        dx = np.arange(nx) - pixel_center[1]\n",
    "        dist2 = dy[:,None]**2 + dx[None,:]**2\n",
    "        self.mask = dist2 > R**2\n",
    "        \n",
    "    def __call__(self, X, step):\n",
    "        X[self.mask] = 0\n",
    "        X[:,:] = np.maximum(X, 0)\n",
    "        return X\n",
    "        \n",
    "    \n",
    "class EELRSource(scarlet.RandomSource):\n",
    "    \"\"\"Source to describe EELR\n",
    "    \n",
    "    It has a free-form morphology, possible constrained to be within R of the center\n",
    "    but its SED can be determined up to a constant.\n",
    "    \"\"\"\n",
    "    def __init__(self, frame, sky_coord, sed=None, R=None):\n",
    "        super().__init__(frame)\n",
    "        \n",
    "        center = np.array(frame.get_pixel(sky_coord), dtype=\"float\")\n",
    "        self.pixel_center = tuple(np.round(center).astype(\"int\"))\n",
    "        \n",
    "        if sed is not None:\n",
    "            self._parameters[0].constraint = SEDConstraint(sed)\n",
    "            self._parameters[0][:] = self._parameters[0].constraint(self._parameters[0], 0)\n",
    "        if R is not None:\n",
    "            self._parameters[1].constraint = RadialMaskConstraint(frame.shape, self.pixel_center, R)\n",
    "            self._parameters[1][:,:] = self._parameters[1].constraint(self._parameters[1], 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_center_source(catalog, dim):\n",
    "    eelr_host_ind = -1\n",
    "    closest_distsq = dim[0]**2 + dim[1]**2\n",
    "    center = (dim[0] / 2, dim[1] / 2)\n",
    "    for k, src in enumerate(catalog):\n",
    "        distsq = (src['y'] - center[0])**2 + (src['x'] - center[1])**2\n",
    "        if distsq < closest_distsq:\n",
    "            eelr_host_ind = k\n",
    "            closest_distsq = distsq\n",
    "    return eelr_host_ind\n",
    "\n",
    "def create_sources(catalog, eelr_host_ind, mags, mask, frame, observation):\n",
    "    sources = []\n",
    "    for k, src in enumerate(catalog):\n",
    "        if k == eelr_host_ind:\n",
    "            sources.append(scarlet.MultiComponentSource(frame, (src['y'], src['x']), observation, thresh=0.2, shifting=True))\n",
    "\n",
    "            # set mag for EELR source\n",
    "            mags = np.ma.masked_array(mags, mask=mask)\n",
    "            eelr_sed = mag2amplitude(mags)\n",
    "            sources.append(EELRSource(frame, (src['y'],src['x']), sed=eelr_sed, R=None))\n",
    "        else:\n",
    "            sources.append(scarlet.ExtendedSource(frame, (src['y'],src['x']), observation, shifting=True, thresh=0.5))\n",
    "    return sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eelr_host_ind = get_center_source(catalog, (images.shape[1], images.shape[2]))\n",
    "print(eelr_host_ind)\n",
    "print(catalog[eelr_host_ind]['x'], catalog[eelr_host_ind]['y'])\n",
    "sources = create_sources(catalog, eelr_host_ind, mags, band_mask, frame, observation)\n",
    "blend = scarlet.Blend(sources, observation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run the fitter\n",
    "%time blend.fit(200, e_rel=1e-5)\n",
    "print(\"scarlet ran for {0} iterations to logL = {1}\".format(len(blend.loss), -blend.loss[-1]))\n",
    "plt.plot(-np.array(blend.loss))\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('log-Likelihood')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = scarlet.display.show_scene(sources, observation=observation, norm=norm, show_observed=True, show_rendered=True, show_residual=True)\n",
    "fig.savefig(f\"{source_id}_scarlet_scene.png\", dpi=200, bbox_inches=\"tight\")\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = scarlet.display.show_sources(sources, observation, show_observed=True, show_rendered=True, norm=norm)\n",
    "fig.savefig(f\"{source_id}_scarlet_sources.png\", dpi=200, bbox_inches=\"tight\")\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = f\"{source_id:05d}_\"\n",
    "print(prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_asinh_stretch(img):\n",
    "    stretch = 0.05\n",
    "    Q = 5\n",
    "    norm = scarlet.display.AsinhMapping(minimum=0, stretch=stretch, Q=Q)\n",
    "    plt.imshow(scarlet.display.img_to_rgb(img, norm=norm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_one_eelr_sample(catalog, eelr_host_ind, mags, mask, frame, observation):\n",
    "    sources = create_sources(catalog, eelr_host_ind, mags, mask, frame, observation)\n",
    "    blend = scarlet.Blend(sources, observation)\n",
    "    blend.fit(200, e_rel=1e-5)\n",
    "#     print(\"scarlet ran for {0} iterations to logL = {1}\".format(len(blend.loss), -blend.loss[-1]))\n",
    "    \n",
    "    # EELR host and source, full model, and logL\n",
    "    return sources[eelr_host_ind], sources[eelr_host_ind+1], blend.get_model(), -blend.loss[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eelr_samples = []\n",
    "for _ in range(2):\n",
    "    eelr_samples.append(do_one_eelr_sample(catalog, eelr_host_ind, mags, band_mask, frame, observation)[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scarlet.display.show_sources(eelr_samples, observation, show_observed=True, show_rendered=True, norm=norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_snr(X):\n",
    "    sig = (X**2).mean(axis=0)\n",
    "    noise = X.var(axis=0)\n",
    "    snr = sig / noise\n",
    "    snr[sig==0] = 0\n",
    "    snr = np.nan_to_num(snr, nan=np.nanmax(snr))\n",
    "    return snr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create heatmap of average (since multiple bands) SNR of each EELR pixel\n",
    "\n",
    "eelr_models = np.stack([sample.get_model() for sample in eelr_samples])\n",
    "\n",
    "snr_per_band = compute_snr(eelr_models)\n",
    "snr = snr_per_band.mean(axis=0)\n",
    "\n",
    "plt.imshow(snr)\n",
    "plt.colorbar()\n",
    "plt.savefig(f\"{prefix}model_snr.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create heatmap of average morphology\n",
    "\n",
    "eelr_morphs = np.stack([sample.morph for sample in eelr_samples])\n",
    "mean_morph = eelr_morphs.mean(axis=0)\n",
    "\n",
    "plt.imshow(mean_morph)\n",
    "plt.colorbar()\n",
    "plt.savefig(f\"{prefix}morph_avg.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create heatmap of morphology SNR\n",
    "\n",
    "morph_snr = compute_snr(eelr_morphs)\n",
    "\n",
    "plt.imshow(morph_snr)\n",
    "plt.colorbar()\n",
    "plt.savefig(f\"{prefix}morph_snr.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unmasked Y band"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_unmasked_y_samples(catalog, eelr_host_ind, mags, num_samples, frame, observation):\n",
    "    unmasked_y_host_samples = []\n",
    "    unmasked_y_eelr_samples = []\n",
    "    unmasked_y_model_samples = []\n",
    "    unmasked_y_logL = []\n",
    "    for _ in range(num_samples):\n",
    "        y = np.random.uniform(22, 27)\n",
    "#         print(f\"y = {y}:\")\n",
    "        this_mags = mags\n",
    "        this_mags[4] = y\n",
    "        host_sample, eelr_sample, model_sample, logL = do_one_eelr_sample(catalog, eelr_host_ind, this_mags, [0, 0, 0, 0, 0], frame, observation)\n",
    "        unmasked_y_host_samples.append(host_sample)\n",
    "        unmasked_y_eelr_samples.append(eelr_sample)\n",
    "        unmasked_y_model_samples.append(model_sample)\n",
    "        unmasked_y_logL.append(logL)\n",
    "    return unmasked_y_host_samples, unmasked_y_eelr_samples, unmasked_y_model_samples, unmasked_y_logL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 10\n",
    "unmasked_y_host_samples, unmasked_y_eelr_samples, unmasked_y_model_samples, unmasked_y_logL = get_unmasked_y_samples(catalog, eelr_host_ind, mags, num_samples, frame, observation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "scarlet.display.show_sources(unmasked_y_eelr_samples[:5], observation, show_observed=True, show_rendered=True, norm=norm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Averaging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_mean_and_var(samples, weights):\n",
    "    # Computes likelihood-weighted mean and variance\n",
    "#     samples = np.stack(samples)\n",
    "    mean = np.average(samples, axis=0, weights=weights)\n",
    "    w_normed = weights / weights.sum()\n",
    "    # variance with reliability weights\n",
    "    var_nonadj = (np.square(w_normed) * np.moveaxis((samples - mean)**2, 0, -1)).sum(axis=-1)\n",
    "    v2 = np.square(w_normed).sum(axis=-1)\n",
    "    var = var_nonadj / (1 - v2)\n",
    "    return mean, var\n",
    "\n",
    "def get_outliers(arr, thresh=5):\n",
    "    # Returns mask of outliers that are more than thresh below the median\n",
    "    mask = np.zeros(arr.shape)\n",
    "    last_size = -1\n",
    "    cur_size = 0\n",
    "    while cur_size > last_size:\n",
    "        last_size = cur_size\n",
    "        med = np.median(arr)\n",
    "        mask = arr < med - thresh\n",
    "        cur_size = np.count_nonzero(mask)\n",
    "    return mask\n",
    "\n",
    "def zero_borders(X):\n",
    "    X[0, :] = 0\n",
    "    X[-1, :] = 0\n",
    "    X[:, 0] = 0\n",
    "    X[:, -1] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "host_morphs = np.array([sample.components[1].morph for sample in unmasked_y_host_samples])\n",
    "# host_morphs = np.array([sample.morph for sample in unmasked_y_host_samples])\n",
    "eelr_morphs = np.array([sample.morph for sample in unmasked_y_eelr_samples])\n",
    "models = np.array(unmasked_y_model_samples)\n",
    "unmasked_y_logL = np.array(unmasked_y_logL)\n",
    "\n",
    "# Drop likelihoods that are more than 7 orders of magnitude (in units of e) less than the max\n",
    "mask = unmasked_y_logL > (unmasked_y_logL.max() - 7)\n",
    "print(f\"Using {np.count_nonzero(mask)} samples.\")\n",
    "scaled_likelihoods = np.exp(unmasked_y_logL[mask] - max(unmasked_y_logL[mask]))\n",
    "host_morphs = host_morphs[mask]\n",
    "eelr_morphs = eelr_morphs[mask]\n",
    "models = models[mask]\n",
    "\n",
    "host_avg, host_var = weighted_mean_and_var(host_morphs, scaled_likelihoods)\n",
    "eelr_avg, eelr_var = weighted_mean_and_var(eelr_morphs, scaled_likelihoods)\n",
    "model_avg, model_var = weighted_mean_and_var(models, scaled_likelihoods)\n",
    "\n",
    "# zero out borders to remove artifacts due to PSF\n",
    "zero_borders(host_avg)\n",
    "zero_borders(eelr_avg)\n",
    "\n",
    "# low pass filter\n",
    "eelr_avg_filtered = gaussian_filter(eelr_avg, sigma=2)\n",
    "\n",
    "# plt.imshow(host_avg)\n",
    "plot_asinh_stretch(host_avg)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "# plt.imshow(host_var)\n",
    "plot_asinh_stretch(host_var)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "display_img(model_avg, norm, catalog)\n",
    "plt.show()\n",
    "\n",
    "plt.imshow(eelr_avg)\n",
    "plt.colorbar()\n",
    "plt.savefig(f\"{prefix}morph_avg.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "\n",
    "plt.imshow(eelr_avg_filtered)\n",
    "plt.colorbar()\n",
    "plt.savefig(f\"{prefix}morph_avg_filtered.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "\n",
    "plt.imshow(eelr_var)\n",
    "plt.colorbar()\n",
    "plt.savefig(f\"{prefix}morph_var.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, morph in enumerate(eelr_morphs):\n",
    "    print(i)\n",
    "    plt.imshow(morph)\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# griz Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"gp_models.pickle\", \"rb\") as infile:\n",
    "    z_to_griz_models = pickle.load(infile)\n",
    "\n",
    "print(z_to_griz_models.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_one_griz_sample(z_):\n",
    "    z_ = np.array([[z_]])\n",
    "    g = 25\n",
    "    r = g - z_to_griz_models[\"g-r\"].sample_y(z_)[0, 0]\n",
    "    i = r - z_to_griz_models[\"r-i\"].sample_y(z_)[0, 0]\n",
    "    z = i - z_to_griz_models[\"i-z\"].sample_y(z_)[0, 0]\n",
    "    return [g, r, i, z]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Large Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_df = pd.read_csv(\"parent_sample/FIRST.csv\", comment=\"#\")\n",
    "first_ids = first_df[\"objid\"].to_numpy()\n",
    "\n",
    "FIG_DIR = \"eelr_outputs_sdss_fixedvar/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_eelr_fig(result):    \n",
    "    if result is None: return\n",
    "\n",
    "    stretch = 1\n",
    "    Q = 5\n",
    "    norm = scarlet.display.AsinhMapping(minimum=0, stretch=stretch, Q=Q)\n",
    "    bands = ['g', 'r', 'i', 'z', 'y']\n",
    "    model_psf = scarlet.PSF(partial(scarlet.psf.gaussian, sigma=.8), shape=(None, 8, 8))\n",
    "    \n",
    "    # load in source\n",
    "    images, psfs = get_images_and_psfs(result[\"id\"])\n",
    "    catalog, segmap, bg_rms = makeCatalog(images)\n",
    "\n",
    "    # define scarlet frame and observation\n",
    "    frame = scarlet.Frame(images.shape, psfs=model_psf, channels=bands)\n",
    "    observation = scarlet.Observation(images, psfs=psfs, channels=bands).match(frame)\n",
    "    \n",
    "    fig, axs = plt.subplots(1, 3, figsize=(15, 5), gridspec_kw={'width_ratios': [1, 1, 1.1]})\n",
    "#     fig.suptitle(f\"Source {result['id']}\")\n",
    "    axs[0].set_title(\"Observation\")\n",
    "    axs[1].set_title(\"Rendered Model\")\n",
    "    axs[2].set_title(\"EELR Morphology\")\n",
    "    display_img(images, norm, catalog, axs[0])\n",
    "    display_img(observation.render(result[\"model_avg\"]), norm, catalog, axs[1])\n",
    "    im = axs[2].imshow(result[\"eelr_avg\"])\n",
    "    fig.colorbar(im, ax=axs[2], fraction=0.04)\n",
    "\n",
    "    # footnote\n",
    "#     if result['id'] in first_ids:\n",
    "#         footnote = \"In FIRST catalog.\"\n",
    "#     else:\n",
    "#         footnote = \"Not in FIRST catalog.\"\n",
    "#     plt.figtext(0.11, 0.07, footnote, fontsize=\"small\", fontstyle=\"italic\", ha=\"left\", va=\"bottom\")\n",
    "\n",
    "    prefix = f\"{result['id']:05d}_\"\n",
    "    plt.savefig(f\"{FIG_DIR}{prefix}eelr.png\", dpi=200, bbox_inches=\"tight\")\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "import traceback\n",
    "\n",
    "def safe_func(func):\n",
    "    @functools.wraps(func)\n",
    "    def wrapper(*args, **kwargs):\n",
    "        try:\n",
    "            return func(*args, **kwargs)\n",
    "        except Exception:\n",
    "            traceback.print_exc()\n",
    "            return None\n",
    "    return wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def model_eelr(objid, num_samples=10, gen_fig=True, use_sdss=False, z_noise=0):\n",
    "\n",
    "    # FIXME: scarlet breaks on 178\n",
    "#     if objid == 178: return None\n",
    "    \n",
    "    source_info = source_list[source_list[\"OBJID\"].astype(int) == objid]\n",
    "    source_info = {k: source_info[k][0] for k in source_info.columns}\n",
    "    \n",
    "    bands = ['g', 'r', 'i', 'z', 'y']\n",
    "    model_psf = scarlet.PSF(partial(scarlet.psf.gaussian, sigma=.8), shape=(None, 8, 8))\n",
    "    \n",
    "    # load in source\n",
    "    images, psfs = get_images_and_psfs(objid)\n",
    "    catalog, segmap, bg_rms = makeCatalog(images)\n",
    "    if use_sdss:\n",
    "        mags = np.nan_to_num(source_info['MAG_AB_LINEONLY'], nan=30.0)\n",
    "    else:\n",
    "        z = source_info[\"Z\"]\n",
    "        z += np.random.uniform(-z_noise, z_noise)\n",
    "\n",
    "    # define scarlet frame and observation\n",
    "    frame = scarlet.Frame(images.shape, psfs=model_psf, channels=bands)\n",
    "    observation = scarlet.Observation(images, psfs=psfs, channels=bands).match(frame)\n",
    "\n",
    "    # sampling\n",
    "    eelr_host_ind = get_center_source(catalog, (images.shape[1], images.shape[2]))\n",
    "    if use_sdss:\n",
    "        unmasked_y_host_samples, unmasked_y_eelr_samples, unmasked_y_model_samples, unmasked_y_logL = get_unmasked_y_samples(catalog, eelr_host_ind, mags, num_samples, frame, observation)\n",
    "    else:\n",
    "        unmasked_y_host_samples = []\n",
    "        unmasked_y_eelr_samples = []\n",
    "        unmasked_y_model_samples = []\n",
    "        mags_samples = []\n",
    "        unmasked_y_logL = []\n",
    "        for _ in range(num_samples):\n",
    "            y = np.random.uniform(18, 26)\n",
    "            mags = do_one_griz_sample(z)\n",
    "            mags.append(y)\n",
    "            mags_samples.append(mags)\n",
    "            host_sample, eelr_sample, model_sample, logL = do_one_eelr_sample(catalog, eelr_host_ind, mags, [0, 0, 0, 0, 0], frame, observation)\n",
    "            unmasked_y_host_samples.append(host_sample)\n",
    "            unmasked_y_eelr_samples.append(eelr_sample)\n",
    "            unmasked_y_model_samples.append(model_sample)\n",
    "            unmasked_y_logL.append(logL)\n",
    "    \n",
    "    # morphology likelihood-weighted mean and variance\n",
    "    host_morphs = np.array([sample.components[1].morph for sample in unmasked_y_host_samples])\n",
    "    # host_morphs = np.array([sample.morph for sample in unmasked_y_host_samples])\n",
    "    eelr_morphs = np.array([sample.morph for sample in unmasked_y_eelr_samples])\n",
    "    models = np.array(unmasked_y_model_samples)\n",
    "    unmasked_y_logL = np.array(unmasked_y_logL)\n",
    "    # Drop likelihoods that are more than 7 orders of magnitude (in units of e) less than the max\n",
    "    mask = unmasked_y_logL > (unmasked_y_logL.max() - 7)\n",
    "    used_samples = np.count_nonzero(mask)\n",
    "    print(f\"Using {used_samples} samples.\")\n",
    "    scaled_likelihoods = np.exp(unmasked_y_logL[mask] - max(unmasked_y_logL[mask]))\n",
    "    host_morphs = host_morphs[mask]\n",
    "    eelr_morphs = eelr_morphs[mask]\n",
    "    models = models[mask]\n",
    "    host_avg, host_var = weighted_mean_and_var(host_morphs, scaled_likelihoods)\n",
    "    eelr_avg, eelr_var = weighted_mean_and_var(eelr_morphs, scaled_likelihoods)\n",
    "    model_avg, model_var = weighted_mean_and_var(models, scaled_likelihoods)\n",
    "    \n",
    "    result = {\"id\": objid,\n",
    "              \"host_avg\": host_avg,\n",
    "              \"host_var\": host_var,\n",
    "              \"eelr_avg\": eelr_avg,\n",
    "              \"eelr_var\": eelr_var,\n",
    "              \"model_avg\": model_avg,\n",
    "              \"model_var\": model_var,\n",
    "              \"num_samples\": used_samples,\n",
    "           }\n",
    "    if not use_sdss:\n",
    "        result[\"z\"] = z\n",
    "        mags_samples = np.array(mags_samples)[mask]\n",
    "        result[\"mags_avg\"], result[\"mags_var\"] = weighted_mean_and_var(mags_samples, scaled_likelihoods)\n",
    "    \n",
    "    if gen_fig:\n",
    "        gen_eelr_fig(result)\n",
    "    \n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sample = \"full\"\n",
    "\n",
    "if sample == \"first\":\n",
    "    source_inds = first_ids\n",
    "elif sample == \"full\":\n",
    "    source_inds = source_list[\"OBJID\"].astype(int)\n",
    "\n",
    "results = Parallel(n_jobs=3, verbose=11)(delayed(safe_func(model_eelr))(i, num_samples=50, use_sdss=True, z_noise=0) for i in source_inds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save/Load Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"results_sdss_fixedvar.pickle\", \"wb\") as outfile:\n",
    "    pickle.dump(results, outfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"results_sdss.pickle\", \"rb\") as infile:\n",
    "    results = pickle.load(infile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Regenerate EELR Figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sample = \"full\"\n",
    "\n",
    "if sample == \"first\":\n",
    "    source_inds = first_ids\n",
    "elif sample == \"full\":\n",
    "    source_inds = source_list[\"OBJID\"].astype(int)\n",
    "\n",
    "Parallel(n_jobs=6, verbose=11)(delayed(gen_eelr_fig)(res) for res in results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Differential Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"results_sdss_fixedvar.pickle\", \"rb\") as infile:\n",
    "    sdss_results = pickle.load(infile)\n",
    "\n",
    "with open(\"results_gp_fixedvar.pickle\", \"rb\") as infile:\n",
    "    gp_results = pickle.load(infile)\n",
    "\n",
    "with open(\"results_gp_noisyz_fixedvar.pickle\", \"rb\") as infile:\n",
    "    gp_noisyz_results = pickle.load(infile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import pearsonr\n",
    "\n",
    "sdss_mse, sdss_totinten, sdss_meanintenvar = [], [], []\n",
    "gp_mse, gp_totinten, gp_meanintenvar = [], [], []\n",
    "gp_noisyz_mse, gp_noisyz_totinten, gp_noisyz_meanintenvar = [], [], []\n",
    "\n",
    "# correlations = []\n",
    "\n",
    "for i in range(len(sdss_results)):\n",
    "    if sdss_results[i] is None: continue\n",
    "    assert(sdss_results[i][\"id\"] == gp_results[i][\"id\"])\n",
    "    \n",
    "    images, psfs = get_images_and_psfs(sdss_results[i][\"id\"])\n",
    "    sdss_mse.append(np.square(sdss_results[i][\"model_avg\"] - images).sum() / images.size)\n",
    "    gp_mse.append(np.square(gp_results[i][\"model_avg\"] - images).sum() / images.size)\n",
    "    gp_noisyz_mse.append(np.square(gp_noisyz_results[i][\"model_avg\"] - images).sum() / images.size)\n",
    "    \n",
    "    mid_r = sdss_results[i][\"eelr_avg\"].shape[0] // 2\n",
    "    mid_c = sdss_results[i][\"eelr_avg\"].shape[1] // 2\n",
    "    win_half = 15\n",
    "    cropped_sdss_eelr_avg = sdss_results[i][\"eelr_avg\"][mid_r-win_half:mid_r+win_half, mid_c-win_half:mid_c+win_half]\n",
    "    cropped_gp_eelr_avg = gp_results[i][\"eelr_avg\"][mid_r-win_half:mid_r+win_half, mid_c-win_half:mid_c+win_half]\n",
    "    cropped_gp_noisyz_eelr_avg = gp_noisyz_results[i][\"eelr_avg\"][mid_r-win_half:mid_r+win_half, mid_c-win_half:mid_c+win_half]\n",
    "    \n",
    "    sdss_totinten.append(cropped_sdss_eelr_avg.sum())\n",
    "    gp_totinten.append(cropped_gp_eelr_avg.sum())\n",
    "    gp_noisyz_totinten.append(cropped_gp_noisyz_eelr_avg.sum())\n",
    "    \n",
    "    sdss_meanintenvar.append(sdss_results[i][\"eelr_var\"].mean())\n",
    "    gp_meanintenvar.append(gp_results[i][\"eelr_var\"].mean())\n",
    "    gp_noisyz_meanintenvar.append(gp_noisyz_results[i][\"eelr_var\"].mean())\n",
    "    \n",
    "#     correlations.append(pearsonr(cropped_sdss_eelr_avg.ravel(), cropped_gp_eelr_avg.ravel())[0])\n",
    "#     correlations.append(pearsonr(sdss_results[i][\"eelr_avg\"].ravel(), gp_results[i][\"eelr_avg\"].ravel())[0]\n",
    "\n",
    "sdss_mse, sdss_totinten, sdss_meanintenvar = np.array(sdss_mse), np.array(sdss_totinten), np.array(sdss_meanintenvar)\n",
    "gp_mse, gp_totinten, gp_meanintenvar = np.array(gp_mse), np.array(gp_totinten), np.array(gp_meanintenvar)\n",
    "gp_noisyz_mse, gp_noisyz_totinten, gp_noisyz_meanintenvar = np.array(gp_noisyz_mse), np.array(gp_noisyz_totinten), np.array(gp_noisyz_meanintenvar)\n",
    "# correlations = np.array(correlations)\n",
    "\n",
    "var_mask = (~np.isnan(sdss_meanintenvar)) & (~np.isnan(gp_meanintenvar)) & (~np.isnan(gp_noisyz_meanintenvar))\n",
    "sdss_meanintenvar = sdss_meanintenvar[var_mask]\n",
    "gp_meanintenvar = gp_meanintenvar[var_mask]\n",
    "gp_noisyz_meanintenvar = gp_noisyz_meanintenvar[var_mask]\n",
    "print(len(sdss_meanintenvar))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mse1, totinten1, meanintenvar1 = sdss_mse, sdss_totinten, sdss_meanintenvar\n",
    "\n",
    "# mse2, totinten2, meanintenvar2 = gp_mse, gp_totinten, gp_meanintenvar\n",
    "# gp_name = \"GP\"\n",
    "\n",
    "mse2, totinten2, meanintenvar2 = gp_noisyz_mse, gp_noisyz_totinten, gp_noisyz_meanintenvar\n",
    "gp_name = \"Noisy GP\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(mse1, mse2)\n",
    "plt.loglog([mse1.min(), mse2.max()], [mse1.min(), mse2.max()], color=\"red\")\n",
    "plt.xlabel(\"SDSS MSE\")\n",
    "plt.ylabel(f\"{gp_name} MSE\")\n",
    "plt.savefig(\"sdss_vs_gp_noisyz_mse.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "\n",
    "print(\"SDSS Median MSE:\", np.median(mse1))\n",
    "print(f\"{gp_name} Median MSE:\", np.median(mse2))\n",
    "\n",
    "print(f\"Proportion where SDSS MSE < {gp_name} MSE:\", np.count_nonzero(mse1 < mse2) / len(mse1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(totinten1, totinten2)\n",
    "plt.loglog([totinten2.min(), totinten1.max()], [totinten2.min(), totinten1.max()], color=\"red\")\n",
    "plt.xlabel(\"SDSS Total Intensity\")\n",
    "plt.ylabel(f\"{gp_name} Total Intensity\")\n",
    "plt.savefig(\"sdss_vs_gp_noisyz_intensity.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "\n",
    "print(\"SDSS Median Total Intensity:\", np.median(totinten1))\n",
    "print(f\"{gp_name} Median Total Intensity:\", np.median(totinten2))\n",
    "\n",
    "print(f\"Proportion where SDSS intensity > {gp_name} intensity:\", np.count_nonzero(totinten1 > totinten2) / len(totinten1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(meanintenvar1, meanintenvar2)\n",
    "plt.loglog([meanintenvar1.min(), meanintenvar2.max()], [meanintenvar1.min(), meanintenvar2.max()], color=\"red\")\n",
    "plt.xlabel(\"SDSS MPIV\")\n",
    "plt.ylabel(f\"{gp_name} MPIV\")\n",
    "plt.savefig(\"sdss_vs_gp_noisyz_intenvar.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "\n",
    "print(\"SDSS Median MPIV:\", np.median(meanintenvar1))\n",
    "print(f\"{gp_name} Median MPIV:\", np.median(meanintenvar2))\n",
    "\n",
    "print(f\"Proportion where SDSS intensity var < {gp_name} intensity var:\", np.count_nonzero(meanintenvar1 < meanintenvar2) / len(meanintenvar1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.hist(correlations)\n",
    "# plt.xlabel(\"Correlation\")\n",
    "# plt.ylabel(\"Frequency\")\n",
    "# # plt.savefig(\"sdss_vs_gp_noisyz_correlation.png\", dpi=200, bbox_inches=\"tight\")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EELR Classification (deprecated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_ellipse(morph):\n",
    "    bkg = sep.Background(morph)\n",
    "    objects = sep.extract(morph, 1.2, err=bkg.globalrms)\n",
    "    if len(objects) != 1:\n",
    "        print(f\"Detected {len(objects)} objects in host!\")\n",
    "    i = max(list(range(len(objects))), key=lambda i: objects[\"cflux\"][i])\n",
    "    return objects[\"x\"][i], objects[\"y\"][i], objects[\"a\"][i], objects[\"b\"][i], objects[\"theta\"][i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.patches import Ellipse\n",
    "\n",
    "host_ellipse = fit_ellipse(host_avg)\n",
    "print(host_ellipse)\n",
    "\n",
    "plot_asinh_stretch(host_avg)\n",
    "e = Ellipse(xy=(host_ellipse[0], host_ellipse[1]),\n",
    "                width=6*host_ellipse[2],\n",
    "                height=6*host_ellipse[3],\n",
    "                angle=host_ellipse[4] * 180. / np.pi)\n",
    "e.set_facecolor('none')\n",
    "e.set_edgecolor('red')\n",
    "plt.gca().add_artist(e)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(eelr_avg.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "# TODO: dynamically set this threshold\n",
    "thresh = 1\n",
    "center_x = host_avg.shape[1] / 2 - 0.5\n",
    "center_y = host_avg.shape[0] / 2 - 0.5\n",
    "thetas = []\n",
    "for r in range(eelr_avg.shape[0]):\n",
    "    for c in range(eelr_avg.shape[1]):\n",
    "        if eelr_avg[r, c] > thresh:\n",
    "            x = c - center_x\n",
    "            y = r - center_y\n",
    "            theta = math.atan2(y, x) - host_ellipse[4]\n",
    "            theta = theta * 180 / math.pi\n",
    "            if theta < 0:\n",
    "                theta += 360\n",
    "            thetas.append(theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(thetas)\n",
    "plt.xlabel(\"Theta (°)\")\n",
    "plt.ylabel(\"Number of Pixels\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def polar_projection(img, center, Rmax=None, resolution=100):\n",
    "    \"\"\"Evaluate img at the location of polar grid coordinates\n",
    "    This method doesn't resample `img` on the polar grid, it merely\n",
    "    transforms the coordinates and picks the nearest pixel.\n",
    "    For resolved features, this is an acceptable approximation\n",
    "    \"\"\"\n",
    "    lims = img.shape\n",
    "    if Rmax is None:\n",
    "        Rmax = np.sqrt(lims[0]**2 + lims[1]**2)\n",
    "    R, P = np.meshgrid(np.linspace(0, Rmax, resolution, dtype=np.float), np.linspace(-np.pi, np.pi, resolution))\n",
    "    Y = np.round(R * np.sin(P)).astype('int') + center[0]\n",
    "    X = np.round(R * np.cos(P)).astype('int') + center[1]\n",
    "    YX = np.dstack((Y,X))\n",
    "    polar = np.array([[ img[tuple(coord)] for coord in YX[i]] for i in range(len(YX))]).T\n",
    "    return polar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_bins = 12  # number of angle bins\n",
    "angles_per_bin = 5  # number of angles to sample per bin\n",
    "\n",
    "eelr_polar = polar_projection(eelr_avg,\n",
    "                              (host_avg.shape[0] // 2, host_avg.shape[1] // 2),\n",
    "                              min(eelr_avg.shape[0] // 2 - 1, eelr_avg.shape[1] // 2 - 1),\n",
    "                              resolution=num_bins * angles_per_bin)\n",
    "\n",
    "# only consider pixels r_thresh or further from the center (units depend on polar resolution)\n",
    "r_thresh = 2\n",
    "angle_intensities = eelr_polar[r_thresh:].sum(axis=0)\n",
    "angle_intensities = np.array([sum(angle_intensities[i:i+angles_per_bin]) for i in range(0, num_bins*angles_per_bin, angles_per_bin)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cyclic_peak_inds(arr):\n",
    "    \"\"\"Finds peaks in a cyclic array.\"\"\"\n",
    "    \n",
    "    prom_thresh = 0.2 * (arr.max() - arr.min())  # min prominence of peaks\n",
    "#     dist_thresh = len(arr) // 2 - 1  # min distance between peaks\n",
    "    peaks, _ = find_peaks(arr, prominence=prom_thresh)\n",
    "\n",
    "    # find peaks that were at the edges\n",
    "    roll = len(arr) // 2\n",
    "    rolled = np.roll(arr, roll)\n",
    "    rolled_peaks, _ = find_peaks(rolled, prominence=prom_thresh)\n",
    "    unrolled_peaks = (rolled_peaks - roll) % len(arr)\n",
    "    \n",
    "    return np.sort(np.union1d(peaks, unrolled_peaks))\n",
    "\n",
    "def plot_with_peaks(arr, peak_inds, ax=None):\n",
    "    bar_colors = np.array([\"blue\"] * len(arr))\n",
    "    if len(peak_inds) > 0:\n",
    "        bar_colors[peak_inds] = \"red\"\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    ax.bar(np.linspace(-np.pi, np.pi, len(arr)), arr, color=bar_colors, width=2*np.pi/len(arr))\n",
    "    ax.set_xlabel(\"Theta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cyclic_peak_inds = get_cyclic_peak_inds(angle_intensities)\n",
    "print(cyclic_peak_inds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_with_peaks(angle_intensities, cyclic_peak_inds)\n",
    "plt.savefig(f\"{prefix}angle_intensities.png\", dpi=200, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison with FIRST (deprecated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: consider distance between peaks?\n",
    "\n",
    "jetlikes = [source_inds[i] for i in range(len(source_inds)) if results[i] and len(results[i][\"peaks\"]) == 2]\n",
    "print(len(jetlikes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_detected, _, first_detected_inds = np.intersect1d(jetlikes, first_ids, return_indices=True)\n",
    "print(f\"Detected {len(first_detected)} of {len(first_ids)} FIRST AGNs\")\n",
    "print(first_detected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_missed = np.delete(first_ids, first_detected_inds)\n",
    "print(first_missed)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}