Skip to content
GitLab
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
e26715b6
Commit
e26715b6
authored
Mar 19, 2020
by
Chris Jewell
Browse files
Implemented Haario and Sachs adaptive Multi-site RWM
parent
5e05d847
Changes
4
Hide whitespace changes
Inline
Side-by-side
covid/model.py
View file @
e26715b6
...
...
@@ -82,7 +82,7 @@ class Homogeneous:
def
dense_to_block_diagonal
(
A
,
n_blocks
):
A_dense
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
A
)
eye
=
tf
.
linalg
.
LinearOperatorIdentity
(
n_blocks
)
eye
=
tf
.
linalg
.
LinearOperatorIdentity
(
n_blocks
,
dtype
=
tf
.
float64
)
A_block
=
tf
.
linalg
.
LinearOperatorKronecker
([
eye
,
A_dense
])
return
A_block
...
...
@@ -101,24 +101,25 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
self
.
n_ages
=
M_tt
.
shape
[
0
]
self
.
n_lads
=
C
.
shape
[
0
]
self
.
Kbar
=
tf
.
reduce_mean
(
M_tt
)
self
.
M
=
tf
.
tuple
([
dense_to_block_diagonal
(
M_tt
,
self
.
n_lads
),
dense_to_block_diagonal
(
M_hh
,
self
.
n_lads
)])
self
.
Kbar
=
tf
.
reduce_mean
(
tf
.
cast
(
M_tt
,
tf
.
float64
)
)
self
.
M
=
tf
.
tuple
([
dense_to_block_diagonal
(
tf
.
cast
(
M_tt
,
tf
.
float64
)
,
self
.
n_lads
),
dense_to_block_diagonal
(
tf
.
cast
(
M_hh
,
tf
.
float64
)
,
self
.
n_lads
)])
C
=
tf
.
cast
(
C
,
tf
.
float64
)
self
.
C
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
C
+
tf
.
transpose
(
C
))
shp
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
np
.
ones_like
(
M_tt
,
dtype
=
np
.
float
32
))
shp
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
np
.
ones_like
(
M_tt
,
dtype
=
np
.
float
64
))
self
.
C
=
tf
.
linalg
.
LinearOperatorKronecker
([
self
.
C
,
shp
])
self
.
N
=
tf
.
constant
(
N
,
dtype
=
tf
.
float
32
)
self
.
N
=
tf
.
constant
(
N
,
dtype
=
tf
.
float
64
)
N_matrix
=
tf
.
reshape
(
self
.
N
,
[
self
.
n_lads
,
self
.
n_ages
])
N_sum
=
tf
.
reduce_sum
(
N_matrix
,
axis
=
1
)
N_sum
=
N_sum
[:,
None
]
*
tf
.
ones
([
1
,
self
.
n_ages
])
N_sum
=
N_sum
[:,
None
]
*
tf
.
ones
([
1
,
self
.
n_ages
]
,
dtype
=
tf
.
float64
)
self
.
N_sum
=
tf
.
reshape
(
N_sum
,
[
-
1
])
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
,
dtype
=
tf
.
int64
)
self
.
bg_select
=
tf
.
constant
(
np
.
less
(
bg_max_t
,
self
.
times
),
dtype
=
tf
.
int64
)
self
.
bg_select
=
tf
.
constant
(
np
.
less
(
self
.
times
,
bg_max_t
),
dtype
=
tf
.
int64
)
self
.
solver
=
tode
.
DormandPrince
()
def
make_h
(
self
,
param
):
...
...
@@ -128,7 +129,7 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
S
,
E
,
I
,
R
=
state
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
.
float
32
)
epsilon
=
param
[
'epsilon'
]
*
tf
.
cast
(
tf
.
gather
(
self
.
bg_select
,
t
),
tf
.
float
64
)
if
m_switch
==
0
:
infec_rate
=
param
[
'beta1'
]
*
tf
.
linalg
.
matvec
(
self
.
M
[
0
],
I
)
else
:
...
...
@@ -148,7 +149,7 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
def
create_initial_state
(
self
,
init_matrix
=
None
):
if
init_matrix
is
None
:
I
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float
32
)
I
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float
64
)
I
[
149
*
17
+
10
]
=
30.
# Middle-aged in Surrey
else
:
np
.
testing
.
assert_array_equal
(
init_matrix
.
shape
,
[
self
.
n_lads
,
self
.
n_ages
],
...
...
@@ -156,8 +157,8 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
(
{
self
.
n_lads
}
,
{
self
.
n_ages
}
)"
)
I
=
init_matrix
.
flatten
()
S
=
self
.
N
-
I
E
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float
32
)
R
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float
32
)
E
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float
64
)
R
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float
64
)
return
np
.
stack
([
S
,
E
,
I
,
R
])
@
tf
.
function
...
...
doc/age_space_model.tex
View file @
e26715b6
...
...
@@ -25,29 +25,28 @@ of Covid-19 across the UK, respecting the availability of human mobility
data as well as known contact behaviour between individuals of different
ages.
A deterministic SEIR model is posited in which infection rate is written
A deterministic SEIR state transition model is posited in which individuals
transition from Susceptible to Exposed (i.e. infected but not yet
infectious) to Infectious to Removed (i.e. quarantined, got better,
or died).
We model the infection rate (rate of S
$
\rightarrow
$
E transition)
as a function of known age-structured contact from Polymod, known
human mobility between MSOAs (Middle Super Output Area), and Census-derived
age structured population density in regions across the UK.
Currently, this model is populated with data for England only, though
we are in the process of extending this to Scotland and Wales.
Noise in daily case numbers
$
y
_{
it
}$
for age-group
$
i
$
in location
$
k
$
on day
$
t
$
is assumed to be Poisson-distributed such that
\[
y
_{
ikt
}
\sim\mbox
{
Poisson
}
(
R
_{
ikt
}
-
R
_{
ikt
-
1
}
)
\]
we are in the process of extending this to Scotland, Wales, and Northern
Ireland.
\section
{
Data
}
\subsection
{
Age-mixing
}
Standard Polymod data for the UK are used, with 17 5-year
age groups
$
[
0
-
5
)
,
[
5
-
10
)
,
\dots
,
[
75
-
80
)
,
[
80
-
\infty
)
$
. Estimated contact
matrices
for term-time
$
M
_{
tt
}$
and school-holidays
$
M
_{
hh
}$
were
extracted
of dimension
$
n
_{
m
}
\times
n
_{
m
}$
where
$
n
_{
m
}
=
17
$
.
Standard Polymod
social mixing
data for the UK are used, with 17 5-year
age groups
$
[
0
-
5
)
,
[
5
-
10
)
,
\dots
,
[
75
-
80
)
,
[
80
-
\infty
)
$
. Estimated contact
matrices
for term-time
$
M
_{
tt
}$
and school-holidays
$
M
_{
hh
}$
were
extracted
of dimension
$
n
_{
m
}
\times
n
_{
m
}$
where
$
n
_{
m
}
=
17
$
.
\subsection
{
Human mobility
}
...
...
@@ -67,8 +66,8 @@ with 0 diagonal.
Age-structured population size within each LAD is taken from publicly
available 2019 Local Authority data giving a vector
$
N
$
of length
$
n
_{
m
}
n
_{
c
}$
, i.e. population for each of
$
n
_{
m
}$
age groups
and
$
n
_{
c
}$
LADs.
$
n
_{
m
}
n
_{
c
}
=
2584
$
, i.e. population for each of
$
n
_{
m
}$
age groups
and
$
n
_{
c
}$
LADs.
\section
{
Model
}
...
...
@@ -77,16 +76,16 @@ $n_{c}$ LADs.
We assemble a country-wide connectivity matrices as Kronecker products,
such that
\[
K
^{
\star
}
=
I
_{
n
_{
c
}}
\bigotimes
M
M
^{
\star
}
=
I
_{
n
_{
c
}}
\bigotimes
M
\]
and
\[
T
^{
\star
}
=
C
\bigotimes\bm
{
1
}_{
n
_{
m
}
\times
n
_{
c
}}
C
^{
\star
}
=
C
\bigotimes\bm
{
1
}_{
n
_{
m
}
\times
n
_{
c
}}
\]
giving two matrices of dimension
$
n
_{
m
}
n
_{
c
}$
.
$
K
^{
\star
}$
is block
diagonal with Polymod mixing matrices.
$
T
^{
\star
}$
expands
the mobility
matrix
$
C
$
such that a block structure of connectivity
between LADs
results.
giving two matrices of dimension
$
n
_{
m
}
n
_{
c
}
\times
n
_{
m
}
n
_{
c
}$
.
$
M
^{
\star
}$
is block
diagonal with Polymod mixing matrices.
$
C
^{
\star
}$
expands
the mobility
matrix
$
C
$
such that a block structure of connectivity
between LADs
results.
\subsection
{
Disease progression model
}
...
...
@@ -95,16 +94,19 @@ number of individual in each age-group-LAD combination at time $t$
by the vectors
$
\vec
{
S
}
(
t
)
,
\vec
{
E
}
(
t
)
,
\vec
{
I
}
(
t
)
,
\vec
{
R
}
(
t
)
$
. We
therefore have
\begin{align*}
\frac
{
\mathrm
{
d
\vec
{
S
}
(t)
}}{
dt
}
&
=-
\beta
_{
1
}
\left
[
K
^{\star}\vec{I}(t)+\beta_{2}\bar{
K}T
^{\star}\frac{{\vec{I}(t)}}{N}\right]
\frac
{
\vec
{
S
}
(t)
}{
N
}
\\
\frac
{
\mathrm
{
d
}
\vec
{
E
}
(t)
}{
dt
}
&
=
\beta
_{
1
}
\left
[
K
^{\star}\vec{I}(t)+\beta_{2}\bar{
K}T
^{\star}\frac{{\vec{I}(t)}}{N}\right]
\frac
{
\vec
{
S
}
(t)
}{
N
}
-
\nu\vec
{
E
}
(t)
\\
\frac
{
\mathrm
{
d
\vec
{
S
}
(t)
}}{
dt
}
&
=-
\beta
_{
1
}
\left
[
M
^{\star}\vec{I}(t)+\beta_{2}\bar{
M}C
^{\star}\frac{{\vec{I}(t)}}{N}\right]
\frac
{
\vec
{
S
}
(t)
}{
N
}
\\
\frac
{
\mathrm
{
d
}
\vec
{
E
}
(t)
}{
dt
}
&
=
\beta
_{
1
}
\left
[
M
^{\star}\vec{I}(t)+\beta_{2}\bar{
M}C
^{\star}\frac{{\vec{I}(t)}}{N}\right]
\frac
{
\vec
{
S
}
(t)
}{
N
}
-
\nu\vec
{
E
}
(t)
\\
\frac
{
\mathrm
{
d
}
\vec
{
I
}
(t)
}{
dt
}
&
=
\nu\vec
{
E
}
(t)-
\gamma\vec
{
I
}
(t)
\\
\frac
{
\mathrm
{
d
}
\vec
{
R
}
(t)
}{
dt
}
&
=
\gamma\vec
{
I
}
(t)
\end{align*}
with parameters: baseline infection rate
$
\beta
_{
1
}$
, commuting infection
ratio
$
\beta
_{
2
}$
, latent period
$
\frac
{
1
}{
\nu
}$
, and infectious
period
$
\frac
{
1
}{
\gamma
}$
.
where
$
\bar
{
M
}$
is the global mean person-person contact rate. Parameters
are: baseline infection rate
$
\beta
_{
1
}$
, commuting infection ratio
$
\beta
_{
2
}$
, latent period
$
\frac
{
1
}{
\nu
}$
, and infectious period
$
\frac
{
1
}{
\gamma
}$
. Typically, we assume that contact with commuters
is
$
\beta
_{
2
}
=
\frac
{
1
}{
3
}$
of that between members of the same age-LAD
combination assuming an 8 hour working day.
\subsection
{
Noise model
}
\subsection
{
Noise model
(under construction)
}
Currently, and subject to discussion, we assume that all detected
cases are synonymous with individuals transitioning
$
I
\rightarrow
R
$
.
...
...
@@ -122,9 +124,7 @@ for Poisson overdispersion.
The model is currently implemented in Python3, using Tensorflow 2.1
with the RK5 differential equation solver implemented in the
\texttt
{
DormandPrince
}
class provided by Tensorflow Probability 0.9. Code may be found at
\url
{
http://fhm-chicas-code.lancs.ac.uk/jewell/covid19uk
}
.
\section
{
Quick results
}
\url
{
http://fhm-chicas-code.lancs.ac.uk/jewell/covid19uk
}
.
See attached SPI-M report
.
The code is a work in progress
\textendash
{}
see README
.
\end{document}
mcmc.py
View file @
e26715b6
...
...
@@ -8,6 +8,7 @@ import tensorflow_probability as tfp
from
tensorflow_probability
import
distributions
as
tfd
from
tensorflow_probability
import
bijectors
as
tfb
import
matplotlib.pyplot
as
plt
from
tensorflow_probability.python.util
import
SeedStream
from
covid.rdata
import
load_population
,
load_age_mixing
,
load_mobility_matrix
from
covid.model
import
CovidUKODE
,
covid19uk_logp
...
...
@@ -36,6 +37,18 @@ def plotting(dates, sim):
plt
.
show
()
def
random_walk_mvnorm_fn
(
covariance
,
name
=
None
):
"""Returns callable that adds Multivariate Normal noise to the input"""
rv
=
tfp
.
distributions
.
MultivariateNormalTriL
(
loc
=
tf
.
zeros
(
covariance
.
shape
[
0
],
dtype
=
tf
.
float64
),
scale_tril
=
tf
.
linalg
.
cholesky
(
tf
.
convert_to_tensor
(
covariance
,
dtype
=
tf
.
float64
)))
def
_fn
(
state_parts
,
seed
):
with
tf
.
name_scope
(
name
or
'random_walk_mvnorm_fn'
):
seed_stream
=
SeedStream
(
seed
,
salt
=
'RandomWalkNormalFn'
)
return
rv
.
sample
()
+
state_parts
if
__name__
==
'__main__'
:
parser
=
optparse
.
OptionParser
()
...
...
@@ -68,16 +81,18 @@ if __name__ == '__main__':
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
(
epsilon
,
beta
):
@
tf
.
function
def
logp
(
par
):
p
=
param
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'
])
p
[
'epsilon'
]
=
par
[
0
]
p
[
'beta1'
]
=
par
[
1
]
p
[
'gamma'
]
=
par
[
2
]
epsilon_logp
=
tfd
.
Gamma
(
concentration
=
tf
.
constant
(
1.
,
tf
.
float64
),
rate
=
tf
.
constant
(
1.
,
tf
.
float64
)).
log_prob
(
p
[
'epsilon'
])
beta_logp
=
tfd
.
Gamma
(
concentration
=
tf
.
constant
(
1.
,
tf
.
float64
),
rate
=
tf
.
constant
(
1.
,
tf
.
float64
)).
log_prob
(
p
[
'beta1'
])
gamma_logp
=
tfd
.
Gamma
(
concentration
=
tf
.
constant
(
100.
,
tf
.
float64
),
rate
=
tf
.
constant
(
400.
,
tf
.
float64
)).
log_prob
(
p
[
'gamma'
])
t
,
sim
,
solve
=
simulator
.
simulate
(
p
,
state_init
)
y_logp
=
covid19uk_logp
(
y_incr
,
sim
,
0.1
)
logp
=
epsilon_logp
+
beta_logp
+
tf
.
reduce_sum
(
y_logp
)
logp
=
epsilon_logp
+
beta_logp
+
gamma_logp
+
tf
.
reduce_sum
(
y_logp
)
return
logp
def
trace_fn
(
_
,
pkr
):
...
...
@@ -87,48 +102,51 @@ if __name__ == '__main__':
pkr
.
inner_results
.
accepted_results
.
step_size
)
unconstraining_bijector
=
[
tfb
.
Exp
(),
tfb
.
Exp
()]
initial_mcmc_state
=
[
0.00
0
0
1
,
0.03
]
print
(
"Initial log likelihood:"
,
logp
(
*
initial_mcmc_state
))
unconstraining_bijector
=
[
tfb
.
Exp
()]
initial_mcmc_state
=
np
.
array
(
[
0.001
,
0.03
6
,
0.25
],
dtype
=
np
.
float64
)
print
(
"Initial log likelihood:"
,
logp
(
initial_mcmc_state
))
@
tf
.
function
def
sample
(
n_samples
,
step_size
=
[
0.01
,
0.1
]
):
def
sample
(
n_samples
,
init_state
,
scale
):
return
tfp
.
mcmc
.
sample_chain
(
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
(
current_state
=
init_state
,
kernel
=
tfp
.
mcmc
.
TransformedTransitionKernel
(
inner_kernel
=
tfp
.
mcmc
.
RandomWalkMetropolis
(
target_log_prob_fn
=
logp
,
step_size
=
step_size
,
num_leapfrog_steps
=
5
),
new_state_fn
=
random_walk_mvnorm_fn
(
scale
)
),
bijector
=
unconstraining_bijector
),
num_adaptation_steps
=
500
),
trace_fn
=
lambda
_
,
pkr
:
pkr
.
inner_results
.
inner_results
.
accepted_results
.
step_size
)
trace_fn
=
lambda
_
,
pkr
:
pkr
.
inner_results
.
is_accepted
)
with
tf
.
device
(
"/CPU:0"
):
cov
=
np
.
diag
([
0.001
,
0.03
,
1.02
])
*
2.38
**
2
/
3
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)
joint_posterior
,
results
=
sample
(
100
,
init_state
=
initial_mcmc_state
,
scale
=
cov
)
for
i
in
range
(
10
):
cov
=
tfp
.
stats
.
covariance
(
tf
.
math
.
log
(
joint_posterior
))
*
2.38
**
2
/
joint_posterior
.
shape
[
0
]
+
tf
.
eye
(
cov
.
shape
[
0
],
dtype
=
tf
.
float64
)
*
1.e-7
print
(
cov
.
numpy
())
posterior_new
,
results
=
sample
(
100
,
joint_posterior
[
-
1
,
:],
cov
)
joint_posterior
=
tf
.
concat
([
joint_posterior
,
posterior_new
],
axis
=
0
)
posterior_new
,
results
=
sample
(
2000
,
init_state
=
joint_posterior
[
-
1
,
:],
scale
=
cov
)
joint_posterior
=
tf
.
concat
([
joint_posterior
,
posterior_new
],
axis
=
0
)
end
=
time
.
perf_counter
()
print
(
f
"Simulation complete in
{
end
-
start
}
seconds"
)
print
(
results
)
print
(
"Acceptance: "
,
np
.
mean
(
results
.
numpy
()))
print
(
tfp
.
stats
.
covariance
(
tf
.
math
.
log
(
joint_posterior
[
1000
:,
:])))
fig
,
ax
=
plt
.
subplots
(
1
,
2
)
ax
[
0
].
plot
(
pi_beta
[
0
])
ax
[
1
].
plot
(
pi_beta
[
1
])
fig
,
ax
=
plt
.
subplots
(
1
,
3
)
ax
[
0
].
plot
(
joint_posterior
[:,
0
])
ax
[
1
].
plot
(
joint_posterior
[:,
1
])
ax
[
2
].
plot
(
joint_posterior
[:,
2
])
plt
.
show
()
print
(
f
"Posterior mean:
{
np
.
mean
(
pi_beta
[
0
]
)
}
"
)
print
(
f
"Posterior mean:
{
np
.
mean
(
joint_posterior
,
axis
=
0
)
}
"
)
with
open
(
'pi_beta_2020-03-15.pkl'
,
'wb'
)
as
f
:
pkl
.
dump
(
pi_beta
,
f
)
pkl
.
dump
(
joint_posterior
,
f
)
#dates = settings['start'] + t.numpy().astype(np.timedelta64)
#plotting(dates, sim)
prediction.py
View file @
e26715b6
"""Prediction functions"""
import
optparse
import
seaborn
import
yaml
import
pickle
as
pkl
import
numpy
as
np
...
...
@@ -7,10 +9,22 @@ import pandas as pd
import
tensorflow
as
tf
from
tensorflow_probability
import
stats
as
tfs
import
matplotlib.pyplot
as
plt
import
h5py
from
covid.model
import
CovidUKODE
from
covid.rdata
import
load_age_mixing
,
load_mobility_matrix
,
load_population
from
covid.util
import
sanitise_settings
,
sanitise_parameter
,
seed_areas
from
covid.util
import
sanitise_settings
,
sanitise_parameter
,
seed_areas
,
doubling_time
def
save_sims
(
sims
,
la_names
,
age_groups
,
filename
):
f
=
h5py
.
File
(
filename
,
'w'
)
dset_sim
=
f
.
create_dataset
(
'prediction'
,
data
=
sims
)
la_long
=
np
.
repeat
(
la_names
,
age_groups
.
shape
[
0
]).
astype
(
np
.
string_
)
age_long
=
np
.
tile
(
age_groups
,
la_names
.
shape
[
0
]).
astype
(
np
.
string_
)
dset_dims
=
f
.
create_dataset
(
"dimnames"
,
data
=
[
b
'sim_id'
,
b
't'
,
b
'state'
,
b
'la_age'
])
dset_la
=
f
.
create_dataset
(
'la_names'
,
data
=
la_long
)
dset_age
=
f
.
create_dataset
(
'age_names'
,
data
=
age_long
)
f
.
close
()
if
__name__
==
'__main__'
:
...
...
@@ -48,25 +62,42 @@ 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
5
-01'
),
settings
[
'holiday'
],
settings
[
'bg_max_time'
],
1
)
np
.
datetime64
(
'2020-0
9
-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
def
prediction
(
beta
):
def
prediction
(
epsilon
,
beta
):
sims
=
tf
.
TensorArray
(
tf
.
float32
,
size
=
beta
.
shape
[
0
])
R0
=
tf
.
TensorArray
(
tf
.
float32
,
size
=
beta
.
shape
[
0
])
#d_time = tf.TensorArray(tf.float32, size=beta.shape[0])
for
i
in
tf
.
range
(
beta
.
shape
[
0
]):
p
=
param
p
[
'epsilon'
]
=
epsilon
[
i
]
p
[
'beta1'
]
=
beta
[
i
]
t
,
sim
,
solver_results
=
simulator
.
simulate
(
p
,
state_init
)
sim_aggr
=
tf
.
reduce_sum
(
sim
,
axis
=
2
)
sims
=
sims
.
write
(
i
,
sim_aggr
)
return
sims
.
gather
(
range
(
beta
.
shape
[
0
]))
r
=
simulator
.
eval_R0
(
p
)
R0
=
R0
.
write
(
i
,
r
[
0
])
#d_time = d_time.write(i, doubling_time(t, sim, '2002-03-01', '2002-04-01'))
#sim_aggr = tf.reduce_sum(sim, axis=2)
sims
=
sims
.
write
(
i
,
sim
)
return
sims
.
gather
(
range
(
beta
.
shape
[
0
])),
R0
.
gather
(
range
(
beta
.
shape
[
0
]))
draws
=
[
pi_beta
[
0
].
numpy
()[
np
.
arange
(
500
,
pi_beta
[
0
].
shape
[
0
],
10
)],
pi_beta
[
1
].
numpy
()[
np
.
arange
(
500
,
pi_beta
[
1
].
shape
[
0
],
10
)]]
with
tf
.
device
(
'/CPU:0'
):
sims
,
R0
=
prediction
(
draws
[
0
],
draws
[
1
])
sims
=
tf
.
stack
(
sims
)
# shape=[n_sims, n_times, n_states, n_metapops]
save_sims
(
sims
,
la_names
,
age_groups
,
'pred_2020-03-15.h5'
)
draws
=
pi_beta
[
0
].
numpy
()[
np
.
arange
(
0
,
pi_beta
[
0
].
shape
[
0
],
20
)]
sims
=
prediction
(
draws
)
dub_time
=
[
doubling_time
(
simulator
.
times
,
sim
,
'2020-03-01'
,
'2020-04-01'
)
for
sim
in
sims
.
numpy
()]
dates
=
np
.
arange
(
date_range
[
0
]
-
np
.
timedelta64
(
1
,
'D'
),
np
.
datetime64
(
'2020-05-01'
),
# Sum over country
sims
=
tf
.
reduce_sum
(
sims
,
axis
=
3
)
print
(
"Plotting..."
,
flush
=
True
)
dates
=
np
.
arange
(
date_range
[
0
]
-
np
.
timedelta64
(
1
,
'D'
),
np
.
datetime64
(
'2020-09-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
)
...
...
@@ -98,3 +129,15 @@ if __name__ == '__main__':
plt
.
ylabel
(
"Incidence per 10,000"
)
fig
.
autofmt_xdate
()
plt
.
show
()
# R0
R0_ci
=
tfs
.
percentile
(
R0
,
q
=
[
2.5
,
50
,
97.5
])
print
(
"R0:"
,
R0_ci
)
fig
=
plt
.
figure
()
seaborn
.
kdeplot
(
R0
.
numpy
(),
ax
=
fig
.
gca
())
plt
.
title
(
"R0"
)
plt
.
show
()
# Doubling time
dub_ci
=
tfs
.
percentile
(
dub_time
,
q
=
[
2.5
,
50
,
97.5
])
print
(
"Doubling time:"
,
dub_ci
)
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new 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