Skip to content

Commit

Permalink
feat: visit specification on mwmVisit load
Browse files Browse the repository at this point in the history 
- readded print -> warnings conversion
- can specify the visit to load on mwmVisit load.
- added relevant tests for the new mwmVisit case
  • Loading branch information
@rileythai
rileythai committed Oct 18, 2024
1 parent 122b368 commit 0c5fe8f
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 62 deletions.
121 changes: 86 additions & 35 deletions specutils/io/default_loaders/sdss_v.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
"""Register reader functions for various spectral formats."""
import warnings
from typing import Optional

import numpy as np
from astropy.units import Unit, Quantity, Angstrom
from astropy.nddata import StdDevUncertainty, InverseVariance
from astropy.io.fits import HDUList, BinTableHDU, ImageHDU
from astropy.utils.exceptions import AstropyUserWarning

from ...spectra import Spectrum1D, SpectrumList
from ..registers import data_loader
Expand Down Expand Up @@ -313,8 +315,6 @@ def load_sdss_apVisit_list(file_obj, **kwargs):


# BOSS REDUX products (specLite, specFull, custom coadd files, etc)


@data_loader(
"SDSS-V spec",
identifier=spec_sdss5_identify,
Expand All @@ -340,7 +340,8 @@ def load_sdss_spec_1D(file_obj, *args, hdu: Optional[int] = None, **kwargs):
"""
if hdu is None:
# TODO: how should we handle this -- multiple things in file, but the user cannot choose.
print('HDU not specified. Loading coadd spectrum (HDU1)')
warnings.warn('HDU not specified. Loading coadd spectrum (HDU1)',
AstropyUserWarning)
hdu = 1 # defaulting to coadd
# raise ValueError("HDU not specified! Please specify a HDU to load.")
elif hdu in [2, 3, 4]:
Expand Down Expand Up @@ -438,16 +439,21 @@ def _load_BOSS_HDU(hdulist: HDUList, hdu: int, **kwargs):
priority=20,
extensions=["fits"],
)
def load_sdss_mwm_1d(file_obj, hdu: Optional[int] = None, **kwargs):
def load_sdss_mwm_1d(file_obj,
hdu: Optional[int] = None,
visit: Optional[int] = None,
**kwargs):
"""
Load an unspecified spec file as a Spectrum1D.
Parameters
----------
file_obj : str, file-like, or HDUList
FITS file name, file object, or HDUList..
hdu : int
hdu : Optional[int]
Specified HDU to load.
visit : Optional[int]
Specified visit index to load.
Returns
-------
Expand All @@ -467,11 +473,22 @@ def load_sdss_mwm_1d(file_obj, hdu: Optional[int] = None, **kwargs):
for i in range(len(hdulist)):
if hdulist[i].header.get("DATASUM") != "0":
hdu = i
print('HDU not specified. Loading spectrum at (HDU{})'.
format(i))
warnings.warn(
'HDU not specified. Loading spectrum at (HDU{})'.
format(i), AstropyUserWarning)
break

return _load_mwmVisit_or_mwmStar_hdu(hdulist, hdu, **kwargs)
# load spectra
spectra_list = _load_mwmVisit_or_mwmStar_hdu(hdulist, hdu)

# if mwmVisit AND visit not specified, load first visit in the HDU
if visit is None:
if spectra_list[0].meta['datatype'].lower() == 'mwmvisit':
warnings.warn(
'Visit to load not specified. Loading spectrum at index 0. Specify with (visit=...)',
AstropyUserWarning)
visit = 0
return spectra_list[visit]


@data_loader(
Expand All @@ -493,12 +510,8 @@ def load_sdss_mwm_list(file_obj, **kwargs):
Returns
-------
SpectrumList
The spectra contained in the file, where:
Spectrum1D
A given spectra of nD flux
None
If there are no spectra for that spectrograph/observatory
SpectrumList[Spectrum1D]
A list spectra from each visit with each instrument at each observatory (mwmVisit), or the coadd from each instrument/observatory (mwmStar).
"""
spectra = SpectrumList()
with read_fileobj_or_hdulist(file_obj, memmap=False, **kwargs) as hdulist:
Expand All @@ -513,12 +526,12 @@ def load_sdss_mwm_list(file_obj, **kwargs):
for hdu in range(1, len(hdulist)):
if hdulist[hdu].header.get("DATASUM") == "0":
# Skip zero data HDU's
# TODO: validate if we want this printed warning or not.
# it might get annoying & fill logs with useless alerts.
print("WARNING: HDU{} ({}) is empty.".format(
hdu, hdulist[hdu].name))
warnings.warn(
"WARNING: HDU{} ({}) is empty.".format(
hdu, hdulist[hdu].name), AstropyUserWarning)
continue
spectra.append(_load_mwmVisit_or_mwmStar_hdu(hdulist, hdu))
spectra_list = _load_mwmVisit_or_mwmStar_hdu(hdulist, hdu)
[spectra.append(spectra_list[i]) for i in range(len(spectra_list))]
return spectra


Expand All @@ -535,8 +548,8 @@ def _load_mwmVisit_or_mwmStar_hdu(hdulist: HDUList, hdu: int, **kwargs):
Returns
-------
Spectrum1D
The spectrum with nD flux contained in the HDU.
list[Spectrum1D]
List of spectrum with 1D flux contained in the HDU.
"""
if hdulist[hdu].header.get("DATASUM") == "0":
Expand Down Expand Up @@ -571,12 +584,7 @@ def _load_mwmVisit_or_mwmStar_hdu(hdulist: HDUList, hdu: int, **kwargs):
mask = mask != 0

# collapse shape if 1D spectra in 2D array
# NOTE: this fixes a jdaviz handling bug for 2D of shape 1,
# it could be that it's expected to be parsed this way.
if flux.shape[0] == 1:
flux = flux[0]
e_flux = e_flux[0]
mask = mask[0]

# Create metadata
meta = dict()
Expand All @@ -597,18 +605,61 @@ def _load_mwmVisit_or_mwmStar_hdu(hdulist: HDUList, hdu: int, **kwargs):
try:
meta['mjd'] = hdulist[hdu].data['mjd']
meta["datatype"] = "mwmVisit"

except KeyError:
meta["mjd"] = (str(hdulist[hdu].data["min_mjd"][0]) + "->" +
str(hdulist[hdu].data["max_mjd"][0]))
meta["min_mjd"] = str(hdulist[hdu].data["min_mjd"][0])
meta["max_mjd"] = str(hdulist[hdu].data["max_mjd"][0])
meta["datatype"] = "mwmStar"
finally:
meta["name"] = hdulist[hdu].name
meta["sdss_id"] = hdulist[hdu].data['sdss_id']

return Spectrum1D(
spectral_axis=spectral_axis,
flux=flux,
uncertainty=e_flux,
mask=mask,
meta=meta,
)
# drop back a list of Spectrum1Ds to unpack
metadicts = _split_mwm_meta_dict(meta)
return [
Spectrum1D(
spectral_axis=spectral_axis,
flux=flux[i],
uncertainty=e_flux[i],
mask=mask[i],
meta=metadicts[i],
) for i in range(flux.shape[0])
]


def _split_mwm_meta_dict(d):
"""
Metadata sub-loader subfunction for MWM files.
Parameters
----------
d: dict
Initial metadata dictionary.
Returns
-------
dicts: list[dict]
List of dicts with unpacked metadata for length > 1 array objects.
"""
# find the length of entries
N = max(len(v) if isinstance(v, np.ndarray) else 1 for v in d.values())

# create N dictionaries to hold the split results
dicts = [{} for _ in range(N)]

for key, value in d.items():
if isinstance(value, np.ndarray):
# Ensure that the length of the list matches N
if len(value) != N:
# an error case we ignore
continue
# distribute each element to the corresponding metadict
for i in range(N):
dicts[i][key] = value[i]
else:
# if it's a single object, copy it to each metadict
for i in range(N):
dicts[i][key] = value

return dicts
65 changes: 38 additions & 27 deletions specutils/io/default_loaders/tests/test_sdss_v.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,16 @@
from specutils import Spectrum1D, SpectrumList


def generate_apogee_hdu(observatory="APO", with_wl=True, datasum="0"):
def generate_apogee_hdu(observatory="APO",
with_wl=True,
datasum="0",
nvisits=1):
wl = (10**(4.179 + 6e-6 * np.arange(8575))).reshape((1, -1))
flux = np.zeros_like(wl)
ivar = np.zeros_like(wl)
pixel_flags = np.zeros_like(wl)
continuum = np.zeros_like(wl)
nmf_rectified_model_flux = np.zeros_like(wl)
flux = np.array([np.zeros_like(wl)] * nvisits)
ivar = np.array([np.zeros_like(wl)] * nvisits)
pixel_flags = np.array([np.zeros_like(wl)] * nvisits)
continuum = np.array([np.zeros_like(wl)] * nvisits)
nmf_rectified_model_flux = np.array([np.zeros_like(wl)] * nvisits)

columns = [
fits.Column(name="spectrum_pk_id", array=[159783564], format="K"),
Expand Down Expand Up @@ -79,10 +82,13 @@ def generate_apogee_hdu(observatory="APO", with_wl=True, datasum="0"):
return fits.BinTableHDU.from_columns(columns, header=header)


def generate_boss_hdu(observatory="APO", with_wl=True, datasum="0"):
def generate_boss_hdu(observatory="APO", with_wl=True, datasum="0", nvisits=1):
wl = (10**(3.5523 + 1e-4 * np.arange(4648))).reshape((1, -1))
flux = ivar = continuum = pixel_flags = nmf_rectified_model_flux = np.zeros_like(
wl)
flux = np.array([np.zeros_like(wl)] * nvisits)
ivar = np.array([np.zeros_like(wl)] * nvisits)
pixel_flags = np.array([np.zeros_like(wl)] * nvisits)
continuum = np.array([np.zeros_like(wl)] * nvisits)
nmf_rectified_model_flux = np.array([np.zeros_like(wl)] * nvisits)
columns = [
fits.Column(name="spectrum_pk_id", array=[0], format="K"),
fits.Column(name="release", array=["sdss5"], format="5A"),
Expand Down Expand Up @@ -285,20 +291,22 @@ def fake_primary_hdu():
]))


def mwm_HDUList(hduflags, with_wl):
def mwm_HDUList(hduflags, with_wl, **kwargs):
hdulist = [fake_primary_hdu()]
for i, flag in enumerate(hduflags):
obs = ["APO", "LCO"]
if i <= 1:
hdulist.append(
generate_boss_hdu(obs[i % 2],
with_wl=with_wl,
datasum=str(flag)))
datasum=str(flag),
**kwargs))
else:
hdulist.append(
generate_apogee_hdu(obs[i % 2],
with_wl=with_wl,
datasum=str(flag)))
datasum=str(flag),
**kwargs))

