work on fixing test for added spectra

desihub · Nov 20, 2023 · 0bd8f63 · 0bd8f63
1 parent 30f0569
commit 0bd8f63
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/py/desispec/spectra.py b/py/desispec/spectra.py
@@ -19,11 +19,11 @@
 
 import numpy as np
 from astropy.table import Table
+from astropy.units import Unit
 
 _specutils_imported = True
 try:
     from specutils import SpectrumList, Spectrum1D
-    from astropy.units import Unit
     from astropy.nddata import InverseVariance, StdDevUncertainty
 except ImportError:
     _specutils_imported = False
@@ -79,6 +79,9 @@ class Spectra(object):
         optional table of metadata, rowmatched to fibermap,
         e.g. a redshift catalog for these spectra
     """
+    wavelength_unit = Unit('Angstrom')
+    flux_density_unit = Unit('10-17 erg cm-2 s-1 AA-1')
+
     def __init__(self, bands=[], wave={}, flux={}, ivar={}, mask=None,
             resolution_data=None, fibermap=None, exp_fibermap=None,
             meta=None, extra=None,
@@ -682,13 +685,11 @@ def to_specutils(self):
         if not _specutils_imported:
             raise NameError("specutils is not available in the environment.")
         sl = SpectrumList()
-        AA = Unit('Angstrom')
-        specunit = Unit('10-17 erg cm-2 s-1 AA-1')
         for i, band in enumerate(self.bands):
             meta = {'band': band}
-            spectral_axis = self.wave[band] * AA
-            flux = self.flux[band] * specunit
-            uncertainty = InverseVariance(self.ivar[band] * (specunit**-2))
+            spectral_axis = self.wave[band] * self.wavelength_unit
+            flux = self.flux[band] * self.flux_density_unit
+            uncertainty = InverseVariance(self.ivar[band] * (self.flux_density_unit**-2))
             mask = self.mask[band] != 0
             meta['int_mask'] = self.mask[band]
             meta['resolution_data'] = self.resolution_data[band]
@@ -774,16 +775,14 @@ def from_specutils(cls, spectra):
         mask = dict()
         resolution_data = None
         extra = None
-        AA = Unit('Angstrom')
-        specunit = Unit('10-17 erg cm-2 s-1 AA-1')
         for i, band in enumerate(bands):
-            wave[band] = sl[i].spectral_axis.to(AA).value
-            flux[band] = sl[i].flux.to(specunit).value
+            wave[band] = sl[i].spectral_axis.to(cls.wavelength_unit).value
+            flux[band] = sl[i].flux.to(cls.flux_density_unit).value
             if isinstance(sl[i].uncertainty, InverseVariance):
-                ivar[band] = sl[i].uncertainty.quantity.to(specunit**-2).value
+                ivar[band] = sl[i].uncertainty.quantity.to(cls.flux_density_unit**-2).value
             elif isinstance(sl[i].uncertainty, StdDevUncertainty):
                 # Future: may need np.isfinite() here?
-                ivar[band] = (sl[i].uncertainty.quantity.to(specunit).value)**-2
+                ivar[band] = (sl[i].uncertainty.quantity.to(cls.flux_density_unit).value)**-2
             else:
                 raise ValueError("Unknown uncertainty type!")
             try:

diff --git a/py/desispec/test/test_spectra.py b/py/desispec/test/test_spectra.py
@@ -225,7 +225,7 @@ def test_read_targetids(self):
         self.assertTrue(np.all(spec_subset.mask['z'] == comp_subset.mask['z']))
 
         # read subset in different order than original file, with repeats and missing targetids
-        spec.fibermap['TARGETID'] = (np.arange(self.nspec) // 2) * 2 # [0,0,2,2,4] for nspec=5
+        spec.fibermap['TARGETID'] = (np.arange(self.nspec) // 2) * 2 # [0, 0, 2, 2, 4, 4] for nspec=6
         write_spectra(self.fileio, spec)
         targetids = [2,10,4,4,4,0,0]
         comp_subset = read_spectra(self.fileio, targetids=targetids)