def mpi_smoothing(input_map, fwhm, map_dist):
import libsharp
beam = hp.gauss_beam(fwhm=fwhm.to(u.rad).value,
lmax=map_dist.smoothing_lmax,
pol=True)
input_map_I = input_map if input_map.ndim == 1 else input_map[0]
input_map_I_contig = np.ascontiguousarray(input_map_I.reshape(
(1, 1, -1))).astype(np.float64, copy=False)
alm_sharp_I = libsharp.analysis(
map_dist.libsharp_grid,
map_dist.libsharp_order,
input_map_I_contig,
spin=0,
comm=map_dist.mpi_comm,
)
map_dist.libsharp_order.almxfl(alm_sharp_I,
np.ascontiguousarray(beam[:, 0:1]))
out = libsharp.synthesis(
map_dist.libsharp_grid,
map_dist.libsharp_order,
alm_sharp_I,
spin=0,
comm=map_dist.mpi_comm,
)[0]
assert np.isnan(out).sum() == 0
if utils.has_polarization(input_map):
alm_sharp_P = libsharp.analysis(
map_dist.libsharp_grid,
map_dist.libsharp_order,
np.ascontiguousarray(input_map[1:3, :].reshape(
(1, 2, -1))).astype(np.float64, copy=False),
spin=2,
comm=map_dist.mpi_comm,
)
map_dist.libsharp_order.almxfl(alm_sharp_P,
np.ascontiguousarray(beam[:, (1, 2)]))
signal_map_P = libsharp.synthesis(
map_dist.libsharp_grid,
map_dist.libsharp_order,
alm_sharp_P,
spin=2,
comm=map_dist.mpi_comm,
)[0]
out = np.vstack((out, signal_map_P))
return out
def mpi_smoothing(self, sky, fwhm):
import libsharp
beam = hp.gauss_beam(fwhm=fwhm.to(u.rad).value,
lmax=self.smoothing_lmax,
pol=True)
sky_I = sky if sky.ndim == 1 else sky[0]
sky_I_contig = np.ascontiguousarray(sky_I.reshape(
(1, 1, -1))).astype(np.float64, copy=False)
alm_sharp_I = libsharp.analysis(
self.libsharp_grid,
self.libsharp_order,
sky_I_contig,
spin=0,
comm=self.mpi_comm,
)
self.libsharp_order.almxfl(alm_sharp_I, np.ascontiguousarray(beam[:, 0:1]))
out = libsharp.synthesis(
self.libsharp_grid,
self.libsharp_order,
alm_sharp_I,
spin=0,
comm=self.mpi_comm,
)[0]
assert np.isnan(out).sum() == 0
if utils.has_polarization(sky):
alm_sharp_P = libsharp.analysis(
self.libsharp_grid,
self.libsharp_order,
np.ascontiguousarray(sky[1:3, :].reshape(
(1, 2, -1))).astype(np.float64, copy=False),
spin=2,
comm=self.mpi_comm,
)
self.libsharp_order.almxfl(alm_sharp_P,
np.ascontiguousarray(beam[:, (1, 2)]))
signal_map_P = libsharp.synthesis(
self.libsharp_grid,
self.libsharp_order,
alm_sharp_P,
spin=2,
comm=self.mpi_comm,
)[0]
out = np.vstack((out, signal_map_P))
return out
def exec(self, data):
"""
Create the timestreams...
This loops over all observations and detectors and uses the pointing
matrix to ...
Args:
data (toast.Data): The distributed data.
"""
if libsharp is None:
raise RuntimeError('libsharp not available')
autotimer = timing.auto_timer(type(self).__name__)
has_pol = len(data[self.signal_map]) > 1
alm_sharp_I = libsharp.analysis(
self.grid,
self.order,
np.ascontiguousarray(data[self.signal_map][0].reshape((1, 1, -1))),
spin=0,
comm=self.comm)
self.order.almxfl(alm_sharp_I, np.ascontiguousarray(self.beam[:, 0:1]))
out = libsharp.synthesis(self.grid,
self.order,
alm_sharp_I,
spin=0,
comm=self.comm)[0]
if has_pol:
alm_sharp_P = libsharp.analysis(
self.grid,
self.order,
np.ascontiguousarray(data[self.signal_map][1:3, :].reshape(
(1, 2, -1))),
spin=2,
comm=self.comm)
self.order.almxfl(alm_sharp_P,
np.ascontiguousarray(self.beam[:, (1, 2)]))
signal_map_P = libsharp.synthesis(self.grid,
self.order,
alm_sharp_P,
spin=2,
comm=self.comm)[0]
out = np.vstack((out, signal_map_P))
data[self.out] = out
assert data[self.signal_map].shape == data[self.out].shape
def smooth_pol(self, fwhm_in, nside_in, maps_in):
if fwhm_in > .9 * self.fwhm and nside_in == self.nside:
return maps_in
if fwhm_in > .9 * self.fwhm and self.nside < nside_in:
# Simple ud_grade
if self.global_rank == 0 and self.verbose:
print('Downgrading Nside {} -> {}'
''.format(nside_in, self.nside),
flush=True)
maps_out = []
npix_in = hp.nside2npix(nside_in)
my_pix_in = self.get_my_pix(nside_in)
for (qmap, umap) in maps_in:
my_mapout = np.zeros(npix_in, dtype=np.float)
qmapout = np.zeros(npix_in, dtype=np.float)
umapout = np.zeros(npix_in, dtype=np.float)
my_mapout[my_pix_in] = qmap
self.comm.Allreduce(my_mapout, qmapout)
my_mapout[my_pix_in] = umap
self.comm.Allreduce(my_mapout, umapout)
del my_mapout
maps_out.append((hp.ud_grade(qmapout, self.nside)[self.my_pix],
hp.ud_grade(umapout,
self.nside)[self.my_pix]))
else:
# Full smoothing
lmax = self.optimal_lmax(fwhm_in, nside_in)
total_beam = self.total_beam(fwhm_in, lmax, pol=True)
if self.global_rank == 0 and self.verbose:
print('Smoothing {} -> {}. lmax = {}. Nside {} -> {}'
''.format(fwhm_in, self.fwhm, lmax, nside_in,
self.nside),
flush=True)
local_m = np.arange(self.rank,
lmax + 1,
self.ntask,
dtype=np.int32)
alminfo = packed_real_order(lmax, ms=local_m)
grid_in = self.get_grid(nside_in)
grid_out = self.get_grid(self.nside)
maps_out = []
for (local_map_Q, local_map_U) in maps_in:
map_P = np.ascontiguousarray(np.vstack(
[local_map_Q, local_map_U]).reshape((1, 2, -1)),
dtype=np.float64)
alm_P = analysis(grid_in,
alminfo,
map_P,
spin=2,
comm=self.comm)
if total_beam is not None:
alminfo.almxfl(alm_P, total_beam)
map_P = synthesis(grid_out,
alminfo,
alm_P,
spin=2,
comm=self.comm)[0]
maps_out.append(map_P)
return maps_out
def smooth(self, fwhm_in, nside_in, maps_in):
""" Smooth a distributed map and change the resolution.
"""
if fwhm_in > .9 * self.fwhm and nside_in == self.nside:
return maps_in
if fwhm_in > .9 * self.fwhm and self.nside < nside_in:
# Simple ud_grade
if self.global_rank == 0 and self.verbose:
print('Downgrading Nside {} -> {}'
''.format(nside_in, self.nside),
flush=True)
maps_out = []
npix_in = hp.nside2npix(nside_in)
my_pix_in = self.get_my_pix(nside_in)
for m in maps_in:
my_outmap = np.zeros(npix_in, dtype=np.float)
outmap = np.zeros(npix_in, dtype=np.float)
my_outmap[my_pix_in] = m
self.comm.Allreduce(my_outmap, outmap)
del my_outmap
maps_out.append(hp.ud_grade(outmap, self.nside)[self.my_pix])
else:
# Full smoothing
lmax = self.optimal_lmax(fwhm_in, nside_in)
total_beam = self.total_beam(fwhm_in, lmax, pol=False)
if self.global_rank == 0 and self.verbose:
print('Smoothing {} -> {}. lmax = {}. Nside {} -> {}'
''.format(fwhm_in, self.fwhm, lmax, nside_in,
self.nside),
flush=True)
local_m = np.arange(self.rank,
lmax + 1,
self.ntask,
dtype=np.int32)
alminfo = packed_real_order(lmax, ms=local_m)
grid_in = self.get_grid(nside_in)
grid_out = self.get_grid(self.nside)
maps_out = []
for local_map in maps_in:
map_I = np.ascontiguousarray(local_map.reshape([1, 1, -1]),
dtype=np.float64)
alm_I = analysis(grid_in,
alminfo,
map_I,
spin=0,
comm=self.comm)
if total_beam is not None:
alminfo.almxfl(alm_I, total_beam)
map_I = synthesis(grid_out,
alminfo,
alm_I,
spin=0,
comm=self.comm)[0][0]
maps_out.append(map_I)
return maps_out
def exec(self, data):
"""
Create the timestreams...
This loops over all observations and detectors and uses the pointing
matrix to ...
Args:
data (toast.Data): The distributed data.
"""
if libsharp is None:
raise RuntimeError('libsharp not available')
autotimer = timing.auto_timer(type(self).__name__)
has_pol = len(data[self.signal_map]) > 1
alm_sharp_I = libsharp.analysis(
self.grid, self.order,
np.ascontiguousarray(data[self.signal_map][0].reshape((1, 1, -1))),
spin=0, comm=self.comm)
self.order.almxfl(alm_sharp_I, np.ascontiguousarray(self.beam[:, 0:1]))
out = libsharp.synthesis(self.grid, self.order, alm_sharp_I, spin=0,
comm=self.comm)[0]
if has_pol:
alm_sharp_P = libsharp.analysis(
self.grid, self.order,
np.ascontiguousarray(
data[self.signal_map][1:3, :].reshape((1, 2, -1))),
spin=2, comm=self.comm)
self.order.almxfl(alm_sharp_P,
np.ascontiguousarray(self.beam[:, (1, 2)]))
signal_map_P = libsharp.synthesis(
self.grid, self.order, alm_sharp_P, spin=2, comm=self.comm)[0]
out = np.vstack((
out,
signal_map_P))
data[self.out] = out
assert data[self.signal_map].shape == data[self.out].shape
local_m_indices = np.arange(rank, lmax + 1, MPI.COMM_WORLD.Get_size(), dtype=np.int32)
if not mpi:
local_m_indices = None
order = libsharp.packed_real_order(lmax, ms=local_m_indices)
local_nl = order.local_size()
print("rank", rank, "local_nl", local_nl, "mval", order.mval())
mpi_comm = MPI.COMM_WORLD if mpi else None
# map2alm
# maps in libsharp are 3D, 2nd dimension is IQU, 3rd is pixel
alm_sharp_I = libsharp.analysis(grid, order,
np.ascontiguousarray(local_map[0].reshape((1, 1, -1))),
spin=0, comm=mpi_comm)
alm_sharp_P = libsharp.analysis(grid, order,
np.ascontiguousarray(local_map[1:].reshape((1, 2, -1))),
spin=2, comm=mpi_comm)
beam = hp.gauss_beam(fwhm=np.radians(fwhm_deg), lmax=lmax, pol=True)
print("Smooth")
# smooth in place (zonca implemented this function)
order.almxfl(alm_sharp_I, np.ascontiguousarray(beam[:, 0:1]))
order.almxfl(alm_sharp_P, np.ascontiguousarray(beam[:, (1, 2)]))
# alm2map
new_local_map_I = libsharp.synthesis(grid, order, alm_sharp_I, spin=0, comm=mpi_comm)
dtype=np.int32)
if not mpi:
local_m_indices = None
order = libsharp.packed_real_order(lmax, ms=local_m_indices)
local_nl = order.local_size()
print("rank", rank, "local_nl", local_nl, "mval", order.mval())
mpi_comm = MPI.COMM_WORLD if mpi else None
# map2alm
# maps in libsharp are 3D, 2nd dimension is IQU, 3rd is pixel
alm_sharp_I = libsharp.analysis(grid,
order,
np.ascontiguousarray(local_map[0].reshape(
(1, 1, -1))),
spin=0,
comm=mpi_comm)
alm_sharp_P = libsharp.analysis(grid,
order,
np.ascontiguousarray(local_map[1:].reshape(
(1, 2, -1))),
spin=2,
comm=mpi_comm)
beam = hp.gauss_beam(fwhm=np.radians(fwhm_deg), lmax=lmax, pol=True)
print("Smooth")
# smooth in place (zonca implemented this function)
order.almxfl(alm_sharp_I, np.ascontiguousarray(beam[:, 0:1]))
order.almxfl(alm_sharp_P, np.ascontiguousarray(beam[:, (1, 2)]))