Adds some functions for visualizing spectra and labeling spectral lines

README.md

@@ -36,6 +36,7 @@ $ snakemake ../results/ar-run/detector_images/all_components.fits \
                      exposure_time=1 \
                      cadence=1 \
                      instrument_design='moxsi_slot' \
-                     convert_to_dn='False' \
+                     apply_gain_conversion='False' \
+                     apply_electron_conversion='False' \
             --cores 1
 ```

mocksipipeline/instrument/optics/design.py

@@ -49,7 +49,7 @@ class OpticalDesign:
     pixel_size_x: u.Quantity[u.micron] = 7 * u.micron
     pixel_size_y: u.Quantity[u.micron] = 7 * u.micron
     detector_shape: tuple = (1504, 2000)
-    camera_gain: u.Quantity[u.ct / u.electron] = 1.8 * u.ct / u.electron
+    camera_gain: u.Quantity[u.DN / u.electron] = 1.8 * u.DN / u.electron
 
     def __eq__(self, value):
         if not isinstance(value, OpticalDesign):

mocksipipeline/instrument/optics/response.py

@@ -15,21 +15,7 @@
 
 from mocksipipeline.instrument.optics.filter import ThinFilmFilter
 
-ALLOWED_SPECTRAL_ORDERS = [
-    -6,
-    -5,
-    -4,
-    -3,
-    -2,
-    -1,
-    0,
-    1,
-    2,
-    3,
-    4,
-    5,
-    6,
-]
+ALLOWED_SPECTRAL_ORDERS = np.arange(-11, 12, 1)
 
 __all__ = [
     'Channel',
@@ -292,7 +278,7 @@ def effective_area(self) -> u.cm ** 2:
 
     @property
     @u.quantity_input
-    def camera_gain(self) -> u.Unit('ct / electron'):
+    def camera_gain(self) -> u.Unit('DN / electron'):
         return self.design.camera_gain
 
     @property
@@ -307,11 +293,11 @@ def electron_per_photon(self) -> u.electron / u.photon:
 
     @property
     @u.quantity_input
-    def wavelength_response(self) -> u.Unit('cm^2 ct / photon'):
+    def wavelength_response(self) -> u.Unit('cm^2 DN / photon'):
         wave_response = (self.effective_area *
                          self.electron_per_photon *
                          self.camera_gain)
-        return np.where(self._energy_out_of_bounds, 0 * u.Unit('cm2 ct /ph'), wave_response)
+        return np.where(self._energy_out_of_bounds, 0 * u.Unit('cm2 DN /ph'), wave_response)
 
     def get_wcs(self, observer, roll_angle=90 * u.deg):
         pc_matrix = pcij_matrix(roll_angle,

mocksipipeline/instrument/optics/tests/test_channel.py

@@ -50,7 +50,7 @@ def test_effective_area(channel):
 def test_wavelength_response(channel):
     assert isinstance(channel.wavelength_response, u.Quantity)
     assert channel.wavelength_response.shape == channel.wavelength.shape
-    assert channel.wavelength_response.unit.is_equivalent('cm2 ct / ph')
+    assert channel.wavelength_response.unit.is_equivalent('cm2 DN / ph')
 
 
 @pytest.mark.parametrize('channel', moxsi_short.channel_list[:4])

mocksipipeline/modeling.py

@@ -97,7 +97,8 @@ def project_spectral_cube(instr_cube,
                           observer,
                           dt=1*u.s,
                           interval=20*u.s,
-                          convert_to_dn=False,
+                          apply_gain_conversion=False,
+                          apply_electron_conversion=False,
                           chunks=None,
                           **wcs_kwargs):
     """
@@ -109,10 +110,16 @@ def project_spectral_cube(instr_cube,
     channel
     dt
     interval
-    convert_to_dn: `bool`, optional
+    apply_electron_conversion: `bool`, optional
+        If True, sum counts in electrons. Poisson sampling will still be done in
+        photon space.
+    apply_gain_conversion: `bool`, optional
         If True, sum counts in DN. Poisson sampling will still be done
-        in photon space
+        in photon space. This only has an effect if ``apply_electron_conversion``
+        is also True.
     """
+    if apply_gain_conversion and not apply_electron_conversion:
+        raise ValueError('Cannot convert to DN without also setting apply_electron_conversion=True')
     # Sample distribution
     lam = (instr_cube.data * instr_cube.unit * u.pix * dt).to_value('photon')
     if chunks is None:
@@ -127,13 +134,18 @@ def project_spectral_cube(instr_cube,
     weights = samples[idx_nonzero]
     # (Optionally) Convert to DN
     unit = 'photon'
-    if convert_to_dn:
+    if apply_electron_conversion:
         # NOTE: we can select the relevant conversion factors this way because the wavelength
         # axis of lam is the same as channel.wavelength and thus their indices are aligned
-        ct_per_photon = channel.electron_per_photon * channel.camera_gain
-        ct_per_photon = np.where(channel._energy_out_of_bounds, 0*u.Unit('ct /ph'), ct_per_photon)
-        weights = weights * ct_per_photon.to_value('ct / ph')[idx_nonzero[0]]
-        unit = 'ct'
+        unit = 'electron'
+        conversion_factor = channel.electron_per_photon.to('electron / ph')
+        if apply_gain_conversion:
+            unit = 'DN'
+            conversion_factor = (conversion_factor * channel.camera_gain).to('DN / ph')
+        conversion_factor = np.where(channel._energy_out_of_bounds,
+                                     0*conversion_factor.unit,
+                                     conversion_factor)
+        weights = weights * conversion_factor.value[idx_nonzero[0]]
     # Map counts to detector coordinates
     overlap_wcs = channel.get_wcs(observer, **wcs_kwargs)
     n_rows = channel.detector_shape[0]

mocksipipeline/visualization.py

@@ -0,0 +1,120 @@
+"""
+Visualization convenience functions for MOXSI data
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+__all__ = ['annotate_lines', 'plot_labeled_spectrum']
+
+
+def annotate_lines(axis, cube, source_location, line_list, component_wcs, cube_total=None, threshold=None, **kwargs):
+    """
+    Add labels for spectral line identifications on MOXSI spectra plots
+
+    Parameters
+    ----------
+    axis
+    cube
+    source_location
+    line_list
+    component_wcs
+    """
+    draggable = kwargs.pop('draggable', False)
+    color = kwargs.get('color', 'k')
+    annotate_kwargs = {
+        'xytext': (0, 150),
+        'textcoords': 'offset points',
+        'rotation': 90,
+        'color': color,
+        'horizontalalignment': 'center',
+        'verticalalignment':'center',
+        'arrowprops': dict(color=color, arrowstyle='-', ls='--'),
+    }
+    annotate_kwargs.update(kwargs)
+
+    line_pos, _, _ = component_wcs.world_to_pixel(source_location, line_list['wavelength'])
+    _, _, line_index = component_wcs.world_to_array_index(source_location, line_list['wavelength'])
+    line_index = np.array(line_index)
+
+    # NOTE: we cannot label lines for which we do not have corresponding data values
+    in_bounds = np.where(np.logical_and(line_index>=0, line_index<cube.data.shape[0]))
+    line_pos = line_pos[in_bounds]
+    line_index = line_index[in_bounds]
+    line_list = line_list[in_bounds]
+
+    # NOTE: we don't want to waste time labeling lines that are below a certain threshold
+    if cube_total is not None and threshold is not None:
+        ratio = cube.data / cube_total.data
+        above_thresh = np.where(ratio[line_index]>threshold)
+        line_pos = line_pos[above_thresh]
+        line_index = line_index[above_thresh]
+        line_list = line_list[above_thresh]
+
+    for pos, index, row in zip(line_pos, line_index, line_list):
+        wave_label = row["wavelength"].to_string(format="latex_inline", precision=5)
+        _annotate = axis.annotate(f'{row["ion name"]}, {wave_label}', (pos, cube.data[index]), **annotate_kwargs)
+        _annotate.draggable(draggable)
+
+    return axis
+
+
+def plot_labeled_spectrum(spectra_components, spectra_total, line_list, source_location, labels=None, **kwargs):
+    """
+    Plot a 1D spectra across the MOXSI detector, including components from multiple orders with selected spectral lines
+    labeled
+
+    Parameters
+    ----------
+    spectra_components: `list`
+        A list of `ndcube.NDCube` objects containing the spectra for each order component we want to plot.
+        Each cube should have a 3D WCS but the data should only have the last dimension of length greater
+        than 1.
+    spectra_total: `ndcube.NDCube
+        The sum of all orders. The shape constraints are the same as the components.
+    line_list: `astropy.table.Table`
+        A list of all the lines to label. At a minimum, there must be a column labeled "wavelength" and "ion name"
+    source_location: `astropy.coordinates.SkyCoord`
+        The location of the source of the emission.
+    labels: `list`, optional
+        Labels for each of the entries in ``spectra_components``
+    kwargs: optional
+        Other options for configuring axes limits and line annotations
+    """
+    x_lim = kwargs.pop('x_lim', None)
+    y_lim = kwargs.pop('y_lim', None)
+    log_y = kwargs.pop('log_y', True)
+    threshold = kwargs.pop('threshold', None)
+    figsize = kwargs.pop('figsize', (20,10))
+
+    fig = kwargs.pop('figure', None)
+    if fig is None:
+        fig = plt.figure(figsize=figsize)  # Avoid creating figure if not needed
+    ax = kwargs.pop('axes', None)
+    for i, component in enumerate(spectra_components):
+        label = labels[i] if labels else None
+        if ax is None:
+            ax = component[0,0,:].plot(label=label)
+        else:
+            component[0,0,:].plot(axes=ax, label=label)
+        annotate_lines(ax,
+                       component[0,0,:],
+                       source_location,
+                       line_list,
+                       component.wcs,
+                       cube_total=spectra_total[0,0,:],
+                       threshold=threshold,
+                       color=ax.lines[-1].get_color(),
+                       **kwargs)
+    spectra_total[0,0,:].plot(axes=ax, color='k', label='Total')
+
+    if log_y:
+        ax.set_yscale('log')
+    ax.set_xlim((1000, 2000) if x_lim is None else x_lim)
+    if y_lim is None:
+        y_lim = (0, spectra_total.data.max()*1.025)
+    ax.set_ylim(y_lim)
+    ax.set_ylabel(f"Flux [{spectra_total.unit:latex}]")
+    ax.set_xlabel('Dispersion direction')
+    ax.legend(loc=4 if log_y else 1, ncol=5)
+
+    return fig, ax

pipeline/config/config.yml

@@ -8,4 +8,5 @@ log_t_right_edge: 7.5
 delta_log_t: 0.1
 exposure_time: 3600  # in s
 cadence: 3600  # in s
-convert_to_dn: True
+apply_electron_conversion: True
+apply_gain_conversion: True

pipeline/workflow/scripts/project_intensity_cubes.py

@@ -31,7 +31,8 @@ def apply_psf(cube, channel):
     # Read in config options
     dt = float(snakemake.config['exposure_time']) * u.s
     interval = float(snakemake.config['cadence']) * u.s
-    convert_to_dn = bool(snakemake.config['convert_to_dn'])
+    apply_electron_conversion = bool(snakemake.config['apply_electron_conversion'])
+    apply_gain_conversion = bool(snakemake.config['apply_gain_conversion'])
     # Load instrument configuration
     instrument_design = getattr(mocksipipeline.instrument.configuration, snakemake.config['instrument_design'])
     # Select channel
@@ -49,6 +50,7 @@ def apply_psf(cube, channel):
                                      wcs_to_celestial_frame(spec_cube.wcs).observer,
                                      dt=dt,
                                      interval=interval,
-                                     convert_to_dn=convert_to_dn)
+                                     apply_gain_conversion=apply_gain_conversion,
+                                     apply_electron_conversion=apply_electron_conversion)
     # Save to file
     write_overlappogram(det_cube, snakemake.output[0])