def apply_apod(div): if apod_params is None: return div weight = div.preflat[0] moo = enmap.downgrade(weight,50) maxval = np.max(enmap.downgrade(weight,50)) apod = np.minimum(1,weight/maxval/apod_params[0])**apod_params[1] return div*apod
def apply_apod(div): if apod_params is None: return div weight = div.preflat[0] moo = enmap.downgrade(weight, 50) maxval = np.max(enmap.downgrade(weight, 50)) apod = np.minimum(1, weight / maxval / apod_params[0])**apod_params[1] return div * apod
def highResPlot2d(array, outPath, down=None, verbose=True, overwrite=True, crange=None): if not (overwrite): if os.path.isfile(outPath): return try: from enlib import enmap, enplot except: traceback.print_exc() printC( "Could not produce plot " + outPath + ". High resolution plotting requires enlib, which couldn't be imported. Continuing without plotting.", color='fail') return if (down is not None) and (down != 1): downmap = enmap.downgrade(enmap.enmap(array)[None], down) else: downmap = enmap.enmap(array)[None] img = enplot.draw_map_field(downmap, enplot.parse_args("-vvvg moo"), crange=crange) img.save(outPath) if verbose: print(bcolors.OKGREEN + "Saved high-res plot to", outPath + bcolors.ENDC)
def high_res_plot_img(array, filename=None, down=None, verbose=True, overwrite=True, crange=None, cmap="planck"): if not (overwrite): if os.path.isfile(filename): return try: from enlib import enmap, enplot except: traceback.print_exc() cprint( "Could not produce plot " + filename + ". High resolution plotting requires enlib, which couldn't be imported. Continuing without plotting.", color='fail') return if (down is not None) and (down != 1): downmap = enmap.downgrade(enmap.enmap(array)[None], down) else: downmap = enmap.enmap(array)[None] img = enplot.draw_map_field(downmap, enplot.parse_args("-c " + cmap + " -vvvg moo"), crange=crange) #img = enplot.draw_map_field(downmap,enplot.parse_args("--grid 1"),crange=crange) if filename is None: img.show() else: img.save(filename) if verbose: print(bcolors.OKGREEN + "Saved high-res plot to", filename + bcolors.ENDC)
def flatFitsToHealpix(fitsFile, nside, downgrade=1): from enlib import enmap imap = enmap.read_map(fitsFile) if downgrade > 1: imap = enmap.downgrade(imap, args.downgrade) omap = imap.to_healpix(nside=args.nside) return omap
def prepare(map, hitmap=False): """Prepare a map for input by cutting off one pixel along each edge, as out-of-bounds data accumulates there, and downgrading to the target resolution.""" # Get rid of polarization for now. Remove this later. if map.ndim == 3: map = map[:1] # Cut off edge pixels map[..., :1, :] = 0 map[..., -1:, :] = 0 map[..., :, :1] = 0 map[..., :, -1:] = 0 # Downsample map = enmap.downgrade(map, args.downgrade) if hitmap: map *= args.downgrade**2 # Pad to fourier-friendly size. Because no cropping # is used, this will result in the same padding for # all maps. map = map.autocrop(method="fft", value="none") return map
def prepare(map, hitmap=False): """Prepare a map for input by cutting off one pixel along each edge, as out-of-bounds data accumulates there, and downgrading to the target resolution.""" # Get rid of polarization for now. Remove this later. if map.ndim == 3: map = map[:1] # Cut off edge pixels map[...,:1,:] = 0 map[...,-1:,:] = 0 map[...,:,:1] = 0 map[...,:,-1:] = 0 # Downsample map = enmap.downgrade(map, args.downgrade) if hitmap: map *= args.downgrade**2 # Pad to fourier-friendly size. Because no cropping # is used, this will result in the same padding for # all maps. map = map.autocrop(method="fft", value="none") return map
parser.add_argument("-d", "--downgrade", type=int, default=1) parser.add_argument("-t", "--thin", type=int, default=1000) parser.add_argument( "-A", "--area-model", type=str, default="exact", help= "How to model pixel area. exact: Compute shape of each pixel. average: Use a single average number for all" ) parser.add_argument("--already-arcmin", action="store_true") args = parser.parse_args() div = enmap.read_fits(args.div) if args.downgrade: div = enmap.downgrade(div, args.downgrade) div *= args.downgrade**2 div = div.reshape((-1, ) + div.shape[-2:])[0] # Individual pixel area if args.area_model == "average": pix_area = div * 0 + div.area() / div.size * (180 * 60 / np.pi)**2 else: pos = div.posmap() diffs = utils.rewind(pos[:, 1:, 1:] - pos[:, :-1, :-1], 0) pix_area = np.abs(diffs[0] * diffs[1]) * np.cos(pos[0, :-1, :-1]) del diffs # Go to square arcmins pix_area /= utils.arcmin**2 # Pad to recover edge pixels pix_area = np.concatenate([pix_area, pix_area[-1:]], 0)
def analyze(self, ref_beam=None, mode="weight", map_max=1e8, div_tol=20, apod_val=0.2, apod_alpha=5, apod_edge=120, beam_tol=1e-4, ps_spec_tol=0.5, ps_smoothing=10, filter_kxrad=20, filter_highpass=200, filter_kx_ymax_scale=1): # Find the typical noise levels. We will use this to decide where # divs and beams etc. can be truncated to improve convergence. datasets = self.datasets ncomp = max([ split.data.map.preflat.shape[0] for dataset in datasets for split in dataset.splits ]) for dataset in datasets: for split in dataset.splits: split.ref_div = robust_ref(split.data.div) # Avoid single, crazy pixels split.data.div = np.minimum(split.data.div, split.ref_div * div_tol) split.data.div = filter_div(split.data.div) split.data.map = np.maximum( -map_max, np.minimum(map_max, split.data.map)) # Expand map to ncomp components split.data.map = add_missing_comps(split.data.map, ncomp) # Build apodization apod = np.minimum(split.data.div / (split.ref_div * apod_val), 1.0)**apod_alpha apod = apod.apod(apod_edge) split.data.div *= apod split.data.H = split.data.div**0.5 dataset.ref_div = np.sum( [split.ref_div for split in dataset.splits]) tot_ref_div = np.sum([dataset.ref_div for dataset in datasets]) ly, lx = enmap.laxes(self.shape, self.wcs) lr = (ly[:, None]**2 + lx[None, :]**2)**0.5 bmin = np.min([beam_size(dataset.beam) for dataset in datasets]) # Deconvolve all the relative beams. These should ideally include pixel windows. # This could matter for planck if ref_beam is not None: for dataset in datasets: rel_beam = beam_ratio(dataset.beam, ref_beam) # Avoid division by zero bspec = np.maximum(eval_beam(rel_beam, lr), 1e-10) # We don't want to divide by tiny numbers, so we will cap the relative # beam. The goal is just to make sure that the deconvolved noise ends up # sufficiently high that anything beyond that is negligible. This will depend # on the div ratios between the different datasets. We can stop deconvolving # when beam*my_div << (tot_div-my_div). But deconvolving even by a factor # 1000 leads to strange numberical errors bspec = np.maximum( bspec, beam_tol * (tot_ref_div / dataset.ref_div - 1)) bspec_dec = np.maximum(bspec, 0.1) for split in dataset.splits: split.data.map = map_ifft( map_fft(split.data.map) / bspec_dec) # In theory we don't need to worry about the beam any more by this point. # But the pixel window might be unknown or missing. So we save the beam so # we can make sure the noise model makes sense dataset.bspec = bspec # We classify this as a low-resolution dataset if we did an appreciable amount of # deconvolution dataset.lowres = np.min(bspec) < 0.5 # Can now build the noise model and rhs for each dataset. # The noise model is N = HCH, where H = div**0.5 and C is the mean 2d noise spectrum # of the whitened map, after some smoothing. for dataset in datasets: nsplit = 0 dset_map, dset_div = None, None for split in dataset.splits: if dset_map is None: dset_map = split.data.map * 0 dset_div = split.data.div * 0 dset_map += split.data.map * split.data.div dset_div += split.data.div # Form the mean map for this dataset dset_map[:, dset_div > 0] /= dset_div[dset_div > 0] # Then use it to build the diff maps and noise spectra dset_ps = None #i=0 for split in dataset.splits: if split.data.empty: continue diff = split.data.map - dset_map wdiff = diff * split.data.H #i+=1 # What is the healthy area of wdiff? Wdiff should have variance # 1 or above. This tells us how to upweight the power spectrum # to take into account missing regions of the diff map. ndown = 10 wvar = enmap.downgrade(wdiff**2, ndown) goodfrac = np.sum(wvar > 1e-3) / float(wvar.size) if goodfrac < 0.1: goodfrac = 0 ps = np.abs(map_fft(wdiff))**2 # correct for unhit areas, which can't be whitened with utils.nowarn(): ps /= goodfrac if dset_ps is None: dset_ps = enmap.zeros(ps.shape, ps.wcs, ps.dtype) dset_ps += ps nsplit += 1 if nsplit < 2: continue # With n splits, mean map has var 1/n, so diff has var (1-1/n) + (n-1)/n = 2*(n-1)/n # Hence tot-ps has var 2*(n-1) dset_ps /= 2 * (nsplit - 1) dset_ps = smooth_pix(dset_ps, ps_smoothing) # Use the beam we saved from earlier to make sure we don't have a remaining # pixel window giving our high-l parts too high weight. If everything has # been correctly deconvolved, we expect high-l dset_ps to go as # 1/beam**2. The lower ls will realistically be no lower than this either. # So we can simply take the max dset_ps_ref = np.min( np.maximum(dset_ps, dataset.bspec**-2 * ps_spec_tol * 0.1)) dset_ps = np.maximum(dset_ps, dset_ps_ref * dataset.bspec**-2 * ps_spec_tol) # Our fourier-space inverse noise matrix is the inverse of this if np.all(np.isfinite(dset_ps)): iN = 1 / dset_ps else: iN = enmap.zeros(dset_ps.shape, dset_ps.wcs, dset_ps.dtype) # Add any fourier-space masks to this if dataset.highpass: kxmask = butter(lx, filter_kxrad, -5) kxmask = 1 - (1 - kxmask[None, :]) * ( np.abs(ly) < bmin * filter_kx_ymax_scale)[:, None] highpass = butter(lr, filter_highpass, -10) filter = highpass * kxmask del kxmask, highpass else: filter = 1 if mode != "filter": iN *= filter dataset.iN = iN dataset.filter = filter self.mode = mode
profiles[polcomb] = [] else: apowers[polcomb] = [] cpowers[polcomb] = [] for i in range(Nsims): print(i) unlensed = enmap.rand_map(shape_sim, wcs_sim, ps) lensed = lensing.lens_map_flat_pix(unlensed, alpha_pix, order=lens_order) klteb = enmap.map2harm(lensed) klteb_beam = klteb * kbeam_sim lteb_beam = enmap.ifft(klteb_beam).real noise = enmap.rand_map(shape_sim, wcs_sim, ps_noise, scalar=True) observed = lteb_beam + noise measured = enmap.downgrade(observed, analysis_pixel_scale / sim_pixel_scale) if i == 0: #debug() shape_dat, wcs_dat = measured.shape, measured.wcs lxmap_dat, lymap_dat, modlmap_dat, angmap_dat, lx_dat, ly_dat = fmaps.get_ft_attributes_enmap( shape_dat, wcs_dat) nT = ntfunc(modlmap_dat) nP = npfunc(modlmap_dat) kbeam_dat = cmb.gauss_beam(modlmap_dat, beam_arcmin) if cluster: modr_dat = enmap.modrmap(shape_dat, wcs_dat) * 180. * 60. / np.pi bin_edges_dat = np.arange(0., modr_dat.max(), 1.0) binner_dat = stats.bin2D(modr_dat, bin_edges_dat)
ftot += m ps_auto += m[:, None] * np.conj(m[None, :]) print("Computing cross spectrum") # Compute auto spectrum ps_auto /= nfile**2 ftot /= nfile # Compute total spectrum ps_tot = ftot[:, None] * np.conj(ftot[None, :]) if len(args.ifiles) > 1: # Subtract to get cross spectrum ps_cross = ps_tot - ps_auto else: ps_cross = ps_tot del ps_tot, ps_auto, ftot ps_cross = enmap.downgrade(ps_cross, args.pregrade) print(ps_cross.shape) print("Normalizing") l = np.sum(ps_cross.lmap()**2, 0)**0.5 l = np.minimum(l, args.lmax_scale) ps_cross *= ps_cross.area() / ps_cross.npix ps_cross *= l**2 / (2 * np.pi) print("Recentering") # Center l=0 def recenter(m, shape): return np.roll(np.roll(m, -shape[-2] // 2, -2), -shape[-1] // 2, -1)
for i in range(Nsims): print(i) unlensed = parray_sim.get_unlensed_cmb(seed=(200 + i)) lensed = lensing.lens_map_flat_pix(unlensed.copy(), alpha_pix.copy(), order=lens_order) #m, = lensing.rand_map(shape, wcs, ps, lmax=lmax, maplmax=maplmax, seed=(seed,i)) klteb = enmap.map2harm(lensed.copy()) klteb_beam = klteb * kbeam_sim lteb_beam = enmap.ifft(klteb_beam).real noise = 0. #parray_sim.get_noise_sim(seed=(300+i)) observed = lteb_beam + noise measured = enmap.downgrade(observed, pixratio) if i == 0: shape_dat, wcs_dat = measured.shape, measured.wcs modr_dat = parray_dat.modrmap * 180. * 60. / np.pi # === ESTIMATOR === template_dat = fmaps.simple_flipper_template_from_enmap( shape_dat, wcs_dat) lxmap_dat, lymap_dat, modlmap_dat, angmap_dat, lx_dat, ly_dat = fmaps.get_ft_attributes_enmap( shape_dat, wcs_dat) taper_percent = 15.0 Ny, Nx = shape_dat taper = fmaps.cosineWindow(Ny,
rhs = enmap.zeros((ncomp,)+shape, area.wcs, dtype) div = enmap.zeros((ncomp,ncomp)+shape, area.wcs, dtype) junk = np.zeros(pcut.njunk, dtype) with bench.show("rhs"): tod *= ivar[:,None] pcut.backward(tod, junk) pmap.backward(tod, rhs) with bench.show("hits"): for i in range(ncomp): div[i,i] = 1 pmap.forward(tod, div[i]) tod *= ivar[:,None] pcut.backward(tod, junk) div[i] = 0 pmap.backward(tod, div[i]) with bench.show("map"): idiv = array_ops.eigpow(div, -1, axes=[0,1], lim=1e-5) map = enmap.map_mul(idiv, rhs) # Estimate central amplitude c = np.array(map.shape[-2:])/2 crad = 50 mcent = map[:,c[0]-crad:c[0]+crad,c[1]-crad:c[1]+crad] mcent = enmap.downgrade(mcent, 4) amp = np.max(mcent) print("%s amp %7.3f asens %7.3f" % (id, amp/1e6, asens)) with bench.show("write"): enmap.write_map("%s%s_map.fits" % (prefix, bid), map) enmap.write_map("%s%s_rhs.fits" % (prefix, bid), rhs) enmap.write_map("%s%s_div.fits" % (prefix, bid), div) del d, scan, pmap, pcut, tod, map, rhs, div, idiv, junk
def get_map(ifile, args, return_info=False): """Read the specified map, and massage it according to the options in args. Relevant ones are sub, autocrop, slice, op, downgrade, scale, mask. Retuns with shape [:,ny,nx], where any extra dimensions have been flattened into a single one.""" with warnings.catch_warnings(): warnings.filterwarnings("ignore") toks = ifile.split(":") ifile, slice = toks[0], ":".join(toks[1:]) m0 = enmap.read_map(ifile) if args.fix_wcs: m0.wcs = enwcs.fix_wcs(m0.wcs) # Save the original map, so we can compare its wcs later m = m0 # Submap slicing currently has wrapping issues if args.sub is not None: default = [[-90,-180],[90,180]] sub = np.array([[(default[j][i] if q == '' else float(q))*np.pi/180 for j,q in enumerate(w.split(":"))]for i,w in enumerate(args.sub.split(","))]).T m = m.submap(sub) # Perform a common autocrop across all fields if args.autocrop: m = enmap.autocrop(m) # If necessary, split into stamps. If no stamp splitting occurs, # a list containing only the original map is returned mlist = extract_stamps(m, args) # The stamp stuff is a bit of an ugly hack. This loop and wcslist # are parts of that hack. for i, m in enumerate(mlist): # Downgrade downgrade = [int(w) for w in args.downgrade.split(",")] m = enmap.downgrade(m, downgrade) # Slicing, either at the file name level or though the slice option m = eval("m"+slice) if args.slice is not None: m = eval("m"+args.slice) flip = (m.wcs.wcs.cdelt*m0.wcs.wcs.cdelt)[::-1]<0 assert m.ndim >= 2, "Image must have at least 2 dimensions" # Apply arbitrary map operations m1 = m if args.op is not None: m = eval(args.op, {"m":m},np.__dict__) # Scale if requested scale = [int(w) for w in args.scale.split(",")] if np.any(np.array(scale)>1): m = enmap.upgrade(m, scale) # Flip such that pixels are in PIL or matplotlib convention, # such that RA increases towards the left and dec upwards in # the final image. Unless a slicing operation on the image # overrrode this. if m.wcs.wcs.cdelt[1] > 0: m = m[...,::-1,:] if m.wcs.wcs.cdelt[0] > 0: m = m[...,:,::-1] if flip[0]: m = m[...,::-1,:] if flip[1]: m = m[...,:,::-1] # Update stamp list mlist[i] = m wcslist = [m.wcs for m in mlist] m = enmap.samewcs(np.asarray(mlist),mlist[0]) if args.stamps is None: m, wcslist = m[0], None # Flatten pre-dimensions mf = m.reshape((-1,)+m.shape[-2:]) # Stack if args.tile is not None: toks = [int(i) for i in args.tile.split(",")] nrow = toks[0] if len(toks) > 0 else -1 ncol = toks[1] if len(toks) > 1 else -1 mf = hwstack(hwexpand(mf, nrow, ncol, args.tile_transpose))[None] # Mask bad data if args.mask is not None: if not np.isfinite(args.mask): mf[np.abs(mf)==args.mask] = np.nan else: mf[np.abs(mf-args.mask)<=args.mask_tol] = np.nan # Done if not return_info: return mf else: info = bunch.Bunch(fname=ifile, ishape=m.shape, wcslist=wcslist) return mf, info
imap = enmap.read_map(imapfiles[i]) if args.slice: imap = eval("imap" + args.slice) # We want y,x-ordering off = off[:, ::-1] box = utils.minmax(off, 0) dets.append(det) offs.append(off) boxes.append(box) imaps.append(imap) box = utils.bounding_box(boxes) box = utils.widen_box(box, rad * 5, relative=False) # We assume that the two maps have the same pixelization imaps = enmap.samewcs(np.array(imaps), imaps[0]) # Downsample by averaging imaps = enmap.downgrade(imaps, (1, args.step)) naz = imaps.shape[-1] # Ok, build our output geometry shape, wcs = enmap.geometry(pos=box, res=args.res * utils.arcmin, proj="car", pre=(naz, )) omap = enmap.zeros(shape, wcs, dtype=dtype) # Normalization norm = enmap.zeros(shape[-2:], wcs) norm[0, 0] = 1 norm = enmap.smooth_gauss(norm, rad)[0, 0] # Loop through slices and populate
import numpy as np, argparse from enlib import enmap, utils, bench parser = argparse.ArgumentParser() parser.add_argument("div") parser.add_argument("ofile", nargs="?", default="/dev/stdout") parser.add_argument("-d", "--downgrade", type=int, default=1) parser.add_argument("-t", "--thin", type=int, default=1000) parser.add_argument("-A", "--area-model", type=str, default="exact", help="How to model pixel area. exact: Compute shape of each pixel. average: Use a single average number for all") args = parser.parse_args() div = enmap.read_fits(args.div) if args.downgrade: div = enmap.downgrade(div, args.downgrade) div *= args.downgrade**2 div = div.reshape((-1,)+div.shape[-2:])[0] # Individual pixel area if args.area_model == "average": pix_area = div*0 + div.area()/div.size*(180*60/np.pi)**2 else: pos = div.posmap() diffs = utils.rewind(pos[:,1:,1:]-pos[:,:-1,:-1],0) pix_area = np.abs(diffs[0]*diffs[1])*np.cos(pos[0,:-1,:-1]) del diffs # Go to square arcmins pix_area /= utils.arcmin**2 # Pad to recover edge pixels pix_area = np.concatenate([pix_area,pix_area[-1:]],0) pix_area = np.concatenate([pix_area,pix_area[:,-1:]],1) # Flatten everything
rhs = enmap.zeros((ncomp,)+shape, area.wcs, dtype) div = enmap.zeros((ncomp,ncomp)+shape, area.wcs, dtype) junk = np.zeros(pcut.njunk, dtype) with bench.show("rhs"): tod *= ivar[:,None] pcut.backward(tod, junk) pmap.backward(tod, rhs) with bench.show("hits"): for i in range(ncomp): div[i,i] = 1 pmap.forward(tod, div[i]) tod *= ivar[:,None] pcut.backward(tod, junk) div[i] = 0 pmap.backward(tod, div[i]) with bench.show("map"): idiv = array_ops.eigpow(div, -1, axes=[0,1], lim=1e-5) map = enmap.map_mul(idiv, rhs) # Estimate central amplitude c = np.array(map.shape[-2:])/2 crad = 50 mcent = map[:,c[0]-crad:c[0]+crad,c[1]-crad:c[1]+crad] mcent = enmap.downgrade(mcent, 4) amp = np.max(mcent) print "%s amp %7.3f asens %7.3f" % (id, amp/1e6, asens) with bench.show("write"): enmap.write_map("%s%s_map.fits" % (prefix, bid), map) enmap.write_map("%s%s_rhs.fits" % (prefix, bid), rhs) enmap.write_map("%s%s_div.fits" % (prefix, bid), div) del d, scan, pmap, pcut, tod, map, rhs, div, idiv, junk
imap = enmap.read_map(imapfiles[i]) if args.slice: imap = eval("imap"+args.slice) # We want y,x-ordering off = off[:,::-1] box = utils.minmax(off,0) dets.append(det) offs.append(off) boxes.append(box) imaps.append(imap) box = utils.bounding_box(boxes) box = utils.widen_box(box, rad*5, relative=False) # We assume that the two maps have the same pixelization imaps = enmap.samewcs(np.array(imaps), imaps[0]) # Downsample by averaging imaps = enmap.downgrade(imaps, (1,args.step)) naz = imaps.shape[-1] # Ok, build our output geometry shape, wcs = enmap.geometry(pos=box, res=args.res*utils.arcmin, proj="car", pre=(naz,)) omap = enmap.zeros(shape, wcs, dtype=dtype) # Normalization norm = enmap.zeros(shape[-2:],wcs) norm[0,0] = 1 norm = enmap.smooth_gauss(norm, rad)[0,0] # Loop through slices and populate bazs = [] for iaz in range(naz): # Get our boresight az
import numpy as np, argparse, healpy from enlib import enmap #import enmap parser = argparse.ArgumentParser() parser.add_argument("ifile") parser.add_argument("ofile") parser.add_argument("-d", "--downgrade", type=int, default=1) parser.add_argument("-N", "--nside", type=int, default=0) args = parser.parse_args() imap = enmap.read_map(args.ifile) if args.downgrade > 1: imap = enmap.downgrade(imap, args.downgrade) omap = imap.to_healpix(nside=args.nside) healpy.write_map(args.ofile, omap)
## Generate scalar-only lensed and unlensed map #m_scal_u, m_scal_l = lensing.rand_map(shape, wcs, cl_scal, cl_phi, seed=args.seed, output="ul", verbose=args.verbose) ## Generate tensor-only lensed and unlensed map #m_tens_u, m_tens_l = lensing.rand_map(shape, wcs, cl_tens, cl_phi, seed=args.seed, output="ul", verbose=args.verbose) np.random.seed(args.seed) phi = enmap.rand_map(shape[-2:], wcs, cl_phi) m_scal_u = enmap.rand_map(shape, wcs, cl_scal) m_tens_u = enmap.rand_map(shape, wcs, cl_tens) m_scal_l = lensing.lens_map_flat(m_scal_u, phi) m_tens_l = lensing.lens_map_flat(m_tens_u, phi) # And the sums m_tot_u = m_scal_u + m_tens_u m_tot_l = m_scal_l + m_tens_l # Convert from TQU to TEB and downgrade def to_eb(m): return enmap.ifft(enmap.map2harm(m)).real m_scal_u, m_scal_l, m_tens_u, m_tens_l, m_tot_u, m_tot_l = [enmap.downgrade(to_eb(i),os) for i in [m_scal_u, m_scal_l, m_tens_u, m_tens_l, m_tot_u, m_tot_l]] # And output utils.mkdir(args.odir) enmap.write_map(args.odir + "/map_scalar.fits", m_scal_u) enmap.write_map(args.odir + "/map_tensor.fits", m_tens_u) enmap.write_map(args.odir + "/map_tot.fits", m_tot_u) enmap.write_map(args.odir + "/map_scalar_lensed.fits", m_scal_l) enmap.write_map(args.odir + "/map_tensor_lensed.fits", m_tens_l) enmap.write_map(args.odir + "/map_tot_lensed.fits", m_tot_l)
def coadd_tile_data(datasets, box, odir, ps_smoothing=10, pad=0, ref_beam=None, cg_tol=1e-6, dump=False, verbose=False, read_cache=False, write_cache=False, div_max_tol=100, div_div_tol=1e-10): # Load data for this box for each dataset datasets, ffpad = read_data(datasets, box, odir, pad=pad, verbose=verbose, read_cache=read_cache, write_cache=write_cache) # We might not find any data if len(datasets) == 0: return None # Find the smallest beam size of the datasets bmin = np.min([beam_size(dataset.beam) for dataset in datasets]) # Subtract mean map from each split to get noise maps. Our noise # model is HNH, where H is div**0.5 and N is the mean 2d noise spectrum # after some smoothing rhs, tot_div = None, None tot_iN, tot_udiv = None, 0 for dataset in datasets: nsplit = 0 dset_map, dset_div = None, None for split in dataset.splits: if dset_map is None: dset_map = split.data.map * 0 dset_div = split.data.div * 0 dset_map += split.data.map * split.data.div dset_div += split.data.div # Form the mean map for this dataset dset_map[:, dset_div > 0] /= dset_div[dset_div > 0] if tot_div is None: tot_div = dset_div * 0 tot_div += dset_div tshape, twcs, tdtype = dset_map.shape, dset_div.wcs, dset_div.dtype # Then use it to build the diff maps and noise spectra dset_ps = None for split in dataset.splits: if split.data.empty: continue diff = split.data.map - dset_map wdiff = diff * split.data.H # What is the healthy area of wdiff? Wdiff should have variance # 1 or above. This tells us how to upweight the power spectrum # to take into account missing regions of the diff map. ndown = 10 wvar = enmap.downgrade(wdiff**2, ndown) goodfrac = np.sum(wvar > 1e-3) / float(wvar.size) if goodfrac < 0.1: goodfrac = 0 #opre = odir + "/" + os.path.basename(split.map)[:-5] #enmap.write_map(opre + "_diff.fits", diff) #enmap.write_map(opre + "_wdiff.fits", wdiff) #enmap.write_map(opre + "_wvar.fits", wvar) ps = np.abs(map_fft(wdiff))**2 #enmap.write_map(opre + "_ps1.fits", ps) # correct for unhit areas, which can't be whitened #print "A", dataset.name, np.median(ps[ps>0]), medloop(ps), goodfrac with utils.nowarn(): ps /= goodfrac #print "B", dataset.name, np.median(ps[ps>0]), medloop(ps), goodfrac #enmap.write_map(opre + "_ps2.fits", ps) #enmap.write_map(opre + "_ps2d.fits", ps) if dset_ps is None: dset_ps = enmap.zeros(ps.shape, ps.wcs, ps.dtype) dset_ps += ps nsplit += 1 if nsplit < 2: continue # With n splits, mean map has var 1/n, so diff has var (1-1/n) + (n-1)/n = 2*(n-1)/n # Hence tot-ps has var 2*(n-1) dset_ps /= 2 * (nsplit - 1) #enmap.write_map(opre + "_ps2d_tot.fits", dset_ps) dset_ps = smooth_pix(dset_ps, ps_smoothing) #enmap.write_map(opre + "_ps2d_smooth.fits", dset_ps) if np.all(np.isfinite(dset_ps)): # Super-low values of the spectrum are not realistic. These appear # due to beam/pixel smoothing in the planck maps. This will be # mostly taken care of when processing the beams, as long as we don't # let them get too small dset_ps = np.maximum(dset_ps, 1e-7) # Optionally cap the max dset_ps, this is mostly to speed up convergence if args.max_ps: dset_ps = np.minimum(dset_ps, args.max_ps) # Our fourier-space inverse noise matrix is based on the inverse noise spectrum iN = 1 / dset_ps #enmap.write_map(opre + "_iN_raw.fits", iN) else: print "Setting weight of dataset %s to zero" % dataset.name #print np.all(np.isfinite(dset_ps)), np.all(dset_ps>0) iN = enmap.zeros(dset_ps.shape, dset_ps.wcs, dset_ps.dtype) # Add any fourier-space masks to this ly, lx = enmap.laxes(tshape, twcs) lr = (ly[:, None]**2 + lx[None, :]**2)**0.5 if dataset.highpass: kxmask = butter(lx, args.kxrad, -3) kxmask = 1 - (1 - kxmask[None, :]) * ( np.abs(ly) < bmin * args.kx_ymax_scale)[:, None] highpass = butter(lr, args.highpass, -10) filter = highpass * kxmask #print "filter weighting", dataset.name del kxmask, highpass else: filter = 1 if not args.filter: iN *= filter # We should deconvolve the relative beam from the maps, # but that's numerically nasty. But it can be handled # inversely. We want (BiNB + ...)x = (BiNB iB m + ...) # where iB is the beam deconvolution operation in map space. # Instead of actually doing that operation, we can compute two # inverse noise matrixes: iN_A = BiNB for the left hand # side and iN_b = BiN for the right hand side. That way we # avoid dividing by any huge numbers. # Add the relative beam iN_A = iN.copy() iN_b = iN.copy() if ref_beam is not None: rel_beam = beam_ratio(dataset.beam, ref_beam) bspec = eval_beam(rel_beam, lr) iN_A *= bspec**2 iN_b *= bspec #moo = iN*0+filter #enmap.write_map(opre + "_filter.fits", moo) # Add filter to noise model if we're downweighting # rather than filtering. dataset.iN_A = iN_A dataset.iN_b = iN_b dataset.filter = filter #print "A", opre #enmap.write_map(opre + "_iN_A.fits", iN_A) #enmap.write_map(opre + "_iN.fits", iN) # Cap to avoid single crazy pixels tot_div = np.maximum(tot_div, np.median(tot_div[tot_div > 0]) * 0.01) tot_idiv = tot_div * 0 tot_idiv[tot_div > div_div_tol] = 1 / tot_div[tot_div > div_div_tol] # Build the right-hand side. The right-hand side is # sum(HNHm) if rhs is None: rhs = enmap.zeros(tshape, twcs, tdtype) for dataset in datasets: i = 0 for split in dataset.splits: if split.data.empty: continue #print "MOO", dataset.name, np.max(split.data.map), np.min(split.data.map), np.max(split.data.div), np.min(split.data.div) w = split.data.H * split.data.map fw = map_fft(w) fw *= dataset.iN_b if args.filter: fw *= dataset.filter w = map_ifft(fw) * split.data.H #enmap.write_map(odir + "/%s_%02d_rhs.fits" % (dataset.name, i), w) rhs += w i += 1 del w, iN, iN_A, iN_b, filter # Now solve the equation def A(x): global times m = enmap.samewcs(x.reshape(rhs.shape), rhs) res = m * 0 times[:] = 0 ntime = 0 for dataset in datasets: for split in dataset.splits: if split.data.empty: continue t = [time.time()] w = split.data.H * m t.append(time.time()) fw = map_fft(w) t.append(time.time()) fw *= dataset.iN_A t.append(time.time()) w = map_ifft(fw) t.append(time.time()) w *= split.data.H t.append(time.time()) res += w for i in range(1, len(t)): times[i - 1] += t[i] - t[i - 1] ntime += 1 #w = enmap.harm2map(dataset.iN_A*enmap.map2harm(w)) #w *= split.data.H #res += w del w times /= ntime return res.reshape(-1) def M(x): m = enmap.samewcs(x.reshape(rhs.shape), rhs) res = m * tot_idiv return res.reshape(-1) solver = cg.CG(A, rhs.reshape(-1), M=M) for i in range(1000): t1 = time.time() solver.step() t2 = time.time() if verbose: print "%5d %15.7e %5.2f: %4.2f %4.2f %4.2f %4.2f %4.2f" % ( solver.i, solver.err, t2 - t1, times[0], times[1], times[2], times[3], times[4]), np.std(solver.x) if dump and solver.i in [1, 2, 5, 10, 20, 50] + range(100, 10000, 100): m = enmap.samewcs(solver.x.reshape(rhs.shape), rhs) enmap.write_map(odir + "/step%04d.fits" % solver.i, m) if solver.err < cg_tol: if dump: m = enmap.samewcs(solver.x.reshape(rhs.shape), rhs) enmap.write_map(odir + "/step_final.fits", m) break tot_map = enmap.samewcs(solver.x.reshape(rhs.shape), rhs) # Get rid of the fourier padding ny, nx = tot_map.shape[-2:] tot_map = tot_map[..., :ny - ffpad[0], :nx - ffpad[1]] tot_div = tot_div[..., :ny - ffpad[0], :nx - ffpad[1]] return bunch.Bunch(map=tot_map, div=tot_div)
# pl.add(ells,theory.uCl('TT',ells)*ells**2.*TCMB**2.) # pl.done(outDir+"cls.png") # sys.exit() for i in range(N): map = enmap.rand_map(shape, wcs, ps) / TCMB # massIndex = massIndices[i] # inputKappaMap, szMap,M500,z = getKappaSZ(b,snap,massIndex,px,thetaMap.shape) # avgM500 += M500 # avgz += z inputKappaMap = kappaMap szMap = kappaMap.copy() * 0. if int(pxDown / px) > 1: inpDown = enmap.downgrade(inputKappaMap, pxDown / px) szDown = enmap.downgrade(szMap, pxDown / px) else: inpDown = inputKappaMap szDown = szMap trueKappaStack += inpDown szStack += szDown # === DEFLECTION MAP === a = alphaMaker(thetaMap) alpha = a.kappaToAlpha(inputKappaMap, test=False) alphamod = 180. * 60. * np.sum(alpha**2, 0)**0.5 / np.pi # print "alphaint ", alphamod[thetaMap*60.*180./np.pi<10.].mean() pos = thetaMap.posmap() + alpha pix = thetaMap.sky2pix(pos, safe=False)
def read_map(fname): m = enmap.read_map(fname) m = m[..., 1:-1, 1:-1] #m = eval("m" + args.slice) m = enmap.downgrade(m, args.downgrade) return m
], nowrap=True) tshape, twcs = enmap.downgrade_geometry(tshape, twcs, args.downgrade) #print ty, tx, np.mean(enmap.box(tshape, twcs),0)/utils.degree # Read in our tile data frhss, kmaps, mjds = [], [], [] for ind in range(comm_intra.rank, len(idirs), comm_intra.size): if not overlaps(pboxes[ind], pixbox, nphi): #print idirs[ind], "does not overlap", pixbox.reshape(-1), pboxes[ind].reshape(-1) continue idir = idirs[ind] lshape, lwcs = enmap.read_map_geometry(idir + "/frhs.fits") pixbox_loc = pixbox - enmap.pixbox_of(wcs, lshape, lwcs)[0] kmap = enmap.read_map(idir + "/kmap.fits", pixbox=pixbox_loc).astype(dtype) if args.downgrade > 1: kmap = enmap.downgrade(kmap, args.downgrade) # Skip tile if it's empty if np.any(~np.isfinite(kmap)) or np.all(kmap < 1e-10): if args.verbose: print "%3d skipping %s (nan or unexposed)" % (comm.rank, idir) continue else: if args.verbose: print "%3d read %s" % (comm.rank, idir) frhs = enmap.read_map(idir + "/frhs.fits", pixbox=pixbox_loc).astype(dtype) if args.downgrade > 1: frhs = enmap.downgrade(frhs, args.downgrade) with h5py.File(idir + "/info.hdf", "r") as hfile: mjd = hfile["mjd"].value frhss.append(frhs) kmaps.append(kmap)
def combine_tiles(ipathfmt, opathfmt, combine=2, downsample=2, itile1=(None, None), itile2=(None, None), tyflip=False, txflip=False, pad_to=None, comm=None, verbose=False): """Given a set of tiles on disk at locaiton ipathfmt % {"y":...,"x"...}, combine them into larger tiles, downsample and write the result to opathfmt % {"y":...,"x":...}. x and y must be contiguous and start at 0. reftile[2] indicates the tile coordinates of the first valid input tile. This needs to be specified if not all tiles of the logical tiling are physically present. tyflip and txflip indicate if the tiles coordinate system is reversed relative to the pixel coordinates or not." """ # Expand combine and downsample to 2d combine = np.zeros(2, int) + combine downsample = np.zeros(2, int) + downsample if pad_to is not None: pad_to = np.zeros(2, int) + pad_to # Handle optional mpi rank, size = (comm.rank, comm.size) if comm is not None else (0, 1) # Find the range of input tiles itile1, itile2 = find_tile_range(ipathfmt, itile1, itile2) # Read the first tile to get its size information ibase = enmap.read_map(ipathfmt % {"y": itile1[0], "x": itile1[1]}) * 0 # Find the set of output tiles we need to consider otile1 = itile1 / combine otile2 = (itile2 - 1) / combine + 1 # And loop over them oyx = [(oy, ox) for oy in range(otile1[0], otile2[0]) for ox in range(otile1[1], otile2[1])] for i in range(rank, len(oyx), size): oy, ox = oyx[i] # Read in all associated tiles into a list of lists rows = [] for dy in range(combine[0]): iy = oy * combine[0] + dy if iy >= itile2[0]: continue cols = [] for dx in range(combine[1]): ix = ox * combine[1] + dx if ix >= itile2[1]: continue if iy < itile1[0] or ix < itile1[1]: # The first tiles are missing on disk, but are # logically a part of the tiling. Use ibase, # which has been zeroed out. cols.append(ibase) else: itname = ipathfmt % {"y": iy, "x": ix} cols.append(enmap.read_map(itname)) if txflip: cols = cols[::-1] rows.append(cols) # Stack them next to each other into a big tile if tyflip: rows = rows[::-1] omap = enmap.tile_maps(rows) # Downgrade if necessary if np.any(downsample > 1): omap = enmap.downgrade(omap, downsample) if pad_to is not None: # Padding happens towards the end of the tiling, # which depends on the flip status padding = np.array( [[0, 0], [pad_to[0] - omap.shape[-2], pad_to[1] - omap.shape[-1]]]) if tyflip: padding[:, 0] = padding[::-1, 0] if txflip: padding[:, 1] = padding[::-1, 1] omap = enmap.pad(omap, padding) # And output otname = opathfmt % {"y": oy, "x": ox} utils.mkdir(os.path.dirname(otname)) enmap.write_map(otname, omap) if verbose: print otname
def read_map(fname): m = enmap.read_map(fname) m = m[...,1:-1,1:-1] #m = eval("m" + args.slice) m = enmap.downgrade(m, args.downgrade) return m