def rvCorrectedArray(self): """ Return the wavelength array corrected for the star's radial velocity. """ if not hasattr(self, '_rvCorrectedArray'): tqdm.write('Creating RV-corrected array.') self._rvCorrectedArray = shift_wavelength(self.barycentricArray, -self.radialVelocity) return self._rvCorrectedArray
def shiftWavelengthArray(self, wavelength_array, shift_velocity): """Doppler shift a wavelength array by an amount equivalent to a given velocity. Parameters ---------- wavelength_array : `unyt.unyt_array` An array containing wavelengths to be Doppler shifted. Needs units of dimension length. velocity : `unyt.unyt_quantity` A Unyt quantity with dimensions length/time to shift the wavelength array by. Returns ------- `unyt.unyt_array` An array of the same shape as the given array, Doppler shifted by the given radial velocity. """ return shift_wavelength(wavelength_array, shift_velocity)
def testShiftMultipleWithSuperluminal(self): wavelengths = [400, 500, 600] * u.nm with pytest.raises(AssertionError): vcl.shift_wavelength(wavelengths, 1e8 * u.km / u.s)
def testShiftMultipleWavelengths(self, wavelength_array): assert pytest.approx( vcl.shift_wavelength(wavelength_array[:2], 100 * u.km / u.s), [4001.33425638, 4501.50103843] * u.angstrom)
def testShiftSingleWavelength(self, wavelength): assert pytest.approx( vcl.shift_wavelength(wavelength, 500 * u.km / u.s), 5008.3991 * u.angstrom)
def testSuperluminalVelocity(self, wavelength, shift_velocity): with pytest.raises(AssertionError): vcl.shift_wavelength(wavelength, shift_velocity)
def __init__(self, transition, observation, order, radial_velocity=None, close_up_plot_path=None, context_plot_path=None, integrated=True, verbose=False): """Construct a fit to an absorption feature using a Gaussian or integrated Gaussian. Parameters ---------- transition : `transition_line.Transition` object A `Transition` object representing the absorption feature to fit. observation : `obs2d.HARPSFile2DScience` object A `HARPSFile2DScience` object to find the absorption feature in. order : int The order in the e2ds file to fit the transition in. Zero-indexed, so ranging from [0-71]. Optional -------- radial_velocity : `unyt.unyt_quantity` A radial velocity (dimensions of length / time) for the object in the observation. Most of the time the radial velocity should be picked up from the observation itself, but for certain objects such as asteroids the supplied radial velocity may not be correct. In such cases, this parameter can be used to override the given radial velocity. close_up_plot_path : string or `pathlib.Path` The file name to save a close-up plot of the fit to. context_plot_path : string or `pathlib.Path` The file name to save a wider context plot (±25 km/s) around the fitted feature to. integrated : bool, Default : True Controls whether to attempt to fit a feature with an integrated Gaussian instead of a Gaussian. verbose : bool, Default : False Whether to print out extra diagnostic information while running the function. """ # Store the transition. self.transition = transition # Grab some observation-specific information from the observation. self.dateObs = observation.dateObs self.BERV = observation.BERV self.airmass = observation.airmass self.exptime = observation.exptime self.calibrationFile = observation.calibrationFile self.calibrationSource = observation.calibrationSource self.order = int(order) # Store the plot paths. self.close_up_plot_path = close_up_plot_path self.context_plot_path = context_plot_path # Define some useful numbers and variables. # The ranges in velocity space to search around to find the minimum of # an absorption line. search_range_vel = 5 * u.km / u.s # The range in velocity space to consider to find the continuum. continuum_range_vel = 25 * u.km / u.s # The number of pixels either side of the flux minimum to use in the # fit. pixel_range = 3 # If no radial velocity is given, use the radial velocity from the # supplied observation. This is mostly for use with things like # asteroids that might not have a radial velocity assigned. if radial_velocity is None: radial_velocity = observation.radialVelocity # Shift the wavelength being searched for to correct for the radial # velocity of the star. nominal_wavelength = self.transition.wavelength.to(u.angstrom) self.correctedWavelength = shift_wavelength(nominal_wavelength, radial_velocity) if verbose: tqdm.write( 'Given RV {:.2f}: line {:.3f} should be at {:.3f}'.format( radial_velocity, nominal_wavelength.to(u.angstrom), self.correctedWavelength.to(u.angstrom))) self.baryArray = observation.barycentricArray[self.order] self.fluxArray = observation.photonFluxArray[self.order] self.errorArray = observation.errorArray[self.order] # Figure out the range in wavelength space to search around the nominal # wavelength for the flux minimum, as well as the range to take for # measuring the continuum. search_range = velocity2wavelength(search_range_vel, self.correctedWavelength) self.continuumRange = velocity2wavelength(continuum_range_vel, self.correctedWavelength) low_search_index = wavelength2index( self.correctedWavelength - search_range, self.baryArray) high_search_index = wavelength2index( self.correctedWavelength + search_range, self.baryArray) self.lowContinuumIndex = wavelength2index( self.correctedWavelength - self.continuumRange, self.baryArray) self.highContinuumIndex = wavelength2index( self.correctedWavelength + self.continuumRange, self.baryArray) self.centralIndex = low_search_index + \ self.fluxArray[low_search_index:high_search_index].argmin() self.continuumLevel = self.fluxArray[self.lowContinuumIndex:self. highContinuumIndex].max() self.fluxMinimum = self.fluxArray[self.centralIndex] self.lowFitIndex = self.centralIndex - pixel_range self.highFitIndex = self.centralIndex + pixel_range + 1 # Grab the wavelengths, fluxes, and errors from the region to be fit. self.wavelengths = self.baryArray[self.lowFitIndex:self.highFitIndex] self.fluxes = self.fluxArray[self.lowFitIndex:self.highFitIndex] self.errors = self.errorArray[self.lowFitIndex:self.highFitIndex] self.lineDepth = self.continuumLevel - self.fluxMinimum self.normalizedLineDepth = self.lineDepth / self.continuumLevel self.initial_guess = (self.lineDepth * -1, self.correctedWavelength.to(u.angstrom).value, 0.05, self.continuumLevel) if verbose: tqdm.write( 'Attempting to fit line at {:.4f} with initial guess:'.format( self.correctedWavelength)) if verbose: tqdm.write('Initial parameters are:\n{}\n{}\n{}\n{}'.format( *self.initial_guess)) # Do the fitting: try: if integrated: wavelengths_lower = observation.pixelLowerArray wavelengths_upper = observation.pixelUpperArray pixel_edges_lower = wavelengths_lower[ self.order, self.lowFitIndex:self.highFitIndex] pixel_edges_upper = wavelengths_upper[ self.order, self.lowFitIndex:self.highFitIndex] self.popt, self.pcov = curve_fit( integrated_gaussian, (pixel_edges_lower.value, pixel_edges_upper.value), self.fluxes, sigma=self.errors, absolute_sigma=True, p0=self.initial_guess, method='lm', maxfev=10000) else: self.popt, self.pcov = curve_fit(gaussian, self.wavelengths.value, self.fluxes, sigma=self.errors, absolute_sigma=True, p0=self.initial_guess, method='lm', maxfev=10000) except (OptimizeWarning, RuntimeError): print(self.continuumLevel) print(self.lineDepth) print(self.initial_guess) self.plotFit(close_up_plot_path, context_plot_path, plot_fit=False, verbose=True) raise if verbose: print(self.popt) print(self.pcov) # Recover the fitted values for the parameters: self.amplitude = self.popt[0] self.mean = self.popt[1] * u.angstrom self.sigma = self.popt[2] * u.angstrom if self.amplitude > 0: err_msg = ('Fit for' f' {self.transition.wavelength.to(u.angstrom)}' ' has a positive amplitude.') tqdm.write(err_msg) self.plotFit(close_up_plot_path, context_plot_path, plot_fit=True, verbose=verbose) raise PositiveAmplitudeError(err_msg) # Find 1-σ errors from the covariance matrix: self.perr = np.sqrt(np.diag(self.pcov)) self.amplitudeErr = self.perr[0] self.meanErr = self.perr[1] * u.angstrom self.meanErrVel = abs( wavelength2velocity(self.mean, self.mean + self.meanErr)) self.sigmaErr = self.perr[2] * u.angstrom if (self.chiSquaredNu > 1): self.meanErr *= np.sqrt(self.chiSquaredNu) if verbose: tqdm.write('χ^2_ν = {}'.format(self.chiSquaredNu)) # Find the full width at half max. # 2.354820 ≈ 2 * sqrt(2 * ln(2)), the relationship of FWHM to the # standard deviation of a Gaussian. self.FWHM = 2.354820 * self.sigma self.FWHMErr = 2.354820 * self.sigmaErr self.velocityFWHM = wavelength2velocity(self.mean, self.mean + self.FWHM).to(u.km / u.s) self.velocityFWHMErr = wavelength2velocity(self.mean, self.mean + self.FWHMErr).to(u.km / u.s) # Compute the offset between the input wavelength and the wavelength # found in the fit. self.offset = self.correctedWavelength - self.mean self.offsetErr = self.meanErr self.velocityOffset = wavelength2velocity(self.correctedWavelength, self.mean) self.velocityOffsetErr = wavelength2velocity( self.mean, self.mean + self.offsetErr) if verbose: print(self.continuumLevel) print(self.fluxMinimum) print(self.wavelengths)
def find_transitions_in_obs(obs_path): """ Find transitions in a given observation file. Parameters ---------- obs_path : `pathlib.Path` or str A file path to an observation file. Returns ------- None. """ # tqdm.write('-' * 40) vprint('Fitting {}...'.format(obs_path.name)) try: obs = obs2d.HARPSFile2DScience(obs_path, pixel_positions=pix_pos, new_coefficients=new_coeffs, update=args.update) # We need to test if new calibration coefficients are available or not, # but if the wavelenth array isn't updated it won't call the function # that checks for them, so call it manually in that case. if set(['ALL', 'WAVE', 'BARY']).isdisjoint(set(args.update)): obs.getWavelengthCalibrationFile() except BlazeFileNotFoundError: err_msg = f'Blaze file not found for {obs_path.name}, continuing.' vprint(err_msg) logger.warning(err_msg) return except NewCoefficientsNotFoundError: err_msg = f'New coefficients not found for {obs_path.name},'\ ' continuing.' vprint(err_msg) logger.warning(err_msg) return obs_dir = data_dir / obs_path.stem if not obs_dir.exists(): os.mkdir(obs_dir) ccd_positions_dir = obs_dir.parent / 'ccd_positions' if not ccd_positions_dir.exists(): os.mkdir(ccd_positions_dir) # Define directory for output pickle files: output_pickle_dir = obs_dir / '_'.join(['pickles', suffix]) if not output_pickle_dir.exists(): os.mkdir(output_pickle_dir) # Define paths for plots to go in: output_plots_dir = obs_dir / '_'.join(['plots', suffix]) if not output_plots_dir.exists(): os.mkdir(output_plots_dir) # Create the plot sub-directories. closeup_dir = output_plots_dir / 'close_up' if not closeup_dir.exists(): os.mkdir(closeup_dir) context_dir = output_plots_dir / 'context' if not context_dir.exists(): os.mkdir(context_dir) # Create some lists to hold x,y coordinates for the CCD position plot. transitions_x, transitions_y, labels = [], [], [] # Create a second set for failed measurments. bad_x, bad_y, bad_labels = [], [], [] fits_list = [] fit_transitions = 0 # tqdm.write('Fitting transitions...') for transition in transitions_list: for order_num in transition.ordersToFitIn: vprint(f'Attempting fit of {transition} in order' f' {order_num}') plot_closeup = closeup_dir / '{}_{}_{}_close.png'.format( obs_path.stem, transition.label, order_num) plot_context = context_dir / '{}_{}_{}_context.png'.format( obs_path.stem, transition.label, order_num) try: fit = GaussianFit(transition, obs, order_num, radial_velocity=rv, verbose=args.verbose, integrated=args.integrated_gaussian, close_up_plot_path=plot_closeup, context_plot_path=plot_context) except RuntimeError: err_msg = ('Unable to fit' f' {transition}_{order_num} for' f' {obs_path.name}!') vprint(err_msg) logger.warning(err_msg) # Append None to fits list to signify that no fit exists for # this transition. fits_list.append(None) # Fit is plotted automatically upon failing, move on to next # transition. continue except PositiveAmplitudeError: err_msg = (f'Fit of {transition} {order_num} failed with' ' PositiveAmplitudeError in' f' {obs_path.name}!') vprint(err_msg) logger.warning(err_msg) fits_list.append(None) if args.create_ccd_plots: bad_y.append(order_num + 1) if rv is None: corrected_wavelength = shift_wavelength( transition.wavelength, obs.radialVelocity) else: corrected_wavelength = shift_wavelength( transition.wavelength, rv) bad_x.append( wavelength2index(corrected_wavelength, obs.barycentricArray[order_num])) bad_labels.append('_'.join( (transition.label, str(order_num)))) continue # Assuming the fit didn't fail, continue on: fits_list.append(fit) fit_transitions += 1 if args.create_plots: # Plot the fit. fit.plotFit(plot_closeup, plot_context) if args.create_ccd_plots: # Save the x and y pixel positions of the transitions on the CCD transitions_y.append(fit.order + 1) transitions_x.append( wavelength2index(fit.correctedWavelength, obs.barycentricArray[order_num])) labels.append(fit.label) # Pickle the list of fits, then compress them to save space before writing # them out. info_msg = f'Fit {fit_transitions}/{len(fits_list)} transitions' +\ f' in {obs_path.name}.' vprint(info_msg) logger.info(info_msg) outfile = output_pickle_dir / '{}_gaussian_fits.lzma'.format( obs._filename.stem) if not outfile.parent.exists(): os.mkdir(outfile.parent) with lzma.open(outfile, 'wb') as f: vprint(f'Pickling and compressing list of fits at {outfile}') f.write(pickle.dumps(fits_list)) # Create a plot to show locations of transitions on the CCD for this # observation. if args.create_ccd_plots: tqdm.write('Creating plot of transition CCD locations...') fig = plt.figure(figsize=(15, 10), tight_layout=True) ax = fig.add_subplot(1, 1, 1) ax.set_xlim(left=0, right=4097) ax.set_ylim(bottom=16, top=73) ax.xaxis.set_major_locator(ticks.MultipleLocator(base=512)) ax.xaxis.set_minor_locator(ticks.MultipleLocator(base=64)) ax.yaxis.set_major_locator(ticks.MultipleLocator(base=2)) ax.set_xlabel('Position across CCD (pixels)') ax.set_ylabel('Order number') ax.grid(which='major', axis='x', color='Gray', alpha=0.8) ax.grid(which='minor', axis='x', color='LightGray', alpha=0.9) for i in range(17, 73, 1): ax.axhline(i, linestyle='--', color='Gray', alpha=0.7) ax.axhline(46.5, linestyle='-.', color='Peru', alpha=0.8) ax.plot(transitions_x, transitions_y, marker='+', color='RoyalBlue', linestyle='', markersize=8) ax.plot(bad_x, bad_y, marker='+', color='FireBrick', linestyle='', markersize=10) texts = [ plt.text(transitions_x[i], transitions_y[i], labels[i], ha='center', va='center') for i in range(len(labels)) ] bad_texts = [ plt.text(bad_x[i], bad_y[i], bad_labels[i], ha='center', va='center', color='DarkRed') for i in range(len(bad_labels)) ] tqdm.write('Adjusting label text positions...') # Add all the text labels together to adjust them. texts.extend(bad_texts) adjust_text(texts, arrowprops=dict(arrowstyle='-', color='OliveDrab'), lim=1000, fontsize=9) ccd_position_filename = ccd_positions_dir /\ '{}_CCD_positions.png'.format(obs_path.stem) fig.savefig(str(ccd_position_filename)) plt.close(fig)