def rand_map_flat(shape,
wcs,
ps,
lmax=None,
lens=True,
aberrate=True,
beta=None,
dir=None,
seed=None,
dtype=None,
verbose=False,
recenter=False,
pad=0):
"""Simulate a random flat-sky map. The input spectrum should be
[{phi,T,E,B},{phi,T,E,b},nl] of lens is True, and just [{T,E,B},{T,E,B},nl]
otherwise."""
if dtype is None: dtype = np.float64
if dir is None: dir = aberration.dir_equ
if beta is None: beta = aberration.beta
ctype = np.result_type(dtype, 0j)
if verbose: print "Generating unlensed cmb"
# No position calculation necessary if we're not lensing or aberrating.
if not lens and not aberrate:
return enmap.rand_map(shape, wcs, ps, seed=seed)
# Otherwise we must deal with various displacements
if aberrate: pad += np.pi * beta * 1.2
pad_pix = int(pad / enmap.pixsize(shape, wcs)**0.5)
if pad_pix > 0:
if verbose: print "Padding"
template = enmap.zeros(shape, wcs, np.int16)
template, pslice = enmap.pad(template, pad_pix, return_slice=True)
pshape, pwcs = template.shape, template.wcs
else:
pshape, pwcs = shape, wcs
# Simulate (padded) lensing map
if lens:
maps = enmap.rand_map((ps.shape[0], ) + pshape[-2:], pwcs, ps)
phi, unlensed = maps[0], maps[1:]
if verbose: print "Lensing"
m = lensing.lens_map_flat(unlensed, phi)
else:
m = enmap.rand_map((ps.shape[0], ) + pshape[-2:], pwcs, ps)
# Then handle aberration if necessary
if aberrate:
if verbose: print "Computing aberration displacement"
pos = m.posmap()
pos = enmap.samewcs(
aberration.remap(pos[1::-1], dir=dir, beta=beta,
recenter=recenter), pos)
amp = pos[3]
pos = pos[1::-1]
if verbose: print "Interpolating aberration"
m = enmap.samewcs(m.at(pos, mask_nan=False), m)
if verbose: print "Applying modulation"
m *= amp
if pad_pix > 0:
if verbose: print "Unpadding"
m = m[pslice]
return m
def getKappaSZ(bSims,snap,massIndex,px,thetaMapshape):
b = bSims
PIX = 2048
maps, z, kappaSimDat, szMapuKDat, projectedM500, trueM500, trueR500, pxInRad, pxInRad = b.getMaps(snap,massIndex,freqGHz=150.)
pxIn = pxInRad * 180.*60./np.pi
hwidth = PIX*pxIn/2.
print "At redshift ", z , " stamp is ", hwidth ," arcminutes wide."
# input pixelization
shapeSim, wcsSim = enmap.geometry(pos=[[-hwidth*arcmin,-hwidth*arcmin],[hwidth*arcmin,hwidth*arcmin]], res=pxIn*arcmin, proj="car")
kappaSim = enmap.enmap(kappaSimDat,wcsSim)
szMapuK = enmap.enmap(szMapuKDat,wcsSim)
# downgrade to native
shapeOut, wcsOut = enmap.geometry(pos=[[-hwidth*arcmin,-hwidth*arcmin],[hwidth*arcmin,hwidth*arcmin]], res=px*arcmin, proj="car")
kappaMap = enmap.project(kappaSim,shapeOut,wcsOut)
szMap = enmap.project(szMapuK,shapeOut,wcsOut)
print thetaMapshape
print szMap.shape
if szMap.shape[0]>thetaMapshape[0]:
kappaMap = enmap.project(kappaSim,thetaMapshape,wcsOut)
szMap = enmap.project(szMapuK,thetaMapshape,wcsOut)
else:
diffPad = ((np.array(thetaMapshape) - np.array(szMap.shape))/2.+0.5).astype(int)
apodWidth = 25
kappaMap = enmap.pad(enmap.apod(kappaMap,apodWidth),diffPad)[:-1,:-1]
szMap = enmap.pad(enmap.apod(szMap,apodWidth),diffPad)[:-1,:-1]
# print szMap.shape
assert szMap.shape==thetaMap.shape
# print z, projectedM500
# pl = Plotter()
# pl.plot2d(kappaMap)
# pl.done("output/kappasim.png")
# pl = Plotter()
# pl.plot2d(szMap)
# pl.done("output/szsim.png")
# sys.exit()
# print "kappaint ", kappaMap[thetaMap*60.*180./np.pi<10.].mean()
return kappaMap,szMap
def sim_srcs(shape, wcs, srcs, beam, omap=None, dtype=None, nsigma=5, rmax=None, method="loop", mmul=1,
return_padded=False):
"""Simulate a point source map in the geometry given by shape, wcs
for the given srcs[nsrc,{dec,ra,T...}], using the beam[{r,val},npoint],
which must be equispaced. If omap is specified, the sources will be
added to it in place. All angles are in radians. The beam is only evaluated up to
the point where it reaches exp(-0.5*nsigma**2) unless rmax is specified, in which
case this gives the maximum radius. mmul gives a factor to multiply the resulting
source model by. This is mostly useful in conction with omap. method can be
"loop" or "vectorized", but "loop" is both faster and uses less memory, so there's
no point in using the latter.
The source simulation is sped up by using a source lookup grid.
"""
if omap is None: omap = enmap.zeros(shape, wcs, dtype)
ishape = omap.shape
omap = omap.preflat
ncomp = omap.shape[0]
# In keeping with the rest of the functions here, srcs is [nsrc,{dec,ra,T,Q,U}].
# The beam parameters are ignored - the beam argument is used instead
amps = srcs[:,2:2+ncomp]
poss = srcs[:,:2].copy()
# Rewind positions to let us use flat-sky approximation for distance calculations
#wcs = enmap.enlib.wcs.fix_wcs(wcs)
ref = np.mean(enmap.box(shape, wcs, corner=False)[:,1])
poss[:,1] = utils.rewind(poss[:,1], ref)
beam = expand_beam(beam, nsigma, rmax)
rmax = nsigma2rmax(beam, nsigma)
# Pad our map by rmax, so we get the contribution from sources
# just ourside our area. We will later split our map into cells of size cres. Let's
# adjust the padding so we have a whole number of cells
cres = utils.nint(rmax/omap.pixshape())
epix = cres-(omap.shape[-2:]+2*cres)%cres
padding = [cres,cres+epix]
wmap, wslice = enmap.pad(omap, padding, return_slice=True)
# Overall we will have this many grid cells
cshape = wmap.shape[-2:]/cres
# Find out which sources matter for which cells
srcpix = wmap.sky2pix(poss.T).T
pixbox= np.array([[0,0],wmap.shape[-2:]],int)
nhit, cell_srcs = build_src_cells(pixbox, srcpix, cres)
posmap = wmap.posmap()
model = eval_srcs_loop(posmap, poss, amps, beam, cres, nhit, cell_srcs)
# Update our work map, through our view
if mmul != 1: model *= mmul
wmap += model
if not return_padded:
# Copy out
omap[:] = wmap[wslice]
# Restore shape
omap = omap.reshape(ishape)
return omap
else:
return wmap.reshape(ishape[:-2]+wmap.shape[-2:]), wslice
def lens_map_flat(cmb_map, phi_map):
raw_pix = cmb_map.pixmap() + enmap.grad_pix(phi_map)
# And extract the interpolated values. Because of a bug in map_pixels with
# mode="wrap", we must handle wrapping ourselves.
npad = int(
np.ceil(
max(np.max(-raw_pix),
np.max(raw_pix -
np.array(cmb_map.shape[-2:])[:, None, None]))))
pmap = enmap.pad(cmb_map, npad, wrap=True)
return enmap.samewcs(
utils.interpol(pmap, raw_pix + npad, order=4, mode="wrap"), cmb_map)
def init_geometry(ishape,iwcs):
modlmap = enmap.modlmap(ishape,iwcs)
bin_edges = np.arange(args.kellmin,args.kellmax,args.dell)
binner = stats.bin2D(modlmap,bin_edges)
if args.beam<1e-5:
kbeam = None
else:
kbeam = maps.gauss_beam(modlmap,args.beam)
lmax = modlmap.max()
ells = np.arange(2,lmax,1)
wnoise_TT = ells*0.+(args.noise*(np.pi/180./60.))**2.
wnoise_PP = 2.*wnoise_TT
nT = modlmap*0.+(args.noise*(np.pi/180./60.))**2.
nP = 2.*nT
ncomp = 3 if pol else 1
ps = np.zeros((ncomp,ncomp,ells.size))
ps[0,0] = wnoise_TT
if pol:
ps[1,1] = wnoise_PP
ps[2,2] = wnoise_PP
oshape = (3,)+ishape if pol else ishape
if not(args.flat) and args.noise_pad>1.e-5:
# Pad noise sim geometry
pad_width_deg = args.noise_pad
pad_width = pad_width_deg * np.pi/180.
res = maps.resolution(oshape[-2:],iwcs)
pad_pixels = int(pad_width/res)
template = enmap.zeros(oshape,iwcs)
btemplate = enmap.pad(template,pad_pixels)
bshape,bwcs = btemplate.shape,btemplate.wcs
del template
del btemplate
ngen = maps.MapGen(bshape,bwcs,ps)
else:
ngen = maps.MapGen(oshape,iwcs,ps)
tmask = maps.mask_kspace(ishape,iwcs,lmin=args.tellmin,lmax=args.tellmax)
pmask = maps.mask_kspace(ishape,iwcs,lmin=args.pellmin,lmax=args.pellmax)
kmask = maps.mask_kspace(ishape,iwcs,lmin=args.kellmin,lmax=args.kellmax)
qest = lensing.qest(ishape,iwcs,theory,noise2d=nT,beam2d=kbeam,kmask=tmask,noise2d_P=nP,kmask_P=pmask,kmask_K=kmask,pol=pol,grad_cut=None,unlensed_equals_lensed=True)
taper,w2 = maps.get_taper_deg(ishape,iwcs,taper_width_degrees = args.taper_width,pad_width_degrees = args.pad_width)
fc = maps.FourierCalc(oshape,iwcs,iau=args.iau)
purifier = maps.Purify(ishape,iwcs,taper) if args.purify else None
return qest,ngen,kbeam,binner,taper,fc,purifier
def draw(self, shape, wcs, window=False, nsigma=10, pad=False):
m = enmap.zeros((3, ) + shape[-2:], wcs)
w = np.max(self.widths) * nsigma
# Find all sources that enter even partially into our area
ipos = np.array([
enmap.sky2pix(shape, wcs, self.pos - w, corner=True, safe=True),
enmap.sky2pix(shape, wcs, self.pos + w, corner=True, safe=True)
]).astype(int)
ipos = np.sort(ipos, 0)
ashape = np.array(shape[-2:])
mask = np.all(ipos[1] >= 0, 0) & np.all(
ipos[0] < ashape[:, None], 0) & np.all(
ipos[1] - ipos[0] < ashape[:, None], 0)
nmask = np.sum(mask)
if nmask == 0: return m
if pad:
# We will include enough pixels on each side that we avoid
# truncating any part of the source
ipos2 = np.array([
enmap.sky2pix(shape,
wcs,
self.pos[:, mask] - w * 2,
corner=True,
safe=True),
enmap.sky2pix(shape,
wcs,
self.pos[:, mask] + w * 2,
corner=True,
safe=True)
]).astype(int)
npad = np.maximum(
0, [-np.min(ipos2, (0, 2)),
np.max(ipos2, (0, 2)) - ashape])
m, mslice = enmap.pad(m, npad, return_slice=True)
for i, (pos, amp, width, ibeam) in enumerate(
zip(self.pos.T[mask], self.amps.T[mask], self.widths.T[mask],
self.ibeam.T[mask])):
# Find necessary bounding box
sub = m.submap([pos - w, pos + w])
if sub.size == 0: continue
add_gauss(sub, pos, amp, ibeam)
# Optinally apply window function
if window: m = enmap.apply_window(m)
if pad: m = m[mslice]
return m
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