From dc2bc464aa62ee114b50749a9f3176d8ff13af78 Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 17:36:29 +0200 Subject: [PATCH 01/12] wcurl recs and parfile --- delensalot/core/iterator/cs_iterator_multi.py | 591 ++++++++++++++++++ delensalot/scripts/run_fromparfile_wcurl.py | 340 ++++++++++ delensalot/utility/utils_steps.py | 177 ++++++ 3 files changed, 1108 insertions(+) create mode 100644 delensalot/core/iterator/cs_iterator_multi.py create mode 100644 delensalot/scripts/run_fromparfile_wcurl.py create mode 100644 delensalot/utility/utils_steps.py diff --git a/delensalot/core/iterator/cs_iterator_multi.py b/delensalot/core/iterator/cs_iterator_multi.py new file mode 100644 index 00000000..0f653f34 --- /dev/null +++ b/delensalot/core/iterator/cs_iterator_multi.py @@ -0,0 +1,591 @@ +"""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, np.zeros_like(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 + assert soltn.ndim == 1, 'Fix following lines' + 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/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py new file mode 100644 index 00000000..ca5fef6e --- /dev/null +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -0,0 +1,340 @@ +"""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 '' + + '' version is standard gradient reconstruction + 'wcurl' version reconstructs both gradient and curl on gradient-only input map + 'wcurlin' version reconstructs both gradient and curl on gradient and curl input map + (starting point still the curl-free QE though) + +""" +import os +from os.path import join as opj +import numpy as np +from multiprocessing 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.core.iterator import steps +from delensalot.core import mpi +from delensalot.core.opfilt.MAP_opfilt_iso_p import alm_filter_nlev_wl +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 + """ + assert k in ['p_eb', 'p_p'], k + 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() // 2)) + 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) + + + if k in ['p_p']: + ffi = deflection(lenjob_geometry, np.zeros_like(plm0), mmax_qlm, numthreads=tr, epsilon=1e-7) + # 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 + 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 = steps.harmonicbump(lmax_qlm, mmax_qlm, 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=lambda : alm_copy(ivfs.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl)) + + 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=1, 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 + 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(sims.sims_cmb_len.get_sim_plm(int(idx)), 5120, 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, plm in zip(itrs, plms): + 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 close 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) + From 84ad5f526814f52b6a866bbbb6641cac4c2fe72d Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 17:56:36 +0200 Subject: [PATCH 02/12] minor fix --- delensalot/core/opfilt/MAP_opfilt_iso_p.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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: From 7f21255daf8316521210703a1f5d8a1205e98ec8 Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 18:21:08 +0200 Subject: [PATCH 03/12] minor fixes --- delensalot/scripts/run_fromparfile_wcurl.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index ca5fef6e..28b05655 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -2,7 +2,7 @@ tests joint lensing gradient and curl potential reconstruction - e.g. python ./run_fromparfile_wcurl.py -itmax 0 -v '' + e.g. python ./run_fromparfile_wcurl.py -itmax 0 -v 'wcurlin' '' version is standard gradient reconstruction 'wcurl' version reconstructs both gradient and curl on gradient-only input map @@ -31,6 +31,7 @@ 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_p import alm_filter_nlev_wl @@ -250,7 +251,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): if 'wcurl' in version: # Sets to zero all L-modes below Lmin in the iterations: assert 'wmf1' not in version - stepper = steps.harmonicbump(lmax_qlm, mmax_qlm, xa=400, xb=1500) + 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, @@ -301,7 +302,6 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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)) @@ -328,7 +328,8 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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, plm in zip(itrs, plms): + for itr, plms in zip(itrs, plms): + plm = np.atleast_2d(plms)[0] # 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)) From 8fe9a8e12596efb20e401fd9b173b9ea6b571984 Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 18:43:09 +0200 Subject: [PATCH 04/12] minor fixes --- delensalot/scripts/run_fromparfile_wcurl.py | 53 +++++++++++++++++---- 1 file changed, 43 insertions(+), 10 deletions(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index 28b05655..8455f678 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -271,7 +271,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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=1, help='maximal 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') @@ -311,7 +311,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): import pylab as pl pl.ion() version = args.v - input_sims = sims + 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") @@ -320,7 +320,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): if args.itmax not in itrs: itrs = np.concatenate([itrs, [args.itmax]]) plms = Rec.load_plms(lib_dir_iterator, itrs) - plm_in = alm_copy(sims.sims_cmb_len.get_sim_plm(int(idx)), 5120, lmax_qlm, mmax_qlm) + 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) @@ -329,13 +329,46 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): axes[2].set_title('cross-corr. coeff.') axes[0].loglog(ls, wls * cpp_in[ls], c='k') for itr, plms in zip(itrs, plms): - plm = np.atleast_2d(plms)[0] # 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)) + 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() - k = input("press close to exit") \ No newline at end of file + 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(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 [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 From 443ecca216093fbbebf6b9e424b0599644d510b2 Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 18:46:54 +0200 Subject: [PATCH 05/12] minor fix --- delensalot/scripts/run_fromparfile_wcurl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index 8455f678..1f96655e 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -255,7 +255,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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=lambda : alm_copy(ivfs.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl)) + wflm0=lambda : alm_copy(ivfs_wcurl.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl)) 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 From 1e7ea21ee6ba913d063b24c72d35d75ac1e29d66 Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 19:05:29 +0200 Subject: [PATCH 06/12] minor fix --- delensalot/scripts/run_fromparfile_wcurl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index 1f96655e..b786aaa4 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -352,7 +352,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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(1, lmax_qlm + 1) + 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') From baf2b9ec47999d33e97afed30693b9a3e843b599 Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 20:21:29 +0200 Subject: [PATCH 07/12] minor fix --- delensalot/scripts/run_fromparfile_wcurl.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index b786aaa4..8c925365 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -13,8 +13,7 @@ import os from os.path import join as opj import numpy as np -from multiprocessing import cpu_count - +from psutil import cpu_count import plancklens from plancklens import utils from plancklens import qresp @@ -186,7 +185,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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() // 2)) + 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. From 6f13ad9634dcc5b685d932ce1fbcb5c2cd71e05c Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 22:04:11 +0200 Subject: [PATCH 08/12] minor fix --- delensalot/scripts/run_fromparfile_wcurl.py | 1 - 1 file changed, 1 deletion(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index 8c925365..8acf27ff 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -7,7 +7,6 @@ '' version is standard gradient reconstruction 'wcurl' version reconstructs both gradient and curl on gradient-only input map 'wcurlin' version reconstructs both gradient and curl on gradient and curl input map - (starting point still the curl-free QE though) """ import os From 8f1e6e7d1eb68c66ee7aea5d9693e8dee643e17d Mon Sep 17 00:00:00 2001 From: jcarron Date: Mon, 4 Sep 2023 22:19:38 +0200 Subject: [PATCH 09/12] minor --- delensalot/scripts/run_fromparfile_wcurl.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/delensalot/scripts/run_fromparfile_wcurl.py b/delensalot/scripts/run_fromparfile_wcurl.py index 8acf27ff..bde22984 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -4,9 +4,9 @@ e.g. python ./run_fromparfile_wcurl.py -itmax 0 -v 'wcurlin' - '' version is standard gradient reconstruction - 'wcurl' version reconstructs both gradient and curl on gradient-only input map - 'wcurlin' version reconstructs both gradient and curl on gradient and curl input map + -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 e8dc4a65a97b2091b031aff41c5d36a340b8cc11 Mon Sep 17 00:00:00 2001 From: jcarron Date: Tue, 5 Sep 2023 09:53:54 +0200 Subject: [PATCH 10/12] ignore .key --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index c1dc62f9..1d00a461 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ outputs/* *.o *.obj +*.key # slurm slurm* From 79cbcaa2b44073639493b87996ef25ea00f1af9d Mon Sep 17 00:00:00 2001 From: jcarron Date: Tue, 5 Sep 2023 12:57:36 +0200 Subject: [PATCH 11/12] minor --- delensalot/core/iterator/cs_iterator_multi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/delensalot/core/iterator/cs_iterator_multi.py b/delensalot/core/iterator/cs_iterator_multi.py index 0f653f34..5fc0f667 100644 --- a/delensalot/core/iterator/cs_iterator_multi.py +++ b/delensalot/core/iterator/cs_iterator_multi.py @@ -300,7 +300,7 @@ def get_template_blm(self, it, it_e, lmaxb=1024, lmin_plm=1, elm_wf:None or np.n 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, np.zeros_like(elm_wf)], self.mmax_filt, 2, lmaxb, mmaxb) + elm, blm = ffi.lensgclm(elm_wf, self.mmax_filt, 2, lmaxb, mmaxb) if cache_cond: self.wf_cacher.cache(fn, blm) return blm From c9b6a19886bfb9471d5cf47d0c762dd7c8801e48 Mon Sep 17 00:00:00 2001 From: jcarron Date: Fri, 29 Sep 2023 21:33:00 +0200 Subject: [PATCH 12/12] adding T-only and T+P recs --- delensalot/core/iterator/cs_iterator_multi.py | 1 - delensalot/core/opfilt/MAP_opfilt_iso_t.py | 2 +- delensalot/core/opfilt/MAP_opfilt_iso_tp.py | 2 +- delensalot/scripts/run_fromparfile_wcurl.py | 32 ++++++++++++++++--- 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/delensalot/core/iterator/cs_iterator_multi.py b/delensalot/core/iterator/cs_iterator_multi.py index 5fc0f667..aa3b7d09 100644 --- a/delensalot/core/iterator/cs_iterator_multi.py +++ b/delensalot/core/iterator/cs_iterator_multi.py @@ -510,7 +510,6 @@ def calc_gradlik(self, itr, key, iwantit=False): soltn, it_soltn = self.load_soltn(itr, key) if it_soltn < itr - 1: soltn *= self.soltn_cond - assert soltn.ndim == 1, 'Fix following lines' 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) 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 index bde22984..b4dd4a16 100644 --- a/delensalot/scripts/run_fromparfile_wcurl.py +++ b/delensalot/scripts/run_fromparfile_wcurl.py @@ -32,7 +32,10 @@ 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 @@ -180,7 +183,6 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): (here if 'noMF' is in version, will not use any mean-fied at the very first step) cg_tol: tolerance of conjugate-gradient filter """ - assert k in ['p_eb', 'p_p'], k libdir_iterator = libdir_iterators(k, simidx, version) if not os.path.exists(libdir_iterator): os.makedirs(libdir_iterator) @@ -231,9 +233,8 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): {'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']: - ffi = deflection(lenjob_geometry, np.zeros_like(plm0), mmax_qlm, numthreads=tr, epsilon=1e-7) # 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) @@ -241,6 +242,29 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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 @@ -253,7 +277,7 @@ def get_itlib(k:str, simidx:int, version:str, cg_tol:float): 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=lambda : alm_copy(ivfs_wcurl.get_sim_emliklm(simidx), None, lmax_unl, mmax_unl)) + 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