# Importing Modules
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
sns.set_context("notebook")

import astropy
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy import constants as const

"""
Explanation: Week 14 - Astropy
Today's Agenda

Useful functions of Astropy
Units
Time
Coordinates
FITS files
Analytic functions
AstroPy Tables and different formats

Astropy is a package that is meant to provide a lot of basic functionality for astronomy work in Python
This can be roughly broken up into two areas. One is astronomical calculations:
* unit and physical quantity conversions
* physical constants specific to astronomy
* celestial coordinate and time transformations
The other is file type and structures:
* FITS files, implementing the former standalone PyFITS interface
* Virtual Observatory (VO) tables
* common ASCII table formats, e.g. for online catalogues or data supplements of scientific publications
* Hierarchical Data Format (HDF5) files
AstroPy normallly comes with the Anaconda installation. But in case you happen to not have it installed it on your computer, you can simply do a 
sh
    pip install --no-deps astropy
You can always update it via 
sh
    conda update astropy
This is just a glimpse of all the features that AstroPy has:
<img src="./images/astropy_sections.png" alt="Astropy Features" width="600">
For purposes of today, we'll focus just on what astropy can do for units, time, coordinates, image manipulation, and more.
End of explanation
"""


d=42*u.meter
t=6*u.second
v=d/t
print v
print v.unit

"""
Explanation: Units
Astropy.units introduces units and allows for unit conversions. It doesn't, however, correctly handle spherical coordinates, but the astropy.coordinates package will address this later.
These units can be used to create objects that are made up of both a value and a unit, and basic math can be easily carried out with these. We can add the .unit and .value properties to get the units and numerical values, respectively.
End of explanation
"""


from astropy.units import imperial
imperial.enable()
print( u.s.find_equivalent_units() )
print( u.m.find_equivalent_units() )

"""
Explanation: Astropy includes a large number of units, and this can include imperial units as well if desired by importing and enabling imperial units. The .find_equivalent_units() function will also return all the other units that are already defined in astropy. Below we do a quick list of the units that are defined for time and length units
End of explanation
"""


print const.c
print const.G

"""
Explanation: The package also provides constants, with the units included. The full list of units can be found here. We can take a quick look at c and G below, and see that these are objects which have value, uncertainty, and units.
End of explanation
"""


Mars=1.5*u.AU
print Mars.to('kilometer')
print Mars.to('mile')
print u.AU.to('kilometer')
print const.c.to('pc/yr')
print const.c.to('fur/fortnight')

"""
Explanation: Astropy has an aditional function that will allow for unit conversions. So we can, for example, create an object that is the distance to Mars, and then convert that to kilometers or miles. A brief note is that if you try to convert a pure unit (like the 4th line below) into another unit, you'll get a unitless value representing the conversion between the two.
This can also be used to convert constants into other units, so we can convert the speed of light to the somewhat useful pc/yr or the entirely unuseful furlong/fortnight
End of explanation
"""


time=1*u.AU/const.c
print(time)

time_s=time.decompose()
print(time_s)

time_min=time_s.to(u.minute)
print(time_min)

"""
Explanation: To use this more practically, we can calculate the time it will take for light to reach the earth just by dividing 1 AU by the speed of light, as done below. Since AU is a unit, and c is in m/s, we end up with an answer that is (AU*m/s). By using .decompose() we can simplify that expression, which in this case will end up with an answer that is just in seconds. Finally, we can then convert that answer to minutes to get the answer of about 8 1/3 minutes that is commonly used. None of this required our doing the conversions where we might've slipped up.
End of explanation
"""


from astropy.time import Time
t=Time(57867.346424, format='mjd', scale='utc')
t1=Time(58867.346424, format='mjd', scale='utc')
print t.mjd
print t.iso
print t.jyear
t1-t

"""
Explanation: Time
Astropy handles time in a similar way to units, with creating Time objects. These objects have two main properties.
The format is simply how the time is displayed. This is the difference between, for example, Julian Date, Modified Julian Date, and ISO time (YYYY-MM-DD HH:MM:SS). The second is the scale, and is the difference between terrestrial time vs time at the barycenter of the solar system.
We can start off by changing a time from one format to many others. We can also subtract times and we will get a timedelta unit.
End of explanation
"""


c = SkyCoord(ra=10.68458*u.degree, dec=41.26917*u.degree, frame='icrs')
print c
print c.ra
print c.dec
print c.ra.hour
print c.ra.hms
print c.galactic

