From 418019c123372e3ce5d8b6332363f970ab92aca1 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Tue, 14 Feb 2023 17:51:29 -0500
Subject: [PATCH 1/7] Add cube writer to properly write out cube generated by
 model fitting.

---
 CHANGES.rst                                   |  2 +
 docs/cubeviz/export_data.rst                  |  7 ++
 jdaviz/configs/cubeviz/helper.py              | 58 ++++++++++++++
 .../model_fitting/tests/test_fitting.py       | 79 +++++++++++++------
 4 files changed, 120 insertions(+), 26 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index b9143919fc..202fa790f4 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -7,6 +7,8 @@ New Features
 Cubeviz
 ^^^^^^^
 
+- Custom Spectrum1D writer for spectral cube generated by Cubeviz. [#2012]
+
 Imviz
 ^^^^^
 
diff --git a/docs/cubeviz/export_data.rst b/docs/cubeviz/export_data.rst
index d889266023..38e0424c67 100644
--- a/docs/cubeviz/export_data.rst
+++ b/docs/cubeviz/export_data.rst
@@ -94,6 +94,13 @@ Alternatively, you can wrap this all into a single command:
 
     mydata = cubeviz.app.get_data_from_viewer("uncert-viewer", "data_name")
 
+To write out a `specutils.Spectrum1D` cube from Cubeviz
+(e.g., a fitted cube from :ref:`model-fitting`):
+
+.. code-block:: python
+
+    mydata.write("mydata.fits", format="jdaviz-cube-writer")
+
 Data can also be accessed directly from ``data_collection`` using the following code:
 
 .. code-block:: python
diff --git a/jdaviz/configs/cubeviz/helper.py b/jdaviz/configs/cubeviz/helper.py
index 1e7502c22d..5caaba8ee6 100644
--- a/jdaviz/configs/cubeviz/helper.py
+++ b/jdaviz/configs/cubeviz/helper.py
@@ -1,6 +1,10 @@
 import numpy as np
+from astropy.io import fits
+from astropy.io import registry as io_registry
 from astropy.utils.decorators import deprecated
 from glue.core import BaseData
+from specutils import Spectrum1D
+from specutils.io.registers import _astropy_has_priorities
 
 from jdaviz.core.helpers import ImageConfigHelper
 from jdaviz.configs.default.plugins.line_lists.line_list_mixin import LineListMixin
@@ -126,3 +130,57 @@ class CubeViz(Cubeviz):
 
 def layer_is_cube_image_data(layer):
     return isinstance(layer, BaseData) and layer.ndim in (2, 3)
+
+
+# TODO: We can remove this when specutils supports it, i.e.,
+#       https://github.com/astropy/specutils/issues/592 and
+#       https://github.com/astropy/specutils/pull/1009
+# NOTE: Cannot use custom_write decorator from specutils because
+#       that involves asking user to manually put something in
+#       their ~/.specutils directory.
+
+def jdaviz_cube_fitswriter(spectrum, file_name, **kwargs):
+    """This is a custom writer for Spectrum1D data cube.
+    This writer is specifically targetting data cube
+    generated from Cubeviz plugins (e.g., cube fitting)
+    with FITS WCS. It writes out data in the following format
+    (with MASK only exist when applicable)::
+
+        No.    Name      Ver    Type
+          0  PRIMARY       1 PrimaryHDU
+          1  SCI           1 ImageHDU (float32)
+          2  MASK          1 ImageHDU (uint16)
+
+    Examples
+    --------
+    To write out a Spectrum1D cube using this writer::
+
+    >>> spec.write("my_output.fits", format="jdaviz-cube-writer", overwrite=True)
+
+    """
+    pri_hdu = fits.PrimaryHDU()
+
+    flux = spectrum.flux
+    sci_hdu = fits.ImageHDU(flux.value.astype(np.float32))
+    sci_hdu.name = "SCI"
+    sci_hdu.header.update(spectrum.meta)
+    sci_hdu.header.update(spectrum.wcs.to_header())
+    sci_hdu.header['BUNIT'] = flux.unit.to_string(format='fits')
+
+    hlist = [pri_hdu, sci_hdu]
+
+    if spectrum.mask is not None:
+        mask_hdu = fits.ImageHDU(spectrum.mask.astype(np.uint16))
+        mask_hdu.name = "MASK"
+        hlist.append(mask_hdu)
+
+    hdulist = fits.HDUList(hlist)
+    hdulist.writeto(file_name, **kwargs)
+
+
+if _astropy_has_priorities():
+    kwargs = {"priority": 0}
+else:
+    kwargs = {}
+io_registry.register_writer(
+    "jdaviz-cube-writer", Spectrum1D, jdaviz_cube_fitswriter, force=False, **kwargs)
diff --git a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
index f8736df915..2b4347fd6b 100644
--- a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
@@ -1,10 +1,13 @@
-import astropy.modeling.models as models
-import astropy.modeling.parameters as params
-from astropy.nddata import StdDevUncertainty
-import astropy.units as u
+import warnings
+
 import numpy as np
 import pytest
-
+from astropy import units as u
+from astropy.io import fits
+from astropy.modeling import models, parameters as params
+from astropy.nddata import StdDevUncertainty
+from astropy.wcs import WCS
+from numpy.testing import assert_allclose, assert_array_equal
 from specutils.spectra import Spectrum1D
 
 from jdaviz.configs.default.plugins.model_fitting import fitting_backend as fb
@@ -100,7 +103,7 @@ def test_fitting_backend(unc):
 
     parameters_expected = np.array([0.7, 4.65, 0.3, 2., 5.55, 0.3, -2.,
                                     8.15, 0.2, 1.])
-    assert np.allclose(fm.parameters, parameters_expected, atol=1e-5)
+    assert_allclose(fm.parameters, parameters_expected, atol=1e-5)
 
     # Returns the fitted model
     fm, fitted_spectrum = fb.fit_model_to_spectrum(spectrum, model_list, expression,
@@ -109,14 +112,11 @@ def test_fitting_backend(unc):
     parameters_expected = np.array([1.0104705, 4.58956282, 0.19590464, 2.39892026,
                                     5.49867754, 0.10834472, -1.66902953, 8.19714439,
                                     0.09535613, 3.99125545])
-    assert np.allclose(fm.parameters, parameters_expected, atol=1e-5)
+    assert_allclose(fm.parameters, parameters_expected, atol=1e-5)
 
 
-# When pytest turns warnings into errors, this silently fails with
-# len(fitted_parameters) == 0
-@pytest.mark.filterwarnings('ignore')
 @pytest.mark.parametrize('unc', ('zeros', None))
-def test_cube_fitting_backend(unc):
+def test_cube_fitting_backend(unc, tmp_path):
     np.random.seed(42)
 
     SIGMA = 0.1  # noise in data
@@ -140,7 +140,14 @@ def test_cube_fitting_backend(unc):
         flux_cube[:, spx[0], spx[1]] = build_spectrum(sigma=SIGMA)[1]
 
     # Transpose so it can be packed in a Spectrum1D instance.
-    flux_cube = flux_cube.transpose(1, 2, 0)
+    flux_cube = flux_cube.transpose(1, 2, 0)  # (15, 14, 200)
+    cube_wcs = WCS({
+        'WCSAXES': 3, 'RADESYS': 'ICRS', 'EQUINOX': 2000.0,
+        'CRPIX3': 38.0, 'CRPIX2': 38.0, 'CRPIX1': 1.0,
+        'CRVAL3': 205.4384, 'CRVAL2': 27.004754, 'CRVAL1': 0.0,
+        'CDELT3': 0.01, 'CDELT2': 0.01, 'CDELT1': 0.05,
+        'CUNIT3': 'deg', 'CUNIT2': 'deg', 'CUNIT1': 'um',
+        'CTYPE3': 'RA---TAN', 'CTYPE2': 'DEC--TAN', 'CTYPE1': 'WAVE'})
 
     # Mask part of the spectral axis to later ensure that it gets propagated through:
     mask = np.zeros_like(flux_cube).astype(bool)
@@ -152,7 +159,7 @@ def test_cube_fitting_backend(unc):
     elif unc is None:
         uncertainties = None
 
-    spectrum = Spectrum1D(flux=flux_cube*u.Jy, spectral_axis=x*u.um,
+    spectrum = Spectrum1D(flux=flux_cube*u.Jy, spectral_axis=x*u.um, wcs=cube_wcs,
                           uncertainty=uncertainties, mask=mask)
 
     # Initial model for fit.
@@ -168,8 +175,10 @@ def test_cube_fitting_backend(unc):
     # n_cpu = 1  # NOTE: UNCOMMENT TO DEBUG LOCALLY, AS NEEDED
 
     # Fit to all spaxels.
-    fitted_parameters, fitted_spectrum = fb.fit_model_to_spectrum(
-        spectrum, model_list, expression, n_cpu=n_cpu)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message=r"The fit may be unsuccessful.*")
+        fitted_parameters, fitted_spectrum = fb.fit_model_to_spectrum(
+            spectrum, model_list, expression, n_cpu=n_cpu)
 
     # Check that parameter results are formatted as expected.
     assert type(fitted_parameters) == list
@@ -198,19 +207,37 @@ def test_cube_fitting_backend(unc):
     # interested here in checking the correctness of the data
     # packaging into the output products.
 
-    assert np.allclose(fitted_model[0].amplitude.value, 1.09, atol=TOL)
-    assert np.allclose(fitted_model[1].amplitude.value, 2.4, atol=TOL)
-    assert np.allclose(fitted_model[2].amplitude.value, -1.7, atol=TOL)
+    assert_allclose(fitted_model[0].amplitude.value, 1.09, atol=TOL)
+    assert_allclose(fitted_model[1].amplitude.value, 2.4, atol=TOL)
+    assert_allclose(fitted_model[2].amplitude.value, -1.7, atol=TOL)
 
-    assert np.allclose(fitted_model[0].mean.value, 4.6, atol=TOL)
-    assert np.allclose(fitted_model[1].mean.value, 5.5, atol=TOL)
-    assert np.allclose(fitted_model[2].mean.value, 8.2, atol=TOL)
+    assert_allclose(fitted_model[0].mean.value, 4.6, atol=TOL)
+    assert_allclose(fitted_model[1].mean.value, 5.5, atol=TOL)
+    assert_allclose(fitted_model[2].mean.value, 8.2, atol=TOL)
 
-    assert np.allclose(fitted_model[0].stddev.value, 0.2, atol=TOL)
-    assert np.allclose(fitted_model[1].stddev.value, 0.1, atol=TOL)
-    assert np.allclose(fitted_model[2].stddev.value, 0.1, atol=TOL)
+    assert_allclose(fitted_model[0].stddev.value, 0.2, atol=TOL)
+    assert_allclose(fitted_model[1].stddev.value, 0.1, atol=TOL)
+    assert_allclose(fitted_model[2].stddev.value, 0.1, atol=TOL)
 
-    assert np.allclose(fitted_model[3].amplitude.value, 4.0, atol=TOL)
+    assert_allclose(fitted_model[3].amplitude.value, 4.0, atol=TOL)
 
     # Check that the fitted spectrum is masked correctly:
-    assert np.all(fitted_spectrum.mask == mask)
+    assert_array_equal(fitted_spectrum.mask, mask)
+
+    # Check I/O roundtrip.
+    out_fn = tmp_path / "fitted_cube.fits"
+    fitted_spectrum.write(out_fn, format="jdaviz-cube-writer", overwrite=True)
+    coo_expected = fitted_spectrum.wcs.pixel_to_world(1, 0, 2)
+    with fits.open(out_fn) as pf:
+        assert len(pf) == 3
+        assert pf[0].name == "PRIMARY"
+        assert pf[1].name == "SCI"
+        assert pf[1].header["BUNIT"] == fitted_spectrum.flux.unit.to_string(format="fits")
+        assert_allclose(pf[1].data, fitted_spectrum.flux.value)
+        assert pf[2].name == "MASK"
+        assert_array_equal(pf[2].data, mask)
+        w = WCS(pf[1].header)
+        coo = w.pixel_to_world(1, 0, 2)
+        assert_allclose(coo[0].value, coo_expected[0].value)  # SpectralCoord
+        assert_allclose([coo[1].ra.deg, coo[1].dec.deg],
+                        [coo_expected[1].ra.deg, coo_expected[1].dec.deg])

From ba2408d97660ee967b6f3731085ef92c2fdd9bb0 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Tue, 14 Feb 2023 17:58:48 -0500
Subject: [PATCH 2/7] Fix docstring formatting

---
 jdaviz/configs/cubeviz/helper.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jdaviz/configs/cubeviz/helper.py b/jdaviz/configs/cubeviz/helper.py
index 5caaba8ee6..a6e0d8154a 100644
--- a/jdaviz/configs/cubeviz/helper.py
+++ b/jdaviz/configs/cubeviz/helper.py
@@ -153,7 +153,7 @@ def jdaviz_cube_fitswriter(spectrum, file_name, **kwargs):
 
     Examples
     --------
-    To write out a Spectrum1D cube using this writer::
+    To write out a Spectrum1D cube using this writer:
 
     >>> spec.write("my_output.fits", format="jdaviz-cube-writer", overwrite=True)
 

From 6602a73317de8673ddc931a2893d7e9d38fa4364 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Tue, 14 Feb 2023 18:54:17 -0500
Subject: [PATCH 3/7] Fix WCS and loader roundtripping

---
 jdaviz/configs/cubeviz/helper.py                  |  2 +-
 jdaviz/configs/cubeviz/plugins/parsers.py         | 15 +++++++++++----
 .../plugins/model_fitting/model_fitting.py        |  9 ++++++---
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/jdaviz/configs/cubeviz/helper.py b/jdaviz/configs/cubeviz/helper.py
index a6e0d8154a..016dc7b59d 100644
--- a/jdaviz/configs/cubeviz/helper.py
+++ b/jdaviz/configs/cubeviz/helper.py
@@ -155,7 +155,7 @@ def jdaviz_cube_fitswriter(spectrum, file_name, **kwargs):
     --------
     To write out a Spectrum1D cube using this writer:
 
-    >>> spec.write("my_output.fits", format="jdaviz-cube-writer", overwrite=True)
+    >>> spec.write("my_output.fits", format="jdaviz-cube-writer", overwrite=True)  # doctest: +SKIP
 
     """
     pri_hdu = fits.PrimaryHDU()
diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
index fac74702ec..8cb455367a 100644
--- a/jdaviz/configs/cubeviz/plugins/parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -140,13 +140,20 @@ def _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, target_w
             category=UserWarning)
         sc = Spectrum1D(flux=flux, wcs=wcs)
 
-    # TODO: Can the unit be defined in a different keyword, e.g., CUNIT1?
     if target_wave_unit is None and hdulist is not None:
-        cunit_key = 'CUNIT3'
+        found_target = False
         for ext in ('SCI', 'FLUX', 'PRIMARY'):  # In priority order
-            if ext in hdulist and cunit_key in hdulist[ext].header:
-                target_wave_unit = u.Unit(hdulist[ext].header[cunit_key])
+            if found_target:
                 break
+            if ext not in hdulist:
+                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]:
+                    target_wave_unit = u.Unit(hdr[cunit_key])
+                    found_target = True
+                    break
 
     if (data_type == 'flux' and target_wave_unit is not None
             and target_wave_unit != sc.spectral_axis.unit):
diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
index bc04566ad5..e937d0e4ce 100644
--- a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
@@ -782,7 +782,10 @@ def _fit_model_to_cube(self, add_data):
             return
 
         # Get the primary data component
-        spec = data.get_object(cls=Spectrum1D, statistic=None)
+        if "_orig_spec" in data.meta:
+            spec = data.meta["_orig_spec"]
+        else:
+            spec = data.get_object(cls=Spectrum1D, statistic=None)
 
         snackbar_message = SnackbarMessage(
             "Fitting model to cube...",
@@ -825,8 +828,8 @@ def _fit_model_to_cube(self, add_data):
 
         # Create new glue data object
         output_cube = Data(label=label,
-                           coords=data.coords)
-        output_cube['flux'] = fitted_spectrum.flux.value
+                           coords=fitted_spectrum.wcs,
+                           flux=fitted_spectrum.flux.value)
         output_cube.get_component('flux').units = fitted_spectrum.flux.unit.to_string()
 
         if add_data:

From a2e083acb8111fbcb93eba541821e7ee2d59ef82 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Wed, 15 Feb 2023 10:14:50 -0500
Subject: [PATCH 4/7] Clarify what mask values mean.

[ci skip]
---
 docs/cubeviz/export_data.rst     | 4 +++-
 jdaviz/configs/cubeviz/helper.py | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/docs/cubeviz/export_data.rst b/docs/cubeviz/export_data.rst
index 38e0424c67..230f64944b 100644
--- a/docs/cubeviz/export_data.rst
+++ b/docs/cubeviz/export_data.rst
@@ -95,7 +95,9 @@ Alternatively, you can wrap this all into a single command:
     mydata = cubeviz.app.get_data_from_viewer("uncert-viewer", "data_name")
 
 To write out a `specutils.Spectrum1D` cube from Cubeviz
-(e.g., a fitted cube from :ref:`model-fitting`):
+(e.g., a fitted cube from :ref:`model-fitting`),
+where the mask (if available) is as defined in
+`Spectrum1D masks <https://specutils.readthedocs.io/en/latest/spectrum1d.html#including-masks>`_:
 
 .. code-block:: python
 
diff --git a/jdaviz/configs/cubeviz/helper.py b/jdaviz/configs/cubeviz/helper.py
index 016dc7b59d..fb8fbf4e36 100644
--- a/jdaviz/configs/cubeviz/helper.py
+++ b/jdaviz/configs/cubeviz/helper.py
@@ -169,6 +169,9 @@ def jdaviz_cube_fitswriter(spectrum, file_name, **kwargs):
 
     hlist = [pri_hdu, sci_hdu]
 
+    # https://specutils.readthedocs.io/en/latest/spectrum1d.html#including-masks
+    # Good: False or 0
+    # Bad: True or non-zero
     if spectrum.mask is not None:
         mask_hdu = fits.ImageHDU(spectrum.mask.astype(np.uint16))
         mask_hdu.name = "MASK"

From f8a0bd827b3af0c843cc27dff8e35865a711a578 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Wed, 15 Feb 2023 13:48:37 -0500
Subject: [PATCH 5/7] Fix bug in cunit lookup

---
 jdaviz/configs/cubeviz/plugins/parsers.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
index 8cb455367a..331d062459 100644
--- 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

From ad0d3e0e146001f6496a01cf5a730ccc79e41510 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Wed, 22 Feb 2023 14:30:45 -0500
Subject: [PATCH 6/7] Address review comments from kecnry

---
 docs/cubeviz/export_data.rst                  |  2 +-
 jdaviz/configs/cubeviz/helper.py              |  9 ++++--
 jdaviz/configs/cubeviz/plugins/parsers.py     |  2 +-
 .../model_fitting/tests/test_fitting.py       | 28 +++++++++++++++++--
 4 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/docs/cubeviz/export_data.rst b/docs/cubeviz/export_data.rst
index 230f64944b..0b7010d447 100644
--- a/docs/cubeviz/export_data.rst
+++ b/docs/cubeviz/export_data.rst
@@ -101,7 +101,7 @@ where the mask (if available) is as defined in
 
 .. code-block:: python
 
-    mydata.write("mydata.fits", format="jdaviz-cube-writer")
+    mydata.write("mydata.fits", format="jdaviz-cube")
 
 Data can also be accessed directly from ``data_collection`` using the following code:
 
diff --git a/jdaviz/configs/cubeviz/helper.py b/jdaviz/configs/cubeviz/helper.py
index fb8fbf4e36..973bbc4d33 100644
--- a/jdaviz/configs/cubeviz/helper.py
+++ b/jdaviz/configs/cubeviz/helper.py
@@ -151,11 +151,14 @@ def jdaviz_cube_fitswriter(spectrum, file_name, **kwargs):
           1  SCI           1 ImageHDU (float32)
           2  MASK          1 ImageHDU (uint16)
 
+    The FITS file generated by this writer does not need a
+    custom reader to be read back into Spectrum1D.
+
     Examples
     --------
     To write out a Spectrum1D cube using this writer:
 
-    >>> spec.write("my_output.fits", format="jdaviz-cube-writer", overwrite=True)  # doctest: +SKIP
+    >>> spec.write("my_output.fits", format="jdaviz-cube", overwrite=True)  # doctest: +SKIP
 
     """
     pri_hdu = fits.PrimaryHDU()
@@ -183,7 +186,7 @@ def jdaviz_cube_fitswriter(spectrum, file_name, **kwargs):
 
 if _astropy_has_priorities():
     kwargs = {"priority": 0}
-else:
+else:  # pragma: no cover
     kwargs = {}
 io_registry.register_writer(
-    "jdaviz-cube-writer", Spectrum1D, jdaviz_cube_fitswriter, force=False, **kwargs)
+    "jdaviz-cube", Spectrum1D, jdaviz_cube_fitswriter, force=False, **kwargs)
diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
index 331d062459..4f8b02eb60 100644
--- a/jdaviz/configs/cubeviz/plugins/parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -152,7 +152,7 @@ def _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, target_w
             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]:
