In [ ]:
from os import path
from astropy.io import fits
import astropy.units as u
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('apw-notebook')
%matplotlib inline

from comoving_rv.db import Session, Base, db_connect
from comoving_rv.db.model import (Run, Observation, TGASSource, SimbadInfo,
                                  SpectralLineInfo, SpectralLineMeasurement)
from comoving_rv.plot import colors

In [ ]:
base_path = '/Volumes/ProjectData/gaia-comoving-followup/'
db_path = path.join(base_path, 'db.sqlite')
engine = db_connect(db_path)
session = Session()

In [ ]:
observations = session.query(Observation).join(Run).filter(Run.name == 'mdm-spring-2017')\
                      .group_by(Observation.object).all()

In [ ]:
n_spectra = 10

np.random.seed(48151623)
idx = np.random.choice(len(observations), size=n_spectra, replace=False)

In [ ]:
from scipy.ndimage import gaussian_filter1d

In [ ]:
fig,ax = plt.subplots(1, 1, figsize=(5, 8))

wvln_range = (5500, 6800)

for j,i in enumerate(idx):
    obs = observations[i]
    data = fits.getdata(obs.path_1d(base_path))
    data = data[np.isfinite(data['wavelength'])]
    data = data[np.argsort(data['wavelength'])]
    
    i1 = np.abs(data['wavelength']-wvln_range[0]).argmin()
    i2 = np.abs(data['wavelength']-wvln_range[1]).argmin()
    data = data[i1:i2]
    
    wvln = data['wavelength']
    flux = data['source_flux'] / data['source_flux'][0]
    ax.plot(wvln, gaussian_filter1d(flux + j/2.9, 1.), marker='', 
            color='k', lw=1)
    # ax.plot(wvln, flux + j/2.9, marker='', 
    #         drawstyle='steps-mid', color='k', lw=1)
    
    # err = 1 / np.sqrt(ivar) / # TODO: normalize by same factor
    # ax.errorbar(wvln, flux, yerr=err,
    #             ecolor='#777777', marker='', linestyle='')

ax.set_xticks(np.arange(5500, 6750+1, 250))
ax.set_yticks([])
ax.set_xlim(wvln_range)
ax.set_ylim(0, 4.3)

# line markers
line_style = dict(linewidth=2, marker='', color=colors['line_marker'], alpha=0.75)
ax.plot([6562.8]*2, [3.5, 3.6], **line_style)
ax.plot([5893]*2, [3.85, 3.95], **line_style)

text_style = dict(ha='center', va='bottom', fontsize=18, color=colors['line_marker'])
ax.text(6562.8, 3.65, r'${\rm H}\alpha$', **text_style)
ax.text(5893, 4.0, r'${\rm Na}\,{\rm D}$', **text_style)

ax.set_xlabel(r'wavelength [${\rm \AA}$]')
ax.set_ylabel('normalized flux + offset')

fig.savefig('sample_spectra.pdf')

In [ ]: