-
Notifications
You must be signed in to change notification settings - Fork 219
/
covasim
executable file
·78 lines (63 loc) · 2.5 KB
/
covasim
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
#!/usr/bin/env python3
'''
Simple script for running the Covid-19 agent-based model
'''
print('Importing...')
import sciris as sc
import covasim as cv
try:
import fire # Optional import
except ModuleNotFoundError as E:
errormsg = f'Please install the "fire" module first, e.g. pip install fire' # Also caught in make_people()
raise ModuleNotFoundError(errormsg) from E
# Set filename if saving
version = 'v0'
date = '2020apr06'
folder = 'results'
basename = f'{folder}/covasim_run_{date}_{version}'
fig_path = f'{basename}.png'
# Configure the sim -- can also just use a normal dictionary
pars = sc.objdict(
pop_size = 20000, # Population size
pop_infected = 1, # Number of initial infections
n_days = 180, # Number of days to simulate
rand_seed = 1, # Random seed
)
def run(pars=pars,
interv=False,
do_plot=True,
do_save=False,
do_show=False,
fig_path=fig_path,
verbose=False
):
'''
Run simulation, optionally with an intervention.
To run simulation with default parameters:
> python sim.py
To run simulation with an intervention:
> python sim.py --interv=True
To run simulation with an intervention and save plot to disk:
> python sim.py --interv=True --do_save=True
To run simulation with custom configuration
> python sim.py --pars "{pop_size:20000, pop_infected:1, n_days:360, rand_seed:1}"
Args:
pars: (dict): configuration for the sim. See description for example.
interv: (bool): whether or not to add an intervention, specified by cv.change_beta(days=45, changes=0.5)
do_plot: (bool): whether or not to generate a plot. Defaults to True.
do_save: (bool): If a plot is generated, whether or not to save it. Defaults to False.
do_show: (bool): If a plot is generated, whether or not to show it. Defaults to False.
fig_path: (str): Path to which save filename. Defaults to results/covasim_run_{date}_{version}.png
verbose: (bool): whether or not turn verbose mode while running simulation. Defaults to False.
'''
if interv:
pars.interventions = cv.change_beta(days=45, changes=0.5) # Optionally add an intervention
print('Making sim...')
sim = cv.Sim(pars=pars)
print('Running...')
sim.run(verbose=verbose)
if do_plot:
print('Plotting...')
fig = sim.plot(do_save=do_save, do_show=do_show, fig_path=fig_path)
if __name__ == '__main__':
fire.Fire(run)