inference.py 12.7 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),
            ),
Chris Jewell's avatar
Chris Jewell committed
283
            "events": (samples[2], (NUM_BURST_SAMPLES, 32, 32, 1)),
284
        },
Chris Jewell's avatar
Chris Jewell committed
285
286
        results_dict=results,
        num_samples=NUM_SAVED_SAMPLES,
287
    )
Chris Jewell's avatar
Chris Jewell committed
288
    posterior._file.create_dataset("initial_state", data=initial_state)
289
290
291
292
    posterior._file.create_dataset(
        "date_range",
        data=np.array(data["date_range"]).astype(h5py.string_dtype()),
    )
293
294
295
    # 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"):
296
    final_results = None
297
    for i in tqdm.tqdm(
298
        range(NUM_BURSTS), unit_scale=NUM_BURST_SAMPLES * config["thin"]
299
    ):
300
        samples, results, final_results = sample(
301
302
            NUM_BURST_SAMPLES,
            init_state=current_state,
303
            thin=config["thin"] - 1,
304
            previous_results=final_results,
305
306
307
308
309
        )
        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
310
        posterior.write_samples(
311
312
313
314
            {
                "beta2": samples[0][:, 0],
                "gamma0": samples[0][:, 1],
                "gamma1": samples[0][:, 2],
Chris Jewell's avatar
Chris Jewell committed
315
                "sigma": samples[0][:, 3],
Chris Jewell's avatar
Chris Jewell committed
316
317
318
                "beta3": tf.zeros(
                    [samples[0].shape[0], 5], dtype=DTYPE
                ),  # samples[0][:, 4:],
319
320
321
322
                "beta1": samples[1][:, 0],
                "xi": samples[1][:, 1:],
                "events": samples[2],
            },
Chris Jewell's avatar
Chris Jewell committed
323
324
325
            first_dim_offset=i * NUM_BURST_SAMPLES,
        )
        posterior.write_results(results, first_dim_offset=i * NUM_BURST_SAMPLES)
326
327
328
        end = perf_counter()

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

Chris Jewell's avatar
Chris Jewell committed
335
    print(
336
        f"Acceptance theta: {posterior['results/block0/is_accepted'][:].mean()}"
Chris Jewell's avatar
Chris Jewell committed
337
    )
338
    print(f"Acceptance xi: {posterior['results/block1/is_accepted'][:].mean()}")
Chris Jewell's avatar
Chris Jewell committed
339
340
341
342
343
344
345
346
347
348
349
350
    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
351

Chris Jewell's avatar
Chris Jewell committed
352
    del posterior
353
354
355
356


if __name__ == "__main__":

357
358
359
360
361
362
363
364
365
366
367
368
369
    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()
370
371
372
373

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

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