simulate.py 1.76 KB
Newer Older
Chris Jewell's avatar
Chris Jewell committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
"""Test simulation for COVID-19 UK model"""

import optparse
import yaml
import numpy as np
import tensorflow as tf

from covid import config
from covid.model import load_data
from covid.pydata import phe_case_data
from covid.util import sanitise_settings, impute_previous_cases
from covid.impl.util import compute_state

import model_spec

DTYPE = config.floatX

# Read in data
parser = optparse.OptionParser()
parser.add_option(
    "--config",
    "-c",
    dest="config",
    default="example_config.yaml",
    help="configuration file",
)
options, cmd_args = parser.parse_args(["-c", "example_config.yaml"])
print("Loading config file:", options.config)

with open(options.config, "r") as f:
    config = yaml.load(f, Loader=yaml.FullLoader)


# Load in covariate data
Chris Jewell's avatar
Chris Jewell committed
35
covar_data = model_spec.read_covariates(config["data"])
Chris Jewell's avatar
Chris Jewell committed
36
37
38

# We load in cases and impute missing infections first, since this sets the
# time epoch which we are analysing.
Chris Jewell's avatar
Chris Jewell committed
39
cases = model_spec.read_cases(config["data"]["reported_cases"])
Chris Jewell's avatar
Chris Jewell committed
40
41

# Single imputation of censored data
Chris Jewell's avatar
Chris Jewell committed
42
events = model_spec.impute_censored_events(cases)
Chris Jewell's avatar
Chris Jewell committed
43

Chris Jewell's avatar
Chris Jewell committed
44
45
# Initial conditions S(0), E(0), I(0), R(0) are calculated
# by calculating the state at the beginning of the inference period
Chris Jewell's avatar
Chris Jewell committed
46
state = compute_state(
Chris Jewell's avatar
Chris Jewell committed
47
    initial_state=tf.concat([covar_data["N"], tf.zeros_like(events[:, 0, :])], axis=-1),
Chris Jewell's avatar
Chris Jewell committed
48
49
50
51
52
53
54
55
56
    events=events,
    stoichiometry=model_spec.STOICHIOMETRY,
)
start_time = state.shape[1] - cases.shape[1]
initial_state = state[:, start_time, :]
events = events[:, start_time:, :]

# Build model and sample
full_probability_model = model_spec.CovidUK(
Chris Jewell's avatar
Chris Jewell committed
57
    covariates=covar_data, initial_state=initial_state, initial_step=0, num_steps=80,
Chris Jewell's avatar
Chris Jewell committed
58
59
)
seir = full_probability_model.model["seir"](
Chris Jewell's avatar
Chris Jewell committed
60
    beta1=0.35, beta2=0.65, xi=[0.0] * 5, gamma=0.49
Chris Jewell's avatar
Chris Jewell committed
61
62
)
sim = seir.sample()