def test_mpi_with_dill(): pool = MPIPool(use_dill=True) pool.wait(lambda: sys.exit(0)) all_tasks = [[random.random() for i in range(1000)]] # test map alone for tasks in all_tasks: results = pool.map(_function, tasks) for r1, r2 in zip(results, [_function(x) for x in tasks]): assert isclose(r1, r2) assert len(results) == len(tasks) # test map with callback for tasks in all_tasks: results = pool.map(_function, tasks, callback=_callback) for r1, r2 in zip(results, [_function(x) for x in tasks]): assert isclose(r1, r2) assert len(results) == len(tasks) pool.close()
import sys # Project from schwimmbad.mpi import MPIPool, MPI from schwimmbad.tests.test_pools import isclose, _function def callback(x): assert MPI.COMM_WORLD.Get_rank() == 0 if MPI is not None and MPI.COMM_WORLD.Get_size() > 1: pool = MPIPool() pool.wait(lambda: sys.exit(0)) all_tasks = [[random.random() for i in range(1000)]] # test map alone for tasks in all_tasks: results = pool.map(_function, tasks) for r1,r2 in zip(results, [_function(x) for x in tasks]): assert isclose(r1, r2) # test map with callback for tasks in all_tasks: results = pool.map(_function, tasks, callback=callback) for r1,r2 in zip(results, [_function(x) for x in tasks]): assert isclose(r1, r2) pool.close()
pars.NonLinearModel.set_params(halofit_version='takahashi') camb.set_feedback_level(level=100) results = camb.get_results(pars) if shellnums is None: shellnum_min = int( results.comoving_radial_distance(zmin) * h // shellwidth) shellnum_max = int( results.comoving_radial_distance(zmax) * h // shellwidth + 1) shellnums = list(range(shellnum_min, shellnum_max + 1)) else: shellnums = list(map(int, shellnums.split(","))) try: pool = MPIPool() except: pool = None if is_cutsky: args = zip(product([1, 2, 3], shellnums), repeat(snapshot_cutsky)) else: args = zip(product([1, 2, 3], shellnums), repeat(None)) if pool is not None: if not pool.is_master(): pool.wait() sys.exit(0) pool.map(generate_lightcone_shell, args) else: map(generate_lightcone_shell, args)
try: pool = MPIPool() except: pool = None ### Loop over galtypes and shells for galaxies and randoms args = product(shellnums, nz_pars, [nrandoms], [dir_out], [lightcone_name_template], [output_name_template], [random_name_template]) if pool is not None: if not pool.is_master(): pool.wait() sys.exit(0) else: pool.map(process_shell, args) else: map(process_shell, args) ### Merge the processed shells flag = pool is None if not flag: flag = pool.is_master() if flag: for nz_par in nz_pars: galtype = nz_par["galtype_index"] out_fname = dir_out + merged_shells_name_template.format( nz_par['galtype'].upper()) if not os.path.exists(out_fname):