diff --git a/docs/custom_loading.rst b/docs/custom_loading.rst index ea706aaab..24af69a6f 100644 --- a/docs/custom_loading.rst +++ b/docs/custom_loading.rst @@ -131,8 +131,61 @@ file. For the general case where none of the spectra are assumed to be the same length, the loader should return a `~specutils.SpectrumList`. Consider the custom JWST data loader as an example: -.. literalinclude:: ../specutils/io/default_loaders/jwst_reader.py - :language: python +.. code-block:: python + + from specutils import Spectrum, SpectrumList + from specutils.io.registers import data_loader + + @data_loader( + "JWST x1d multi", identifier=identify_jwst_x1d_multi_fits, + dtype=SpectrumList, extensions=['fits'], priority=10, + ) + def jwst_x1d_multi_loader(file_obj, **kwargs): + """Loader for JWST x1d 1-D spectral data in FITS format""" + return _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', **kwargs) + + def _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', flux_col=None, **kwargs): + """Implementation of loader for JWST x1d 1-D spectral data in FITS format""" + + if extname not in ['COMBINE1D', 'EXTRACT1D']: + raise ValueError('Incorrect extname given for 1d spectral data.') + + spectra = [] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + + primary_header = hdulist["PRIMARY"].header + + for hdu in hdulist: + # Read only the BinaryTableHDUs named COMBINE1D/EXTRACT1D and SCI + if hdu.name != extname: + continue + + header = hdu.header + + # Correct some known bad unit strings before reading the table + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for c in hdu.columns: + if c.unit in bad_units: + c.unit = bad_units[c.unit] + + data = QTable.read(hdu) + + if data[0]['WAVELENGTH'].shape != (): + # In this case we have multiple spectra packed into a single extension, one target + # per row of the table + for row in data: + if hasattr(row['WAVELENGTH'], 'mask') and np.all(row['WAVELENGTH'].mask): + # If everything is masked out we don't bother to read it in at all + continue + srctype = row['SOURCE_TYPE'] + spec = _jwst_spectrum_from_table(row, header, primary_header, flux_col, srctype) + spectra.append(spec) + else: + # Otherwise the whole table is defining a single spectrum + spec = _jwst_spectrum_from_table(data, header, primary_header, flux_col) + spectra.append(spec) + + return SpectrumList(spectra) Note that by default, any loader that uses ``dtype=Spectrum`` will also automatically add a reader for `~specutils.SpectrumList`. This enables user @@ -142,6 +195,229 @@ many `~specutils.Spectrum` objects. This method is available since `~specutils.SpectrumList` makes use of the Astropy IO registry (see `astropy.io.registry.read`). +Lazy Loading +^^^^^^^^^^^^ + +By default, `~specutils.SpectrumList` data loaders will load all spectra eagerly. Loaders optionally support +lazy loading so that individual spectra are only loaded into memory when accessed. This can be useful for lists with a +large number of spectra. + +Implementation +~~~~~~~~~~~~~~ +Lazy loading is opt-in per data loader. To implement lazy loading for a given data loader, define a custom function to be +passed into the ``lazy_loader`` argument of the ``@data_loader`` decorator. The loader function should return a ``SpectrumList`` built using +the :meth:`specutils.SpectrumList.from_lazy` class method. The function should: + +* Determine the total number of spectra. +* Define an index-based loader function that returns a single ``Spectrum``. +* Return the resulting ``SpectrumList``. + +See the following example for the Roman 1d spectra asdf data loader. + +.. code-block:: python + + def _lazy_loader(file_obj, **kwargs): + """Lazy loader for Roman spectra""" + # read in the input file + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + # get the roman spectral source ids + sources = list(roman["data"].keys()) + + def _loader(i: int) -> Spectrum: + """Function to load a single spectra from the input file given a list index""" + # select the proper source + source = sources[i] + with read_fileobj_or_asdftree(file_obj, **kwargs) as af2: + roman2 = af2["roman"] + # load a single Spectrum + return _load_roman_spectrum(roman2, source) + + # create the lazy SpectrumList, pass in the number of spectra and the individual spectrum loader + sl = SpectrumList.from_lazy(length=len(sources), loader=_loader) + return sl + + + @data_loader( + "Roman 1d combined", + identifier=identify_1d_combined, # standard function for format identification + dtype=SpectrumList, + extensions=["asdf"], + priority=10, + force=True, + lazy_loader=_lazy_loader, # function to handle lazy loading + ) + def roman_1d_combined_list(file_obj, **kwargs): + """Load all Roman 1d combined extracted spectra""" + # standard eager loading of all spectra + spectra = SpectrumList() + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + meta = roman["meta"] + # load the spectra + for source in roman["data"]: + # load single spectrum + spectrum = _load_roman_spectrum(roman, source) + spectra.append(spectrum) + + return spectra + + + def _load_roman_spectrum(roman: dict, source: str) -> Spectrum: + """Load a single Roman spectrum""" + meta = copy.deepcopy(roman["meta"]) + meta['source_id'] = source + data = roman["data"][source] if source else roman["data"] + flux = data['flux'] * u.Unit(meta["unit_flux"]) + flux_err = StdDevUncertainty(data['flux_error']) + wavelength = data['wl'] * u.Unit(meta["unit_wl"]) + return Spectrum(spectral_axis=wavelength, flux=flux, uncertainty=flux_err, meta=meta) + +Usage +~~~~~ + +Once implmemented, lazy loading can be activated by passing ``lazy_load=True`` to ``SpectrumList.read``. +This creates a list of placeholder objects of length equal to the number of spectra loaded into the list. +The ``repr`` indicates a lazy list with how many spectra are currently loaded into memory + +.. code-block:: python + + # example file with 6 spectral sources + speclist = SpectrumList.read("/path/to/roman.asdf", format="Roman 1d combined", lazy_load=True) + + speclist + lazy list: 0 items loaded; access an index to load a spectrum: + [, , + , , + , ] + + # inspect the lazy list + len(speclist) + 6 + + # verify it is lazy + speclist.is_lazy + True + + # check how many are loaded + speclist.n_loaded + 0 + +Accessing a list item will lazily load the corresponding spectrum into memory. + +.. code-block:: python + + # access the first spectrum + speclist[0] + (length=275); uncertainty=StdDevUncertainty)> + + # check the repr again + speclist + lazy list: 1 items loaded; access an index to load a spectrum: + [ (length=275); uncertainty=StdDevUncertainty)>, + , , , + , ] + + # check how many are loaded + speclist.n_loaded + 1 + +**Optional Labels** +You can optionally pass a list of labels to use as placeholder values in the lazy list repr, instead of +the default pointer ````. This can be done by passing a list of strings to the ``labels`` argument +of ``SpectrumList.from_lazy`` in the lazy loader function. + +.. code-block:: python + + def _lazy_load_roman(file_obj, **kwargs): + """Lazy loader for SpectrumList""" + + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + # create a list of roman source ids + sources = list(roman["data"].keys()) + + def _loader(i: int) -> Spectrum: + ... + + # pass the source ids as placeholder labels + sl = SpectrumList.from_lazy( + length=len(sources), loader=_loader, labels=sources + ) + return sl + +Loading the lazy list with display these labels instead: + +.. code-block:: python + + speclist = SpectrumList.read("/path/to/roman.asdf", format="Roman 1d combined", lazy_load=True) + + speclist + lazy list: 0 items loaded; access an index to load a spectrum: + ['402849', '403613', '403686', '404935', '404979', '414981'] + + +.. note:: + + Lazy loaders can be outfitted to any existing data loader. See the example data loader for loading + ``SDSS-V spec`` formatted FITS files. + + +Alternate List Indexing +^^^^^^^^^^^^^^^^^^^^^^^ + +Alternate ID labels allow string indexing of a ``SpectrumList``. This is useful for long lists of +spectra that can be more easily identified by a name or ID rather than a list index. Alternate IDs +are optional, and can be added to any ``SpectrumList`` data loader with the :meth:`specutils.SpectrumList.set_id_map` +class method. This method accepts a dictionary mapping of string labels to list indices. + +For example, + +.. code-block:: python + + from specutils import SpectrumList + + # instantiate a SpectrumList + ss = SpectrumList(['a', 'b', 'c']) + ss + ['a', 'b', 'c'] + + # set an alternate id indexing + ss.set_id_map({'spec1': 0, 'spec2': 1, 'spec3': 2}) + + # access an item with a list index + ss[0] + 'a' + + # access an item with an alternate id + ss['spec1'] + 'a' + +The example Roman data loaders uses a string target source id as alternate ids. + +.. code-block:: python + + def _load_roman_multisource(file_obj, **kwargs): + """Load all Roman spectra into a SpectrumList""" + + spectra = SpectrumList() + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + meta = roman["meta"] + sources = list(roman['data'].keys()) + + # set the alternate ids to roman source ids + source_idx_map = dict(zip(sources, range(len(sources)))) + spectra.set_id_map(source_idx_map) + + # load the spectra + for source in roman["data"]: + spectrum = _load_roman_spectrum(roman, source) + spectra.append(spectrum) + + return spectra .. _custom_writer: diff --git a/specutils/io/default_loaders/roman.py b/specutils/io/default_loaders/roman.py new file mode 100644 index 000000000..3f4b8ab0e --- /dev/null +++ b/specutils/io/default_loaders/roman.py @@ -0,0 +1,313 @@ +import copy +import warnings + +import astropy.units as u +from astropy.nddata import StdDevUncertainty +from astropy.utils.exceptions import AstropyUserWarning + +from ...spectra import Spectrum, SpectrumList +from ..parsing_utils import read_fileobj_or_asdftree +from ..registers import data_loader + +# Optional Roman GDPS package +# if the user has this installed, roman spectral files serialize into pydantic datamodels +try: + import roman_gdps +except ImportError: + roman_gdps = None + + +def _identify_roman_1d(*args, mode: str, units: str=None, **kwargs) -> bool: + """Identify a Roman SSC spectral file + + Parameters + ---------- + mode : str + the spectral mode + units : str, optional + the flux units substring to check, by default None + + Returns + ------- + bool + whether the input matches the given conditions + """ + with read_fileobj_or_asdftree(*args, **kwargs) as af: + # fail if not a roman asdf file + if "asdf_library" not in af.tree or "roman" not in af.tree: + return False + + # normalize to standard python dicts + if roman_gdps: + roman = af.tree["roman"].model_dump() + else: + roman = af.tree.get("roman", {}) + + # set base condition + meta = roman.get("meta", {}) + base = (isinstance(roman, dict) and + (meta.get('optical_element') or meta.get('instr_optical_element')) in ("GRISM", "PRISM") and + "data" in roman and isinstance(roman["data"], dict) + ) + + # check specific conditions + if mode == 'single': + return base and "flux" in roman["data"] + if mode == 'multi': + return base and "flux" not in roman["data"] and "1d_spectra" in roman["meta"].get("title", "") and units in roman['meta']['unit_flux'] + + +def identify_1d_single_source(origin, *args, **kwargs): + """Check if input is a Roman single source 1d extracted spectrum""" + + return _identify_roman_1d(*args, mode='single', **kwargs) + +def identify_1d_combined(origin, *args, **kwargs): + """Check if input is a set of Roman 1d combined extracted spectra""" + + return _identify_roman_1d(*args, mode='multi', units='W m**(-2)', **kwargs) + + +def identify_1d_individual(origin, *args, **kwargs): + """Check if input is a set of Roman 1d individual extracted spectra""" + + return _identify_roman_1d(*args, mode='multi', units='DN/s', **kwargs) + + +def _load_roman_spectrum(roman: dict, source: str) -> Spectrum: + """Load a single Roman spectrum + + Create and return a Spectrum object from a Roman SSC input dictionary. + + Parameters + ---------- + roman : dict + A collection of Roman spectral data and metadata. + source : str + A source identifier string. + + Returns + ------- + Spectrum + the Roman Spectrum object + """ + + meta = copy.deepcopy(roman["meta"]) + meta['source_id'] = source + data = roman["data"][source] if source else roman["data"] + flux = data['flux'] * u.Unit(meta["unit_flux"]) + flux_err = StdDevUncertainty(data['flux_error']) + wavelength = data['wl'] * u.Unit(meta["unit_wl"]) + return Spectrum(spectral_axis=wavelength, flux=flux, uncertainty=flux_err, meta=meta) + + +def _load_roman_first(file_obj, source=None, **kwargs): + """Load the Roman spectrum for the first source found""" + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + if source is None: + source = next(iter(roman['data'])) + warnings.warn(f"source not specified. Loading first source found: {source}.", + AstropyUserWarning) + elif source not in roman['data']: + raise ValueError(f"Invalid source! Source {source} is not spectra.") + + return _load_roman_spectrum(roman, source) + + +def _lazy_load_roman(file_obj, **kwargs): + """Lazy loader for SpectrumList""" + + # read the file and get the sources + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + sources = list(roman["data"].keys()) + source_idx_map = dict(zip(sources, range(len(sources)))) + + # define the lazy loader + def _loader(i: int) -> Spectrum: + source = sources[i] + with read_fileobj_or_asdftree(file_obj, **kwargs) as af2: + roman2 = af2["roman"] + return _load_roman_spectrum(roman2, source) + + # create the lazy object + cache_size = kwargs.get("cache_size", None) + sl = SpectrumList.from_lazy( + length=len(sources), loader=_loader, cache_size=cache_size, labels=sources + ) + + # set the sources ids as alternate indices + sl.set_id_map(source_idx_map) + + return sl + + +def _load_roman_multisource(file_obj, **kwargs) -> SpectrumList: + """Load all Roman spectra into a SpectrumList""" + + spectra = SpectrumList() + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + meta = roman["meta"] + sources = list(roman['data'].keys()) + + # set the alternate ids to roman source ids + source_idx_map = dict(zip(sources, range(len(sources)))) + meta['source_map'] = source_idx_map + spectra.set_id_map(source_idx_map) + # load the spectra + for source in roman["data"]: + spectrum = _load_roman_spectrum(roman, source) + spectra.append(spectrum) + + return spectra + + +@data_loader( + "Roman 1d spectra", identifier=identify_1d_single_source, dtype=Spectrum, + extensions=['asdf'], priority=10, +) +def roman_1d_spectrum_loader(file_obj, **kwargs): + """Load a Roman single source 1d extracted spectrum + + Loader for Roman 1d extracted spectral data in ASDF format. + + Parameters + ---------- + 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()). + + Returns + ------- + Spectrum + The spectrum contained in the file. + """ + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: + roman = af["roman"] + meta = roman["meta"] + data = roman["data"] + flux = data['flux'] * u.Unit(meta["unit_flux"]) + flux_err = StdDevUncertainty(data['flux_error']) + wavelength = data['wl'] * u.Unit(meta["unit_wl"]) + return Spectrum(spectral_axis=wavelength, flux=flux, uncertainty=flux_err, meta=meta) + + + +@data_loader( + "Roman 1d combined", + identifier=identify_1d_combined, + dtype=Spectrum, + extensions=["asdf"], + priority=10, + autogenerate_spectrumlist=False, +) +def roman_1d_combined(file_obj, source: str = None, **kwargs): + """Load Roman 1d combined extracted spectra + + Loads a set of extracted spectra from a Roman 1d combined ASDF file. + If not source is specified, it loads the first source found in the data dict. + + Parameters + ---------- + file_obj + the input file or file object + source : str, optional + the source id string, by default None + + Returns + ------- + Spectrum + The Roman spectrum object + """ + return _load_roman_first(file_obj, source=source, **kwargs) + + +@data_loader( + "Roman 1d combined", + identifier=identify_1d_combined, + dtype=SpectrumList, + extensions=["asdf"], + priority=10, + force=True, + lazy_loader=_lazy_load_roman, +) +def roman_1d_combined_list(file_obj, **kwargs): + """Load all Roman 1d combined extracted spectra + + Loads a set of extracted spectra from a Roman 1d combined ASDF file, + as a list of Spectrum objects. + + Parameters + ---------- + file_obj + the input file or file object + + Returns + ------- + SpectrumList + A set of Roman spectrum object + """ + return _load_roman_multisource(file_obj, **kwargs) + + +@data_loader( + "Roman 1d individual", + identifier=identify_1d_individual, + dtype=Spectrum, + extensions=["asdf"], + priority=10, + autogenerate_spectrumlist=False, +) +def roman_1d_individual(file_obj, source: str=None, **kwargs): + """Load Roman 1d individual extracted spectra + + Loads a set of extracted spectra from a Roman 1d individual ASDF file. + If not source is specified, it loads the first source found in the data dict. + + Parameters + ---------- + file_obj + the input file or file object + source : str, optional + the source id string, by default None + + Returns + ------- + Spectrum + The Roman spectrum object + """ + return _load_roman_first(file_obj, source=source, **kwargs) + + +@data_loader( + "Roman 1d individual", + identifier=identify_1d_individual, + dtype=SpectrumList, + extensions=["asdf"], + priority=10, + force=True, + lazy_loader=_lazy_load_roman, +) +def roman_1d_individual_list(file_obj, **kwargs): + """Load all Roman 1d individual extracted spectra + + Loads a set of extracted spectra from a Roman 1d individual ASDF file, + as a list of Spectrum objects. + + Parameters + ---------- + file_obj + the input file or file object + source : str, optional + the source id string, by default None + + Returns + ------- + SpectrumList + A set of Roman spectrum object + """ + return _load_roman_multisource(file_obj, **kwargs) + + diff --git a/specutils/io/default_loaders/sdss_v.py b/specutils/io/default_loaders/sdss_v.py index d3575f3d9..ca32ba566 100644 --- a/specutils/io/default_loaders/sdss_v.py +++ b/specutils/io/default_loaders/sdss_v.py @@ -450,6 +450,33 @@ def load_sdss_spec_1D(file_obj, *args, hdu: Optional[int] = None, **kwargs): return _load_BOSS_HDU(hdulist, hdu, **kwargs) +def _lazy_sdss_spec_loader(fileobj, **kwargs): + """Lazy loader example for SDSS-V spec files""" + # Build list of HDU indices + labels once + with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: + hdu_indices = [] + labels = [] + for idx in range(1, len(hdulist)): + name = hdulist[idx].name + if name in ["SPALL", "ZALL", "ZLINE"]: + continue + hdu_indices.append(idx) + labels.append(name) + + def _loader(i: int) -> Spectrum: + hdu_idx = hdu_indices[i] + with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: + return _load_BOSS_HDU(hdulist, hdu_idx, **kwargs) + + sl = SpectrumList.from_lazy( + length=len(hdu_indices), + loader=_loader, + labels=labels, # optional + ) + sl.set_id_map(dict(zip(labels, range(len(labels))))) # optional + return sl + + @data_loader( "SDSS-V spec", identifier=spec_sdss5_identify, @@ -457,6 +484,7 @@ def load_sdss_spec_1D(file_obj, *args, hdu: Optional[int] = None, **kwargs): force=True, priority=5, extensions=["fits"], + lazy_loader=_lazy_sdss_spec_loader ) def load_sdss_spec_list(file_obj, **kwargs): """ diff --git a/specutils/io/default_loaders/tests/test_roman.py b/specutils/io/default_loaders/tests/test_roman.py new file mode 100644 index 000000000..eacc19eff --- /dev/null +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -0,0 +1,253 @@ +import astropy.units as u +import numpy as np +import pytest +from astropy.utils.exceptions import AstropyUserWarning + +from specutils import Spectrum, SpectrumList +from specutils.io.default_loaders import roman as roman_loaders # noqa: F401 + +asdf = pytest.importorskip("asdf") + + +def _roman_tree(mode='single', optel="GRISM", unit="W m**(-2) nm**(-1)", nsrc=1): + """Function to create a fake asdf tree for roman spectral data""" + wl = np.linspace(1.0, 2.0, 5, dtype=float) + flux = np.arange(5, dtype=float) + 1.0 + ferr = np.full_like(flux, 0.1) + + meta = { + "optical_element": optel, + "title": "my_1d_spectra_product", + "unit_wl": "nm", + "unit_flux": unit, + } + + if mode == 'single': + data = { + "wl": wl, + "flux": flux, + "flux_error": ferr, + } + elif mode == 'multi': + data = {} + for i in range(nsrc): + data[f"40{i}"] = { + "wl": wl, + "flux": (np.arange(5, dtype=float) + 1.0) * (i + 1), + "flux_error": np.full_like(flux, 0.1 * (i + 1)), + } + + return { + "roman": { + "meta": meta, + "data": data + } + } + +@pytest.fixture() +def roman_single(tmp_path): + """Fixture to create a single source roman spectral asdf file""" + path = tmp_path / "roman_test_single.asdf" + af = asdf.AsdfFile(_roman_tree(mode='single')) + af.write_to(path) + yield str(path) + af.close() + + +@pytest.fixture() +def roman_multi(tmp_path): + """Fixture to create a roman ssc 1d_combined spectral file""" + path = tmp_path / "roman_test_multi.asdf" + af = asdf.AsdfFile(_roman_tree(mode='multi', nsrc=3)) + af.write_to(path) + yield str(path) + af.close() + + +@pytest.fixture() +def roman_file(tmp_path): + """Fixture factory to create variations on roman spectral products""" + def _roman_file(mode='single', optel='GRISM', unit="W m**(-2) nm**(-1)", nsrc=1): + path = tmp_path / "roman_test.asdf" + af = asdf.AsdfFile(_roman_tree(mode=mode, optel=optel, unit=unit, nsrc=nsrc)) + af.write_to(path) + af.close() + return str(path) + yield _roman_file + + +def test_identify_roman_single(roman_single): + """test we can identify a roman single source spectrum""" + assert roman_loaders.identify_1d_single_source(None, roman_single) is True + assert roman_loaders.identify_1d_combined(None, roman_single) is False + assert roman_loaders.identify_1d_individual(None, roman_single) is False + +def test_identify_roman_combined(roman_multi): + """test we can identify a roman 1d combined spectra file""" + assert roman_loaders.identify_1d_single_source(None, roman_multi) is False + assert roman_loaders.identify_1d_combined(None, roman_multi) is True + assert roman_loaders.identify_1d_individual(None, roman_multi) is False + +def test_identify_roman_individual(roman_file): + """test we can identify a roman 1d individual spectra file""" + roman_multi = roman_file(mode='multi', unit="DN/s") + assert roman_loaders.identify_1d_single_source(None, roman_multi) is False + assert roman_loaders.identify_1d_combined(None, roman_multi) is False + assert roman_loaders.identify_1d_individual(None, roman_multi) is True + +def test_identify_not_spectra(roman_file): + """test that another roman file is not mis-identified""" + roman_wfi = roman_file(mode='single', optel='WFI', unit="DN/s") + assert roman_loaders.identify_1d_single_source(None, roman_wfi) is False + assert roman_loaders.identify_1d_combined(None, roman_wfi) is False + assert roman_loaders.identify_1d_individual(None, roman_wfi) is False + + +def test_roman_1d_spectrum(roman_single): + """test loading a roman 1d single source spectrum""" + spec = Spectrum.read(roman_single, format="Roman 1d spectra") + + assert isinstance(spec, Spectrum) + assert spec.spectral_axis.unit == u.nm + assert spec.unit == u.Unit("W m**(-2) nm**(-1)") + assert spec.shape == (5,) + assert spec.uncertainty is not None + assert np.allclose(spec.uncertainty.array, 0.1) + + +def test_roman_1d_combined_load_specific_source(roman_multi): + """test loading a roman 1d combined spectrum for a specific source""" + spec = Spectrum.read(roman_multi, format="Roman 1d combined", source="402") + assert isinstance(spec, Spectrum) + assert spec.meta["source_id"] == "402" + assert spec.spectral_axis.unit == u.nm + assert spec.unit == u.Unit("W m**(-2) nm**(-1)") + assert spec.shape == (5,) + + +def test_roman_1d_combined_load_first_source(roman_multi): + """test loading without a source throws the warning and loads first source""" + with pytest.warns(AstropyUserWarning, match="source not specified"): + spec = Spectrum.read(roman_multi, format="Roman 1d combined") + + assert spec.meta["source_id"] == "400" + + +def test_roman_1d_combined_invalid_source(roman_multi): + """test loading with a bad source id raises error""" + with pytest.raises(ValueError, match="Invalid source"): + Spectrum.read(roman_multi, format="Roman 1d combined", source="nope") + + +def test_roman_1d_combined_list(roman_multi): + """test we can load a 1d combined list""" + speclist = SpectrumList.read(roman_multi, format="Roman 1d combined") + + assert isinstance(speclist, SpectrumList) + assert speclist.is_lazy is False + assert speclist.n_loaded == 3 + assert len(speclist) == 3 + for i, sp in enumerate(speclist): + assert isinstance(sp, Spectrum) + assert sp.meta["source_id"] == f"40{i}" + assert "source_map" in sp.meta + assert sp.meta["source_map"][f"40{i}"] == i + +def test_roman_altid_index(roman_multi): + """test we can select on alternate id""" + speclist = SpectrumList.read(roman_multi, format="Roman 1d combined") + spec1 = speclist[1] + spec2 = speclist["401"] + spec3 = speclist["1"] + assert spec1 == spec2 == spec3 + +def test_roman_alid_fails(roman_multi): + """test we fail on altid""" + speclist = SpectrumList.read(roman_multi, format="Roman 1d combined") + with pytest.raises(IndexError, match="list index out of range"): + speclist[10] + + with pytest.raises(IndexError, match="list index out of range"): + speclist["410"] + + with pytest.raises(KeyError, match="not found in id mapping, and cannot resolve to a list"): + speclist["ab"] + + speclist._id_map = None + with pytest.raises(KeyError, match="No id mapping provided for alternate indexing"): + speclist["ab"] + + +def test_roman_nonunique_altid_fails(): + """test we fail on non-unique altid""" + with pytest.raises(ValueError, match="Labels must be unique! Non-unique labels:"): + a = ['400', '401', '402', '403', '401'] + SpectrumList.from_lazy(len(a), lambda x: x, labels=a) + + +def test_roman_lazy_loaded_spectrum(roman_multi): + """test we can lazy load spectra""" + speclist = SpectrumList.read(roman_multi, format="Roman 1d combined", lazy_load=True, cache_asdf=True) + assert speclist.is_lazy is True + assert speclist.n_loaded == 0 + spec = speclist[0] + assert isinstance(spec, Spectrum) + assert speclist.n_loaded == 1 + + +def test_roman_repr(roman_multi): + """test we get a lazy repr""" + speclist = SpectrumList.read(roman_multi, format="Roman 1d combined", lazy_load=True, cache_asdf=True) + assert speclist.is_lazy is True + assert speclist.n_loaded == 0 + assert "lazy list: 0 items loaded; access an index to load a spectrum:" in repr(speclist) + assert "['400'" in repr(speclist) + + # load 1 + speclist[0] + assert "lazy list: 1 items loaded;" in repr(speclist) + assert "[ bool: + """Whether the SpectrumList is in lazy-loading mode""" + return self._lazy_loader is not None + + @property + def n_loaded(self) -> int: + """Number of spectra in the list currently loaded""" + if not self.is_lazy: + return len(self) + + return sum(1 for x in super().__iter__() if x is not _placeholder) + + def set_id_map(self, id_map: dict[str, int]): + """Set a mapping of alternate string labels to list indices. + + This allows accessing items in the list using string labels, e.g. + source ids, in addition to int indices. + + Parameters + ---------- + id_map : dict[str, int] + Mapping of string keys to list indices + """ + if not isinstance(id_map, dict): + raise TypeError("input id map must be a dictionary") + + # sanity check the keys are unique + self.check_unique_labels(list(id_map.keys())) + + self._id_map = dict(id_map) + + def _resolve_key(self, key: str) -> int: + """Resolve a string key to a list index""" + + # return normal list index + if key.isdigit() and (self._id_map is None or key not in self._id_map): + return int(key) + + # otherwise it must be provided by the mapping + if self._id_map is None: + raise KeyError("No id mapping provided for alternate indexing") + + if key not in self._id_map: + raise KeyError(f"Key '{key}' not found in id mapping, and cannot resolve to a list index.") + + return self._id_map[key] + + def __getitem__(self, value: str | int): + """Retrieve items from the list normally or lazily""" + + # resolve any string key + if isinstance(value, str): + value = self._resolve_key(value) + + # preserve original slice behaviour + if isinstance(value, slice): + return SpectrumList([self[i] for i in range(*value.indices(len(self)))]) + + # use lazy item getter + if self.is_lazy: + return self._lazy_get(int(value)) + + # use normal item getter + return super().__getitem__(value) + + def _lazy_get(self, idx: int): + """Lazily retrieve an item from the list""" + + # preserve original reverse slicing + if idx < 0: + idx = len(self) + idx + if idx < 0 or idx >= len(self): + raise IndexError("list index out of range") + + # get the current item and check if it's a placeholder object + # if not, then return it + current = super().__getitem__(idx) + if current is not _placeholder: + return current + + # use the lazy loader to get the spectrum object + # replace the placeholder item with it + val = self._lazy_loader(idx) + self[idx] = val + + return val + + def __repr__(self) -> str: + """Build string repr the list + + Non-lazy lists have normal reprs. Lazy lists use placeholder values, + or optional labels if provided. Once the item is loaded, the normal + item repr is used. + """ + # use normal repr + if not self.is_lazy or self.n_loaded == len(self): + return super().__repr__() + + # build the lazy repr + prefix = f"lazy list: {self.n_loaded} items loaded; access an index to load a spectrum:\n" + return prefix + f"[{', '.join(self._lazy_repr(i) for i in range(len(self)))}]" + + def _lazy_repr(self, ii: int): + """Return item repr without triggering a lazy load""" + + # get item + item = super().__getitem__(ii) + + # use normal item repr + if item is not _placeholder: + return repr(item) + + # use placeholder or optional label repr + labels = self._lazy_labels + if labels is not None and ii < len(labels): + return repr(labels[ii]) + + return repr(item) + + @classmethod + def check_unique_labels(cls, labels: list[str]): + """Check that labels are unique""" + nonuniq = {k for k,v in Counter(labels).items() if v > 1} + if nonuniq: + raise ValueError(f"Labels must be unique! Non-unique labels: {nonuniq}") + + @classmethod + def from_lazy( + cls, length: int, loader: Callable, cache_size: Optional[int] = None, labels: list = None + ) -> "SpectrumList": + """Construct a lazy-loading SpectrumList. + + Constructs a Spectrumlist using placeholder objects and sets a + loader callable used to instantiate Spectrum objects lazily on item + get. Once a Spectrum is loaded, it replaces the placeholder item, and + repeated access does not re-run the loader. + + If cache_size is specified, then the loader is also cached with + an lru_cache. + + Parameters + ---------- + length : int + Total number of spectra in the list. + loader : Callable + Callable taking an int index and returns a Spectrum. + cache_size : int or None, optional + If provided, wraps the loader in an lru_cache of this + size. + labels : list, optional + Optional list of placeholder display labels shown by the repr + before spectra are materialized. + + Returns + ------- + SpectrumList + A lazy-loadable SpectrumList + + """ + # create placeholder list + speclist = cls([_placeholder] * length) + + # optionally cache the loader + if cache_size: + cache_size = max(int(cache_size), 0) + loader = lru_cache(maxsize=cache_size)(loader) + + # check labels + if labels: + cls.check_unique_labels(labels) + + # set lazy parameters + speclist._lazy_loader = loader + speclist._lazy_labels = labels + return speclist