inference.py 12.8 KB
Newer Older
Chris Jewell's avatar
Chris Jewell committed
1
"""MCMC Test Rig for COVID-19 UK model"""
2
3
# pylint: disable=E402

4
5
import h5py
import pickle as pkl
6
from time import perf_counter
7
8
import tqdm
import yaml
9
import numpy as np
10
11
import tensorflow as tf
import tensorflow_probability as tfp
Chris Jewell's avatar
Chris Jewell committed
12

13
14
15
16
17
18
from gemlib.util import compute_state
from gemlib.mcmc import UncalibratedEventTimesUpdate
from gemlib.mcmc import UncalibratedOccultUpdate, TransitionTopology
from gemlib.mcmc import GibbsKernel
from gemlib.mcmc import MultiScanKernel
from gemlib.mcmc import AdaptiveRandomWalkMetropolis
Chris Jewell's avatar
Chris Jewell committed
19
from gemlib.mcmc import Posterior
20

21
import covid.model_spec as model_spec
22

23
24
tfd = tfp.distributions
tfb = tfp.bijectors
25
DTYPE = model_spec.DTYPE
26
27


28
def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
29
    """Constructs and runs the MCMC"""
30

31
32
33
34
    if tf.test.gpu_device_name():
        print("Using GPU")
    else:
        print("Using CPU")
35

36
37
    with open(data_file, "rb") as f:
        data = pkl.load(f)
38
39
40
41

    # We load in cases and impute missing infections first, since this sets the
    # time epoch which we are analysing.
    # Impute censored events, return cases
Chris Jewell's avatar
Chris Jewell committed
42
    print("Data shape:", data["cases"].shape)
43
    events = model_spec.impute_censored_events(data["cases"].astype(DTYPE))
44
45
46
47
48
49
50
51
52

    # Initial conditions are calculated by calculating the state
    # at the beginning of the inference period
    #
    # Imputed censored events that pre-date the first I-R events
    # in the cases dataset are discarded.  They are only used to
    # to set up a sensible initial state.
    state = compute_state(
        initial_state=tf.concat(
53
            [data["N"][:, tf.newaxis], tf.zeros_like(events[:, 0, :])],
54
            axis=-1,
55
        ),
56
        events=events,
57
        stoichiometry=model_spec.STOICHIOMETRY,
58
    )
59
    start_time = state.shape[1] - data["cases"].shape[1]
60
61
62
63
    initial_state = state[:, start_time, :]
    events = events[:, start_time:, :]

    ########################################################
64
    # Construct the MCMC kernels #
65
66
    ########################################################
    model = model_spec.CovidUK(
67
        covariates=data,
68
69
70
        initial_state=initial_state,
        initial_step=0,
        num_steps=events.shape[1],
71
    )
72

73
    def joint_log_prob(block0, block1, events):
74
        return model.log_prob(
Chris Jewell's avatar
Chris Jewell committed
75
            dict(
76
                beta2=block0[0],
77
78
                gamma0=block0[1],
                gamma1=block0[2],
Chris Jewell's avatar
Chris Jewell committed
79
                sigma=block0[3],
Chris Jewell's avatar
Chris Jewell committed
80
                beta1=block1[0],
Chris Jewell's avatar
Chris Jewell committed
81
                xi=block1[1:],
Chris Jewell's avatar
Chris Jewell committed
82
83
                seir=events,
            )
84
        )
85

86
    # Build Metropolis within Gibbs sampler
87
    def make_blk0_kernel(shape, name):
Chris Jewell's avatar
Chris Jewell committed
88
        def fn(target_log_prob_fn, _):
89
            return tfp.mcmc.TransformedTransitionKernel(
90
                inner_kernel=AdaptiveRandomWalkMetropolis(
91
                    target_log_prob_fn=target_log_prob_fn,
92
93
                    initial_covariance=np.eye(shape[0], dtype=model_spec.DTYPE)
                    * 1e-1,
94
95
                    covariance_burnin=200,
                ),
96
                bijector=tfp.bijectors.Blockwise(
Chris Jewell's avatar
Chris Jewell committed
97
98
99
100
                    bijectors=[
                        tfp.bijectors.Exp(),
                        tfp.bijectors.Identity(),
                        tfp.bijectors.Exp(),
Chris Jewell's avatar
Chris Jewell committed
101
                        # tfp.bijectors.Identity(),
Chris Jewell's avatar
Chris Jewell committed
102
                    ],
Chris Jewell's avatar
Chris Jewell committed
103
                    block_sizes=[1, 2, 1],  # , 5],
104
                ),
105
106
                name=name,
            )