"""
Explanation: Coordinates
Coordinates again work by using an object time defined for this purpose. We can establish a point in the ICRS frame (this is approximately the equatorial coordinate) by defining the ra and dec. Note that here we are using u.degree in specifying the coordinates.
We can then print out the RA and dec, as well as change the units displayed. In the last line, we can also convert from ICRS equatorial coordinates to galactic coordinates.
End of explanation
"""


hosts={}
data=np.loadtxt('./data/planets.tab', dtype='str', delimiter='\t')
print data[0]
hosts['ra_hours']=data[1:,9].astype(float)
hosts['ra']=data[1:,6].astype(float)
hosts['dec']=data[1:,8].astype(float)
#print hosts['ra_hours']
#print hosts['dec']

import astropy.units as u
import astropy.coordinates as coord
from astropy.coordinates import SkyCoord
ra = coord.Angle(hosts['ra']*u.degree)
ra = ra.wrap_at(180*u.degree)
dec = coord.Angle(hosts['dec']*u.degree)

fig = plt.figure(figsize=(12,8))
ax = fig.add_subplot(111, projection="mollweide")
plt.title('Map of Exoplanets')
ax.scatter(ra.radian, dec.radian)
ax.set_xticklabels(['14h','16h','18h','20h','22h','0h','2h','4h','6h','8h','10h'])
ax.grid(True)
plt.show()

"""
Explanation: Slightly practical application of this
Using some of these astropy functions, we can do some fancier applications. Starting off, we import a listing of stars with RA and dec from the attached table, and store them in the coordinate formats that are used by astropy. We then use matplotlib to plot this, and are able to easily convert them into radians thanks to astropy. This plot is accurate, but it lacks reference for where these points are.
End of explanation
"""


a=coord.Angle((np.arange(361)-180)*u.degree)
b=coord.Angle(np.zeros(len(a))*u.degree)
numpoints=360
galaxy=SkyCoord(l=coord.Angle((np.arange(numpoints))*u.degree), b=coord.Angle(np.zeros(numpoints)*u.degree), frame='galactic')
ecliptic=SkyCoord(lon=coord.Angle((np.arange(numpoints))*u.degree), lat=coord.Angle(np.zeros(numpoints)*u.degree), frame='geocentrictrueecliptic')
ecl_eq=ecliptic.icrs
gal_eq=galaxy.icrs
#print gal_eq
fixed_ra=[]
for item in gal_eq.ra.radian:
   if item < np.pi:
      fixed_ra.append(item)
   else:
      fixed_ra.append(item-2*np.pi)
i=np.argmin(fixed_ra)
fixed_dec=[x for x in gal_eq.dec.radian]

fixed_ra_eq=[]
for item in ecl_eq.ra.radian:
   if item < np.pi:
      fixed_ra_eq.append(item)
   else:
      fixed_ra_eq.append(item-2*np.pi)
j=np.argmin(fixed_ra_eq)
fixed_dec_eq=[x for x in ecl_eq.dec.radian]

fig = plt.figure(figsize=(12,8))
ax = fig.add_subplot(111, projection="mollweide")
plt.title('Map of Exoplanets')
ax.scatter(ra.radian, dec.radian)
ax.plot(a.radian, b.radian, color='r', lw=2)
#ax.scatter(gal_eq.ra.radian, gal_eq.dec.radian, color='g')
ax.plot(fixed_ra[i:]+fixed_ra[:i], fixed_dec[i:]+fixed_dec[:i], color='g', lw=2)
ax.plot(fixed_ra_eq[j:]+fixed_ra_eq[:j], fixed_dec_eq[j:]+fixed_dec_eq[:j], color='m', lw=2)
ax.set_xticklabels(['14h','16h','18h','20h','22h','0h','2h','4h','6h','8h','10h'])
ax.grid(True)
plt.show()

"""
Explanation: To fix this, we will add some references to this by adding a few more sets of data points. The first is relatively simple, we put in a line at the celestial equator. This just has to be a set of points that are all at declination of 0, and from -180 to +180 degrees in RA. These are a and b in the below.
We also want to add the planes of the ecliptic and the galaxy on this. For both, we use coordinate objects and provide numpy arrays where one coordinate is at zero, and the other goes from 0 to 360. With astropy we can then easily convert from each coordinate system to ICRS. There's some for loops to modify the plotting, but the important thing is that this will give us a plot that has not just the locations of all the planets that we've plotted, but will also include the celestial equator, galactic plane, and ecliptic plane on it.
End of explanation
"""


# We will use  `wget` to download the necessary file to the `data` folder.
!wget 'http://star.herts.ac.uk/~gb/python/656nmos.fits' -O ./data/hst_image.fits

