Skip to content

Commit 6a37dbf

Browse files
msyriacmsyriac
committed
Merge branch 'master' of github.com:simonsobs/pixell
* 'master' of github.com:simonsobs/pixell: More flexibility in wavelets (#270) enmap.downgrade flexibility and kappa to phi (#268)
2 parents 2eb78ef + 0d500c3 commit 6a37dbf

File tree

3 files changed

+114
-46
lines changed

3 files changed

+114
-46
lines changed

‎pixell/lensing.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,26 @@ def phi_to_kappa(phi_alm,phi_ainfo=None):
9191
kappa_alm: The filtered alms phi_alm * l * (l+1) / 2
9292
"""
9393
from . import curvedsky
94-
return curvedsky.almxfl(alm=phi_alm,lfilter=lambda x: x*(x+1)/2,ainfo=phi_ainfo)
94+
return curvedsky.almxfl(alm=phi_alm,lfilter=lambda l: l*(l+1)/2,ainfo=phi_ainfo)
95+
96+
def kappa_to_phi(kappa_alm,kappa_ainfo=None):
97+
"""Convert lensing convergence alms kappa_alm to
98+
lensing potential alms phi_alm, i.e.
99+
kappa_alm / ( l * (l+1) / 2 )
100+
101+
Args:
102+
kappa_alm: (...,N) ndarray of spherical harmonic alms of lensing convergence
103+
kappa_ainfo: If ainfo is provided, it is an alm_info describing the layout
104+
of the input alm. Otherwise it will be inferred from the alm itself.
105+
106+
Returns:
107+
phi_alm: The filtered alms kappa_alm / ( l * (l+1) / 2 )
108+
"""
109+
from . import curvedsky
110+
with utils.nowarn():
111+
oalm = curvedsky.almxfl(alm=kappa_alm,lfilter=lambda l: 1./(l*(l+1)/2) ,ainfo=kappa_ainfo)
112+
return utils.remove_nan(oalm)
113+
95114

96115
def kappa_to_phi(kappa_alm,kappa_ainfo=None):
97116
"""Convert lensing convergence alms kappa_alm to

‎pixell/wavelets.py

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ def __init__(self, lpeaks):
137137
"""
138138
self.lpeaks = lpeaks
139139
self.lmaxs = np.append(self.lpeaks[1:],self.lpeaks[-1])
140+
self.lmins = np.append(self.lpeaks[0],self.lpeaks[:-1])
140141
self.lmin = self.lpeaks[0]
141142
self.lmax = self.lpeaks[-1]
142143
@property
@@ -303,12 +304,22 @@ def wcs(self): return self.uht.shape
303304
def geometry(self): return self.shape, self.wcs
304305
@property
305306
def nlevel(self): return len(self.geometries)
306-
def map2wave(self, map, owave=None):
307+
def map2wave(self, map, owave=None, fl = None, scales=None, fill_value=None):
307308
"""Transform from an enmap map[...,ny,nx] to a multimap of wavelet coefficients,
308309
which is effectively a group of enmaps with the same pre-dimensions but varying shape.
309310
If owave is provided, it should be a multimap with the right shape (compatible with
310311
the .geometries member of this class), and will be overwritten with the result. In
311312
any case the resulting wavelet coefficients are returned.
313+
314+
A filter fl (either an array or a function; see curvedsky.almxfl) can be
315+
provided that pre-filters the map in spherical harmonic space, e.g. to
316+
convolve maps to a common beam.
317+
318+
A list of the indices of wavelet coefficients to be calculated can be provided
319+
in scales; None defaults to all scales. For wavelet coefficients that are not
320+
calculated, a map of zeros wil be provided instead of performing the corresponding
321+
harmonic to real space transform. Alternatively, a fill_value different from zero
322+
can be specified.
312323
"""
313324
filters, norms, lmids = self.filters, self.norms, self.lmids
314325
# Output geometry. Can't just use our existing one because it doesn't know about the
@@ -317,18 +328,31 @@ def map2wave(self, map, owave=None):
317328
if owave is None: owave = multimap.zeros(geos, map.dtype)
318329
if self.uht.mode == "flat":
319330
fmap = enmap.fft(map, normalize=False)
331+
if fl is not None:
332+
raise NotImplementedError("Pre-filtering not yet implemented for flat-sky wavelets.")
320333
for i, (shape, wcs) in enumerate(self.geometries):
321-
fsmall = enmap.resample_fft(fmap, shape, norm=None, corner=True)
322-
fsmall *= filters[i] / (norms[i]*fmap.npix)
323-
owave.maps[i] = enmap.ifft(fsmall, normalize=False).real
334+
if i in scales:
335+
fsmall = enmap.resample_fft(fmap, shape, norm=None, corner=True)
336+
fsmall *= filters[i] / (norms[i]*fmap.npix)
337+
owave.maps[i] = enmap.ifft(fsmall, normalize=False).real
338+
else:
339+
owave.maps[i] = enmap.zeros(shape,wcs)
340+
if fill_value is not None: owave.maps[i][:] = np.nan
341+
324342
else:
325343
ainfo = curvedsky.alm_info(lmax=self.basis.lmax)
326344
alm = curvedsky.map2alm(map, ainfo=ainfo)
345+
if fl is not None:
346+
alm = curvedsky.almxfl(alm,fl)
327347
for i, (shape, wcs) in enumerate(self.geometries):
328-
smallinfo = curvedsky.alm_info(lmax=self.basis.lmaxs[i])
329-
asmall = curvedsky.transfer_alm(ainfo, alm, smallinfo)
330-
smallinfo.lmul(asmall, filters[i]/norms[i], asmall)
331-
curvedsky.alm2map(asmall, owave.maps[i])
348+
if i in scales:
349+
smallinfo = curvedsky.alm_info(lmax=self.basis.lmaxs[i])
350+
asmall = curvedsky.transfer_alm(ainfo, alm, smallinfo)
351+
smallinfo.lmul(asmall, filters[i]/norms[i], asmall)
352+
curvedsky.alm2map(asmall, owave.maps[i])
353+
else:
354+
owave.maps[i] = enmap.zeros(shape,wcs)
355+
if fill_value is not None: owave.maps[i][:] = fill_value
332356
return owave
333357
def wave2map(self, wave, omap=None):
334358
"""Transform from the wavelet coefficients wave (multimap), to the corresponding enmap.

‎tests/test_pixell.py

Lines changed: 62 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -735,6 +735,30 @@ def test_prepare_alm_mmax(self):
735735
self.assertEqual(ainfo_out.mmax, ainfo_in.mmax)
736736
self.assertEqual(ainfo_out.nelem, ainfo_in.nelem)
737737

738+
739+
def test_lens_alms(self):
740+
# We generate phi alms and convert them to kappa and back
741+
lmax = 100
742+
ps = np.zeros(lmax+1)
743+
ls = np.arange(lmax+1)
744+
ps[ls>=2] = 1./ls[ls>=2]
745+
phi_alm = curvedsky.rand_alm(ps,lmax=lmax)
746+
kappa_alm = lensing.phi_to_kappa(phi_alm)
747+
phi_alm2 = lensing.kappa_to_phi(kappa_alm)
748+
np.testing.assert_array_almost_equal(phi_alm, phi_alm2)
749+
750+
def test_downgrade(self):
751+
shape,wcs = enmap.geometry(pos=(0,0),shape=(100,100),res=0.01)
752+
imap = enmap.ones(shape,wcs)
753+
for dfact in [None,1]:
754+
omap = enmap.downgrade(imap,dfact)
755+
np.testing.assert_equal(imap,omap)
756+
757+
dfact = 2
758+
omap = enmap.downgrade(imap,dfact,op=np.sum)
759+
np.testing.assert_equal(omap,np.ones(enmap.scale_geometry(shape,wcs,1./dfact)[0])*4)
760+
761+
738762
def test_almxfl(self):
739763
# We try to filter alms of shape (nalms,) and (ncomp,nalms) with
740764
# a filter of shape (nells,)
@@ -901,43 +925,44 @@ def test_alm2map_healpix_roundtrip(self):
901925
alm_out = curvedsky.map2alm_healpix(omap, spin=spin, ainfo=ainfo, niter=niter)
902926
np.testing.assert_array_almost_equal(alm_out, alm)
903927

904-
# Disabled for now because the version of ducc currently on pypi
905-
# has an adjointness bug. It's fixed in the ducc git repo.
906-
#def test_adjointness(self):
907-
# # This tests if alm2map_adjoint is the adjoint of alm2map,
908-
# # and if map2alm_adjoint is the adjoint of map2alm.
909-
# # (This doesn't test if they're correct, just that they're
910-
# # consistent with each other). This test is a bit slow, taking
911-
# # 5 s or so. It would be much faster if we dropped the ncomp=3 case.
912-
# for dtype in [np.float32, np.float64]:
913-
# for ncomp in [1,3]:
914-
# # Define our geometries
915-
# geos = []
916-
# res = 30*utils.degree
917-
# shape, wcs = enmap.fullsky_geometry(res=res, variant="fejer1")
918-
# geos.append(("fullsky_fejer1", shape, wcs))
919-
920-
# shape, wcs = enmap.fullsky_geometry(res=res, variant="cc")
921-
# geos.append(("fullsky_cc", shape, wcs))
922-
# lmax = shape[-2]-2
923-
924-
# shape, wcs = enmap.Geometry(shape, wcs)[3:-3,3:-3]
925-
# geos.append(("patch_cc", shape, wcs))
926-
927-
# wcs = wcs.deepcopy()
928-
# wcs.wcs.crpix += 0.123
929-
# geos.append(("patch_gen_cyl", shape, wcs))
930-
931-
# shape, wcs = enmap.geometry(np.array([[-45,45],[45,-45]])*utils.degree, res=res, proj="tan")
932-
# geos.append(("patch_tan", shape, wcs))
933-
934-
# for gi, (name, shape, wcs) in enumerate(geos):
935-
# mat1 = alm_bash(curvedsky.alm2map, shape, wcs, ncomp, lmax, dtype)
936-
# mat2 = map_bash(curvedsky.alm2map_adjoint, shape, wcs, ncomp, lmax, dtype)
937-
# np.testing.assert_array_almost_equal(mat1, mat2)
938-
# mat1 = map_bash(curvedsky.map2alm, shape, wcs, ncomp, lmax, dtype)
939-
# mat2 = alm_bash(curvedsky.map2alm_adjoint, shape, wcs, ncomp, lmax, dtype)
940-
# np.testing.assert_array_almost_equal(mat1, mat2)
928+
# MM: Re-enabled 09/17/2024
929+
# --Disabled for now because the version of ducc currently on pypi
930+
# has an adjointness bug. It's fixed in the ducc git repo.--
931+
def test_adjointness(self):
932+
# This tests if alm2map_adjoint is the adjoint of alm2map,
933+
# and if map2alm_adjoint is the adjoint of map2alm.
934+
# (This doesn't test if they're correct, just that they're
935+
# consistent with each other). This test is a bit slow, taking
936+
# 5 s or so. It would be much faster if we dropped the ncomp=3 case.
937+
for dtype in [np.float32, np.float64]:
938+
for ncomp in [1,3]:
939+
# Define our geometries
940+
geos = []
941+
res = 30*utils.degree
942+
shape, wcs = enmap.fullsky_geometry(res=res, variant="fejer1")
943+
geos.append(("fullsky_fejer1", shape, wcs))
944+
945+
shape, wcs = enmap.fullsky_geometry(res=res, variant="cc")
946+
geos.append(("fullsky_cc", shape, wcs))
947+
lmax = shape[-2]-2
948+
949+
shape, wcs = enmap.Geometry(shape, wcs)[3:-3,3:-3]
950+
geos.append(("patch_cc", shape, wcs))
951+
952+
wcs = wcs.deepcopy()
953+
wcs.wcs.crpix += 0.123
954+
geos.append(("patch_gen_cyl", shape, wcs))
955+
956+
shape, wcs = enmap.geometry(np.array([[-45,45],[45,-45]])*utils.degree, res=res, proj="tan")
957+
geos.append(("patch_tan", shape, wcs))
958+
959+
for gi, (name, shape, wcs) in enumerate(geos):
960+
mat1 = alm_bash(curvedsky.alm2map, shape, wcs, ncomp, lmax, dtype)
961+
mat2 = map_bash(curvedsky.alm2map_adjoint, shape, wcs, ncomp, lmax, dtype)
962+
np.testing.assert_array_almost_equal(mat1, mat2)
963+
mat1 = map_bash(curvedsky.map2alm, shape, wcs, ncomp, lmax, dtype)
964+
mat2 = alm_bash(curvedsky.map2alm_adjoint, shape, wcs, ncomp, lmax, dtype)
965+
np.testing.assert_array_almost_equal(mat1, mat2)
941966

942967

943968
#def test_sharp_alm2map_der1(self):

0 commit comments

Comments
 (0)