diff --git a/.gitignore b/.gitignore index c1dc62f9..1d00a461 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ outputs/* *.o *.obj +*.key # slurm slurm* diff --git a/delensalot/core/iterator/cs_iterator_multi.py b/delensalot/core/iterator/cs_iterator_multi.py new file mode 100644 index 00000000..aa3b7d09 --- /dev/null +++ b/delensalot/core/iterator/cs_iterator_multi.py @@ -0,0 +1,590 @@ +"""Module for curved-sky iterative lensing estimation + + + In contrast to cs_iterator.py, this module attemps to reconstruct jointly several fields + (for example the lensing gradient and curl potential jointly) + + +""" +from __future__ import annotations + +import os +from os.path import join as opj +import shutil +import time +import sys +import numpy as np + +from plancklens.qcinv import multigrid + +from lenspyx.utils_hp import Alm, almxfl, alm2cl, alm_copy +from lenspyx.remapping.utils_geom import pbdGeometry, pbounds, Geom +from lenspyx import cachers + +from delensalot.core.iterator import bfgs, loggers +from delensalot.core.opfilt import opfilt_base +from delensalot.utility import utils_qe +from delensalot.utils import cli + + +import logging +log = logging.getLogger(__name__) +from logdecorator import log_on_start, log_on_end + +from delensalot.utility.utils_steps import nrstep, gradient, gradient_dotop + +@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 logger_norms(loggers.logger_norms): + def __init__(self, txt_file): + super().__init__(txt_file) + self.txt_file = txt_file + self.ti = None + + + def on_iterdone(self, itr:int, key:str, iterator:gclm_iterator): + incr = iterator.hess_cacher.load('rlm_sn_%s_%s' % (itr-1, key)) + norm_inc = iterator.calc_norm(incr) / iterator.calc_norm(iterator.get_hlm(0)) + norms = [iterator.calc_norm(iterator.load_gradient(itr - 1))] + norm_grad_0 = iterator.calc_norm(iterator.load_gradient(0)) + for i in [0]: norms[i] = norms[i] / norm_grad_0 + + with open(opj(iterator.lib_dir, 'history_increment.txt'), 'a') as file: + file.write('%03d %.1f %.6f %.6f \n' + % (itr, time.time() - self.ti, norm_inc, norms[0])) + file.close() + + +class gclm_iterator(object): + def __init__(self, lib_dir:str, h:str, lm_max_dlm:list[tuple], + dat_maps:list or np.ndarray, plm0s:list or np.ndarray, pp_h0s:list[np.ndarray], + cpp_priors:list[np.ndarray], labels:tuple[str], cls_filt:dict, + ninv_filt:opfilt_base.alm_filter_wl, + k_geom:Geom, + chain_descr, stepper:nrstep, + lm_maxee:tuple[int] or None=None, + logger=None, + NR_method=100, tidy=0, verbose=True, soltn_cond=True, wflm0=None, _usethisE=None): + """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 + plm0s: starting point for each field to be reconstructed + pp_h0s: the starting hessian estimate for each field. (cl array, ~ 1 / N0 of the lensing potential) + cpp_priors: fiducial lensing potential spectrum used for the prior term for each component + labels: components identification string (e.g. ('p', 'x') for joint lensing gradient and curl rec) + cls_filt (dict): dictionary containing the filter cmb unlensed spectra (here, only 'ee' is required) + k_geom: scarf geometry for once-per-iterations opertations (like cehcking for invertibility etc) + stepper: custom calculation of NR-step + wflm0(optional): callable with Wiener-filtered CMB map search starting point + + + """ + assert h in ['k', 'p', 'd'] + lmax_filt, mmax_filt = ninv_filt.lmax_sol, ninv_filt.mmax_sol + plm0s = plm0s if isinstance(plm0s, list) else [plm0s] + pp_h0s = pp_h0s if isinstance(pp_h0s, list) else [pp_h0s] + cpp_priors = cpp_priors if isinstance(cpp_priors, list) else [cpp_priors] + + assert len(lm_max_dlm) == len(plm0s) + assert len(plm0s) == len(pp_h0s) and len(plm0s) == len(cpp_priors) + assert len(labels) >= len(plm0s) + + + for plm, (lmax_qlm, mmax_qlm) in zip(plm0s, lm_max_dlm): + assert len(pp_h0s[0]) > lmax_qlm + assert Alm.getlmax(plm0s[0].size, mmax_qlm) == lmax_qlm + if mmax_qlm is None: mmax_qlm = lmax_qlm + + # lmax'es: here same for all, but easy to change + self.lmaxs_qlm = [lmax_qlm for lmax_qlm, mmax_qlm in lm_max_dlm] + self.mmaxs_qlm = [mmax_qlm for lmax_qlm, mmax_qlm in lm_max_dlm] + + self.h = h + + self.lib_dir = lib_dir + self.cacher = cachers.cacher_npy(lib_dir) + self.hess_cacher = cachers.cacher_npy(opj(self.lib_dir, 'hessian')) + self.wf_cacher = cachers.cacher_npy(opj(self.lib_dir, 'wflms')) + if logger is None: + logger = logger_norms(opj(lib_dir, 'history_increment.txt')) + self.logger = logger + + self.chain_descr = chain_descr + self.opfilt = sys.modules[ninv_filt.__module__] # filter module containing the ch-relevant info + self.stepper = stepper + self.soltn_cond = soltn_cond + + self.dat_maps = dat_maps + + chhs = [] + hh_h0s = [] + for cpp_prior, pp_h0, lmax_qlm in zip(cpp_priors, pp_h0s, self.lmaxs_qlm): + chh_p = cpp_prior[:lmax_qlm+1] * self._p2h(lmax_qlm) ** 2 + hh_h0_p = cli(pp_h0[:lmax_qlm + 1] * self._h2p(lmax_qlm) ** 2 + cli(chh_p)) #~ (1/Cpp + 1/N0)^-1 + hh_h0_p *= (cpp_prior> 0) + chhs.append(chh_p) + hh_h0s.append(hh_h0_p) + + self.chhs = chhs # (rescaled) isotropic approximation to the likelihood curvature for each component + self.hh_h0s = hh_h0s + + self.NR_method = NR_method + self.tidy = tidy + self.verbose = verbose + + self.cls_filt = cls_filt + self.lmax_filt = lmax_filt + self.mmax_filt = mmax_filt + + self.filter = ninv_filt + self.k_geom = k_geom + # Defining a trial newton step length : + + self.wflm0 = wflm0 + gclm_fname = '%s_%slm_it%03d' % ({'p': 'phi', 'o': 'om'}['p'], self.h, 0) + if not self.cacher.is_cached(gclm_fname): + glm0 = gradient(plm0s, self.mmaxs_qlm, labels=labels) + self.cacher.cache(gclm_fname, glm0.almxfl([self._p2h(self.lmaxs_qlm[0]),self._p2h(self.lmaxs_qlm[0]) ], False).getarray()) + self.logger.startup(self) + self.labels = labels[:len(plm0s)] + self.gradlm_size = np.sum([Alm.getsize(lmax_q, mmax_q) for lmax_q, mmax_q in zip(self.lmaxs_qlm, self.mmaxs_qlm)]) + # Size of total gradient array + + self._usethisE = _usethisE + + self.lm_maxee = lm_maxee + + def _p2h(self, lmax): + if self.h == 'p': + return np.ones(lmax + 1, dtype=float) + elif self.h == 'k': + return 0.5 * np.arange(lmax + 1, dtype=float) * np.arange(1, lmax + 2, dtype=float) + elif self.h == 'd': + return np.sqrt(np.arange(lmax + 1, dtype=float) * np.arange(1, lmax + 2), dtype=float) + else: + assert 0, self.h + ' not implemented' + + def _h2p(self, lmax): return cli(self._p2h(lmax)) + + def hlm2dlm(self, hlm:gradient, inplace): + """Rescaling of the gradient if desired """ + h2ds = [] + for lmax, mmax in zip(hlm.lmaxs, hlm.mmaxs): #Fix h values + if self.h == 'd': + h2d = np.ones(lmax + 1, dtype=float) + elif self.h == 'p': + h2d = np.sqrt(np.arange(lmax + 1, dtype=float) * np.arange(1, lmax + 2, dtype=float)) + elif self.h == 'k': + h2d = cli(0.5 * np.sqrt(np.arange(lmax + 1, dtype=float) * np.arange(1, lmax + 2, dtype=float))) + else: + assert 0, self.h + ' not implemented' + h2ds.append(h2d) + if inplace: + hlm.almxfl(h2ds, True) + else: + return hlm.almxfl(h2ds, False) + + + def _sk2plm(self, itr): + sk_fname = lambda k: 'rlm_sn_%s_%s' % (k, 'p') + rlm = self.cacher.load('phi_%slm_it000'%self.h) + for i in range(itr): + rlm += self.hess_cacher.load(sk_fname(i)) + return rlm + + def _yk2grad(self, itr): + yk_fname = lambda k: 'rlm_yn_%s_%s' % (k, 'p') + rlm = self.load_gradient(0).getarray() + for i in range(itr): + rlm += self.hess_cacher.load(yk_fname(i)) + return rlm + + def is_iter_done(self, itr): + """Returns True if the iteration 'itr' has been performed already and False if not + + """ + if itr <= 0: + return self.cacher.is_cached('%s_%slm_it000' % ('phi', self.h)) + sk_fname = lambda k: 'rlm_sn_%s_%s' % (k, 'p') + return self.hess_cacher.is_cached(sk_fname(itr - 1)) + + def _is_qd_grad_done(self, itr, key): + if itr <= 0: + return self.cacher.is_cached('%slm_grad%slik_it%03d' % (self.h, key.lower(), 0)) + yk_fname = lambda k: 'rlm_yn_%s_%s' % (k, 'p') + for i in range(itr): + if not self.hess_cacher.is_cached(yk_fname(i)): + return False + return True + + + @log_on_start(logging.INFO, "get_template_blm() started: it={it}, calc={calc}") + @log_on_end(logging.INFO, "get_template_blm() finished: it={it}") + def get_template_blm(self, it, it_e, lmaxb=1024, lmin_plm=1, elm_wf:None or np.ndarray=None, dlm_mod=None, calc=False, Nmf=None, + perturbative=False): + """Builds a template B-mode map with the iterated phi and input elm_wf + + Args: + it: iteration index of lensing tracer + it_e: iteration index of E-tracer (for best results use it_e = it + 1) + elm_wf: Wiener-filtered E-mode (healpy alm array), if not an iterated solution (it_e will ignored if set) + lmin_plm: the lensing tracer is zeroed below lmin_plm + lmaxb: the B-template is calculated up to lmaxb (defaults to lmax elm_wf) + perturbative: use pertubative instead of full remapping if set (may be useful for QE) + + Returns: + blm healpy array + + Note: + It can be a real lot better to keep the same L range as the iterations + + """ + cache_cond = (lmin_plm == 1) and (elm_wf is None) + # TODO this needs a cleaner implementation. Duplicate in map_delenser + if dlm_mod is not None: + dlm_mod_string = '_dlmmod' + else: + dlm_mod_string = '' + if Nmf == None: + pass + else: + dlm_mod_string += "{:03d}".format(Nmf) + fn = 'btempl_p%03d_e%03d_lmax%s%s' % (it, it_e, lmaxb, dlm_mod_string) + fn += 'perturbative' * perturbative + if not calc: + if self.wf_cacher.is_cached(fn): + return self.wf_cacher.load(fn) + if elm_wf is None: + if it_e > 0: + e_fname = 'wflm_%s_it%s' % ('p', it_e - 1) + assert self.wf_cacher.is_cached(e_fname) + elm_wf = self.wf_cacher.load(e_fname) + elif it_e == 0: + elm_wf = self.wflm0() + else: + assert 0,'dont know what to do with it_e = ' + str(it_e) + assert Alm.getlmax(elm_wf.size, self.mmax_filt) == self.lmax_filt + mmaxb = lmaxb + dlm = self.get_hlm(it) + + # subtract field from phi + if dlm_mod is not None: + dlm -= dlm_mod + self.hlm2dlm(dlm, inplace=True) + assert self.lmaxs_qlm[0] == self.lmaxs_qlm[1] + assert self.mmaxs_qlm[0] == self.mmaxs_qlm[1] + dlm.almxfl([np.arange(self.lmaxs_qlm[0] + 1, dtype=int) >= lmin_plm] * 2, True) + if perturbative: # Applies perturbative remapping + assert dlm.labels in [('p', 'x'), ('p',)], 'not implemented' + get_alm = lambda a: elm_wf if a == 'e' else np.zeros_like(elm_wf) + geom, sht_tr = self.filter.ffi.geom, self.filter.ffi.sht_tr + d1 = geom.alm2map_spin([dlm.get_comp('p'), dlm.get_comp('x')], 1, self.lmaxs_qlm[0], self.mmaxs_qlm[0], sht_tr, [-1., 1.]) + dp = utils_qe.qeleg_multi([2], +3, [utils_qe.get_spin_raise(2, self.lmax_filt)])(get_alm, geom, sht_tr) + dm = utils_qe.qeleg_multi([2], +1, [utils_qe.get_spin_lower(2, self.lmax_filt)])(get_alm, geom, sht_tr) + dlens = -0.5 * ((d1[0] - 1j * d1[1]) * dp + (d1[0] + 1j * d1[1]) * dm) + del dp, dm, d1 + elm, blm = geom.map2alm_spin([dlens.real, dlens.imag], 2, lmaxb, mmaxb, sht_tr, [-1., 1.]) + else: # Applies full remapping + assert dlm.labels in [('p', 'x'), ('p',)], 'not implemented' + ffi = self.filter.ffi.change_dlm([dlm.get_comp('p'), dlm.get_comp('x')], self.mmaxs_qlm[0]) + elm, blm = ffi.lensgclm(elm_wf, self.mmax_filt, 2, lmaxb, mmaxb) + if cache_cond: + self.wf_cacher.cache(fn, blm) + return blm + + def _get_ffi(self, itr): + dlm = self.hlm2dlm(self.get_hlm(itr), False) + if dlm.labels in [('p', 'x'), ('p',)]: + glm, clm = np.copy(dlm.get_comp('p')), np.copy(dlm.get_comp('x')) + assert self.lmaxs_qlm[0] == self.lmaxs_qlm[1] + assert self.mmaxs_qlm[0] == self.mmaxs_qlm[1] + ffi = self.filter.ffi.change_dlm([glm, clm], self.mmaxs_qlm[0], cachers.cacher_mem()) + return ffi + elif 'pee' in dlm.labels and 'p_eb' in dlm.labels: # EE and EB lensing components + ffi_ee = self.filter.ffi.change_dlm([dlm.get_comp('pee'), dlm.get_comp('xee')], self.mmaxs_qlm[0], cachers.cacher_mem()) + ffi_eb = self.filter.ffi.change_dlm([dlm.get_comp('p_eb'), dlm.get_comp('x_eb')], self.mmaxs_qlm[1], cachers.cacher_mem()) + return [ffi_ee, ffi_eb] + else: + assert 0, ('dont know what to do with labels ', dlm.labels) + + def get_hlm(self, itr): + """Loads current estimate of the anistropy sources. It is a complex array""" + if itr < 0: + return np.zeros(self.gradlm_size, dtype=complex) + fn = '%s_%slm_it%03d' % ({'p': 'phi', 'o': 'om'}['p'.lower()], self.h, itr) + if self.cacher.is_cached(fn): + return gradient.fromarray(self.cacher.load(fn), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + return gradient.fromarray(self._sk2plm(itr), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + + + def load_soltn(self, itr, key): + """Load starting point for the conjugate gradient inversion. + + """ + assert key.lower() in ['p', 'o'] + for i in np.arange(itr - 1, -1, -1): + fname = 'wflm_%s_it%s' % (key.lower(), i) + if self.wf_cacher.is_cached(fname): + return self.wf_cacher.load(fname), i + if callable(self.wflm0): + return self.wflm0(), -1 + # TODO: for MV this need a change + return np.zeros((1, Alm.getsize(self.lmax_filt, self.mmax_filt)), dtype=complex).squeeze(), -1 + + + def load_graddet(self, itr): + fn= '%slm_grad%sdet_it%03d' % (self.h, 'p'.lower(), itr) + return gradient.fromarray(self.cacher.load(fn), self.lmaxs_qlm, self.mmaxs_qlm) + + def load_gradpri(self, itr): + assert self.is_iter_done(itr -1) + ret = self.get_hlm(itr) + ret.almxfl([cli(chh) for chh in self.chhs], True) + return ret + + def load_gradquad(self, k): + fn = '%slm_grad%slik_it%03d' % (self.h, 'p'.lower(), k) + return gradient.fromarray(self.cacher.load(fn), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + + def load_gradient(self, itr): + """Loads the total gradient at iteration iter. + + All necessary alm's must have been calculated previously + + """ + if itr == 0: + g = self.load_gradpri(0) + g += self.load_graddet(0) + g += self.load_gradquad(0) + return g + return gradient.fromarray(self._yk2grad(itr), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + + def dotop(self, glms1:np.ndarray or gradient, glms2:np.ndarray or gradient): + if isinstance(glms1, gradient) and isinstance(glms2, gradient): + return gradient_dotop(glms1, glms2) + ret = 0. + N = 0 + for lmax, mmax in zip(self.lmaxs_qlm, self.mmaxs_qlm): + siz = Alm.getsize(lmax, mmax) + cl = alm2cl(glms1[N:N+siz], glms2[N:N+siz], None, mmax, None) + ret += np.sum(cl * (2 * np.arange(len(cl)) + 1)) + N += siz + return ret + + def calc_norm(self, qlm:np.ndarray): + return np.sqrt(self.dotop(qlm, qlm)) + + + def apply_H0k(self, grad_lm:np.ndarray, kr): + ret = np.empty_like(grad_lm) + N = 0 + for lmax, mmax, h0 in zip(self.lmaxs_qlm, self.mmaxs_qlm, self.hh_h0s): + siz = Alm.getsize(lmax, mmax) + ret[N:N+siz] = almxfl(grad_lm[N:N+siz], h0, mmax, False) + N += siz + return ret + + def apply_B0k(self, grad_lm:np.ndarray, kr): + ret = np.empty_like(grad_lm) + N = 0 + for lmax, mmax, h0 in zip(self.lmaxs_qlm, self.mmaxs_qlm, self.hh_h0s): + siz = Alm.getsize(lmax, mmax) + ret[N:N+siz] = almxfl(grad_lm[N:N+siz], cli(h0), mmax, False) #TOD0 this assumes >= 0 + N += siz + return ret + + @log_on_start(logging.INFO, "get_hessian() started: k={k}, key={key}") + @log_on_end(logging.INFO, "get_hessian() finished: k={k}, key={key}") + + def get_hessian(self, k, key): + """Inverse hessian that will produce phi_iter. + + + """ + # Zeroth order inverse hessian : + BFGS_H = bfgs.BFGS_Hessian(self.hess_cacher, self.apply_H0k, {}, {}, self.dotop, + L=self.NR_method, verbose=self.verbose, apply_B0k=self.apply_B0k) + # Adding the required y and s vectors : + for k_ in range(np.max([0, k - BFGS_H.L]), k): + BFGS_H.add_ys('rlm_yn_%s_%s' % (k_, key), 'rlm_sn_%s_%s' % (k_, key), k_) + return BFGS_H + + + @log_on_start(logging.INFO, "build_incr() started: it={it}, key={key}") + @log_on_end(logging.INFO, "build_incr() finished: it={it}, key={key}") + def build_incr(self, it, key, gradn:gradient): + """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 + :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 = gradn - self.load_gradient(k) + self.hess_cacher.cache(yk_fname, yk.getarray()) + 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 = gradient.fromarray(BFGS.get_mHkgk(gradn.getarray(), k), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + # giving the invertibility check that up at least for the moment + #incr = self.ensure_invertibility(self.get_hlm(it - 1), self.stepper.build_incr(incr, it)) + incr = self.stepper.build_incr(incr, it) + self.hess_cacher.cache(sk_fname, incr.getarray()) + prt_time(time.time() - t0, label=' Exec. time for descent direction calculation') + assert self.hess_cacher.is_cached(sk_fname), sk_fname + + + @log_on_start(logging.INFO, "iterate() started: it={itr}, key={key}") + @log_on_end(logging.INFO, "iterate() finished: it={itr}, key={key}") + 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): + assert self.is_iter_done(itr - 1), 'previous iteration not done' + self.logger.on_iterstart(itr, key, self) + # Calculation in // of lik and det term : + glm = self.calc_gradlik(itr, key) + glm += self.calc_graddet(itr) + glm += self.load_gradpri(itr - 1) + glm.almxfl([chh > 0 for chh in self.chhs], True) # kills all modes where priors are set to zero + 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))) + + + @log_on_start(logging.INFO, "calc_gradlik() started: it={itr}, key={key}") + @log_on_end(logging.INFO, "calc_gradlik() finished: it={itr}, key={key}") + 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) + 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' + self.filter.set_ffi(self._get_ffi(itr - 1)) + 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() + q_geom = pbdGeometry(self.k_geom, pbounds(0., 2 * np.pi)) + Gs, Cs = self.filter.get_qlms(self.dat_maps, soltn, q_geom) + + # GC can either G, C for a simple component gradient, or (G1, G2), (C1, C2) etc for a multicomponent + # TODO match G and Cs to label components in general case + if self.labels in [('p', 'x'), ('p',), ('x',), ('x', 'p')]: + grad_lm = gradient( [-Gs* ('p' in self.labels), -Cs * ('x' in self.labels)], self.mmaxs_qlm, labels=self.labels) + grad_lm.almxfl([self._h2p(self.lmaxs_qlm[0]), self._h2p(self.lmaxs_qlm[1])], True) + elif self.labels == ('pee', 'p_eb'): + if self.lm_maxee is not None and self.lm_maxee[0] < self.lmaxs_qlm[0]: + print("seeing smaller lmax for ee, patching") + # ee only up to some lmax, then total gradient + felp = np.ones(self.lmaxs_qlm[0] + 1) * (np.arange(self.lmaxs_qlm[0] + 1) > self.lm_maxee[0]) + felm = np.ones(self.lmaxs_qlm[0] + 1) * (np.arange(self.lmaxs_qlm[0] + 1) <= self.lm_maxee[0]) + Gee = almxfl(Gs[0], felm, self.mmaxs_qlm[0], False) + G_p = Gs[1] + alm_copy(almxfl(Gs[0], felp, self.mmaxs_qlm[0], False), self.mmaxs_qlm[0], self.lmaxs_qlm[1], self.mmaxs_qlm[1]) + G_ee =Gee + almxfl(alm_copy(G_p, self.mmaxs_qlm[1], self.lmaxs_qlm[0], self.mmaxs_qlm[0]), felp, self.mmaxs_qlm[0], False) + grad_lm = gradient([-G_ee, -G_p], self.mmaxs_qlm, labels=self.labels) + + else: + grad_lm = gradient([-Gs[0], -Gs[1]], self.mmaxs_qlm, labels=self.labels) + del Cs + else: + assert 0, ('dont know what to do with ',self.labels) + 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, grad_lm.getarray()) + return grad_lm + + @log_on_start(logging.INFO, "calc_graddet() started: it={itr}. subclassed") + @log_on_end(logging.INFO, "calc_graddet() finished: it={itr}. subclassed") + def calc_graddet(self, itr): + assert 0, 'subclass this' + + +class iterator_cstmf(gclm_iterator): + """Constant mean-field + + + """ + + def __init__(self, lib_dir:str, h:str, lm_max_dlm:list[tuple], + dat_maps:list or np.ndarray, plm0s:list or np.ndarray, mf0s:list or np.ndarray, pp_h0s:list[np.ndarray], + cpp_priors:list[np.ndarray], labels:tuple[str], cls_filt:dict, + ninv_filt:opfilt_base.alm_filter_wl, + k_geom:Geom, + chain_descr, stepper:nrstep,**kwargs): + super(iterator_cstmf, self).__init__(lib_dir, h, lm_max_dlm, dat_maps, plm0s, pp_h0s, cpp_priors, labels, cls_filt, + ninv_filt, k_geom, chain_descr, stepper, **kwargs) + + if not self.cacher.is_cached('mf'): + if not isinstance(mf0s, list): + mf0s = [mf0s] + mf0 = gradient(mf0s, self.mmaxs_qlm) + self.cacher.cache('mf', mf0.almxfl([self._h2p(self.lmaxs_qlm[0]), self._h2p(self.lmaxs_qlm[1])],False).getarray()) + + + @log_on_start(logging.INFO, "load_graddet() started: it={k}") + @log_on_end(logging.INFO, "load_graddet() finished: it={k}") + def load_graddet(self, k): + return gradient.fromarray(self.cacher.load('mf'), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + + @log_on_start(logging.INFO, "calc_graddet() started: it={k}") + @log_on_end(logging.INFO, "calc_graddet() finished: it={k}") + def calc_graddet(self, k): + return gradient.fromarray(self.cacher.load('mf'), self.lmaxs_qlm, self.mmaxs_qlm, labels=self.labels) + + +# TODO add visitor pattern if desired \ No newline at end of file diff --git a/delensalot/core/opfilt/MAP_opfilt_iso_p.py b/delensalot/core/opfilt/MAP_opfilt_iso_p.py index 6a91be41..b5577b0d 100644 --- a/delensalot/core/opfilt/MAP_opfilt_iso_p.py +++ b/delensalot/core/opfilt/MAP_opfilt_iso_p.py @@ -295,7 +295,7 @@ def _get_gpmap(self, elm_wf:np.ndarray, spin:int, q_pbgeom:pbdGeometry): fl[:spin] *= 0. fl = np.sqrt(fl) elm = np.atleast_2d(almxfl(elm_wf, fl, self.mmax_sol, False)) - ffi = self.ffi.change_geom(q_pbgeom.geom) if q_pbgeom is not self.ffi.pbgeom else self.ffi + ffi = self.ffi.change_geom(q_pbgeom.geom) if q_pbgeom.geom is not self.ffi.pbgeom.geom else self.ffi return ffi.gclm2lenmap(elm, self.mmax_sol, spin, False) class pre_op_diag: diff --git a/delensalot/core/opfilt/MAP_opfilt_iso_t.py b/delensalot/core/opfilt/MAP_opfilt_iso_t.py index 246e2873..31b75a58 100644 --- a/delensalot/core/opfilt/MAP_opfilt_iso_t.py +++ b/delensalot/core/opfilt/MAP_opfilt_iso_t.py @@ -202,7 +202,7 @@ def _get_gtmap(self, tlm_wf:np.ndarray, q_pbgeom:utils_geom.pbdGeometry): """ assert Alm.getlmax(tlm_wf.size, self.mmax_sol) == self.lmax_sol, ( Alm.getlmax(tlm_wf.size, self.mmax_sol), self.lmax_sol) fl = -np.sqrt(np.arange(self.lmax_sol + 1) * np.arange(1, self.lmax_sol + 2)) - ffi = self.ffi.change_geom(q_pbgeom) if q_pbgeom is not self.ffi.pbgeom else self.ffi + ffi = self.ffi.change_geom(q_pbgeom.geom) if q_pbgeom.geom is not self.ffi.pbgeom.geom else self.ffi return ffi.gclm2lenmap([almxfl(tlm_wf, fl, self.mmax_sol, False), np.zeros_like(tlm_wf)], self.mmax_sol, 1, False) diff --git a/delensalot/core/opfilt/MAP_opfilt_iso_tp.py b/delensalot/core/opfilt/MAP_opfilt_iso_tp.py index 0813fdc9..fdd3cc3e 100644 --- a/delensalot/core/opfilt/MAP_opfilt_iso_tp.py +++ b/delensalot/core/opfilt/MAP_opfilt_iso_tp.py @@ -298,7 +298,7 @@ def _get_gpmap(self, elm_wf:np.ndarray, spin:int, q_pbgeom:pbdGeometry): fl[:spin] *= 0. fl = np.sqrt(fl) elm = np.atleast_2d(almxfl(elm_wf, fl, self.mmax_sol, False)) - ffi = self.ffi.change_geom(q_pbgeom.geom) if q_pbgeom is not self.ffi.pbgeom else self.ffi + ffi = self.ffi.change_geom(q_pbgeom.geom) if q_pbgeom.geom is not self.ffi.pbgeom.geom else self.ffi return ffi.gclm2lenmap(elm, self.mmax_sol, spin, False) def _get_irestmap(self, tlm_dat:np.ndarray, tlm_wf:np.ndarray, q_pbgeom: pbdGeometry): diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py new file mode 100644 index 00000000..b4dd4a16 --- /dev/null +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -0,0 +1,396 @@ +"""Iterative reconstruction for masked polarization CMB data + + tests joint lensing gradient and curl potential reconstruction + + e.g. python ./run_fromparfile_wcurl.py -itmax 0 -v 'wcurlin' + + -v '' version is standard gradient reconstruction + -v 'wcurl' version reconstructs both gradient and curl on gradient-only input map + -v 'wcurlin' version reconstructs both gradient and curl on gradient and curl input map + +""" +import os +from os.path import join as opj +import numpy as np +from psutil import cpu_count +import plancklens +from plancklens import utils +from plancklens import qresp +from plancklens import qest, qecl +from plancklens.qcinv import cd_solve +from plancklens.sims import phas, maps +from plancklens.sims.cmbs import sims_cmb_unl +from plancklens.filt import filt_simple, filt_util + +from lenspyx.remapping.deflection import deflection +from lenspyx.remapping.utils_geom import Geom, pbdGeometry, pbounds +from lenspyx.utils import cli +from lenspyx.utils_hp import gauss_beam, almxfl, alm2cl, alm_copy +from lenspyx import cachers +from lenspyx.sims import sims_cmb_len + +from delensalot.utility import utils_steps +from delensalot.core.iterator import steps +from delensalot.core import mpi +from delensalot.core.opfilt.MAP_opfilt_iso_t import alm_filter_nlev_wl as alm_filter_nlev_wl_t +from delensalot.core.opfilt.MAP_opfilt_iso_p import alm_filter_nlev_wl +from delensalot.core.opfilt.MAP_opfilt_iso_tp import alm_filter_nlev_wl as alm_filter_nlev_wl_tp + +from delensalot.core.iterator.cs_iterator import iterator_cstmf as iterator_cstmf +from delensalot.core.iterator.cs_iterator_multi import iterator_cstmf as iterator_cstmf_wcurl + +suffix = 'delensalot_idealized_curly' # descriptor to distinguish this parfile from others... +TEMP = opj(os.environ['SCRATCH'], suffix) +DATDIR = opj(os.environ['SCRATCH'], suffix, 'sims') +DATDIRwcurl = opj(os.environ['SCRATCH'],suffix, 'simswcurl') + +if not os.path.exists(DATDIR): + os.makedirs(DATDIR) +# harmonic space noise phas down to 4096 +noise_phas = phas.lib_phas(opj(os.environ['HOME'], 'noisephas_lmax%s'%4096), 3, 4096) # T, E, and B noise phases +cmb_phas = phas.lib_phas(opj(os.environ['HOME'], 'cmbphas_lmax%s'%5120), 5, 5120) # unlensed T E B P O CMB phases + +lmax_ivf, mmax_ivf, beam, nlev_t, nlev_p = (4096, 4096, 1., 0.5 / np.sqrt(2), 0.5) +lmin_tlm, lmin_elm, lmin_blm = (1, 2, 2) # The fiducial transfer functions are set to zero below these lmins +# for delensing useful to cut much more B. It can also help since the cg inversion does not have to reconstruct those. + +lmax_qlm, mmax_qlm = (5120, 5120) # Lensing map is reconstructed down to this lmax and mmax +# NB: the QEs from plancklens does not support mmax != lmax, but the MAP pipeline does +lmax_unl, mmax_unl = (5120, 5120) # Delensed CMB is reconstructed down to this lmax and mmax + + +#----------------- pixelization and geometry info for the input maps and the MAP pipeline and for lensing operations +lenjob_geometry = Geom.get_thingauss_geometry(lmax_unl * 2, 2) +lenjob_pbgeometry = pbdGeometry(lenjob_geometry, pbounds(0., 2 * np.pi)) +Lmin = 1 # The reconstruction of all lensing multipoles below that will not be attempted +mc_sims_mf_it0 = np.array([]) # sims to use to build the very first iteration mean-field (QE mean-field) Here 0 since idealized + + +# Multigrid chain descriptor +# The hard coded number nside 2048 here is irrelevant for diagonal preconditioner +chain_descrs = lambda lmax_sol, cg_tol : [[0, ["diag_cl"], lmax_sol, 2048, np.inf, cg_tol, cd_solve.tr_cg, cd_solve.cache_mem()]] +libdir_iterators = lambda qe_key, simidx, version: opj(TEMP,'%s_sim%04d'%(qe_key, simidx) + version) +#------------------ + +# Fiducial CMB spectra for QE and iterative reconstructions +# (here we use very lightly suboptimal lensed spectra QE weights) +cls_path = opj(os.path.dirname(plancklens.__file__), 'data', 'cls') +cls_unl = utils.camb_clfile(opj(cls_path, 'FFP10_wdipole_lenspotentialCls.dat')) +cls_len = utils.camb_clfile(opj(cls_path, 'FFP10_wdipole_lensedCls.dat')) +cls_unl_wcurl = utils.camb_clfile(opj(cls_path, 'FFP10_wdipole_lenspotentialCls.dat')) +cls_unl_wcurl['oo'] = np.loadtxt(opj(cls_path, 'FFP10_fieldrotationCls.dat')) # lensing curl potential + +gradcls = utils.camb_clfile(opj(cls_path, 'FFP10_wdipole_gradlensedCls.dat')) + +# Fiducial model of the transfer function +transf_tlm = gauss_beam(beam/180 / 60 * np.pi, lmax=lmax_ivf) * (np.arange(lmax_ivf + 1) >= lmin_tlm) +transf_elm = gauss_beam(beam/180 / 60 * np.pi, lmax=lmax_ivf) * (np.arange(lmax_ivf + 1) >= lmin_elm) +transf_blm = gauss_beam(beam/180 / 60 * np.pi, lmax=lmax_ivf) * (np.arange(lmax_ivf + 1) >= lmin_blm) +transf_d = {'t':transf_tlm, 'e':transf_elm, 'b':transf_blm} +# Isotropic approximation to the filtering (used eg for response calculations) +ftl = cli(cls_len['tt'][:lmax_ivf + 1] + (nlev_t / 180 / 60 * np.pi) ** 2 * cli(transf_tlm ** 2)) * (transf_tlm > 0) +fel = cli(cls_len['ee'][:lmax_ivf + 1] + (nlev_p / 180 / 60 * np.pi) ** 2 * cli(transf_elm ** 2)) * (transf_elm > 0) +fbl = cli(cls_len['bb'][:lmax_ivf + 1] + (nlev_p / 180 / 60 * np.pi) ** 2 * cli(transf_blm ** 2)) * (transf_blm > 0) + +# Same using unlensed spectra (used for unlensed response used to initiate the MAP curvature matrix) +ftl_unl = cli(cls_unl['tt'][:lmax_ivf + 1] + (nlev_t / 180 / 60 * np.pi) ** 2 * cli(transf_tlm ** 2)) * (transf_tlm > 0) +fel_unl = cli(cls_unl['ee'][:lmax_ivf + 1] + (nlev_p / 180 / 60 * np.pi) ** 2 * cli(transf_elm ** 2)) * (transf_elm > 0) +fbl_unl = cli(cls_unl['bb'][:lmax_ivf + 1] + (nlev_p / 180 / 60 * np.pi) ** 2 * cli(transf_blm ** 2)) * (transf_blm > 0) + +# ------------------------- +# ---- Input simulation libraries. Here we use the NERSC FFP10 CMBs with homogeneous noise and consistent transfer function +# We define explictly the phase library such that we can use the same phases for for other purposes in the future as well if needed +# I am putting here the phases in the home directory such that they dont get NERSC auto-purged +# actual data transfer function for the sim generation: +transf_dat = gauss_beam(beam / 180 / 60 * np.pi, lmax=4096) # (taking here full FFP10 cmb's which are given to 4096) +cls_noise = {'t': np.full(4097, (nlev_t /180 / 60 * np.pi) ** 2) * (cls_len['tt'][:4097] > 0), + 'e': np.full(4097, (nlev_p / 180 / 60 * np.pi) ** 2) * (cls_len['ee'][:4097] > 0), + 'b': np.full(4097, (nlev_p / 180 / 60 * np.pi) ** 2) * (cls_len['bb'][:4097] > 0),} +cls_transf = {f: transf_dat for f in ['t', 'e', 'b']} +if mpi.rank ==0: + # Problem of creating dir in parallel if does not exist + cacher = cachers.cacher_npy(DATDIR) + cacher_wcurl = cachers.cacher_npy(DATDIRwcurl) +mpi.barrier() + +cacher = cachers.cacher_npy(DATDIR) +cacher_wcurl = cachers.cacher_npy(DATDIRwcurl) + +cmb_unl = sims_cmb_unl(cls_unl, cmb_phas) +cmb_unl_wcurl = sims_cmb_unl(cls_unl_wcurl, cmb_phas) + +cmb_len = sims_cmb_len(4096, cmb_unl, cache=cacher, epsilon=1e-7) +cmb_len_wcurl = sims_cmb_len(4096, cmb_unl_wcurl, cache=cacher_wcurl, epsilon=1e-7) + +sims = maps.cmb_maps_harmonicspace(cmb_len, cls_transf, cls_noise, noise_phas) +sims_wcurl = maps.cmb_maps_harmonicspace(cmb_len_wcurl, cls_transf, cls_noise, noise_phas) +# ------------------------- + +ivfs = filt_simple.library_fullsky_alms_sepTP(opj(TEMP, 'ivfs'), sims, transf_d, cls_len, ftl, fel, fbl, cache=True) +ivfs_wcurl = filt_simple.library_fullsky_alms_sepTP(opj(TEMP, 'ivfs_wcurl'), sims_wcurl, transf_d, cls_len, ftl, fel, fbl, cache=True) + +# ---- QE libraries from plancklens to calculate unnormalized QE (qlms) and their spectra (qcls) +mc_sims_bias = np.arange(0, dtype=int) +mc_sims_var = np.arange(0, 60, dtype=int) +qlms_dd = qest.library_sepTP(opj(TEMP, 'qlms_dd'), ivfs, ivfs, cls_len['te'], 2048, lmax_qlm=lmax_qlm) +qcls_dd = qecl.library(opj(TEMP, 'qcls_dd'), qlms_dd, qlms_dd, mc_sims_bias) + + +qlms_dd_wcurl = qest.library_sepTP(opj(TEMP, 'qlms_dd_wcurl'), ivfs_wcurl, ivfs_wcurl, cls_len['te'], 2048, lmax_qlm=lmax_qlm) +qcls_dd_wcurl = qecl.library(opj(TEMP, 'qcls_dd_wcurl'), qlms_dd_wcurl, qlms_dd_wcurl, mc_sims_bias) + +# ------------------------- +# This following block is only necessary if a full, Planck-like QE lensing power spectrum analysis is desired +# This uses 'ds' and 'ss' QE's, crossing data with sims and sims with other sims. + +# This remaps idx -> idx + 1 by blocks of 60 up to 300. This is used to remap the sim indices for the 'MCN0' debiasing term in the QE spectrum +ss_dict = { k : v for k, v in zip( np.concatenate( [ range(i*60, (i+1)*60) for i in range(0,5) ] ), + np.concatenate( [ np.roll( range(i*60, (i+1)*60), -1 ) for i in range(0,5) ] ) ) } +ds_dict = { k : -1 for k in range(300)} # This remap all sim. indices to the data maps to build QEs with always the data in one leg + +ivfs_d = filt_util.library_shuffle(ivfs, ds_dict) +ivfs_s = filt_util.library_shuffle(ivfs, ss_dict) + +qlms_ds = qest.library_sepTP(opj(TEMP, 'qlms_ds'), ivfs, ivfs_d, cls_len['te'], 2048, lmax_qlm=lmax_qlm) +qlms_ss = qest.library_sepTP(opj(TEMP, 'qlms_ss'), ivfs, ivfs_s, cls_len['te'], 2048, lmax_qlm=lmax_qlm) + +qcls_ds = qecl.library(opj(TEMP, 'qcls_ds'), qlms_ds, qlms_ds, np.array([])) # for QE RDN0 calculations +qcls_ss = qecl.library(opj(TEMP, 'qcls_ss'), qlms_ss, qlms_ss, np.array([])) # for QE RDN0 / MCN0 calculations + +def get_n0_iter(k='p_p'): + from plancklens import n0s + fnN0s = 'N0siter' + k * (k != 'p_p') + fndelcls = 'delcls'+ k * (k != 'p_p') + cachecond = True + if not cacher_wcurl.is_cached(fnN0s) or not cacher_wcurl.is_cached(fndelcls): + _, N0sg, _, N0c, _, delcls = n0s.get_N0_iter(k, nlev_t, nlev_p, beam, cls_unl, {'t':lmin_tlm, 'e':lmin_elm, 'b':lmin_blm}, lmax_ivf,10, ret_delcls=True, ret_curl=True, lmax_qlm=lmax_qlm) + if cachecond: + cacher_wcurl.cache(fnN0s, np.array([N0sg, N0c])) + cacher_wcurl.cache(fndelcls, np.array([delcls[-1][spec] for spec in ['ee', 'bb', 'pp']])) + return np.array([N0sg, N0c]), delcls + delcls = cacher_wcurl.load(fndelcls) + delclsdict = {'ee': delcls[0], 'bb':delcls[1], 'pp':delcls[2]} + return cacher_wcurl.load(fnN0s), delclsdict + +# ------------------------- + +def get_itlib(k:str, simidx:int, version:str, cg_tol:float): + """Return iterator instance for simulation idx and qe_key type k + Args: + k: 'p_p' for Pol-only, 'ptt' for T-only, 'p_eb' for EB-only, etc + simidx: simulation index to build iterative lensing estimate on + version: string to use to test variants of the iterator with otherwise the same parfile + (here if 'noMF' is in version, will not use any mean-fied at the very first step) + cg_tol: tolerance of conjugate-gradient filter + """ + libdir_iterator = libdir_iterators(k, simidx, version) + if not os.path.exists(libdir_iterator): + os.makedirs(libdir_iterator) + tr = int(os.environ.get('OMP_NUM_THREADS', cpu_count(logical=False))) + print("Using %s threads"%tr) + cpp = np.copy(cls_unl_wcurl['pp'][:lmax_qlm + 1]) + cpp[:Lmin] *= 0. + coo = np.copy(cls_unl_wcurl['oo'][:lmax_qlm + 1]) + coo[:Lmin] *= 0. + # QE mean-field fed in as constant piece in the iteration steps: + if 'wcurlin' in version: + qlms_dd_QE = qlms_dd_wcurl + sims_MAP = sims_wcurl + else: + qlms_dd_QE = qlms_dd + sims_MAP = sims + + mf_sims = np.unique(mc_sims_mf_it0 if not 'noMF' in version else np.array([])) + mf0_p = qlms_dd_QE.get_sim_qlm_mf('p' + k[1:], mf_sims) # Mean-field to subtract on the first iteration: + mf0_o = qlms_dd_QE.get_sim_qlm_mf('x' + k[1:], mf_sims) # Mean-field to subtract on the first iteration: + + if simidx in mf_sims: # We dont want to include the sim we consider in the mean-field... + Nmf = len(mf_sims) + mf0_p = (mf0_p - qlms_dd_QE.get_sim_qlm('p' + k[1:], int(simidx)) / Nmf) * (Nmf / (Nmf - 1)) + mf0_o = (mf0_o - qlms_dd_QE.get_sim_qlm('x' + k[1:], int(simidx)) / Nmf) * (Nmf / (Nmf - 1)) + + plm0 = qlms_dd_QE.get_sim_qlm('p' + k[1:], int(simidx)) - mf0_p # Unormalized quadratic estimate: + olm0 = qlms_dd_QE.get_sim_qlm('x' + k[1:], int(simidx)) - mf0_o # Unormalized quadratic estimate: + + # Isotropic normalization of the QE + Rpp, Roo = qresp.get_response(k, lmax_ivf, 'p', cls_len, cls_len, {'e': fel, 'b': fbl, 't': ftl}, + lmax_qlm=lmax_qlm)[0:2] + # Isotropic Wiener-filter (here assuming for simplicity N0 ~ 1/R) + WF_p = cpp * utils.cli(cpp + utils.cli(Rpp)) + WF_o = coo * utils.cli(coo + utils.cli(Roo)) + + plm0 = alm_copy(plm0, None, lmax_qlm, mmax_qlm) # Just in case the QE and MAP mmax'es were not consistent + almxfl(plm0, utils.cli(Rpp), mmax_qlm, True) # Normalized QE + almxfl(plm0, WF_p, mmax_qlm, True) # Wiener-filter QE + almxfl(plm0, cpp > 0, mmax_qlm, True) + + olm0 = alm_copy(olm0, None, lmax_qlm, mmax_qlm) # Just in case the QE and MAP mmax'es were not consistent + almxfl(olm0, utils.cli(Roo), mmax_qlm, True) # Normalized QE + almxfl(olm0, WF_o, mmax_qlm, True) # Wiener-filter QE assuming the curl signal is the expected one + almxfl(olm0, coo > 0, mmax_qlm, True) + + Rpp_unl, Roo_unl = qresp.get_response(k, lmax_ivf, 'p', cls_unl, cls_unl, + {'e': fel_unl, 'b': fbl_unl, 't': ftl_unl}, lmax_qlm=lmax_qlm)[0:2] + # Lensing deflection field instance (initiated here with zero deflection) + + ffi = deflection(lenjob_geometry, np.zeros_like(plm0), mmax_qlm, numthreads=tr, epsilon=1e-7) + if k in ['p_p']: + # Here multipole cuts are set by the transfer function (those with 0 are not considered) + filtr = alm_filter_nlev_wl(nlev_p, ffi, transf_elm, (lmax_unl, mmax_unl), (lmax_ivf, mmax_ivf), + transf_b=transf_blm, nlev_b=nlev_p) + # dat maps must now be given in harmonic space in this idealized configuration + eblm = np.array(sims_MAP.get_sim_pmap(int(simidx))) + datmaps = np.array([alm_copy(eblm[0], None, lmax_ivf, mmax_ivf), alm_copy(eblm[1], None, lmax_ivf, mmax_ivf) ]) + del eblm + wflm0 = lambda: alm_copy(ivfs_wcurl.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl) + elif k in ['p']: + filtr = alm_filter_nlev_wl_tp(nlev_t, nlev_p, ffi, transf_tlm, (lmax_unl, mmax_unl), (lmax_ivf, mmax_ivf), + transf_b=transf_blm, transf_e=transf_elm) + # dat maps must now be given in harmonic space in this idealized configuration + eblm = np.array(sims_MAP.get_sim_pmap(int(simidx))) + tlm = sims_MAP.get_sim_tmap(int(simidx)) + datmaps = np.array([alm_copy(tlm, None, lmax_ivf, mmax_ivf), + alm_copy(eblm[0], None, lmax_ivf, mmax_ivf), + alm_copy(eblm[1], None, lmax_ivf, mmax_ivf)]) + del tlm, eblm + wflm0 = lambda: np.array([alm_copy(ivfs_wcurl.get_sim_tmliklm(simidx), None, lmax_unl, mmax_unl), + alm_copy(ivfs_wcurl.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl)]) + + elif k in ['ptt']: + # Here multipole cuts are set by the transfer function (those with 0 are not considered) + filtr = alm_filter_nlev_wl_t(nlev_t, ffi, transf_tlm, (lmax_unl, mmax_unl), (lmax_ivf, mmax_ivf)) + # dat maps must now be given in harmonic space in this idealized configuration + tlm = sims_MAP.get_sim_tmap(int(simidx)) + datmaps = alm_copy(tlm, None, lmax_ivf, mmax_ivf) + del tlm + wflm0 = lambda: alm_copy(ivfs_wcurl.get_sim_tmliklm(simidx), None, lmax_unl, mmax_unl) + + else: + assert 0 + + + k_geom = filtr.ffi.geom # Customizable Geometry for position-space operations in calculations of the iterated QEs etc + if 'wcurl' in version: + # Sets to zero all L-modes below Lmin in the iterations: + assert 'wmf1' not in version + stepper = utils_steps.harmonicbump(xa=400, xb=1500) + iterator = iterator_cstmf_wcurl(libdir_iterator, 'p', [(lmax_qlm, mmax_qlm), (lmax_qlm, mmax_qlm)], datmaps, + [plm0, olm0], [mf0_p, mf0_o], [Rpp_unl, Roo_unl], [cpp, coo], ('p', 'x'), cls_unl, filtr, k_geom, + chain_descrs(lmax_unl, cg_tol), stepper, + wflm0=wflm0) + + else: # standard gradient only + stepper = steps.harmonicbump(lmax_qlm, mmax_qlm, xa=400, xb=1500) # reduce the gradient by 0.5 for large scale and by 0.1 for small scales to improve convergence in regimes where the deflection field is not invertible + iterator = iterator_cstmf(libdir_iterator, 'p', (lmax_qlm, mmax_qlm), datmaps, + plm0, plm0 * 0, Rpp_unl, cpp, cls_unl, filtr, k_geom, chain_descrs(lmax_unl, cg_tol), stepper + , wflm0=lambda : alm_copy(ivfs.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl)) + return iterator + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser(description='test iterator full-sky with pert. resp.') + parser.add_argument('-k', dest='k', type=str, default='p_p', help='rec. type') + parser.add_argument('-itmax', dest='itmax', type=int, default=-1, help='maximal iter index') + parser.add_argument('-tol', dest='tol', type=float, default=7., help='-log10 of cg tolerance default') + parser.add_argument('-imin', dest='imin', type=int, default=0, help='minimal sim index') + parser.add_argument('-imax', dest='imax', type=int, default=0, help='maximal sim index') + parser.add_argument('-v', dest='v', type=str, default='', help='iterator version') + parser.add_argument('-p', dest='plot', action='store_true', help='make some plots on the fly') + + + args = parser.parse_args() + tol_iter = lambda it : 10 ** (- args.tol) # tolerance a fct of iterations ? + soltn_cond = lambda it: True # Uses (or not) previous E-mode solution as input to search for current iteration one + + + mpi.barrier = lambda : 1 # redefining the barrier (Why ? ) + from delensalot.core.iterator.statics import rec as Rec + jobs = [] + for idx in np.arange(args.imin, args.imax + 1): + lib_dir_iterator = libdir_iterators(args.k, idx, args.v) + if Rec.maxiterdone(lib_dir_iterator) < args.itmax or args.plot: + jobs.append(idx) + + if mpi.rank ==0: + print("Caching things in " + TEMP) + + + for idx in jobs[mpi.rank::mpi.size]: + lib_dir_iterator = libdir_iterators(args.k, idx, args.v) + print("iterator folder: " + lib_dir_iterator) + + if args.itmax >= 0 and Rec.maxiterdone(lib_dir_iterator) < args.itmax: + itlib = get_itlib(args.k, idx, args.v, 1.) + for i in range(args.itmax + 1): + print("****Iterator: setting cg-tol to %.4e ****"%tol_iter(i)) + print("****Iterator: setting solcond to %s ****"%soltn_cond(i)) + itlib.chain_descr = chain_descrs(lmax_unl, tol_iter(i)) + itlib.soltn_cond = soltn_cond(i) + print("doing iter " + str(i)) + itlib.iterate(i, 'p') + + if args.plot and mpi.rank == 0: + import pylab as pl + pl.ion() + version = args.v + input_sims = sims if not 'wcurlin' in version else sims_wcurl + fig, axes = pl.subplots(1, 3, figsize=(15, 5)) + for idx in jobs[0:1]: # only first + print("plots") + lib_dir_iterator = libdir_iterators(args.k, idx, version) + itrs = np.unique(np.linspace(0, args.itmax + 1, 5, dtype=int)) # plotting max 5 curves + if args.itmax not in itrs: + itrs = np.concatenate([itrs, [args.itmax]]) + plms = Rec.load_plms(lib_dir_iterator, itrs) + plm_in = alm_copy(input_sims.sims_cmb_len.get_sim_plm(int(idx)), None, lmax_qlm, mmax_qlm) + cpp_in = alm2cl(plm_in, plm_in, lmax_qlm, mmax_qlm, lmax_qlm) + ls = np.arange(1, lmax_qlm + 1) + wls = ls ** 2 * (ls + 1) ** 2 / (2 * np.pi) + axes[0].set_title('auto-spectra') + axes[1].set_title('cross-spectra') + axes[2].set_title('cross-corr. coeff.') + axes[0].loglog(ls, wls * cpp_in[ls], c='k') + for itr, plms in zip(itrs, plms): + ncomp = 1 + ('wcurl' in args.v) + plm_size = len(plms) // ncomp + for icomp in [0]: + plm = plms[icomp * plm_size : (icomp + 1) * plm_size] # In the curly case there are two components + cxx = alm2cl(plm, plm_in, lmax_qlm, mmax_qlm, lmax_qlm) + cpp = alm2cl(plm, plm, lmax_qlm, mmax_qlm, lmax_qlm) + axes[0].loglog(ls, wls * cpp[ls], label='itr ' + str(itr)) + axes[1].loglog(ls, wls * cxx[ls], label='itr ' + str(itr)) + axes[2].semilogx(ls, cxx[ls] / np.sqrt(cpp * cpp_in)[ls], label='itr ' + str(itr)) + + for ax in axes: + ax.legend() + if 'wcurlin' in args.v: + fig, axes = pl.subplots(1, 3, figsize=(15, 5)) + for idx in jobs[0:1]: # only first + print("plots") + lib_dir_iterator = libdir_iterators(args.k, idx, version) + itrs = np.unique(np.linspace(0, args.itmax + 1, 5, dtype=int)) # plotting max 5 curves + if args.itmax not in itrs: + itrs = np.concatenate([itrs, [args.itmax]]) + plms = Rec.load_plms(lib_dir_iterator, itrs) + plm_in = alm_copy(input_sims.sims_cmb_len.get_sim_olm(int(idx)), None, lmax_qlm, mmax_qlm) + cpp_in = alm2cl(plm_in, plm_in, lmax_qlm, mmax_qlm, lmax_qlm) + ls = np.arange(2, lmax_qlm + 1) + wls = ls ** 2 * (ls + 1) ** 2 / (2 * np.pi) + axes[0].set_title('auto-spectra') + axes[1].set_title('cross-spectra') + axes[2].set_title('cross-corr. coeff.') + axes[0].loglog(ls, wls * cpp_in[ls], c='k') + for itr, plms in zip(itrs, plms): + ncomp = 1 + ('wcurl' in args.v) + plm_size = len(plms) // ncomp + for icomp in [1]: + plm = plms[icomp * plm_size : (icomp + 1) * plm_size] # In the curly case there are two components + cxx = alm2cl(plm, plm_in, lmax_qlm, mmax_qlm, lmax_qlm) + cpp = alm2cl(plm, plm, lmax_qlm, mmax_qlm, lmax_qlm) + axes[0].loglog(ls, wls * cpp[ls], label='itr ' + str(itr)) + axes[1].loglog(ls, wls * cxx[ls], label='itr ' + str(itr)) + axes[2].semilogx(ls, cxx[ls] / np.sqrt(cpp * cpp_in)[ls], label='itr ' + str(itr)) + + for ax in axes: + ax.legend() + k = input("press a key to exit") \ No newline at end of file diff --git a/delensalot/utility/utils_steps.py b/delensalot/utility/utils_steps.py new file mode 100644 index 00000000..7dadad57 --- /dev/null +++ b/delensalot/utility/utils_steps.py @@ -0,0 +1,177 @@ +"""Module to handle steps in abstract 'gradient' format + + +""" +from __future__ import annotations +from lenspyx.utils_hp import Alm, almxfl, alm2cl +import numpy as np + + +class gradient: + """Class to abstract away the handling of the joint reconstruction of several fields + + """ + def __init__(self, componentlist:list[np.ndarray], mmax_list:list[int], labels:tuple[str]=('p', 'x')): + """ + Args: + componentlist: list of the gradient components, all healpy-like alm arrays + mmax_list: mmax's of the arrays + labels: strings to distinguish components if needed + + + """ + self.comps = componentlist + self.mmaxs = mmax_list + self.lmaxs = [Alm.getlmax(alm.size, mmax) for alm, mmax in zip(componentlist, mmax_list)] + self.labels = labels[:len(self.comps)] + + def almxfl(self, cl_list:list[np.ndarray], inplace): + assert len(cl_list) == len(self.comps), (len(cl_list), len(self.comps)) + if inplace: + for alm, cl, mmax in zip(self.comps, cl_list, self.mmaxs): + almxfl(alm, cl, mmax, True) + else: + comps = [almxfl(alm, cl, mmax, False) for alm, cl, mmax in zip(self.comps, cl_list, self.mmaxs)] + return gradient(comps, self.mmaxs, labels=self.labels) + + def alm2cl(self): + return [alm2cl(alm, alm, None, mmax, None) for alm, mmax in zip(self.comps, self.mmaxs)] + + def get_comp(self, label:str): + if label in self.labels: + return self.comps[self.labels.index(label)] + elif label[0] == 'x' and 'p' + label[1:] in self.labels: + return np.zeros_like(self.get_comp('p' + label[1:])) + else: + assert 0, 'no ' + label + ' component in this gradient' + + def getarray(self): + return np.concatenate(self.comps) + + @staticmethod + def fromarray(arr, lmaxs:list[int], mmaxs:list[int], labels=None): + """Builds instance from its array 'getarray' """ + N = 0 + comps = [] + for lmax, mmax in zip(lmaxs, mmaxs): + size = Alm.getsize(lmax, mmax) + comps.append(arr[N:N+size]) + N += size + return gradient(comps, mmaxs, labels=labels) + + def copy(self): + return gradient([np.copy(alm) for alm in self.comps], self.mmaxs, labels=self.labels) + + def __mul__(self, other): + if np.isscalar(other): + return gradient([alm * other for alm in self.comps], self.mmaxs, labels=self.labels) + else: + assert 0, 'not implemented' + + def __truediv__(self, other): + if np.isscalar(other): + return gradient([alm / other for alm in self.comps], self.mmaxs, labels=self.labels) + else: + assert 0, 'not implemented' + + def __sub__(self, other): + if isinstance(other, gradient): + assert other.lmaxs == self.lmaxs # not necessary but hard to think when this should not be a bug + assert other.mmaxs == self.mmaxs + return gradient([alm - blm for alm, blm in zip(self.comps, other.comps)], self.mmaxs, labels=self.labels) + else: + assert 0, 'not implemented' + + def __add__(self, other): + if isinstance(other, gradient): + assert other.lmaxs == self.lmaxs # not necessary but hard to think when this should not be a bug + assert other.mmaxs == self.mmaxs + return gradient([alm + blm for alm, blm in zip(self.comps, other.comps)], self.mmaxs, labels=self.labels) + else: + assert 0, 'not implemented' + + def __iadd__(self, other): + if isinstance(other, gradient): + assert other.lmaxs == self.lmaxs # not necessary but hard to think when this should not be a bug + assert other.mmaxs == self.mmaxs + for alm, blm in zip(self.comps, other.comps): + alm += blm + return self + else: + assert 0, 'not implemented' + + def __isub__(self, other): + if isinstance(other, gradient): + assert other.lmaxs == self.lmaxs # not necessary but hard to think when this should not be a bug + assert other.mmaxs == self.mmaxs + for alm, blm in zip(self.comps, other.comps): + alm -= blm + return self + else: + assert 0, 'not implemented' + + def __imul__(self, other): + if np.isscalar(other): + for alm in self.comps: + alm *= other + return self + else: + assert 0, 'not implemented' + + def __itruediv__(self, other): + if np.isscalar(other): + for alm in self.comps: + alm /= other + return self + else: + assert 0, 'not implemented' + + def __neg__(self): + return gradient([-alm for alm in self.comps], self.mmaxs, labels=self.labels) + + def __pos__(self): + return self.copy() + +def gradient_dotop(g1:gradient, g2:gradient): + assert g1.mmaxs == g2.mmaxs + assert g1.lmaxs == g2.lmaxs + ret = 0. + for alm1, alm2, mmax in zip(g1.comps, g2.comps, g1.mmaxs): + cl = alm2cl(alm1, alm2, None, mmax, None) + ret += np.sum(cl * (2 * np.arange(len(cl)) + 1 )) + return ret + +class nrstep(object): + def __init__(self, val=1.): + self.val = val + + + def build_incr(self, incrlm:gradient, itr:int): + return incrlm * self.val + +class harmonicbump(nrstep): + def __init__(self, xa=400, xb=1500, a=0.5, b=0.1, scale=50): + """Harmonic bumpy step that were useful for s06b and s08b + + """ + self.scale = scale + self.bump_params = (xa, xb, a, b) + + def steplen(self, lmax_qlm): + xa, xb, a, b = self.bump_params + return self.bp(np.arange(lmax_qlm + 1),xa, a, xb, b, scale=self.scale) + + + def build_incr(self, incrlm:gradient, itr:int): + incrlm.almxfl([self.steplen(lmax) for lmax in incrlm.lmaxs], True) + return incrlm + + @staticmethod + def bp(x, xa, a, xb, b, scale=50): + """Bump function with f(xa) = a and f(xb) = b with transition at midpoint over scale scale + + """ + x0 = (xa + xb) * 0.5 + r = lambda x_: np.arctan(np.sign(b - a) * (x_ - x0) / scale) + np.sign(b - a) * np.pi * 0.5 + return a + r(x) * (b - a) / r(xb) +