Repository: RadioAstronomySoftwareGroup/pyuvsim

Issue:
pyuvsim ignoring (and not updating) `UVData.uvw_array`
I stumbled upon an interesting feature when helping out a user who was running into an issue taking a `UVData` object from a prior observation, passing it through pyuvsim, and then writing out the result as an MS file (via the `UVData.write_ms` method). Long story short, what we found was that when passing a `UVData` object, `run_uvdata_uvsim` does not use the `uvw_array` values for calculating the interferometer response, favoring instead using the antenna positions to internally generate ENUs (and then calculating fringe response from there). This would be fine, except that `run_uvdata_uvsim` does _not_ appear to update `uvw_array` when it returns the object, which means that an object with incorrect uvw coordinates gets passed back.

At a minimum it would be helpful to raise an error when the values in `uvw_array` are not consistent w/ the antenna position-based values, though I think a better solution would be to simply overwrite the `uvw_array` values and wipe out any info on the phase centers (though I don't know the downstream ramifications of this, and figure this is a question for a more experienced dev/user to address). Happy to a write up a PR if someone wants to recommend which path to take.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/RadioAstronomySoftwareGroup/pyuvsim/pull/476?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=RadioAstronomySoftwareGroup) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 98.71%. Comparing base [(`d24b1c6`)](https://app.codecov.io/gh/RadioAstronomySoftwareGroup/pyuvsim/commit/d24b1c68235f87ae88e3fcdc072c1556f9432bd0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=RadioAstronomySoftwareGroup) to head [(`f74fad9`)](https://app.codecov.io/gh/RadioAstronomySoftwareGroup/pyuvsim/commit/f74fad9537c086d7b3658eb39aff11ca81c1a049?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=RadioAstronomySoftwareGroup).

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##              main     #476      +/-   ##
===========================================
- Coverage   100.00%   98.71%   -1.29%     
===========================================
  Files           11       11              
  Lines         2166     2172       +6     
===========================================
- Hits          2166     2144      -22     
- Misses           0       28      +28     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/RadioAstronomySoftwareGroup/pyuvsim/pull/476?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=RadioAstronomySoftwareGroup).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=RadioAstronomySoftwareGroup).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 547878b..9eb1eb7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,10 @@ changes for numpy 2.0 compatibility.
 - Use UVData.new method to do the UVData object initialization. This leads to a
 change in the default conjugation convention from `"ant2<ant1"` to `"ant1<ant2"`.
 
+### Fixed
+- Fixed a bug where `run_uvdata_uvsim` does not update uvw_array to what it actually
+used for the simulation, resulting in output objects with incorrect uvw arrays.
+
 ### Deprecated
 - The `"reorder_blt_kw"` parameter in the `simsetup.initialize_uvdata_from_params`
 function in favor of setting the ordering via the new "ordering" section in the
