def _calc_Mscal(comm, wEBlm, lmax=None, verbose=True):
'''Calculate the temperature/V polarization mode-mixing matrix
'''
rank = comm.Get_rank()
size = comm.Get_size()
if lmax is None:
lmax = H.Alm.getlmax(len(wEBlm[0, :]))
Mscal = np.zeros([lmax+1, lmax+1])
wl = H.alm2cl(wEBlm[0, :])
for l1 in range(2+rank, lmax+1, size):
if verbose:
print("Scal l1 = ", l1)
for l2 in range(2, lmax+1):
l3min = np.abs(l1-l2)
l3max = np.abs(l1+l2)
l3vals = np.arange(l3min, l3max+1)
# Calculates Wigner 3j symbol for all valid l3
JT = wc.wigner3j_vect(2*l1, 2*l2, 0, 0)
# Since we only have wl up to lmax we must ignore all terms that
# have l3 > lmax
idx = l3vals <= lmax
Mscal[l1, l2] = np.sum((2*l3vals[idx]+1)*wl[l3vals[idx]] *
JT[idx]**2)
normfact = (2.0*l2+1.0) / (4.0*np.pi)
Mscal[l1, l2] *= normfact
Mscal = comm.allreduce(Mscal)
return Mscal
def _correct_TE(apfunction, x, w):
'''Calculate an amplitude correction of the TE cross-correlation
due to the apodization of the correlation function.
Parameters
----------
apfunction : array
The function dependent on angle used to apodize the correlation
function.
x : array
cos(angle), the Legendre-Gauss quadrature points
w : array
the weights needed for the Legendre-Gauss quadrature
Returns
-------
fact : array
TE correction function due to the fact that the correlation function
is apodized
'''
nell = len(apfunction)
lmax = nell - 1
#Calculate f_l from apodizing function in real space
fl = np.zeros(nell)
for i in range(nell):
Pl0 = scipy.special.lpmv(0, i, x)
fl[i] = np.sum(
w * apfunction * Pl0) * 2 * np.pi * (2 * i + 1) / (4 * np.pi)
kcross = np.zeros([nell, nell])
rank = comm.rank
size = comm.size
for l1 in range(2 + rank, nell, size):
for l2 in range(l1, nell):
l3min = np.abs(l1 - l2)
l3max = np.abs(l1 + l2)
#runs from l3min to l3max
wigner00 = wc.wigner3j_vect(2 * l1, 2 * l2, 2 * 0, 2 * 0)
wigner22 = wc.wigner3j_vect(2 * l1, 2 * l2, 2 * 2, -2 * 2)
if l3max > lmax:
tmp = lmax - l3max
wigner00 = wigner00[:tmp]
wigner22 = wigner22[:tmp]
kcross[l1, l2] = np.sum(wigner00 * wigner22 * fl[l3min:l3max + 1])
kcross[l2, l1] = kcross[l1, l2]
kcross[l1, l2] *= 2 * l2 + 1
if l2 != l1:
kcross[l2, l1] *= 2 * l1 + 1
kcross = comm.allreduce(kcross)
fact = np.sum(kcross, axis=1)
return fact
def _calc_Mcross(comm, wEBlm, lmax=None, cl_type='pseudo', verbose=True):
'''Calculate the temp-pol mode-mixing matrix.
'''
rank = comm.Get_rank()
size = comm.Get_size()
if lmax is None:
lmax = H.Alm.getlmax(len(wEBlm[0, :]))
l, m = H.Alm.getlm(lmax)
Mcross = np.zeros([2*lmax+2, 2*lmax+2])
a0 = 1.0
a1 = 2.0
a2 = 1.0
l1_vals = range(2+rank, lmax+1, size)
# Don't need to calculate this if we are looking at pseudo Cls, but it
# is not the bottleneck in the calculation so I don't care
Nl12 = H_ext._nfunc(l1_vals, 2)
Nl11 = H_ext._nfunc(l1_vals, 1)
Nl10 = H_ext._nfunc(l1_vals, 0)
fact0s = Nl10 / Nl12
fact1s = Nl11 / Nl12
for l1, fact0, fact1 in zip(l1_vals, fact0s, fact1s):
if verbose:
print("Cross l1 = ", l1)
for l2 in range(2, lmax+1):
l3min = np.abs(l1-l2)
l3max = np.abs(l1+l2)
# Minimum l3 for case when m3=1,2
l3min2 = np.max([l3min, 2])
l3min1 = np.max([l3min, 1])
wc_0 = wc.wigner3j_vect(2*l1, 2*l2, -2*2, 2*2)
wc_1 = wc.wigner3j_vect(2*l1, 2*l2, 2*2, -2*2)
Jp0 = wc_0 + wc_1
# Jm0 = wc_0 - wc_1
JT = wc.wigner3j_vect(2*l1, 2*l2, 0, 0)
if cl_type == 'pure' or cl_type == 'hybrid':
wc_0 = wc.wigner3j_vect(2*l1, 2*l2, (-2+1)*2, 2*2)
wc_1 = wc.wigner3j_vect(2*l1, 2*l2, (2-1)*2, -2*2)
Jp1 = wc_0 + wc_1
# Jm1 = wc_0 - wc_1
wc_0 = wc.wigner3j_vect(2*l1, 2*l2, (-2+2)*2, 2*2)
wc_1 = wc.wigner3j_vect(2*l1, 2*l2, (2-2)*2, -2*2)
Jp2 = wc_0 + wc_1
# Jm2 = wc_0 - wc_1
idx = np.all([l >= l3min, l <= l3max], axis=0)
l_tmp = l[idx]
m_tmp = m[idx]
m0 = m_tmp == 0
# sub array of wEBlm that have valid l (for l = l3)
wEBlm_tmp = wEBlm[:, idx]
idx_l3_0 = np.array(l_tmp - l3min, dtype=np.int) # Jp0,Jm0
idx_l3_1 = np.array(l_tmp - l3min1, dtype=np.int) # Jp1,Jm1
idx_l3_2 = np.array(l_tmp - l3min2, dtype=np.int) # Jp2,Jm2
# Subsets of the subset. When l3min is 0 or 1 and take the whole
# range, Jp1(Jm1) and/or Jp2(Jm2) get non-zero values when
# they should be zero (when l3 = 0 or 1) because negative indices
# in idx_l3_1 and idx_l3_2 wrap around to the end
idx1 = l_tmp >= l3min1
idx2 = l_tmp >= l3min2
# TE,TE, and TB,TB (2/4)
termE = a2 * JT[idx_l3_0]*Jp0[idx_l3_0] * np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[0, :]))
if cl_type == 'hybrid' or cl_type == 'pseudo':
Mcross[l1, l2] = 2*np.sum(termE) - np.sum(termE[m0])
if cl_type == 'pure' or cl_type == 'hybrid':
termE[idx1] += a1*fact1 * (JT[idx_l3_0]*Jp1[idx_l3_1] * np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[1, :])))[idx1]
termE[idx2] += a0*fact0 * (JT[idx_l3_0]*Jp2[idx_l3_2] * np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[2, :])))[idx2]
Mcross[l1+lmax+1, l2+lmax+1] = 2*np.sum(termE) - np.sum(termE[m0])
if cl_type == 'pure':
Mcross[l1, l2] = Mcross[l1+lmax+1, l2+lmax+1]
# TE,TB and TB,TE (4/4)
if cl_type == 'pure' or cl_type == 'hybrid':
termE = np.zeros_like(JT[idx_l3_0], dtype=np.float64)
termE[idx1] += a1*fact1 * (JT[idx_l3_0]*Jp1[idx_l3_1] * np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[-1, :])))[idx1]
termE[idx2] += a0*fact0 * (JT[idx_l3_0]*Jp2[idx_l3_2] * np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[-2, :])))[idx2]
Mcross[l1+lmax+1, l2] = 2*np.sum(termE) - np.sum(termE[m0])
if cl_type == 'pure':
Mcross[l1, l2+lmax+1] = Mcross[l1+lmax+1, l2]
normfact = (2.0*l2+1.0) / (4.0*np.pi)
Mcross[l1, l2] *= normfact / 2.0
Mcross[l1+lmax+1, l2+lmax+1] *= normfact / 2.0
Mcross[l1, l2+lmax+1] *= normfact / 2.0
Mcross[l1+lmax+1, l2] *= -normfact / 2.0
Mcross = comm.allreduce(Mcross)
return Mcross
def _calc_Mpol(comm, wEBlm, lmax=None, cl_type='pseudo', verbose=True):
'''Calculate the polarization mode mixing matrix
'''
rank = comm.Get_rank()
size = comm.Get_size()
if lmax is None:
lmax = H.Alm.getlmax(len(wEBlm[0, :]))
l, m = H.Alm.getlm(lmax)
Mpol = np.zeros([3*lmax+3, 3*lmax+3])
a0 = 1.0
a1 = 2.0
a2 = 1.0
l1_vals = range(2+rank, lmax+1, size)
# Don't need to calculate this if we are looking at pseudo Cls, but it
# is not the bottleneck in the calculation so I don't care
Nl12 = H_ext._nfunc(l1_vals, 2)
Nl11 = H_ext._nfunc(l1_vals, 1)
Nl10 = H_ext._nfunc(l1_vals, 0)
fact0s = Nl10 / Nl12
fact1s = Nl11 / Nl12
for l1, fact0, fact1 in zip(l1_vals, fact0s, fact1s):
if verbose:
print("Pol l1 = ", l1)
for l2 in range(2, lmax+1):
l3min = np.abs(l1-l2)
l3max = np.abs(l1+l2)
# Minimum l3 for when m3 is 1 or 2
l3min2 = np.max([l3min, 2])
l3min1 = np.max([l3min, 1])
# Wigner Symbols that we need are
# (l1 l2 l3 ) and (l1 l2 l3) for m=0 (pseudo)
# (-2+m 2 -m ) (2-m -2 m ) and m=1,2 (pure)
# m=0 term needed for pseudo and pure
wc_0 = wc.wigner3j_vect(2*l1, 2*l2, -2*2, 2*2)
wc_1 = wc.wigner3j_vect(2*l1, 2*l2, 2*2, -2*2)
Jp0 = wc_0 + wc_1
Jm0 = wc_0 - wc_1
# m=1,2 terms are only needed for pure modes
if (cl_type == 'pure') or (cl_type == 'hybrid'):
wc_0 = wc.wigner3j_vect(2*l1, 2*l2, (-2+1)*2, 2*2)
wc_1 = wc.wigner3j_vect(2*l1, 2*l2, (2-1)*2, -2*2)
Jp1 = wc_0 + wc_1
Jm1 = wc_0 - wc_1
wc_0 = wc.wigner3j_vect(2*l1, 2*l2, (-2+2)*2, 2*2)
wc_1 = wc.wigner3j_vect(2*l1, 2*l2, (2-2)*2, -2*2)
Jp2 = wc_0 + wc_1
Jm2 = wc_0 - wc_1
# This allows us to remove any sum over m or l3
idx = np.all([l >= l3min, l <= l3max], axis=0)
l_tmp = l[idx]
m_tmp = m[idx]
m0 = m_tmp == 0
wEBlm_tmp = wEBlm[:, idx]
# Calculate the correct index in the Jp and Jm terms
idx_l3_0 = np.array(l_tmp - l3min, dtype=np.int) # Jp0, Jm0
idx_l3_1 = np.array(l_tmp - l3min1, dtype=np.int) # Jp1, Jm1
idx_l3_2 = np.array(l_tmp - l3min2, dtype=np.int) # Jp2, Jm2
# When l3min is 0 or 1 and take the whole range, Jp1(Jm1) and/or
# Jp2(Jm2) get non-zero values when they should be zero (when l3 =
# 0 or 1) because negative indices in idx_l3_1 and idx_l3_2 wrap
# around to the end
idx1 = l_tmp >= l3min1
idx2 = l_tmp >= l3min2
# EE,EE and BB,BB (2/9)
termE = a2 * wEBlm_tmp[0, :] * Jp0[idx_l3_0]
termB = np.zeros_like(termE)
if cl_type == 'pseudo' or cl_type == 'hybrid':
Mpol[l1, l2] = 2*np.sum(termE*np.conj(termE))
Mpol[l1, l2] -= np.sum(termE[m0]*np.conj(termE[m0]))
if cl_type == 'pure' or cl_type == 'hybrid':
termE[idx1] += a1*fact1*wEBlm_tmp[1, idx1]*Jp1[idx_l3_1][idx1]
termE[idx2] += a0*fact0*wEBlm_tmp[2, idx2]*Jp2[idx_l3_2][idx2]
termB[idx1] += a1*fact1*wEBlm_tmp[-1, idx1]*Jm1[idx_l3_1][idx1]
termB[idx2] += a0*fact0*wEBlm_tmp[-2, idx2]*Jm2[idx_l3_2][idx2]
Mpol[l1+lmax+1, l2+lmax+1] = 2*np.sum(termE*np.conj(termE) +
termB*np.conj(termB))
Mpol[l1+lmax+1, l2+lmax+1] -= np.sum(termE[m0]*np.conj(termE[m0]) +
termB[m0]*np.conj(termB[m0]))
if cl_type == 'pure':
Mpol[l1, l2] = Mpol[l1+lmax+1, l2+lmax+1]
# EE,BB and BB,EE (4/9)
termE = a2 * wEBlm_tmp[0, :] * Jm0[idx_l3_0]
termB = np.zeros_like(termE)
if cl_type == 'pseudo' or cl_type == 'hybrid':
Mpol[l1, l2+lmax+1] = 2*np.sum(termE*np.conj(termE))
Mpol[l1, l2+lmax+1] -= np.sum(termE[m0]*np.conj(termE[m0]))
if cl_type == 'pseudo':
Mpol[l1+lmax+1, l2] = Mpol[l1, l2+lmax+1]
if cl_type == 'pure' or cl_type == 'hybrid':
termE[idx1] += a1*fact1*wEBlm_tmp[1, idx1]*Jm1[idx_l3_1][idx1]
termE[idx2] += a0*fact0*wEBlm_tmp[2, idx2]*Jm2[idx_l3_2][idx2]
termB[idx1] += a1*fact1*wEBlm_tmp[-1, idx1]*Jp1[idx_l3_1][idx1]
termB[idx2] += a0*fact0*wEBlm_tmp[-2, idx2]*Jp2[idx_l3_2][idx2]
Mpol[l1+lmax+1, l2] = 2*np.sum(termE*np.conj(termE) +
termB*np.conj(termB))
Mpol[l1+lmax+1, l2] -= np.sum(termE[m0]*np.conj(termE[m0]) +
termB[m0]*np.conj(termB[m0]))
if cl_type == 'pure':
Mpol[l1, l2+lmax+1] = Mpol[l1+lmax+1, l2]
# EB,EB (5/9)
if cl_type == 'pseudo' or cl_type == 'pure':
# EE,EE - EE,BB
Mpol[l1+2*lmax+2, l2+2*lmax+2] = Mpol[l1, l2] - Mpol[l1, l2+lmax+1]
else:
termE = np.zeros_like(Jp0[idx_l3_0], dtype=np.float64)
termE[idx2] += a0*fact0*(Jp0[idx_l3_0]*Jp2[idx_l3_2]-Jm0[idx_l3_0]*Jm2[idx_l3_2])[idx2] * np.real(wEBlm_tmp[0, idx2]*np.conj(wEBlm_tmp[2, idx2]))
termE[idx1] += a1*fact1*(Jp0[idx_l3_0]*Jp1[idx_l3_1]-Jm0[idx_l3_0]*Jm1[idx_l3_1])[idx1] * np.real(wEBlm_tmp[0, idx1]*np.conj(wEBlm_tmp[1, idx1]))
termE += a2 * (Jp0[idx_l3_0]*Jp0[idx_l3_0] - Jm0[idx_l3_0]*Jm0[idx_l3_0]) * np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[0, :]))
Mpol[l1+2*lmax+2, l2+2*lmax+2] = 2*np.sum(termE) - np.sum(termE[m0])
# EE,EB and BB,EB (7/9)
if cl_type == 'pure' or cl_type == 'hybrid':
termE = np.zeros_like(Jp0[idx_l3_0], dtype=np.float64)
# only non-zero for pure (and B is pure in hybrid).
termE[idx2] += a0*a0*fact0*fact0*(Jp2[idx_l3_2]*Jp2[idx_l3_2]*np.real(wEBlm_tmp[2, :]*np.conj(wEBlm_tmp[-2, :]))
- Jm2[idx_l3_2]*Jm2[idx_l3_2]*np.real(wEBlm_tmp[-2, :]*np.conj(wEBlm_tmp[2, :])))[idx2]
termE[idx2] += a0*a1*fact0*fact1*(Jp2[idx_l3_2]*Jp1[idx_l3_1]*np.real(wEBlm_tmp[2, :]*np.conj(wEBlm_tmp[-1, :]))
- Jm2[idx_l3_2]*Jm1[idx_l3_1]*np.real(wEBlm_tmp[-2, :]*np.conj(wEBlm_tmp[1, :])))[idx2]
termE[idx2] += a0*a2*fact0*(-Jm2[idx_l3_2]*Jm0[idx_l3_0]*np.real(wEBlm_tmp[-2, :]*np.conj(wEBlm_tmp[0, :])))[idx2]
termE[idx2] += a1*a0*fact1*fact0*(Jp1[idx_l3_1]*Jp2[idx_l3_2]*np.real(wEBlm_tmp[1, :]*np.conj(wEBlm_tmp[-2, :]))
- Jm1[idx_l3_1]*Jm2[idx_l3_2]*np.real(wEBlm_tmp[-1, :]*np.conj(wEBlm_tmp[2, :])))[idx2]
termE[idx1] += a1*a1*fact1*fact1*(Jp1[idx_l3_1]*Jp1[idx_l3_1]*np.real(wEBlm_tmp[1, :]*np.conj(wEBlm_tmp[-1, :]))
- Jm1[idx_l3_1]*Jm1[idx_l3_1]*np.real(wEBlm_tmp[-1, :]*np.conj(wEBlm_tmp[1, :])))[idx1]
termE[idx1] += a1*a2*fact1*(-Jm1[idx_l3_1]*Jm0[idx_l3_0]*np.real(wEBlm_tmp[-1, :]*np.conj(wEBlm_tmp[0, :])))[idx1]
termE[idx2] += a2*a0*fact0*(Jp0[idx_l3_0]*Jp2[idx_l3_2]*np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[-2, :])))[idx2]
termE[idx1] += a2*a1*fact1*(Jp0[idx_l3_0]*Jp1[idx_l3_1]*np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[-1, :])))[idx1]
Mpol[l1+lmax+1, l2+2*lmax+2] = 2*np.sum(termE) - np.sum(termE[m0])
if cl_type == 'pure':
Mpol[l1, l2+2*lmax+2] = Mpol[l1+lmax+1, l2+2*lmax+2]
# EB,EE and EB,BB (9/9)
if cl_type == 'pure':
Mpol[l1+2*lmax+2, l2] = Mpol[l1, l2+2*lmax+2]
Mpol[l1+2*lmax+2, l2+lmax+1] = Mpol[l1, l2+2*lmax+2]
elif cl_type == 'hybrid':
termE = np.zeros_like(Jp0[idx_l3_0], dtype=np.float64)
termE[idx2] += a0*fact0*(Jp0[idx_l3_0]*Jp2[idx_l3_2]*np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[-2, :])))[idx2]
termE[idx1] += a1*fact1*(Jp0[idx_l3_0]*Jp1[idx_l3_1]*np.real(wEBlm_tmp[0, :]*np.conj(wEBlm_tmp[-1, :])))[idx1]
termB = np.zeros_like(Jm0[idx_l3_0], dtype=np.float64)
termB[idx2] += a0*fact0*(Jm0[idx_l3_0]*Jm2[idx_l3_2]*np.real(wEBlm_tmp[-2, :]*np.conj(wEBlm_tmp[0, :])))[idx2]
termB[idx1] += a1*fact1*(Jm0[idx_l3_0]*Jm1[idx_l3_1]*np.real(wEBlm_tmp[-1, :]*np.conj(wEBlm_tmp[0, :])))[idx1]
Mpol[l1+2*lmax+2, l2] = 2*np.sum(termE) - np.sum(termE[m0])
Mpol[l1+2*lmax+2, l2+lmax+1] = 2*np.sum(termB) - np.sum(termB[m0])
normfact = (2.0*l2+1.0) / (4.0*np.pi)
Mpol[l1, l2] *= normfact / 4.0
Mpol[l1+lmax+1, l2+lmax+1] *= normfact / 4.0
Mpol[l1, l2+lmax+1] *= normfact / 4.0
Mpol[l1+lmax+1, l2] *= normfact / 4.0
Mpol[l1+2*lmax+2, l2+2*lmax+2] *= normfact / 4.0
Mpol[l1, l2+2*lmax+2] *= normfact / 2.0
Mpol[l1+lmax+1, l2+2*lmax+2] *= -normfact / 2.0
Mpol[l1+2*lmax+2, l2] *= -normfact / 4.0
Mpol[l1+2*lmax+2, l2+lmax+1] *= normfact / 4.0
Mpol = comm.allreduce(Mpol)
return Mpol