"""
Generic loader definitions
"""
import os
from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS
from astropy.units import Unit
from astropy.nddata import StdDevUncertainty
from ...interfaces.decorators import data_loader, data_writer
from ...core.data import Spectrum1DRef
__all__ = ['fits_identify', 'simple_generic_loader', 'simple_generic_writer']
[docs]def fits_identify(*args, **kwargs):
"""
Check whether given filename is FITS. This is used for Astropy I/O
Registry.
"""
return (isinstance(args[0], str) and
args[0].lower().split('.')[-1] in ['fits', 'fit', 'fits.gz'])
[docs]def simple_generic_writer(data, file_name, **kwargs):
"""
Basic `Spectrum1DRef` FITS writer.
"""
# Create fits columns
flux_col = fits.Column(name='FLUX', format='E', array=data.data.data.value)
disp_col = fits.Column(name='WAVE', format='E', array=data.dispersion.data.value)
uncert_col = fits.Column(name='UNCERT', format='E', array=data.uncertainty.array)
mask_col = fits.Column(name='MASK', format='L', array=data.mask)
cols = fits.ColDefs([flux_col, disp_col, uncert_col, mask_col])
# Create the bin table
tbhdu = fits.BinTableHDU.from_columns(cols)
# Create header
prihdu = fits.PrimaryHDU(header=data.meta.get('header')
or data.wcs.to_header())
# Compose
thdulist = fits.HDUList([prihdu, tbhdu])
thdulist.writeto("{}.fits".format(file_name), overwrite=True)
@data_loader(label="Simple Fits", identifier=fits_identify, extensions=["fits", "fit"],
writer=simple_generic_writer)
[docs]def simple_generic_loader(file_name, **kwargs):
"""
Basic FITS file loader
Presumption is the primary data is a table with columns 'flux'
and 'err'. The dispersion information is encoded in the FITS
header keywords.
Parameters
----------
file_name: str
The path to the FITS file
Returns
-------
data: Spectrum1DRef
The data.
"""
name = os.path.basename(file_name.rstrip(os.sep)).rsplit('.', 1)[0]
hdulist = fits.open(file_name, **kwargs)
header = hdulist[0].header
tab = Table.read(file_name)
meta = {'header': header}
wcs = WCS(hdulist[0].header)
unit = Unit('Jy')
uncertainty = StdDevUncertainty(tab["err"])
data = tab["flux"]
hdulist.close()
return Spectrum1DRef(data=data, name=name, wcs=wcs,
uncertainty=uncertainty, unit=unit, meta=meta)