diff --git a/src/pyuvsim/uvsim.py b/src/pyuvsim/uvsim.py
index 8706499..d9afe85 100644
--- a/src/pyuvsim/uvsim.py
+++ b/src/pyuvsim/uvsim.py
@@ -492,8 +492,7 @@ def uvdata_to_task_iter(
     else:
         src_iter = [range(Nsrcs_total)]
     # Build the antenna list.
-    # remove the pragma below after pyuvdata v3.0 is released
-    if hasattr(input_uv, "telescope"):  # pragma: nocover
+    if hasattr(input_uv, "telescope"):
         antenna_names = input_uv.telescope.antenna_names
         antpos_enu = input_uv.telescope.get_enu_antpos()
         antenna_numbers = input_uv.telescope.antenna_numbers
@@ -511,7 +510,8 @@ def uvdata_to_task_iter(
         if world.lower() == "earth":
             tel_loc = EarthLocation.from_geocentric(*tloc, unit="m")
         elif world.lower() == "moon":
-            if not hasmoon:
+            if not hasmoon:  # pragma: nocover
+                # can only get here if pyuvdata < 3.0 and no lunarsky installed
                 raise ValueError(
                     "Need lunarsky module to simulate an array on the Moon."
                 )
@@ -671,6 +671,86 @@ def _set_nsky_parts(Nsrcs, cat_nfreqs, Nsky_parts):
     return Nsky_parts
 
 
+def _update_uvd(uv_container, *, input_uv, catalog, input_order):
+    """
+    Apply final updates to the simulated uvdata object.
+
+    Parameters
+    ----------
+    uv_container : UVData object
+        The uvdata object with the simulated visibilities that will be updated.
+    input_uv : UVData object
+        The input UVData object that the simulation ran from which may contain
+        some useful metadata that we want to get into the history of the
+        simulated UVData object.
+    catalog : SkyModel object
+        The input SkyModel object that was used in the simulation which may contain
+        some useful metadata that we want to get into the history of the
+        simulated UVData object.
+    input_order : tuple of str
+        The baseline-time ordering of the input_uv when it was passed in.
+
+    """
+    if input_order is not None:
+        if len(input_order) < 2:
+            input_order = (input_order[0], None)
+        # Don't call the function if we're already expecting (time, baseline)
+        if input_order != ("time", "baseline"):
+            uv_container.reorder_blts(order=input_order[0], minor_order=input_order[1])
+    else:
+        warnings.warn(
+            "The parameter `blt_order` could not be identified for input_uv "
+            " so the ordering cannot be restored."
+            "The output UVData object will have (time, baseline) ordering."
+        )
+
+    # set the uvws and phase center properly
+    # the sim uses the antenna locations and applies no projection, so
+    # if an object is passed in with projected phase centers or uvws not
+    # calculated from the antenna locations in ENU are wrong.
+    if np.any(~uv_container._check_for_cat_type("unprojected")):
+        # If there are any projected phase centers, just replace the whole
+        # catalog with an unprojected one.
+        uv_container.phase_center_catalog = {
+            0: {"cat_name": "unprojected", "cat_type": "unprojected"}
+        }
+    uv_container.set_uvws_from_antenna_positions(update_vis=False)
+
+    # Updating file history.
+    history = simutils.get_version_string()
+    if catalog.filename is not None:
+        history += (
+            " Sources from source list(s): [" + ", ".join(catalog.filename) + "]."
+        )
+    if "obsparam" in input_uv.extra_keywords:
+        obs_param_file = input_uv.extra_keywords["obsparam"]
+        telescope_config_file = input_uv.extra_keywords["telecfg"]
+        antenna_location_file = input_uv.extra_keywords["layout"]
+        history += (
+            " Based on config files: "
+            + obs_param_file
+            + ", "
+            + telescope_config_file
+            + ", "
+            + antenna_location_file
+        )
+    elif hasattr(uv_container, "filename") and uv_container.filename is not None:
+        history += (
+            " Based on uvdata file(s): [" + ", ".join(uv_container.filename) + "]."
+        )
+    history += " Npus = " + str(mpi.Npus) + "."
+
+    # add pyradiosky version
+    history += catalog.pyradiosky_version_str
+
+    # add pyuvdata version info
+    history += uv_container.pyuvdata_version_str
+
+    uv_container.history = history
+
+    return uv_container
+
+
 def run_uvdata_uvsim(
     input_uv,
     beam_list,
@@ -890,52 +970,9 @@ def run_uvdata_uvsim(
         vis_data.Free()
 
     if rank == 0:
-        if input_order is not None:
-            if len(input_order) < 2:
-                input_order = (input_order[0], None)
-            # Don't call the function if we're already expecting (time, baseline)
-            if input_order != ("time", "baseline"):
-                uv_container.reorder_blts(
-                    order=input_order[0], minor_order=input_order[1]
-                )
-        else:
-            warnings.warn(
-                "The parameter `blt_order` could not be identified for input_uv "
-                " so the ordering cannot be restored."
-                "The output UVData object will have (time, baseline) ordering."
-            )
-
-        # Updating file history.
-        history = simutils.get_version_string()
-        if catalog.filename is not None:
-            history += (
-                " Sources from source list(s): [" + ", ".join(catalog.filename) + "]."
-            )
-        if "obsparam" in input_uv.extra_keywords:
-            obs_param_file = input_uv.extra_keywords["obsparam"]
-            telescope_config_file = input_uv.extra_keywords["telecfg"]
-            antenna_location_file = input_uv.extra_keywords["layout"]
-            history += (
-                " Based on config files: "
-                + obs_param_file
-                + ", "
-                + telescope_config_file
-                + ", "
-                + antenna_location_file
-            )
-        elif hasattr(uv_container, "filename") and uv_container.filename is not None:
-            history += (
-                " Based on uvdata file(s): [" + ", ".join(uv_container.filename) + "]."
-            )
-        history += " Npus = " + str(mpi.Npus) + "."
-
-        # add pyradiosky version
-        history += catalog.pyradiosky_version_str
-
-        # add pyuvdata version info
-        history += uv_container.pyuvdata_version_str
-
-        uv_container.history = history
+        uv_container = _update_uvd(
+            uv_container, input_uv=input_uv, catalog=catalog, input_order=input_order
+        )
 
         return uv_container
 


Test patch:
diff --git a/tests/test_run.py b/tests/test_run.py
index 7854b3b..22ce158 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -201,7 +201,7 @@ def test_powerbeam_sim(cst_beam):
         cfg, return_catname=False
     )
     sky_model = pyuvsim.simsetup.SkyModelData(sky_model)
-    beam_dict = {n: 0 for n in range(4)}
+    beam_dict = dict.fromkeys(range(4), 0)
 
     with pytest.raises(ValueError, match="Beam type must be efield!"):
         pyuvsim.run_uvdata_uvsim(input_uv, beams, beam_dict, catalog=sky_model)
diff --git a/tests/test_uvsim.py b/tests/test_uvsim.py
index a723da9..3bb8020 100644
--- a/tests/test_uvsim.py
+++ b/tests/test_uvsim.py
@@ -177,7 +177,7 @@ def uvdata_two_redundant_bls_triangle_sources():
     beam_list[0] = pyuvsim.AnalyticBeam("airy", diameter=13.0)
 
     time = Time(uv_obj.time_array[0], format="jd", scale="utc")
-    sources, kwds = pyuvsim.create_mock_catalog(
+    sources, _ = pyuvsim.create_mock_catalog(
         time, arrangement="triangle", Nsrcs=3, return_data=True
     )
 
@@ -734,7 +734,7 @@ def test_nsky_parts_large(capsys):
         hera_uv.use_future_array_shapes()
     hera_uv.select(times=np.unique(hera_uv.time_array)[0:3], freq_chans=range(3))
     time = Time(hera_uv.time_array[0], scale="utc", format="jd")
-    sources, kwds = pyuvsim.create_mock_catalog(
+    sources, _ = pyuvsim.create_mock_catalog(
         time, arrangement="random", Nsrcs=25, return_data=True, rseed=100
     )
 
@@ -1134,20 +1134,23 @@ def test_fullfreq_check(uvobj_beams_srcs):
     next(taskiter1)
 
 
-@pytest.mark.skipif("hasmoon")
 def test_moonloc_error(uvobj_beams_srcs):
-    # Break if the uvobj indicates that the sim is on the Moon, but lunarsky is not available.
 
     uv_obj, beam_list, beam_dict, sources = uvobj_beams_srcs
+    if hasattr(uv_obj, "telescope"):
+        # This whole test can go away once we require pyuvdata >= 3.0
+        pytest.skip()
 
     uv_obj.extra_keywords["world"] = "moon"
 
-    with pytest.raises(ValueError, match="Need lunarsky module to simulate"):
-        next(
-            pyuvsim.uvsim.uvdata_to_task_iter(
-                range(5), uv_obj, sources, beam_list, beam_dict
+    if not hasmoon:
+        # Break if the uvobj indicates that the sim is on the Moon, but lunarsky is not installed.
+        with pytest.raises(ValueError, match="Need lunarsky module to simulate"):
+            next(
+                pyuvsim.uvsim.uvdata_to_task_iter(
+                    range(5), uv_obj, sources, beam_list, beam_dict
+                )
             )
-        )
 
     uv_obj.extra_keywords["world"] = "gallifrey"
 
@@ -1293,3 +1296,21 @@ def test_nblts_not_square(uvdata_two_redundant_bls_triangle_sources, cut_beam):
             out_uv.Nfreqs,
             out_uv.Npols,
         )
+
+
+@pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected")
+def test_uvdata_uvsim_uvw_setting(uvdata_two_redundant_bls_triangle_sources):
+    pytest.importorskip("mpi4py")
+    uvd, beam_list, beam_dict, sky_model = uvdata_two_redundant_bls_triangle_sources
+    uvd2 = uvd.copy()
+    uvd2.uvw_array = uvd2.uvw_array * 1.2
+
+    out_uv1 = pyuvsim.uvsim.run_uvdata_uvsim(
+        input_uv=uvd, beam_list=beam_list, beam_dict=beam_dict, catalog=sky_model
+    )
+
+    out_uv2 = pyuvsim.uvsim.run_uvdata_uvsim(
+        input_uv=uvd2, beam_list=beam_list, beam_dict=beam_dict, catalog=sky_model
+    )
+
+    assert out_uv1 == out_uv2