Fix bug in cunit lookup

spacetelescope · Feb 27, 2023 · f8a0bd8 · f8a0bd8
1 parent a2e083a
commit f8a0bd8
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -149,8 +149,10 @@ def _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, target_w
                 continue
             hdr = hdulist[ext].header
             # The WCS could be swapped or unswapped.
-            for cunit_key in ('CUNIT3', 'CUNIT1'):
-                if not found_target and cunit_key in hdr and 'WAVE' in hdr[cunit_key]:
+            for cunit_num in (3, 1):
+                cunit_key = f"CUNIT{cunit_num}"
+                ctype_key = f"CTYPE{cunit_num}"
+                if not found_target and cunit_key in hdr and 'WAVE' in hdr[ctype_key]:
                     target_wave_unit = u.Unit(hdr[cunit_key])
                     found_target = True
                     break