-
Notifications
You must be signed in to change notification settings - Fork 0
/
plot_maps_all.py
executable file
·88 lines (68 loc) · 2.27 KB
/
plot_maps_all.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
#!/usr/bin/env python3
import argparse
import numpy as np
import multiprocessing
import time
import utils
import plot_maps
def add_arguments(parser):
parser.add_argument("-v", "--verbose", action="store_true", help="verbose mode")
parser.add_argument("-r", "--root", type=utils.dir_path, default="out", help="root directory")
parser.add_argument(
"-n",
"--ncpu",
type=int,
default=4,
help="number of CPUs to use (default: 4)",
)
def process(run):
# start timer
tic = time.perf_counter()
# check if directory is complete
if not utils.is_complete(run):
return run, None
# create directory if needed
plotdir = run / "plots"
plotdir.mkdir(exist_ok=True)
# get last ref
ref = utils.get_last_ref(run)
# read data
maps = utils.read_maps(run, ref=ref)
hits, cond = utils.read_hits_cond(run, ref=ref)
residuals = utils.read_residuals(run, ref=ref)
sky_in = utils.read_input_sky()
# define a mask for pixels outside the solved patch
mask = hits < 1
for m in maps.values():
m[mask] = np.nan
cond[mask] = np.nan
plot_maps.plot_hits_cond(hits, cond, plotdir)
plot_maps.plot_res_hist(maps, sky_in, plotdir)
plot_maps.plot_maps(maps, sky_in, plotdir, map_range_P=10)
plot_maps.plot_residuals(residuals, plotdir)
elapsed = time.perf_counter() - tic
return run, elapsed
def main(args):
runs = list(utils.get_all_runs(args.root))
if len(runs) == 0:
return
if args.ncpu > 0:
ncpu = args.ncpu
else:
ncpu = multiprocessing.cpu_count()
# Don't use more CPUs than runs to process
ncpu = min(ncpu, len(runs))
with multiprocessing.Pool(processes=ncpu) as pool:
if args.verbose:
print(f"Using {ncpu} CPU")
for run, elapsed in pool.imap_unordered(process, runs):
if elapsed is None:
print(f"Could not plot maps for '{run}' (missing files)")
continue
if args.verbose:
print(f"Processed '{run}' in {elapsed:.3f} seconds")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Produce plots of output maps for all runs.")
add_arguments(parser)
args = parser.parse_args()
main(args)