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)