107

108
109
        return fn

110
    def make_blk1_kernel(shape, name):
Chris Jewell's avatar
Chris Jewell committed
111
        def fn(target_log_prob_fn, _):
112
            return AdaptiveRandomWalkMetropolis(
113
                target_log_prob_fn=target_log_prob_fn,
114
115
                initial_covariance=np.eye(shape[0], dtype=model_spec.DTYPE)
                * 1e-1,
116
117
118
119
120
                covariance_burnin=200,
                name=name,
            )

        return fn
Chris Jewell's avatar
Chris Jewell committed
121

122
123
124
    def make_partially_observed_step(
        target_event_id, prev_event_id=None, next_event_id=None, name=None
    ):
Chris Jewell's avatar
Chris Jewell committed
125
        def fn(target_log_prob_fn, _):
126
            return tfp.mcmc.MetropolisHastings(
127
                inner_kernel=UncalibratedEventTimesUpdate(
128
                    target_log_prob_fn=target_log_prob_fn,
129
130
131
132
                    target_event_id=target_event_id,
                    prev_event_id=prev_event_id,
                    next_event_id=next_event_id,
                    initial_state=initial_state,
133
134
135
                    dmax=config["dmax"],
                    mmax=config["m"],
                    nmax=config["nmax"],
136
137
138
139
140
                ),
                name=name,
            )

        return fn
141

142
    def make_occults_step(prev_event_id, target_event_id, next_event_id, name):
Chris Jewell's avatar
Chris Jewell committed
143
        def fn(target_log_prob_fn, _):
144
            return tfp.mcmc.MetropolisHastings(
145
                inner_kernel=UncalibratedOccultUpdate(
146
                    target_log_prob_fn=target_log_prob_fn,
147
148
                    topology=TransitionTopology(
                        prev_event_id, target_event_id, next_event_id
149
                    ),
150
                    cumulative_event_offset=initial_state,
151
                    nmax=config["occult_nmax"],
152
153
                    t_range=(events.shape[1] - 21, events.shape[1]),
                    name=name,
154
155
156
157
158
159
                ),
                name=name,
            )

        return fn

Chris Jewell's avatar
Chris Jewell committed
160
    def make_event_multiscan_kernel(target_log_prob_fn, _):
161
        return MultiScanKernel(
162
            config["num_event_time_updates"],
163
164
165
166
167
168
169
170
171
            GibbsKernel(
                target_log_prob_fn=target_log_prob_fn,
                kernel_list=[
                    (0, make_partially_observed_step(0, None, 1, "se_events")),
                    (0, make_partially_observed_step(1, 0, 2, "ei_events")),
                    (0, make_occults_step(None, 0, 1, "se_occults")),
                    (0, make_occults_step(0, 1, 2, "ei_occults")),
                ],
                name="gibbs1",
172
            ),
173
174
        )

175
176
    # MCMC tracing functions
    def trace_results_fn(_, results):
Chris Jewell's avatar
Chris Jewell committed
177
178
179
180
        """Packs results into a dictionary"""
        results_dict = {}
        res0 = results.inner_results

181
        results_dict["block0"] = {
Chris Jewell's avatar
Chris Jewell committed
182
183
184
185
186
            "is_accepted": res0[0].inner_results.is_accepted,
            "target_log_prob": res0[
                0
            ].inner_results.accepted_results.target_log_prob,
        }
187
        results_dict["block1"] = {
Chris Jewell's avatar
Chris Jewell committed
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
            "is_accepted": res0[1].is_accepted,
            "target_log_prob": res0[1].accepted_results.target_log_prob,
        }

        def get_move_results(results):
            return {
                "is_accepted": results.is_accepted,
                "target_log_prob": results.accepted_results.target_log_prob,
                "proposed_delta": tf.stack(
                    [
                        results.accepted_results.m,
                        results.accepted_results.t,
                        results.accepted_results.delta_t,
                        results.accepted_results.x_star,
                    ]
                ),
            }
