Fix a bug in properly handling diffuse maps with units of Jy/sr

CHANGELOG.md

@@ -15,6 +15,7 @@ string/object mode switch.
 creating the UVData object, rather than down-selecting afterwards, saving memory and time.
 
 ### Fixed
+- A bug in handling diffuse maps with units of Jy/sr.
 - A bug when specifying a frequency buffer for beam frequency selects in telescope
 config files.
 

src/pyuvsim/simsetup.py

@@ -673,12 +673,8 @@ def __init__(self, sky_in=None, filename=None):
             stokes_in = sky_in.stokes
 
             if isinstance(stokes_in, units.Quantity):
-                if stokes_in.unit.is_equivalent("Jy"):
-                    stokes_in = stokes_in.to_value("Jy")
-                    self.flux_unit = "Jy"
-                elif stokes_in.unit.is_equivalent("K"):
-                    stokes_in = stokes_in.to_value("K")
-                    self.flux_unit = "K"
+                self.flux_unit = stokes_in.unit.to_string()
+                stokes_in = stokes_in.value
 
             self.stokes_I = stokes_in[0, ...]
 

tests/test_run.py

@@ -13,7 +13,7 @@
 import yaml
 from astropy import units
 from astropy.constants import c as speed_of_light
-from astropy.coordinates import Latitude, Longitude
+from astropy.coordinates import EarthLocation, Latitude, Longitude
 from astropy.time import Time
 from pyradiosky.utils import jy_to_ksr, stokes_to_coherency
 from pyuvdata import ShortDipoleBeam, UniformBeam, UVData
@@ -167,6 +167,86 @@ def test_analytic_diffuse(model, tol, tmpdir):
     np.testing.assert_allclose(ana / 2, dat, atol=tol, rtol=0)
 
 
+def test_diffuse_units(tmpdir):
+    pytest.importorskip("analytic_diffuse")
+    pytest.importorskip("astropy_healpix")
+
+    sky_k, _, _, _, _ = pyuvsim.simsetup._create_catalog_diffuse(
+        map_nside=128,
+        diffuse_model="polydome",
+        diffuse_params={"n": 2},
+        time=Time(2457458.1738949567, format="jd"),
+        array_location=EarthLocation.from_geodetic(
+            lat=-30.72153, lon=21.42830, height=1073.0
+        ),
+    )
+
+    sky_k.spectral_type = "spectral_index"
+    sky_k.reference_frequency = np.full(sky_k.Ncomponents, 1e8) * units.Hz
+    sky_k.spectral_index = np.full(sky_k.Ncomponents, -0.8)
+    sky_k.check()
+
+    sky_jy = sky_k.copy()
+    sky_jy.kelvin_to_jansky()
+    assert sky_jy.stokes.unit == units.Unit("Jy/sr")
+
+    catpath_k = str(tmpdir.join("diffuse_units_k.skyh5"))
+    sky_k.write_skyh5(catpath_k)
+
+    catpath_jy = str(tmpdir.join("diffuse_units_jy.skyh5"))
+    sky_jy.write_skyh5(catpath_jy)
+
+    # Making configuration files for this simulation.
+    template_path = os.path.join(
+        SIM_DATA_PATH, "test_config", "obsparam_diffuse_sky.yaml"
+    )
+    obspar_path = str(tmpdir.join("obsparam_diffuse_sky.yaml"))
+    layout_path = str(tmpdir.join("threeant_layout.csv"))
+    herauniform_path = str(tmpdir.join("hera_uniform.yaml"))
+
+    teleconfig = {
+        "beam_paths": {0: UniformBeam()},
+        "telescope_location": "(-30.72153, 21.42830, 1073.0)",
+        "telescope_name": "HERA",
+    }
+    antpos_enu = np.array([[0, 0, 0], [0, 3, 0], [5, 0, 0]], dtype=float)
+
+    pyuvsim.simsetup._write_layout_csv(
+        layout_path, antpos_enu, np.arange(3).astype(str), np.arange(3)
+    )
+    with open(herauniform_path, "w") as ofile:
+        yaml.dump(teleconfig, ofile, default_flow_style=False)
+
+    with open(template_path) as yfile:
+        obspar = yaml.safe_load(yfile)
+    obspar["telescope"]["array_layout"] = layout_path
+    obspar["telescope"]["telescope_config_name"] = herauniform_path
+    obspar["sources"] = {"catalog": catpath_k}
+
+    obspar["filing"]["outfile_name"] = "diffuse_units_k_sim.uvh5"
+    obspar["filing"]["output_format"] = "uvh5"
+    obspar["filing"]["outdir"] = str(tmpdir)
+
+    with open(obspar_path, "w") as ofile:
+        yaml.dump(obspar, ofile, default_flow_style=False)
+
+    uv_out_k = pyuvsim.run_uvsim(obspar_path, return_uv=True)
+
+    obspar["sources"] = {"catalog": catpath_jy}
+    obspar["filing"]["outfile_name"] = "diffuse_units_jy_sim.uvh5"
+    with open(obspar_path, "w") as ofile:
+        yaml.dump(obspar, ofile, default_flow_style=False)
+
+    uv_out_jy = pyuvsim.run_uvsim(obspar_path, return_uv=True)
+
+    # make the histories the same for comparison
+    uv_out_jy.history = uv_out_k.history
+    # harmonize phase center catalog info
+    uv_out_jy._consolidate_phase_center_catalogs(other=uv_out_k, ignore_name=True)
+
+    assert uv_out_k == uv_out_jy
+
+
 @pytest.mark.filterwarnings("ignore:Fixing auto polarization power beams")
 @pytest.mark.parametrize(
     ("problem", "err_msg"),

tests/test_simsetup.py

@@ -1927,24 +1927,34 @@ def test_mock_catalog_moon():
 
 
 @pytest.mark.filterwarnings("ignore:Input ra and dec parameters are being used instead")
-@pytest.mark.parametrize("unit", ["Jy", "K"])
+@pytest.mark.parametrize("unit", ["Jy", "K", "Jy/sr", "K sr"])
 def test_skymodeldata_with_quantity_stokes(unit, cat_with_some_pols):
     # Support for upcoming pyradiosky change setting SkyModel.stokes
     # to an astropy Quantity.
-    if unit == "Jy":
+    if unit in ["Jy", "K sr"]:
         sky = cat_with_some_pols
+        if unit != "Jy":
+            sky.jansky_to_kelvin()
     else:
         pytest.importorskip("analytic_diffuse")
         pytest.importorskip("astropy_healpix")
         sky, _ = simsetup.create_mock_catalog(
             Time.now(), arrangement="diffuse", diffuse_model="monopole", map_nside=16
         )
+        sky.spectral_type = "spectral_index"
+        sky.reference_frequency = np.full(sky.Ncomponents, 1e8) * units.Hz
+        sky.spectral_index = np.full(sky.Ncomponents, -0.8)
+        sky.check()
+
+        if unit != "K":
+            sky.kelvin_to_jansky()
+
     if not isinstance(sky.stokes, units.Quantity):
         sky.stokes *= units.Unit(unit)
 
     smd = simsetup.SkyModelData(sky)
     assert np.all(sky.stokes.to_value(unit)[0] == smd.stokes_I)
-    assert smd.flux_unit == unit
+    assert units.Unit(smd.flux_unit) == units.Unit(unit)
 
     sky2 = smd.get_skymodel()
     assert sky2 == sky