covariance in tabular fits

Author: kbwestfall

specutils/io/default_loaders/tabular_fits.py

@@ -47,9 +47,9 @@ def tabular_fits_loader(file_obj, column_mapping=None, hdu=1, store_data_header=
 
     Parameters
     ----------
-    file_obj : str, file-like, or :class:`~astropy.io.fits.HDUList`
-            FITS file name, object (provided from name by Astropy I/O Registry),
-            or HDU list (as resulting from `~astropy.io.fits.open`).
+    file_obj : str, file-like, or HDUList
+        FITS file name, object (provided from name by Astropy I/O Registry),
+        or HDUList (as resulting from astropy.io.fits.open()).
     hdu : int
         The HDU of the fits file (default: 1st extension) to read from
     store_data_header : bool
@@ -73,8 +73,8 @@ def tabular_fits_loader(file_obj, column_mapping=None, hdu=1, store_data_header=
 
     Returns
     -------
-    data : :class:`Spectrum1D`
-        The spectrum that is represented by the data in the input table.
+    data : Spectrum1D
+        The spectrum that is represented by the data in this table.
     """
     # Parse the wcs information. The wcs will be passed to the column finding
     # routines to search for spectral axis information in the file.
@@ -86,6 +86,12 @@ def tabular_fits_loader(file_obj, column_mapping=None, hdu=1, store_data_header=
         else:
             tab.meta = hdulist[0].header
 
+        # Determine if there is a correlation matrix
+        correl = None
+        if 'CORREL' in [h.name for h in hdulist]:
+            correl = Table.read(hdulist['CORREL'])
+            correl.meta = hdulist['CORREL'].header
+
     # Minimal checks for wcs consistency with table data -
     # assume 1D spectral axis (having shape (0, NAXIS1),
     # or alternatively compare against shape of 1st column.
@@ -96,9 +102,9 @@ def tabular_fits_loader(file_obj, column_mapping=None, hdu=1, store_data_header=
     # If no column mapping is given, attempt to parse the file using
     # unit information
     if column_mapping is None:
-        return generic_spectrum_from_table(tab, wcs=wcs)
+        return generic_spectrum_from_table(tab, wcs=wcs, correl=correl, **kwargs)
 
-    return spectrum_from_column_mapping(tab, column_mapping, wcs=wcs)
+    return spectrum_from_column_mapping(tab, column_mapping, wcs=wcs, correl=correl)
 
 
 @custom_writer("tabular-fits")
@@ -173,7 +179,7 @@ def tabular_fits_writer(spectrum, file_name, hdu=1, update_header=False, store_d
     if spectrum.uncertainty is not None:
         if isinstance(spectrum.uncertainty, Covariance):
             var, correl = spectrum.uncertainty.to_tables()
-            columns.append(np.sqrt(var))
+            columns.append(np.sqrt(var) * funit)
             colnames.append("uncertainty")
         else:
             try:

specutils/io/parsing_utils.py

@@ -6,7 +6,7 @@
 import contextlib
 
 from astropy.io import fits
-from astropy.nddata import StdDevUncertainty
+from astropy.nddata import StdDevUncertainty, Covariance
 from astropy.utils.exceptions import AstropyUserWarning
 import astropy.units as u
 import warnings
@@ -52,7 +52,7 @@ def read_fileobj_or_hdulist(*args, **kwargs):
                 hdulist.close()
 
 
-def spectrum_from_column_mapping(table, column_mapping, wcs=None, verbose=False):
+def spectrum_from_column_mapping(table, column_mapping, wcs=None, correl=None, verbose=False):
     """
     Given a table and a mapping of the table column names to attributes
     on the Spectrum1D object, parse the information into a Spectrum1D.
@@ -76,6 +76,11 @@ def spectrum_from_column_mapping(table, column_mapping, wcs=None, verbose=False)
     wcs : :class:`~astropy.wcs.WCS` or :class:`gwcs.WCS`
         WCS object passed to the Spectrum1D initializer.
 
+    correl : `astropy.table.Table`, optional
+        Table with correlation matrix for the uncertainties in coordinate
+        format; see `~astropy.nddata.Covariance`. If None, uncertainties are
+        assumed to be uncorrelated.
+
     verbose : bool
         Print extra info.
 
@@ -131,14 +136,28 @@ def spectrum_from_column_mapping(table, column_mapping, wcs=None, verbose=False)
         spec_kwargs.setdefault(kwarg_name, kwarg_val)
 
     # Ensure that the uncertainties are a subclass of NDUncertainty
+    if spec_kwargs.get('uncertainty') is None and correl is not None:
+        warnings.warn('Unable to parse uncertainty from provided table.  Ignoring provided '
+                      'correlation matrix data.')
+        correl = None
     if spec_kwargs.get('uncertainty') is not None:
-        spec_kwargs['uncertainty'] = StdDevUncertainty(
-            spec_kwargs.get('uncertainty'))
+        if correl is not None:
+            err = spec_kwargs.get('uncertainty')
+            try:
+                spec_kwargs['uncertainty'] = Covariance.from_tables(err**2, correl, quiet=True)
+            except (ValueError, TypeError):
+                warnings.warn('Unable to parse correlation table into a Covariance object.  '
+                              'Ignoring correlation matrix data.')
+                correl = None
+        # NOTE: This is not an `else` or `elif` block in order to catch the
+        # change to correl=None when handling the excemption above.
+        if correl is None:
+            spec_kwargs['uncertainty'] = StdDevUncertainty(spec_kwargs.get('uncertainty'))
 
     return Spectrum1D(**spec_kwargs, wcs=wcs, meta={'header': table.meta})
 
 
-def generic_spectrum_from_table(table, wcs=None):
+def generic_spectrum_from_table(table, wcs=None, correl=None, **kwargs):
     """
     Load spectrum from an Astropy table into a Spectrum1D object.
     Uses the following logic to figure out which column is which:
@@ -155,12 +174,15 @@ def generic_spectrum_from_table(table, wcs=None):
 
     Parameters
     ----------
-    table : :class:`~astropy.table.Table`
-        Table containing a column of ``flux``, and optionally ``spectral_axis``
-        and ``uncertainty`` as defined above.
-    wcs : :class:`~astropy.wcs.WCS`
+    table : `astropy.table.QTable`
+        Tabulated data with units
+    wcs : :class:`~astropy.wcs.WCS`, optional
         A FITS WCS object. If this is present, the machinery will fall back
-        and default to using the ``wcs`` to find the dispersion information.
+        to using the wcs to find the dispersion information.
+    correl : `astropy.table.Table`, optional
+        Table with correlation matrix for the uncertainties in coordinate
+        format; see `~astropy.nddata.Covariance`. If None, uncertainties are
+        assumed to be uncorrelated.
 
     Returns
     -------
@@ -275,12 +297,27 @@ def _find_spectral_column(table, columns_to_search, spectral_axis):
         if table[err_column].ndim > 1:
             err = table[err_column].T
         elif flux.ndim > 1:  # Repeat uncertainties over all flux columns
+            if correl is not None:
+                warnings.warn('When applying correlated errors, the dimensionality of the error '
+                              'array must match the dimensionality of the flux array.  Ignoring '
+                              'correlated errors.')
+                correl = None
             err = np.tile(table[err_column], flux.shape[0], 1)
         else:
             err = table[err_column]
-        err = StdDevUncertainty(err.to(err.unit))
-        if np.min(table[err_column]) <= 0.:
-            warnings.warn("Standard Deviation has values of 0 or less", AstropyUserWarning)
+        if correl is not None:
+            try:
+                err = Covariance.from_tables(err**2, correl, quiet=True)
+            except (ValueError, TypeError):
+                warnings.warn('Unable to parse correlation table into a Covariance object.  '
+                              'Ignoring correlation matrix data.')
+                correl = None
+        # NOTE: This is not an `else` or `elif` block in order to catch the
+        # change to correl=None when handling the excemption above.
+        if correl is None:
+            err = StdDevUncertainty(err.to(err.unit))
+            if np.min(table[err_column]) <= 0.:
+                warnings.warn("Standard Deviation has values of 0 or less", AstropyUserWarning)
     else:
         err = None
 

specutils/manipulation/extract_spectral_region.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from astropy import units as u
+from astropy.nddata import Covariance
 from ..spectra import Spectrum1D, SpectralRegion
 
 __all__ = ['extract_region', 'extract_bounding_spectral_region', 'spectral_slab']
@@ -189,6 +190,9 @@ def _get_joined_value(sps, key, unique_inds=None):
                 uncert = sps[0].uncertainty
                 if uncert is None:
                     return None
+                if isinstance(uncert, Covariance):
+                    raise NotImplementedError("Cannot yet combine spectral regions with "
+                                              "covariant uncertainties.")
                 uncert._array = np.concatenate([sp.uncertainty._array for sp in sps])
                 return uncert[unique_inds] if unique_inds is not None else uncert
             elif key in concat_keys or key == 'spectral_axis':

specutils/spectra/spectrum1d.py

@@ -5,7 +5,7 @@
 from astropy import units as u
 from astropy.utils.decorators import lazyproperty
 from astropy.utils.decorators import deprecated
-from astropy.nddata import NDUncertainty, NDIOMixin, NDArithmeticMixin
+from astropy.nddata import NDUncertainty, NDIOMixin, NDArithmeticMixin, Covariance
 
 from .spectral_axis import SpectralAxis
 from .spectrum_mixin import OneDSpectrumMixin
@@ -36,8 +36,10 @@ class Spectrum1D(OneDSpectrumMixin, NDCube, NDIOMixin, NDArithmeticMixin):
     Parameters
     ----------
     flux : `~astropy.units.Quantity`
-        The flux data for this spectrum. This can be a simple `~astropy.units.Quantity`,
-        or an existing `~Spectrum1D` or `~ndcube.NDCube` object.
+        The flux data for this spectrum. This can be a simple
+        `~astropy.units.Quantity`, or an existing `~Spectrum1D` or
+        `~ndcube.NDCube` object.  If an `~ndcube.NDCube` object, all other
+        arguments are ignored.
     spectral_axis : `~astropy.units.Quantity` or `~specutils.SpectralAxis`
         Dispersion information with the same shape as the last (or only)
         dimension of flux, or one greater than the last dimension of flux
@@ -60,8 +62,10 @@ class Spectrum1D(OneDSpectrumMixin, NDCube, NDIOMixin, NDArithmeticMixin):
         values represent edges of the wavelength bin, or centers of the bin.
     uncertainty : `~astropy.nddata.NDUncertainty`
         Contains uncertainty information along with propagation rules for
-        spectrum arithmetic. Can take a unit, but if none is given, will use
-        the unit defined in the flux.
+        spectrum arithmetic. Can take a unit, but if none is given, will use the
+        unit defined in the flux.  Note that functionality is limited for
+        `~astropy.nddata.Covariance` instances, particularly for
+        multidimensional data.
     mask : `~numpy.ndarray`-like
         Array where values in the flux to be masked are those that
         ``astype(bool)`` converts to True. (For example, integer arrays are not
@@ -211,7 +215,11 @@ def __init__(self, flux=None, spectral_axis=None, wcs=None,
                                                 len(kwargs["mask"].shape)-temp_axes[0]-1, -1)
                     if "uncertainty" in kwargs:
                         if kwargs["uncertainty"] is not None:
-                            if isinstance(kwargs["uncertainty"], NDUncertainty):
+                            if isinstance(kwargs["uncertainty"], Covariance):
+                                raise NotImplementedError("Cannot yet handle covariance for "
+                                                          "multidimensional Spectrum1D objects "
+                                                          "with a WCS coordinate system.")
+                            elif isinstance(kwargs["uncertainty"], NDUncertainty):
                                 # Account for Astropy uncertainty types
                                 unc_len = len(kwargs["uncertainty"].array.shape)
                                 temp_unc = np.swapaxes(kwargs["uncertainty"].array,
@@ -302,10 +310,13 @@ def __init__(self, flux=None, spectral_axis=None, wcs=None,
             raise ValueError('Spectral axis must be strictly increasing or decreasing.')
 
         if hasattr(self, 'uncertainty') and self.uncertainty is not None:
-            if not flux.shape == self.uncertainty.array.shape:
-                raise ValueError(
-                    "Flux axis ({}) and uncertainty ({}) shapes must be the "
-                    "same.".format(flux.shape, self.uncertainty.array.shape))
+            if isinstance(self.uncertainty, Covariance):
+                uncertainty_shape = self.uncertainty.data_shape
+            else:
+                uncertainty_shape = self.uncertainty.array.shape
+            if not flux.shape == uncertainty_shape:
+                raise ValueError(f"Flux axis ({flux.shape}) and uncertainty ({uncertainty_shape}) "
+                                 "shapes must be the same.")
 
     def __getitem__(self, item):
         """
@@ -322,7 +333,6 @@ def __getitem__(self, item):
         The first case is handled by the parent class, while the second is
         handled here.
         """
-
         if self.flux.ndim > 1 or (isinstance(item, tuple) and item[0] is Ellipsis):
             if isinstance(item, tuple):
                 if len(item) == len(self.flux.shape) or item[0] is Ellipsis:
@@ -371,11 +381,17 @@ def __getitem__(self, item):
             else:
                 new_meta = deepcopy(self.meta)
 
+            if isinstance(self.uncertainty, Covariance):
+                new_unc = self.uncertainty.sub_matrix(item)
+            elif self.uncertainty is not None:
+                new_unc = self.uncertainty[item]
+            else:
+                new_unc = None
+
             return self._copy(
                 flux=self.flux[item],
                 spectral_axis=self.spectral_axis[spec_item],
-                uncertainty=self.uncertainty[item]
-                if self.uncertainty is not None else None,
+                uncertainty=new_unc,
                 mask=self.mask[item] if self.mask is not None else None,
                 meta=new_meta, wcs=None)
 
@@ -748,8 +764,13 @@ def __str__(self):
 
         # Add information about uncertainties if available
         if self.uncertainty:
+            _arr = (
+                self.uncertainty.toarray()
+                if isinstance(self.uncertainty, Covariance)
+                else self.uncertainty.array
+            )
             result += (f'\nUncertainty={type(self.uncertainty).__name__} '
-                       f'({np.array2string(self.uncertainty.array, threshold=8)}'
+                       f'({np.array2string(_arr, threshold=8)}'
                        f' {self.uncertainty.unit})')
 
         return result

specutils/tests/test_loaders.py

@@ -6,14 +6,15 @@
 import pytest
 import astropy.units as u
 import numpy as np
+from scipy import sparse
 from astropy.io import fits
 from astropy.io.fits.verify import VerifyWarning
 from astropy.table import Table
 from astropy.units import UnitsWarning
 from astropy.wcs import FITSFixedWarning, WCS
 from astropy.io.registry import IORegistryError
 from astropy.modeling import models
-from astropy.nddata import StdDevUncertainty, InverseVariance, VarianceUncertainty
+from astropy.nddata import StdDevUncertainty, InverseVariance, VarianceUncertainty, Covariance
 from astropy.tests.helper import quantity_allclose
 from astropy.utils.exceptions import AstropyUserWarning
 
@@ -584,6 +585,45 @@ def test_tabular_fits_writer(tmp_path, spectral_axis):
                              spectrum.uncertainty.quantity)
 
 
+def test_tabular_fits_cov_io(tmp_path):
+    # Create a fake spectrum
+    wave = np.arange(4500., 5500., 1.) * u.AA
+    flux = np.full(len(wave), 10) * u.Jy
+    # And covariance
+    cov_diags = [
+        np.ones(1000, dtype=float),
+        np.full(1000-1, 0.5, dtype=float),
+        np.full(1000-2, 0.2, dtype=float),
+    ]
+    cov = Covariance(sparse.diags(cov_diags, [0, 1, 2]), unit=u.Jy**2)
+
+    # Create the Spectrum1D
+    spectrum = Spectrum1D(flux=flux, spectral_axis=wave, uncertainty=cov)
+    # Simple test of ingestion
+    assert np.array_equal(spectrum.uncertainty.toarray(), cov.toarray()), \
+            'Covariance not included correctly'
+    # Write it
+    tmpfile = str(tmp_path / '_tst.fits')
+    spectrum.write(tmpfile, format='tabular-fits')
+
+    # Check the output file
+    with fits.open(tmpfile) as hdu:
+        assert len(hdu) == 3, 'Should contain 3 extensions, primary, DATA, CORREL'
+        assert np.array_equal([h.name for h in hdu], ['PRIMARY', 'DATA', 'CORREL']), \
+                'Extension names are wrong'
+        assert all([h.__class__.__name__ == 'BinTableHDU' for h in hdu[1:]]), \
+                'Data extensions should both be BinTableHDU'
+        assert len(hdu['CORREL'].data) == cov.nnz, 'Number of non-zero cov elements mismatch'
+
+    # Read it
+    _spectrum = Spectrum1D.read(tmpfile)
+    assert spectrum.flux.unit == _spectrum.flux.unit
+    assert spectrum.spectral_axis.unit == _spectrum.spectral_axis.unit
+    assert quantity_allclose(spectrum.spectral_axis, _spectrum.spectral_axis)
+    assert quantity_allclose(spectrum.flux, _spectrum.flux)
+    assert np.array_equal(spectrum.uncertainty.toarray(), _spectrum.uncertainty.toarray())
+
+
 @pytest.mark.parametrize("ndim", range(1, 4))
 @pytest.mark.parametrize("spectral_axis",
                          ['wavelength', 'frequency', 'energy', 'wavenumber'])