import numpy as np
from astropy.io import ascii
import matplotlib.pyplot as plt
#Gamma Ray Burst Data
GRB = ascii.read(u'GRB_afterglow.dat',data_start=2,header_start=1)
GRB.colnames
#t is time since GRB trigger event (i.e. the GRB was detected)
#f is the flux measured between 2-10 keV wavelength regions
#(Where must this data be taken from?)
#s is the error in the flux measurement
fig,ax = plt.subplots()
ax.scatter(np.log10(GRB['t']),np.log10(GRB['f'])-11.,color='black')
ax.set_ylabel(r'$\log$(Flux) (ergs cm$^{-2}$ s$^{-1}$)',fontsize=24)
ax.set_xlabel(r'$\log$(Time since trigger) (s)',fontsize=24)
#Fit with covariance matrix
p,cov = np.polyfit(np.log10(GRB['t']),np.log10(GRB['f'])-11.,1,cov=True)
#Find the rms
n = np.float(GRB['f'].size)
rms = np.sqrt(np.sum((p[0]*np.log10(GRB['t'])+p[1]-np.log10(GRB['f'])+11.)**2.)/n)
xs = np.arange(1.5,6.5,.5)
ax.plot(xs,p[0]*xs+p[1],'--',color='red',linewidth=3)
print('Fit parameters')
print(u'm '+str(p[0])+'+/-'+str(np.sqrt(cov[0,0])))
print(u'b '+str(p[1])+'+/-'+str(np.sqrt(cov[1,1])))
print(u'RMS = '+str(rms))
#Is a linear fit the best?
#or putting things back we can show measurements errors
fig,ax = plt.subplots()
ax.errorbar(np.log10(GRB['t']),GRB['f'],yerr=GRB['s'],fmt='o',color='black')
ax.set_ylabel(r'Flux (10$^{-11}$ ergs cm$^{-2}$ s$^{-1}$)',fontsize=24)
ax.set_xlabel(r'$\log$(Time since trigger) (s)',fontsize=24)
xs = np.arange(2.0,6.5,.01)
ax.plot(xs,10.**(p[0]*xs+p[1]+11.),'--',color='red',linewidth=3)
ax.set_ylim([0,140.])