polarization covariances (all)

pspy/mcm_fortran/mcm_fortran.f90

@@ -341,4 +341,235 @@ subroutine binning_matrix(mcm, binLo, binHi, binsize, bbl, doDl)
 end subroutine
 
 
+! sub-blocks of the coupling matrix
+
+subroutine calc_coupling_elem_spin2_pp(wcl, l1, l2, elem)
+    implicit none
+    integer, intent(in)   :: l1, l2
+    real(8), intent(in)   :: wcl(:)
+    real(8), intent(inout):: elem
+    real(8) :: thrcof1(2*size(wcl)), l1f, l2f
+    integer :: info, l3, wlmin, wlmax, i, lmax
+    lmax = size(wcl)-1 ! wcl starts at 0
+
+    call drc3jj(dble(l1), dble(l2), -2d0, 2d0, l1f, l2f, thrcof1, size(thrcof1), info)
+
+    wlmin = int(l1f)
+    wlmax = min(lmax, int(l2f))
+    elem = 0
+    do l3 = wlmin, wlmax
+        i = l3 - wlmin + 1
+        elem = elem + wcl(l3 + 1) * thrcof1(i)**2 * (1 + (-1)**(l1 + l2 + l3)) / 2
+    end do
+end subroutine
+
+
+subroutine calc_coupling_elem_spin02(wcl, l1, l2, elem)
+    implicit none
+    integer, intent(in)   :: l1, l2
+    real(8), intent(in)   :: wcl(:)
+    real(8), intent(inout):: elem
+    real(8) :: thrcof0(2*size(wcl)), thrcof1(2*size(wcl)), l1f, l2f
+    integer :: info, l3, wlmin, wlmax, i, lmax
+    lmax = size(wcl)-1 ! wcl starts at 0
+
+    call drc3jj(dble(l1), dble(l2), 0d0, 0d0, l1f, l2f, thrcof0, size(thrcof0), info)
+    call drc3jj(dble(l1), dble(l2), -2d0, 2d0, l1f, l2f, thrcof1, size(thrcof1), info)
+
+    wlmin = int(l1f)
+    wlmax = min(lmax, int(l2f))
+    elem = 0
+    do l3 = wlmin, wlmax
+        i = l3 - wlmin + 1
+        elem = elem + wcl(l3 + 1) * thrcof0(i) * thrcof1(i)
+    end do
+end subroutine
+
+
+subroutine calc_coupling_elem_spin2_mm(wcl, l1, l2, elem)
+    implicit none
+    integer, intent(in)   :: l1, l2
+    real(8), intent(in)   :: wcl(:)
+    real(8), intent(inout):: elem
+    real(8) :: thrcof1(2*size(wcl)), l1f, l2f
+    integer :: info, l3, wlmin, wlmax, i, lmax
+    lmax = size(wcl)-1 ! wcl starts at 0
+
+    call drc3jj(dble(l1), dble(l2), -2d0, 2d0, l1f, l2f, thrcof1, size(thrcof1), info)
+
+    wlmin = int(l1f)
+    wlmax = min(lmax, int(l2f))
+    elem = 0
+    do l3 = wlmin, wlmax
+        i = l3 - wlmin + 1
+        elem = elem + wcl(l3 + 1) * thrcof1(i)**2 * (1 - (-1)**(l1 + l2 + l3)) / 2
+    end do
+end subroutine
+
+
+! unlike calc_coupling_spin0, this keeps ell=0,1,lmax
+subroutine calc_coupling_block_spin0(wcl, l_exact, l_band, l_toeplitz, coupling)
+    implicit none
+    real(8), intent(in)    :: wcl(:)
+    integer, intent(in)    :: l_exact , l_band, l_toeplitz
+    real(8), intent(inout) :: coupling(:,:)
+    integer :: l1, l2, nlmax, lmax_band
+
+    nlmax = size(coupling,1) - 1
+
+    !$omp parallel do private(l2, l1) schedule(dynamic)
+    do l1 = 2, min(nlmax, l_exact)
+        do l2 = l1, nlmax
+            call calc_coupling_elem_spin0(wcl, l1, l2, coupling(l1-1, l2-1))
+        end do
+    end do
+
+    if (l_exact .lt. nlmax) then
+        !$omp parallel do private(l2, l1, lmax_band) schedule(dynamic)
+        do l1 = l_exact + 1, l_toeplitz
+            if (l1 .lt. l_toeplitz) then
+                lmax_band = min(l1 + l_band, nlmax)
+            else
+                lmax_band = nlmax
+            end if
+
+            do l2 = l1, lmax_band
+                call calc_coupling_elem_spin0(wcl, l1, l2, coupling(l1-1, l2-1))
+            end do
+        end do
+
+        if (l_toeplitz .lt. nlmax) then
+            !$omp parallel do
+            do l1 = l_toeplitz + 1, nlmax
+                call calc_coupling_elem_spin0(wcl, l1, l1, coupling(l1-1, l1-1))
+            end do
+        end if
+    end if
+
+end subroutine
+
+
+subroutine calc_coupling_block_spin02(wcl, l_exact, l_band, l_toeplitz, coupling)
+    implicit none
+    real(8), intent(in)    :: wcl(:)
+    integer, intent(in)    :: l_exact , l_band, l_toeplitz
+    real(8), intent(inout) :: coupling(:,:)
+    integer :: l1, l2, nlmax, lmax_band
+
+    nlmax = size(coupling,1) - 1
+
+    !$omp parallel do private(l2, l1) schedule(dynamic)
+    do l1 = 2, min(nlmax, l_exact)
+        do l2 = l1, nlmax
+            call calc_coupling_elem_spin02(wcl, l1, l2, coupling(l1-1, l2-1))
+        end do
+    end do
+
+    if (l_exact .lt. nlmax) then
+        !$omp parallel do private(l2, l1, lmax_band) schedule(dynamic)
+        do l1 = l_exact + 1, l_toeplitz
+            if (l1 .lt. l_toeplitz) then
+                lmax_band = min(l1 + l_band, nlmax)
+            else
+                lmax_band = nlmax
+            end if
+
+            do l2 = l1, lmax_band
+                call calc_coupling_elem_spin02(wcl, l1, l2, coupling(l1-1, l2-1))
+            end do
+        end do
+
+        if (l_toeplitz .lt. nlmax) then
+            !$omp parallel do
+            do l1 = l_toeplitz + 1, nlmax
+                call calc_coupling_elem_spin02(wcl, l1, l1, coupling(l1-1, l1-1))
+            end do
+        end if
+    end if
+
+end subroutine
+
+
+subroutine calc_coupling_block_spin2_pp(wcl, l_exact, l_band, l_toeplitz, coupling)
+    implicit none
+    real(8), intent(in)    :: wcl(:)
+    integer, intent(in)    :: l_exact , l_band, l_toeplitz
+    real(8), intent(inout) :: coupling(:,:)
+    integer :: l1, l2, nlmax, lmax_band
+
+    nlmax = size(coupling,1) - 1
+
+    !$omp parallel do private(l2, l1) schedule(dynamic)
+    do l1 = 2, min(nlmax, l_exact)
+        do l2 = l1, nlmax
+            call calc_coupling_elem_spin2_pp(wcl, l1, l2, coupling(l1-1, l2-1))
+        end do
+    end do
+
+    if (l_exact .lt. nlmax) then
+        !$omp parallel do private(l2, l1, lmax_band) schedule(dynamic)
+        do l1 = l_exact + 1, l_toeplitz
+            if (l1 .lt. l_toeplitz) then
+                lmax_band = min(l1 + l_band, nlmax)
+            else
+                lmax_band = nlmax
+            end if
+
+            do l2 = l1, lmax_band
+                call calc_coupling_elem_spin2_pp(wcl, l1, l2, coupling(l1-1, l2-1))
+            end do
+        end do
+
+        if (l_toeplitz .lt. nlmax) then
+            !$omp parallel do
+            do l1 = l_toeplitz + 1, nlmax
+                call calc_coupling_elem_spin2_pp(wcl, l1, l1, coupling(l1-1, l1-1))
+            end do
+        end if
+    end if
+
+end subroutine
+
+
+subroutine calc_coupling_block_spin2_mm(wcl, l_exact, l_band, l_toeplitz, coupling)
+    implicit none
+    real(8), intent(in)    :: wcl(:)
+    integer, intent(in)    :: l_exact , l_band, l_toeplitz
+    real(8), intent(inout) :: coupling(:,:)
+    integer :: l1, l2, nlmax, lmax_band
+
+    nlmax = size(coupling,1) - 1
+
+    !$omp parallel do private(l2, l1) schedule(dynamic)
+    do l1 = 2, min(nlmax, l_exact)
+        do l2 = l1, nlmax
+            call calc_coupling_elem_spin2_mm(wcl, l1, l2, coupling(l1-1, l2-1))
+        end do
+    end do
+
+    if (l_exact .lt. nlmax) then
+        !$omp parallel do private(l2, l1, lmax_band) schedule(dynamic)
+        do l1 = l_exact + 1, l_toeplitz
+            if (l1 .lt. l_toeplitz) then
+                lmax_band = min(l1 + l_band, nlmax)
+            else
+                lmax_band = nlmax
+            end if
+
+            do l2 = l1, lmax_band
+                call calc_coupling_elem_spin2_mm(wcl, l1, l2, coupling(l1-1, l2-1))
+            end do
+        end do
+
+        if (l_toeplitz .lt. nlmax) then
+            !$omp parallel do
+            do l1 = l_toeplitz + 1, nlmax
+                call calc_coupling_elem_spin2_mm(wcl, l1, l1, coupling(l1-1, l1-1))
+            end do
+        end if
+    end if
+
+end subroutine
+
+
 end module

