-
Notifications
You must be signed in to change notification settings - Fork 7
/
auto_coadd.py
93 lines (89 loc) · 4.43 KB
/
auto_coadd.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
import numpy as np, argparse, os, imp, time
from scipy import ndimage
from enlib import enmap, retile, utils, jointmap, bunch, cg, fft, mpi
parser = argparse.ArgumentParser()
parser.add_argument("config")
parser.add_argument("sel")
parser.add_argument("template", nargs="?")
parser.add_argument("odir")
# Let's use a 4x4 degree test patch with 1 degree of padding on each side and
# another 1 degree of apodization
parser.add_argument("-b", "--box", type=str, default="-5:-1,36:32")
parser.add_argument("-t", "--tilesize", type=int, default=480)
parser.add_argument("-p", "--pad", type=int, default=240)
parser.add_argument("-a", "--apod-val", type=float, default=2e-1)
parser.add_argument("-A", "--apod-alpha", type=float, default=5)
parser.add_argument("-E", "--apod-edge", type=float, default=120)
parser.add_argument( "--kxrad", type=float, default=90)
parser.add_argument( "--kx-ymax-scale", type=float, default=1)
parser.add_argument( "--highpass", type=float, default=200)
parser.add_argument( "--cg-tol", type=float, default=1e-6)
parser.add_argument( "--max-ps", type=float, default=0)
parser.add_argument("-M", "--mode", type=str, default="weight")
parser.add_argument("-c", "--cont", action="store_true")
parser.add_argument("-v", "--verbose", action="store_true")
parser.add_argument("-i", "--maxiter", type=int, default=100)
parser.add_argument( "--ncomp", type=int, default=3)
parser.add_argument("-s", "--slice", type=str, default=None)
args = parser.parse_args()
# Why did I use this?
enmap.extent_model.append("intermediate")
dtype = np.float64
comm = mpi.COMM_WORLD
ncomp = args.ncomp
utils.mkdir(args.odir)
config = jointmap.read_config(args.config)
datasets = jointmap.get_datasets(config, args.sel)
# Get the reference beam, which is the highest beam at each l
ref_beam = datasets[0].beam
for dataset in datasets[1:]:
ref_beam = np.maximum(ref_beam, dataset.beam)
if args.template is None:
# Process a single box
box = np.array([[float(w) for w in fromto.split(":")] for fromto in args.box.split(",")]).T*utils.degree
utils.mkdir(args.odir)
coadder = jointmap.AutoCoadder.read(datasets, box, pad=args.pad, verbose=True, dtype=dtype,
cache_dir=args.odir, read_cache=args.cont)
coadder.analyze(ref_beam, mode=args.mode,
apod_val=args.apod_val, apod_alpha=args.apod_alpha, apod_edge=args.apod_edge,
filter_kxrad=args.kxrad, filter_kx_ymax_scale=args.kx_ymax_scale, filter_highpass=args.highpass)
if args.slice: coadder.set_slice(args.slice)
coadder.calc_precon()
coadder.calc_rhs()
res = coadder.solve(verbose=True, cg_tol=args.cg_tol, dump_dir=args.odir, maxiter=args.maxiter)
if res is None: print "No data found"
else:
enmap.write_map(args.odir + "/map.fits", res.map)
enmap.write_map(args.odir + "/div.fits", res.div)
else:
# We will loop over tiles in the area defined by template
shape, wcs = jointmap.read_geometry(args.template)
pre = (ncomp,)
shape = pre + shape
tshape = np.array([args.tilesize,args.tilesize])
ntile = np.floor((shape[-2:]+tshape-1)/tshape).astype(int)
tyx = [(y,x) for y in range(ntile[0]-1,-1,-1) for x in range(ntile[1])]
for i in range(comm.rank, len(tyx), comm.size):
y, x = tyx[i]
if args.cont and os.path.isfile(args.odir + "/map_padtile%(y)03d_%(x)03d.fits" % {"y":y,"x":x}):
print "%3d skipping %3d %3d" % (comm.rank, y, x)
continue
print "%3d processing %3d %3d" % (comm.rank, y, x)
tpos = np.array(tyx[i])
pbox = np.array([tpos*tshape,np.minimum((tpos+1)*tshape,shape[-2:])])
box = enmap.pix2sky(shape, wcs, pbox.T).T
# Set up the coadder
coadder = jointmap.AutoCoadder.read(datasets, box, pad=args.pad, verbose=True, dtype=dtype, ncomp=ncomp)
if coadder is not None:
coadder.analyze(ref_beam, mode=args.mode,
apod_val=args.apod_val, apod_alpha=args.apod_alpha, apod_edge=args.apod_edge,
filter_kxrad=args.kxrad, filter_kx_ymax_scale=args.kx_ymax_scale, filter_highpass=args.highpass)
if args.slice: coadder.set_slice(args.slice)
coadder.calc_precon()
coadder.calc_rhs()
res = coadder.solve(verbose=True, cg_tol=args.cg_tol, maxiter=args.maxiter,
dump_dir = args.odir if comm.rank==0 else None)
if coadder is None or res is None:
res = jointmap.make_dummy_tile(shape, wcs, box, pad=args.pad)
enmap.write_map(args.odir + "/map_padtile%(y)03d_%(x)03d.fits" % {"y":y,"x":x}, res.map)
enmap.write_map(args.odir + "/div_padtile%(y)03d_%(x)03d.fits" % {"y":y,"x":x}, res.div)