diff --git a/CHANGES.rst b/CHANGES.rst index eedf081154..cfd6de6749 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -6,6 +6,8 @@ New Features * Model fitting results are logged in a table within the plugin [#2093]. +* Auto-identify a configuration/helper for a given data file [#2124] + Cubeviz ^^^^^^^ diff --git a/jdaviz/core/data_formats.py b/jdaviz/core/data_formats.py index e1e317badc..7caa29253b 100644 --- a/jdaviz/core/data_formats.py +++ b/jdaviz/core/data_formats.py @@ -2,12 +2,22 @@ import pathlib import astropy.io +from astropy.io import registry, fits +from astropy.nddata import CCDData +from astropy.wcs import WCS + from specutils.io.registers import identify_spectrum_format -from specutils import SpectrumList +from specutils import Spectrum1D, SpectrumList, SpectrumCollection +from stdatamodels import asdf_in_fits from jdaviz.core.config import list_configurations -__all__ = ['guess_dimensionality', 'get_valid_format', 'identify_data'] +__all__ = [ + 'guess_dimensionality', + 'get_valid_format', + 'identify_data', + 'identify_helper' +] # create a default file format to configuration mapping default_mapping = {'JWST x1d': 'specviz', 'JWST s2d': 'specviz2d', @@ -112,3 +122,153 @@ def identify_data(filename, current=None): raise ValueError('Mismatch between input file format and loaded configuration.') return valid_format, config + + +def _get_wcs(filename, header): + """ + Get gwcs.wcs.WCS or astropy.wcs.WCS from FITS file. + """ + try: + with asdf_in_fits.open(filename) as af: + wcs = af.tree['meta']['wcs'] + + # if the file doesn't have ASDF-in-FITS, then + # the 'meta' key doesn't exist, yielding a KeyError: + except KeyError: + # fall back on using astropy WCS: + wcs = WCS(header) + + return wcs + + +def identify_helper(filename, ext=1): + """ + Guess the appropriate viz helper for a data file. + + Parameters + ---------- + filename : str (path-like) + Name for a local data file. + ext : int + Extension from the FITS file. + + Returns + ------- + helper_name : str + Name of the best-guess helper for ``filename``. + """ + supported_dtypes = [ + Spectrum1D, + SpectrumList, + SpectrumCollection, + CCDData + ] + + if filename.lower().endswith('asdf'): + # ASDF files are only supported in jdaviz for + # Roman WFI 2D images, so suggest imviz: + return 'imviz' + + header = fits.getheader(filename, ext=ext) + data = fits.getdata(filename, ext=ext) + wcs = _get_wcs(filename, header) + has_spectral_axis = 'spectral' in wcs.world_axis_object_classes + + n_axes = ( + int(has_spectral_axis) + + + sum([component[0] in ['celestial', 'angle'] + for component in wcs.world_axis_object_components]) - + + # remove any slit_frame axis from the count + (0 if not hasattr(wcs, 'available_frames') else + int('slit_frame' in wcs.available_frames)) + ) + + # use astropy to recognize some data formats: + possible_formats = {} + for cls in supported_dtypes: + fmt = registry.identify_format( + 'read', cls, filename, None, {}, {} + ) + if fmt: + possible_formats[cls] = fmt + + # If CCDData is the only match: + if len(possible_formats) == 1: + only_key, only_value = possible_formats.popitem() + if only_key == CCDData: + # could be 2D spectrum or 2D image. break tie with WCS: + if has_spectral_axis: + if n_axes > 1: + return 'specviz2d' + return 'specviz' + elif not isinstance(data, fits.BinTableHDU): + return 'imviz' + + # Ensure specviz is chosen when ``data`` is a table or recarray + # and there's a "known" spectral column name: + if isinstance(data, (fits.BinTableHDU, fits.fitsrec.FITS_rec)): + # now catch spectra in FITS tables, looking for + # columns with "wave" or "flux" in the names: + table_columns = [getattr(col, 'name', col).lower() for col in data.columns] + + # these are "known" prefixes for column names + # in FITS tables of spectral observations + known_spectral_columns = [ + 'wave', + 'flux' + ] + + # this list of bools indicates any + # spectral column names found: + found_spectral_columns = [ + found_col.startswith(known_col) + for known_col in known_spectral_columns + for found_col in table_columns + ] + + # if at least one spectral column is found: + if sum(found_spectral_columns): + return 'specviz' + + # If the data could be spectral: + for cls in [Spectrum1D, SpectrumList]: + if cls in possible_formats.keys(): + recognized_spectrum_format = possible_formats[cls][0].lower() + + # first catch known JWST spectrum types: + if (n_axes == 3 and + recognized_spectrum_format.find('s3d') > -1): + return 'cubeviz' + elif (n_axes == 2 and + recognized_spectrum_format.find('x1d') > -1): + return 'specviz' + + # we intentionally don't choose specviz2d for + # data recognized as 's2d' as we did with the cases above, + # because 2D data products could be 2D spectra *or* 2D images + # that the registry recognizes as s2d. + + # Use WCS to break the tie below: + elif n_axes == 2: + if has_spectral_axis: + return 'specviz2d' + return 'imviz' + + elif n_axes == 1: + return 'specviz' + + try: + # try using the specutils registry: + valid_format, config = identify_data(filename) + return config + except ValueError: + # if file type not recognized: + pass + + if n_axes == 2 and not has_spectral_axis: + # at this point, non-spectral 2D data are likely images: + return 'imviz' + + raise ValueError(f"No helper could be auto-identified for {filename}.") diff --git a/jdaviz/core/tests/test_data_menu.py b/jdaviz/core/tests/test_data_menu.py index 2486303138..c48e5f3ec7 100644 --- a/jdaviz/core/tests/test_data_menu.py +++ b/jdaviz/core/tests/test_data_menu.py @@ -1,6 +1,30 @@ +import pytest import numpy as np from glue.core.roi import XRangeROI +from astropy.utils.data import download_file + +from jdaviz.core.data_formats import identify_helper + +# URIs to example JWST/HST files on MAST, and their +# corresponding jdaviz helpers: +example_uri_helper = [ + ['mast:HST/product/id4301ouq_drz.fits', 'imviz'], + ['mast:HST/product/ldq601030_x1dsum.fits', 'specviz'], + ['mast:HST/product/o4xw01dkq_flt.fits', 'specviz2d'], + ['mast:JWST/product/jw01324-o001_s00094_niriss_f200w-gr150c-gr150r_x1d.fits', + 'specviz'], + ['mast:JWST/product/jw01324-o006_s00005_nirspec_f100lp-g140h_s2d.fits', + 'specviz2d'], + ['mast:JWST/product/jw01345-o001_t021_nircam_clear-f200w_i2d.fits', 'imviz'], + ['mast:JWST/product/jw01373-o028_t001_nirspec_g395h-f290lp_s3d.fits', + 'cubeviz'], + ['mast:JWST/product/jw01373-o031_t007_miri_ch1-shortmediumlong_s3d.fits', + 'cubeviz'], + ['mast:JWST/product/jw01783-o004_t008_nircam_clear-f444w_i2d.fits', 'imviz'], + ['mast:JWST/product/jw02732-o004_t004_miri_ch1-shortmediumlong_x1d.fits', + 'specviz'] +] def test_data_menu_toggles(specviz_helper, spectrum1d): @@ -75,3 +99,27 @@ def test_visibility_toggle(imviz_helper): visible=True) assert iv.layers[0].visible is True assert po.stretch_preset.value == 90 + + +@pytest.mark.remote_data +@pytest.mark.filterwarnings(r"ignore::astropy.wcs.wcs.FITSFixedWarning") +@pytest.mark.parametrize( + "uri, expected_helper", example_uri_helper +) +def test_auto_config_detection(uri, expected_helper): + url = f'https://mast.stsci.edu/api/v0.1/Download/file/?uri={uri}' + fn = download_file(url, cache=True) + helper_name = identify_helper(fn) + assert helper_name == expected_helper + + +@pytest.mark.remote_data +@pytest.mark.filterwarnings(r"ignore::astropy.wcs.wcs.FITSFixedWarning") +def test_auto_config_manga(): + # Check that MaNGA cubes go to cubeviz. This file is + # originally from + # https://data.sdss.org/sas/dr14/manga/spectro/redux/v2_1_2/7495/stack/manga-7495-12704-LOGCUBE.fits.gz + URL = 'https://stsci.box.com/shared/static/28a88k1qfipo4yxc4p4d40v4axtlal8y.fits' + fn = download_file(URL, cache=True) + helper_name = identify_helper(fn) + assert helper_name == 'cubeviz'