def kappa_to_phi(kappa, modlmap, return_fphi=False):
fphi = enmap.samewcs(kappa_to_fphi(kappa, modlmap), kappa)
phi = enmap.samewcs(ifft(fphi).real, kappa)
if return_fphi:
return phi, fphi
else:
return phi
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 dump_maps(ofile, tod, data, pos, amp, rad=args.radius*m2r, res=args.resolution*m2r): ncomp = amp.shape[1] dmaps, drhs, ddiv = make_maps(tod.astype(dtype), data, pos, ncomp, rad, res) dstack = stack_maps(drhs, ddiv, amp[:,0]) dmaps = np.concatenate((dmaps,[dstack]),0) dmaps = enmap.samewcs(dmaps, ddiv) enmap.write_map(ofile, dmaps[::-1])
def __call__(self, x): xmap = self.dof.unzip(x) res = xmap*0 for info in self.infos: t = [time.time()] work = xmap*info.H t.append(time.time()) umap = info.U.apply(work) t.append(time.time()) fmap = fft.fft(umap+0j, axes=[-2,-1]) t.append(time.time()) fmap = info.N.apply(fmap, exp=0.5) t.append(time.time()) if info.W is not None: fmap = info.W.apply(fmap) t.append(time.time()) fmap = info.N.apply(fmap, exp=0.5) t.append(time.time()) umap = fft.ifft(fmap, umap+0j, axes=[-2,-1], normalize=True).real t.append(time.time()) work = enmap.samewcs(info.U.trans(umap, work),work) t.append(time.time()) work *= info.H t.append(time.time()) t = np.array(t) print " %4.2f"*(len(t)-1) % tuple(t[1:]-t[:-1]) res += work res = utils.allreduce(res,comm) return self.dof.zip(res)
def __call__(self, x):
xmap = self.dof.unzip(x)
res = xmap * 0
for info in self.infos:
t = [time.time()]
work = xmap * info.H
t.append(time.time())
umap = info.U.apply(work)
t.append(time.time())
fmap = fft.fft(umap + 0j, axes=[-2, -1])
t.append(time.time())
fmap = info.N.apply(fmap, exp=0.5)
t.append(time.time())
if info.W is not None:
fmap = info.W.apply(fmap)
t.append(time.time())
fmap = info.N.apply(fmap, exp=0.5)
t.append(time.time())
umap = fft.ifft(fmap, umap + 0j, axes=[-2, -1],
normalize=True).real
t.append(time.time())
work = enmap.samewcs(info.U.trans(umap, work), work)
t.append(time.time())
work *= info.H
t.append(time.time())
t = np.array(t)
print " %4.2f" * (len(t) - 1) % tuple(t[1:] - t[:-1])
res += work
res = utils.allreduce(res, comm)
return self.dof.zip(res)
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 alm2map_pos(alm,
pos,
ainfo=None,
oversample=2.0,
spin=2,
deriv=False,
verbose=False):
"""Projects the given alms (with layout) on the specified pixel positions.
alm[ncomp,nelem], pos[2,...] => res[ncomp,...]. It projects on a large
cylindrical grid and then interpolates to the actual pixels. This is the
general way of doing things, but not the fastest. Computing pos and
interpolating takes a significant amount of time."""
alm_full = np.atleast_2d(alm)
if ainfo is None: ainfo = sharp.alm_info(nalm=alm_full.shape[-1])
ashape, ncomp = alm_full.shape[:-2], alm_full.shape[-2]
if deriv:
# If we're computing derivatives, spin isn't allowed.
# alm must be either [ntrans,nelem] or [nelem],
# and the output will be [ntrans,2,ny,nx] or [2,ny,nx]
ashape = ashape + (ncomp, )
ncomp = 2
tmap = make_projectable_map_by_pos(pos, ainfo.lmax, ashape + (ncomp, ),
oversample, alm.real.dtype)
alm2map_cyl(alm,
tmap,
ainfo=ainfo,
spin=spin,
deriv=deriv,
direct=True,
verbose=verbose)
# Project down on our final pixels. This will result in a slight smoothing
res = enmap.samewcs(tmap.at(pos[:2], mode="wrap"), pos)
# Remove any extra dimensions we added
if alm.ndim == alm_full.ndim - 1: res = res[0]
return res
def rand_map(shape,
wcs,
ps_lensinput,
lmax=None,
maplmax=None,
dtype=np.float64,
seed=None,
oversample=2.0,
spin=2,
output="l",
geodesic=True,
verbose=False):
ctype = np.result_type(dtype, 0j)
# First draw a random lensing field, and use it to compute the undeflected positions
if verbose: print "Computing observed coordinates"
obs_pos = enmap.posmap(shape, wcs)
if verbose: print "Generating alms"
alm = curvedsky.rand_alm(ps_lensinput, lmax=lmax, seed=seed, dtype=ctype)
phi_alm, cmb_alm = alm[0], alm[1:]
# Truncate alm if we want a smoother map. In taylens, it was necessary to truncate
# to a lower lmax for the map than for phi, to avoid aliasing. The appropriate lmax
# for the cmb was the one that fits the resolution. FIXME: Can't slice alm this way.
#if maplmax: cmb_alm = cmb_alm[:,:maplmax]
del alm
if "p" in output:
if verbose: print "Computing phi map"
phi_map = curvedsky.alm2map(phi_alm,
enmap.zeros(shape[-2:], wcs, dtype=dtype))
if verbose: print "Computing grad map"
grad = curvedsky.alm2map(phi_alm,
enmap.zeros((2, ) + shape[-2:], wcs, dtype=dtype),
deriv=True)
if verbose: print "Computing alpha map"
raw_pos = enmap.samewcs(
offset_by_grad(obs_pos, grad, pol=True, geodesic=geodesic), obs_pos)
del obs_pos, phi_alm
if "a" not in output: del grad
if "u" in output:
if verbose: print "Computing unlensed map"
cmb_raw = curvedsky.alm2map(cmb_alm,
enmap.zeros(shape, wcs, dtype=dtype),
spin=spin)
if verbose: print "Computing lensed map"
cmb_obs = curvedsky.alm2map_pos(cmb_alm,
raw_pos[:2],
oversample=oversample,
spin=spin)
if raw_pos.shape[0] > 2 and np.any(raw_pos[2]):
if verbose: print "Rotating polarization"
cmb_obs = enmap.rotate_pol(cmb_obs, raw_pos[2])
del cmb_alm, raw_pos
# Output in same order as specified in output argument
res = []
for c in output:
if c == "l": res.append(cmb_obs)
elif c == "u": res.append(cmb_raw)
elif c == "p": res.append(phi_map)
elif c == "a": res.append(grad)
return tuple(res)
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 = array_ops.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")
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 = array_ops.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")
def build_single(ifile, srcs, beam, ofile, mask_level=0, apod_size=16): imap = enmap.read_map(ifile) omap, oslice = pointsrcs.sim_srcs(imap.shape[-2:], imap.wcs, srcs, beam, return_padded=True) if mask_level: mask = omap > mask_level omap = 1-np.cos(np.minimum(1,ndimage.distance_transform_edt(1-mask)/16.0)*np.pi) omap = enmap.samewcs(omap, imap) omap = omap[oslice] enmap.write_map(ofile, omap)
def make_thumbs(self, off, amps): self.P.set_offset(off) tod2 = self.tod*0 self.P.forward(tod2, amps, pmul=1) data = self.thumb_mapper.map(self.tod.copy()) model = self.thumb_mapper.map(tod2) resid = data-model thumbs = enmap.samewcs([data,model,resid],data) return thumbs
def make_thumbs(self, off, amps): self.P.set_offset(off) tod2 = self.tod*0 self.P.forward(tod2, amps, pmul=1) data = self.thumb_mapper.map(self.tod.copy()) model = self.thumb_mapper.map(tod2) resid = data-model thumbs = enmap.samewcs([data,model,resid],data) return thumbs
def sim_scan_fast(imap, sigma, rad, n): mask = np.zeros(imap.shape[-2:]) mask[rad:-rad,rad:-rad] = 1 w = np.maximum(1 - ndimage.distance_transform_edt(1-mask)/rad,0)**3 w = w[None,:,:]*np.array([1,0.5,0.5])[:,None,None]*sigma**-2*n w = enmap.samewcs(w, imap) m = imap + np.random.standard_normal(imap.shape)*w**-0.5 m[~np.isfinite(m)] = 0 return m, w
def project_mat(pix, template, mat=None):
if mat is None: mat = np.full([pix.shape[-1], pix.shape[-1]], 1.0)
pix = np.asarray(pix)
off = np.asarray(template.shape[-2:]) / 2
rpix = (pix[:, :, None] - pix[:, None, :]) + off[:, None, None]
# Flatten
rpix = rpix.reshape(2, -1)
mat = np.asarray(mat).reshape(-1)
res = utils.bin_multi(rpix, template.shape[-2:], weights=mat)
return enmap.samewcs(res, template)
def project_mat(pix, template, mat=None): if mat is None: mat = np.full([pix.shape[-1],pix.shape[-1]],1.0) pix = np.asarray(pix) off = np.asarray(template.shape[-2:])/2 rpix = (pix[:,:,None] - pix[:,None,:]) + off[:,None,None] # Flatten rpix = rpix.reshape(2,-1) mat = np.asarray(mat).reshape(-1) res = utils.bin_multi(rpix, template.shape[-2:], weights=mat) return enmap.samewcs(res, template)
def map_to_color(map, crange, args):
"""Compute an [{R,G,B},ny,nx] color map based on a map[1 or 3, ny,nx]
map and a corresponding color range crange[{min,max}]. Relevant args
fields: color, method, rgb. If rgb is not true, only the first element
of the input map will be used. Otherwise 3 will be used."""
map = ((map.T-crange[0])/(crange[1]-crange[0])).T # .T ensures broadcasting for rgb case
if args.reverse_color: map = 1-map
if args.rgb: m_color = colorize.colorize(map, desc=args.color, driver=args.method, mode="direct")
else: m_color = colorize.colorize(map[0], desc=args.color, driver=args.method)
m_color = enmap.samewcs(np.rollaxis(m_color,2), map)
return m_color
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 eval(self, rvec):
"""Evaluate beam at positions rvec[{dra,dec},...] relative to beam center"""
# Decompose into parallel and orthogonal parts
rvec = np.asanyarray(rvec).copy()
rvec[0] *= np.cos(self.dec_ref)
rpara = np.sum(rvec*self.e_para[:,None,None],0)
rorto = np.sum(rvec*self.e_orto[:,None,None],0)
# Evaluate each beam component
ipara = rpara/self.res+self.vbeam.size/2
bpara = utils.interpol(self.vbeam, ipara[None], mask_nan=False, order=self.order, prefilter=False)
borto = np.exp(-0.5*rorto**2/self.sigma**2)
res = enmap.samewcs(bpara*borto, rvec)
return res
def dump_maps(ofile,
tod,
data,
pos,
amp,
rad=args.radius * m2r,
res=args.resolution * m2r):
ncomp = amp.shape[1]
dmaps, drhs, ddiv = make_maps(tod.astype(dtype), data, pos, ncomp, rad,
res)
dstack = stack_maps(drhs, ddiv, amp[:, 0])
dmaps = np.concatenate((dmaps, [dstack]), 0)
dmaps = enmap.samewcs(dmaps, ddiv)
enmap.write_map(ofile, dmaps[::-1])
def build_single(ifile, srcs, beam, ofile, mask_level=0, apod_size=16):
imap = enmap.read_map(ifile)
omap, oslice = pointsrcs.sim_srcs(imap.shape[-2:],
imap.wcs,
srcs,
beam,
return_padded=True)
if mask_level:
mask = omap > mask_level
omap = 1 - np.cos(
np.minimum(1,
ndimage.distance_transform_edt(1 - mask) / 16.0) *
np.pi)
omap = enmap.samewcs(omap, imap)
omap = omap[oslice]
enmap.write_map(ofile, omap)
def eval(self, rvec):
"""Evaluate beam at positions rvec[{dra,dec},...] relative to beam center"""
# Decompose into parallel and orthogonal parts
rvec = np.asanyarray(rvec).copy()
rvec[0] *= np.cos(self.dec_ref)
rpara = np.sum(rvec * self.e_para[:, None, None], 0)
rorto = np.sum(rvec * self.e_orto[:, None, None], 0)
# Evaluate each beam component
ipara = rpara / self.res + self.vbeam.size / 2
bpara = utils.interpol(self.vbeam,
ipara[None],
mask_nan=False,
order=self.order,
prefilter=False)
borto = np.exp(-0.5 * rorto**2 / self.sigma**2)
res = enmap.samewcs(bpara * borto, rvec)
return res
def read(self,y,x):
if self.nphi: x = x % self.nphi
for c in self.cache:
if c[0] == (y,x): return c[1]
else:
while len(self.cache) >= self.ncache:
del self.cache[0]
fname = self.pathfmt % {"y":y,"x":x}
if os.path.isfile(fname):
m = enmap.read_map(fname)
if self.crop:
m = m[...,self.crop:-self.crop,self.crop:-self.crop]
if self.ncomp:
m = m.preflat
extra = np.tile(m[:1]*0, (self.ncomp-len(m),1,1))
m = enmap.samewcs(np.concatenate([m,extra],0),m)
else:
m = None
self.cache.append([(y,x),m])
return m
def solve(self, b, x0=None, verbose=False, nmin=0):
if x0 is None: x0 = b * 0
def wrap(fun):
def foo(x):
xmap = enmap.samewcs(x.reshape(b.shape), b)
return fun(xmap).reshape(-1)
return foo
solver = CG(wrap(self.A),
b.reshape(-1),
x0=x0.reshape(-1),
M=wrap(self.M))
#for i in range(50):
#while solver.err > 1e-6:
while solver.err > 1e-2 or solver.i < nmin:
solver.step()
if verbose:
print "%5d %15.7e %15.7e" % (solver.i, solver.err,
solver.err_true)
return enmap.samewcs(solver.x.reshape(b.shape), b)
def build_gauss(posmap, sigma):
r2 = np.sum(posmap**2,0)
return np.exp(-0.5*r2/sigma**2)
# Ok, read in the maps
maps, divs, ids = [], [], []
for i, (rfile,dfile) in enumerate(zip(mapfiles, divfiles)):
print "Reading %s" % rfile
map = read_map(rfile)
print "Reading %s" % dfile
div = read_map(dfile).preflat[0]
maps.append(map)
divs.append(div)
ids.append(os.path.basename(rfile)[:-9])
maps = enmap.samewcs(np.asarray(maps),maps[0])
divs = enmap.samewcs(np.asarray(divs),divs[0])
nmap = maps.shape[0]
ncomp= maps.shape[1]
ref, refdiv = maps[0], divs[0]
ref_small = eval("ref"+args.slice)
refdiv_small = eval("refdiv"+args.slice)
# Fit gaussian
fitter = SingleFitter(build_gauss(ref_small.posmap(), beam_sigma)[None], ref_small[:1], refdiv_small)
p_ref = fitter.fit(verbose=True)
# Don't override amplitude or offset
p_ref[2:] = [1,0]
ref_small = apply_params(ref_small, p_ref)
# Choose a reference map. Fit all maps to it. Coadd to find
def mul(w, m):
if w.ndim < 4: return m * w
elif w.ndim == 4:
return enmap.samewcs(array_ops.matmul(w, m, axes=[0, 1]), m)
else:
raise NotImplementedError("Only 2d, 3d or 4d weight maps understood")
def calc_profile(self, off): pos = self.pos + off[:,None,None] r = np.sum(pos**2,0)**0.5 pix = r/self.dr return enmap.samewcs(utils.interpol(self.beam_pre, pix[None], prefilter=False, mask_nan=False), pos)
import numpy as np, argparse, os
from enlib import enmap, utils
from enact import files
parser = argparse.ArgumentParser()
parser.add_argument("ifiles", nargs="+")
parser.add_argument("ofile")
args = parser.parse_args()
maps = []
for ifile in args.ifiles:
ibase = os.path.basename(ifile)
print ibase
data = enmap.read_map(ifile)
maps.append(data)
maps = enmap.samewcs(np.array(maps),maps[0])
print np.sum(maps**2)
print maps.shape
masked = np.ma.MaskedArray(maps, mask=maps==0)
print np.sum(masked**2)
medmap = np.ma.median(masked,0)
medmap = enmap.samewcs(np.asarray(medmap),maps)
# And run through again, subtracting it
enmap.write_map(args.ofile, medmap)
def mapdiag(map):
if map.ndim < 4: return map
elif map.ndim == 4: enmap.samewcs(np.einsum("iiyx->iyx", map), map)
else: raise NotImplementedError
TOnly=not (pol),
gradCut=grad_cut,
uEqualsL=not (cluster))
fkmaps = fftfast.fft(measured, axes=[-2, -1])
if pol:
qest.updateTEB_X(fkmaps[0], fkmaps[1], fkmaps[2], alreadyFTed=True)
else:
qest.updateTEB_X(fkmaps, alreadyFTed=True)
qest.updateTEB_Y()
for polcomb in pol_list:
print(("Reconstructing", polcomb, " for ", i, " ..."))
kappa_recon = enmap.samewcs(qest.getKappa(polcomb).real, measured)
if i == 0:
io.quickPlot2d(kappa_recon, out_dir + "kappa_recon_single.png")
kappa_recon -= kappa_recon.mean()
if cluster:
cents_prof, prof = binner_dat.bin(kappa_recon)
profiles[polcomb].append(prof)
else:
downk = enmap.downgrade(kappa_map,
analysis_pixel_scale / sim_pixel_scale)
kpower = fmaps.get_simple_power_enmap(kappa_recon)
cents_pwr, aclkk = dbinner_dat.bin(kpower)
cpower = fmaps.get_simple_power_enmap(enmap1=kappa_recon,
enmap2=downk)
cents_pwr, cclkk = dbinner_dat.bin(cpower)
apowers[polcomb].append(aclkk)
def divdiag(div):
if div.ndim == 2: return div.preflat
elif div.ndim == 3: return div
elif div.ndim == 4: return enmap.samewcs(np.einsum("aayx->ayx",div),div)
else: raise ValueError("Invalid div shape: %s" % div.shape)
def divdiag(div):
if div.ndim == 2: return div.preflat
elif div.ndim == 3: return div
elif div.ndim == 4: return enmap.samewcs(np.einsum("aayx->ayx",div),div)
else: raise ValueError("Invalid div shape: %s" % div.shape)
def calc_map_block_ivar(map, n): m = map_to_blocks(map, n) vmap = np.var(m, axis=(-3,-1)) vmap[vmap!=0] = 1/vmap[vmap!=0] return enmap.samewcs(vmap, map[...,::n,::n])
def calc_map_block_mean(map, n): m = map_to_blocks(map, n) return enmap.samewcs(np.mean(m, axis=(-3,-1)), map[...,::n,::n])
def reorder(map, nrow, ncol, dets): return enmap.samewcs(map[utils.transpose_inds(dets,nrow,ncol)],map)
import numpy as np, argparse
from scipy import ndimage
from enlib import enmap
parser = argparse.ArgumentParser()
parser.add_argument("ifile")
parser.add_argument("ofile")
parser.add_argument("-r", "--apod-radius", type=int, default=64)
args = parser.parse_args()
def make_apod(shape, rad):
mask = np.zeros(shape[-2:])
mask[rad:-rad, rad:-rad] = 1
w = np.maximum(1 - ndimage.distance_transform_edt(1 - mask) / rad, 0) ** 3
return w
teb = enmap.read_map(args.ifile)
tqu = enmap.harm2map(enmap.fft(teb))
mask = make_apod(teb.shape, args.apod_radius)
tqu_mask = tqu * mask[None]
teb_mask = enmap.ifft(enmap.map2harm(tqu_mask)).real
res = enmap.samewcs([teb, tqu, tqu_mask, teb_mask], teb)
enmap.write_map(args.ofile, res)
for i in range(nfile):
det, off = files.read_point_template(ilayfiles[i])
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 = []
scan.noise.apply(tod) scan.pcut.backward(tod, ojunk[scan.cut_range[0]:scan.cut_range[1]]) scan.pmap.backward(tod, omap) del tod omap = utils.allreduce(omap, comm) return np.concatenate([omap.reshape(-1),ojunk],0) def M(x): map = x[:area.size].reshape(area.shape) junk = x[area.size:] omap = map*0 omap[:] = enmap.map_mul(idiv, map) ojunk= junk/jdiv return np.concatenate([omap.reshape(-1),ojunk],0) def dot(x,y): mprod = np.sum(x[:area.size]*y[:area.size]) jprod = np.sum(x[area.size:]*y[area.size:]) return mprod + comm.allreduce(jprod) bin = enmap.map_mul(idiv, cg_rhs) enmap.write_map(args.odir + "/map_bin.fits", bin) b = np.concatenate([cg_rhs.reshape(-1),cg_rjunk],0) solver = cg.CG(A, b, M=M, dot=dot) for i in range(nstep): solver.step() if comm.rank == 0: print "%5d %15.7e" % (solver.i, solver.err) if solver.i % args.ostep == 0: map = enmap.samewcs(solver.x[:area.size].reshape(area.shape),area) enmap.write_map(args.odir + "/map%04d.fits" % solver.i, map)
print fbin f1,f2 = [min(nfreq-1,int(i*fmax/dfreq/nbin)) for i in [fbin,fbin+1]] fsub = ft[:,f1:f2] cov = array_ops.measure_cov(fsub) std = np.diag(cov)**0.5 corr = cov / std[:,None] / std[None,:] myrhs = project_mat(pix, template, corr) mydiv = project_mat(pix, template) return fbin, myrhs, mydiv def collect(args): fbin, myrhs, mydiv = args rhs[fbin] += myrhs div[fbin] += mydiv p = multiprocessing.Pool(args.nmulti) for fbin in range(nbin): p.apply_async(handle_bin, [fbin], callback=collect) p.close() p.join() del ft # Collect the results if comm.rank == 0: print "Reducing" rhs = enmap.samewcs(utils.allreduce(rhs, comm), rhs) div = enmap.samewcs(utils.allreduce(div, comm), div) with utils.nowarn(): map = rhs/div if comm.rank == 0: print "Writing" enmap.write_map(args.ofile, map)
def mapdiag(map):
if map.ndim < 4: return map
elif map.ndim == 4: enmap.samewcs(np.einsum("iiyx->iyx",map),map)
else: raise NotImplementedError
def rand_map(shape,
wcs,
ps_lensinput,
lmax=None,
maplmax=None,
dtype=np.float64,
seed=None,
oversample=2.0,
spin=2,
output="l",
geodesic=True,
verbose=False,
delta_theta=None):
import curvedsky, sharp
ctype = np.result_type(dtype, 0j)
# Restrict to target number of components
oshape = shape[-3:]
if len(oshape) == 2: shape = (1, ) + tuple(shape)
# First draw a random lensing field, and use it to compute the undeflected positions
if verbose: print("Generating alms")
alm, ainfo = curvedsky.rand_alm(ps_lensinput,
lmax=lmax,
seed=seed,
dtype=ctype,
return_ainfo=True)
phi_alm, cmb_alm = alm[0], alm[1:1 + shape[-3]]
# Truncate alm if we want a smoother map. In taylens, it was necessary to truncate
# to a lower lmax for the map than for phi, to avoid aliasing. The appropriate lmax
# for the cmb was the one that fits the resolution. FIXME: Can't slice alm this way.
#if maplmax: cmb_alm = cmb_alm[:,:maplmax]
del alm
if delta_theta is None: bsize = shape[-2]
else:
bsize = utils.nint(abs(delta_theta / utils.degree / wcs.wcs.cdelt[1]))
# Adjust bsize so we don't get any tiny blocks at the end
nblock = shape[-2] // bsize
bsize = int(shape[-2] / (nblock + 0.5))
# Allocate output maps
if "p" in output: phi_map = enmap.empty(shape[-2:], wcs, dtype=dtype)
if "k" in output:
kappa_map = enmap.empty(shape[-2:], wcs, dtype=dtype)
l = np.arange(ainfo.lmax + 1.0)
kappa_alm = ainfo.lmul(phi_alm, l * (l + 1) / 2)
for i1 in range(0, shape[-2], bsize):
curvedsky.alm2map(kappa_alm, kappa_map[..., i1:i1 + bize, :])
del kappa_alm
if "a" in output:
grad_map = enmap.empty((2, ) + shape[-2:], wcs, dtype=dtype)
if "u" in output: cmb_raw = enmap.empty(shape, wcs, dtype=dtype)
if "l" in output: cmb_obs = enmap.empty(shape, wcs, dtype=dtype)
# Then loop over dec bands
for i1 in range(0, shape[-2], bsize):
i2 = min(i1 + bsize, shape[-2])
lshape, lwcs = enmap.slice_geometry(shape, wcs,
(slice(i1, i2), slice(None)))
if "p" in output:
if verbose: print("Computing phi map")
curvedsky.alm2map(phi_alm, phi_map[..., i1:i2, :])
if verbose: print("Computing grad map")
if "a" in output: grad = grad_map[..., i1:i2, :]
else: grad = enmap.zeros((2, ) + lshape[-2:], lwcs, dtype=dtype)
curvedsky.alm2map(phi_alm, grad, deriv=True)
if "l" not in output: continue
if verbose: print("Computing observed coordinates")
obs_pos = enmap.posmap(lshape, lwcs)
if verbose: print("Computing alpha map")
raw_pos = enmap.samewcs(
offset_by_grad(obs_pos, grad, pol=shape[-3] > 1,
geodesic=geodesic), obs_pos)
del obs_pos, grad
if "u" in output:
if verbose: print("Computing unlensed map")
curvedsky.alm2map(cmb_alm, cmb_raw[..., i1:i2, :], spin=spin)
if verbose: print("Computing lensed map")
cmb_obs[..., i1:i2, :] = curvedsky.alm2map_pos(cmb_alm,
raw_pos[:2],
oversample=oversample,
spin=spin)
if raw_pos.shape[0] > 2 and np.any(raw_pos[2]):
if verbose: print("Rotating polarization")
cmb_obs[..., i1:i2, :] = enmap.rotate_pol(cmb_obs[..., i1:i2, :],
raw_pos[2])
del raw_pos
del cmb_alm, phi_alm
# Output in same order as specified in output argument
res = []
for c in output:
if c == "l": res.append(cmb_obs.reshape(oshape))
elif c == "u": res.append(cmb_raw.reshape(oshape))
elif c == "p": res.append(phi_map)
elif c == "k": res.append(kappa_map)
elif c == "a": res.append(grad_map)
return tuple(res)
pl.add(fine_ells, lclbb * fine_ells**2., color="C2", ls="--")
pl.done(out_dir + "lccomp.png")
pl = io.Plotter(scaleX='log')
pl.add(cents, lte * cents**2., color="C0", ls="-")
pl.add(fine_ells, lclte * fine_ells**2., color="C0", ls="--")
pl.done(out_dir + "lccompte.png")
fkmaps = fftfast.fft(measured, axes=[-2, -1])
if deconvolve_beam: fkmaps = np.nan_to_num(fkmaps / kbeam_dat)
if maxlike and cluster:
polcomb = "TT"
fkmapsdc = np.nan_to_num(fkmaps / kbeam_dat)
maps = enmap.samewcs(
fftfast.ifft(fkmapsdc * fMaskCMB_T, normalize=True,
axes=[-2, -1]).real, measured)
#kappa_model = init_kappa_model
k = 0
io.quickPlot2d(maps, out_dir + "map_iter_" + str(k).zfill(3) + ".png")
from scipy.integrate import simps
Ny, Nx = shape_dat[-2:]
pixScaleY, pixScaleX = enmap.pixshape(shape_dat, wcs_dat)
Ukappa = init_kappa_model
Uft = fftfast.fft(Ukappa, axes=[-2, -1])
Upower = np.real(Uft * Uft.conjugate())
Nl2d = qest_maxlike.N.Nlkk[polcomb]
area = Nx * Ny * pixScaleX * pixScaleY
Upower = Upower * area / (Nx * Ny)**2
wfilter = np.nan_to_num(Upower / Nl2d)
def mul(mat, vec, axes=[0, 1]):
return enmap.samewcs(
array_ops.matmul(mat.astype(vec.dtype), vec, axes=axes), mat, vec)
import numpy as np, argparse
from enlib import enmap, log
parser = argparse.ArgumentParser()
parser.add_argument("ifiles", nargs="+")
parser.add_argument("ofile")
parser.add_argument("-v", "--verbose", action="store_true")
args = parser.parse_args()
L = log.init(level=log.DEBUG if args.verbose else log.ERROR)
maps = []
for ifile in args.ifiles:
L.info("Reading %s" % ifile)
maps.append(enmap.read_map(ifile))
L.info("Stacking")
maps = enmap.samewcs(maps, maps[0])
L.info("Writing %s" % args.ofile)
enmap.write_map(args.ofile, maps)
def filter_div(div):
"""Downweight very thin stripes in the div - they tend to be problematic single detectors"""
return enmap.samewcs(ndimage.minimum_filter(div, size=2), div)
def diag_plot(self, off): profile = self.calc_profile(off) amp, amp_div = self.calc_amp(profile) model_rhs = self.nmat(profile*amp[:,None,None]) omap = enmap.samewcs(np.concatenate([self.data.rhs, model_rhs, self.data.rhs-model_rhs][::-1],-1),profile) return omap
return enmap.samewcs(map[utils.transpose_inds(dets, nrow, ncol)], map)
# Read our map, and give each row a weight
pickup = enmap.read_map(args.pickup_map)
pickup = reorder(pickup, nrow, ncol, d.dets)
weight = np.median((pickup[:, 1:] - pickup[:, :-1])**2, -1)
weight[weight > 0] = 1 / weight[weight > 0]
# Find the output pixel for each input pixel
baz = pickup[:1].posmap()[1, 0]
bel = baz * 0 + args.el * utils.degree
ipoint = np.array([baz, bel])
opoint = ipoint[:, None, :] + d.point_offset.T[:, :, None]
opix = template.sky2pix(opoint[::-1]).astype(int) # [{y,x},ndet,naz]
opix = np.rollaxis(opix, 1) # [ndet,{y,x},naz]
omap = enmap.zeros((3, ) + template.shape[-2:], template.wcs)
odiv = enmap.zeros((3, 3) + template.shape[-2:], template.wcs)
for det in range(d.ndet):
omap += utils.bin_multi(opix[det], template.shape[-2:], weight[det] *
pickup[det]) * d.det_comps[det, :, None, None]
odiv += utils.bin_multi(
opix[det], template.shape[-2:], weight[det]) * d.det_comps[
det, :, None, None, None] * d.det_comps[det, None, :, None, None]
odiv = enmap.samewcs(array_ops.eigpow(odiv, -1, axes=[0, 1]), odiv)
omap = enmap.samewcs(array_ops.matmul(odiv, omap, axes=[0, 1]), omap)
enmap.write_map(args.ofile, omap)
def pow(mat, exp, axes=[0, 1]):
return enmap.samewcs(array_ops.eigpow(mat, exp, axes=axes), mat, exp)
def mul(w,m):
if w.ndim < 4: return m*w
elif w.ndim == 4: return enmap.samewcs(array_ops.matmul(w,m, axes=[0,1]),m)
else: raise NotImplementedError("Only 2d, 3d or 4d weight maps understood")
parser = argparse.ArgumentParser()
parser.add_argument("ifiles", nargs=2)
parser.add_argument("ofile")
parser.add_argument("-b", "--binsize", type=int, default=3)
parser.add_argument("-s", "--smooth", type=float, default=30)
parser.add_argument("--div", type=float, default=1.0)
args = parser.parse_args()
b = args.binsize
smooth = args.smooth * np.pi/180/60/(8*np.log(2))
m = [enmap.read_map(f) for f in args.ifiles]
dm = (m[1]-m[0])/2
pixarea = dm.area()/np.product(dm.shape[-2:])*(180*60/np.pi)**2
# Compute standard deviation in bins
dm = dm[...,:dm.shape[-2]/b*b,:dm.shape[-1]/b*b]
dm_blocks = dm.reshape(dm.shape[:-2]+(dm.shape[-2]/b,b,dm.shape[-1]/b,b))
var = np.std(dm_blocks,axis=(-3,-1))**2*pixarea/args.div
# This reshaping stuff messes up the wcs, which doesn't notice
# that we now have bigger pixels. So correct that.
var = enmap.samewcs(var, dm[...,::b,::b])
typ = np.median(var[var!=0])
var[~np.isfinite(var)] = 0
var = np.minimum(var,typ*1e6)
svar = enmap.smooth_gauss(var, smooth)
sigma = svar**0.5
enmap.write_map(args.ofile, sigma)
def foo(x):
xmap = enmap.samewcs(x.reshape(b.shape), b)
return fun(xmap).reshape(-1)
def reorder(map, nrow, ncol, dets):
return enmap.samewcs(map[utils.transpose_inds(dets, nrow, ncol)], map)
for i in range(nfile):
det, off = files.read_point_template(ilayfiles[i])
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]
w = w[None, :, :] * np.array([1, 0.5, 0.5])[:, None, None] * sigma ** -2 * n
w = enmap.samewcs(w, imap)
m = imap + np.random.standard_normal(imap.shape) * w ** -0.5
m[np.isnan(m)] = 0
return m, w
imap = enmap.read_map(args.ifile)
edge_width = 256
scan_shape = [d - 2 * edge_width for d in imap.shape[-2:]]
mtot, wtot = imap * 0, imap * 0
n = 0
outmaps = []
for ostep in range(0, int(np.log((args.scan_noise / args.target_noise) ** 2) / np.log(2)) + 1):
nloc = 2 ** ostep
if nloc < 512:
for i in range(0, nloc):
m, mw = sim_scan(imap, args.scan_noise, scan_shape, edge_width)
mtot += m * mw
wtot += mw
else:
m, mw = sim_scan_fast(imap, args.scan_noise, edge_width, nloc)
mtot += m * mw
wtot += mw
outmaps.append(mtot / wtot)
print "%5d %9.3f" % (2 ** (ostep + 1) - 1, np.max(wtot) ** -0.5)
outmaps = enmap.samewcs(outmaps, imap)
enmap.write_map(args.ofile, outmaps)
chits = np.zeros(nsrc)
ref_pixs = find_ref_pixs(divs)
for i in range(nref):
corr = measure_corr(pmaps, nmat, divs, tod, ref_pixs[i])
for i in range(nsrc):
if corr[i, 0,0] < 0.1: continue
corrs[i] += corr[i]
chits[i] += 1
if np.any(chits==0):
skip("Failed to measure correlations")
continue
corrs /= chits[:,None,None]
del tod
# Write as enmap + fits table
omap = enmap.samewcs([rhss, divs, corrs], rhss)
header = omap.wcs.to_header(relax=True)
header['NAXIS'] = omap.ndim
for i,n in enumerate(omap.shape[::-1]):
header['NAXIS%d'%(i+1)] = n
header['id'] = id
header['off_x'] = d.point_correction[0]/utils.arcmin
header['off_y'] = d.point_correction[1]/utils.arcmin
header['bore_el'] = np.mean(d.boresight[2])/utils.degree
header['bore_az1'] = np.min(d.boresight[1])/utils.degree
header['bore_az2'] = np.max(d.boresight[1])/utils.degree
header['ctime'] = np.mean(d.boresight[0])
map_hdu = fits.PrimaryHDU(omap, header)
srcinfo = np.zeros(nsrc, [('sid','i'),('ra','f'),('dec','f'),('amp','f')])
srcinfo["sid"] = sids
d.boresight[2] = args.el # In degrees, calibrated in next step
d = actdata.calibrate(d, exclude=["autocut"])
def reorder(map, nrow, ncol, dets):
return enmap.samewcs(map[utils.transpose_inds(dets,nrow,ncol)],map)
# Read our map, and give each row a weight
pickup = enmap.read_map(args.pickup_map)
pickup = reorder(pickup, nrow, ncol, d.dets)
weight = np.median((pickup[:,1:]-pickup[:,:-1])**2,-1)
weight[weight>0] = 1/weight[weight>0]
# Find the output pixel for each input pixel
baz = pickup[:1].posmap()[1,0]
bel = baz*0 + args.el * utils.degree
ipoint = np.array([baz,bel])
opoint = ipoint[:,None,:] + d.point_offset.T[:,:,None]
opix = template.sky2pix(opoint[::-1]).astype(int) # [{y,x},ndet,naz]
opix = np.rollaxis(opix, 1) # [ndet,{y,x},naz]
omap = enmap.zeros((3,)+template.shape[-2:], template.wcs)
odiv = enmap.zeros((3,3)+template.shape[-2:], template.wcs)
for det in range(d.ndet):
omap += utils.bin_multi(opix[det], template.shape[-2:], weight[det]*pickup[det]) * d.det_comps[det,:,None,None]
odiv += utils.bin_multi(opix[det], template.shape[-2:], weight[det]) * d.det_comps[det,:,None,None,None] * d.det_comps[det,None,:,None,None]
odiv = enmap.samewcs(array_ops.eigpow(odiv, -1, axes=[0,1]), odiv)
omap = enmap.samewcs(array_ops.matmul(odiv, omap, axes=[0,1]), omap)
enmap.write_map(args.ofile, omap)
# Add some convenience data hour = sdata[i].ctime/3600.%24 ostr += " | %5.2f %9.4f %9.4f | %7.4f %2d | %9.4f %9.4f" % (hour, fit.poss_hor[i,0]/utils.degree, fit.poss_hor[i,1]/utils.degree, fit.time/fit.nsrc, fit.nsrc, sdata[i].srcpos[0]/utils.degree, sdata[i].srcpos[1]/utils.degree) print ostr sys.stdout.flush() f.write(ostr + "\n") f.flush() if args.minimaps: # Build shifted models smap = project_maps([s.map.preflat[0] for s in sdata], fit.poss_cel, shape, wcs) sdiv = project_maps([s.div for s in sdata], fit.poss_cel, shape, wcs) smod = project_maps(fit.models, fit.poss_cel, shape, wcs) smap = enmap.samewcs([smap,smod,smap-smod],smap) # And build scaled coadd. When map is divided by a, div = ivar # is multiplied by a**2. Points in very noisy regions can have # large error bars in the amplitude, and hence might randomly # appear to have a very strong signal. We don't want to let these # dominate, so downweight points with atypically high variance. # FIXME: Need a better approach than this. This fixed some things, # but broke others. #beam_exposure = [] #for i, s in enumerate(sdata): # med_div = np.median(s.div[s.div!=0]) # profile = fit.models[i]/fit.amps[i] # avg_div = np.sum(s.div*profile)/np.sum(profile) # beam_exposure.append(avg_div/med_div) #beam_exposure = np.array(beam_exposure) #weight = np.minimum(1,beam_exposure*1.25)**4