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):