# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ***
# # HCI Week 9 - Principal Component Analysis
# *<NAME> // Leiden Observatory*
# ***
# Today we cover the implementation of Principal Component Analysis as described in [Amara and Quanz (2012)](http://adsabs.harvard.edu/abs/2012MNRAS.427..948A). Go and read Section 3.1.
#
# ## Principle Component Analysis
#
# PCA is 'shift, scale, rotate'. The `first principal component` is the direction in which the transformed data has **MAXIMUM** variance. The second and higher components are *orthogonal to the first principal component*.
#
# Any single science image can be thought of as a Iinear combination of an orthogonal basis set formed from all the science images. We can write this as:
#
# $$
# I(\vec{x})=\sum a_{i} \phi_{i}(\vec{x})
# $$
#
# where $I(\vec{x})$ is the image of the PSF, $\phi(\vec{x})$ is a given basis and $a_{i}$ is the coefficient for each basis function.
#
# Instead of small wedges, the whole stack of images is taken and a **Principal Component Analysis** (PCA) is performed on it.

# +
import numpy as np
import matplotlib.pyplot as plt
# %config InlineBackend.figure_format = 'retina' 
from astropy.io import fits

import matplotlib as mpl
mpl.rc('image', interpolation='nearest', origin='lower')

def wfits(im, fname):
    """wfits - write im to file fname, automatically overwriting any old file"""
    from astropy.io import fits
    hea = fits.PrimaryHDU(im)
    hea.writeto(fname, overwrite=True)

def r_theta(im, xc, yc):
    """r_theta - make a radius mask and return the radius rr and the angle phi for point (xc,yc)"""
    ny, nx = im.shape
    yp, xp = np.mgrid[0:ny,0:nx]
    yp = yp - yc
    xp = xp - xc
    rr = np.sqrt(np.power(yp,2.) + np.power(xp,2.))
    phi = np.arctan2(yp, xp)
    return(rr, phi)

def cen_rot2 (im, rot, dim_out, offset1=(0,0), offset2=(0,0), order=2):
    """cen_rot2 - takes a cube of images im, and a set of rotation angles in rot,
                and translates the middle of the frame with a size dim_out to the middle of
                a new output frame with an additional rotation of rot.
                """
    from scipy.ndimage import affine_transform
    a = rot * np.pi / 180.
    
    # make a rotation matrix
    transform=np.array([[np.cos(a),-np.sin(a)],[np.sin(a),np.cos(a)]])
    
    # calculate total offset for image output

    # determine centre of input image

    # -0.5 is there for fencepost counting error
    c_in = np.array(offset1) - 0.5
    c_out = 0.5 * np.array(dim_out) - 0.5

    # c_out has to be pre-rotated to make offset correct
    offset = c_in - c_out.dot(transform) - np.array(offset2).dot(transform)
    
    # perform the transformation
    dst=affine_transform( \
        im,transform.T, order=order,offset=offset, \
        output_shape=dim_out, cval=0.0)
    return(dst)

def derot_cube(imin, rotangle, out_dims, star_cen):
    """derot_cube - takes a cube of images im, and a set of rotation angles in rot,
                and translates the middle of the frame with a size dim_out to the middle of
                a new output frame with an additional rotation of rot degrees.
                """
    # get shape of input image
    nim, nx, ny = imin.shape
    
    # make a zero value numpy array for the output
    im_out = np.zeros((nim, out_dims[0], out_dims[1]))
    
    # for every frame in imin perform the derotation of rotangle
    for j, frame in enumerate(np.split(imin, nim, axis=0)):
        frame_sq = np.squeeze(frame)
        frame_rot = cen_rot2(frame_sq, rotangle[j], out_dims, offset1=star_cen)
        im_out[j,:,:] = frame_rot
        
    return(im_out)    


# -

# ## Data of the planetary system HR 8799

# The data is from the [NIRC2](https://www2.keck.hawaii.edu/inst/nirc2/) camera on the Keck II 10m telescope, from observing program U159N2 by <NAME>. This data is courtesy of <NAME>. Image registration by <NAME>. 
#
# The data is taken in the L prime (Lp) filter (covers 3.42 to 4.12 microns in wavelength), in the narrow camera pixel scale (10mas/pixel)
#
# https://www2.keck.hawaii.edu/inst/nirc2/genspecs.html

