#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SPECEX - SPECtra EXtractor.
Extract spectra from spectral data cubes.
Copyright (C) 2022-2023 Maurizio D'Addona <mauritiusdadd@gmail.com>
"""
from __future__ import absolute_import, division, print_function
import os
import sys
import json
import argparse
import logging
from typing import Optional, Any, List, Union
import numpy as np
import multiprocessing
from scipy.signal import savgol_filter
from astropy import wcs
from astropy.table import Table, Column
from astropy.io import fits
from astropy.io.registry import IORegistryError
from astropy.coordinates import SkyCoord
from astropy.time import Time
from astropy import units
import matplotlib.pyplot as plt
from specex.utils import (
plot_zfit_check, get_log_img, stack, plot_scandata,
get_spectrum_snr, get_spectrum_snr_emission, parse_regionfile,
HAS_REGION, HAS_REPROJECT, simple_pbar_callback, do_reproject
)
from specex.sources import get_spectrum
from specex.cube import get_hdu, SpectralCube
from specex.exceptions import exception_handler
from specex.cube import (
KNOWN_SPEC_EXT_NAMES,
KNOWN_VARIANCE_EXT_NAMES,
KNOWN_INVAR_EXT_NAMES,
KNOWN_MASK_EXT_NAMES,
KNOWN_RCURVE_EXT_NAMES
)
try:
from specex.rrspecex import rrspecex, get_templates
except Exception:
HAS_RR = False
else:
HAS_RR = True
try:
from IPython.core import ultratb
from IPython import embed
except Exception:
HAS_IPYTHON = False
else:
HAS_IPYTHON = True
SPECTRA_WEIGHTING_MODES = {
'none': 'no weighting',
'whitelight': 'weighting using cube whitelight image',
'log-whitelight': 'weighting using cube whitelight image in log10 scale',
}
if HAS_REPROJECT:
SPECTRA_WEIGHTING_MODES['<file url>'] = 'weighting using an external image'
def __argshandler(options: Optional[List[str]] = None) -> Any:
"""
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='Extract spectra from a datacube and optionally compute '
'the redshift using rrspecex backend.'
)
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(
'--mode', metavar='WORKING_MODE', type=str, default='auto',
help='Set the working mode of the program. Can be one of the following'
' modes: "auto", "kron_ellipse", "kron_circle", "circular_aperture". '
'When using "kron_ellipse" mode, the input catalog must contain the '
'column for objetcs kron radius, semi-major and semi-minor axes and '
'for the rotation angle of the objects. In this mode all the pixels '
'whitin this ellipse will be added together to compute the spectra. '
'"kron_circular" mode requires the same parameters of the previous '
'mode, but a circular aperture is used instead of an ellipse and the '
'radius of this aperture is the circularized kron radius (i.e. the '
'kron darius multiplied by the axis ratio q = minor-axis/major-axis. '
'In circular_aperture mode, no morphological parameters are required '
'to be in the input catalog and the aperture size is set using the '
'input parameter --aperture-size. In "auto" mode, the best working '
'mode is selected based on the parameters available in the input '
'catalog.'
)
parser.add_argument(
'--catalog', metavar='CATALOG', type=str, default=None,
help='A catalog containing the position and the (optional) morphology '
'of the objects of which you want to extract the spectra. The catalog '
'must contain at least the RA and DEC coordinates and optionally the '
'dimensions of the isophotal axes and the rotation angle of the '
' isophotal ellipse, as well as the kron radius. The way how these '
'parameter are used depends on the program working mode (see the '
'--mode section for more information). The catalog format is'
' by default compatible by the one generated by Sextractor, so the '
'column name used to read the aformentioned qauntities are '
'respectively ALPHA_J2000, DELTA_J2000, A_IMAGE, B_IMAGE, THETA_IMAGE '
'and KRON_RADIUS. If the catalog has different column names for these '
'quantities, they can be overridden with a corresponding --key-* '
'argument (see below). The sky frame of the RA and DEC coordinates can'
' be set using --cat-skyframe.'
)
parser.add_argument(
'--regionfile', metavar='REGFILE', type=str, default=None,
help='A region file containing regions from which extract spectra. '
'The name of the refion will be used as ID for the spectra.'
)
parser.add_argument(
'--nspectra', metavar='N', type=int, default=0,
help='Number of spectra to extract (usefull for debugging purposes).'
)
parser.add_argument(
'--sn-threshold', metavar='SN_THRESH', type=float, default=1.5,
help='Set the signal to Noise ratio threshold: objects with their '
'spectrum having a SN ratio lower than threshold will be ignored. '
'If SN_THRESH is set to a value equal or lower than 0, SN filtering '
'is disabled'
)
parser.add_argument(
'--no-nans', action='store_true', default=False,
help="This option tells the program to replace eventual NaNs in the"
"extracted spectrum with linarly interpolated data. In this case the "
"output fits files contain an additional extension named 'NAN_MASK' "
"that identifies the previously bad pixels in the spectrum."
)
parser.add_argument(
'--info-hdu', metavar='INFO_HDU', type=int, default=0,
help='The HDU containing cube metadata. If this argument '
'Set this to -1 to automatically detect the HDU containing the info. '
'NOTE that this value is zero indexed (i.e. firts HDU has index 0).'
)
parser.add_argument(
'--spec-hdu', metavar='SPEC_HDU', type=int, default=None,
help='The HDU containing the spectral data to use. If this argument '
'Set this to -1 to automatically detect the HDU containing spectra. '
'NOTE that this value is zero indexed (i.e. second HDU has index 1).'
)
parser.add_argument(
'--var-hdu', metavar='VAR_HDU', type=int, default=None,
help='The HDU containing the variance of the spectral data. '
'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=int, default=None,
help='The HDU containing the valid pixel mask of the spectral data. '
'Set this to -1 if no mask 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(
'--invert-mask', action='store_true', default=False,
help='Set whether to use the inverse of the cube mask.'
)
parser.add_argument(
'--aperture-size', metavar='APERTURE', type=str, default='0.8mas',
help='Set the diameter of a fixed circula aperture used in the '
'"fixed-aperture" mode. The pixels within this aperture will be '
'used for spectra extraction. Units can be specified using the names '
'compatible with astropy.units (eg. 5mas, 1arcsec, etc.).'
'The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--anulus-size', metavar='R1,R2', type=str, default='None',
help='Set the inner and outer radii in arcsecs of a fixed anular '
'aperture used in the "fixed-aperture" mode. The pixels within this '
'anulus will be used to compute the background nearby the objects. If '
'The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--weighting', metavar='MODE', type=str, default='None',
help='Set the spatial weighting for the spectra. Available modes are:'
+ ';'.join([
f"{mode_name} for {mode_desc}"
for mode_name, mode_desc in SPECTRA_WEIGHTING_MODES.items()
]) +
". Note that in order to use extrenal images as weightmap, the python "
"package 'reproject' is needed. (Is reproject installed: " +
"YES" if HAS_REPROJECT else "NO" +
")."
)
parser.add_argument(
'--cat-pixelscale', metavar='PIXEL_SCALE', type=str, default='1mas',
help='Set the scale size in mas of the major and minor axes. Useful if'
'A and B are given in pixels instead of physical units (like in the '
'case of catalogs extracted by sextractor). Units can be specified '
'using the astropy.units names (ie. 5mas, 1arcsec, etc.). If it is '
'adimensional, then the value is interpreted as mas per pixel. '
'NOTE that this is not the pixel scale of the cube but it is the size '
' of a pixel of the image used to generate the input catalog. '
'This parameter will be ignored if using a region file as input.'
'The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--cat-format', metavar='FORMAT', type=str, default=None,
help='Manually set the format of the input catalog to %(metavar)s. '
'If not specified the program will try to automatically detect the '
'format.'
)
parser.add_argument(
'--cutouts-image', metavar='IMAGE', type=str, default=None,
help='The path of a FITS image (grayscale or single/multiext RGB) from'
'which to extract cutout of the objects. If this option is not '
'specified then the cutouts will be extracted directly from the '
'input cube.'
)
parser.add_argument(
'--cutouts-size', metavar='SIZE', type=float, default=10,
help='Size of the cutouts of the object in pixel or arcseconds. '
'The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-ra', metavar='RA_KEY', type=str, default='ALPHA_J2000',
help='Set the name of the column from which to read the RA of the '
'objects. The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-dec', metavar='DEC_KEY', type=str, default='DELTA_J2000',
help='Set the name of the column from which to read the RA of the '
'objects. The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-a', metavar='A_KEY', type=str, default='A_IMAGE',
help='Set the name of the column from which to read the size of the'
'semi-major isophotal axis. The default value is '
'%(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-b', metavar='B_KEY', type=str, default='B_IMAGE',
help='Set the name of the column from which to read the size of the'
'semi-minor isophotal axis. The default value is '
'%(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-theta', metavar='THETA_KEY', type=str, default='THETA_IMAGE',
help='Set the name of the column from which to read the rotation angle'
' of the isophotal ellipse. The default value is '
'%(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-kron', metavar='KRON_KEY', type=str, default='KRON_RADIUS',
help='Set the name of the column from which to read the size of the '
'Kron radius (in units of A or B). The default value is '
'%(metavar)s=%(default)s.'
)
parser.add_argument(
'--key-id', metavar='OBJID_KEY', type=str, default=None,
help='Set the name of the column from which to read the id of the '
'object to which the spectral data belong. If this argument is not '
'specified then a suitable column will be selected according to its '
'name or, if none is found, a new progressive id will be generated. '
'This id will be included in the fits header.'
)
parser.add_argument(
'--cat-skyframe', metavar='SKY_FRAME', type=str, default='icrs',
help='Set the sky frame RA and DEC values of the input catalog. This '
'option is ignored if a region file is used for the extraction. '
'The default value is %(metavar)s=%(default)s.'
)
parser.add_argument(
'--out-name', metavar='NAME', type=str, default=None,
help='Set the name format of the extracted spectra. If not specified '
'the default format spec_%%d.fits is used (where %%d is a '
'progressive number). Value from the columns of input zcat can be '
'specified using the parenthesis {COL_NAME}, similar to python '
'f-string. For example, in the format string spec_{RA}_{DEC}.fits '
'the {RA} and {DEC} will be replaced, for each object, to the '
'corresponding values in the input zcat.'
)
parser.add_argument(
'--outdir', metavar='DIR', type=str, default=None,
help='Set the directory where extracted spectra will be output. '
'If this parameter is not specified, then a new directory will be '
'created based on the name of the input cube.'
)
parser.add_argument(
'--check-images', action='store_true', default=False,
help='Whether or not to generate check images (cutouts, etc.).'
)
parser.add_argument(
"--checkimg-outdir", type=str, default=None, required=False,
help='Set the directory where check images are saved (when they are '
'enabled thorugh the appropriate parameter).'
)
if HAS_RR:
rr_help_note = ''
else:
rr_help_note = ' (WARNING: this option is unavailable since redrock' \
' seems not to be installed).'
parser.add_argument(
"--zbest", type=str, default=None, required=False,
help='Whether to find redshift from the spectra. ' + rr_help_note
)
parser.add_argument(
"--priors", type=str, default=None, required=False,
help='optional redshift prior file. ' + rr_help_note
)
parser.add_argument(
"--nminima", type=int, default=3, required=False,
help="the number of redshift minima to search. " + rr_help_note
)
parser.add_argument(
"--mp", type=int, default=0, required=False,
help="if not using MPI, the number of multiprocessing processes to use"
" (defaults to half of the hardware threads). " + rr_help_note
)
parser.add_argument(
"-t", "--templates", type=str, default=None, required=False,
help="template file or directory. " + rr_help_note
)
parser.add_argument(
"--debug", action='store_true', default=False,
help="Enable some debug otupout on plots and on the console."
)
args = None
if options is None:
args = parser.parse_args()
else:
args = parser.parse_args(options)
if args.cutouts_image and not os.path.exists(args.cutouts_image):
print(f"The file {args.cutouts_image} does not exist!")
sys.exit(1)
if args.mp == 0:
args.mp = int(multiprocessing.cpu_count() / 4)
return args
# TODO: resample to constant log(lambda)
[docs]
def get_spplate_fits(cube_header, spec_header, obj_ids, spec_data,
var_data=None, and_mask_data=None, or_mask_data=None,
wdisp_data=None, sky_data=None):
"""
Generate a plate fits file similar to the ones used in SDSS.
Parameters
----------
cube_header : TYPE
DESCRIPTION.
spec_header : TYPE
DESCRIPTION.
obj_ids : TYPE
DESCRIPTION.
spec_data : TYPE
DESCRIPTION.
var_data : TYPE, optional
DESCRIPTION. The default is None.
and_mask_data : TYPE, optional
DESCRIPTION. The default is None.
or_mask_data : TYPE, optional
DESCRIPTION. The default is None.
wdisp_data : TYPE, optional
DESCRIPTION. The default is None.
sky_data : TYPE, optional
DESCRIPTION. The default is None.
Returns
-------
hdul : TYPE
DESCRIPTION.
"""
spec_hdu = fits.PrimaryHDU()
ivar_hdu = fits.ImageHDU()
andmsk_hdu = fits.ImageHDU()
ormsk_hdu = fits.ImageHDU()
wdisp_hdu = fits.ImageHDU()
tbl_hdu = fits.BinTableHDU()
sky_hdu = fits.ImageHDU()
if isinstance(spec_data, np.ma.MaskedArray):
spec_data = spec_data.filled(np.nan)
if isinstance(var_data, np.ma.MaskedArray):
var_data = var_data.filled(np.nan)
spec_hdu.data = spec_data.T
spec_hdu.header['PLATEID'] = 0
try:
spec_hdu.header['MJD'] = cube_header['MJD']
except KeyError:
try:
spec_hdu.header['MJD'] = cube_header['MJD-OBS']
except KeyError:
print(
"WARNING: no MJD found in the cube metadata, using extraction "
"time instead!",
file=sys.stderr
)
spec_hdu.header['MJD'] = Time.now().mjd
# See the following link for COEFF0 and COEFF1 definition
# https://classic.sdss.org/dr7/products/spectra/read_spSpec.html
try:
spec_hdu.header["COEFF0"] = spec_header["COEFF0"]
spec_hdu.header["COEFF1"] = spec_header["COEFF1"]
except KeyError:
print(
"WARING: wavelength binning of the input cube is not compatible "
"with boss log10 binning data model.\n Computing the best "
"values of COEFF0 and COEFF1 for this wavelengt range...",
file=sys.stderr
)
# Compute coeff0 and coeff1 that best approximate wavelength range
# of the datacube
crpix = spec_header["CRPIX3"] - 1
crval = spec_header["CRVAL3"]
cd1 = spec_header["CD3_3"]
w_end = spec_data.shape[1] - crpix
coeff0 = np.log10(crval + cd1*(crpix))
coeff1 = np.log10(1 + (w_end*cd1)/crval) / w_end
spec_hdu.header['COEFF0'] = coeff0
spec_hdu.header['COEFF1'] = coeff1
spec_hdu.header["CRVAL1"] = spec_header["CRVAL3"]
spec_hdu.header["CD1_1"] = spec_header["CD3_3"]
spec_hdu.header["CDELT1"] = spec_header["CD3_3"]
spec_hdu.header["CRPIX1"] = spec_header["CRPIX3"]
spec_hdu.header['BUNIT'] = spec_header['BUNIT']
spec_hdu.header["CUNIT1"] = spec_header["CUNIT3"]
spec_hdu.header["CTYPE1"] = spec_header["CTYPE3"]
if var_data is not None:
ivar_hdu.data = 1 / var_data.T
if wdisp_data is None:
wdisp_hdu.data = np.zeros_like(spec_data.T)
else:
wdisp_hdu.data = wdisp_data.T
info_table = Table()
info_table.add_column(
Column(data=obj_ids, name='FIBERID', dtype='>i4')
)
tbl_hdu.data = info_table.as_array()
hdul = fits.HDUList([
spec_hdu, ivar_hdu, andmsk_hdu, ormsk_hdu, wdisp_hdu, tbl_hdu, sky_hdu
])
return hdul
[docs]
def get_spsingle_fits(main_header, spec_wcs_header, obj_spectrum,
spec_hdu_header, obj_spectrum_var=None, no_nans=False):
"""
Generate a fits containing the spectral data.
Parameters
----------
main_header : dict
Dictionary containing the base information of the spectrum extraction.
The keys of the dictionaries are the names of the fits CARDS that will
be written, while the values can be either strings or tuples in the
form of (cadr value, card comment).
spec_wcs_header : astropy.io.fits.Header object
Header containing the WCS information of the spectrum.
obj_spectrum : numpy.ndarray
The actual spectral data.
spec_hdu_header : astropy.io.fits.Header object
Header of the HDU of the original cube containing the spectral data.
obj_spectrum_var : numpy.ndarray (default=None)
The variance of the spectral data.
no_nans : bool, optional
If true, nan values in the spectrum are replaced by a linear
interpolation. In this case an additional extension is added to the
fits file, containing a mask for identifying the previously bax pixels.
Returns
-------
hdul : astropy.io.fits.HDUList object
The fits HDUL containing the extracted spectra.
"""
my_time = Time.now()
my_time.format = 'isot'
if isinstance(obj_spectrum, np.ma.MaskedArray):
obj_spectrum = obj_spectrum.filled(np.nan)
if isinstance(obj_spectrum_var, np.ma.MaskedArray):
obj_spectrum_var = obj_spectrum_var.filled(np.nan)
# The primary hdu contains no data but a header with information on
# how the spectrum was extracted
main_hdu = fits.PrimaryHDU()
for key, val in main_header.items():
try:
main_hdu.header[key] = val
except ValueError:
print(
f"Invalid value for spec header: {key} = {val}",
file=sys.stderr
)
spec_header = fits.Header()
spec_header['BUNIT'] = spec_hdu_header['BUNIT']
spec_header["OBJECT"] = spec_hdu_header["OBJECT"]
spec_header["CUNIT1"] = spec_hdu_header["CUNIT3"]
spec_header["CTYPE1"] = spec_hdu_header["CTYPE3"]
spec_header["OBJECT"] = spec_hdu_header["OBJECT"]
spec_header["CRPIX1"] = spec_hdu_header["CRPIX3"]
spec_header["CRVAL1"] = spec_hdu_header["CRVAL3"]
spec_header["CDELT1"] = spec_hdu_header["CD3_3"]
nanmask = np.isnan(obj_spectrum)
if no_nans:
obj_spectrum[nanmask] = 0
if obj_spectrum_var is not None:
obj_spectrum_var[nanmask] = np.nanmax(obj_spectrum_var) * 100.0
hdu_list = [
main_hdu,
fits.ImageHDU(
data=obj_spectrum,
header=spec_header,
name='SPECTRUM'
)
]
if obj_spectrum_var is not None:
var_header = fits.Header()
var_header['BUNIT'] = spec_hdu_header['BUNIT']
var_header["OBJECT"] = spec_hdu_header["OBJECT"]
var_header["CUNIT1"] = spec_hdu_header["CUNIT3"]
var_header["CTYPE1"] = spec_hdu_header["CTYPE3"]
var_header["OBJECT"] = spec_hdu_header["OBJECT"]
var_header["CRPIX1"] = spec_hdu_header["CRPIX3"]
var_header["CRVAL1"] = spec_hdu_header["CRVAL3"]
var_header["CDELT1"] = spec_hdu_header["CD3_3"]
var_header["OBJECT"] = spec_hdu_header["OBJECT"]
hdu_list.append(
fits.ImageHDU(
data=obj_spectrum_var,
header=var_header,
name='VARIANCE'
)
)
hdu_list.append(
fits.ImageHDU(
data=nanmask.astype('uint8'),
name='NAN_MASK'
)
)
hdul = fits.HDUList(hdu_list)
return hdul
[docs]
def load_specex_file(filename, spec_hdu_index=None, var_hdu_index=None,
wd_hdu_index=None, mask_hdu_index=None, smoothing=0):
"""
Load a spectrum file generated by specex (or any compatible file).
Parameters
----------
filename : str
The path of the fits file containing the spectrum.
spec_hdu_index : int or str, optional
The index or name of the HDU containing the flux.
The default is None.
var_hdu_index : int or str, optional
The index or name of the HDU containing the variance.
The default is None.
wd_hdu_index : int or str, optional
The index or name of the HDU containing the resolution curve.
The default is None.
mask_hdu_index : int or str, optional
The index or name of the HDU containing the flux nan mask.
The default is None.
smoothing : int, optional
If greather than zero, the the spectrum is smoothed with a window of
this size. The default is 0.
Raises
------
ValueError
When flux and variance have different shape.
Returns
-------
dict
A dict containing the following keys:
* 'id': the ID of the spectrum,
* 'main_header': the FITS header of the primary HDU,
* 'wavelength': the wavelengths of the dispersion grid,
* 'flux': the fluxes,
* 'variance': the variance of the fluxes,
* 'r_curve': the resolution curve,
* 'nan_mask': a mask indicating invalid fluxes,
* 'flux_units': the units of the fluxes,
* 'wavelength_units': the units of the wavelengths
"""
with fits.open(filename) as hdul:
valid_id_keys = [
f"{i}{j}"
for i in ['', 'OBJ', 'OBJ_', 'TARGET', 'TARGET_']
for j in ['ID', 'NUMBER', 'UID', 'UUID']
]
spec_id = None
for hdu in hdul:
for key in valid_id_keys:
try:
spec_id = hdu.header[key]
except KeyError:
continue
else:
break
main_header = get_hdu(
hdul,
hdu_index=0,
valid_names=['PRIMARY', 'primary'],
msg_index_error="WARNING: No Primary HDU",
exit_on_errors=False
).header.copy()
spec_hdu = get_hdu(
hdul,
hdu_index=spec_hdu_index,
valid_names=KNOWN_SPEC_EXT_NAMES,
msg_err_notfound="WARNING: No spectrum HDU",
exit_on_errors=False
)
var_hdu = get_hdu(
hdul,
hdu_index=var_hdu_index,
valid_names=KNOWN_VARIANCE_EXT_NAMES,
msg_err_notfound="WARNING: No variance HDU",
exit_on_errors=False
)
if var_hdu is None:
ivar_hdu = get_hdu(
hdul,
hdu_index=var_hdu_index,
valid_names=KNOWN_INVAR_EXT_NAMES,
msg_err_notfound="WARNING: No inverse variance HDU",
exit_on_errors=False
)
mask_hdu = get_hdu(
hdul,
hdu_index=mask_hdu_index,
valid_names=KNOWN_MASK_EXT_NAMES,
msg_err_notfound="WARNING: No nan mask HDU",
exit_on_errors=False
)
wd_hdu = get_hdu(
hdul,
hdu_index=wd_hdu_index,
valid_names=KNOWN_RCURVE_EXT_NAMES,
msg_err_notfound="WARNING: No r curve HDU",
exit_on_errors=False
)
smoothing_window_size = 4*smoothing + 1
if spec_hdu is None:
flux = None
flux_units = None
lam = None
wavelength_units = None
else:
flux = spec_hdu.data.copy()
spec_wcs = wcs.WCS(spec_hdu.header)
if smoothing > 0:
flux = savgol_filter(flux, smoothing_window_size, 3)
flux = flux.astype('float32')
pixel = np.arange(len(flux))
if spec_wcs.has_spectral:
lam = (
spec_wcs.pixel_to_world(pixel)
.Angstrom
.astype('float32')
)
else:
try:
coeff0 = spec_hdu.header["COEFF0"]
coeff1 = spec_hdu.header["COEFF1"]
except KeyError:
lam = None
else:
lam = (10**(coeff0 + coeff1*pixel)).astype('float32')
try:
flux_units = spec_hdu.header['BUNIT']
except KeyError:
flux_units = None
try:
wavelength_units = spec_hdu.header['CUNIT1']
except KeyError:
wavelength_units = None
if mask_hdu is None:
flux_nan_mask = None
else:
flux_nan_mask = mask_hdu.data.astype(bool)
if flux is not None:
flux_nan_mask |= np.isnan(flux)
if var_hdu is None:
if ivar_hdu is None:
var = None
else:
var = 1 / ivar_hdu.data.copy()
if smoothing > 0:
var = savgol_filter(var, smoothing_window_size, 3)
var = var.astype('float32')
else:
var = var_hdu.data.copy()
if smoothing > 0:
var = savgol_filter(var, smoothing_window_size, 3)
var = var.astype('float32')
if wd_hdu is None:
wd = None
else:
wd = wd_hdu.data.copy()
if smoothing > 0:
wd = savgol_filter(wd, smoothing_window_size, 3)
wd = wd.astype('float32')
if (
(flux is not None) and
(var is not None) and
flux.shape != var.shape
):
raise ValueError(
f"'{filename}' - "
"Spectral data invalid or corruptede: Flux data shape "
"do not match variance data one!"
)
return {
'id': spec_id,
'main_header': main_header,
'wavelength': lam,
'flux': flux,
'variance': var,
'r_curve': wd,
'nan_mask': flux_nan_mask,
'flux_units': flux_units,
'wavelength_units': wavelength_units
}
[docs]
def speclist2ez(spec_files, spec_hdu_name='SPECTRUM', var_hdu_name='VARIANCE'):
"""
Generate a list of spectra that can be imported into pandora EZ.
Parameters
----------
spec_files : list of str
List of file paths.
spec_hdu_name : str, optional
The name of the HDU containing the spectral data.
The default is 'SPECTRUM'.
var_hdu_name : str, optional
The name of the HDU containing the variance data.
The default is 'VARIANCE'.
Returns
-------
txt_data : TYPE
DESCRIPTION.
"""
txt_data = ""
for file in spec_files:
txt_data += f'{file}[{spec_hdu_name}](1) - '
txt_data += f'{file}[{var_hdu_name}](1) -\n'
return txt_data
[docs]
def specex(options=None):
"""
Run the main program.
Returns
-------
None.
"""
# Get user input from argument parsers helper
args = __argshandler(options)
if args.debug and HAS_IPYTHON:
sys.excepthook = exception_handler
# Working mode
mode = None
if args.catalog is not None:
if args.cat_format is not None:
sources = Table.read(args.catalog, format=args.cat_format)
else:
try:
sources = Table.read(args.catalog)
except IORegistryError:
try:
sources = Table.read(args.catalog, format='ascii')
except IORegistryError:
print(
"ERROR: cannot determine the format of the input "
"catalog. Try to specify it manually using the "
"option --cat-format.",
file=sys.stderr
)
sys.exit(1)
basename_with_ext = os.path.basename(args.catalog)
pixelscale = units.Quantity(args.cat_pixelscale)
coord_frame = args.cat_skyframe # frame of source coordinates
if str(pixelscale.unit.physical_type) == 'dimensionless':
pixelscale = pixelscale * units.mas
elif args.regionfile is not None:
if not HAS_REGION:
logging.error(
"astropy regions package is needed to handle regionfiles!"
)
sys.exit(1)
sources, coord_frame = parse_regionfile(
args.regionfile,
key_ra=args.key_ra,
key_dec=args.key_dec,
key_id=args.key_id,
file_format=args.cat_format,
)
basename_with_ext = os.path.basename(args.regionfile)
pixelscale = 1.0 * units.arcsec
mode = 'regions'
else:
logging.error(
"You must provide at least an input catalog or a region file!",
)
sys.exit(1)
basename = os.path.splitext(basename_with_ext)[0]
with SpectralCube.open(args.input_cube[0]) as my_cube:
if args.debug:
print(f"Flux HDU: {my_cube.spec_hdu.name}")
if my_cube.var_hdu is not None:
print(f"Variance HDU: {my_cube.var_hdu.name}")
if my_cube.mask_hdu is not None:
print(f"Mask HDU: {my_cube.mask_hdu.name}")
# Get spectral and celstial WCS from the cube
celestial_wcs = my_cube.getSpecWCS().celestial
spectral_wcs = my_cube.getSpecWCS().spectral
# If weighting with an external image, read that
valid_wl_modes = ['none', 'whitelight', 'log-whitelight']
weight_img: Union[None, np.ndarray] = None
if args.weighting.lower() not in valid_wl_modes:
if not os.path.isfile(args.weighting):
logging.error(
f"weighting image does not exist '{args.weighting}' "
)
sys.exit(1)
with fits.open(args.weighting) as hdul:
wimg_hdu = hdul[0]
wimg_shape = wimg_hdu.data.shape
if len(wimg_shape) != 2:
logging.error(
"ERROR: only grayscale images are supported as "
"weightmap!"
)
sys.exit(1)
elif HAS_REPROJECT:
# Reproject the external weight image to the
# datacube footprint
weight_img, wimg_mask = do_reproject(
wimg_hdu, my_cube
)
weight_img[wimg_mask == 0] = np.nan
elif wimg_shape != my_cube.spec_hdu.data.shape[1:]:
logging.error(
"The shape of the external weightmap image does "
"not match the datacube and reproject is not "
"installed! Use --help for more information."
)
sys.exit(1)
else:
weight_img = wimg_hdu.data.copy()
# TODO: we assume here that the distortion across the sky plane is
# negligible, but it is very bold assumption.
# A more general approach that takes into account also distortion
# in function of the position on the sky is needed.
# Get the pixelscales in arcsec/pixel.
cube_pixelscale_x, cube_pixelscale_y = [
units.Quantity(sc_val, u)
for sc_val, u in zip(
wcs.utils.proj_plane_pixel_scales(celestial_wcs),
celestial_wcs.wcs.cunit
)
]
if args.debug:
print(
"Cube X and Y pixelscales are: "
f"{cube_pixelscale_x.to(units.arcsec)}, "
f"{cube_pixelscale_y.to(units.arcsec)}",
file=sys.stderr
)
ra_unit = sources[args.key_ra].unit
if ra_unit is None:
if args.debug:
logging.warning(
"RA data has no units, assuming values are in degrees!",
)
ra_unit = 'deg'
dec_unit = sources[args.key_dec].unit
if dec_unit is None:
if args.debug:
logging.warning(
"DEC data has no units, assuming values are in degrees!",
)
dec_unit = 'deg'
# Get the objects positions in the cube (in pixels)
obj_sky_coords = SkyCoord(
sources[args.key_ra], sources[args.key_dec],
unit=(ra_unit, dec_unit),
frame=coord_frame
)
obj_x, obj_y = wcs.utils.skycoord_to_pixel(
coords=obj_sky_coords,
wcs=celestial_wcs
)
sources.add_column(obj_x, name='CUBE_X_IMAGE')
sources.add_column(obj_y, name='CUBE_Y_IMAGE')
img_height, img_width = my_cube.getSpatialSizePixels()
yy, xx = np.meshgrid(
np.arange(img_height),
np.arange(img_width),
indexing='ij'
)
extracted_data = np.zeros((img_height, img_width))
# Chech if there is a footprint mask in the cube
if my_cube.mask_hdu is not None:
cube_footprint = my_cube.mask_hdu.data.prod(axis=0) == 0
if args.invert_mask:
cube_footprint = ~cube_footprint
else:
cube_footprint = None
if args.outdir is None:
outdir = f'{basename}_spectra'
else:
outdir = args.outdir
if not os.path.isdir(outdir):
os.mkdir(outdir)
if mode is None:
if args.mode == 'auto':
if (
args.key_a in sources.colnames and
args.key_b in sources.colnames and
args.key_kron in sources.colnames
):
mode = 'kron_ellipse'
else:
mode = 'circular_aperture'
else:
mode = args.mode
print(f"Extracting spectra with strategy '{mode}'", file=sys.stderr)
n_objects = len(sources)
if args.nspectra:
n_objects = int(args.nspectra)
# Try to identify the column with object IDs
if args.key_id is None:
valid_id_keys = [
f"{i}{j}"
for i in ['', 'OBJ', 'OBJ_', 'TARGET', 'TARGET_']
for j in ['ID', 'NUMBER', 'UID', 'UUID']
]
for key in valid_id_keys:
if key.lower() in sources.colnames:
key_id = key.lower()
break
elif key.upper() in sources.colnames:
key_id = key.upper()
break
else:
key_id = None
else:
key_id = args.key_id
out_specfiles = []
specex_apertures = {} # these values are in physical units (arcsecs)
specex_anuli = {} # these values are in physical units (arcsecs)
source_ids = []
valid_sources_mask = np.zeros(len(sources), dtype=bool)
# Check wheter to use anuli to compute background
anulus_str = args.anulus_size.lower().strip().split(',')
anulus_r_in = None
anulus_r_out = None
if anulus_str and len(anulus_str) == 2:
try:
anulus_r1 = float(anulus_str[0])
anulus_r2 = float(anulus_str[1])
except ValueError:
pass
else:
anulus_r_in = min(anulus_r1, anulus_r2)
anulus_r_out = max(anulus_r1, anulus_r2)
if args.debug:
print(
f"Read {len(sources[:n_objects])} sources from "
"input catalog/regionfile...",
file=sys.stderr
)
for i, source in enumerate(sources[:n_objects]):
simple_pbar_callback(i + 1, n_objects)
if key_id is not None:
obj_id = source[key_id]
else:
obj_id = f"{i:06}"
obj_ra = source[args.key_ra]
obj_dec = source[args.key_dec]
xx_tr = xx - source['CUBE_X_IMAGE']
yy_tr = yy - source['CUBE_Y_IMAGE']
anulus_mask = None
if mode == 'regions':
anulus_mask = None
pix_reg = source['region'].to_pixel(celestial_wcs)
reg_mask = pix_reg.to_mask()
mask = reg_mask.to_image((img_height, img_width)) > 0
specex_apertures[obj_id] = (
1 * pixelscale,
1 * pixelscale,
0 * units.rad
)
elif mode == 'kron_ellipse':
ang = np.deg2rad(source[args.key_theta])
x_over_a = xx_tr*np.cos(ang) + yy_tr*np.sin(ang)
x_over_a /= source[args.key_a]
y_over_b = xx_tr*np.sin(ang) - yy_tr*np.cos(ang)
y_over_b /= source[args.key_b]
specex_apertures[obj_id] = (
source[args.key_a] * pixelscale,
source[args.key_b] * pixelscale,
source[args.key_theta] * units.rad
)
mask = (x_over_a**2 + y_over_b**2) < (1.0/source[args.key_kron])
elif mode == 'kron_circular':
kron_circular = source[args.key_kron] * source[args.key_b]
kron_circular /= source[args.key_a]
specex_apertures[obj_id] = (
kron_circular * pixelscale,
kron_circular * pixelscale,
0 * units.rad
)
mask = (xx_tr**2 + yy_tr**2) < (kron_circular ** 2)
elif mode == 'circular_aperture':
aperture_size = units.Quantity(args.aperture_size)
aperture_radius = aperture_size / 2
if str(aperture_size.unit.physical_type) == 'dimensionless':
aperture_radius = aperture_radius * units.mas
try:
specex_apertures[obj_id] = (
aperture_radius,
aperture_radius,
0 * units.rad
)
except TypeError:
logging.warning(
f"object {obj_id} has an invalid ID, skipping..."
)
continue
xx_scaled = xx_tr * cube_pixelscale_x
yy_scaled = yy_tr * cube_pixelscale_y
mask = (xx_scaled**2 + yy_scaled**2) < (aperture_radius)**2
# TODO: fix here
"""
if anulus_r_in and anulus_r_out:
specex_anuli[obj_id] = (
anulus_r_in,
anulus_r_out
)
anulus_mask_out = rad_coords <= anulus_r_out / cube_pixelscale
anulus_mask_in = rad_coords > anulus_r_in / cube_pixelscale
anulus_mask = anulus_mask_in & anulus_mask_out
else:
specex_anuli[obj_id] = None
anulus_mask = None
"""
if cube_footprint is not None:
mask &= cube_footprint
if np.sum(mask) == 0:
logging.warning(
f"object {obj_id} has no footprint, skipping..."
)
continue
obj_extraction_mask = np.zeros_like(mask, dtype='float32')
if args.weighting.lower() == 'none':
weights = None
obj_extraction_mask[mask] = 1.0
else:
if args.weighting.lower() == 'whitelight':
# summing selected spaxels along spectral axis to get a
# whitelight weightmap
w_arr = np.ma.array(
my_cube.spec_hdu.data[:, mask],
mask=~np.isfinite(my_cube.spec_hdu.data[:, mask]),
dtype='float32'
)
weights = np.ma.mean(w_arr, axis=0)
weights -= np.ma.min(weights)
elif args.weighting.lower() == 'log-whitelight':
w_arr = np.ma.array(
my_cube.spec_hdu.data[:, mask],
mask=~np.isfinite(my_cube.spec_hdu.data[:, mask]),
dtype='float32'
)
weights = np.ma.mean(w_arr, axis=0)
weights -= np.ma.min(weights)
weights = np.ma.log10(1 + weights)
elif weight_img is not None:
weights = weight_img[mask]
weights -= np.ma.min(weights)
obj_extraction_mask[mask] = weights
var_spaxels = None
if my_cube.var_hdu is not None:
var_spaxels = my_cube.var_hdu.data[:, mask]
obj_spectrum, obj_spectrum_var = get_spectrum(
flux_spaxels=my_cube.spec_hdu.data[:, mask],
var_spaxels=var_spaxels,
weights=weights
)
obj_spectrum = obj_spectrum.filled(np.nan)
obj_spectrum_var = obj_spectrum_var.filled(np.nan)
if np.sum(
np.isnan(obj_spectrum_var[~np.isnan(obj_spectrum)])
) != 0:
logging.warning(
f"\nWARNING: object {obj_id} is weird! Skipping...\n"
)
continue
if np.sum(~np.isnan(obj_spectrum)) < 100:
logging.warning(
f"WARNING: object {obj_id} has invalid data..."
)
if anulus_mask is not None:
# TODO: fix here!
# Here we suppose that background vaies very slowly within
# object aperture. In this case we can copute the mean
# background and subtract it total contribution within the
# aperture.
bkg_spectrum = np.median(
my_cube.spec_hdu.data[:, mask], axis=0
)
smoothed_bkg_spectrum = savgol_filter(bkg_spectrum, 51, 11)
obj_spectrum -= smoothed_bkg_spectrum*np.sum(mask)
if np.sum(~np.isnan(obj_spectrum)) == 0:
if args.debug:
logging.warning(
f"WARNING: object {obj_id} has no spectral data, "
"skipping..."
)
continue
sn_spec = get_spectrum_snr(obj_spectrum, obj_spectrum_var)
sn_em = get_spectrum_snr_emission(obj_spectrum, obj_spectrum_var)
if (
(args.sn_threshold > 0) and
(max(sn_spec, sn_em) < args.sn_threshold)
):
if args.debug:
logging.warning(
f"WARNING: object {obj_id} has SNR={sn_spec:.2f} < "
f"{args.sn_threshold}, skipping..."
)
continue
my_time = Time.now()
my_time.format = 'isot'
if args.out_name is None:
outname = f"spec_{obj_id}.fits"
else:
outname = args.out_name.replace('%d', obj_id)
try:
outname = outname.format(**source)
except KeyError as exc:
logging.error(
f"ERROR: output name formatting error: {exc}"
)
apertures_info = [
x.to_string() for x in specex_apertures[obj_id]
]
main_header = {
'CUBE': (
basename_with_ext, "Spectral cube used for extraction"
),
'RA': (
obj_ra, "Ra of the center of the object"
),
'DEC': (
obj_dec, "Dec of the center of the object"
),
'FRAME': (
coord_frame, "Ra and Dec sky frame"
),
'NPIX': (
np.sum(mask), 'Number of pixels used for this spectra'
),
'HISTORY': f'Extracted on {str(my_time)}',
'ID': obj_id,
'SN': (
sn_spec, "SNR of the total spectrum"
),
'SN_EMISS': (
sn_em, "SNR due to emission lines only"
),
'EXT_MODE': (
mode, "Spex extraction mode"
),
'EXT_APER': (
json.dumps(apertures_info),
"Apertures used in specex"
)
}
# Copy the spectral part of the WCS into the new FITS
spec_wcs_header = spectral_wcs.to_header()
hdul = get_spsingle_fits(
main_header, spec_wcs_header,
obj_spectrum, my_cube.spec_hdu.header,
obj_spectrum_var, no_nans=args.no_nans
)
out_file_name = os.path.join(outdir, outname)
hdul.writeto(
out_file_name,
overwrite=True
)
# If we arrived here, then the object has a vail spectrum and can
# be saved. Let's also mark the object as valid and store its ID
out_specfiles.append(os.path.realpath(out_file_name))
source_ids.append(obj_id)
valid_sources_mask[i] = True
# Add also the extracted pixels to the extraction map
extracted_data += obj_extraction_mask
if anulus_mask is not None:
extracted_data -= anulus_mask
# Discard all invalid sources
sources = sources[valid_sources_mask]
source_ids = np.array(source_ids)
with open("spec_list_pandora_ez.txt", 'w') as f:
f.write(speclist2ez(out_specfiles))
if args.debug:
print(
f"\nExtracted {len(sources)} valid sources\n",
file=sys.stderr
)
if args.checkimg_outdir is not None:
check_images_outdir = args.checkimg_outdir
else:
check_images_outdir = os.path.join(outdir, 'checkimages')
if not os.path.isdir(check_images_outdir):
os.mkdir(check_images_outdir)
# NOTE: for some weird reason redrock hangs if any figure has been
# created using matplotlib. Do all the plottings only after
# redrock has finished!
if args.zbest:
rrspecex_options = [
'--zbest', args.zbest,
'--checkimg-outdir', check_images_outdir
]
if args.priors is not None:
rrspecex_options += ['--priors', args.priors]
if args.nminima is not None:
rrspecex_options += ['--nminima', f'{args.nminima:d}']
if args.mp is not None:
rrspecex_options += ['--mp', f'{args.mp:d}']
if args.templates is not None:
rrspecex_options += ['--templates', f'{args.templates}']
rrspecex_options += out_specfiles
targets, zfit, scandata = rrspecex(options=rrspecex_options)
if args.debug or not args.cutouts_image:
stacked_cube = stack(my_cube.spec_hdu.data)
if args.debug:
fig = plt.figure(figsize=(10, 10))
ax = plt.subplot(projection=celestial_wcs)
logimg, cvmin, cvmax = get_log_img(stacked_cube)
ax.imshow(
logimg,
origin='lower',
vmin=cvmin,
vmax=cvmax,
cmap='jet'
)
fig.savefig("stacked_cube.png")
plt.tight_layout()
plt.close(fig)
if args.debug:
plot_templates = get_templates(templates=args.templates)
else:
plot_templates = None
print("", file=sys.stderr)
# Creating directory to save Chi-2 plots
if args.debug:
x2_plots_out_dir = os.path.join(
check_images_outdir,
"X2_plots"
)
if not os.path.isdir(x2_plots_out_dir):
os.mkdir(x2_plots_out_dir)
for i, target in enumerate(targets):
logging.info("Generating previews...")
simple_pbar_callback(i, n_objects-1)
fig, axs = plot_zfit_check(
target,
zfit,
plot_template=plot_templates,
wavelengt_units=my_cube.spec_hdu.header['CUNIT3'],
flux_units=my_cube.spec_hdu.header['BUNIT'],
)
figname = f'spectrum_{target.spec_id}.png'
figname = os.path.join(check_images_outdir, figname)
fig.savefig(figname, dpi=150)
plt.close(fig)
if args.debug:
figname = f'scandata_{target.spec_id}.png'
figname = os.path.join(x2_plots_out_dir, figname)
fig, axs = plot_scandata(target, scandata)
fig.savefig(figname, dpi=150)
plt.close(fig)
specex_map = fits.PrimaryHDU(
data=extracted_data,
header=celestial_wcs.to_header()
)
specex_map.writeto(
f"{my_cube.getBaseName()}_extmap.fits",
overwrite=True
)
if args.check_images:
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
ax.imshow(extracted_data, origin='lower')
fig.savefig("sext_extraction_map.png", dpi=150)
plt.close(fig)
if __name__ == '__main__':
specex()