hotspot_detection.py 4.65 KB
Newer Older
Chris Jewell's avatar
Chris Jewell committed
1
"""Hotspot detection given a posterior"""
2

Chris Jewell's avatar
Chris Jewell committed
3
4
import os
import yaml
5
import pickle as pkl
Chris Jewell's avatar
Chris Jewell committed
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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import h5py
import numpy as np
import geopandas as gp

import tensorflow as tf
from gemlib.util import compute_state

from covid.cli_arg_parse import cli_args

import model_spec

DTYPE = model_spec.DTYPE


@tf.function
def predicted_incidence(param, init_state, init_step, num_steps, priors):
    """Runs the simulation forward in time from `init_state` at time `init_time`
       for `num_steps`.
    :param theta: a tensor of batched theta parameters [B] + theta.shape
    :param xi: a tensor of batched xi parameters [B] + xi.shape
    :param events: a [B, M, S] batched state tensor
    :param init_step: the initial time step
    :param num_steps: the number of steps to simulate
    :param priors: the priors for gamma
    :returns: a tensor of srt_quhape [B, M, num_steps, X] where X is the number of state
              transitions
    """

    def sim_fn(args):
        beta1_, beta2_, beta3_, sigma_, xi_, gamma0_, gamma1_, init_ = args

        par = dict(
            beta1=beta1_,
            beta2=beta2_,
            beta3=beta3_,
            gamma0=gamma0_,
            gamma1=gamma1_,
            xi=xi_,
        )

        model = model_spec.CovidUK(
            covar_data,
            initial_state=init_,
            initial_step=init_step,
            num_steps=num_steps,
            priors=priors,
        )
        sim = model.sample(**par)
        return sim["seir"]

    events = tf.map_fn(
        sim_fn,
        elems=(
            param["beta1"],
            param["beta2"],
            param["beta3"],
            param["sigma"],
            param["xi"],
            param["gamma0"],
            param["gamma1"],
            init_state,
        ),
        fn_output_signature=(tf.float64),
    )
    return events


def quantile_observed(observed, prediction):
    predicted_sum = tf.reduce_sum(prediction, axis=-1)
    observed_sum = tf.reduce_sum(observed, axis=-1)
    q = tf.reduce_mean(
        tf.cast(predicted_sum < observed_sum, tf.float32), axis=0
    )
    return q


if __name__ == "__main__":

    args = cli_args()

    # Get general config
    with open(args.config, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    inference_period = [
        np.datetime64(x) for x in config["Global"]["inference_period"]
    ]

    # Load covariate data
    covar_data = model_spec.read_covariates(config)
96
97
98
99

    output_folder_path = config["output"]["results_dir"]
    geopackage_path = os.path.expandvars(
        os.path.join(output_folder_path, config["output"]["geopackage"])
Chris Jewell's avatar
Chris Jewell committed
100
    )
101
102
103
104

    # Load geopackage
    geo = gp.read_file(geopackage_path)

Chris Jewell's avatar
Chris Jewell committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
    geo = geo[geo["lad19cd"].str.startswith("E")]  # England only, for now.
    geo = geo.sort_values("lad19cd")

    # Load posterior file
    posterior_path = os.path.join(
        config["output"]["results_dir"], config["output"]["posterior"]
    )
    print("Using posterior:", posterior_path)
    posterior = h5py.File(
        os.path.expandvars(
            posterior_path,
        ),
        "r",
        rdcc_nbytes=1024 ** 3,
        rdcc_nslots=1e6,
    )

    # Pre-determined thinning of posterior (better done in MCMC?)
    idx = range(6000, 10000, 10)
    param = dict(
        beta1=posterior["samples/beta1"][idx],
        beta2=posterior["samples/beta2"][idx],
        beta3=posterior["samples/beta3"][idx],
        sigma=posterior["samples/sigma"][
            idx,
        ],
        xi=posterior["samples/xi"][idx],
        gamma0=posterior["samples/gamma0"][idx],
        gamma1=posterior["samples/gamma1"][idx],
    )
    events = posterior["samples/events"][idx]
    init_state = posterior["initial_state"][:]
    state_timeseries = compute_state(
        init_state, events, model_spec.STOICHIOMETRY
    )

    # Prediction requires simulation from 2 weeks ago
    prediction = predicted_incidence(
        param,
        init_state=state_timeseries[..., -14, :],
        init_step=state_timeseries.shape[-2] - 14,
146
        num_steps=56,
Chris Jewell's avatar
Chris Jewell committed
147
148
149
150
151
152
153
154
155
        priors=config["mcmc"]["prior"],
    )

    # prediction quantiles
    q_obs7 = quantile_observed(events[0, :, -7:, 2], prediction[..., 7:14, 2])
    q_obs14 = quantile_observed(events[0, :, -14:, 2], prediction[..., :14, 2])

    geo["Pr(pred<obs)_7"] = q_obs7.numpy()
    geo["Pr(pred<obs)_14"] = q_obs14.numpy()
156

Chris Jewell's avatar
Chris Jewell committed
157
158
159
160
161
162
    geo.to_file(
        os.path.join(
            config["output"]["results_dir"], config["output"]["geopackage"]
        ),
        driver="GPKG",
    )
163
164
165

    with open(
        os.path.expandvars(
166
167
168
169
170
            os.path.join(
                output_folder_path, config["output"]["posterior_predictive"]
            )
        ),
        "wb",
171
172
    ) as f:
        pkl.dump(prediction, f)