-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathcovasim_utils.py
199 lines (170 loc) · 6.53 KB
/
covasim_utils.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
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
""" Utilities for running Covasim. """
import numpy as np
import covasim as cv
import sciris as sc
import pandas as pd
import os
def avg_age(location):
ages = cv.data.loaders.map_entries(cv.data.country_age_data.data, location)[location]
total_pop = sum(ages.values())
avg = 0
for age in ages:
prob = ages[age] / total_pop
if "-" in age:
midpoint = np.mean([int(x) for x in age.split("-")])
else:
midpoint = 80
avg += midpoint * prob
return round(avg)
# return avg
def household_size(location):
return cv.data.loaders.get_household_size(location)
# Instantiate named covasim interventions with parameters
intervention = {
"standardTest": cv.test_prob(
symp_prob=0.2, asymp_prob=0.001, symp_quar_prob=1, asymp_quar_prob=1
),
"noTest": cv.test_prob(
symp_prob=0, asymp_prob=0, symp_quar_prob=0, asymp_quar_prob=0
),
"optimalTest": cv.test_prob(
symp_prob=1, asymp_prob=1, symp_quar_prob=1, asymp_quar_prob=1
),
"standardTrace": cv.contact_tracing(
trace_probs={"h": 1, "w": 0.5, "s": 0.5, "c": 0.3}, quar_period=14
),
"noTrace": cv.contact_tracing(
trace_probs={"h": 0, "w": 0, "s": 0, "c": 0}, quar_period=14
),
"optimalTrace": cv.contact_tracing(
trace_probs={"h": 1, "w": 1, "s": 1, "c": 1}, quar_period=14
),
}
# Provide a list of interventions with which to run covasim
interventions = {
"baseline": [],
"standardTest": [intervention["standardTest"]],
"noTest": [intervention["noTest"]],
"optimalTest": [intervention["optimalTest"]],
"standardTrace": [intervention["standardTest"], intervention["standardTrace"]],
"noTrace": [intervention["standardTest"], intervention["noTrace"]],
"optimalTrace": [intervention["standardTest"], intervention["optimalTrace"]],
"traceNoTest": [intervention["standardTrace"]],
}
def preprocess_data(data):
assert "intervention" not in data, "intervention is a column in data"
data["interventions"] = [str(x) for x in data["interventions"]]
data["interventions"] = data["interventions"].astype("category")
data["pop_type"] = data["pop_type"].astype("category")
data["location"] = data["location"].astype("category")
if "start_day" in data:
data["start_day"] = data["start_day"].astype("category")
return data
def dict_plus(dic1, dic2):
for field in dic2:
if field not in dic1:
dic1[field] = []
dic1[field].append(dic2[field])
def chunks(lst, n):
"""Yield successive n-sized chunks from lst."""
for i in range(0, len(lst), n):
yield lst[i : i + n]
def aggregate_by_week(data, desired_outputs=None):
if desired_outputs is None:
desired_outputs = data.columns
week_by_week = {k: [] for k in desired_outputs}
for c in chunks(data, 7):
for k in week_by_week:
if k.startswith("new_"):
week_by_week[k].append(c[k].sum())
elif k.startswith("n_"):
week_by_week[k].append(c[k].iloc[-1])
else:
week_by_week[k].append(c[k].iloc[0])
return pd.DataFrame(week_by_week)
def run_covasim_by_week(label, params, desired_outputs, n_runs=30):
print("Params", params)
print("Desired outputs", desired_outputs)
intervention_sim = cv.MultiSim(
cv.Sim(pars=params, label=label, verbose=0), keep_people=True
)
intervention_sim.run(n_runs=n_runs, verbose=0)
temporal = []
for sim in intervention_sim.sims:
df = sim.to_df()
quar_period = 14
for i in sim["interventions"]:
if hasattr(i, "quar_period"):
quar_period = i.quar_period
df = df[desired_outputs]
week_by_week = pd.DataFrame(aggregate_by_week(df, desired_outputs))
dic = week_by_week.to_dict(orient="list")
week_dic = {
f"{k}_{w+1}": item for k in desired_outputs for w, item in enumerate(dic[k])
}
week_dic["quar_period"] = quar_period
params["interventions"] = label
for k, v in params.items():
week_dic[k] = v
temporal.append(week_dic)
# add avg. age and contacts to data
week_dic["avg_age"] = np.average(sim.people["age"])
week_dic["contacts_h"] = sim.pars["contacts"]["h"]
week_dic["contacts_s"] = sim.pars["contacts"]["s"]
week_dic["contacts_w"] = sim.pars["contacts"]["w"]
week_dic["contacts_c"] = sim.pars["contacts"]["c"]
return pd.DataFrame(temporal)
def run_covasim_basic(label, params, desired_outputs, n_runs=30):
intervention_sim = cv.MultiSim(cv.Sim(pars=params, label=label, verbose=0))
intervention_sim.run(n_runs=n_runs, verbose=0)
results = {k: [] for k in desired_outputs}
results["quar_period"] = []
for sim in intervention_sim.sims:
df = sim.to_df()
quar_period = 14
for i in sim["interventions"]:
if hasattr(i, "quar_period"):
quar_period = i.quar_period
df = df[desired_outputs]
for k in desired_outputs:
results[k].append(df[k].iloc[-1])
results["quar_period"].append(quar_period)
data = pd.DataFrame(results)
params["interventions"] = label
for k, v in params.items():
data[k] = [v for _ in range(n_runs)]
return data
def msims(default, investigate, include_baseline=True):
"""
Generate a list of MultiSims to be run to investigate a series of parameter
changes.
Parameters
----------
default : dict
The default parameters with which to run covasim.
investigate : [dict]
A list of dictionaries with the investigation parameters.
Returns
-------
[MultiSim]
A list of MultiSims with the specified parameters.
"""
sims = [
cv.MultiSim(cv.Sim(pars={**default, **d}, label=l, verbose=0))
for d, l in investigate
]
if include_baseline:
sims += [cv.MultiSim(cv.Sim(pars=default, label="Baseline"))]
return sims
def save_results_df(results_df, out_path, file_name):
"""
Given a data frame containing Covasim results, save the dataframe to a specified path.
If the out_path directory doesn't exist, this method will create it.
:param results_df: (pandas data frame) containing Covasim results.
:param out_path: (string) path to save location.
:param file_name: (string) name of file excluding extension.
"""
if not os.path.exists(out_path):
os.mkdir(out_path)
out_path = os.path.join(out_path, f"{file_name}.csv")
results_df.to_csv(out_path)