#!/usr/bin/env python
# coding: utf-8

# In[1]:


from IPython.core.display import SVG
url = 'https://stsci.box.com/shared/static/p2pmd5pzmmhs8hscoojnvlb4bcvnstck.svg'
SVG(url=url)


# ## Photutils
# 
# - Code: https://github.com/astropy/photutils
# - Documentation: http://photutils.readthedocs.org/en/latest/
# - Issue Tracker:  https://github.com/astropy/photutils/issues

# ## Photutils Overview
# 
# - Background and background RMS estimation
# - Aperture Photometry
# - PSF Photometry (needs work)
# - Source Detection and Extraction
#   - DAOFIND and IRAF's starfind
#   - image segmentation based on (local) threshold
#   - local peak finder
# - Segmentation-based source photometry and morphological properties
# - Centroids
# - More to come!

# In[2]:


# initial imports
import numpy as np
import matplotlib.pyplot as plt

# run the %matplotlib magic command to enable inline plotting
# in the current Notebook
get_ipython().run_line_magic('matplotlib', 'inline')


# ## Part 1:  Perform aperture photometry on some sources in the XDF

# ### Load the data

# In[3]:


import os
import os.path as op
from astropy.io import fits

#path = /home/analyst/data/glue'
path = op.join(os.environ['HOME'], 'data/glue')
sci_fn = 'hlsp_xdf_hst_wfc3ir-60mas_hudf_f160w_v1_sci.fits'
sci_fn = op.join(path, sci_fn)

sci_hdulist = fits.open(sci_fn)


# In[4]:


# define the data array
data = sci_hdulist[0].data.astype(np.float)
print(data.shape)
print(data.dtype)


# In[5]:


# display the data
from astropy.visualization import scale_image
plt.imshow(scale_image(data, scale='sqrt', percent=98.),
           origin='lower', cmap='Greys_r')


# In[6]:


# load the weight file
wht_fn = 'hlsp_xdf_hst_wfc3ir-60mas_hudf_f160w_v1_wht.fits'
wht_fn = op.join(path, wht_fn)
#wht_fn = 'https://archive.stsci.edu/pub/hlsp/xdf/hlsp_xdf_hst_wfc3ir-60mas_hudf_f160w_v1_wht.fits'

wht = fits.getdata(wht_fn).astype(np.float)

# create an RMS map from the inverse-variance WHT map
err = 1. / np.sqrt(wht)


# In[7]:


# extract the data header and create a WCS object
from astropy.wcs import WCS
hdr = sci_hdulist[0].header
wcs = WCS(hdr)


# In[8]:


# plot a histogram of data values
idx_good = np.isfinite(err)  # exclude regions with no coverage
r = plt.hist(data[idx_good].ravel(), bins=50, range=(-0.01, 0.01))


# ## Aperture photometry

# In[9]:


url = 'http://photutils.readthedocs.org/en/latest/_images/inheritance-41bd5c4e22ff4e2f493f41359f854c831e8e9c86.svg'
SVG(url=url)


# ## Methods for handling aperture/pixel intersection

# In[10]:


url = 'https://stsci.box.com/shared/static/bhe6eyer7w5pmn9spxfjdx3celj9irzt.svg'
SVG(url=url)


# In[11]:


from photutils import CircularAperture, aperture_photometry


# In[12]:


# define the aperture
position = (2815.73, 4209.43)
radius = 5.
aperture = CircularAperture(position, r=radius)


# In[13]:


# perform the photometry; the default method is 'exact'
phot = aperture_photometry(data, aperture)


# In[14]:


phot


# In[15]:


# center method
phot = aperture_photometry(data, aperture,  method='center')
phot


# In[16]:


# subpixel method, subpixels=5 (same as SExtractor)
phot = aperture_photometry(data, aperture,  method='subpixel', subpixels=5)
phot


# In[17]:


# include units
import astropy.units as u
unit = u.electron / u.s
phot = aperture_photometry(data, aperture, unit=unit)
phot


