Skip to content

Commit

Permalink
algorithm tested, fixed bugs
Browse files Browse the repository at this point in the history
  • Loading branch information
damonge committed Oct 6, 2024
1 parent 03807c4 commit e09a690
Show file tree
Hide file tree
Showing 3 changed files with 159 additions and 31 deletions.
115 changes: 115 additions & 0 deletions pymaster/tests/test_master_anisotropic.py
Original file line number Diff line number Diff line change
@@ -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])
7 changes: 4 additions & 3 deletions pymaster/workspaces.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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],
Expand Down
68 changes: 40 additions & 28 deletions src/nmt_master.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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++) {
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down

0 comments on commit e09a690

Please sign in to comment.