def update_plots(self): """ This redraws the plot planels, taking care to reselect axis ranges and such. """ import numpy as np import tayph.drag_colour as dcb import tayph.functions as fun import copy import pdb import tayph.plotting as plotting if self.npx != len(self.x_axis): print('--------- Redrawing to account for order width mismatch') 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. self.xrange = [0,self.npx-1] # self.yrange=[0,self.nexp-1]#Should not be needed as self.y_axis cant change! 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.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.ax[2].set_xlim((min(self.x_axis),max(self.x_axis))) # self.ax[2].set_ylim(0,self.img_max) self.ax[2].clear() self.img3=self.ax[2].plot(self.x_axis,self.meanspec) # self.cbar.remove() # self.cbar = self.fig.colorbar(self.img2, ax=self.ax.ravel().tolist(),aspect = 20) # self.cbarD = dcb.DraggableColorbar_fits(self.cbar,[self.img2],'hot') # self.cbarD.connect() else: array1 = copy.deepcopy(self.order.ravel()) array2 = copy.deepcopy(self.residual.ravel()) 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. self.img1.set_array(array1) self.img1.set_clim(vmin=0,vmax=self.img_max) self.img2.set_array(array2) self.img2.set_clim(vmin=self.vmin,vmax=self.vmax) self.img3[0].set_ydata(self.meanspec) 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[2].set_ylim(0,self.img_max) self.draw_masked_areas() self.fig.canvas.draw_idle()
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 construct_gaussian_model(rv, ccf, dp, shadowname, xrange=[-200, 200], Nxticks=20.0, Nyticks=10.0): #Now we do the whole thing again for pulsations # nexp = np.shape(ccf)[0] # yrange=[0,nexp-1] # y_axis = 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_pulsations(fig, ax, rv_sel, z, 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.45, 0.15, 0.15]) 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) #Then the offset slider: rax_slider2 = plt.axes([0.8, 0.35, 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.28, 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.