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 [ ]:

    

files = ["../data/apVisit-r5-6994-56770-261.fits", "../data/apVisit-r5-6994-56794-177.fits"]


    

In [ ]:

    

def load_file(filename, chip=None):
    hdulist1 = fits.open(filename)
    wvln = hdulist1[4].data
    flux = hdulist1[1].data
    flux_err = hdulist1[2].data
    
    if chip is None:
        return {'wvln': wvln, 'flux': flux, 'flux_err': flux_err}
    else:
        return {'wvln': wvln[chip], 'flux': flux[chip], 'flux_err': flux_err[chip]}


    

In [ ]:

    

d['wvln'][2].min(), d['wvln'][0].max()


    

In [ ]:

    

fig,ax = pl.subplots(1,1,figsize=(12,6))
for fn in files:
    d = load_file(fn)
#     for i in range(3):
    i = 2
    ax.plot(d['wvln'][i], d['flux'][i], drawstyle='steps', marker=None)


    

In [ ]:

    

step = ((1*u.km/u.s) / c).decompose()
step


    

In [ ]:

    

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


    

In [ ]:

    

M = int((np.log(all_spectra[0]['wvln'].max()) - np.log(all_spectra[0]['wvln'].min())) / step) + 1
grid = np.exp(np.log(all_spectra[0]['wvln'].min()) + np.arange(M) * step)
grid.size


    

In [ ]:

    

for spec in all_spectra:
    interp_f = interp1d(spec['wvln'], spec['flux'], kind='cubic', bounds_error=False)
    spec['interp_flux'] = interp_f(grid)


    

In [ ]:

    

for spec in all_spectra:
    print(np.where(np.isnan(spec['interp_flux'])))


    

In [ ]:

    

fig,axes = pl.subplots(2,1,figsize=(12,12),sharex=True)

axes[0].plot(all_spectra[0]['wvln'], all_spectra[0]['flux'], linestyle='none', marker='o')
axes[0].plot(grid, all_spectra[0]['interp_flux'], linestyle='none', marker='.')

axes[1].plot(all_spectra[1]['wvln'], all_spectra[1]['flux'], linestyle='none', marker='o')
axes[1].plot(grid, all_spectra[1]['interp_flux'], linestyle='none', marker='.')

axes[0].set_xlim(15500, 15504)
axes[0].set_ylim(13000, 16000)


    

In [ ]:

    

def weight_function(size, edge_buffer=64):
    w = np.ones(size)
    
    w[:edge_buffer] = 0.
    w[edge_buffer:2*edge_buffer] = np.linspace(0.,1,edge_buffer)
    w[-edge_buffer:] = 0.
    w[-2*edge_buffer:-edge_buffer] = np.linspace(1,0.,edge_buffer)
    
    return w

edge_buffer = 64
weight = weight_function(grid.size, edge_buffer)


    

In [ ]:

    

for spec in all_spectra:
    idx = np.where(np.isnan(spec['interp_flux']))
    assert np.all(weight[idx] == 0)
    spec['interp_flux'][idx] = 0.


    

In [ ]:

    

fig,axes = pl.subplots(2,1,figsize=(12,12),sharex=True)

axes[0].plot(all_spectra[0]['interp_flux']*weight_function(grid.size), linestyle='none', marker='o')
axes[1].plot(all_spectra[1]['interp_flux']*weight_function(grid.size), linestyle='none', marker='o')

# axes[0].set_xlim(0, 150)
axes[0].set_xlim(all_spectra[0]['interp_flux'].size-150, all_spectra[0]['interp_flux'].size)


    

In [ ]:

    

for spec in all_spectra:
    spec['mean_flux'] = np.sum(weight * spec['interp_flux']) / np.sum(weight)
    print(spec['mean_flux'])


    

In [ ]:

    

def xcor_spectra(spec1, spec2, weight, d_index):
    assert d_index > 0
    ws1 = weight * (spec1['interp_flux'] - spec1['mean_flux'])
    ws2 = weight * (spec2['interp_flux'] - spec2['mean_flux'])
    denom = np.sum(weight[:-d_index] * weight[d_index:])
    return np.sum(ws1[:-d_index] * ws2[d_index:]) / denom


    

In [ ]:

    

shifts = np.arange(1, 64+1, 1)


    

In [ ]:

    

# -- spec1, spec1
# -- spec2, spec2
for spec in all_spectra:
    acors = np.zeros_like(shifts)
    for j,i in enumerate(shifts):
        acors[j] = xcor_spectra(spec, spec, weight, i)

    pl.plot(shifts, acors)
    
# -- spec1, spec2
xcors = np.zeros_like(shifts)
for j,i in enumerate(shifts):
    xcors[j] = xcor_spectra(all_spectra[0], all_spectra[1], weight, i)

pl.plot(shifts, xcors)

# -- spec2, spec1
xcors = np.zeros_like(shifts)
for j,i in enumerate(shifts):
    xcors[j] = xcor_spectra(all_spectra[1], all_spectra[0], weight, i)

pl.plot(shifts, xcors)


    

In [ ]: