#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SPECEX - SPECtra EXtractor.
This module provides functions to stack spectral cubes.
Copyright (C) 2022-2023 Maurizio D'Addona <mauritiusdadd@gmail.com>
"""
import os
import argparse
from typing import Union
import numpy as np
from astropy.wcs import WCS
from astropy.io import fits
import matplotlib.pyplot as plt
from specex.utils import stack
from specex.cube import get_hdu
[docs]
def stack_and_plot(ext, basename, suffix="", is_mask=False, override_wcs=None,
dpi=150, wave_ranges=None):
"""
Stack and plot a spectral datacube.
Stack a datacube along the spectral axis and plot the result as a
PNG and a FITS file.
Parameters
----------
ext : astropy.io.fits.ImageHDU
The fits exension containing the datacube.
basename : str
The name of the output images.
suffix : str, optional
An optional suffix for the output image names. The default is "".
is_mask : bool, optional
If True, the stacked image is trated as a boolean mask where every
pixel having non-zero value will be considered as True.
The default is False.
override_wcs : astropy.wcs.WCS, optional
An optional WCS to override the one present in the datacube.
The default is None.
dpi : int, optional
The resolution of the output PNG file. The default is 150.
wave_range : TYPE, optional
The optional wavelength range to use, if None all the available
wavelenghts are used. The default is None.
Returns
-------
new_data : numpy.ndarray
The stacked image.
img_wcs : astropy.wcs.WCS
The WCS of the stacked image.
"""
if ext.data is None:
return None
img_wcs = WCS(ext.header)
wave_mask: Union[np.ndarray, None]
if wave_ranges is None:
wave_mask = None
elif img_wcs.has_spectral:
wave_mask = np.zeros(ext.data.shape[0], dtype=bool)
for wave_range in wave_ranges:
wave_index = np.arange(ext.data.shape[0])
wave_angstrom = img_wcs.spectral.pixel_to_world(wave_index)
wave_angstrom = wave_angstrom.Angstrom
mask = wave_angstrom >= np.nanmin(wave_range)
mask &= wave_angstrom <= np.nanmax(wave_range)
wave_mask |= mask
img_height, img_width = ext.data.shape[1], ext.data.shape[2]
img_figsize = (
img_width / dpi,
img_height / dpi
)
new_data = stack(ext.data, wave_mask=wave_mask)
if is_mask:
new_data = ~(new_data == new_data.max()) * 1.0
levels = np.percentile(new_data[np.isfinite(new_data)], [3, 97])
fig, ax = plt.subplots(
1, 1,
figsize=img_figsize,
subplot_kw={'projection': img_wcs.celestial}
)
ax.imshow(new_data, vmin=np.min(levels), vmax=np.max(levels))
out_name_png = f"{basename}_{suffix}.png"
fig.savefig(out_name_png, dpi=dpi, bbox_inches="tight")
plt.close(fig)
if override_wcs:
new_header = override_wcs.celestial.to_header()
else:
new_header = img_wcs.celestial.to_header()
out_name_fits = f"{basename}_{suffix}.fits"
hdu = fits.PrimaryHDU(data=new_data, header=new_header)
hdu.writeto(out_name_fits, overwrite=True)
return new_data, img_wcs
def __argshandler(options=None):
"""
Parse the arguments given by the user.
Returns
-------
args: Namespace
A Namespace containing the parsed arguments. For more information see
the python documentation of the argparse module.
"""
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument(
'input_cube', metavar='SPEC_CUBE', type=str, nargs=1,
help='The spectral cube in fits format from which spectra will be '
'extracted.'
)
parser.add_argument(
'--dpi', metavar='DPI', type=int, default=150,
help='Set the DPI of the output figures. '
'The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--spec-hdu', metavar='SPEC_HDU', type=str, default=-1,
help='Index or name of the HDU containing the spectral data to '
'use. Set this to -1 to automatically detect the HDU containing '
'the fluxes. '
'The default value is %(metavar)s=%(default)s. '
'NOTE that this value is zero indexed (i.e. second HDU has index 1).'
)
parser.add_argument(
'--var-hdu', metavar='VAR_HDU', type=str, default=-1,
help='Index or name of the HDU containing the variance. '
'Set this to -1 if no variance data is present in the cube. '
'The default value is %(metavar)s=%(default)s. '
'NOTE that this value is zero indexed (i.e. third HDU has index 2).'
)
parser.add_argument(
'--mask-hdu', metavar='MASK_HDU', type=str, default=-1,
help='Index or name of the HDU containing the data mask. '
'Set this to -1 if no variance data is present in the cube. '
'The default value is %(metavar)s=%(default)s. '
'NOTE that this value is zero indexed (i.e. fourth HDU has index 3).'
)
parser.add_argument(
'--outdir', metavar='DIR', type=str, default=None,
help='Set the directory where extracted spectra will be outputed. '
'If this parameter is not specified, then a new directory will be '
'created based on the name of the input cube.'
)
parser.add_argument(
'--wave-ranges', metavar='WAVE_RANGE', type=str, default=None,
help='Set the wavelength range(s) to stack, in the format '
'WAVE_RANGE=RANGE1_START-RANGE1_STOP,RANGE2_START-RANGE2_STOP,...'
' If not specified, the whole wavelength range is used.'
)
args = None
if options is None:
args = parser.parse_args()
else:
args = parser.parse_args(options)
return args
[docs]
def cube_stack(options=None):
"""
Run the main program of this module.
Returns
-------
None.
"""
args = __argshandler(options)
input_file = args.input_cube[0]
basename = os.path.basename(args.input_cube[0])
basename = os.path.splitext(basename)[0]
if args.outdir is not None:
basename = os.path.join(args.outdir, basename)
with fits.open(input_file) as hdul:
try:
spec_hdu_index = int(args.spec_hdu)
except ValueError:
spec_hdu_index = args.spec_hdu
else:
if spec_hdu_index < 0:
spec_hdu_index = None
try:
var_hdu_index = int(args.var_hdu)
except ValueError:
var_hdu_index = args.var_hdu
else:
if var_hdu_index < 0:
var_hdu_index = None
try:
mask_hdu_index = int(args.mask_hdu)
except ValueError:
mask_hdu_index = args.mask_hdu
else:
if mask_hdu_index < 0:
mask_hdu_index = None
spec_hdu = get_hdu(
hdul,
hdu_index=spec_hdu_index,
valid_names=['data', 'spec', 'spectrum', 'spectra'],
msg_err_notfound="ERROR: Cannot determine which HDU contains "
"spectral data, try to specify it manually!",
msg_index_error="ERROR: Cannot open HDU {} to read specra!"
)
if args.wave_ranges is not None:
wave_ranges = [
[float(k) for k in x.split('-')]
for x in args.wave_ranges.split(',')
]
else:
wave_ranges = None
dat, dat_wcs = stack_and_plot(
spec_hdu,
basename,
'data',
wave_ranges=wave_ranges,
dpi=args.dpi
)
if var_hdu_index is not None:
var_hdu = get_hdu(
hdul,
hdu_index=var_hdu_index,
valid_names=['stat', 'var', 'variance', 'noise'],
msg_err_notfound="WARNING: Cannot determine which HDU contains "
"the variance data, try to specify it manually!",
msg_index_error="ERROR: Cannot open HDU {} to read the "
"variance!",
exit_on_errors=False
)
var, var_wcs = stack_and_plot(
var_hdu,
basename,
'variance',
override_wcs=dat_wcs,
wave_ranges=wave_ranges,
dpi=args.dpi
)
if mask_hdu_index is not None:
mask_hdu = get_hdu(
hdul,
hdu_index=mask_hdu_index,
valid_names=['mask', 'platemask', 'footprint', 'dq'],
msg_err_notfound="WARNING: Cannot determine which HDU contains "
"the mask data, try to specify it manually!",
msg_index_error="ERROR: Cannot open HDU {} to read the mask!",
exit_on_errors=False
)
mask, mask_wcs = stack_and_plot(
mask_hdu,
basename,
'mask',
is_mask=True,
override_wcs=dat_wcs,
wave_ranges=wave_ranges,
dpi=args.dpi
)
if __name__ == '__main__':
cube_stack()