# In[18]:


help(aperture_photometry)


# ## Perform aperture photometry at 3 positions
# ## Also input the error array

# In[19]:


positions = [(2815.73, 4209.43), (2798.63, 4289.41), (2768.62, 4211.63)]
radius = 5.
apertures = CircularAperture(positions, r=radius)

phot = aperture_photometry(data, apertures, error=err)
phot


# In[20]:


# this time include the Poisson error of the sources
eff_gain = hdr['TEXPTIME']

#toterr = np.sqrt(err**2 + (data / eff_gain))
#phot = aperture_photometry(data, apertures, error=toterr)
phot = aperture_photometry(data, apertures, error=err, effective_gain=eff_gain)
phot


# ## Bad pixel masking

# In[21]:


# create a bad pixel
data2 = data.copy()
y, x = 4209, 2816
data2[y, x] = 100.
aperture_photometry(data2, apertures, error=err)


# In[22]:


# mask the bad pixel
mask = np.zeros_like(data2, dtype=bool)
mask[y, x] = True
aperture_photometry(data2, apertures, error=err, mask=mask)


# ## Add columns to the photometry table

# In[23]:


# Add an ID column
phot['id'] = np.arange(len(phot)) + 1
phot


# In[24]:


# rearrange the column order to put 'id' first
colnames = phot.colnames
colnames.insert(0, colnames.pop(colnames.index('id')))
phot = phot[colnames]
phot


# In[25]:


# calculate the SNR and add it to the table
snr = phot['aperture_sum'] / phot['aperture_sum_err']
phot['snr'] = snr
phot


# In[26]:


# calculate the F160W AB magnitude and add it to the table
f160w_zpt = 25.9463
abmag = -2.5 * np.log10(phot['aperture_sum']) + f160w_zpt
phot['abmag'] = abmag
phot


# In[27]:


# calculate the ICRS RA and Dec and add them to the table
from astropy.wcs.utils import pixel_to_skycoord
x, y = np.transpose(positions)
coord = pixel_to_skycoord(x, y, wcs)
phot['ra_icrs'] = coord.icrs.ra
phot['dec_icrs'] = coord.icrs.dec
phot


# ## Aperture photometry using sky apertures

# In[28]:


coord


# In[29]:


# define circular aperture using sky coordinates
from photutils import SkyCircularAperture
radius2 = 5. * u.pix   # sky aperture radius must be a Quantity
sky_apers = SkyCircularAperture(coord, r=radius2)


# In[30]:


# aperture_photometry needs either an wcs object
phot2 = aperture_photometry(data, sky_apers, wcs=wcs)
phot2


# In[31]:


# ...or an hdu (i.e. header and data)
sci_hdulist[0].header['BUNIT'] = 'electron/s'
phot2 = aperture_photometry(sci_hdulist[0], sky_apers)
phot2


# In[32]:


# SkyApertures can take angular radii
radius3 = 5. * u.arcsec
sky_apers3 = SkyCircularAperture(coord, r=radius3)
phot3 = aperture_photometry(data, sky_apers3, wcs=wcs)
phot3


# In[33]:


# can convert SkyApertures to PixelApertures
apers4 = sky_apers3.to_pixel(wcs)
apers4


# In[34]:


# get the radius in pixels
print(sky_apers3.r)
print(apers4.r)   # radius in pixels

print(sky_apers3.r) / (0.060 * u.arcsec / u.pix)


# ## Apertures can plot themselves

# In[35]:


plt.imshow(scale_image(data, scale='sqrt', percent=98.),
           origin='lower', cmap='Greys_r')
apertures.plot(color='blue')
plt.xlim(2700, 2900)
plt.ylim(4150, 4350)


# ## Multiple Apertures at Each Position

# In[36]:


unit = u.electron / u.s
radii = [3, 4, 5]
tbl_list = []
for radius in radii:
    tbl_list.append(aperture_photometry(data,
        CircularAperture(positions, radius),
        error=err, unit=unit))


