-
Notifications
You must be signed in to change notification settings - Fork 12
/
mapmaking.py
2098 lines (2012 loc) · 82.7 KB
/
mapmaking.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
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Separable map making. Defines a set of Signals, which contain forward
# and backaward projection as well as a preconditioner. These can then be
# used via
#
# A matrix (cut is one of the signals here)
# for each scan
# for each (signal,sflat)
# signal.forward(scan,tod,sflat)
# scan.noise.apply(tod)
# for each (signal,sflat) reverse
# signal.backward(scan,tod,sflat)
#
# M matrix
# for each (signal,sflat)
# signal.precon(sflat)
#
# Init: Because we want to treat cuts as a signal, but also want cuts
# to be available for other signals to use, this is a bit cumbersome.
# We make sure to initialize the cut signal first, and pass it as an
# extra argument to the other signals
# for each scan
# signal_cut.add(scan)
# for each other signal
# signal.add(scan, signal_cut)
#
# Or perhaps one should add the cut signal reference when *initializing*
# the signal objects, and then require cuts to be passed first in the
# signal array (in general passing signals other signals depend on first).
# Then it would look like
#
# for each scan
# for each signal
# signal.add(scan)
#
# signal_cut = SignalCut(...)
# signal_map = SignalMap(..., cut=signal_cut)
# signal_phase = SignalPhase(..., cut=signal_cut)
# signals = [signal_cut, signal_map, signal_phase]
from __future__ import division, print_function
import numpy as np, h5py, logging, gc
from . import enmap, dmap, array_ops, pmat, utils, todfilter, pointsrcs, zipper
from . import config, nmat, bench, gapfill, mpi, sampcut, fft, bunch
from .cg import CG
L = logging.getLogger(__name__)
def dump(fname, d):
print("dumping " + fname)
with h5py.File(fname, "w") as hfile:
hfile["data"] = d
config.default("debug_raw",False,"moo")
######## Signals ########
class Signal:
"""General signal interface."""
def __init__(self, name, ofmt, output, ext):
self.ext = ext
self.dof = zipper.ArrayZipper(np.zeros([0]))
self.precon = PreconNull()
self.prior = PriorNull()
self.post = []
self.filters= []
self.name = name
self.ofmt = ofmt
self.output = output
def prepare (self, x): return x.copy()
def forward (self, scan, tod, x): pass
def backward(self, scan, tod, x): pass
def precompute(self, scan): pass
def free(self): pass
def finish (self, x, y): x[:] = y
def polmat(self): return self.zeros(mat=True)
def transform(self, m, op): return op(m)
def polinv(self, m): return 1/m
def polmul(self, mat, m, inplace=False): return mat*m
# This one is a potentially cheaper version of self.prepare(self.zeros())
def work (self, mat=False): return self.prepare(self.zeros(mat=mat))
def write (self, prefix, tag, x): pass
def postprocess(self, x):
for p in self.post: x = p(x)
return x
def filter(self, x):
for f in self.filters:
x = f(x)
return x
class SignalMap(Signal):
def __init__(self, scans, area, comm, cuts=None, name="main", ofmt="{name}", output=True,
ext="fits", pmat_order=None, sys=None, nuisance=False, data=None, extra=[], split=False):
Signal.__init__(self, name, ofmt, output, ext)
self.area = area
self.cuts = cuts
self.dof = zipper.ArrayZipper(area, comm=comm)
self.dtype= area.dtype
self.split= split
if data is not None:
self.data = data
else:
self.data = {scan: pmat.PmatMap(scan, area, order=pmat_order, sys=sys, extra=extra, split=scan.split if split else None) for scan in scans}
def forward(self, scan, tod, work, tmul=1, mmul=1):
if scan not in self.data: return
self.data[scan].forward(tod, work, tmul=tmul, mmul=mmul)
def backward(self, scan, tod, work, tmul=1, mmul=1):
if scan not in self.data: return
self.data[scan].backward(tod, work, tmul=tmul, mmul=mmul)
def finish(self, m, work):
if m.flags["C_CONTIGUOUS"]:
self.dof.comm.Allreduce(work, m)
else:
tmp = np.ascontiguousarray(m)
self.dof.comm.Allreduce(work, tmp)
m[:] = tmp
def zeros(self, mat=False):
shape = self.area.shape
if mat: shape = shape[:-2]+shape[-3:]
return enmap.zeros(shape, self.area.wcs, self.area.dtype)
def polinv(self, m): return array_ops.eigpow(m, -1, axes=[-4,-3], lim=config.get("eig_limit"), fallback="scalar")
def polmul(self, mat, m, inplace=False):
if not inplace: m = m.copy()
m[:] = array_ops.matmul(mat, m, axes=[-4,-3])
return m
def write(self, prefix, tag, m):
if not self.output: return
oname = self.ofmt.format(name=self.name)
oname = "%s%s_%s.%s" % (prefix, oname, tag, self.ext)
if self.dof.comm.rank == 0:
enmap.write_map(oname, m)
class SignalMapFast(SignalMap):
def __init__(self, scans, area, comm, cuts=None, name="main", ofmt="{name}", output=True,
ext="fits", sys=None, nuisance=False, data=None, extra=[]):
if data is None:
data = {scan: pmat.PmatMapFast(scan, area, sys=sys, extra=extra) for scan in scans}
SignalMap.__init__(self, scans, area, comm=comm, cuts=cuts, name=name, ofmt=ofmt,
output=output, ext=ext, sys=sys, nuisance=nuisance, data=data)
def precompute(self, scan):
if scan not in self.data: return
self.pix, self.phase = self.data[scan].get_pix_phase()
def free(self):
del self.pix, self.phase
def forward(self, scan, tod, work):
if scan not in self.data: return
self.data[scan].forward(tod, work, self.pix, self.phase)
def backward(self, scan, tod, work):
if scan not in self.data: return
self.data[scan].backward(tod, work, self.pix, self.phase)
config.default("dmap_format","merged","How to store dmaps on disk. 'merged': combine into a single fits file before writing. This is memory intensive. 'tiles': Write the tiles that make up the dmap directly to disk.")
class SignalDmap(Signal):
def __init__(self, scans, subinds, area, cuts=None, name="main", ofmt="{name}", output=True,
ext="fits", pmat_order=None, sys=None, nuisance=False, data=None, extra=[]):
Signal.__init__(self, name, ofmt, output, ext)
self.area = area
self.cuts = cuts
self.dof = dmap.DmapZipper(area)
self.dtype= area.dtype
self.subs = subinds
if data is None:
data = {}
work = area.tile2work()
for scan, subind in zip(scans, subinds):
data[scan] = [pmat.PmatMap(scan, work.maps[subind], order=pmat_order, sys=sys, extra=extra), subind]
self.data = data
def prepare(self, m): return m.tile2work()
def forward(self, scan, tod, work, tmul=1, mmul=1):
if scan not in self.data: return
mat, ind = self.data[scan]
mat.forward(tod, work.maps[ind], tmul=tmul, mmul=mmul)
def backward(self, scan, tod, work, tmul=1, mmul=1):
if scan not in self.data: return
mat, ind = self.data[scan]
mat.backward(tod, work.maps[ind], tmul=tmul, mmul=mmul)
def finish(self, m, work):
m.work2tile(work)
def filter(self, work):
res = []
for w in work.maps:
for filter in self.filters:
w = filter(w)
res.append(w)
return dmap.Workspace(res)
def zeros(self, mat=False):
geom = self.area.geometry
if mat:
geom = geom.copy()
geom.pre = geom.pre + geom.pre[-1:]
return dmap.zeros(geom)
def transform(self, imap, op):
omap = imap.copy()
for dtile in omap.tiles:
dtile[:] = op(dtile)
return omap
def polinv(self, idiv):
odiv = idiv.copy()
for dtile in odiv.tiles:
dtile[:] = array_ops.eigpow(dtile, -1, axes=[-4,-3], lim=config.get("eig_limit"), fallback="scalar")
return odiv
def polmul(self, mat, m, inplace=False):
if not inplace: m = m.copy()
for idtile, mtile in zip(mat.tiles, m.tiles):
mtile[:] = array_ops.matmul(idtile, mtile, axes=[-4,-3])
return m
def work(self): return self.area.geometry.build_work()
def write(self, prefix, tag, m):
if not self.output: return
oname = self.ofmt.format(name=self.name)
oname = "%s%s_%s.%s" % (prefix, oname, tag, self.ext)
merged= config.get("dmap_format") == "merged"
dmap.write_map(oname, m, merged=merged)
class SignalDmapFast(SignalDmap):
def __init__(self, scans, subinds, area, cuts=None, name="main", ofmt="{name}", output=True,
ext="fits", sys=None, nuisance=False, data=None, extra=[]):
if data is None:
data = {}
work = area.tile2work()
for scan, subind in zip(scans, subinds):
data[scan] = [pmat.PmatMapFast(scan, work.maps[subind], sys=sys, extra=extra), subind]
SignalDmap.__init__(self, scans, subinds, area, cuts=cuts, name=name, ofmt=ofmt,
output=output, ext=ext, sys=sys, nuisance=nuisance, data=data)
def precompute(self, scan):
if scan not in self.data: return
mat, ind = self.data[scan]
self.pix, self.phase = mat.get_pix_phase()
def free(self):
del self.pix, self.phase
def forward(self, scan, tod, work):
if scan not in self.data: return
mat, ind = self.data[scan]
mat.forward(tod, work.maps[ind], self.pix, self.phase)
def backward(self, scan, tod, work):
if scan not in self.data: return
mat, ind = self.data[scan]
mat.backward(tod, work.maps[ind], self.pix, self.phase)
class SignalCut(Signal):
def __init__(self, scans, dtype, comm, name="cut", ofmt="{name}_{rank:02}", output=False, cut_type=None, cuts=None, keep=False):
Signal.__init__(self, name, ofmt, output, ext="hdf")
self.data = {}
self.dtype = dtype
cutrange = [0,0]
for si, scan in enumerate(scans):
# Use cuts from scan by default, but allow override
cut = None if cuts is None else cuts[si]
mat = pmat.PmatCut(scan, cut_type, tmul=keep*1.0, cut=cut)
cutrange = [cutrange[1], cutrange[1]+mat.njunk]
self.data[scan] = [mat, cutrange]
self.njunk = cutrange[1]
self.dof = zipper.ArrayZipper(np.zeros(self.njunk, self.dtype), shared=False, comm=comm)
def forward(self, scan, tod, junk):
if scan not in self.data: return
mat, cutrange = self.data[scan]
mat.forward(tod, junk[cutrange[0]:cutrange[1]])
def backward(self, scan, tod, junk):
if scan not in self.data: return
mat, cutrange = self.data[scan]
mat.backward(tod, junk[cutrange[0]:cutrange[1]])
def zeros(self, mat=False): return np.zeros(self.njunk, self.dtype)
def write(self, prefix, tag, m):
if not self.output: return
oname = self.ofmt.format(name=self.name, rank=self.dof.comm.rank)
oname = "%s%s_%s.%s" % (prefix, oname, tag, self.ext)
with h5py.File(oname, "w") as hfile:
hfile["data"] = m
class SignalNoiseRect(Signal):
def __init__(self, scans, area, yspeed, yoffs, comm, cuts=None, mode="keepaz", name="noiserect", ofmt="{name}", output=True,
ext="fits", pmat_order=None, sys=None, nuisance=False, data=None, extra=[]):
Signal.__init__(self, name, ofmt, output, ext)
self.area = area
self.cuts = cuts
self.yspeed = yspeed
self.yoffs = yoffs
self.mode = mode
self.dof = zipper.ArrayZipper(area, comm=comm)
self.dtype= area.dtype
if data is not None:
self.data = data
else:
self.data = {}
for i, scan in enumerate(scans):
self.data[scan] = pmat.PmatNoiseRect(scan, area, yspeed, yoffs[i], mode=mode, extra=extra)
def forward(self, scan, tod, work):
if scan not in self.data: return
self.data[scan].forward(tod, work)
def backward(self, scan, tod, work):
if scan not in self.data: return
self.data[scan].backward(tod, work)
def finish(self, m, work):
self.dof.comm.Allreduce(work, m)
def zeros(self, mat=False):
shape = self.area.shape
if mat: shape = shape[-2:] + shape
return enmap.zeros(shape, self.area.wcs, self.area.dtype)
def polinv(self, m): return array_ops.eigpow(m, -1, axes=[0,1], lim=config.get("eig_limit"), fallback="scalar")
def polmul(self, mat, m, inplace=False):
if not inplace: m = m.copy()
m[:] = array_ops.matmul(mat, m, axes=[0,1])
return m
def write(self, prefix, tag, m):
if not self.output: return
oname = self.ofmt.format(name=self.name)
oname = "%s%s_%s.%s" % (prefix, oname, tag, self.ext)
if self.dof.comm.rank == 0:
enmap.write_map(oname, m)
class PhaseMap:
"""Helper class for SignalPhase. Represents a set of [...,ndet,naz] phase "maps",
one per scanning pattern."""
def __init__(self, patterns, dets, maps):
self.patterns = patterns
self.dets = dets
self.maps = maps
@staticmethod
def read(dirname, rewind=False):
dets = np.loadtxt(dirname + "/dets.txt").astype(int)
patterns = []
maps = []
with open(dirname + "/info.txt", "r") as f:
for line in f:
line = line.strip()
if len(line) == 0 or line.startswith("#"): continue
toks = line.split()
el, az1, az2 = [float(w)*utils.degree for w in toks[:3]]
map = enmap.read_map(dirname + "/" + toks[3])
if rewind:
off = utils.rewind(az1)-az1
az1, az2 = az1+off, az2+off
maz = map.pix2sky([0,0])[1]
off2 = utils.rewind(maz)-maz
map.wcs.wcs.crval[0] += off2/utils.degree
maz2 = map.pix2sky([0,0])[1]
print("pattern rewind off1 %8.3f off2 %8.3f diff %8.3f maz %8.3f maz2 %8.3f" % (off/utils.degree, off2/utils.degree, (off-off2)/utils.degree, maz/utils.degree, maz2/utils.degree))
patterns.append([[el,az1],[el,az2]])
maps.append(map)
return PhaseMap(patterns, dets, maps)
def write(self, dirname, fmt="{pid:02}_{az0:.0f}_{az1:.0f}_{el:.0f}"):
utils.mkdir(dirname)
np.savetxt(dirname + "/dets.txt", self.dets, fmt="%5d")
with open(dirname + "/info.txt", "w") as f:
for pid, (pattern, m) in enumerate(zip(self.patterns, self.maps)):
oname = ("scan_" + fmt + ".fits").format(pid=pid, el=pattern[0,0]/utils.degree, az0=pattern[0,1]/utils.degree, az1=pattern[1,1]/utils.degree)
enmap.write_map(dirname + "/" + oname, m)
f.write("%8.3f %8.3f %8.3f %s\n" % (pattern[0,0]/utils.degree, pattern[0,1]/utils.degree,
pattern[1,1]/utils.degree, oname))
@staticmethod
def zeros(patterns, dets, res=1*utils.arcmin, det_unit=1, hysteresis=True, dtype=np.float64):
maps = []
ndet = len(dets)
for pattern in patterns:
az0, az1 = utils.widen_box(pattern, margin=2*res, relative=False)[:,1]
naz = int(np.ceil((az1-az0)/res))
az1 = az0 + naz*res
shape, wcs = enmap.geometry(pos=[[0,az0],[float(ndet)/det_unit*utils.degree,az1]], shape=(ndet,naz), proj="car")
if hysteresis: shape = (2,)+tuple(shape)
maps.append(enmap.zeros(shape, wcs, dtype=dtype))
return PhaseMap(patterns, dets, maps)
def copy(self):
patterns = np.array(self.patterns)
dets = np.array(self.dets)
maps = [m.copy() for m in self.maps]
return PhaseMap(patterns, dets, maps)
class SignalPhase(Signal):
def __init__(self, scans, areas, pids, comm, cuts=None, name="phase", ofmt="{name}", output=True, ext="fits"):
Signal.__init__(self, name, ofmt, output, ext)
self.areas = areas # template PhaseMaps defining scan patterns and pixelizations
self.pids = pids # index of each scan into scanning patterns
self.comm = comm
self.cuts = cuts
self.dtype = areas.maps[0].dtype
self.data = {}
# Argh! I should hever have switched to string-based detector names!
def get_uids(detnames): return np.char.partition(detnames, "_")[:,2].astype(int)
# Set up the detector mapping for each scan
for pid, scan in zip(pids,scans):
dets = get_uids(scan.dets)
inds = utils.find(areas.dets, dets)
mat = pmat.PmatScan(scan, areas.maps[pid], inds)
self.data[scan] = [pid, mat]
self.dof = zipper.MultiZipper([
zipper.ArrayZipper(map, comm=comm) for map in self.areas.maps],
comm=comm)
def prepare(self, ms):
return [m.copy() for m in ms]
def forward(self, scan, tod, work):
if scan not in self.data: return
pid, mat = self.data[scan]
mat.forward(tod, work[pid])
def backward(self, scan, tod, work):
if scan not in self.data: return
pid, mat = self.data[scan]
mat.backward(tod, work[pid])
def finish(self, ms, work):
for m, w in zip(ms, work):
self.dof.comm.Allreduce(w,m)
def zeros(self, mat=False):
return [area*0 for area in self.areas.maps]
def write(self, prefix, tag, ms):
if not self.output: return
if self.comm.rank != 0: return
omaps = self.areas.copy()
omaps.maps = ms
oname = self.ofmt.format(name=self.name)
oname = "%s%s_%s" % (prefix, oname, tag)
omaps.write(oname)
class SignalMapBuddies(SignalMap):
def __init__(self, scans, area, comm, cuts=None, name="main", ofmt="{name}", output=True, ext="fits", pmat_order=None, sys=None, nuisance=False):
Signal.__init__(self, name, ofmt, output, ext)
self.area = area
self.cuts = cuts
self.dof = zipper.ArrayZipper(area, comm=comm)
self.dtype= area.dtype
self.comm = comm
self.data = {scan: [
pmat.PmatMap(scan, area, order=pmat_order, sys=sys),
pmat.PmatMapMultibeam(scan, area, scan.buddy_offs,
scan.buddy_comps, order=pmat_order, sys=sys)
] for scan in scans}
def forward(self, scan, tod, work):
if scan not in self.data: return
for pmat in self.data[scan]:
pmat.forward(tod, work)
def backward(self, scan, tod, work):
if scan not in self.data: return
for pmat in self.data[scan]:
pmat.backward(tod, work)
# Ugly hack
def get_nobuddy(self):
data = {k:v[0] for k,v in self.data.iteritems()}
return SignalMap(self.data.keys(), self.area, self.comm, cuts=self.cuts, output=False, data=data)
class SignalGain(Signal):
def __init__(self, scans, signals, dtype, comm, model=None, name="gain", ofmt="{name}_{rank:02}", output=True):
Signal.__init__(self, name, ofmt, output, ext="txt")
self.data = {}
self.dtype = dtype
self.signals = signals
# The sky model to assume. Must be compatible with the signals
# passed in. This can be freely updated by assigning to .model
self.model = model
# Calculate the gain degrees of freedom mapping
self.data = {}
i1 = 0
for si, scan in enumerate(scans):
i2 = i1+scan.ndet
self.data[scan] = (i1,i2)
i1 = i2
self.ntot = i1
self.dof = zipper.ArrayZipper(np.zeros(self.ntot, self.dtype), shared=False, comm=comm)
def precompute(self, scan):
self.profile = np.zeros((scan.ndet,scan.nsamp),self.dtype)
self.signals.forward(scan, self.profile, self.model)
def free(self):
self.profile = None
def forward(self, scan, tod, gains):
if scan not in self.data: return
i1,i2 = self.data[scan]
tod += self.profile*gains[i1:i2,None]
def backward(self, scan, tod, gains):
if scan not in self.data: return
i1,i2 = self.data[scan]
gains[i1:i2] += np.sum(self.profile*tod,1)
def zeros(self, mat=False): return np.zeros(self.ntot, self.dtype)
def write(self, prefix, tag, m):
if not self.output: return
oname = self.ofmt.format(name=self.name, rank=self.dof.comm.rank)
oname = "%s%s_%s.%s" % (prefix, oname, tag, self.ext)
# We want to output the tods in sorted order
scans = [scan for scan in self.data]
ids = [scan.id for scan in scans]
order = np.argsort(ids)
from enact import actdata
with open(oname, "w") as ofile:
for ind in order:
scan = scans[ind]
i1, i2= self.data[scan]
gvals = m[i1:i2]
dets = actdata.split_detname(scan.dets)[1]
msg = [scan.id]
for det,gval in zip(dets, gvals):
msg.append("%s:%.3f" % (det,gval))
msg = " ".join(msg)
ofile.write(msg + "\n")
# Special source handling v2:
# Solve for the map and the local model errors around strong sources jointly. The
# equation system will be degenerate, so the map will not contain a complete solution
# in these regions. So before outputting, we need do do
# tod = Pmap map + Psrcsamp srcssamps
# totmap = map_white(tod)
# So this requires a new signal that's very similar to our cut signal, as well as a
# postprocessing operation that adds up the values in the src pixels.
class SignalSrcSamp(SignalCut):
def __init__(self, scans, dtype, comm, mask=None, name="srcsamp", ofmt="{name}_{rank:02}", output=False, cut_type="full", sys="cel", keep=True):
Signal.__init__(self, name, ofmt, output, ext="hdf")
self.data = {}
self.dtype = dtype
self.keep = keep
self.sys = sys
self.cut_type = cut_type
cutrange = [0,0]
self.mask = mask
for scan in scans:
# Project our sample density into time domain
tod = np.zeros((scan.ndet,scan.nsamp), np.float32)
pmap = pmat.PmatMap(scan, mask, sys=sys, order=1)
pmap.forward(tod, mask)
tod[:] = tod > 0.5
# Accumulate density and use this to define a sample mask
cumsamps = np.concatenate([[0],utils.floor(np.cumsum(tod))])
# minimum shouldn't be necessary, but present to guard against floating point issues
tod_mask = np.minimum(cumsamps[1:]-cumsamps[:-1],1).reshape(tod.shape)
del cumsamps, tod
cut = sampcut.from_mask(tod_mask); del tod_mask
mat = pmat.PmatCut(scan, cut_type, tmul=keep*1.0, cut=cut)
cutrange = [cutrange[1], cutrange[1]+mat.njunk]
self.data[scan] = [mat, cutrange]
self.njunk = cutrange[1]
self.dof = zipper.ArrayZipper(np.zeros(self.njunk, self.dtype), shared=False, comm=comm)
class SignalTemplate(Signal):
def __init__(self, scans, template, comm, name="template", ofmt="{name}", output=True,
ext="txt", pmat_order=None, sys=None, extra=[]):
Signal.__init__(self, name, ofmt, output, ext)
self.template = template
self.dof = zipper.ArrayZipper(np.zeros(len(scans), dtype=template.dtype), shared=False, comm=comm)
self.dtype= template.dtype
self.data = {scan: pmat.PmatMap(scan, template, order=pmat_order, sys=sys, extra=extra) for scan in scans}
self.inds = {scan:i for i, scan in enumerate(scans)}
self.scans= scans
self.n = len(scans)
def forward(self, scan, tod, work, tmul=1, mmul=1):
if scan not in self.data: return
self.data[scan].forward(tod, self.template*work[self.inds[scan]], tmul=tmul, mmul=mmul)
def backward(self, scan, tod, work, tmul=1, mmul=1):
if scan not in self.data: return
Ptd = self.template*0
self.data[scan].backward(tod, Ptd, tmul=tmul)
work[self.inds[scan]] = work[self.inds[scan]]*mmul + np.sum(Ptd)
def zeros(self, mat=False): return np.zeros(self.n, self.dtype)
def write(self, prefix, tag, m):
if not self.output: return
oname = self.ofmt.format(name=self.name, rank=self.dof.comm.rank)
oname = "%s%s_%s.%s" % (prefix, oname, tag, self.ext)
with open(oname, "w") as ofile:
for i in range(self.n):
ofile.write("%s %8.3f\n" % (self.scans[i].id, m[i]))
class Signals(Signal):
"""Class for handling multiple signals jointly."""
def __init__(self, signals, comm, name="signals", ofmt="{name}_{rank:02}", ext="hdf", output=False):
Signal.__init__(self, name, ofmt, output, ext)
self.signals = signals
self.dof = zipper.MultiZipper([signal.dof for signal in signals], comm=comm)
self.precon = PreconSignals(signals)
def zeros(self, mat=False): return [signal.zeros(mat=mat) for signal in self.signals]
def prepare(self, maps): return [signal.prepare(map) for signal,map in zip(self.signals, maps)]
def filter(self, maps): return [signal.filter(map) for signal,map in zip(self.signals, maps)]
def precompute(self, scan):
for signal in self.signals:
signal.precompute(scan)
def forward (self, scan, tod, works):
for signal, work in list(zip(self.signals, works))[::-1]:
signal.forward(scan, tod, work)
def backward(self, scan, tod, works):
for signal, work in zip(self.signals, works):
signal.backward(scan, tod, work)
def free(self):
for signal in self.signals:
signal.free()
def finish (self, maps, works):
for signal, map, work in zip(self.signals, maps, works):
signal.finish(map, work)
def transform(self, maps, op): return [signal.transform(map, op) for signal,map in zip(self.signals, maps)]
def polinv(self, maps): return [signal.polinv(map) for signal,map in zip(self.signals, maps)]
def polmul(self, mats, maps, inplace=False): return [signal.polmul(mat, map, inplace=inplace) for signal,mat,map in zip(self.signals, mats, maps)]
# This one is a potentially cheaper version of self.prepare(self.zeros())
def write (self, prefix, tag, maps):
for signal, map in zip(self.signals, maps):
signal.write(prefix, tag, map)
def postprocess(self, maps): return [signal.postprocess(map) for signal,map in zip(self.signals, maps)]
def prior(self, scans, xins, xouts):
for signal, scans, xin, xout in zip(self.signals, xins, xouts):
signal.prior(scans, xin, xout)
######## Preconditioners ########
# Preconditioners have a lot of overlap with Eqsys.A. That's not
# really surprising, as their job is to approximate A, but it does
# lead to duplicate code. The preconditioners differ by:
# 1. Only invovling [signal,cut]
# 2. Looping over components of the signal (if it has components)
# We can do this via a loop of the type
# for each compmap:
# eqsys.A(compmap)
# e.g. one sends in a reduced equation system to the preconditioner:
# PreconMapBinned(eqsys)
# instead of
# PreconMapBinned(signal, signal_cut, scans, weights)
# But th eqsys above would not be the real eqsys. So wouldn't it really be
# PreconMapBinned(Eqsys(scans, [signal, signal_cut], weights=weights))
# which is pretty long. The internal eqsys use is an implementation
# detail. Pass the arguments we need, and use an eqsys internally if
# that makes things simpler. Eqsyses are cheap to construct.
class PreconNull:
def __init__(self, delayed=False): pass
def build(self, scans): pass
def __call__(self, m): pass
def write(self, prefix="", ext="fits"): pass
config.default("eig_limit", 1e-3, "Smallest relative eigenvalue to invert in eigenvalue inversion. Ones smaller than this are set to zero.")
# The binned preconditioners use a 3x3 covmat per pixel. This matrix is sometimes
# poorly conditioned, and can't be inverted. I currently use eigenvalue inversion
# for this, which works by setting bad eigenvalues to zero. But is this sufficient?
# Consider a pixel that is hit by a single sample with phase [1,1,0], resulting in
# the covmat [110,110,000]. This has a single nonzero eigenvalue, and after
# eigenvalue inversion it will still be proportional to [110,110,000]. This preconditioner
# forces U to be zero, but it allows T,Q to float. This means that we can't distinguish
# [0,0,0] from [1,-1,0] or [1e9,-1e9,0], which can lead to runaway pixel values.
# To avoid this problem we can make the poorly conditioned pixels T-only by setting
# all but [0,0] of icov to 0.
#
# I've moved to this workaround, even though it appears to perform marginally worse
# at for my one-tod 10-cg test. The main thing that helps in both cases is to use an
# eig_limit like 1e-3 instead of 1e-6, which I used before
# Preconditioners now have a build(self, scans) method that does the actual building.
# To keep backward compatibility this is called automatically in the constructor
# unless delayed=True is passed (ideally this would default to True, but that would
# break things, so I don't do that right away). build() should be safe to call multiple
# times, which would update the preconditioners based on the scans without leaving it
# in an inconsistent state
class PreconMapBinned:
def __init__(self, signal, scans, weights, noise=True, hits=True, mask=None, delayed=False):
"""Binned preconditioner: (P'W"P)", where W" is a white
nosie approximation of N". If noise=False, instead computes
(P'P)". If hits=True, also computes a hitcount map."""
self.signal = signal
self.weights= weights
self.noise = noise
self.mask = mask
self.calc_hits = hits
# Backwards compatibility
if not delayed: self.build(scans)
def build(self, scans):
self.div = self.signal.zeros(mat=True)
calc_div_map(self.div, self.signal, scans, self.weights, noise=self.noise)
if self.mask is not None: self.div *= self.mask
self.idiv = array_ops.eigpow(self.div, -1, axes=[-4,-3], lim=config.get("eig_limit"), fallback="scalar")
if self.calc_hits:
# Build hitcount map too
self.hits = self.signal.area.copy()
self.hits = calc_hits_map(self.hits, self.signal, scans)
if self.mask is not None: self.hits *= self.mask
else: self.hits = None
def __call__(self, m):
m[:] = array_ops.matmul(self.idiv, m, axes=[-4,-3])
def write(self, prefix):
self.signal.write(prefix, "div", self.div)
if self.hits is not None:
self.signal.write(prefix, "hits", self.hits)
class PreconDmapBinned:
def __init__(self, signal, scans, weights, noise=True, hits=True, delayed=False):
"""Binned preconditioner: (P'W"P)", where W" is a white
nosie approximation of N". If noise=False, instead computes
(P'P)". If hits=True, also computes a hitcount map."""
self.signal = signal
self.weights = weights
self.noise = noise
self.calc_hits = hits
if not delayed: self.build(scans)
def build(self, scans):
self.div = self.signal.zeros(mat=True)
calc_div_map(self.div, self.signal, scans, self.weights, noise=self.noise)
self.idiv = self.div.copy()
for dtile in self.idiv.tiles:
dtile[:] = array_ops.eigpow(dtile, -1, axes=[-4,-3], lim=config.get("eig_limit"), fallback="scalar")
if self.calc_hits:
# Build hitcount map too
self.hits = self.signal.area.copy()
self.hits = calc_hits_map(self.hits, self.signal, scans)
else: self.hits = None
def __call__(self, m):
for idtile, mtile in zip(self.idiv.tiles, m.tiles):
mtile[:] = array_ops.matmul(idtile, mtile, axes=[-4,-3])
def write(self, prefix):
self.signal.write(prefix, "div", self.div)
if self.hits is not None:
self.signal.write(prefix, "hits", self.hits)
class PreconMapHitcount:
def __init__(self, signal, scans, delayed=False):
"""Hitcount preconditioner: (P'P)_TT."""
self.signal = signal
if not delayed: self.build(scans)
def build(self, scans):
self.hits = self.signal.area.copy()
self.hits = calc_hits_map(self.hits, self.signal, scans)
self.ihits= 1.0/np.maximum(self.hits,1)
def __call__(self, m):
m *= self.ihits[...,None,:,:]
def write(self, prefix):
self.signal.write(prefix, "hits", self.hits)
class PreconDmapHitcount:
def __init__(self, signal, scans, delayed=False):
"""Hitcount preconditioner: (P'P)_TT"""
self.signal = signal
if not delayed: self.build(scans)
def build(self, scans):
self.hits = self.signal.area.copy()
self.hits = calc_hits_map(self.hits, self.signal, scans)
def __call__(self, m):
for htile, mtile in zip(self.hits.tiles, m.tiles):
htile = np.maximum(htile, 1)
mtile /= htile[...,None,:,:]
def write(self, prefix):
self.signal.write(prefix, "hits", self.hits)
class PreconMapTod:
def __init__(self, signal, scans, weights, delayed=False):
"""Preconditioner based on inverting the tod noise model independently per TOD,
along with a P'P inversion: precon = P"'NP", where P" = P(P'P)" """
self.signal = signal
self.weights = weights
self.scans = scans
self.maxnoise = 10
if not delayed: self.build(scans)
def build(self, scans):
# First get P'P, which is the standard div but with no noise model applied
ncomp = self.signal.area.shape[0]
self.ptp = enmap.zeros((ncomp,)+self.signal.area.shape, self.signal.area.wcs, self.signal.area.dtype)
calc_div_map(self.ptp, self.signal, scans, self.weights, noise=False)
self.iptp = array_ops.eigpow(self.ptp, -1, axes=[0,1], lim=config.get("eig_limit"))
self.hits = self.ptp[0,0].astype(np.int32)
def __call__(self, m):
# This is pretty slow. Let's see if it is worth it. I fear the discontinuities
# will be just as slow to resolve
m[:] = array_ops.matmul(self.iptp, m, axes=[0,1])
omap = m*0
for scan in self.scans:
tod = np.zeros([scan.ndet, scan.nsamp], self.signal.dtype)
self.signal.precompute(scan)
self.signal.forward (scan, tod, m)
for weight in self.weights: weight(scan, tod)
noise_ref = np.median(scan.noise.D)*self.maxnoise
scan.noise.D = np.minimum(scan.noise.D, noise_ref)
scan.noise.E = np.minimum(scan.noise.E, noise_ref)
scan.noise.apply(tod, inverse=True)
for weight in self.weights[::-1]: weight(scan, tod)
sampcut.gapfill_const(scan.cut, tod, inplace=True)
self.signal.backward(scan, tod, omap)
self.signal.free()
m[:] = 0
self.signal.finish(m, omap)
m[:] = array_ops.matmul(self.iptp, m, axes=[0,1])
return m
def write(self, prefix):
self.signal.write(prefix, "ptp", self.ptp)
self.signal.write(prefix, "hits", self.hits)
class PreconCut:
def __init__(self, signal, scans, weights=[], delayed=False):
self.signal = signal
self.weights = weights
if not delayed: self.build(scans)
def build(self, scans):
# I'm ignoring weights in this preconditioner after the performance
# invesstigation below:
# With 200 steps max:
# full+white+nowin: super 1e-40 @ CG 3
# full+white+igwin: good 1e-12
# full+white+win: super 1e-35 @ CG 9
# full+jon +nowin: medium 1e-7
# full+jon +igwin: medium 1e-7
# full+jon +win: good 1e-11
# poly+white+nowin: super 1e-32 @ CG 16
# poly+white+igwin: good 1e-11
# poly+white+win: great 1e-31 @ CG 162
# poly+jon +nowin: medium 1e-8
# poly+jon +igwin: medium 1e-8
# poly+jon +win: bad 1e-4
# So overall that's bad is anything with poly+win.
# The best would be to not do windowing. I'll consider that. But if one
# has windowing, ignoring it in the preconditioner seems best. The issue
# is that windowing breaks the orthogonality of the legendre polynomials,
# and the preconditioner relies on that orthogonality. Ignoring the windowing
# keeps the preconditioner self-consistent even though it's no longer fully
# consistent with the actual data.
junk = self.signal.zeros()
iwork = self.signal.prepare(junk)
owork = self.signal.prepare(junk)
iwork[:] = 1
for scan in scans:
with bench.mark("div_" + self.signal.name):
tod = np.zeros((scan.ndet, scan.nsamp), iwork.dtype)
self.signal.forward (scan, tod, iwork)
#for weight in self.weights: weight(scan, tod)
scan.noise.white(tod)
#for weight in self.weights[::-1]: weight(scan, tod)
self.signal.backward(scan, tod, owork)
times = [bench.stats[s]["time"].last for s in ["div_"+self.signal.name]]
L.debug("div %s %6.3f %s" % ((self.signal.name,)+tuple(times)+(scan.id,)))
self.signal.finish(junk, owork)
junk = np.abs(junk)
# Avoid negative numbers or division by zero. Should not be necessary
# if we ignore windowing, but doesn't hurt much.
ref = np.max(junk)/1e7
self.idiv = 1/np.maximum(junk,ref)
def __call__(self, m):
m *= self.idiv
def write(self, prefix):
self.signal.write(prefix, "idiv", self.idiv)
class PreconPhaseBinned:
def __init__(self, signal, scans, weights, delayed=False):
self.signal = signal
self.weights = weights
if not delayed: self.build(scans)
def build(self, scans):
div = [area*0+1 for area in self.signal.areas.maps]
# The cut samples are included here becuase they must be avoided, but the
# actual computation of the junk sample preconditioner happens elsewhere.
# This is a bit redundant, but should not cost much time since this only happens
# in the beginning.
iwork = self.signal.prepare(div)
owork = self.signal.prepare(self.signal.zeros())
prec_div_helper(self.signal, scans, self.weights, iwork, owork)
self.signal.finish(div, owork)
hits = self.signal.zeros()
owork = self.signal.prepare(hits)
for scan in scans:
tod = np.full((scan.ndet, scan.nsamp), 1, self.signal.dtype)
self.signal.backward(scan, tod, owork)
self.signal.finish(hits, owork)
for i in range(len(hits)):
hits[i] = hits[i].astype(np.int32)
self.div = div
self.hits = hits
def __call__(self, ms):
for d, m in zip(self.div, ms):
m[d!=0] /= d[d!=0]
def write(self, prefix):
self.signal.write(prefix, "div", self.div)
class PreconSignals:
def __init__(self, signals, delayed=False):
self.signals = signals
# If we're not delayed, then our signals' precons will already
# have been initialized, so don't do anything in either case
def build(self, scans):
for signal in self.signals:
signal.precon.build(scans)
def __call__(self, maps):
for signal, map in zip(self.signals, maps):
signal.precon(map)
def write(self, prefix="", ext=None):
for signal in self.signals:
signal.write(prefix=prefix)
def prec_div_helper(signal, scans, weights, iwork, owork, cuts=None, noise=True):
# The argument list of this one is so long that it almost doesn't save any
# code.
if cuts is None: cuts = [scan.cut for scan in scans]
for si, scan in enumerate(scans):
with bench.mark("div_Pr_" + signal.name):
signal.precompute(scan)
with bench.mark("div_P_" + signal.name):
tod = np.zeros((scan.ndet, scan.nsamp), signal.dtype)
signal.forward(scan, tod, iwork)
#signal_cut.forward (scan, tod, ijunk)
with bench.mark("div_white"):
for weight in weights: weight(scan, tod)
if noise: scan.noise.white(tod)
for weight in weights[::-1]: weight(scan, tod)
with bench.mark("div_PT_" + signal.name):
#signal_cut.backward(scan, tod, ojunk)
sampcut.gapfill_const(cuts[si], tod, inplace=True)
signal.backward(scan, tod, owork)
with bench.mark("div_Fr_" + signal.name):
signal.free()
times = [bench.stats[s]["time"].last for s in ["div_P_"+signal.name, "div_white", "div_PT_" + signal.name]]
L.debug("div %s %6.3f %6.3f %6.3f %s" % ((signal.name,)+tuple(times)+(scan.id,)))
def calc_div_map(div, signal, scans, weights, cuts=None, noise=True):
# The cut samples are included here becuase they must be avoided, but the
# actual computation of the junk sample preconditioner happens elsewhere.
# This is a bit redundant, but should not cost much time since this only happens
# in the beginning.
for i in range(div.shape[-4]):
div[...,i,i,:,:] = 1
iwork = signal.prepare(div[...,i,:,:,:])
owork = signal.prepare(signal.zeros())
prec_div_helper(signal, scans, weights, iwork, owork, cuts=cuts, noise=noise)
signal.finish(div[...,i,:,:,:], owork)
def calc_crosslink_map(signal, scans, weights, cuts=None, noise=True):
saved_comps = [scan.comps.copy() for scan in scans]
if cuts is None: cuts = [scan.cut for scan in scans]
for scan in scans: scan.comps[:] = np.array([1,1,0])
cmap = signal.zeros()
owork = signal.prepare(cmap)
for si, scan in enumerate(scans):
with bench.mark("cmap_Pr_" + signal.name): signal.precompute(scan)
with bench.mark("cmap_tod"):
tod = np.full((scan.ndet, scan.nsamp), 1.0, signal.dtype)
with bench.mark("cmap_white"):
for weight in weights: weight(scan, tod)
if noise: scan.noise.white(tod)
for weight in weights: weight(scan, tod)
with bench.mark("cmap_PT_" + signal.name):
#signal_cut.backward(scan, tod, ojunk)
sampcut.gapfill_const(cuts[si], tod, inplace=True)
signal.backward(scan, tod, owork)
signal.free()
times = [bench.stats[s]["time"].last for s in ["cmap_white", "cmap_PT_" + signal.name]]
L.debug("crossmap %s %6.3f %6.3f %s" % ((signal.name,)+tuple(times)+(scan.id,)))
signal.finish(cmap, owork)
# Restore saved components
for scan, comp in zip(scans, saved_comps): scan.comps = comp
return cmap
def calc_ptsrc_map(signal, scans, src_filters):
# First compute P'W"srcs
cmap = signal.zeros()
owork = signal.prepare(cmap)
for scan in scans:
tod = np.zeros((scan.ndet, scan.nsamp), signal.dtype)
with bench.mark("srcmap_srcs"):
for src_filter in src_filters:
src_filter(scan, tod)
with bench.mark("srcmap_PT_" + signal.name):
sampcut.gapfill_const(scan.cut, tod, inplace=True)
signal.backward(scan, tod, owork)
signal.free()
times = [bench.stats[s]["time"].last for s in ["srcmap_srcs", "srcmap_PT_" + signal.name]]
L.debug("srcmap %s %6.3f %6.3f %s" % ((signal.name,)+tuple(times)+(scan.id,)))
signal.finish(cmap, owork)
# Then divide out the hits
with utils.nowarn():
cmap /= signal.precon.hits
cmap.fillbad(0, inplace=True)
# We want to know what the source model is, which is the opposite of what was
# subtracted
cmap *= -1
return cmap
def calc_icov_map(signal, scans, pos, weights):
"""Compute a map containing the inverse covariance structure around the set
of positions pos[:,{y,x}] in pixels. This will be computed in one
operation, so if the points are too close to each other their covariance
information will interfere with each other."""
icov = signal.zeros()
ocov = icov.copy()
# Set the chosen positions to one. This is very inelegant. Should have
# global pixel setting methods in dmap
pos = np.zeros([1,2],int)+pos # broadcast to correct shape
if isinstance(icov, enmap.ndmap):
icov[0,pos[:,0],pos[:,1]] = 1
elif isinstance(icov, dmap.Dmap):
for tile, ind in zip(icov.tiles, icov.loc_inds):
pbox = icov.geometry.tile_boxes[ind]
for pix in pos:
y,x = pix - pbox[0]
if y >= 0 and x >= 0 and y < tile.shape[-2] and x < tile.shape[-1]:
tile[0,y,x] = 1
# The rest proceeds similarly to the crosslinking map
iwork = signal.prepare(icov)
owork = signal.prepare(ocov)
for scan in scans:
with bench.mark("icov_Pr_" + signal.name): signal.precompute(scan)
with bench.mark("icov_P_" + signal.name):
tod = np.zeros((scan.ndet, scan.nsamp), signal.dtype)
signal.forward(scan, tod, iwork)
with bench.mark("icov_nmat"):
for weight in weights: weight(scan, tod)
scan.noise.apply(tod)
for weight in weights: weight(scan, tod)
with bench.mark("icov_PT_" + signal.name):
sampcut.gapfill_const(scan.cut, tod, inplace=True)
signal.backward(scan, tod, owork)
with bench.mark("icov_Fr_" + signal.name): signal.free()
times = [bench.stats[s]["time"].last for s in ["icov_P_" + signal.name, "icov_nmat", "icov_PT_" + signal.name]]
L.debug("icov %s %6.3f %6.3f %6.3f %s" % ((signal.name,)+tuple(times)+(scan.id,)))
signal.finish(ocov, owork)
return ocov[0]
def calc_hits_map(hits, signal, scans, cuts=None):
hits = hits*0
work = signal.prepare(hits)
if cuts is None: cuts = [scan.cut for scan in scans]
for si, scan in enumerate(scans):
with bench.mark("hits_Pr_" + signal.name):
signal.precompute(scan)