205

Chris Jewell's avatar
Chris Jewell committed
206
207
208
209
210
        res1 = res0[2].inner_results
        results_dict["move/S->E"] = get_move_results(res1[0])
        results_dict["move/E->I"] = get_move_results(res1[1])
        results_dict["occult/S->E"] = get_move_results(res1[2])
        results_dict["occult/E->I"] = get_move_results(res1[3])
211

Chris Jewell's avatar
Chris Jewell committed
212
        return results_dict
213
214

    # Build MCMC algorithm here.  This will be run in bursts for memory economy
215
    @tf.function(autograph=use_autograph, experimental_compile=use_xla)
Chris Jewell's avatar
Chris Jewell committed
216
    def sample(n_samples, init_state, thin=0, previous_results=None):
217
218
219
220
        with tf.name_scope("main_mcmc_sample_loop"):

            init_state = init_state.copy()

221
            gibbs_schema = GibbsKernel(
222
                target_log_prob_fn=joint_log_prob,
223
                kernel_list=[
224
225
                    (0, make_blk0_kernel(init_state[0].shape, "block0")),
                    (1, make_blk1_kernel(init_state[1].shape, "block1")),
226
                    (2, make_event_multiscan_kernel),
227
                ],
228
                name="gibbs0",
229
            )
Chris Jewell's avatar
Chris Jewell committed
230

231
232
233
234
            samples, results, final_results = tfp.mcmc.sample_chain(
                n_samples,
                init_state,
                kernel=gibbs_schema,
Chris Jewell's avatar
Chris Jewell committed
235
                num_steps_between_results=thin,
236
237
238
                previous_kernel_results=previous_results,
                return_final_kernel_results=True,
                trace_fn=trace_results_fn,
239
240
            )

241
            return samples, results, final_results
242

243
244
245
    ###############################
    # Construct bursted MCMC loop #
    ###############################
246
247
    NUM_BURSTS = int(config["num_bursts"])
    NUM_BURST_SAMPLES = int(config["num_burst_samples"])
248
    NUM_SAVED_SAMPLES = NUM_BURST_SAMPLES * NUM_BURSTS
249
250
251
252
253

    # RNG stuff
    tf.random.set_seed(2)

    current_state = [
Chris Jewell's avatar
Chris Jewell committed
254
255
256
        np.array(
            [0.6, 0.0, 0.0, 0.1], dtype=DTYPE
        ),  # , 0.0, 0.0, 0.0, 0.0, 0.0], dtype=DTYPE),
257
        np.zeros(
258
            model.model["xi"](0.0, 0.1).event_shape[-1] + 1,
259
260
            dtype=DTYPE,
        ),
261
262
        events,
    ]
263
    print("Initial logpi:", joint_log_prob(*current_state))
264

Chris Jewell's avatar
Chris Jewell committed
265
266
267
    # Output file
    samples, results, _ = sample(1, current_state)
    posterior = Posterior(
268
        output_file,
269
270
271
272
        sample_dict={
            "beta2": (samples[0][:, 0], (NUM_BURST_SAMPLES,)),
            "gamma0": (samples[0][:, 1], (NUM_BURST_SAMPLES,)),
            "gamma1": (samples[0][:, 2], (NUM_BURST_SAMPLES,)),
Chris Jewell's avatar
Chris Jewell committed
273
            "sigma": (samples[0][:, 3], (NUM_BURST_SAMPLES,)),
Chris Jewell's avatar
Chris Jewell committed
274
275
276
277
            "beta3": (
                tf.zeros([1, 5], dtype=DTYPE),
                (NUM_BURST_SAMPLES, 2),
            ),  # (samples[0][:, 4:], (NUM_BURST_SAMPLES, 2)),
278
279
280
281
282
            "beta1": (samples[1][:, 0], (NUM_BURST_SAMPLES,)),
            "xi": (
                samples[1][:, 1:],
                (NUM_BURST_SAMPLES, samples[1].shape[1] - 1),
            ),
283
284
            "events": (
                samples[2],
285
                (NUM_BURST_SAMPLES, min(32, events.shape[0]), 32, 1),
286
            ),
287
        },
Chris Jewell's avatar
Chris Jewell committed
288
289
        results_dict=results,
        num_samples=NUM_SAVED_SAMPLES,
290
    )
