From 3ea11dc41594d0d4e999005f83a3b0046bd98c3b Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Fri, 6 Sep 2024 13:40:15 -0400
Subject: [PATCH] fixed equivs, more tests passing?

---
 jdaviz/tests/test_utils.py | 10 ++++++++++
 jdaviz/utils.py            | 26 ++++++++++++++++++--------
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/jdaviz/tests/test_utils.py b/jdaviz/tests/test_utils.py
index 86294bcd12..e7122917c0 100644
--- a/jdaviz/tests/test_utils.py
+++ b/jdaviz/tests/test_utils.py
@@ -15,6 +15,8 @@
 
 PHOTUTILS_LT_1_12_1 = not minversion(photutils, "1.12.1.dev")
 
+PIX2 = u.pix * u.pix
+
 
 def test_spec_sb_flux_conversion():
     # Actual spectrum content does not matter, just the meta is used here.
@@ -27,14 +29,22 @@ def test_spec_sb_flux_conversion():
     spec.meta["_pixel_scale_factor"] = 0.1
     assert_allclose(flux_conversion(values, u.Jy / u.sr, u.Jy, spec), [1, 2, 3])
     assert_allclose(flux_conversion(values, u.Jy, u.Jy / u.sr, spec), [100, 200, 300])
+    
+    # conversions with eq pixels
+    assert_allclose(flux_conversion(values, u.Jy / PIX2, u.Jy, spec), [10, 20, 30])
+    assert_allclose(flux_conversion(values, u.Jy, u.Jy / PIX2, spec), [10, 20, 30])
 
     # complex translation Jy / sr -> erg / (Angstrom s cm2 sr)
     targ = [2.99792458e-12, 1.49896229e-12, 9.99308193e-13] * (u.erg / (u.Angstrom * u.s * u.cm**2 * u.sr))  # noqa: E501
     assert_allclose(flux_conversion(values, u.Jy / u.sr, u.erg / (u.Angstrom * u.s * u.cm**2 * u.sr), spec), targ.value)  # noqa: E501
+    # swap sr for pix2 and check conversion
+    assert_allclose(flux_conversion(values, u.Jy / PIX2, u.erg / (u.Angstrom * u.s * u.cm**2 * PIX2), spec), targ.value)  # noqa: E501
 
     # complex translation erg / (Angstrom s cm2 sr) -> Jy / sr
     targ = [3.33564095e+13, 2.66851276e+14, 9.00623057e+14] * (u.Jy / u.sr)
     assert_allclose(flux_conversion(values, u.erg / (u.Angstrom * u.s * u.cm**2 * u.sr), u.Jy / u.sr, spec), targ.value)  # noqa: E501
+    # swap sr for pix2 and check conversion
+    assert_allclose(flux_conversion(values, u.erg / (u.Angstrom * u.s * u.cm**2 * PIX2), u.Jy / PIX2, spec), targ.value)  # noqa: E501
 
     spectral_values = spec.spectral_axis
     spec_unit = u.MJy
diff --git a/jdaviz/utils.py b/jdaviz/utils.py
index ce3dba9034..be0bc4fdf5 100644
--- a/jdaviz/utils.py
+++ b/jdaviz/utils.py
@@ -468,16 +468,20 @@ def flux_conversion(values, original_units, target_units, spec=None, eqv=None, s
 def _indirect_conversion(values, orig_units, targ_units, eqv,
                          spec_unit=None, image_data=None):
 
-    # make these an input parameter since they're already determined
-    # here for now until i make sure this function isn't called elsewhere
+    # Note: is there a way we could write this to not require 'spec_unit'? It 
+    # seems like it falls back on this to get a solid angle unit, but can we
+    # assume pix2 now if there are none? or use the display units?
+
     solid_angle_in_orig = check_if_unit_is_per_solid_angle(orig_units, return_unit=True)
     solid_angle_in_targ = check_if_unit_is_per_solid_angle(targ_units, return_unit=True)
-    solid_angle_in_spec = check_if_unit_is_per_solid_angle(spec_unit, return_unit=True)
+    if spec_unit is not None:
+        solid_angle_in_spec = check_if_unit_is_per_solid_angle(spec_unit, return_unit=True)
+    else:
+        solid_angle_in_spec = None
 
     # indirect units cannot be directly converted, and require
     # additional conversions to reach the desired end unit.
-    if (spec_unit and spec_unit in [orig_units, targ_units]
-            and not check_if_unit_is_per_solid_angle(spec_unit)):
+    if (spec_unit and spec_unit in [orig_units, targ_units] and not solid_angle_in_spec):
         if u.Unit(targ_units) in indirect_units():
             temp_targ = targ_units * solid_angle_in_targ
             values = (values * orig_units).to_value(temp_targ, equivalencies=eqv)
@@ -511,8 +515,6 @@ def _indirect_conversion(values, orig_units, targ_units, eqv,
 
         return values, orig_units
 
-# Note: should unify how these custom equivs are written, either they should take in any flux unit
-# or return multiple equivalencies for all non-equivalent flux unit types.
 
 
 def _eqv_pixar_sr(pixar_sr):
@@ -543,7 +545,15 @@ def _eqv_flux_to_sb_pixel(flux_unit):
     """
 
     pix2 = u.pix * u.pix
-    equiv = [(flux_unit, flux_unit / pix2, lambda x: x, lambda x: x)]
+
+    # generate an equivalency for each flux type that would need
+    # another equivalency for converting to/from
+    flux_units = [u.MJy, u.erg / (u.s * u.cm**2 * u.Angstrom),
+                  u.ph / (u.Angstrom * u.s * u.cm**2),
+                  u.ph / (u.Hz * u.s * u.cm**2)]
+    equiv = []
+    for flux_unit in flux_units:
+        equiv.append((flux_unit, flux_unit / pix2, lambda x: x, lambda x: x))
 
     return equiv