def read_shadow(dp, shadowname, rv, ccf):
"""
This reads shadow parameters from a pickle file generated by the fit_doppler_
model class below, and evaluates it on the user-supplied rv and ccf (the latter of which is
just for shape-purposes and the mask defined during the creation of the shadow model.
"""
import pickle
import os.path
import sys
from pathlib import Path
import tayph.util as ut
from tayph.vartests import typetest, dimtest
import tayph.system_parameters as sp
import numpy as np
typetest(shadowname, str, 'shadowname in read_shadow().')
typetest(rv, np.ndarray, 'rv in read_shadow()')
typetest(ccf, np.ndarray, 'ccf in read_shadow()')
dp = ut.check_path(dp, exists=True)
inpath = dp / (shadowname + '.pkl')
ut.check_path(inpath, exists=True)
RVp = sp.RV(dp)
pickle_in = open(inpath, "rb")
d = pickle.load(pickle_in)
params = d["fit_params"]
T = d["Transit"]
p = d["Phases"]
aRstar = d["aRstar"]
vsys = d["vsys"]
i = d["inclination"]
n_c = d["n_components"]
maskHW = d["maskHW"]
offset = d["offset"] #Offset of second component.
S = d["S"]
D = d["D"]
RVp = sp.RV(dp) + vsys
mask = mask_ccf_near_RV(rv, ccf, RVp, maskHW)
return (evaluate_shadow(params,
rv,
ccf,
T,
p,
aRstar,
vsys,
i,
n_c,
offset,
S,
D,
leastsq=False), mask)
def __init__(self, fig, ax, rv, ccf, primer, dp, outname):
"""We initialize with a figure object, three axis objects (in a list)
an rv axis, the CCF, the user-generated prior on the doppler shadow (the two)
points, and the pointer to the data in order to load physics."""
import tayph.system_parameters as sp
import numpy as np
import pdb
import scipy.interpolate as interpol
import tayph.functions as fun
import tayph.util as ut
import sys
#Upon initialization, we pass the keywords onto self.
nexp = np.shape(ccf)[0]
nrv = np.shape(ccf)[1]
self.rv = rv
self.ccf = ccf
self.p = sp.phase(dp)
if len(self.p) != nexp:
print('ERROR IN FIT_DOPPLER_MODEL __INIT__:')
print('The height of the CCF does not match nexp.')
sys.exit()
transit = sp.transit(dp) - 1.0
self.T = abs(transit / max(abs(transit)))
self.ax = ax
self.aRstar = sp.paramget('aRstar', dp)
self.vsys = sp.paramget('vsys', dp)
self.RVp = sp.RV(dp) + self.vsys
self.inclination = sp.paramget('inclination', dp)
self.n_components = 1
self.maskHW = 10.0 #Default masking half-width
self.offset = 0.0
self.D = 0 #Initial degree of polynomial fitting.
self.S = 0
self.dp = ut.check_path(dp, exists=True)
self.outpath = self.dp / (outname + '.pkl')
#Translate the pivot to RV-phase points.
#Need to interpolate on phase only, because the primer was already defined in RV space.
p_i = interpol.interp1d(np.arange(nexp, dtype=float), self.p)
p1 = float(p_i(primer[0][1]))
p2 = float(p_i(primer[1][1]))
v1 = primer[0][0]
v2 = primer[1][0]
self.primer = [[v1, p1], [v2, p2]]
self.fit_spinorbit()
# self.primer = a*self.rv+b
self.v_star_primer = fun.local_v_star(self.p, self.aRstar,
self.inclination, self.vsini_fit,
self.l_fit) + self.vsys
# ax[0].plot(v_star_primer,fun.findgen(nexp),'--',color='black',label='Spin-orbit fit to primer')
self.mask_ccf()
self.fit_model()
def construct_KpVsys(rv, ccf, ccf_e, dp, kprange=[0, 300], dkp=1.0):
"""The name says it all. Do good tests."""
import tayph.functions as fun
import tayph.operations as ops
import numpy as np
import tayph.system_parameters as sp
import matplotlib.pyplot as plt
import astropy.io.fits as fits
import tayph.util as ut
import sys
import pdb
Kp = fun.findgen((kprange[1] - kprange[0]) / dkp + 1) * dkp + kprange[0]
n_exp = np.shape(ccf)[0]
KpVsys = np.zeros((len(Kp), len(rv)))
KpVsys_e = np.zeros((len(Kp), len(rv)))
transit = sp.transit(dp) - 1.0
transit /= np.nansum(transit)
transitblock = fun.rebinreform(transit, len(rv)).T
j = 0
ccfs = []
for i in Kp:
dRV = sp.RV(dp, vorb=i) * (-1.0)
ccf_shifted = shift_ccf(rv, ccf, dRV)
ccf_e_shifted = shift_ccf(rv, ccf_e, dRV)
ccfs.append(ccf_shifted)
KpVsys[j, :] = np.nansum(transitblock * ccf_shifted, axis=0)
KpVsys_e[j, :] = (np.nansum((transitblock * ccf_e_shifted)**2.0,
axis=0))**0.5
# plt.plot(rv,KpVsys[j,:])
# plt.fill_between(rv, KpVsys[j,:]-KpVsys_e[j,:], KpVsys[j,:]+KpVsys_e[j,:],alpha=0.5)
# plt.show()
# pdb.set_trace()
j += 1
ut.statusbar(i, Kp)
return (Kp, KpVsys, KpVsys_e)
def inject_model(list_of_wls,
list_of_orders,
dp,
modelname,
model_library='library/models'):
"""This function takes a list of spectral orders and injects a model with library
identifier modelname, and system parameters as defined in dp. The model is blurred taking into
account spectral resolution and rotation broadening (with an LSF as per Brogi et al.) and
finite-exposure broadening (with a box LSF).
It returns a copy of the list of orders with the model injected."""
import tayph.util as ut
import tayph.system_parameters as sp
import tayph.models
import astropy.constants as const
import numpy as np
import scipy
import tayph.operations as ops
from tayph.vartests import typetest, dimtest
import pdb
import copy
import matplotlib.pyplot as plt
# dimtest(order,[0,len(wld)])
dp = ut.check_path(dp)
typetest(modelname, str, 'modelname in models.inject_model()')
typetest(model_library, str, 'model_library in models.inject_model()')
c = const.c.to('km/s').value #In km/s
Rd = sp.paramget('resolution', dp)
planet_radius = sp.paramget('Rp', dp)
inclination = sp.paramget('inclination', dp)
P = sp.paramget('P', dp)
transit = sp.transit(dp)
n_exp = len(transit)
vsys = sp.paramget('vsys', dp)
rv = sp.RV(dp) + vsys
dRV = sp.dRV(dp)
phi = sp.phase(dp)
dimtest(transit, [n_exp])
dimtest(rv, [n_exp])
dimtest(phi, [n_exp])
dimtest(dRV, [n_exp])
mask = (transit - 1.0) / (np.min(transit - 1.0))
wlm, fxm = get_model(modelname, library=model_library)
if wlm[-1] <= wlm[0]: #Reverse the wl axis if its sorted the wrong way.
wlm = np.flipud(wlm)
fxm = np.flipud(fxm)
#With the model and the revelant parameters in hand, now only select that
#part of the model that covers the wavelengths of the order provided.
#A larger wavelength range would take much extra time because the convolution
#is a slow operation.
N_orders = len(list_of_wls)
if N_orders != len(list_of_orders):
raise RuntimeError(
f'in models.inject_model: List_of_wls and list_of_orders are not of the '
f'same length ({N_orders} vs {len(list_of_orders)})')
if np.min(wlm) > np.min(list_of_wls) - 1.0 or np.max(
wlm) < np.max(list_of_wls) + 1.0:
ut.tprint(
'WARNING in model injection: Data grid falls (partly) outside of model range. '
'Setting missing area to 1.0. (meaning, no planet absorption.)')
modelsel = [
(wlm >= np.min(list_of_wls) - 1.0) & (wlm <= np.max(list_of_wls) + 1.0)
]
wlm = wlm[tuple(modelsel)]
fxm = fxm[tuple(modelsel)]
fxm_b = ops.blur_rotate(wlm, fxm, c / Rd, planet_radius, P,
inclination) #Only do this once per dataset.
oversampling = 2.5
wlm_cv, fxm_bcv, vstep = ops.constant_velocity_wl_grid(
wlm, fxm_b, oversampling=oversampling)
if np.min(dRV) < c / Rd / 10.0:
dRV_min = c / Rd / 10.0 #If the minimum dRV is less than 10% of the spectral
#resolution, we introduce a lower limit to when we are going to blur, because the effect
#becomes insignificant.
else:
dRV_min = np.min(dRV)
if dRV_min / vstep < 3: #Meaning, if the smoothing is going to be undersampled by this choice
#in v_step, it means that the oversampling parameter in ops.constant_velocity_wl_grid was
#not high enough. Then we adjust it. I choose a value of 3 here to be safe, even though
#ops.smooth below accepts values as low as 2.
oversampling_new = 3.0 / (
dRV_min / vstep) * oversampling #scale up the oversampling factor.
wlm_cv, fxm_bcv, vstep = ops.constant_velocity_wl_grid(
wlm, fxm_b, oversampling=oversampling_new)
list_of_orders_injected = copy.deepcopy(list_of_orders)
for i in range(n_exp):
if dRV[i] >= c / Rd / 10.0:
fxm_b2 = ops.smooth(fxm_bcv, dRV[i] / vstep, mode='box')
else:
fxm_b2 = copy.deepcopy(fxm_bcv)
shift = 1.0 + rv[i] / c
fxm_i = scipy.interpolate.interp1d(
wlm_cv * shift, fxm_b2, fill_value=1.0,
bounds_error=False) #This is a class that can be called.
#Fill_value = 1 because if the model does not fully cover the order, it will be padded with 1.0s,
#assuming that we are dealing with a model that is in transmission.
for j in range(len(list_of_orders)):
list_of_orders_injected[j][i] *= (
1.0 + mask[i] * (fxm_i(list_of_wls[j]) - 1.0)) #This assumes
#that the model is in transit radii. This can definitely be vectorised!
#These are for checking that the broadening worked as expected:
# injection_total[i,:]= scipy.interpolate.interp1d(wlm_cv*shift,fxm_b2)(wld)
# injection_rot_only[i,:]=scipy.interpolate.interp1d(wlm*shift,fxm_b)(wld)
# injection_pure[i,:]=scipy.interpolate.interp1d(wlm*shift,fxm)(wld)
# ut.save_stack('test.fits',[injection_pure,injection_rot_only,injection_total])
# pdb.set_trace()
# ut.writefits('test.fits',injection)
# pdb.set_trace()
return (list_of_orders_injected)
def plot_ccf(rv,
ccf,
dp,
xrange=[-200, 200],
yrange=[0, 0],
Nxticks=10.0,
Nyticks=10.0,
title='',
doppler_model=[],
i_legend=True,
show=True):
"""
This is a routine that does all the plotting of the cleaned 2D CCF. It overplots
the expected planet velocity, modified by the systemic velocity given in the config file.
Optionally, a trace of the doppler model is added as a list (of y-values) in the
doppler_model parameter. Set i_legend to False if you wish to remove the interactive legend.
"""
import numpy as np
import matplotlib.pyplot as plt
import pdb
import tayph.drag_colour as dcb
import tayph.functions as fun
import pylab as pl
import tayph.system_parameters as sp
#Load necessary physics for overplotting planet velocity.
vsys = sp.paramget('vsys', dp)
RVp = sp.RV(dp) + vsys
nexp = np.shape(ccf)[0]
#Default to the entire y-axis if yrange = [0,0]
if all(v == 0 for v in yrange):
yrange = [0, nexp - 1]
x2, y2, z, rv_sel, y_sel, xticks, yticks, vmin, vmax = plotting_scales_2D(
rv,
fun.findgen(nexp),
ccf,
xrange,
yrange,
Nxticks=Nxticks,
Nyticks=Nyticks,
nsigma=3.0)
#The plotting
fig, ax = plt.subplots(figsize=(12, 6))
img = ax.pcolormesh(x2, y2, z, vmin=vmin, vmax=vmax, cmap='hot')
ax.axis([x2.min(), x2.max(), y2.min(), y2.max()])
line1, = ax.plot(RVp,
fun.findgen(nexp),
'--',
color='black',
label='Planet rest-frame')
if len(doppler_model) > 0:
line2, = ax.plot(doppler_model + vsys,
fun.findgen(nexp),
'--',
color='black',
label='Doppler shadow')
ax.set_xticks(xticks)
ax.set_yticks(yticks)
ax.set_title(title)
ax.set_xlabel('Radial velocity (km/s)')
ax.set_ylabel('Exposure')
#The colourbar
cbar = plt.colorbar(img, format='%05.4f', aspect=15)
# cbar.set_norm(dcb.MyNormalize(vmin=vmin,vmax=vmax,stretch='linear'))
cbar = dcb.DraggableColorbar_fits(cbar, img, 'hot')
cbar.connect()
#The clickable legend.
if len(doppler_model) > 0:
lines = [line1, line2]
else:
lines = [line1]
if i_legend == True:
interactive_legend(fig, ax, lines)
if show == True:
plt.show()
return (fig, ax, cbar)
def construct_doppler_model(rv,
ccf,
dp,
shadowname,
xrange=[-200, 200],
Nxticks=20.0,
Nyticks=10.0):
"""This is the the main function to construct a doppler model. The above are mostly dependencies."""
import numpy as np
import matplotlib.pyplot as plt
import tayph.drag_colour as dcb
import tayph.functions as fun
import tayph.system_parameters as sp
import tayph.plotting as fancyplots
import sys
from matplotlib.widgets import Slider, Button, RadioButtons, CheckButtons
#This is for setting plot axes in the call to plotting_scales_2D below.
nexp = np.shape(ccf)[0]
yrange = [0, nexp - 1]
y_axis = np.arange(nexp, dtype=float) #fun.findgen(nexp)
#And for adding the planet line:
vsys = sp.paramget('vsys', dp)
vsini = sp.paramget('vsini', dp)
RVp = sp.RV(dp) + vsys
transit = sp.transit(dp)
sel_transit = y_axis[(transit < 1.0)]
transit_start = min(sel_transit)
transit_end = max(sel_transit)
fig, ax, cbar = fancyplots.plot_ccf(
rv,
ccf,
dp,
xrange=xrange,
Nxticks=Nxticks,
Nyticks=Nyticks,
i_legend=False,
show=False) #Initiates the plot for the primer.
primer = prime_doppler_model(fig, ax, cbar) #We use the active Figure
#to let the user indicate where the doppler shadow is located (the primer). This calls the plt.show()
#which was not called by plot_ccf. To proceed, we would like to model the shadow using the
#primer, along with fancy visualization of the ccf.
#We first re-instate a plot. This plot is a bit more complex than the primer so we don't
#use plot_ccf anymore; though the philosophy is similar.
x2, y2, z, rv_sel, y_sel, xticks, yticks, vmin, vmax = fancyplots.plotting_scales_2D(
rv,
y_axis,
ccf,
xrange,
yrange,
Nxticks=Nxticks,
Nyticks=Nyticks,
nsigma=3.0)
#Create empty plots.
fig, ax = plt.subplots(3, 1, sharex=True, figsize=(13, 6))
plt.subplots_adjust(right=0.75)
#Here we initiate the model instance that does the fitting and handles the GUI.
#This does an initial fit based on the primer.
model_callback = fit_doppler_model(fig, ax, rv_sel, z, primer, dp,
shadowname)
s_init = model_callback.maskHW #Masking half-width around the planet RV to start the slider with.
#We continue by overplotting the velocty traces and transit markers onto the 3 subplots, saving the references to the lines.
#These lines can be set to be visible/invisible using the checkbuttons below, and are set to be invisible
#when the plot opens.
l_planet = []
t_start = []
t_end = []
l_primer = []
l_vfit = []
for sub_ax in ax:
sub_ax.axis([x2.min(), x2.max(), y2.min(), y2.max()])
sub_ax.set_xticks(xticks)
sub_ax.set_yticks(yticks)
sub_ax.set_ylabel('Exposure')
l1 = sub_ax.plot(RVp,
y_axis,
'--',
color='black',
label='Planet rest-frame',
visible=False)[0]
l2 = sub_ax.plot(rv,
rv * 0.0 + transit_start,
'--',
color='white',
label='Transit start',
visible=False)[0]
l3 = sub_ax.plot(rv,
rv * 0.0 + transit_end,
'--',
color='white',
label='Transit end',
visible=False)[0]
l4 = sub_ax.plot(model_callback.v_star_primer,
y_axis,
'--',
color='black',
label='Local velocity (primer)',
visible=False)[0]
l5 = sub_ax.plot(model_callback.v_star_fit,
y_axis,
'--',
color='black',
label='Local velocity (fit)',
visible=False)[0]
l_planet.append(l1)
t_start.append(l2)
t_end.append(l3)
l_primer.append(l4)
l_vfit.append(l5)
ax[0].set_title('Data')
ax[1].set_title('Model shadow')
ax[2].set_title('Residual')
ax[2].set_xlabel('Radial velocity (km/s)')
#Here we actually plot the initial fit, which will be modified each time the parameters are changed
#using the GUI buttons/sliders.
img1 = ax[0].pcolormesh(x2, y2, z, vmin=vmin, vmax=vmax, cmap='hot')
img2 = ax[1].pcolormesh(x2,
y2,
model_callback.model,
vmin=vmin,
vmax=vmax,
cmap='hot')
img3 = ax[2].pcolormesh(x2,
y2,
z - model_callback.model,
vmax=vmax,
cmap='hot')
#This trick to associate a single CB to multiple axes comes from
#https://stackoverflow.com/questions/13784201/matplotlib-2-subplots-1-colorbar
cbar = fig.colorbar(img1,
ax=ax.ravel().tolist(),
format='%05.4f',
aspect=15)
cbar = dcb.DraggableColorbar_fits(cbar, [img1, img2, img3], 'hot')
cbar.connect()
#We define the interface and the bahaviour of button/slider callbacks.
#First the check buttons for showing the lines defined above.
rax_top = plt.axes([0.8, 0.65, 0.15, 0.25])
rax_top.set_title('Plot:')
labels = [
'Planet velocity', 'Transit start/end', 'Shadow v$_c$ primer',
'Shadow $v_c$ fit', 'Masked area'
]
start = [False, False, False, False, False,
False] #Start with none of these actually visible.
check = CheckButtons(rax_top, labels, start)
def func(label):
index = labels.index(label)
lines = [l_planet, np.append(t_end, t_start), l_primer, l_vfit]
if index < len(lines):
for l in lines[index]:
l.set_visible(not l.get_visible())
if index == len(
lines
): #I.e. if we are on the last element, which is not a line an option for SHOWING the masked area:
status = check.get_status()[-1]
if status == True: #If true, mask the image.
data = z * model_callback.ccf_mask
data[np.isnan(
data
)] = np.inf #The colobar doesn't eat NaNs, set them to inf instead for the plot. Makes them white, too.
img1.set_array((data).ravel())
img3.set_array((data - model_callback.model).ravel())
if status == False: #If false (unclicked), then just the data w/o mask.
img1.set_array(z.ravel())
img3.set_array((z - model_callback.model).ravel())
plt.draw()
check.on_clicked(func)
#Then the choice for 1 or 2 Gaussian fitting components:
rax_middle = plt.axes([0.8, 0.53, 0.15, 0.10])
clabels = ['1 component', '2 components']
radio = RadioButtons(rax_middle, clabels)
def cfunc(label):
index = clabels.index(label)
model_callback.n_components = index + 1
model_callback.fit_model() #Each time we change the choice, refit.
status = check.get_status()[-1]
if status == True: #If true, mask the image.
data = z * model_callback.ccf_mask
data[np.isnan(
data
)] = np.inf #The colobar doesn't eat NaNs, set them to inf instead for the plot.
img2.set_array(model_callback.model.ravel())
img3.set_array((data - model_callback.model).ravel())
if status == False: #If false (unclicked), then just the data w/o mask.
img2.set_array(model_callback.model.ravel())
img3.set_array((z - model_callback.model).ravel())
plt.draw()
radio.on_clicked(cfunc)
rax_deg = plt.axes([0.8, 0.35, 0.07, 0.13])
plt.title('Polynomial degree', fontsize=8)
rax_poly = plt.axes([0.88, 0.35, 0.07, 0.13])
dlabels = ['0 (off)', '2', '4', '6']
plabels = ['Single', 'Even', 'Full']
dradio = RadioButtons(rax_deg, dlabels)
pradio = RadioButtons(rax_poly, plabels)
def update_degree(label):
model_callback.D = dlabels.index(label) * 2
model_callback.fit_model() #Each time we change the choice, refit.
status = check.get_status()[-1]
if status == True: #If true, mask the image.
data = z * model_callback.ccf_mask
data[np.isnan(
data
)] = np.inf #The colobar doesn't eat NaNs, set them to inf instead for the plot.
img2.set_array(model_callback.model.ravel())
img3.set_array((data - model_callback.model).ravel())
if status == False: #If false (unclicked), then just the data w/o mask.
img2.set_array(model_callback.model.ravel())
img3.set_array((z - model_callback.model).ravel())
plt.draw()
def update_poly(label):
model_callback.S = plabels.index(label)
model_callback.fit_model() #Each time we change the choice, refit.
status = check.get_status()[-1]
if status == True: #If true, mask the image.
data = z * model_callback.ccf_mask
data[np.isnan(
data
)] = np.inf #The colobar doesn't eat NaNs, set them to inf instead for the plot.
img2.set_array(model_callback.model.ravel())
img3.set_array((data - model_callback.model).ravel())
if status == False: #If false (unclicked), then just the data w/o mask.
img2.set_array(model_callback.model.ravel())
img3.set_array((z - model_callback.model).ravel())
plt.draw()
dradio.on_clicked(update_degree)
pradio.on_clicked(update_poly)
#Then the offset slider:
rax_slider2 = plt.axes([0.8, 0.25, 0.15, 0.02])
rax_slider2.set_title('Offset 2nd component')
offset_slider = Slider(rax_slider2,
'',
-1.0 * np.ceil(vsini),
np.ceil(vsini),
valinit=0.0,
valstep=1.0)
def update_offset(val):
model_callback.offset = offset_slider.val
status = radio.value_selected
if status == clabels[
1]: #Only update the fit if we are actually asked to do 2 components.
model_callback.fit_model()
# data = z*model_callback.ccf_mask
# data[np.isnan(data)] = np.inf#The colobar doesn't eat NaNs...
# img1.set_array((data).ravel())
img2.set_array(model_callback.model.ravel())
img3.set_array((z - model_callback.model).ravel())
# if status == False:#If false (unclicked), then just the data w/o mask.
# img1.set_array(z.ravel())
# img2.set_array(model_callback.model.ravel())
# img3.set_array((z-model_callback.model).ravel())
plt.draw()
offset_slider.on_changed(update_offset)
#Then the slider:
rax_slider = plt.axes([0.8, 0.18, 0.15, 0.02])
rax_slider.set_title('Mask width (km/s)')
mask_slider = Slider(rax_slider,
'',
0.0,
30.0,
valinit=s_init,
valstep=1.0)
def update(val):
model_callback.maskHW = mask_slider.val
model_callback.mask_ccf()
model_callback.fit_model()
status = check.get_status()[-1]
if status == True: #If true, mask the image.
data = z * model_callback.ccf_mask
data[np.isnan(data)] = np.inf #The colobar doesn't eat NaNs...
img1.set_array((data).ravel())
img2.set_array(model_callback.model.ravel())
img3.set_array((data - model_callback.model).ravel())
if status == False: #If false (unclicked), then just the data w/o mask.
img1.set_array(z.ravel())
img2.set_array(model_callback.model.ravel())
img3.set_array((z - model_callback.model).ravel())
plt.draw()
mask_slider.on_changed(update)
#And finally the save button.
rax_save = plt.axes([0.875, 0.1, 0.06, 0.05])
bsave = Button(rax_save, 'Save')
bsave.on_clicked(model_callback.save)
rax_cancel = plt.axes([0.8, 0.1, 0.06, 0.05])
bcancel = Button(rax_cancel, 'Cancel')
bcancel.on_clicked(model_callback.cancel)
plt.show() #All fitting is done before this line through event handling.
def Kp_parallel(i):
dRV = sp.RV(dp, vorb=i) * (-1.0)
ccf_shifted = shift_ccf(rv, ccf, dRV)
ccf_e_shifted = shift_ccf(rv, ccf_e, dRV)
return (np.nansum(transitblock * ccf_shifted, axis=0), (np.nansum(
(transitblock * ccf_e_shifted)**2.0, axis=0))**0.5)