Source code for pypeit.gui.binospec_ifu_extract
"""
Matplotlib GUI for interactively selecting fibers from a fiber-fed IFU
spec1d file and extracting a combined 1D spectrum.
.. include:: ../include/links.rst
"""
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Circle, RegularPolygon, Rectangle
from matplotlib.ticker import FuncFormatter
from matplotlib.widgets import Button, CheckButtons, RadioButtons, RangeSlider
from pypeit import log, PypeItError
from pypeit.core import datacube
# Bright optical sky lines (vacuum, Angstrom). Used only for the integrated
# per-fiber flux shown in the fiber-IFU extractor display, so that residual
# sky-subtraction errors do not bias the per-fiber flux estimate. Extraction
# itself uses the full wavelength range.
#
# TODO: Replace this hand-rolled list with PypeIt's built-in sky-line
# resources (e.g. ``data/sky_spec/sky_single_mg.dat`` or the
# ``OH_*_lines.dat`` line lists; cf. the matching TODO in pypeit/qa.py),
# and expose the mask half-width as a GUI control rather than a module
# constant.
_SKY_LINE_WAVELENGTHS = (
5577.34, # [OI]
5890.00, 5896.00, # Na I D
6300.30, 6363.78, # [OI]
7340.0, 7821.0, 8344.0, 8430.0, 8761.0, 8867.0, # OH band centers (approx)
)
_SKY_LINE_MASK_HALFWIDTH = 10.0 # Angstrom
[docs]
def sky_line_mask(wave):
"""Return a boolean mask that is ``True`` inside any sky-line region.
Parameters
----------
wave : `numpy.ndarray`_
1D array of wavelengths in Angstrom.
Returns
-------
`numpy.ndarray`_
1D boolean array, same shape as ``wave``, ``True`` where the
wavelength falls within ``_SKY_LINE_MASK_HALFWIDTH`` of any line
in ``_SKY_LINE_WAVELENGTHS``.
"""
mask = np.zeros_like(wave, dtype=bool)
for lam in _SKY_LINE_WAVELENGTHS:
mask |= np.abs(wave - lam) < _SKY_LINE_MASK_HALFWIDTH
return mask
[docs]
def compute_fiber_fluxes(waves, fluxes, wave_min, wave_max):
"""Per-fiber integrated flux with sky-line masking.
Each fiber is integrated over its **own** native wavelength grid, with
pixels in any sky-line region (see ``_SKY_LINE_WAVELENGTHS``) excluded.
Used to colour the fiber-IFU extractor display; the actual extracted
spectrum (:func:`pypeit.core.datacube.resample_and_combine`) is
unaffected.
Parameters
----------
waves, fluxes : :obj:`list` of `numpy.ndarray`_
Per-fiber native wavelengths and fluxes (parallel lists).
wave_min, wave_max : :obj:`float`
Integration range, Angstrom.
Returns
-------
`numpy.ndarray`_
1D array of integrated fluxes, one per fiber.
"""
out = np.zeros(len(waves), dtype=float)
for i, (w, f) in enumerate(zip(waves, fluxes)):
in_range = (w >= wave_min) & (w <= wave_max)
sky = sky_line_mask(w)
mask = in_range & ~sky
out[i] = np.nansum(f[mask])
return out
[docs]
class BinospecIFUExtractGUI:
"""Matplotlib widget GUI for selecting fibers and extracting spectra.
Displays the IFU as a hexagonal flux map, lets the user select fibers by
drawing a region or clicking, and writes the combined selection as a
:class:`~pypeit.onespec.OneSpec` FITS file.
Parameters
----------
fibers : :obj:`list` of :obj:`dict`
Per-fiber records (see :func:`pypeit.core.datacube.load_fibers`), each
carrying ``wave``, ``flux``, ``ivar``, ``ra``, and ``dec``.
raw_hdr : `astropy.io.fits.Header`_
Primary header of the source spec1d, copied to the output.
pyp_spec : :obj:`str`
Spectrograph short name, written as ``PYP_SPEC`` on the output.
outfile : :obj:`str`
Default output OneSpec FITS path.
"""
def __init__(self, fibers, raw_hdr, pyp_spec, outfile):
self.fibers = fibers
self.raw_hdr = raw_hdr
self.pyp_spec = pyp_spec
self.outfile = outfile
self.waves = [f['wave'] for f in fibers]
self.fluxes = [f['flux'] for f in fibers]
self.ivars = [f['ivar'] for f in fibers]
self.ra = np.array([f['ra'] for f in fibers])
self.dec = np.array([f['dec'] for f in fibers])
self.dec_center = float(np.median(self.dec))
self.cos_dec = np.cos(np.radians(self.dec_center))
self.ra_scaled = self.ra * self.cos_dec
self.nspec = len(fibers)
# Slider extents. Filter to fibers with at least one positive
# wavelength so an all-zero/masked fiber does not crash np.min.
valid_waves = [w[w > 0] for w in self.waves if np.any(w > 0)]
if not valid_waves:
raise PypeItError("No valid wavelength data in any fiber")
self.wave_min = float(np.min([w.min() for w in valid_waves]))
self.wave_max = float(np.max([w.max() for w in valid_waves]))
self.current_wave_min = self.wave_min
self.current_wave_max = self.wave_max
self.fiber_fluxes = compute_fiber_fluxes(
self.waves, self.fluxes, self.wave_min, self.wave_max)
# Selection state
self.y_scale = 'linear'
self.mask_sky = False
self.selection_shape = 'rectangle'
self.selection_start = None
self.selection_patch = None
self.selected_region = None
self.clicked_fibers = set()
self.fiber_mask = None
self.extracted_wave = None
self.extracted_flux = None
self.extracted_ivar = None
self._create_gui()
[docs]
def _create_gui(self):
"""Build all GUI axes, widgets, and event connections."""
self.fig = plt.figure(figsize=(16, 8))
self.ax_image = plt.subplot2grid((4, 3), (0, 0), rowspan=4,
colspan=2)
self.ax_spectrum = plt.subplot2grid((4, 3), (0, 2), rowspan=2,
colspan=1)
self.ax_controls = plt.subplot2grid((4, 3), (2, 2), colspan=1)
self.ax_coords = plt.subplot2grid((4, 3), (3, 2), colspan=1)
self.ax_controls.axis('off')
self.ax_coords.axis('off')
self._create_hexagon_plot()
self.ax_spectrum.set_xlabel('Wavelength (Å)')
self.ax_spectrum.set_ylabel('Flux')
self.ax_spectrum.set_title('Extracted Spectrum')
self.ax_spectrum.grid(True, alpha=0.3)
# Buttons
self.btn_extract = Button(
plt.axes([0.77, 0.34, 0.15, 0.05]), 'Extract Spectrum')
self.btn_extract.on_clicked(self._on_extract)
self.btn_reset = Button(
plt.axes([0.77, 0.28, 0.15, 0.05]), 'Reset Display')
self.btn_reset.on_clicked(self._on_reset)
self.btn_save = Button(
plt.axes([0.77, 0.22, 0.15, 0.05]), 'Save Spectrum')
self.btn_save.on_clicked(self._on_save)
# Mode radio
self.radio = RadioButtons(
plt.axes([0.68, 0.22, 0.08, 0.10]),
('Rectangle', 'Circle', 'Click Fibers'))
self.radio.on_clicked(self._on_shape_change)
# Y-axis scale radio for the extracted-spectrum panel
self.yscale_radio = RadioButtons(
plt.axes([0.77, 0.13, 0.08, 0.08]),
('Linear', 'Log'))
self.yscale_radio.on_clicked(self._on_yscale_change)
# Sky-line masking toggle for the extracted-spectrum panel
self.skymask_check = CheckButtons(
plt.axes([0.86, 0.13, 0.12, 0.08]),
['Mask Sky Lines'], [self.mask_sky])
self.skymask_check.on_clicked(self._on_skymask_change)
# Wavelength range slider
self.wave_slider = RangeSlider(
plt.axes([0.1, 0.05, 0.45, 0.03]), 'Wavelength (Å)',
self.wave_min, self.wave_max,
valinit=(self.wave_min, self.wave_max),
valstep=((self.wave_max - self.wave_min) / 100) or None)
self.wave_slider.on_changed(self._on_wave_change)
# Mouse events
self.fig.canvas.mpl_connect('button_press_event', self._on_press)
self.fig.canvas.mpl_connect('motion_notify_event', self._on_motion)
self.fig.canvas.mpl_connect('button_release_event', self._on_release)
plt.tight_layout()
plt.subplots_adjust(bottom=0.10)
[docs]
def _create_hexagon_plot(self):
"""Populate ``ax_image`` with a hexagonal fiber array flux map."""
# Hexagon size from median nearest-neighbour separation
distances = []
for i in range(min(50, self.nspec)):
dra = self.ra_scaled - self.ra_scaled[i]
ddec = self.dec - self.dec[i]
d = np.sqrt(dra**2 + ddec**2)
d = d[d > 0]
if d.size:
distances.append(d.min())
sep = float(np.median(distances)) if distances else 1.0 / 3600.0
hex_radius = sep * 0.58
valid = self.fiber_fluxes[np.isfinite(self.fiber_fluxes)
& (self.fiber_fluxes != 0)]
vmin = float(np.percentile(valid, 5)) if valid.size else 0.0
vmax = float(np.percentile(valid, 95)) if valid.size else 1.0
patches = [RegularPolygon((self.ra_scaled[i], self.dec[i]),
numVertices=6, radius=hex_radius,
orientation=np.pi / 2)
for i in range(self.nspec)]
self.hex_collection = PatchCollection(patches, cmap='viridis',
edgecolors='gray',
linewidths=0.5)
self.hex_collection.set_array(self.fiber_fluxes)
self.hex_collection.set_clim(vmin, vmax)
self.ax_image.add_collection(self.hex_collection)
# nspec >= 1 is guaranteed by load_fibers raising PypeItError on
# empty input; otherwise ra_scaled.min() would crash here.
pad = 5 * hex_radius
self.ax_image.set_xlim(self.ra_scaled.min() - pad,
self.ra_scaled.max() + pad)
self.ax_image.set_ylim(self.dec.min() - pad, self.dec.max() + pad)
self.ax_image.set_aspect('equal')
self.ax_image.xaxis.set_major_formatter(FuncFormatter(
lambda x, pos: f'{x / self.cos_dec:.5f}'))
self.ax_image.set_xlabel('RA (deg)')
self.ax_image.set_ylabel('Dec (deg)')
self.ax_image.set_title('Fiber Array (drag to select region)')
plt.colorbar(self.hex_collection, ax=self.ax_image,
label='Integrated Flux')
# ------------------------------------------------------------------
# Callbacks
# ------------------------------------------------------------------
[docs]
def _on_wave_change(self, val):
wave_min, wave_max = val
self.current_wave_min = float(wave_min)
self.current_wave_max = float(wave_max)
self.fiber_fluxes = compute_fiber_fluxes(
self.waves, self.fluxes,
self.current_wave_min, self.current_wave_max)
valid = self.fiber_fluxes[np.isfinite(self.fiber_fluxes)
& (self.fiber_fluxes != 0)]
if valid.size:
self.hex_collection.set_clim(np.percentile(valid, 5),
np.percentile(valid, 95))
self.hex_collection.set_array(self.fiber_fluxes)
if self.extracted_wave is not None:
self._apply_spectrum_xlim()
self.fig.canvas.draw()
[docs]
def _apply_spectrum_xlim(self):
"""Sync the spectrum axis to the current slider window."""
if self.extracted_wave is None:
return
lo, hi = self.current_wave_min, self.current_wave_max
self.ax_spectrum.set_xlim(lo, hi)
in_window = ((self.extracted_wave >= lo)
& (self.extracted_wave <= hi))
if self.mask_sky:
in_window &= ~sky_line_mask(self.extracted_wave)
flux = self.extracted_flux[in_window]
flux = flux[np.isfinite(flux)]
if self.y_scale == 'log':
# Log axis cannot represent non-positive flux; ignore it when
# setting the limits (it is also masked from the plotted data).
flux = flux[flux > 0]
if flux.size:
fmin, fmax = float(flux.min()), float(flux.max())
self.ax_spectrum.set_ylim(fmin / 1.5, fmax * 1.5)
elif flux.size:
fmin, fmax = float(flux.min()), float(flux.max())
pad = 0.05 * (fmax - fmin) if fmax > fmin else max(abs(fmax), 1.0)
self.ax_spectrum.set_ylim(fmin - pad, fmax + pad)
[docs]
def _selected_indices(self):
mask = self._get_fiber_mask()
for i in self.clicked_fibers:
mask[i] = True
return np.where(mask)[0]
[docs]
def _on_extract(self, event):
idxs = self._selected_indices()
if idxs.size == 0:
log.info("No fibers selected")
return
log.info(f"Extracting from {idxs.size} fibers")
try:
wave, flux, ivar = datacube.resample_and_combine(
[self.waves[i] for i in idxs],
[self.fluxes[i] for i in idxs],
[self.ivars[i] for i in idxs])
except PypeItError as e:
log.warning(str(e))
return
self.extracted_wave = wave
self.extracted_flux = flux
self.extracted_ivar = ivar
self.fiber_mask = np.zeros(self.nspec, dtype=bool)
self.fiber_mask[idxs] = True
self._draw_spectrum()
self.fig.canvas.draw()
[docs]
def _draw_spectrum(self):
"""Redraw the extracted spectrum honoring the current display toggles.
On a log scale, non-positive flux values (e.g. sky-subtraction
residuals) are masked so they neither appear in the trace nor
distort the axis limits. When sky-line masking is enabled, pixels
within ``_SKY_LINE_MASK_HALFWIDTH`` of a line in
``_SKY_LINE_WAVELENGTHS`` are likewise masked.
"""
if self.extracted_wave is None:
return
wave = self.extracted_wave
flux = self.extracted_flux
ivar = self.extracted_ivar
# Clamp tiny-ivar pixels so the envelope doesn't blow up the y-axis.
sigma_raw = np.where(ivar > 0,
1.0 / np.sqrt(np.maximum(ivar, 1e-300)),
np.nan)
med_sig = float(np.nanmedian(sigma_raw)) if np.any(ivar > 0) else 0.0
sigma = np.where(np.isfinite(sigma_raw)
& (sigma_raw < 10.0 * med_sig),
sigma_raw, 0.0)
plot_flux = flux
lower = flux - sigma
upper = flux + sigma
if self.y_scale == 'log':
# Mask non-positive flux: a log axis cannot represent it. The
# lower error envelope is clamped to a small positive value so
# matplotlib does not drop the band where flux > 0 but
# flux - sigma <= 0.
bad = ~(flux > 0)
plot_flux = np.where(bad, np.nan, flux)
lower = np.where(bad, np.nan, np.maximum(lower, 1e-300))
upper = np.where(bad, np.nan, upper)
if self.mask_sky:
# Blank out sky-line regions so residuals there don't clutter
# the displayed trace. The saved spectrum is unaffected.
sky = sky_line_mask(wave)
plot_flux = np.where(sky, np.nan, plot_flux)
lower = np.where(sky, np.nan, lower)
upper = np.where(sky, np.nan, upper)
self.ax_spectrum.clear()
self.ax_spectrum.set_yscale(self.y_scale)
self.ax_spectrum.plot(wave, plot_flux, lw=0.8)
self.ax_spectrum.fill_between(wave, lower, upper, alpha=0.3)
self.ax_spectrum.set_xlabel('Wavelength (Å)')
self.ax_spectrum.set_ylabel('Flux')
n_fib = int(self.fiber_mask.sum()) if self.fiber_mask is not None else 0
self.ax_spectrum.set_title(f'Extracted Spectrum ({n_fib} fibers)')
self.ax_spectrum.grid(True, alpha=0.3)
self._apply_spectrum_xlim()
[docs]
def _on_yscale_change(self, label):
self.y_scale = label.lower()
if self.extracted_wave is None:
self.ax_spectrum.set_yscale(self.y_scale)
else:
self._draw_spectrum()
self.fig.canvas.draw()
[docs]
def _on_skymask_change(self, label):
self.mask_sky = bool(self.skymask_check.get_status()[0])
if self.extracted_wave is not None:
self._draw_spectrum()
self.fig.canvas.draw()
[docs]
def _on_save(self, event):
if self.extracted_wave is None:
log.info("Nothing to save -- run Extract Spectrum first")
return
outfile = self._default_save_path()
datacube.write_onespec(self.extracted_wave, self.extracted_flux,
self.extracted_ivar, self.raw_hdr,
self.pyp_spec, outfile)
log.info(f"Wrote {outfile}")
[docs]
def _default_save_path(self):
"""Return a non-clobbering default output path.
Starts from ``self.outfile`` (``..._<base>.fits``) and inserts a
two-digit counter before the extension (``..._<base>_00.fits``),
incrementing the counter until a path that does not already exist on
disk is found.
Returns
-------
:obj:`str`
The first ``..._<NN>.fits`` path (NN starting at ``00``) that does
not yet exist.
"""
path = Path(self.outfile)
n = 0
while True:
candidate = path.with_name(f"{path.stem}_{n:02d}{path.suffix}")
if not candidate.exists():
return str(candidate)
n += 1
[docs]
def _on_reset(self, event):
self.selected_region = None
self.clicked_fibers.clear()
self.selection_start = None
if self.selection_patch:
self.selection_patch.remove()
self.selection_patch = None
# Restore wavelength range to full extent.
self.current_wave_min = self.wave_min
self.current_wave_max = self.wave_max
self.wave_slider.set_val((self.wave_min, self.wave_max))
self.fiber_fluxes = compute_fiber_fluxes(
self.waves, self.fluxes, self.wave_min, self.wave_max)
valid = self.fiber_fluxes[np.isfinite(self.fiber_fluxes)
& (self.fiber_fluxes != 0)]
vmin = float(np.percentile(valid, 5)) if valid.size else 0.0
vmax = float(np.percentile(valid, 95)) if valid.size else 1.0
self.hex_collection.set_array(self.fiber_fluxes)
self.hex_collection.set_clim(vmin, vmax)
# Reset highlighting.
self.hex_collection.set_edgecolors('gray')
self.hex_collection.set_linewidths(0.5)
# Clear spectrum panel and extracted data.
self.ax_spectrum.clear()
self.ax_spectrum.set_yscale(self.y_scale)
self.ax_spectrum.set_xlabel('Wavelength (Å)')
self.ax_spectrum.set_ylabel('Flux')
self.ax_spectrum.set_title('Extracted Spectrum')
self.ax_spectrum.grid(True, alpha=0.3)
self.extracted_wave = None
self.extracted_flux = None
self.extracted_ivar = None
self.fiber_mask = None
self.fig.canvas.draw()
[docs]
def _on_shape_change(self, label):
self.selection_shape = label.lower()
if self.selection_patch:
self.selection_patch.remove()
self.selection_patch = None
if self.selection_shape == 'click fibers':
self.ax_image.set_title('Fiber Array (click hexagons to add/remove)')
else:
self.ax_image.set_title('Fiber Array (drag to select region)')
self.fig.canvas.draw()
[docs]
def _on_press(self, event):
if event.inaxes != self.ax_image:
return
if self.selection_shape == 'click fibers':
dra = self.ra_scaled - event.xdata
ddec = self.dec - event.ydata
d = np.sqrt(dra**2 + ddec**2)
i = int(np.argmin(d))
if d[i] < 3.0 / 3600.0:
if i in self.clicked_fibers:
self.clicked_fibers.remove(i)
else:
self.clicked_fibers.add(i)
self._update_selection_display()
return
self.selection_start = (event.xdata, event.ydata)
if self.selection_patch:
self.selection_patch.remove()
self.selection_patch = None
[docs]
def _on_motion(self, event):
if (self.selection_shape == 'click fibers'
or self.selection_start is None
or event.inaxes != self.ax_image):
return
if self.selection_patch:
self.selection_patch.remove()
self.selection_patch = None
x0, y0 = self.selection_start
x1, y1 = event.xdata, event.ydata
if self.selection_shape == 'rectangle':
self.selection_patch = Rectangle(
(x0, y0), x1 - x0, y1 - y0,
fill=False, edgecolor='cyan', linewidth=2, linestyle='--')
else:
r = np.sqrt((x1 - x0)**2 + (y1 - y0)**2)
self.selection_patch = Circle(
(x0, y0), r, fill=False, edgecolor='cyan',
linewidth=2, linestyle='--')
self.ax_image.add_patch(self.selection_patch)
# Live preview: temporarily compute the would-be `selected_region`
# so `_update_selection_display` lights up the fibers under the
# current rubber-band. Restore the prior region so release is the
# only place that commits.
old_region = self.selected_region
if self.selection_shape == 'rectangle':
x0_true, x1_true = x0 / self.cos_dec, x1 / self.cos_dec
self.selected_region = {
'type': 'rectangle',
'ra_min': min(x0_true, x1_true),
'ra_max': max(x0_true, x1_true),
'dec_min': min(y0, y1),
'dec_max': max(y0, y1),
}
else:
r_scaled = np.sqrt((x1 - x0)**2 + (y1 - y0)**2)
self.selected_region = {
'type': 'circle',
'ra_center': x0 / self.cos_dec,
'dec_center': y0,
'radius_scaled': r_scaled,
}
self._update_selection_display()
self.selected_region = old_region
[docs]
def _on_release(self, event):
if (self.selection_shape == 'click fibers'
or self.selection_start is None
or event.inaxes != self.ax_image):
self.selection_start = None
if self.selection_patch:
self.selection_patch.remove()
self.selection_patch = None
self.fig.canvas.draw()
return
x0_s, y0 = self.selection_start
x1_s, y1 = event.xdata, event.ydata
x0 = x0_s / self.cos_dec
x1 = x1_s / self.cos_dec
if self.selection_shape == 'rectangle':
self.selected_region = {
'type': 'rectangle',
'ra_min': min(x0, x1), 'ra_max': max(x0, x1),
'dec_min': min(y0, y1), 'dec_max': max(y0, y1),
}
else:
r_scaled = np.sqrt((x1_s - x0_s)**2 + (y1 - y0)**2)
self.selected_region = {
'type': 'circle',
'ra_center': x0, 'dec_center': y0,
'radius_scaled': r_scaled,
}
self.selection_start = None
self._update_selection_display()
[docs]
def _get_fiber_mask(self):
"""Return a boolean mask of fibers inside ``self.selected_region``.
Returns
-------
`numpy.ndarray`_
1D boolean array of length ``self.nspec``, ``True`` for fibers
whose sky position falls within the currently drawn selection
region. Returns all-``False`` if no region is defined.
"""
mask = np.zeros(self.nspec, dtype=bool)
r = self.selected_region
if r is None:
return mask
if r['type'] == 'rectangle':
mask = ((self.ra >= r['ra_min']) & (self.ra <= r['ra_max'])
& (self.dec >= r['dec_min']) & (self.dec <= r['dec_max']))
else: # circle (radius is in scaled coordinates)
dra = self.ra_scaled - (r['ra_center'] * self.cos_dec)
ddec = self.dec - r['dec_center']
mask = (dra**2 + ddec**2) <= r['radius_scaled']**2
return mask
[docs]
def _update_selection_display(self):
"""Redraw hexagon edges to highlight selected fibers.
Fibers selected by the drawn region (``_get_fiber_mask``) **or** by
individual clicks (``self.clicked_fibers``) are drawn with a thick
cyan edge; all others revert to a thin gray edge.
"""
region = self._get_fiber_mask()
edges = []
widths = []
for i in range(self.nspec):
if region[i] or i in self.clicked_fibers:
edges.append('cyan')
widths.append(3.0)
else:
edges.append('gray')
widths.append(0.5)
self.hex_collection.set_edgecolors(edges)
self.hex_collection.set_linewidths(widths)
self.fig.canvas.draw()