In [ ]:

    

# Third-party
from astropy.io import ascii, fits
import astropy.coordinates as coord
import astropy.units as u
from astropy.constants import c
import matplotlib as mpl
import matplotlib.pyplot as pl
import numpy as np
from scipy.interpolate import interp1d

pl.style.use('apw-notebook')
%matplotlib inline
# pl.style.use('classic')
# %matplotlib notebook


    

In [ ]:

    

data_files = ["../data/apVisit-r5-6994-56770-261.fits", "../data/apVisit-r5-6994-56794-177.fits"]
model_file = "../data/apStar-r5-2M00004994+1621552.fits"


    

In [ ]:

    

min_wvln = 15329
max_wvln = 15359


    

In [ ]:

    

def load_file(filename, chip):
    hdulist1 = fits.open(filename)
    wvln = hdulist1[4].data[chip]
    ix = (wvln >= min_wvln) & (wvln <= max_wvln)
    
    wvln = wvln[ix]
    flux = hdulist1[1].data[chip,ix]
    flux_err = hdulist1[2].data[chip,ix]
    
    return {'wvln': wvln, 'flux': flux, 'flux_err': flux_err}
    
def load_model_file(filename):
    hdulist1 = fits.open(filename)
    flux = hdulist1[1].data[0]
    flux_err = hdulist1[2].data[0]
    wvln = 10**(hdulist1[0].header['CRVAL1'] + np.arange(flux.size) * hdulist1[0].header['CDELT1'])
    
#     ix = (wvln >= min_wvln) & (wvln <= max_wvln)
    ix = (wvln < 15750) & (wvln > 15150) # HACK: magic numbers
    return {'wvln': wvln[ix], 'flux': flux[ix], 'flux_err': flux_err[ix]}


    

In [ ]:

    

d = load_file(fn, chip=2)
d['wvln'].shape


    

In [ ]:

    

chip = 2

fig,ax = pl.subplots(1,1,figsize=(12,6))

for fn in data_files:
    d = load_file(fn, chip=chip)
    ax.plot(d['wvln'], d['flux'], drawstyle='steps', marker=None)
    
ref_spec = load_model_file(model_file)
ax.plot(ref_spec['wvln'], 3.2*ref_spec['flux'], drawstyle='steps', marker=None, lw=2.) # HACK: scale up

# _d = 175
# ax.set_xlim(15150.+_d, 15175.+_d)
# ax.set_ylim(10000, 20000)


    

In [ ]:

    

all_spectra = [load_file(f, chip=2) for f in files]


    

In [ ]:

    

ref_spec['interp'] = interp1d(ref_spec['wvln'], ref_spec['flux'], kind='cubic', bounds_error=False)


    

In [ ]:

    

def get_design_matrix(data, ref_spec, v1, v2):
    """
    Note: Positive velocity is a redshift.
    """
    X = np.ones((3, data['wvln'].shape[0]))
    X[1] = ref_spec['interp'](data['wvln'] * (1 + v1/c)) # this is only good to first order in (v/c)
    X[2] = ref_spec['interp'](data['wvln'] * (1 + v2/c))
    return X


    

In [ ]:

    

def get_optimal_chisq(data, ref_spec, v1, v2):
    X = get_design_matrix(data, ref_spec, v1, v2)
    return np.linalg.solve( X.dot(X.T), X.dot(data['flux']) )


    

In [ ]:

    

spec_i = 1
v1 = 35 * u.km/u.s
v2 = -5 * u.km/u.s
X = get_design_matrix(all_spectra[spec_i], ref_spec, v1, v2)
opt_pars = get_optimal_chisq(all_spectra[spec_i], ref_spec, v1, v2)
opt_pars


    

In [ ]:

    

def make_synthetic_spectrum(X, pars):
    return X.T.dot(pars)

def compute_chisq(data, X, opt_pars):
    synth_spec = make_synthetic_spectrum(X, opt_pars)
    return -np.sum((synth_spec - data['flux'])**2)


    

In [ ]:

    

# opt_pars = np.array([1.1E+4, 0.5,  0.5])
synth_spec = make_synthetic_spectrum(X, opt_pars)


    

In [ ]:

    

pl.plot(all_spectra[spec_i]['wvln'], all_spectra[spec_i]['flux'], marker=None, drawstyle='steps')
pl.plot(all_spectra[spec_i]['wvln'], synth_spec, marker=None, drawstyle='steps')


    

In [ ]:

    

_v1_grid = np.linspace(25, 45, 32)
_v2_grid = np.linspace(-15, 5, 32)
shp = (_v1_grid.size, _v2_grid.size)
v_grid = np.vstack(map(np.ravel, np.meshgrid(_v1_grid, _v2_grid))).T
v_grid.shape


    

In [ ]:

    

chisq = np.zeros(v_grid.shape[0])
for i in range(v_grid.shape[0]):
    v1,v2 = v_grid[i]
    opt_pars = get_optimal_chisq(all_spectra[spec_i], ref_spec, 
                                 v1*u.km/u.s, v2*u.km/u.s)
    chisq[i] = compute_chisq(all_spectra[spec_i], X, opt_pars)


    

In [ ]:

    

fig,ax = pl.subplots(1,1,figsize=(9,8))

cb = ax.pcolormesh(v_grid[:,0].reshape(shp), v_grid[:,1].reshape(shp), 
                   chisq.reshape(shp), cmap='magma')

fig.colorbar(cb)


    

In [ ]:

    

fig,ax = pl.subplots(1,1,figsize=(9,8))

cb = ax.pcolormesh(v_grid[:,0].reshape(shp), v_grid[:,1].reshape(shp), 
                   np.exp(chisq-chisq.max()).reshape(shp), cmap='magma')

fig.colorbar(cb)


    

In [ ]:

    

fig,ax = pl.subplots(1,1,figsize=(9,8))

cb = ax.pcolormesh(v_grid[:,0].reshape(shp), v_grid[:,1].reshape(shp), 
                   chisq.reshape(shp), cmap='magma')

fig.colorbar(cb)


    

In [ ]:

    

fig,ax = pl.subplots(1,1,figsize=(9,8))

cb = ax.pcolormesh(v_grid[:,0].reshape(shp), v_grid[:,1].reshape(shp), 
                   np.exp(chisq-chisq.max()).reshape(shp), cmap='magma')

fig.colorbar(cb)


    

In [ ]:

    

from scipy.optimize import leastsq


    

In [ ]:

    

def errfunc(pars, data_spec, ref_spec):
    v1,v2,a,b,c = pars
    X = get_design_matrix(data_spec, ref_spec, v1*u.km/u.s, v2*u.km/u.s)
    synth_spec = make_synthetic_spectrum(X, [a,b,c])
    return (synth_spec - data_spec['flux'])


    

In [ ]:

    

levmar_opt_pars,ier = leastsq(errfunc, x0=[35,-5]+opt_pars.tolist(), args=(all_spectra[0], ref_spec))


    

In [ ]:

    

levmar_opt_pars


    

In [ ]:

    

data_spec = all_spectra[0]
X = get_design_matrix(data_spec, ref_spec, levmar_opt_pars[0]*u.km/u.s, levmar_opt_pars[1]*u.km/u.s)
synth_spec = make_synthetic_spectrum(X, levmar_opt_pars[2:])
pl.plot(data_spec['wvln'], data_spec['flux'], marker=None, drawstyle='steps')
pl.plot(data_spec['wvln'], synth_spec, marker=None, drawstyle='steps')


    

In [ ]: