add tests for multiple iterations

Author: michaelbuehlmann

tests/test_distribute.py

@@ -8,7 +8,7 @@
 from mpipartition import Partition, distribute
 
 
-def create_and_distribute(dimensions: int, N: int) -> None:
+def create_and_distribute(dimensions: int, N: int, iterations: int) -> None:
     partition = Partition(dimensions)
     assert dimensions < 7
     assert dimensions > 0
@@ -17,7 +17,7 @@ def create_and_distribute(dimensions: int, N: int) -> None:
     coord_keys = [x for x in labels]
 
     data = {x: np.random.uniform(0, 1, N) for x in labels}
-    data = distribute(partition, 1.0, data, coord_keys)
+    data = distribute(partition, 1.0, data, coord_keys, all2all_iterations=iterations)
 
     valid = np.ones(len(data[labels[0]]), dtype=np.bool_)
     for i, label in enumerate(labels):
@@ -33,19 +33,27 @@ def create_and_distribute(dimensions: int, N: int) -> None:
 
 @pytest.mark.mpi
 def test_1d() -> None:
-    create_and_distribute(1, 1000)
+    create_and_distribute(1, 1000, 1)
+    create_and_distribute(1, 1000, 2)
+    create_and_distribute(1, 1000, 4)
 
 
 @pytest.mark.mpi
 def test_2d() -> None:
-    create_and_distribute(2, 100)
+    create_and_distribute(2, 100, 1)
+    create_and_distribute(2, 100, 2)
+    create_and_distribute(2, 100, 4)
 
 
 @pytest.mark.mpi
 def test_3d() -> None:
-    create_and_distribute(3, 10)
+    create_and_distribute(3, 10, 1)
+    create_and_distribute(3, 10, 2)
+    create_and_distribute(3, 10, 4)
 
 
 @pytest.mark.mpi
 def test_4d() -> None:
-    create_and_distribute(4, 5)
+    create_and_distribute(4, 5, 1)
+    create_and_distribute(4, 5, 2)
+    create_and_distribute(4, 5, 4)

tests/test_s2_distribute.py

@@ -8,7 +8,7 @@
 from mpipartition import S2Partition, s2_distribute
 
 
-def create_and_distribute(N: int, equal_area: bool) -> None:
+def create_and_distribute(N: int, equal_area: bool, iterations: int) -> None:
     partition = S2Partition(equal_area=equal_area)
     np.random.seed(partition.rank)
 
@@ -17,7 +17,7 @@ def create_and_distribute(N: int, equal_area: bool) -> None:
     id = np.arange(N, dtype=np.int64) + partition.rank * N
 
     data = {"phi": phi, "theta": theta, "id": id}
-    data = s2_distribute(partition, data)
+    data = s2_distribute(partition, data, all2all_iterations=iterations)
 
     valid = np.ones(len(data["phi"]), dtype=np.bool_)
     valid &= data["phi"] >= partition.phi_extent[0]
@@ -30,9 +30,13 @@ def create_and_distribute(N: int, equal_area: bool) -> None:
 
 @pytest.mark.mpi
 def test_s2_distribute_equal_area() -> None:
-    create_and_distribute(100, equal_area=True)
+    create_and_distribute(100, equal_area=True, iterations=1)
+    create_and_distribute(100, equal_area=True, iterations=2)
+    create_and_distribute(100, equal_area=True, iterations=4)
 
 
 @pytest.mark.mpi
 def test_s2_distribute_unequal_area() -> None:
-    create_and_distribute(100, equal_area=False)
+    create_and_distribute(100, equal_area=False, iterations=1)
+    create_and_distribute(100, equal_area=False, iterations=2)
+    create_and_distribute(100, equal_area=False, iterations=4)