Temporary fix for two failing MPI tests related to chunk partition (N…

…anoComp#2972)

* fix failing MPI tests

* specify chunk layout in test_fragment_stats.py

* specify chunk layout in test_dump_load.py

* more fixes

* fixes

* fixes for test_fragment_stats.py

* fixes for test_dump_load.py

python/tests/test_dump_load.py

@@ -153,6 +153,11 @@ def _load_dump_structure_3d(
         boundary_layers = [mp.Absorber(0.2)]
         k_point = mp.Vector3(0.4, -1.3, 0.7)
 
+        if mp.count_processors() == 2:
+            chunk_layout = mp.BinaryPartition(data=[(mp.Z, 0), 0, 1])
+        else:
+            chunk_layout = None
+
         sim1 = mp.Simulation(
             resolution=resolution,
             cell_size=cell,
@@ -161,6 +166,7 @@ def _load_dump_structure_3d(
             boundary_layers=boundary_layers,
             default_material=default_material,
             sources=[sources],
+            chunk_layout=chunk_layout,
         )
 
         sample_point = mp.Vector3(0.73, -0.33, 0.61)
@@ -189,6 +195,9 @@ def get_ref_field_point(sim):
         if chunk_sim:
             chunk_layout = sim1
 
+        if chunk_layout is None and mp.count_processors() == 2:
+            chunk_layout = mp.BinaryPartition(data=[(mp.Z, 0), 0, 1])
+
         sim = mp.Simulation(
             resolution=resolution,
             cell_size=cell,

python/tests/test_fragment_stats.py

@@ -24,12 +24,18 @@ def make_dft_vecs(
 
 
 def make_sim(cell, res, pml, dims, create_gv=True, k_point=False):
+    if mp.count_processors() == 2:
+        chunk_layout = mp.BinaryPartition(data=[(mp.X, 0), 0, 1])
+    else:
+        chunk_layout = None
+
     sim = mp.Simulation(
         cell_size=cell,
         resolution=res,
         boundary_layers=pml,
         dimensions=dims,
         k_point=k_point,
+        chunk_layout=chunk_layout,
     )
     if create_gv:
         sim._create_grid_volume(False)