print(hdulist)
return fits.HDUList(hdulist)
Expand Down Expand Up @@ -426,23 +434,25 @@ def spec_HDUList(n_spectra):

# TEST MWM loaders
@pytest.mark.parametrize(
"file_obj, hdu, with_wl, hduflags",
"file_obj, hdu, visit, with_wl, hduflags, nvisits",
[
("mwm-temp", None, False, [0, 0, 1, 0]),
("mwm-temp", 3, False, [0, 0, 1, 0]),
("mwm-temp", None, True, [0, 1, 1, 0]),
("mwm-temp", 2, True, [0, 1, 1, 0]),
("mwm-temp", None, None, False, [0, 0, 1, 0], 1), # visit
("mwm-temp", 3, None, False, [0, 0, 1, 0], 1),
("mwm-temp", None, 2, False, [0, 1, 1, 0], 3), # multi-ext visits
("mwm-temp", 2, 2, False, [0, 1, 1, 0], 3),
("mwm-temp", None, None, True, [0, 0, 1, 0], 1), # star
("mwm-temp", 3, None, True, [0, 0, 1, 0], 1),
("mwm-temp", None, None, True, [0, 1, 1, 0], 1),
("mwm-temp", 2, None, True, [0, 1, 1, 0], 1),
],
)
def test_mwm_1d(file_obj, hdu, with_wl, hduflags):
def test_mwm_1d(file_obj, hdu, visit, with_wl, hduflags, nvisits):
"""Test mwm Spectrum1D loader"""
tmpfile = str(file_obj) + ".fits"
mwm_HDUList(hduflags, with_wl).writeto(tmpfile, overwrite=True)
mwm_HDUList(hduflags, with_wl, nvisits=nvisits).writeto(tmpfile,
overwrite=True)

