update tests to run in paralell. do gymnastics for warnings

Author: mkolopanis

pyuvsim/tests/test_uvsim.py

@@ -918,6 +918,7 @@ def test_run_mpierr():
             pyuvsim.run_uvdata_uvsim(UVData(), ['beamlist'])
 
 
+@pytest.mark.parallel(2)
 @pytest.mark.parametrize("order", [("bda",), ("baseline", "time"), ("ant2", "time")])
 def test_ordering(uvdata_two_redundant_bls_triangle_sources, order):
     pytest.importorskip('mpi4py')
@@ -931,41 +932,62 @@ def test_ordering(uvdata_two_redundant_bls_triangle_sources, order):
         beam_dict=beam_dict,
         catalog=sky_model,
     )
-    assert out_uv.blt_order == order
-    assert out_uv.blt_order == uvdata_linear.blt_order
+    rank = pyuvsim.mpi.get_rank()
+    if rank == 0:
+        print(rank, out_uv)
+        assert out_uv.blt_order == order
+        assert out_uv.blt_order == uvdata_linear.blt_order
 
-    uvdata_linear.data_array = out_uv.data_array
+        uvdata_linear.data_array = out_uv.data_array
 
-    uvdata_linear.reorder_blts(order="time", minor_order="baseline", conj_convention="ant1<ant2")
+        uvdata_linear.reorder_blts(
+            order="time", minor_order="baseline", conj_convention="ant1<ant2"
+        )
 
-    assert np.allclose(
-        uvdata_linear.get_data((0, 1)), uvdata_linear.get_data((1, 2))
-    )
-    assert not np.allclose(
-        uvdata_linear.get_data((0, 1)), uvdata_linear.get_data((0, 2))
-    )
+        assert np.allclose(
+            uvdata_linear.get_data((0, 1)), uvdata_linear.get_data((1, 2))
+        )
+        assert not np.allclose(
+            uvdata_linear.get_data((0, 1)), uvdata_linear.get_data((0, 2))
+        )
 
 
+@pytest.mark.parallel(2)
 @pytest.mark.parametrize("order", [("bda",), ("baseline", "time"), ("ant2", "time")])
 def test_order_warning(uvdata_two_redundant_bls_triangle_sources, order):
     pytest.importorskip('mpi4py')
+    # need to get the mpi initialized
+    # now that simulations require at least 2 PUs
+    pyuvsim.mpi.start_mpi()
+    rank = pyuvsim.mpi.get_rank()
     uvdata_linear, beam_list, beam_dict, sky_model = uvdata_two_redundant_bls_triangle_sources
 
     uvdata_linear.reorder_blts(*order)
     # delete the order like we forgot to set it
     uvdata_linear.blt_order = None
-    with uvtest.check_warnings(
-        UserWarning, match="The parameter `blt_order` could not be identified."
-    ):
+    if rank == 0:
+        with uvtest.check_warnings(
+            UserWarning, match="The parameter `blt_order` could not be identified."
+        ):
+
+            out_uv = pyuvsim.uvsim.run_uvdata_uvsim(
+                input_uv=uvdata_linear.copy(),
+                beam_list=beam_list,
+                beam_dict=beam_dict,
+                catalog=sky_model,
+            )
+
+        assert out_uv.blt_order == ("time", "baseline")
+    else:
         out_uv = pyuvsim.uvsim.run_uvdata_uvsim(
             input_uv=uvdata_linear.copy(),
             beam_list=beam_list,
             beam_dict=beam_dict,
             catalog=sky_model,
         )
-    assert out_uv.blt_order == ("time", "baseline")
 
 
+@pytest.mark.parallel(2)
 def test_nblts_not_square(uvdata_two_redundant_bls_triangle_sources):
     pytest.importorskip('mpi4py')
     uvdata_linear, beam_list, beam_dict, sky_model = uvdata_two_redundant_bls_triangle_sources
@@ -977,7 +999,7 @@ def test_nblts_not_square(uvdata_two_redundant_bls_triangle_sources):
     indices = np.nonzero(
         uvdata_linear.baseline_array == uvdata_linear.antnums_to_baseline(0, 2)
     )[0]
-    print(uvdata_linear.antnums_to_baseline(0, 2), indices)
+
     # discard half of them
     indices = indices[::2]
     blt_inds = np.delete(np.arange(uvdata_linear.Nblts), indices)
@@ -991,12 +1013,13 @@ def test_nblts_not_square(uvdata_two_redundant_bls_triangle_sources):
         beam_dict=beam_dict,
         catalog=sky_model,
     )
-
-    assert np.allclose(
-        out_uv.get_data((0, 1)), out_uv.get_data((1, 2))
-    )
-    # make sure (0, 2) has fewer times
-    assert out_uv.get_data((0, 2)).shape == (out_uv.Ntimes // 2, out_uv.Nfreqs, out_uv.Npols)
+    rank = pyuvsim.mpi.get_rank()
+    if rank == 0 :
+        assert np.allclose(
+            out_uv.get_data((0, 1)), out_uv.get_data((1, 2))
+        )
+        # make sure (0, 2) has fewer times
+        assert out_uv.get_data((0, 2)).shape == (out_uv.Ntimes // 2, out_uv.Nfreqs, out_uv.Npols)
 
 
 def test_tqdm_import_error():