def testWavelengthOutOfRange(self, wavelength_array):
with pytest.raises(RuntimeError):
vcl.wavelength2index(3000 * u.angstrom,
wavelength_array,
reverse=False)
vcl.wavelength2index(8000 * u.angstrom,
wavelength_array,
reverse=False)
def testWavelengthUnyt(self, wavelength_array):
assert vcl.wavelength2index(4001 * u.angstrom,
wavelength_array,
reverse=False) == 0
assert vcl.wavelength2index(4499 * u.angstrom,
wavelength_array,
reverse=False) == 1
assert vcl.wavelength2index(6400 * u.angstrom,
wavelength_array,
reverse=False) == 5
def testFindWavelengthInTwoOrders(self, s):
assert s.findWavelength(5039 * u.angstrom,
s.barycentricArray,
mid_most=False) == (39, 40)
index1 = wavelength2index(5039 * u.angstrom, s.barycentricArray[39])
index2 = wavelength2index(5039 * u.angstrom, s.barycentricArray[40])
assert abs(index1 - 2047.5) > abs(index2 - 2047.5)
assert s.findWavelength(5034 * u.angstrom,
s.barycentricArray,
mid_most=False) == (39, 40)
def testBadWavelengthValue(self, wavelength_array):
with pytest.raises(AssertionError):
vcl.wavelength2index('6521', wavelength_array, reverse=False)
with pytest.raises(AssertionError):
vcl.wavelength2index(6521, wavelength_array, reverse=False)
with pytest.raises(AssertionError):
vcl.wavelength2index(6521.0, wavelength_array, reverse=False)
def testReversedWavelengthArray(self, wavelength_array):
reversed_wavelengths = [x for x in reversed(wavelength_array)]
assert vcl.wavelength2index(4001 * u.angstrom,
reversed_wavelengths,
reverse=True) == 0
def findWavelength(self,
wavelength,
wavelength_array,
mid_most=True,
verbose=False):
"""Find which orders contain a given wavelength.
This function will return the indices of the wavelength orders that
contain the given wavelength. The result will be a tuple of length 1 or
2 containing integers in the range [0, 71].
Parameters
----------
wavelength : unyt_quantity
The wavelength to find in the wavelength array. This should be a
unyt_quantity object of length 1.
wavelength_array : `unyt.unyt_array`
An array of wavelengths in the shape of a HARPS extracted spectrum
(72, 4096) to be searched.
mid_most : bool, Default : *True*, optional
In a 2D extracted echelle spectrograph like HARPS, a wavelength
near the ends of an order can appear a second time in an adjacent
order. By default `findWavelength` will return only the single
order where the wavelength is closest to the geometric center of
the CCD, which corresponds to the point where the signal-to-noise
ratio is highest. Setting this to *False* will allow for the
possibility of a length-2 tuble being returned containing the
numbers of both orders a wavelength is found in.
Returns
-------
If mid_most is false: tuple
A tuple of ints of length 1 or 2, representing the indices of
the orders in which the input wavelength is found.
If mid_most is true: int
An int representing the order in which the wavelength found is
closest to the geometrical center.
In both cases the integers returned will be in the range [0, 71].
"""
wavelength_to_find = wavelength.to(u.angstrom)
# Make sure the wavelength to find is in the array in the first place.
err_str = "Given wavelength not in array limits: {} ({}, {})".format(
wavelength_to_find, wavelength_array[0, 0], wavelength_array[-1,
-1])
if not (wavelength_array[0, 0] <= wavelength_to_find <=
wavelength_array[-1, -1]):
raise WavelengthNotFoundInArrayError(err_str)
# Set up a list to hold the indices of the orders where the wavelength
# is found.
orders_wavelength_found_in = []
for order in range(0, 72):
if (wavelength_array[order, 0] <= wavelength_to_find <=
wavelength_array[order, -1]):
orders_wavelength_found_in.append(order)
if len(orders_wavelength_found_in) == 1:
continue
elif len(orders_wavelength_found_in) == 2:
break
assert len(orders_wavelength_found_in) > 0, 'Wavelength not found'
' in array.'
if mid_most:
# If only one array: great, return it.
if len(orders_wavelength_found_in) == 1:
return orders_wavelength_found_in[0]
# Found in two arrays: figure out which is closer to the geometric
# center of the CCD, which conveiently falls around the middle
# of the 4096-element array.
elif len(orders_wavelength_found_in) == 2:
order1, order2 = orders_wavelength_found_in
index1 = wavelength2index(wavelength_to_find,
wavelength_array[order1])
index2 = wavelength2index(wavelength_to_find,
wavelength_array[order2])
# Check which index is closest to the pixel in the geometric
# center of the 4096-length array, given 0-indexing in Python.
if abs(index1 - 2047.5) > abs(index2 - 2047.5):
return order2
else:
return order1
else:
return tuple(orders_wavelength_found_in)
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)