# ## Image data and position angle data
# The data is approximately 100 images taken over the course of one night on Keck in 2009. Keck is an Alt-az telescope, so the sky rotates in the camera field of view. To derotate the images so that they are all aligned on the sky, you need a list of derotation angles, provided in `rotnth.fits`. You can read it in the same way as the data cube file.
#
# To save storage sapce, this cube is a cropped version of the data from the earlier Practicum. The sky background is **already subtracted** from this data set, so you do not need to do it again.
#

# +
# read in the data cube
im, he = fits.getdata('im_pca.fits.gz', header=True)

# read in the rotation angles
rot = fits.getdata('rotnth.fits')

# -

# ## Performing regular ADI
#
# First, we will perform a regular ADI reduction, but this time using an annular mask (Question 9.1) and [numpy Broadcasting](https://numpy.org/devdocs/user/theory.broadcasting.html) to pull out the required pixels from the data cube (Question 9.2).
#
# ### Defining an annulus
#
# Assume that the star is in the centre of each image frame.
#
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.1</span> <span style="  font-size:1em; color: SteelBlue;"> (2 points): </span>
#
# 1.  **Use `r_theta` to make a mask `m` that has the same dimensions as a single image frame, with inner radius of 10 pixels and outer radius of 120 pixels. Display this image with `ax.imshow()`.**
# * **Print out the numer of pixels in the annular region and print out the shape of mask `m`.**
#
# ***
#
#

# +
# Q1 answer here
xc,yc = 150,150
m = r_theta(im[0],xc,yc)[0]
m[(m>120)] = 0
m[(m<10)] = 0
m[(m>=10)*(m<=120)] = 1
m = m.astype(int)
fig, ax = plt.subplots(1,1, figsize=(6,6))
ax = ax.imshow(m)
fig.colorbar(ax)
plt.title('Annular mask')
plt.show()
# Q1 end of answer

print('Number of pixels in annular region is {}'.format(len(m[m>0])))
print('The shape of the mask is {}'.format(m.shape))
# -

# ## Using broadcasting to select annular regions through the whole data cube
#
# You've used broadcasting during this course to select columns and lines through data cubes and to generate functions. You can also use 2D masks to select data from a 3D data cube.
#
#
# <div class="alert alert-block alert-info">
# <b>The Broadcasting Rule: </b> <br>In order to broadcast, the size of the trailing axes for both arrays in an operation must either be the same size, or one of them must be one. <a href="https://numpy.org/devdocs/user/theory.broadcasting.html"> Explanation</a></div>
#
#
# You can use a mask to select pixels from a 2D image like so - in this case, we are looking at frame 5 in the image cube `im`:
#     
#     ann_pix = im[5,:][m]
#     print ann_pix
#     (44928,)
#    
# Notice that the mask selection has flattened the output array into a 1D numpy vector.
#
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.2</span> <span style="  font-size:1em; color: SteelBlue;"> (2 points): </span>
#
# If your data cube has a shape of `[X,Y,N]` and your mask has a shape of `[X,Y]`, then `image[mask]` will broadcast across `N` to output a 2D array containing `N` columns with `1D` vectors.
#
# The dimensions of `im` are in the wrong order, though. We need to reorder them!
#
# 1. **Use `np.swapaxes()` to rearrange the order of dimensions in `im`.**
# * **Use the annulus mask to pull out all annulus pixels in the data cube into an array `A`.**
# * **Print the shape of `A`.**
#
#
# <div class="alert alert-block alert-info">
# <b>Tip: </b> Using <code>print(x.shape)</code> a lot will help you understand what's going on. </div>
#
#
# ***

# +
print("The shape of the image im is:     {}".format(im.shape))

# Q2 answer here
im_swap = np.swapaxes(im,0,2)
im_swap = np.swapaxes(im_swap,0,1)
A = im_swap[m==1]


# Q2 end of answer

print("The shape of the array A is:      {}".format(A.shape))
# -

# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.3</span> <span style="  font-size:1em; color: SteelBlue;"> (2 point): </span>
#
# 1. **Calculate the median annulus of all the annuli combined and subtract it off all the individual annuli. (This is the equivalent of calculating the median image of the cube, and then subtracting this off all the individual frames in the cube.) Call it `Amedsub`.**
#
#
# <div class="alert alert-block alert-info">
# <b>Note: </b> You should NOT use a <code>for</code> loop to do this, only do it using broadcasting! Ask us if you are still unsure how to do this. </div>
#
# * **Put `Amedsub` back into an empty data cube and derotate the image frames into 300x300 pixel images.  They should be aligned with the sky. Then take the median of this image stack. Call the new array `ADImed`. Use `ax.imshow(ADImed)` to plot it from a range of -500 to +500 counts.**
#
#
# <div class="alert alert-block alert-info">
# <b>Note 1: </b> After you have put <code>Amedsub</code> into the data cube, you will need to rearrange the axes back when you do the derotation of the image frames! <br>
# <b>Note 2: </b> You should see three planets, with the innermost one showing lots of speckle noise next to it.</div>
#
#
# ***

# +
# Q3 answer here

# subtract the median annulus off of the individual annulus
Amedsub = A - np.median(A,axis=1)[:,np.newaxis]
# put annuli back to empty image cube
im_empty = np.zeros_like(im_swap)
im_empty[m==1] = Amedsub
# swap axes to original setting
im_back = np.swapaxes(im_empty,0,1)
im_back = np.swapaxes(im_back,0,2)
# derotate the image cube
outdim = 300,300
im_rot = derot_cube(im_back,rot,outdim,(150,150))
ADImed = np.median(im_rot,axis=0)

fig, ax = plt.subplots(1,1, figsize=(6,6))
ax.imshow(ADImed,vmin=-500,vmax=500)
ax.set_xlabel('X [pixels]')
ax.set_ylabel('Y [pixels]')
plt.show()
# Q3 end of answer
# -

# ## Computing the SVD
#
# The next step is to now calculate the Singular Value Decomposition of `A`.
#
# To do this, we also need to make sure that the data in each annulus has a mean of zero - check how to take the mean of one specific axis by reading the documentation on [`np.mean(axis= ... `](https://numpy.org/doc/stable/reference/generated/numpy.mean.html)
#
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.4</span> <span style="  font-size:1em; color: SteelBlue;"> (1 point): </span>
#
# **For `Amedsub`, calculate and plot the mean flux in each individual annulus. Subtract it off the `Amedsub`. Call this array `Ain`.** 
#
# ***

# Q4 answer here
mean_flux = np.mean(Amedsub,axis=0)
plt.figure(figsize=(6,6))
plt.plot(mean_flux)
plt.xlabel('Frame number')
plt.ylabel('Mean flux')
plt.show()
Ain = Amedsub - mean_flux[np.newaxis,:]
# Q4 end of answer

# ## Performing the SVD
#
# We now decompose `Ain` into three separate matrices using `np.linalg.svd`, so that:
#
# $$ \texttt{Ain} = \texttt{U} \cdot \texttt{S} \cdot \texttt{V} $$
#
# Where `S` is a diagonal matrix where the diagonal elements are the eigenvalues, `U` contains columns of eigenvectors, and `V` contains rows of eigenvectors.
#

# +
U, s, V = np.linalg.svd(Ain, full_matrices=False)
print("Shape of U:       {}".format(U.shape))
print("Shape of s:       {}".format(s.shape))
print("Shape of V:       {}".format(V.shape))

S = np.diag(s)
# -

# ## The size of the eigenvalues
#
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.5</span> <span style="  font-size:1em; color: SteelBlue;"> (2 points): </span>
#
# **Make a `log10` plot of the eigenvalues `s`. Label the x axis 'PCA number' and the y-axis 'log10 eigenvalues'.** 
#
# ***

# Q5 answer here
plt.figure(figsize=(6,6))
plt.semilogy(s)
plt.xlabel('PCA number')
plt.ylabel('$\log$ 10 eigenvalues')
plt.show()
# Q5 end answer 

#
#  
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.6</span> <span style="  font-size:1em; color: SteelBlue;"> (2 points): </span>
#
# **Check that $\texttt{U} \cdot \texttt{S} \cdot \texttt{V} $ is close to `Ain` by using the `np.isclose()` function and printing the fraction of elements that return `True`.**
#
#
# <div class="alert alert-block alert-info">
# <b>Note: </b> Using Numpy the matrix multiplication: $C = A\cdot B$ can be done with: <code>C = np.dot(A,B)</code> </div>
#
# ***

