Commit 744ee455 authored by Chris Jewell's avatar Chris Jewell
Browse files

Removed redundant code from `inference.py` (duplicated in `model_spec.py`)....

Removed redundant code from `inference.py` (duplicated in `model_spec.py`). Small change to `xi_freq` parameter in `model_spec.py`.
parent 440ce766
%% LyX 2.2.4 created this file. For more info, see http://www.lyx.org/.
%% Do not edit unless you really know what you are doing.
\documentclass[english]{article}
\usepackage[T1]{fontenc}
\usepackage[latin9]{inputenc}
\usepackage{geometry}
\geometry{verbose,tmargin=2.5cm,bmargin=2.5cm,lmargin=2.5cm,rmargin=2.5cm}
\setlength{\parskip}{\medskipamount}
\setlength{\parindent}{0pt}
\PassOptionsToPackage{lined, algonl, boxed}{algorithm2e}
\usepackage{url}
\usepackage{bm}
\usepackage{algorithm2e}
\usepackage{amsmath}
\usepackage[numbers]{natbib}
\usepackage{babel}
\begin{document}
\title{UK Age and Space structured Covid-19 model\\
Technical Description}
\author{CPJ}
\maketitle
\part{Model Concept}
We wish to develop a model that will enable us to assess spatial spread
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 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) aggegated
to Upper Tier Local Authority (UTLA) regions, and Census 2011-derived
age structured population density in UTLA 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, Wales, and Northern
Ireland.
\section{Age-mixing}
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$.
\section{Human mobility}
2011 Census data from ONS on daily mean numbers of commuters moving
from each Residential MSOA to Workplace MSOA. MSOAs are aggregated
to UTLA regions\footnote{City of Westminster and City of London are aggregated, as are Cornwall
and Scilly to allow mapping onto MSOAs.} for which we have age-structured population density. The resulting
matrix $C$ is of dimension $n_{c}\times n_{c}$ where $n_{c}=149$.
Since this matrix is for Residence to Workplace movement only, we
assume that the mean number of journeys between each UTLA is given
by
\[
T=C+C^{T}
\]
with 0 diagonal.
\section{Population size}
Age-structured population size within each UTLA is taken from publicly
available 2019 UTLA data giving a vector $N$ of length $n_{m}n_{c}=2533$,
i.e. population for each of $n_{m}$ age groups and $n_{c}$ UTLAs.
\section{Connectivity matrix}
We assemble a country-wide connectivity matrices as Kronecker products,
such that
\[
M^{\star}=I_{n_{c}}\bigotimes M
\]
and
\[
C^{\star}=C\bigotimes\bm{1}_{n_{m}\times n_{c}}
\]
giving two matrices of dimension $n_{m}n_{c}\times n_{m}n_{c}$. $M^{\star}$
is block diagonal with Polymod mixing matrices, performed for both
$M_{tt}$ and $M_{hh}$. $C^{\star}$ expands the mobility matrix
$C$ such that a block structure of connectivity between UTLAs results.
\part{ODE-based Model}
\section{Disease progression model}
We assume an SEIR model described as a system of ODEs. We denote the
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_{t}\left[M^{\star}\vec{I}(t)+\beta_{2}w_{t}\bar{M}C^{\star}\frac{{\vec{I}(t)}}{N}\right]\frac{\vec{S}(t)}{N}\\
\frac{\mathrm{d}\vec{E}(t)}{dt} & =\beta_{t}\left[M^{\star}\vec{I}(t)+\beta_{2}w_{t}\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*}
where $\bar{M}$ is the global mean person-person contact rate, and
$w_{t}$ is the total rail ticket sales in the UK expressed as a fraction
of the 2019 mean (a proxy for reduction in travel). Parameters are:
\begin{align*}
\beta_{t} & =\begin{cases}
\beta_{1} & \mbox{if }t<T_{L}\\
\beta_{1}\beta_{3} & \mbox{otherwise }
\end{cases}
\end{align*}
with $T_{L}$ the date of lock-down restrictions, $\beta_{1}$ a baseline
transmisison rate, and $\beta_{3}$ giving the ratio of post-lockdown
to pre-lockdown transmission; commuting infection ratio $\beta_{2}$;
latent period $\frac{1}{\nu}=4$days; and infectious period $\frac{1}{\gamma}$.
We assume that contact with commuters is $\beta_{2}=\frac{1}{3}$
of that between members of the same age-UTLA combination assuming
an 8 hour working day.
\section{Noise model}
Currently, and subject to discussion, we assume that all detected
cases are synonymous with individuals transitioning $I\rightarrow R$.
We assume the number of new cases in each age-LAD combination are
given by
\[
y_{ik}(t)\sim\mbox{Negative Binomial}\left(r,\phi(R_{ik}(t)-R_{ik}(t-1))\right)
\]
where $\phi$ is the case reporting fraction (i.e. proportion of infections
that are eventually detected) and $r$ is an overdispersion parameter.
\section{Inference}
We are interested in making inference on $\beta_{1},\beta_{3},\gamma,I_{0}$,
and $r$. Prior distributions are chosen to be
\begin{align*}
\beta_{1} & \sim\mbox{Gamma}(1,1)\\
\beta_{3} & \sim\mbox{Gamma}(20,20)\\
\gamma & \sim\mbox{Gamma}(100,400)\\
I_{0} & \sim\mbox{Gamma}(1.5,0.05)\\
r & \sim\mbox{\mbox{Gamma}(}0.1,0.1)
\end{align*}
specified to express \emph{a priori }relative ignorance about $\beta_{1}$
and $r$, stronger information on $\beta_{3}$ (post-lockdown transmission
is the same as pre-lockdown transmission but could be either greater
or smaller), and $I_{0}$ (initial number of individuals infected)
and strong information about $\gamma$ with the\emph{ }belief that
the infectious period is approximately 4 days.
Bayesian inference is performed by MCMC to estimate the joint posterior
distribution of the parameters conditional on the observed data up
to the analysis time.
\subsection{ODE Implementation}
The model is currently implemented in Python3, using Tensorflow 2.2.0
with the RK45 differential equation solver implemented in the \texttt{DormandPrince}
class provided by Tensorflow Probability 0.9. The MCMC is implemented
using Tensorflow Probability's Metropolis-Hastings framework with
additional adaptation steps according to \citet{HaarEtAl01}. The
code implementation may be found at \url{http://github.com/chrism0dwk/covid19uk/tree/space_age_fitting}.
\part{Stochastic implementation}
The stochastic version of the model takes the form of a discrete-time
Markov process. Since we may wish to develop beyond the simple linear
SEIR state transition model, we construct a general framework for
an arbitrary STM.
At timepoint $t$, we define a Markov transition rate matrix over
the state vector $X(t)=(S(t),E(t),I(t),R(t))^{T}$ as
\[
Q(t)=\left(\begin{array}{cccc}
-\lambda(t) & \lambda(t) & 0 & 0\\
0 & -\nu & \nu & 0\\
0 & 0 & -\gamma & \gamma\\
0 & 0 & 0 & 0
\end{array}\right)
\]
where $\lambda(t)$ is the transition rate for $S\rightarrow E$,
$\nu$ is for $E\rightarrow I$ and $\gamma$ is for $I\rightarrow R$.
The Markov transition matrix may then be calculated as
\[
P(t)=e^{Q(t)\delta t}
\]
where $\delta t$ is the size of the timestep used in the model.
Since $P(t)$ is expensive to compute, we take the approximation
\[
p_{ij}(t)=\frac{q_{ij}(t)}{\sum_{j}q_{ij}(t)}\left[1-e^{-\sum_{j}q_{ij}(t)\delta t}\right],\;i\ne j
\]
and
\[
p_{ii}(t)=1-\sum_{j\ne i}p_{ij}(t)
\]
where $p_{ij}(t)$ and $q_{ij}(t)$ are the $i,j$th elements of $P(t)$
and $Q(t)$.
\section{Simulation}
Simulating from the model involves propagating the system state $X_{0},X_{1},\dots,X_{T}$
from initial conditions $X_{0}$ and maximum time $T$ in units of
$\delta t$. We denote by $Z_{t}$ the \emph{event matrix} at time
$t$ which gives the number of individuals transitioning from each
epidemiological state in $X_{t}$ to every epidemiological state at
the next timestep $X_{t+1}$ (including individuals who remain in
a state). We denote the $i$th epidemiological state at time $t$
by $\bm{x}_{i}(t)$. We further denote the $i$th row of a matrix
$M$ as $\bm{m}_{i\cdot}$ and likewise the $j$th column as $\bm{m}_{\cdot j}$.
The method for evolving the state is shown in Algorithm \ref{alg:Discrete-time-simulation}.
\begin{algorithm}[h]
\SetAlgoLined
Initialise $t=0,\; X(t)$\;
\While{t<T}{
Calculate $P(t)$ given $X(t)$\;
\For{$i$ in 1 to $nrow(P(t))$}{
Draw $\bm{z}_{i\cdot}(t) \sim \mbox{Multinomial}(\bm{x}_{i}(t),\bm{p}_{i\cdot}(t))$
}
Let $X(t+1)=\sum_i \bm{z}_{i\cdot}(t)$ \tcp*[h]{column sums of $Z(t)$}\;
Let $t=t+1$\;
}
\caption{\label{alg:Discrete-time-simulation}Discrete time simulation algorithm}
\end{algorithm}
\section{Inference}
\subsection{Likelihood function}
A realisation of the epidemic process from Algorithm \ref{alg:Discrete-time-simulation}
is characterised by a sequence of (realised) event matrices $Z(0),\dots,Z(T-1)$,
conditional on an initial state $X(0)$ and model parameters $\bm{\theta}$.
Using the Markovian property of the model the full data joint likelihood
is therefore
\begin{equation}
L(Z(0),\dots,Z(T-1);X(0),\bm{\theta})\propto\prod_{t=0}^{T-1}\prod_{i}\prod_{j}p_{ij}(t)^{z_{ij}(t)}\label{eq:likelihood}
\end{equation}
being the product of Multinomial PMFs for each state update at each
timepoint in the epidemic.
\textbf{Remarks}
\begin{enumerate}
\item The disarmingly simple form of the likelihood in Equation \ref{eq:likelihood}
belies the complexity of having to calculate $P(t)$ for each term
over $t$. To do this requires calculation of $X(t)$ since $Q(t)$
typically depends on $X(t)$. In practice, this means that it is convenient
to implement the leftmost product (over $t$) as a loop within code,
updating $X(t)$ according to $Z(t)$ at each iteration. The products
$i$ and $j$ are vectorised in code so as to make use of vectorised
CPU instructions or GPU acceleration.
\item Typically, only certain elements of $Z(t)$ are observed as data within
an epidemic. For the SEIR example above, typically only $I\rightarrow R$
transitions are observed, this being assumed synonymous with detection
(and subsequent isolation/recovery with immunity/death).
\end{enumerate}
\subsection{MCMC fitting}
The stochastic model is fitted using Bayesian statistics, with missing
data treated as part of the joint posterior distribution. The MCMC
scheme proceeds according to
\begin{enumerate}
\item Update $\bm{\theta}|Z(0),\dots,Z(T-1)$;
\item Update $Z(0),\dots,Z(T-1)|\bm{\theta}$.
\end{enumerate}
\subsubsection{Updating event times}
For any unobserved transition $ij$, do the following:
\begin{enumerate}
\item Choose a timepoint to update $t^{\star}\sim\mbox{UniformInteger}(0,T-1)$
\item Choose a metapopulation to update $m^{\star}\sim\mbox{UniformInteger(0,M-1) }$
\item Choose a number of events to move $x^{\star}\sim\mbox{Binomial}(Z_{ij}(t^{\star}),p)$
\item Choose where to move the events to as follows
\[
s^{\star}=u^{\star}d^{\star}
\]
with
\begin{align*}
u^{\star} & \sim\mbox{DiscreteUniform}((-1,1))\\
d^{\star} & \sim\mbox{Poisson}(\alpha)
\end{align*}
\end{enumerate}
Tuning parameters are $p$, the proportion of events on any particular
day/metapopulation to move, and $\alpha$ which determines the distance
in time to move the events.
This proposal is potentially asymmetric due to the Binomial sampling
of $x^{\star}|z_{ij}(t^{\star}),p$ so requires a ratio of PMFs of
Binomial distributions in the MH accept/reject ratio.
\textbf{Extensions}
\begin{enumerate}
\item An obvious extension is to try to move multiple metapopulations at
once \textendash{} easy to implement, adds another tuning parameter;
\item Non-trivially, we could propose a joint proposal, which makes use
of the model structure.
\end{enumerate}
\bibliographystyle{plain}
\bibliography{bibliography}
\end{document}
......@@ -3,17 +3,14 @@ import optparse
import os
from time import perf_counter
import tqdm
import yaml
import h5py
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow_probability as tfp
import tqdm
import yaml
from covid import config
from covid.model import load_data, DiscreteTimeStateTransitionModel
from covid.util import impute_previous_cases
from covid.impl.util import compute_state
from covid.impl.mcmc import UncalibratedLogRandomWalk, random_walk_mvnorm_fn
from covid.impl.event_time_mh import UncalibratedEventTimesUpdate
......@@ -34,28 +31,6 @@ else:
tfd = tfp.distributions
tfb = tfp.bijectors
DTYPE = config.floatX
STOICHIOMETRY = tf.constant([[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]], dtype=DTYPE)
def impute_censored_events(cases):
"""Imputes censored S->E and E->I events using geometric
sampling algorithm in `impute_previous_cases`
There are application-specific magic numbers hard-coded below,
which reflect experimentation to get the right lag between EI and
IR events, and SE and EI events respectively. These were chosen
by experimentation and examination of the resulting epidemic
trajectories.
:param cases: a MxT matrix of case numbers (I->R)
:returns: a MxTx3 tensor of events where the first two indices of
the right-most dimension contain the imputed event times.
"""
ei_events, lag_ei = impute_previous_cases(cases, 0.44)
se_events, lag_se = impute_previous_cases(ei_events, 2.0)
ir_events = np.pad(cases, ((0, 0), (lag_ei + lag_se - 2, 0)))
ei_events = np.pad(ei_events, ((0, 0), (lag_se - 1, 0)))
return tf.stack([se_events, ei_events, ir_events], axis=-1)
if __name__ == "__main__":
......@@ -82,7 +57,7 @@ if __name__ == "__main__":
cases = model_spec.read_cases(config["data"]["reported_cases"])
# Impute censored events, return cases
events = impute_censored_events(cases)
events = model_spec.impute_censored_events(cases)
# Initial conditions are calculated by calculating the state
# at the beginning of the inference period
......@@ -95,13 +70,11 @@ if __name__ == "__main__":
[covar_data["N"], tf.zeros_like(events[:, 0, :])], axis=-1
),
events=events,
stoichiometry=STOICHIOMETRY,
stoichiometry=model_spec.STOICHIOMETRY,
)
start_time = state.shape[1] - cases.shape[1]
initial_state = state[:, start_time, :]
events = events[:, start_time:, :]
xi_freq = 14
num_xi = events.shape[1] // xi_freq
num_metapop = covar_data["N"].shape[0]
########################################################
......@@ -109,7 +82,6 @@ if __name__ == "__main__":
########################################################
model = model_spec.CovidUK(
covariates=covar_data,
xi_freq=14,
initial_state=initial_state,
initial_step=0,
num_steps=events.shape[1],
......@@ -271,7 +243,7 @@ if __name__ == "__main__":
current_state = [
np.array([0.85, 0.3, 0.25], dtype=DTYPE),
np.zeros(num_xi, dtype=DTYPE),
np.zeros(model.model['xi']().event_shape[-1], dtype=DTYPE),
events,
]
......
......@@ -14,7 +14,7 @@ DTYPE = floatX
STOICHIOMETRY = tf.constant([[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]])
TIME_DELTA = 1.0
XI_FREQ = 14
def read_covariates(paths):
"""Loads covariate data
......@@ -63,7 +63,8 @@ def impute_censored_events(cases):
return tf.stack([se_events, ei_events, ir_events], axis=-1)
def CovidUK(covariates, xi_freq, initial_state, initial_step, num_steps):
def CovidUK(covariates, initial_state, initial_step, num_steps):
def beta1():
return tfd.Gamma(
concentration=tf.constant(1.0, dtype=DTYPE),
......@@ -78,9 +79,9 @@ def CovidUK(covariates, xi_freq, initial_state, initial_step, num_steps):
def xi():
sigma = tf.constant(0.01, dtype=DTYPE)
phi = tf.constant(12.0, dtype=DTYPE)
phi = tf.constant(24.0, dtype=DTYPE)
kernel = tfp.math.psd_kernels.MaternThreeHalves(sigma, phi)
idx_pts = tf.cast(tf.range(num_steps // xi_freq) * xi_freq, dtype=DTYPE)
idx_pts = tf.cast(tf.range(num_steps // XI_FREQ) * XI_FREQ, dtype=DTYPE)
return tfd.GaussianProcess(kernel, index_points=idx_pts[:, tf.newaxis])
def nu():
......@@ -114,7 +115,7 @@ def CovidUK(covariates, xi_freq, initial_state, initial_step, num_steps):
w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, W.shape[0] - 1)
commute_volume = tf.gather(W, w_idx)
xi_idx = tf.cast(
tf.clip_by_value(t // 14, 0, xi.shape[0] - 1), dtype=tf.int64,
tf.clip_by_value(t // XI_FREQ, 0, xi.shape[0] - 1), dtype=tf.int64,
)
xi_ = tf.gather(xi, xi_idx)
beta = beta1 * tf.math.exp(xi_)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment