Skip to content
GitLab
Menu
Projects
Groups
Snippets
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Chris Jewell
covid19uk
Commits
d2e2af66
Commit
d2e2af66
authored
Mar 19, 2020
by
Chris Jewell
Browse files
Bivariate MCMC algorithm
parent
a82e7c8f
Changes
7
Hide whitespace changes
Inline
Side-by-side
covid/model.py
View file @
d2e2af66
...
...
@@ -88,7 +88,7 @@ def dense_to_block_diagonal(A, n_blocks):
class
CovidUKODE
:
# TODO: add background case importation rate to the UK, e.g. \epsilon term.
def
__init__
(
self
,
M_tt
,
M_hh
,
C
,
N
,
start
,
end
,
holidays
,
t_step
):
def
__init__
(
self
,
M_tt
,
M_hh
,
C
,
N
,
start
,
end
,
holidays
,
bg_max_t
,
t_step
):
"""Represents a CovidUK ODE model
:param K_tt: a MxM matrix of age group mixing in term time
...
...
@@ -117,7 +117,8 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
self
.
times
=
np
.
arange
(
start
,
end
,
np
.
timedelta64
(
t_step
,
'D'
))
m_select
=
(
np
.
less_equal
(
holidays
[
0
],
self
.
times
)
&
np
.
less
(
self
.
times
,
holidays
[
1
])).
astype
(
np
.
int64
)
self
.
m_select
=
tf
.
constant
(
m_select
)
self
.
m_select
=
tf
.
constant
(
m_select
,
dtype
=
tf
.
int64
)
self
.
bg_select
=
tf
.
constant
(
np
.
less
(
bg_max_t
,
self
.
times
),
dtype
=
tf
.
int64
)
self
.
solver
=
tode
.
DormandPrince
()
def
make_h
(
self
,
param
):
...
...
@@ -125,13 +126,15 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
def
h_fn
(
t
,
state
):
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
m_switch
=
tf
.
gather
(
self
.
m_select
,
tf
.
cast
(
t
,
tf
.
int64
))
t
=
tf
.
clip_by_value
(
tf
.
cast
(
t
,
tf
.
int64
),
0
,
self
.
m_select
.
shape
[
0
]
-
1
)
m_switch
=
tf
.
gather
(
self
.
m_select
,
t
)
epsilon
=
param
[
'epsilon'
]
*
tf
.
cast
(
tf
.
gather
(
self
.
bg_select
,
t
),
tf
.
float32
)
if
m_switch
==
0
:
infec_rate
=
param
[
'beta1'
]
*
tf
.
linalg
.
matvec
(
self
.
M
[
0
],
I
)
else
:
infec_rate
=
param
[
'beta1'
]
*
tf
.
linalg
.
matvec
(
self
.
M
[
1
],
I
)
infec_rate
+=
param
[
'beta1'
]
*
param
[
'beta2'
]
*
self
.
Kbar
*
tf
.
linalg
.
matvec
(
self
.
C
,
I
/
self
.
N_sum
)
infec_rate
=
S
/
self
.
N
*
infec_rate
infec_rate
=
S
/
self
.
N
*
(
infec_rate
+
epsilon
)
dS
=
-
infec_rate
dE
=
infec_rate
-
param
[
'nu'
]
*
E
...
...
covid/util.py
View file @
d2e2af66
...
...
@@ -2,9 +2,10 @@
import
numpy
as
np
def
sanitise_parameter
(
par_dict
):
"""Sanitises a dictionary of parameters"""
par
=
[
'beta1'
,
'beta2'
,
'nu'
,
'gamma'
]
par
=
[
'epsilon'
,
'beta1'
,
'beta2'
,
'nu'
,
'gamma'
]
d
=
{
key
:
np
.
float64
(
par_dict
[
key
])
for
key
in
par
}
return
d
...
...
@@ -13,7 +14,8 @@ def sanitise_settings(par_dict):
d
=
{
'start'
:
np
.
datetime64
(
par_dict
[
'start'
]),
'end'
:
np
.
datetime64
(
par_dict
[
'end'
]),
'time_step'
:
float
(
par_dict
[
'time_step'
]),
'holiday'
:
np
.
array
([
np
.
datetime64
(
date
)
for
date
in
par_dict
[
'holiday'
]])}
'holiday'
:
np
.
array
([
np
.
datetime64
(
date
)
for
date
in
par_dict
[
'holiday'
]]),
'bg_max_time'
:
np
.
datetime64
(
par_dict
[
'bg_max_time'
])}
return
d
...
...
covid_ode.py
View file @
d2e2af66
...
...
@@ -12,7 +12,7 @@ from covid.rdata import *
def
sanitise_parameter
(
par_dict
):
"""Sanitises a dictionary of parameters"""
par
=
[
'beta1'
,
'beta2'
,
'nu'
,
'gamma'
]
par
=
[
'epsilon'
,
'beta1'
,
'beta2'
,
'nu'
,
'gamma'
]
d
=
{
key
:
np
.
float64
(
par_dict
[
key
])
for
key
in
par
}
return
d
...
...
@@ -21,7 +21,8 @@ def sanitise_settings(par_dict):
d
=
{
'start'
:
np
.
datetime64
(
par_dict
[
'start'
]),
'end'
:
np
.
datetime64
(
par_dict
[
'end'
]),
'time_step'
:
float
(
par_dict
[
'time_step'
]),
'holiday'
:
np
.
array
([
np
.
datetime64
(
date
)
for
date
in
par_dict
[
'holiday'
]])}
'holiday'
:
np
.
array
([
np
.
datetime64
(
date
)
for
date
in
par_dict
[
'holiday'
]]),
'bg_max_time'
:
np
.
datetime64
(
par_dict
[
'bg_max_time'
])}
return
d
...
...
@@ -211,33 +212,17 @@ if __name__ == '__main__':
K1_hh
[
14
:,
:]
*=
cocooning_ratio
K1_hh
[:,
14
:]
*=
cocooning_ratio
model
_term
=
CovidUKODE
(
K1_tt
,
T
,
N
)
model_holiday
=
CovidUKODE
(
K1_hh
,
T
,
N
)
model
=
CovidUKODE
(
K1_tt
,
K_hh
,
T
,
N
,
settings
[
'start'
],
settings
[
'end'
],
settings
[
'holiday'
],
settings
[
'bg_max_time'
],
1
)
seeding
=
seed_areas
(
N
,
n_names
)
# Seed 40-44 age group, 30 seeds by popn size
state_init
=
model_term
.
create_initial_state
(
init_matrix
=
seeding
)
print
(
'R_term='
,
model_term
.
eval_R0
(
param
))
print
(
'R_holiday='
,
model_holiday
.
eval_R0
(
param
))
# School holidays and closures
@
tf
.
function
()
def
simulate
():
t0
,
sim_0
,
solve0
=
model_term
.
simulate
(
param
,
state_init
,
settings
[
'start'
],
settings
[
'holiday'
][
0
],
settings
[
'time_step'
])
t1
,
sim_1
,
solve1
=
model_holiday
.
simulate
(
param
,
sim_0
[
-
1
,
:,
:],
settings
[
'holiday'
][
0
],
settings
[
'holiday'
][
1
],
settings
[
'time_step'
],
solver_state
=
None
)
t2
,
sim_2
,
_
=
model_term
.
simulate
(
param
,
sim_1
[
-
1
,
:,
:],
settings
[
'holiday'
][
1
],
settings
[
'end'
],
settings
[
'time_step'
],
solver_state
=
None
)
t
=
tf
.
concat
([
t0
,
t1
+
t0
[
-
1
],
t2
+
t0
[
-
1
]
+
t1
[
-
1
]],
axis
=
0
)
sim
=
tf
.
concat
([
tf
.
expand_dims
(
state_init
,
axis
=
0
),
sim_0
,
sim_1
,
sim_2
],
axis
=
0
)
return
t
,
sim
state_init
=
model
.
create_initial_state
(
init_matrix
=
seeding
)
print
(
'R_term='
,
model
.
eval_R0
(
param
))
#print('R_holiday=', model_holiday.eval_R0(param))
start
=
time
.
perf_counter
()
t
,
sim
=
simulate
(
)
t
,
sim
,
_
=
model
.
simulate
(
param
,
state_init
)
end
=
time
.
perf_counter
()
print
(
f
'Complete in
{
end
-
start
}
seconds'
)
...
...
@@ -252,11 +237,10 @@ if __name__ == '__main__':
for
i
,
r
in
enumerate
(
cocoon_ratios
):
print
(
f
"Simulation, r=
{
r
}
"
,
flush
=
True
)
t
,
sim
=
cocooning_model
(
r
)
plot_age_attack_rate
(
fig_attack
.
gca
(),
sim
,
N
,
f
"
{
1
-
r
}
"
,
f
"C
{
i
}
"
)
#fig_attack.gca().legend(title="Contact ratio")
#
fig_attack.gca().legend(title="Contact ratio")
plot_infec_curve
(
fig_uk
.
gca
(),
sim
.
numpy
(),
f
"
{
1
-
r
}
"
,
f
"C
{
i
}
"
)
#fig_uk.gca().legend(title="Contact ratio")
#
fig_uk.gca().legend(title="Contact ratio")
fig_attack
.
autofmt_xdate
()
fig_uk
.
autofmt_xdate
()
...
...
data/DailyConfirmedCases.csv
View file @
d2e2af66
...
...
@@ -42,4 +42,5 @@ DateVal,CMODateCount,CumCases
2020-03-11,83,456
2020-03-12,139,590
2020-03-13,207,797
2020-03-14,343,1140
2020-03-14,264,1061
2020-03-15,330,1391
mcmc.py
View file @
d2e2af66
...
...
@@ -63,54 +63,72 @@ if __name__ == '__main__':
y_incr
=
np
.
round
((
y
[
1
:]
-
y
[:
-
1
])
*
0.8
)
simulator
=
CovidUKODE
(
K_tt
,
K_hh
,
T
,
N
,
date_range
[
0
]
-
np
.
timedelta64
(
1
,
'D'
),
date_range
[
1
],
settings
[
'holiday'
],
int
(
settings
[
'time_step'
]))
date_range
[
1
],
settings
[
'holiday'
],
settings
[
'bg_max_time'
],
int
(
settings
[
'time_step'
]))
seeding
=
seed_areas
(
N
,
n_names
)
# Seed 40-44 age group, 30 seeds by popn size
state_init
=
simulator
.
create_initial_state
(
init_matrix
=
seeding
)
#@tf.function
def
logp
(
beta
1
):
def
logp
(
epsilon
,
beta
):
p
=
param
p
[
'beta1'
]
=
beta1
p
[
'epsilon'
]
=
epsilon
p
[
'beta1'
]
=
beta
epsilon_logp
=
tfd
.
Gamma
(
concentration
=
1.
,
rate
=
1.
).
log_prob
(
p
[
'epsilon'
])
beta_logp
=
tfd
.
Gamma
(
concentration
=
1.
,
rate
=
1.
).
log_prob
(
p
[
'beta1'
])
t
,
sim
,
solve
=
simulator
.
simulate
(
p
,
state_init
)
y_logp
=
covid19uk_logp
(
y_incr
,
sim
,
0.1
)
return
beta_logp
+
tf
.
reduce_sum
(
y_logp
)
logp
=
epsilon_logp
+
beta_logp
+
tf
.
reduce_sum
(
y_logp
)
return
logp
unconstraining_bijector
=
[
tfb
.
Log
()]
initial_mcmc_state
=
[
0.03
]
print
(
"Initial log likelihood:"
,
logp
(
0.03
))
def
trace_fn
(
_
,
pkr
):
return
(
pkr
.
inner_results
.
log_accept_ratio
,
pkr
.
inner_results
.
accepted_results
.
target_log_prob
,
pkr
.
inner_results
.
accepted_results
.
step_size
)
unconstraining_bijector
=
[
tfb
.
Exp
(),
tfb
.
Exp
()]
initial_mcmc_state
=
[
0.00001
,
0.03
]
print
(
"Initial log likelihood:"
,
logp
(
*
initial_mcmc_state
))
@
tf
.
function
def
sample
():
def
sample
(
n_samples
,
step_size
=
[
0.01
,
0.1
]
):
return
tfp
.
mcmc
.
sample_chain
(
num_results
=
2000
,
num_burnin_steps
=
50
0
,
num_results
=
n_samples
,
num_burnin_steps
=
0
,
current_state
=
initial_mcmc_state
,
kernel
=
tfp
.
mcmc
.
SimpleStepSizeAdaptation
(
tfp
.
mcmc
.
TransformedTransitionKernel
(
inner_kernel
=
tfp
.
mcmc
.
HamiltonianMonteCarlo
(
target_log_prob_fn
=
logp
,
step_size
=
0.0001
,
step_size
=
step_size
,
num_leapfrog_steps
=
5
),
bijector
=
unconstraining_bijector
),
num_adaptation_steps
=
400
),
trace_fn
=
lambda
_
,
pkr
:
pkr
.
inner_results
.
inner_results
.
is_accepted
)
start
=
time
.
perf_counter
()
pi_beta
,
accept
=
sample
()
end
=
time
.
perf_counter
()
print
(
f
"Simulation complete in
{
end
-
start
}
seconds"
)
plt
.
plot
(
pi_beta
[
0
])
num_adaptation_steps
=
500
),
trace_fn
=
lambda
_
,
pkr
:
pkr
.
inner_results
.
inner_results
.
accepted_results
.
step_size
)
with
tf
.
device
(
"/CPU:0"
):
start
=
time
.
perf_counter
()
pi_beta
,
results
=
sample
(
1000
,
[
0.1
,
0.1
])
# sd = tfp.stats.stddev(pi_beta, sample_axis=1)
# sd = sd / tf.reduce_sum(sd)
# new_step_size = [results[-1] * sd[0], results[-1] * sd[1]]
# print("New step size:",new_step_size)
# pi_beta, results = sample(1000, step_size=new_step_size)
end
=
time
.
perf_counter
()
print
(
f
"Simulation complete in
{
end
-
start
}
seconds"
)
print
(
results
)
fig
,
ax
=
plt
.
subplots
(
1
,
2
)
ax
[
0
].
plot
(
pi_beta
[
0
])
ax
[
1
].
plot
(
pi_beta
[
1
])
plt
.
show
()
print
(
f
"Posterior mean:
{
np
.
mean
(
pi_beta
[
0
])
}
"
)
with
open
(
'pi_beta_2020-03-15.pkl'
,
'wb'
)
as
f
:
with
open
(
'pi_beta_2020-03-15.pkl'
,
'wb'
)
as
f
:
pkl
.
dump
(
pi_beta
,
f
)
#dates = settings['start'] + t.numpy().astype(np.timedelta64)
#plotting(dates, sim)
ode_config.yaml
View file @
d2e2af66
...
...
@@ -8,18 +8,19 @@ data:
reported_cases
:
data/DailyConfirmedCases.csv
parameter
:
beta1
:
0
.0
35
# R0 2.4
#
beta1: 0.04
epsilon
:
1
.0
e-6
beta1
:
0.04
# R0 2.4
beta2
:
0.33
# Contact with commuters 1/3rd of the time
nu
:
0.25
gamma
:
0.25
settings
:
start
:
2020-02-04
end
:
2020-0
5
-01
end
:
2020-0
4
-01
holiday
:
-
2020-04-06
-
2020-04-17
bg_max_time
:
2020-03-01
time_step
:
1.
output
:
...
...
prediction.py
View file @
d2e2af66
...
...
@@ -31,6 +31,7 @@ if __name__ == '__main__':
N
,
n_names
=
load_population
(
config
[
'data'
][
'population_size'
])
param
=
sanitise_parameter
(
config
[
'parameter'
])
param
[
'epsilon'
]
=
0.0
settings
=
sanitise_settings
(
config
[
'settings'
])
case_reports
=
pd
.
read_csv
(
config
[
'data'
][
'reported_cases'
])
...
...
@@ -47,11 +48,11 @@ if __name__ == '__main__':
date_range
[
1
]
+
np
.
timedelta64
(
1
,
'D'
),
np
.
timedelta64
(
1
,
'D'
))
simulator
=
CovidUKODE
(
K_tt
,
K_hh
,
T
,
N
,
date_range
[
0
]
-
np
.
timedelta64
(
1
,
'D'
),
np
.
datetime64
(
'2020-0
7
-01'
),
settings
[
'holiday'
],
1
)
np
.
datetime64
(
'2020-0
5
-01'
),
settings
[
'holiday'
],
settings
[
'bg_max_time'
],
1
)
seeding
=
seed_areas
(
N
,
n_names
)
# Seed 40-44 age group, 30 seeds by popn size
state_init
=
simulator
.
create_initial_state
(
init_matrix
=
seeding
)
#
@tf.function
@
tf
.
function
def
prediction
(
beta
):
sims
=
tf
.
TensorArray
(
tf
.
float32
,
size
=
beta
.
shape
[
0
])
for
i
in
tf
.
range
(
beta
.
shape
[
0
]):
...
...
@@ -65,7 +66,7 @@ if __name__ == '__main__':
draws
=
pi_beta
[
0
].
numpy
()[
np
.
arange
(
0
,
pi_beta
[
0
].
shape
[
0
],
20
)]
sims
=
prediction
(
draws
)
dates
=
np
.
arange
(
date_range
[
0
]
-
np
.
timedelta64
(
1
,
'D'
),
np
.
datetime64
(
'2020-0
7
-01'
),
dates
=
np
.
arange
(
date_range
[
0
]
-
np
.
timedelta64
(
1
,
'D'
),
np
.
datetime64
(
'2020-0
5
-01'
),
np
.
timedelta64
(
1
,
'D'
))
total_infected
=
tfs
.
percentile
(
tf
.
reduce_sum
(
sims
[:,
:,
1
:
3
],
axis
=
2
),
q
=
[
2.5
,
50
,
97.5
],
axis
=
0
)
removed
=
tfs
.
percentile
(
sims
[:,
:,
3
],
q
=
[
2.5
,
50
,
97.5
],
axis
=
0
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment