def running_MAD_2D(z, w, verbose=False, parallel=False):
"""Computers a running standard deviation of a 2-dimensional array z.
The stddev is evaluated over the vertical block with width w pixels.
The output is a 1D array with length equal to the width of z.
This is very slow on arrays that are wide in x (hundreds of thousands of points)."""
import astropy.stats as stats
import numpy as np
from tayph.vartests import typetest, dimtest, postest
import tayph.util as ut
if parallel: from joblib import Parallel, delayed
typetest(z, np.ndarray, 'z in fun.running_MAD_2D()')
dimtest(z, [0, 0], 'z in fun.running_MAD_2D()')
typetest(w, [int, float], 'w in fun.running_MAD_2D()')
postest(w, 'w in fun.running_MAD_2D()')
size = np.shape(z)
ny = size[0]
nx = size[1]
s = np.arange(0, nx, dtype=float) * 0.0
dx1 = int(0.5 * w)
dx2 = int(int(0.5 * w) + (w % 2)) #To deal with odd windows.
for i in range(nx):
minx = max([0, i - dx1]) #This here is only a 3% slowdown.
maxx = min([nx, i + dx2])
s[i] = stats.mad_std(
z[:, minx:maxx],
ignore_nan=True) #This is what takes 97% of the time.
if verbose: ut.statusbar(i, nx)
return (s)
def selmax(y_in,p,s=0.0):
"""This program returns the p (fraction btw 0 and 1) highest points in y,
ignoring the very top s % (default zero, i.e. no points ignored), for the
purpose of outlier rejection."""
import tayph.util as ut
from tayph.vartests import postest,dimtest
import numpy as np
import copy
postest(p)
y=copy.deepcopy(y_in)#Copy itself, or there are pointer-related troubles...
if s < 0.0:
raise Exception("ERROR in selmax: s should be zero or positive.")
if p >= 1.0:
raise Exception("ERROR in selmax: p should be strictly between 0.0 and 1.0.")
if s >= 1.0:
raise Exception("ERROR in selmax: s should be strictly less than 1.0.")
postest(-1.0*p+1.0)
#ut.nantest('y in selmax',y)
dimtest(y,[0])#Test that it is one-dimensional.
y[np.isnan(y)]=np.nanmin(y)#set all nans to the minimum value such that they will not be selected.
y_sorting = np.flipud(np.argsort(y))#These are the indices in descending order (thats why it gets a flip)
N=len(y)
if s == 0.0:
max_index = np.max([int(round(N*p)),1])#Max because if the fraction is 0 elements, then at least it should contain 1.0
return y_sorting[0:max_index]
if s > 0.0:
min_index = np.max([int(round(N*s)),1])#If 0, then at least it should be 1.0
max_index = np.max([int(round(N*(p+s))),2]) #If 0, then at least it should be 1+1.
return y_sorting[min_index:max_index]
def telluric_correction_order(i):
order = list_of_orders[i]
order_cor = order * 0.0
error = list_of_sigmas[i]
error_cor = error * 0.0
wl = copy.deepcopy(list_of_wls[i]) #input wl axis, either 1D or 2D.
#If it is 1D, we make it 2D by tiling it vertically:
if wl.ndim == 1: wl = np.tile(wl, (Nexp, 1)) #Tile it into a 2D thing.
#If read (2D) or tiled (1D) correctly, wl and order should have the same shape:
dimtest(wl, np.shape(order),
f'Wl axis of order {i}/{No} in apply_telluric_correction()')
dimtest(error, np.shape(order),
f'errors {i}/{No} in apply_telluric_correction()')
for j in range(Nexp):
T_i = interp.interp1d(wlT[j], fxT[j],
fill_value="extrapolate")(wl[j])
postest(T_i,
f'T-spec of exposure {j} in apply_telluric_correction()')
nantest(T_i,
f'T-spec of exposure {j} in apply_telluric_correction()')
order_cor[j] = order[j] / T_i
error_cor[j] = error[
j] / T_i #I checked that this works because the SNR before and after
#telluric correction is identical.
return (order_cor, error_cor)
def envelope(wlm, fxm, binsize, selfrac=0.05, mode='top', threshold=''):
"""
This program measures the top or bottom envelope of a spectrum (wl,fx), by
chopping it up into bins of size binsze (unit of wl), and measuring the mean
of the top n % of values in that bin. Setting the mode to 'bottom' will do the
oppiste: The mean of the bottom n% of values. The output is the resulting wl
and flux points of these bins.
Example: wle,fxe = envelope(wl,fx,1.0,selfrac=3.0,mode='top')
"""
import numpy as np
import tayph.util as ut
import tayph.functions as fun
from tayph.vartests import typetest, dimtest, nantest, postest
from matplotlib import pyplot as plt
typetest(wlm, np.ndarray, 'wlm in ops.envelope()')
typetest(fxm, np.ndarray, 'fxm in ops.envelope()')
dimtest(wlm, [len(fxm)])
typetest(binsize, [int, float], 'binsize in ops.envelope()')
typetest(selfrac, float, 'percentage in ops.envelope()')
typetest(mode, str, 'mode in envelope')
nantest(fxm, 'fxm in ops.envelope()')
nantest(wlm, 'wlm in ops.envelope()')
postest(wlm, 'wlm in ops.envelope()')
postest(binsize, 'binsize in ops.envelope()')
if mode not in ['top', 'bottom']:
raise Exception(
f'RuntimeError in ops.envelope(): Mode should be set to "top" or "bottom" ({mode}).'
)
binsize = float(binsize)
if mode == 'bottom':
fxm *= -1.0
wlcs = np.array([]) #Center wavelengths
fxcs = np.array([]) #Flux at center wavelengths
i_start = 0
wlm_start = wlm[i_start]
for i in range(0, len(wlm) - 1):
if wlm[i] - wlm_start >= binsize:
wlsel = wlm[i_start:i]
fxsel = fxm[i_start:i]
maxsel = fun.selmax(fxsel, selfrac)
wlcs = np.append(wlcs, np.mean(wlsel[maxsel]))
fxcs = np.append(fxcs, np.mean(fxsel[maxsel]))
i_start = i + 1
wlm_start = wlm[i + 1]
if isinstance(threshold, float) == True:
#This means that the threshold value is set, and we set all bins less than
#that threshold to the threshold value:
if mode == 'bottom':
threshold *= -1.0
fxcs[(fxcs < threshold)] = threshold
if mode == 'bottom':
fxcs *= -1.0
fxm *= -1.0
return wlcs, fxcs
def running_mean_2D(D, w):
"""This computes a running mean on a 2D array in a window with width w that
slides over the array in the horizontal (x) direction."""
import numpy as np
from tayph.vartests import typetest, dimtest, postest
typetest(D, np.ndarray, 'z in fun.running_mean_2D()')
typetest(w, [int, float], 'w in fun.running_mean_2D()')
postest(w, 'w in fun.running_mean_2D()')
ny, nx = D.shape
m2 = strided_window(D, w, pad=True)
s = np.nanmean(m2, axis=(1, 2))
return (s)
def findgen(n,integer=False):
"""This is basically IDL's findgen function.
a = findgen(5) will return an array with 5 elements from 0 to 4:
[0,1,2,3,4]
"""
import numpy as np
from tayph.vartests import typetest,postest
typetest(n,[int,float],'n in findgen()')
typetest(integer,bool,'integer in findgen()')
postest(n,'n in findgen()')
n=int(n)
if integer:
return np.linspace(0,n-1,n).astype(int)
else:
return np.linspace(0,n-1,n)
def running_MAD(z,w):
"""Computers a running standard deviation of a 1-dimensional array z.
The stddev is evaluated over a range with width w pixels.
The output is a 1D array with length equal to the width of z."""
import astropy.stats as stats
import numpy as np
from tayph.vartests import typetest,dimtest,postest
typetest(z,np.ndarray,'z in fun.running_MAD()')
typetest(w,[int,float],'w in fun.running_MAD()')
postest(w,'w in fun.running_MAD_2D()')
nx = len(z)
s = np.arange(0,nx,dtype=float)*0.0
dx1=int(0.5*w)
dx2=int(int(0.5*w)+(w%2))#To deal with odd windows.
for i in range(nx):
minx = max([0,i-dx1])
maxx = min([nx,i+dx2])
s[i] = stats.mad_std(z[minx:maxx],ignore_nan=True)
return(s)
def running_MAD_2D(z,w):
"""Computers a running standard deviation of a 2-dimensional array z.
The stddev is evaluated over the vertical block with width w pixels.
The output is a 1D array with length equal to the width of z."""
import astropy.stats as stats
import numpy as np
from tayph.vartests import typetest,dimtest,postest
typetest(z,np.ndarray,'z in fun.running_MAD_2D()')
dimtest(z,[0,0],'z in fun.running_MAD_2D()')
typetest(w,[int,float],'w in fun.running_MAD_2D()')
postest(w,'w in fun.running_MAD_2D()')
size = np.shape(z)
ny = size[0]
nx = size[1]
s = findgen(nx)*0.0
for i in range(nx):
minx = max([0,i-int(0.5*w)])
maxx = min([nx-1,i+int(0.5*w)])
s[i] = stats.mad_std(z[:,minx:maxx],ignore_nan=True)
return(s)
def v_orb(dp):
"""
This program calculates the orbital velocity in km/s for the planet in the
data sequence provided in dp, the data-path. dp starts in the root folder,
i.e. it starts with data/projectname/. This assumes a circular orbit.
The output is a number in km/s.
Parameters
----------
dp : str, path like
The path to the dataset containing the config file.
Returns
-------
v_orb : float
The planet's orbital velocity.
list_of_sigmas_corrected : list
List of 2D error matrices, telluric corrected.
"""
import numpy as np
import pdb
import astropy.units as u
from tayph.vartests import typetest,postest
dp=check_dp(dp)#Path object
P=paramget('P',dp)
r=paramget('a',dp)
typetest(P,float,'P in sp.v_orb()')
typetest(r,float,'r in sp.v_orb()')
postest(P,'P in sp.v_orb()')
postest(r,'r in sp.v_orb()')
return (2.0*np.pi*r*u.AU/(P*u.d)).to('km/s').value
def mask_orders(list_of_wls,list_of_orders,dp,maskname,w,c_thresh,manual=False):
"""
This code takes the list of orders and masks out bad pixels.
It combines two steps, a simple sigma clipping step and a manual step, where
the user can interactively identify bad pixels in each order. The sigma
clipping is done on a threshold of c_thresh, using a rolling standard dev.
with a width of w pixels. Manual masking is a big routine needed to support
a nice GUI to do that.
If c_thresh is set to zero, sigma clipping is skipped. If manual=False, the
manual selection of masking regions (which is manual labour) is turned off.
If both are turned off, the list_of_orders is returned unchanged.
If either or both are active, the routine will output 1 or 2 FITS files that
contain a stack (cube) of the masks for each order. The first file is the mask
that was computed automatically, the second is the mask that was constructed
manually. This is done so that the manual mask can be transplanted onto another
dataset, or saved under a different file-name, to limit repetition of work.
At the end of the routine, the two masks are merged into a single list, and
applied to the list of orders.
"""
import tayph.operations as ops
import numpy as np
import tayph.functions as fun
import tayph.plotting as plotting
import sys
import matplotlib.pyplot as plt
import tayph.util as ut
import warnings
from tayph.vartests import typetest,dimtest,postest
ut.check_path(dp)
typetest(maskname,str,'maskname in mask_orders()')
typetest(w,[int,float],'w in mask_orders()')
typetest(c_thresh,[int,float],'c_thresh in mask_orders()')
postest(w,'w in mask_orders()')
postest(c_thresh,'c_thresh in mask_orders()')
typetest(list_of_wls,list,'list_of_wls in mask_orders()')
typetest(list_of_orders,list,'list_of_orders in mask_orders()')
typetest(manual,bool,'manual keyword in mask_orders()')
dimtest(list_of_wls,[0,0],'list_of_wls in mask_orders()')
dimtest(list_of_orders,[len(list_of_wls),0,0],'list_of_orders in mask_orders()')
if c_thresh <= 0 and manual == False:
print('---WARNING in mask_orders: c_thresh is set to zero and manual masking is turned off.')
print('---Returning orders unmasked.')
return(list_of_orders)
N = len(list_of_orders)
void = fun.findgen(N)
list_of_orders = ops.normalize_orders(list_of_orders,list_of_orders)[0]#first normalize. Dont want outliers to
#affect the colour correction later on, so colour correction cant be done before masking, meaning
#that this needs to be done twice; as colour correction is also needed for proper maskng. The second variable is
#a dummy to replace the expected list_of_sigmas input.
N_NaN = 0
list_of_masked_orders = []
for i in range(N):
list_of_masked_orders.append(list_of_orders[i])
list_of_masks = []
if c_thresh > 0:#Check that c_thresh is positive. If not, skip sigma clipping.
print('------Sigma-clipping mask')
for i in range(N):
order = list_of_orders[i]
N_exp = np.shape(order)[0]
N_px = np.shape(order)[1]
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
meanspec = np.nanmean(order,axis = 0)
meanblock = fun.rebinreform(meanspec,N_exp)
res = order / meanblock - 1.0
sigma = fun.running_MAD_2D(res,w)
with np.errstate(invalid='ignore'):#https://stackoverflow.com/questions/25345843/inequality-comparison-of-numpy-array-with-nan-to-a-scalar
sel = np.abs(res) >= c_thresh*sigma
N_NaN += np.sum(sel)#This is interesting because True values count as 1, and False as zero.
order[sel] = np.nan
list_of_masks.append(order*0.0)
ut.statusbar(i,void)
print(f'%s outliers identified and set to NaN ({N_NaN}/{round(N_NaN/np.size(list_of_masks)*100.0,3)}).')
else:
print('------Skipping sigma-clipping (c_thres <= 0)')
#Do nothing to list_of_masks. It is now an empty list.
#We now automatically proceed to manual masking, because at this point
#it has already been established that it must have been turned on.
list_of_masks_manual = []
if manual == True:
previous_list_of_masked_columns = load_columns_from_file(dp,maskname,mode='relaxed')
list_of_masked_columns = manual_masking(list_of_wls,list_of_orders,list_of_masks,saved = previous_list_of_masked_columns)
print('------Successfully concluded manual mask.')
write_columns_to_file(dp,maskname,list_of_masked_columns)
print('------Building manual mask from selected columns')
for i in range(N):
order = list_of_orders[i]
N_exp = np.shape(order)[0]
N_px = np.shape(order)[1]
list_of_masks_manual.append(np.zeros((N_exp,N_px)))
for j in list_of_masked_columns[i]:
list_of_masks_manual[i][:,int(j)] = np.nan
#We write 1 or 2 mask files here. The list of manual masks
#and list_of_masks (auto) are either filled, or either is an emtpy list if
#c_thresh was set to zero or manual was set to False (because they were defined
#as empty lists initially, and then not filled with anything).
write_mask_to_file(dp,maskname,list_of_masks,list_of_masks_manual)
return(0)
def manual_masking(list_of_wls,list_of_orders,list_of_masks,Nxticks = 20,Nyticks = 10,saved = []):
import numpy as np
import matplotlib.pyplot as plt
import pdb
import tayph.drag_colour as dcb
import tayph.util as ut
from tayph.vartests import typetest,postest,dimtest
import tayph.functions as fun
import sys
from matplotlib.widgets import Slider, Button, RadioButtons, CheckButtons
"""
This brings the user into a GUI in which he/she/they can both visualize
spectral orders and mask regions of bad data; such as occur at edges
of orders, inside deep spectral lines (stellar or telluric) or other places.
In the standard workflow, this routine succeeds an automatic masking step that has
performed a sigma clipping using a rolling standard deviation. This procedure has
added a certain number of pixels to the mask. If this routine is subsequently
called, the user can mask out bad *columns* by selecting them in the GUI. These
are then added to that mask as well.
"""
typetest(Nxticks,int,'Nxticks in manual_masking()')
typetest(Nyticks,int,'Nyticks in manual_masking()')
postest(Nxticks,'Nxticks in manual_masking()')
postest(Nyticks,'Nyticks in manual_masking()')
typetest(list_of_wls,list,'list_of_wls in manual_masking()')
typetest(list_of_orders,list,'list_of_orders in manual_masking()')
typetest(list_of_masks,list,'list_of_masks in manual_masking()')
typetest(saved,list,'saved list_of_masks in manual_masking()')
dimtest(list_of_wls,[0,0],'list_of_wls in manual_masking()')
dimtest(list_of_orders,[len(list_of_wls),0,0],'list_of_orders in manual_masking()')
dimtest(list_of_masks,[len(list_of_wls),0,0],'list_of_masks in manual_masking()')
print('------Entered manual masking mode')
M = mask_maker(list_of_wls,list_of_orders,saved,Nxticks=Nxticks,Nyticks=Nyticks,nsigma=3.0) #Mask callback
#This initializes all the parameters of the plot. Which order it is plotting, what
#dimensions these have, the actual data arrays to plot; etc. Initialises on order 56.
#I also dumped most of the buttons and callbacks into there; so this thing
#is probably a bit of a sphaghetti to read.
#Here we only need to define the buttons and add them to the plot. In fact,
#I could probably have shoved most of this into the class as well, (see
#me defining all these buttons as class attributes? But, I
#suppose that the structure is a bit more readable this way.
#The button to add a region to the mask:
rax_sub = plt.axes([0.8, 0.4, 0.14, 0.05])
M.bsub = Button(rax_sub, ' Unmask columns ')
M.bsub.on_clicked(M.subtract)
#The button to add a region to the mask:
rax_add = plt.axes([0.8, 0.48, 0.14, 0.05])
M.badd = Button(rax_add, ' Mask columns ')
M.badd.color=M.col_passive[0]
M.badd.hovercolor=M.col_passive[1]
M.add_connector = M.badd.on_clicked(M.add)
#The button to add a region to the mask:
rax_all = plt.axes([0.8, 0.56, 0.14, 0.05])
M.ball = Button(rax_all, ' Apply to all ')
M.ball.on_clicked(M.applyall)
rax_future = plt.axes([0.8, 0.64, 0.14, 0.05])
M.ballf = Button(rax_future, ' Apply to all above')
M.ballf.on_clicked(M.applyallfuture)
rax_clearentire = plt.axes([0.8, 0.72, 0.14, 0.05])
M.bcent = Button(rax_clearentire, ' Unmask order ')
M.bcent.on_clicked(M.cleartotal)
#The button to add a region to the mask:
rax_entire = plt.axes([0.8, 0.8, 0.14, 0.05])
M.bent = Button(rax_entire, ' Mask order ')
M.bent.on_clicked(M.applytotal)
# rax_atv = plt.axes([0.9, 0.45, 0.1, 0.15])
# clabels = ['To air', 'No conversion','To vaccuum']
# radio = RadioButtons(rax_atv,clabels)
# radio.on_clicked(M.atv)
#The mask width:
rax_slider = plt.axes([0.8, 0.3, 0.14, 0.02])
rax_slider.set_title('Mask width')
M.MW_slider = Slider(rax_slider,'', 1,200,valinit=M.MW,valstep=1)#Store the slider in the model class
M.MW_slider.on_changed(M.slide_maskwidth)
#The slider to cycle through orders:
rax_slider = plt.axes([0.8, 0.2, 0.14, 0.02])
rax_slider.set_title('Order')
M.mask_slider = Slider(rax_slider,'', 0,M.N_orders-1,valinit=M.N,valstep=1)#Store the slider in the model class
M.mask_slider.on_changed(M.slide_order)
#The Previous order button:
rax_prev = plt.axes([0.8, 0.1, 0.04, 0.05])
bprev = Button(rax_prev, ' <<< ')
bprev.on_clicked(M.previous)
#The Next order button:
rax_next = plt.axes([0.86, 0.1, 0.04, 0.05])
bnext = Button(rax_next, ' >>> ')
bnext.on_clicked(M.next)
#The save button:
rax_save = plt.axes([0.92, 0.1, 0.07, 0.05])
bsave = Button(rax_save, 'Save')
bsave.on_clicked(M.save)
#The cancel button:
rax_cancel = plt.axes([0.92, 0.025, 0.07, 0.05])
bcancel = Button(rax_cancel, 'Cancel')
bcancel.on_clicked(M.cancel)
plt.show()
#When pressing "save", the figure is closed, the suspesion caused by plt.show() is
#lifted and we continue to exit this GUI function, as its job has been done:
#return all the columns that were selected by the user.
return(M.list_of_selected_columns)
def __init__(self,list_of_wls,list_of_orders,list_of_saved_selected_columns,Nxticks,Nyticks,nsigma=3.0):
"""
We initialize with a figure object, three axis objects (in a list)
the wls, the orders, the masks already made; and we do the first plot.
NOTICE: Anything that is potted in these things as INF actually used to be
a NaN that was masked out before.
"""
import numpy as np
import pdb
import tayph.functions as fun
import tayph.plotting as plotting
import tayph.drag_colour as dcb
import matplotlib.pyplot as plt
import itertools
from matplotlib.widgets import MultiCursor
import tayph.util as ut
from tayph.vartests import typetest,postest
import copy
#Upon initialization, we raise the keywords onto self.
self.N_orders = len(list_of_wls)
if len(list_of_wls) < 1 or len(list_of_orders) < 1:# or len(list_of_masks) <1:
raise Exception('Runtime Error in mask_maker init: lists of WLs, orders and/or masks have less than 1 element.')
if len(list_of_wls) != len(list_of_orders):# or len(list_of_wls) != len(list_of_masks):
raise Exception('Runtime Error in mask_maker init: List of wls and list of orders have different length (%s & %s).' % (len(list_of_wls),len(list_of_orders)))
typetest(Nxticks,int,'Nxticks in mask_maker init',)
typetest(Nyticks,int,'Nyticks in mask_maker init',)
typetest(nsigma,float,'Nsigma in mask_maker init',)
postest(Nxticks,varname='Nxticks in mask_maker init')
postest(Nyticks,varname='Nyticks in mask_maker init')
postest(nsigma,varname='Nsigma in mask_maker init')
self.list_of_wls = list_of_wls
self.list_of_orders = list_of_orders
self.list_of_selected_columns = list(list_of_saved_selected_columns)
#Normally, if there are no saved columns to load, list_of_saved_selected_columns is an empty list. However if
#it is set, then its automatically loaded into self.list_of_selected_columns upon init.
#Below there is a check to determine whether it was empty or not, and whether the list of columns
#has the same length as the list of orders.
if len(self.list_of_selected_columns) == 0:
for i in range(self.N_orders):
self.list_of_selected_columns.append([])#Make a list of empty lists.
#This will contain all columns masked by the user, on top of the things
#that are already masked by the program.
else:
if len(self.list_of_selected_columns) != self.N_orders:
raise Exception('Runtime Error in mask_maker init: Trying to restore previously saved columns but the number of orders in the saved column file does not match the number of orders provided.')
print('------Restoring previously saved columns in mask-maker')
#All checks are now complete. Lets prepare to do the masking.
# self.N = min([56,self.N_orders-1])#We start on order 56, or the last order if order 56 doesn't exist.
self.N=0
#Set the current active order to order , and calculate the meanspec
#and residuals to be plotted, which are saved in self.
self.set_order(self.N)
#Sorry for the big self.spaghetti of code. This initializes the plot.
#Functions and vars further down in the class will deal with updating the plots
#as buttons are pressed. Some of this is copied from the construct_doppler_model
#function; but this time I made it part of the class.
#First define plotting and axis parameters for the colormesh below.
self.Nxticks = Nxticks
self.Nyticks = Nyticks
self.nsigma = nsigma
self.xrange = [0,self.npx-1]
self.yrange=[0,self.nexp-1]
self.x_axis=fun.findgen(self.npx).astype(int)
self.y_axis = fun.findgen(self.nexp).astype(int)
self.x2,self.y2,self.z,self.wl_sel,self.y_axis_sel,self.xticks,self.yticks,void1,void2= plotting.plotting_scales_2D(self.x_axis,self.y_axis,self.residual,self.xrange,self.yrange,Nxticks=self.Nxticks,Nyticks=self.Nyticks,nsigma=self.nsigma)
self.fig,self.ax = plt.subplots(3,1,sharex=True,figsize=(14,6))#Initialize the figure and 3 axes.
plt.subplots_adjust(left=0.05)#Make them more tight, we need all the space we can get.
plt.subplots_adjust(right=0.85)
self.ax[0].set_title('Spectral order %s (%s - %s nm)' % (self.N,round(np.min(self.wl),1),round(np.max(self.wl),1)))
self.ax[1].set_title('Residual of time-average')
self.ax[2].set_title('Time average 1D spectrum')
array1 = copy.deepcopy(self.order)
array2 = copy.deepcopy(self.residual)
array1[np.isnan(array1)] = np.inf#The colobar doesn't eat NaNs, so now set them to inf just for the plot.
array2[np.isnan(array2)] = np.inf#And here too.
#The previous three lines are repeated in self.update_plots()
self.img1=self.ax[0].pcolormesh(self.x2,self.y2,array1,vmin=0,vmax=self.img_max,cmap='hot')
self.img2=self.ax[1].pcolormesh(self.x2,self.y2,array2,vmin=self.vmin,vmax=self.vmax,cmap='hot')
self.img3=self.ax[2].plot(self.x_axis,self.meanspec)
self.ax[2].set_xlim((min(self.x_axis),max(self.x_axis)))
self.ax[2].set_ylim(0,self.img_max)
#This trick to associate a single CB to multiple axes comes from
#https://stackoverflow.com/questions/13784201/matplotlib-2-subplots-1-colorbar
self.cbar = self.fig.colorbar(self.img2, ax=self.ax.ravel().tolist(),aspect = 20)
self.cbar = dcb.DraggableColorbar_fits(self.cbar,[self.img2],'hot')
self.cbar.connect()
#The rest is for dealing with the masking itself; the behaviour of the
#add/subtact buttons, the cursor and the saving of the masked columns.
self.col_active = ['coral','mistyrose']#The colours for the ADD and SUBTRACT buttons that
#can be activated.
self.col_passive = ['lightgrey','whitesmoke']#Colours when they are not active.
self.MW = 50#The default masking width.
self.addstatus = 0#The status for adding-to-mask mode starts as zero; i.e. it starts inactive.
self.substatus = 0#Same for subtraction.
self.list_of_polygons = []#This stores the polygons that are currently plotted. When initializing, none are plotted.
#However when proceeding through draw_masked_areas below, this variable could become populated if previously selected
#column were loaded from file.
self.multi = MultiCursor(self.fig.canvas, (self.ax[0],self.ax[1],self.ax[2]), color='g', lw=1, horizOn=False, vertOn=True)
self.multi.set_active(False)#The selection cursor starts deactivated as well, and is activated and deactivated
#further down as the buttons are pressed.
self.apply_to_all = False
self.apply_to_all_future = False
#The following show the z value of the plotted arrays in the statusbar,
#taken from https://matplotlib.org/examples/api/image_zcoord.html
numrows, numcols = self.order.shape
def format_coord_order(x, y):
col = int(x + 0.5)
row = int(y + 0.5)
if col >= 0 and col < numcols and row >= 0 and row < numrows:
z = self.order[row, col]
return 'x=%1.4f, y=%1.4f, z=%1.4f' % (x, y, z)
else:
return 'x=%1.4f, y=%1.4f' % (x, y)
def format_coord_res(x, y):
col = int(x + 0.5)
row = int(y + 0.5)
if col >= 0 and col < numcols and row >= 0 and row < numrows:
z = self.residual[row, col]
return 'x=%1.4f, y=%1.4f, z=%1.4f' % (x, y, z)
else:
return 'x=%1.4f, y=%1.4f' % (x, y)
self.ax[0].format_coord = format_coord_order
self.ax[1].format_coord = format_coord_res
self.draw_masked_areas()
def run_instance(p):
"""This runs the entire cross correlation analysis cascade."""
import numpy as np
from astropy.io import fits
import astropy.constants as const
import astropy.units as u
from matplotlib import pyplot as plt
import os.path
import scipy.interpolate as interp
import pylab
import pdb
import os.path
import os
import sys
import glob
import distutils.util
import pickle
import copy
from pathlib import Path
import tayph.util as ut
import tayph.operations as ops
import tayph.functions as fun
import tayph.system_parameters as sp
import tayph.tellurics as telcor
import tayph.masking as masking
import tayph.models as models
from tayph.ccf import xcor, clean_ccf, filter_ccf, construct_KpVsys
from tayph.vartests import typetest, notnegativetest, nantest, postest, typetest_array, dimtest
import tayph.shadow as shadow
# from lib import analysis
# from lib import cleaning
# from lib import masking as masking
# from lib import shadow as shadow
# from lib import molecfit as telcor
#First parse the parameter dictionary into required variables and test them.
typetest(p, dict, 'params in run_instance()')
dp = Path(p['dp'])
ut.check_path(dp, exists=True)
modellist = p['modellist']
templatelist = p['templatelist']
model_library = p['model_library']
template_library = p['template_library']
typetest(modellist, [str, list], 'modellist in run_instance()')
typetest(templatelist, [str, list], 'templatelist in run_instance()')
typetest(model_library, str, 'model_library in run_instance()')
typetest(template_library, str, 'template_library in run_instance()')
ut.check_path(model_library, exists=True)
ut.check_path(template_library, exists=True)
if type(modellist) == str:
modellist = [modellist] #Force this to be a list
if type(templatelist) == str:
templatelist = [templatelist] #Force this to be a list
typetest_array(modellist, str, 'modellist in run_instance()')
typetest_array(templatelist, str, 'modellist in run_instance()')
shadowname = p['shadowname']
maskname = p['maskname']
typetest(shadowname, str, 'shadowname in run_instance()')
typetest(maskname, str, 'shadowname in run_instance()')
RVrange = p['RVrange']
drv = p['drv']
f_w = p['f_w']
resolution = sp.paramget('resolution', dp)
typetest(RVrange, [int, float], 'RVrange in run_instance()')
typetest(drv, [int, float], 'drv in run_instance()')
typetest(f_w, [int, float], 'f_w in run_instance()')
typetest(resolution, [int, float], 'resolution in run_instance()')
nantest(RVrange, 'RVrange in run_instance()')
nantest(drv, 'drv in run_instance()')
nantest(f_w, 'f_w in run_instance()')
nantest(resolution, 'resolution in run_instance()')
postest(RVrange, 'RVrange in run_instance()')
postest(drv, 'drv in run_instance()')
postest(resolution, 'resolution in run_instance()')
notnegativetest(f_w, 'f_w in run_instance()')
do_colour_correction = p['do_colour_correction']
do_telluric_correction = p['do_telluric_correction']
do_xcor = p['do_xcor']
plot_xcor = p['plot_xcor']
make_mask = p['make_mask']
apply_mask = p['apply_mask']
c_subtract = p['c_subtract']
do_berv_correction = p['do_berv_correction']
do_keplerian_correction = p['do_keplerian_correction']
make_doppler_model = p['make_doppler_model']
skip_doppler_model = p['skip_doppler_model']
typetest(do_colour_correction, bool,
'do_colour_correction in run_instance()')
typetest(do_telluric_correction, bool,
'do_telluric_correction in run_instance()')
typetest(do_xcor, bool, 'do_xcor in run_instance()')
typetest(plot_xcor, bool, 'plot_xcor in run_instance()')
typetest(make_mask, bool, 'make_mask in run_instance()')
typetest(apply_mask, bool, 'apply_mask in run_instance()')
typetest(c_subtract, bool, 'c_subtract in run_instance()')
typetest(do_berv_correction, bool, 'do_berv_correction in run_instance()')
typetest(do_keplerian_correction, bool,
'do_keplerian_correction in run_instance()')
typetest(make_doppler_model, bool, 'make_doppler_model in run_instance()')
typetest(skip_doppler_model, bool, 'skip_doppler_model in run_instance()')
#We start by defining constants and preparing for generating output.
c = const.c.value / 1000.0 #in km/s
colourdeg = 3 #A fitting degree for the colour correction.
print(
f'---Passed parameter input tests. Initiating output folder tree in {Path("output")/dp}.'
)
libraryname = str(template_library).split('/')[-1]
if str(dp).split('/')[0] == 'data':
dataname = str(dp).replace('data/', '')
print(
f'------Data is located in data/ folder. Assuming output name for this dataset as {dataname}'
)
else:
dataname = dp
print(
f'------Data is NOT located in data/ folder. Assuming output name for this dataset as {dataname}'
)
list_of_wls = [] #This will store all the data.
list_of_orders = [] #All of it needs to be loaded into your memory.
list_of_sigmas = []
trigger2 = 0 #These triggers are used to limit the generation of output in the forloop.
trigger3 = 0
n_negative_total = 0 #This will hold the total number of pixels that were set to NaN because they were zero when reading in the data.
air = sp.paramget('air', dp) #Read bool from str in config file.
typetest(air, bool, 'air in run_instance()')
filelist_orders = [str(i) for i in Path(dp).glob('order_*.fits')]
if len(filelist_orders) == 0:
raise Exception(
f'Runtime error: No orders_*.fits files were found in {dp}.')
try:
order_numbers = [
int(i.split('order_')[1].split('.')[0]) for i in filelist_orders
]
except:
raise Exception(
'Runtime error: Failed casting fits filename numerals to ints. Are the filenames of the spectral orders correctly formatted?'
)
order_numbers.sort() #This is the ordered list of numerical order IDs.
n_orders = len(order_numbers)
if n_orders == 0:
raise Exception(
f'Runtime error: n_orders may not have ended up as zero. ({n_orders})'
)
#Loading the data from the datafolder.
if do_xcor == True or plot_xcor == True or make_mask == True:
print(f'---Loading orders from {dp}.')
# for i in range(startorder,endorder+1):
for i in order_numbers:
wavepath = dp / f'wave_{i}.fits'
orderpath = dp / f'order_{i}.fits'
sigmapath = dp / f'sigma_{i}.fits'
ut.check_path(wavepath, exists=True)
ut.check_path(orderpath, exists=True)
ut.check_path(sigmapath, exists=False)
wave_axis = fits.getdata(wavepath)
dimtest(wave_axis, [0], 'wavelength grid in run_instance()')
n_px = len(wave_axis) #Pixel width of the spectral order.
if air == False:
if i == np.min(order_numbers):
print("------Assuming wavelengths are in vaccuum.")
list_of_wls.append(1.0 * wave_axis)
else:
if i == np.min(order_numbers):
print("------Applying airtovac correction.")
list_of_wls.append(ops.airtovac(wave_axis))
order_i = fits.getdata(orderpath)
if i == np.min(order_numbers):
dimtest(
order_i, [0, n_px], f'order {i} in run_instance()'
) #For the first order, check that it is 2D and that is has a width equal to n_px.
n_exp = np.shape(
order_i
)[0] #then fix n_exp. All other orders should have the same n_exp.
print(f'------{n_exp} exposures recognised.')
else:
dimtest(order_i, [n_exp, n_px], f'order {i} in run_instance()')
#Now test for negatives, set them to NaN and track them.
n_negative = len(order_i[order_i <= 0])
if trigger3 == 0 and n_negative > 0:
print("------Setting negative values to NaN.")
trigger3 = -1
n_negative_total += n_negative
order_i[order_i <= 0] = np.nan
postest(order_i, f'order {i} in run_instance().'
) #make sure whatever comes out here is strictly positive.
list_of_orders.append(order_i)
try: #Try to get a sigma file. If it doesn't exist, we raise a warning. If it does, we test its dimensions and append it.
sigma_i = fits.getdata(sigmapath)
dimtest(sigma_i, [n_exp, n_px],
f'order {i} in run_instance().')
list_of_sigmas.append(sigma_i)
except FileNotFoundError:
if trigger2 == 0:
print(
'------WARNING: Sigma (flux error) files not provided. Assuming sigma = sqrt(flux). This is standard practise for HARPS data, but e.g. ESPRESSO has a pipeline that computes standard errors on each pixel for you.'
)
trigger2 = -1
list_of_sigmas.append(np.sqrt(order_i))
print(
f"------{n_negative_total} negative values set to NaN ({np.round(100.0*n_negative_total/n_exp/n_px/len(order_numbers),2)}% of total spectral pixels in dataset.)"
)
if len(list_of_orders) != n_orders:
raise Exception(
'Runtime error: n_orders is not equal to the length of list_of_orders. Something went wrong when reading them in?'
)
print('---Finished loading dataset to memory.')
#Apply telluric correction file or not.
# plt.plot(list_of_wls[60],list_of_orders[60][10],color='red')
# plt.plot(list_of_wls[60],list_of_orders[60][10]+list_of_sigmas[60][10],color='red',alpha=0.5)#plot corrected spectra
# plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='red',alpha=0.5)#plot SNR
if do_telluric_correction == True and n_orders > 0:
print('---Applying telluric correction')
telpath = dp / 'telluric_transmission_spectra.pkl'
list_of_orders, list_of_sigmas = telcor.apply_telluric_correction(
telpath, list_of_wls, list_of_orders, list_of_sigmas)
# plt.plot(list_of_wls[60],list_of_orders[60][10],color='blue')
# plt.plot(list_of_wls[60],list_of_orders[60][10]+list_of_sigmas[60][10],color='blue',alpha=0.5)#plot corrected spectra
# plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='blue',alpha=0.5) #plot SNR
# plt.show()
# pdb.set_trace()
#Do velocity correction of wl-solution. Explicitly after telluric correction
#but before masking. Because the cross-correlation relies on columns being masked.
#Then if you start to move the CCFs around before removing the time-average,
#each masked column becomes slanted. Bad deal.
rv_cor = 0
if do_berv_correction == True:
rv_cor += sp.berv(dp)
if do_keplerian_correction == True:
rv_cor -= sp.RV_star(dp) * (1.0)
if type(rv_cor) != int and len(list_of_orders) > 0:
print('---Reinterpolating data to correct velocities')
list_of_orders_cor = []
list_of_sigmas_cor = []
for i in range(len(list_of_wls)):
order = list_of_orders[i]
sigma = list_of_sigmas[i]
order_cor = order * 0.0
sigma_cor = sigma * 0.0
for j in range(len(list_of_orders[0])):
wl_i = interp.interp1d(list_of_wls[i],
order[j],
bounds_error=False)
si_i = interp.interp1d(list_of_wls[i],
sigma[j],
bounds_error=False)
wl_cor = list_of_wls[i] * (
1.0 - (rv_cor[j] * u.km / u.s / const.c)
) #The minus sign was tested on a slow-rotator.
order_cor[j] = wl_i(wl_cor)
sigma_cor[j] = si_i(
wl_cor
) #I checked that this works because it doesn't affect the SNR, apart from wavelength-shifting it.
list_of_orders_cor.append(order_cor)
list_of_sigmas_cor.append(sigma_cor)
ut.statusbar(i, fun.findgen(len(list_of_wls)))
# plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='blue')
# plt.plot(list_of_wls[60],list_of_orders_cor[60][10]/list_of_sigmas_cor[60][10],color='red')
# plt.show()
# sys.exit()
list_of_orders = list_of_orders_cor
list_of_sigmas = list_of_sigmas_cor
if len(list_of_orders) != n_orders:
raise RuntimeError(
'n_orders is no longer equal to the length of list_of_orders, though it was before. Something went wrong during telluric correction or velocity correction.'
)
#Compute / create a mask and save it to file (or not)
if make_mask == True and len(list_of_orders) > 0:
if do_colour_correction == True:
print(
'---Constructing mask with intra-order colour correction applied'
)
masking.mask_orders(list_of_wls,
ops.normalize_orders(list_of_orders,
list_of_sigmas,
colourdeg)[0],
dp,
maskname,
40.0,
5.0,
manual=True)
else:
print(
'---Constructing mask WITHOUT intra-order colour correction applied.'
)
print(
'---Switch on colour correction if you see colour variations in the 2D spectra.'
)
masking.mask_orders(list_of_wls,
list_of_orders,
dp,
maskname,
40.0,
5.0,
manual=True)
if apply_mask == False:
print(
'---WARNING in run_instance: Mask was made but is not applied to data (apply_mask == False)'
)
#Apply the mask that was previously created and saved to file.
if apply_mask == True:
print('---Applying mask')
list_of_orders = masking.apply_mask_from_file(dp, maskname,
list_of_orders)
list_of_sigmas = masking.apply_mask_from_file(dp, maskname,
list_of_sigmas)
#Interpolate over all isolated NaNs and set bad columns to NaN (so that they are ignored in the CCF)
if do_xcor == True:
print('---Healing NaNs')
list_of_orders = masking.interpolate_over_NaNs(
list_of_orders
) #THERE IS AN ISSUE HERE: INTERPOLATION SHOULD ALSO HAPPEN ON THE SIGMAS ARRAY!
list_of_sigmas = masking.interpolate_over_NaNs(list_of_sigmas)
#Normalize the orders to their average flux in order to effectively apply a broad-band colour correction (colour is typically a function of airmass and seeing).
if do_colour_correction == True:
print('---Normalizing orders to common flux level')
# plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='blue',alpha=0.4)
list_of_orders_normalised, list_of_sigmas_normalised, meanfluxes = ops.normalize_orders(
list_of_orders, list_of_sigmas, colourdeg
) #I tested that this works because it doesn't alter the SNR.
meanfluxes_norm = meanfluxes / np.nanmean(meanfluxes)
else:
meanfluxes_norm = fun.findgen(len(
list_of_orders[0])) * 0.0 + 1.0 #All unity.
# plt.plot(list_of_wls[60],list_of_orders_normalised[60][10]/list_of_sigmas[60][10],color='red',alpha=0.4)
# plt.show()
# sys.exit()
if len(list_of_orders) != n_orders:
raise RuntimeError(
'n_orders is no longer equal to the length of list_of_orders, though it was before. Something went wrong during masking or colour correction.'
)
#Construct the cross-correlation templates in case we will be computing or plotting the CCF.
if do_xcor == True or plot_xcor == True:
list_of_wlts = []
list_of_templates = []
outpaths = []
for templatename in templatelist:
print(f'---Building template {templatename}')
wlt, T = models.build_template(templatename,
binsize=0.5,
maxfrac=0.01,
resolution=resolution,
template_library=template_library,
c_subtract=c_subtract)
T *= (-1.0)
if np.mean(wlt) < 50.0: #This is likely in microns:
print(
'------WARNING: The loaded template has a mean wavelength less than 50.0, meaning that it is very likely not in nm, but in microns. I have divided by 1,000 now and hope for the best...'
)
wlt *= 1000.0
list_of_wlts.append(wlt)
list_of_templates.append(T)
outpath = Path('output') / Path(dataname) / Path(
libraryname) / Path(templatename)
if not os.path.exists(outpath):
print(
f"------The output location ({outpath}) didn't exist, I made it now."
)
os.makedirs(outpath)
outpaths.append(outpath)
#Perform the cross-correlation on the entire list of orders.
for i in range(len(list_of_wlts)):
templatename = templatelist[i]
wlt = list_of_wlts[i]
T = list_of_templates[i]
outpath = outpaths[i]
if do_xcor == True:
print(
f'---Cross-correlating spectra with template {templatename}.')
t1 = ut.start()
rv, ccf, ccf_e, Tsums = xcor(
list_of_wls,
list_of_orders_normalised,
np.flipud(np.flipud(wlt)),
T,
drv,
RVrange,
list_of_errors=list_of_sigmas_normalised)
ut.end(t1)
print(f'------Writing CCFs to {str(outpath)}')
ut.writefits(outpath / 'ccf.fits', ccf)
ut.writefits(outpath / 'ccf_e.fits', ccf_e)
ut.writefits(outpath / 'RV.fits', rv)
ut.writefits(outpath / 'Tsum.fits', Tsums)
else:
print(
f'---Reading CCFs with template {templatename} from {str(outpath)}.'
)
if os.path.isfile(outpath / 'ccf.fits') == False:
raise FileNotFoundError(
f'CCF output not located at {outpath}. Rerun with do_xcor=True to create these files?'
)
rv = fits.getdata(outpath / 'rv.fits')
ccf = fits.getdata(outpath / 'ccf.fits')
ccf_e = fits.getdata(outpath / 'ccf_e.fits')
Tsums = fits.getdata(outpath / 'Tsum.fits')
ccf_cor = ccf * 1.0
ccf_e_cor = ccf_e * 1.0
print('---Cleaning CCFs')
ccf_n, ccf_ne, ccf_nn, ccf_nne = clean_ccf(rv, ccf_cor, ccf_e_cor, dp)
if make_doppler_model == True and skip_doppler_model == False:
shadow.construct_doppler_model(rv,
ccf_nn,
dp,
shadowname,
xrange=[-200, 200],
Nxticks=20.0,
Nyticks=10.0)
make_doppler_model = False # This sets it to False after it's been run once, for the first template.
if skip_doppler_model == False:
print('---Reading doppler shadow model from ' + shadowname)
doppler_model, dsmask = shadow.read_shadow(
dp, shadowname, rv, ccf
) #This returns both the model evaluated on the rv,ccf grid, as well as the mask that blocks the planet trace.
ccf_clean, matched_ds_model = shadow.match_shadow(
rv, ccf_nn, dsmask, dp, doppler_model
) #THIS IS AN ADDITIVE CORRECTION, SO CCF_NNE DOES NOT NEED TO BE ALTERED AND IS STILL VALID VOOR CCF_CLEAN
else:
print('---Not performing shadow correction')
ccf_clean = ccf_nn * 1.0
matched_ds_model = ccf_clean * 0.0
if f_w > 0.0:
print('---Performing high-pass filter on the CCF')
ccf_clean_filtered, wiggles = filter_ccf(
rv, ccf_clean, v_width=f_w
) #THIS IS ALSO AN ADDITIVE CORRECTION, SO CCF_NNE IS STILL VALID.
else:
print('---Skipping high-pass filter')
ccf_clean_filtered = ccf_clean * 1.0
wiggles = ccf_clean * 0.0 #This filtering is additive so setting to zero is accurate.
print('---Weighing CCF rows by mean fluxes that were normalised out')
ccf_clean_weighted = np.transpose(
np.transpose(ccf_clean_filtered) * meanfluxes_norm
) #MULTIPLYING THE AVERAGE FLUXES BACK IN! NEED TO CHECK THAT THIS ALSO GOES PROPERLY WITH THE ERRORS!
ccf_nne = np.transpose(np.transpose(ccf_nne) * meanfluxes_norm)
ut.save_stack(outpath / 'cleaning_steps.fits', [
ccf, ccf_cor, ccf_nn, ccf_clean, matched_ds_model,
ccf_clean_filtered, wiggles, ccf_clean_weighted
])
ut.writefits(outpath / 'ccf_cleaned.fits', ccf_clean_weighted)
ut.writefits(outpath / 'ccf_cleaned_error.fits', ccf_nne)
print('---Constructing KpVsys')
Kp, KpVsys, KpVsys_e = construct_KpVsys(rv, ccf_clean_weighted,
ccf_nne, dp)
ut.writefits(outpath / 'KpVsys.fits', KpVsys)
ut.writefits(outpath / 'KpVsys_e.fits', KpVsys_e)
ut.writefits(outpath / 'Kp.fits', Kp)
return
sys.exit()
if plot_xcor == True and inject_model == False:
print('---Plotting KpVsys')
analysis.plot_KpVsys(rv, Kp, KpVsys, dp)
#Now repeat it all for the model injection.
if inject_model == True:
for modelname in modellist:
outpath_i = outpath + modelname + '/'
if do_xcor == True:
print('---Injecting model ' + modelname)
list_of_orders_injected = models.inject_model(
list_of_wls,
list_of_orders,
dp,
modelname,
model_library=model_library
) #Start with the unnormalised orders from before.
#Normalize the orders to their average flux in order to effectively apply
#a broad-band colour correction (colour is a function of airmass and seeing).
if do_colour_correction == True:
print(
'------Normalizing injected orders to common flux level'
)
list_of_orders_injected, list_of_sigmas_injected, meanfluxes_injected = ops.normalize_orders(
list_of_orders_injected, list_of_sigmas, colourdeg)
meanfluxes_norm_injected = meanfluxes_injected / np.mean(
meanfluxes_injected)
else:
meanfluxes_norm_injected = fun.findgen(
len(list_of_orders_injected[0])
) * 0.0 + 1.0 #All unity.
print('------Cross-correlating injected orders')
rv_i, ccf_i, ccf_e_i, Tsums_i = analysis.xcor(
list_of_wls,
list_of_orders_injected,
np.flipud(np.flipud(wlt)),
T,
drv,
RVrange,
list_of_errors=list_of_sigmas_injected)
print('------Writing injected CCFs to ' + outpath_i)
if not os.path.exists(outpath_i):
print("---------That path didn't exist, I made it now.")
os.makedirs(outpath_i)
ut.writefits(outpath_i + '/' + 'ccf_i_' + modelname + '.fits',
ccf_i)
ut.writefits(
outpath_i + '/' + 'ccf_e_i_' + modelname + '.fits',
ccf_e_i)
else:
print('---Reading injected CCFs from ' + outpath_i)
if os.path.isfile(outpath_i + 'ccf_i_' + modelname +
'.fits') == False:
print('------ERROR: Injected CCF not located at ' +
outpath_i + 'ccf_i_' + modelname + '.fits' +
'. Set do_xcor and inject_model to True?')
sys.exit()
if os.path.isfile(outpath_i + 'ccf_e_i_' + modelname +
'.fits') == False:
print('------ERROR: Injected CCF error not located at ' +
outpath_i + 'ccf_e_i_' + modelname + '.fits' +
'. Set do_xcor and inject_model to True?')
sys.exit()
# f.close()
# f2.close()
ccf_i = fits.getdata(outpath_i + 'ccf_i_' + modelname +
'.fits')
ccf_e_i = fits.getdata(outpath_i + 'ccf_e_i_' + modelname +
'.fits')
print('---Cleaning injected CCFs')
ccf_n_i, ccf_ne_i, ccf_nn_i, ccf_nne_i = cleaning.clean_ccf(
rv, ccf_i, ccf_e_i, dp)
ut.writefits(outpath_i + 'ccf_normalized_i.fits', ccf_nn_i)
ut.writefits(outpath_i + 'ccf_ne_i.fits', ccf_ne_i)
# if make_doppler_model == True and skip_doppler_model == False:
# shadow.construct_doppler_model(rv,ccf_nn,dp,shadowname,xrange=[-200,200],Nxticks=20.0,Nyticks=10.0)
if skip_doppler_model == False:
# print('---Reading doppler shadow model from '+shadowname)
# doppler_model,maskHW = shadow.read_shadow(dp,shadowname,rv,ccf)
ccf_clean_i, matched_ds_model_i = shadow.match_shadow(
rv, ccf_nn_i, dp, doppler_model, maskHW)
else:
print(
'---Not performing shadow correction on injected spectra either.'
)
ccf_clean_i = ccf_nn_i * 1.0
matched_ds_model_i = ccf_clean_i * 0.0
if f_w > 0.0:
ccf_clean_i_filtered, wiggles_i = cleaning.filter_ccf(
rv, ccf_clean_i, v_width=f_w)
else:
ccf_clean_i_filtered = ccf_clean_i * 1.0
ut.writefits(outpath_i + 'ccf_cleaned_i.fits',
ccf_clean_i_filtered)
ut.writefits(outpath + 'ccf_cleaned_i_error.fits', ccf_nne)
print(
'---Weighing injected CCF rows by mean fluxes that were normalised out'
)
ccf_clean_i_filtered = np.transpose(
np.transpose(ccf_clean_i_filtered) * meanfluxes_norm_injected
) #MULTIPLYING THE AVERAGE FLUXES BACK IN! NEED TO CHECK THAT THIS ALSO GOES PROPERLY WITH THE ERRORS!
ccf_nne_i = np.transpose(
np.transpose(ccf_nne_i) * meanfluxes_norm_injected)
print('---Constructing injected KpVsys')
Kp, KpVsys_i, KpVsys_e_i = analysis.construct_KpVsys(
rv, ccf_clean_i_filtered, ccf_nne_i, dp)
ut.writefits(outpath_i + 'KpVsys_i.fits', KpVsys_i)
# ut.writefits(outpath+'KpVsys_e_i.fits',KpVsys_e_i)
if plot_xcor == True:
print('---Plotting KpVsys with ' + modelname + ' injected.')
analysis.plot_KpVsys(rv, Kp, KpVsys, dp, injected=KpVsys_i)
def build_template(templatename,
binsize=1.0,
maxfrac=0.01,
mode='top',
resolution=0.0,
c_subtract=True,
twopass=False,
template_library='models/library'):
"""This routine reads a specified model from the library and turns it into a
cross-correlation template by subtracting the top-envelope (or bottom envelope),
if c_subtract is set to True."""
import tayph.util as ut
from tayph.vartests import typetest, postest, notnegativetest
import numpy as np
import tayph.operations as ops
import astropy.constants as const
from astropy.io import fits
from matplotlib import pyplot as plt
from scipy import interpolate
from pathlib import Path
typetest(templatename, str, 'templatename mod.build_template()')
typetest(binsize, [int, float], 'binsize mod.build_template()')
typetest(maxfrac, [int, float], 'maxfrac mod.build_template()')
typetest(
mode,
str,
'mode mod.build_template()',
)
typetest(resolution, [int, float], 'resolution in mod.build_template()')
typetest(twopass, bool, 'twopass in mod.build_template()')
binsize = float(binsize)
maxfrac = float(maxfrac)
resolution = float(resolution)
postest(binsize, 'binsize mod.build_template()')
postest(maxfrac, 'maxfrac mod.build_template()')
notnegativetest(resolution, 'resolution in mod.build_template()')
template_library = ut.check_path(template_library, exists=True)
c = const.c.to('km/s').value
if mode not in ['top', 'bottom']:
raise Exception(
f'RuntimeError in build_template: Mode should be set to "top" or "bottom" ({mode}).'
)
wlt, fxt = get_model(templatename, library=template_library)
if wlt[-1] <= wlt[0]: #Reverse the wl axis if its sorted the wrong way.
wlt = np.flipud(wlt)
fxt = np.flipud(fxt)
if c_subtract == True:
wle, fxe = ops.envelope(
wlt, fxt - np.median(fxt), binsize, selfrac=maxfrac,
mode=mode) #These are binpoints of the top-envelope.
#The median of fxm is first removed to decrease numerical errors, because the spectrum may
#have values that are large (~1.0) while the variations are small (~1e-5).
e_i = interpolate.interp1d(wle, fxe, fill_value='extrapolate')
envelope = e_i(wlt)
T = fxt - np.median(fxt) - envelope
absT = np.abs(T)
T[(absT < 1e-4 * np.max(absT)
)] = 0.0 #This is now continuum-subtracted and binary-like.
#Any values that are small are taken out.
#This therefore assumes that the model has lines that are deep compared to the numerical
#error of envelope subtraction (!).
else:
T = fxt * 1.0
if resolution != 0.0:
dRV = c / resolution
print(
f'------Blurring template to resolution of data ({round(resolution,0)}, {round(dRV,2)} km/s)'
)
wlt_cv, T_cv, vstep = ops.constant_velocity_wl_grid(wlt,
T,
oversampling=2.0)
print(f'---------v_step is {np.round(vstep,3)} km/s')
print(
f'---------So the resolution blurkernel has an avg width of {np.round(dRV/vstep,3)} px.'
)
T_b = ops.smooth(T_cv, dRV / vstep, mode='gaussian')
wlt = wlt_cv * 1.0
T = T_b * 1.0
return (wlt, T)
def apply_telluric_correction(inpath, list_of_wls, list_of_orders,
list_of_sigmas):
"""
This applies a set of telluric spectra (computed by molecfit) for each exposure
in our time series that were written to a pickle file by write_telluric_transmission_to_file.
List of errors are provided to propagate telluric correction into the error array as well.
Parameters
----------
inpath : str, path like
The path to the pickled transmission spectra.
list_of_wls : list
List of wavelength axes.
list_of_orders :
List of 2D spectral orders, matching to the wavelength axes in dimensions and in number.
list_of_isgmas :
List of 2D error matrices, matching dimensions and number of list_of_orders.
Returns
-------
list_of_orders_corrected : list
List of 2D spectral orders, telluric corrected.
list_of_sigmas_corrected : list
List of 2D error matrices, telluric corrected.
"""
import scipy.interpolate as interp
import numpy as np
import tayph.util as ut
import tayph.functions as fun
from tayph.vartests import dimtest, postest, typetest, nantest
wlT, fxT = read_telluric_transmission_from_file(inpath)
typetest(list_of_wls, list, 'list_of_wls in apply_telluric_correction()')
typetest(list_of_orders, list,
'list_of_orders in apply_telluric_correction()')
typetest(list_of_sigmas, list,
'list_of_sigmas in apply_telluric_correction()')
typetest(wlT, list,
'list of telluric wave-axes in apply_telluric_correction()')
typetest(
fxT, list,
'list of telluric transmission spectra in apply_telluric_correction()')
No = len(list_of_wls)
x = fun.findgen(No)
if No != len(list_of_orders):
raise Exception(
'Runtime error in telluric correction: List of data wls and List of orders do not have the same length.'
)
Nexp = len(wlT)
if Nexp != len(fxT):
raise Exception(
'Runtime error in telluric correction: List of telluric wls and telluric spectra read from file do not have the same length.'
)
if Nexp != len(list_of_orders[0]):
raise Exception(
f'Runtime error in telluric correction: List of telluric spectra and data spectra read from file do not have the same length ({Nexp} vs {len(list_of_orders[0])}).'
)
list_of_orders_cor = []
list_of_sigmas_cor = []
# ut.save_stack('test.fits',list_of_orders)
# pdb.set_trace()
for i in range(No): #Do the correction order by order.
order = list_of_orders[i]
order_cor = order * 0.0
error = list_of_sigmas[i]
error_cor = error * 0.0
wl = list_of_wls[i]
dimtest(order, [0, len(wl)],
f'order {i}/{No} in apply_telluric_correction()')
dimtest(error, np.shape(order),
f'errors {i}/{No} in apply_telluric_correction()')
for j in range(Nexp):
T_i = interp.interp1d(wlT[j], fxT[j], fill_value="extrapolate")(wl)
postest(T_i,
f'T-spec of exposure {j} in apply_telluric_correction()')
nantest(T_i,
f'T-spec of exposure {j} in apply_telluric_correction()')
order_cor[j] = order[j] / T_i
error_cor[j] = error[
j] / T_i #I checked that this works because the SNR before and after telluric correction is identical.
list_of_orders_cor.append(order_cor)
list_of_sigmas_cor.append(error_cor)
ut.statusbar(i, x)
return (list_of_orders_cor, list_of_sigmas_cor)
def apply_telluric_correction(inpath, list_of_wls, list_of_orders,
list_of_sigmas):
"""
This applies a set of telluric spectra (computed by molecfit) for each exposure
in our time series that were written to a pickle file by write_telluric_transmission_to_file.
List of errors are provided to propagate telluric correction into the error array as well.
Parameters
----------
inpath : str, path like
The path to the pickled transmission spectra.
list_of_wls : list
List of wavelength axes.
list_of_orders :
List of 2D spectral orders, matching to the wavelength axes in dimensions and in number.
list_of_isgmas :
List of 2D error matrices, matching dimensions and number of list_of_orders.
Returns
-------
list_of_orders_corrected : list
List of 2D spectral orders, telluric corrected.
list_of_sigmas_corrected : list
List of 2D error matrices, telluric corrected.
"""
import scipy.interpolate as interp
import numpy as np
import tayph.util as ut
import tayph.functions as fun
from tayph.vartests import dimtest, postest, typetest, nantest
import copy
T = read_telluric_transmission_from_file(inpath)
wlT = T[0]
fxT = T[1]
typetest(list_of_wls, list, 'list_of_wls in apply_telluric_correction()')
typetest(list_of_orders, list,
'list_of_orders in apply_telluric_correction()')
typetest(list_of_sigmas, list,
'list_of_sigmas in apply_telluric_correction()')
typetest(wlT, list,
'list of telluric wave-axes in apply_telluric_correction()')
typetest(
fxT, list,
'list of telluric transmission spectra in apply_telluric_correction()')
No = len(list_of_wls) #Number of orders.
x = fun.findgen(No)
Nexp = len(wlT)
#Test dimensions
if No != len(list_of_orders):
raise Exception(
'Runtime error in telluric correction: List of wavelength axes and List '
'of orders do not have the same length.')
if Nexp != len(fxT):
raise Exception(
'Runtime error in telluric correction: List of telluric wls and telluric '
'spectra read from file do not have the same length.')
if Nexp != len(list_of_orders[0]):
raise Exception(
f'Runtime error in telluric correction: List of telluric spectra and data'
f'spectra read from file do not have the same length ({Nexp} vs {len(list_of_orders[0])}).'
)
#Corrected orders will be stored here.
list_of_orders_cor = []
list_of_sigmas_cor = []
#Do the correction order by order:
for i in range(No):
order = list_of_orders[i]
order_cor = order * 0.0
error = list_of_sigmas[i]
error_cor = error * 0.0
wl = copy.deepcopy(list_of_wls[i]) #input wl axis, either 1D or 2D.
#If it is 1D, we make it 2D by tiling it vertically:
if wl.ndim == 1: wl = np.tile(wl, (Nexp, 1)) #Tile it into a 2D thing.
#If read (2D) or tiled (1D) correctly, wl and order should have the same shape:
dimtest(wl, np.shape(order),
f'Wl axis of order {i}/{No} in apply_telluric_correction()')
dimtest(error, np.shape(order),
f'errors {i}/{No} in apply_telluric_correction()')
for j in range(Nexp):
T_i = interp.interp1d(wlT[j], fxT[j],
fill_value="extrapolate")(wl[j])
postest(T_i,
f'T-spec of exposure {j} in apply_telluric_correction()')
nantest(T_i,
f'T-spec of exposure {j} in apply_telluric_correction()')
order_cor[j] = order[j] / T_i
error_cor[j] = error[
j] / T_i #I checked that this works because the SNR before and after
#telluric correction is identical.
list_of_orders_cor.append(order_cor)
list_of_sigmas_cor.append(error_cor)
ut.statusbar(i, x)
return (list_of_orders_cor, list_of_sigmas_cor)
def normalize_orders(list_of_orders, list_of_sigmas, deg=1, nsigma=4):
"""
If deg is set to 1, this function will normalise based on the mean flux in each order.
If set higher, it will remove the average spectrum in each order and fit a polynomial
to the residual. This means that in the presence of spectral lines, the fluxes will be
slightly lower than if def=1 is used. nsigma is only used if deg > 1, and is used to
throw away outliers from the polynomial fit. The program also computes the total
mean flux of each exposure in the time series - totalled over all orders. These
are important to correctly weigh the cross-correlation functions later. The
inter-order colour correction is assumed to be an insignificant modification to
these weights.
Parameters
----------
list_of_orders : list
The list of 2D orders that need to be normalised.
list_of_sigmas : list
The list of 2D error matrices corresponding to the 2D orders that need to be normalised.
deg : int
The polynomial degree to remove. If set to 1, only the average flux is removed. If higher,
polynomial fits are made to the residuals after removal of the average spectrum.
nsigma : int, float
The number of sigmas beyond which outliers are rejected from the polynomial fit.
Only used when deg > 1.
Returns
-------
out_list_of_orders : list
The normalised 2D orders.
out_list_of_sigmas : list
The corresponding errors.
meanfluxes : np.array
The mean flux of each exposure in the time series, averaged over all orders.
"""
import numpy as np
import tayph.functions as fun
from tayph.vartests import dimtest, postest, typetest
import warnings
typetest(list_of_orders, list, 'list_of_orders in ops.normalize_orders()')
typetest(list_of_sigmas, list, 'list_of_sigmas in ops.normalize_orders()')
dimtest(list_of_orders[0], [0, 0]) #Test that the first order is 2D.
dimtest(list_of_sigmas[0],
[0, 0]) #And that the corresponding sigma array is, as well.
n_exp = np.shape(list_of_orders[0])[0] #Get the number of exposures.
for i in range(len(list_of_orders)): #Should be the same for all orders.
dimtest(list_of_orders[i], [n_exp, 0])
dimtest(list_of_sigmas[i], np.shape(list_of_orders[i]))
typetest(deg, int, 'degree in ops.normalize_orders()')
typetest(nsigma, [int, float], 'nsigma in ops.normalize_orders()')
postest(deg, 'degree in ops.normalize_orders()')
postest(nsigma, 'degree in ops.normalize_orders()')
N = len(list_of_orders)
out_list_of_orders = []
out_list_of_sigmas = []
#First compute the exposure-to-exposure flux variations to be used as weights.
meanfluxes = fun.findgen(n_exp) * 0.0
N_i = 0
for i in range(N):
m = np.nanmedian(list_of_orders[i], axis=1) #Median or mean?
if np.sum(np.isnan(m)) > 0:
print(
'---Warning in normalise_orders: Skipping order %s because many nans are present.'
% i)
else:
N_i += 1
meanfluxes += m #These contain the exposure-to-exposure variability of the time-series.
meanfluxes /= N_i #These are the weights.
if deg == 1:
for i in range(N):
#What I'm doing here is probably stupid and numpy division will probably work just fine without
#IDL-relics.
n_px = np.shape(list_of_orders[i])[1]
meanflux = np.nanmedian(
list_of_orders[i],
axis=1) #Average flux in each order. Median or mean?
meanblock = fun.rebinreform(
meanflux / np.nanmean(meanflux), n_px
).T #This is a slow operation. Row-by-row division is better done using a double-transpose...
out_list_of_orders.append(list_of_orders[i] / meanblock)
out_list_of_sigmas.append(list_of_sigmas[i] / meanblock)
else:
for i in range(N):
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
meanspec = np.nanmean(list_of_orders[i],
axis=0) #Average spectrum in each order.
x = np.array(range(len(meanspec)))
poly_block = list_of_orders[
i] * 0.0 #Array that will host the polynomial fits.
colour = list_of_orders[
i] / meanspec #What if there are zeroes? I.e. padding around the edges of the order?
for j, s in enumerate(list_of_orders[i]):
idx = np.isfinite(colour[j])
if np.sum(idx) > 0:
p = np.poly1d(np.polyfit(x[idx], colour[j][idx], deg))(
x) #Polynomial fit to the colour variation.
res = colour[
j] / p - 1.0 #The residual, which is flat around zero if it's a good fit. This has all sorts of line residuals that we need to throw out.
#We do that using the weight keyword of polyfit, and just set all those weights to zero.
sigma = np.nanstd(res)
w = x * 0.0 + 1.0 #Start with a weight function that is 1.0 everywhere.
with warnings.catch_warnings():
warnings.simplefilter("ignore",
category=RuntimeWarning)
w[np.abs(res) > nsigma * sigma] = 0.0
w = x * 0.0 + 1.0 #Start with a weight function that is 1.0 everywhere.
p2 = np.poly1d(
np.polyfit(x[idx], colour[j][idx], deg, w=w[idx])
)(
x
) #Second, weighted polynomial fit to the colour variation.
poly_block[j] = p2
out_list_of_orders.append(list_of_orders[i] / poly_block)
out_list_of_sigmas.append(list_of_sigmas[i] / poly_block)
return (out_list_of_orders, out_list_of_sigmas, meanfluxes)
def xcor(list_of_wls,
list_of_orders,
wlm,
fxm,
drv,
RVrange,
plot=False,
list_of_errors=None):
"""
This routine takes a combined dataset (in the form of lists of wl spaces,
spectral orders and possible a matching list of errors on those spectal orders),
as well as a template (wlm,fxm) to cross-correlate with, and the cross-correlation
parameters (drv,RVrange). The code takes on the order of ~10 minutes for an entire
HARPS dataset, which appears to be superior to my old IDL pipe.
The CCF used is the Geneva-style weighted average; not the Pearson CCF. Therefore
it measures true 'average' planet lines, with flux on the y-axis of the CCF.
The template must therefore be (something close to) a binary mask, with values
inside spectral lines (the CCF is scale-invariant so their overall scaling
doesn't matter),
It returns the RV axis and the resulting CCF in a tuple.
Thanks to Brett Morris (bmorris3), this code now implements a clever numpy broadcasting trick to
instantly apply and interpolate the wavelength shifts of the model template onto
the data grid in 2 dimensions. The matrix multiplication operator (originally
recommended to me by Matteo Brogi) allowed this 2D template matrix to be multiplied
with a 2D spectral order. np.hstack() is used to concatenate all orders end to end,
effectively making a giant single spectral order (with NaNs in between due to masking).
All these steps have eliminated ALL the forloops from the equation, and effectuated a
speed gain of a factor between 2,000 and 3,000. The time to do cross correlations is now
typically measured in 100s of milliseconds rather than minutes.
This way of calculation does impose some strict rules on NaNs, though. To keep things fast,
NaNs are now used to set the interpolated template matrix to zero wherever there are NaNs in the data.
These NaNs are found by looking at the first spectrum in the stack, with the assumption that
every NaN is in an all-NaN column. In the standard cross-correlation work-flow, isolated
NaNs are interpolated over (healed), after all.
The places where there are NaN columns in the data are therefore set to 0 in the template matrix.
The NaN values themselves are then set to to an arbitrary value, since they will never
weigh into the average by construction.
Parameters
----------
list_of_wls : list
List of wavelength axes of the data.
list_of_orders : list
List of corresponding 2D orders.
list_of_errors : list
Optional, list of corresponding 2D error matrices.
wlm : np.ndarray
Wavelength axis of the template.
fxm : np.ndarray
Weight-axis of the template.
drv : int,float
The velocity step onto which the CCF is computed. Typically ~1 km/s.
RVrange : int,float
The velocity range in the positive and negative direction over which to
evaluate the CCF. Typically >100 km/s.
plot : bool
Set to True for diagnostic plotting.
Returns
-------
RV : np.ndarray
The radial velocity grid over which the CCF is evaluated.
CCF : np.ndarray
The weighted average flux in the spectrum as a function of radial velocity.
CCF_E : np.ndarray
Optional. The error on each CCF point propagated from the error on the spectral values.
Tsums : np.ndarray
The sum of the template for each velocity step. Used for normalising the CCFs.
"""
import tayph.functions as fun
import astropy.constants as const
import tayph.util as ut
from tayph.vartests import typetest, dimtest, postest, nantest
import numpy as np
import scipy.interpolate
import astropy.io.fits as fits
import matplotlib.pyplot as plt
import sys
import pdb
#===FIRST ALL SORTS OF TESTS ON THE INPUT===
if len(list_of_wls) != len(list_of_orders):
raise ValueError(
f'In xcor(): List of wls and list of orders have different length ({len(list_of_wls)} & {len(list_of_orders)}).'
)
dimtest(wlm, [len(fxm)], 'wlm in ccf.xcor()')
typetest(wlm, np.ndarray, 'wlm in ccf.xcor')
typetest(fxm, np.ndarray, 'fxm in ccf.xcor')
typetest(drv, [int, float], 'drv in ccf.xcor')
typetest(
RVrange,
float,
'RVrange in ccf.xcor()',
)
postest(RVrange, 'RVrange in ccf.xcor()')
postest(drv, 'drv in ccf.xcor()')
nantest(wlm, 'fxm in ccf.xcor()')
nantest(fxm, 'fxm in ccf.xcor()')
nantest(drv, 'drv in ccf.xcor()')
nantest(RVrange, 'RVrange in ccf.xcor()')
drv = float(drv)
N = len(list_of_wls) #Number of orders.
if np.ndim(list_of_orders[0]) == 1.0:
n_exp = 1
else:
n_exp = len(list_of_orders[0][:, 0]) #Number of exposures.
#===Then check that all orders indeed have n_exp exposures===
for i in range(N):
if len(list_of_orders[i][:, 0]) != n_exp:
raise ValueError(
f'In xcor(): Not all orders have {n_exp} exposures.')
#===END OF TESTS. NOW DEFINE CONSTANTS===
c = const.c.to('km/s').value #In km/s
RV = fun.findgen(
2.0 * RVrange / drv +
1) * drv - RVrange #..... CONTINUE TO DEFINE THE VELOCITY GRID
beta = 1.0 - RV / c #The doppler factor with which each wavelength is to be shifted.
n_rv = len(RV)
#===STACK THE ORDERS IN MASSIVE CONCATENATIONS===
stack_of_orders = np.hstack(list_of_orders)
stack_of_wls = np.concatenate(list_of_wls)
if list_of_errors is not None:
stack_of_errors = np.hstack(list_of_errors) #Stack them horizontally
#Check that the number of NaNs is the same in the orders as in the errors on the orders;
#and that they are in the same place; meaning that if I add the errors to the orders, the number of
#NaNs does not increase (NaN+value=NaN).
if (np.sum(np.isnan(stack_of_orders)) != np.sum(
np.isnan(stack_of_errors + stack_of_orders))) and (np.sum(
np.isnan(stack_of_orders)) != np.sum(
np.isnan(stack_of_errors))):
raise ValueError(
f"in CCF: The number of NaNs in list_of_orders and list_of_errors is not equal ({np.sum(np.isnan(list_of_orders))},{np.sum(np.isnan(list_of_errors))})"
)
#===HERE IS THE JUICY BIT===
shifted_wavelengths = stack_of_wls * beta[:, np.
newaxis] #2D broadcast of wl_data, each row shifted by beta[i].
T = scipy.interpolate.interp1d(wlm, fxm, bounds_error=False, fill_value=0)(
shifted_wavelengths) #...making this a 2D thing.
T[:, np.isnan(
stack_of_orders[0]
)] = 0.0 #All NaNs are assumed to be in all-NaN columns. If that is not true, the below nantest will fail.
T_sums = np.sum(T, axis=1)
#We check whether there are isolated NaNs:
n_nans = np.sum(np.isnan(stack_of_orders),
axis=0) #This is the total number of NaNs in each column.
n_nans[n_nans == len(
stack_of_orders
)] = 0 #Whenever the number of NaNs equals the length of a column, set the flag to zero.
if np.max(
n_nans
) > 0: #If there are any columns which still have NaNs in them, we need to crash.
raise ValueError(
f"in CCF: Not all NaN values are purely in columns. There are still isolated NaNs. Remove those."
)
stack_of_orders[np.isnan(
stack_of_orders)] = 47e20 #Set NaNs to arbitrarily high values.
CCF = stack_of_orders @ T.T / T_sums #Here it the entire cross-correlation. Over all orders and velocity steps. No forloops.
CCF_E = CCF * 0.0
#If the errors were provided, we do the same to those:
if list_of_errors is not None:
stack_of_errors[np.isnan(
stack_of_errors
)] = 42e20 #we have already tested that these NaNs are in the same place.
CCF_E = stack_of_errors**2 @ (
T.T / T_sums)**2 #This has been mathematically proven.
#===THAT'S ALL. TEST INTEGRITY AND RETURN THE RESULT===
nantest(
CCF, 'CCF in ccf.xcor()'
) #If anything went wrong with NaNs in the data, these tests will fail because the matrix operation @ is non NaN-friendly.
nantest(CCF_E, 'CCF_E in ccf.xcor()')
if list_of_errors != None:
return (RV, CCF, np.sqrt(CCF_E), T_sums)
return (RV, CCF, T_sums)
def convolve(array, kernel, edge_degree=1, fit_width=2):
"""It's unbelievable, but I could not find the python equivalent of IDL's
/edge_truncate keyword, which truncates the kernel at the edge of the convolution.
Therefore, unfortunately, I need to code a convolution operation myself.
Stand by to be slowed down by an order of magnitude #thankspython.
Nope! Because I can just use np.convolve for most of the array, just not the edge...
So the strategy is to extrapolate the edge of the array using a polynomial fit
to the edge elements. By default, I fit over a range that is twice the length of the kernel; but
this value can be modified using the fit_width parameter.
Parameters
----------
array : list, np.ndarray
The horizontal axis.
kernel : list, np.ndarray
The convolution kernel. It is required to have a length that is less than 25% of the size of the array.
edge_degree : int
The polynomial degree by which the array is extrapolated in order to
fit_width : int
The length of the area at the edges of array used to fit the polynomial, in units of the length of the kernel.
Increase this number for small kernels or noisy arrays.
Returns
-------
array_convolved : np.array
The input array convolved with the kernel
Example
-------
>>> import numpy as np
>>> a=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
>>> b=[-0.5,0,0.5]
>>> c=convolve(a,b,edge_degree=1)
"""
import numpy as np
import tayph.functions as fun
from tayph.vartests import typetest, postest, dimtest
typetest(edge_degree, int, 'edge_degree in ops.convolve()')
typetest(fit_width, int, 'edge_degree in ops.convolve()')
typetest(array, [list, np.ndarray])
typetest(kernel, [list, np.ndarray])
dimtest(array, [0], 'array in ops.convolve()')
dimtest(kernel, [0], 'array in ops.convolve()')
postest(edge_degree, 'edge_degree in ops.convolve()')
postest(fit_width, 'edge_degree in ops.convolve()')
array = np.array(array)
kernel = np.array(kernel)
if len(kernel) >= len(array) / 4.0:
raise Exception(
f"Error in ops.convolve(): Kernel length is larger than a quarter of the array ({len(kernel)}, {len(array)}). Can't extrapolate over that length. And you probably don't want to be doing a convolution like that, anyway."
)
if len(kernel) % 2 != 1:
raise Exception(
'Error in ops.convolve(): Kernel needs to have an odd number of elements.'
)
#Perform polynomial fits at the edges.
x = fun.findgen(len(array))
fit_left = np.polyfit(x[0:len(kernel) * 2], array[0:len(kernel) * 2],
edge_degree)
fit_right = np.polyfit(x[-2 * len(kernel) - 1:-1],
array[-2 * len(kernel) - 1:-1], edge_degree)
#Pad both the x-grid (onto which the polynomial is defined)
#and the data array.
pad = fun.findgen((len(kernel) - 1) / 2)
left_pad = pad - (len(kernel) - 1) / 2
right_pad = np.max(x) + pad + 1
left_array_pad = np.polyval(fit_left, left_pad)
right_array_pad = np.polyval(fit_right, right_pad)
#Perform the padding.
x_padded = np.append(left_pad, x)
x_padded = np.append(
x_padded, right_pad
) #Pad the array with the missing elements of the kernel at the edge.
array_padded = np.append(left_array_pad, array)
array_padded = np.append(array_padded, right_array_pad)
#Reverse the kernel because np.convol does that automatically and I don't want that.
#(Imagine doing a derivative with a kernel [-1,0,1] and it gets reversed...)
kr = kernel[::-1]
#The valid keyword effectively undoes the padding, leaving only those values for which the kernel was entirely in the padded array.
#This thus again has length equal to len(array).
return np.convolve(array_padded, kr, 'valid')
def smooth(fx, w, mode='box', edge_degree=1):
"""
This function takes a spectrum, and blurs it using either a
Gaussian kernel or a box kernel, which have a FWHM width of w px everywhere.
Meaning that the width changes dynamically on a constant d-lambda grid.
Set the mode to gaussian or box. Because in box, care is taken to correctly
interpolate the edges, it is about twice slower than the Gaussian.
This interpolation is done manually in the fun.box function.
"""
import numpy as np
import tayph.util as ut
from tayph.vartests import typetest, postest
import tayph.functions as fun
from matplotlib import pyplot as plt
typetest(w, [int, float], 'w in ops.smooth()')
typetest(fx, np.ndarray, 'fx in ops.smooth()')
typetest(mode, str, 'mode in ops.smooth()')
typetest(edge_degree, int, 'edge_degree in ops.smooth()')
postest(w, 'w in ops.smooth()')
if mode not in ['box', 'gaussian']:
raise Exception(
f'RuntimeError in ops.smooth(): Mode should be set to "top" or "bottom" ({mode}).'
)
truncsize = 4.0 #The gaussian is truncated at 4 sigma.
shape = np.shape(fx)
sig_w = w / 2 * np.sqrt(
2.0 * np.log(2)) #Transform FWHM to Gaussian sigma. In km/s.
trunc_dist = np.round(sig_w * truncsize).astype(int)
#First define the kernel.
kw = int(np.round(truncsize * sig_w * 2.0))
if kw % 2.0 != 1.0: #This is to make sure that the kernel has an odd number of
#elements, and that it is symmetric around zero.
kw += 1
kx = fun.findgen(kw)
kx -= np.mean(
kx) #This must be centered around zero. Doing a hardcoded check:
if (-1.0) * kx[-1] != kx[0]:
print(kx)
raise Exception(
"ERROR in box_smooth: Kernel could not be made symmetric somehow. Attempted kernel grid is printed above. Kernel width is %s pixels."
% kw)
if mode == 'gaussian':
k = fun.gaussian(kx, 1.0, 0.0, sig_w)
if mode == 'box':
k = fun.box(kx, 1.0, 0.0, w)
kx = kx[k > 0.0]
k = k[k > 0.0]
if (-1.0) * kx[-1] != kx[0]:
print(kx)
raise Exception(
"ERROR in box_smooth: Kernel could not be made symmetric AFTER CROPPING OUT THE BOX, somehow. Attempted kernel grid is printed above. Kernel width is %s pixels."
% kw)
k /= np.sum(k)
return (convolve(fx, k, edge_degree))
def constant_velocity_wl_grid(wl, fx, oversampling=1.0):
"""This function will define a constant-velocity grid that is (optionally)
sampled a number of times finer than the SMALLEST velocity difference that is
currently in the grid.
Example: wl_cv,fx_cv = constant_velocity_wl_grid(wl,fx,oversampling=1.5).
This function is hardcoded to raise an exception if wl or fx contain NaNs,
because interp1d does not handle NaNs.
Parameters
----------
wl : list, np.ndarray
The wavelength array to be resampled.
fx : list, np.ndarray
The flux array to be resampled.
oversampling : float
The factor by which the wavelength array is *minimally* oversampled.
Returns
-------
wl : np.array
The new wavelength grid.
fx : np.array
The interpolated flux values.
a : float
The velocity step in km/s.
"""
import astropy.constants as consts
import numpy as np
import tayph.functions as fun
from tayph.vartests import typetest, nantest, dimtest, postest
from scipy import interpolate
import pdb
import matplotlib.pyplot as plt
typetest(
oversampling,
[int, float],
'oversampling in constant_velocity_wl_grid()',
)
typetest(wl, [list, np.ndarray], 'wl in constant_velocity_wl_grid()')
typetest(fx, [list, np.ndarray], 'fx in constant_velocity_wl_grid()')
nantest(wl, 'wl in in constant_velocity_wl_grid()')
nantest(fx, 'fx in constant_velocity_wl_grid()')
dimtest(wl, [0], 'wl in constant_velocity_wl_grid()')
dimtest(fx, [len(wl)], 'fx in constant_velocity_wl_grid()')
postest(oversampling, 'oversampling in constant_velocity_wl_grid()')
oversampling = float(oversampling)
wl = np.array(wl)
fx = np.array(fx)
c = consts.c.to('km/s').value
dl = derivative(wl)
dv = dl / wl * c
a = np.min(dv) / oversampling
wl_new = 0.0
#The following while loop will define the new pixel grid.
#It starts trying 100,000 points, and if that's not enough to cover the entire
#range from min(wl) to max(wl), it will add 100,000 more; until it's enough.
n = len(wl)
while np.max(wl_new) < np.max(wl):
x = fun.findgen(n)
wl_new = np.exp(a / c * x) * np.min(wl)
n += len(wl)
wl_new[0] = np.min(
wl) #Artificially set to zero to avoid making a small round
#off error in that exponent.
#Then at the end we crop the part that goes too far:
wl_new_cropped = wl_new[(wl_new <= np.max(wl))]
x_cropped = x[(wl_new <= np.max(wl))]
i_fx = interpolate.interp1d(wl, fx)
fx_new_cropped = i_fx(wl_new_cropped)
return (wl_new_cropped, fx_new_cropped, a)
