-
Notifications
You must be signed in to change notification settings - Fork 7
/
build_todinfo.py
50 lines (44 loc) · 1.74 KB
/
build_todinfo.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
from __future__ import division, print_function
import numpy as np, os, h5py
from enlib import config, mpi, coordinates, utils, errors, tagdb
from enact import filedb, actdata, todinfo
parser = config.ArgumentParser(os.environ["HOME"] + "/.enkirc")
parser.add_argument("tagfile")
parser.add_argument("sel", nargs="?", default="")
parser.add_argument("ofile")
args = parser.parse_args()
file_db = filedb.setup_filedb()
scan_db = todinfo.read(args.tagfile, vars={"root":config.get("root")})
comm = mpi.COMM_WORLD
# We want to process *all* tods, not just selected ones. Could also
# achieve this by adding /all to the selector.
scan_db = scan_db.select(scan_db.query(args.sel, apply_default_query=False))
ids = scan_db.ids
stats = []
for ind in range(comm.rank, len(ids), comm.size):
id = ids[ind]
entry = file_db[id]
try:
stats.append(todinfo.build_tod_stats(entry))
except (errors.DataMissing,AttributeError) as e:
print("%3d %4d/%d %5.1f%% Skipping %s (%s)" % (comm.rank, ind+1, len(ids), (ind+1)/float(len(ids))*100, id, str(e)))
continue
print("%3d %4d/%d %5.1f%% %s" % (comm.rank, ind+1, len(ids), (ind+1)/float(len(ids))*100, id))
stats = todinfo.merge_tod_stats(stats)
if comm.rank == 0: print("Reducing")
comm.Barrier()
for key in stats:
stats[key] = utils.allgatherv(stats[key],comm)
# Sort by id and move id index last
inds = np.argsort(stats["id"])
for key in stats:
stats[key] = utils.moveaxis(stats[key][inds],0,-1)
stat_db = todinfo.Todinfo(stats)
# Merge with original tags. Rightmost overrides for overridable fields.
# For tags, we get the union. This means that stat_db can't override incorrect
# tags in scan_db, just add to them.
stat_db = scan_db + stat_db
if comm.rank == 0:
print("Writing")
stat_db.write(args.ofile)
print("Done")