covid-sim
run_sample.py
1 #!/usr/bin/env python3
2 """Run the sample data.
3 
4 See README.md in this directory for more information.
5 """
6 
7 import argparse
8 import gzip
9 import multiprocessing
10 import os
11 import shutil
12 import subprocess
13 import sys
14 
15 def try_remove(f):
16  try:
17  os.remove(f)
18  except OSError as e:
19  pass
20 
21 def parse_args():
22  """Parse the arguments.
23 
24  On exit: Returns the result of calling argparse.parse()
25 
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
31  """
32  parser = argparse.ArgumentParser()
33  try:
34  cpu_count = len(os.sched_getaffinity(0))
35  except AttributeError:
36  # os.sched_getaffinity isn't available
37  cpu_count = multiprocessing.cpu_count()
38  if cpu_count is None or cpu_count == 0:
39  cpu_count = 2
40 
41  script_path = os.path.dirname(os.path.realpath(__file__))
42 
43  # Default values
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)
48 
49  parser.add_argument(
50  "country",
51  help="Country to run sample for")
52  parser.add_argument(
53  "--covidsim",
54  help="Location of CovidSim binary, if none specified will build")
55  parser.add_argument(
56  "--datadir",
57  help="Directory at root of input data",
58  default=script_path)
59  parser.add_argument(
60  "--paramdir",
61  help="Directory with input parameter files",
62  default=param_dir)
63  parser.add_argument(
64  "--srcdir",
65  help="Directory with source in - needed if --covidsim isn't specified",
66  default=src_dir)
67  parser.add_argument(
68  "--outputdir",
69  help="Directory to store output data",
70  default=output_dir)
71  parser.add_argument(
72  "--threads",
73  help="Number of threads to use",
74  default=cpu_count
75  )
76  args = parser.parse_args()
77 
78  return args
79 
80 args = parse_args()
81 
82 # Lists of places that need to be handled specially
83 united_states = [ "United_States" ]
84 canada = [ "Canada" ]
85 usa_territories = ["Alaska", "Hawaii", "Guam", "Virgin_Islands_US", "Puerto_Rico", "American_Samoa"]
86 nigeria = ["Nigeria"]
87 
88 # Determine whether we need to build the tool or use a user supplied one:
89 if args.covidsim is not None:
90  exe = args.covidsim
91 else:
92  build_dir = os.path.join(args.outputdir, "build")
93 
94  # Ensure we do a clean build
95  shutil.rmtree(build_dir, ignore_errors=True)
96  os.makedirs(build_dir, exist_ok=False)
97  cwd = os.getcwd()
98  os.chdir(build_dir)
99 
100  # Build
101  subprocess.run(['cmake', args.srcdir], check=True)
102  subprocess.run(['cmake', '--build', '.'], check=True)
103 
104  # Where the exe ends up depends on the OS.
105  if os.name == 'nt':
106  exe = os.path.join(build_dir, "Debug", "src", "CovidSim.exe")
107  else:
108  exe = os.path.join(build_dir, "src", "CovidSim")
109 
110  os.chdir(cwd)
111 
112 # Ensure output directory exists
113 os.makedirs(args.outputdir, exist_ok=True)
114 
115 # The admin file to use
116 admin_file = os.path.join(args.datadir, "admin_units",
117  "{0}_admin.txt".format(args.country))
118 
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))
123  exit(1)
124 
125 # Population density file in gziped form, text file, and binary file as
126 # processed by CovidSim
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"
133 else:
134  wpop_file_root = "eur"
135 
136 wpop_file_gz = os.path.join(
137  args.datadir,
138  "populations",
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))
144  exit(1)
145 
146 wpop_file = os.path.join(
147  args.outputdir,
148  "wpop_{0}.txt".format(wpop_file_root))
149 wpop_bin = os.path.join(
150  args.outputdir,
151  "{0}_pop_density.bin".format(args.country))
152 
153 # gunzip wpop fie
154 try_remove(wpop_file)
155 try_remove(wpop_bin)
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)
159 
160 # Configure pre-parameter file. This file doesn't change between runs:
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")
165 else:
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))
171  exit(1)
172 
173 # Configure No intervention parameter file. This is run first
174 # and provides a baseline
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))
180  exit(1)
181 
182 # Configure an intervention (controls) parameter file.
183 # In reality you will run CovidSim many times with different parameter
184 # controls.
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))
192  exit(1)
193 
194 school_file = None
195 if args.country in united_states:
196  school_file = os.path.join(args.datadir, "populations", "USschools.txt")
197 
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))
202  exit(1)
203 
204 # Some command_line settings
205 r = 3.0
206 rs = r/2
207 
208 # This is the temporary network that represents initial state of the
209 # simulation
210 network_bin = os.path.join(
211  args.outputdir,
212  "Network_{0}_T{1}_R{2}.bin".format(args.country, args.threads, r))
213 try_remove(network_bin)
214 
215 # Run the no intervention sim. This also does some extra setup which is one
216 # off for each R.
217 print("No intervention: {0} NoInt {1}".format(args.country, r))
218 cmd = [
219  exe,
220  "/c:{0}".format(args.threads),
221  "/A:" + admin_file
222 ]
223 if school_file:
224  cmd.extend(["/s:" + school_file])
225 cmd.extend([
226  "/PP:" + pp_file, # Preparam file
227  "/P:" + no_int_file, # Param file
228  "/O:" + os.path.join(args.outputdir,
229  "{0}_NoInt_R0={1}".format(args.country, r)), # Output
230  "/D:" + wpop_file, # Input (this time text) pop density
231  "/M:" + wpop_bin, # Where to save binary pop density
232  "/S:" + network_bin, # Where to save binary net setup
233  "/R:{0}".format(rs),
234  "98798150", # These four numbers are RNG seeds
235  "729101",
236  "17389101",
237  "4797132"
238  ])
239 print("Command line: " + " ".join(cmd))
240 process = subprocess.run(cmd, check=True)
241 
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))
245  cmd = [
246  exe,
247  "/c:{0}".format(args.threads),
248  "/A:" + admin_file
249  ]
250  if school_file:
251  cmd.extend(["/s:" + school_file])
252  cmd.extend([
253  "/PP:" + pp_file,
254  "/P:" + cf,
255  "/O:" + os.path.join(args.outputdir,
256  "{0}_{1}_R0={2}".format(args.country, root, r)),
257  "/D:" + wpop_bin, # Binary pop density file (speedup)
258  "/L:" + network_bin, # Network to load
259  "/R:{0}".format(rs),
260  "98798150",
261  "729101",
262  "17389101",
263  "4797132"
264  ])
265  print("Command line: " + " ".join(cmd))
266  process = subprocess.run(cmd, check=True)
def parse_args()
Definition: run_sample.py:21