Пример #1
0
def resolution(shape,wcs):
    res = np.min(np.abs(enmap.extent(shape,wcs))/shape[-2:])
    return res
Пример #2
0
def make_geometry(shape=None,wcs=None,hole_radius=None,cmb2d_TEB=None,n2d_IQU=None,context_width=None,n=None,beam2d=None,deproject=True,iau=False,res=None,tot_pow2d=None,store_pcov=False,pcov=None):

    """
    Make covariances for brute force maxlike inpainting of CMB maps.
    Eq 3 of arXiv:1109.0286

    shape,wcs -- enmap geometry of big map
    cmb2d_TEB -- (ncomp,ncomp,Ny,Nx) 2D CMB power in physical units
    n2d_IQU -- (ncomp,ncomp,Ny,Nx) 2D noise power in physical units
    hole_radius in radians
    context_width in radians or n as number of pixels
    beam2d -- (Ny,Nx) 2D beam template
    tot_pow2d -- (ncomp,ncomp,Ny,Nx) 2D total IQU power in physical units (includes CMB, beam and noise). If this argument is provided
                 cmb2d_TEB, n2d_IQU and beam2d will be ignored.
    deproject -- whether to deproject common mode
    iau -- whether to use IAU convention for polarization
    res -- specify resolution in radians instead of inferring from enmap geometry


    """

    # Make a flat stamp geometry
    if res is None: res = np.min(np.abs(enmap.extent(shape,wcs))/shape[-2:])
    if n is None: n = int(context_width/res)

    # Get the pix-pix covariance on the stamp geometry given CMB theory, beam and 2D noise on the big map
    if pcov is None:
        if tot_pow2d is not None:
                pcov = fcov_to_rcorr(shape,wcs,tot_pow2d,n)
        else:
                pcov = stamp_pixcov_from_theory(n,cmb2d_TEB,n2d,beam2d=beam2d,iau=iau)


    # Do we have polarization?
    ncomp = pcov.shape[0]
    assert ncomp==1 or ncomp==2 or ncomp==3

    # Select the hole (m1) and context(m2) across all components
    m1,m2 = get_geometry_regions(ncomp,n,res,hole_radius)

    # --- Make sure that the pcov is in the right order vector(I,Q,U) ---
    # It is currently in (ncomp,ncomp,n,n) order
    # We transpose it to (ncomp,n,ncomp,n) order
    # so that when it is reshaped into a 2D array, a row/column will correspond to an (I,Q,U) vector
    pcov = np.transpose(pcov,(0,2,1,3))
    pcov = pcov.reshape((ncomp*n**2,ncomp*n**2))

    # Invert
    Cinv = np.linalg.inv(pcov)
    
    # Woodbury deproject common mode
    if deproject:
        # Deproject I,Q,U common mode separately
        #u = (np.zeros((n*n,ncomp,ncomp))+np.eye(ncomp)).reshape(n*n*ncomp,ncomp)
        u = np.zeros((n*n*ncomp,ncomp))
        for i in range(ncomp):
            u[i*n*n:(i+1)*n*n,i] = 1
        #u = np.ones((ncomp*n**2,1)) # Deproject mode common to all of I,Q,U
        Cinvu = np.linalg.solve(pcov,u)
        precalc = np.dot(Cinvu,np.linalg.solve(np.dot(u.T,Cinvu),u.T))
        correction = np.dot(precalc,Cinv)
        Cinv -= correction
    
    # Get matrices for maxlike solution Eq 3 of arXiv:1109.0286
    cslice = Cinv[m1][:,m1]
    mul2 = Cinv[m1][:,m2]
    mean_mul = -np.linalg.solve(cslice,mul2)
    cov = np.linalg.inv(Cinv[m1][:,m1])
    cov_root = utils.eigpow(cov,0.5)

    geometry = {}
    geometry['covsqrt'] = cov_root
    geometry['meanmul'] = mean_mul
    geometry['n'] = n
    geometry['res'] = res
    geometry['m1'] = m1
    geometry['m2'] = m2
    geometry['ncomp'] = ncomp
    geometry['hole_radius'] = hole_radius
    if store_pcov: geometry['pcov'] = pcov

    return geometry