"""
Explanation: Reading in FITS files
One of the useful things with Astropy is that you can use it for reading in FITS files, and extracting info such as bands, exposure times, intrument information, etc.
In this example, we will read in a FITS image file, and extract its information
End of explanation
"""


from astropy.io import fits
filename = './data/hst_image.fits'
hdulist = fits.open(filename)

"""
Explanation: Now we can extract some of the information stored in the FITS file.
End of explanation
"""


hdulist.info()

"""
Explanation: The returned object, hdulist, (an instance of the HDUList class) behaves like a Python list, and each element maps to a Header-Data Unit (HDU) in the FITS file. You can view more information about the FITS file with:
End of explanation
"""


hdu = hdulist[0]

"""
Explanation: As we can see, this file contains two HDUs. The first contains the image, the second a data table. To access the primary HDU, which contains the main data, you can then do:
End of explanation
"""


np.sort(hdulist[0].header.keys())

"""
Explanation: To read the header of the FITS file, you can read hdulist. The following shows the different keys for the header
End of explanation
"""


hdu.data.shape

"""
Explanation: As we can see, this file contains two HDUs. The first contains the image, the second a data table.
Let's look at the image of the FITS file.
The hdu object then has two important attributes: data, which behaves like a Numpy array, can be used to access the data, and header, which behaves like a dictionary, can be used to access the header information. First, we can take a look at the data:
End of explanation
"""


hdu.header

"""
Explanation: This tells us that it is a 1600-by-1600 pixel image. We can now take a peak at the header. To access the primary HDU, which contains the main data, you can then do:
End of explanation
"""


hdu.header['INSTRUME']

hdu.header['EXPTIME']

"""
Explanation: We can access individual header keywords using standard item notation:
End of explanation
"""


plt.figure(figsize=(10,10))
plt.imshow(np.log10(hdu.data), origin='lower', cmap='gray', vmin=1.5, vmax=3)

"""
Explanation: We can plot the image using matplotlib:
End of explanation
"""


hdu.header['MODIFIED'] = '2014-12-01'  # adds a new keyword

"""
Explanation: You can also add new fields to the FITS file
End of explanation
"""


hdu.data = hdu.data - 0.5

"""
Explanation: and we can also change the data, for example subtracting a background value:
End of explanation
"""


hdu.writeto('./data/hubble-image-background-subtracted.fits', overwrite=True)

!ls ./data

"""
Explanation: This only changes the FITS file in memory. You can write to a file with:
End of explanation
"""


from astropy.analytic_functions import blackbody_lambda, blackbody_nu

def Planck_func(temp, lam_arr, opt='lam'):
    """
    Computes the Blackbody radiation curve of a blackbody of a given temperature `temp`.
    
    Parameters
    ----------
    temp: float or array-like
        temperature(s) of the blackbody
    
    lam_arr: float or array_like
        aray of wavelenths to evaluate the Planck function.
    
    opt: str, optional (default = 'lam')
        Option for returning either the flux of `lambda` (wavelength) or `nu` (frequency).
        Options:
            - `lam`: Return flux for `lambda' or wavelength
            - `nu` : Returns flux for `nu` (frequency)
    """
    wavelengths = lam_arr * u.AA
    temperature = temp * u.K
    with np.errstate(all='ignore'):
        flux_lam = blackbody_lambda(wavelengths, temperature)
        flux_nu  = blackbody_nu(wavelengths, temperature)
    
    if opt=='lam':
        return flux_lam
    if opt=='nu':
        return flux_nu

"""
Explanation: Analytic Functions
Astropy comes with some built-in analytic functions, e.g. the blackbody radiation function.
Blackbody Radiation
Blackbody flux is calculated with Planck law (Rybicki & Lightman 1979)
$$B_{\lambda}(T) = \frac{2 h c^{2} / \lambda^{5}}{exp(h c / \lambda k T) - 1}$$
$$B_{\nu}(T) = \frac{2 h \nu^{3} / c^{2}}{exp(h \nu / k T) - 1}$$
End of explanation
"""


lam_arr = np.arange(1e2, 2e4)
nu_arr  = (const.c/(lam_arr * u.AA)).to(1./u.s).value

fig = plt.figure(figsize=(15,8))
ax1  = fig.add_subplot(121, axisbg='white')
ax2  = fig.add_subplot(122, axisbg='white')
ax1.set_xlabel(r'$\lambda$ (Ansgstrom)', fontsize=25)
ax1.set_ylabel(r'$B_{\lambda}(T)$', fontsize=25)
ax2.set_xlabel(r'$\nu$ (\textrm{s}^{-1})', fontsize=25)
ax2.set_ylabel(r'$B_{\nu}(T)$', fontsize=25)
ax2.set_xscale('log')

