{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3e52c34d",
   "metadata": {},
   "source": [
    "# Galaxy Modeling\n",
    "\n",
    "In this section, we demonstrate the basics of light profile modeling on a galaxy using a single component Sérsic profile. We also demonstrate how the `photometry` and `petrosian` PetroFit sub-models can be used to improve the initial guesses of the Sérsic parameters.\n",
    "\n",
    "**Please note that this is an extensive analysis to demonstrate PetroFit's fitting workflow. Most practical applications of PetroFit will not require all of the steps discussed in this section.**\n",
    "\n",
    "\n",
    "To start with PetroFit, simply import it as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa40d0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import petrofit as pf "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a68cbd8e",
   "metadata": {},
   "source": [
    "## Loading Example Data\n",
    "\n",
    "The following data is a cutout of a group of bright galaxies in Abell 2744 (located at `(3.596248, -30.388517)`). 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). The cutout image used in this documentation can be found in the git repository at the following path `petrofit/docs/data/abell_2744_dwarf_galaxy_f105w.fits.gz`.\n",
    "\n",
    "### Loading Data and RMS Images\n",
    "We first use `astropy`'s ``CCDData`` to load the example data and visualize it through `matplotlib`. The rms image is loaded using `astropy`'s ``fits`` sub-module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91575418",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.nddata import CCDData\n",
    "from astropy.io import fits \n",
    "\n",
    "image = CCDData.read('data/abell_2744_dwarf_galaxy_f105w.fits.gz')\n",
    "rms = fits.getdata('data/abell_2744_dwarf_galaxy_f105w_rms.fits.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c75fd46",
   "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": "ffb27660",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "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.05 # Use the image std as max and min of all plots \n",
    "vmin = - vmax \n",
    "\n",
    "fig, axs = plt.subplots(1,2, figsize=[12, 6])\n",
    "plt.sca(axs[0])\n",
    "plt.imshow(image.data, vmin=vmin, vmax=vmax)\n",
    "plt.title(\"Mock Galaxy\")\n",
    "plt.xlabel(\"Pixels\")\n",
    "plt.ylabel(\"Pixels\")\n",
    "\n",
    "plt.sca(axs[1])\n",
    "plt.imshow(rms)\n",
    "plt.title(\"RMS Image\")\n",
    "plt.xlabel(\"Pixels\")\n",
    "plt.ylabel(\"Pixels\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24044134",
   "metadata": {},
   "source": [
    "## Image Errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f16d0ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from photutils.utils import calc_total_error\n",
    "err = calc_total_error(\n",
    "    data=image.data, # Input Image\n",
    "    bkg_error=rms, # All sources of background error except source Poisson error\n",
    "    effective_gain=image.header['EXPTIME'] # Factor to convert data units to counts\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b76dbf6",
   "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": "b5d8bba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.stats import sigma_clipped_stats\n",
    "image_mean, image_median, image_stddev = sigma_clipped_stats(image.data, sigma=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f99306",
   "metadata": {},
   "source": [
    "Here we identity sources in the input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "648a41a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = image_stddev * 3\n",
    "\n",
    "# Min Source size (area)\n",
    "npixels = 2**2\n",
    "\n",
    "cat, segm, segm_deblend = pf.make_catalog(    \n",
    "    image.data, \n",
    "    threshold=threshold, \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": "8d1050fc",
   "metadata": {},
   "source": [
    "## Single Source Photometry \n",
    "\n",
    "\n",
    "Here we pick a galaxy that can be modeled using a single Sérsic model, the galaxy in the middle looks like an elliptical galaxy that can be well described by such a profile. To make sure we selected the right galaxy, we use the `plot_target` function to plot a cutout of the source. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c2384e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sort and select object of interest in the catalog\n",
    "sorted_idx_list = pf.order_cat(cat, key='area', reverse=True)\n",
    "idx = sorted_idx_list[2] # index 0 is largest \n",
    "source = cat[idx]  # get source from the catalog \n",
    "\n",
    "pf.plot_target(\n",
    "    position=(source.maxval_xindex, source.maxval_yindex), \n",
    "    image=image.data, \n",
    "    size=100, \n",
    "    vmax=vmax, \n",
    "    vmin=vmin\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d172029",
   "metadata": {},
   "source": [
    "After selecting the source, we define a radius list for the aperture photometry. Once `r_list` is defined, we run the `source_photometry` step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c1a6837",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Max aperture radius \n",
    "max_pix = 70\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 (i.e 1 aperture per pixel increase in r) \n",
    ")\n",
    "\n",
    "print(\"len(r_list) = {}\".format(len(r_list)))\n",
    "\n",
    "\n",
    "# Photomerty \n",
    "flux_arr, area_arr, error_arr = pf.source_photometry(\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",
    "    error=err, # Error image (optional)\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=True, vmax=vmax, vmin=vmin, # Show plot with max and min defined above\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d167e5fa",
   "metadata": {},
   "source": [
    "## Petrosian Profile \n",
    "\n",
    " \n",
    "We use the photometry that we obtained to construct a `Petrosian` profile of the galaxy. This will allow us to estimate the parameters of the galaxies Sérsic profile. We also initialize a `PetrosianCorrection` which we can use to estimate the Sérsic index (`n`) and `r_eff` (half-light radius). The correction grid we use for `PetrosianCorrection` was created specifically for this dataset using the PSF above. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2161ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pf.Petrosian(r_list, area_arr, flux_arr, flux_err=error_arr)\n",
    "\n",
    "pc = pf.PetrosianCorrection.read(\"data/f105w_psf_corr.ecsv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdcaebb5",
   "metadata": {},
   "source": [
    "Next we compute and apply the epsilon that produces the correct `r_total_flux` and creates a new corrected Petrosian profile (`corrected_p`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42ff7344",
   "metadata": {},
   "outputs": [],
   "source": [
    "pc.plot_correction(p)\n",
    "plt.show()\n",
    "corrected_p = pc.correct(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cb7fca8",
   "metadata": {},
   "source": [
    "Galaxies with a high Sérsic index (high concentration) have `r_total_flux`s that extend much further than we can possibly measure. If the corrected `r_total_flux` is out of the photometric range (i.e ` r_total_flux > max_pix `), we would be unable to measure the flux at `r_total_flux` (the total flux), which also means that we would be unable to measure the Petrosian radius (`np.nan` is returned). If the corrected ` r_total_flux` is out of range and we feel like the photometry data range goes far out enough to estimate the total flux, we can set `epsilon` to a value that will result in a ` r_total_flux`  equal to the last photometry radius."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0d4ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_photometry_r =  max(corrected_p.r_list)\n",
    "\n",
    "if corrected_p.r_total_flux > max_photometry_r:\n",
    "    print(\"Truncation was applied\")\n",
    "    truncated_epsilon = max_photometry_r / corrected_p.r_petrosian\n",
    "    corrected_p.epsilon = truncated_epsilon\n",
    "else:\n",
    "    print(\"Truncation was not needed\")\n",
    "\n",
    "# Print Radii\n",
    "print('total_flux = ', corrected_p.total_flux)\n",
    "print('max_photometry_r = ', max_photometry_r)\n",
    "\n",
    "# Plot Petrosian \n",
    "corrected_p.plot()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab39545f",
   "metadata": {},
   "source": [
    "## Estimating Sérsic Parameters \n",
    "\n",
    "`astropy`’s `Sersic2D` implements this model and allows us to provide initial guesses for the Sérsic parameters. Getting good estimates of these parameters is very important because the Levenberg-Marquardt algorithm is sensitive and may return parameters that correspond to a local minimum. Because of this, we follow the steps below to estimate the Sérsic parameters.\n",
    "\n",
    "The `Sersic2D` model has the following parameters:\n",
    "\n",
    "- [**amplitude**: Surface brightness at r_eff.](#Amplitude-at-r_eff)\n",
    "- [**r_eff**: Effective (half-light) radius.](#Half-Light-Radius-(r_eff))\n",
    "- [**n**: Sérsic Index.](#Sérsic-Index)\n",
    "- [**x_0 and y_0**: x and y position of the center.](#Center-Pixel)\n",
    "- [**ellip**:  Ellipticity of the profile.](#Ellipticity-and-Elongation)\n",
    "- [**theta**: Rotation angle in radians, counterclockwise from the positive x-axis.](#Source-Orientation)\n",
    "\n",
    "\n",
    "\n",
    "### Center Pixel\n",
    "\n",
    "We can use the `get_source_position` to find the max pixel position of the source in the image using the `SourceCatalog` object. Since this center is relative to the image, we save the result in `image_x_0` and `image_y_0`."
   ]
  },
  {
   "cell_type": "raw",
   "id": "cb0c268a",
   "metadata": {},
   "source": [
    "<div class=\"admonition note\">\n",
    "<p class=\"admonition-title\">Note</p>\n",
    "<p>The <code class=\"docutils literal notranslate\"><span class=\"pre\">x_0</span></code> and <code class=\"docutils literal notranslate\"><span class=\"pre\">y_0</span></code> vlaues we use in the model in this example are defined in the <a class=\"reference external\" href=\"#Zoom-and-Mask-Image\">Zoom and Mask Image section</a>. This is because we make a cutout of the target at <code class=\"docutils literal notranslate\"><span class=\"pre\">(image_x_0,</span> <span class=\"pre\">image_y_0)</span></code> which shifts the coordinate space. If no cutouts are made, then we can use <code class=\"docutils literal notranslate\"><span class=\"pre\">(image_x_0,</span> <span class=\"pre\">image_y_0)</span></code> as the center of the <code class=\"docutils literal notranslate\"><span class=\"pre\">astropy</span></code> model.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "508cb6a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_x_0, image_y_0 = pf.get_source_position(source)\n",
    "\n",
    "print(image_x_0, image_y_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29805c6d",
   "metadata": {},
   "source": [
    "###  Ellipticity and Elongation\n",
    "\n",
    "We pass the source's `SourceCatalog` object to the `get_source_ellip` and `get_source_elong` functions to find the source's light profile ellipticity and elongation respectively. These values are derived from the segmentation footprint of the source. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c3178e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ellip = pf.get_source_ellip(source)\n",
    "elong = pf.get_source_elong(source)\n",
    "\n",
    "print(ellip)\n",
    "print(elong)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2655cec5",
   "metadata": {},
   "source": [
    "### Source Orientation\n",
    "\n",
    "We pass the source's `SourceCatalog` object to the `get_source_theta` function to find the orientation of the light profile in radians, counterclockwise from the positive x-axis. In some catalogs, this value is known as the position angle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c98ec68",
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = pf.get_source_theta(source)\n",
    "\n",
    "print(theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dfa5518",
   "metadata": {},
   "source": [
    "### Sérsic Index\n",
    "\n",
    "The PetroFit `PetrosianCorrection` contains a grid that maps the uncorrected (PSF convolved) Petrosian half-light radius and concentration index to an epsilon value that gives the correct Petrosian radii. This grid can also be used to map the Petrosian radius, uncorrected half-light radius, and concentration index to a Sérsic index. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57ef0d9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = pc.estimate_n(p)\n",
    "\n",
    "print(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881f2ffe",
   "metadata": {},
   "source": [
    "### Half-Light Radius (r_eff)\n",
    "\n",
    "We use the corrected Petrosian radius to estimate the half-light radius (`r_eff`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0dec7fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "r_eff = corrected_p.r_half_light\n",
    "\n",
    "print(r_eff)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1189983",
   "metadata": {},
   "source": [
    "### Amplitude at r_eff\n",
    "\n",
    "To find the flux at the half-light radius, we fit an elliptical isophote at `r_eff` using `photutils`’s `Ellipse` model using the `get_amplitude_at_r` function. If this value can not be computed, `np.nan` is returned so make sure to check using `np.isnan(amplitude)`\n",
    "\n",
    "```\n",
    "from photutils.isophote import EllipseGeometry, Ellipse\n",
    "\n",
    "# Define EllipseGeometry using ellip and theta\n",
    "g = EllipseGeometry(image_x_0, image_y_0, 1., ellip, theta)\n",
    "\n",
    "# Create Ellipse model \n",
    "ellipse = Ellipse(image.data, geometry=g)\n",
    "\n",
    "# Fit isophote at r_eff\n",
    "iso = ellipse.fit_isophote(r_eff)\n",
    "\n",
    "# Get flux at r_eff\n",
    "amplitude = iso.intens\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca1b19a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from petrofit.segmentation import get_amplitude_at_r\n",
    "amplitude = get_amplitude_at_r(r_eff, image, image_x_0, image_y_0, ellip, theta)\n",
    "\n",
    "print(amplitude)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9e59d78",
   "metadata": {},
   "source": [
    "## Zoom and Mask Image\n",
    "\n",
    "Before we fit the image, it is important to mask all nearby sources to avoid interference and zoom in to avoid fitting irrelevant pixels (especially those outside of `r_total_flux`). For this example, we will mask nearby sources using the `masked_segm_image` function that masks pixels using source segmentation footprints and make a zoomed cutout using `astropy`’s `Cutout2D` function. We also make a second cutout (`fitting_image_unmasked`) of the same dimensions as the fitting image, but without masking nearby sources. We will subtract the fitted model from this unmasked image once the fitting is complete. \n",
    "\n",
    "We make a cutout of size `cutout_size` centered at the source, as we identified in the [finding center section](#Center-Pixel), and define a new center (`x_0` and `y_0`) that we will use in the `astropy` model. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbe30014",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.nddata import Cutout2D\n",
    "\n",
    "# Define cutout size \n",
    "cutout_size = 30\n",
    "\n",
    "# Make an image with all sources masked except `source`\n",
    "masked_image = pf.masked_segm_image(source, image.data, segm_deblend, fill=None, mask_background=False)\n",
    "\n",
    "# Make cutouts \n",
    "fitting_image = Cutout2D(masked_image, (image_x_0, image_y_0), cutout_size, mode='partial', copy=True)\n",
    "fitting_image_unmasked = Cutout2D(image.data, (image_x_0, image_y_0), cutout_size, mode='partial', copy=True)\n",
    "\n",
    "# Define new center \n",
    "x_0 = y_0 = cutout_size / 2\n",
    "\n",
    "# Plot cutout that will be fit \n",
    "plt.imshow(fitting_image.data, vmin=vmin, vmax=vmax)\n",
    "plt.title(\"Galaxy in Abell 2744\")\n",
    "plt.xlabel(\"Pixels\")\n",
    "plt.ylabel(\"Pixels\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33f984f4",
   "metadata": {},
   "source": [
    "## Make Sérsic Model\n",
    "\n",
    "### Astropy Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cb0021c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.modeling import models \n",
    "\n",
    "center_slack = 4\n",
    "\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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84dec3c3",
   "metadata": {},
   "source": [
    "### Load PSF "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc9f08a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.io import fits\n",
    "\n",
    "# Load PSF image (2D array)\n",
    "PSF = fits.getdata('data/f105w_psf.fits.gz')\n",
    "\n",
    "# Normalize PSF \n",
    "PSF = PSF / PSF.sum()\n",
    "\n",
    "# Note that the PSF shape is odd on all sides\n",
    "print(\"PSF Shape = {}\".format(PSF.shape))\n",
    "\n",
    "# Plot PSF and use vmax and vmin to show difraction spikes\n",
    "plt.imshow(PSF, vmin=0, vmax=PSF.std()/10)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "817344e0",
   "metadata": {},
   "source": [
    "### Oversampling "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d6d88a",
   "metadata": {},
   "source": [
    "**Oversample Entire Model Image**\n",
    "\n",
    "To oversample the image by a factor, you can pass a single integer value. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7d7d660",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Oversample the entire image by a factor of 5\n",
    "oversample = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8a178f5",
   "metadata": {},
   "source": [
    "### Create PetroFit Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef24ad08",
   "metadata": {},
   "outputs": [],
   "source": [
    "psf_sersic_model = pf.PSFConvolvedModel2D(sersic_model, psf=PSF, oversample=oversample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8014d5a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disable the PSF rotation, \n",
    "psf_sersic_model.fixed['psf_pa'] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f53b0ee0",
   "metadata": {},
   "source": [
    "## Fitting the PSF Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "646afb23",
   "metadata": {},
   "outputs": [],
   "source": [
    "fitting_err =  Cutout2D(err, (image_x_0, image_y_0), cutout_size, mode='partial', copy=True)\n",
    "fitting_weights = 1 / fitting_err.data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aff6a5c9",
   "metadata": {},
   "source": [
    "To fit the galaxy we prepared with the `PSFConvolvedModel2D` we constructed, we call the `fit_model` as follows: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "990fe9a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "from astropy.modeling import fitting\n",
    "fitted_model, fit_info = pf.fit_model(\n",
    "    fitting_image.data, psf_sersic_model,\n",
    "    weights=fitting_weights,\n",
    "    calc_uncertainties=True,\n",
    "    maxiter=10000,\n",
    "    epsilon=1.4901161193847656e-08,\n",
    "    acc=1e-09,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "544b714d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pf.print_model_params(fitted_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92341da0",
   "metadata": {},
   "source": [
    "## Model Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73fd5f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "pf.plot_fit(fitted_model, fitting_image_unmasked, vmax=vmax, vmin=vmin, figsize=[3*6, 6])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5d03654",
   "metadata": {},
   "source": [
    "If we want to make a model image that we can subtract from the original data image, we just update the center of the model and generate a model image the same dimensions as the original image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a09353ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy \n",
    "\n",
    "# Reposition the Model\n",
    "# --------------------\n",
    "\n",
    "# Copy the model before changing the parameters\n",
    "fitted_model_copy = deepcopy(fitted_model)\n",
    "\n",
    "# Change the frame of reference from the fitting \n",
    "# image to the original data coordinates. \n",
    "# Remember that (x_0, y_0) -> (image_x_0, image_y_0) \n",
    "fitted_model_copy.x_0 =  image_x_0 + (fitted_model.x_0 - x_0)\n",
    "fitted_model_copy.y_0 =  image_y_0 + (fitted_model.y_0 - y_0)\n",
    "\n",
    "# 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",
    "full_fitted_model_image = pf.model_to_image(\n",
    "    fitted_model_copy, \n",
    "    full_fitted_image_size, \n",
    "    \n",
    ")\n",
    "\n",
    "# Plot Model Image\n",
    "# ----------------\n",
    "pf.plot_fit(fitted_model_copy, image.data,vmin=vmin, vmax=vmax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be820294",
   "metadata": {},
   "source": [
    "## Estimated vs. Fitted Parameters\n",
    "\n",
    "In this section we compare the initial estimates that we derived purely from photometry vs. the parameters that were fit to the data. In the printout below, we see the fitted values in the first column and our initial guesses in the second column for each parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "badba372",
   "metadata": {},
   "outputs": [],
   "source": [
    "assumptions = [\n",
    "    amplitude,\n",
    "    r_eff,\n",
    "    n,\n",
    "    x_0,\n",
    "    y_0,\n",
    "    ellip,\n",
    "    theta,\n",
    "    0 # psf_pa\n",
    "]\n",
    "\n",
    "param_stds = fitted_model.stds\n",
    "\n",
    "print(\"assum\\tfit\\t\\tparam_name\")\n",
    "for param_name, param_val, param_std, assumption in zip(fitted_model.param_names, fitted_model.parameters,  param_stds.stds, assumptions):\n",
    "    print(\"{:0.2f}\\t{:0.2f} ± {:0.3f}\\t{}\".format(assumption, param_val, param_std, param_name))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ef8b94",
   "metadata": {},
   "source": [
    "## Galaxy vs. Model Profiles \n",
    "\n",
    "Now that we have a model image, we can do photometry and compare the Petrosian measurements of the model to that of the original galaxy image. Since we know the exact position of the target in the model image and because there are no intruding sources, we can perform direct photometry using `photometry_step` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7f36d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Photomerty \n",
    "model_flux_arr, model_area_arr, model_error_arr = pf.source_photometry(\n",
    "    \n",
    "    # Inputs \n",
    "    source, # Source (`photutils.segmentation.catalog.SourceCatalog`)\n",
    "    full_fitted_model_image, # 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=True, vmax=vmax, vmin=vmin, # Show plot with max and min defined above\n",
    ")\n",
    "\n",
    "plt.title(\"Model Image\")\n",
    "plt.xlabel(\"Pixels\")\n",
    "plt.ylabel(\"Pixels\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f0a018",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot Light profile \n",
    "plt.plot(r_list, flux_arr, label='Data', lw=5, alpha=0.7)\n",
    "plt.plot(r_list, model_flux_arr, label='Model', lw=5, alpha=0.7)\n",
    "\n",
    "plt.title(\"Light Profiles\")\n",
    "plt.xlabel(\"Aperture Radius [Pixels]\")\n",
    "plt.ylabel(\"Cumulative Flux  [{}]\".format(image.unit))\n",
    "pf.mpl_tick_frame()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d874013",
   "metadata": {},
   "source": [
    "### Petrosian Radii Comparison\n",
    "\n",
    "We compare the Petrosian profiles and radii of the two images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "350d3197",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_p = pf.Petrosian(r_list, model_area_arr, model_flux_arr)\n",
    "\n",
    "model_corrected_p = pc.correct(model_p)\n",
    "\n",
    "\n",
    "# Plot Image \n",
    "plt.imshow(full_fitted_model_image, vmax=vmax, vmin=vmin)\n",
    "\n",
    "model_position = (\n",
    "    fitted_model_copy.x_0.value,\n",
    "    fitted_model_copy.y_0.value)\n",
    "\n",
    "model_position = (image_x_0, image_y_0)\n",
    "model_elong = 1 / (1 - fitted_model_copy.ellip.value)\n",
    "\n",
    "model_theta = fitted_model_copy.theta.value\n",
    "\n",
    "model_corrected_p.imshow(\n",
    "    position=model_position,\n",
    "    elong=model_elong, \n",
    "    theta=model_theta\n",
    ")\n",
    "\n",
    "plt.title(\"Model Image\")\n",
    "plt.xlabel(\"Pixels\")\n",
    "plt.ylabel(\"Pixels\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce40f7bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "corrected_p.plot(plot_r=True)\n",
    "model_corrected_p.plot(plot_r=True, color='tab:orange')\n",
    "\n",
    "plt.gca().get_legend().remove()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a443d341",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Data r_total_flux = {:0.2f}\".format(corrected_p.r_total_flux))\n",
    "print(\"Model r_total_flux = {:0.2f}\".format(model_corrected_p.r_total_flux))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b865eade",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Data r_half_light = {:0.2f}\".format(corrected_p.r_half_light))\n",
    "print(\"Model r_half_light= {:0.2f}\".format(model_corrected_p.r_half_light))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556d1281",
   "metadata": {},
   "source": [
    "### Total Flux Comparison\n",
    "\n",
    "Finally we compare the total Petrosian flux of the data vs the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ccf7c61",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Data Corrected Total Flux = {:0.2f}\".format(corrected_p.total_flux * image.unit))\n",
    "print(\"Model Corrected Total Flux = {:0.2f}\".format(model_corrected_p.total_flux * image.unit))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ad7946b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Data Corrected AB mag = {:0.2f} mag\".format(pf.hst_flux_to_abmag(corrected_p.total_flux, image.header) ))\n",
    "print(\"Model Corrected AB mag = {:0.2f} mag\".format(pf.hst_flux_to_abmag(model_corrected_p.total_flux, image.header) ))"
   ]
  }
 ],
 "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
}