{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d6541c4",
   "metadata": {},
   "source": [
    "# Multi-object Photometry\n",
    "\n",
    "In this section we demonstrate how to perform photometry on multiple objects in an image. We choose a faint group of galaxies and run the photometry steps as described in the [Photometry](./photometry.ipynb#Photometry) section.\n",
    "\n",
    "To start with PetroFit, simply import it as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2de4bfde",
   "metadata": {},
   "outputs": [],
   "source": [
    "import petrofit as pf "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0d11bba",
   "metadata": {},
   "source": [
    "## Loading Example Data\n",
    "\n",
    "The following data is a cutout of a group of faint galaxies in Abell 2744. The original data was acquired by the [Hubble Frontier Fields](https://frontierfields.org) team via the WFC3 instrument in the `F105W` filter and can be directly downloaded from the [Mikulski Archive for Space Telescopes](https://archive.stsci.edu/pub/hlsp/frontier/abell2744/images/hst/v1.0/hlsp_frontier_hst_wfc3-60mas_abell2744_f105w_v1.0_drz.fits). \n",
    "\n",
    "We first use `astropy`'s ``CCDData`` to load the example data and visualize it through `matplotlib`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c73113d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.nddata import CCDData\n",
    "\n",
    "image = CCDData.read('data/abell_2744_group_of_galaxies_f105w.fits.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e867d428",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# Hidden cell\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "# Stop Fit Model to Data section warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore', append=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1de6d092",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "plt.rcParams['figure.figsize'] = [6, 6]\n",
    "plt.rcParams['image.origin'] = 'lower'\n",
    "plt.rcParams['font.size'] = 12\n",
    "\n",
    "vmax =  0.01 # Use the image std as max and min of all plots \n",
    "vmin = - vmax \n",
    "\n",
    "plt.imshow(image.data, vmin=vmin, vmax=vmax)\n",
    "plt.title(\"Galaxy Group in Abell 2744 Frontier Field\")\n",
    "plt.xlabel(\"Pixels\")\n",
    "plt.ylabel(\"Pixels\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3abd847f",
   "metadata": {},
   "source": [
    "## Make Source Catalog \n",
    "\n",
    "We will use the sigma clipped std as a threshold at the segmentation and deblending steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "205cf01e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.stats import sigma_clipped_stats\n",
    "\n",
    "# Compute stats for threshold\n",
    "image_mean, image_median, image_stddev = sigma_clipped_stats(image.data, sigma=3)\n",
    "\n",
    "# Define threshold\n",
    "threshold = image_stddev \n",
    "\n",
    "# Min Source size (area)\n",
    "npixels = 5**2\n",
    "\n",
    "\n",
    "cat, segm, segm_deblend = pf.make_catalog(    \n",
    "    image.data, \n",
    "    threshold=threshold,\n",
    "    wcs=image.wcs,\n",
    "    deblend=True,                 \n",
    "    npixels=npixels,\n",
    "    contrast=0.00,\n",
    "    plot=True, vmax=vmax, vmin=vmin,\n",
    "    figsize=[12, 6]\n",
    ")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# Display source properties\n",
    "print(\"Num of Targets:\", len(cat))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d8493b8",
   "metadata": {},
   "source": [
    "## Photometry Loop\n",
    "\n",
    "We define the list of aperture radii and proceed to the photometry step. In this case, instead of selecting a source, we loop through the source catalog and perform photometry on each object. After constructing the photometry we create a ` Petrosian` object for the source. We save the `Petrosian` in a python dictionary (`petrosian_properties`) for later use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e31fff82",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_pix=35\n",
    "\n",
    "r_list = pf.make_radius_list(\n",
    "    max_pix=max_pix, # Max pixel to go up to\n",
    "    n=max_pix # the number of radii to produce \n",
    ")\n",
    "\n",
    "\n",
    "petrosian_properties = {}\n",
    "\n",
    "for idx, source in enumerate(cat):\n",
    "    \n",
    "    # Photomerty \n",
    "    flux_arr, area_arr, error_arr = pf.source_photometry(\n",
    "\n",
    "        # Inputs \n",
    "        source, # Source (`photutils.segmentation.catalog.SourceCatalog`)\n",
    "        image.data, # Image as 2D array \n",
    "        segm_deblend, # Deblended segmentation map of image\n",
    "        r_list, # list of aperture radii  \n",
    "\n",
    "        # Options \n",
    "        cutout_size=max(r_list)*2, # Cutout out size, set to double the max radius  \n",
    "        bg_sub=True, # Subtract background  \n",
    "        sigma=3, sigma_type='clip', # Fit a 2D plane to pixels within 3 sigma of the mean\n",
    "        plot=False, vmax=vmax, vmin=vmin, # Show plot with max and min defined above\n",
    "    )\n",
    "    plt.show()\n",
    "    \n",
    "    p = pf.Petrosian(r_list, area_arr, flux_arr)\n",
    "    \n",
    "    petrosian_properties[source] = p\n",
    "    \n",
    "print(\"Completed for {} Sources\".format(len(petrosian_properties)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7646c36d",
   "metadata": {},
   "source": [
    "## Compute Total Mags \n",
    "\n",
    "We compute the total magnitudes of the sources by looping through their ` Petrosian` objects. We use the ` flux_to_abmag` function to convert the total flux to AB mags of each object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6dfa0ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "\n",
    "mag_list = []\n",
    "\n",
    "for source in petrosian_properties:\n",
    "    \n",
    "    # Get Petrosian\n",
    "    p = petrosian_properties[source]\n",
    "    \n",
    "    # Compute HST Flux -> mags for total_flux\n",
    "    mag = pf.hst_flux_to_abmag(p.total_flux, image.header)\n",
    "    \n",
    "    # Add to mag list\n",
    "    mag_list.append(mag)\n",
    "\n",
    "# Convert mag_list to array \n",
    "mag_list = np.array(mag_list)\n",
    "\n",
    "mag_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a29a8bc9",
   "metadata": {},
   "source": [
    "## Photometry Catalog\n",
    "\n",
    "We construct and save a photometry catalog with the magnitudes we computed. To construct the table, we first use the `SourceCatalog.to_table()` function that returns an `astropy` `table`. This will include important info about each source. We then add a new column (`MAG_F105W`) with the total mags we computed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2d71dd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Segmentation catalog to astropy table.\n",
    "photo_cat = cat.to_table()\n",
    "\n",
    "# Add new column with mags.\n",
    "photo_cat.add_column(mag_list, index=4, name='MAG_F105W')\n",
    "\n",
    "# Save to file.\n",
    "photo_cat.write('temp/example_photo_cat.csv', overwrite=True)\n",
    "\n",
    "photo_cat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06535ee7",
   "metadata": {},
   "source": [
    "# Simultaneous Fitting\n",
    "\n",
    "In this section, we explore how to make compound models that can be used to describe multiple objects in an image. We use the same dataset as the [Multi-object Photometry section](#Multi-object-Photometry) to fit the nine faint sources with Sérsic profiles. \n",
    "\n",
    "## Make Individual Models\n",
    "\n",
    "To do this we loop through the sources and construct `astropy` ` Sersic2D` models for each source. We also make initial guesses of the paramters as described in the Estimating Sérsic Parameters section. At the end of each iteration, we add the newly constructed model in a list `model_list`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "031f39c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from photutils.isophote import EllipseGeometry, Ellipse\n",
    "from astropy.modeling import models\n",
    "\n",
    "# AstroPy Model List\n",
    "model_list = []\n",
    "\n",
    "# For each source\n",
    "for source in list(petrosian_properties.keys()):\n",
    "    \n",
    "    # Get Petrosian\n",
    "    p = petrosian_properties[source]\n",
    "    \n",
    "    # Estimate center \n",
    "    position = pf.get_source_position(source)    \n",
    "    x_0, y_0 = position\n",
    "    \n",
    "    # Estimate shape \n",
    "    elong = pf.get_source_elong(source)\n",
    "    ellip = pf.get_source_ellip(source)\n",
    "    theta = pf.get_source_theta(source)\n",
    "    \n",
    "    # Estimate Sersic index\n",
    "    n = 1\n",
    "    \n",
    "    # Estimate r_half_light\n",
    "    r_eff = p.r_half_light\n",
    "    \n",
    "    # Estimate amplitude\n",
    "    amplitude = pf.get_amplitude_at_r(r_eff, image, x_0, y_0, ellip, theta)\n",
    "    \n",
    "    # Allow for 4 pixel center slack \n",
    "    center_slack = 4\n",
    "    \n",
    "    # Make astropy model\n",
    "    sersic_model = models.Sersic2D(\n",
    "\n",
    "            amplitude=amplitude,\n",
    "            r_eff=r_eff,\n",
    "            n=n,\n",
    "            x_0=x_0,\n",
    "            y_0=y_0,\n",
    "            ellip=ellip, \n",
    "            theta=theta,\n",
    "\n",
    "            bounds = pf.get_default_sersic_bounds({\n",
    "                'x_0': (x_0 - center_slack/2, x_0 + center_slack/2),\n",
    "                'y_0': (y_0 - center_slack/2, y_0 + center_slack/2),\n",
    "            }),\n",
    "    )\n",
    "    \n",
    "    # Add to model list \n",
    "    model_list.append(sersic_model)\n",
    "    \n",
    "    # Over-plot Petrosian radii \n",
    "    p.imshow(position=position, elong=elong, theta=theta, lw=1.25)\n",
    "\n",
    "# Plot image of sources \n",
    "plt.imshow(image.data, vmax=vmax, vmin=vmin)\n",
    "pf.mpl_tick_frame()\n",
    "plt.xlabel('Pixels')\n",
    "plt.ylabel('Pixels')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd6fba37",
   "metadata": {},
   "source": [
    "## Make Compound Model\n",
    "\n",
    "To make a single compound model that represents all the sources of interest, we add up all the models. `astropy` models can be added like numbers or arrays, so we convert the model list to a `numpy` array and sum it. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62a78c7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "compound_model = np.array(model_list).sum()\n",
    "\n",
    "compound_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa7f1f8",
   "metadata": {},
   "source": [
    "## Make PSFConvolvedModel2D\n",
    "\n",
    "Now that we have a single model that represents all the sources, we can create a `PSFConvolvedModel2D` with the appropriate parameters. After we load the PSF, we wrap the compound model and PSF using `PSFConvolvedModel2D`. We specify an oversampling factor 4 to account for poor CCD sampling.\n",
    "\n",
    "**Load and Normalize PSF**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f37cac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.io import fits\n",
    "\n",
    "PSF = fits.getdata('data/f105w_psf.fits.gz')\n",
    "PSF = PSF / PSF.sum()\n",
    "\n",
    "plt.imshow(PSF, vmin=0, vmax=PSF.std()/10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d813bcc4",
   "metadata": {},
   "source": [
    "**PSFConvolvedModel2D**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dede573",
   "metadata": {},
   "outputs": [],
   "source": [
    "psf_sersic_model = pf.PSFConvolvedModel2D(compound_model, psf=PSF, oversample=4)\n",
    "\n",
    "psf_sersic_model.fixed['psf_pa'] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a51a810",
   "metadata": {},
   "source": [
    "## Fit Model to Data\n",
    "\n",
    "We fit the compound model using a Levenberg-Marquardt algorithm and save the returned optimized copy of the fitted model in  `fitted_model`. Since this the compound model is composed of many parameters, we may see `astropy` warnings when the fitter explores parameters that cause issues, such as division by zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ae27f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "fitted_model, _ = pf.fit_model(\n",
    "    image.data, psf_sersic_model,\n",
    "    maxiter=10000,\n",
    "    epsilon=1.4901161193847656e-08,\n",
    "    acc=1e-09,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c25efce0",
   "metadata": {},
   "source": [
    "## Generate Model Image\n",
    "\n",
    "To generate a model image we use the `model_to_image` utility function. This function allows us to define the center of the model image and the side length of the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13ec334d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make Model Image\n",
    "# ----------------\n",
    "\n",
    "# Set the size of the model image equal to the original image\n",
    "full_fitted_image_size = image.data.shape[0]\n",
    "\n",
    "# Center the model image at the center of the original image\n",
    "# so the two images cover the same window\n",
    "full_fitted_image_center = full_fitted_image_size // 2\n",
    "\n",
    "# Generate a model image from the model\n",
    "fitted_model_image = pf.model_to_image(\n",
    "    fitted_model,\n",
    "    full_fitted_image_size, \n",
    ")\n",
    "\n",
    "# Plot Model Image\n",
    "# ----------------\n",
    "\n",
    "fig, ax = plt.subplots(2, 2, figsize=[12, 12])\n",
    "\n",
    "# Plot raw data\n",
    "im = ax[0, 0].imshow(image.data, vmin=vmin, vmax=vmax)\n",
    "ax[0, 0].set_title(\"Hubble F105W Image\")\n",
    "ax[0, 0].set_xlabel(\"Pixels\")\n",
    "ax[0, 0].set_ylabel(\"Pixels\")\n",
    "#ax[0, 0].axis('off')\n",
    "\n",
    "# Plot Petrosian radii\n",
    "plt.sca(ax[0, 1])\n",
    "for i, source in enumerate(petrosian_properties):\n",
    "    p = petrosian_properties[source]\n",
    "    \n",
    "    position = pf.get_source_position(source)    \n",
    "    x_0, y_0 = position\n",
    "    \n",
    "    elong = pf.get_source_elong(source)\n",
    "    ellip = pf.get_source_ellip(source)\n",
    "    theta = pf.get_source_theta(source)\n",
    "    \n",
    "    p.imshow(position=position, elong=elong, theta=theta, lw=1.25)\n",
    "    if i == 0:\n",
    "        plt.legend()\n",
    "ax[0, 1].imshow(image.data, vmin=vmin, vmax=vmax)\n",
    "ax[0, 1].set_title(\"Petrosian Radii\")\n",
    "ax[0, 1].set_xlabel(\"Pixels\")\n",
    "ax[0, 1].set_ylabel(\"Pixels\")\n",
    "# ax[0, 1].axis('off')\n",
    "\n",
    "# Plot Model Image\n",
    "ax[1, 0].imshow(fitted_model_image, vmin=vmin, vmax=vmax)\n",
    "ax[1, 0].set_title(\"Simultaneously Fitted Sersic Models\")\n",
    "ax[1, 0].set_xlabel(\"Pixels\")\n",
    "ax[1, 0].set_ylabel(\"Pixels\")\n",
    "# ax[1, 0].axis('off')\n",
    "\n",
    "# Plot Residual\n",
    "ax[1, 1].imshow(image.data - fitted_model_image, vmin=vmin, vmax=vmax)\n",
    "ax[1, 1].set_title(\"Residual\")\n",
    "ax[1, 1].set_xlabel(\"Pixels\")\n",
    "ax[1, 1].set_ylabel(\"Pixels\")\n",
    "# ax[1, 1].axis('off')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf62ad6f",
   "metadata": {},
   "source": [
    "## Analysis of Background\n",
    "\n",
    "We can now create a background image using the residual and perform some statistical  analysis on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deba9c81",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define background image\n",
    "background_image = image.data - fitted_model_image\n",
    "\n",
    "# Compute stats\n",
    "# -------------\n",
    "\n",
    "noise_mean = background_image.mean()\n",
    "noise_sigma = background_image.std()\n",
    "noise_3_sigma = noise_sigma * 3.\n",
    "noise_8_sigma = noise_sigma * 8.\n",
    "\n",
    "print(\"noise_mean = {}\".format(noise_mean))\n",
    "print(\"noise_sigma = {}\".format(noise_sigma))\n",
    "print(\"noise_3_sigma = {}\".format(noise_3_sigma))\n",
    "\n",
    "# Plots\n",
    "# -----\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=[16, 8])\n",
    "plt.sca(axs[0])\n",
    "plt.imshow(image.data - fitted_model_image, vmin=vmin, vmax=vmax)\n",
    "plt.title('Residual Image')\n",
    "\n",
    "plt.sca(axs[1])\n",
    "n, bins, patches = plt.hist(background_image.flatten(), bins=35, align='left', \n",
    "                            color='black', label=\"Binned Residual Image Pixel Values\")\n",
    "plt.plot(bins[:-1], n, c='r', linewidth=3)\n",
    "plt.axvline(image.data.mean(), label=\"Raw Input Image Mean\", c='g',linestyle=\"--\")\n",
    "plt.axvline(noise_mean, label=\"Residual Image Mean\", linestyle=\"--\")\n",
    "\n",
    "\n",
    "plt.xlabel('Flux Bins [{}]'.format(str(image.unit)))\n",
    "plt.ylabel('Count')\n",
    "plt.title('Residual Image Pixel Value Histogram')\n",
    "plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bb1e3ce",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "Now we see the residual image mean is near the mean of the noise distribution, we can make a segmentation map using the residual image 3-sigma as the detection threshold. Notice how some of the sources we were able to fit were below the 3-sigma estimate of the background (residual image). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a50d96b",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_threshold = noise_3_sigma\n",
    "\n",
    "new_cat, new_segm, new_segm_deblend = pf.make_catalog(    \n",
    "    image.data, \n",
    "    threshold=new_threshold,\n",
    "    wcs=image.wcs,\n",
    "    deblend=True,                 \n",
    "    npixels=npixels,\n",
    "    contrast=0.00,\n",
    "    plot=True, vmax=vmax, vmin=vmin,\n",
    "    figsize=[12, 6]\n",
    ")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# Display source properties\n",
    "print(\"Num of Targets:\", len(new_cat))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}