diff --git a/python/tests/test_dump_load.py b/python/tests/test_dump_load.py index c97712733..cf429fc70 100644 --- a/python/tests/test_dump_load.py +++ b/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, diff --git a/python/tests/test_fragment_stats.py b/python/tests/test_fragment_stats.py index 2bbfe0bb3..e56148530 100644 --- a/python/tests/test_fragment_stats.py +++ b/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)