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)