fix read_spectra skipping exp_fibermap

desihub · Nov 3, 2023 · 7a05933 · 7a05933
1 parent d34aa51
commit 7a05933
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 31 deletions.
diff --git a/py/desispec/io/spectra.py b/py/desispec/io/spectra.py
@@ -310,36 +310,40 @@ def read_spectra(
     for h in range(1, nhdu):
         name = hdus[h].read_header()["EXTNAME"]
         log.debug('Reading %s', name)
-        if name == "FIBERMAP" and name not in skip_hdus:
-            fmap = encode_table(
-                Table(
-                    hdus[h].read(rows=rows, columns=select_columns["FIBERMAP"]),
-                    copy=True,
-                ).as_array()
-            )
-        elif name == "EXP_FIBERMAP" and name not in skip_hdus:
-            expfmap = encode_table(
-                Table(
-                    hdus[h].read(rows=exp_rows, columns=select_columns["EXP_FIBERMAP"]),
-                    copy=True,
-                ).as_array()
-            )
-        elif name == "SCORES" and name not in skip_hdus:
-            scores = encode_table(
-                Table(
-                    hdus[h].read(rows=rows, columns=select_columns["SCORES"]),
-                    copy=True,
-                ).as_array()
-            )
-        elif name == "EXTRA_CATALOG" and name not in skip_hdus:
-            extra_catalog = encode_table(
-                Table(
-                    hdus[h].read(
-                        rows=rows, columns=select_columns["EXTRA_CATALOG"]
-                    ),
-                    copy=True,
-                ).as_array()
-            )
+        if name == "FIBERMAP":
+            if name not in skip_hdus:
+                fmap = encode_table(
+                    Table(
+                        hdus[h].read(rows=rows, columns=select_columns["FIBERMAP"]),
+                        copy=True,
+                    ).as_array()
+                )
+        elif name == "EXP_FIBERMAP":
+            if name not in skip_hdus:
+                expfmap = encode_table(
+                    Table(
+                        hdus[h].read(rows=exp_rows, columns=select_columns["EXP_FIBERMAP"]),
+                        copy=True,
+                    ).as_array()
+                )
+        elif name == "SCORES":
+            if name not in skip_hdus:
+                scores = encode_table(
+                    Table(
+                        hdus[h].read(rows=rows, columns=select_columns["SCORES"]),
+                        copy=True,
+                    ).as_array()
+                )
+        elif name == "EXTRA_CATALOG":
+            if name not in skip_hdus:
+                extra_catalog = encode_table(
+                    Table(
+                        hdus[h].read(
+                            rows=rows, columns=select_columns["EXTRA_CATALOG"]
+                        ),
+                        copy=True,
+                    ).as_array()
+                )
         else:
             # Find the band based on the name
             mat = re.match(r"(.*)_(.*)", name)
@@ -374,6 +378,7 @@ def read_spectra(
                 res[band] = _read_image(hdus, h, ftype, rows=rows)
             elif type != "MASK" and type != "RESOLUTION" and type not in skip_hdus:
                 # this must be an "extra" HDU
+                log.debug('Reading extra HDU %s', name)
                 if extra is None:
                     extra = {}
                 if band not in extra:

diff --git a/py/desispec/test/test_spectra.py b/py/desispec/test/test_spectra.py
@@ -281,7 +281,7 @@ def test_read_skip_hdus(self):
         # manually create the spectra and write
         spec = Spectra(bands=self.bands, wave=self.wave, flux=self.flux,
             ivar=self.ivar, mask=self.mask, resolution_data=self.res,
-            fibermap=self.fmap1, meta=self.meta)
+            fibermap=self.fmap1, meta=self.meta, exp_fibermap=self.fmap1)
 
         write_spectra(self.fileio, spec)
 
@@ -290,6 +290,11 @@ def test_read_skip_hdus(self):
         self.assertIsNone(test.R)
         self.assertIsNotNone(test.fibermap) #- fibermap not skipped
 
+        test = read_spectra(self.fileio, skip_hdus=('EXP_FIBERMAP', 'SCORES', 'RESOLUTION'))
+        self.assertIsNone(test.exp_fibermap)
+        self.assertIsNone(test.scores)
+        self.assertIsNone(test.R)
+        self.assertIsNotNone(test.fibermap) #- fibermap not skipped
 
     def test_empty(self):