def rate(self): """Calculates the rate parameter""" state = compute_state(self._initial_state, self._events, self._stoichiometry, closed=True) tms_state = tf.transpose(state, perm=(1, 0, 2)) def fn(elems): return tf.stack(self._transition_rate_fn(*elems), axis=-1) rates = tf.vectorized_map( fn=fn, elems=( self._initial_step + tf.range(tms_state.shape[0]), tms_state, ), ) def integrated_rate_fn(): """Use mid-point integration to estimate the constant rate over time. """ integrated_rates = tms_state[..., :-1] * rates return (integrated_rates[:-1, ...] + integrated_rates[1:, ...]) / 2.0 return (tf.reduce_sum(integrated_rate_fn(), axis=(-3, -2)) + self._baseline_hazard_rate_priors["rate"])
def within_between(input_files, output_file): """Calculates PAF for within- and between-location infection. :param input_files: a list of [data pickle, posterior samples pickle] :param output_file: a csv with within/between summary """ with open(input_files[0], "rb") as f: covar_data = pkl.load(f) with open(input_files[1], "rb") as f: samples = pkl.load(f) beta2 = samples["beta2"] events = samples["seir"] init_state = samples["init_state"] state_timeseries = compute_state(init_state, events, model_spec.STOICHIOMETRY) within, between = calc_pressure_components(covar_data, beta2, state_timeseries[..., -1, :]) df = pd.DataFrame( dict( within_mean=np.mean(within, axis=0), between_mean=np.mean(between, axis=0), p_within_gt_between=np.mean(within > between), ), index=pd.Index(covar_data["locations"]["lad19cd"], name="location"), ) df.to_csv(output_file)
def r_fn(args): beta1_, beta2_, beta3_, sigma_, xi_, gamma0_, events_ = args t = events_.shape[-2] - 1 state = compute_state(init_state, events_, model_spec.STOICHIOMETRY) state = tf.gather(state, t, axis=-2) # State on final inference day model = model_spec.CovidUK( covariates=covar_data, initial_state=init_state, initial_step=0, num_steps=events_.shape[-2], priors=priors, ) xi_pred = model_spec.conditional_gp( model.model["xi"](beta1_, sigma_), xi_, tf.constant([events.shape[-2] + model_spec.XI_FREQ], dtype=model_spec.DTYPE)[:, tf.newaxis], ) par = dict( beta1=beta1_, beta2=beta2_, beta3=beta3_, sigma=sigma_, gamma0=gamma0_, xi=xi_, ) print("xi shape:", par["xi"].shape) ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par) ngm = ngm_fn(t, state) return ngm
def predicted_incidence( posterior_samples, init_state, covar_data, init_step, num_steps, out_of_sample=False, ): """Runs the simulation forward in time from `init_state` at time `init_time` for `num_steps`. :param param: a dictionary of model parameters :covar_data: a dictionary of model covariate data :param init_step: the initial time step :param num_steps: the number of steps to simulate :returns: a tensor of srt_quhape [B, M, num_steps, X] where X is the number of state transitions """ posterior_state = compute_state( init_state, posterior_samples["seir"], model_spec.STOICHIOMETRY, ) posterior_samples["new_init_state"] = posterior_state[..., init_step, :] del posterior_samples["seir"] # For out-of-sample prediction, we have to re-simulate the # alpha_t trajectory given the starting point. if out_of_sample is True: alpha_t = posterior_samples["alpha_0"][:, tf.newaxis] + tf.cumsum( posterior_samples["alpha_t"], axis=-1) if init_step > 0: posterior_samples["alpha_0"] = alpha_t[:, init_step - 1] # Remove alpha_t from the posterior to make TFP re-simulate it. del posterior_samples["alpha_t"] @tf.function def do_sim(): def sim_fn(args): par = tf.nest.pack_sequence_as(posterior_samples, args) init_ = par["new_init_state"] del par["new_init_state"] model = model_spec.CovidUK( covar_data, initial_state=init_, initial_step=init_step, num_steps=num_steps, ) sim = model.sample(**par) return sim["seir"] return tf.map_fn( sim_fn, elems=tf.nest.flatten(posterior_samples), fn_output_signature=(tf.float64), ) return posterior_samples["new_init_state"], do_sim()
def body(state_, occults_): state_t1 = tf.roll(state_, shift=-1, axis=-2) neg_state_idx = tf.where(state_t1 < 0) first_neg_state_idx = tf.gather( neg_state_idx, tf.concat( [ [[0]], tf.where(neg_state_idx[:-1, 0] - neg_state_idx[1:, 0]) + 1, ], axis=0, ), ) mask = tf.scatter_nd( first_neg_state_idx, tf.ones([first_neg_state_idx.shape[0], 1], dtype=state_t1.dtype), state_t1.shape, ) delta_occults = tf.einsum("mts,xs->mtx", state_t1 * mask, stoichiometry) new_occults = tf.clip_by_value( occults_ - delta_occults, clip_value_min=0.0, clip_value_max=1.0e6 ) new_state = compute_state( init_state, events + new_occults, stoichiometry ) return new_state, new_occults
def prevalence(input_files, output_file): """Reconstruct predicted prevalence from original data and projection. :param input_files: a list of [data pickle, samples pickle, prediction pickle] :param output_file: a csv containing prevalence summary """ offset = 4 # Account for recording lag timepoints = np.array([0, 7, 14, 28, 56], np.int32) + offset with open(input_files[0], "rb") as f: data = pkl.load(f) with open(input_files[1], "rb") as f: samples = pkl.load(f) with open(input_files[2], "rb") as f: prediction = pkl.load(f) insample_state = compute_state(samples["init_state"], samples["seir"], STOICHIOMETRY) predicted_state = compute_state(insample_state[..., -1, :], prediction, STOICHIOMETRY) def calc_prev(state, name=None): prev = np.sum(state[..., 1:3], axis=-1) / np.squeeze(data["N"]) return mean_and_ci(prev, name=name) idx = prediction.coords["location"] prev = pd.DataFrame( calc_prev(predicted_state[..., timepoints[0], :], name="prev"), index=idx, ) for t in timepoints[1:]: tmp = pd.DataFrame( calc_prev(predicted_state[..., t, :], name=f"prev{t-offset}"), index=idx, ) prev = pd.concat([prev, tmp], axis="columns") prev.to_csv(output_file)
def predicted_incidence(posterior_samples, covar_data, init_step, num_steps): """Runs the simulation forward in time from `init_state` at time `init_time` for `num_steps`. :param param: a dictionary of model parameters :covar_data: a dictionary of model covariate data :param init_step: the initial time step :param num_steps: the number of steps to simulate :returns: a tensor of srt_quhape [B, M, num_steps, X] where X is the number of state transitions """ @tf.function def sim_fn(args): beta1_, beta2_, sigma_, xi_, gamma0_, gamma1_, init_ = args par = dict( beta1=beta1_, beta2=beta2_, sigma=sigma_, xi=xi_, gamma0=gamma0_, gamma1=gamma1_, ) model = model_spec.CovidUK( covar_data, initial_state=init_, initial_step=init_step, num_steps=num_steps, ) sim = model.sample(**par) return sim["seir"] posterior_state = compute_state( posterior_samples["init_state"], posterior_samples["seir"], model_spec.STOICHIOMETRY, ) init_state = posterior_state[..., init_step, :] events = tf.map_fn( sim_fn, elems=( posterior_samples["beta1"], posterior_samples["beta2"], posterior_samples["sigma"], posterior_samples["xi"], posterior_samples["gamma0"], posterior_samples["gamma1"], init_state, ), fn_output_signature=(tf.float64), ) return init_state, events
def regularize_occults(events, occults, init_state, stoichiometry): """Regularizes an occult matrix such that counting processes are valid :param events: a [M, T, X] events tensor :param occults: a [M, T, X] occults tensor :param init_state: a [M, S] initial state tensor :param stoichiometry: a [X, S] stoichiometry matrix :returns: an tuple containing updated (state, occults) tensors """ from gemlib.util import compute_state def body(state_, occults_): state_t1 = tf.roll(state_, shift=-1, axis=-2) neg_state_idx = tf.where(state_t1 < 0) first_neg_state_idx = tf.gather( neg_state_idx, tf.concat( [ [[0]], tf.where(neg_state_idx[:-1, 0] - neg_state_idx[1:, 0]) + 1, ], axis=0, ), ) mask = tf.scatter_nd( first_neg_state_idx, tf.ones([first_neg_state_idx.shape[0], 1], dtype=state_t1.dtype), state_t1.shape, ) delta_occults = tf.einsum("mts,xs->mtx", state_t1 * mask, stoichiometry) new_occults = tf.clip_by_value( occults_ - delta_occults, clip_value_min=0.0, clip_value_max=1.0e6 ) new_state = compute_state( init_state, events + new_occults, stoichiometry ) return new_state, new_occults def cond(state_, _): return tf.reduce_any(state_ < 0) state = compute_state(init_state, events + occults, stoichiometry) new_state, new_occults = tf.while_loop(cond, body, (state, occults)) return new_state, new_occults
def _log_prob(self, y, **kwargs): """Calculates the log probability of observing epidemic events y :param y: a list of tensors. The first is of shape [n_times] containing times, the second is of shape [n_times, n_states, n_states] containing event matrices. :param param: a list of parameters :returns: a scalar giving the log probability of the epidemic """ dtype = dtype_util.common_dtype([y, self.initial_state], dtype_hint=self.dtype) events = tf.convert_to_tensor(y, dtype) with tf.name_scope("DiscreteApproxContStateTransitionModel.log_prob"): state_timeseries = compute_state( self.initial_state, events, self.stoichiometry, closed=True, ) tms_timeseries = tf.transpose(state_timeseries, perm=(1, 0, 2)) tmr_events = tf.transpose(events, perm=(1, 0, 2)) def fn(elems): return tf.stack(self.transition_rates(*elems), axis=-1) rates = tf.vectorized_map( fn=fn, elems=( self.initial_step + tf.range(tms_timeseries.shape[0]), tms_timeseries, ), ) def integrated_rates(): """Use mid-point integration to estimate the constant rate over time """ integrated_rates = tms_timeseries[..., :-1] * rates return (integrated_rates[:-1, ...] + integrated_rates[1:, ...]) / 2.0 log_hazard_rate = tf.reduce_sum(tmr_events * tf.math.log(integrated_rates())) log_survival = tf.reduce_sum(integrated_rates()) * self.time_delta log_denom = tf.reduce_sum(tf.math.lgamma(tmr_events + 1.0)) return log_hazard_rate - log_survival - log_denom
def prevalence(prediction, popsize): prev = compute_state(prediction["initial_state"], prediction["events"], STOICHIOMETRY) prev = xarray.DataArray( prev.numpy(), coords=[ np.arange(prev.shape[0]), prediction.coords["location"], prediction.coords["time"], np.arange(prev.shape[-1]), ], dims=["iteration", "location", "time", "state"], ) prev_per_1e5 = (prev[..., 1:3].sum(dim="state").reset_coords(drop=True) / np.array(popsize)[np.newaxis, :, np.newaxis] * 100000) return xarray2summarydf(prev_per_1e5)
def r_fn(args): par = tf.nest.pack_sequence_as(samples, args) state = compute_state(initial_state, par["seir"], model_spec.STOICHIOMETRY) del par["seir"] def fn(t): state_ = tf.gather(state, t, axis=-2) # State on final inference day ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par) ngm = ngm_fn(t, state_) return ngm ngm = tf.vectorized_map(fn, elems=times) return tf.reduce_sum(ngm, axis=-2) # sum over destinations
def discrete_markov_log_prob(events, init_state, init_step, time_delta, hazard_fn, stoichiometry): """Calculates an unnormalised log_prob function for a discrete time epidemic model. :param events: a `[M, T, X]` batch of transition events for metapopulation M, times `T`, and transitions `X`. :param init_state: a vector of shape `[M, S]` the initial state of the epidemic for `M` metapopulations and `S` states :param init_step: the initial time step, as an offset to `range(events.shape[-2])` :param time_delta: the size of the time step. :param hazard_fn: a function that takes a state and returns a matrix of transition rates. :param stoichiometry: a `[X, S]` matrix describing the state update for each transition. :return: a scalar log probability for the epidemic. """ num_meta = events.shape[-3] num_times = events.shape[-2] num_states = stoichiometry.shape[-1] state_timeseries = compute_state(init_state, events, stoichiometry) # MxTxS tms_timeseries = tf.transpose(state_timeseries, perm=(1, 0, 2)) def fn(elems): return hazard_fn(*elems) tx_coords = transition_coords(stoichiometry) rates = tf.vectorized_map(fn=fn, elems=[tf.range(num_times), tms_timeseries]) rate_matrix = make_transition_matrix(rates, tx_coords, tms_timeseries.shape) probs = approx_expm(rate_matrix * time_delta) # [T, M, S, S] to [M, T, S, S] probs = tf.transpose(probs, perm=(1, 0, 2, 3)) event_matrix = make_transition_matrix(events, tx_coords, [num_meta, num_times, num_states]) event_matrix = tf.linalg.set_diag( event_matrix, state_timeseries - tf.reduce_sum(event_matrix, axis=-1)) logp = tfd.Multinomial( tf.cast(state_timeseries, dtype=tf.float32), probs=tf.cast(probs, dtype=tf.float32), name="log_prob", ).log_prob(tf.cast(event_matrix, dtype=tf.float32)) return tf.cast(tf.reduce_sum(logp), dtype=events.dtype)
def r_fn(args): beta1_, beta2_, beta3_, sigma_, xi_, gamma0_, events_ = args t = events_.shape[-2] - 1 state = compute_state(samples["init_state"], events_, model_spec.STOICHIOMETRY) state = tf.gather(state, t, axis=-2) # State on final inference day par = dict( beta1=beta1_, beta2=beta2_, beta3=beta3_, sigma=sigma_, gamma0=gamma0_, xi=xi_, ) ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par) ngm = ngm_fn(t, state) return ngm
def predicted_incidence(posterior_samples, init_state, covar_data, init_step, num_steps): """Runs the simulation forward in time from `init_state` at time `init_time` for `num_steps`. :param param: a dictionary of model parameters :covar_data: a dictionary of model covariate data :param init_step: the initial time step :param num_steps: the number of steps to simulate :returns: a tensor of srt_quhape [B, M, num_steps, X] where X is the number of state transitions """ posterior_state = compute_state( init_state, posterior_samples["seir"], model_spec.STOICHIOMETRY, ) posterior_samples["new_init_state"] = posterior_state[..., init_step, :] del posterior_samples["seir"] @tf.function def do_sim(): def sim_fn(args): par = tf.nest.pack_sequence_as(posterior_samples, args) init_ = par["new_init_state"] del par["new_init_state"] model = model_spec.CovidUK( covar_data, initial_state=init_, initial_step=init_step, num_steps=num_steps, ) sim = model.sample(**par) return sim["seir"] return tf.map_fn( sim_fn, elems=tf.nest.flatten(posterior_samples), fn_output_signature=(tf.float64), ) return posterior_samples["new_init_state"], do_sim()
def r_fn(args): beta1_, beta2_, beta_3, sigma_, xi_, gamma0_, events_ = args t = events_.shape[-2] - 1 state = compute_state(init_state, events_, model_spec.STOICHIOMETRY) state = tf.gather(state, t, axis=-2) # State on final inference day model = model_spec.CovidUK( covariates=covar_data, initial_state=init_state, initial_step=0, num_steps=events_.shape[-2], priors=priors, ) xi_pred = model_spec.conditional_gp( model.model["xi"](beta1_, sigma_), xi_, tf.constant([events.shape[-2] + model_spec.XI_FREQ], dtype=model_spec.DTYPE)[:, tf.newaxis], ) # FNC NOTE: # adding another 0.0 to beta3 as TF complains of dimension mismatch otherwise par = dict( beta1=beta1_, beta2=beta2_, beta3=tf.concat([beta_3, [0.0, 0.0]], axis=-1), sigma=sigma_, gamma0=gamma0_, xi=xi_, # tf.reshape(xi_pred.sample(), [1]), ) print("xi shape:", par["xi"].shape) ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par) ngm = ngm_fn(t, state) return ngm
def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True): """Constructs and runs the MCMC""" if tf.test.gpu_device_name(): print("Using GPU") else: print("Using CPU") data = xarray.open_dataset(data_file, group="constant_data") cases = xarray.open_dataset(data_file, group="observations")[ "cases" ].astype(DTYPE) dates = cases.coords["time"] # Impute censored events, return cases # Take the last week of data, and repeat a further 3 times # to get a better occult initialisation. extra_cases = tf.tile(cases[:, -7:], [1, 3]) cases = tf.concat([cases, extra_cases], axis=-1) events = model_spec.impute_censored_events(cases).numpy() # 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( [ tf.constant(data["N"], DTYPE)[:, tf.newaxis], tf.zeros_like(events[:, 0, :]), ], axis=-1, ), events=events, stoichiometry=model_spec.STOICHIOMETRY, ) start_time = state.shape[1] - cases.shape[1] initial_state = state[:, start_time, :] events = events[:, start_time:-21, :] # Clip off the "extra" events ######################################################## # Construct the MCMC kernels # ######################################################## model = model_spec.CovidUK( covariates=data, initial_state=initial_state, initial_step=0, num_steps=events.shape[1], ) param_bij = tfb.Invert( # Forward transform unconstrains params tfb.Blockwise( [ tfb.Softplus(low=dtype_util.eps(DTYPE)), tfb.Identity(), tfb.Identity(), ], block_sizes=[1, 3, events.shape[1]], ) ) def joint_log_prob(unconstrained_params, events): params = param_bij.inverse(unconstrained_params) return model.log_prob( dict( psi=params[0], beta_area=params[1], gamma0=params[2], gamma1=params[3], alpha_0=params[4], alpha_t=params[5:], seir=events, ) ) + param_bij.inverse_log_det_jacobian( unconstrained_params, event_ndims=1 ) # MCMC tracing functions ############################### # Construct bursted MCMC loop # ############################### current_chain_state = [ tf.concat( [ np.array([0.1, 0.0, 0.0, 0.0], dtype=DTYPE), np.full(events.shape[1], -1.75, dtype=DTYPE,), ], axis=0, ), events, ] print("Num time steps:", events.shape[1], flush=True) print("alpha_t shape", model.event_shape["alpha_t"], flush=True) print("Initial chain state:", current_chain_state[0], flush=True) print("Initial logpi:", joint_log_prob(*current_chain_state), flush=True) # Output file posterior = run_mcmc( joint_log_prob_fn=joint_log_prob, current_state=current_chain_state, param_bijector=param_bij, initial_conditions=initial_state, config=config, output_file=output_file, ) posterior._file.create_dataset("initial_state", data=initial_state) posterior._file.create_dataset( "time", data=np.array(dates).astype(str).astype(h5py.string_dtype()), ) print(f"Acceptance theta: {posterior['results/hmc/is_accepted'][:].mean()}") 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()}" ) del posterior
) # Pre-determined thinning of posterior (better done in MCMC?) idx = range(6000, 10000, 10) param = dict( beta1=posterior["samples/beta1"][idx], beta2=posterior["samples/beta2"][idx], beta3=posterior["samples/beta3"][idx, ], sigma=posterior["samples/sigma"][idx, ], xi=posterior["samples/xi"][idx], gamma0=posterior["samples/gamma0"][idx], gamma1=posterior["samples/gamma1"][idx], ) events = posterior["samples/events"][idx] init_state = posterior["initial_state"][:] state_timeseries = compute_state(init_state, events, model_spec.STOICHIOMETRY) # Build model model = model_spec.CovidUK( covar_data, initial_state=init_state, initial_step=0, num_steps=events.shape[1], priors=config["mcmc"]["prior"], ) ngms = calc_R_it(param, events, init_state, covar_data, config["mcmc"]["prior"]) b, _ = power_iteration(ngms) rt = rayleigh_quotient(ngms, b) q = np.arange(0.05, 1.0, 0.05)
def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True): """Constructs and runs the MCMC""" if tf.test.gpu_device_name(): print("Using GPU") else: print("Using CPU") data = xarray.open_dataset(data_file, group="constant_data") cases = xarray.open_dataset(data_file, group="observations")["cases"] # We load in cases and impute missing infections first, since this sets the # time epoch which we are analysing. # Impute censored events, return cases print("Data shape:", cases.shape) events = model_spec.impute_censored_events(cases.astype(DTYPE)) # 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( [ tf.constant(data["N"], DTYPE)[:, tf.newaxis], tf.zeros_like(events[:, 0, :]), ], axis=-1, ), events=events, stoichiometry=model_spec.STOICHIOMETRY, ) start_time = state.shape[1] - cases.shape[1] initial_state = state[:, start_time, :] events = events[:, start_time:, :] ######################################################## # Construct the MCMC kernels # ######################################################## model = model_spec.CovidUK( covariates=data, initial_state=initial_state, initial_step=0, num_steps=events.shape[1], ) def joint_log_prob(block0, block1, events): return model.log_prob( dict( beta2=block0[0], gamma0=block0[1], gamma1=block0[2], sigma=block0[3], beta1=block1[0], xi=block1[1:], seir=events, )) # Build Metropolis within Gibbs sampler def make_blk0_kernel(shape, name): def fn(target_log_prob_fn, _): return tfp.mcmc.TransformedTransitionKernel( inner_kernel=AdaptiveRandomWalkMetropolis( target_log_prob_fn=target_log_prob_fn, initial_covariance=np.eye(shape[0], dtype=model_spec.DTYPE) * 1e-1, covariance_burnin=200, ), bijector=tfp.bijectors.Blockwise( bijectors=[ tfp.bijectors.Exp(), tfp.bijectors.Identity(), tfp.bijectors.Exp(), # tfp.bijectors.Identity(), ], block_sizes=[1, 2, 1], # , 5], ), name=name, ) return fn def make_blk1_kernel(shape, name): def fn(target_log_prob_fn, _): return AdaptiveRandomWalkMetropolis( target_log_prob_fn=target_log_prob_fn, initial_covariance=np.eye(shape[0], dtype=model_spec.DTYPE) * 1e-1, covariance_burnin=200, name=name, ) return fn def make_partially_observed_step(target_event_id, prev_event_id=None, next_event_id=None, name=None): def fn(target_log_prob_fn, _): return tfp.mcmc.MetropolisHastings( inner_kernel=UncalibratedEventTimesUpdate( target_log_prob_fn=target_log_prob_fn, target_event_id=target_event_id, prev_event_id=prev_event_id, next_event_id=next_event_id, initial_state=initial_state, dmax=config["dmax"], mmax=config["m"], nmax=config["nmax"], ), name=name, ) return fn def make_occults_step(prev_event_id, target_event_id, next_event_id, name): def fn(target_log_prob_fn, _): return tfp.mcmc.MetropolisHastings( inner_kernel=UncalibratedOccultUpdate( target_log_prob_fn=target_log_prob_fn, topology=TransitionTopology(prev_event_id, target_event_id, next_event_id), cumulative_event_offset=initial_state, nmax=config["occult_nmax"], t_range=(events.shape[1] - 21, events.shape[1]), name=name, ), name=name, ) return fn def make_event_multiscan_kernel(target_log_prob_fn, _): return MultiScanKernel( config["num_event_time_updates"], 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", ), ) # MCMC tracing functions def trace_results_fn(_, results): """Packs results into a dictionary""" results_dict = {} res0 = results.inner_results results_dict["block0"] = { "is_accepted": res0[0].inner_results.is_accepted, "target_log_prob": res0[0].inner_results.accepted_results.target_log_prob, } results_dict["block1"] = { "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, ]), } 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]) return results_dict # Build MCMC algorithm here. This will be run in bursts for memory economy @tf.function(autograph=use_autograph, experimental_compile=use_xla) def sample(n_samples, init_state, thin=0, previous_results=None): with tf.name_scope("main_mcmc_sample_loop"): init_state = init_state.copy() gibbs_schema = GibbsKernel( target_log_prob_fn=joint_log_prob, kernel_list=[ (0, make_blk0_kernel(init_state[0].shape, "block0")), (1, make_blk1_kernel(init_state[1].shape, "block1")), (2, make_event_multiscan_kernel), ], name="gibbs0", ) samples, results, final_results = tfp.mcmc.sample_chain( n_samples, init_state, kernel=gibbs_schema, num_steps_between_results=thin, previous_kernel_results=previous_results, return_final_kernel_results=True, trace_fn=trace_results_fn, ) return samples, results, final_results ############################### # Construct bursted MCMC loop # ############################### NUM_BURSTS = int(config["num_bursts"]) NUM_BURST_SAMPLES = int(config["num_burst_samples"]) NUM_SAVED_SAMPLES = NUM_BURST_SAMPLES * NUM_BURSTS # RNG stuff tf.random.set_seed(2) current_state = [ tf.constant([0.6, 0.0, 0.0, 0.1], dtype=DTYPE), # , 0.0, 0.0, 0.0, 0.0, 0.0], dtype=DTYPE), tf.zeros( model.model["xi"](0.0, 0.1).event_shape[-1] + 1, dtype=DTYPE, ), events, ] print("Initial logpi:", joint_log_prob(*current_state)) # Output file samples, results, _ = sample(1, current_state) posterior = Posterior( output_file, sample_dict={ "beta2": samples[0][:, 0], "gamma0": samples[0][:, 1], "gamma1": samples[0][:, 2], "sigma": samples[0][:, 3], "beta1": samples[1][:, 0], "xi": samples[1][:, 1:], "seir": samples[2], }, results_dict=results, num_samples=NUM_SAVED_SAMPLES, ) posterior._file.create_dataset("initial_state", data=initial_state) # 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"): final_results = None for i in tqdm.tqdm(range(NUM_BURSTS), unit_scale=NUM_BURST_SAMPLES * config["thin"]): samples, results, final_results = sample( NUM_BURST_SAMPLES, init_state=current_state, thin=config["thin"] - 1, previous_results=final_results, ) current_state = [s[-1] for s in samples] print(current_state[0].numpy(), flush=True) start = perf_counter() posterior.write_samples( { "beta2": samples[0][:, 0], "gamma0": samples[0][:, 1], "gamma1": samples[0][:, 2], "sigma": samples[0][:, 3], "beta1": samples[1][:, 0], "xi": samples[1][:, 1:], "seir": samples[2], }, first_dim_offset=i * NUM_BURST_SAMPLES, ) posterior.write_results(results, first_dim_offset=i * NUM_BURST_SAMPLES) end = perf_counter() print("Storage time:", end - start, "seconds") for k, v in results.items(): print( f"Acceptance {k}:", tf.reduce_mean(tf.cast(v["is_accepted"], tf.float32)), ) print( f"Acceptance theta: {posterior['results/block0/is_accepted'][:].mean()}" ) print( f"Acceptance xi: {posterior['results/block1/is_accepted'][:].mean()}") 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()}" ) del posterior
def runInference(pipelineData): # Read in settings config = pipelineData['config'] # Extract data if 'covar_data' in pipelineData: covar_data = pipelineData['covar_data'] else: covar_data, data = GetData.CovarData(config) pipelineData['covar_data'] = covar_data pipelineData['data'] = data # inference_period = config['dates']['inference_period'] # date_low = config['dates']['low'] # date_high = config['dates']['high'] # weekday = config['dates']['weekday'] # We load in cases and impute missing infections first, since this sets the # time epoch which we are analysing. cases = pipelineData['data']['cases_wide'] # Impute censored events, return cases events = model_spec.impute_censored_events(cases) # 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. _initial_state = tf.concat( [covar_data["N"], tf.zeros_like(events[:, 0, :])], axis=-1) state = compute_state( initial_state=_initial_state, events=events, stoichiometry=model_spec.STOICHIOMETRY, ) start_time = state.shape[1] - cases.shape[1] initial_state = state[:, start_time, :] events = events[:, start_time:, :] num_metapop = covar_data["N"].shape[0] ######################################################## # Build the model, and then construct the MCMC kernels # ######################################################## def convert_priors(node): if isinstance(node, dict): for k, v in node.items(): node[k] = convert_priors(v) return node return float(node) model = model_spec.CovidUK( covariates=covar_data, initial_state=initial_state, initial_step=0, num_steps=events.shape[1], priors=convert_priors(config["mcmc"]["prior"]), ) # Full joint log posterior distribution # $\pi(\theta, \xi, y^{se}, y^{ei} | y^{ir})$ # FNC NOTE: # adding another 0.0 to beta3 as TF complains of dimension mismatch otherwise def logp(block0, block1, events): return model.log_prob( dict( beta2=block0[0], gamma0=block0[1], gamma1=block0[2], sigma=block0[3], beta3=tf.concat([block0[4:6], [0.0, 0.0]], axis=-1), beta1=block1[0], xi=block1[1:], seir=events, )) # Build Metropolis within Gibbs sampler # # Kernels are: # Q(\theta, \theta^\prime) # Q(\xi, \xi^\prime) # Q(Z^{se}, Z^{se\prime}) (partially-censored) # Q(Z^{ei}, Z^{ei\prime}) (partially-censored) # Q(Z^{se}, Z^{se\prime}) (occult) # Q(Z^{ei}, Z^{ei\prime}) (occult) def make_blk0_kernel(shape, name): def fn(target_log_prob_fn, _): return tfp.mcmc.TransformedTransitionKernel( inner_kernel=AdaptiveRandomWalkMetropolis( target_log_prob_fn=target_log_prob_fn, initial_covariance=np.eye(shape[0], dtype=model_spec.DTYPE) * 1e-1, covariance_burnin=200, ), bijector=tfp.bijectors.Blockwise( bijectors=[ tfp.bijectors.Exp(), tfp.bijectors.Identity(), tfp.bijectors.Exp(), tfp.bijectors.Identity(), ], block_sizes=[1, 2, 1, 2], ), name=name, ) return fn def make_blk1_kernel(shape, name): def fn(target_log_prob_fn, _): return AdaptiveRandomWalkMetropolis( target_log_prob_fn=target_log_prob_fn, initial_covariance=np.eye(shape[0], dtype=model_spec.DTYPE) * 1e-1, covariance_burnin=200, name=name, ) return fn def make_partially_observed_step(target_event_id, prev_event_id=None, next_event_id=None, name=None): def fn(target_log_prob_fn, _): return tfp.mcmc.MetropolisHastings( inner_kernel=UncalibratedEventTimesUpdate( target_log_prob_fn=target_log_prob_fn, target_event_id=target_event_id, prev_event_id=prev_event_id, next_event_id=next_event_id, initial_state=initial_state, dmax=config["mcmc"]["dmax"], mmax=config["mcmc"]["m"], nmax=config["mcmc"]["nmax"], ), name=name, ) return fn def make_occults_step(prev_event_id, target_event_id, next_event_id, name): def fn(target_log_prob_fn, _): return tfp.mcmc.MetropolisHastings( inner_kernel=UncalibratedOccultUpdate( target_log_prob_fn=target_log_prob_fn, topology=TransitionTopology(prev_event_id, target_event_id, next_event_id), cumulative_event_offset=initial_state, nmax=config["mcmc"]["occult_nmax"], t_range=(events.shape[1] - 21, events.shape[1]), name=name, ), name=name, ) return fn def make_event_multiscan_kernel(target_log_prob_fn, _): return MultiScanKernel( config["mcmc"]["num_event_time_updates"], 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", ), ) # MCMC tracing functions def trace_results_fn(_, results): """Packs results into a dictionary""" results_dict = {} res0 = results.inner_results results_dict["block0"] = { "is_accepted": res0[0].inner_results.is_accepted, "target_log_prob": res0[0].inner_results.accepted_results.target_log_prob, } results_dict["block1"] = { "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, ]), } 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]) return results_dict # Build MCMC algorithm here. This will be run in bursts for memory economy @tf.function(autograph=False, experimental_compile=True) def sample(n_samples, init_state, thin=0, previous_results=None): with tf.name_scope("main_mcmc_sample_loop"): init_state = init_state.copy() gibbs_schema = GibbsKernel( target_log_prob_fn=logp, kernel_list=[ (0, make_blk0_kernel(init_state[0].shape, "block0")), (1, make_blk1_kernel(init_state[1].shape, "block1")), (2, make_event_multiscan_kernel), ], name="gibbs0", ) samples, results, final_results = tfp.mcmc.sample_chain( n_samples, init_state, kernel=gibbs_schema, num_steps_between_results=thin, previous_kernel_results=previous_results, return_final_kernel_results=True, trace_fn=trace_results_fn, ) return samples, results, final_results #################################### # Construct bursted MCMC loop here # #################################### # MCMC Control NUM_BURSTS = int(config["mcmc"]["num_bursts"]) NUM_BURST_SAMPLES = int(config["mcmc"]["num_burst_samples"]) NUM_EVENT_TIME_UPDATES = int(config["mcmc"]["num_event_time_updates"]) NUM_SAVED_SAMPLES = NUM_BURST_SAMPLES * NUM_BURSTS # RNG stuff tf.random.set_seed(2) current_state = [ np.array([0.2, 0.0, 0.0, 0.1, 0.0, 0.0], dtype=DTYPE), np.zeros( model.model["xi"](0.0, 0.1).event_shape[-1] # + model.model["beta3"]().event_shape[-1] + 1, dtype=DTYPE, ), events, ] print("Initial logpi:", logp(*current_state)) # Output file samples, results, _ = sample(1, current_state) print('Storing posterior data at', config["PosteriorData"]["address"]) posterior = Posterior( config["PosteriorData"]["address"], sample_dict={ "beta2": (samples[0][:, 0], (NUM_BURST_SAMPLES, )), "gamma0": (samples[0][:, 1], (NUM_BURST_SAMPLES, )), "gamma1": (samples[0][:, 2], (NUM_BURST_SAMPLES, )), "sigma": (samples[0][:, 3], (NUM_BURST_SAMPLES, )), "beta3": (samples[0][:, 4:], (NUM_BURST_SAMPLES, 2)), "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, 43, 43, 1)), # Change so it adapts size correctly }, results_dict=results, num_samples=NUM_SAVED_SAMPLES, ) posterior._file.create_dataset("initial_state", data=initial_state) posterior._file.create_dataset("config", data=yaml.dump(config)) # 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"): final_results = None for i in tqdm.tqdm(range(NUM_BURSTS), unit_scale=NUM_BURST_SAMPLES * config["mcmc"]["thin"]): samples, results, final_results = sample( NUM_BURST_SAMPLES, init_state=current_state, thin=config["mcmc"]["thin"] - 1, previous_results=final_results, ) current_state = [s[-1] for s in samples] print(current_state[0].numpy(), flush=True) start = perf_counter() posterior.write_samples( { "beta2": samples[0][:, 0], "gamma0": samples[0][:, 1], "gamma1": samples[0][:, 2], "sigma": samples[0][:, 3], "beta3": samples[0][:, 4:], "beta1": samples[1][:, 0], "xi": samples[1][:, 1:], "events": samples[2], }, first_dim_offset=i * NUM_BURST_SAMPLES, ) posterior.write_results(results, first_dim_offset=i * NUM_BURST_SAMPLES) end = perf_counter() print("Storage time:", end - start, "seconds") print( "Acceptance theta:", tf.reduce_mean( tf.cast(results["block0"]["is_accepted"], tf.float32)), ) print( "Acceptance xi:", tf.reduce_mean( tf.cast(results["block1"]["is_accepted"], tf.float32), ), ) print( "Acceptance move S->E:", tf.reduce_mean( tf.cast(results["move/S->E"]["is_accepted"], tf.float32)), ) print( "Acceptance move E->I:", tf.reduce_mean( tf.cast(results["move/E->I"]["is_accepted"], tf.float32)), ) print( "Acceptance occult S->E:", tf.reduce_mean( tf.cast(results["occult/S->E"]["is_accepted"], tf.float32)), ) print( "Acceptance occult E->I:", tf.reduce_mean( tf.cast(results["occult/E->I"]["is_accepted"], tf.float32)), ) print( f"Acceptance theta: {posterior['results/block0/is_accepted'][:].mean()}" ) print( f"Acceptance xi: {posterior['results/block1/is_accepted'][:].mean()}") 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()}" ) del posterior
def runSummary(pipelineData): # Pipeline data should contain config at a minimium config = pipelineData['config'] settings = config['SummaryData'] if settings['input'] == 'processed': summaryData = GetData.SummaryData.process(config) pipelineData['summary'] = summaryData return pipelineData # as we're running in a function, we need to assign covar_data before defining # functions that call it in order for it to be in scope # previously, covar_dict was defined in the __name__ == 'main' portion of this script # moving to a pipeline necessitates this change. # grab all data from dicts # inference_period = config['dates']['inference_period'] # date_low = config['dates']['low'] # date_high = config['dates']['high'] # weekday = config['dates']['weekday'] if 'covar_data' in pipelineData: covar_data = pipelineData['covar_data'] else: covar_data, tmp = GetData.CovarData(config) # Reproduction number calculation def calc_R_it(param, events, init_state, covar_data, priors): """Calculates effective reproduction number for batches of metapopulations :param theta: a tensor of batched theta parameters [B] + theta.shape :param xi: a tensor of batched xi parameters [B] + xi.shape :param events: a [B, M, T, X] batched events tensor :param init_state: the initial state of the epidemic at earliest inference date :param covar_data: the covariate data :return a batched vector of R_it estimates """ def r_fn(args): beta1_, beta2_, beta_3, sigma_, xi_, gamma0_, events_ = args t = events_.shape[-2] - 1 state = compute_state(init_state, events_, model_spec.STOICHIOMETRY) state = tf.gather(state, t, axis=-2) # State on final inference day model = model_spec.CovidUK( covariates=covar_data, initial_state=init_state, initial_step=0, num_steps=events_.shape[-2], priors=priors, ) xi_pred = model_spec.conditional_gp( model.model["xi"](beta1_, sigma_), xi_, tf.constant([events.shape[-2] + model_spec.XI_FREQ], dtype=model_spec.DTYPE)[:, tf.newaxis], ) # FNC NOTE: # adding another 0.0 to beta3 as TF complains of dimension mismatch otherwise par = dict( beta1=beta1_, beta2=beta2_, beta3=tf.concat([beta_3, [0.0, 0.0]], axis=-1), sigma=sigma_, gamma0=gamma0_, xi=xi_, # tf.reshape(xi_pred.sample(), [1]), ) print("xi shape:", par["xi"].shape) ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par) ngm = ngm_fn(t, state) return ngm return tf.vectorized_map( r_fn, elems=( param["beta1"], param["beta2"], param["beta3"], param["sigma"], param["xi"], param["gamma0"], events, ), ) @tf.function def predicted_incidence(param, init_state, init_step, num_steps, priors): """Runs the simulation forward in time from `init_state` at time `init_time` for `num_steps`. :param theta: a tensor of batched theta parameters [B] + theta.shape :param xi: a tensor of batched xi parameters [B] + xi.shape :param events: a [B, M, S] batched state tensor :param init_step: the initial time step :param num_steps: the number of steps to simulate :param priors: the priors for gamma :returns: a tensor of srt_quhape [B, M, num_steps, X] where X is the number of state transitions """ def sim_fn(args): beta1_, beta2_, beta3_, sigma_, xi_, gamma0_, gamma1_, init_ = args # FNC NOTE: # adding another 0.0 to beta3 as TF complains of dimension mismatch otherwise par = dict( beta1=beta1_, beta2=beta2_, beta3=tf.concat([beta3_, [0.0, 0.0]], axis=-1), gamma0=gamma0_, gamma1=gamma1_, xi=xi_, ) model = model_spec.CovidUK( covar_data, initial_state=init_, initial_step=init_step, num_steps=num_steps, priors=priors, ) sim = model.sample(**par) return sim["seir"] events = tf.map_fn( sim_fn, elems=( param["beta1"], param["beta2"], param["beta3"], param["sigma"], param["xi"], param["gamma0"], param["gamma1"], init_state, ), fn_output_signature=(tf.float64), ) return events # Today's prevalence def prevalence(predicted_state, population_size, name=None): """Computes prevalence of E and I individuals :param state: the state at a particular timepoint [batch, M, S] :param population_size: the size of the population :returns: a dict of mean and 95% credibility intervals for prevalence in units of infections per person """ prev = tf.reduce_sum(predicted_state[:, :, 1:3], axis=-1) / tf.squeeze(population_size) return mean_and_ci(prev, name=name) def predicted_events(events, name=None): num_events = tf.reduce_sum(events, axis=-1) return mean_and_ci(num_events, name=name) # Load posterior file posterior_path = config['PosteriorData']['address'] print("Using posterior:", posterior_path) posterior = h5py.File( os.path.expandvars(posterior_path, ), "r", rdcc_nbytes=1024**3, rdcc_nslots=1e6, ) # Pre-determined thinning of posterior (better done in MCMC?) if posterior["samples/beta1"].size >= 10000: idx = range(6000, 10000, 10) else: print('Using smaller MCMC sample range') print('Size of posterior["samples/beta1"] is', posterior["samples/beta1"].size) idx = range(600, 1000, 10) param = dict( beta1=posterior["samples/beta1"][idx], beta2=posterior["samples/beta2"][idx], beta3=posterior["samples/beta3"][idx, ], sigma=posterior["samples/sigma"][idx, ], xi=posterior["samples/xi"][idx], gamma0=posterior["samples/gamma0"][idx], gamma1=posterior["samples/gamma1"][idx], ) events = posterior["samples/events"][idx] init_state = posterior["initial_state"][:] state_timeseries = compute_state(init_state, events, model_spec.STOICHIOMETRY) # Build model model = model_spec.CovidUK( covar_data, initial_state=init_state, initial_step=0, num_steps=events.shape[1], priors=config["mcmc"]["prior"], ) ngms = calc_R_it(param, events, init_state, covar_data, config["mcmc"]["prior"]) b, _ = power_iteration(ngms) rt = rayleigh_quotient(ngms, b) q = np.arange(0.05, 1.0, 0.05) # FNC Note: removed dict from this and # instead added Rt as a sheet name in the excel writer rt_quantiles = pd.DataFrame(np.quantile(rt, q, axis=-1), index=q).T rt_quantiles.to_excel(config['RtQuantileData']['address'], sheet_name='Rt') # Prediction requires simulation from the last available timepoint for 28 + 4 + 1 days # Note a 4 day recording lag in the case timeseries data requires that # now = state_timeseries.shape[-2] + 4 prediction = predicted_incidence( param, init_state=state_timeseries[..., -1, :], init_step=state_timeseries.shape[-2] - 1, num_steps=70, priors=config["mcmc"]["prior"], ) predicted_state = compute_state(state_timeseries[..., -1, :], prediction, model_spec.STOICHIOMETRY) # Prevalence now prev_now = prevalence(predicted_state[..., 4, :], covar_data["N"], name="prev") # Incidence of detections now cases_now = predicted_events(prediction[..., 4:5, 2], name="cases") # Incidence from now to now+7 cases_7 = predicted_events(prediction[..., 4:11, 2], name="cases7") cases_14 = predicted_events(prediction[..., 4:18, 2], name="cases14") cases_21 = predicted_events(prediction[..., 4:25, 2], name="cases21") cases_28 = predicted_events(prediction[..., 4:32, 2], name="cases28") cases_56 = predicted_events(prediction[..., 4:60, 2], name="cases56") # Prevalence at day 7 prev_7 = prevalence(predicted_state[..., 11, :], covar_data["N"], name="prev7") prev_14 = prevalence(predicted_state[..., 18, :], covar_data["N"], name="prev14") prev_21 = prevalence(predicted_state[..., 25, :], covar_data["N"], name="prev21") prev_28 = prevalence(predicted_state[..., 32, :], covar_data["N"], name="prev28") prev_56 = prevalence(predicted_state[..., 60, :], covar_data["N"], name="prev56") # Package up summary data # this will be saved into a pickle # Add LADs in for later reference summaryData = { 'cases': { 'now': cases_now, '7': cases_7, '14': cases_14, '21': cases_21, '28': cases_28, '56': cases_56 }, 'prev': { 'now': prev_now, '7': prev_7, '14': prev_14, '21': prev_21, '28': prev_28, '56': prev_56 }, 'metrics': { 'ngms': ngms, 'b': b, 'rt': rt, 'q': q }, 'LADs': config['lad19cds'] } # Save and pass on the output data if config['GenerateOutput']['summary']: settings = config['SummaryData'] if settings['format'] == 'pickle': fn = settings['address'] with open(fn, 'wb') as file: pickle.dump(summaryData, file) pipelineData['summary'] = summaryData return pipelineData
def _log_prob(self, y, **kwargs): """Calculates the log probability of observing epidemic events y :param y: a list of tensors. The first is of shape [n_times] containing times, the second is of shape [n_times, n_states, n_states] containing event matrices. :param param: a list of parameters :returns: a scalar giving the log probability of the epidemic """ dtype = dtype_util.common_dtype([y, self.initial_state], dtype_hint=self.dtype) events = tf.convert_to_tensor(y, dtype) with tf.name_scope("StateTransitionMarginalModel.log_prob"): state_timeseries = compute_state( initial_state=self.initial_state, events=events, stoichiometry=self.stoichiometry, closed=True, ) tms_timeseries = tf.transpose(state_timeseries, perm=(1, 0, 2)) tmr_events = tf.transpose(events, perm=(1, 0, 2)) def fn(elems): return tf.stack(self.transition_rates(*elems), axis=-1) rates = tf.vectorized_map( fn=fn, elems=( self._initial_step + tf.range(tms_timeseries.shape[0]), tms_timeseries, ), ) def integrated_rate_fn(): """Use mid-point integration to estimate the constant rate over time. """ integrated_rates = tms_timeseries[..., :-1] * rates return (integrated_rates[:-1, ...] + integrated_rates[1:, ...]) / 2.0 integrated_rates = integrated_rate_fn() log_norm_constant = tf.reduce_sum( tf.math.multiply_no_nan(tf.math.log(integrated_rates), tmr_events) - tf.math.lgamma(tmr_events + 1.0), axis=(0, 1), ) pi_concentration = ( tf.reduce_sum(tmr_events, axis=(0, 1)) + self.baseline_hazard_rate_priors["concentration"]) pi_rate = (tf.reduce_sum(integrated_rates * self.time_delta, axis=(0, 1)) + self.baseline_hazard_rate_priors["rate"]) log_prob = (log_norm_constant + tf.math.lgamma(pi_concentration) - (pi_concentration) * tf.math.log(pi_rate)) return tf.reduce_sum(log_prob)
) # Pre-determined thinning of posterior (better done in MCMC?) idx = range(6000, 10000, 10) param = dict( beta1=posterior["samples/beta1"][idx], beta2=posterior["samples/beta2"][idx], beta3=posterior["samples/beta3"][idx], sigma=posterior["samples/sigma"][idx, ], xi=posterior["samples/xi"][idx], gamma0=posterior["samples/gamma0"][idx], gamma1=posterior["samples/gamma1"][idx], ) events = posterior["samples/events"][idx] init_state = posterior["initial_state"][:] state_timeseries = compute_state(init_state, events, model_spec.STOICHIOMETRY) # Prediction requires simulation from 2 weeks ago prediction = predicted_incidence( param, init_state=state_timeseries[..., -14, :], init_step=state_timeseries.shape[-2] - 14, num_steps=56, priors=config["mcmc"]["prior"], ) # prediction quantiles q_obs7 = quantile_observed(events[0, :, -7:, 2], prediction[..., 7:14, 2]) q_obs14 = quantile_observed(events[0, :, -14:, 2], prediction[..., :14, 2]) geo["Pr(pred<obs)_7"] = q_obs7.numpy()