From fdd00a52a90c3ba079de3d3b9a022e66fc57a8b5 Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Tue, 13 Feb 2024 11:34:45 -0500
Subject: [PATCH 1/2] update y axis label based on unit type

Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com>

Update jdaviz/configs/specviz/plugins/viewers.py

Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com>
---
 CHANGES.rst                                  |  4 +++
 jdaviz/configs/specviz/plugins/viewers.py    | 18 +++++++++++++-
 jdaviz/configs/specviz/tests/test_viewers.py | 26 ++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 jdaviz/configs/specviz/tests/test_viewers.py

diff --git a/CHANGES.rst b/CHANGES.rst
index e4b968fa46..304ef8a0a8 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -99,6 +99,10 @@ Mosviz
 Specviz
 ^^^^^^^
 
+- Check unit type (e.g., flux density, surface brightness, counts, etc) for generating
+  display label for the y axis in spectral viewer. Previously it was hard coded
+  to always display ``flux density`` no matter the input unit. [#2703]
+
 Specviz2d
 ^^^^^^^^^
 
diff --git a/jdaviz/configs/specviz/plugins/viewers.py b/jdaviz/configs/specviz/plugins/viewers.py
index 7486551be1..e282b19dd9 100644
--- a/jdaviz/configs/specviz/plugins/viewers.py
+++ b/jdaviz/configs/specviz/plugins/viewers.py
@@ -545,8 +545,24 @@ def _plot_uncertainties(self):
                 self.figure.marks = list(self.figure.marks) + [error_line_mark]
 
     def set_plot_axes(self):
-        # Set axes labels for the spectrum viewer
+        # Set y axes labels for the spectrum viewer
+        y_display_unit = self.state.y_display_unit
+        y_unit = u.Unit(y_display_unit) if y_display_unit else u.dimensionless_unscaled
+
+        if y_unit.is_equivalent(u.Jy / u.sr):
+            flux_unit_type = "Surface brightness"
+        elif y_unit.is_equivalent(u.erg / (u.s * u.cm**2)):
+            flux_unit_type = 'Flux'
+        elif y_unit.is_equivalent(u.electron / u.s) or y_unit.physical_type == 'dimensionless':
+            # electron / s or 'dimensionless_unscaled' should be labeled counts
+            flux_unit_type = "Counts"
+        elif y_unit.is_equivalent(u.W):
+            flux_unit_type = "Luminosity"
+
+        # default to Flux Density for flux density or uncaught types
         flux_unit_type = "Flux density"
+
+        # Set x axes labels for the spectrum viewer
         x_disp_unit = self.state.x_display_unit
         x_unit = u.Unit(x_disp_unit) if x_disp_unit else u.dimensionless_unscaled
         if x_unit.is_equivalent(u.m):
diff --git a/jdaviz/configs/specviz/tests/test_viewers.py b/jdaviz/configs/specviz/tests/test_viewers.py
new file mode 100644
index 0000000000..883ebc86b7
--- /dev/null
+++ b/jdaviz/configs/specviz/tests/test_viewers.py
@@ -0,0 +1,26 @@
+import astropy.units as u
+import numpy as np
+import pytest
+from specutils import Spectrum1D
+
+
+@pytest.mark.parametrize(
+    ('input_unit', 'y_axis_label'),
+    [(u.MJy, 'Flux density'),
+     (u.MJy / u.sr, 'Surface brightness'),
+     (u.electron / u.s, 'Counts'),
+     (u.dimensionless_unscaled, 'Counts'),
+     (u.erg / (u.s * u.cm ** 2), 'Flux'),
+     (u.erg / u.s, 'Luminosity')])
+def test_spectrum_viewer_axis_labels(specviz_helper, input_unit, y_axis_label):
+
+    flux = np.arange(1, 10) * input_unit
+    spectral_axis = np.arange(1, 10) * u.um
+
+    spec = Spectrum1D(flux, spectral_axis)
+
+    specviz_helper.load_data(spec)
+
+    label = specviz_helper.app.get_viewer_by_id('specviz-0').figure.axes[1].label
+
+    assert (y_axis_label in label)

From 04a29edab55f1478fefef8f2d985b960820ca68a Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Wed, 14 Feb 2024 10:26:49 -0500
Subject: [PATCH 2/2] .

---
 jdaviz/configs/specviz/plugins/viewers.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/viewers.py b/jdaviz/configs/specviz/plugins/viewers.py
index e282b19dd9..d23e1e8c86 100644
--- a/jdaviz/configs/specviz/plugins/viewers.py
+++ b/jdaviz/configs/specviz/plugins/viewers.py
@@ -558,9 +558,9 @@ def set_plot_axes(self):
             flux_unit_type = "Counts"
         elif y_unit.is_equivalent(u.W):
             flux_unit_type = "Luminosity"
-
-        # default to Flux Density for flux density or uncaught types
-        flux_unit_type = "Flux density"
+        else:
+            # default to Flux Density for flux density or uncaught types
+            flux_unit_type = "Flux density"
 
         # Set x axes labels for the spectrum viewer
         x_disp_unit = self.state.x_display_unit