temp_arr = [6e3, 8e3, 1e4, 1.2e4]
for temp in temp_arr:
    ax1.plot(lam_arr, Planck_func(temp, lam_arr=lam_arr, opt='lam'), label='T = {0} K'.format(int(temp)))
    ax2.plot(nu_arr , Planck_func(temp, lam_arr=lam_arr, opt='nu' ), label='T = {0} K'.format(int(temp)))
    ax1.legend(loc=1, prop={'size':20})

"""
Explanation: Let's plot the Planck function for two bodies with temperatures $T_1 = 8000\ K$ and $T_2 = 6000\ K$
End of explanation
"""


!head ./data/sources.dat

from astropy.io import ascii
sources_tb = ascii.read('./data/sources.dat')

print( sources_tb )

"""
Explanation: AstroPy Tables
Read files
You can use Astropy to read tables from data files. We'll use it to read the sources.dat file, which contains columns and rows of data
End of explanation
"""


from astropy.table import Table, Column, MaskedColumn
x = np.random.uniform(low=10, high=20, size=(1000,))
y = np.random.uniform(low=100, high=50, size=(x.size,))
z = np.random.uniform(low=30, high=50, size=(x.size,))
data = Table([x, y], names=['x', 'y'])
print(data)

ascii.write(data, './data/astropy_data.tb', overwrite=True)

"""
Explanation: Write to files
You can also write directoy to a file using the data in the AstroPy table.
Let's create a new AstroPy Table:
End of explanation
"""


!head ./data/astropy_data.tb

"""
Explanation: Let's see what's in the astropy_data.tb file
End of explanation
"""


ascii.write(data, './data/astropy_data_2.tb', delimiter=',', overwrite=True)

!head ./data/astropy_data_2.tb

"""
Explanation: You can also specify the delimiter of the file. For example, we can separate it with a comma.
End of explanation
"""


df = data.to_pandas()
df.head()

"""
Explanation: AstroPy Tables to other Formats
The AstroPy tables can also be converted to multiple formats
to Pandas DataFrames
A nice feature of AstroPy Tables is that you can export your data into different formats. 
For example, you can export it as a Pandas Dataframe. 
See here for more info on how to use pandas with Astropy: http://docs.astropy.org/en/stable/table/pandas.html
End of explanation
"""


data

"""
Explanation: And to compare, let's see the AstroPy Tables format
End of explanation
"""


import sys
ascii.write(data[0:10], sys.stdout, format='latex')

"""
Explanation: to LaTeX tables
A nice thing about AstroPy is that you can convert your data into LaTeX tables. This is easily done with writing it to a file. You can then copy it and use it on your next publication
End of explanation
"""


ascii.write(data, './data/astropy_data_latex.tex', format='latex')

# I'm only showing the first 10 lines
!head ./data/astropy_data_latex.tex

"""
Explanation: To save it as a file, you can do this:
End of explanation
"""


ascii.write(data, './data/astropy_data_csv.csv', format='csv', fast_writer=False) 

!head ./data/astropy_data_csv.csv

"""
Explanation: to CSV files
End of explanation
"""


t = Table(masked=True)
t['x'] = MaskedColumn([1.0, 2.0], unit='m', dtype='float32')
t['x'][1] = np.ma.masked
t['y'] = MaskedColumn([False, True], dtype='bool')

t

"""
Explanation: Other formats
AstroPy tables come with a great support for many different types of files.
This is a list of the supported files that you can import/export AstroPy tables.

Data tables and Column types
You can also use AstroPy tables to preserve the metadata of a column. For example, you can keep the units of each column, so that you use the data later on, and still be able to use unit conversions, etc. for this.
End of explanation
"""


from astropy.extern.six.moves import StringIO
fh = StringIO()
t.write(fh, format='ascii.ecsv')  
table_string = fh.getvalue()      
print(table_string) 

Table.read(table_string, format='ascii')  

"""
Explanation: Now we can save it into a ecsv file. This type of file will preserve the type of units, and more, for each of the columns
End of explanation
"""


t.write('./data/astropy_data_ecsv.ecsv', format='ascii.ecsv', overwrite=True)

"""
Explanation: Or you can dump it into a file
End of explanation
"""


data_ecsv = ascii.read('./data/astropy_data_ecsv.ecsv', format='ecsv')
data_ecsv

data_ecsv['x']

"""
Explanation: And you can now read it in
End of explanation
"""


