From e09a690c31bc262f11cbc723c3cbf25dcce247cb Mon Sep 17 00:00:00 2001 From: damonge Date: Sun, 6 Oct 2024 19:49:15 +0100 Subject: [PATCH] algorithm tested, fixed bugs --- pymaster/tests/test_master_anisotropic.py | 115 ++++++++++++++++++++++ pymaster/workspaces.py | 7 +- src/nmt_master.c | 68 +++++++------ 3 files changed, 159 insertions(+), 31 deletions(-) create mode 100644 pymaster/tests/test_master_anisotropic.py diff --git a/pymaster/tests/test_master_anisotropic.py b/pymaster/tests/test_master_anisotropic.py new file mode 100644 index 0000000..c23abd9 --- /dev/null +++ b/pymaster/tests/test_master_anisotropic.py @@ -0,0 +1,115 @@ +import numpy as np +import healpy as hp +import pymaster as nmt + + +def test_anisotropic_weighting(): + # Test parameters + nside = 64 + spin = 2 + nlb = 4 + nsims = 100 + + # Create anisotropic weights + mask = hp.read_map("test/benchmarks/msk.fits") + mask = hp.ud_grade(mask, nside_out=nside) + mask = hp.smoothing(mask, sigma=np.radians(1.0)) + mask[mask > 1] = 1 + mask[mask < 0.001] = 0 + + delta_m = 0.9 + r_m = 0.5 + w11 = (1+delta_m)*mask + w22 = (1-delta_m)*mask + w12 = r_m*np.sqrt(w11*w22) + + # Input power spectra + ls = np.arange(3*nside) + cl_temp = 1/(ls+10) + cl_tt = 1.5*cl_temp + cl_te = 0.6*cl_temp + cl_tb = 0.3*cl_temp + cl_ee = 1.0*cl_temp + cl_eb = 0.2*cl_temp + cl_bb = 0.4*cl_temp + + # Workspaces + b = nmt.NmtBin.from_nside_linear(nside, nlb=nlb) + leff = b.get_effective_ells() + f0 = nmt.NmtField(mask, None, spin=0, n_iter=0) + fs = nmt.NmtField(mask, None, spin=spin, n_iter=0) + fsa = nmt.NmtField(w11, None, spin=spin, n_iter=0, + mask_12=w12, mask_22=w22) + w0sa = nmt.NmtWorkspace.from_fields(f0, fsa, b) + wsa0 = nmt.NmtWorkspace.from_fields(fsa, f0, b) + wsasa = nmt.NmtWorkspace.from_fields(fsa, fsa, b) + wssa = nmt.NmtWorkspace.from_fields(fs, fsa, b) + wsas = nmt.NmtWorkspace.from_fields(fsa, fs, b) + + # Run simulations + cl0sa_s = [] + clsa0_s = [] + clsasa_s = [] + clssa_s = [] + clsas_s = [] + for i in range(nsims): + almt, alme, almb = hp.synalm([cl_tt, cl_ee, cl_bb, + cl_te, cl_eb, cl_tb], new=True) + map_t = hp.alm2map(almt, nside, lmax=3*nside-1) + map_q, map_u = hp.alm2map_spin([alme, almb], nside, spin, + lmax=3*nside-1, mmax=3*nside-1) + + f0 = nmt.NmtField(mask, [map_t], n_iter=0) + fs = nmt.NmtField(mask, [map_q, map_u], spin=spin, n_iter=0) + fsa = nmt.NmtField(w11, [map_q, map_u], spin=spin, n_iter=0, + mask_12=w12, mask_22=w22) + + cl0sa_s.append(w0sa.decouple_cell(nmt.compute_coupled_cell(f0, fsa))) + clsa0_s.append(wsa0.decouple_cell(nmt.compute_coupled_cell(fsa, f0))) + clsasa_s.append(wsasa.decouple_cell(nmt.compute_coupled_cell(fsa, + fsa))) + clssa_s.append(wssa.decouple_cell(nmt.compute_coupled_cell(fs, fsa))) + clsas_s.append(wsas.decouple_cell(nmt.compute_coupled_cell(fsa, fs))) + cl0sa_s = np.array(cl0sa_s) + clsa0_s = np.array(clsa0_s) + clsasa_s = np.array(clsasa_s) + clssa_s = np.array(clssa_s) + clsas_s = np.array(clsas_s) + # Mean + cl0sa_m = np.mean(cl0sa_s, axis=0) + clsa0_m = np.mean(clsa0_s, axis=0) + clsasa_m = np.mean(clsasa_s, axis=0) + clssa_m = np.mean(clssa_s, axis=0) + clsas_m = np.mean(clsas_s, axis=0) + # STD + cl0sa_e = np.std(cl0sa_s, axis=0) + clsa0_e = np.std(clsa0_s, axis=0) + clsasa_e = np.std(clsasa_s, axis=0) + clssa_e = np.std(clssa_s, axis=0) + clsas_e = np.std(clsas_s, axis=0) + # Truth + cl0sa_t = w0sa.decouple_cell(w0sa.couple_cell([cl_te, cl_tb])) + clsa0_t = wsa0.decouple_cell(wsa0.couple_cell([cl_te, cl_tb])) + clsasa_t = wsasa.decouple_cell(wsasa.couple_cell([cl_ee, cl_eb, + cl_eb, cl_bb])) + clssa_t = wssa.decouple_cell(wssa.couple_cell([cl_ee, cl_eb, + cl_eb, cl_bb])) + clsas_t = wsas.decouple_cell(wsas.couple_cell([cl_ee, cl_eb, + cl_eb, cl_bb])) + + # Compare all power spectra and check for > 6sigma deviations + def comp_cl(clm, cle, clt, nsigma=6): + lgood = leff < 2*nside + for m, e, t in zip(clm, cle, clt): + r = ((m-t)*np.sqrt(nsims)/e)[lgood] + if np.any(np.fabs(r) > nsigma): + return False + return True + + test_0sa = comp_cl(cl0sa_m, cl0sa_e, cl0sa_t) + test_sa0 = comp_cl(clsa0_m, clsa0_e, clsa0_t) + test_sasa = comp_cl(clsasa_m, clsasa_e, clsasa_t) + test_ssa = comp_cl(clssa_m, clssa_e, clssa_t) + test_sas = comp_cl(clsas_m, clsas_e, clsas_t) + + assert np.all([test_0sa, test_sa0, test_sasa, test_ssa, test_sas]) diff --git a/pymaster/workspaces.py b/pymaster/workspaces.py index 4933476..4946bae 100644 --- a/pymaster/workspaces.py +++ b/pymaster/workspaces.py @@ -263,9 +263,10 @@ def compute_coupling_matrix(self, fl1, fl2, bins, is_teb=False, alm2 = fl2.get_mask_alms() pcl_mask = hp.alm2cl(alm1, alm2, lmax=fl1.ainfo_mask.lmax) if anisotropic_mask_any: + pcl0 = pcl_mask * 0 pclm_00 = pcl_mask - pclm_0e = pclm_0b = pclm_e0 = pclm_b0 = None - pclm_ee = pclm_eb = pclm_be = pclm_bb = None + pclm_0e = pclm_0b = pclm_e0 = pclm_b0 = pcl0 + pclm_ee = pclm_eb = pclm_be = pclm_bb = pcl0 if fl1.anisotropic_mask: alm1a = fl1.get_anisotropic_mask_alms() if fl2.anisotropic_mask: @@ -276,7 +277,7 @@ def compute_coupling_matrix(self, fl1, fl2, bins, is_teb=False, if fl1.anisotropic_mask: pclm_e0 = hp.alm2cl(alm1a[0], alm2, lmax=fl1.ainfo_mask.lmax) pclm_b0 = hp.alm2cl(alm1a[1], alm2, lmax=fl1.ainfo_mask.lmax) - if fl1.anisotropic_mask: + if fl2.anisotropic_mask: pclm_ee = hp.alm2cl(alm1a[0], alm2a[0], lmax=fl1.ainfo_mask.lmax) pclm_eb = hp.alm2cl(alm1a[0], alm2a[1], diff --git a/src/nmt_master.c b/src/nmt_master.c index 2bf64cc..750ed64 100644 --- a/src/nmt_master.c +++ b/src/nmt_master.c @@ -497,30 +497,34 @@ nmt_workspace *nmt_compute_coupling_matrix_anisotropic(int spin1, int spin2, int max_spin=NMT_MAX(spin1, spin2); int lstart=max_spin; - int has_ss2=(spin1!=0) && (spin2!=0) && (spin1!=spin2); + int reuse_ss1=(spin1==spin2); int is_0s=(spin1==0) && (spin2!=0); int is_s0=(spin1!=0) && (spin2==0); int is_ss=(spin1!=0) && (spin2!=0); #pragma omp parallel default(none) \ - shared(pcl_masks, lstart, has_ss2, lmax, lmax_mask) \ + shared(pcl_masks, lstart, reuse_ss1, lmax, lmax_mask) \ shared(mask_aniso_1, mask_aniso_2, is_0s, is_s0, is_ss) \ shared(spin1, spin2, i_00, i_0e, i_0b, i_e0, i_b0) \ shared(i_ee, i_eb, i_be, i_bb, w) { int ll2,ll3; double *wigner_ss1=NULL, *wigner_ss2=NULL, *wigner_ss1a=NULL, *wigner_ss2a=NULL; + // s -s 0 terms wigner_ss1=my_malloc(2*(lmax_mask+1)*sizeof(double)); + if(reuse_ss1) + wigner_ss2=wigner_ss1; + else + wigner_ss2=my_malloc(2*(lmax_mask+1)*sizeof(double)); + + // s s -2s terms if(mask_aniso_1) wigner_ss1a=my_malloc(2*(lmax_mask+1)*sizeof(double)); - if(has_ss2) { - wigner_ss2=my_malloc(2*(lmax_mask+1)*sizeof(double)); - if(mask_aniso_2) + if(mask_aniso_2) { + if((!reuse_ss1) || (mask_aniso_1 == 0)) wigner_ss2a=my_malloc(2*(lmax_mask+1)*sizeof(double)); - } - else { - wigner_ss2=wigner_ss1; - wigner_ss2a=wigner_ss1a; + else + wigner_ss2a=wigner_ss1a; } #pragma omp for schedule(dynamic) @@ -535,19 +539,25 @@ nmt_workspace *nmt_compute_coupling_matrix_anisotropic(int spin1, int spin2, lmin_here=abs(ll2-ll3); lmax_here=ll2+ll3; + // s -s 0 terms drc3jj(ll2,ll3,spin1,-spin1,&lmin_ss1,&lmax_ss1,wigner_ss1,2*(lmax_mask+1)); + if(reuse_ss1) { + lmin_ss2=lmin_ss1; + lmax_ss2=lmax_ss1; + } + else + drc3jj(ll2,ll3,spin2,-spin2,&lmin_ss2,&lmax_ss2,wigner_ss2,2*(lmax_mask+1)); + + // s s -2s terms if(mask_aniso_1) drc3jj(ll2,ll3,spin1,spin1,&lmin_ss1a,&lmax_ss1a,wigner_ss1a,2*(lmax_mask+1)); - if(has_ss2) { - drc3jj(ll2,ll3,spin2,-spin2,&lmin_ss2,&lmax_ss2,wigner_ss2,2*(lmax_mask+1)); - if(mask_aniso_2) + if(mask_aniso_2) { + if((!reuse_ss1) || (mask_aniso_1 == 0)) drc3jj(ll2,ll3,spin2,spin2,&lmin_ss2a,&lmax_ss2a,wigner_ss2a,2*(lmax_mask+1)); - } - else { - lmin_ss2=lmin_ss1; - lmax_ss2=lmax_ss1; - lmin_ss2a=lmin_ss1a; - lmax_ss2a=lmax_ss1a; + else { + lmin_ss2a=lmin_ss1a; + lmax_ss2a=lmax_ss1a; + } } for(l1=lmin_here;l1<=lmax_here;l1++) { @@ -583,18 +593,17 @@ nmt_workspace *nmt_compute_coupling_matrix_anisotropic(int spin1, int spin2, mbb = wss1a*wss2a*pcl_masks[i_bb][l1]; } } - if(is_0s) { w->coupling_matrix_unbinned[2*ll2+0][2*ll3+0] += l3fac*(m00-m0e); //0E,0E w->coupling_matrix_unbinned[2*ll2+0][2*ll3+1] += l3fac*(-m0b); //0E,0B - w->coupling_matrix_unbinned[4*ll2+1][4*ll3+0] += l3fac*(-m0b); //0B,0E - w->coupling_matrix_unbinned[4*ll2+1][4*ll3+1] += l3fac*(m00+m0e); //0B,0B + w->coupling_matrix_unbinned[2*ll2+1][2*ll3+0] += l3fac*(-m0b); //0B,0E + w->coupling_matrix_unbinned[2*ll2+1][2*ll3+1] += l3fac*(m00+m0e); //0B,0B } if(is_s0) { - w->coupling_matrix_unbinned[4*ll2+0][4*ll3+0] += l3fac*(m00-me0); //E0,E0 - w->coupling_matrix_unbinned[4*ll2+0][4*ll3+2] += l3fac*(-mb0); //E0,B0 - w->coupling_matrix_unbinned[4*ll2+2][4*ll3+0] += l3fac*(-mb0); //B0,E0 - w->coupling_matrix_unbinned[4*ll2+2][4*ll3+2] += l3fac*(m00+me0); //B0,B0 + w->coupling_matrix_unbinned[2*ll2+0][2*ll3+0] += l3fac*(m00-me0); //E0,E0 + w->coupling_matrix_unbinned[2*ll2+0][2*ll3+1] += l3fac*(-mb0); //E0,B0 + w->coupling_matrix_unbinned[2*ll2+1][2*ll3+0] += l3fac*(-mb0); //B0,E0 + w->coupling_matrix_unbinned[2*ll2+1][2*ll3+1] += l3fac*(m00+me0); //B0,B0 } if(is_ss) { w->coupling_matrix_unbinned[4*ll2+0][4*ll3+0] += l3fac*(splus * (m00-m0e-me0+mee) + sminus * mbb); //EE,EE @@ -619,12 +628,15 @@ nmt_workspace *nmt_compute_coupling_matrix_anisotropic(int spin1, int spin2, } } //end omp for + // s -s 0 terms free(wigner_ss1); + if(!reuse_ss1) + free(wigner_ss2); + // s s -2s terms if(mask_aniso_1) free(wigner_ss1a); - if(has_ss2) { - free(wigner_ss2); - if(mask_aniso_2) + if(mask_aniso_2) { + if((!reuse_ss1) || (mask_aniso_1 == 0)) free(wigner_ss2a); } } //end omp parallel