From cb484df8edbf5f8af71efb68266e1a42dc6f7ae5 Mon Sep 17 00:00:00 2001 From: Mathew Madhavacheril Date: Mon, 23 Sep 2024 14:00:05 -0400 Subject: [PATCH 1/4] more flexibility in wavelets --- pixell/wavelets.py | 38 ++++++++++++++++++++++++++++++-------- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/pixell/wavelets.py b/pixell/wavelets.py index dd3db95..e68f236 100644 --- a/pixell/wavelets.py +++ b/pixell/wavelets.py @@ -303,12 +303,21 @@ def wcs(self): return self.uht.shape def geometry(self): return self.shape, self.wcs @property def nlevel(self): return len(self.geometries) - def map2wave(self, map, owave=None): + def map2wave(self, map, owave=None, fl = None, skip_coeffs=[]): """Transform from an enmap map[...,ny,nx] to a multimap of wavelet coefficients, which is effectively a group of enmaps with the same pre-dimensions but varying shape. If owave is provided, it should be a multimap with the right shape (compatible with the .geometries member of this class), and will be overwritten with the result. In any case the resulting wavelet coefficients are returned. + + A filter fl (either an array or a function; see curvedsky.almxfl) can be + provided that pre-filters the map in spherical harmonic space, e.g. to + convolve maps to a common beam. + + A list of the indices of wavelet coefficients to be skipped can be provided + in skip_coeffs. For these wavelet coefficients, a map of nans wil be + provided instead of performing the corresponding harmonic to real + space transform. """ filters, norms, lmids = self.filters, self.norms, self.lmids # Output geometry. Can't just use our existing one because it doesn't know about the @@ -317,18 +326,31 @@ def map2wave(self, map, owave=None): if owave is None: owave = multimap.zeros(geos, map.dtype) if self.uht.mode == "flat": fmap = enmap.fft(map, normalize=False) + if not(fl is None): + raise NotImplementedError("Pre-filtering not yet implemented for flat-sky wavelets.") for i, (shape, wcs) in enumerate(self.geometries): - fsmall = enmap.resample_fft(fmap, shape, norm=None, corner=True) - fsmall *= filters[i] / (norms[i]*fmap.npix) - owave.maps[i] = enmap.ifft(fsmall, normalize=False).real + if not(i in skip_coeffs): + fsmall = enmap.resample_fft(fmap, shape, norm=None, corner=True) + fsmall *= filters[i] / (norms[i]*fmap.npix) + owave.maps[i] = enmap.ifft(fsmall, normalize=False).real + else: + owave.maps[i] = enmap.empty(shape,wcs) + owave.maps[i][:] = np.nan + else: ainfo = curvedsky.alm_info(lmax=self.basis.lmax) alm = curvedsky.map2alm(map, ainfo=ainfo) + if not(fl is None): + alm = curvedsky.almxfl(alm,fl) for i, (shape, wcs) in enumerate(self.geometries): - smallinfo = curvedsky.alm_info(lmax=self.basis.lmaxs[i]) - asmall = curvedsky.transfer_alm(ainfo, alm, smallinfo) - smallinfo.lmul(asmall, filters[i]/norms[i], asmall) - curvedsky.alm2map(asmall, owave.maps[i]) + if not(i in skip_coeffs): + smallinfo = curvedsky.alm_info(lmax=self.basis.lmaxs[i]) + asmall = curvedsky.transfer_alm(ainfo, alm, smallinfo) + smallinfo.lmul(asmall, filters[i]/norms[i], asmall) + curvedsky.alm2map(asmall, owave.maps[i]) + else: + owave.maps[i] = enmap.empty(shape,wcs) + owave.maps[i][:] = np.nan return owave def wave2map(self, wave, omap=None): """Transform from the wavelet coefficients wave (multimap), to the corresponding enmap. From 056bf4c9ab79172eba1cea21dc77ac4a052d849f Mon Sep 17 00:00:00 2001 From: Mathew Madhavacheril Date: Mon, 23 Sep 2024 14:58:43 -0400 Subject: [PATCH 2/4] change arguments --- pixell/wavelets.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/pixell/wavelets.py b/pixell/wavelets.py index e68f236..3a45080 100644 --- a/pixell/wavelets.py +++ b/pixell/wavelets.py @@ -303,7 +303,7 @@ def wcs(self): return self.uht.shape def geometry(self): return self.shape, self.wcs @property def nlevel(self): return len(self.geometries) - def map2wave(self, map, owave=None, fl = None, skip_coeffs=[]): + def map2wave(self, map, owave=None, fl = None, scales=None, fill_value=None): """Transform from an enmap map[...,ny,nx] to a multimap of wavelet coefficients, which is effectively a group of enmaps with the same pre-dimensions but varying shape. If owave is provided, it should be a multimap with the right shape (compatible with @@ -314,10 +314,11 @@ def map2wave(self, map, owave=None, fl = None, skip_coeffs=[]): provided that pre-filters the map in spherical harmonic space, e.g. to convolve maps to a common beam. - A list of the indices of wavelet coefficients to be skipped can be provided - in skip_coeffs. For these wavelet coefficients, a map of nans wil be - provided instead of performing the corresponding harmonic to real - space transform. + A list of the indices of wavelet coefficients to be calculated can be provided + in scales; None defaults to all scales. For wavelet coefficients that are not + calculated, a map of zeros wil be provided instead of performing the corresponding + harmonic to real space transform. Alternatively, a fill_value different from zero + can be specified. """ filters, norms, lmids = self.filters, self.norms, self.lmids # Output geometry. Can't just use our existing one because it doesn't know about the @@ -329,13 +330,13 @@ def map2wave(self, map, owave=None, fl = None, skip_coeffs=[]): if not(fl is None): raise NotImplementedError("Pre-filtering not yet implemented for flat-sky wavelets.") for i, (shape, wcs) in enumerate(self.geometries): - if not(i in skip_coeffs): + if i in scales: fsmall = enmap.resample_fft(fmap, shape, norm=None, corner=True) fsmall *= filters[i] / (norms[i]*fmap.npix) owave.maps[i] = enmap.ifft(fsmall, normalize=False).real else: - owave.maps[i] = enmap.empty(shape,wcs) - owave.maps[i][:] = np.nan + owave.maps[i] = enmap.zeros(shape,wcs) + if fill_value is not None: owave.maps[i][:] = np.nan else: ainfo = curvedsky.alm_info(lmax=self.basis.lmax) @@ -343,14 +344,14 @@ def map2wave(self, map, owave=None, fl = None, skip_coeffs=[]): if not(fl is None): alm = curvedsky.almxfl(alm,fl) for i, (shape, wcs) in enumerate(self.geometries): - if not(i in skip_coeffs): + if i in scales: smallinfo = curvedsky.alm_info(lmax=self.basis.lmaxs[i]) asmall = curvedsky.transfer_alm(ainfo, alm, smallinfo) smallinfo.lmul(asmall, filters[i]/norms[i], asmall) curvedsky.alm2map(asmall, owave.maps[i]) else: - owave.maps[i] = enmap.empty(shape,wcs) - owave.maps[i][:] = np.nan + owave.maps[i] = enmap.zeros(shape,wcs) + if fill_value is not None: owave.maps[i][:] = fill_value return owave def wave2map(self, wave, omap=None): """Transform from the wavelet coefficients wave (multimap), to the corresponding enmap. From 1ad28432a413c6f9dbe7b858118886162834dc36 Mon Sep 17 00:00:00 2001 From: Mathew Madhavacheril Date: Mon, 23 Sep 2024 16:54:07 -0400 Subject: [PATCH 3/4] lmins and lmaxs for cosine needlets --- pixell/wavelets.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pixell/wavelets.py b/pixell/wavelets.py index 3a45080..adf18d8 100644 --- a/pixell/wavelets.py +++ b/pixell/wavelets.py @@ -137,6 +137,7 @@ def __init__(self, lpeaks): """ self.lpeaks = lpeaks self.lmaxs = np.append(self.lpeaks[1:],self.lpeaks[-1]) + self.lmins = np.append(self.lpeaks[0],self.lpeaks[:-1]) self.lmin = self.lpeaks[0] self.lmax = self.lpeaks[-1] @property From c693e5fc90c355476850f637e152493e7e384832 Mon Sep 17 00:00:00 2001 From: Mathew Madhavacheril Date: Tue, 24 Sep 2024 00:01:31 -0400 Subject: [PATCH 4/4] style change --- pixell/wavelets.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pixell/wavelets.py b/pixell/wavelets.py index adf18d8..baa61ac 100644 --- a/pixell/wavelets.py +++ b/pixell/wavelets.py @@ -328,7 +328,7 @@ def map2wave(self, map, owave=None, fl = None, scales=None, fill_value=None): if owave is None: owave = multimap.zeros(geos, map.dtype) if self.uht.mode == "flat": fmap = enmap.fft(map, normalize=False) - if not(fl is None): + if fl is not None: raise NotImplementedError("Pre-filtering not yet implemented for flat-sky wavelets.") for i, (shape, wcs) in enumerate(self.geometries): if i in scales: @@ -342,7 +342,7 @@ def map2wave(self, map, owave=None, fl = None, scales=None, fill_value=None): else: ainfo = curvedsky.alm_info(lmax=self.basis.lmax) alm = curvedsky.map2alm(map, ainfo=ainfo) - if not(fl is None): + if fl is not None: alm = curvedsky.almxfl(alm,fl) for i, (shape, wcs) in enumerate(self.geometries): if i in scales: