Merge pull request #621 from havok2063/multiformat

Multi-extension data format identification

jdaviz/core/data_formats.py

@@ -4,18 +4,21 @@
 
 import astropy.io
 from specutils.io.registers import identify_spectrum_format
-from specutils import Spectrum1D
+from specutils import SpectrumList
 
 from jdaviz.core.config import list_configurations
 
 
 # create a default file format to configuration mapping
 default_mapping = {'JWST x1d': 'specviz', 'JWST s2d': 'specviz2d',
                    'JWST s3d': 'cubeviz', 'MaNGA cube': 'cubeviz',
-                   'MaNGA rss': 'imviz'}
+                   'MaNGA rss': 'specviz'}
 
-formats_table = astropy.io.registry.get_formats(data_class=Spectrum1D,
-                                                readwrite='Read')
+# get formats table for specutils objects
+formats_table = astropy.io.registry.get_formats(readwrite='Read')
+formats_table.add_index('Data class')
+formats_table = formats_table.loc[['Spectrum1D', 'SpectrumList']]
+formats_table.sort(['Data class', 'Format'])
 
 file_to_config_mapping = {i: default_mapping.get(
     i, 'specviz') for i in formats_table['Format']}
@@ -63,7 +66,7 @@ def get_valid_format(filename):
             The recommended application configuration
     """
 
-    valid_file_format = identify_spectrum_format(filename)
+    valid_file_format = identify_spectrum_format(filename, SpectrumList)
     ndim = guess_dimensionality(filename)
 
     if valid_file_format:

jdaviz/tests/test_data_formats.py

@@ -6,54 +6,69 @@
 from jdaviz.core.config import list_configurations
 
 
-def create_sdss(n_dim=2):
+def create_sdss(n_dim=2, **kwargs):
     """ create a fake SDSS single fiber spectral fits file """
     primary = fits.PrimaryHDU(header=fits.Header({'TELESCOP': "SDSS 2.5-M", 'FIBERID': 555}))
     ext = fits.BinTableHDU(name='DATA', header=fits.Header({'TTYPE3': "ivar"}))
-    hdulist = fits.HDUList([primary, ext])
-    return hdulist
+    return fits.HDUList([primary, ext])
 
 
-def create_manga(n_dim):
+def create_manga(n_dim=None, **kwargs):
     """ create a fake SDSS MaNGA fits file """
     primary = fits.PrimaryHDU(header=fits.Header({'TELESCOP': "SDSS 2.5-M"}))
     ext = fits.ImageHDU(name='FLUX', header=fits.Header({'INSTRUME': "MaNGA"}))
     shape = (1,) if n_dim == 1 else (1, 2,) if n_dim == 2 else (1, 2, 3) if n_dim == 3 else (1,)
     ext.data = np.empty(shape)
-    hdulist = fits.HDUList([primary, ext])
-    return hdulist
+    return fits.HDUList([primary, ext])
 
 
-def create_jwst(n_dim):
+def _create_bintable(name, ver=1):
+    return fits.BinTableHDU.from_columns(
+                [fits.Column(name='target', format='20A', array=np.ones(3)),
+                 fits.Column(name='V_mag', format='E', array=np.ones(3))],
+                name=name, ver=ver
+            )
+
+
+def create_jwst(n_dim=None, multi=None, ext=None):
     """ create a fake JWST fits file """
     primary = fits.PrimaryHDU(header=fits.Header({'TELESCOP': "JWST"}))
-    ext = fits.ImageHDU(name='ASDF')
+    asdf = fits.ImageHDU(name='ASDF')
     sci = fits.ImageHDU(name='SCI')
-    exts = [primary, ext, sci]
+    exts = [primary]
     if n_dim == 1:
-        ex1d = fits.ImageHDU(name='EXTRACT1D')
+        ex1d = _create_bintable(ext, ver=1)
         exts.append(ex1d)
-    shape = (1,) if n_dim == 1 else (1, 2,) if n_dim == 2 else (1, 2, 3) if n_dim == 3 else (1,)
+        # create two more extensions
+        if multi:
+            for i in range(1, 3):
+                ex1d = _create_bintable(ext, ver=i)
+                exts.append(ex1d)
+    shape = (1,) if n_dim == 1 else (1, 2,) if n_dim == 2 else (1, 2, 3) if n_dim == 3 else (1, )
     sci.data = np.empty(shape)
-    hdulist = fits.HDUList(exts)
-    return hdulist
+    if n_dim > 1:
+        exts.append(sci)
+    exts.append(asdf)
+    return fits.HDUList(exts)
 
 
-def create_generic(n_dim=1):
+def create_generic(n_dim=1, **kwargs):
     """ create a generic fits file of dimension N """
     primary = fits.PrimaryHDU()
     ext = fits.ImageHDU(name='DATA')
     shape = (1,) if n_dim == 1 else (1, 2,) if n_dim == 2 else (1, 2, 3) if n_dim == 3 else (1,)
     ext.data = np.empty(shape)
-    hdulist = fits.HDUList([primary, ext])
-    return hdulist
+    return fits.HDUList([primary, ext])
 
 
 data = {'MaNGA cube': {'ndim': 3, 'fxn': create_manga},
         'MaNGA rss': {'ndim': 2, 'fxn': create_manga},
         'JWST s3d': {'ndim': 3, 'fxn': create_jwst},
         'JWST s2d': {'ndim': 2, 'fxn': create_jwst},
         'JWST x1d': {'ndim': 1, 'fxn': create_jwst},
+        'JWST c1d': {'ndim': 1, 'fxn': create_jwst, 'ext': 'COMBINE1D'},
+        'JWST x1d multi': {'ndim': 1, 'fxn': create_jwst, 'multi': True},
+        'JWST c1d multi': {'ndim': 1, 'fxn': create_jwst, 'ext': 'COMBINE1D', 'multi': True},
         'SDSS-III/IV spec': {'ndim': 2, 'fxn': create_sdss},
         'generic 3d': {'ndim': 3, 'fxn': create_generic},
         'generic 1d': {'ndim': 1, 'fxn': create_generic}
@@ -72,7 +87,9 @@ def _create_fits(name):
         # generate fake fits.HDUList
         ndim = data[name]['ndim']
         fxn = data[name]['fxn']
-        hdulist = fxn(n_dim=ndim)
+        multi = data[name].get('multi', None)
+        ext = data[name].get('ext', 'EXTRACT1D')
+        hdulist = fxn(n_dim=ndim, ext=ext, multi=multi)
         hdulist.writeto(filepath)
         return str(filepath)
 
@@ -81,10 +98,13 @@ def _create_fits(name):
 
 @pytest.mark.parametrize('name, expconf',
                          [('MaNGA cube', 'cubeviz'),
-                          ('MaNGA rss', 'imviz'),
+                          ('MaNGA rss', 'specviz'),
                           ('JWST s3d', 'cubeviz'),
                           ('JWST s2d', 'specviz2d'),
                           ('JWST x1d', 'specviz'),
+                          ('JWST c1d', 'specviz'),
+                          ('JWST x1d multi', 'specviz'),
+                          ('JWST c1d multi', 'specviz'),
                           ('SDSS-III/IV spec', 'specviz'),
                           ('generic 3d', 'cubeviz'),
                           ('generic 1d', 'specviz')])
@@ -108,7 +128,6 @@ def test_list_configurations():
 
 @pytest.mark.parametrize('name, expconf, expstat',
                          [('MaNGA cube', 'cubeviz', 'Success: Valid Format'),
-                          ('MaNGA rss', 'imviz', 'Success: Valid Format'),
                           ('generic 1d', 'specviz', 'Error: Cannot determine format of '
                                                     'the file to load.')])
 def test_identify_data(create_fake_fits, name, expconf, expstat):