2 """Run the sample data. 4 See README.md in this directory for more information. 22 """Parse the arguments. 24 On exit: Returns the result of calling argparse.parse() 26 args.covidsim is the name of the CovidSim executable 27 args.datadir is the directory with the input data 28 args.paramdir is the directory with the parameters in it 29 args.outputdir is the directory where output will be stored 30 args.threads is the number of threads to use 32 parser = argparse.ArgumentParser()
34 cpu_count = len(os.sched_getaffinity(0))
35 except AttributeError:
37 cpu_count = multiprocessing.cpu_count()
38 if cpu_count
is None or cpu_count == 0:
41 script_path = os.path.dirname(os.path.realpath(__file__))
44 data_dir = script_path
45 param_dir = os.path.join(script_path,
"param_files")
46 output_dir = os.getcwd()
47 src_dir = os.path.join(data_dir, os.pardir)
51 help=
"Country to run sample for")
54 help=
"Location of CovidSim binary, if none specified will build")
57 help=
"Directory at root of input data",
61 help=
"Directory with input parameter files",
65 help=
"Directory with source in - needed if --covidsim isn't specified",
69 help=
"Directory to store output data",
73 help=
"Number of threads to use",
76 args = parser.parse_args()
83 united_states = [
"United_States" ]
85 usa_territories = [
"Alaska",
"Hawaii",
"Guam",
"Virgin_Islands_US",
"Puerto_Rico",
"American_Samoa"]
89 if args.covidsim
is not None:
92 build_dir = os.path.join(args.outputdir,
"build")
95 shutil.rmtree(build_dir, ignore_errors=
True)
96 os.makedirs(build_dir, exist_ok=
False)
101 subprocess.run([
'cmake', args.srcdir], check=
True)
102 subprocess.run([
'cmake',
'--build',
'.'], check=
True)
106 exe = os.path.join(build_dir,
"Debug",
"src",
"CovidSim.exe")
108 exe = os.path.join(build_dir,
"src",
"CovidSim")
113 os.makedirs(args.outputdir, exist_ok=
True)
116 admin_file = os.path.join(args.datadir,
"admin_units",
117 "{0}_admin.txt".format(args.country))
119 if not os.path.exists(admin_file):
120 print(
"Unable to find admin file for country: {0}".format(args.country))
121 print(
"Data directory: {0}".format(args.datadir))
122 print(
"Looked for: {0}".format(admin_file))
127 if args.country
in united_states + canada:
128 wpop_file_root =
"usacan" 129 elif args.country
in usa_territories:
130 wpop_file_root =
"us_terr" 131 elif args.country
in nigeria:
132 wpop_file_root =
"nga_adm1" 134 wpop_file_root =
"eur" 136 wpop_file_gz = os.path.join(
139 "wpop_{0}.txt.gz".format(wpop_file_root))
140 if not os.path.exists(wpop_file_gz):
141 print(
"Unable to find population file for country: {0}".format(args.country))
142 print(
"Data directory: {0}".format(args.datadir))
143 print(
"Looked for: {0}".format(wpop_file_gz))
146 wpop_file = os.path.join(
148 "wpop_{0}.txt".format(wpop_file_root))
149 wpop_bin = os.path.join(
151 "{0}_pop_density.bin".format(args.country))
154 try_remove(wpop_file)
156 with gzip.open(wpop_file_gz,
'rb')
as f_in:
157 with open(wpop_file,
'wb')
as f_out:
158 shutil.copyfileobj(f_in, f_out)
161 if args.country
in united_states:
162 pp_file = os.path.join(args.paramdir,
"preUS_R0=2.0.txt")
163 elif args.country
in nigeria:
164 pp_file = os.path.join(args.paramdir,
"preNGA_R0=2.0.txt")
166 pp_file = os.path.join(args.paramdir,
"preUK_R0=2.0.txt")
167 if not os.path.exists(pp_file):
168 print(
"Unable to find pre-parameter file")
169 print(
"Param directory: {0}".format(args.paramdir))
170 print(
"Looked for: {0}".format(pp_file))
175 no_int_file = os.path.join(args.paramdir,
"p_NoInt.txt")
176 if not os.path.exists(no_int_file):
177 print(
"Unable to find parameter file")
178 print(
"Param directory: {0}".format(args.paramdir))
179 print(
"Looked for: {0}".format(no_int_file))
185 control_roots = [
"PC7_CI_HQ_SD" ]
186 for root
in control_roots:
187 cf = os.path.join(args.paramdir,
"p_{0}.txt".format(root))
188 if not os.path.exists(cf):
189 print(
"Unable to find parameter file")
190 print(
"Param directory: {0}".format(args.paramdir))
191 print(
"Looked for: {0}".format(cf))
195 if args.country
in united_states:
196 school_file = os.path.join(args.datadir,
"populations",
"USschools.txt")
198 if not os.path.exists(school_file):
199 print(
"Unable to find school file for country: {0}".format(args.country))
200 print(
"Data directory: {0}".format(args.datadir))
201 print(
"Looked for: {0}".format(school_file))
210 network_bin = os.path.join(
212 "Network_{0}_T{1}_R{2}.bin".format(args.country, args.threads, r))
213 try_remove(network_bin)
217 print(
"No intervention: {0} NoInt {1}".format(args.country, r))
220 "/c:{0}".format(args.threads),
224 cmd.extend([
"/s:" + school_file])
228 "/O:" + os.path.join(args.outputdir,
229 "{0}_NoInt_R0={1}".format(args.country, r)),
239 print(
"Command line: " +
" ".join(cmd))
240 process = subprocess.run(cmd, check=
True)
242 for root
in control_roots:
243 cf = os.path.join(args.paramdir,
"p_{0}.txt".format(root))
244 print(
"Intervention: {0} {1} {2}".format(args.country, root, r))
247 "/c:{0}".format(args.threads),
251 cmd.extend([
"/s:" + school_file])
255 "/O:" + os.path.join(args.outputdir,
256 "{0}_{1}_R0={2}".format(args.country, root, r)),
265 print(
"Command line: " +
" ".join(cmd))
266 process = subprocess.run(cmd, check=
True)