-
Notifications
You must be signed in to change notification settings - Fork 7
/
submit_map
executable file
·116 lines (107 loc) · 4.88 KB
/
submit_map
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
#!/usr/bin/env python
# Script for quickly submitting tod2map runs
import argparse, os, subprocess, re, sys
parser = argparse.ArgumentParser()
parser.add_argument("sel")
parser.add_argument("odir")
parser.add_argument("-n", "--nsplit", type=int, default=4)
parser.add_argument("-p", "--patch", type=str, default=None)
parser.add_argument("-N", "--nnode", type=int, default=1)
parser.add_argument("-T", "--tblock", type=float, default=4)
parser.add_argument("-d", "--dry-run", action="store_true")
parser.add_argument("-s", "--sky", type=str, default="sky")
parser.add_argument( "--sys", type=str, default=None)
parser.add_argument("-t", "--time", type=str, default="24:00:00")
parser.add_argument("--split-mode", type=str, default="jon")
parser.add_argument("--no-prune", action="store_true")
args, unknown = parser.parse_known_args()
from enlib import utils, config, enmap
config.init()
root = config.get("root")
utils.mkdir(args.odir)
rundir = args.odir + "/runs"
utils.mkdir(rundir)
array_freqs = {"pa1":["f150"], "pa2":["f150"], "pa3":["f090","f150"],
"pa4":["f150","f220"], "pa5":["f090","f150"], "pa6":["f090","f150"]}
season_arrays = {"s13":["pa1"], "s14":["pa1","pa2"], "s15":["pa1","pa2","pa3"],
"s16":["pa2","pa3","pa4"], "s17":["pa4","pa5","pa6"], "s18":["pa4","pa5","pa6"],
"s19":["pa4","pa5","pa6"], "s20":["pa4","pa5","pa7"], "s21":["pa4","pa5","pa7"]}
tag_map = {"+":"", "~":"no", ">":"_gt_", "<":"_lt_"}
def remap(toks, tag_map):
toks = list(toks)
for key in tag_map:
for i, tok in enumerate(toks):
toks[i] = tok.replace(key, tag_map[key])
return toks
# Sel format is normal, except there can be :-separated parts, which will
# be iterated over
toks = utils.split_outside(args.sel,",")
alts = [utils.split_outside(tok,":") for tok in toks]
def get_patch_file(toks, override=None):
if override is not None: return override
pfmt = root + "/area/%s.fits"
for tok in toks:
fname = pfmt % tok
if os.path.isfile(fname): return fname
raise ValueError("Can infer patch file")
def invalid_combination(toks):
# Find seasons, arrays and freqs
seasons, arrays, freqs = [], [], []
for tok in toks:
m = re.match(r"(s\d\d)", tok)
if m: seasons.append(m.group(1))
m = re.match(r"(pa\d)", tok)
if m: arrays.append(m.group(1))
m = re.match(r"\+?(f\d\d\d)", tok)
if m: freqs.append(m.group(1))
# We only handle the simple case of one array and one freq
if len(seasons) != 1 or len(arrays) != 1 or len(freqs) != 1:
return False
season, array, freq = seasons[0], arrays[0], freqs[0]
if season in season_arrays and array not in season_arrays[season]: return True
if array not in array_freqs: return False
return freq not in array_freqs[array]
for toks in utils.list_combination_iter(alts):
# Skip invalid array-frequency combinations
if not args.no_prune and invalid_combination(toks): continue
# Expand variables in toks, based on the toks themselves. Only positional expansion supported
otoks = remap(toks, tag_map)
for i in range(args.nsplit):
if args.split_mode == "jon":
osel = ",".join(toks) + ",int32(jon/%f)%%%d==%d" % (args.tblock,args.nsplit,i)
elif args.split_mode == "baz":
osel = ",".join(toks) + ",int32(((baz+180)%%360-180+200)/400.*%d)==%d" % (args.nsplit,i)
elif args.split_mode.startswith("file:"):
fname = args.split_mode[5:]
fname = fname.format(*otoks, i=i)
osel = ",".join(toks) + ",@" + fname
else: raise ValueError(args.split_mode)
otag = "_".join(otoks) + "_%dway_%d" % (args.nsplit, i)
#otag = "_".join(tok.replace("+","").replace("~","no") for tok in toks) + "_%dway_%d" % (args.nsplit,i)
patch_file = get_patch_file(otoks, args.patch)
shape, wcs = enmap.read_map_geometry(patch_file)
npix = shape[-2]*shape[-1]
map_type = "map" if npix < 4e7 else "dmap"
# Allow variable replacement in the unknown arguments
unknown_expanded = [a.format(*otoks, i=i) for a in unknown]
batch = """#!/bin/bash
#SBATCH --nodes %(nnode)d --ntasks-per-node=10 --cpus-per-task=4 --time=%(time)s
#SBATCH --job-name %(name)s
cd "%(cdir)s"
OMP_NUM_THREADS=4 mpirun -bind-to none -npernode 10 --oversubscribe python -u ~/local/tenki/tod2map2.py --dmap_format=tiles -S %(sky)s:%(patch)s,type=%(map_type)s%(sys)s "%(osel)s" "%(odir)s" "%(otag)s" %(extra_args)s""" % {
"nnode": args.nnode, "name": otag, "cdir": os.getcwd(),
"patch": patch_file, "map_type": map_type, "osel": osel,
"sky": args.sky, "time": args.time, "sys": ",sys="+args.sys if args.sys else "",
"otag": otag, "odir": args.odir, "extra_args": " ".join(unknown_expanded)}
runfile = rundir + "/%s.txt" % otag
with open(runfile, "w") as f:
f.write(batch + "\n")
if not args.dry_run:
subprocess.call(["sbatch",runfile])
# Copy our command line argument to the run directory, but avoid clobbering
for i in range(100):
ofile = rundir + "/submit%d.txt" % i
if os.path.isfile(ofile): continue
with open(ofile, "w") as f:
f.write("submit_map " + " ".join(sys.argv[1:]) + "\n")
break