def test_ifft_input_shape(self):
# Tests ifft for various shapes and choices of axes.
# 1D IFFT over last axis for 3d array.
fsignal = np.ones((1, 2, 5), dtype=np.complex128)
fsignal[0, 1, :] = 10.
out_exp = np.zeros((1, 2, 5))
out_exp[0, 0, 0] = 5
out_exp[0, 1, 0] = 50
out = fft.ifft(fsignal)
np.testing.assert_allclose(out, out_exp, atol=1e-12)
self.assertTrue(out.flags['C_CONTIGUOUS'])
# 1D IFFT over middle axis for 3d array.
fsignal = np.ones((1, 5, 2), dtype=np.complex128)
fsignal[0, :, 1] = 10.
out_exp = np.zeros((1, 5, 2))
out_exp[0, 0, 0] = 5
out_exp[0, 0, 1] = 50
out = fft.ifft(fsignal, axes=[-2])
np.testing.assert_allclose(out, out_exp, atol=1e-12)
self.assertTrue(out.flags['C_CONTIGUOUS'])
# 2D IFFT over last 2 axes of 4d array.
fsignal = np.ones((1, 2, 5, 10), dtype=np.complex128)
fsignal[0, 1, :] = 10.
out_exp = np.zeros((1, 2, 5, 10))
out_exp[0, 0, 0, 0] = 50
out_exp[0, 1, 0, 0] = 500
out = fft.ifft(fsignal, axes=[-2, -1])
np.testing.assert_allclose(out, out_exp, atol=1e-12)
self.assertTrue(out.flags['C_CONTIGUOUS'])
# 2D IFFT over last 2 axes of 4d non-contiguous array.
fsignal = np.ones((1, 2, 5, 10), dtype=np.complex128)
fsignal[0, 1, :] = 10.
tod = np.zeros((5, 10, 1, 2),
dtype=np.complex128).transpose(2, 3, 0, 1)
out_exp = np.zeros_like(tod)
out_exp[0, 0, 0, 0] = 50
out_exp[0, 1, 0, 0] = 500
out = fft.ifft(fsignal, tod=tod, axes=[-2, -1])
self.assertTrue(np.shares_memory(tod, out))
np.testing.assert_allclose(out, out_exp, atol=1e-12)
self.assertFalse(out.flags['C_CONTIGUOUS'])
# 2D IFFT over middle 2 axes of 4d array.
fsignal = np.ones((1, 5, 10, 2), dtype=np.complex128)
fsignal[0, :, :, 1] = 10.
out_exp = np.zeros((1, 5, 10, 2))
out_exp[0, 0, 0, 0] = 50
out_exp[0, 0, 0, 1] = 500
out = fft.ifft(fsignal, axes=[-3, -2])
np.testing.assert_allclose(out, out_exp, atol=1e-12)
self.assertTrue(out.flags['C_CONTIGUOUS'])
def __signal_postprocessing__(self,
patch,
signal_idx,
alm_patch,
save_map,
oshape,
owcs,
apply_window=True):
signal = self.get_template(patch, shape=oshape, wcs=owcs)
signal = signal if len(alm_patch.shape) > 1 else signal[0, ...]
curvedsky.alm2map(alm_patch, signal, spin=[0, 2], verbose=True)
if apply_window:
print('apply window')
axes = [-2, -1]
for idx in range(signal.shape[0]):
kmap = pfft.fft(signal[idx], axes=axes)
wy, wx = enmap.calc_window(kmap.shape)
wind = wy[:, None]**1 * wx[None, :]**1
kmap *= wind
signal[idx] = (pfft.ifft(kmap, axes=axes, normalize=True)).real
del kmap
if save_map:
self.signals[signal_idx] = signal.copy()
self.manage_cache(self.signals, self.max_cached)
return signal
def process(kouts, name="default", ellmax=None, y=False, y_ellmin=400):
ellmax = lmax if ellmax is None else ellmax
ksilc = enmap.zeros((Ny, Nx), wcs, dtype=np.complex128).reshape(-1)
ksilc[modlmap.reshape(-1) < lmax] = np.nan_to_num(kouts.copy())
ksilc = enmap.enmap(ksilc.reshape((Ny, Nx)), wcs)
ksilc[modlmap > ellmax] = 0
if y: ksilc[modlmap < y_ellmin] = 0
msilc = np.nan_to_num(
fft.ifft(ksilc, axes=[-2, -1], normalize=True).real * bmask)
enmap.write_map(proot + "outmap_%s.fits" % name, enmap.enmap(msilc, wcs))
p2d = fc.f2power(ksilc, ksilc)
bin_edges = np.arange(100, 3000, 40)
binner = stats.bin2D(modlmap, bin_edges)
cents, p1d = binner.bin(p2d)
try:
# io.plot_img(np.log10(np.fft.fftshift(p2d)),proot+"ksilc_%s.png" % name,aspect='auto')
# io.plot_img(msilc,proot+"msilc_%s.png" % name)
# io.plot_img(msilc,proot+"lmsilc_%s.png" % name,lim=300)
tmask = maps.mask_kspace(shape,
wcs,
lmin=300,
lmax=5000 if not (y) else 1500)
fmap = maps.filter_map(msilc, tmask) * bmask
io.plot_img(fmap, proot + "hmsilc_%s.png" % name, high_res=True)
except:
pass
return cents, p1d
Ldl2 = (Lx * l2x + Ly * l2y) # \vec{L}.\vec{l}_2
Lxl1 = (Ly * l1x - Lx * l1y) # \vec{L} x \vec{l}_1 for curl
Lxl2 = (Ly * l2x - Lx * l2y) # \vec{L} x \vec{l}_2 for curl
# More built-in special symbols
# cos(2\theta_{12}), sin(2\theta_{12}) for polarization
cos2t12, sin2t12 = _helpers.substitute_trig(l1x, l1y, l2x, l2y, l1, l2)
# Custom symbol wrapper
def e(symbol):
# TODO: add exceptions if symbol doesn't correspond to key structure
return Symbol(symbol)
ifft = lambda x: efft.ifft(x, axes=[-2, -1], normalize=True)
fft = lambda x: efft.fft(x, axes=[-2, -1])
evaluate = _helpers.evaluate
def factorize_2d_convolution_integral(expr,
l1funcs=None,
l2funcs=None,
groups=None,
validate=True):
"""Reduce a sympy expression of variables l1x,l1y,l2x,l2y,l1,l2 into a sum of
products of factors that depend only on vec(l1) and vec(l2) and neither, each. If the expression
appeared as the integrand in an integral over vec(l1), where
vec(l2) = vec(L) - vec(l1) then this reduction allows one to evaluate the
integral as a function of vec(L) using FFTs instead of as a convolution.
def getActpolCmbFgSim(beamfileDict,
shape, wcs,
iterationNum, cmbDir, freqs, psa,
cmbSet = 0, \
doBeam = True, applyWindow = True, verbose = True, cmbMaptype = 'LensedCMB',
foregroundSeed = 0, simType = 'cmb', foregroundPowerFile = None, applyModulation = True):
nTQUs = len('TQU')
firstTime = True
output = enmap.empty((
len(freqs),
nTQUs,
) + shape[-2:], wcs)
if simType == 'cmb':
filename = cmbDir + "/fullsky%s_alm_set%02d_%05d.fits" % (
cmbMaptype, cmbSet, iterationNum)
if verbose:
print('getActpolCmbFgSim(): loading CMB a_lms from %s' % filename)
import healpy
almTebFullskyOnecopy = np.complex128(
healpy.fitsfunc.read_alm(filename, hdu=(1, 2, 3)))
#Now tile the same for all the frequencies requested (i.e. CMB is same at all frequencies).
#The beam convolution happens below.
almTebFullsky = np.tile(almTebFullskyOnecopy, (len(freqs), 1, 1))
if verbose:
print('getActpolCmbFgSim(): done')
elif simType == 'foregrounds':
outputFreqs = ['f090', 'f150']
foregroundPowers \
= powspec.read_spectrum(os.path.join(os.path.dirname(os.path.abspath(__file__)),
'../data/',
foregroundPowerFile),
ncol = 3,
expand = 'row')
if verbose:
print('getActpolCmbFgSim(): getting foreground a_lms')
almTFullsky90and150 = curvedsky.rand_alm(foregroundPowers,
seed=foregroundSeed)
almTebFullsky = np.zeros((
len(freqs),
nTQUs,
) + (len(almTFullsky90and150[-1]), ),
dtype=np.complex128)
for fi, freq in enumerate(freqs):
if freq in outputFreqs:
almTebFullsky[fi, 'TQU'.index('T'), :] \
= almTFullsky90and150[outputFreqs.index(freq), :]
#Convolve with beam on full sky
for fi, freq in enumerate(freqs):
if doBeam:
beamFile = beamfileDict[psa + '_' + freq]
if verbose:
print('getActpolCmbFgSim(): applying beam from %s' % beamFile)
beamData = (np.loadtxt(beamFile))[:, 1]
else:
if verbose:
print('getActpolCmbFgSim(): not convolving with beam')
beamData = np.repeat(1., almTebFullsky.shape[-1])
import healpy
#couldn't quickly figure out how to vectorize this so loop from 0 to 2.
for tqui in range(nTQUs):
almTebFullsky[fi, tqui] = healpy.sphtfunc.almxfl(
almTebFullsky[fi, tqui].copy(), beamData)
#These lines stolen from curvedsky.rand_map
#doing all freqs at once gives error:
#sharp.pyx in sharp.execute_dp (cython/sharp.c:12118)()
#ValueError: ndarray is not C-contiguous
#so loop over all freqs once for now.
curvedsky.alm2map(almTebFullsky, output, spin=[0, 2], verbose=True)
# outputThisfreq = enmap.empty(( nTQUs,)+shape[-2:], wcs)
# curvedsky.alm2map(almTebFullsky[fi,:,:], outputThisfreq, spin = [0,2], verbose = True)
# output[fi,...] = outputThisfreq
# curvedsky.alm2map(almTebFullsky[fi,:,:], output[fi,:,:,:], spin = [0,2], verbose = True)
if applyModulation:
from pixell import aberration
for fi, freq in enumerate(freqs):
print('applying modulation for frequency %s' % freq)
output, A = aberration.boost_map(output,
aberration.dir_equ,
aberration.beta,
return_modulation=True,
aberrate=False,
freq=freqStrToValGhz[freq] * 1e9)
if applyWindow:
from pixell import fft
#The axes along which to FFT
axes = [-2, -1]
if verbose:
print('getActpolCmbFgSim(): applying pixel window function')
nfreq = len(freqs)
for idx in range(nfreq):
fd = fft.fft(output[idx], axes=axes)
wy, wx = enmap.calc_window(fd.shape)
twoDWindow = wy[:, None]**1 * wx[None, :]**1
#Careful, this is quietly multiplying an array with shape [N_freq, N_TQU, N_y, N_x] with one of shape [N_y, N_x]
fd *= twoDWindow
output[idx] = (fft.ifft(fd, axes=axes, normalize=True)).real
del fd
if verbose:
print('getActpolCmbFgSim(): done')
return enmap.ndmap(output, wcs)