# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Environment (ebak)
#     language: ''
#     name: ebak
# ---

# +
import astropy.units as u
import matplotlib.pyplot as pl
pl.style.use('apw-notebook')
# %matplotlib inline
import numpy as np
from scipy.optimize import leastsq
from scipy.interpolate import interp1d
import matplotlib.gridspec as gridspec

import ebak.inference as einf
import ebak.io as eio
# -

data_file = "../data/apVisit-r5-6994-56770-261.fits"
ref_file = "../data/apStar-r5-2M00004994+1621552.fits"

chip_idx = 2
data_spec = eio.load_apVisit(data_file, chip=chip_idx)

# +
ref_spec = eio.load_apStar(ref_file)

# only need the reference spectrum over the range of wavelengths considered (here, chip=2)
ix = (ref_spec['wvln'] > data_spec['wvln'].min()) & (ref_spec['wvln'] < data_spec['wvln'].max())
for key in ref_spec.keys():
    ref_spec[key] = ref_spec[key][ix][32:-32]

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

# trim data spectrum to have ~same wavelength coverage
for key in data_spec.keys():
    ix = (data_spec['wvln'] > ref_spec['wvln'].min()) & (data_spec['wvln'] < ref_spec['wvln'].max())
    data_spec[key] = data_spec[key][ix]
# -

pl.plot(ref_spec['wvln'], ref_spec['flux'], drawstyle='steps', marker=None)
pl.plot(data_spec['wvln'], data_spec['flux']/3, drawstyle='steps', marker=None)

# Velocities are in km/s -- ln likelihood parameter vector just takes [v1, v2] and the two spectra

einf.ln_likelihood([35.,-5], data_spec, ref_spec)



StarcoderdataJupyterScriptsDedupFiltered