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