In [35]:
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from scipy.optimize import fmin_cg, minimize
import h5py
c = 2.99792458e8 # m/s
In [36]:
def doppler(v):
frac = (1. - v/c) / (1. + v/c)
return np.sqrt(frac)
def gamma(v):
return 1. / np.sqrt(1. - (v/c) ** 2)
def dlndopplerdv(v):
dv = doppler(v)
return -1. * gamma(v) / (2. * c) * (1. / dv + dv)
def state(v, xs, xps):
'''
outputs: (M, Mp, v, xs, ms, mps, ehs, bees, seas)
M and Mp are the lengths of data and model wavelength grids
v is the RV
xs is the wavelength values of the data grid
ms is the data index m at which there is an interpolated model value
mps is the model index m' from which we interpolate to ms
ehs, bees, and seas go into the coefficients of interpolation
'''
# every input must be 1-d
M = len(xs)
Mp = len(xps)
xps_shifted = xps + np.log(doppler(v))
ms = np.arange(M)
mps = np.searchsorted(xps_shifted, xs, side='left')
good = (mps > 0) * (mps < Mp)
ms = ms[good]
mps = mps[good]
ehs = xps_shifted[mps] - xs[ms]
bees = xs[ms] - xps_shifted[mps - 1]
seas = ehs + bees
return (M, Mp, v, xs, ms, mps, ehs, bees, seas)
def Pdot(state, vec):
# takes state and model flux vector, returns (shifted) model interpolated into data space
# unpack state
M, Mp, v, xs, ms, mps, ehs, bees, seas = state
# do shit
result = np.zeros(M)
result[ms] = vec[mps - 1] * ehs / seas + vec[mps] * bees / seas
return result
def dotP(state, vec):
# takes state and data flux vector, returns data interpolated into (shifted) model space
# unpack state
M, Mp, v, xs, ms, mps, ehs, bees, seas = state
# do shit
result = np.zeros(Mp)
result[mps - 1] += vec[ms] * ehs / seas
result[mps] += vec[ms] * bees / seas
return result
def dotdPdv(state, vec):
# unpack state
M, Mp, v, xs, ms, mps, ehs, bees, seas = state
# do shit
result = np.zeros(Mp)
foos = vec[ms] / seas * dlndopplerdv(v) # * xs[ms] ??
result[mps - 1] += foos
result[mps] -= foos
return result
def dPdotdv(state, vec):
# unpack state
M, Mp, v, xs, ms, mps, ehs, bees, seas = state
# do shit
result = np.zeros(M)
result[ms] = (vec[mps - 1] - vec[mps]) * dlndopplerdv(v) / seas
return result
In [37]:
with h5py.File('../data/hip54287.hdf5', 'r') as f:
N = 75
inds = (f['xs'][:] > 5900.) & (f['xs'][:] < 6000.)
data = np.copy(f['data'])[:N,inds]
data_xs = np.log(np.copy(f['xs'][inds]))
ivars = np.copy(f['ivars'])[:N,inds]
true_rvs = np.copy(f['true_rvs'])[:N]
bervs = np.copy(f['berv'])[:N] * -1.e3
for i in xrange(len(data)):
data[i] /= np.median(data[i])
data = np.log(data)
In [47]:
plt.plot(data_xs, data[0,:])
plt.plot(data_xs, data[10,:])
plt.ylim([-0.05, 0.05])
In [39]:
def make_template(all_data, rvs, xs, dx):
"""
`all_data`: `[N, M]` array of pixels
`rvs`: `[N]` array of RVs
`xs`: `[M]` array of wavelength values
`dx`: linear spacing desired for template wavelength grid (A)
"""
(N,M) = np.shape(all_data)
all_xs = np.empty_like(all_data)
for i in range(N):
all_xs[i,:] = xs - np.log(doppler(rvs[i])) # shift to rest frame
all_data, all_xs = np.ravel(all_data), np.ravel(all_xs)
tiny = 10.
template_xs = np.arange(min(all_xs)-tiny*dx, max(all_xs)+tiny*dx, dx)
template_ys = np.nan + np.zeros_like(template_xs)
for i,t in enumerate(template_xs):
ind = (all_xs >= t-dx/2.) & (all_xs < t+dx/2.)
if np.sum(ind) > 0:
template_ys[i] = np.nanmedian(all_data[ind])
ind_nan = np.isnan(template_ys)
template_ys.flat[ind_nan] = np.interp(template_xs[ind_nan], template_xs[~ind_nan], template_ys[~ind_nan]) #np.interp(template_xs[ind_nan], template_xs[~ind_nan], template_ys[~ind_nan])
return template_xs, template_ys
def subtract_template(data_xs, data, model_xs_t, model_ys_t, rvs_t):
(N,M) = np.shape(data)
data_sub = np.copy(data)
for n,v in enumerate(rvs_t):
s = state(v, data_xs, model_xs_t)
model_ys_t_shifted = Pdot(s, model_ys_t)
data_sub[n,:] -= np.ravel(model_ys_t_shifted)
if n == 0:
plt.plot(data_xs, data[n,:], color='k')
plt.plot(data_xs, data_sub[n,:], color='blue')
plt.plot(data_xs, np.ravel(model_ys_t_shifted), color='red')
return data_sub
In [40]:
x0_star = -np.copy(bervs)
x0_star -= np.mean(x0_star)
x0_t = np.zeros(N)
model_xs_star, model_ys_star = make_template(data, x0_star, data_xs, np.log(6000.01) - np.log(6000.))
model_xs_t, model_ys_t = make_template(data, x0_t, data_xs, np.log(6000.01) - np.log(6000.))
In [41]:
plt.plot(model_xs_star, model_ys_star, color='k', alpha=0.5)
plt.plot(model_xs_t, model_ys_t, color='red', alpha=0.5)
plt.plot(data_xs, data[0,:], color='blue', alpha=0.5)
plt.ylim([-0.1, 0.1])
Out[41]:
In [42]:
def rv_lnprior(rvs):
return -0.5 * np.mean(rvs)**2/1.**2
def drv_lnprior_dv(rvs):
return np.zeros_like(rvs) - np.mean(rvs)/1.**2/len(rvs)
def lnlike_star(rvs_star, rvs_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t, model_ys_t):
try:
N = len(rvs_star)
except:
N = 1
lnlike = 0.
dlnlike_dv = np.zeros(N)
for n in range(N):
state_star = state(rvs_star[n], data_xs, model_xs_star)
pd_star = Pdot(state_star, model_ys_star)
state_t = state(rvs_t[n], data_xs, model_xs_t)
pd_t = Pdot(state_t, model_ys_t)
pd = pd_star + pd_t
lnlike += -0.5 * np.sum((data[n,:] - pd)**2 * ivars[n,:])
dpd_dv = dPdotdv(state_star, model_ys_star)
dlnlike_dv[n] = np.sum((data[n,:] - pd) * ivars[n,:] * dpd_dv)
lnpost = lnlike + rv_lnprior(rvs_star)
dlnpost_dv = dlnlike_dv + drv_lnprior_dv(rvs_star)
return -1 * lnpost, -1 * dlnpost_dv
def lnlike_t(rvs_t, rvs_star, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t, model_ys_t):
try:
N = len(rvs_t)
except:
N = 1
lnlike = 0.
dlnlike_dv = np.zeros(N)
for n in range(N):
state_star = state(rvs_star[n], data_xs, model_xs_star)
pd_star = Pdot(state_star, model_ys_star)
state_t = state(rvs_t[n], data_xs, model_xs_t)
pd_t = Pdot(state_t, model_ys_t)
pd = pd_star + pd_t
lnlike += -0.5 * np.sum((data[n,:] - pd)**2 * ivars[n,:])
dpd_dv = dPdotdv(state_t, model_ys_t)
dlnlike_dv[n] = np.sum((data[n,:] - pd) * ivars[n,:] * dpd_dv)
lnpost = lnlike + rv_lnprior(rvs_t)
dlnpost_dv = dlnlike_dv + drv_lnprior_dv(rvs_t)
return -1 * lnpost, -1 * dlnpost_dv
In [43]:
def model_ys_lnprior(w):
return -0.5 * np.sum(w**2)/100.**2
def dmodel_ys_lnprior_dw(w):
return -1.*w / 100.**2
def dlnlike_star_dw_star(model_ys_star, rvs_star, rvs_t, data_xs, data, ivars, model_xs_star, model_xs_t, model_ys_t):
try:
N = len(rvs_star)
except:
N = 1
lnlike = 0.
Mp = len(model_xs_star)
dlnlike_dw = np.zeros(Mp)
for n in range(N):
state_star = state(rvs_star[n], data_xs, model_xs_star)
pd_star = Pdot(state_star, model_ys_star)
state_t = state(rvs_t[n], data_xs, model_xs_t)
pd_t = Pdot(state_t, model_ys_t)
pd = pd_star + pd_t
dp_star = dotP(state_star, (data[n,:] - pd)*ivars[n,:])
lnlike += -0.5 * np.sum((data[n,:] - pd)**2 * ivars[n,:])
dlnlike_dw += dp_star
lnprior = model_ys_lnprior(model_ys_star)
dlnprior = dmodel_ys_lnprior_dw(model_ys_star)
return -lnlike - lnprior, -dlnlike_dw - dlnprior
def dlnlike_t_dw_t(model_ys_t, rvs_star, rvs_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t):
try:
N = len(rvs_t)
except:
N = 1
lnlike = 0.
Mp = len(model_xs_t)
dlnlike_dw = np.zeros(Mp)
for n in range(N):
state_star = state(rvs_star[n], data_xs, model_xs_star)
pd_star = Pdot(state_star, model_ys_star)
state_t = state(rvs_t[n], data_xs, model_xs_t)
pd_t = Pdot(state_t, model_ys_t)
pd = pd_star + pd_t
dp_t = dotP(state_t, (data[n,:] - pd)*ivars[n,:])
lnlike += -0.5 * np.sum((data[n,:] - pd)**2 * ivars[n,:])
dlnlike_dw += dp_t
lnprior = model_ys_lnprior(model_ys_t)
dlnprior = dmodel_ys_lnprior_dw(model_ys_t)
return -lnlike - lnprior, -dlnlike_dw - dlnprior
In [44]:
def improve_telluric_model(model_ys_t, x0_star, x0_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t):
w = np.copy(model_ys_t)
for i in range(50):
lnlike, dlnlike_dw = dlnlike_t_dw_t(w, x0_star, x0_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t)
w -= 5e-7 * dlnlike_dw
return w
def improve_star_model(model_ys_star, x0_star, x0_t, data_xs, data, ivars, model_xs_star, model_xs_t, model_ys_t):
w = np.copy(model_ys_star + np.random.normal(0, 0.01, len(model_ys_star)))
for i in range(50):
lnlike, dlnlike_dw = dlnlike_star_dw_star(w, x0_star, x0_t, data_xs, data, ivars, model_xs_star, model_xs_t, model_ys_t)
w -= 5e-7 * dlnlike_dw
return w
In [45]:
for n in range(5):
print "Fitting stellar RVs..."
soln_star = minimize(lnlike_star, x0_star, args=(x0_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t, model_ys_t),
method='BFGS', jac=True, options={'disp':True, 'gtol':1.e-2, 'eps':1.5e-5})['x']
model_ys_t = improve_telluric_model(model_ys_t, soln_star, x0_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t)
model_ys_star = improve_star_model(model_ys_star, soln_star, x0_t, data_xs, data, ivars, model_xs_star, model_xs_t, model_ys_t)
print "Star model improved. Fitting telluric RVs..."
soln_t = minimize(lnlike_t, x0_t, args=(soln_star, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t, model_ys_t),
method='BFGS', jac=True, options={'disp':True, 'gtol':1.e-2, 'eps':1.5e-5})['x']
model_ys_t = improve_telluric_model(model_ys_t, soln_star, soln_t, data_xs, data, ivars, model_xs_star, model_ys_star, model_xs_t)
model_ys_star = improve_star_model(model_ys_star, soln_star, soln_t, data_xs, data, ivars, model_xs_star, model_xs_t, model_ys_t)
x0_star = soln_star
x0_t = soln_t
print "iter {0}: star std = {1:.2f}, telluric std = {2:.2f}".format(n, np.std(soln_star + true_rvs), np.std(soln_t))
plt.scatter(np.arange(N), soln_star + true_rvs, color='k')
plt.scatter(np.arange(N), soln_t, color='red')
plt.show()
In [ ]:
plt.scatter(bervs, soln_star+true_rvs)
In [ ]:
plt.plot(data_xs, data[0,:], color='k')
plt.plot(data_xs, data_star[0,:] + data_t[0,:], color='red')
In [ ]:
plt.plot(data_xs, data[0,:] - data_star[0,:] - data_t[0,:])
In [ ]:
plt.hist(data[0,:] - data_star[0,:] - data_t[0,:])
In [46]:
plt.plot(model_xs_star, model_ys_star, 'k', alpha=0.5)
plt.plot(model_xs_t, model_ys_t, 'r', alpha=0.5)
plt.ylim([-0.05, 0.05])
Out[46]:
In [48]:
print len(data_xs)
print len(model_xs_star)
print len(model_xs_t)
In [ ]: