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