+                if cunit_key in hdr and 'WAVE' in hdr[ctype_key]:
                     target_wave_unit = u.Unit(hdr[cunit_key])
                     found_target = True
                     break
diff --git a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
index 2b4347fd6b..4b67e82534 100644
--- a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
@@ -4,6 +4,7 @@
 import pytest
 from astropy import units as u
 from astropy.io import fits
+from astropy.io.registry.base import IORegistryError
 from astropy.modeling import models, parameters as params
 from astropy.nddata import StdDevUncertainty
 from astropy.wcs import WCS
@@ -116,7 +117,7 @@ def test_fitting_backend(unc):
 
 
 @pytest.mark.parametrize('unc', ('zeros', None))
-def test_cube_fitting_backend(unc, tmp_path):
+def test_cube_fitting_backend(cubeviz_helper, unc, tmp_path):
     np.random.seed(42)
 
     SIGMA = 0.1  # noise in data
@@ -226,13 +227,14 @@ def test_cube_fitting_backend(unc, tmp_path):
 
     # Check I/O roundtrip.
     out_fn = tmp_path / "fitted_cube.fits"
-    fitted_spectrum.write(out_fn, format="jdaviz-cube-writer", overwrite=True)
+    fitted_spectrum.write(out_fn, format="jdaviz-cube", overwrite=True)
+    flux_unit_str = fitted_spectrum.flux.unit.to_string(format="fits")
     coo_expected = fitted_spectrum.wcs.pixel_to_world(1, 0, 2)
     with fits.open(out_fn) as pf:
         assert len(pf) == 3
         assert pf[0].name == "PRIMARY"
         assert pf[1].name == "SCI"
