def _map2gclm(spin, lmax, _map, tht, wg, threads=1):
"""
input _map is (Ntheta,Nphi)-shaped.
Output vtm is (Ntheta,2 * lmax + 1) shaped, with v[t,lmax + m] = vtm.
This assume the _map samples uniformly the longitude [0,2pi) with Nphi = _map.shape[1] points.
tht and wg : colatitude and integration weights (e.g. GL zeroes and weights)
Here is a trick to save fac of 2 time in map2alm:
\int dcost f(t) sLlm(t) = int_(upperhalf) f+(t) sLlm + int_(upperhalf) f-(t) sLlm
where f+,f_ are the symmmetric / antisymmetric part, and the first term non zero only for l-m of the same partity and vice vers.
"""
vtm = shts.map2vtm(spin, lmax, _map, pfftwthreads=threads)
vlm = shts.vtm2vlm(spin, tht, vtm * np.outer(wg, np.ones(vtm.shape[1])))
return shts.util.vlm2alm(vlm)
def lens_glm_GLth_sym_timed(spin,
dlm,
glm,
lmax_target,
nband=16,
facres=0,
clm=None,
olm=None,
rotpol=True):
"""
Same as lens_alm but lens simultnously a North and South colatitude band,
to make profit of the symmetries of the spherical harmonics.
"""
assert spin >= 0, spin
times = {}
t0 = time.time()
tGL, wg = gauleg.get_xgwg(lmax_target + 2)
times['GL points and weights'] = time.time() - t0
target_nt = 3**1 * 2**(11 + facres
) # on one hemisphere (0.87 arcmin spacing)
th1s = np.arange(nband) * (np.pi * 0.5 / nband)
th2s = np.concatenate((th1s[1:], [np.pi * 0.5]))
Nt = target_nt / nband
tGL = np.arccos(tGL)
tGL = np.sort(tGL)
wg = wg[np.argsort(tGL)]
times['pol. rot.'] = 0.
times['vtm2defl2ang'] = 0.
times['vtmdefl'] = 0.
def coadd_times(tim):
for _k, _t in tim.iteritems():
if _k not in times:
times[_k] = _t
else:
times[_k] += _t
shapes = []
shapes_d = []
tGLNs = []
tGLSs = []
wgs = []
# Collects (Nt,Nphi) per band and prepare wisdom
wisdomhash = str(lmax_target) + '_' + str(nband) + '_' + str(facres +
1000) + '.npy'
assert os.path.exists(
os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/')
t0 = time.time()
print "building and caching FFTW wisdom, this might take a while"
for ib, th1, th2 in zip(range(nband), th1s, th2s):
Np = get_Nphi(th1,
th2,
facres=facres,
target_amin=60. * 90. /
target_nt) # same spacing as theta grid
Np_d = min(get_Nphi(th1, th2, target_amin=180. * 60. / lmax_target),
2 * lmax_target) #Equator point density
pixN, = np.where((tGL >= th1) & (tGL <= th2))
pixS, = np.where((tGL >= (np.pi - th2)) & (tGL <= (np.pi - th1)))
assert np.all(pixN[::-1] == len(tGL) - 1 -
pixS), 'symmetry of GL points'
shapes_d.append((len(pixN), Np_d))
shapes.append((Nt, Np))
tGLNs.append(tGL[pixN])
tGLSs.append(tGL[pixS])
wgs.append(np.concatenate([wg[pixN], wg[pixS]]))
print "BAND %s in %s. deflection (%s x %s) pts " % (ib, nband,
len(pixN), Np_d)
print " interpolation (%s x %s) pts " % (Nt, Np)
#==== For each block we have the following ffts:
# (Np_d) complex to complex (deflection map) BACKWARD (vtm2map)
# (Nt,Np) complex to complex (bicubic prefiltering) BACKWARD (vt2mmap) (4 threads)
# (Nt) complex to complex (bicubic prefiltering) FORWARD (vt2map)
# (Np_d) complex to complex FORWARD (map2vtm)
# Could rather do a try with FFTW_WISDOM_ONLY
if not os.path.exists(
os.path.dirname(os.path.realpath(__file__)) +
'/pyfftw_cache/' + wisdomhash):
a = pyfftw.empty_aligned(Np_d, dtype='complex128')
b = pyfftw.empty_aligned(Np_d, dtype='complex128')
fft = pyfftw.FFTW(a, b, direction='FFTW_FORWARD', threads=1)
fft = pyfftw.FFTW(a, b, direction='FFTW_BACKWARD', threads=1)
a = pyfftw.empty_aligned(Nt, dtype='complex128')
b = pyfftw.empty_aligned(Nt, dtype='complex128')
fft = pyfftw.FFTW(a, b, direction='FFTW_FORWARD', threads=1)
a = pyfftw.empty_aligned((Nt, Np), dtype='complex128')
b = pyfftw.empty_aligned((Nt, Np), dtype='complex128')
fft = pyfftw.FFTW(a,
b,
direction='FFTW_BACKWARD',
axes=(0, 1),
threads=4)
if not os.path.exists(
os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' +
wisdomhash):
np.save(
os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' +
wisdomhash, pyfftw.export_wisdom())
pyfftw.import_wisdom(
np.load(
os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' +
wisdomhash))
shts.PYFFTWFLAGS = ['FFTW_WISDOM_ONLY']
times['pyfftw_caches'] = time.time() - t0
print "Total number of interpo points: %s = %s ** 2" % (np.sum(
[np.prod(s)
for s in shapes]), np.sqrt(1. * np.sum([np.prod(s) for s in shapes])))
print "Total number of deflect points: %s = %s ** 2" % (np.sum([
np.prod(s) for s in shapes_d
]), np.sqrt(1. * np.sum([np.prod(s) for s in shapes_d])))
glmout = np.zeros(shts.util.lmax2nlm(lmax_target), dtype=np.complex)
clmout = np.zeros(shts.util.lmax2nlm(lmax_target), dtype=np.complex)
for ib, th1, th2 in zip(range(nband), th1s, th2s):
Nt_d, Np_d = shapes_d[ib]
Nt, Np = shapes[ib]
t0 = time.time()
vtm_def = shts.vlm2vtm_sym(1, _th2colat(tGLNs[ib]),
shts.util.alm2vlm(dlm, clm=olm))
times['vtmdefl'] += time.time() - t0
#==== gettting deflected positions
# NB: forward slice to keep theta -> pi - theta correspondance.
t0 = time.time()
dmapN = shts.vtm2map(1, vtm_def[:Nt_d, :], Np_d).flatten()
dmapS = shts.vtm2map(1, vtm_def[slice(Nt_d, 2 * Nt_d), :],
Np_d).flatten()
told = np.outer(tGLNs[ib], np.ones(Np_d)).flatten()
phiold = np.outer(np.ones(Nt_d),
np.arange(Np_d) * (2. * np.pi / Np_d)).flatten()
tnewN, phinewN = _buildangles((told, phiold), dmapN.real, dmapN.imag)
tnewS, phinewS = _buildangles(((np.pi - told)[::-1], phiold),
dmapS.real, dmapS.imag)
del vtm_def
times['vtm2defl2ang'] += time.time() - t0
#===== Adding a 10 pixels buffer for new angles to be safely inside interval.
# th1,th2 is mapped onto pi - th2,pi -th1 so we need to make sure to cover both buffers
matnewN = np.max(tnewN)
mitnewN = np.min(tnewN)
matnewS = np.max(tnewS)
mitnewS = np.min(tnewS)
buffN = 10 * (matnewN - mitnewN) / (Nt - 1) / (1. - 2. * 10. /
(Nt - 1))
buffS = 10 * (matnewS - mitnewS) / (Nt - 1) / (1. - 2. * 10. /
(Nt - 1))
_thup = min(np.pi - (matnewS + buffS), mitnewN - buffN)
_thdown = max(np.pi - (mitnewS - buffS), matnewN + buffN)
#==== these are the theta and limits. It is ok to go negative or > 180
dphi_patch = (2. * np.pi) / Np * max(np.sin(_thup), np.sin(_thdown))
dth_patch = (_thdown - _thup) / (Nt - 1)
print 'input t1,t2 %.3f %.3f in degrees' % (_thup / np.pi * 180,
_thdown / np.pi * 180.)
print 'North %.3f and South %.3f buffers in amin' % (
buffN / np.pi * 180 * 60, buffS / np.pi * 180. * 60.)
print "cell (theta,phi) in amin (%.3f,%.3f)" % (
dth_patch / np.pi * 60. * 180, dphi_patch / np.pi * 60. * 180)
if spin == 0:
lenN, lenS, tim = lens_band_sym_timed(glm,
_thup,
_thdown,
Nt, (tnewN, phinewN),
(tnewS, phinewS),
Nphi=Np)
ret = np.zeros((2 * Nt_d, Np_d), dtype=complex)
ret[:Nt_d, :] = lenN.reshape((Nt_d, Np_d))
ret[Nt_d:, :] = lenS.reshape((Nt_d, Np_d))
vtm = shts.map2vtm(spin, lmax_target, ret)
glmout -= shts.vtm2tlm_sym(
np.concatenate([tGLNs[ib], tGLSs[ib]]),
vtm * np.outer(wgs[ib], np.ones(vtm.shape[1])))
else:
assert 0, 'fix this'
lenNR, lenNI, lenSR, lenSI, tim = gclm2lensmap_symband_timed(
spin,
glm,
_thup,
_thdown,
Nt, (tnewN, phinewN), (tnewS, phinewS),
Nphi=Nphi,
clm=clm)
retN = (lenNR + 1j * lenNI).reshape((len(pixN), Np_d))
retS = (lenSR + 1j * lenSI).reshape((len(pixN), Np_d))
glm, clm = shts.util.vlm2alm(
shts.vtm2vlm(spin, tGL,
vtm * np.outer(wg, np.ones(vtm.shape[1]))))
t0 = time.time()
if rotpol and spin > 0:
ret[pixN, :] *= polrot(spin, retN.flatten(), tnewN, dmapN.real,
dmapN.imag)
ret[pixS, :] *= polrot(spin, retS.flatten(), tnewS, dmapS.real,
dmapS.imag)
times['pol. rot.'] += time.time() - t0
coadd_times(tim)
t0 = time.time()
print "STATS for lmax tlm %s lmax dlm %s" % (hp.Alm.getlmax(
glm.size), hp.Alm.getlmax(dlm.size))
tot = 0.
for _k, _t in times.iteritems():
print '%20s: %.2f' % (_k, _t)
tot += _t
print "%20s: %2.f sec." % ('tot', tot)
return glmout, clmout, ret