# +
# Q6 answer here

Aclose = np.dot(np.dot(U,S),V)
print("The fraction of elements return 'True' is: ",len(np.isclose(Aclose,Ain))/len(Ain))

# Q6 end answer 
# -

# ## Approximate the PSF with $k$ numbers of eigenmodes
#
# We can now use all the eigenmodes to *precisely* reconstruct each PSF in the data cube - but of course, that's NOT what we want, because all the images have the planets in them.
#
# Instead, we want to take the first `k` eigenmodes and fit them to each input PSF frame individually. The first `k` modes should not have any planet flux in them, because `k` modes will not be enough to encode a planet moving in an arc through the whole data cube. 
#
# Starting with `k=10`, make `Uk` , `sk` and `Vk` such that you clip out the first `k` columns of `U`, the first `k` values of `s` and the first `k` rows of `V`.
#

# +
k=10

Uk = U[:,:k]
sk = s[:k]
Vk = V[:k,:]

print("Shape of Uk:       {}".format(Uk.shape))
print("Shape of sk:       {}".format(sk.shape))
print("Shape of Vk:       {}".format(Vk.shape))
# -

# ## Reconstructing the PCA PSF and subtracting it
#
# We can now use the first $k$ eigenimages and fit them to each of the individual input images, creating a `reconstructed` image which can then be subtracted off `Ain` to produce the PCA subtracted annuli.
#
# After calculating `residuals` below, put it into an image array called `resid`.

# +
transformed = np.dot(Vk, Ain.T)

reconstructed = np.dot(transformed.T, Vk)

residuals = Ain - reconstructed

resid = np.zeros_like(im_swap)

resid[m==1] = residuals

wfits(np.swapaxes(resid,0,-1), 'resid.fits')
# -

# ## Derotating the residuals and stacking together
#
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.7</span> <span style="  font-size:1em; color: SteelBlue;"> (4 points): </span>
#
# 1. **Finally, take the data cube `resid` and derotate to the sky.** 
# *  **Take its median and combine it in a file called `PCAmed`.**
# *  **Display `ADImed` from -400 to 400 and `PCAmed` from -200 to 200 using `ax.imshow()`.** 
#
# ***

# +
# Q7 answer here
resid = np.swapaxes(resid,1,2)
resid = np.swapaxes(resid,0,1)
resid_rot = derot_cube(resid,rot,outdim,(150,150))
PCAmed = np.median(resid_rot,axis=0)


fig, (ax1,ax2) = plt.subplots(1,2, figsize=(12,6))
ax1.imshow(ADImed,vmin=-400,vmax=400)
ax1.set_xlabel('X [pixels]')
ax1.set_ylabel('Y [pixels]')
ax1.set_title('ADImed')
ax2.imshow(PCAmed,vmin=-200,vmax=200)
ax2.set_xlabel('X [pixels]')
ax2.set_ylabel('Y [pixels]')
ax2.set_title('PCAmed')
plt.show()

# Q7 end answer 
# -

# ## Comparing median ADI and PCA
#
#
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.8</span> <span style="  font-size:1em; color: SteelBlue;"> (4 points): </span>
#
# 1. **Define annuli from 0.2 arcsec through to 2.3 arcsec in 0.1 arcsec steps.**
# * **Plot the radius of the middle of width of the ring versus the r.m.s. in each the ring using ADI and for PCA with k = 5, 10, and 20.**
# * **Plot them all on the same graph. Use `ax.set_xlim([0,2.3])` to set the axis limit. Add labels, legend, title, ect.** 
#
# ***
#
#
#

# +
def get_PCAmed(k=10):
    
    Uk = U[:,:k]
    sk = s[:k]
    Vk = V[:k,:]

    transformed = np.dot(Vk, Ain.T)
    reconstructed = np.dot(transformed.T, Vk)
    residuals = Ain - reconstructed
    resid = np.zeros_like(im_swap)
    resid[m==1] = residuals
    resid = np.swapaxes(resid,1,2)
    resid = np.swapaxes(resid,0,1)
    resid_rot = derot_cube(resid,rot,outdim,(150,150))
    PCAmed = np.median(resid_rot,axis=0)
    return PCAmed



