diff --git a/doc/source/index.rst b/doc/source/index.rst index d72779f5..c574b557 100644 --- a/doc/source/index.rst +++ b/doc/source/index.rst @@ -51,6 +51,7 @@ for getting started. ocrreject radialvel r_util + splice sshift stisnoise tastis diff --git a/doc/source/splice.rst b/doc/source/splice.rst new file mode 100644 index 00000000..ab39e50d --- /dev/null +++ b/doc/source/splice.rst @@ -0,0 +1,11 @@ +.. _splice: + +************************** +splice +************************** + +.. currentmodule:: stistools.splice + +.. automodule:: stistools.splice + :members: + :undoc-members: \ No newline at end of file diff --git a/stistools/__init__.py b/stistools/__init__.py index 65864d48..c0962db0 100644 --- a/stistools/__init__.py +++ b/stistools/__init__.py @@ -8,6 +8,7 @@ from . import x1d from . import x2d from . import mktrace +from . import splice from . import sshift from . import stisnoise from . import wx2d diff --git a/stistools/splice.py b/stistools/splice.py new file mode 100644 index 00000000..225cfa9a --- /dev/null +++ b/stistools/splice.py @@ -0,0 +1,886 @@ +#! /usr/bin/env python + +import numpy as np +from astropy.io import fits +from astropy.table import Table +from functools import reduce + +__doc__ = """ +The ``splice`` module concatenates the several orders contained in a STIS +Echelle ``_x1d`` spectrum while co-adding the overlapping sections. This code +emulates the splice module previously implemented on the STSDAS IRAF package. + +The function :func:`splice` takes as input an ``_x1d`` Echelle spectrum +(including path) and outputs an Astropy Table containing the wavelength, flux, +flux uncertainty and data-quality flags. The spliced spectrum can be saved as +an additional extension to the ``_x1d`` file by setting ``update_fits`` to +``True``. The spliced spectrum can also be exported to an ascii file by setting +a path and filename to ``output_file``. + +The codebase from the ULLYSES project, which is also sponsored by STScI, possess +routines that can be used to co-add and splice STIS echelle spectra. The main +differences between these codes are the following: 1) ``splice`` does not change +the spectra in non-overlap regions; 2) ``splice`` takes into account +data-quality (DQ) flags and assignes specific DQ values to co-added pixels; +3) ``splice`` only works for STIS spectra and does not concatenate or co-add +COS spectra or multiple STIS spectra. + +Examples +-------- + +Read a spectrum with :func:`read_spectrum` and plot all the different orders +using ``matplotlib`` to visualize how an echelle extracted spectrum looks like: + +>>> import matplotlib.pyplot as plt +>>> from stistools import splice +>>> spectrum = splice.read_spectrum('oblh01040_x1d.fits') +>>> for s in spectrum: +>>> plt.plot(s['wavelength'], s['flux']) +>>> _ = plt.xlabel(r'Wavelength (${\rm \AA}$)') +>>> _ = plt.ylabel(r'Flux density (erg s$^{-1}$ cm$^{-2}$ ${\rm \AA}^{-1}$)') + +Splice the Echelle spectrum orders with :func:`splice` and plot it +using ``matplotlib``: + +>>> import matplotlib.pyplot as plt +>>> from stistools import splice +>>> spliced_spectrum = splice.splice('oblh01040_x1d.fits') +>>> plt.errorbar(spliced_spectrum['WAVELENGTH'], spliced_spectrum['FLUX'], +>>> yerr=spliced_spectrum['ERROR']) +>>> _ = plt.xlabel(r'Wavelength (${\rm \AA}$)') +>>> _ = plt.ylabel(r'Flux density (erg s$^{-1}$ cm$^{-2}$ ${\rm \AA}^{-1}$)') + +""" +__taskname__ = "splice" +__version__ = "1.0" +__vdate__ = "28-February-2023" +__author__ = "Leonardo Dos Santos" +__all__ = ["nearest_index", "read_spectrum", "find_overlap", "merge_overlap", + "concatenate_sections", "splice"] + + +# Useful tool when dealing with binned data +def nearest_index(array, target_value): + """ + Finds the index of a value in ``array`` that is closest to ``target_value``. + + Parameters + ---------- + array : ``numpy.array`` + Target array. + target_value : ``float`` + Target value. + + Returns + ------- + index : ``int`` + Index of the value in ``array`` that is closest to ``target_value``. + """ + index = array.searchsorted(target_value) + index = np.clip(index, 1, len(array) - 1) + left = array[index - 1] + right = array[index] + index -= target_value - left < right - target_value + return index + + +# Read the Echelle spectrum based on dataset name and a prefix for file location +def read_spectrum(x1d_input, truncate_edge_left=None, truncate_edge_right=None): + """ + This is a fairly straightforward function to read the spectrum from a `x1d` + FITS file. + + Parameters + ---------- + x1d_input : ``str`` + Path and name of the ``*_x1d.fits`` file containing the spectrum. + + truncate_edge_left : ``int``, optional + Set the number of low-resolution pixels at the left edge of the detector + where the spectra should be truncated. If ``None``, then no truncation + is applied. Default is ``None``. + + truncate_edge_right : ``int``, optional + Set the number of low-resolution pixels at the right edge of the + detector where the spectra should be truncated. If ``None``, then no + truncation is applied. Default is ``None``. + + Returns + ------- + spectrum : ``list`` + List of all the orders contained in the Echelle spectrum and their + respective fluxes. + """ + with fits.open(x1d_input) as hdu: + header = hdu[0].header + data = hdu['SCI'].data + optical_element = header['OPT_ELEM'] + if optical_element[0] != 'E': + raise TypeError("This is not an Echelle spectrum.") + wavelength = data['WAVELENGTH'] + flux = data['FLUX'] + uncertainty = data['ERROR'] + data_quality = data['DQ'] + gross_counts = data['GROSS'] + net_counts = data['NET'] + n_orders = len(wavelength) + + # We index the spectral regions as a `dict` in order to avoid confusion with + # too many numerical indexes. Also, since the orders are in reverse order, + # we index them in the opposite way + spectrum = [{'wavelength': wavelength[-i - 1], + 'flux': flux[-i - 1], + 'uncertainty': uncertainty[-i - 1], + 'data_quality': data_quality[-i - 1], + 'gross': gross_counts[-i - 1], + 'net': net_counts[-i - 1]} + for i in range(n_orders)] + + # If edge truncation is passed, we simply assign a different DQ flags to the + # pixels to be truncated. We choose the flag 4096, which corresponds to + # "Extracted flux affected by bad input data." + if truncate_edge_right is not None: + for sk in spectrum: + sk['data_quality'][-truncate_edge_right:] = 4096 + if truncate_edge_left is not None: + for sk in spectrum: + sk['data_quality'][:truncate_edge_left] = 4096 + + return spectrum + + +# Identify overlaps in the whole spectrum +def find_overlap(spectrum): + """ + Find and return the overlapping sections of the Echelle spectrum. + + Parameters + ---------- + spectrum : ``list`` + List of dictionaries containing the orders of the Echelle spectrum. + It should resemble the output of ``read_spectrum()``. + + Returns + ------- + unique_sections : ``list`` + List containing the unique sections of the spectrum. + + overlap_pair_sections : ``list`` + List containing the overlapping pairs of the spectrum. + + overlap_trio_sections : ``list`` + List containing the overlapping trios of the spectrum. + """ + n_orders = len(spectrum) + + # Identify the wavelength borders of each order + borders = [] + for order in spectrum: + borders.append([min(order['wavelength']), max(order['wavelength'])]) + + unique_sections = [] + first_trio = [] + second_trio = [] + third_trio = [] + first_pair = [] + second_pair = [] + + # The following code is hacky and not pretty to look at, but it works. Sorry + # for the mess! + + # First we deal with the first orders + order = spectrum[0] + wl = order['wavelength'] + idx = [nearest_index(wl, borders[1][0]), + nearest_index(wl, borders[2][0])] + + # There is always a unique section here + unique_idx = np.arange(0, idx[0], 1) + unique_sections.append({'wavelength': order['wavelength'][unique_idx], + 'flux': order['flux'][unique_idx], + 'uncertainty': order['uncertainty'][unique_idx], + 'data_quality': order['data_quality'][unique_idx], + 'gross': order['gross'][unique_idx], + 'net': order['net'][unique_idx]}) + + # There is also a pair overlap. But since it's with the next order, it's + # considered a "second pair" + overlap_01 = np.arange(idx[0], idx[1], 1) + second_pair.append({'wavelength': order['wavelength'][overlap_01], + 'flux': order['flux'][overlap_01], + 'uncertainty': order['uncertainty'][overlap_01], + 'data_quality': order['data_quality'][overlap_01], + 'gross': order['gross'][overlap_01], + 'net': order['net'][overlap_01]}) + + if idx[1] < 1023: + # There is a trio overlap + overlap_012 = np.arange(idx[1], 1023, 1) + third_trio.append( + {'wavelength': order['wavelength'][overlap_012], + 'flux': order['flux'][overlap_012], + 'uncertainty': order['uncertainty'][overlap_012], + 'data_quality': order['data_quality'][overlap_012], + 'gross': order['gross'][overlap_012], + 'net': order['net'][overlap_012]} + ) + else: + pass + + # Now the second order + order = spectrum[1] + wl = order['wavelength'] + idx = [nearest_index(wl, borders[2][0]), + nearest_index(wl, borders[0][1]), + nearest_index(wl, borders[3][0])] + + # There are two pairs, potentially one or two trios, and potentially a + # unique section + if idx[1] > idx[0]: + # There is a trio overlap + overlap_01 = np.arange(0, idx[0], 1) + overlap_012 = np.arange(idx[0], idx[1], 1) + unique_1 = None + + if idx[2] < 1023: + # There is another trio overlap + overlap_12 = np.arange(idx[1], idx[2], 1) + overlap_123 = np.arange(idx[2], 1023, 1) + else: + overlap_123 = None + overlap_12 = np.arange(idx[0], idx[2], 1) + + else: + # No trio overlap + overlap_01 = np.arange(0, idx[1], 1) + overlap_012 = None + unique_1 = np.arange(idx[1], idx[0], 1) + overlap_123 = None + overlap_12 = np.arange(idx[0], idx[2], 1) + + # Add the to the lists + if unique_1 is not None: + unique_sections.append( + {'wavelength': order['wavelength'][unique_1], + 'flux': order['flux'][unique_1], + 'uncertainty': order['uncertainty'][unique_1], + 'data_quality': order['data_quality'][unique_1], + 'gross': order['gross'][unique_1], + 'net': order['net'][unique_1]} + ) + else: + pass + + first_pair.append({'wavelength': order['wavelength'][overlap_01], + 'flux': order['flux'][overlap_01], + 'uncertainty': order['uncertainty'][overlap_01], + 'data_quality': order['data_quality'][overlap_01], + 'gross': order['gross'][overlap_01], + 'net': order['net'][overlap_01]}) + second_pair.append({'wavelength': order['wavelength'][overlap_12], + 'flux': order['flux'][overlap_12], + 'uncertainty': order['uncertainty'][overlap_12], + 'data_quality': order['data_quality'][overlap_12], + 'gross': order['gross'][overlap_12], + 'net': order['net'][overlap_12]}) + + if overlap_012 is not None: + second_trio.append({'wavelength': order['wavelength'][overlap_012], + 'flux': order['flux'][overlap_012], + 'uncertainty': order['uncertainty'][overlap_012], + 'data_quality': order['data_quality'][overlap_012], + 'gross': order['gross'][overlap_012], + 'net': order['net'][overlap_012]}) + else: + pass + + if overlap_123 is not None: + third_trio.append({'wavelength': order['wavelength'][overlap_123], + 'flux': order['flux'][overlap_123], + 'uncertainty': order['uncertainty'][overlap_123], + 'data_quality': order['data_quality'][overlap_123], + 'gross': order['gross'][overlap_123], + 'net': order['net'][overlap_123]}) + + # Now we deal with the third to third before last orders in a loop + for i in range(n_orders - 4): + order = spectrum[i + 2] + wl = order['wavelength'] + idx = [nearest_index(wl, borders[i][1]), + nearest_index(wl, borders[i + 3][0]), + nearest_index(wl, borders[i + 1][1]), + nearest_index(wl, borders[i + 4][0])] + if idx[0] > 0: + overlap_idx_012 = np.arange(0, idx[0], 1) + else: + overlap_idx_012 = None + if idx[2] < idx[1]: + overlap_idx_12 = np.arange(idx[0], idx[2], 1) + overlap_idx_123 = None + unique_idx_2 = np.arange(idx[2], idx[1], 1) + overlap_idx_23 = np.arange(idx[1], idx[3], 1) + elif idx[2] == idx[1]: + overlap_idx_12 = np.arange(idx[0], idx[2], 1) + overlap_idx_123 = None + unique_idx_2 = np.array([idx[2], ]) + overlap_idx_23 = np.arange(idx[1], idx[3], 1) + else: + overlap_idx_12 = np.arange(idx[0], idx[1], 1) + overlap_idx_123 = np.arange(idx[1], idx[2], 1) + unique_idx_2 = None + overlap_idx_23 = np.arange(idx[2], idx[3], 1) + if idx[3] < 1023: + overlap_idx_234 = np.arange(idx[3], 1023, 1) + else: + overlap_idx_234 = None + + if len(overlap_idx_12) > 0: + first_pair.append( + {'wavelength': order['wavelength'][overlap_idx_12], + 'flux': order['flux'][overlap_idx_12], + 'uncertainty': order['uncertainty'][overlap_idx_12], + 'data_quality': order['data_quality'][overlap_idx_12], + 'gross': order['gross'][overlap_idx_12], + 'net': order['net'][overlap_idx_12]} + ) + else: + pass + + if len(overlap_idx_23) > 0: + second_pair.append( + {'wavelength': order['wavelength'][overlap_idx_23], + 'flux': order['flux'][overlap_idx_23], + 'uncertainty': order['uncertainty'][overlap_idx_23], + 'data_quality': order['data_quality'][overlap_idx_23], + 'gross': order['gross'][overlap_idx_23], + 'net': order['net'][overlap_idx_23]} + ) + else: + pass + + if overlap_idx_012 is not None: + first_trio.append( + {'wavelength': order['wavelength'][overlap_idx_012], + 'flux': order['flux'][overlap_idx_012], + 'uncertainty': order['uncertainty'][overlap_idx_012], + 'data_quality': order['data_quality'][overlap_idx_012], + 'gross': order['gross'][overlap_idx_012], + 'net': order['net'][overlap_idx_012]} + ) + else: + pass + + if overlap_idx_123 is not None: + second_trio.append( + {'wavelength': order['wavelength'][overlap_idx_123], + 'flux': order['flux'][overlap_idx_123], + 'uncertainty': order['uncertainty'][overlap_idx_123], + 'data_quality': order['data_quality'][overlap_idx_123], + 'gross': order['gross'][overlap_idx_123], + 'net': order['net'][overlap_idx_123]} + ) + else: + pass + + if overlap_idx_234 is not None: + third_trio.append( + {'wavelength': order['wavelength'][overlap_idx_234], + 'flux': order['flux'][overlap_idx_234], + 'uncertainty': order['uncertainty'][overlap_idx_234], + 'data_quality': order['data_quality'][overlap_idx_234], + 'gross': order['gross'][overlap_idx_234], + 'net': order['net'][overlap_idx_234]} + ) + else: + pass + + if unique_idx_2 is not None: + unique_sections.append( + {'wavelength': order['wavelength'][unique_idx_2], + 'flux': order['flux'][unique_idx_2], + 'uncertainty': order['uncertainty'][unique_idx_2], + 'data_quality': order['data_quality'][unique_idx_2], + 'gross': order['gross'][unique_idx_2], + 'net': order['net'][unique_idx_2]} + ) + else: + pass + + # Now we deal with the last orders. Almost there! + order = spectrum[-2] + wl = order['wavelength'] + idx = [nearest_index(wl, borders[-4][1]), + nearest_index(wl, borders[-1][0]), + nearest_index(wl, borders[-3][1])] + + # There are two pairs, potentially one or two trios, and potentially a + # unique section + if idx[0] > 0: + # There is a trio overlap + overlap_012 = np.arange(0, idx[0], 1) + if idx[2] > idx[1]: + # There is another trio overlap + overlap_123 = np.arange(idx[1], idx[2], 1) + overlap_12 = np.arange(idx[0], idx[1], 1) + overlap_23 = np.arange(idx[2], 1023, 1) + unique_2 = None + else: + overlap_123 = None + unique_2 = np.arange(idx[2], idx[1], 1) + overlap_12 = np.arange(idx[0], idx[2], 1) + overlap_23 = np.arange(idx[1], 1023, 1) + else: + overlap_012 = None + overlap_123 = None + unique_2 = np.arange(idx[2], idx[1], 1) + overlap_12 = np.arange(idx[0], idx[2], 1) + overlap_23 = np.arange(idx[1], 1023, 1) + + # Add the to the lists + if unique_2 is not None: + unique_sections.append( + {'wavelength': order['wavelength'][unique_2], + 'flux': order['flux'][unique_2], + 'uncertainty': order['uncertainty'][unique_2], + 'data_quality': order['data_quality'][unique_2], + 'gross': order['gross'][unique_2], + 'net': order['net'][unique_2]} + ) + else: + pass + + if len(overlap_12) > 0: + first_pair.append({'wavelength': order['wavelength'][overlap_12], + 'flux': order['flux'][overlap_12], + 'uncertainty': order['uncertainty'][overlap_12], + 'data_quality': order['data_quality'][overlap_12], + 'gross': order['gross'][overlap_12], + 'net': order['net'][overlap_12]}) + else: + pass + + if len(overlap_23) > 0: + second_pair.append({'wavelength': order['wavelength'][overlap_23], + 'flux': order['flux'][overlap_23], + 'uncertainty': order['uncertainty'][overlap_23], + 'data_quality': order['data_quality'][overlap_23], + 'gross': order['gross'][overlap_23], + 'net': order['net'][overlap_23]}) + else: + pass + + if overlap_012 is not None: + first_trio.append({'wavelength': order['wavelength'][overlap_012], + 'flux': order['flux'][overlap_012], + 'uncertainty': order['uncertainty'][overlap_012], + 'data_quality': order['data_quality'][overlap_012], + 'gross': order['gross'][overlap_012], + 'net': order['net'][overlap_012]}) + else: + pass + + if overlap_123 is not None: + second_trio.append({'wavelength': order['wavelength'][overlap_123], + 'flux': order['flux'][overlap_123], + 'uncertainty': order['uncertainty'][overlap_123], + 'data_quality': order['data_quality'][overlap_123], + 'gross': order['gross'][overlap_123], + 'net': order['net'][overlap_123]}) + + # Finally deal with the last order + order = spectrum[-1] + wl = order['wavelength'] + idx = [nearest_index(wl, borders[-3][1]), + nearest_index(wl, borders[-2][1])] + + # There is always a unique section here + unique_idx = np.arange(idx[1], 1023, 1) + unique_sections.append( + {'wavelength': order['wavelength'][unique_idx], + 'flux': order['flux'][unique_idx], + 'uncertainty': order['uncertainty'][unique_idx], + 'data_quality': order['data_quality'][unique_idx], + 'gross': order['gross'][unique_idx], + 'net': order['net'][unique_idx]} + ) + + # There is also a pair overlap. But since it's with the previous order, it's + # considered a "first pair" + overlap_23 = np.arange(idx[0], idx[1], 1) + + if len(overlap_23) > 0: + first_pair.append( + {'wavelength': order['wavelength'][overlap_23], + 'flux': order['flux'][overlap_23], + 'uncertainty': order['uncertainty'][overlap_23], + 'data_quality': order['data_quality'][overlap_23], + 'gross': order['gross'][overlap_23], + 'net': order['net'][overlap_23]} + ) + else: + pass + + if idx[0] > 0: + # There is a trio overlap + overlap_123 = np.arange(0, idx[0], 1) + first_trio.append( + {'wavelength': order['wavelength'][overlap_123], + 'flux': order['flux'][overlap_123], + 'uncertainty': order['uncertainty'][overlap_123], + 'data_quality': order['data_quality'][overlap_123], + 'gross': order['gross'][overlap_123], + 'net': order['net'][overlap_123]} + ) + else: + pass + + # With all that done, we assemble the overlap sections into a large list + overlap_pair_sections = [] + overlap_trio_sections = [] + n_pairs = len(first_pair) + n_trios = len(first_trio) + for i in range(n_pairs): + overlap_pair_sections.append([first_pair[i], second_pair[i]]) + if n_trios > 0: + for i in range(n_trios): + overlap_trio_sections.append([first_trio[i], second_trio[i], + third_trio[i]]) + + return unique_sections, overlap_pair_sections, overlap_trio_sections + + +# Merge overlapping sections +def merge_overlap(overlap_sections, + acceptable_dq_flags=(0, 64, 128, 1024, 2048), + weight='sensitivity'): + """ + Merges overlapping spectral regions. The basic workflow of this function + is to interpolate the sections into a common wavelength table and calculate + the weighted mean flux for each wavelength bin. If the fluxes are + inconsistent between each other, the code can use the flux with higher SNR + instead of the mean. If there are still outlier fluxes (compared to + neighboring pixels), the code uses the flux from the lower SNR section + instead. Co-added (merged) pixels will have their DQ flag set to `32768` if + they are the result of combining good pixels (according to the list of + acceptable flags). Their DQ flag will be set to `65536` if the combined + pixels do not have an acceptable DQ flag. + + Parameters + ---------- + overlap_sections : ``list`` + List of dictionaries containing the overlapping spectra of neighboring + orders. + + acceptable_dq_flags : array-like, optional + Data-quality flags that are acceptable when co-adding overlapping + spectra. The default values are (0, 64, 128, 1024, 2048), which + correspond to: 0 = regular pixel, 64 = vignetted pixel, 128 = pixel in + overscan region, 1024 = small blemish, 2048 = more than 30% of + background pixels rejected by sigma-clipping in the data reduction. + + weight : ``str``, optional + Defines how to merge the overlapping sections. The options currently + implemented are ``'sensitivity'`` and ``'snr'`` (inverse square of the + uncertainties). Default is ``'sensitivity'``. + + Returns + ------- + overlap_merged : ``dict`` + Dictionary containing the merged overlapping spectrum. + """ + n_overlaps = len(overlap_sections) + bitwise_or_acceptable_dq_flags = reduce(np.bitwise_or, acceptable_dq_flags) + + # First we need to determine which spectrum has a higher sensitivity + avg_sensitivity = np.array([np.nanmean(ok['net'] / ok['flux']) + for ok in overlap_sections]) + + # We interpolate the lower-SNR spectra to the wavelength bins of the higher + # SNR spectrum. + max_sens_idx = np.where(avg_sensitivity == np.nanmax(avg_sensitivity))[0][0] + overlap_ref = overlap_sections.pop(np.copy(max_sens_idx)) + + f_interp = [] + err_interp = [] + net_interp = [] + for i in range(n_overlaps - 1): + # Since we cannot really interpolate DQ flags, before we do the + # interpolation, we need to identify the pixels where + # the DQ flags are not acceptable, and assign them a NaN value + overlap_s_f = np.copy(overlap_sections[i]['flux']) + overlap_s_u = np.copy(overlap_sections[i]['uncertainty']) + overlap_s_n = np.copy(overlap_sections[i]['net']) + overlap_s_dq = overlap_sections[i]['data_quality'] + where_dq_bad = np.where(overlap_s_dq & bitwise_or_acceptable_dq_flags) + overlap_s_f[where_dq_bad] = np.nan + overlap_s_u[where_dq_bad] = np.nan + overlap_s_n[where_dq_bad] = np.nan + + # Now we perform the interpolation + f_interp.append(np.interp(overlap_ref['wavelength'], + overlap_sections[i]['wavelength'], + overlap_s_f)) + err_interp.append(np.interp(overlap_ref['wavelength'], + overlap_sections[i]['wavelength'], + overlap_s_u)) + net_interp.append(np.interp(overlap_ref['wavelength'], + overlap_sections[i]['wavelength'], + overlap_s_n)) + f_interp = np.array(f_interp) + err_interp = np.array(err_interp) + net_interp = np.array(net_interp) + sens_interp = net_interp / f_interp # This is a good estimate of the + # sensitivity. If there were NaNs, however, we set the sensitivity to an + # arbitraty value; these interpolated pixels will be ignored anyway during + # merging + sens_interp[np.where(np.isnan(sens_interp))] = 0.0 + sens_ref = overlap_ref['net'] / overlap_ref['flux'] + + # Merge the spectra. We will take the weighted averaged, with weights equal + # to the inverse of the uncertainties squared multiplied by a scale factor + # to avoid numerical overflows. + if weight == 'sensitivity': + scale = 1E-10 + weights_interp = sens_interp * scale + weights_ref = sens_ref * scale + elif weight == 'snr': + scale = 1E-20 + weights_interp = (1 / err_interp) ** 2 * scale + weights_ref = (1 / overlap_ref['uncertainty']) ** 2 * scale + elif weight == 'binary': + weights_interp = np.zeros_like(sens_interp) + weights_ref = np.ones_like(sens_ref) + else: + raise ValueError( + 'The weighting option "{}" is not implemented.'.format(weight)) + + # Here we deal with the data-quality flags. We only accept flags that are + # listed in `acceptable_dq_flags`. Let's initialize the dq flag arrays + dq_ref = overlap_ref['data_quality'] + # We start assuming that all the dq weights are zero + dq_weights_ref = np.zeros_like(dq_ref) + # And then for each acceptable dq, if the element of the dq array is one + # of the acceptable flags, we set its dq weight to one + dq_weights_ref[np.where(dq_ref & bitwise_or_acceptable_dq_flags)] = 1 + # The lines above do not catch a DQ flag of zero, so we have to manually + # add them in case they are an acceptable DQ flag + dq_weights_ref_0 = np.zeros_like(dq_weights_ref) + if any(0 in acceptable_dq_flags for it in range(len(acceptable_dq_flags))) \ + is True: + dq_weights_ref_0[np.where(dq_ref == 0)] = 1 + dq_weights_ref += dq_weights_ref_0 + + # For the merged section, we create a new data-quality array filled with + # 32768, which is what we establish as the flag for co-added pixels + dq_merge = np.ones_like(dq_ref, dtype=int) * 32768 + # And the pixel-by-pixel weights are all one, except where there were set to + # Nan + dq_weights_interp = np.ones_like(f_interp) + dq_weights_interp[np.where(np.isnan(f_interp))] = 0 + + # Now we need to get rid of the NaNs to avoid numerical problems. We will + # assign them an arbitrary value of -99, but the value does not matter + # because their weight is zero anyway + f_interp[np.where(np.isnan(f_interp))] = -99 + err_interp[np.where(np.isnan(err_interp))] = -99 + net_interp[np.where(np.isnan(net_interp))] = -99 + sens_interp[np.where(np.isnan(sens_interp))] = -99 + + # Now we need to verify if we are setting the dq weighting to zero in both + # the reference and the interpolated dqs. If this is the case, we will + # set their weights to one and then flag these pixels + sum_dq_weights = np.copy(dq_weights_ref + np.sum(dq_weights_interp, axis=0)) + dq_weights_ref[sum_dq_weights < 1] = 1 + for i in range(n_overlaps - 1): + dq_weights_interp[i][sum_dq_weights < 1] = 1 + dq_merge[sum_dq_weights < 1] = 65536 + + # And then we multiply the original weights by the dq weights + weights_interp *= dq_weights_interp + weights_ref *= dq_weights_ref + + # This following array will be important later + sum_weights = np.sum(weights_interp, axis=0) + weights_ref + + # Finally co-add the overlaps + wl_merge = np.copy(overlap_ref['wavelength']) + + f_merge = np.zeros_like(overlap_ref['flux']) + err_merge = np.zeros_like(overlap_ref['uncertainty']) + for i in range(n_overlaps - 1): + f_merge += f_interp[i] * weights_interp[i] + err_merge += err_interp[i] ** 2 * weights_interp[i] ** 2 + f_merge += overlap_ref['flux'] * weights_ref + err_merge += overlap_ref['uncertainty'] ** 2 * weights_ref ** 2 + f_merge = f_merge / sum_weights + err_merge = err_merge ** 0.5 / sum_weights + + overlap_merged = {'wavelength': wl_merge, 'flux': f_merge, + 'uncertainty': err_merge, 'data_quality': dq_merge} + + return overlap_merged + + +# Splice the spectra +def concatenate_sections(unique_spectra_list, merged_pair_list, + merged_trio_list): + """ + Concatenate the unique and the (merged) overlapping spectra. + + Parameters + ---------- + unique_spectra_list : ``list`` + List of unique spectra. + + merged_pair_list : ``list`` + List of merged overlapping pair spectra. + + merged_trio_list : ``list`` + List of merged overlapping trio spectra. + + Returns + ------- + spliced_wavelength : ``numpy.ndarray`` + Array containing the wavelengths in the entire spectrum. + + spliced_flux : ``numpy.ndarray`` + Array containing the fluxes in the entire spectrum. + + spliced_uncertainty : ``numpy.ndarray`` + Array containing the flux uncertainties in the entire spectrum. + """ + n_pair_overlap = len(merged_pair_list) + all_spectra = [] + + # We always start with the first unique spectrum + all_spectra.append(unique_spectra_list[0]) + k = 1 + for i in range(n_pair_overlap): + all_spectra.append(merged_pair_list[i]) + try: + all_spectra.append(merged_trio_list[i]) + pass + except IndexError: + all_spectra.append(unique_spectra_list[k]) + k += 1 + pass + + # There may still be some unique spectra remaining + unique_remaining = len(unique_spectra_list) - k + for ik in range(unique_remaining): + all_spectra.append(unique_spectra_list[k]) + k += 1 + + spliced_wavelength = \ + np.concatenate([spectrum['wavelength'] for spectrum in all_spectra]) + spliced_flux = \ + np.concatenate([spectrum['flux'] for spectrum in all_spectra]) + spliced_uncertainty = \ + np.concatenate([spectrum['uncertainty'] for spectrum in all_spectra]) + spliced_data_quality = \ + np.concatenate([spectrum['data_quality'] for spectrum in all_spectra]) + + return spliced_wavelength, spliced_flux, spliced_uncertainty, \ + spliced_data_quality + + +# The splice pipeline does everything +def splice(x1d_input, update_fits=False, output_file=None, weight='sensitivity', + acceptable_dq_flags=(0, 64, 128, 1024, 2048), + truncate_edge_left=None, truncate_edge_right=None): + """ + The main workhorse of the package. This pipeline performs all the steps + necessary to merge overlapping spectral sections and splice them with the + unique sections. + + Parameters + ---------- + x1d_input : ``str`` + Path and name of the ``*_x1d.fits`` file containing the spectrum. + + update_fits : ``bool``, optional + Use carefully, since it can modify fits files permanently. Parameter + that decides whether to update the ``*_x1d.fits`` file with a new + extension containing the spliced spectrum. + + output_file : ``str`` or ``None``, optional + String containing the location to save the output spectrum as an ascii + file. If ``None``, no output file is saved and the code returns an + Astropy Table instead. Default is ``None``. + + weight : ``str``, optional + Defines how to merge the overlapping sections. The options currently + implemented are ``'sensitivity'`` and ``'snr'`` (inverse square of the + uncertainties). Default is ``'sensitivity'``. + + acceptable_dq_flags : array-like, optional + Data-quality flags that are acceptable when co-adding overlapping + spectra. The default values are (0, 64, 128, 1024, 2048), which + correspond to: 0 = regular pixel, 64 = vignetted pixel, 128 = pixel in + overscan region, 1024 = small blemish, 2048 = more than 30% of + background pixels rejected by sigma-clipping in the data reduction. + + truncate_edge_left (``int``, optional): + Set the number of low-resolution pixels at the left edge of the detector + where the spectra should be truncated. If ``None``, then no truncation + is applied. Default is ``None``. + + truncate_edge_right (``int``, optional): + Set the number of low-resolution pixels at the right edge of the + detector where the spectra should be truncated. If ``None``, then no + truncation is applied. Default is ``None``. + + Returns + ------- + spliced_spectrum_table : ``astropy.Table`` object + Astropy Table containing the spliced spectrum. Only returned if + ``output_file`` is ``None``. + """ + # Read the data + sections = read_spectrum(x1d_input, truncate_edge_left=truncate_edge_left, + truncate_edge_right=truncate_edge_right) + + unique_sections, overlap_pair_sections, overlap_trio_sections = \ + find_overlap(sections) + + # Merge the overlapping spectral sections + merged_pairs = [ + merge_overlap(overlap_pair_sections[k], acceptable_dq_flags, weight) + for k in range(len(overlap_pair_sections)) + ] + + if len(overlap_trio_sections) > 0: + merged_trios = [ + merge_overlap(overlap_trio_sections[k], acceptable_dq_flags, weight) + for k in range(len(overlap_trio_sections)) + ] + else: + merged_trios = [] + + # By now we have three lists: unique_sections, merged_pairs and + # merged_trios. The next step is to concatenate everything in the correct + # order. + + # Finally splice the unique and merged sections + wavelength, flux, uncertainty, dq = concatenate_sections(unique_sections, + merged_pairs, + merged_trios) + + # Instantiate the spectrum dictionary + spectrum_dict = \ + {'WAVELENGTH': wavelength, 'FLUX': flux, 'ERROR': uncertainty, 'DQ': dq} + spliced_spectrum_table = Table(spectrum_dict) + table_hdu = fits.BinTableHDU(spliced_spectrum_table) + + # This feature modifies the fits input file! Use carefully! + if update_fits is True: + with fits.open(x1d_input, mode='update') as hdul: + hdul.append(table_hdu) + else: + pass + + # Return or output the result + if output_file is None: + return spliced_spectrum_table + else: + spliced_spectrum_table.write(output_file, format='ascii') diff --git a/tests/test_splice.py b/tests/test_splice.py new file mode 100644 index 00000000..71377df4 --- /dev/null +++ b/tests/test_splice.py @@ -0,0 +1,30 @@ +from .resources import BaseSTIS +import pytest +import numpy as np +from stistools.splice import splice + + +@pytest.mark.bigdata +@pytest.mark.slow +class TestSplice(BaseSTIS): + # Test the entire splicing pipeline + def test_pipeline(self, precision_threshold=1E-6): + dataset = 'oblh01040_x1d.fits' + truth = 'spliced_spectrum_truth.dat' + path_input = self.get_data("splice/input", dataset) + path_truth = self.get_data("splice/truth", truth) + + spectrum_table = splice(path_input, weight='sensitivity') + + truth = np.loadtxt(path_truth, skiprows=1) + wl_truth = truth[:, 0] + f_truth = truth[:, 1] + u_truth = truth[:, 2] + + f_scale = 1E-12 + wl_diff = abs(np.sum(spectrum_table['WAVELENGTH'].data - wl_truth)) + f_diff = abs(np.sum(spectrum_table['FLUX'].data - f_truth)) / f_scale + u_diff = abs(np.sum(spectrum_table['ERROR'].data - u_truth)) / f_scale + diff = wl_diff + f_diff + u_diff + + assert(diff < precision_threshold)