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

# Astronomical data analysis using Python
# =======
# 
# Lecture 10
# -----------------
# 
# 

# In[1]:


from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.wcs import WCS
from astropy.visualization import ZScaleInterval
from astropy.coordinates import SkyCoord
hdulist = fits.open('h_n4603_f555_mosaic.fits')
wcs = WCS(hdulist[0].header)
interval = ZScaleInterval()
vmin,vmax = interval.get_limits(hdulist[0].data)
ax = plt.subplot(111, projection=wcs)
ax.imshow(hdulist[0].data, cmap='gray_r', vmin=vmin, vmax=vmax, interpolation=None, origin='lower')
ax.set_xlabel("Right Ascension"); ax.set_ylabel("Declination")
ax.coords.grid(color='red', alpha=0.5, linestyle='solid')
ax.plot_coord(SkyCoord("12h40m57s","-40d57m33s", frame="fk5"), "ro")
hdulist.close()


# In[2]:


import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import make_lupton_rgb # see Lupton et al. (2004)
# These images are available at https://archive.stsci.edu/prepds/appp/lmc.html
image_r = fits.getdata('h_pu4k2bs01_f814w_sci_v20.fits')/2
image_g = fits.getdata('h_pu4k2bs01_f606w_sci_v20.fits')/3
image_b = fits.getdata('h_pu4k2bs01_f450w_sci_v20.fits')
image = make_lupton_rgb(image_r, image_g, image_b, stretch = 0.03, Q=10, filename='lmc.jpg')
plt.imshow(image,origin='lower')


# # The composite RGB image
# 
# <IMG SRC="lmc.jpg" width=450px>

# # astropy.stats
# 
# The `astropy.stats` package holds statistical functions or algorithms used in astronomy. While the `scipy.stats` and `statsmodels` packages contains a wide range of statistical tools, they are general-purpose packages and are missing **some tools that are particularly useful or specific to astronomy** . This package provides this missing functionality, but will not replace scipy.stats if its implementation.
# 
# You will find, in practice, that `scipy.stats` which is a huge package, contains almost everything you will need. `astropy.stats` adds a few missing pieces. We will discuss two topics from `astropy.stats` - sigma clipping and jackknife errors - which I have found useful. 
# 
# 

# # Sigma Clipping

# In[3]:


from astropy import stats
import numpy as np
data = np.array([1, 5, 6, 8, 100, 150, 3, 2])
clipped_data = stats.sigma_clip(data, sigma=2, maxiters=5)
print (clipped_data)
print (data.mean())
print (clipped_data.mean())


# # The jackknife statistic

# In[4]:


from astropy.stats import jackknife_stats
data = np.array([1,2,3,4,5,6,7,8,9,0]) 
test_statistic = np.mean
estimate, bias, stderr, conf_interval = jackknife_stats(data, test_statistic, 0.95)
print (estimate)
print (bias)
print (stderr)  
print (conf_interval)


# Cosmology
# =========
# 
# This submodule of astropy allows you to do various cosmological calculations based on a model of cosmology.
# 
# We begin by importing the cosmology sub-module.

# In[5]:


from astropy import cosmology


# Now, before we can make do any cosmological calculations, we need to choose a model. Let's do that.

# In[6]:


print (cosmology.parameters.available)


# The above are the various models of cosmology available, you can choose one of them by saying, 

# In[7]:


from astropy.cosmology import WMAP9


# Or you could define your own cosmology by saying,
# 
#     from astropy.comsology import FlatLambdaCDM
#     mycosmo = FlatLambdaCDM(..., ...., ....) 
#     
# Refer documentation for more details. From Astropy 5.0 onwards, you can read or write a cosmology from a file.

# Performing Cosmological Calculations
# -----------------

# In[8]:


WMAP9.Ode(3)  # density parameter for dark energy at redshift z=3   (in units of critical density)


# In[9]:


WMAP9.critical_density(3)   # critical density at z=3


# In[10]:


WMAP9.Tcmb(1100)    # CMB temperature at z=1100


# In[11]:


WMAP9.angular_diameter_distance(2)  # Angular diameter distance in Mpc at z=2.


# In[12]:


WMAP9.arcsec_per_kpc_comoving(3)  # Angular separation in arcsec corresponding to a comoving kpc at z=3


# In[13]:


WMAP9.scale_factor(4)     # a = 1/(1+z)


# In[14]:


WMAP9.age(1100)       # Age of universe at z=1100


# In[15]:


print (dir(WMAP9))


# # astropy.modeling
# 
# astropy.modeling is a module in Python designed to give you
# 
# * Access to commonly used models.
# * As well as fit them to various data.
# 

# In[16]:


from astropy.modeling import models
import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(9, 14, 100)
gauss_example1 = models.Gaussian1D(amplitude=1.0, mean=12, stddev=0.5)
gauss_example2 = models.Gaussian1D(amplitude=2.0, mean=13.5, stddev=0.2)
gauss_total = gauss_example1 + gauss_example2
y = gauss_total(x)

plt.scatter(x,y)


# In[17]:


import numpy.random as nr

y_noise = nr.normal(0, 0.1, len(x))
y_obs = 12 + 0.01*x**2 + y + y_noise
plt.scatter(x, y_obs)


# We saw how trivial it is to use a model and actually evaluate it over a range. But a more useful thing we need to do with models is to fit some data. Now, pretend that x and y_obs are two arrays that contain our observed data as in the plot above.

