def _get_response(qes, source, cls_cmb, fal_leg1, lmax_qlm, fal_leg2=None):
fal_leg2 = fal_leg1 if fal_leg2 is None else fal_leg2
RGG = np.zeros(lmax_qlm + 1, dtype=float)
RCC = np.zeros(lmax_qlm + 1, dtype=float)
RGC = np.zeros(lmax_qlm + 1, dtype=float)
RCG = np.zeros(lmax_qlm + 1, dtype=float)
Ls = np.arange(lmax_qlm + 1, dtype=int)
for qe in qes:
si, ti = (qe.leg_a.spin_in, qe.leg_b.spin_in)
so, to = (qe.leg_a.spin_ou, qe.leg_b.spin_ou)
for s2 in ([0, -2, 2]):
FA = uspin.get_spin_matrix(si, s2, fal_leg1)
if np.any(FA):
for t2 in ([0, -2, 2]):
FB = uspin.get_spin_matrix(ti, t2, fal_leg2)
if np.any(FB):
rW_st, prW_st, mrW_st, s_cL_st = get_covresp(source, -s2, t2, cls_cmb, len(FB) - 1)
clA = ut.joincls([qe.leg_a.cl, FA])
clB = ut.joincls([qe.leg_b.cl, FB, mrW_st.conj()])
Rpr_st = uspin.wignerc(clA, clB, so, s2, to, -s2 + rW_st, lmax_out=lmax_qlm) * s_cL_st(Ls)
rW_ts, prW_ts, mrW_ts, s_cL_ts = get_covresp(source, -t2, s2, cls_cmb, len(FA) - 1)
clA = ut.joincls([qe.leg_a.cl, FA, mrW_ts.conj()])
clB = ut.joincls([qe.leg_b.cl, FB])
Rpr_st = Rpr_st + uspin.wignerc(clA, clB, so, -t2 + rW_ts, to, t2, lmax_out=lmax_qlm) * s_cL_ts(Ls)
assert rW_st == rW_ts and rW_st >= 0, (rW_st, rW_ts)
if rW_st > 0:
clA = ut.joincls([qe.leg_a.cl, FA])
clB = ut.joincls([qe.leg_b.cl, FB, prW_st.conj()])
Rmr_st = uspin.wignerc(clA, clB, so, s2, to, -s2 - rW_st, lmax_out=lmax_qlm) * s_cL_st(Ls)
clA = ut.joincls([qe.leg_a.cl, FA, prW_ts.conj()])
clB = ut.joincls([qe.leg_b.cl, FB])
Rmr_st = Rmr_st + uspin.wignerc(clA, clB, so, -t2 - rW_ts, to, t2, lmax_out=lmax_qlm) * s_cL_ts(Ls)
else:
Rmr_st = Rpr_st
prefac = qe.cL(Ls)
RGG += prefac * ( Rpr_st.real + Rmr_st.real * (-1) ** rW_st)
RCC += prefac * ( Rpr_st.real - Rmr_st.real * (-1) ** rW_st)
RGC += prefac * (-Rpr_st.imag + Rmr_st.imag * (-1) ** rW_st)
RCG += prefac * ( Rpr_st.imag + Rmr_st.imag * (-1) ** rW_st)
return RGG, RCC, RGC, RCG
def _get_nhl(qes1, qes2, cls_ivfs, lmax_out, cls_ivfs_bb=None, cls_ivfs_ab=None, ret_terms=False):
GG_N0 = np.zeros(lmax_out + 1, dtype=float)
CC_N0 = np.zeros(lmax_out + 1, dtype=float)
GC_N0 = np.zeros(lmax_out + 1, dtype=float)
CG_N0 = np.zeros(lmax_out + 1, dtype=float)
cls_ivfs_aa = cls_ivfs
cls_ivfs_bb = cls_ivfs if cls_ivfs_bb is None else cls_ivfs_bb
cls_ivfs_ab = cls_ivfs if cls_ivfs_ab is None else cls_ivfs_ab
cls_ivfs_ba = cls_ivfs_ab
if ret_terms:
terms = []
for qe1 in qes1:
cL1 = qe1.cL(np.arange(lmax_out + 1))
for qe2 in qes2:
cL2 = qe2.cL(np.arange(lmax_out + 1))
si, ti, ui, vi = (qe1.leg_a.spin_in, qe1.leg_b.spin_in, qe2.leg_a.spin_in, qe2.leg_b.spin_in)
so, to, uo, vo = (qe1.leg_a.spin_ou, qe1.leg_b.spin_ou, qe2.leg_a.spin_ou, qe2.leg_b.spin_ou)
assert so + to >= 0 and uo + vo >= 0, (so, to, uo, vo)
clsu = utils.joincls([qe1.leg_a.cl, qe2.leg_a.cl.conj(), uspin.spin_cls(si, ui, cls_ivfs_aa)])
cltv = utils.joincls([qe1.leg_b.cl, qe2.leg_b.cl.conj(), uspin.spin_cls(ti, vi, cls_ivfs_bb)])
R_sutv = utils.joincls([uspin.wignerc(clsu, cltv, so, uo, to, vo, lmax_out=lmax_out), cL1, cL2])
clsv = utils.joincls([qe1.leg_a.cl, qe2.leg_b.cl.conj(), uspin.spin_cls(si, vi, cls_ivfs_ab)])
cltu = utils.joincls([qe1.leg_b.cl, qe2.leg_a.cl.conj(), uspin.spin_cls(ti, ui, cls_ivfs_ba)])
R_sutv = R_sutv + utils.joincls([uspin.wignerc(clsv, cltu, so, vo, to, uo, lmax_out=lmax_out), cL1, cL2])
# we now need -s-t uv
sgnms = (-1) ** (si + so)
sgnmt = (-1) ** (ti + to)
clsu = utils.joincls([sgnms * qe1.leg_a.cl.conj(), qe2.leg_a.cl.conj(), uspin.spin_cls(-si, ui, cls_ivfs_aa)])
cltv = utils.joincls([sgnmt * qe1.leg_b.cl.conj(), qe2.leg_b.cl.conj(), uspin.spin_cls(-ti, vi, cls_ivfs_bb)])
R_msmtuv = utils.joincls([uspin.wignerc(clsu, cltv, -so, uo, -to, vo, lmax_out=lmax_out), cL1, cL2])
clsv = utils.joincls([sgnms * qe1.leg_a.cl.conj(), qe2.leg_b.cl.conj(), uspin.spin_cls(-si, vi, cls_ivfs_ab)])
cltu = utils.joincls([sgnmt * qe1.leg_b.cl.conj(), qe2.leg_a.cl.conj(), uspin.spin_cls(-ti, ui, cls_ivfs_ba)])
R_msmtuv = R_msmtuv + utils.joincls([uspin.wignerc(clsv, cltu, -so, vo, -to, uo, lmax_out=lmax_out), cL1, cL2])
GG_N0 += 0.5 * R_sutv.real
GG_N0 += 0.5 * (-1) ** (to + so) * R_msmtuv.real
CC_N0 += 0.5 * R_sutv.real
CC_N0 -= 0.5 * (-1) ** (to + so) * R_msmtuv.real
GC_N0 -= 0.5 * R_sutv.imag
GC_N0 -= 0.5 * (-1) ** (to + so) * R_msmtuv.imag
CG_N0 += 0.5 * R_sutv.imag
CG_N0 -= 0.5 * (-1) ** (to + so) * R_msmtuv.imag
if ret_terms:
terms += [0.5 * R_sutv, 0.5 * (-1) ** (to + so) * R_msmtuv]
return (GG_N0, CC_N0, GC_N0, CG_N0) if not ret_terms else (GG_N0, CC_N0, GC_N0, CG_N0, terms)