within_between.py 3.1 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

    C = tf.convert_to_tensor(covariates["C"], dtype=model_spec.DTYPE)
19
20
    C = tf.linalg.set_diag(C, tf.zeros(C.shape[0], dtype=model_spec.DTYPE))

Chris Jewell's avatar
Chris Jewell committed
21
22
23
24
25
26
    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
    )
27
28

    def within_fn(t, state):
29
30
31
32
33
        w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, W.shape[0] - 1)
        commute_volume = tf.gather(W, w_idx)
        rate = state[..., 2] - beta2 * state[
            ..., 2
        ] / N * commute_volume * tf.reduce_sum(C, axis=-2)
34
35
36
37
38
        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)
39
40
41
42
43
        rate = (
            beta2
            * commute_volume
            * tf.linalg.matvec(C + tf.transpose(C), state[..., 2] / N)
        )
44
45
46
47
48
        return rate

    return within_fn, between_fn


Chris Jewell's avatar
Chris Jewell committed
49
# @tf.function
50
def calc_pressure_components(covariates, beta2, state):
Chris Jewell's avatar
Chris Jewell committed
51
    def atomic_fn(args):
52
53
        beta2_, state_ = args
        within_fn, between_fn = make_within_rate_fns(covariates, beta2_)
Chris Jewell's avatar
Chris Jewell committed
54
55
        within = within_fn(covariates["W"].shape[0], state_)
        between = between_fn(covariates["W"].shape[0], state_)
56
57
58
        total = within + between
        return within / total, between / total

59
    return tf.vectorized_map(atomic_fn, elems=(beta2, state))
60
61


62
args = cli_args()
63
64
65
66
67

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

    inference_period = [
68
        np.datetime64(x) for x in config["Global"]["inference_period"]
69
70
71
    ]

# Load covariate data
72
covar_data = model_spec.read_covariates(config)
73
74
75
76

# Load posterior file
posterior = h5py.File(
    os.path.expandvars(
77
78
79
        os.path.join(
            config["output"]["results_dir"], config["output"]["posterior"]
        )
80
81
82
83
84
85
86
87
    ),
    "r",
    rdcc_nbytes=1024 ** 3,
    rdcc_nslots=1e6,
)

# Pre-determined thinning of posterior (better done in MCMC?)
idx = range(6000, 10000, 10)
88
beta2 = posterior["samples/beta2"][idx]
89
90
91
92
93
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(
94
    covar_data, beta2, state_timeseries[..., -1, :]
95
96
)

97
98
99
100
101
gpkg = gp.read_file(
    os.path.join(
        config["output"]["results_dir"], config["output"]["geopackage"]
    )
)
102
103
104
gpkg = gpkg[gpkg["lad19cd"].str.startswith("E")]
gpkg = gpkg.sort_values("lad19cd")

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

107
108
109
110
111
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(
112
113
114
    os.path.join(
        config["output"]["results_dir"], config["output"]["geopackage"]
    ),
115
116
    driver="GPKG",
)