# In[37]:


# combine the table list using hstack
from astropy.table import hstack
phot = hstack(tbl_list)
phot


# In[38]:


# instead of 1, 2, 3 names, use r3, r4, r5
from astropy.table import hstack
names = ['r3', 'r4', 'r5']
phot2 = hstack(tbl_list, table_names=names)
phot2


# In[39]:


# do some table cleanup
for key in ['xcenter_r4', 'ycenter_r4', 'xcenter_r5', 'ycenter_r5']:
    del phot2[key]
phot2.rename_column('xcenter_r3', 'xcenter')
phot2.rename_column('ycenter_r3', 'ycenter')
colnames = phot2.colnames
colnames.insert(0, colnames.pop(colnames.index('ycenter')))
colnames.insert(0, colnames.pop(colnames.index('xcenter')))
phot2 = phot2[colnames]
phot2


# In[40]:


# add the aperture radii to the table
phot2['aperture_radius_r3'] = radii[0] * u.pix
phot2['aperture_radius_r4'] = radii[1] * u.pix
phot2['aperture_radius_r5'] = radii[2] * u.pix
phot2


# ## Encircled flux

# In[41]:


# multiple apertures at each position
radii = np.linspace(0.1, 20, 100)
flux = []
for r in radii:
    ap = CircularAperture(positions[0], r=r)
    phot = aperture_photometry(data, ap)
    flux.append(phot['aperture_sum'][0])


# In[42]:


plt.plot(radii, flux)
plt.title('Encircled Flux')
plt.xlabel('Radius (pixels)')
plt.ylabel('Aperture Sum ($e^{-1}/s$)')


# ## Local background estimation

# In[43]:


# define the annular apertures
from photutils import CircularAnnulus
apertures = CircularAperture(positions, r=3)
annulus_apertures = CircularAnnulus(positions, r_in=10., r_out=15.)


# In[44]:


# plot the apertures
from astropy.visualization import scale_image
plt.imshow(scale_image(data, scale='sqrt', percent=98.), 
           origin='lower', cmap='Greys_r')
apertures.plot(color='blue')
annulus_apertures.plot(color='green', hatch='//', alpha=0.8)
plt.xlim(2700, 2900)
plt.ylim(4150, 4350)


# In[45]:


# measure the aperture sums
phot = aperture_photometry(data, apertures)
bkg = aperture_photometry(data, annulus_apertures)


# In[46]:


# calculate the mean background level (per pixel) in the annuli
bkg_mean = bkg['aperture_sum'] / annulus_apertures.area()
bkg_mean


# In[47]:


# now calculate the total background in the circular aperture
bkg_sum = bkg_mean * apertures.area()

phot['bkg_sum'] = bkg_sum
phot


# In[48]:


# subtract the background
flux_bkgsub = phot['aperture_sum'] - bkg_sum

phot['aperture_sum_bkgsub'] = flux_bkgsub
phot


# ## Part 2:  Image Segmentation

# In[49]:


# define a threshold image
bkg = 0.
threshold = bkg + (3.0 * err)   # 3-sigma (per pixel!) above the background


# In[50]:


# detect sources
from photutils import detect_sources
segm = detect_sources(data, threshold, npixels=5)


# In[51]:


# object minimum snr
obj_min_snr = np.sqrt(5) * 3.
obj_min_snr


# ### The segmentation image is the isophotal footprint of each source above the threshold

# In[52]:


# plot the segmentation image
plt.imshow(segm, origin='lower')
plt.colorbar()
print('Found {0} sources'.format(np.max(segm)))


# ## Filter (smooth) the data prior to source detection

# In[53]:


# define the kernel
from astropy.convolution import Gaussian2DKernel
from astropy.stats import gaussian_fwhm_to_sigma
sigma = 2.0 * gaussian_fwhm_to_sigma    # FWHM = 2 pixels
kernel = Gaussian2DKernel(sigma, x_size=5, y_size=5)
kernel.array