pspy/so_cov.py

@@ -1084,142 +1084,3 @@ def measure_white_noise_level(var, mask, dtype=np.float64):
     buffer *= var.pixsizemap().flatten()
     noise_tot = np.sum(buffer)
     return noise_tot / mask_tot
-
-# only works for TTTT for now
-def generate_aniso_couplings(survey_name, win, var, lmax, niter=0):
-    """
-    survey_name : list of names of the four splits
-    win : dict with windows with keys 'Ta', "Tb', ...
-    var : dict with variance maps, with keys 'Ta', 'Tb'
-    """
-    survey_id = ["Ta", "Tb", "Tc", "Td"]
-    # set up some tools to work with the survey names
-    id2name = dict(zip(survey_id, survey_name))
-
-    var_a, var_b, var_c, var_d = (
-        win['Ta'].copy(), win['Tb'].copy(), win['Tc'].copy(), win['Td'].copy())
-
-    pixsizes = win['Ta'].data.pixsizemap()  # assuming they're the same footprint
-    var_a.data *= np.sqrt(pixsizes * var['Ta'].data)
-    var_b.data *= np.sqrt(pixsizes * var['Tb'].data)
-    var_c.data *= np.sqrt(pixsizes * var['Tc'].data)
-    var_d.data *= np.sqrt(pixsizes * var['Td'].data)
-
-    zero_var = win['Ta'].copy()
-    zero_var.data *= 0
-
-    # allows indexing with the delta function
-    v_a = (zero_var, var_a)
-    v_b = (zero_var, var_b)
-    v_c = (zero_var, var_c)
-    v_d = (zero_var, var_d)
-
-    coupling_1 = cov_coupling_spin0(
-        {'Ta': win['Ta'],
-        'Tb': win['Tb'],
-        'Tc': win['Tc'],
-        'Td': win['Td']},
-        lmax, niter)['TaTcTbTd']
-
-    coupling_2 = cov_coupling_spin0(
-        {'Ta': win['Ta'],
-        'Tb': win['Tb'],
-        'Tc': win['Tc'],
-        'Td': win['Td']},
-        lmax, niter)['TaTdTbTc']
-
-    coupling_3 = cov_coupling_spin0(
-        {'Ta': win['Ta'],
-        'Tb': v_b[delta2(id2name['Tb'], id2name['Td'])],
-        'Tc': win['Tc'],
-        'Td': v_d[delta2(id2name['Tb'], id2name['Td'])]},
-        lmax, niter)['TaTcTbTd']
-
-    coupling_4 = cov_coupling_spin0(
-        {'Ta': v_a[delta2(id2name['Ta'], id2name['Tc'])],
-        'Tb': win['Tb'],
-        'Tc': v_c[delta2(id2name['Ta'], id2name['Tc'])],
-        'Td': win['Td']},
-        lmax, niter)['TaTcTbTd']
-
-    coupling_5 = cov_coupling_spin0(
-        {'Ta': win['Ta'],
-        'Tb': v_a[delta2(id2name['Tb'], id2name['Tc'])],
-        'Tc': v_c[delta2(id2name['Tb'], id2name['Tc'])],
-        'Td': win['Td']},
-        lmax, niter)['TaTdTbTc']
-
-    coupling_6 = cov_coupling_spin0(
-        {'Ta': v_a[delta2(id2name['Ta'], id2name['Td'])],
-        'Tb': win['Tb'],
-        'Tc': win['Tc'],
-        'Td': v_d[delta2(id2name['Ta'], id2name['Td'])]},
-        lmax, niter)['TaTdTbTc']
-
-    coupling_7 = cov_coupling_spin0(
-        {'Ta': v_a[delta2(id2name['Ta'], id2name['Tc'])],
-        'Tb': v_b[delta2(id2name['Tb'], id2name['Td'])],
-        'Tc': v_c[delta2(id2name['Ta'], id2name['Tc'])],
-        'Td': v_d[delta2(id2name['Tb'], id2name['Td'])]},
-        lmax, niter)['TaTcTbTd']
-
-    coupling_8 = cov_coupling_spin0(
-        {'Ta': v_a[delta2(id2name['Ta'], id2name['Td'])],
-        'Tb': v_b[delta2(id2name['Tb'], id2name['Tc'])],
-        'Tc': v_c[delta2(id2name['Tb'], id2name['Tc'])],
-        'Td': v_d[delta2(id2name['Ta'], id2name['Td'])]},
-        lmax, niter)['TaTcTbTd']
-
-    return [coupling_1, coupling_2, coupling_3, coupling_4,
-            coupling_5, coupling_6, coupling_7, coupling_8]
-
-
-def cov_spin0_aniso1(Clth, Rl, couplings, binning_file, lmax,
-                     mbb_inv_ab, mbb_inv_cd, binned_mcm=True):
-    """Estimate the analytic covariance in the case of anisotropic noise pixel variance,
-       using the theoretical Cls, the ratios of noise to white noise power spectra, the
-       coupling matrices estimated from the masks and variance maps, and the mode-coupling
-       matrices.
-
-    Parameters
-    ----------
-
-    Clth: dictionary
-      A dictionary of theoretical power spectrum (auto and cross) for the different split
-      combinations ('TaTb' etc)
-    Rl: dictionary
-      a dictionary containing the ratios of noise power spectra to white noise
-    couplings: list of arrays
-      a list containing the eight coupling matrices involved in this covariance
-    binning_file: data file
-      a binning file with format bin low, bin high, bin mean
-    lmax: int
-      the maximum multipole to consider
-    mbb_inv_ab: 2d array
-      the inverse mode coupling matrix for the 'TaTb' power spectrum
-    mbb_inv_cd: 2d array
-      the inverse mode coupling matrix for the 'TcTd' power spectrum
-    binned_mcm: boolean
-      specify if the mode coupling matrices are binned or not
-
-    """
-
-    geom = lambda cl : symmetrize(cl, mode="geo")
-
-    cov =  geom(Clth["TaTc"]) * geom(Clth["TbTd"]) * couplings[0],
-    cov += geom(Clth["TaTd"]) * geom(Clth["TbTc"]) * couplings[1]
-    cov += geom(Rl['Tb']) * geom(Rl['Td']) * geom(Clth["TaTc"]) * couplings[2]
-    cov += geom(Rl['Ta']) * geom(Rl['Tc']) * geom(Clth["TbTd"]) * couplings[3]
-    cov += geom(Rl['Tb']) * geom(Rl['Tc']) * geom(Clth["TaTd"]) * couplings[4]
-    cov += geom(Rl['Ta']) * geom(Rl['Td']) * geom(Clth["TbTc"]) * couplings[5]
-    cov += geom(Rl['Ta']) * geom(Rl['Tb']) * geom(Rl['Tc']) * geom(Rl['Td']) * couplings[6]
-    cov += geom(Rl['Ta']) * geom(Rl['Tb']) * geom(Rl['Tc']) * geom(Rl['Td']) * couplings[7]
-
-    if binned_mcm == True:
-        analytic_cov = bin_mat(cov, binning_file, lmax)
-        analytic_cov = mbb_inv_ab @ analytic_cov @ mbb_inv_cd.T
-    else:
-        full_analytic_cov = mbb_inv_ab @ cov @ mbb_inv_cd.T
-        analytic_cov = bin_mat(full_analytic_cov, binning_file, lmax)
-
-    return cov