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
3e5268ca
Commit
3e5268ca
authored
Mar 02, 2020
by
Chris Jewell
Browse files
Added ODE model.
parent
7c8712ef
Changes
3
Hide whitespace changes
Inline
Side-by-side
covid/model.py
View file @
3e5268ca
"""Functions for infection rates"""
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
...
...
@@ -56,3 +59,49 @@ class Homogeneous:
self
.
param
=
param
return
chain_binomial_simulate
(
self
.
h
,
initial_state
,
time_lims
[
0
],
time_lims
[
1
],
1.
,
self
.
stoichiometry
)
class
CovidUKODE
:
def
__init__
(
self
,
K
,
T
,
N
,
n_lads
=
152
):
"""Represents a CovidUK ODE model
:param K: a MxM matrix of age group mixing
:param T: a n_ladsxn_lads matrix of inter-LAD connectivity
: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
])
T
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
T
)
shp
=
tf
.
linalg
.
LinearOperatorFullMatrix
(
np
.
ones
([
n_lads
,
n_lads
]))
self
.
T
=
tf
.
linalg
.
LinearOperatorKronecker
([
T
,
shp
])
def
make_h
(
self
,
param
):
@
tf
.
function
def
h_fn
(
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
dS
=
-
infec_rate
dE
=
infec_rate
-
param
[
'nu'
]
*
E
dI
=
param
[
'nu'
]
*
E
-
param
[
'gamma'
]
*
I
dR
=
param
[
'gamma'
]
*
I
df
=
tf
.
stack
([
dS
,
dE
,
dI
,
dR
])
return
df
return
h_fn
@
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
covid_ode.py
View file @
3e5268ca
...
...
@@ -12,15 +12,16 @@ state = np.array([np.full([popsize], 999.),
K
=
np
.
random
.
uniform
(
size
=
[
popsize
,
popsize
])
param
=
{
'beta'
:
0.2
,
'nu'
:
0.14
,
'gamma'
:
0.14
}
param
=
{
'beta'
:
0.0002
,
'nu'
:
0.14
,
'gamma'
:
0.14
}
@
tf
.
function
def
h
(
t
,
state
):
print
(
state
)
state
=
tf
.
unstack
(
state
,
axis
=
0
)
S
,
E
,
I
,
R
=
state
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
...
...
@@ -31,9 +32,11 @@ def h(t, state):
@
tf
.
function
def
solve_ode
(
rates
,
t_init
,
state_init
,
t
):
return
tode
.
DormandPrince
().
solve
(
rates
,
t_init
,
state_init
,
solution_times
=
t
)
return
tode
.
DormandPrince
(
first_step_size
=
1.
,
max_num_steps
=
5000
).
solve
(
rates
,
t_init
,
state_init
,
solution_times
=
t
)
solution_times
=
np
.
arange
(
0.
,
365.
,
1.
)
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
scratch_covid.py
View file @
3e5268ca
...
...
@@ -11,12 +11,12 @@ state = np.array([np.full([popsize], 999.),
np
.
full
([
popsize
],
0.
),
np
.
full
([
popsize
],
1.
),
np
.
full
([
popsize
],
0.
)],
dtype
=
np
.
float32
)
print
(
state
)
print
(
"Running..."
,
flush
=
True
,
sep
=
''
)
tf
.
summary
.
trace_on
(
graph
=
True
,
profiler
=
True
)
t
,
sim_state
=
model
.
sample
(
state
,
[
0.
,
5
.
],
[
0.
,
10
.
],
{
'beta'
:
0.2
,
'nu'
:
0.14
,
'gamma'
:
0.14
})
print
(
"Done"
,
flush
=
True
)
print
(
sim_state
)
...
...
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