within_between.py 2.94 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
"""Creates a medium term prediction"""

import os
import yaml
import numpy as np
import h5py
import geopandas as gp

import tensorflow as tf

from covid.cli_arg_parse import cli_args
Chris Jewell's avatar
Chris Jewell committed
12
from gemlib.util import compute_state
13
14
15
import model_spec


16
def make_within_rate_fns(covariates, beta2):
17
18
19
20
21

    C = tf.convert_to_tensor(covariates["C"], dtype=model_spec.DTYPE)
    C = tf.linalg.set_diag(
        C + tf.transpose(C), tf.zeros(C.shape[0], dtype=model_spec.DTYPE)
    )
Chris Jewell's avatar
Chris Jewell committed
22
23
24
25
26
27
    W = tf.convert_to_tensor(
        tf.squeeze(covariates["W"]), dtype=model_spec.DTYPE
    )
    N = tf.convert_to_tensor(
        tf.squeeze(covariates["N"]), dtype=model_spec.DTYPE
    )
28
29

    def within_fn(t, state):
Chris Jewell's avatar
Chris Jewell committed
30
        rate = state[..., 2]
31
32
33
34
35
        return rate

    def between_fn(t, state):
        w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, W.shape[0] - 1)
        commute_volume = tf.gather(W, w_idx)
36
        rate = beta2 * commute_volume * tf.linalg.matvec(C, state[..., 2] / N)
37
38
39
40
41
        return rate

    return within_fn, between_fn


Chris Jewell's avatar
Chris Jewell committed
42
# @tf.function
43
def calc_pressure_components(covariates, beta2, state):
Chris Jewell's avatar
Chris Jewell committed
44
    def atomic_fn(args):
45
46
        beta2_, state_ = args
        within_fn, between_fn = make_within_rate_fns(covariates, beta2_)
Chris Jewell's avatar
Chris Jewell committed
47
48
        within = within_fn(covariates["W"].shape[0], state_)
        between = between_fn(covariates["W"].shape[0], state_)
49
50
51
        total = within + between
        return within / total, between / total

52
    return tf.vectorized_map(atomic_fn, elems=(beta2, state))
53
54


55
args = cli_args()
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

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

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

# Load covariate data
covar_data = model_spec.read_covariates(
    config["data"], date_low=inference_period[0], date_high=inference_period[1]
)

# Load posterior file
posterior = h5py.File(
    os.path.expandvars(
72
73
74
        os.path.join(
            config["output"]["results_dir"], config["output"]["posterior"]
        )
75
76
77
78
79
80
81
82
    ),
    "r",
    rdcc_nbytes=1024 ** 3,
    rdcc_nslots=1e6,
)

# Pre-determined thinning of posterior (better done in MCMC?)
idx = range(6000, 10000, 10)
83
beta2 = posterior["samples/beta2"][idx]
84
85
86
87
88
events = posterior["samples/events"][idx]
init_state = posterior["initial_state"][:]
state_timeseries = compute_state(init_state, events, model_spec.STOICHIOMETRY)

within, between = calc_pressure_components(
89
    covar_data, beta2, state_timeseries[..., -1, :]
90
91
)

92
93
94
95
96
gpkg = gp.read_file(
    os.path.join(
        config["output"]["results_dir"], config["output"]["geopackage"]
    )
)
97
98
99
gpkg = gpkg[gpkg["lad19cd"].str.startswith("E")]
gpkg = gpkg.sort_values("lad19cd")

Chris Jewell's avatar
Chris Jewell committed
100
101
print("Within shape:", within.shape)

102
103
104
105
106
gpkg["within_mean"] = np.mean(within, axis=0)
gpkg["between_mean"] = np.mean(between, axis=0)
gpkg["p_within_gt_between"] = np.mean(within > between)

gpkg.to_file(
107
108
109
    os.path.join(
        config["output"]["results_dir"], config["output"]["geopackage"]
    ),
110
111
    driver="GPKG",
)