From 3f577359f6f970119cacb43ccb1a319336d3ac8a Mon Sep 17 00:00:00 2001 From: Sebastian-Belkner Date: Wed, 6 Mar 2024 13:11:53 -0800 Subject: [PATCH] fix zbound bug in MAP rec (niv_geomlib didnt agree with truncated obs maps), small refactors --- .../config/transformer/lerepi2dlensalot.py | 22 ++++++----- delensalot/core/handler.py | 39 ++++++++++++------- delensalot/core/iterator/cs_iterator.py | 1 - delensalot/core/opfilt/bmodes_ninv.py | 31 ++++++++++++--- 4 files changed, 63 insertions(+), 30 deletions(-) diff --git a/delensalot/config/transformer/lerepi2dlensalot.py b/delensalot/config/transformer/lerepi2dlensalot.py index b6a62799..f7aeaa7d 100644 --- a/delensalot/config/transformer/lerepi2dlensalot.py +++ b/delensalot/config/transformer/lerepi2dlensalot.py @@ -368,8 +368,9 @@ def _process_Itrec(dl, it): dl.lenjob_geominfo = it.lenjob_geominfo dl.lenjob_geomlib = get_geom(it.lenjob_geominfo) - thtbounds = (np.arccos(dl.zbounds[1]), np.arccos(dl.zbounds[0])) - dl.lenjob_geomlib.restrict(*thtbounds, northsouth_sym=False, update_ringstart=True) + # FIXME not sure if we should restrict here + # thtbounds = (np.arccos(dl.zbounds[1]), np.arccos(dl.zbounds[0])) + # dl.lenjob_geomlib.restrict(*thtbounds, northsouth_sym=False, update_ringstart=True) # TODO this needs cleaner implementation if dl.version == '' or dl.version == None: @@ -558,6 +559,8 @@ def _process_Itrec(dl, it): dl.lenjob_geominfo = it.lenjob_geominfo dl.lenjob_geomlib = get_geom(it.lenjob_geominfo) + thtbounds = (np.arccos(dl.zbounds[1]), np.arccos(dl.zbounds[0])) + dl.lenjob_geomlib.restrict(*thtbounds, northsouth_sym=False, update_ringstart=True) if dl.version == '' or dl.version == None: dl.mf_dirname = opj(dl.TEMP, l2T_Transformer.ofj('mf', {'Nmf': dl.Nmf})) @@ -796,14 +799,13 @@ def _process_Madel(dl, ma): dl.masks[maskflavour][maskid] = dl.masks_fromfn[maskid] dl.binmasks[maskflavour][maskid] = np.where(dl.masks[maskflavour][maskid]>0,1,0) - ## Binning and power spectrum calculator specific preparation - if ma.Cl_fid == 'ffp10': - dl.cls_unl = camb_clfile(cf.analysis.cls_unl) - dl.cls_len = camb_clfile(cf.analysis.cls_len) - dl.clg_templ = dl.cls_len['ee'] - dl.clc_templ = dl.cls_len['bb'] - dl.clg_templ[0] = 1e-32 - dl.clg_templ[1] = 1e-32 + ## Binning and power spectrum calculator specific preparation. + dl.cls_unl = camb_clfile(cf.analysis.cls_unl) + dl.cls_len = camb_clfile(cf.analysis.cls_len) + dl.clg_templ = dl.cls_len['ee'] + dl.clc_templ = dl.cls_len['bb'] + dl.clg_templ[0] = 1e-32 + dl.clg_templ[1] = 1e-32 dl.binning = ma.binning if dl.binning == 'binned': diff --git a/delensalot/core/handler.py b/delensalot/core/handler.py index b409049d..692bfc8c 100644 --- a/delensalot/core/handler.py +++ b/delensalot/core/handler.py @@ -498,7 +498,9 @@ def postrun_sky(self): class Noise_modeller(Basejob): - + ''' + CURRENTLY NOT USED + ''' def __init__(self, dlensalot_model): super().__init__(dlensalot_model) @@ -547,13 +549,15 @@ def __init__(self, dlensalot_model): marge_monopole=True, marge_dipole=True, marge_maps=[]) transf_elm_loc = gauss_beam(self.beam / 180 / 60 * np.pi, lmax=self.lm_max_ivf[0]) if self.OBD == 'OBD': - nivjob_geomlib_ = get_geom(self.nivjob_geominfo) + # FIXME not sure which nivjob_geomlib to pass here, restricted or not? + # nivjob_geomlib_ = get_geom(self.nivjob_geominfo) self.cinv_p = cinv_p_OBD.cinv_p(opj(self.libdir_QE, 'cinv_p'), self.lm_max_ivf[0], self.nivjob_geominfo[1]['nside'], self.cls_len, - transf_elm_loc[:self.lm_max_ivf[0]+1], self.nivp_desc, geom=nivjob_geomlib_, #self.nivjob_geomlib, + transf_elm_loc[:self.lm_max_ivf[0]+1], self.nivp_desc, geom=self.nivjob_geomlib, #self.nivjob_geomlib, chain_descr=self.chain_descr(self.lm_max_ivf[0], self.cg_tol), bmarg_lmax=self.lmin_teb[2], - zbounds=(-1,1), _bmarg_lib_dir=self.obd_libdir, _bmarg_rescal=self.obd_rescale, + zbounds=self.zbounds, _bmarg_lib_dir=self.obd_libdir, _bmarg_rescal=self.obd_rescale, sht_threads=self.tr) + # (-1,1) else: self.cinv_p = filt_cinv.cinv_p(opj(self.TEMP, 'cinv_p'), self.lm_max_ivf[0], self.nivjob_geominfo[1]['nside'], self.cls_len, @@ -749,19 +753,21 @@ def run(self, task=None): if task == 'calc_phi': for idx in self.jobs[taski][mpi.rank::mpi.size]: + log.info("get_sim_qlm..") self.qlms_dd.get_sim_qlm(self.k, int(idx)) + log.info("get_sim_qlm done.") if np.all(self.simulationdata.obs_lib.maps == DEFAULT_NotAValue): self.simulationdata.purgecache() mpi.barrier() + for idx in self.jobs[taski][mpi.rank::mpi.size]: - ## If meanfield subtraction is done, only one task must calculate the meanfield first before get_plm() is called, otherwise read-errors because all tasks try calculating/accessing it at once. + ## If meanfield subtraction is requested, only one task must calculate the meanfield first before get_plm() is called, otherwise read-errors because all tasks try calculating/accessing it at once. ## The way I fix this (the next two lines) is a bit unclean. if self.QE_subtract_meanfield: self.qlms_dd.get_sim_qlm_mf(self.k, [int(simidx_mf) for simidx_mf in self.simidxs_mf]) self.get_plm(idx, self.QE_subtract_meanfield) if np.all(self.simulationdata.obs_lib.maps == DEFAULT_NotAValue): - self.simulationdata.purgecache() - + self.simulationdata.purgecache() if task == 'calc_meanfield': if len(self.jobs[taski])>0: @@ -824,6 +830,7 @@ def get_mchain(self, simidx, key, it=0): def get_meanfield(self, simidx): # Either return MC MF, filter.qlms_mf, or mfvar ret = np.zeros_like(self.qlms_dd.get_sim_qlm(self.k, 0)) + fn_mf = opj(self.libdir_QE, 'mf_allsims.npy') if np.dtype(self.mfvar) == str: if self.mfvar == 'qlms_mf': # calculate MF estimate using Lewis&Carron trick @@ -833,6 +840,7 @@ def get_meanfield(self, simidx): if self.mfvar == None: # MC MF, and exclude the current simidx ret = self.qlms_dd.get_sim_qlm_mf(self.k, [int(simidx_mf) for simidx_mf in self.simidxs_mf]) + np.save(fn_mf, ret) # plancklens already stores that in qlms_dd/ but I want to have this more conveniently without the naming gibberish if simidx in self.simidxs_mf: ret = (ret - self.qlms_dd.get_sim_qlm(self.k, int(simidx)) / self.Nmf) * (self.Nmf / (self.Nmf - 1)) else: @@ -840,6 +848,7 @@ def get_meanfield(self, simidx): ret = hp.read_alm(self.mfvar) if simidx in self.simidxs_mf: ret = (ret - self.qlms_dd_mfvar.get_sim_qlm(self.k, int(simidx)) / self.Nmf) * (self.Nmf / (self.Nmf - 1)) + return ret return ret @@ -993,11 +1002,10 @@ def __init__(self, dlensalot_model): self.simulationdata = self.simgen.simulationdata self.qe = QE_lr(dlensalot_model, caller=self) self.qe.simulationdata = self.simgen.simulationdata # just to be sure, so we have a single truth in MAP_lr. - - if self.OBD == 'OBD': - nivjob_geomlib_ = get_geom(self.nivjob_geominfo) - self.tpl = template_dense(self.lmin_teb[2], nivjob_geomlib_, self.tr, _lib_dir=self.obd_libdir, rescal=self.obd_rescale) + # FIXME not sure why this was here.. that caused mismatch for calc_gradlik niv_job geom and qumaps when truncated + # nivjob_geomlib_ = get_geom(self.nivjob_geominfo) + self.tpl = template_dense(self.lmin_teb[2], self.nivjob_geomlib, self.tr, _lib_dir=self.obd_libdir, rescal=self.obd_rescale) else: self.tpl = None @@ -1017,6 +1025,7 @@ def __init__(self, dlensalot_model): elif self.it_filter_directional == 'isotropic': self.sims_MAP = self.simulationdata self.filter = self.get_filter() + log.info('------ init done ----') # # @base_exception_handler @log_on_start(logging.DEBUG, "MAP.map.collect_jobs() started") @@ -1250,11 +1259,16 @@ def _prepare_job(self): # @log_on_start(logging.DEBUG, "get_basemap() started") # @log_on_end(logging.DEBUG, "get_basemap() finished") def get_basemap(self, simidx): + ''' + Return a B-map to be delensed. Can be the map handled in the sims_lib library (basemap='lens'), 'lens_ffp10' (this potentially is the same as sims_lib currently uses ffp10 as baseline), + or the observed map itself, in which case the residual foregrounds and noise will still be in there. + + ''' # TODO depends if data comes from delensalot simulations or from external.. needs cleaner implementation if self.basemap == 'lens': return almxfl(alm_copy(self.simulationdata.get_sim_sky(simidx, space='alm', spin=0, field='polarization')[1], self.simulationdata.lmax, *self.lm_max_blt), self.ttebl['e'], self.lm_max_blt[0], inplace=False) elif self.basemap == 'lens_ffp10': - return almxfl(alm_copy(planck2018_sims.cmb_len_ffp10.get_sim_blm(simidx), None, lmaxout=self.lm_max_blt[0], mmaxout=self.lm_max_blt[1]), gauss_beam(2.3 / 180 / 60 * np.pi, lmax=self.lm_max_blt[1])) + return almxfl(alm_copy(planck2018_sims.cmb_len_ffp10.get_sim_blm(simidx), None, lmaxout=self.lm_max_blt[0], mmaxout=self.lm_max_blt[1]), gauss_beam(self.beam / 180 / 60 * np.pi, lmax=self.lm_max_blt[1])) else: # only checking for map to save some memory.. if np.all(self.simulationdata.maps == DEFAULT_NotAValue): @@ -1263,7 +1277,6 @@ def get_basemap(self, simidx): return hp.map2alm_spin(self.simulationdata.get_sim_obs(simidx, space='map', spin=2, field='polarization'), spin=2, lmax=self.lm_max_blt[0], mmax=self.lm_max_blt[1])[1] - @log_on_start(logging.DEBUG, "_delens() started") @log_on_end(logging.DEBUG, "_delens() finished") def delens(self, simidx, outputdata): diff --git a/delensalot/core/iterator/cs_iterator.py b/delensalot/core/iterator/cs_iterator.py index 8e432d65..477187c5 100644 --- a/delensalot/core/iterator/cs_iterator.py +++ b/delensalot/core/iterator/cs_iterator.py @@ -478,7 +478,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 - 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/bmodes_ninv.py b/delensalot/core/opfilt/bmodes_ninv.py index e0f074f2..ad87df99 100644 --- a/delensalot/core/opfilt/bmodes_ninv.py +++ b/delensalot/core/opfilt/bmodes_ninv.py @@ -305,12 +305,31 @@ def __init__(self, geom, n_inv, b_transf, lmax_marg=0, zbounds=(-1., 1.), blm_ra self.templates.append(template_bfilt(lmax_marg=lmax_marg, geom=geom, sht_threads=sht_threads, _lib_dir=_bmarg_lib_dir)) if len(self.templates) > 0: if _bmarg_lib_dir is not None and os.path.exists( os.path.join(_bmarg_lib_dir, 'tniti.npy')): - if mpi.rank == 0: - log.info("Loading " + os.path.join(_bmarg_lib_dir, 'tniti.npy')) - self.tniti = np.load(os.path.join(_bmarg_lib_dir, 'tniti.npy')) - else: - pass - self.tniti = mpi.bcast(self.tniti, root=0) + nmodes_ = int((lmax_marg + 1) * lmax_marg + lmax_marg + 1 - 4) + log.info("Loading " + os.path.join(_bmarg_lib_dir, 'tniti.npy')) + self.tniti = np.load(os.path.join(_bmarg_lib_dir, 'tniti.npy')) + log.info("done") + # from mpi4py import MPI + # # np.frombuffer(buf, dtype=np.double, count=data_size) + # if mpi.rank == 0: + # win = MPI.Win.Allocate_shared(nmodes_*nmodes_, MPI.DOUBLE, comm=MPI.COMM_WORLD) + # self.tniti, _ = win.Shared_query(0) + # log.info("Loading " + os.path.join(_bmarg_lib_dir, 'tniti.npy')) + # self.tniti[:] = np.load(os.path.join(_bmarg_lib_dir, 'tniti.npy')) + # print("tniti shape: {}".format(self.tniti.shape)) + # log.info('sharing tniti across slurm workers..') + # win.Fence() + # else: + # self.tniti = np.zeros(shape=(nmodes_,nmodes_)) + # win = MPI.Win.Allocate_shared(0, MPI.DOUBLE, comm=MPI.COMM_WORLD) + # size = win.Shared_query(0)[0] + # buf, _ = win.Shared_query(0) + # win.Fence() + # self.tniti[:] = buf[:] + + # for rowi, row in enumerate(self.tniti): + # self.tniti[rowi] = mpi.bcast(self.tniti[rowi], root=0) + # log.info('sharing tniti across slurm workers.. DONE.') if _bmarg_rescal != 1.: log.info("**** RESCALING tiniti with %.4f"%_bmarg_rescal) self.tniti *= _bmarg_rescal