From 3c1400179c7ab2989ef77fa5198cecc0882768c7 Mon Sep 17 00:00:00 2001 From: Ricky O'Steen <39831871+rosteen@users.noreply.github.com> Date: Fri, 1 Nov 2024 14:04:01 -0400 Subject: [PATCH] Add ability to load 2D flux arrays in Specviz (#3229) * Add ability to load 2D flux arrays in Specviz by splitting into separate spectra * Changelog * Handle this in file load case * Fix typo * Add test and debug * Codestyle again, whoops * Add keyword to force Specviz parser to use SpectrumList.read * Add new keyword to helper load_data * Implement suggestions from @rileythai Co-authored-by: rileythai * Add more test coverage, debug * Codestyle Codestyle * Apparently need this to standardize the input label * One more label handling fix * Apply suggestions from code review Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com> * Add clarity to test and changelog * Add meta, debug --------- Co-authored-by: rileythai Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com> --- CHANGES.rst | 2 + jdaviz/configs/specviz/helper.py | 6 +- jdaviz/configs/specviz/plugins/parsers.py | 91 ++++++++++++++++----- jdaviz/configs/specviz/tests/test_helper.py | 21 +++++ 4 files changed, 98 insertions(+), 22 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 131aafbd33..2b1a782015 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -17,6 +17,8 @@ Mosviz Specviz ^^^^^^^ +- Specviz parser will now split a spectrum with a 2D flux array into multiple spectra on load + (useful for certain SDSS file types). [#3229] Specviz2d ^^^^^^^^^ diff --git a/jdaviz/configs/specviz/helper.py b/jdaviz/configs/specviz/helper.py index 7eefe9e212..c13ac67175 100644 --- a/jdaviz/configs/specviz/helper.py +++ b/jdaviz/configs/specviz/helper.py @@ -41,7 +41,8 @@ def __init__(self, *args, **kwargs): handler=self._redshift_listener) def load_data(self, data, data_label=None, format=None, show_in_viewer=True, - concat_by_file=False, cache=None, local_path=None, timeout=None): + concat_by_file=False, cache=None, local_path=None, timeout=None, + load_as_list=False): """ Load data into Specviz. @@ -80,7 +81,8 @@ def load_data(self, data, data_label=None, format=None, show_in_viewer=True, concat_by_file=concat_by_file, cache=cache, local_path=local_path, - timeout=timeout) + timeout=timeout, + load_as_list=load_as_list) def get_spectra(self, data_label=None, spectral_subset=None, apply_slider_redshift="Warn"): """Returns the current data loaded into the main viewer diff --git a/jdaviz/configs/specviz/plugins/parsers.py b/jdaviz/configs/specviz/plugins/parsers.py index 51fbfc1ee6..8af49c2bc3 100644 --- a/jdaviz/configs/specviz/plugins/parsers.py +++ b/jdaviz/configs/specviz/plugins/parsers.py @@ -18,7 +18,8 @@ @data_parser_registry("specviz-spectrum1d-parser") def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_viewer=True, - concat_by_file=False, cache=None, local_path=os.curdir, timeout=None): + concat_by_file=False, cache=None, local_path=os.curdir, timeout=None, + load_as_list=False): """ Loads a data file or `~specutils.Spectrum1D` object into Specviz. @@ -46,18 +47,32 @@ def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_v remote requests in seconds (passed to `~astropy.utils.data.download_file` or `~astroquery.mast.Conf.timeout`). + load_as_list : bool, optional + Force the parser to load the input file with the `~specutils.SpectrumList` read function + instead of `~specutils.Spectrum1D`. """ spectrum_viewer_reference_name = app._jdaviz_helper._default_spectrum_viewer_reference_name # If no data label is assigned, give it a unique name if not data_label: data_label = app.return_data_label(data, alt_name="specviz_data") + # Still need to standardize the label + elif isinstance(data_label, (list, tuple)): + data_label = [app.return_data_label(label, alt_name="specviz_data") for label in data_label] # noqa + else: + data_label = app.return_data_label(data_label, alt_name="specviz_data") + if isinstance(data, SpectrumCollection): raise TypeError("SpectrumCollection detected." " Please provide a Spectrum1D or SpectrumList") elif isinstance(data, Spectrum1D): - data_label = [app.return_data_label(data_label, alt_name="specviz_data")] - data = [data] + # Handle the possibility of 2D spectra by splitting into separate spectra + if data.flux.ndim == 1: + data_label = [data_label] + data = [data] + elif data.flux.ndim == 2: + data = split_spectrum_with_2D_flux_array(data) + data_label = [f"{data_label} [{i}]" for i in range(len(data))] # No special processing is needed in this case, but we include it for completeness elif isinstance(data, SpectrumList): pass @@ -65,40 +80,48 @@ def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_v # special processing for HDUList if isinstance(data, fits.HDUList): data = [Spectrum1D.read(data)] - data_label = [app.return_data_label(data_label, alt_name="specviz_data")] + data_label = [data_label] else: # list treated as SpectrumList if not an HDUList data = SpectrumList.read(data, format=format) else: # try parsing file_obj as a URI/URL: - data = download_uri_to_path( - data, cache=cache, local_path=local_path, timeout=timeout - ) + data = download_uri_to_path(data, cache=cache, local_path=local_path, timeout=timeout) path = pathlib.Path(data) - if path.is_file(): + if path.is_dir() or load_as_list: + data = SpectrumList.read(str(path), format=format) + if data == []: + raise ValueError(f"`specutils.SpectrumList.read('{str(path)}')` " + "returned an empty list") + elif path.is_file(): try: - data = [Spectrum1D.read(str(path), format=format)] - data_label = [app.return_data_label(data_label, alt_name="specviz_data")] + data = Spectrum1D.read(str(path), format=format) + if data.flux.ndim == 2: + data = split_spectrum_with_2D_flux_array(data) + else: + data = [data] + data_label = [app.return_data_label(data_label, alt_name="specviz_data")] except IORegistryError: # Multi-extension files may throw a registry error data = SpectrumList.read(str(path), format=format) - elif path.is_dir(): - data = SpectrumList.read(str(path), format=format) - if data == []: - raise ValueError(f"`specutils.SpectrumList.read('{str(path)}')` " - "returned an empty list") else: raise FileNotFoundError("No such file: " + str(path)) + # step through SpectrumList and convert any 2D spectra to 1D spectra if isinstance(data, SpectrumList): + new_data = [] + for spec in data: + if spec.flux.ndim == 2: + new_data.extend(split_spectrum_with_2D_flux_array(spec)) + else: + new_data.append(spec) + data = SpectrumList(new_data) + if not isinstance(data_label, (list, tuple)): - temp_labels = [] - for i in range(len(data)): - temp_labels.append(f"{data_label} {i}") - data_label = temp_labels + data_label = [f"{data_label} [{i}]" for i in range(len(data))] elif len(data_label) != len(data): raise ValueError(f"Length of data labels list ({len(data_label)}) is different" f" than length of list of data ({len(data)})") @@ -121,7 +144,7 @@ def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_v # if the concatenated list of uncertanties is all nan (meaning # they were all nan to begin with, or all None), it will be set # to None on the final Spectrum1D - if spec.uncertainty[wlind] is not None: + if spec.uncertainty is not None and spec.uncertainty[wlind] is not None: dfnuallorig.append(spec.uncertainty[wlind].array) else: dfnuallorig.append(np.nan) @@ -246,3 +269,31 @@ def combine_lists_to_1d_spectrum(wl, fnu, dfnu, wave_units, flux_units): spec = Spectrum1D(flux=fnuall * flux_units, spectral_axis=wlall * wave_units, uncertainty=unc) return spec + + +def split_spectrum_with_2D_flux_array(data): + """ + Helper function to split Spectrum1D of 2D flux to a SpectrumList of nD objects. + + Parameters + ---------- + data : `~specutils.Spectrum1D` + Spectrum with 2D flux array + + Returns + ------- + new_data : `~specutils.SpectrumList` + List of unpacked spectra. + """ + new_data = [] + for i in range(data.flux.shape[0]): + unc = None + mask = None + if data.uncertainty is not None: + unc = data.uncertainty[i, :] + if data.mask is not None: + mask = data.mask[i, :] + new_data.append(Spectrum1D(flux=data.flux[i, :], spectral_axis=data.spectral_axis, + uncertainty=unc, mask=mask, meta=data.meta)) + + return new_data diff --git a/jdaviz/configs/specviz/tests/test_helper.py b/jdaviz/configs/specviz/tests/test_helper.py index 6ca8c67064..8e87125071 100644 --- a/jdaviz/configs/specviz/tests/test_helper.py +++ b/jdaviz/configs/specviz/tests/test_helper.py @@ -405,6 +405,27 @@ def test_load_spectrum_list_directory_concat(tmpdir, specviz_helper): assert len(specviz_helper.app.data_collection) == 41 +def test_load_2d_flux(specviz_helper): + # Test loading a spectrum with a 2D flux, which should be split into separate + # 1D Spectrum1D objects to load in Specviz. + spec = Spectrum1D(spectral_axis=np.linspace(4000, 6000, 10)*u.Angstrom, + flux=np.ones((4, 10))*u.Unit("1e-17 erg / (Angstrom cm2 s)")) + specviz_helper.load_data(spec, data_label="test") + + assert len(specviz_helper.app.data_collection) == 4 + assert specviz_helper.app.data_collection[0].label == "test [0]" + + spec2 = Spectrum1D(spectral_axis=np.linspace(4000, 6000, 10)*u.Angstrom, + flux=np.ones((2, 10))*u.Unit("1e-17 erg / (Angstrom cm2 s)")) + + # Make sure 2D spectra in a SpectrumList also get split properly. + spec_list = SpectrumList([spec, spec2]) + specviz_helper.load_data(spec_list, data_label="second test") + + assert len(specviz_helper.app.data_collection) == 10 + assert specviz_helper.app.data_collection[-1].label == "second test [5]" + + def test_plot_uncertainties(specviz_helper, spectrum1d): specviz_helper.load_data(spectrum1d)