diff --git a/pspy/mcm_fortran/mcm_fortran.f90 b/pspy/mcm_fortran/mcm_fortran.f90 index 34aa8a9..8bd9d51 100644 --- a/pspy/mcm_fortran/mcm_fortran.f90 +++ b/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 diff --git a/pspy/so_cov.py b/pspy/so_cov.py index d46ee7b..73e377e 100644 --- a/pspy/so_cov.py +++ b/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 diff --git a/pspy/so_map_preprocessing.py b/pspy/so_map_preprocessing.py index 0fe34b5..739161b 100644 --- a/pspy/so_map_preprocessing.py +++ b/pspy/so_map_preprocessing.py @@ -1,12 +1,43 @@ import numpy as np, pylab as plt +from scipy.integrate import quad from pspy import flat_tools, pspy_utils from pixell import enmap, utils - -def build_std_filter(shape, wcs, vk_mask, hk_mask, dtype=np.float64): +from itertools import product + +def build_std_filter(shape, wcs, vk_mask, hk_mask, dtype=np.float64, + shape_y=None, wcs_y=None): + """Construct a "standard," binary-cross style filter. + + Parameters + ---------- + shape : (ny, nx) tuple + Shape of map and therefore shape of DFT. + wcs : astropy.wcs.WCS + wcs describing the map, used in enmap.laxes + vk_mask : (kx0, kx1) tuple + Modes between kx0 and kx1 (exclusive) are set to 0, otherwise 1. + hk_mask : (ky0, ky1) tuple + Modes between ky0 and ky1 (exclusive) are set to 0, otherwise 1. + dtype : np.dtype, optional + Type of output filter array, by default np.float64. + shape_y : (ny, nx) tuple, optional + If provided, together with wcs_y, used to override the y-spacing and + shape in Fourier space in enmap.laxes, by default None. + wcs_y : astropy.wcs.WCS, optional + If provided, together with shape_y, used to override the y-spacing and + shape in Fourier space in enmap.laxes, by default None. + + Returns + ------- + (ny, nx) enmap.ndmap + The binary-cross filter. + """ ly, lx = enmap.laxes(shape, wcs, method="auto") + if shape_y is not None and wcs_y is not None: + ly, _ = enmap.laxes(shape_y, wcs_y, method="auto") ly = ly.astype(dtype) lx = lx.astype(dtype) - filter = enmap.ones(shape[-2:], wcs, dtype) + filter = enmap.ones((ly.size, lx.size), wcs, dtype) if vk_mask is not None: id_vk = np.where((lx > vk_mask[0]) & (lx < vk_mask[1])) filter[:, id_vk] = 0 @@ -25,8 +56,11 @@ def apply_std_filter(imap, filter): return filtered_map -def build_sigurd_filter(shape, wcs, lbounds, dtype=np.float64): +def build_sigurd_filter(shape, wcs, lbounds, dtype=np.float64, + shape_y=None, wcs_y=None): ly, lx = enmap.laxes(shape, wcs, method="auto") + if shape_y is not None and wcs_y is not None: + ly, _ = enmap.laxes(shape_y, wcs_y, method="auto") ly = ly.astype(dtype) lx = lx.astype(dtype) lbounds = np.asarray(lbounds) @@ -34,7 +68,7 @@ def build_sigurd_filter(shape, wcs, lbounds, dtype=np.float64): lbounds = np.broadcast_to(lbounds, (1,2)) if lbounds.ndim > 2 or lbounds.shape[-1] != 2: raise ValueError("lbounds must be [:,{ly,lx}]") - filter = enmap.ones(shape[-2:], wcs, dtype) + filter = enmap.ones((ly.size, lx.size), wcs, dtype) # Apply the filters for i , (ycut, xcut) in enumerate(lbounds): filter *= 1-(np.exp(-0.5*(ly/ycut)**2)[:,None]*np.exp(-0.5*(lx/xcut)**2)[None,:]) @@ -68,7 +102,51 @@ def analytical_tf(map, filter, binning_file, lmax): return bin_c, tf - +def analytical_tf_vkhk(vk_mask, hk_mask, ells, dtype=np.float64): + if vk_mask is None: + vk_mask = [0, 0] + if hk_mask is None: + hk_mask = [0, 0] + assert len(vk_mask) == 2 and len(hk_mask) == 2, \ + f'{len(vk_mask)=} and {len(hk_mask)=} should both be 2' + + out = np.zeros(ells.shape, dtype=dtype) + for vk, hk in product(vk_mask, hk_mask): + rk = np.sqrt(vk**2 + hk**2) + out[ells > rk] += .25 - (np.abs(np.arcsin(vk/ells[ells > rk])) + np.abs(np.arcsin(hk/ells[ells > rk]))) / (2*np.pi) + return out + +def analytical_tf_yxint(yxfilter_func, ells, x_fac=1, y_fac=1, dtype=np.float64): + def x(r, t): + return r*x_fac*np.cos(t) + def y(r, t): + return r*y_fac*np.sin(t) + + if x_fac == y_fac: + def v(r, t): + return r*x_fac + else: + def v(r, t): + return r*np.sqrt(x_fac**2 * np.sin(t)**2 + y_fac**2 * np.cos(t)**2) + + assert ells.ndim == 1, 'expect 1d ells' + assert np.all(ells >= 0), 'expect non-negative ells' + out = np.zeros(ells.shape, dtype=dtype) + for i, ell in enumerate(ells): + if i % 100 == 0: + print(ell) + + def vell(t): + return v(ell, t) + def rphifilter_func(t): + return yxfilter_func(y(ell, t), x(ell, t)) * vell(t) + + if x_fac == y_fac: + out[i] = quad(rphifilter_func, 0, 2*np.pi)[0] / (2*np.pi*ell*x_fac) + else: + out[i] = quad(rphifilter_func, 0, 2*np.pi)[0] / quad(vell, 0, 2*np.pi)[0] + + return out #def analytical_tf_old(map, binning_file, lmax, vk_mask=None, hk_mask=None): # import time diff --git a/pspy/so_mcm.py b/pspy/so_mcm.py index 576ede6..4e9f494 100644 --- a/pspy/so_mcm.py +++ b/pspy/so_mcm.py @@ -473,3 +473,86 @@ def read_coupling(prefix, spin_pairs=None): Bbl = np.load(prefix + "_Bbl.npy") return mode_coupling_inv, Bbl + + +def coupling_block( + mode, + win1, + lmax, + niter=None, + win2=None, + input_alm=False, + l_exact=None, + l_toep=None, + l_band=None, + l3_pad=2000, + legacy_ell_range=False): + + """Get the coupling matrix for two fields + Parameters + ---------- + mode: string + one of "00", "02", "++" or "--" + win1: so_map (or alm) + the window function of survey 1, if input_alm=True, expect wlm1 + lmax: integer + the maximum multipole to consider for the spectra computation + win2: so_map (or alm) + the window function of survey 2, if input_alm=True, expect wlm2 + niter: int + specify the number of iteration in map2alm + l_toep: int + l_band: int + l_exact: int + """ + + if input_alm == False: + if niter is None: + raise ValueError("niter required") + l_max_limit = win1.get_lmax_limit() + if lmax > l_max_limit: raise ValueError("the requested lmax is too high with respect to the map pixellisation") + maxl = np.minimum(lmax + l3_pad, l_max_limit) + + win1 = sph_tools.map2alm(win1, niter=niter, lmax=maxl) + if win2 is not None: + win2 = sph_tools.map2alm(win2, niter=niter, lmax=maxl) + + if win2 is None: + wcl = hp.alm2cl(win1) + else: + wcl = hp.alm2cl(win1, win2) + + l = np.arange(len(wcl)) + wcl *= (2 * l + 1) + + mcm = np.zeros((lmax + 1, lmax + 1)) + + if l_toep is None: l_toep = lmax + 1 + if l_band is None: l_band = lmax + 1 + if l_exact is None: l_exact = lmax + 1 + + if mode == "00": + mcm_fortran.calc_coupling_block_spin0(wcl, l_exact, l_band, l_toep, mcm.T) + elif mode == "02": + mcm_fortran.calc_coupling_block_spin02(wcl, l_exact, l_band, l_toep, mcm.T) + elif mode == "++": + mcm_fortran.calc_coupling_block_spin2_pp(wcl, l_exact, l_band, l_toep, mcm.T) + elif mode == "--": + mcm_fortran.calc_coupling_block_spin2_mm(wcl, l_exact, l_band, l_toep, mcm.T) + else: + raise ValueError(f"mode \"{mode}\" not one of (\"00\", \"02\", \"++\", \"--\")") + + if l_toep < lmax: + mcm = format_toepliz_fortran2(mcm, l_toep, l_exact, lmax) + + mcm_fortran.fill_upper(mcm.T) + + mcm_full = np.zeros((lmax + 1, lmax + 1)) + mcm_full[2:, 2:] = mcm[:-2, :-2] + mcm_full[0, 0] = 1.0 + mcm_full[1, 1] = 1.0 + + if legacy_ell_range: + return mcm_full[2:-1, 2:-1] + else: + return mcm_full diff --git a/pspy/tests/test_so_cov.py b/pspy/tests/test_so_cov.py index 5146983..a7da0e9 100644 --- a/pspy/tests/test_so_cov.py +++ b/pspy/tests/test_so_cov.py @@ -2,87 +2,41 @@ import tempfile, unittest, os import numpy as np -from pspy import so_cov, so_map, so_mcm, pspy_utils +from pspy import so_cov, so_map, so_mcm, pspy_utils, so_window TEST_DATA_PREFIX = os.path.join(os.path.dirname(__file__), "data", "example_newcov") class SOCovmatTests(unittest.TestCase): - def setUp(self, verbose=False): - # read in downgraded windows and ivar - self.lmax = 1350 - self.window = so_map.read_map(os.path.join(TEST_DATA_PREFIX, "window_adjusted.fits")) - var0 = so_map.read_map(os.path.join(TEST_DATA_PREFIX, "ivar1_adjusted.fits")) - var1 = so_map.read_map(os.path.join(TEST_DATA_PREFIX, "ivar2_adjusted.fits")) - self.cl_dict = np.load(os.path.join(TEST_DATA_PREFIX, "cl_dict.npy")) - self.nl_dict = np.load(os.path.join(TEST_DATA_PREFIX, "nl_dict.npy")) - self.cov_ref = np.load(os.path.join(TEST_DATA_PREFIX, "analytic_cov_ll2.npy")) + def test_custom_couplings(self, verbose=False): + ra0, ra1, dec0, dec1 = -25, 25, -15, 15 + res = 2 + lmax = 500 + l_exact, l_band, l_toep = 50, 100, 200 - var0.data = 1 / var0.data - var1.data = 1 / var1.data - var0.data[np.isinf(var0.data)] = 0. - var1.data[np.isinf(var1.data)] = 0. - self.var0, self.var1 = var0, var1 + binary_car = so_map.car_template(1, ra0, ra1, dec0, dec1, res) + binary_car.data[:] = 0 + binary_car.data[1:-1, 1:-1] = 1 + window = so_window.create_apodization(binary_car, apo_type="C1", apo_radius_degree=2) + window2 = so_window.create_apodization(binary_car, apo_type="C2", apo_radius_degree=4) - self.binning_file = os.path.join(TEST_DATA_PREFIX, "binning.dat") # ignored - # pspy_utils.create_binning_file(bin_size=40, n_bins=100, file_name=self.binning_file) + default_coup_02 = so_mcm.mcm_and_bbl_spin0and2( + (window, window2), "", lmax=lmax, niter=0, l_exact=l_exact, l_band=l_band, l_toep=l_toep, return_coupling_only=True) - + xi = so_mcm.coupling_block("00", window, lmax, niter=0, l_exact=l_exact, l_band=l_band, l_toep=l_toep) + np.testing.assert_almost_equal( default_coup_02[0,:,:], xi[2:-1,2:-1]) - def test_covmat(self, verbose=False): - lmax, binning_file = self.lmax, self.binning_file - window, var0, var1 = self.window, self.var0, self.var1 - cl_dict, nl_dict = self.cl_dict, self.nl_dict + xi = so_mcm.coupling_block("02", window, lmax, niter=0, win2=window2, l_exact=l_exact, l_band=l_band, l_toep=l_toep) + np.testing.assert_almost_equal( default_coup_02[1,:,:], xi[2:-1,2:-1]) - # set up scenario - survey_id = ["Ta", "Tb", "Tc", "Td"] - survey_name = ["split_0", "split_1", "split_0", "split_1"] - win = {'Ta': window, 'Tb': window, 'Tc': window, 'Td': window} - var = {'Ta': var0, 'Tb': var1, 'Tc': var0, 'Td': var1} + xi = so_mcm.coupling_block("02", window2, lmax, niter=0, win2=window, l_exact=l_exact, l_band=l_band, l_toep=l_toep) + np.testing.assert_almost_equal( default_coup_02[2,:,:], xi[2:-1,2:-1]) - # generate mcm - mbb_inv, Bbl = so_mcm.mcm_and_bbl_spin0(window, binning_file, lmax=lmax, - type="Dl", niter=0, binned_mcm = False) + xi = so_mcm.coupling_block("++", window2, lmax, niter=0, win2=window2, l_exact=l_exact, l_band=l_band, l_toep=l_toep) + np.testing.assert_almost_equal( default_coup_02[3,:,:], xi[2:-1,2:-1]) - couplings = so_cov.generate_aniso_couplings(survey_name, win, var, lmax) - - id2spec = {'Ta': 'dr6&pa6_f150_s1', 'Tb': 'dr6&pa6_f150_s4', - 'Tc': 'dr6&pa6_f150_s1', 'Td': 'dr6&pa6_f150_s4'} - - white_noise = { - 'Ta': so_cov.measure_white_noise_level(var0.data, window.data), - 'Tb': so_cov.measure_white_noise_level(var1.data, window.data), - 'Tc': so_cov.measure_white_noise_level(var0.data, window.data), - 'Td': so_cov.measure_white_noise_level(var1.data, window.data) - } - - # compute Rl, could put this in a function maybe - Clth_dict = {} - Rl_dict = {} - for name1, id1 in zip(survey_name, survey_id): - Rl_dict[id1] = np.sqrt(nl_dict[id1[0] + id1[0] + id2spec[id1] + id2spec[id1]][:lmax] - / white_noise[id1])[2:] - for name2, id2 in zip(survey_name, survey_id): - spec = id1[0] + id2[0] - key12 = spec + id2spec[id1] + id2spec[id2] - if name1 == name2: - Clth_dict[id1 + id2] = cl_dict[key12][:lmax] - else: - Clth_dict[id1 + id2] = cl_dict[key12][:lmax] - Clth_dict[id1 + id2] = Clth_dict[id1 + id2][2:] - - cov_e = so_cov.cov_spin0_aniso1( - Clth_dict, Rl_dict, couplings, - binning_file, lmax, mbb_inv, mbb_inv, binned_mcm=False) - - num_diag = 30 # check first 30 off-diagonals - diag_max_errors = [ - np.max(np.abs(np.diag(self.cov_ref[2:,2:] / cov_e[0], k) - 1)) - for k in range(0,num_diag) - ] - np.testing.assert_almost_equal(np.zeros(num_diag), - diag_max_errors, 7, - err_msg="aniso covmats don't match") + xi = so_mcm.coupling_block("--", window2, lmax, niter=0, win2=window2, l_exact=l_exact, l_band=l_band, l_toep=l_toep) + np.testing.assert_almost_equal( default_coup_02[4,:,:], xi[2:-1,2:-1]) if __name__ == "__main__":