Chris Jewell's avatar
Chris Jewell committed
291
    posterior._file.create_dataset("initial_state", data=initial_state)
292
293
294
295
    posterior._file.create_dataset(
        "date_range",
        data=np.array(data["date_range"]).astype(h5py.string_dtype()),
    )
296
297
298
    # We loop over successive calls to sample because we have to dump results
    #   to disc, or else end OOM (even on a 32GB system).
    # with tf.profiler.experimental.Profile("/tmp/tf_logdir"):
299
    final_results = None
300
    for i in tqdm.tqdm(
301
        range(NUM_BURSTS), unit_scale=NUM_BURST_SAMPLES * config["thin"]
302
    ):
303
        samples, results, final_results = sample(
304
305
            NUM_BURST_SAMPLES,
            init_state=current_state,
306
            thin=config["thin"] - 1,
307
            previous_results=final_results,
308
309
310
311
312
        )
        current_state = [s[-1] for s in samples]
        print(current_state[0].numpy(), flush=True)

        start = perf_counter()
Chris Jewell's avatar
Chris Jewell committed
313
        posterior.write_samples(
314
315
316
317
            {
                "beta2": samples[0][:, 0],
                "gamma0": samples[0][:, 1],
                "gamma1": samples[0][:, 2],
Chris Jewell's avatar
Chris Jewell committed
318
                "sigma": samples[0][:, 3],
Chris Jewell's avatar
Chris Jewell committed
319
320
321
                "beta3": tf.zeros(
                    [samples[0].shape[0], 5], dtype=DTYPE
                ),  # samples[0][:, 4:],
322
323
324
325
                "beta1": samples[1][:, 0],
                "xi": samples[1][:, 1:],
                "events": samples[2],
            },
Chris Jewell's avatar
Chris Jewell committed
326
327
328
            first_dim_offset=i * NUM_BURST_SAMPLES,
        )
        posterior.write_results(results, first_dim_offset=i * NUM_BURST_SAMPLES)
329
330
331
        end = perf_counter()

        print("Storage time:", end - start, "seconds")
332
        for k, v in results.items():
333
334
335
336
            print(
                f"Acceptance {k}:",
                tf.reduce_mean(tf.cast(v["is_accepted"], tf.float32)),
            )
Chris Jewell's avatar
Chris Jewell committed
337

Chris Jewell's avatar
Chris Jewell committed
338
    print(
339
        f"Acceptance theta: {posterior['results/block0/is_accepted'][:].mean()}"
Chris Jewell's avatar
Chris Jewell committed
340
    )
341
    print(f"Acceptance xi: {posterior['results/block1/is_accepted'][:].mean()}")
Chris Jewell's avatar
Chris Jewell committed
342
343
344
345
346
347
348
349
350
351
352
353
    print(
        f"Acceptance move S->E: {posterior['results/move/S->E/is_accepted'][:].mean()}"
    )
    print(
        f"Acceptance move E->I: {posterior['results/move/E->I/is_accepted'][:].mean()}"
    )
    print(
        f"Acceptance occult S->E: {posterior['results/occult/S->E/is_accepted'][:].mean()}"
    )
    print(
        f"Acceptance occult E->I: {posterior['results/occult/E->I/is_accepted'][:].mean()}"
    )
Chris Jewell's avatar
Chris Jewell committed
354

Chris Jewell's avatar
Chris Jewell committed
355
    del posterior
356
357
358
359


if __name__ == "__main__":

360
361
362
363
364
365
366
367
368
369
370
371
372
    from argparse import ArgumentParser

    parser = ArgumentParser(description="Run MCMC inference algorithm")
    parser.add_argument(
        "-c", "--config", type=str, help="Config file", required=True
    )
    parser.add_argument(
        "-o", "--output", type=str, help="Output file", required=True
    )
    parser.add_argument(
        "data_file", type=str, help="Data pickle file", required=True
    )
    args = parser.parse_args()
373
374
375
376

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

377
    mcmc(args.data_file, args.output, config["Mcmc"])