if hdu is None:
data = Spectrum1D.read(tmpfile)
else:
data = Spectrum1D.read(tmpfile, hdu=hdu)
data = Spectrum1D.read(tmpfile, hdu=hdu, visit=visit)
assert isinstance(data, Spectrum1D)
assert isinstance(data.meta["header"], fits.Header)
if data.meta["instrument"].lower() == "apogee":
Expand All @@ -463,19 +473,20 @@ def test_mwm_1d(file_obj, hdu, with_wl, hduflags):
"file_obj, with_wl, hduflags",
[
("mwm-temp", False, [0, 0, 1, 1]),
("mwm-temp", True, [0, 0, 1, 1]),
("mwm-temp", False, [0, 1, 1, 0]),
("mwm-temp", True, [0, 1, 1, 0]),
("mwm-temp", False, [1, 1, 0, 0]),
("mwm-temp", True, [1, 1, 0, 0]),
("mwm-temp", False, [1, 1, 1, 1]),
("mwm-temp", True, [0, 0, 1, 1]),
("mwm-temp", True, [0, 1, 1, 0]),
("mwm-temp", True, [1, 1, 0, 0]),
("mwm-temp", True, [1, 1, 1, 1]),
],
)
def test_mwm_list(file_obj, with_wl, hduflags):
"""Test mwm SpectrumList loader"""
tmpfile = str(file_obj) + ".fits"
mwm_HDUList(hduflags, with_wl).writeto(tmpfile, overwrite=True)
mwm_HDUList(hduflags, with_wl,
nvisits=1 if with_wl else 3).writeto(tmpfile, overwrite=True)

data = SpectrumList.read(tmpfile)
assert isinstance(data, SpectrumList)
Expand Down

0 comments on commit 0c5fe8f

Please sign in to comment.