def solve(w, m):
if w.ndim < 4: return m / w
elif w.ndim == 4:
# This is slower, but handles low-hit areas near the edge better
iw = utils.eigpow(w, -1, axes=[0, 1])
return enmap.samewcs(array_ops.matmul(iw, m, axes=[0, 1]), m)
#return array_ops.solve_masked(w,m,axes=[0,1])
else:
raise NotImplementedError("Only 2d, 3d or 4d weight maps understood")
dscov /= np.nan_to_num(
maps.gauss_beam(modlmap, ifwhm) * maps.gauss_beam(modlmap, jfwhm))
# dncov[sel] = 1e90 #np.inf # inf gives invertible matrix but nonsensical output, 1e90 gives noninvertible, but with eigpow sensible
# dscov[sel] = 1e90 #np.inf
# io.plot_img((np.fft.fftshift(ncov)),"tuncov%d%d.png"%(aindex1,aindex2),aspect='auto')
# io.plot_img(np.log10(np.fft.fftshift(scov+ncov)),"udsncov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
# io.plot_img(np.log10(np.fft.fftshift(ncov)),"uncov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
# io.plot_img(np.log10(np.fft.fftshift(scov)),"usncov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
# io.plot_img(np.log10(np.fft.fftshift(dscov+dncov)),"dsncov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
# io.plot_img(np.log10(np.fft.fftshift(dncov)),"dncov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
# io.plot_img(np.log10(np.fft.fftshift(dscov)),"dsncov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
# io.plot_img(np.log10(np.fft.fftshift(ncov/dncov)),"rcov%d%d.png"%(aindex1,aindex2),aspect='auto',lim=[-5,1])
Scov[aindex1, aindex2] = dscov[modlmap < lmax].reshape(-1)
Ncov[aindex1, aindex2] = dncov[modlmap < lmax].reshape(-1)
if aindex1 != aindex2:
Scov[aindex2, aindex1] = Scov[aindex1, aindex2].copy()
if aindex1 != aindex2:
Ncov[aindex2, aindex1] = Ncov[aindex1, aindex2].copy()
Cov = Scov + Ncov
np.save("cov.npy", Cov)
#Cov = np.load("cov.npy")
Cov = np.rollaxis(Cov, 2)
icinv = utils.eigpow(Cov, -1)
np.save("icinv.npy", icinv)
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
if bcov[i, j] == 0.: bcov[i, j] = 700000
print(np.diagonal(bcov))
print(bcov)
np.savetxt(f"bcov_{seed}.txt", bcov, delimiter=',')
np.savetxt(f"gcov_{seed}.txt", gcov, delimiter=',')
gcorr = stats.cov2corr(gcov)
bcorr = stats.cov2corr(bcov)
print(gcorr.min(), gcorr.max())
print(bcorr.min(), bcorr.max())
io.plot_img(gcorr, f"det_gcov_{seed}.png", flip=False, lim=[0.5, 1])
io.plot_img(bcorr, f"det_bcov_{seed}.png", flip=False, lim=[0.5, 1])
print(seed)
gi = np.linalg.inv(gcov)
bi = np.linalg.inv(bcov)
bi2 = utils.eigpow(bcov, -1)
print(np.linalg.eigh(gcov)[0])
print(np.linalg.eigh(bcov)[0])
pl.done("detscatter.png")
pl = io.Plotter(xyscale='linlin', xlabel='a', ylabel='r')
c = 0
for i in range(narrays):
for j in range(i, narrays):
v1 = vals[12][c]
v2 = vals[13][c]
c = c + 1
qid1 = qids[i]
qid2 = qids[j]
def inpaint_uncorrelated_save_geometries(coords,hole_radius,ivar,output_dir,
theory_fn=None,beam_fn=None,include_signal=True,
pol=True,context_fraction=2./3.,
deproject=True,verbose_every_nsrcs=1,comm=None):
"""
This MPI-parallelized function will pre-calculate quantities required to inpaint a map with circular holes.
The results will be saved to disk in output_dir. The MPI parallelization is done over the number of sources.
Arguments
---------
coords: (N,2) array-like, float
Array containing the (dec,ra) coordinates in radians for N sources to be inpainted.
hole_radius: float
Each source will be inpainted with a hole of this radius in radians.
ivar: ndmap
(Ny,Nx) ndmap containing II inverse variance per pixel.
QQ and UU will be assumed to be half of II, and QU to be zero.
output_dir: string
Directory to save results to. Specify this for the inpaint_from_saved_geometries function.
theory_fn: function
Function such that theory_fn(string,ells) returns the CMB+foreground power spectrum where string
can be TT,TE,EE,BB.
beam_fn: function
Function such that beam_fn(ells) returns the beam transfer function.
include_signal: bool, optional
Whether to assume the CMB+foregrounds are present in the map. Set to False when inpainting
split differences.
pol: bool, optional
Whether to jointly inpaint I/Q/U maps or only assume I maps.
context_fraction: float
The ratio of the width of the context region around the hole to the diameter of the hole.
Larger fractions result in more informed inpainting, but takes longer time and memory.
comm: MPI communicator object
"""
import h5py
from . import mpi,io
if not(coords.ndim==2): raise ValueError
if not(coords.shape[1]==2): raise ValueError
if not(ivar.ndim==2): raise ValueError
nsrcs = coords.shape[0]
if not(comm is None):
rank = comm.Get_rank()
numcores = comm.Get_size()
num_each,each_tasks = mpi.mpi_distribute(nsrcs,numcores)
my_tasks = each_tasks[rank]
else:
comm = mpi.MPI.COMM_WORLD
my_tasks = range(nsrcs)
if pol:
ncomp = 3
else:
ncomp = 1
rtot = hole_radius * (1 + context_fraction)
np.savetxt(f'{output_dir}/source_inpaint_attributes.txt',np.asarray([[ncomp,hole_radius,context_fraction],]),fmt='%d,%.15f,%.15f',header='ncomp,hole_radius,context_fraction')
pixboxes = enmap.neighborhood_pixboxes(ivar.shape[-2:], ivar.wcs, coords, rtot)
def skip_warning(task): print(f"Skipped source {task}.")
ocoords = []
oinds = []
for ind,task in enumerate(my_tasks):
pixbox = pixboxes[task]
ithumb = ivar.extract_pixbox(pixbox)
if ithumb is None:
skip_warning(task)
continue
if not(ithumb.ndim==2): raise ValueError
if np.any(np.asarray(ithumb.shape)<=1):
skip_warning(task)
continue
if np.all(ithumb<=0):
skip_warning(task)
continue
Ny,Nx = ithumb.shape
modlmap = ithumb.modlmap()
modrmap = ithumb.modrmap()
if include_signal:
scov = scov_from_theory(modlmap,theory_fn,beam_fn,iau=False,ncomp=ncomp)
else:
scov = 0
ncov = ncov_from_ivar(ithumb,ncomp=ncomp)
pcov = scov + ncov
m1,m2 = get_regions(ncomp,modrmap,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*Ny*Nx,ncomp*Ny*Nx))
# Invert
Cinv = np.linalg.inv(pcov)
# Woodbury deproject common mode
if deproject:
# Deproject I,Q,U common mode separately
u = np.zeros((Ny*Nx*ncomp,ncomp))
for i in range(ncomp):
u[i*Ny*Nx:(i+1)*Ny*Nx,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['shape'] = np.asarray((Ny,Nx))
geometry['m1'] = m1
geometry['m2'] = m2
with h5py.File(f'{output_dir}/source_inpaint_geometry_{task}.hdf','w') as f:
for key in geometry:
f.create_dataset(key,data=geometry[key])
ocoords.append(coords[task])
oinds.append(task)
if verbose_every_nsrcs:
if (ind+1)%verbose_every_nsrcs==0: print(f"Done with {ind+1} / {len(my_tasks)}...")
ocoords = utils.allgatherv(ocoords,comm)
oinds = utils.allgatherv(oinds,comm)
io.save_cols(f'{output_dir}/source_inpaint_coords.txt',ocoords.swapaxes(0,1),header='Dec,RA')
np.savetxt(f'{output_dir}/source_inpaint_task_indices.txt',oinds,fmt='%d')
def get_mv(nls):
ninv = utils.eigpow(nls, -1., axes=[0, 1])
ncoadd = 1. / ninv.sum(axis=(0, 1))
return ncoadd