-        assert pf[1].header["BUNIT"] == fitted_spectrum.flux.unit.to_string(format="fits")
+        assert pf[1].header["BUNIT"] == flux_unit_str
         assert_allclose(pf[1].data, fitted_spectrum.flux.value)
         assert pf[2].name == "MASK"
         assert_array_equal(pf[2].data, mask)
@@ -241,3 +243,23 @@ def test_cube_fitting_backend(unc, tmp_path):
         assert_allclose(coo[0].value, coo_expected[0].value)  # SpectralCoord
         assert_allclose([coo[1].ra.deg, coo[1].dec.deg],
                         [coo_expected[1].ra.deg, coo_expected[1].dec.deg])
+
+    # Our custom format is not registered to readers, just writers.
+    # You can read it back in without custom read. See "Cubeviz roundtrip" below.
+    with pytest.raises(IORegistryError, match="No reader defined"):
+        Spectrum1D.read(out_fn, format="jdaviz-cube")
+
+    # Check Cubeviz roundtrip.
+    cubeviz_helper.load_data(out_fn)
+    assert len(cubeviz_helper.app.data_collection) == 2
+    data_sci = cubeviz_helper.app.data_collection["fitted_cube.fits[SCI]"]
+    flux_sci = data_sci.get_component("flux")
+    assert_allclose(flux_sci.data, fitted_spectrum.flux.value)
+    assert flux_sci.units == flux_unit_str
+    coo = data_sci.coords.pixel_to_world(1, 0, 2)
+    assert_allclose(coo[0].value, coo_expected[0].value)  # SpectralCoord
+    assert_allclose([coo[1].ra.deg, coo[1].dec.deg],
+                    [coo_expected[1].ra.deg, coo_expected[1].dec.deg])
+    data_mask = cubeviz_helper.app.data_collection["fitted_cube.fits[MASK]"]
+    flux_mask = data_mask.get_component("flux")
+    assert_array_equal(flux_mask.data, mask)

From e05119dc022ce12c4c06a30e7cd5e0f43fcc7de6 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Mon, 27 Feb 2023 11:23:25 -0500
Subject: [PATCH 7/7] Remove redundant spectral_axis in test case

---
 .../configs/default/plugins/model_fitting/tests/test_fitting.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
index 4b67e82534..9b06640887 100644
--- a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
@@ -160,7 +160,7 @@ def test_cube_fitting_backend(cubeviz_helper, unc, tmp_path):
     elif unc is None:
         uncertainties = None
 
-    spectrum = Spectrum1D(flux=flux_cube*u.Jy, spectral_axis=x*u.um, wcs=cube_wcs,
+    spectrum = Spectrum1D(flux=flux_cube*u.Jy, wcs=cube_wcs,
                           uncertainty=uncertainties, mask=mask)
 
     # Initial model for fit.