update tests to run in paralell. do gymnastics for warnings

Author: mkolopanis

pyuvsim/tests/test_run.py

@@ -144,10 +144,10 @@ def test_analytic_diffuse(model, tmpdir, backend, progbar):
     )
     if pyuvsim.mpi.rank == 0:
         # Convert from Jy to K sr
-        dat = uv_out.data_array[:, 0, 0, 0] * jy_to_ksr(uv_out.freq_array[0, 0]).value
+        dat = uv_out.data_array[:, 0, 0] * jy_to_ksr(uv_out.freq_array[0]).value
         # Evaluate the solution and compare to visibilities.
         soln = analytic_diffuse.get_solution(modname)
-        uvw_lam = uv_out.uvw_array * uv_out.freq_array[0, 0] / c_ms
+        uvw_lam = uv_out.uvw_array * uv_out.freq_array[0] / c_ms
         ana = soln(uvw_lam, **params)
         assert np.allclose(ana / 2, dat, atol=1e-2)
 
@@ -252,7 +252,7 @@ def test_zenith_spectral_sim(spectral_type, tmpdir, backend, progbar):
     uv_out = pyuvsim.run_uvsim(params, return_uv=True, backend=backend, progbar=progbar)
     if pyuvsim.mpi.rank == 0:
         for ii in range(uv_out.Nbls):
-            assert np.allclose(uv_out.data_array[ii, 0, :, 0], spectrum / 2)
+            assert np.allclose(uv_out.data_array[ii, :, 0], spectrum / 2)
 
 
 def test_pol_error():
@@ -293,7 +293,7 @@ def test_sim_on_moon(future_shapes):
     )
     if pyuvsim.mpi.rank == 0:
 
-        assert np.allclose(uv_out.data_array[:, 0, :, 0], 0.5)
+        assert np.allclose(uv_out.data_array[:, :, 0], 0.5)
         assert uv_out.extra_keywords['world'] == 'moon'
 
 

pyuvsim/tests/test_uvsim.py

@@ -921,6 +921,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')
@@ -934,41 +935,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
@@ -980,7 +1002,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)
@@ -994,12 +1016,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():