# +
# Q8 answer here
def get_PCArms(PCAmed):
    rmask = r_theta(ADImed,150,150)[0]
    annuli = np.arange(.2,2.3+.1,.1)
    pscale = 10e-3 
    annuli_px = annuli/pscale
    rms_PCA = [np.std(PCAmed[(rmask<annuli_px[i+1])*(rmask>annuli_px[i])]) for i in range(len(mid_rad))]
    return np.asarray(rms_PCA)
    
k = [5,10,20]
PCAmed1 = get_PCAmed(k[0])
PCAmed2 = get_PCAmed(k[1])
PCAmed3 = get_PCAmed(k[2])

rmask = r_theta(ADImed,150,150)[0]
annuli = np.arange(.2,2.3+.1,.1)
pscale = 10e-3 
annuli_px = annuli/pscale
mid_rad = .5*(annuli[:-1]+annuli[1:])

rms_ADI = [np.std(ADImed[(rmask<annuli_px[i+1])*(rmask>annuli_px[i])]) for i in range(len(mid_rad))]
rms_PCA1 = get_PCArms(PCAmed1)
rms_PCA2 = get_PCArms(PCAmed2)
rms_PCA3 = get_PCArms(PCAmed3)
rms_PCA = [rms_PCA1,rms_PCA2,rms_PCA3]

fig, ax = plt.subplots(1,1, figsize=(6,6))
ax.plot(mid_rad,rms_ADI,label='RMS ADI')
[ax.plot(mid_rad,rms_PCA[i],label='RMS PCA, k = '+str(k[i])) for i in range(3)]
ax.set_xlabel('$radius [arcsec]$')
ax.set_ylabel('RMS')
ax.set_xlim([0,2.3])
ax.legend()
ax.set_title('RMS in annuli using ADI and PCA methods')
plt.show()
# Q8 end of answer
# -

#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Question 9.9</span> <span style="  font-size:1em; color: SteelBlue;"> (4 points): </span>
#
# **Measure the total flux from the planet at (188,95) using a 10 pixel diameter aperture for the four images above.** 
#
# **Answer the question: What happens to the planet flux as you go to higher values of $k$?**
#
# ***
#
#

# +
# sanity check that the planet is indeed at about (188,95)
def zoom(im,x,y,bb):
    '''zoom - returns a square 2D image cut out from im centered on location (x,y) and with half-width bb'''
    return(im[y-bb:y+bb,x-bb:x+bb])

plt.imshow(zoom(ADImed,188,95,10))
plt.show()

# +
# Q9 answer here
planet_mask = r_theta(ADImed,188,95)[0]
planet_mask[planet_mask<=10] = 1
planet_mask[planet_mask>10] = 0
flux_ADI = np.sum(ADImed[planet_mask==1])
flux_PCA5 = np.sum(PCAmed1[planet_mask==1])
flux_PCA10 = np.sum(PCAmed2[planet_mask==1])
flux_PCA20 = np.sum(PCAmed3[planet_mask==1])
flux_tot = [flux_ADI,flux_PCA5,flux_PCA10,flux_PCA20]
label = ['ADI','k=5','k=10','k=20']

plt.figure(figsize=(5,5))
plt.plot(label,flux_tot)
plt.xlabel('Methods')
plt.ylabel('Flux')
plt.title('Flux inside 10 pixel diameter of the planet')
plt.show()
print("It can be seen, when going to high k value, the planet flux becomes lower.")
# Q9 end of answer
# -

#
# <div class="alert alert-block alert-info">
# <b>REMEMBER:</b> to make sure your code compiles cleanly before submitting it! Do <b>NOT</b> send your fits cubes back! Only this Jupyter notebook!</div>
#
# **Thank you for submitting all the Jupyter notebooks from the past 9 weeks. We hope you enjoyed the tutorials and that you now understand (how to apply) different techniques used in High Contrast Imaging.**
#
# ---
# <span style="  font-size:2em; color: SteelBlue;">Your time has a value as well</span> 
#  </span>
#
# ---
#
# ### How long did this Practicum take?
# Please enter the number of hours you spent on this Practicum (including the time spend besides the time spend during the Lecture) below:
#

# Number of hours taken for this practicum:
# ~ 3 hours.



StarcoderdataJupyterScriptsDedupFiltered