# In[54]:


# detect the sources
segm = detect_sources(data, threshold, npixels=5, filter_kernel=kernel)
plt.imshow(segm)
plt.colorbar()
print('Found {0} sources'.format(np.max(segm)))


# ### IMPORTANT NOTE:
# ### Source deblending is not yet available

# ## Measure the photometry and morphological properties of detected sources

# In[55]:


from photutils import segment_properties, properties_table
props = segment_properties(data, segm, error=err)
tbl = properties_table(props)


# ### props is a list of SegmentProperties objects

# ### tbl is a Table of isophotal photometry and morphological properties

# In[56]:


tbl


# In[57]:


snr = tbl['segment_sum'] / tbl['segment_sum_err']
print('Minimum SNR: {0}'.format(np.min(snr)))


# In[58]:


# a subset of sources can be specified,
# defined by their labels in the segmentation image
labels = [1, 5, 20, 50, 75, 80]
props2 = segment_properties(data, segm, error=err, wcs=wcs, labels=labels)
tbl = properties_table(props2)
tbl


# In[59]:


# a subset of columns can be specified
columns = ['id', 'xcentroid', 'ycentroid', 'segment_sum', 'area']
tbl2 = properties_table(props2, columns=columns)
tbl2


# In[60]:


# props is a list of SegmentProperties objects
obj = props[0]   # id = 1 (segmentation label)
print(obj.id)
print(obj.xcentroid)
print(obj.area)
print(obj.ellipticity)
print(obj.covariance)
print(obj.orientation)
print(obj.orientation.to(u.deg))


# In[61]:


# plot the first object
fig, ax = plt.subplots(figsize=(4, 6), ncols=2)
origin ='lower'
cmap = 'hot'
ax[0].imshow(obj.data_cutout_ma, origin=origin, cmap=cmap)
ax[1].imshow(obj.error_cutout_ma, origin=origin, cmap=cmap)


# In[62]:


f160w_zpt = 25.9463
abmag = -2.5 * np.log10(obj.segment_sum) + f160w_zpt
print(abmag)

print(obj.segment_sum / obj.segment_sum_err)


# ## For the currently defined source properties, see
# http://photutils.readthedocs.org/en/latest/api/photutils.segmentation.SegmentProperties.html#photutils.segmentation.SegmentProperties

# ## Define the approximate isophotal ellipses for each object

# In[63]:


y0, y1 = 2700, 3000
x0, x1 = 3350, 3650

from photutils import EllipticalAperture
r = 3.    # approximate isophotal extent
apertures = []
for prop in props:
    position = (prop.xcentroid.value - x0, prop.ycentroid.value - y0)
    a = prop.semimajor_axis_sigma.value * r
    b = prop.semiminor_axis_sigma.value * r
    theta = prop.orientation.value
    apertures.append(EllipticalAperture(position, a, b, theta=theta))


# In[64]:


fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8))
ax1.imshow(scale_image(data[y0:y1, x0:x1], scale='sqrt', percent=98.),
           origin='lower', cmap='Greys_r')
ax2.imshow(segm[y0:y1, x0:x1], origin='lower', cmap='jet')
for aperture in apertures:
    aperture.plot(color='red', lw=1.5, alpha=0.5, ax=ax1)
    aperture.plot(color='white', lw=1.5, alpha=1.0, ax=ax2)


# ## The segmentation image can be reused on other data:

# ### - define the source segmentation image from a detection image and then use it to perform photometry in individual bands)

# ### - data in individual bands must be registered pixelwise

# In[65]:


# read the F105W XDF data
path = '/Users/lbradley/Box Sync/workshop/data/glue'
fn = 'hlsp_xdf_hst_wfc3ir-60mas_hudf_f105w_v1_sci.fits'
data_f105w = fits.getdata(op.join(path, fn))


# In[66]:


# extract and measure sources in the F105W data
props_f105w = segment_properties(data_f105w, segm, wcs=wcs, labels=labels)
tbl_f105w = properties_table(props_f105w, columns=columns)
tbl_f105w


