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>

.

.

.

added test

.

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
 ^^^^^^^^^
 

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
+
+        # default to Flux Density for flux density or uncaught types
         flux_unit_type = "Flux density"
+
+        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"
+
+        # 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):

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.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)
+
+    spec = Spectrum1D(flux, spectral_axis * u.nm)
+
+    specviz_helper.load_data(spec)
+
+    label = specviz_helper.app.get_viewer_by_id('specviz-0').figure.axes[1].label
+    label = label.split('[')[0].strip()
+
+    assert (label == y_axis_label)