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
750c80c7
Commit
750c80c7
authored
Mar 03, 2020
by
Chris Jewell
Browse files
Implemented country-wide ODE model.
parent
3e5268ca
Changes
4
Hide whitespace changes
Inline
Side-by-side
covid/model.py
View file @
750c80c7
...
...
@@ -2,11 +2,13 @@
import
tensorflow
as
tf
import
tensorflow_probability
as
tfp
tode
=
tfp
.
math
.
ode
import
numpy
as
np
from
covid.impl.chainbinom_simulate
import
chain_binomial_simulate
class
CovidUK
:
def
__init__
(
self
,
K
,
T
,
W
):
self
.
K
=
K
...
...
@@ -38,17 +40,16 @@ class CovidUK:
class
Homogeneous
:
def
__init__
(
self
):
self
.
stoichiometry
=
tf
.
constant
([[
-
1
,
1
,
0
,
0
],
[
0
,
-
1
,
1
,
0
],
[
0
,
0
,
-
1
,
1
]],
dtype
=
tf
.
float32
)
[
0
,
-
1
,
1
,
0
],
[
0
,
0
,
-
1
,
1
]],
dtype
=
tf
.
float32
)
def
h
(
self
,
state
):
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
hazard_rates
=
tf
.
stack
([
self
.
param
[
'beta'
]
*
I
/
tf
.
reduce_sum
(
state
),
self
.
param
[
'beta'
]
*
I
/
tf
.
reduce_sum
(
state
),
self
.
param
[
'nu'
]
*
tf
.
ones_like
(
I
),
self
.
param
[
'gamma'
]
*
tf
.
ones_like
(
I
)
])
...
...
@@ -63,7 +64,7 @@ class Homogeneous:
class
CovidUKODE
:
def
__init__
(
self
,
K
,
T
,
N
,
n_lads
=
152
):
def
__init__
(
self
,
K
,
T
,
N
,
n_lads
):
"""Represents a CovidUK ODE model
:param K: a MxM matrix of age group mixing
...
...
@@ -71,23 +72,26 @@ class CovidUKODE:
:param N: a vector of population sizes in each LAD
:param n_lads: the number of LADS
"""
K
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
K
)
eye
=
tf
.
linalg
.
LinearOperatorDiag
(
n_lads
)
self
.
K
=
tf
.
linalg
.
LinearOperatorKronecker
([
K
,
eye
])
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
])
T
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
T
)
shp
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
np
.
ones
([
n_lads
,
n_lads
]))
self
.
T
=
tf
.
linalg
.
LinearOperatorKronecker
([
T
,
shp
])
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
def
make_h
(
self
,
param
):
@
tf
.
function
def
h_fn
(
state
):
def
h_fn
(
t
,
state
):
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
infec_rate
=
tf
.
matvec
(
self
.
K
,
I
)
infec_rate
+=
tf
.
reduce_mean
(
self
.
K
)
*
tf
.
matvec
(
self
.
T
,
I
/
self
.
N
)
infec_rate
=
infec_rate
*
S
/
self
.
N
infec_rate
=
tf
.
linalg
.
matvec
(
self
.
K
,
I
)
infec_rate
+=
self
.
K
bar
*
tf
.
linalg
.
matvec
(
self
.
T
,
I
/
self
.
N
)
infec_rate
=
param
[
'beta'
]
*
S
/
self
.
N
*
infec_rate
dS
=
-
infec_rate
dE
=
infec_rate
-
param
[
'nu'
]
*
E
...
...
@@ -96,12 +100,22 @@ class CovidUKODE:
df
=
tf
.
stack
([
dS
,
dE
,
dI
,
dR
])
return
df
return
h_fn
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
])
@
tf
.
function
def
simulate
(
self
,
param
,
state_init
,
t_start
,
t_end
,
t_step
=
1.
):
h
=
self
.
make_h
(
param
)
t
=
np
.
linspace
(
t_start
,
t_end
,
t_step
)
sim
=
tode
.
DormandPrince
(
first_step_size
=
1.
,
max_num_steps
=
5000
).
solve
(
h
,
t_start
,
state_init
,
solution_times
=
t
)
return
sim
\ No newline at end of file
t
=
np
.
arange
(
start
=
t_start
,
stop
=
t_end
,
step
=
t_step
)
print
(
f
"Running simulation with times
{
t_start
}
-
{
t_end
}
:
{
t_step
}
"
)
solver
=
tode
.
DormandPrince
()
results
=
solver
.
solve
(
ode_fn
=
h
,
initial_time
=
t_start
,
initial_state
=
state_init
,
solution_times
=
t
)
return
results
.
times
,
results
.
states
covid/rdata.py
0 → 100644
View file @
750c80c7
"""Loads R data.frame data structures"""
import
numpy
import
pyreadr
as
pyr
def
load_age_mixing
(
rds_file
:
str
):
"""Loads age mixing matrix from R.
: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
]
def
load_mobility_matrix
(
rds_file
:
str
):
"""Loads mobility COO from RDS file.
:param rds_file: a .rds file containing an R data.frame with columns 'Residence',
and 'Workplace' indicating matrix coordinates, and first column containing the value
"""
raw
=
pyr
.
read_r
(
rds_file
)
df
=
list
(
raw
.
values
())[
0
]
colnames
=
df
.
columns
mobility_matrix
=
df
.
pivot
(
index
=
'Workplace'
,
columns
=
'Residence'
,
values
=
colnames
[
0
])
mobility_matrix
[
mobility_matrix
.
isna
()]
=
0.
return
mobility_matrix
def
load_population
(
rds_file
:
str
):
"""Loads population data from RDS file.
:param rds_file: and RDS file containing a data.frame with columns 'age', 'LA.code', 'n'
"""
raw
=
pyr
.
read_r
(
rds_file
)
df
=
list
(
raw
.
values
())[
0
]
df
=
df
.
sort_values
(
by
=
[
'LA.code'
,
'age'
])
return
df
covid_ode.py
View file @
750c80c7
import
optparse
import
numpy
as
np
import
tensorflow
as
tf
import
t
ensorflow_probability
as
tfp
tode
=
tfp
.
math
.
ode
import
yaml
import
t
ime
import
matplotlib.pyplot
as
plt
from
covid.model
import
CovidUKODE
from
covid.rdata
import
*
popsize
=
2500
state
=
np
.
array
([
np
.
full
([
popsize
],
999.
),
np
.
full
([
popsize
],
0.
),
np
.
full
([
popsize
],
1.
),
np
.
full
([
popsize
],
0.
)],
dtype
=
np
.
float64
)
K
=
np
.
random
.
uniform
(
size
=
[
popsize
,
popsize
])
def
sanitise_parameter
(
par_dict
):
"""Sanitises a dictionary of parameters"""
par
=
[
'beta'
,
'nu'
,
'gamma'
]
d
=
{
key
:
np
.
float64
(
par_dict
[
key
])
for
key
in
par
}
return
d
param
=
{
'beta'
:
0.0002
,
'nu'
:
0.14
,
'gamma'
:
0.14
}
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'
])}
return
d
def
h
(
t
,
state
):
print
(
state
)
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
if
__name__
==
'__main__'
:
infec_rate
=
param
[
'beta'
]
*
S
*
tf
.
linalg
.
matvec
(
K
,
I
)
dS
=
-
infec_rate
dE
=
infec_rate
-
param
[
'nu'
]
*
E
dI
=
param
[
'nu'
]
*
E
-
param
[
'gamma'
]
*
I
dR
=
param
[
'gamma'
]
*
I
parser
=
optparse
.
OptionParser
()
parser
.
add_option
(
"--config"
,
"-c"
,
dest
=
"config"
,
help
=
"configuration file"
)
options
,
args
=
parser
.
parse_args
()
with
open
(
options
.
config
,
'r'
)
as
ymlfile
:
config
=
yaml
.
load
(
ymlfile
)
df
=
tf
.
stack
([
dS
,
dE
,
dI
,
dR
])
return
df
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
)
np
.
fill_diagonal
(
T
,
0.
)
N
=
load_population
(
config
[
'data'
][
'population_size'
])[
'n'
].
to_numpy
(
dtype
=
np
.
float32
)
@
tf
.
function
def
solve_ode
(
rates
,
t_init
,
state_init
,
t
):
return
tode
.
DormandPrince
(
first_step_size
=
1.
,
max_num_steps
=
5000
).
solve
(
rates
,
t_init
,
state_init
,
solution_times
=
t
)
param
=
sanitise_parameter
(
config
[
'parameter'
])
settings
=
sanitise_settings
(
config
[
'settings'
])
solution_times
=
np
.
arange
(
0.
,
365.
,
1.
)
model
=
CovidUKODE
(
K
,
T
,
N
,
T
.
shape
[
0
]
)
print
(
"Running..."
,
flush
=
True
)
result
=
solve_ode
(
rates
=
h
,
t_init
=
0.
,
state_init
=
state
,
t
=
solution_times
)
print
(
"Done"
)
print
(
result
)
\ No newline at end of file
state_init
=
model
.
create_initial_state
()
start
=
time
.
perf_counter
()
t
,
sim
=
model
.
simulate
(
param
,
state_init
,
settings
[
'start'
],
settings
[
'end'
],
settings
[
'time_step'
])
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
ode_config.yaml
0 → 100644
View file @
750c80c7
# Covid_ODE model configuration
data
:
age_mixing_matrix
:
data/polymod_normal_df.rds
mobility_matrix
:
data/movement.rds
population_size
:
data/pop.rds
parameter
:
beta
:
0.09
nu
:
0.25
gamma
:
0.65
settings
:
start
:
0.
end
:
365.
time_step
:
1.
output
:
simulation
:
covid_ode.sim
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