def A(self, xmap):
#enmap.write_map(args.odir + "/A1.fits", xmap)
iPx = mul(self.iP, enmap.map2harm(xmap))
#enmap.write_map(args.odir + "/A2.fits", iPx.real)
SiPx = enmap.harm2map(mul(self.S, iPx))
#enmap.write_map(args.odir + "/A3.fits", SiPx)
NiPx = mul(self.noise, enmap.harm2map(iPx))
#enmap.write_map(args.odir + "/A3.fits", NiPx)
res = xmap + SiPx + NiPx
#enmap.write_map(args.odir + "/A4.fits", res)
return res
def A(self, u):
s, a = self.dof.unzip(u)
# U"u = [S"s, 0a]
Uu = self.dof.zip(en.harm2map(en.map_mul(self.iS, en.map2harm(s))),
a * 0)
# P'N"P u
PNPu = self.PT(en.map_mul(self.iN, self.P(u)))
return Uu + PNPu
def calc_b(self):
PNd = self.PT(en.map_mul(self.iN, self.d))
Uw1_s = en.harm2map(
en.map_mul(self.hS,
en.rand_gauss_harm(self.d.shape[-3:], self.d.wcs)))
Uw1_a = np.zeros(self.T.shape[0])
Uw1 = self.dof.zip(Uw1_s, Uw1_a)
PNw2 = self.PT(
en.map_mul(self.hN, en.rand_gauss(self.d.shape, self.d.wcs)))
return PNd + Uw1 + PNw2
def M(self, u):
# Multiplying things out, the full expression for A is:
# [ S" + sum(N") sum(N"T) ]
# [ sum(T'N") sum(T'T) ]
# A reasonable approximation for this is
# [ S" + sum(sigma^{-2}) 0 ]
# [ 0 sum(T'T) ]
# which can be directly inverted.
s, a = self.dof.unzip(u)
# Solve for the cmb signal component
res_s = en.harm2map(en.map_mul(self.S_prec, en.map2harm(s)))
res_a = np.linalg.solve(self.TT, a)
return self.dof.zip(res_s, res_a)
def blur(m, sigma): l = np.sum(m.lmap()**2,0)**0.5 fm = enmap.map2harm(m) fm *= np.exp(-0.5*l**2*sigma**2)[None,:,:] return enmap.harm2map(fm)
def M(self, xmap):
#enmap.write_map(args.odir + "/M1.fits", xmap)
res = enmap.harm2map(mul(self.prec, enmap.map2harm(xmap)))
#enmap.write_map(args.odir + "/M2.fits", res)
return res
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)
