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
996a4afc
Commit
996a4afc
authored
Mar 08, 2020
by
Chris Jewell
Browse files
Implemented household sim
parent
750c80c7
Changes
8
Hide whitespace changes
Inline
Side-by-side
covid/impl/chainbinom_simulate.py
View file @
996a4afc
...
...
@@ -31,7 +31,7 @@ def chain_binomial_simulate(hazard_fn, state, start, end, time_step, stoichiomet
propagate
=
chain_binomial_propagate
(
hazard_fn
,
time_step
,
stoichiometry
)
times
=
tf
.
range
(
start
,
end
,
time_step
)
output
=
tf
.
TensorArray
(
tf
.
float
32
,
size
=
times
.
shape
[
0
])
output
=
tf
.
TensorArray
(
tf
.
float
64
,
size
=
times
.
shape
[
0
])
output
=
output
.
write
(
0
,
state
)
for
i
in
tf
.
range
(
1
,
times
.
shape
[
0
]):
...
...
covid/model.py
View file @
996a4afc
...
...
@@ -8,6 +8,25 @@ import numpy as np
from
covid.impl.chainbinom_simulate
import
chain_binomial_simulate
def
power_iteration
(
A
,
tol
=
1e-3
):
b_k
=
tf
.
random
.
normal
([
A
.
shape
[
1
],
1
])
epsilon
=
1.
i
=
0
while
tf
.
greater
(
epsilon
,
tol
):
b_k1
=
tf
.
matmul
(
A
,
b_k
)
b_k1_norm
=
tf
.
linalg
.
norm
(
b_k1
)
b_k_new
=
b_k1
/
b_k1_norm
epsilon
=
tf
.
reduce_sum
(
tf
.
pow
(
b_k_new
-
b_k
,
2
))
b_k
=
b_k_new
i
+=
1
return
b_k
,
i
#@tf.function
def
rayleigh_quotient
(
A
,
b
):
b
=
tf
.
reshape
(
b
,
[
b
.
shape
[
0
],
1
])
numerator
=
tf
.
matmul
(
tf
.
transpose
(
b
),
tf
.
matmul
(
A
,
b
))
denominator
=
tf
.
matmul
(
tf
.
transpose
(
b
),
b
)
return
numerator
/
denominator
class
CovidUK
:
def
__init__
(
self
,
K
,
T
,
W
):
...
...
@@ -19,19 +38,18 @@ class CovidUK:
[
0
,
-
1
,
1
,
0
],
[
0
,
0
,
-
1
,
1
]]
@
tf
.
function
def
h
(
self
,
state
):
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
hazard_rates
=
tf
.
stack
([
self
.
param
[
'beta1'
]
*
tf
.
dot
(
self
.
T
,
tf
.
dot
(
self
.
K
,
I
)),
self
.
param
[
'beta1'
]
*
tf
.
dot
(
self
.
T
,
tf
.
dot
(
self
.
K
,
I
))
/
self
.
K
.
shape
[
0
]
,
self
.
param
[
'nu'
],
self
.
param
[
'gamma'
]
])
return
hazard_rates
@
tf
.
function
def
sample
(
self
,
initial_state
,
time_lims
,
param
):
self
.
param
=
param
return
chain_binomial_simulate
(
self
.
h
,
initial_state
,
time_lims
[
0
],
...
...
@@ -64,34 +82,42 @@ class Homogeneous:
class
CovidUKODE
:
def
__init__
(
self
,
K
,
T
,
N
,
n_lads
):
def
__init__
(
self
,
K
,
T
,
N
):
"""Represents a CovidUK ODE model
:param K: a MxM matrix of age group mixing
:param T: a n_ladsxn_lads matrix of inter-LAD co
nnectivity
:param K: a MxM matrix of age group mixing
in term time
:param T: a n_ladsxn_lads matrix of inter-LAD co
mmuting
:param N: a vector of population sizes in each LAD
:param n_lads: the number of LADS
"""
self
.
n_ages
=
K
.
shape
[
0
]
self
.
n_lads
=
T
.
shape
[
0
]
self
.
Kbar
=
tf
.
reduce_mean
(
K
)
self
.
K
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
K
)
eye
=
tf
.
linalg
.
LinearOperatorIdentity
(
n_lads
)
self
.
K
=
tf
.
linalg
.
LinearOperatorKronecker
([
self
.
K
,
eye
])
eye
=
tf
.
linalg
.
LinearOperatorIdentity
(
self
.
n_lads
)
self
.
K
=
tf
.
linalg
.
LinearOperatorKronecker
([
eye
,
self
.
K
])
self
.
T
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
T
)
shp
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
np
.
ones_like
(
K
,
dtype
=
np
.
float32
))
self
.
T
=
tf
.
linalg
.
LinearOperatorKronecker
([
self
.
T
,
shp
])
self
.
N
=
N
self
.
n_lads
=
n_lads
self
.
N
=
tf
.
constant
(
N
,
dtype
=
tf
.
float32
)
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
])
self
.
N_sum
=
tf
.
reshape
(
N_sum
,
[
-
1
])
def
make_h
(
self
,
param
):
def
h_fn
(
t
,
state
):
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
infec_rate
=
tf
.
linalg
.
matvec
(
self
.
K
,
I
)
infec_rate
+=
self
.
Kbar
*
tf
.
linalg
.
matvec
(
self
.
T
,
I
/
self
.
N
)
infec_rate
=
param
[
'beta'
]
*
S
/
self
.
N
*
infec_rate
infec_rate
=
param
[
'beta1'
]
*
tf
.
linalg
.
matvec
(
self
.
K
,
I
)
infec_rate
+=
param
[
'beta1'
]
*
param
[
'beta2'
]
*
self
.
Kbar
*
tf
.
linalg
.
matvec
(
self
.
T
,
I
/
self
.
N
_sum
)
infec_rate
=
S
/
self
.
N
*
infec_rate
dS
=
-
infec_rate
dE
=
infec_rate
-
param
[
'nu'
]
*
E
...
...
@@ -105,17 +131,37 @@ class CovidUKODE:
def
create_initial_state
(
self
,
init_matrix
=
None
):
if
init_matrix
is
None
:
I
=
np
.
ones
(
self
.
N
.
shape
,
dtype
=
np
.
float32
)
S
=
self
.
N
-
I
E
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float32
)
R
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float32
)
return
np
.
stack
([
S
,
E
,
I
,
R
])
I
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float32
)
I
[
149
*
17
+
10
]
=
30.
# Middle-aged in Surrey
else
:
np
.
testing
.
assert_array_equal
(
init_matrix
.
shape
,
[
self
.
n_lads
,
self
.
n_ages
],
err_msg
=
f
"init_matrix does not have shape [<num lads>,<num ages>]
\
(
{
self
.
n_lads
}
,
{
self
.
n_ages
}
)"
)
I
=
init_matrix
.
flatten
()
S
=
self
.
N
-
I
E
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float32
)
R
=
np
.
zeros
(
self
.
N
.
shape
,
dtype
=
np
.
float32
)
return
np
.
stack
([
S
,
E
,
I
,
R
])
@
tf
.
function
def
simulate
(
self
,
param
,
state_init
,
t_start
,
t_end
,
t_step
=
1.
):
h
=
self
.
make_h
(
param
)
t
=
np
.
arange
(
start
=
t_start
,
stop
=
t_end
,
step
=
t_step
)
print
(
f
"Running simulation with times
{
t_start
}
-
{
t_end
}
:
{
t_step
}
"
)
t0
=
0.
t1
=
(
t_end
-
t_start
)
/
np
.
timedelta64
(
1
,
'D'
)
t
=
np
.
arange
(
start
=
t0
,
stop
=
t1
,
step
=
t_step
)[
1
:]
solver
=
tode
.
DormandPrince
()
results
=
solver
.
solve
(
ode_fn
=
h
,
initial_time
=
t
_start
,
initial_state
=
state_init
,
solution_times
=
t
)
results
=
solver
.
solve
(
ode_fn
=
h
,
initial_time
=
t
0
,
initial_state
=
state_init
,
solution_times
=
t
)
return
results
.
times
,
results
.
states
def
ngm
(
self
,
param
):
infec_rate
=
param
[
'beta1'
]
*
self
.
K
.
to_dense
()
infec_rate
+=
param
[
'beta1'
]
*
param
[
'beta2'
]
*
self
.
Kbar
*
self
.
T
.
to_dense
()
/
self
.
N_sum
[
None
,
:]
ngm
=
infec_rate
/
param
[
'gamma'
]
return
ngm
def
eval_R0
(
self
,
param
,
tol
=
1e-8
):
ngm
=
self
.
ngm
(
param
)
# Dominant eigen value by power iteration
dom_eigen_vec
,
i
=
power_iteration
(
ngm
,
tol
=
tol
)
R0
=
rayleigh_quotient
(
ngm
,
dom_eigen_vec
)
return
tf
.
squeeze
(
R0
),
i
covid/rdata.py
View file @
996a4afc
...
...
@@ -2,6 +2,7 @@
import
numpy
import
pyreadr
as
pyr
import
numpy
as
np
def
load_age_mixing
(
rds_file
:
str
):
...
...
@@ -10,7 +11,9 @@ def load_age_mixing(rds_file: str):
:param rds_file: a .rds file containing an R data.frame with mixing matrix
"""
raw
=
pyr
.
read_r
(
rds_file
)
return
list
(
raw
.
values
())[
0
]
K
=
list
(
raw
.
values
())[
0
]
age_groups
=
K
.
columns
return
K
.
to_numpy
(
dtype
=
np
.
float32
),
age_groups
def
load_mobility_matrix
(
rds_file
:
str
):
...
...
@@ -24,7 +27,7 @@ def load_mobility_matrix(rds_file: str):
colnames
=
df
.
columns
mobility_matrix
=
df
.
pivot
(
index
=
'Workplace'
,
columns
=
'Residence'
,
values
=
colnames
[
0
])
mobility_matrix
[
mobility_matrix
.
isna
()]
=
0.
return
mobility_matrix
return
mobility_matrix
.
to_numpy
(
dtype
=
np
.
float32
),
mobility_matrix
.
index
.
to_numpy
()
def
load_population
(
rds_file
:
str
):
...
...
@@ -35,4 +38,4 @@ def load_population(rds_file: str):
raw
=
pyr
.
read_r
(
rds_file
)
df
=
list
(
raw
.
values
())[
0
]
df
=
df
.
sort_values
(
by
=
[
'LA.code'
,
'age'
])
return
df
return
df
[
'n'
].
to_numpy
(
dtype
=
np
.
float32
),
df
[[
'name'
,
'Area.name.2'
]]
covid_ode.py
View file @
996a4afc
import
optparse
import
numpy
as
np
import
yaml
import
time
import
h5py
import
matplotlib.pyplot
as
plt
import
tensorflow
as
tf
import
yaml
from
covid.model
import
CovidUKODE
from
covid.rdata
import
*
...
...
@@ -10,18 +12,149 @@ from covid.rdata import *
def
sanitise_parameter
(
par_dict
):
"""Sanitises a dictionary of parameters"""
par
=
[
'beta'
,
'nu'
,
'gamma'
]
par
=
[
'beta
1
'
,
'beta2'
,
'nu'
,
'gamma'
]
d
=
{
key
:
np
.
float64
(
par_dict
[
key
])
for
key
in
par
}
return
d
def
sanitise_settings
(
par_dict
):
d
=
{
'start'
:
np
.
float64
(
par_dict
[
'start'
]),
'end'
:
np
.
float64
(
par_dict
[
'end'
]),
'time_step'
:
np
.
float64
(
par_dict
[
'time_step'
])}
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'
]])}
return
d
def
seed_areas
(
N
,
names
,
age_group
=
8
,
num_la
=
152
,
num_age
=
17
,
n_seed
=
30.
):
areas
=
[
'Inner London'
,
'Outer London'
,
'West Midlands (Met County)'
,
'Greater Manchester (Met County)'
]
names_matrix
=
names
[
'Area.name.2'
].
to_numpy
().
reshape
([
num_la
,
num_age
])
seed_areas
=
np
.
in1d
(
names_matrix
[:,
age_group
],
areas
)
N_matrix
=
N
.
reshape
([
num_la
,
num_age
])
# LA x Age
pop_size_sub
=
N_matrix
[
seed_areas
,
age_group
]
# Gather
n
=
np
.
round
(
n_seed
*
pop_size_sub
/
pop_size_sub
.
sum
())
seeding
=
np
.
zeros_like
(
N_matrix
)
seeding
[
seed_areas
,
age_group
]
=
n
# Scatter
return
seeding
def
sum_age_groups
(
sim
):
infec
=
sim
[:,
2
,
:]
infec
=
infec
.
reshape
([
infec
.
shape
[
0
],
152
,
17
])
infec_uk
=
infec
.
sum
(
axis
=
2
)
return
infec_uk
def
sum_la
(
sim
):
infec
=
sim
[:,
2
,
:]
infec
=
infec
.
reshape
([
infec
.
shape
[
0
],
152
,
17
])
infec_uk
=
infec
.
sum
(
axis
=
1
)
return
infec_uk
def
sum_total_removals
(
sim
):
remove
=
sim
[:,
3
,
:]
return
remove
.
sum
(
axis
=
1
)
def
final_size
(
sim
):
remove
=
sim
[:,
3
,
:]
remove
=
remove
.
reshape
([
remove
.
shape
[
0
],
152
,
17
])
fs
=
remove
[
-
1
,
:,
:].
sum
(
axis
=
0
)
return
fs
def
write_hdf5
(
filename
,
param
,
t
,
sim
):
with
h5py
.
File
(
filename
,
"w"
)
as
f
:
dset_sim
=
f
.
create_dataset
(
"simulation"
,
sim
.
shape
,
dtype
=
'f'
)
dset_sim
[:]
=
sim
dset_t
=
f
.
create_dataset
(
"time"
,
t
.
shape
,
dtype
=
'f'
)
dset_t
[:]
=
t
grp_param
=
f
.
create_group
(
"parameter"
)
for
k
,
v
in
param
.
items
():
d_beta
=
grp_param
.
create_dataset
(
k
,
[
1
],
dtype
=
'f'
)
d_beta
[()]
=
v
def
plot_total_curve
(
sim
):
infec_uk
=
sum_la
(
sim
)
infec_uk
=
infec_uk
.
sum
(
axis
=
1
)
removals
=
sum_total_removals
(
sim
)
times
=
np
.
datetime64
(
'2020-02-20'
)
+
np
.
arange
(
removals
.
shape
[
0
])
plt
.
plot
(
times
,
infec_uk
,
'r-'
,
label
=
'Infected'
)
plt
.
plot
(
times
,
removals
,
'b-'
,
label
=
'Removed'
)
plt
.
title
(
'UK total cases'
)
plt
.
xlabel
(
'Date'
)
plt
.
ylabel
(
'Num infected or removed'
)
plt
.
legend
()
def
plot_by_age
(
sim
,
labels
,
t0
=
np
.
datetime64
(
'2020-02-20'
),
ax
=
None
):
if
ax
is
None
:
ax
=
plt
.
figure
().
gca
()
infec_uk
=
sum_la
(
sim
)
total_uk
=
infec_uk
.
mean
(
axis
=
1
)
t
=
t0
+
np
.
arange
(
infec_uk
.
shape
[
0
])
colours
=
plt
.
cm
.
viridis
(
np
.
linspace
(
0.
,
1.
,
infec_uk
.
shape
[
1
]))
for
i
in
range
(
infec_uk
.
shape
[
1
]):
ax
.
plot
(
t
,
infec_uk
[:,
i
],
'r-'
,
alpha
=
0.4
,
color
=
colours
[
i
],
label
=
labels
[
i
])
ax
.
plot
(
t
,
total_uk
,
'-'
,
color
=
'black'
,
label
=
'Mean'
)
return
ax
def
plot_by_la
(
sim
,
labels
,
t0
=
np
.
datetime64
(
'2020-02-20'
),
ax
=
None
):
if
ax
is
None
:
ax
=
plt
.
figure
().
gca
()
infec_uk
=
sum_age_groups
(
sim
)
total_uk
=
infec_uk
.
mean
(
axis
=
1
)
t
=
t0
+
np
.
arange
(
infec_uk
.
shape
[
0
])
colours
=
plt
.
cm
.
viridis
(
np
.
linspace
(
0.
,
1.
,
infec_uk
.
shape
[
1
]))
for
i
in
range
(
infec_uk
.
shape
[
1
]):
ax
.
plot
(
t
,
infec_uk
[:,
i
],
'r-'
,
alpha
=
0.4
,
color
=
colours
[
i
],
label
=
labels
[
i
])
ax
.
plot
(
t
,
total_uk
,
'-'
,
color
=
'black'
,
label
=
'Mean'
)
return
ax
def
draw_figs
(
sim
,
N
):
# Attack rate
N
=
N
.
reshape
([
152
,
17
]).
sum
(
axis
=
0
)
fs
=
final_size
(
sim
)
attack_rate
=
fs
/
N
print
(
"Attack rate:"
,
attack_rate
)
print
(
"Overall attack rate: "
,
np
.
sum
(
fs
)
/
np
.
sum
(
N
))
# Total UK epidemic curve
plot_total_curve
(
sim
)
plt
.
xticks
(
rotation
=
45
,
horizontalalignment
=
"right"
)
plt
.
savefig
(
'total_uk_curve.pdf'
)
plt
.
show
()
# TotalUK epidemic curve by age-group
fig
,
ax
=
plt
.
subplots
(
1
,
2
,
figsize
=
[
24
,
12
])
plot_by_la
(
sim
,
la_names
,
ax
=
ax
[
0
])
plot_by_age
(
sim
,
age_groups
,
ax
=
ax
[
1
])
ax
[
1
].
legend
()
plt
.
xticks
(
rotation
=
45
,
horizontalalignment
=
"right"
)
fig
.
autofmt_xdate
()
plt
.
savefig
(
'la_age_infec_curves.pdf'
)
plt
.
show
()
# Plot attack rate
plt
.
figure
(
figsize
=
[
4
,
2
])
plt
.
plot
(
age_groups
,
attack_rate
,
'o-'
)
plt
.
xticks
(
rotation
=
90
)
plt
.
title
(
'Age-specific attack rate'
)
plt
.
savefig
(
'age_attack_rate.pdf'
)
plt
.
show
()
if
__name__
==
'__main__'
:
parser
=
optparse
.
OptionParser
()
...
...
@@ -31,22 +164,52 @@ if __name__ == '__main__':
with
open
(
options
.
config
,
'r'
)
as
ymlfile
:
config
=
yaml
.
load
(
ymlfile
)
K
=
load_age_mixing
(
config
[
'data'
][
'age_mixing_matrix'
]).
to_numpy
(
dtype
=
np
.
float32
)
T
=
load_mobility_matrix
(
config
[
'data'
][
'mobility_matrix'
]).
to_numpy
(
dtype
=
np
.
float32
)
K_tt
,
age_groups
=
load_age_mixing
(
config
[
'data'
][
'age_mixing_matrix_term'
])
K_hh
,
_
=
load_age_mixing
(
config
[
'data'
][
'age_mixing_matrix_hol'
])
T
,
la_names
=
load_mobility_matrix
(
config
[
'data'
][
'mobility_matrix'
])
np
.
fill_diagonal
(
T
,
0.
)
N
=
load_population
(
config
[
'data'
][
'population_size'
])[
'n'
].
to_numpy
(
dtype
=
np
.
float32
)
N
,
n_names
=
load_population
(
config
[
'data'
][
'population_size'
])
param
=
sanitise_parameter
(
config
[
'parameter'
])
settings
=
sanitise_settings
(
config
[
'settings'
])
model
=
CovidUKODE
(
K
,
T
,
N
,
T
.
shape
[
0
])
state_init
=
model
.
create_initial_state
()
# Straight, no school closures
model_term
=
CovidUKODE
(
K_tt
,
T
,
N
)
model_holiday
=
CovidUKODE
(
K_hh
/
2.
,
T
/
2.
,
N
)
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
=
model_term
.
simulate
(
param
,
state_init
,
settings
[
'start'
],
settings
[
'holiday'
][
0
],
settings
[
'time_step'
])
t1
,
sim_1
=
model_holiday
.
simulate
(
param
,
sim_0
[
-
1
,
:,
:],
settings
[
'holiday'
][
0
],
settings
[
'holiday'
][
1
],
settings
[
'time_step'
])
t2
,
sim_2
=
model_term
.
simulate
(
param
,
sim_1
[
-
1
,
:,
:],
settings
[
'holiday'
][
1
],
settings
[
'end'
],
settings
[
'time_step'
])
t
=
tf
.
concat
([
t0
,
t1
,
t2
],
axis
=
0
)
sim
=
tf
.
concat
([
tf
.
expand_dims
(
state_init
,
axis
=
0
),
sim_0
,
sim_1
,
sim_2
],
axis
=
0
)
return
t
,
sim
start
=
time
.
perf_counter
()
t
,
sim
=
model
.
simulate
(
param
,
state_init
,
settings
[
'start'
],
settings
[
'end'
],
settings
[
'time_step'
])
t
,
sim
=
simulate
()
end
=
time
.
perf_counter
()
print
(
f
"Complete in
{
end
-
start
}
seconds"
)
plt
.
plot
(
t
,
sim
[:,
2
,
0
],
'r-'
)
plt
.
show
()
\ No newline at end of file
print
(
f
'Complete in
{
end
-
start
}
seconds'
)
draw_figs
(
sim
.
numpy
(),
N
)
if
'simulation'
in
config
[
'output'
]:
write_hdf5
(
config
[
'output'
][
'simulation'
],
param
,
t
,
sim
)
print
(
f
"Complete in
{
end
-
start
}
seconds"
)
\ No newline at end of file
household_sim.py
0 → 100644
View file @
996a4afc
"""Household simulation for SPI-M"""
import
time
import
matplotlib.pyplot
as
plt
import
numpy
as
np
import
pandas
as
pd
import
tensorflow
as
tf
from
covid.model
import
CovidUK
# Household size distribution in the UK, taken from ONS, 2019
hh_size_distr
=
pd
.
DataFrame
({
'size'
:
[
1
,
2
,
3
,
4
,
5
,
6
],
'freq'
:
[
6608
,
8094
,
3918
,
3469
,
1186
,
463
]})
# Draw households
def
draw_households
(
n
:
int
,
dist
:
pd
.
DataFrame
)
->
np
.
array
:
probs
=
dist
[
'freq'
]
/
np
.
sum
(
dist
[
'freq'
])
hh_sizes
=
np
.
random
.
choice
(
dist
[
'size'
],
n
,
p
=
probs
)
#hh_vec = np.repeat(np.arange(hh_sizes.shape[0]), hh_sizes)
return
hh_sizes
# Simulator
class
HHstochastic
(
CovidUK
):
def
__init__
(
self
,
N
):
self
.
stoichiometry
=
tf
.
constant
([[
-
1
,
1
,
0
,
0
],
[
0
,
-
1
,
1
,
0
],
[
0
,
0
,
-
1
,
1
]],
dtype
=
tf
.
float64
)
self
.
N
=
N
self
.
popsize
=
tf
.
reduce_sum
(
N
)
def
h
(
self
,
state
):
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
eye
=
tf
.
eye
(
I
.
shape
[
1
],
dtype
=
tf
.
float64
)
infec_rate
=
tf
.
linalg
.
tensordot
(
I
,
eye
,
axes
=
[[
1
],
[
1
]])
infec_rate
/=
self
.
N
infec_rate
*=
self
.
param
[
'omega'
]
infec_rate
+=
tf
.
reduce_sum
(
I
,
axis
=
1
,
keepdims
=
True
)
/
self
.
popsize
infec_rate
*=
self
.
param
[
'beta'
]
hazard_rates
=
tf
.
stack
([
infec_rate
,
tf
.
constant
(
np
.
full
(
S
.
shape
,
self
.
param
[
'nu'
]),
dtype
=
tf
.
float64
),
tf
.
constant
(
np
.
full
(
S
.
shape
,
self
.
param
[
'gamma'
]),
dtype
=
tf
.
float64
)
])
return
hazard_rates
if
__name__
==
'__main__'
:
hh
=
draw_households
(
100
,
hh_size_distr
)
K
=
(
hh
[:,
None
]
==
hh
[
None
,
:])
nsim
=
120
model
=
HHstochastic
(
hh
.
astype
(
np
.
float64
))
init_state
=
np
.
stack
([
np
.
broadcast_to
(
hh
,
[
nsim
,
hh
.
shape
[
0
]]),
np
.
zeros
([
nsim
,
hh
.
shape
[
0
]]),
np
.
zeros
([
nsim
,
hh
.
shape
[
0
]]),
np
.
zeros
([
nsim
,
hh
.
shape
[
0
]])]).
astype
(
np
.
float64
)
init_state
[
2
,
:,
0
]
=
1.
init_state
[
0
,
:,
0
]
=
init_state
[
0
,
:,
0
]
-
1.
start
=
time
.
perf_counter
()
t
,
sim
=
model
.
sample
(
init_state
,
[
0.
,
365.
],
{
'beta'
:
0.3
,
'omega'
:
1.5
,
'nu'
:
0.25
,
'gamma'
:
0.25
})
end
=
time
.
perf_counter
()
print
(
f
"Completed in
{
end
-
start
}
seconds"
)
r
=
tf
.
reduce_sum
(
sim
[:,
2
,
:,
:],
axis
=
2
)
plt
.
plot
(
t
,
r
,
'r-'
,
alpha
=
0.2
)
plt
.
show
()
ode_config.yaml
View file @
996a4afc
# Covid_ODE model configuration
data
:
age_mixing_matrix
:
data/polymod_normal_df.rds
age_mixing_matrix_term
:
data/polymod_normal_df.rds
age_mixing_matrix_hol
:
data/polymod_no_school_df.rds
mobility_matrix
:
data/movement.rds
population_size
:
data/pop.rds
parameter
:
beta
:
0.09
beta1
:
0.04
# R0 2.4
beta2
:
0.33
# Contact with commuters 1/3rd of the time
nu
:
0.25
gamma
:
0.
6
5
gamma
:
0.
2
5
settings
:
start
:
0.
end
:
365.
start
:
2020-02-04
end
:
2020-12-30
holiday
:
-
2020-04-06
-
2020-05-01
time_step
:
1.
output
:
simulation
:
covid_ode.
sim
simulation
:
covid_ode.
hd5
r0_calc.py
0 → 100644
View file @
996a4afc
import
numpy
as
np
from
scipy
import
optimize
as
opt
from
covid.rdata
import
*
from
covid.model
import
CovidUKODE
if
__name__
==
'__main__'
:
K
,
_
=
load_age_mixing
(
'data/polymod_normal_df.rds'
)
T
,
_
=
load_mobility_matrix
(
'data/movement.rds'
)
np
.
fill_diagonal
(
T
,
0.
)
N
,
_
=
load_population
(
'data/pop.rds'
)
model
=
CovidUKODE
(
K
,
T
,
N
)
print
(
"R0 (beta=0.036) = "
,
model
.
eval_R0
({
'beta1'
:
0.036
,
'beta2'
:
0.33
,
'gamma'
:
0.25
}))
R0
=
2.4
def
opt_fn
(
beta
,
R0
):
r0
,
i
=
model
.
eval_R0
({
'beta1'
:
beta
,
'beta2'
:
0.33
,
'gamma'
:
0.25
},
1e-9
)
return
(
r0
-
R0
)
**
2
res
=
opt
.
minimize_scalar
(
opt_fn
,
args
=
R0
)
print
(
res
[
'x'
])
scratch_covid.py
View file @
996a4afc
import
tensorflow
as
tf
import
numpy
as
np
from
covid.model
import
Homogeneous
from
covid.rdata
import
*
from
covid.model
import
CovidUKODE
sw
=
tf
.
summary
.
create_file_writer
(
'tf_log'
)
K
=
load_age_mixing
(
'data/polymod_normal_df.rds'
).
to_numpy
(
dtype
=
np
.
float32
)
T
=
load_mobility_matrix
(
'data/movement.rds'
).
to_numpy
(
dtype
=
np
.
float32
)
np
.
fill_diagonal
(
T
,
0.
)
N
=
load_population
(
'data/pop.rds'
)[
'n'
].
to_numpy
(
dtype
=
np
.
float32
)
model
=
Homogeneous
()
popsize
=
2500
state
=
np
.
array
([
np
.
full
([
popsize
],
999.
),
np
.
full
([
popsize
],
0.
),
np
.
full
([
popsize
],
1.
),
np
.
full
([
popsize
],
0.
)],
dtype
=
np
.
float32
)
model
=
CovidUKODE
(
K
,
T
,
N
,
T
.
shape
[
0
])
print
(
"Running..."
,
flush
=
True
,
sep
=
''
)
tf
.
summary
.
trace_on
(
graph
=
True
,
profiler
=
True
)
t
,
sim_state
=
model
.
sample
(
state
,
[
0.
,
10.
],
{
'beta'
:
0.2
,
'nu'
:
0.14
,
'gamma'
:
0.14
})
print
(
"Done"
,
flush
=
True
)
print
(
sim_state
)
with
sw
.
as_default
():
tf
.
summary
.
trace_export
(
'profile'
,
step
=
0
,
profiler_outdir
=
'tf_log'
)
sw
.
flush
()
sw
.
close
()
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