def resolution(shape,wcs):
res = np.min(np.abs(enmap.extent(shape,wcs))/shape[-2:])
return res
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