def test_mpi_smoothing(mpi_comm):
nside = 128
lmax = 2 * nside
map_dist = pysm.MapDistribution(pixel_indices=None,
mpi_comm=mpi_comm,
smoothing_lmax=lmax,
nside=nside)
model = pysm.Model(nside, map_dist=map_dist)
distributed_map = model.read_map("pysm_2/dust_temp.fits")
fwhm = 5 * u.deg
smoothed_distributed_map = pysm.mpi_smoothing(distributed_map,
fwhm,
map_dist=map_dist)
full_map_rank0 = pysm.mpi.assemble_map_on_rank0(
mpi_comm,
smoothed_distributed_map,
model.map_dist.pixel_indices,
n_components=1,
npix=hp.nside2npix(nside),
)[0]
if mpi_comm.rank == 0:
np.testing.assert_allclose(
full_map_rank0,
hp.smoothing(
pysm.read_map("pysm_2/dust_temp.fits", nside=nside).value,
fwhm.to(u.rad).value,
iter=0,
lmax=lmax,
use_pixel_weights=False,
),
rtol=1e-5,
)
def test_read_map_mpi_pixel_indices(mpi_comm):
# Reads pixel [0] on rank 0
# pixels [0,1] on rank 1
# pixels [0,1,2] on rank 2 and so on.
map_dist = pysm.MapDistribution(mpi_comm=mpi_comm,
pixel_indices=list(
range(0, mpi_comm.rank + 1)))
m = pysm.read_map("pysm_2/dust_temp.fits",
nside=8,
field=0,
map_dist=map_dist)
assert len(m) == mpi_comm.rank + 1
def test_read_map_mpi_uniform_distribution(mpi_comm):
# Spreads the map equally across processes
map_dist = pysm.MapDistribution(
mpi_comm=mpi_comm,
pixel_indices=pysm.mpi.distribute_pixels_uniformly(mpi_comm, nside=8),
)
m = pysm.read_map("pysm_2/dust_temp.fits",
nside=8,
field=0,
map_dist=map_dist)
npix = hp.nside2npix(8)
assert (
npix % mpi_comm.size == 0
), "This test requires the size of the communicator to divide the number of pixels {}".format(
npix)
assert len(m) == npix / mpi_comm.size
def init_sky(self, pysm_sky_config, pysm_precomputed_cmb_K_CMB):
if pysm is None:
raise RuntimeError("pysm not available")
if pysm_precomputed_cmb_K_CMB is not None:
raise NotImplementedError(
"pysm_precomputed_cmb_K_CMB is not currently supported")
# cmb = {
# "model": "pre_computed",
# "nside": self._nside,
# "pixel_indices": self._pixel_indices,
# }
# # PySM expects uK_CMB
# cmb["A_I"], cmb["A_Q"], cmb["A_U"] = (
# np.array(
# pysm.read_map(
# pysm_precomputed_cmb_K_CMB,
# self._nside,
# field=(0, 1, 2),
# pixel_indices=self._pixel_indices,
# mpi_comm=self._comm,
# )
# )
# * 1e6
# )
# initialized_sky_config["cmb"] = [cmb]
# # remove cmb from the pysm string
# pysm_sky_config.pop("cmb", None)
if self.map_dist is None:
self.map_dist = (None
if self._comm is None else pysm.MapDistribution(
pixel_indices=self._pixel_indices,
nside=self._nside,
mpi_comm=self._comm,
))
return pysm.Sky(
nside=self._nside,
preset_strings=pysm_sky_config,
component_objects=self.pysm_component_objects,
map_dist=self.map_dist,
output_unit=self._units,
)
def test_mpi_assemble(mpi_comm):
nside = 128
lmax = 2 * nside
map_dist = pysm.MapDistribution(pixel_indices=None,
mpi_comm=mpi_comm,
nside=nside)
model = pysm.Model(nside, map_dist=map_dist)
distributed_map = model.read_map("pysm_2/dust_temp.fits")
full_map_rank0 = pysm.mpi.assemble_map_on_rank0(
mpi_comm,
distributed_map,
model.map_dist.pixel_indices,
n_components=1,
npix=hp.nside2npix(nside),
)[0]
if mpi_comm.rank == 0:
np.testing.assert_allclose(
full_map_rank0,
pysm.read_map("pysm_2/dust_temp.fits", nside=nside).value,
rtol=1e-5,
)
def test_read_map_mpi_uniform_distribution(mpi_comm):
# Spreads the map equally across processes
map_dist = pysm.MapDistribution(
mpi_comm=mpi_comm,
pixel_indices=pysm.mpi.distribute_pixels_uniformly(mpi_comm, nside=8),
)
m = pysm.read_map("pysm_2/dust_temp.fits",
nside=8,
field=0,
map_dist=map_dist)
npix = hp.nside2npix(8)
assert (
npix % mpi_comm.size == 0
), "This test requires the size of the communicator to divide the number of pixels {}".format(
npix)
num_local_pix = len(m)
assert num_local_pix == npix / mpi_comm.size
complete_m = pysm.read_map("pysm_2/dust_temp.fits", nside=8, field=0)
np.testing.assert_allclose(
m,
complete_m[num_local_pix * mpi_comm.rank:num_local_pix *
(mpi_comm.rank + 1)],
)
import numpy as np
import pysm
import pysm.units as u
from mpi4py import MPI
nside = 32
map_dist = pysm.MapDistribution(pixel_indices=None,
nside=nside,
mpi_comm=MPI.COMM_WORLD)
sky = pysm.Sky(nside=nside,
preset_strings=["d1", "s1", "f1", "a1"],
map_dist=map_dist)
m = sky.get_emission(freq=np.arange(50, 55) * u.GHz,
weights=np.array([0.1, 0.3, 0.5, 0.3, 0.1]))[0]
print(map_dist.mpi_comm.rank, m.shape, m.min(), m.max())
m_smoothed = pysm.apply_smoothing_and_coord_transform(m,
fwhm=1 * u.deg,
map_dist=map_dist)
print(map_dist.mpi_comm.rank, m_smoothed.shape, m_smoothed.min(),
m_smoothed.max())
def simulate_sky_signal(args, comm, data, focalplanes, signalname=None, mc=0):
""" Use PySM to simulate smoothed sky signal.
"""
log = Logger.get()
timer = Timer()
timer.start()
# Convolve a signal TOD from PySM
if comm.world_rank == 0:
log.info("Simulating sky signal with PySM")
map_dist = (None if comm is None else pysm.MapDistribution(
nside=args.nside, mpi_comm=comm.comm_rank))
pysm_component_objects = []
pysm_model = []
for model_tag in args.pysm_model.split(","):
if not model_tag.startswith("SO"):
pysm_model.append(model_tag)
else:
if so_pysm_models is None:
raise RuntimeError(
"{} requires so_pysm_models".format(model_tag))
if model_tag == "SO_x1_cib":
pysm_component_objects.append(
so_pysm_models.WebSkyCIB(
websky_version="0.3",
interpolation_kind="linear",
nside=args.nside,
map_dist=map_dist,
))
elif model_tag == "SO_x1_ksz":
pysm_component_objects.append(
so_pysm_models.WebSkySZ(
version="0.3",
nside=args.nside,
map_dist=map_dist,
sz_type="kinetic",
))
elif model_tag == "SO_x1_tsz":
pysm_component_objects.append(
so_pysm_models.WebSkySZ(
version="0.3",
nside=args.nside,
map_dist=map_dist,
sz_type="thermal",
))
elif model_tag.startswith("SO_x1_cmb"):
lensed = "unlensed" not in model_tag
include_solar_dipole = "solar" in model_tag
pysm_component_objects.append(
so_pysm_models.WebSkyCMBMap(
websky_version="0.3",
lensed=lensed,
include_solar_dipole=include_solar_dipole,
seed=1,
nside=args.nside,
map_dist=map_dist,
))
else:
if not model_tag.endswith("s") and args.nside > 512:
model_tag += "s"
pysm_component_objects.append(
so_pysm_models.get_so_models(model_tag,
args.nside,
map_dist=map_dist))
if signalname is None:
signalname = "pysmsignal"
op_sim_pysm = OpSimPySM(
data,
comm=comm.comm_rank,
out=signalname,
pysm_model=pysm_model,
pysm_component_objects=pysm_component_objects,
focalplanes=focalplanes,
apply_beam=args.pysm_apply_beam,
coord="G", # setting G doesn't perform any rotation
map_dist=map_dist,
)
assert args.coord in "CQ", "Input SO models are always in Equatorial coordinates"
op_sim_pysm.exec(data)
if comm.comm_world is not None:
comm.comm_world.barrier()
timer.stop()
if comm.world_rank == 0:
timer.report("PySM")
return signalname
