diff --git a/delensalot/core/iterator/cs_iterator_lognormal.py b/delensalot/core/iterator/cs_iterator_lognormal.py new file mode 100644 index 00000000..015485b3 --- /dev/null +++ b/delensalot/core/iterator/cs_iterator_lognormal.py @@ -0,0 +1,375 @@ +"""Module for curved-sky iterative lensing estimation + + Version revised on March 2023 + + Among the changes: + * delensalot'ed this with great improvements in execution time + * novel and more stable way of calculating the delfection angles and inverses + * optionally change main variable from plm to klm or dlm with expected better behavior ? + * rid of alm2rlm which was just wasting a little bit of time and loads of memory + * abstracted bfgs with cacher and dot_op + + + + #FIXME: loading of total gradient seems mixed up with loading of quadratic gradient... + #TODO: make plm0 possibly a path? + #FIXME: Chh = 0 not resulting in 0 estimate +""" + +import os +from os.path import join as opj +import shutil +import time +import sys +import numpy as np + +import logging +log = logging.getLogger(__name__) +from logdecorator import log_on_start, log_on_end + +from plancklens.qcinv import multigrid + +import lenspyx.remapping.utils_geom as utils_geom +from lenspyx.remapping.utils_geom import pbdGeometry, pbounds +from lenspyx.lensing import get_geom + +from delensalot.utils import cli, read_map +from delensalot.utility.utils_hp import Alm, almxfl, alm2cl +from delensalot.utility import utils_qe + +from delensalot.core import cachers +from delensalot.core.opfilt import opfilt_base +from delensalot.core.iterator import bfgs, steps + +from . import cs_iterator as csit + +alm2rlm = lambda alm : alm # get rid of this +rlm2alm = lambda rlm : rlm + + +@log_on_start(logging.INFO, " Start of prt_time()") +@log_on_end(logging.INFO, " Finished prt_time()") +def prt_time(dt, label=''): + dh = np.floor(dt / 3600.) + dm = np.floor(np.mod(dt, 3600.) / 60.) + ds = np.floor(np.mod(dt, 60)) + log.info("\r [" + ('%02d:%02d:%02d' % (dh, dm, ds)) + "] " + label) + return + +typs = ['T', 'QU', 'TQU'] + + +class iterator_pertmf(csit.iterator_pertmf): + def __init__(self, lib_dir:str, h:str, lm_max_dlm:tuple, + dat_maps:list or np.ndarray, plm0:np.ndarray, mf_resp:np.ndarray, pp_h0:np.ndarray, + cpp_prior:np.ndarray, cls_filt:dict, ninv_filt:opfilt_base.alm_filter_wl, k_geom:utils_geom.Geom, + chain_descr, stepper:steps.nrstep, mf0=None, kappa0=None, muG=None, clG=None, **kwargs): + """Lensing map iterator + + The bfgs hessian updates are called 'hlm's and are either in plm, dlm or klm space + + Args: + h: 'k', 'd', 'p' if bfgs updates act on klm's, dlm's or plm's respectively + pp_h0: the starting hessian estimate. (cl array, ~ 1 / N0 of the lensing potential) + cpp_prior: fiducial lensing potential spectrum used for the prior term + cls_filt (dict): dictionary containing the filter cmb unlensed spectra (here, only 'ee' is required) + k_geom: lenspyx geometry for once-per-iterations operations (like checking for invertibility etc, QE evals...) + stepper: custom calculation of NR-step + wflm0(optional): callable with Wiener-filtered CMB map search starting point + + """ + + super(csit.iterator_pertmf, self).__init__(lib_dir, h, lm_max_dlm, dat_maps, plm0, pp_h0, cpp_prior, cls_filt, + ninv_filt, k_geom, chain_descr, stepper, **kwargs) + assert mf_resp.ndim == 1 and mf_resp.size > self.lmax_qlm, mf_resp.shape + if mf0 is not None: + assert self.lmax_qlm == Alm.getlmax(mf0.size, self.mmax_qlm), (self.lmax_qlm, Alm.getlmax(mf0.size, self.lmax_qlm)) + self.cacher.cache('mf', almxfl(mf0, self._h2p(self.lmax_qlm), self.mmax_qlm, False)) + self.p_mf_resp = mf_resp + + nside = 4096 + geominfo = ('healpix', {'nside': nside}) + #geominfo_defl = ('thingauss', {'lmax': 4200 + 300, 'smax': 2}) + self.q_geom = get_geom(geominfo) + + self.ffi = ninv_filt.ffi + + ells = np.arange(0, self.hh_h0.size, 1) + factor = ells*(ells+1)/2 + self.hh_h0 *= factor**2. #transform in kappa space, and assume this is ok for the lognormal Gaussian field + + #kappa0 = 1.0272441232149767 + #ymu = 0.023049319395827456 + #cly = np.loadtxt(direc+"clgaussian.txt") + + """kappa0 = 0.8869370911600852 + ymu = -0.12403122348608665 + cly = np.loadtxt(direc+"clgaussian_5120.txt") + """ + + kappa0 = kappa0 + ymu = muG + cly = clG + + cly[:1] *= 0. + + self.cly = cly + self.kappa0 = kappa0 + self.ymu = ymu #mean of Gaussian filed + + """ #incr is in y map + geominfo = ('healpix', {'nside': nside}) + ninvjob_geometry_new = get_geom(geominfo) + kappa0 = 1.0272441232149767 + ymu = 0.023049319395827456 + ells = np.arange(0, self.lmax_qlm+1, 1) + factor = ells*(ells+1)/2 + klm0 = almxfl(plm0, factor, mmax = self.mmax_qlm, inplace = False) #k0 estimate + kmap = ninvjob_geometry_new.alm2map(klm0, self.lmax_qlm, self.mmax_qlm, 6, (-1., 1.)) #kmap estimate from harmonic to real space + y = np.log(kmap+kappa0) #take log of shifted kappa map to obtain the y field + ylm = ninvjob_geometry_new.map2alm(y, lself.max_qlm, self.mmax_qlm, 6, (-1., 1.)) + plm0 = ylm + print("Done with ylm") + """ + + + def get_hlm(self, itr, key): + """Loads current estimate """ + print("Loading hlm, iteration", itr) + if itr < 0: + return np.zeros(Alm.getsize(self.lmax_qlm, self.mmax_qlm), dtype=complex) + assert key.lower() in ['p', 'o'], key # potential or curl potential. + fn = '%s_%slm_it%03d' % ({'p': 'phi', 'o': 'om'}[key.lower()], self.h, itr) + if self.cacher.is_cached(fn): + return self.cacher.load(fn) + return self._sk2plm(itr) + + def phi_to_kappa(self, phi_lm): + lmax = Alm.getlmax(phi_lm.size, None) + ells = np.arange(0, lmax+1, 1) + factor = ells*(ells+1)/2 + return almxfl(phi_lm, factor, lmax, False) + + + def kappa_to_phi(self, kappa_lm): + ells = np.arange(0, self.lmax_qlm+1, 1) + factor = ells*(ells+1)/2 + return almxfl(kappa_lm, cli(factor), self.mmax_qlm, False) + + def alm2map(self, alm): + return self.q_geom.alm2map(alm.copy(), self.lmax_qlm, self.mmax_qlm, self.ffi.sht_tr, (-1., 1.)) + + def map2alm(self, map): + return self.q_geom.map2alm(map.copy(), self.lmax_qlm, self.mmax_qlm, self.ffi.sht_tr, (-1., 1.)) + + def kappa_shifted(self, kappa): + return kappa+self.kappa0 + + def kappa_to_y_real(self, kappa): + y = np.log(self.kappa_shifted(kappa))-self.ymu + return y + + def y_to_kappa_real(self, y): + kappa = np.exp(y)-self.kappa0 + return kappa + + def load_gradpri(self, itr, key): + """ + Log-normal gradient with respect to p_lm + """ + + assert key in ['p'], key + ' not implemented' + assert self.is_iter_done(itr -1 , key) + ret = self.get_hlm(itr, key) + ret = almxfl(ret, cli(self.cly), self.mmax_qlm, False) + return ret + + + def get_y(self, itr, key): + + assert key in ['p'], key + ' not implemented' + assert self.is_iter_done(itr -1 , key) + ret = self.get_hlm(itr, key) + + kappa_lm = self.phi_to_kappa(ret) + kappa = self.alm2map(kappa_lm) + + return np.log(self.kappa_shifted(kappa)) + + + @log_on_start(logging.INFO, "calc_gradlik(it={itr}, key={key}) started") + @log_on_end(logging.INFO, "calc_gradlik(it={itr}, key={key}) finished") + def calc_gradlik(self, itr, key, iwantit=False): + """Computes the quadratic part of the gradient for plm iteration 'itr' + + """ + assert self.is_iter_done(itr - 1, key) + assert itr > 0, itr + assert key.lower() in ['p', 'o'], key # potential or curl potential. + if not self._is_qd_grad_done(itr, key) or iwantit: + assert key in ['p'], key + ' not implemented' + dlm = self.get_hlm(itr - 1, key) + + #now this is in reality the y field, so you have to get kappa + yreal = self.alm2map(dlm) + kappa = self.y_to_kappa_real(yreal) + kappa_lm = self.map2alm(kappa) + dlm = self.kappa_to_phi(kappa_lm) + + self.hlm2dlm(dlm, True) + ffi = self.filter.ffi.change_dlm([dlm, None], self.mmax_qlm, cachers.cacher_mem(safe=False)) + self.filter.set_ffi(ffi) + mchain = multigrid.multigrid_chain(self.opfilt, self.chain_descr, self.cls_filt, self.filter) + if self._usethisE is not None: + if callable(self._usethisE): + log.info("iterator: using custom WF E") + soltn = self._usethisE(self.filter, itr) + else: + assert 0, 'dont know what to do this with this E input' + else: + soltn, it_soltn = self.load_soltn(itr, key) + if it_soltn < itr - 1: + soltn *= self.soltn_cond + mchain.solve(soltn, self.dat_maps, dot_op=self.filter.dot_op()) + fn_wf = 'wflm_%s_it%s' % (key.lower(), itr - 1) + log.info("caching " + fn_wf) + self.wf_cacher.cache(fn_wf, soltn) + else: + log.info("Using cached WF solution at iter %s "%itr) + + t0 = time.time() + if ffi.pbgeom.geom is self.k_geom and ffi.pbgeom.pbound == pbounds(0., 2 * np.pi): + # This just avoids having to recalculate angles on a new geom etc + q_geom = ffi.pbgeom + else: + q_geom = pbdGeometry(self.k_geom, pbounds(0., 2 * np.pi)) + G, C = self.filter.get_qlms(self.dat_maps, soltn, q_geom) + almxfl(G if key.lower() == 'p' else C, self._h2p(self.lmax_qlm), self.mmax_qlm, True) + log.info('get_qlms calculation done; (%.0f secs)'%(time.time() - t0)) + if itr == 1: #We need the gradient at 0 and the yk's to be able to rebuild all gradients + fn_lik = '%slm_grad%slik_it%03d' % (self.h, key.lower(), 0) + self.cacher.cache(fn_lik, -G if key.lower() == 'p' else -C) + return -G if key.lower() == 'p' else -C + + + @log_on_start(logging.INFO, "iterate(it={itr}, key={key}) started") + @log_on_end(logging.INFO, "iterate(it={itr}, key={key}) finished") + def iterate(self, itr, key): + """Performs iteration number 'itr' + + This is done by collecting the gradients at level iter, and the lower level potential + + """ + assert key.lower() in ['p', 'o'], key # potential or curl potential. + if not self.is_iter_done(itr, key): + assert self.is_iter_done(itr - 1, key), 'previous iteration not done' + self.logger.on_iterstart(itr, key, self) + # Calculation in // of lik and det term : + glm_like = self.calc_gradlik(itr, key) + glm_det = self.calc_graddet(itr, key) + + glm = glm_det+glm_like + + glm_kappa = self.kappa_to_phi(glm) #gradient in kappa + glm_kappa_real = self.alm2map(glm_kappa) + + ylm = self.get_hlm(itr-1, key) + y = self.alm2map(ylm) + glm_y_real = glm_kappa_real*np.exp(y) #gradient in y + glm = self.map2alm(glm_y_real) + + glm_pri = self.load_gradpri(itr - 1, key) #prior in y + glm += glm_pri + almxfl(glm, self.cly > 0, self.mmax_qlm, True) + + + self.build_incr(itr, key, glm) + del glm + + self.logger.on_iterdone(itr, key, self) + if self.tidy > 2: # Erasing deflection databases + if os.path.exists(opj(self.lib_dir, 'ffi_%s_it%s'%(key, itr))): + shutil.rmtree(opj(self.lib_dir, 'ffi_%s_it%s'%(key, itr))) + + + def load_gradquad(self, k, key): + fn = '%slm_grad%slik_it%03d' % (self.h, key.lower(), k) + result = self.cacher.load(fn) + glm_kappa = self.kappa_to_phi(result) #gradient in kappa + glm_kappa_real = self.alm2map(glm_kappa) + ylm = self.get_hlm(k, key) + y = self.alm2map(ylm) + glm_y_real = glm_kappa_real*np.exp(y) #gradient in y + + return self.map2alm(glm_y_real) + + @log_on_start(logging.INFO, "load_graddet(it={itr}, key={key}) started") + @log_on_end(logging.INFO, "load_graddet(it={itr}, key={key}) finished") + def load_graddet(self, itr, key): + assert self.h == 'p', 'check this line is ok for other h' + mf = almxfl(self.get_hlm(itr - 1, key), self.p_mf_resp * self._h2p(self.lmax_qlm), self.mmax_qlm, False) + if self.cacher.is_cached('mf'): + mf += self.cacher.load('mf') + return mf + + + def get_kappa(self, ylm, kappa0): + y = self.alm2map(ylm) + return self.kappa_shifted(np.exp(y), -kappa0) + + def get_plm(self, ylm, kappa0): + kappa = self.get_kappa(ylm, kappa0) + return self.kappa_to_phi(self.map2alm(kappa)) + + + def load_gradient(self, itr, key): + """Loads the total gradient at iteration iter. + + All necessary alm's must have been calculated previously + + """ + if itr == 0: + g = self.load_gradpri(0, key) + g += self.load_graddet(0, key) + g += self.load_gradquad(0, key) + return g + return self._yk2grad(itr) + + + @log_on_start(logging.INFO, "build_incr(it={it}, key={key}) started") + @log_on_end(logging.INFO, "build_incr(it={it}, key={key}) finished") + def build_incr(self, it, key, gradn): + """Search direction + + BGFS method with 'self.NR method' BFGS updates to the hessian. + Initial hessian are built from N0s. + + :param it: current iteration level. Will produce the increment to phi_{k-1}, from gradient est. g_{k-1} + phi_{k_1} + output = phi_k + :param key: 'p' or 'o' + :param gradn: current estimate of the gradient (alm array) + :return: increment for next iteration (alm array) + + s_k = x_k+1 - x_k = - H_k g_k + y_k = g_k+1 - g_k + """ + assert it > 0, it + k = it - 2 + yk_fname = 'rlm_yn_%s_%s' % (k, key) + if k >= 0 and not self.hess_cacher.is_cached(yk_fname): # Caching hessian BFGS yk update : + yk = alm2rlm(gradn - self.load_gradient(k, key)) + self.hess_cacher.cache(yk_fname, yk) + k = it - 1 + BFGS = self.get_hessian(k, key) # Constructing L-BFGS hessian + # get descent direction sk = - H_k gk : (rlm array). Will be cached directly + sk_fname = 'rlm_sn_%s_%s' % (k, key) + if not self.hess_cacher.is_cached(sk_fname): + log.info("calculating descent direction" ) + t0 = time.time() + incr = BFGS.get_mHkgk(alm2rlm(gradn), k) + incr = alm2rlm(self.stepper.build_incr(incr, it)) + self.hess_cacher.cache(sk_fname, incr) + prt_time(time.time() - t0, label=' Exec. time for descent direction calculation') + assert self.hess_cacher.is_cached(sk_fname), sk_fname + + diff --git a/delensalot/core/iterator/statics.py b/delensalot/core/iterator/statics.py index 4d9f7cbf..ba64aade 100644 --- a/delensalot/core/iterator/statics.py +++ b/delensalot/core/iterator/statics.py @@ -9,6 +9,36 @@ alm2rlm = lambda x : x.copy() rlm2alm = lambda x : x.copy() + +from plancklens import shts +from delensalot.utility.utils_hp import Alm, almxfl +from plancklens.utils import cli + +def kappa_to_phi(k_lm): + lmax = Alm.getlmax(k_lm.size, None) + ells = np.arange(0, lmax+1, 1) + factor = ells*(ells+1)/2 + return almxfl(k_lm, cli(factor), lmax, False) + +def kappa_shifted(kappa, kappa0): + return kappa+kappa0 + +def get_kappa(ylm, kappa0, lmax): + y = shts.alm2map(ylm, lmax) + return kappa_shifted(np.exp(y), -kappa0) + +def get_plm(ylm, kappa0, lmax): + kappa = get_kappa(ylm, kappa0, lmax) + return kappa_to_phi(shts.map2alm(kappa, lmax)) + + +def transform(rlm, kappa0): + if kappa0 is not None: + print('kappa0 is not None, transforming to phi') + lmax = Alm.getlmax(rlm.size, None) + rlm = get_plm(rlm, kappa0, lmax) + return rlm + #TODO this looks like a 'query' class to me. May be refactored. class rec: """Static methods to reach for iterated lensing maps etc @@ -27,22 +57,32 @@ def maxiterdone(lib_dir): return itr @staticmethod - def load_plms(lib_dir, itrs): + def load_plms(lib_dir, itrs, kappa0 = None): """Loads plms for the requested itrs""" lib_dir = os.path.abspath(lib_dir) cacher = cachers.cacher_npy(lib_dir) itmax = np.max(itrs) sk_fname = lambda k: os.path.join(lib_dir, 'hessian', 'rlm_sn_%s_%s' % (k, 'p')) rlm = alm2rlm(cacher.load(os.path.join(lib_dir, 'phi_plm_it000'))) - ret = [] if 0 not in itrs else [rlm2alm(rlm)] + + """if kappa0 is not None: + print('kappa0 is not None, transforming to phi') + lmax = Alm.getlmax(rlm.size, None) + rlm = get_plm(rlm, kappa0, lmax) + """ + + ret = [] if 0 not in itrs else [rlm2alm(transform(rlm, kappa0))] for i in range(itmax): if cacher.is_cached(sk_fname(i)): rlm += cacher.load(sk_fname(i)) + if (i + 1) in itrs: - ret.append(rlm2alm(rlm)) + print("Doing for kappa0", kappa0) + ret.append(rlm2alm(transform(rlm, kappa0))) else: log.info("*** Could only build up to itr number %s"%i) return ret + return ret @staticmethod diff --git a/delensalot/core/iterator/steps.py b/delensalot/core/iterator/steps.py index b3866685..8b8abb40 100644 --- a/delensalot/core/iterator/steps.py +++ b/delensalot/core/iterator/steps.py @@ -14,8 +14,9 @@ def steplen(self, itr, incrnorm): return self.val def build_incr(self, incrlm, itr): - print('incr step val %.5f'%self.val) + #print('incr step val %.5f'%self.val) return incrlm * self.val + #return almxfl(incrlm, self.val, self.mmax_qlm, False) class harmonicbump(nrstep): def __init__(self, lmax_qlm, mmax_qlm, xa=400, xb=1500, a=0.5, b=0.1, scale=50, flt=None): diff --git a/delensalot/core/opfilt/MAP_opfilt_aniso_t.py b/delensalot/core/opfilt/MAP_opfilt_aniso_t.py index 2d931627..150968c4 100644 --- a/delensalot/core/opfilt/MAP_opfilt_aniso_t.py +++ b/delensalot/core/opfilt/MAP_opfilt_aniso_t.py @@ -28,7 +28,8 @@ def apply_fini(*args, **kwargs): class alm_filter_ninv_wl(opfilt_base.alm_filter_wl): def __init__(self, ninv_geom:utils_geom.Geom, ninv: np.ndarray, ffi:remapping.deflection, transf:np.ndarray, - unlalm_info:tuple, lenalm_info:tuple, sht_threads:int,verbose=False, lmin_dotop=0, tpl:tni.template_tfilt or None =None): + unlalm_info:tuple, lenalm_info:tuple, sht_threads:int,verbose=False, lmin_dotop=0, tpl:tni.template_tfilt or None =None, + extra_ninv: np.ndarray = None): r"""CMB inverse-variance and Wiener filtering instance, using unlensed E and lensing deflection Args: @@ -40,6 +41,7 @@ def __init__(self, ninv_geom:utils_geom.Geom, ninv: np.ndarray, ffi:remapping.de lenalm_info: tuple of int, lmax and mmax of lensed CMB sht_threads: number of threads for lenspyx SHTs verbose: some printout if set, defaults to False + extra_ninv: np.ndarray, extra map for inverse noise filtering to be applied to the observed map """ lmax_unl, mmax_unl = unlalm_info @@ -75,6 +77,9 @@ def __init__(self, ninv_geom:utils_geom.Geom, ninv: np.ndarray, ffi:remapping.de self.template = tpl + self.extra_ninv = extra_ninv # I do not use this one for the diagonal pre-conditioned CG + ## assume extra_inv will enter in the form: B^t extra_ninv B, then I want to take the inverse of this + def hashdict(self): return {'ninv':self._ninv_hash(), 'transf':clhash(self.b_transf_tlm), 'deflection':self.ffi.hashdict(), @@ -101,6 +106,28 @@ def apply_map(self, tmap): """ tmap *= self.n_inv + + if self.extra_ninv is not None: + """ + Here I am doing this: + N^{-1} - N^{-1}B^t \hat{S}^2 B N^{-1}, applied on X map + + This assumes that we have a small perturbation to the noise. + + Note that for now here the noise is diagonal in pixel space. This is not the case for the extra piece of here. + """ + B = tmap.copy() + B = self.geom_.map2alm(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.)) + B = almxfl(B, self.b_transf_tlm, self.mmax_len, inplace = False) + B = self.geom_.alm2map(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.)) + B *= self.extra_ninv + B = self.geom_.map2alm(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.)) + B = almxfl(B, self.b_transf_tlm, self.mmax_len, inplace = False) + B = self.geom_.alm2map(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.)) + B *= self.n_inv + tmap -= B + + if self.template is not None: ts = [self.template] # Hack, this is only meant for one template coeffs = np.concatenate(([t.dot(tmap) for t in ts])) diff --git a/delensalot/sims/foregrounds/pointsources.py b/delensalot/sims/foregrounds/pointsources.py new file mode 100644 index 00000000..38fdb2eb --- /dev/null +++ b/delensalot/sims/foregrounds/pointsources.py @@ -0,0 +1,110 @@ +""" +Generates simple point sources, can be correlated with an input lensing convergence map. +""" + +import numpy as np +import healpy as hp + + +class Foreground(object): + def __init__(self): + pass + + @staticmethod + def randomizing_fg(mappa: np.ndarray): + """ + Randomizes the phase of the input map, preserving the amplitude. + """ + f = lambda z: np.abs(z) * np.exp(1j*np.random.uniform(0., 2.*np.pi, size = z.shape)) + return f(mappa) + + def randomized_map(self, mappa: np.ndarray, nside: int): + """ + Randomizes the phase of the input map, preserving the amplitude. + """ + alm = hp.map2alm(mappa) + alm = self.randomizing_fg(alm) + return hp.alm2map(alm, nside) + + + @staticmethod + def matched_filter(input_map_alm: np.ndarray, total_cl: np.ndarray, signal_cl: np.ndarray, nside: int): + """ + Returns the matched filter map. + """ + alm = hp.almxfl(input_map_alm, np.nan_to_num(1/total_cl)) + alm = hp.almxfl(alm, signal_cl) + return hp.alm2map(alm, nside) + + def mask_from_matched_filter(self, input_map_alm: np.ndarray, total_cl: np.ndarray, signal_cl: np.ndarray, nside: int, threshold: float = 0.5): + """ + Returns a mask from the matched filter map. + """ + mappa = self.matched_filter(input_map_alm, total_cl, signal_cl, nside) + SN_map = abs(mappa) / np.std(mappa) + mask = np.where(mappa > threshold, 1, 0) + return mask + + @staticmethod + def smooth_map(mappa: np.ndarray, fwhm: float, nside: int): + """ + Smooths the input map. + """ + return hp.smoothing(mappa, fwhm = np.radians(fwhm)) + + +class PointSourcesSimple(Foreground): + + def __init__(self, nside: int = 2048) -> None: + self.nside = nside + + @staticmethod + def phi_lm_to_kappa_lm(plm: np.ndarray) -> np.ndarray: + """ + Converts the input phi_lm to kappa_lm. + """ + lmax = hp.Alm.getlmax(plm.size) + ls = np.arange(0, lmax) + factor = (ls * (ls + 1.)) / 2. + return hp.almxfl(plm, factor) + + + @staticmethod + def _get_position_from_kappa_default(rng, kappa: np.ndarray, factor: float = 0.5) -> np.ndarray: + """ + Returns the positions of the point sources from the input kappa map. + """ + positions = np.where(rng.poisson(1+kappa*factor) > 0)[0] + return positions + + @staticmethod + def _get_position_from_kappa_alternative(rng, kappa: np.ndarray, factor: float = 0.5) -> np.ndarray: + """ + Returns the positions of the point sources from the input kappa map. + """ + positions = np.where(rng.poisson(abs(kappa)*factor) > 0)[0] + return positions + + + def generate_ps(self, nsrc: int, amp: float = 100, seed: int = 0, plm: np.ndarray = None, factor: float = 0.5) -> np.ndarray: + + rng = np.random.default_rng(seed) + + mappa = np.zeros(hp.nside2npix(self.nside)) + + if plm is not None: + + klm = self.phi_lm_to_kappa_lm(plm) + kmap = hp.alm2map(klm, self.nside, verbose = False) + positions = self._get_position_from_kappa_default(rng, kmap, factor) + nsources = len(positions) + + else: + + nsources = rng.poisson(nsrc) + positions = np.random.randint(0, len(mappa), nsources) + + amplitudes = rng.poisson(amp, nsources) + mappa[positions] = amplitudes + + return mappa diff --git a/delensalot/sims/generic.py b/delensalot/sims/generic.py index c57e689a..f7829a20 100644 --- a/delensalot/sims/generic.py +++ b/delensalot/sims/generic.py @@ -51,9 +51,11 @@ class sims_cmb_len(object): verbose(defaults to True): lenspyx timing info printout + extra_tlm: optional extra map to add to the CMB map, e.g. foregrounds + """ def __init__(self, lib_dir, lmax, cls_unl, lib_pha=None, offsets_plm=None, offsets_cmbunl=None, - dlmax=1024, nside_lens=4096, epsilon=1e-7, nbands=8, cache_plm=True, verbose=True): + dlmax=1024, nside_lens=4096, epsilon=1e-7, nbands=8, cache_plm=True, verbose=True, extra_tlm = None): fields = _get_fields(cls_unl) @@ -99,6 +101,8 @@ def __init__(self, lib_dir, lmax, cls_unl, lib_pha=None, offsets_plm=None, offse self.lens_module = lenspyx self.verbose=verbose + self.extra_tlm = extra_tlm + @staticmethod def offset_index(idx, block_size, offset): """Offset index by amount 'offset' cyclically within blocks of size block_size @@ -169,17 +173,32 @@ def _cache_eblm(self, idx): def get_sim_tlm(self, idx, ret=True): fname = os.path.join(self.lib_dir, 'sim_%04d_tlm.fits' % idx) + pfname = os.path.join(self.lib_dir, 'sim_%04d_plm.fits' % idx) if not os.path.exists(fname): tlm = self.unlcmbs.get_sim_tlm(self.offset_index(idx, self.offset_cmb[0], self.offset_cmb[1])) dlm = self.get_sim_plm(idx) + + hp.write_alm(pfname, dlm) + assert 'o' not in self.fields, 'not implemented' lmaxd = hp.Alm.getlmax(dlm.size) hp.almxfl(dlm, np.sqrt(np.arange(lmaxd + 1, dtype=float) * np.arange(1, lmaxd + 2)), inplace=True) Tlen = self.lens_module.alm2lenmap(np.array(tlm), [dlm, None], epsilon=self.epsilon, verbose=self.verbose) hp.write_alm(fname, hp.map2alm(Tlen, lmax=self.lmax, iter=0)) + + if (self.extra_tlm is not None): + """ + Adding an extra CMB component to the lensed CMB. + """ + extrafname = os.path.join(self.lib_dir, f'sim_{idx:04}_{self.extra_tlm.get_name()}lm.fits') + if (not os.path.exists(extrafname)): + extra_tlm = hp.map2alm(self.extra_tlm(idx), lmax=self.lmax, iter=0) + hp.write_alm(extrafname, extra_tlm) + if ret: - return hp.read_alm(fname) + total = hp.read_alm(fname) + (0 if self.extra_tlm is None else hp.read_alm(extrafname)) + return total def get_sim_elm(self, idx, ret=True): fname = os.path.join(self.lib_dir, 'sim_%04d_elm.fits' % idx) diff --git a/delensalot/sims/lognormal/__init__.py b/delensalot/sims/lognormal/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/delensalot/sims/lognormal/lognormal_utils.py b/delensalot/sims/lognormal/lognormal_utils.py new file mode 100644 index 00000000..daaf7afe --- /dev/null +++ b/delensalot/sims/lognormal/lognormal_utils.py @@ -0,0 +1,170 @@ +"""Module with utilities for the lognormal sims module. +""" + +import numpy as np +import flt +import healpy as hp + + +def get_out_quantities_from_a_map(outmap: np.ndarray) -> tuple: + """ + Returns the mean, variance, skewness, lambda, muG and sigmaG from a given map. + """ + skewness = get_skew_from_map(outmap) + variance = get_variance_from_map(outmap) + mean = get_mean_from_map(outmap) + lamb = get_lambda_from_skew(skewness, variance, mean) + + alpha = get_alpha(mean, lamb) + sigmaG = get_sigma_gauss(alpha, variance) + muG = get_mu_gauss(alpha, variance) + return mean, variance, skewness, lamb, muG, sigmaG + + +def create_lognormal_single_map(inputcl: np.ndarray, nside: int, lmax_gen: int, mu: float = 0.0, lamb: float = 0.0): + """ + Creates a lognormal map with a given power spectrum and skewness. These are the parameters of the lognormal distribution that are specified with mu and lamb. + """ + + + alpha = get_alpha(mu, lamb) + xisinput = cl2xi(inputcl)/alpha/alpha + + xigaussian = np.log(xisinput+1) + clgaussian = xi2cl(xigaussian) + + vargauss = np.dot(np.arange(1, 2*len(clgaussian), 2), clgaussian)/(4*np.pi) + + lmax_gen = 2*nside-1 if lmax_gen is None else lmax_gen + almgaussian = hp.synalm(clgaussian, lmax = lmax_gen) #GENERATE TO HIGH LMAX + maps = hp.alm2map(almgaussian, nside = nside, pol = False) + + #vargauss = np.array([xigaussian[i, i][0] for i in range(Nfields)]) + #vargauss = np.array([np.var(m) for m in maps]) + + expmu = (mu+lamb)*np.exp(-vargauss*0.5) + maps = np.array(maps) + maps = np.exp(maps) + maps *= expmu + maps -= lamb + return maps + + +#shifted log-normal distribution +def shifted_lognormal_zero_mean(x, sigmaG, lamb): + #equation (22) of https://www.aanda.org/articles/aa/pdf/2011/12/aa17294-11.pdf + return np.exp(-(np.log(x/lamb+1)+sigmaG**2/2)**2./(2.*sigmaG**2.))/(x+lamb)/sigmaG/np.sqrt(2.*np.pi)*(x>-lamb) + + + +def cl2xi(cl: np.ndarray, closed = False): + """ + This goes from angular power spectrum to a correlation function calculated at theta points. + """ + ls = np.arange(0, len(cl)) + factorcl = (2*ls+1)/(4*np.pi) + coeffs = cl*factorcl + return flt.idlt(coeffs, closed = closed) + + +def theta(n, closed = False): + """ + Returns the theta for which the cl2xi are calculated, for a given n + """ + return flt.theta(n, closed = closed) + + +def xi2cl(xi: np.ndarray, closed = False): + """ + This goes from correlation function calculated at theta points to an angular power spectrum. + """ + ls = np.arange(0, len(xi)) + factorcl = (2*ls+1)/(4*np.pi) + return flt.dlt(xi, closed = closed)/factorcl + + + +def get_mean_from_map(mappa: np.ndarray): + """ + Mean from a map. + """ + return np.mean(mappa) + +def get_variance_from_map(mappa: np.ndarray): + """ + Variance from a map. + """ + return np.mean(mappa**2.)-np.mean(mappa)**2. + +def get_skew_from_map(mappa: np.ndarray): + """ + Skewness from a map. + """ + return np.mean((mappa-get_mean_from_map(mappa))**3.)/np.mean(mappa**2.)**1.5 + +def y_skew(skew): + """ + Formula (12) from https://arxiv.org/pdf/1602.08503.pdf + It relates the skewnees to some factor that is used to get the lambda parameter for the log-normal generation. + """ + result = 2+skew**2.+skew*np.sqrt(4+skew**2.) + result /= 2 + return np.power(result, 1/3.) + +def get_lambda_from_skew(skew, var, mu): + lmbda = np.sqrt(var)/skew*(1+y_skew(skew)+1/y_skew(skew))-mu + return lmbda + +def get_alpha(mu, lmbda): + """ + Below formula (7) from https://arxiv.org/pdf/1602.08503.pdf + """ + return mu+lmbda + +def get_mu_gauss(alpha, var): + """ + Gets the mu parameter for the Gaussian distribution for the log-normal + """ + result = np.log(alpha**2./np.sqrt(var+alpha**2.)) + return result + +def get_sigma_gauss(alpha, var): + """ + Gets the sigma parameter for the Gaussian distribution for the log-normal. + + Here the variance is the variance of the wanted log-normal field. + """ + result = np.log(1+var/alpha**2.) + result = np.sqrt(result) + return result + + +def suppress(l: np.ndarray, lsup: float = 7000, supindex: float = 10): + """ + Suppression factor at high ell. + """ + return np.exp(-1.0*np.power(l/lsup, supindex)) + +def suppress_cls(inputcl: np.ndarray, l: np.ndarray, lsup: float = 7000, supindex: float = 10): + """ + Suppresses an input angular power spectrum at high ell. + """ + return inputcl*suppress(l, lsup, supindex) + +def process_cl(inputcl: np.ndarray, function: callable = suppress_cls, **kwargs): + """ + Processes an input angular power spectrum with a given function. + + The function must take inputcl and l as arguments. + + Args: + inputcl (np.ndarray): Input angular power spectrum. + function (callable) default suppress_cls: Function to process the angular power spectrum. + **kwargs: Keyword arguments for the function. + + Returns: + np.ndarray: Processed angular power spectrum. + """ + ls = np.arange(0, len(inputcl)) + return function(inputcl = inputcl, l = ls, **kwargs) + \ No newline at end of file diff --git a/delensalot/sims/lognormal/sims_lognormal.py b/delensalot/sims/lognormal/sims_lognormal.py new file mode 100644 index 00000000..3385f1ec --- /dev/null +++ b/delensalot/sims/lognormal/sims_lognormal.py @@ -0,0 +1,81 @@ +""" +Generates a log-normal simulation with a given power spectrum and skewness. It uses methods described in e.g. https://arxiv.org/abs/1602.08503 + +NOTE: + An improvement could be done by setting some work at the unlcmb library level. Will leave this for now as an exploration. + The reason is that in the future users might want to set up their own generation of sims. There should be some simple way to ovverride + getting maps, maybe just some class inheritance somewhere. +""" + +import healpy as hp +import numpy as np +import lognormal_utils as lu +from delensalot.sims import sims_gaussian +import os + + +class sims_gaussian(sims_gaussian.sims_gaussian): + """Simulations with lognormal phi + + Args: + lib_dir: the phases of the CMB maps and the lensed CMBs will be stored there + lmax_cmb: cmb maps are generated down to this max multipole + cls_unl: dictionary of unlensed CMB spectra + dlmax, nside_lens, facres, nbands: lenspyx lensing module parameters + wcurl: include field rotation map in the lensing deflection (default to False for historical reasons) + + + This uses the cl_fid phi from sims_postborn to generate new lensing potential fields. + + """ + def __init__(self, lib_dir, lmax_cmb, cls_unl:dict, wcurl=False, + dlmax=1024, nside_lens=4096, epsilon=1e-5, cache_plm=True, mu: float = 0.0, var: float = 1.0, skew: float = 0.0, input_cl: np.ndarray = None, lmax_gen: int = 8000): + + lmax_plm = lmax_cmb + dlmax + mmax_plm = lmax_plm + + self.lmax_plm = lmax_plm + self.mmax_plm = mmax_plm + self.path = None + + self.cache_plm = cache_plm + self.wcurl = wcurl + self.epsilon = epsilon + + cmb_cls = {} + for k in cls_unl.keys(): + cmb_cls[k] = np.copy(cls_unl[k][:lmax_cmb + dlmax + 1]) + + super(sims_gaussian, self).__init__(lib_dir, lmax_cmb, cmb_cls, + dlmax=dlmax, nside_lens=nside_lens, epsilon=self.epsilon) + + if input_cl is None: + self.input_cl = cmb_cls['pp'] + + self.mu = mu + self.lamb = lu.get_lambda_from_skew(skew, var, mu) + self.lmax_gen = lmax_gen + + + @staticmethod + def kappa_lm_to_phi_lm(klm: np.ndarray) -> np.ndarray: + """ + Converts the input kappa_lm to phi_lm. + """ + lmax = hp.Alm.getlmax(klm.size) + ls = np.arange(0, lmax) + factor = np.nan_to_num(1/((ls * (ls + 1.)) / 2.)) + return hp.almxfl(klm, factor) + + def get_sim_plm(self, idx): + """ + Get a simulated lensing potential map + """ + fn = os.path.join(self.lib_dir, 'plm_in_%04d_lmax%s.fits'%(idx, self.lmax_plm)) + if not os.path.exists(fn): + klm = lu.create_lognormal_single_map(inputcl = self.input_cl, nside = self.nnside_lensside, lmax_gen = self.lmax_gen, mu = self.mu, lamb = self.lamb) + plm = self.kappa_lm_to_phi_lm(klm) + if self.cache_plm: + hp.write_alm(fn, plm) + return plm + return hp.read_alm(fn) \ No newline at end of file diff --git a/delensalot/sims/sims_extra.py b/delensalot/sims/sims_extra.py new file mode 100644 index 00000000..e91add59 --- /dev/null +++ b/delensalot/sims/sims_extra.py @@ -0,0 +1,25 @@ +""" +Extra sims utility. +""" + +import os +import healpy as hp +import numpy as np + + +class Extra(object): + """ + Example: extra_tlm = Extra('fgs', fgnames) + """ + + def __init__(self, baseWebsky, name, fgnames): + self.name = name + self.fgnames = fgnames + self.directory = baseWebsky + + def __call__(self, idx): + return np.sum([hp.read_map(opj(self.directory, f'{fgname}.fits')) for fgname in self.fgnames], axis = 0) + + def get_name(self): + return self.name + \ No newline at end of file diff --git a/delensalot/sims/sims_general.py b/delensalot/sims/sims_general.py new file mode 100644 index 00000000..398af364 --- /dev/null +++ b/delensalot/sims/sims_general.py @@ -0,0 +1,75 @@ +"""Module to allow general simulations that include cmb+foregrounds+other at the signal level. +""" + +from plancklens.sims import maps + +import healpy as hp + + +class cmb_maps(maps.cmb_maps): + + """ + Class to handle multiple objects that give a field at the CMB signal level. + """ + + def __init__(self, **kwargs): + """ + Initializes the cmb_maps object. Note, for now you have to initialize it with a sims_cmb_len object. + """ + super(cmb_maps, self).__init__(**kwargs) + self.components = [] + self.lmax = self.sims_cmb_len.lmax + + + def get_sim_tmap(self,idx): + """Returns temperature healpy map for a simulation + + Args: + idx: simulation index + + Returns: + healpy map + """ + tmap = self.get_sim_tlm(idx) + hp.almxfl(tmap,self.cl_transf,inplace=True) + tmap = hp.alm2map(tmap,self.nside) + return tmap + self.get_sim_tnoise(idx) + + def get_sim_tlm(self, idx): + """ + Returns the temperature alm of the sum of the components. + + Args: + idx: simulation index + + Returns: + alm of the sum of the components. + """ + return sum([self.check(c.get_sim_tlm(idx)) for c in self.components]) + + def check(self, alms): + """ + Just a simple check to make sure the alms have the same lmax as the one initialized the object. + + Args: + alms: alms to check + + Returns: + alms + """ + assert hp.Alm.getlmax(alms.size) == self.lmax, "The alms you are trying to add have a different lmax than the one you initialized the object with!" + return alms + + def __add__(self, other): + """ + Adds a component to the cmb_maps object. + """ + return self._update(other) + + def _update(self, other): + """ + Adds a component to the cmb_maps object. + """ + assert callable(other.get_sim_tlm), "The object you are trying to add has to have a get_sim_tlm method!" + self.components.append(other) + return self \ No newline at end of file