-
Notifications
You must be signed in to change notification settings - Fork 7
/
planet_map.py
202 lines (190 loc) · 7.36 KB
/
planet_map.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
from __future__ import division, print_function
import numpy as np, sys, os
from enlib import utils
with utils.nowarn(): import h5py
from enlib import config, pmat, mpi, errors, gapfill, enmap, bench, sampcut, cg
from enlib import fft, array_ops
from enact import filedb, actscan, actdata, cuts, nmat_measure
config.set("pmat_cut_type", "full")
parser = config.ArgumentParser(os.environ["HOME"]+"./enkirc")
parser.add_argument("planet")
parser.add_argument("area")
parser.add_argument("sel")
parser.add_argument("odir")
parser.add_argument("tag", nargs="?")
parser.add_argument("-R", "--dist", type=float, default=0.2)
parser.add_argument("-e", "--equator", action="count", default=0)
parser.add_argument("-z", "--zenith", action="count", default=0)
parser.add_argument("-c", "--cont", action="store_true")
parser.add_argument("--sim", type=str, default=None, help="Passing a sel here sets up simulation mode. The simulations will consist of data from the sim sel TODs with the scanning pattern of the real TODs, and with the signal read off from the area map")
parser.add_argument("--noiseless", action="store_true", help="Replace signal with simulation instead of adding them. This can be used to get noise free transfer functions")
parser.add_argument("--dbox", type=str, default=None, help="Select only detectors in y1:y2,x1:x2 in the focalplane, relative to the center of the array, in degrees.")
parser.add_argument("--tags", type=str, default=None)
parser.add_argument("-m", "--model", type=str, default="joneig")
args = parser.parse_args()
zenith = args.zenith - args.equator
comm = mpi.COMM_WORLD
filedb.init()
ids = filedb.scans[args.sel]
R = args.dist * utils.degree
csize= 100
dtype= np.float32
shape, wcs = enmap.read_map_geometry(filedb.get_patch_path(args.area))
ncomp= 3
shape= shape[-2:]
area = enmap.zeros((ncomp,)+shape, wcs, dtype)
model_fknee = 10
model_alpha = 10
sys = "hor:"+args.planet
if zenith <= 0: sys += "/0_0"
utils.mkdir(args.odir)
prefix = args.odir + "/"
if args.tag: prefix += args.tag + "_"
if args.dbox: dbox = np.array([[float(w) for w in tok.split(":")] for tok in args.dbox.split(",")]).T*utils.degree
else: dbox = None
if args.sim:
sim_ids = filedb.scans[args.sim][:len(ids)]
area = enmap.read_map(args.area).astype(dytpe)
def smooth(tod, srate):
ft = fft.rfft(tod)
freq = fft.rfftfreq(tod.shape[-1])*srate
flt = 1/(1+(freq/model_fknee)**model_alpha)
ft *= flt
fft.ifft(ft, tod, normalize=True)
return tod
def broaden_beam_hor(tod, d, ibeam, obeam):
ft = fft.rfft(tod)
k = 2*np.pi*fft.rfftfreq(d.nsamp, 1/d.srate)
el = np.mean(d.boresight[2,::100])
skyspeed = d.speed*np.cos(el)
sigma = (obeam**2-ibeam**2)**0.5
ft *= np.exp(-0.5*(sigma/skyspeed)**2*k**2)
fft.ifft(ft, tod, normalize=True)
def calc_model_joneig(tod, cut, srate=400):
return smooth(gapfill.gapfill_joneig(tod, cut, inplace=False), srate)
def calc_model_constrained(tod, cut, srate=400, mask_scale=0.3, lim=3e-4, maxiter=50, verbose=False):
# First do some simple gapfilling to avoid messing up the noise model
tod = sampcut.gapfill_linear(cut, tod, inplace=False)
ft = fft.rfft(tod) * tod.shape[1]**-0.5
iN = nmat_measure.detvecs_jon(ft, srate)
del ft
iV = iN.ivar*mask_scale
def A(x):
x = x.reshape(tod.shape)
Ax = iN.apply(x.copy())
Ax += sampcut.gapfill_const(cut, x*iV[:,None], 0, inplace=True)
return Ax.reshape(-1)
b = sampcut.gapfill_const(cut, tod*iV[:,None], 0, inplace=True).reshape(-1)
x0 = sampcut.gapfill_linear(cut, tod).reshape(-1)
solver = cg.CG(A, b, x0)
while solver.i < maxiter and solver.err > lim:
solver.step()
if verbose:
print("%5d %15.7e" % (solver.i, solver.err))
res = solver.x.reshape(tod.shape)
res = smooth(res, srate)
return res
calc_model = {"joneig": calc_model_joneig, "constrained": calc_model_constrained}[args.model]
for ind in range(comm.rank, len(ids), comm.size):
id = ids[ind]
bid = id.replace(":","_")
entry = filedb.data[id]
if args.tags: entry.tag = args.tags
oname = "%s%s_map.fits" % (prefix, bid)
if args.cont and os.path.isfile(oname):
print("Skipping %s (already done)" % (id))
continue
# Read the tod as usual
try:
if not args.sim:
with bench.show("read"):
d = actdata.read(entry)
else:
sim_id = sim_ids[ind]
sim_entry = filedb.data[sim_id]
with bench.show("read"):
d = actdata.read(entry, ["boresight"])
d += actdata.read(sim_entry, exclude=["boresight"])
with bench.show("calibrate"):
d = actdata.calibrate(d, exclude=["autocut"])
if d.ndet == 0 or d.nsamp < 2: raise errors.DataMissing("no data in tod")
# Select detectors if needed
if dbox is not None:
mid = np.mean(utils.minmax(d.point_template, 0), 0)
off = d.point_template-mid
good = np.all((off > dbox[0])&(off < dbox[1]),-1)
d = d.restrict(dets=d.dets[good])
if d.ndet == 0 or d.nsamp < 2: raise errors.DataMissing("No data left")
except errors.DataMissing as e:
print("Skipping %s (%s)" % (id, str(e)))
continue
print("Processing %s" % id, d.ndet, d.nsamp)
#broaden_beam_hor(d.tod, d, 1.35*utils.arcmin*utils.fwhm, 1.57*utils.arcmin*utils.fwhm)
# Very simple white noise model. This breaks if the beam has been tod-smoothed by this point.
with bench.show("ivar"):
tod = d.tod
del d.tod
tod -= np.mean(tod,1)[:,None]
tod = tod.astype(dtype)
diff = tod[:,1:]-tod[:,:-1]
diff = diff[:,:diff.shape[-1]//csize*csize].reshape(d.ndet,-1,csize)
ivar = 1/(np.median(np.mean(diff**2,-1),-1)/2**0.5)
del diff
# Estimate noise level
asens = np.sum(ivar)**-0.5 / d.srate**0.5
with bench.show("actscan"):
scan = actscan.ACTScan(entry, d=d)
with bench.show("pmat"):
pmap = pmat.PmatMap(scan, area, sys=sys)
pcut = pmat.PmatCut(scan)
rhs = enmap.zeros((ncomp,)+shape, wcs, dtype)
div = enmap.zeros((ncomp,ncomp)+shape, wcs, dtype)
junk = np.zeros(pcut.njunk, dtype)
# Generate planet cut
with bench.show("planet cut"):
planet_cut = cuts.avoidance_cut(d.boresight, d.point_offset, d.site,
args.planet, R)
if args.sim:
if args.noiseless: tod_orig = tod.copy()
with bench.show("inject"):
pmap.forward(tod, area)
# Compute atmospheric model
with bench.show("atm model"):
model = calc_model(tod, planet_cut, d.srate)
if args.sim and args.noiseless:
model -= calc_model(tod_orig, planet_cut, d.srate)
tod -= tod_orig
del tod_orig
with bench.show("atm subtract"):
tod -= model
del model
tod = tod.astype(dtype, copy=False)
# Should now be reasonably clean of correlated noise.
# Proceed to make simple binned map
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, fallback="scalar")
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