forked from CoronaNetDataScience/corona_tscs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathcorona_tscs_betab.stan
164 lines (126 loc) · 6 KB
/
corona_tscs_betab.stan
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
// Coronavirus tracking model
// Robert Kubinec
// New York University Abu Dhabi
// March 20, 2020
functions{
vector setstep(vector col_var) {
vector[rows(col_var)] diff_var = col_var + (max(fabs(col_var)) * sqrt(0.0000007)) - col_var;
return diff_var;
}
}
data {
int time_all;
int num_country;
int cases[num_country,time_all];
int tests[num_country,time_all];
int time_outbreak[num_country,time_all];
int S; // number of suppression measures
matrix[time_all,3] ortho_time;
matrix[num_country,S] suppress;
vector[time_all] count_outbreak;
int country_pop[num_country];
real phi_scale; // prior on how much change there could be in infection rate over time
}
transformed data {
matrix[num_country,time_all] time_outbreak_trans1; // convert raw time numbers to ortho-normal polynomials
matrix[num_country,time_all] time_outbreak_trans2; // convert raw time numbers to ortho-normal polynomials
matrix[num_country,time_all] time_outbreak_trans3; // convert raw time numbers to ortho-normal polynomials
matrix[num_country,3] time_array[time_all];
matrix[num_country,time_all] time_outbreak_center;
for(t in 1:time_all) {
for(n in 1:num_country) {
if(time_outbreak[n,t]>0) {
time_outbreak_trans1[n,t] = ortho_time[time_outbreak[n,t],1];
time_outbreak_trans2[n,t] = ortho_time[time_outbreak[n,t],2];
time_outbreak_trans3[n,t] = ortho_time[time_outbreak[n,t],3];
} else {
time_outbreak_trans1[n,t] = 0;
time_outbreak_trans2[n,t] = 0;
time_outbreak_trans3[n,t] = 0;
}
}
}
// make a new array of time indices
for(t in 1:time_all) {
time_array[t] = append_col(time_outbreak_trans1[,t],
append_col(time_outbreak_trans2[,t],
time_outbreak_trans3[,t]));
}
// need a centered time vector
for(n in 1:num_country) {
time_outbreak_center[n,] = (to_vector(time_outbreak[n,])' - mean(to_vector(time_outbreak[n,])))/sd(to_vector(time_outbreak[n,]));
}
}
parameters {
vector[3] poly; // polinomial function of time
real<lower=0> finding; // difficulty of identifying infected cases
real world_infect; // infection rate based on number of travelers
row_vector[S] suppress_effect[2]; // suppression effect of govt. measures, cannot increase virus transmission rate
vector<lower=0>[num_country] country_test_raw; // unobserved rate at which countries are willing to test vs. number of infected
// we assume that as infection rates increase, more tests will be conducted
vector[2] alpha; // other intercepts
vector<lower=0>[2] phi; // shape parameter for infected
vector[num_country-1] country_int_free; // varying intercepts by country - 1 for identification
real<lower=0> sigma_test_raw; // estimate of between-state testing heterogeneity
}
transformed parameters {
vector[num_country] country_int = append_row(0,country_int_free); // not strictly necessary but centers around 0
matrix[num_country,time_all] num_infected_high; // modeled infection rates for domestic transmission
num_infected_high[,1] = rep_vector(0,num_country);
for(t in 1:time_all) {
//real num_low;
num_infected_high[,t] = alpha[2] + country_int +
time_array[t]*poly +
world_infect*count_outbreak[t] +
(suppress_effect[1]*suppress')' +
+ ((suppress_effect[2]*suppress') .* time_outbreak_center[,t]')';
}
}
model {
poly ~ normal(0,10); // could be large
world_infect ~ normal(0,1);
alpha ~ normal(0,10); // this can reach extremely low values
phi ~ exponential(phi_scale);
for(i in 1:2)
suppress_effect[i] ~ normal(0,2);
finding ~ exponential(.1);
country_int_free ~ normal(0,3);
sigma_test_raw ~ exponential(.1);
country_test_raw ~ exponential(sigma_test_raw); // more likely near the middle than the ends
// first model probability of infection
//next model the true infection rate as a function of time since outbreak
for (t in 2:time_all) {
vector[num_country] mix_prop = inv_logit(num_infected_high[,t]);
// locations for cases and tests
vector[num_country] mu_cases = inv_logit(-2.19 + finding*num_infected_high[,t]);
vector[num_country] mu_tests = inv_logit(alpha[1] + country_test_raw .* num_infected_high[,t]);
tests[,t] ~ beta_binomial(country_pop,mu_tests*phi[1],(1-mu_tests)*phi[1]);
cases[,t] ~ beta_binomial(tests[,t],mu_cases*phi[2],(1-mu_cases)*phi[2]);
}
}
generated quantities {
vector[S] suppress_margin;
for(s in 1:S) {
matrix[rows(suppress),cols(suppress)] suppress_high = suppress;
matrix[rows(suppress),cols(suppress)] suppress_low = suppress;
matrix[rows(suppress),rows(count_outbreak)] y_high;
matrix[rows(suppress),rows(count_outbreak)] y_low;
matrix[rows(suppress),rows(count_outbreak)] y_diff;
suppress_high[,s] = suppress[,s] + setstep(suppress[,s]);
suppress_low[,s] = suppress[,s] - setstep(suppress[,s]);
for(t in 1:time_all) {
y_high[,t] = alpha[2] + country_int +
time_array[t]*poly +
world_infect*count_outbreak[t] +
(suppress_effect[1]*suppress_high')' +
((suppress_effect[2]*suppress_high') .* time_outbreak_center[,t]')';
y_low[,t] = alpha[2] + country_int +
time_array[t]*poly +
world_infect*count_outbreak[t] +
(suppress_effect[1]*suppress_low')' +
((suppress_effect[2]*suppress_low') .* time_outbreak_center[,t]')';
y_diff[,t] = (inv_logit(y_high[,t]) - inv_logit(y_low[,t])) ./ (suppress_high[,s] - suppress_low[,s]);
}
suppress_margin[s] = mean(to_vector(y_diff));
}
}