pass
# Read in all our scans
data = scanutils.read_scans_autobalance(ids,
actscan.ACTScan,
comm,
filedb.data,
dets=args.dets,
downsample=down,
sky_local=dist,
osys=msys)
if data.n == 0:
L.info("Giving up")
sys.exit(1)
read_ids = [ids[ind] for ind in utils.allgatherv(data.inds, comm)]
read_ndets = utils.allgatherv([len(scan.dets) for scan in data.scans], comm)
read_nsamp = utils.allgatherv(
[scan.cut.size - scan.cut.sum() for scan in data.scans], comm)
read_dets = utils.uncat(
utils.allgatherv(
np.concatenate([scan.dets for scan in data.scans])
if len(data.scans) > 0 else np.zeros(0, int), comm), read_ndets)
if comm.rank == 0:
# Save accept list
with open(root + "accept.txt", "w") as f:
for id, dets in zip(read_ids, read_dets):
f.write("%s %3d: " % (id, len(dets)) +
" ".join([str(d) for d in dets]) + "\n")
# Save autocuts
if data.autocuts is not None:
def mcrdn0(icov,
get_kmap,
power,
nsims,
qfunc1,
qfunc2=None,
Xdat=None,
use_mpi=True,
verbose=True,
skip_rd=False):
"""
Using Monte Carlo sims, this function calculates
the anisotropic Gaussian N0 bias
between two quadratic estimators. It returns both the
data realization-dependent version (RDN0) and the pure-simulation
version (MCN0).
e.g.
>> mcrdn0(0,qfunc,get_kmap,comm,power)
Parameters
----------
icov: int
The index of the realization passed to get_kmap if performing
a covariance calculation - otherwise, set to zero.
get_kmap: function
Function for getting the filtered a_lms of data and simulation
maps. See notes at top of module.
power: function
Returns C(l) from two maps x,y, as power(x,y).
nsims: int
Number of sims
qfunc1: function
Function for reconstructing lensing from maps x and y,
called as e.g. qfunc(x, y). See e.g. SoLensPipe.qfunc.
The x and y arguments accept a [T_alm,E_alm,B_alm] tuple.
The function should return an (N,...) array where N is typically
two components for the gradient and curl.
qfunc2: function, optional
Same as above, for the third and fourth legs of the 4-point
RDN0.
comm: object, optional
MPI communicator
verbose: bool, optional
Whether to show progress
skip_rd: bool, optional
Whether to skip the RDN0 terms. The first returned component
is then None.
Returns
-------
rdn0: (N*(N+1)/2,...) array
Estimate of the RDN0 bias. If N=2 for gradient and curl,
the three components correspond to the gradient RDN0, the
curl RDN0 and the gradient x curl RDN0. None is returned
if skip_rd is True.
mcn0: (N*(N+1)/2,...) array
Estimate of the MCN0 bias. If N=2 for gradient and curl,
the three components correspond to the gradient RDN0, the
curl RDN0 and the gradient x curl RDN0.
"""
qa = qfunc1
qb = qfunc2
mcn0evals = []
if not (skip_rd):
assert Xdat is not None # Data
rdn0evals = []
if use_mpi:
comm, rank, my_tasks = mpi.distribute(nsims)
else:
comm, rank, my_tasks = FakeCommunicator(), 0, range(nsims)
for i in my_tasks:
if rank == 0 and verbose:
print("MCRDN0: Rank %d doing task %d" % (rank, i))
Xs = get_kmap((icov, 0, i))
if not (skip_rd):
qaXXs = qa(Xdat, Xs)
qbXXs = qb(Xdat, Xs) if qb is not None else qaXXs
qaXsX = qa(Xs, Xdat)
qbXsX = qb(Xs, Xdat) if qb is not None else qaXsX
rdn0_only_term = power(qaXXs,qbXXs) + power(qaXsX,qbXXs) \
+ power(qaXsX,qbXsX) + power(qaXXs,qbXsX)
Xsp = get_kmap((icov, 1, i))
qaXsXsp = qa(Xs, Xsp)
qbXsXsp = qb(Xs, Xsp) if qb is not None else qaXsXsp
qbXspXs = qb(Xsp, Xs) if qb is not None else qa(Xsp, Xs)
mcn0_term = (power(qaXsXsp, qbXsXsp) + power(qaXsXsp, qbXspXs))
mcn0evals.append(mcn0_term.copy())
if not (skip_rd): rdn0evals.append(rdn0_only_term - mcn0_term)
if not (skip_rd):
avgrdn0 = utils.allgatherv(rdn0evals, comm)
else:
avgrdn0 = None
avgmcn0 = utils.allgatherv(mcn0evals, comm)
return avgrdn0, avgmcn0
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 mcn1(icov,
get_kmap,
power,
nsims,
qfunc1,
qfunc2=None,
comm=None,
verbose=True):
"""
MCN1 for alpha=XY cross beta=AB
qfunc(x,y) returns QE reconstruction minus mean-field in fourier space
Parameters
----------
icov: int
The index of the realization passed to get_kmap if performing
a covariance calculation - otherwise, set to zero.
get_kmap: function
Function for getting the filtered a_lms of data and simulation
maps. See notes at top of module.
power: function
Returns C(l) from two maps x,y, as power(x,y).
nsims: int
Number of sims
qfunc1: function
Function for reconstructing lensing from maps x and y,
called as e.g. qfunc(x, y). See e.g. SoLensPipe.qfunc.
The x and y arguments accept a [T_alm,E_alm,B_alm] tuple.
The function should return an (N,...) array where N is typically
two components for the gradient and curl.
qfunc2: function, optional
Same as above, for the third and fourth legs of the 4-point
MCN1.
comm: object, optional
MPI communicator
verbose: bool, optional
Whether to show progress
Returns
-------
mcn1: (N*(N+1)/2,...) array
Estimate of the MCN1 bias. If N=2 for gradient and curl,
the three components correspond to the gradient MCN1, the
curl MCN1 and the gradient x curl MCN1.
Estimate of the MCN1 bias
"""
qa = qfunc1
qb = qfunc2
comm, rank, my_tasks = mpi.distribute(nsims)
n1evals = []
for i in my_tasks:
if rank == 0 and verbose:
print("MCN1: Rank %d doing task %d" % (comm.rank, i))
Xs = get_kmap((icov, 0, i)) # S
Ysp = get_kmap((icov, 1, i)) # S'
Xsk = get_kmap((icov, 2, i)) # Sphi
Yskp = get_kmap((icov, 3, i)) # Sphi'
qa_Xsk_Yskp = qa(Xsk, Yskp)
qb_Xsk_Yskp = qb(Xsk, Yskp) if qb is not None else qa_Xsk_Yskp
qb_Yskp_Xsk = qb(Yskp, Xsk) if qb is not None else qa(Yskp, Xsk)
qa_Xs_Ysp = qa(Xs, Ysp)
qb_Xs_Ysp = qb(Xs, Ysp) if qb is not None else qa_Xs_Ysp
qb_Ysp_Xs = qb(Ysp, Xs) if qb is not None else qa(Ysp, Xs)
qb_Yskp_Xsk = qb(Yskp, Xsk) if qb is not None else qa(Yskp, Xsk)
term = power(qa_Xsk_Yskp,qb_Xsk_Yskp) + power(qa_Xsk_Yskp,qb_Yskp_Xsk) \
- power(qa_Xs_Ysp,qb_Xs_Ysp) - power(qa_Xs_Ysp,qb_Ysp_Xs)
n1evals.append(term.copy())
n1s = utils.allgatherv(n1evals, comm)
return n1s
uxcls_nobh_2.append(uxcl_nobh_2)
uacls_nobh.append(uacl_nobh)
if bh:
r_bh_1 = plensing.phi_to_kappa(q_bh_1(Xdat, Xdat))
r_bh_2 = plensing.phi_to_kappa(q_bh_2(
Xdat, Xdat)) if not (diag) else r_bh_1.copy()
uxcl_bh_1 = cs.alm2cl(r_bh_1[0], ikalm)
uxcl_bh_2 = cs.alm2cl(r_bh_2[0], ikalm)
uacl_bh = cs.alm2cl(r_bh_1[0], r_bh_2[0])
uxcls_bh_1.append(uxcl_bh_1)
uxcls_bh_2.append(uxcl_bh_2)
uacls_bh.append(uacl_bh)
with bench.show("MPI Gather"):
uicls = putils.allgatherv(uicls, comm)
uxcls_nobh_1 = putils.allgatherv(uxcls_nobh_1, comm)
uxcls_nobh_2 = putils.allgatherv(uxcls_nobh_2, comm)
uacls_nobh = putils.allgatherv(uacls_nobh, comm)
if bh:
uxcls_bh_1 = putils.allgatherv(uxcls_bh_1, comm)
uxcls_bh_2 = putils.allgatherv(uxcls_bh_2, comm)
uacls_bh = putils.allgatherv(uacls_bh, comm)
if rank == 0:
with bench.show("Labels"):
labs = solenspipe.get_labels()
with bench.show("Stats"):
suicls = stats.get_stats(uicls)