# Once we have a plot, our next step is to choose a model. How do we choose a model? Let us assume that these are spectral features with some known wavelengths. Each feature can be assumed to be a Gaussian. To minimise the number of independent parameters, we may also consider that the difference in the wavelengths is a constant. 
# 
# Next, these emission features are sitting atop a continuum. Assuming that the continuum is not varying at a fast rate wrt wavelength, we can further assume that a quadratic polynomial suffices in accounting for this variation.
# 
# So, our model could be a second order polynomial plus two Gaussians, with different means but separated by a fixed wavelength. *There are subjective decisions we have made here informed by our knowledge of the physical situation.*
# 

# In[18]:


# So, let us define our model.
model = models.Gaussian1D(amplitude=1.0, mean=12.1, stddev=0.5) +        models.Gaussian1D(amplitude=1.0, mean=13.6, stddev=0.4) +        models.Polynomial1D(degree=2)

print(model.param_names)


# In[19]:


# Our model is not complete. We must supply our constraint.
def constraint_mean(model):
    mean_0 = model.mean_1 - 1.5
    return mean_0

model.mean_0.tied = constraint_mean


# The model is now ready. We have the (simulated)  data. What we now need is a fitting algorithm. Let us choose the Levenberg-Marquardt algorithm. For linear fits, use LinearLSQFitter()

# In[20]:


from astropy.modeling import fitting
fitter = fitting.LevMarLSQFitter()

model_fit = fitter(model, x, y_obs)
print (model_fit.param_names)
print(model_fit.parameters)


# In[21]:


dict(zip(model_fit.param_names, model_fit.parameters))


# In[22]:


plt.scatter(x, y_obs, color='black', alpha=0.3)
plt.plot(x, model_fit(x), color='red')


# # `astropy.convolution`
# 
# `astropy.convolution` provides convolution functions and kernels that offer improvements compared to the SciPy `scipy.ndimage` convolution routines, including:
# 
# * Proper treatment of `NaN` values (ignoring them during convolution and replacing NaN pixels with interpolated values)
# * A single function for 1D, 2D, and 3D convolution
# * Improved options for the treatment of edges
# * Both direct and Fast Fourier Transform (FFT) versions
# * Built-in kernels that are commonly used in Astronomy

# I have tried to cover the most frequently used aspects of `astropy`. But to learn more do check out the excellent tutorials and documentation of astropy available at:
# 
# http://www.astropy.org
# 
# Astropy is still undergoing rapid development and new features are being added continuously. The coordinated packages are also in a state of flux. An excellent resourse is the astropy mailing list:
# 
# 
# https://mail.python.org/pipermail/astropy/
# 
# Do join this list to get answers to any issues you face with astropy and its affiliated packages.
# 

# # Astroquery
# 
# You may have accessed data in some online archives, one search at a time. This may be too slow if your sample contains thousands of objects. In such a situation, you must write a program to automatically access the data that you are interested in.Astroquery is the Python language package that will provide you with this capability and it is very easy to code. It contains a number of subpackages for each data repository.
# 
# 
# There are two other packages with complimentary functionality as Astroquery: pyvo is an Astropy affiliated package, and Simple-Cone-Search-Creator which generates a cone search service complying with the IVOA standard. They are more oriented to general virtual observatory discovery and queries, whereas Astroquery has web service specific interfaces.
# 
# Astroquery follows a **continuous deployment model**.
# 

# In[23]:


import astroquery
print (astroquery.version.version)


# **Use the latest astroquery version available**. Website configurations change constantly and your programs could stop working suddenly because the website changed its interface.
# 
# To help you get started, see the sample queries in the astroquery gallery.
# 
# https://astroquery.readthedocs.io/en/latest/gallery.html
# 

# In[24]:


from astroquery.simbad import Simbad
from astropy import coordinates
import astropy.units as u
# works only for ICRS coordinates:
c = coordinates.SkyCoord("05h35m17.3s -05d23m28s", frame='icrs')
r = 1 * u.arcminute
result_table = Simbad.query_region(c, radius=r)
result_table.pprint(show_unit=True, max_width=100, max_lines=10)


# # Cross match with any Vizier catalog

# In[25]:


get_ipython().system('cat pos_list.csv')


# In[26]:


from astropy import units as u
from astroquery.xmatch import XMatch
table = XMatch.query(cat1=open('pos_list.csv'),
                     cat2='vizier:II/246/out', # 2MASS catalog
                     max_distance=5 * u.arcsec, colRA1='ra',
                     colDec1='dec')
print (type(table))
print (table.colnames)


# In[27]:


print (table)


# # Getting data from the SDSS

# In[28]:


from astroquery.sdss import SDSS
from astropy import coordinates as coords
pos = coords.SkyCoord('0h8m05.63s +14d50m23.3s', frame='icrs')
xid = SDSS.query_region(pos, spectro=True)
print(xid)


# In[29]:


sp = SDSS.get_spectra(matches=xid)
print (sp)
spec = sp[0][1].data


# In[30]:


import matplotlib.pyplot as plt
print (spec.dtype)
flux = spec['flux']
loglam = spec['loglam']
model = spec['model']
plt.scatter(loglam,flux,marker='.',alpha=0.3,color='red')
plt.plot(loglam,model)


# In[31]:


import numpy as np
im = SDSS.get_images(matches=xid,band='r')
print (im)
image = im[0][0].data
print (image.shape)


# # Display the downloaded image

# In[32]:


plt.imshow(np.sqrt(image),origin='lower',vmin=10)


# # Be careful
# 
# Astroquery can easily run amok and download a lot of data that you don't need. So, be a little cautious at the download steps and make sure that you are downloading what you expect. Setting up `assert` statements before you start the download will be very useful.