From 88d82e7740aa1a6a1ecc13d102d320763e5ca8fc Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Tue, 30 Jun 2026 11:45:57 -0400 Subject: [PATCH 01/12] init migration of lazy features --- docs/custom_loading.rst | 280 ++++++++++++++++++++++++- specutils/io/default_loaders/sdss_v.py | 28 +++ specutils/io/registers.py | 55 ++++- specutils/spectra/spectrum_list.py | 195 ++++++++++++++++- 4 files changed, 545 insertions(+), 13 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: 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/registers.py b/specutils/io/registers.py index cdee16f99..d632c801a 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_loader=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_loader : Callable, optional + A loader function to create a lazy-loading SpectrumList. """ def identifier_wrapper(ident): def wrapper(*args, **kwargs): @@ -70,13 +80,40 @@ def wrapper(*args, **kwargs): return wrapper def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + 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_loader is not None: + return lazy_loader(*args, **kwargs) + + # If not lazy, use eager loader + kwargs.pop("cache_size", None) + return func(*args, **kwargs) + + # Spectrum loaders + 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 +140,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 +170,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..76f44d2e1 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -1,8 +1,13 @@ +from functools import lru_cache +from typing import Callable, Optional +from collections import Counter from astropy.nddata import NDIOMixin - __all__ = ['SpectrumList'] +# a temporary placeholder object for lists +_placeholder = object() + class SpectrumList(list, NDIOMixin): """ @@ -15,3 +20,191 @@ 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) + + # 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 + """ + 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 From 5d627b47587346ff85c7b888d5c43659b2c10629 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Tue, 30 Jun 2026 11:46:09 -0400 Subject: [PATCH 02/12] init test of jwst lazy loader --- specutils/io/default_loaders/jwst_reader.py | 191 +++++++++++++------- 1 file changed, 124 insertions(+), 67 deletions(-) diff --git a/specutils/io/default_loaders/jwst_reader.py b/specutils/io/default_loaders/jwst_reader.py index d595eeb06..f6c34d5ac 100644 --- a/specutils/io/default_loaders/jwst_reader.py +++ b/specutils/io/default_loaders/jwst_reader.py @@ -132,6 +132,33 @@ def _identify_jwst_fits(*args): except Exception: return False +# def _lazy_x1d_loader(fileobj, **kwargs): +# """Lazy loader example for JWST x1d 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 _jwst_spec1d_loader(hdulist, hdu_idx, **kwargs) +# return _jwst_spec1d_loader(hdulist, extname="COMBINE1D", **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( "JWST c1d", identifier=identify_jwst_c1d_fits, dtype=Spectrum, @@ -159,9 +186,21 @@ def jwst_c1d_single_loader(file_obj, **kwargs): "Use SpectrumList.read() instead.") +def _lazy_jwst_c1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) + + +def _lazy_jwst_x1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + + @data_loader( - "JWST c1d multi", identifier=identify_jwst_c1d_multi_fits, - dtype=SpectrumList, extensions=['fits'], priority=10, + "JWST c1d multi", + identifier=identify_jwst_c1d_multi_fits, + dtype=SpectrumList, + extensions=["fits"], + priority=10, + lazy_loader=_lazy_jwst_c1d_loader, ) def jwst_c1d_multi_loader(file_obj, **kwargs): """ @@ -208,8 +247,12 @@ def jwst_x1d_single_loader(file_obj, **kwargs): @data_loader( - "JWST x1d multi", identifier=identify_jwst_x1d_multi_fits, - dtype=SpectrumList, extensions=['fits'], priority=10, + "JWST x1d multi", + identifier=identify_jwst_x1d_multi_fits, + dtype=SpectrumList, + extensions=["fits"], + priority=10, + lazy_loader=_lazy_jwst_x1d_loader, ) def jwst_x1d_multi_loader(file_obj, **kwargs): """ @@ -229,68 +272,6 @@ def jwst_x1d_multi_loader(file_obj, **kwargs): return _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', **kwargs) -@data_loader( - "JWST x1d MIRI MRS", identifier=identify_jwst_miri_mrs, dtype=SpectrumList, - extensions=['*'], priority=10, -) -def jwst_x1d_miri_mrs_loader(input, missing="raise", **kwargs): - """ - Loader for JWST x1d MIRI MRS spectral data in FITS format. - - A single data set consists of a bunch of _x1d files corresponding to - a variety of wavelength bands. This reader reads them one by one and packs - the result into a SpectrumList instance. - - Parameters - ---------- - input : list of str or file-like - List of FITS file names, or objects (provided from name by - Astropy I/O Registry). Alternatively, a directory path on - which glob.glob runs with pattern an implicit pattern "_x1d.fits", - or a directory path with a glob pattern already set. - missing : {'warn', 'silent'} - Allows the user to continue loading if one file is missing by setting - the value to "warn" or "silent". In the first case a warning will be issued - to the user, in the latter the file will silently be skipped. Any other - value will result in a FileNotFoundError if any files in the list are missing. - - Returns - ------- - SpectrumList - A list of the spectra that are contained in all the files. - """ - - # If input is a list, go read each file. If directory, glob-expand - # list of file names. - if not isinstance(input, (list, tuple)): - if os.path.isdir(input): - file_list = glob.glob(os.path.join(input, "*_x1d.fits"), recursive=True) - else: - file_list = glob.glob(input, recursive=True) - else: - file_list = input - - spectra = [] - for file_obj in file_list: - try: - sp = _jwst_spec1d_loader(file_obj, **kwargs) - except FileNotFoundError as e: - if missing.lower() == "warn": - warnings.warn(f'Failed to load {file_obj}: {repr(e)}') - continue - elif missing.lower() == "silent": - continue - else: - raise FileNotFoundError(f"Failed to load {file_obj}: {repr(e)}. " - "To suppress this error, set argument missing='warn'") - - spectra.append(sp) - - # the call to `chain.from_iterable` allows us to handle multiple HDUs - # stored within multiple FITS files, all unpacked into one `SpectrumList` - return SpectrumList(chain.from_iterable(spectra)) - - def _jwst_spectrum_from_table(data, hdu_header, primary_header, flux_col=None, srctype=None): # Create a Spectrum from either a table or a single row of a table @@ -356,7 +337,7 @@ def _jwst_spectrum_from_table(data, hdu_header, primary_header, flux_col=None, s meta['source_id'] = data['SOURCE_ID'] if unpadded_indices is not None: - # In this case the spectra arrays have been padded to make them have consistent length + # In this case the spectra arrays have been padded to make their arrays have consistent length flux = flux[unpadded_indices].unmasked uncertainty = uncertainty[unpadded_indices] @@ -454,6 +435,82 @@ def jwst_s2d_single_loader(filename, **kwargs): raise RuntimeError(f"Input data has {len(spectrum_list)} spectra.") +def _lazy_jwst_c1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) + + +def _lazy_jwst_x1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + + +def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwargs): + """Lazy loader for JWST x1d/c1d 1-D spectral data in FITS format.""" + + entries = [] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + for hdu_idx, hdu in enumerate(hdulist): + if hdu.name != extname: + continue + + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for column in hdu.columns: + if column.unit in bad_units: + column.unit = bad_units[column.unit] + + data = QTable.read(hdu) + + if data[0]["WAVELENGTH"].shape != (): + for row_idx, row in enumerate(data): + if hasattr(row["WAVELENGTH"], "mask") and np.all(row["WAVELENGTH"].mask): + print(f"Skipping row {row_idx} of HDU {hdu_idx} because all wavelength values are masked.") + continue + + if "SOURCE_ID" in data.colnames: + label = f"hdu{hdu_idx}_source_{row['SOURCE_ID']}" + else: + label = f"hdu{hdu_idx}_{row_idx}" + + entries.append((hdu_idx, row_idx, label)) + else: + entries.append((hdu_idx, None, f"{hdu.name}_{hdu_idx}")) + + if len(entries) == 0: + raise ValueError("No valid HDU found to load.") + + labels = [label for _, _, label in entries] + + def _loader(i: int) -> Spectrum: + hdu_idx, row_idx, _ = entries[i] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + primary_header = hdulist["PRIMARY"].header + hdu = hdulist[hdu_idx] + + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for column in hdu.columns: + if column.unit in bad_units: + column.unit = bad_units[column.unit] + + data = QTable.read(hdu) + + if row_idx is None: + return _jwst_spectrum_from_table(data, hdu.header, primary_header, flux_col) + + row = data[row_idx] + return _jwst_spectrum_from_table(row, hdu.header, primary_header, flux_col, row["SOURCE_TYPE"]) + + sl = SpectrumList.from_lazy(length=len(entries), loader=_loader, labels=labels) + sl.set_id_map(dict(zip(labels, range(len(labels))))) + return sl + + +def _lazy_jwst_c1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) + + +def _lazy_jwst_x1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + + @data_loader( "JWST s2d multi", identifier=identify_jwst_s2d_multi_fits, dtype=SpectrumList, extensions=['fits'], priority=10, From c431946bdbc49157a5030fe337a4abf3904b9bbc Mon Sep 17 00:00:00 2001 From: havok2063 Date: Wed, 1 Jul 2026 09:21:58 -0400 Subject: [PATCH 03/12] expanding laziness to other sdss; simplifying lazy logic --- specutils/io/default_loaders/sdss_v.py | 120 +++++++++++++++++++------ specutils/spectra/spectrum_list.py | 4 + 2 files changed, 99 insertions(+), 25 deletions(-) diff --git a/specutils/io/default_loaders/sdss_v.py b/specutils/io/default_loaders/sdss_v.py index ca32ba566..e2ff16a07 100644 --- a/specutils/io/default_loaders/sdss_v.py +++ b/specutils/io/default_loaders/sdss_v.py @@ -451,31 +451,17 @@ def load_sdss_spec_1D(file_obj, *args, hdu: Optional[int] = None, **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) + """Lazy loader for SDSS-V spec files.""" - 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) + def _select(ext): + return ext.name not in ["SPALL", "ZALL", "ZLINE"] - 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 + def _load(hdulist, hdu_idx): + return _load_BOSS_HDU(hdulist, hdu_idx, **kwargs) + return _sdss_lazy_loader( + fileobj, select_hdu=_select, load_source=_load, **kwargs + ) @data_loader( "SDSS-V spec", @@ -501,12 +487,15 @@ def load_sdss_spec_list(file_obj, **kwargs): The spectra contained in the file. """ with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: - spectra = list() + spectra = SpectrumList() + labels = [] for hdu in range(1, len(hdulist)): if hdulist[hdu].name in ["SPALL", "ZALL", "ZLINE"]: continue spectra.append(_load_BOSS_HDU(hdulist, hdu, **kwargs)) - return SpectrumList(spectra) + labels.append(hdulist[hdu].name) + _set_labels(spectra, labels) + return spectra def _load_BOSS_HDU(hdulist: HDUList, hdu: int, **kwargs): @@ -557,6 +546,78 @@ def _load_BOSS_HDU(hdulist: HDUList, hdu: int, **kwargs): mask=mask, meta=meta) +def _set_labels(spectra: SpectrumList, labels: list): + """Set the labels for each index in a SpectrumList object""" + spectra.set_id_map(dict(zip(labels, range(len(labels))))) + + +def _sdss_lazy_loader(fileobj: object, select_hdu: callable, load_source: callable, **kwargs) -> SpectrumList: + """_summary_ + + _extended_summary_ + + Parameters + ---------- + fileobj : object + the file object to load + select_hdu : callable + function to determine which HDUs to load + load_source : callable + function to load a specific source from an HDU + + Returns + ------- + SpectrumList + The list of spectra contained in the file + """ + # Build list of HDU indices + labels once + hdu_indices = [] + labels = [] + with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: + for idx in range(1, len(hdulist)): + ext = hdulist[idx] + # skip the HDU + if not select_hdu(ext): + continue + hdu_indices.append(idx) + labels.append(ext.name) + + def _loader(i: int) -> Spectrum: + hdu_idx = hdu_indices[i] + with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: + return load_source(hdulist, hdu_idx, **kwargs) + + sl = SpectrumList.from_lazy(length=len(hdu_indices), loader=_loader, labels=labels) + _set_labels(sl, labels) + return sl + + +def _lazy_sdss_mwm_loader(fileobj, **kwargs): + """Lazy loader example for SDSS-V mwm files""" + + def _select(ext): + return ext.header.get("DATASUM") != "0" and len(ext.data) > 0 + + def _load(hdulist, hdu_idx, **kwargs): + return _load_mwmVisit_or_mwmStar_hdu(hdulist, hdu_idx, **kwargs) + + return _sdss_lazy_loader( + fileobj, select_hdu=_select, load_source=_load, **kwargs + ) + + +def _lazy_sdss_astra_loader(fileobj, **kwargs): + """Lazy loader for SDSS-V astra files.""" + + def _select(ext): + return ext.header.get("DATASUM") != "0" and len(ext.data) > 0 + + def _load(hdulist, hdu_idx, **kwargs): + return _load_astra_hdu(hdulist, hdu_idx, visit=0, **kwargs) + + return _sdss_lazy_loader( + fileobj, select_hdu=_select, load_source=_load, **kwargs + ) # MWM LOADERS @data_loader( @@ -614,6 +675,7 @@ def load_sdss_mwm_1d(file_obj, hdu: Optional[int] = None, **kwargs): dtype=SpectrumList, priority=20, extensions=["fits"], + lazy_loader=_lazy_sdss_mwm_loader ) def load_sdss_mwm_list(file_obj, **kwargs): """ @@ -630,6 +692,7 @@ def load_sdss_mwm_list(file_obj, **kwargs): A list of spectra from each visit with each instrument at each observatory (mwmVisit), or the coadd from each instrument/observatory (mwmStar). """ + labels = [] spectra = SpectrumList() with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: # Check if file is empty first @@ -644,7 +707,10 @@ def load_sdss_mwm_list(file_obj, **kwargs): if hduext.header.get("DATASUM") == "0" or len(hduext.data) == 0: # Skip zero data HDU's continue + labels.append(hduext.name) spectra.append(_load_mwmVisit_or_mwmStar_hdu(hdulist, i)) + # + _set_labels(spectra, labels) return spectra @@ -801,6 +867,7 @@ def load_sdss_astra_1d( dtype=SpectrumList, priority=20, extensions=["fits"], + lazy_loader=_lazy_sdss_astra_loader ) def load_sdss_astra_list(file_obj, **kwargs): """Load an astraStar/astraVisit model spectrum file as a `~specutils.SpectrumList`. @@ -818,7 +885,7 @@ def load_sdss_astra_list(file_obj, **kwargs): """ spectra = SpectrumList() - + labels = [] with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: # Check if file is empty first datasums = [] @@ -832,8 +899,11 @@ def load_sdss_astra_list(file_obj, **kwargs): if hdulist[hdu].header.get("DATASUM") == "0": # Skip zero data HDU's continue + labels.append(hdulist[hdu].name) spectra.extend(_load_astra_hdu(hdulist, hdu)) + _set_labels(spectra, labels) + if len(spectra) == 0: raise ValueError("No valid HDU found to load.") diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index 76f44d2e1..d24cb4231 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -64,6 +64,10 @@ def set_id_map(self, id_map: dict[str, int]): self._id_map = dict(id_map) + @property + def labels(self) -> Optional[dict[str]]: + return self._id_map or self._lazy_labels + def _resolve_key(self, key: str) -> int: """Resolve a string key to a list index""" From fa9fdadff4e8cc746ddf3b76b3391686c1dbc35f Mon Sep 17 00:00:00 2001 From: havok2063 Date: Wed, 1 Jul 2026 09:24:38 -0400 Subject: [PATCH 04/12] adding docstring --- specutils/io/default_loaders/sdss_v.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/specutils/io/default_loaders/sdss_v.py b/specutils/io/default_loaders/sdss_v.py index e2ff16a07..9a73c37cc 100644 --- a/specutils/io/default_loaders/sdss_v.py +++ b/specutils/io/default_loaders/sdss_v.py @@ -552,9 +552,11 @@ def _set_labels(spectra: SpectrumList, labels: list): def _sdss_lazy_loader(fileobj: object, select_hdu: callable, load_source: callable, **kwargs) -> SpectrumList: - """_summary_ + """Make a SDSS lazy loader - _extended_summary_ + Create a lazy loader callable that looks up the individual + spectrum to load from an HDU extension on demand. Also + builds a list of labels based on the HDU extension name. Parameters ---------- @@ -582,6 +584,7 @@ def _sdss_lazy_loader(fileobj: object, select_hdu: callable, load_source: callab hdu_indices.append(idx) labels.append(ext.name) + # create the lazy loader callable for SpectrumList def _loader(i: int) -> Spectrum: hdu_idx = hdu_indices[i] with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: From a7c0681733be15f0cdd13f82d0581d2851db3db3 Mon Sep 17 00:00:00 2001 From: havok2063 Date: Wed, 1 Jul 2026 10:45:10 -0400 Subject: [PATCH 05/12] add lazy tests --- .../io/default_loaders/tests/test_sdss_v.py | 51 +++++++++++ specutils/spectra/spectrum_list.py | 6 +- specutils/tests/test_spectrum_list.py | 86 +++++++++++++++++++ 3 files changed, 142 insertions(+), 1 deletion(-) create mode 100644 specutils/tests/test_spectrum_list.py diff --git a/specutils/io/default_loaders/tests/test_sdss_v.py b/specutils/io/default_loaders/tests/test_sdss_v.py index c14cf3568..0c4d9385f 100644 --- a/specutils/io/default_loaders/tests/test_sdss_v.py +++ b/specutils/io/default_loaders/tests/test_sdss_v.py @@ -706,6 +706,36 @@ def test_mwm_list(file_obj, with_wl, hduflags): os.remove(tmpfile) +def test_mwm_lazy_load_with_labels(): + """test SDSS-V mwm lazy loader with labels""" + tmpfile = "mwm-lazy-temp.fits" + hduflags = [1, 0, 1, 1] + nvisits = 3 + mwm_HDUList(hduflags, with_wl=False, nvisits=nvisits).writeto( + tmpfile, + overwrite=True, + ) + + data = SpectrumList.read(tmpfile, format="SDSS-V mwm", lazy_load=True) + assert isinstance(data, SpectrumList) + assert data.is_lazy + assert data.n_loaded == 0 + + assert isinstance(data.labels, dict) + assert "BOSS/APO" in data.labels + assert "APOGEE/APO" in data.labels + + first = data["BOSS/APO"] + assert isinstance(first, Spectrum) + assert data.n_loaded == 1 + + second = data[1] + assert isinstance(second, Spectrum) + assert data.n_loaded == 2 + + os.remove(tmpfile) + + @pytest.mark.parametrize( "file_obj, with_wl, hduflags, pipeline", [ @@ -918,6 +948,27 @@ def test_spec_list(file_obj, n_spectra): os.remove(tmpfile) +def test_spec_lazy_load_with_labels(): + """test SDSS-V spec lazy loader""" + tmpfile = "spec-lazy-temp.fits" + n_spectra = 5 + spec_HDUList(n_spectra).writeto(tmpfile, overwrite=True) + + data = SpectrumList.read(tmpfile, format="SDSS-V spec", lazy_load=True) + assert isinstance(data, SpectrumList) + assert data.is_lazy + assert data.n_loaded == 0 + + assert isinstance(data.labels, dict) + assert "COADD" in data.labels + + coadd = data["COADD"] + assert isinstance(coadd, Spectrum) + assert data.n_loaded == 1 + + os.remove(tmpfile) + + @pytest.mark.parametrize( "file_obj,hdu", [ diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index d24cb4231..5ee8d4c9e 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -65,7 +65,7 @@ def set_id_map(self, id_map: dict[str, int]): self._id_map = dict(id_map) @property - def labels(self) -> Optional[dict[str]]: + def labels(self) -> Optional[dict[str] | list]: return self._id_map or self._lazy_labels def _resolve_key(self, key: str) -> int: @@ -75,6 +75,10 @@ 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) + # if lazy labels but not id map, set the mapping + if not self._id_map and self._lazy_labels: + self.set_id_map(dict(zip(self._lazy_labels, range(len(self._lazy_labels))))) + # otherwise it must be provided by the mapping if self._id_map is None: raise KeyError("No id mapping provided for alternate indexing") diff --git a/specutils/tests/test_spectrum_list.py b/specutils/tests/test_spectrum_list.py new file mode 100644 index 000000000..a5ca04ad4 --- /dev/null +++ b/specutils/tests/test_spectrum_list.py @@ -0,0 +1,86 @@ +import pytest +import string +from specutils.spectra import SpectrumList + + +labels = list(string.ascii_lowercase[:10]) + + +def test_nonlazy_spectrum_list(): + """test non-lazy SpectrumList behave like a normal list.""" + sl = SpectrumList(range(10)) + + assert not sl.is_lazy + assert len(sl) == 10 + assert sl.n_loaded == 10 + assert sl.labels is None + assert '1, 2, 3' in repr(sl) + +def test_nonlazy_labels(): + """test non-lazy lists can have labels.""" + sl = SpectrumList(range(10)) + labels = {f"item{i}": i for i in range(10)} + sl.set_id_map(labels) + + assert not sl.is_lazy + assert len(sl) == 10 + assert sl.n_loaded == 10 + assert sl.labels == labels + assert sl.labels["item3"] == 3 + # it does not change the repr + assert 'item1, item2, item3' not in repr(sl) + + +def test_lazy_spectrum_list_no_labels(): + """test lazy loading items with no labels.""" + # list of items + items = list(range(10)) + + # define an item loader + def loader(i): + return items[i] + + sl = SpectrumList.from_lazy(length=len(items), loader=loader) + + # check if lazy + assert sl.is_lazy + assert len(sl) == 10 + assert sl.n_loaded == 0 + assert sl.labels is None + + # load an item + assert sl[3] == 3 + assert sl.n_loaded == 1 + assert 'lazy list: 1 items loaded' in repr(sl) + + # check the first item isn't loaded yet + assert "load a spectrum:\n[ Date: Thu, 23 Jul 2026 14:10:56 -0400 Subject: [PATCH 06/12] fixing jwst lazy loader --- specutils/io/default_loaders/jwst_reader.py | 140 +++++++++--------- .../default_loaders/tests/test_jwst_reader.py | 32 ++++ 2 files changed, 100 insertions(+), 72 deletions(-) diff --git a/specutils/io/default_loaders/jwst_reader.py b/specutils/io/default_loaders/jwst_reader.py index f6c34d5ac..bf6ea2d86 100644 --- a/specutils/io/default_loaders/jwst_reader.py +++ b/specutils/io/default_loaders/jwst_reader.py @@ -381,11 +381,7 @@ def _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', flux_col=None, **kwargs): 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] + _normalize_jwst_column_units(hdu) data = QTable.read(hdu) @@ -407,62 +403,30 @@ def _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', flux_col=None, **kwargs): return SpectrumList(spectra) -@data_loader( - "JWST s2d", identifier=identify_jwst_s2d_fits, dtype=Spectrum, - extensions=['fits'], priority=10, -) -def jwst_s2d_single_loader(filename, **kwargs): - """ - Loader for JWST s2d 2D rectified spectral data in FITS format. - - Parameters - ---------- - filename : str - The path to the FITS file - - Returns - ------- - Spectrum - The spectrum contained in the file. - """ - spectrum_list = _jwst_s2d_loader(filename, **kwargs) - if len(spectrum_list) == 1: - return spectrum_list[0] - elif len(spectrum_list) > 1: - raise RuntimeError(f"Input data has {len(spectrum_list)} spectra. " - "Use SpectrumList.read() instead.") - else: - raise RuntimeError(f"Input data has {len(spectrum_list)} spectra.") - - -def _lazy_jwst_c1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) - +def _normalize_jwst_column_units(hdu): + """Normalize known malformed JWST table column unit strings in-place.""" + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for column in hdu.columns: + if column.unit in bad_units: + column.unit = bad_units[column.unit] -def _lazy_jwst_x1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) - - -def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwargs): - """Lazy loader for JWST x1d/c1d 1-D spectral data in FITS format.""" +def _get_jwst_spec1d_labels(file_obj, extname="EXTRACT1D", **kwargs): + """Build lazy-load data entries and labels for JWST 1D spectra.""" entries = [] + labels = [] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: for hdu_idx, hdu in enumerate(hdulist): if hdu.name != extname: continue - bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} - for column in hdu.columns: - if column.unit in bad_units: - column.unit = bad_units[column.unit] - + _normalize_jwst_column_units(hdu) data = QTable.read(hdu) if data[0]["WAVELENGTH"].shape != (): for row_idx, row in enumerate(data): if hasattr(row["WAVELENGTH"], "mask") and np.all(row["WAVELENGTH"].mask): - print(f"Skipping row {row_idx} of HDU {hdu_idx} because all wavelength values are masked.") continue if "SOURCE_ID" in data.colnames: @@ -470,45 +434,77 @@ def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwa else: label = f"hdu{hdu_idx}_{row_idx}" - entries.append((hdu_idx, row_idx, label)) + entries.append((hdu_idx, row_idx)) + labels.append(label) else: - entries.append((hdu_idx, None, f"{hdu.name}_{hdu_idx}")) + entries.append((hdu_idx, None)) + labels.append(f"{hdu.name}_{hdu_idx}") - if len(entries) == 0: - raise ValueError("No valid HDU found to load.") + return entries, labels - labels = [label for _, _, label in entries] - def _loader(i: int) -> Spectrum: - hdu_idx, row_idx, _ = entries[i] - with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: - primary_header = hdulist["PRIMARY"].header - hdu = hdulist[hdu_idx] +def _load_jwst_spec1d_lazy_source(file_obj, entry, flux_col=None, **kwargs): + """Load a single JWST 1D spectrum from a lazy entry tuple.""" + hdu_idx, row_idx = entry - bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} - for column in hdu.columns: - if column.unit in bad_units: - column.unit = bad_units[column.unit] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + primary_header = hdulist["PRIMARY"].header + hdu = hdulist[hdu_idx] - data = QTable.read(hdu) + _normalize_jwst_column_units(hdu) + data = QTable.read(hdu) + + if row_idx is None: + return _jwst_spectrum_from_table(data, hdu.header, primary_header, flux_col) + + row = data[row_idx] + return _jwst_spectrum_from_table(row, hdu.header, primary_header, flux_col, row["SOURCE_TYPE"]) + + +def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwargs): + """Lazy loader for JWST x1d/c1d 1-D spectral data in FITS format.""" - if row_idx is None: - return _jwst_spectrum_from_table(data, hdu.header, primary_header, flux_col) + # get the data labels and rows + entries, labels = _get_jwst_spec1d_labels(file_obj, extname=extname, **kwargs) + if len(entries) == 0: + raise ValueError("No valid HDU found to load.") - row = data[row_idx] - return _jwst_spectrum_from_table(row, hdu.header, primary_header, flux_col, row["SOURCE_TYPE"]) + # the single source loader fxn + def _loader(i: int) -> Spectrum: + return _load_jwst_spec1d_lazy_source(file_obj, entries[i], flux_col=flux_col, **kwargs) sl = SpectrumList.from_lazy(length=len(entries), loader=_loader, labels=labels) - sl.set_id_map(dict(zip(labels, range(len(labels))))) return sl -def _lazy_jwst_c1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) +@data_loader( + "JWST s2d", + identifier=identify_jwst_s2d_fits, + dtype=Spectrum, + extensions=["fits"], + priority=10, +) +def jwst_s2d_single_loader(filename, **kwargs): + """ + Loader for JWST s2d 2D rectified spectral data in FITS format. + Parameters + ---------- + filename : str + The path to the FITS file -def _lazy_jwst_x1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + Returns + ------- + Spectrum + The spectrum contained in the file. + """ + spectrum_list = _jwst_s2d_loader(filename, **kwargs) + if len(spectrum_list) == 1: + return spectrum_list[0] + elif len(spectrum_list) > 1: + raise RuntimeError(f"Input data has {len(spectrum_list)} spectra. Use SpectrumList.read() instead.") + else: + raise RuntimeError(f"Input data has {len(spectrum_list)} spectra.") @data_loader( diff --git a/specutils/io/default_loaders/tests/test_jwst_reader.py b/specutils/io/default_loaders/tests/test_jwst_reader.py index b640d2625..040e3b82f 100644 --- a/specutils/io/default_loaders/tests/test_jwst_reader.py +++ b/specutils/io/default_loaders/tests/test_jwst_reader.py @@ -182,6 +182,38 @@ def test_jwst_wfss_multi_reader(tmp_path, spec_multi_new, format): assert data[2].flux.unit == u.MJy/u.sr +@pytest.mark.parametrize( + "spec_multi_new, format", + [("EXTRACT1D", "JWST x1d multi"), ("COMBINE1D", "JWST c1d multi")], + indirect=["spec_multi_new"], +) +def test_jwst_wfss_multi_lazy_load_with_labels(tmp_path, spec_multi_new, format): + """Test lazy loading and label-based access for JWST c1d/x1d packed multi data.""" + tmpfile = str(tmp_path / "jwst-lazy.fits") + spec_multi_new.writeto(tmpfile) + + data = SpectrumList.read(tmpfile, format=format, lazy_load=True) + assert isinstance(data, SpectrumList) + assert data.is_lazy + assert data.n_loaded == 0 + assert len(data) == 3 + + # Before string-key access, lazy labels are exposed as a list. + assert isinstance(data.labels, list) + assert "hdu1_source_1" in data.labels + assert "hdu1_source_2" in data.labels + assert "hdu1_source_3" in data.labels + + first = data["hdu1_source_1"] + assert isinstance(first, Spectrum) + assert data.n_loaded == 1 + + second = data[1] + assert isinstance(second, Spectrum) + assert data.n_loaded == 2 + assert len(data) == 3 + + @pytest.mark.parametrize('spec_single, format', [('EXTRACT1D', 'JWST x1d'), ('COMBINE1D', 'JWST c1d')], indirect=['spec_single']) From 41606c1f660ad4c8f501c027ea2fad635764e9e6 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Tue, 30 Jun 2026 11:45:57 -0400 Subject: [PATCH 07/12] init migration of lazy features --- docs/custom_loading.rst | 280 ++++++++++++++++++++++++- specutils/io/default_loaders/sdss_v.py | 28 +++ specutils/io/registers.py | 55 ++++- specutils/spectra/spectrum_list.py | 195 ++++++++++++++++- 4 files changed, 545 insertions(+), 13 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: 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/registers.py b/specutils/io/registers.py index cdee16f99..d632c801a 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_loader=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_loader : Callable, optional + A loader function to create a lazy-loading SpectrumList. """ def identifier_wrapper(ident): def wrapper(*args, **kwargs): @@ -70,13 +80,40 @@ def wrapper(*args, **kwargs): return wrapper def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + 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_loader is not None: + return lazy_loader(*args, **kwargs) + + # If not lazy, use eager loader + kwargs.pop("cache_size", None) + return func(*args, **kwargs) + + # Spectrum loaders + 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 +140,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 +170,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..76f44d2e1 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -1,8 +1,13 @@ +from functools import lru_cache +from typing import Callable, Optional +from collections import Counter from astropy.nddata import NDIOMixin - __all__ = ['SpectrumList'] +# a temporary placeholder object for lists +_placeholder = object() + class SpectrumList(list, NDIOMixin): """ @@ -15,3 +20,191 @@ 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) + + # 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 + """ + 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 From e6e37033e33dd670cac764d626e0f76f6b776168 Mon Sep 17 00:00:00 2001 From: Brian Cherinka Date: Tue, 30 Jun 2026 11:46:09 -0400 Subject: [PATCH 08/12] init test of jwst lazy loader --- specutils/io/default_loaders/jwst_reader.py | 191 +++++++++++++------- 1 file changed, 124 insertions(+), 67 deletions(-) diff --git a/specutils/io/default_loaders/jwst_reader.py b/specutils/io/default_loaders/jwst_reader.py index d595eeb06..f6c34d5ac 100644 --- a/specutils/io/default_loaders/jwst_reader.py +++ b/specutils/io/default_loaders/jwst_reader.py @@ -132,6 +132,33 @@ def _identify_jwst_fits(*args): except Exception: return False +# def _lazy_x1d_loader(fileobj, **kwargs): +# """Lazy loader example for JWST x1d 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 _jwst_spec1d_loader(hdulist, hdu_idx, **kwargs) +# return _jwst_spec1d_loader(hdulist, extname="COMBINE1D", **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( "JWST c1d", identifier=identify_jwst_c1d_fits, dtype=Spectrum, @@ -159,9 +186,21 @@ def jwst_c1d_single_loader(file_obj, **kwargs): "Use SpectrumList.read() instead.") +def _lazy_jwst_c1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) + + +def _lazy_jwst_x1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + + @data_loader( - "JWST c1d multi", identifier=identify_jwst_c1d_multi_fits, - dtype=SpectrumList, extensions=['fits'], priority=10, + "JWST c1d multi", + identifier=identify_jwst_c1d_multi_fits, + dtype=SpectrumList, + extensions=["fits"], + priority=10, + lazy_loader=_lazy_jwst_c1d_loader, ) def jwst_c1d_multi_loader(file_obj, **kwargs): """ @@ -208,8 +247,12 @@ def jwst_x1d_single_loader(file_obj, **kwargs): @data_loader( - "JWST x1d multi", identifier=identify_jwst_x1d_multi_fits, - dtype=SpectrumList, extensions=['fits'], priority=10, + "JWST x1d multi", + identifier=identify_jwst_x1d_multi_fits, + dtype=SpectrumList, + extensions=["fits"], + priority=10, + lazy_loader=_lazy_jwst_x1d_loader, ) def jwst_x1d_multi_loader(file_obj, **kwargs): """ @@ -229,68 +272,6 @@ def jwst_x1d_multi_loader(file_obj, **kwargs): return _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', **kwargs) -@data_loader( - "JWST x1d MIRI MRS", identifier=identify_jwst_miri_mrs, dtype=SpectrumList, - extensions=['*'], priority=10, -) -def jwst_x1d_miri_mrs_loader(input, missing="raise", **kwargs): - """ - Loader for JWST x1d MIRI MRS spectral data in FITS format. - - A single data set consists of a bunch of _x1d files corresponding to - a variety of wavelength bands. This reader reads them one by one and packs - the result into a SpectrumList instance. - - Parameters - ---------- - input : list of str or file-like - List of FITS file names, or objects (provided from name by - Astropy I/O Registry). Alternatively, a directory path on - which glob.glob runs with pattern an implicit pattern "_x1d.fits", - or a directory path with a glob pattern already set. - missing : {'warn', 'silent'} - Allows the user to continue loading if one file is missing by setting - the value to "warn" or "silent". In the first case a warning will be issued - to the user, in the latter the file will silently be skipped. Any other - value will result in a FileNotFoundError if any files in the list are missing. - - Returns - ------- - SpectrumList - A list of the spectra that are contained in all the files. - """ - - # If input is a list, go read each file. If directory, glob-expand - # list of file names. - if not isinstance(input, (list, tuple)): - if os.path.isdir(input): - file_list = glob.glob(os.path.join(input, "*_x1d.fits"), recursive=True) - else: - file_list = glob.glob(input, recursive=True) - else: - file_list = input - - spectra = [] - for file_obj in file_list: - try: - sp = _jwst_spec1d_loader(file_obj, **kwargs) - except FileNotFoundError as e: - if missing.lower() == "warn": - warnings.warn(f'Failed to load {file_obj}: {repr(e)}') - continue - elif missing.lower() == "silent": - continue - else: - raise FileNotFoundError(f"Failed to load {file_obj}: {repr(e)}. " - "To suppress this error, set argument missing='warn'") - - spectra.append(sp) - - # the call to `chain.from_iterable` allows us to handle multiple HDUs - # stored within multiple FITS files, all unpacked into one `SpectrumList` - return SpectrumList(chain.from_iterable(spectra)) - - def _jwst_spectrum_from_table(data, hdu_header, primary_header, flux_col=None, srctype=None): # Create a Spectrum from either a table or a single row of a table @@ -356,7 +337,7 @@ def _jwst_spectrum_from_table(data, hdu_header, primary_header, flux_col=None, s meta['source_id'] = data['SOURCE_ID'] if unpadded_indices is not None: - # In this case the spectra arrays have been padded to make them have consistent length + # In this case the spectra arrays have been padded to make their arrays have consistent length flux = flux[unpadded_indices].unmasked uncertainty = uncertainty[unpadded_indices] @@ -454,6 +435,82 @@ def jwst_s2d_single_loader(filename, **kwargs): raise RuntimeError(f"Input data has {len(spectrum_list)} spectra.") +def _lazy_jwst_c1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) + + +def _lazy_jwst_x1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + + +def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwargs): + """Lazy loader for JWST x1d/c1d 1-D spectral data in FITS format.""" + + entries = [] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + for hdu_idx, hdu in enumerate(hdulist): + if hdu.name != extname: + continue + + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for column in hdu.columns: + if column.unit in bad_units: + column.unit = bad_units[column.unit] + + data = QTable.read(hdu) + + if data[0]["WAVELENGTH"].shape != (): + for row_idx, row in enumerate(data): + if hasattr(row["WAVELENGTH"], "mask") and np.all(row["WAVELENGTH"].mask): + print(f"Skipping row {row_idx} of HDU {hdu_idx} because all wavelength values are masked.") + continue + + if "SOURCE_ID" in data.colnames: + label = f"hdu{hdu_idx}_source_{row['SOURCE_ID']}" + else: + label = f"hdu{hdu_idx}_{row_idx}" + + entries.append((hdu_idx, row_idx, label)) + else: + entries.append((hdu_idx, None, f"{hdu.name}_{hdu_idx}")) + + if len(entries) == 0: + raise ValueError("No valid HDU found to load.") + + labels = [label for _, _, label in entries] + + def _loader(i: int) -> Spectrum: + hdu_idx, row_idx, _ = entries[i] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + primary_header = hdulist["PRIMARY"].header + hdu = hdulist[hdu_idx] + + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for column in hdu.columns: + if column.unit in bad_units: + column.unit = bad_units[column.unit] + + data = QTable.read(hdu) + + if row_idx is None: + return _jwst_spectrum_from_table(data, hdu.header, primary_header, flux_col) + + row = data[row_idx] + return _jwst_spectrum_from_table(row, hdu.header, primary_header, flux_col, row["SOURCE_TYPE"]) + + sl = SpectrumList.from_lazy(length=len(entries), loader=_loader, labels=labels) + sl.set_id_map(dict(zip(labels, range(len(labels))))) + return sl + + +def _lazy_jwst_c1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) + + +def _lazy_jwst_x1d_loader(fileobj, **kwargs): + return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + + @data_loader( "JWST s2d multi", identifier=identify_jwst_s2d_multi_fits, dtype=SpectrumList, extensions=['fits'], priority=10, From e15650e659ba479f74fb6cbaddaf9690402a06d9 Mon Sep 17 00:00:00 2001 From: havok2063 Date: Wed, 1 Jul 2026 09:21:58 -0400 Subject: [PATCH 09/12] expanding laziness to other sdss; simplifying lazy logic --- specutils/io/default_loaders/sdss_v.py | 120 +++++++++++++++++++------ specutils/spectra/spectrum_list.py | 4 + 2 files changed, 99 insertions(+), 25 deletions(-) diff --git a/specutils/io/default_loaders/sdss_v.py b/specutils/io/default_loaders/sdss_v.py index ca32ba566..e2ff16a07 100644 --- a/specutils/io/default_loaders/sdss_v.py +++ b/specutils/io/default_loaders/sdss_v.py @@ -451,31 +451,17 @@ def load_sdss_spec_1D(file_obj, *args, hdu: Optional[int] = None, **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) + """Lazy loader for SDSS-V spec files.""" - 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) + def _select(ext): + return ext.name not in ["SPALL", "ZALL", "ZLINE"] - 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 + def _load(hdulist, hdu_idx): + return _load_BOSS_HDU(hdulist, hdu_idx, **kwargs) + return _sdss_lazy_loader( + fileobj, select_hdu=_select, load_source=_load, **kwargs + ) @data_loader( "SDSS-V spec", @@ -501,12 +487,15 @@ def load_sdss_spec_list(file_obj, **kwargs): The spectra contained in the file. """ with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: - spectra = list() + spectra = SpectrumList() + labels = [] for hdu in range(1, len(hdulist)): if hdulist[hdu].name in ["SPALL", "ZALL", "ZLINE"]: continue spectra.append(_load_BOSS_HDU(hdulist, hdu, **kwargs)) - return SpectrumList(spectra) + labels.append(hdulist[hdu].name) + _set_labels(spectra, labels) + return spectra def _load_BOSS_HDU(hdulist: HDUList, hdu: int, **kwargs): @@ -557,6 +546,78 @@ def _load_BOSS_HDU(hdulist: HDUList, hdu: int, **kwargs): mask=mask, meta=meta) +def _set_labels(spectra: SpectrumList, labels: list): + """Set the labels for each index in a SpectrumList object""" + spectra.set_id_map(dict(zip(labels, range(len(labels))))) + + +def _sdss_lazy_loader(fileobj: object, select_hdu: callable, load_source: callable, **kwargs) -> SpectrumList: + """_summary_ + + _extended_summary_ + + Parameters + ---------- + fileobj : object + the file object to load + select_hdu : callable + function to determine which HDUs to load + load_source : callable + function to load a specific source from an HDU + + Returns + ------- + SpectrumList + The list of spectra contained in the file + """ + # Build list of HDU indices + labels once + hdu_indices = [] + labels = [] + with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: + for idx in range(1, len(hdulist)): + ext = hdulist[idx] + # skip the HDU + if not select_hdu(ext): + continue + hdu_indices.append(idx) + labels.append(ext.name) + + def _loader(i: int) -> Spectrum: + hdu_idx = hdu_indices[i] + with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: + return load_source(hdulist, hdu_idx, **kwargs) + + sl = SpectrumList.from_lazy(length=len(hdu_indices), loader=_loader, labels=labels) + _set_labels(sl, labels) + return sl + + +def _lazy_sdss_mwm_loader(fileobj, **kwargs): + """Lazy loader example for SDSS-V mwm files""" + + def _select(ext): + return ext.header.get("DATASUM") != "0" and len(ext.data) > 0 + + def _load(hdulist, hdu_idx, **kwargs): + return _load_mwmVisit_or_mwmStar_hdu(hdulist, hdu_idx, **kwargs) + + return _sdss_lazy_loader( + fileobj, select_hdu=_select, load_source=_load, **kwargs + ) + + +def _lazy_sdss_astra_loader(fileobj, **kwargs): + """Lazy loader for SDSS-V astra files.""" + + def _select(ext): + return ext.header.get("DATASUM") != "0" and len(ext.data) > 0 + + def _load(hdulist, hdu_idx, **kwargs): + return _load_astra_hdu(hdulist, hdu_idx, visit=0, **kwargs) + + return _sdss_lazy_loader( + fileobj, select_hdu=_select, load_source=_load, **kwargs + ) # MWM LOADERS @data_loader( @@ -614,6 +675,7 @@ def load_sdss_mwm_1d(file_obj, hdu: Optional[int] = None, **kwargs): dtype=SpectrumList, priority=20, extensions=["fits"], + lazy_loader=_lazy_sdss_mwm_loader ) def load_sdss_mwm_list(file_obj, **kwargs): """ @@ -630,6 +692,7 @@ def load_sdss_mwm_list(file_obj, **kwargs): A list of spectra from each visit with each instrument at each observatory (mwmVisit), or the coadd from each instrument/observatory (mwmStar). """ + labels = [] spectra = SpectrumList() with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: # Check if file is empty first @@ -644,7 +707,10 @@ def load_sdss_mwm_list(file_obj, **kwargs): if hduext.header.get("DATASUM") == "0" or len(hduext.data) == 0: # Skip zero data HDU's continue + labels.append(hduext.name) spectra.append(_load_mwmVisit_or_mwmStar_hdu(hdulist, i)) + # + _set_labels(spectra, labels) return spectra @@ -801,6 +867,7 @@ def load_sdss_astra_1d( dtype=SpectrumList, priority=20, extensions=["fits"], + lazy_loader=_lazy_sdss_astra_loader ) def load_sdss_astra_list(file_obj, **kwargs): """Load an astraStar/astraVisit model spectrum file as a `~specutils.SpectrumList`. @@ -818,7 +885,7 @@ def load_sdss_astra_list(file_obj, **kwargs): """ spectra = SpectrumList() - + labels = [] with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: # Check if file is empty first datasums = [] @@ -832,8 +899,11 @@ def load_sdss_astra_list(file_obj, **kwargs): if hdulist[hdu].header.get("DATASUM") == "0": # Skip zero data HDU's continue + labels.append(hdulist[hdu].name) spectra.extend(_load_astra_hdu(hdulist, hdu)) + _set_labels(spectra, labels) + if len(spectra) == 0: raise ValueError("No valid HDU found to load.") diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index 76f44d2e1..d24cb4231 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -64,6 +64,10 @@ def set_id_map(self, id_map: dict[str, int]): self._id_map = dict(id_map) + @property + def labels(self) -> Optional[dict[str]]: + return self._id_map or self._lazy_labels + def _resolve_key(self, key: str) -> int: """Resolve a string key to a list index""" From adeb7bca4eba17de968353eb972ed4501698a548 Mon Sep 17 00:00:00 2001 From: havok2063 Date: Wed, 1 Jul 2026 09:24:38 -0400 Subject: [PATCH 10/12] adding docstring --- specutils/io/default_loaders/sdss_v.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/specutils/io/default_loaders/sdss_v.py b/specutils/io/default_loaders/sdss_v.py index e2ff16a07..9a73c37cc 100644 --- a/specutils/io/default_loaders/sdss_v.py +++ b/specutils/io/default_loaders/sdss_v.py @@ -552,9 +552,11 @@ def _set_labels(spectra: SpectrumList, labels: list): def _sdss_lazy_loader(fileobj: object, select_hdu: callable, load_source: callable, **kwargs) -> SpectrumList: - """_summary_ + """Make a SDSS lazy loader - _extended_summary_ + Create a lazy loader callable that looks up the individual + spectrum to load from an HDU extension on demand. Also + builds a list of labels based on the HDU extension name. Parameters ---------- @@ -582,6 +584,7 @@ def _sdss_lazy_loader(fileobj: object, select_hdu: callable, load_source: callab hdu_indices.append(idx) labels.append(ext.name) + # create the lazy loader callable for SpectrumList def _loader(i: int) -> Spectrum: hdu_idx = hdu_indices[i] with read_fileobj_or_hdulist(fileobj, memmap=False, **kwargs) as hdulist: From 4799d8c8c5dc83162448d5e78e1c1a584e2f87fa Mon Sep 17 00:00:00 2001 From: havok2063 Date: Wed, 1 Jul 2026 10:45:10 -0400 Subject: [PATCH 11/12] add lazy tests --- .../io/default_loaders/tests/test_sdss_v.py | 51 +++++++++++ specutils/spectra/spectrum_list.py | 6 +- specutils/tests/test_spectrum_list.py | 86 +++++++++++++++++++ 3 files changed, 142 insertions(+), 1 deletion(-) create mode 100644 specutils/tests/test_spectrum_list.py diff --git a/specutils/io/default_loaders/tests/test_sdss_v.py b/specutils/io/default_loaders/tests/test_sdss_v.py index c14cf3568..0c4d9385f 100644 --- a/specutils/io/default_loaders/tests/test_sdss_v.py +++ b/specutils/io/default_loaders/tests/test_sdss_v.py @@ -706,6 +706,36 @@ def test_mwm_list(file_obj, with_wl, hduflags): os.remove(tmpfile) +def test_mwm_lazy_load_with_labels(): + """test SDSS-V mwm lazy loader with labels""" + tmpfile = "mwm-lazy-temp.fits" + hduflags = [1, 0, 1, 1] + nvisits = 3 + mwm_HDUList(hduflags, with_wl=False, nvisits=nvisits).writeto( + tmpfile, + overwrite=True, + ) + + data = SpectrumList.read(tmpfile, format="SDSS-V mwm", lazy_load=True) + assert isinstance(data, SpectrumList) + assert data.is_lazy + assert data.n_loaded == 0 + + assert isinstance(data.labels, dict) + assert "BOSS/APO" in data.labels + assert "APOGEE/APO" in data.labels + + first = data["BOSS/APO"] + assert isinstance(first, Spectrum) + assert data.n_loaded == 1 + + second = data[1] + assert isinstance(second, Spectrum) + assert data.n_loaded == 2 + + os.remove(tmpfile) + + @pytest.mark.parametrize( "file_obj, with_wl, hduflags, pipeline", [ @@ -918,6 +948,27 @@ def test_spec_list(file_obj, n_spectra): os.remove(tmpfile) +def test_spec_lazy_load_with_labels(): + """test SDSS-V spec lazy loader""" + tmpfile = "spec-lazy-temp.fits" + n_spectra = 5 + spec_HDUList(n_spectra).writeto(tmpfile, overwrite=True) + + data = SpectrumList.read(tmpfile, format="SDSS-V spec", lazy_load=True) + assert isinstance(data, SpectrumList) + assert data.is_lazy + assert data.n_loaded == 0 + + assert isinstance(data.labels, dict) + assert "COADD" in data.labels + + coadd = data["COADD"] + assert isinstance(coadd, Spectrum) + assert data.n_loaded == 1 + + os.remove(tmpfile) + + @pytest.mark.parametrize( "file_obj,hdu", [ diff --git a/specutils/spectra/spectrum_list.py b/specutils/spectra/spectrum_list.py index d24cb4231..5ee8d4c9e 100644 --- a/specutils/spectra/spectrum_list.py +++ b/specutils/spectra/spectrum_list.py @@ -65,7 +65,7 @@ def set_id_map(self, id_map: dict[str, int]): self._id_map = dict(id_map) @property - def labels(self) -> Optional[dict[str]]: + def labels(self) -> Optional[dict[str] | list]: return self._id_map or self._lazy_labels def _resolve_key(self, key: str) -> int: @@ -75,6 +75,10 @@ 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) + # if lazy labels but not id map, set the mapping + if not self._id_map and self._lazy_labels: + self.set_id_map(dict(zip(self._lazy_labels, range(len(self._lazy_labels))))) + # otherwise it must be provided by the mapping if self._id_map is None: raise KeyError("No id mapping provided for alternate indexing") diff --git a/specutils/tests/test_spectrum_list.py b/specutils/tests/test_spectrum_list.py new file mode 100644 index 000000000..a5ca04ad4 --- /dev/null +++ b/specutils/tests/test_spectrum_list.py @@ -0,0 +1,86 @@ +import pytest +import string +from specutils.spectra import SpectrumList + + +labels = list(string.ascii_lowercase[:10]) + + +def test_nonlazy_spectrum_list(): + """test non-lazy SpectrumList behave like a normal list.""" + sl = SpectrumList(range(10)) + + assert not sl.is_lazy + assert len(sl) == 10 + assert sl.n_loaded == 10 + assert sl.labels is None + assert '1, 2, 3' in repr(sl) + +def test_nonlazy_labels(): + """test non-lazy lists can have labels.""" + sl = SpectrumList(range(10)) + labels = {f"item{i}": i for i in range(10)} + sl.set_id_map(labels) + + assert not sl.is_lazy + assert len(sl) == 10 + assert sl.n_loaded == 10 + assert sl.labels == labels + assert sl.labels["item3"] == 3 + # it does not change the repr + assert 'item1, item2, item3' not in repr(sl) + + +def test_lazy_spectrum_list_no_labels(): + """test lazy loading items with no labels.""" + # list of items + items = list(range(10)) + + # define an item loader + def loader(i): + return items[i] + + sl = SpectrumList.from_lazy(length=len(items), loader=loader) + + # check if lazy + assert sl.is_lazy + assert len(sl) == 10 + assert sl.n_loaded == 0 + assert sl.labels is None + + # load an item + assert sl[3] == 3 + assert sl.n_loaded == 1 + assert 'lazy list: 1 items loaded' in repr(sl) + + # check the first item isn't loaded yet + assert "load a spectrum:\n[ Date: Thu, 23 Jul 2026 14:10:56 -0400 Subject: [PATCH 12/12] fixing jwst lazy loader --- specutils/io/default_loaders/jwst_reader.py | 140 +++++++++--------- .../default_loaders/tests/test_jwst_reader.py | 32 ++++ 2 files changed, 100 insertions(+), 72 deletions(-) diff --git a/specutils/io/default_loaders/jwst_reader.py b/specutils/io/default_loaders/jwst_reader.py index f6c34d5ac..bf6ea2d86 100644 --- a/specutils/io/default_loaders/jwst_reader.py +++ b/specutils/io/default_loaders/jwst_reader.py @@ -381,11 +381,7 @@ def _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', flux_col=None, **kwargs): 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] + _normalize_jwst_column_units(hdu) data = QTable.read(hdu) @@ -407,62 +403,30 @@ def _jwst_spec1d_loader(file_obj, extname='EXTRACT1D', flux_col=None, **kwargs): return SpectrumList(spectra) -@data_loader( - "JWST s2d", identifier=identify_jwst_s2d_fits, dtype=Spectrum, - extensions=['fits'], priority=10, -) -def jwst_s2d_single_loader(filename, **kwargs): - """ - Loader for JWST s2d 2D rectified spectral data in FITS format. - - Parameters - ---------- - filename : str - The path to the FITS file - - Returns - ------- - Spectrum - The spectrum contained in the file. - """ - spectrum_list = _jwst_s2d_loader(filename, **kwargs) - if len(spectrum_list) == 1: - return spectrum_list[0] - elif len(spectrum_list) > 1: - raise RuntimeError(f"Input data has {len(spectrum_list)} spectra. " - "Use SpectrumList.read() instead.") - else: - raise RuntimeError(f"Input data has {len(spectrum_list)} spectra.") - - -def _lazy_jwst_c1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) - +def _normalize_jwst_column_units(hdu): + """Normalize known malformed JWST table column unit strings in-place.""" + bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} + for column in hdu.columns: + if column.unit in bad_units: + column.unit = bad_units[column.unit] -def _lazy_jwst_x1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) - - -def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwargs): - """Lazy loader for JWST x1d/c1d 1-D spectral data in FITS format.""" +def _get_jwst_spec1d_labels(file_obj, extname="EXTRACT1D", **kwargs): + """Build lazy-load data entries and labels for JWST 1D spectra.""" entries = [] + labels = [] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: for hdu_idx, hdu in enumerate(hdulist): if hdu.name != extname: continue - bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} - for column in hdu.columns: - if column.unit in bad_units: - column.unit = bad_units[column.unit] - + _normalize_jwst_column_units(hdu) data = QTable.read(hdu) if data[0]["WAVELENGTH"].shape != (): for row_idx, row in enumerate(data): if hasattr(row["WAVELENGTH"], "mask") and np.all(row["WAVELENGTH"].mask): - print(f"Skipping row {row_idx} of HDU {hdu_idx} because all wavelength values are masked.") continue if "SOURCE_ID" in data.colnames: @@ -470,45 +434,77 @@ def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwa else: label = f"hdu{hdu_idx}_{row_idx}" - entries.append((hdu_idx, row_idx, label)) + entries.append((hdu_idx, row_idx)) + labels.append(label) else: - entries.append((hdu_idx, None, f"{hdu.name}_{hdu_idx}")) + entries.append((hdu_idx, None)) + labels.append(f"{hdu.name}_{hdu_idx}") - if len(entries) == 0: - raise ValueError("No valid HDU found to load.") + return entries, labels - labels = [label for _, _, label in entries] - def _loader(i: int) -> Spectrum: - hdu_idx, row_idx, _ = entries[i] - with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: - primary_header = hdulist["PRIMARY"].header - hdu = hdulist[hdu_idx] +def _load_jwst_spec1d_lazy_source(file_obj, entry, flux_col=None, **kwargs): + """Load a single JWST 1D spectrum from a lazy entry tuple.""" + hdu_idx, row_idx = entry - bad_units = {"(MJy/sr)^2": "MJy2 sr-2"} - for column in hdu.columns: - if column.unit in bad_units: - column.unit = bad_units[column.unit] + with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist: + primary_header = hdulist["PRIMARY"].header + hdu = hdulist[hdu_idx] - data = QTable.read(hdu) + _normalize_jwst_column_units(hdu) + data = QTable.read(hdu) + + if row_idx is None: + return _jwst_spectrum_from_table(data, hdu.header, primary_header, flux_col) + + row = data[row_idx] + return _jwst_spectrum_from_table(row, hdu.header, primary_header, flux_col, row["SOURCE_TYPE"]) + + +def _jwst_spec1d_lazy_loader(file_obj, extname="EXTRACT1D", flux_col=None, **kwargs): + """Lazy loader for JWST x1d/c1d 1-D spectral data in FITS format.""" - if row_idx is None: - return _jwst_spectrum_from_table(data, hdu.header, primary_header, flux_col) + # get the data labels and rows + entries, labels = _get_jwst_spec1d_labels(file_obj, extname=extname, **kwargs) + if len(entries) == 0: + raise ValueError("No valid HDU found to load.") - row = data[row_idx] - return _jwst_spectrum_from_table(row, hdu.header, primary_header, flux_col, row["SOURCE_TYPE"]) + # the single source loader fxn + def _loader(i: int) -> Spectrum: + return _load_jwst_spec1d_lazy_source(file_obj, entries[i], flux_col=flux_col, **kwargs) sl = SpectrumList.from_lazy(length=len(entries), loader=_loader, labels=labels) - sl.set_id_map(dict(zip(labels, range(len(labels))))) return sl -def _lazy_jwst_c1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="COMBINE1D", **kwargs) +@data_loader( + "JWST s2d", + identifier=identify_jwst_s2d_fits, + dtype=Spectrum, + extensions=["fits"], + priority=10, +) +def jwst_s2d_single_loader(filename, **kwargs): + """ + Loader for JWST s2d 2D rectified spectral data in FITS format. + Parameters + ---------- + filename : str + The path to the FITS file -def _lazy_jwst_x1d_loader(fileobj, **kwargs): - return _jwst_spec1d_lazy_loader(fileobj, extname="EXTRACT1D", **kwargs) + Returns + ------- + Spectrum + The spectrum contained in the file. + """ + spectrum_list = _jwst_s2d_loader(filename, **kwargs) + if len(spectrum_list) == 1: + return spectrum_list[0] + elif len(spectrum_list) > 1: + raise RuntimeError(f"Input data has {len(spectrum_list)} spectra. Use SpectrumList.read() instead.") + else: + raise RuntimeError(f"Input data has {len(spectrum_list)} spectra.") @data_loader( diff --git a/specutils/io/default_loaders/tests/test_jwst_reader.py b/specutils/io/default_loaders/tests/test_jwst_reader.py index b640d2625..040e3b82f 100644 --- a/specutils/io/default_loaders/tests/test_jwst_reader.py +++ b/specutils/io/default_loaders/tests/test_jwst_reader.py @@ -182,6 +182,38 @@ def test_jwst_wfss_multi_reader(tmp_path, spec_multi_new, format): assert data[2].flux.unit == u.MJy/u.sr +@pytest.mark.parametrize( + "spec_multi_new, format", + [("EXTRACT1D", "JWST x1d multi"), ("COMBINE1D", "JWST c1d multi")], + indirect=["spec_multi_new"], +) +def test_jwst_wfss_multi_lazy_load_with_labels(tmp_path, spec_multi_new, format): + """Test lazy loading and label-based access for JWST c1d/x1d packed multi data.""" + tmpfile = str(tmp_path / "jwst-lazy.fits") + spec_multi_new.writeto(tmpfile) + + data = SpectrumList.read(tmpfile, format=format, lazy_load=True) + assert isinstance(data, SpectrumList) + assert data.is_lazy + assert data.n_loaded == 0 + assert len(data) == 3 + + # Before string-key access, lazy labels are exposed as a list. + assert isinstance(data.labels, list) + assert "hdu1_source_1" in data.labels + assert "hdu1_source_2" in data.labels + assert "hdu1_source_3" in data.labels + + first = data["hdu1_source_1"] + assert isinstance(first, Spectrum) + assert data.n_loaded == 1 + + second = data[1] + assert isinstance(second, Spectrum) + assert data.n_loaded == 2 + assert len(data) == 3 + + @pytest.mark.parametrize('spec_single, format', [('EXTRACT1D', 'JWST x1d'), ('COMBINE1D', 'JWST c1d')], indirect=['spec_single'])