inference.py 12.6 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
        np.array([0.6, 0.0, 0.0, 0.1], dtype=DTYPE), #, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=DTYPE),
255
        np.zeros(
256
            model.model["xi"](0.0, 0.1).event_shape[-1] + 1,
257
258
            dtype=DTYPE,
        ),
259
260
        events,
    ]
261
    print("Initial logpi:", joint_log_prob(*current_state))
262

Chris Jewell's avatar
Chris Jewell committed
263
264
265
    # Output file
    samples, results, _ = sample(1, current_state)
    posterior = Posterior(
266
        output_file,
267
268
269
270
        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
271
            "sigma": (samples[0][:, 3], (NUM_BURST_SAMPLES,)),
Chris Jewell's avatar
Chris Jewell committed
272
            "beta3": (tf.zeros([1,5], dtype=DTYPE), (NUM_BURST_SAMPLES, 2)), #(samples[0][:, 4:], (NUM_BURST_SAMPLES, 2)),
273
274
275
276
277
278
279
            "beta1": (samples[1][:, 0], (NUM_BURST_SAMPLES,)),
            "xi": (
                samples[1][:, 1:],
                (NUM_BURST_SAMPLES, samples[1].shape[1] - 1),
            ),
            "events": (samples[2], (NUM_BURST_SAMPLES, 64, 64, 1)),
        },
Chris Jewell's avatar
Chris Jewell committed
280
281
        results_dict=results,
        num_samples=NUM_SAVED_SAMPLES,
282
    )
Chris Jewell's avatar
Chris Jewell committed
283
    posterior._file.create_dataset("initial_state", data=initial_state)
284
285
286
287
    posterior._file.create_dataset(
        "date_range",
        data=np.array(data["date_range"]).astype(h5py.string_dtype()),
    )
288
289
290
    # 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"):
291
    final_results = None
292
    for i in tqdm.tqdm(
293
        range(NUM_BURSTS), unit_scale=NUM_BURST_SAMPLES * config["thin"]
294
    ):
295
        samples, results, final_results = sample(
296
297
            NUM_BURST_SAMPLES,
            init_state=current_state,
298
            thin=config["thin"] - 1,
299
            previous_results=final_results,
300
301
302
303
304
        )
        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
305
        posterior.write_samples(
306
307
308
309
            {
                "beta2": samples[0][:, 0],
                "gamma0": samples[0][:, 1],
                "gamma1": samples[0][:, 2],
Chris Jewell's avatar
Chris Jewell committed
310
                "sigma": samples[0][:, 3],
Chris Jewell's avatar
Chris Jewell committed
311
                "beta3": tf.zeros([samples[0].shape[0], 5], dtype=DTYPE), #samples[0][:, 4:],
312
313
314
315
                "beta1": samples[1][:, 0],
                "xi": samples[1][:, 1:],
                "events": samples[2],
            },
Chris Jewell's avatar
Chris Jewell committed
316
317
318
            first_dim_offset=i * NUM_BURST_SAMPLES,
        )
        posterior.write_results(results, first_dim_offset=i * NUM_BURST_SAMPLES)
319
320
321
        end = perf_counter()

        print("Storage time:", end - start, "seconds")
322
        for k, v in results.items():
323
324
325
326
            print(
                f"Acceptance {k}:",
                tf.reduce_mean(tf.cast(v["is_accepted"], tf.float32)),
            )
Chris Jewell's avatar
Chris Jewell committed
327

Chris Jewell's avatar
Chris Jewell committed
328
    print(
329
        f"Acceptance theta: {posterior['results/block0/is_accepted'][:].mean()}"
Chris Jewell's avatar
Chris Jewell committed
330
    )
331
    print(f"Acceptance xi: {posterior['results/block1/is_accepted'][:].mean()}")
Chris Jewell's avatar
Chris Jewell committed
332
333
334
335
336
337
338
339
340
341
342
343
    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
344

Chris Jewell's avatar
Chris Jewell committed
345
    del posterior
346
347
348
349


if __name__ == "__main__":

350
351
352
353
354
355
356
357
358
359
360
361
362
    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()
363
364
365
366

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

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