From 5aff91f59ec185b6e0be8782f319cf2ee0f17010 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Mon, 12 Jan 2026 15:31:24 -0500 Subject: [PATCH 01/14] adding init roman ssc loaders --- specutils/io/default_loaders/roman.py | 282 ++++++++++++++++++++++++++ specutils/io/parsing_utils.py | 67 +++++- 2 files changed, 344 insertions(+), 5 deletions(-) create mode 100644 specutils/io/default_loaders/roman.py diff --git a/specutils/io/default_loaders/roman.py b/specutils/io/default_loaders/roman.py new file mode 100644 index 000000000..cbb7abc7c --- /dev/null +++ b/specutils/io/default_loaders/roman.py @@ -0,0 +1,282 @@ +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 + + +# class SpectrumList(SpectrumList): +# def __getitem__(self, value): +# n_items = len(self) +# meta = getattr(super().__getitem__(0), 'meta', None) +# smap = meta.get('source_map', {}) + +# if isinstance(value, int): +# if value < n_items: +# return super().__getitem__(value) +# elif str(value) in smap: +# return super().__getitem__(smap[str(value)]) +# else: +# raise IndexError("list index out of range") + +# if isinstance(value, str): +# if smap and value in smap: +# return super().__getitem__(smap[value]) + + +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 _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()) + source_idx_map = dict(zip(sources, range(len(sources)))) + + meta['source_map'] = source_idx_map + + 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, +) +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 +) +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 +) +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 +) +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/parsing_utils.py b/specutils/io/parsing_utils.py index 6d35028a7..dfded55d3 100644 --- a/specutils/io/parsing_utils.py +++ b/specutils/io/parsing_utils.py @@ -1,18 +1,24 @@ -import numpy as np +import contextlib +import io import os import re import urllib -import io -import contextlib +import warnings +import astropy.units as u +import numpy as np from astropy.io import fits from astropy.nddata import StdDevUncertainty from astropy.utils.exceptions import AstropyUserWarning -import astropy.units as u -import warnings from specutils.spectra import Spectrum +# Optional ASDF support +try: + import asdf +except ImportError: + asdf = None + @contextlib.contextmanager def read_fileobj_or_hdulist(*args, **kwargs): @@ -52,6 +58,57 @@ def read_fileobj_or_hdulist(*args, **kwargs): hdulist.close() +@contextlib.contextmanager +def read_fileobj_or_asdftree(*args, **kwargs): + """Context manager for reading a filename or file object + + Returns + ------- + af : :class:`~asdf._asdf.AsdfFile` + Provides a generator-iterator representing the open file object handle. + """ + if not asdf: + raise ImportError("The 'asdf' package is required to read ASDF files.") + + # Access the fileobj or filename arg + # Do this so identify functions are useable outside of Spectrum.read context + try: + fileobj = args[2] + except IndexError: + fileobj = args[0] + + # astropy's identify registry reuses same open fileobj handle, so try to rewind it + try: + fileobj.seek(0) + except (AttributeError, io.UnsupportedOperation, OSError): + pass + + # read the asdf file + if isinstance(fileobj, asdf.AsdfFile): + if fileobj._fd is None: + # file is closed if _fd is None, otherwise it's a subclass of asdf.generic_io.GenericFile + af = asdf.open(fileobj.uri, **kwargs) + else: + af = fileobj + elif isinstance(fileobj, io.BufferedReader): + # fileobj.name avoids issues caused by the unified reader when no format is specified where + # the binary blob is read twice but not rewound in between, resulting in an + # "invalid ASDF file" error message + af = asdf.open(fileobj.name, **kwargs) + else: + af = asdf.open(fileobj, **kwargs) + + try: + yield af + # Cleanup even after identifier function has thrown an exception: rewind generic file handles. + finally: + if not isinstance(fileobj, asdf.AsdfFile): + try: + fileobj.seek(0) + except (AttributeError, io.UnsupportedOperation): + af.close() + + def spectrum_from_column_mapping(table, column_mapping, wcs=None, verbose=False): """ Given a table and a mapping of the table column names to attributes From 644ef1c8c29ed4c63d6fd7356e037047521b5a62 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Wed, 14 Jan 2026 13:56:49 -0500 Subject: [PATCH 02/14] adding tests for roman --- .../io/default_loaders/tests/test_roman.py | 192 ++++++++++++++++++ specutils/io/parsing_utils.py | 4 + 2 files changed, 196 insertions(+) create mode 100644 specutils/io/default_loaders/tests/test_roman.py 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..6d441b2b8 --- /dev/null +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -0,0 +1,192 @@ +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[str(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="2") + assert isinstance(spec, Spectrum) + assert spec.meta["source_id"] == "2" + 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"] == "0" + + +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 len(speclist) == 3 + for i, sp in enumerate(speclist): + assert isinstance(sp, Spectrum) + assert sp.meta["source_id"] == str(i) + assert "source_map" in sp.meta + assert sp.meta["source_map"][str(i)] == i + + +def test_roman_1d_individual_list(roman_file): + """test we can load a 1d individual list""" + roman_indiv = roman_file(mode='multi', unit="DN/s", nsrc=3) + speclist = SpectrumList.read(roman_indiv, format="Roman 1d individual") + + assert isinstance(speclist, SpectrumList) + assert len(speclist) == 3 + assert speclist[0].unit == u.Unit("DN/s") + assert speclist[0].spectral_axis.unit == u.nm + assert speclist[0].meta["source_id"] == "0" + assert speclist[1].meta["source_id"] == "1" + assert speclist[2].meta["source_id"] == "2" + + +data_input = ["file", "asdf", "blob"] + +@pytest.fixture(params=data_input) +def data(request, roman_single): + if request.param == "file": + tmp = roman_single + elif request.param == "asdf": + tmp = asdf.open(roman_single) + elif request.param == "blob": + tmp = open(roman_single, 'rb') + else: + tmp = roman_single + + yield tmp + + try: + tmp.close() + except AttributeError: + pass + +def test_read_fileobj_or_asdftree(data): + """Test Spectrum read with different data inputs""" + s = Spectrum.read(data) + assert isinstance(s, Spectrum) diff --git a/specutils/io/parsing_utils.py b/specutils/io/parsing_utils.py index dfded55d3..1aba06461 100644 --- a/specutils/io/parsing_utils.py +++ b/specutils/io/parsing_utils.py @@ -107,6 +107,10 @@ def read_fileobj_or_asdftree(*args, **kwargs): fileobj.seek(0) except (AttributeError, io.UnsupportedOperation): af.close() + finally: + af.close() + else: + af.close() def spectrum_from_column_mapping(table, column_mapping, wcs=None, verbose=False): From 20f5ad0a659898692679d7ac5300fa0a772f2a17 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Wed, 21 Jan 2026 14:45:56 -0500 Subject: [PATCH 03/14] initial lazy loader for spectrum list for roman --- specutils/io/default_loaders/roman.py | 64 ++++++++------ specutils/io/parsing_utils.py | 100 +++++++++++++++++----- specutils/io/registers.py | 59 ++++++++++--- specutils/spectra/spectrum_list.py | 118 +++++++++++++++++++++++++- 4 files changed, 283 insertions(+), 58 deletions(-) diff --git a/specutils/io/default_loaders/roman.py b/specutils/io/default_loaders/roman.py index cbb7abc7c..183be1759 100644 --- a/specutils/io/default_loaders/roman.py +++ b/specutils/io/default_loaders/roman.py @@ -17,25 +17,6 @@ roman_gdps = None -# class SpectrumList(SpectrumList): -# def __getitem__(self, value): -# n_items = len(self) -# meta = getattr(super().__getitem__(0), 'meta', None) -# smap = meta.get('source_map', {}) - -# if isinstance(value, int): -# if value < n_items: -# return super().__getitem__(value) -# elif str(value) in smap: -# return super().__getitem__(smap[str(value)]) -# else: -# raise IndexError("list index out of range") - -# if isinstance(value, str): -# if smap and value in smap: -# return super().__getitem__(smap[value]) - - def _identify_roman_1d(*args, mode: str, units: str=None, **kwargs) -> bool: """Identify a Roman SSC spectral file @@ -122,7 +103,7 @@ def _load_roman_spectrum(roman: dict, source: str) -> Spectrum: 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: + with read_fileobj_or_asdftree(file_obj, **kwargs) as af: roman = af["roman"] if source is None: source = next(iter(roman['data'])) @@ -133,6 +114,28 @@ def _load_roman_first(file_obj, source=None, **kwargs): return _load_roman_spectrum(roman, source) + +def _lazy_load_roman(file_obj, *, cache_size=0, **kwargs): + """Lazy loader for SpectrumList""" + 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)))) + + 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) + + sl = SpectrumList.from_lazy(length=len(sources), loader=_loader, cache_size=cache_size, labels=sources) + + # Optional / independent: enable string indexing for Roman IDs + 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() @@ -143,7 +146,7 @@ def _load_roman_multisource(file_obj, **kwargs) -> SpectrumList: source_idx_map = dict(zip(sources, range(len(sources)))) meta['source_map'] = source_idx_map - + spectra.set_id_map(source_idx_map) for source in roman["data"]: spectrum = _load_roman_spectrum(roman, source) spectra.append(spectrum) @@ -183,10 +186,14 @@ def roman_1d_spectrum_loader(file_obj, **kwargs): @data_loader( - "Roman 1d combined", identifier=identify_1d_combined, dtype=Spectrum, - extensions=['asdf'], priority=10, + "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): +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. @@ -208,8 +215,13 @@ def roman_1d_combined(file_obj, source:str=None, **kwargs): @data_loader( - "Roman 1d combined", identifier=identify_1d_combined, dtype=SpectrumList, - extensions=['asdf'], priority=10, force=True + "Roman 1d combined", + identifier=identify_1d_combined, + dtype=SpectrumList, + extensions=["asdf"], + priority=10, + force=True, + lazy_factory=_lazy_load_roman, ) def roman_1d_combined_list(file_obj, **kwargs): """Load all Roman 1d combined extracted spectra diff --git a/specutils/io/parsing_utils.py b/specutils/io/parsing_utils.py index 1aba06461..df531d65e 100644 --- a/specutils/io/parsing_utils.py +++ b/specutils/io/parsing_utils.py @@ -4,6 +4,8 @@ import re import urllib import warnings +import weakref +from functools import lru_cache import astropy.units as u import numpy as np @@ -13,12 +15,43 @@ from specutils.spectra import Spectrum +obj_cache = weakref.WeakKeyDictionary() + + # Optional ASDF support try: import asdf except ImportError: asdf = None +@lru_cache(maxsize=8) +def open_cache(key: str): + return asdf.open(key) + + +def clear_cache(): + # clear object cache + for af in list(obj_cache.values()): + try: + af.close() + except Exception: + pass + obj_cache.clear() + + # clear lru_cache + open_cache.cache_clear() + + +def get_cached_object(fileobj, **kwargs): + if isinstance(fileobj, (str, os.PathLike)): + return open_cache(fileobj) + else: + af = obj_cache.get(fileobj) + if af is None: + af = asdf.open(fileobj, **kwargs) + obj_cache[fileobj] = af + return af + @contextlib.contextmanager def read_fileobj_or_hdulist(*args, **kwargs): @@ -58,6 +91,37 @@ def read_fileobj_or_hdulist(*args, **kwargs): hdulist.close() +def open_input(fileobj, cache_asdf: bool = None, **kwargs): + # Caller passed an AsdfFile: reuse it if open; reopen if closed and uri available. + if isinstance(fileobj, asdf.AsdfFile): + if fileobj._fd is not None: + return fileobj + if getattr(fileobj, "uri", None): + if cache_asdf: + return open_cache(fileobj.uri) + return asdf.open(fileobj.uri, **kwargs) + + # Cache-enabled path: cache by stable key when possible, otherwise by object identity. + if cache_asdf: + if isinstance(fileobj, (str, os.PathLike)): + return open_cache(os.fspath(fileobj)) + + # BufferedReader: prefer reopening via name when available (avoids “blob read twice” issues) + if isinstance(fileobj, io.BufferedReader) and getattr(fileobj, "name", None): + return open_cache(fileobj.name) + + af = obj_cache.get(fileobj) + if af is None: + af = asdf.open(fileobj, **kwargs) + obj_cache[fileobj] = af + return af + + # Non-cached path: always create a transient handle we close on exit. + if isinstance(fileobj, io.BufferedReader) and getattr(fileobj, "name", None): + return asdf.open(fileobj.name, **kwargs) + return asdf.open(fileobj, **kwargs) + + @contextlib.contextmanager def read_fileobj_or_asdftree(*args, **kwargs): """Context manager for reading a filename or file object @@ -83,35 +147,29 @@ def read_fileobj_or_asdftree(*args, **kwargs): except (AttributeError, io.UnsupportedOperation, OSError): pass - # read the asdf file - if isinstance(fileobj, asdf.AsdfFile): - if fileobj._fd is None: - # file is closed if _fd is None, otherwise it's a subclass of asdf.generic_io.GenericFile - af = asdf.open(fileobj.uri, **kwargs) - else: - af = fileobj - elif isinstance(fileobj, io.BufferedReader): - # fileobj.name avoids issues caused by the unified reader when no format is specified where - # the binary blob is read twice but not rewound in between, resulting in an - # "invalid ASDF file" error message - af = asdf.open(fileobj.name, **kwargs) - else: - af = asdf.open(fileobj, **kwargs) + # caching options + # Close only when we are not caching AND the caller did not provide an already-open AsdfFile. + cache_asdf = kwargs.pop("cache_asdf", None) + should_close = (not cache_asdf) and not (isinstance(fileobj, asdf.AsdfFile) and fileobj._fd is not None) + + af = open_input(fileobj, cache_asdf=cache_asdf, **kwargs) try: yield af # Cleanup even after identifier function has thrown an exception: rewind generic file handles. finally: - if not isinstance(fileobj, asdf.AsdfFile): + if should_close: try: - fileobj.seek(0) - except (AttributeError, io.UnsupportedOperation): af.close() - finally: - af.close() - else: - af.close() + except Exception: + pass + # Always best-effort rewind of non-AsdfFile inputs (keep your current logic) + if not isinstance(fileobj, asdf.AsdfFile): + try: + fileobj.seek(0) + except (AttributeError, io.UnsupportedOperation, OSError): + pass def spectrum_from_column_mapping(table, column_mapping, wcs=None, verbose=False): """ diff --git a/specutils/io/registers.py b/specutils/io/registers.py index cdee16f99..0e260cb05 100644 --- a/specutils/io/registers.py +++ b/specutils/io/registers.py @@ -10,7 +10,7 @@ from astropy.io import registry as io_registry -from ..spectra import Spectrum, SpectrumList, SpectrumCollection +from ..spectra import Spectrum, SpectrumCollection, SpectrumList __all__ = ['data_loader', 'custom_writer', 'get_loaders_by_extension', 'identify_spectrum_format'] @@ -25,8 +25,17 @@ def _astropy_has_priorities(): return False -def data_loader(label, identifier=None, dtype=Spectrum, extensions=None, - priority=0, force=False, autogenerate_spectrumlist=True, verbose=False): +def data_loader( + label, + identifier=None, + dtype=Spectrum, + extensions=None, + priority=0, + force=False, + autogenerate_spectrumlist=True, + verbose=False, + lazy_factory=None, +): """ Wraps a function that can be added to an `~astropy.io.registry` for custom file reading. @@ -56,7 +65,8 @@ def data_loader(label, identifier=None, dtype=Spectrum, extensions=None, data_loader that reads Spectrum objects. Default is ``True``. verbose : bool Print extra info. - + lazy_factory : Callable, optional + A factory function to create a lazy-loading SpectrumList. """ def identifier_wrapper(ident): def wrapper(*args, **kwargs): @@ -70,13 +80,44 @@ def wrapper(*args, **kwargs): return wrapper def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + if dtype is SpectrumList: + lazy_load = bool(kwargs.pop("lazy_load", False)) + cache_size = kwargs.pop("cache_size", None) + + if lazy_load and lazy_factory is not None: + # file_obj = args[0] if args else None + # if lazy_requires_path and not isinstance(file_obj, (str, os.PathLike)): + # return func(*args, **kwargs) + + # Let the factory build the list; registry only passes cache_size optionally + if cache_size is not None: + return lazy_factory(*args, cache_size=int(cache_size), **kwargs) + return lazy_factory(*args, **kwargs) + + # If not lazy, don’t leak these into eager loaders + return func(*args, **kwargs) + + # non-SpectrumList readers + kwargs.pop("lazy_load", None) + kwargs.pop("cache_size", None) + return func(*args, **kwargs) + if _astropy_has_priorities(): io_registry.register_reader( - label, dtype, func, priority=priority, force=force, + label, + dtype, + wrapper, + priority=priority, + force=force, ) else: io_registry.register_reader( - label, dtype, func, force=force, + label, + dtype, + wrapper, + force=force, ) if identifier is None: @@ -103,7 +144,7 @@ def decorator(func): ) # Include the file extensions as attributes on the function object - func.extensions = extensions + wrapper.extensions = extensions if verbose: print(f"Successfully loaded reader \"{label}\".") @@ -133,9 +174,7 @@ def load_spectrum_list(*args, **kwargs): if verbose: print(f"Created SpectrumList reader for \"{label}\".") - @wraps(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) + return wrapper return decorator diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index 486a83658..c0887beab 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -1,9 +1,14 @@ -from astropy.nddata import NDIOMixin +from functools import lru_cache +from typing import Callable, Optional +from astropy.nddata import NDIOMixin __all__ = ['SpectrumList'] +_placeholder = object() + + class SpectrumList(list, NDIOMixin): """ A list that is used to hold a list of `~specutils.Spectrum` objects @@ -15,3 +20,114 @@ class SpectrumList(list, NDIOMixin): `~specutils.Spectrum`. For more on this topic, see :ref:`specutils-representation-overview`. """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._id_map: Optional[dict[str, int]] = None + + self._lazy_loader: Optional[Callable] = None + self._lazy_cache_size: int = 0 + self._lazy_labels: Optional[list[str]] = None + + @property + def is_lazy(self) -> bool: + return self._lazy_loader is not None + + @property + def n_loaded(self) -> int: + 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]): + self._id_map = dict(id_map) + + def _resolve_key(self, key: str) -> int: + 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): + if isinstance(value, str): + value = self._resolve_key(value) + + # Preserve normal list slice behavior (and avoid int(slice) errors) + if isinstance(value, slice): + return SpectrumList([self[i] for i in range(*value.indices(len(self)))]) + + if self.is_lazy: + return self._lazy_get(int(value)) + + return super().__getitem__(value) + + def _lazy_get(self, idx: int): + if idx < 0: + idx = len(self) + idx + if idx < 0 or idx >= len(self): + raise IndexError("list index out of range") + + current = super().__getitem__(idx) + if current is not _placeholder: + return current + + val = self._lazy_loader(idx) + self[idx] = val + + return val + + def _lazy_repr(self, ii: int): + """Return repr for item i without triggering lazy loading.""" + item = super().__getitem__(ii) + + if item is not _placeholder: + return repr(item) + + labels = self._lazy_labels + if labels is not None and ii < len(labels): + return repr(labels[ii]) + + return repr(item) + + def __repr__(self) -> str: + if not self.is_lazy: + return super().__repr__() + + return f"[{', '.join(self._lazy_repr(i) for i in range(len(self)))}]" + + @classmethod + def from_lazy(cls, *, length: int, loader: Callable, cache_size: int = 0, labels: list = None) -> "SpectrumList": + """ + Create a lazy-loading `SpectrumList`. + + Parameters + ---------- + length : int + The number of spectra in the list. + loader : Callable + A function that takes a single integer index and returns the + corresponding `~specutils.Spectrum` object. + cache_size : int, optional + The number of spectra to cache in memory, by default 10. + + Returns + ------- + SpectrumList + A lazy-loading `SpectrumList`. + """ + speclist = cls([_placeholder] * length) + cache_size = max(int(cache_size), 0) + if cache_size > 0: + loader = lru_cache(maxsize=cache_size)(loader) + speclist._lazy_loader = loader + + speclist._lazy_labels = labels + return speclist From 353e7030ac215b7fc684afe05a5750a4f4212697 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Thu, 22 Jan 2026 14:24:38 -0500 Subject: [PATCH 04/14] some cleanup and init object caching; tests --- specutils/io/default_loaders/roman.py | 6 +- .../io/default_loaders/tests/test_roman.py | 36 ++++-- specutils/io/parsing_utils.py | 58 ++++++---- specutils/io/registers.py | 22 ++-- specutils/spectra/spectrum_list.py | 103 ++++++++++++++---- 5 files changed, 155 insertions(+), 70 deletions(-) diff --git a/specutils/io/default_loaders/roman.py b/specutils/io/default_loaders/roman.py index 183be1759..24b01d519 100644 --- a/specutils/io/default_loaders/roman.py +++ b/specutils/io/default_loaders/roman.py @@ -115,7 +115,7 @@ def _load_roman_first(file_obj, source=None, **kwargs): return _load_roman_spectrum(roman, source) -def _lazy_load_roman(file_obj, *, cache_size=0, **kwargs): +def _lazy_load_roman(file_obj, **kwargs): """Lazy loader for SpectrumList""" with read_fileobj_or_asdftree(file_obj, **kwargs) as af: roman = af["roman"] @@ -128,7 +128,9 @@ def _loader(i: int) -> Spectrum: roman2 = af2["roman"] return _load_roman_spectrum(roman2, source) - sl = SpectrumList.from_lazy(length=len(sources), loader=_loader, cache_size=cache_size, labels=sources) + sl = SpectrumList.from_lazy( + length=len(sources), loader=_loader, cache_size=kwargs.get("cache_size", None), labels=sources + ) # Optional / independent: enable string indexing for Roman IDs sl.set_id_map(source_idx_map) diff --git a/specutils/io/default_loaders/tests/test_roman.py b/specutils/io/default_loaders/tests/test_roman.py index 6d441b2b8..9a9a3560f 100644 --- a/specutils/io/default_loaders/tests/test_roman.py +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -31,7 +31,7 @@ def _roman_tree(mode='single', optel="GRISM", unit="W m**(-2) nm**(-1)", nsrc=1) elif mode == 'multi': data = {} for i in range(nsrc): - data[str(i)] = { + 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)), @@ -117,9 +117,9 @@ def test_roman_1d_spectrum(roman_single): 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="2") + spec = Spectrum.read(roman_multi, format="Roman 1d combined", source="402") assert isinstance(spec, Spectrum) - assert spec.meta["source_id"] == "2" + 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,) @@ -130,7 +130,7 @@ def test_roman_1d_combined_load_first_source(roman_multi): with pytest.warns(AstropyUserWarning, match="source not specified"): spec = Spectrum.read(roman_multi, format="Roman 1d combined") - assert spec.meta["source_id"] == "0" + assert spec.meta["source_id"] == "400" def test_roman_1d_combined_invalid_source(roman_multi): @@ -144,14 +144,32 @@ def test_roman_1d_combined_list(roman_multi): 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"] == str(i) + assert sp.meta["source_id"] == f"40{i}" assert "source_map" in sp.meta - assert sp.meta["source_map"][str(i)] == i + assert sp.meta["source_map"][f"40{i}"] == i +def test_roman_altid_index(roman_multi): + 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_lazy_loaded_spectrum(roman_multi): + 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_1d_individual_list(roman_file): """test we can load a 1d individual list""" roman_indiv = roman_file(mode='multi', unit="DN/s", nsrc=3) @@ -161,9 +179,9 @@ def test_roman_1d_individual_list(roman_file): assert len(speclist) == 3 assert speclist[0].unit == u.Unit("DN/s") assert speclist[0].spectral_axis.unit == u.nm - assert speclist[0].meta["source_id"] == "0" - assert speclist[1].meta["source_id"] == "1" - assert speclist[2].meta["source_id"] == "2" + assert speclist[0].meta["source_id"] == "400" + assert speclist[1].meta["source_id"] == "401" + assert speclist[2].meta["source_id"] == "402" data_input = ["file", "asdf", "blob"] diff --git a/specutils/io/parsing_utils.py b/specutils/io/parsing_utils.py index df531d65e..5fbb0c737 100644 --- a/specutils/io/parsing_utils.py +++ b/specutils/io/parsing_utils.py @@ -15,21 +15,24 @@ from specutils.spectra import Spectrum +# for caching objects obj_cache = weakref.WeakKeyDictionary() - # Optional ASDF support try: import asdf except ImportError: asdf = None + @lru_cache(maxsize=8) def open_cache(key: str): + """Open asdf file and cache it""" return asdf.open(key) def clear_cache(): + """Clear object and lru caches""" # clear object cache for af in list(obj_cache.values()): try: @@ -43,14 +46,26 @@ def clear_cache(): def get_cached_object(fileobj, **kwargs): + """Get the cached object""" + # cache by path-like input if isinstance(fileobj, (str, os.PathLike)): - return open_cache(fileobj) - else: - af = obj_cache.get(fileobj) - if af is None: - af = asdf.open(fileobj, **kwargs) - obj_cache[fileobj] = af - return af + return open_cache(os.fspath(fileobj)) + + # cache by uri for asdf input + if isinstance(fileobj, asdf.AsdfFile): + return open_cache(fileobj.uri) + + # prefer caching by filename for file-like objects + name = getattr(fileobj, "name", None) + if isinstance(name, str) and name: + return open_cache(name) + + # fallback to cache by object + af = obj_cache.get(fileobj) + if af is None: + af = asdf.open(fileobj, **kwargs) + obj_cache[fileobj] = af + return af @contextlib.contextmanager @@ -92,33 +107,23 @@ def read_fileobj_or_hdulist(*args, **kwargs): def open_input(fileobj, cache_asdf: bool = None, **kwargs): + """Open the asdf info with or without cache""" # Caller passed an AsdfFile: reuse it if open; reopen if closed and uri available. if isinstance(fileobj, asdf.AsdfFile): if fileobj._fd is not None: return fileobj if getattr(fileobj, "uri", None): if cache_asdf: - return open_cache(fileobj.uri) + return get_cached_object(fileobj, **kwargs) return asdf.open(fileobj.uri, **kwargs) - # Cache-enabled path: cache by stable key when possible, otherwise by object identity. if cache_asdf: - if isinstance(fileobj, (str, os.PathLike)): - return open_cache(os.fspath(fileobj)) - - # BufferedReader: prefer reopening via name when available (avoids “blob read twice” issues) - if isinstance(fileobj, io.BufferedReader) and getattr(fileobj, "name", None): - return open_cache(fileobj.name) - - af = obj_cache.get(fileobj) - if af is None: - af = asdf.open(fileobj, **kwargs) - obj_cache[fileobj] = af - return af + return get_cached_object(fileobj, **kwargs) # Non-cached path: always create a transient handle we close on exit. - if isinstance(fileobj, io.BufferedReader) and getattr(fileobj, "name", None): - return asdf.open(fileobj.name, **kwargs) + name = getattr(fileobj, "name", None) + if isinstance(name, str) and name: + return asdf.open(name, **kwargs) return asdf.open(fileobj, **kwargs) @@ -141,6 +146,11 @@ def read_fileobj_or_asdftree(*args, **kwargs): except IndexError: fileobj = args[0] + # If Astropy opened a transient file object, it may be closed by the time a + # lazy loader re-enters here. If we can recover a filename, do so. + if getattr(fileobj, "closed", False) and getattr(fileobj, "name", None): + fileobj = fileobj.name + # astropy's identify registry reuses same open fileobj handle, so try to rewind it try: fileobj.seek(0) diff --git a/specutils/io/registers.py b/specutils/io/registers.py index 0e260cb05..92308d0b3 100644 --- a/specutils/io/registers.py +++ b/specutils/io/registers.py @@ -82,28 +82,24 @@ def wrapper(*args, **kwargs): def decorator(func): @wraps(func) def wrapper(*args, **kwargs): - if dtype is SpectrumList: - lazy_load = bool(kwargs.pop("lazy_load", False)) - cache_size = kwargs.pop("cache_size", None) + lazy_load = bool(kwargs.pop("lazy_load", False)) + if lazy_load and dtype is Spectrum: + raise ValueError("Lazy loading is not supported for Spectrum objects.") + # check SpectrumList loaders for lazy option + if dtype is SpectrumList: if lazy_load and lazy_factory is not None: - # file_obj = args[0] if args else None - # if lazy_requires_path and not isinstance(file_obj, (str, os.PathLike)): - # return func(*args, **kwargs) - - # Let the factory build the list; registry only passes cache_size optionally - if cache_size is not None: - return lazy_factory(*args, cache_size=int(cache_size), **kwargs) return lazy_factory(*args, **kwargs) - # If not lazy, don’t leak these into eager loaders + # If not lazy, use eager loader + kwargs.pop("cache_size", None) return func(*args, **kwargs) - # non-SpectrumList readers - kwargs.pop("lazy_load", None) + # Spectrum loaders kwargs.pop("cache_size", None) return func(*args, **kwargs) + if _astropy_has_priorities(): io_registry.register_reader( label, diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index c0887beab..798ff1ca7 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -5,7 +5,7 @@ __all__ = ['SpectrumList'] - +# a temporary placeholder object for lists _placeholder = object() @@ -23,31 +23,49 @@ class SpectrumList(list, NDIOMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + + # Mapping of alternate string ids to list index self._id_map: Optional[dict[str, int]] = None + # Parameters for lazy loading self._lazy_loader: Optional[Callable] = None self._lazy_cache_size: int = 0 self._lazy_labels: Optional[list[str]] = None @property def is_lazy(self) -> 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 + """ 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. + # otherwise it must be provided by the mapping if self._id_map is None: raise KeyError("No id mapping provided for alternate indexing") @@ -56,78 +74,119 @@ def _resolve_key(self, key: str) -> int: return self._id_map[key] - def __getitem__(self, value): + 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 normal list slice behavior (and avoid int(slice) errors) + # 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: + return super().__repr__() + + # build the lazy repr + return f"[{', '.join(self._lazy_repr(i) for i in range(len(self)))}]" + def _lazy_repr(self, ii: int): - """Return repr for item i without triggering lazy loading.""" + """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) - def __repr__(self) -> str: - if not self.is_lazy: - return super().__repr__() + @classmethod + def from_lazy( + cls, length: int, loader: Callable, cache_size: Optional[int] = None, labels: list = None + ) -> "SpectrumList": + """Construct a lazy-loading SpectrumList. - return f"[{', '.join(self._lazy_repr(i) for i in range(len(self)))}]" + 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. - @classmethod - def from_lazy(cls, *, length: int, loader: Callable, cache_size: int = 0, labels: list = None) -> "SpectrumList": - """ - Create a lazy-loading `SpectrumList`. + If cache_size is specified, then the loader is also cached with + an lru_cache. Parameters ---------- length : int - The number of spectra in the list. + Total number of spectra in the list. loader : Callable - A function that takes a single integer index and returns the - corresponding `~specutils.Spectrum` object. - cache_size : int, optional - The number of spectra to cache in memory, by default 10. + 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-loading `SpectrumList`. + A lazy-loadable SpectrumList + """ + # create placeholder list speclist = cls([_placeholder] * length) - cache_size = max(int(cache_size), 0) - if cache_size > 0: + + # optionally cache the loader + if cache_size: + cache_size = max(int(cache_size), 0) loader = lru_cache(maxsize=cache_size)(loader) - speclist._lazy_loader = loader + # set lazy parameters + speclist._lazy_loader = loader speclist._lazy_labels = labels return speclist From c6eeea946207034b639031abe1f8c77948ddd18c Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Thu, 22 Jan 2026 14:56:58 -0500 Subject: [PATCH 05/14] tweaking tests --- .../io/default_loaders/tests/test_roman.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/specutils/io/default_loaders/tests/test_roman.py b/specutils/io/default_loaders/tests/test_roman.py index 9a9a3560f..7c6af9b4a 100644 --- a/specutils/io/default_loaders/tests/test_roman.py +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -153,16 +153,33 @@ def test_roman_1d_combined_list(roman_multi): 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_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 From e6ac9dfcafd0f8dea0d9233d642232f501e2959f Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Thu, 22 Jan 2026 14:58:36 -0500 Subject: [PATCH 06/14] renaming lazy kwarg --- specutils/io/default_loaders/roman.py | 2 +- specutils/io/registers.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/specutils/io/default_loaders/roman.py b/specutils/io/default_loaders/roman.py index 24b01d519..7e934adee 100644 --- a/specutils/io/default_loaders/roman.py +++ b/specutils/io/default_loaders/roman.py @@ -223,7 +223,7 @@ def roman_1d_combined(file_obj, source: str = None, **kwargs): extensions=["asdf"], priority=10, force=True, - lazy_factory=_lazy_load_roman, + lazy_loader=_lazy_load_roman, ) def roman_1d_combined_list(file_obj, **kwargs): """Load all Roman 1d combined extracted spectra diff --git a/specutils/io/registers.py b/specutils/io/registers.py index 92308d0b3..d632c801a 100644 --- a/specutils/io/registers.py +++ b/specutils/io/registers.py @@ -34,7 +34,7 @@ def data_loader( force=False, autogenerate_spectrumlist=True, verbose=False, - lazy_factory=None, + lazy_loader=None, ): """ Wraps a function that can be added to an `~astropy.io.registry` for custom @@ -65,8 +65,8 @@ def data_loader( data_loader that reads Spectrum objects. Default is ``True``. verbose : bool Print extra info. - lazy_factory : Callable, optional - A factory function to create a lazy-loading SpectrumList. + lazy_loader : Callable, optional + A loader function to create a lazy-loading SpectrumList. """ def identifier_wrapper(ident): def wrapper(*args, **kwargs): @@ -88,8 +88,8 @@ def wrapper(*args, **kwargs): # check SpectrumList loaders for lazy option if dtype is SpectrumList: - if lazy_load and lazy_factory is not None: - return lazy_factory(*args, **kwargs) + if lazy_load and lazy_loader is not None: + return lazy_loader(*args, **kwargs) # If not lazy, use eager loader kwargs.pop("cache_size", None) From 86d8473474c152cac080b9f163731ff5eee8ef1f Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Thu, 22 Jan 2026 15:18:32 -0500 Subject: [PATCH 07/14] cleanup --- specutils/io/default_loaders/roman.py | 31 +++++++++++++++++++++------ 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/specutils/io/default_loaders/roman.py b/specutils/io/default_loaders/roman.py index 7e934adee..3f4b8ab0e 100644 --- a/specutils/io/default_loaders/roman.py +++ b/specutils/io/default_loaders/roman.py @@ -117,22 +117,27 @@ def _load_roman_first(file_obj, source=None, **kwargs): 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=kwargs.get("cache_size", None), labels=sources + length=len(sources), loader=_loader, cache_size=cache_size, labels=sources ) - # Optional / independent: enable string indexing for Roman IDs + # set the sources ids as alternate indices sl.set_id_map(source_idx_map) return sl @@ -140,15 +145,18 @@ def _loader(i: int) -> Spectrum: 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()) - source_idx_map = dict(zip(sources, range(len(sources)))) + # 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) @@ -245,8 +253,12 @@ def roman_1d_combined_list(file_obj, **kwargs): @data_loader( - "Roman 1d individual", identifier=identify_1d_individual, dtype=Spectrum, - extensions=['asdf'], priority=10 + "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 @@ -270,8 +282,13 @@ def roman_1d_individual(file_obj, source: str=None, **kwargs): @data_loader( - "Roman 1d individual", identifier=identify_1d_individual, dtype=SpectrumList, - extensions=['asdf'], priority=10, force=True + "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 From 448a210f8632dee70496ea559c0ba6bab492a8d9 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Thu, 22 Jan 2026 16:11:45 -0500 Subject: [PATCH 08/14] adding example lazy loader for sdss-v spec files --- specutils/io/default_loaders/sdss_v.py | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) 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): """ From 95474f7fc25ede3210dc83fe6512ac4f29204923 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Mon, 26 Jan 2026 13:51:45 -0500 Subject: [PATCH 09/14] cleanup --- specutils/io/parsing_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/specutils/io/parsing_utils.py b/specutils/io/parsing_utils.py index 5fbb0c737..3ef9828ad 100644 --- a/specutils/io/parsing_utils.py +++ b/specutils/io/parsing_utils.py @@ -108,7 +108,7 @@ def read_fileobj_or_hdulist(*args, **kwargs): def open_input(fileobj, cache_asdf: bool = None, **kwargs): """Open the asdf info with or without cache""" - # Caller passed an AsdfFile: reuse it if open; reopen if closed and uri available. + # handle an input open or closed asdf instance if isinstance(fileobj, asdf.AsdfFile): if fileobj._fd is not None: return fileobj @@ -120,7 +120,7 @@ def open_input(fileobj, cache_asdf: bool = None, **kwargs): if cache_asdf: return get_cached_object(fileobj, **kwargs) - # Non-cached path: always create a transient handle we close on exit. + # non-cached; return a temp handle we can close name = getattr(fileobj, "name", None) if isinstance(name, str) and name: return asdf.open(name, **kwargs) @@ -158,7 +158,7 @@ def read_fileobj_or_asdftree(*args, **kwargs): pass # caching options - # Close only when we are not caching AND the caller did not provide an already-open AsdfFile. + # close when we aren't caching or when an already-open AsdfFile is provided. cache_asdf = kwargs.pop("cache_asdf", None) should_close = (not cache_asdf) and not (isinstance(fileobj, asdf.AsdfFile) and fileobj._fd is not None) @@ -174,7 +174,7 @@ def read_fileobj_or_asdftree(*args, **kwargs): except Exception: pass - # Always best-effort rewind of non-AsdfFile inputs (keep your current logic) + # cleanup if not isinstance(fileobj, asdf.AsdfFile): try: fileobj.seek(0) From ae72ef166282acd6167591f218439b2a13b2555d Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Wed, 4 Feb 2026 10:05:54 -0500 Subject: [PATCH 10/14] updating lazy repr with inline note --- specutils/spectra/spectrum_list.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index 798ff1ca7..6a24b6d47 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -122,11 +122,12 @@ def __repr__(self) -> str: item repr is used. """ # use normal repr - if not self.is_lazy: + if not self.is_lazy or self.n_loaded == len(self): return super().__repr__() # build the lazy repr - return f"[{', '.join(self._lazy_repr(i) for i in range(len(self)))}]" + 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""" From 44e117bc25627979da2c2f3fbac35bb7cbe8d667 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Wed, 4 Feb 2026 10:12:06 -0500 Subject: [PATCH 11/14] adding lazy repr test --- .../io/default_loaders/tests/test_roman.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/specutils/io/default_loaders/tests/test_roman.py b/specutils/io/default_loaders/tests/test_roman.py index 7c6af9b4a..895761ad2 100644 --- a/specutils/io/default_loaders/tests/test_roman.py +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -187,6 +187,25 @@ def test_roman_lazy_loaded_spectrum(roman_multi): 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 "['402849'" in repr(speclist) + + # load 1 + speclist[0] + assert "lazy list: 1 items loaded;" in repr(speclist) + assert "[ Date: Wed, 4 Feb 2026 13:54:56 -0500 Subject: [PATCH 12/14] fixing test --- specutils/io/default_loaders/tests/test_roman.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/specutils/io/default_loaders/tests/test_roman.py b/specutils/io/default_loaders/tests/test_roman.py index 895761ad2..46f4492bd 100644 --- a/specutils/io/default_loaders/tests/test_roman.py +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -194,7 +194,7 @@ def test_roman_repr(roman_multi): 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 "['402849'" in repr(speclist) + assert "['400'" in repr(speclist) # load 1 speclist[0] From 4e0ee662a28a14a9165536e33bdf6c9c11d00668 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Wed, 4 Feb 2026 16:02:52 -0500 Subject: [PATCH 13/14] add checks for unique labels --- .../io/default_loaders/tests/test_roman.py | 7 +++++++ specutils/spectra/spectrum_list.py | 19 ++++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/specutils/io/default_loaders/tests/test_roman.py b/specutils/io/default_loaders/tests/test_roman.py index 46f4492bd..eacc19eff 100644 --- a/specutils/io/default_loaders/tests/test_roman.py +++ b/specutils/io/default_loaders/tests/test_roman.py @@ -178,6 +178,13 @@ def test_roman_alid_fails(roman_multi): 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) diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index 6a24b6d47..76f44d2e1 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -1,6 +1,6 @@ from functools import lru_cache from typing import Callable, Optional - +from collections import Counter from astropy.nddata import NDIOMixin __all__ = ['SpectrumList'] @@ -56,6 +56,12 @@ def set_id_map(self, id_map: dict[str, int]): 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: @@ -146,6 +152,13 @@ def _lazy_repr(self, ii: int): 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 @@ -187,6 +200,10 @@ def from_lazy( 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 From 2627845b4fa6e80ac17f8fd33ffae5f0edc92542 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Fri, 6 Feb 2026 15:25:10 -0500 Subject: [PATCH 14/14] updating docs --- docs/custom_loading.rst | 280 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 278 insertions(+), 2 deletions(-) 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: