def gather_patches_plain(self):
"""Assemble patch statistics relevant to lensing at small scales.
Compute the small scale (ell > lmin) temperature power at different
patches across the sky as well as the average amplitude of the
background temperature gradient (ell < ldT).
"""
self._edge = 0 # Not throwing away edge pixels
m_fft = enmap.fft(self.map_in)
hp = np.zeros(self.map_in.shape)
hp[np.where((self._lmod > self.lmin) & (self._lmod < self.lmax))] = 1.
self._Tss = enmap.ifft(m_fft * hp)
self._dTy, self._dTx = gradient_flat(self.map_in, self.ldT)
# Scale geometry for lower res map of patches
pshp, pwcs = enmap.scale_geometry(self.map_in.shape, self.map_in.wcs,
1. / self._p)
if not self.savesteps:
del self.map_in, m_fft
self._T2patch = enmap.zeros(pshp, pwcs)
self._dTxpatch = enmap.zeros(pshp, pwcs)
self._dTypatch = enmap.zeros(pshp, pwcs)
Trs = self._Tss[:pshp[-2] * self._p, :pshp[-1] * self._p].reshape(
[pshp[-2], self._p, pshp[-1], self._p])
dTxrs = self._dTx[:pshp[-2] * self._p, :pshp[-1] * self._p].reshape(
[pshp[-2], self._p, pshp[-1], self._p])
dTyrs = self._dTy[:pshp[-2] * self._p, :pshp[-1] * self._p].reshape(
[pshp[-2], self._p, pshp[-1], self._p])
self._T2patch[:, :] = np.var(Trs, axis=(1, 3))
self._dTypatch[:, :] = np.mean(dTyrs, axis=(1, 3))
self._dTxpatch[:, :] = np.mean(dTxrs, axis=(1, 3))
self._dT2patch = self._dTxpatch**2 + self._dTypatch**2
if not self.savesteps:
del self._dTypatch, self._dTxpatch, self._dTy, self._dTx, self._Tss
def numgradient_flat(map_in):
"""Return the gradient maps of map_in.
Perform the central difference derivative on map_in. The values at
boundaries contain left/right derivatives.
Parameters
----------
map_in: ndmap, ndarray
The input map.
Returns
-------
dy: ndmap, ndarray
The gradient in the y-direction
dx: ndmap, ndarray
The gradient in the x-direction
"""
yy, xx = map_in.posmap()
dy = enmap.zeros(map_in.shape, map_in.wcs)
dx = enmap.zeros(map_in.shape, map_in.wcs)
dy[1:-1,:] = (map_in[:-2,:] - map_in[2:,:]) / (yy[:-2,:] - yy[2:,:])
dx[:,1:-1] = (map_in[:,:-2] - map_in[:,2:]) / (xx[:,:-2] - xx[:,2:])
dy[0,:] = (map_in[1,:] - map_in[0,:]) / (yy[1,:] - yy[0,:])
dy[-1,:] = (map_in[-1,:] - map_in[-2,:]) / (yy[-1,:] - yy[-2,:])
dx[:,0] = (map_in[:,1] - map_in[:,0]) / (xx[:,1] - xx[:,0])
dx[:,-1] = (map_in[:,-1] - map_in[:,-2]) / (xx[:,-1] - xx[:,-2])
return dy, dx
def gather_patches_cos(self):
"""Assemble patch statistics for small scale lensing with cos filter.
Compute the small scale (ell > 3000) temperature power at different
patches across the sky as well as the average amplitude of the
background temperature gradient (ell < 2000). For the small scale
statistics, also apply a filter in Fourier space such that:
.. math::
f_\\ell = \\cos(\\hat{\\ell}\\cdot\\hat{\\nabla T})
"""
self._edge = 5 # Edge pixels to throw away
p = self._p
m_fft = enmap.fft(self.map_in)
hp = np.zeros(self.map_in.shape)
hp[np.where((self._lmod > self.lmin) & (self._lmod < self.lmax))] = 1.
# Apply pre-whitening or Wiener/inverse variance filters, then top hat
if self._aW and self.fid is not None:
cs = CubicSpline(self.fid[0], self.fid[1]) # (ell, Cl)
m_fft = m_fft / cs(self._lmod)
self._Tss = enmap.ifft(m_fft * hp)
self._dTy, self._dTx = gradient_flat(self.map_in, self.ldT)
# Scale geometry for lower res map of patches
pshp, pwcs = enmap.scale_geometry(self.map_in.shape, self.map_in.wcs,
1. / self._p)
if not self.savesteps:
del self.map_in, m_fft
self._T2patch = enmap.zeros(pshp, pwcs)
self._dTxpatch = enmap.zeros(pshp, pwcs)
self._dTypatch = enmap.zeros(pshp, pwcs)
self._T_sub = np.zeros((pshp[-2], pshp[-1], p, p))
for i in range(self._T2patch.shape[-2]):
for j in range(self._T2patch.shape[-1]):
self._dTypatch[i, j] = np.mean(self._dTy[i * p:(i + 1) * p,
j * p:(j + 1) * p])
self._dTxpatch[i, j] = np.mean(self._dTx[i * p:(i + 1) * p,
j * p:(j + 1) * p])
Tp = self._Tss[i * p:(i + 1) * p, j * p:(j + 1) * p]
lsuby, lsubx = Tp.lmap()
lsubmod = Tp.modlmap()
lsubmod[0, 0] = 1.
fl = 1.j*(lsubx*self._dTxpatch[i,j] + \
lsuby*self._dTypatch[i,j]) / \
(lsubmod * np.sqrt(self._dTxpatch[i,j]**2 + \
self._dTypatch[i,j]**2))
fl[0, 0] = 0.
self._T_sub[i, j, :, :] = enmap.ifft(enmap.fft(Tp) * fl).real
# Throw away pixels with edge effects
self._T2patch[i, j] = np.var(self._T_sub[i, j, 5:-5, 5:-5])
self._dT2patch = self._dTxpatch**2 + self._dTypatch**2
if not self.savesteps:
del self._dTypatch, self._dTxpatch, self._dTy, self._dTx, self._Tss
del self._T_sub
def test_30_supergeom_translate(self):
"""Check that coords.get_supergeom does sensible thing for maps in
cylindrical projections with compatible but not identical
crval.
"""
proj = 'CAR'
ra0, dec0 = CRVAL
res = 0.01 * DEG
wcs = coords.get_wcs_kernel(proj, ra0, dec0, res)
wcs.wcs.crpix = (60, 70)
map0 = enmap.zeros((100, 200), wcs=wcs)
map0[2, 3] = 10.
map0[90, 192] = 11.
# Extracts.
m1 = map0[:10, :10]
m2 = map0[-10:, -10:]
# In simple cylindrical projections, there's a degeneracy
# between crval and crpix in the longitude component -- crval
# can be anywhere on the equator. It is useful to be able to
# join maps even if they have different crval[0], provided the
# pixel centers line up. (The same is not true of crval[1],
# which tips the native equator relative to the celestial
# equator.)
for axis, should_work in [(0, True), (1, False)]:
dpix = 10.5
m2 = map0[-10:, -10:]
m2.wcs.wcs.crpix[axis] += dpix
m2.wcs.wcs.crval[axis] += dpix * m2.wcs.wcs.cdelt[axis]
if should_work:
sg = coords.get_supergeom((m1.shape, m1.wcs),
(m2.shape, m2.wcs))
mapx = enmap.zeros(*sg)
mapx.insert(m1)
mapx.insert(m2)
self.assertTupleEqual(map0.shape,
mapx.shape,
msg="Reconstructed map shape.")
self.assertTrue(np.all(mapx == map0),
msg="Reconstructed map data.")
else:
msg = "Translating crval in dec should cause "\
"coord consistency check failure."
with self.assertRaises(ValueError, msg=msg):
sg = coords.get_supergeom((m1.shape, m1.wcs),
(m2.shape, m2.wcs))
def gather_patches_cos2(self):
"""Assemble patch statistics for small scale lensing with cos^2 filter.
Compute the small scale (ell > 3000) temperature power at different
patches across the sky as well as the average amplitude of the
background temperature gradient (ell < 2000). For the small scale
statistics, also apply a filter in Fourier space such that:
.. math::
f_\\ell = \\cos^2(\\hat{\\ell}\\cdot\\hat{\\nabla T})
"""
self._edge = 3 # Edge pixels to throw away
p = self._p
m_fft = enmap.fft(self.map_in)
hp = np.zeros(self.map_in.shape)
hp[np.where((self._lmod > self.lmin) & (self._lmod < self.lmax))] = 1.
self._Tss = enmap.ifft(m_fft * hp)
self._dTy, self._dTx = gradient_flat(self.map_in, self.ldT)
# Scale geometry for lower res map of patches
pshp, pwcs = enmap.scale_geometry(self.map_in.shape, self.map_in.wcs,
1. / self._p)
if not self.savesteps:
del self.map_in, m_fft
self._T2patch = enmap.zeros(pshp, pwcs)
self._dTxpatch = enmap.zeros(pshp, pwcs)
self._dTypatch = enmap.zeros(pshp, pwcs)
self._T_sub = np.zeros((pshp[-2], pshp[-1], p, p))
for i in range(self._T2patch.shape[-2]):
for j in range(self._T2patch.shape[-1]):
self._dTypatch[i, j] = np.mean(self._dTy[i * p:(i + 1) * p,
j * p:(j + 1) * p])
self._dTxpatch[i, j] = np.mean(self._dTx[i * p:(i + 1) * p,
j * p:(j + 1) * p])
Tp = self._Tss[i * p:(i + 1) * p, j * p:(j + 1) * p]
lsuby, lsubx = Tp.lmap()
lsubmod = Tp.modlmap()
lsubmod[0, 0] = 1. # Avoid divide by 0; set fl here to 0 later
fl = (lsubx*self._dTxpatch[i,j] + \
lsuby*self._dTypatch[i,j])**2 / \
(lsubmod * np.sqrt(self._dTxpatch[i,j]**2 + \
self._dTypatch[i,j]**2))**2
fl[0, 0] = 0.
self._T_sub[i, j, :, :] = enmap.ifft(enmap.fft(Tp) * fl).real
# Throw away pixels with edge effects
self._T2patch[i, j] = np.var(self._T_sub[i, j, 3:-3, 3:-3])
self._dT2patch = self._dTxpatch**2 + self._dTypatch**2
if not self.savesteps:
del self._dTypatch, self._dTxpatch, self._dTy, self._dTx, self._Tss
del self._T_sub
def process(kouts, name="default", ellmax=None, y=False, y_ellmin=400):
ellmax = lmax if ellmax is None else ellmax
ksilc = enmap.zeros((Ny, Nx), wcs, dtype=np.complex128).reshape(-1)
ksilc[modlmap.reshape(-1) < lmax] = np.nan_to_num(kouts.copy())
ksilc = enmap.enmap(ksilc.reshape((Ny, Nx)), wcs)
ksilc[modlmap > ellmax] = 0
if y: ksilc[modlmap < y_ellmin] = 0
msilc = np.nan_to_num(
fft.ifft(ksilc, axes=[-2, -1], normalize=True).real * bmask)
enmap.write_map(proot + "outmap_%s.fits" % name, enmap.enmap(msilc, wcs))
p2d = fc.f2power(ksilc, ksilc)
bin_edges = np.arange(100, 3000, 40)
binner = stats.bin2D(modlmap, bin_edges)
cents, p1d = binner.bin(p2d)
try:
# io.plot_img(np.log10(np.fft.fftshift(p2d)),proot+"ksilc_%s.png" % name,aspect='auto')
# io.plot_img(msilc,proot+"msilc_%s.png" % name)
# io.plot_img(msilc,proot+"lmsilc_%s.png" % name,lim=300)
tmask = maps.mask_kspace(shape,
wcs,
lmin=300,
lmax=5000 if not (y) else 1500)
fmap = maps.filter_map(msilc, tmask) * bmask
io.plot_img(fmap, proot + "hmsilc_%s.png" % name, high_res=True)
except:
pass
return cents, p1d
def load_ivar(path, box=None, mJy=True, fcode=None):
files = glob.glob(path)
if len(files) == 1:
ivar = enmap.read_map(files[0])
elif len(files) == 2:
ivar1 = enmap.read_map(files[0])
ivar2 = enmap.read_map(files[1])
ivar = (1 / 4 * ivar1**-1 + 1 / 4 * ivar2**-1)**-1
else:
raise ValueError("Unknown format")
# ugly hack: e.g. f220 ivar is actually div which has shape
# 3x3xpixells what we really need is the diagonal components
if len(ivar.shape) == 4:
ivar_correct = enmap.zeros(ivar.shape[1:], ivar.wcs)
ivar_correct[0] = ivar[0, 0]
ivar_correct[1] = ivar[1, 1]
ivar_correct[2] = ivar[2, 2]
del ivar
ivar = ivar_correct
if mJy:
if fcode == 'f090':
ivar /= (244.1 * 1e3)**2
elif fcode == 'f150':
ivar /= (371.74 * 1e3)**2
elif fcode == 'f220':
ivar /= (483.69 * 1e3)**2
elif fcode == 'f350':
ivar /= (287.5 * 1e3)**2
else:
raise ValueError("Unknown fcode")
if box is not None: return ivar.submap(box)
else: return ivar
def get_lensed_cmb(self, seed=None, kappa=None, save_output=True, verbose=True, overwrite=False, dtype=np.float64):
try:
assert(seed is not None)
assert(not overwrite)
if verbose:
print(f"trying to load saved lensed cmb. sim idx: {seed}")
lmaps = enmap.empty((3,)+self.shape, self.wcs, dtype=dtype)
for i, polidx in enumerate(['T','Q','U']):
fname = self.get_output_file_name("lensed_cmb", seed, polidx=polidx)
lmaps[i] = enmap.read_map(fname)
except:
ualm = self._generate_unlensed_cmb_alm(seed=seed)
if kappa is None:
kappa = self._get_kappa(seed, dtype=np.float64)
lmax = 10000
l = np.arange(lmax+1)
l_fact = 1/((l*(l+1))/2)
l_fact[0] = 0
kalm = curvedsky.map2alm(kappa, lmax=lmax)
kalm = curvedsky.map2alm(kappa, lmax=lmax); del kappa
kalm = hp.almxfl(kalm, l_fact)
tshape, twcs = enmap.fullsky_geometry(res=1*utils.arcmin)
print("start lensing")
lmaps = lensing.lens_map_curved((3,)+tshape, twcs, kalm, ualm)[0]
lalms = curvedsky.map2alm(lmaps, lmax=lmax, spin=0)
lmaps = curvedsky.alm2map(lalms, enmap.zeros((3,)+self.shape, self.wcs), spin=0)
print("finish lensing")
if save_output:
for i, polidx in enumerate(['T','Q','U']):
fname = self.get_output_file_name("lensed_cmb", seed, polidx=polidx)
os.makedirs(os.path.dirname(fname), exist_ok=True)
enmap.write_map(fname, lmaps[i].astype(np.float32))
return lmaps.astype(dtype)
def alm2stripe(alm, width, resol, proj='car'):
"""Return the stripe centered at dec=0 of the map corresponding to alm.
Return a slice of the map corresponding to alm from declination -width/2
to +width/2, and right ascension -180deg to +180deg.
Parameters
----------
alm: array
The set of alms to convert to map-space.
width: value
The width of the map (centered at dec=0) in degrees.
resol: value
The resolution of the map in arcmin.
proj: str, default='car'
The projection of the map. Must be compatible with pixell.
Returns
-------
stmap: array
The output map.
"""
box = np.array([[-width/2,180], [width/2,-180]]) * utils.degree
shape, wcs = enmap.geometry(pos=box, res=resol*utils.arcmin, proj=proj)
cmap = enmap.zeros(shape, wcs)
return curvedsky.alm2map(alm, cmap, method='cyl')
def get_n2d_data(splits,ivars,mask_a,coadd_estimator=False,flattened=False,plot_fname=None,dtype=None):
assert np.all(np.isfinite(splits))
assert np.all(np.isfinite(ivars))
assert np.all(np.isfinite(mask_a))
if coadd_estimator:
coadd,civars = get_coadd(splits,ivars,axis=1)
data = splits - coadd[:,None,...]
del coadd
else:
data = splits
assert np.all(np.isfinite(data))
if flattened:
# sivars = np.sqrt(ivars)
sivars = ((1./ivars) - (1./civars[:,None,...]))**-0.5
sivars[~np.isfinite(sivars)] = 0
ffts = enmap.fft(data*mask_a*sivars,normalize="phys")
if plot_fname is not None: plot(plot_fname+"_fft_maps",data*mask_a*sivars,quantile=0)
wmaps = mask_a + enmap.zeros(ffts.shape,mask_a.wcs,dtype=dtype) # WARNING: type
del ivars, data, splits
else:
assert np.all(np.isfinite(data*mask_a*ivars))
ffts = enmap.fft(data*mask_a*ivars,normalize="phys")
if plot_fname is not None: plot(plot_fname+"_fft_maps",data*mask_a*ivars,quantile=0)
wmaps = ivars * mask_a
del ivars, data, splits
n2d = get_n2d(ffts,wmaps,coadd_estimator=coadd_estimator,plot_fname=plot_fname,dtype=dtype)
assert np.all(np.isfinite(n2d))
return n2d
def corr_to_mat(corr, n):
res = enmap.zeros([n,n,n,n],dtype=corr.dtype)
for i in range(n):
tmp = np.roll(corr, i, 0)[:n,:]
for j in range(n):
res[i,j] = np.roll(tmp, j, 1)[:,:n]
return res
def gradient_flat(map_in, lmax=2000):
"""Return the gradient maps of map_in.
Compute the gradient of the map_in in Fourier space. Simultaneously
low-pass the maps such that they only include modes ell < lmax.
Parameters
----------
map_in: ndmap, ndarray
The input map.
lmax: int
The cutoff ell for low-passing the maps.
Returns
-------
dy: ndmap, ndarray
The gradient in the y-direction
dx: ndmap, ndarray
The gradient in the x-direction
"""
ly, lx = map_in.lmap()
lmod = map_in.modlmap()
lp = enmap.zeros(map_in.shape)
lp[np.where(lmod < lmax)] = 1.
map_fft = enmap.fft(map_in)
dx = enmap.ifft(map_fft*lp*lx*1j).real
dy = enmap.ifft(map_fft*lp*ly*1j).real
return dy, dx
def build_iN_constcorr_prior(data,
cmb=None,
lknee0=2000,
ivars=None,
constcov=False):
if ivars is None: ivars = data.ivars
N = enmap.zeros((data.n, data.n) + data.maps.shape[-2:], data.maps.wcs,
data.maps.dtype)
ref = np.mean(ivars, (-2, -1))
ref = np.maximum(ref, np.max(ref) * 1e-4)
norm = 1 / (ref / ivars.pixsize())
for i, freq in enumerate(data.freqs):
lknee = lknee0 * freq / 100
N[i, i] = (1 + (np.maximum(0.5, data.l) / lknee)**-3.5) * norm[i]
# Deconvolve the pixel window from the theoretical flat-at-high-l spectrum
N = N / data.wy[:, None]**2 / data.wx[None, :]**2
# Apply the beam-convolved cmb if available
if cmb is not None:
Bf = data.fconvs[:, None, None] * data.beams
N += cmb * Bf[:, None] * Bf[None, :]
del Bf
if not constcov:
N /= norm[:, None, None, None]**0.5 * norm[None, :, None, None]**0.5
iN = array_ops.eigpow(N, -1, axes=[0, 1])
return iN
def read_monolithic(idir, verbose=True, slice=None, dtype=None):
# Find the range of input tiles
ipathfmt = idir + "/tile%(y)03d_%(x)03d.fits"
itile1, itile2 = find_tile_range(ipathfmt)
def read(fname):
m = enmap.read_map(fname)
if slice: m = eval("m" + slice)
return m
# Read the first and last tile to get the total dimensions
m1 = read(ipathfmt % {"y": itile1[0], "x": itile1[1]})
m2 = read(ipathfmt % {"y": itile2[0] - 1, "x": itile2[1] - 1})
wy, wx = m1.shape[-2:]
oshape = tuple(
np.array(m1.shape[-2:]) * (itile2 - itile1 - 1) +
np.array(m2.shape[-2:]))
if dtype is None: dtype = m1.dtype
omap = enmap.zeros(m1.shape[:-2] + oshape, m1.wcs, dtype)
del m1, m2
# Now loop through all tiles and copy them in to the correct position
for ty in range(itile1[0], itile2[0]):
for tx in range(itile1[1], itile2[1]):
m = read(ipathfmt % {"y": ty, "x": tx})
oy = ty - itile1[0]
ox = tx - itile1[1]
omap[..., oy * wy:(oy + 1) * wy, ox * wx:(ox + 1) * wx] = m
if verbose: print(ipathfmt % {"y": ty, "x": tx})
return omap
def test_b_sign(self):
"""
We generate a random IQU map with geometry such that cdelt[0]<0
We transform this to TEB with map2harm and map2alm followed by
scalar harm2map and alm2map and use these as reference T,E,B maps.
We flip the original map along the RA direction.
We transform this to TEB with map2harm and map2alm followed by
scalar harm2map and alm2map and use these as comparison T,E,B maps.
We compare these maps.
"""
ells,cltt,clee,clbb,clte = np.loadtxt(DATA_PREFIX+"cosmo2017_10K_acc3_lensedCls.dat",unpack=True)
ps_cmb = np.zeros((3,3,ells.size))
ps_cmb[0,0] = cltt
ps_cmb[1,1] = clee
ps_cmb[2,2] = clbb
ps_cmb[1,0] = clte
ps_cmb[0,1] = clte
np.random.seed(100)
# Curved-sky is fine
lmax = 1000
alm = curvedsky.rand_alm_healpy(ps_cmb,lmax=lmax)
shape,iwcs = enmap.fullsky_geometry(res=np.deg2rad(10./60.))
wcs = enmap.empty(shape,iwcs)[...,::-1].wcs
shape = (3,) + shape
imap = curvedsky.alm2map(alm,enmap.empty(shape,wcs))
oalm = curvedsky.map2alm(imap.copy(),lmax=lmax)
rmap = curvedsky.alm2map(oalm,enmap.empty(shape,wcs),spin=0)
imap2 = imap.copy()[...,::-1]
oalm = curvedsky.map2alm(imap2.copy(),lmax=lmax)
rmap2 = curvedsky.alm2map(oalm,enmap.empty(shape,wcs),spin=0)
assert np.all(np.isclose(rmap[0],rmap2[0]))
assert np.all(np.isclose(rmap[1],rmap2[1]))
assert np.all(np.isclose(rmap[2],rmap2[2]))
# Flat-sky
px = 2.0
N = 300
shape,iwcs = enmap.geometry(pos=(0,0),res=np.deg2rad(px/60.),shape=(300,300))
shape = (3,) + shape
a = enmap.zeros(shape,iwcs)
a = a[...,::-1]
wcs = a.wcs
seed = 100
imap = enmap.rand_map(shape,wcs,ps_cmb,seed=seed)
kmap = enmap.map2harm(imap.copy())
rmap = enmap.harm2map(kmap,spin=0) # reference map
imap = imap[...,::-1]
kmap = enmap.map2harm(imap.copy())
rmap2 = enmap.harm2map(kmap,spin=0)[...,::-1] # comparison map
assert np.all(np.isclose(rmap[0],rmap2[0]))
assert np.all(np.isclose(rmap[1],rmap2[1],atol=1e0))
assert np.all(np.isclose(rmap[2],rmap2[2],atol=1e0))
def corr_to_mat(corr, ny,nx=None):
if nx is None: nx = ny
res = enmap.zeros([ny,nx,ny,nx],dtype=corr.dtype)
for i in range(ny):
tmp = np.roll(corr, i, 0)[:ny,:]
for j in range(nx):
res[i,j] = np.roll(tmp, j, 1)[:,:nx]
return res
def load_geometries(qids):
imap = enmap.zeros(pshape, pwcs)
imap2 = enmap.pad(imap, 900)
shape, wcs = imap2.shape, imap2.wcs
gs = {}
for qid in qids:
gs[qid] = shape, wcs
return gs
def make_hybrid(tfile, pfile):
# Make a hybrid TQU map from specified temperature and polarization files
T = enmap.read_map(tfile, sel=(0, ))
if ("545" in tfile) or ("857" in tfile): npol = 1
else: npol = 3
omap = enmap.zeros((npol, ) + T.shape, T.wcs, T.dtype)
omap[0] = T
if npol == 3: omap[1:] = enmap.read_map(pfile, sel=(slice(1, 3), ))
return omap
def __init__(self,
shape,
wcs,
comps="T",
noise_model=None,
dtype_tod=np.float32,
dtype_map=np.float64,
comm=mpi.COMM_WORLD,
recenter=None,
verbose=False):
if shape is not None:
shape = shape[-2:]
if noise_model is None:
noise_model = NmatUncorr()
self.shape = shape
self.wcs = wcs
self.comm = comm
self.comps = comps
self.ncomp = len(comps)
self.dtype_tod = dtype_tod
self.dtype_map = dtype_map
self.verbose = verbose
if shape is None:
self.map_rhs = None
self.map_div = None
self.auto_grow = True
else:
self.map_rhs = enmap.zeros((self.ncomp, ) + self.shape, self.wcs,
self.dtype_map)
self.map_div = enmap.zeros((self.ncomp, self.ncomp) + self.shape,
self.wcs, self.dtype_map)
self.auto_grow = False
self.map_idiv = None
self.noise_model = noise_model
self.observations = []
# Cut-related stuff
self.junk_offs = None
self.junk_rhs = []
self.junk_div = []
# Set up the degrees of freedom we will solve for
self.dof = MultiZipper(comm=comm)
#self.dof.add(MapZipper(self.map_rhs.shape, self.map_rhs.wcs))
self.ready = False
self.recenter = recenter
def read_map(fname, shape=None, wcs=None, ncomp=3):
m = nonan(read_helper(fname, shape, wcs))
#return m.preflat[:1]
m = m.reshape(-1, m.shape[-2], m.shape[-1])
if ncomp == 0: return m[0]
if len(m) == 1:
res = enmap.zeros((ncomp, ) + m.shape[1:], m.wcs, m.dtype)
res[0] = m
return res
else:
return m
def fcov_to_rcorr(shape,wcs,p2d,N):
"""Convert a 2D PS into a pix-pix covariance
"""
ncomp = p2d.shape[0]
p2d *= np.prod(shape[-2:])/enmap.area(shape,wcs)
ocorr = enmap.zeros((ncomp,ncomp,N*N,N*N),wcs)
for i in range(ncomp):
for j in range(i,ncomp):
dcorr = ps2d_to_mat(p2d[i,j].copy(), N).reshape((N*N,N*N))
ocorr[i,j] = dcorr.copy()
if i!=j: ocorr[j,i] = dcorr.copy()
return ocorr
def read_div(fname, shape=None, wcs=None, ncomp=3):
m = nonan(read_helper(fname, shape, wcs)) * 1.0
if ncomp == 0: return m.preflat[0]
#return m.preflat[:1][None]
if m.ndim == 2:
res = enmap.zeros((ncomp, ncomp) + m.shape[-2:], m.wcs, m.dtype)
for i in range(ncomp):
res[i, i] = m
return res
elif m.ndim == 4:
return m
else:
raise ValueError("Wrong number of dimensions in div %s" % fname)
def test_20_supergeom_simple(self):
"""Check that coords.get_supergeom does sensible things in simple cases."""
for proj in ['TAN', 'CEA']:
ra0, dec0 = CRVAL
res = 0.01 * DEG
wcs = coords.get_wcs_kernel(proj, ra0, dec0, res)
wcs.wcs.crpix = (60, 70)
map0 = enmap.zeros((100,200), wcs=wcs)
map0[2, 3] = 10.
map0[90, 192] = 11.
# Extracts.
m1 = map0[:10,:10]
m2 = map0[-10:,-10:]
# Reconstruct.
sg = coords.get_supergeom((m1.shape, m1.wcs), (m2.shape, m2.wcs))
mapx = enmap.zeros(*sg)
mapx.insert(m1)
mapx.insert(m2)
self.assertTupleEqual(map0.shape, mapx.shape)
self.assertTrue(np.all(mapx==map0))
def get_n2d(ffts, wmaps, plot_fname=None, coadd_estimator=False, dtype=None):
assert np.all(np.isfinite(ffts))
assert np.all(np.isfinite(wmaps))
shape, wcs = ffts.shape[-2:], ffts.wcs
modlmap = enmap.modlmap(shape, wcs)
Ny, Nx = shape[-2:]
nfreqs = ffts.shape[0]
npol = ffts.shape[2]
ncomps = nfreqs * npol
def ncomp_to_freq_pol(index):
ifreq = index // 3
ipol = index % 3
return ifreq, ipol
n2d = enmap.zeros((ncomps, ncomps, Ny, Nx), wcs,
dtype=dtype) # WARNING: type)
pols = ['I', 'Q', 'U']
for i in range(ncomps):
for j in range(i, ncomps):
ifreq, ipol = ncomp_to_freq_pol(i)
jfreq, jpol = ncomp_to_freq_pol(j)
isplits = ffts[ifreq, :, ipol]
iwts = wmaps[ifreq, :, 0]
if i != j:
jsplits = ffts[jfreq, :, jpol]
jwts = wmaps[jfreq, :, 0]
else:
jsplits = None
jwts = None
n2d[i, j] = noise_power(isplits,
iwts,
jsplits,
jwts,
coadd_estimator=coadd_estimator,
pfunc=naive_power)
if i != j: n2d[j, i] = n2d[i, j]
if plot_fname is not None:
plot("%s_%d_%s_%d_%s" \
% (plot_fname,ifreq,pols[ipol],
jfreq,pols[jpol]),
enmap.enmap(np.arcsinh(np.fft.fftshift(n2d[i,j])),wcs),dg=1,quantile=0)
return n2d
def allocate_output(area):
'''
Parameters
----------
area : str
Absolute path to footprint map.
Returns
-------
hitmap : (1, ny, nx) enmap
cutmap : (1, ny, nx) enmap
'''
shape, wcs = enmap.read_map_geometry(area)
hitmap = enmap.zeros((1, ) + shape[-2:], wcs, dtype)
cutmap = hitmap * 0
return hitmap, cutmap
def webget(tilepath,
box,
tsize=675,
res=0.5 * utils.arcmin,
dtype=np.float64,
verbose=False):
# Build the geometry representing the tiles web map. This is a normal
# fullsky geometry, except that the origin is in the top-left corner
# instead of the bottom-left, so we flip the y axis.
geo = enmap.Geometry(*enmap.fullsky_geometry(res=res))[::-1]
ntile = np.array(geo.shape) // tsize
pbox = enmap.subinds(*geo, box, cap=False, noflip=True)
# Set up output geometry with the same ordering as the input one. We will
# flip it to the final ordering at the end
ogeo = geo[pbox[0, 0]:pbox[1, 0], pbox[0, 1]:pbox[1, 1]]
omap = enmap.zeros(*ogeo, dtype=dtype)
# Loop through tiles
t1 = pbox[0] // tsize
t2 = (pbox[1] + tsize - 1) // tsize
urls = []
for ty in range(t1[0], t2[0]):
ty = ty % ntile[0]
for tx in range(t1[1], t2[1]):
tx = tx % ntile[1]
url = tilepath.format(y=ty, x=tx)
urls.append(url)
datas = download_all(urls, verbose=verbose)
di = 0
for ty in range(t1[0], t2[0]):
ty = ty % ntile[0]
y1, y2 = ty * tsize, (ty + 1) * tsize
for tx in range(t1[1], t2[1]):
tx = tx % ntile[1]
x1, x2 = tx * tsize, (tx + 1) * tsize
tgeo = geo[y1:y2, x1:x2]
data = datas[di]
di += 1
imgdata = np.array(Image.open(io.BytesIO(data)))
mapdata, mask = unpack(imgdata)
map = enmap.enmap(mapdata, tgeo.wcs)
omap.insert(map)
# Flip omap to normal ordering
omap = omap[::-1]
return omap
def ivar_hp_to_cyl(hmap, shape, wcs):
comm = mpi.COMM_WORLD
rstep = 100
dtype = np.float32
nside = hp.npix2nside(hmap.size)
dec, ra = enmap.posaxes(shape, wcs)
pix = np.zeros(shape, np.int32)
psi = np.zeros(shape, dtype)
# Get the pixel area. We assume a rectangular pixelization, so this is just
# a function of y
ipixsize = 4 * np.pi / (12 * nside**2)
opixsize = get_pixsize_rect(shape, wcs)
nblock = (shape[-2] + rstep - 1) // rstep
for bi in range(comm.rank, nblock, comm.size):
if bi % comm.size != comm.rank:
continue
i = bi * rstep
rdec = dec[i:i + rstep]
opos = np.zeros((2, len(rdec), len(ra)))
opos[0] = rdec[:, None]
opos[1] = ra[None, :]
ipos = opos[::-1]
pix[i:i + rstep, :] = hp.ang2pix(nside, np.pi / 2 - ipos[1], ipos[0])
del ipos, opos
for i in range(0, shape[-2], rstep):
pix[i:i + rstep] = utils.allreduce(pix[i:i + rstep], comm)
omap = enmap.zeros((1, ) + shape, wcs, dtype)
imap = np.array(hmap).astype(dtype)
imap = imap[None]
bad = hp.mask_bad(imap)
bad |= imap <= 0
imap[bad] = 0
del bad
# Read off the nearest neighbor values
omap[:] = imap[:, pix]
omap *= opixsize[:, None] / ipixsize
# We ignore QU mixing during rotation for the noise level, so
# it makes no sense to maintain distinct levels for them
mask = omap[1:] > 0
omap[1:] = np.mean(omap[1:], 0)
omap[1:] *= mask
del mask
return omap[0]
def merge_tiles(shape, wcs, paths, margin=100, dtype=np.float32):
# Get the pre-dimensions from the first tile
shape = enmap.read_map_geometry(paths[0][0])[0][:-2] + shape[-2:]
omap = enmap.zeros(shape, wcs, dtype)
ny = len(paths)
nx = len(paths[0])
for ty in range(ny):
for tx in range(nx):
fname = paths[ty][tx]
imap = enmap.read_map(fname)
h, w = imap.shape[-2:]
wy = make_edge_interp(h - 2 * margin,
margin,
ty > 0,
ty < ny - 1,
dtype=imap.dtype)
wx = make_edge_interp(w - 2 * margin, margin, dtype=imap.dtype)
imap = imap * wy[:, None] * wx[None, :]
omap.insert(imap, op=np.ndarray.__iadd__)
return omap
def eval_srcs_loop(posmap,
poss,
amps,
beam,
cres,
nhit,
cell_srcs,
dtype=np.float64,
op=np.add,
verbose=False):
# Loop through each cell
ncy, ncx = nhit.shape
model = enmap.zeros(amps.shape[-1:] + posmap.shape[-2:], posmap.wcs, dtype)
for cy in range(ncy):
for cx in range(ncx):
nsrc = nhit[cy, cx]
if verbose:
print("map cell %5d/%d with %5d srcs" %
(cy * ncx + cx + 1, ncy * ncx, nsrc))
if nsrc == 0: continue
srcs = cell_srcs[cy, cx, :nsrc]
y1, y2 = (cy + 0) * cres[0], (cy + 1) * cres[0]
x1, x2 = (cx + 0) * cres[1], (cx + 1) * cres[1]
pixpos = posmap[:, y1:y2, x1:x2]
srcpos = poss[srcs].T # [2,nsrc]
srcamp = amps[srcs].T # [ncomp,nsrc]
r = utils.angdist(pixpos[::-1, None, :, :], srcpos[::-1, :, None,
None])
bpix = (r - beam[0, 0]) / (beam[0, 1] - beam[0, 0])
# Evaluate the beam at these locations
bval = utils.interpol(beam[1],
bpix[None],
mode="constant",
order=1,
mask_nan=False) # [nsrc,ry,rx]
cmodel = srcamp[:, :, None, None] * bval
cmodel = op.reduce(cmodel, -3)
op(model[:, y1:y2, x1:x2], cmodel, model[:, y1:y2, x1:x2])
return model
def get_input(input_name,set_idx,sim_idx,shape,wcs):
if input_name=='CMB':
cmb_type = 'LensedUnabberatedCMB'
signal_path = sints.dconfig['actsims']['signal_path']
cmb_file = os.path.join(signal_path, 'fullsky%s_alm_set%02d_%05d.fits' %(cmb_type, set_idx, sim_idx))
alms = hp.fitsfunc.read_alm(cmb_file, hdu = 1)
elif input_name=='tSZ':
ellmax = 8101
ells = np.arange(ellmax)
cyy = fgs.power_y(ells)[None,None]
cyy[0,0][ells<2] = 0
comptony_seed = seed_tracker.get_fg_seed(set_idx, sim_idx, 'comptony')
alms = curvedsky.rand_alm(cyy, seed = comptony_seed,lmax=ellmax)
elif input_name=='kappa':
signal_path = sints.dconfig['actsims']['signal_path']
k_file = os.path.join(signal_path, 'fullskyPhi_alm_%05d.fits' %(sim_idx))
palms = hp.fitsfunc.read_alm(k_file)
from pixell import lensing as plensing
alms = plensing.phi_to_kappa(palms)
omap = enmap.zeros((1,)+shape[-2:],wcs)
omap = curvedsky.alm2map(np.complex128(alms),omap,spin=0)[0]
return omap