def build_cmb_2d(shape, wcs, cl_cmb, dtype=np.float32):
lmap = enmap.lmap(shape, wcs)
l = np.sum(lmap**2, 0)**0.5
cmb = enmap.samewcs(utils.interp(l, np.arange(cl_cmb.shape[-1]), cl_cmb),
l).astype(dtype)
# Rotate [TEB,EB] -> [TQU,TQU]. FIXME: not a perfect match
R = enmap.queb_rotmat(lmap, spin=2, inverse=True)
cmb[1:, :] = np.einsum("abyx,bcyx->acyx", R, cmb[1:, :])
cmb[:, 1:] = np.einsum("abyx,cbyx->acyx", cmb[:, 1:], R)
return cmb
def rotate_pol_power(shape,wcs,cov,iau=False,inverse=False):
"""Rotate a 2D power spectrum from TQU to TEB (inverse=False) or
back (inverse=True). cov is a (3,3,Ny,Nx) 2D power spectrum.
WARNING: This function is duplicated from orphics.maps to make
this module independent. Ideally, it should be implemented in
enlib.enmap.
"""
rot = np.zeros((3,3,cov.shape[-2],cov.shape[-1]))
rot[0,0,:,:] = 1
prot = enmap.queb_rotmat(enmap.lmap(shape,wcs), inverse=inverse, iau=iau)
rot[1:,1:,:,:] = prot
Rt = np.transpose(rot, (1,0,2,3))
tmp = np.einsum("ab...,bc...->ac...",rot,cov)
rp2d = np.einsum("ab...,bc...->ac...",tmp,Rt)
return rp2d
def __init__(self,shape,wcs,groups=None):
# Symbolic
self.l1x,self.l1y,self.l2x,self.l2y,self.l1,self.l2 = get_ells()
self.Lx,self.Ly,self.L = get_Ls()
if groups is None: groups = [self.Lx*self.Lx,self.Ly*self.Ly,self.Lx*self.Ly]
self._default_groups = groups
self.integrands = {}
self.ul1s = {}
self.ul2s = {}
self.ogroups = {}
self.ogroup_weights = {}
self.ogroup_symbols = {}
self.l1funcs = []
self.l2funcs = []
# Diagnostic
self.nfft = 0
self.nifft = 0
# Numeric
self.shape,self.wcs = shape,wcs
self.modlmap = enmap.modlmap(shape,wcs)
self.lymap,self.lxmap = enmap.lmap(shape,wcs)
self.pixarea = np.prod(enmap.pixshape(shape,wcs))
def integrate(shape,
wcs,
feed_dict,
expr,
xmask=None,
ymask=None,
cache=True,
validate=True,
groups=None,
pixel_units=False):
"""
Integrate an arbitrary expression after factorizing it.
Parameters
----------
shape : tuple
The shape of the array for the geometry of the footprint. Typically
(...,Ny,Nx) for Ny pixels in the y-direction and Nx in the x-direction.
wcs : :obj:`astropy.wcs.wcs.WCS`
The wcs object completing the specification of the geometry of the footprint.
feed_dict: dict
Mapping from names of custom symbols to numpy arrays.
expr: :obj:`sympy.core.symbol.Symbol`
A sympy expression containing recognized symbols (see docs)
xmask: (Ny,Nx) ndarray,optional
Fourier space 2D mask for the l1 part of the integral. Defaults to ones.
ymask: (Ny,Nx) ndarray, optional
Fourier space 2D mask for the l2 part of the integral. Defaults to ones.
cache: boolean, optional
Whether to store in memory and reuse repeated terms. Defaults to true.
validate: boolean,optional
Whether to check that the final expression and the original agree. Defaults to True.
groups: list,optional
Group all terms such that they have common factors of the provided list of
expressions to reduce the number of FFTs.
pixel_units: boolean,optional
Whether the input is in pixel units or not.
Returns
-------
result : (Ny,Nx) ndarray
The numerical result of the integration of the expression after factorization.
"""
# Geometry
modlmap = enmap.modlmap(shape, wcs)
lymap, lxmap = enmap.lmap(shape, wcs)
pixarea = np.prod(enmap.pixshape(shape, wcs))
feed_dict['L'] = modlmap
feed_dict['Ly'] = lymap
feed_dict['Lx'] = lxmap
shape = shape[-2:]
ones = np.ones(shape, dtype=np.float32)
val = 0.
if xmask is None: xmask = ones
if ymask is None: ymask = ones
# Expression
syms = expr.free_symbols
l1funcs = []
l2funcs = []
for sym in syms:
strsym = str(sym)
if strsym[-3:] == "_l1": l1funcs.append(sym)
elif strsym[-3:] == "_l2": l2funcs.append(sym)
integrands,ul1s,ul2s, \
ogroups,ogroup_weights, \
ogroup_symbols = factorize_2d_convolution_integral(expr,l1funcs=l1funcs,l2funcs=l2funcs,
validate=validate,groups=groups)
def _fft(x):
return fft(x + 0j)
def _ifft(x):
return ifft(x + 0j)
if cache:
cached_u1s = []
cached_u2s = []
for u1 in ul1s:
l12d = evaluate(u1, feed_dict) * ones
cached_u1s.append(_ifft(l12d * xmask))
for u2 in ul2s:
l22d = evaluate(u2, feed_dict) * ones
cached_u2s.append(_ifft(l22d * ymask))
# For each term, the index of which group it belongs to
def get_l1l2(term):
if cache:
ifft1 = cached_u1s[term['l1index']]
ifft2 = cached_u2s[term['l2index']]
else:
l12d = evaluate(term['l1'], feed_dict) * ones
ifft1 = _ifft(l12d * xmask)
l22d = evaluate(term['l2'], feed_dict) * ones
ifft2 = _ifft(l22d * ymask)
return ifft1, ifft2
if ogroups is None:
for i, term in enumerate(integrands):
ifft1, ifft2 = get_l1l2(term)
ot2d = evaluate(term['other'], feed_dict) * ones
ffft = _fft(ifft1 * ifft2)
val += ot2d * ffft
else:
vals = np.zeros((len(ogroup_symbols), ) + shape, dtype=np.float32) + 0j
for i, term in enumerate(integrands):
ifft1, ifft2 = get_l1l2(term)
gindex = ogroups[i]
vals[gindex, ...] += ifft1 * ifft2 * ogroup_weights[i]
for i, group in enumerate(ogroup_symbols):
ot2d = evaluate(ogroup_symbols[i], feed_dict) * ones
ffft = _fft(vals[i, ...])
val += ot2d * ffft
mul = 1 if pixel_units else 1. / pixarea
return val * mul
def kfilter_map(m, apo, kx_cut, ky_cut, unpixwin=True, legacy_steve=False):
r"""Apply a k-space filter on a map.
By default, we do not reproduce the output of Steve's code. We do offer
this functionality: set the optional flag `legacy_steve=True` to offset
the mask and the map by one pixel in each dimension.
You should filter both in temperature and polarization.
Parameters
----------
m : enmap
Input map which this function will filter.
apo : enmap
This map is a smooth tapering of the edges of the map, to be multiplied
into the map prior to filtering. The filtered map is divided by the
nonzero pixels of this map at the end. This is required because
maps of actual data are unlikely to be periodic, which will induce
ringing when one applies a k-space filter. To solve this, we taper the
edges of the map to zero prior to filtering.
See :py:func:`nawrapper.ps.rectangular_apodization`
kx_cut : float
We cut modes with wavenumber :math:`|k_x| < k_x^{\mathrm{cut}}`.
ky_cut : float
We cut modes with wavenumber :math:`|k_y| < k_y^{\mathrm{cut}}`.
unpixwin : bool
Correct for the CAR pixel window if True.
legacy_steve : bool
Use a slightly different filter if True, to reproduce Steve's pipeline.
Steve's k-space filter as of June 2019 had a bug where two k-space
modes (the most positive cut mode in x and most negative cut mode in y)
were not cut. This has a very small effect on the spectrum, but for
reproducibility purposes we offer this behavior. By default we do not
use this. To reproduce Steve's code behavior you should set `legacy_steve=True`.
Returns
-------
result : enmap
The map with the specified k-space filter applied.
"""
alm = enmap.fft(m * apo, normalize=True)
if unpixwin: # remove pixel window in Fourier space
wy, wx = enmap.calc_window(m.shape)
alm /= wy[:, np.newaxis]
alm /= wx[np.newaxis, :]
ly, lx = enmap.lmap(alm.shape, alm.wcs)
kfilter_x = np.abs(lx) >= kx_cut
kfilter_y = np.abs(ly) >= ky_cut
if legacy_steve: # Steve's kspace filter appears to do this
cut_x_k = np.unique(lx[(np.abs(lx) <= kx_cut)])
cut_y_k = np.unique(ly[(np.abs(ly) <= ky_cut)])
# keep most negative kx and most positive ky
kfilter_x[np.isclose(lx, cut_x_k[0])] = True
kfilter_y[np.isclose(ly, cut_y_k[-1])] = True
result = enmap.ifft(alm * kfilter_x * kfilter_y, normalize=True).real
result[apo > 0.0] = result[apo > 0.0] / apo[apo > 0.0]
return result