# In[67]:


# combine the F160W and F105W tables
from astropy.table import hstack
phot = hstack([tbl, tbl_f105w], table_names=['f160w', 'f105w'])
phot


# ## Segmentation image can be modified before getting photometry/properties
# 
#  - remove source segments (artifacts, diffraction spikes, etc.)
#  - combine segments
#  - mask regions of segmentation image (e.g. near image borders)
# 
# See http://photutils.readthedocs.org/en/latest/photutils/segmentation.html
#  
# ### A SourceExtractor segmentation image can be input to segment_properties()
# ```
# CHECKIMAGE_TYPE   SEGMENTATION
# CHECKIMAGE_NAME   segmentation.fits
# ```

# ## Part 3: Detecting Stars (DAOFIND)

# In[68]:


# load a star image
# https://github.com/astropy/photutils-datasets/blob/master/data/M6707HH.fits.gz
from photutils.datasets import load_star_image
hdu = load_star_image()    
star_data = hdu.data[0:400, 0:400]


# In[69]:


# estimate background and background rms from
# sigma-clipped statistics
from astropy.stats import sigma_clipped_stats
mean, median, std = sigma_clipped_stats(star_data, sigma=3.0)


# In[70]:


# find stars using daofind
from photutils import daofind
sources = daofind(star_data - median, fwhm=3.0, threshold=5.*std)    
sources


# In[71]:


# plot the results
from photutils import CircularAperture
from astropy.visualization import SqrtStretch
from astropy.visualization.mpl_normalize import ImageNormalize
positions = (sources['xcentroid'], sources['ycentroid'])
apertures = CircularAperture(positions, r=4.)
norm = ImageNormalize(stretch=SqrtStretch())
plt.imshow(star_data, cmap='Greys', origin='lower', norm=norm)
apertures.plot(color='blue', lw=1.5, alpha=0.5)


# ## Part 4: 2D Background estimation

# In[72]:


# lets add a large background gradient
ny, nx = star_data.shape
y, x = np.mgrid[:ny, :nx]
gradient =  x * y * 1.e-1
data2 = star_data + gradient
plt.imshow(scale_image(data2, scale='sqrt', percent=99.), origin='lower', cmap='Greys')


# In[73]:


from photutils.background import Background
bkg = Background(data2, (10, 10), filter_shape=(3, 3), method='median')


# In[74]:


# plot the estimated background
plt.imshow(bkg.background, origin='lower', cmap='Greys')


# In[75]:


# plot the background-subtracted data
plt.imshow(scale_image(data2 - bkg.background, scale='sqrt', percent=99.), origin='lower', cmap='Greys')


# In[76]:


threshold = median + 2.*std
segm_stars = detect_sources(data2 - bkg.background, threshold, npixels=5)
plt.imshow(segm_stars)


# In[77]:


# create a source mask from the segmentation image
mask = (segm_stars > 0)
bkg2 = Background(data2, (10, 10), filter_shape=(3, 3),
                  method='sextractor', mask=mask)


# In[78]:


plt.imshow(scale_image(data2 - bkg2.background, scale='sqrt', percent=99.),
           origin='lower', cmap='Greys')


# In[79]:


y0 = 200
data2_sub = data2 - bkg2.background
plt.plot(star_data[y0, :])
plt.plot(data2[y0, :])
plt.plot(data2_sub[y0, :], color='cyan')
plt.plot(star_data[y0, :] - median)
plt.axhline(0.0)


# ## imutils
# ###Image Statistics and Arithmetic Convenience Tools
# 
# - Experimental development code from agile sprints
# - Hope to move functionality to the astropy core or affiliated packages (or possibly upstream to numpy, scipy, or scikit-image)
# 
# - Code: https://github.com/spacetelescope/imutils
# - (very preliminary) Documentation: http://imutils.readthedocs.org/en/latest/
# 
# 

# In[79]: