household_sim.py 2.37 KB
 Chris Jewell committed Mar 08, 2020 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 ``````"""Household simulation for SPI-M""" import time import matplotlib.pyplot as plt import numpy as np import pandas as pd import tensorflow as tf from covid.model import CovidUK # Household size distribution in the UK, taken from ONS, 2019 hh_size_distr = pd.DataFrame({'size': [1, 2, 3, 4, 5, 6], 'freq': [6608, 8094, 3918, 3469, 1186, 463]}) # Draw households def draw_households(n: int, dist: pd.DataFrame) -> np.array: probs = dist['freq'] / np.sum(dist['freq']) hh_sizes = np.random.choice(dist['size'], n, p=probs) #hh_vec = np.repeat(np.arange(hh_sizes.shape[0]), hh_sizes) return hh_sizes # Simulator class HHstochastic(CovidUK): def __init__(self, N): self.stoichiometry = tf.constant([[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]], dtype=tf.float64) self.N = N self.popsize = tf.reduce_sum(N) def h(self, state): state = tf.unstack(state, axis=0) S, E, I, R = state eye = tf.eye(I.shape[1], dtype=tf.float64) infec_rate = tf.linalg.tensordot(I, eye, axes=[[1], [1]]) infec_rate /= self.N infec_rate *= self.param['omega'] infec_rate += tf.reduce_sum(I, axis=1, keepdims=True) / self.popsize infec_rate *= self.param['beta'] hazard_rates = tf.stack([ infec_rate, tf.constant(np.full(S.shape, self.param['nu']), dtype=tf.float64), tf.constant(np.full(S.shape, self.param['gamma']), dtype=tf.float64) ]) return hazard_rates if __name__=='__main__': hh = draw_households(100, hh_size_distr) K = (hh[:, None] == hh[None, :]) nsim = 120 model = HHstochastic(hh.astype(np.float64)) init_state = np.stack([np.broadcast_to(hh, [nsim, hh.shape[0]]), np.zeros([nsim, hh.shape[0]]), np.zeros([nsim, hh.shape[0]]), np.zeros([nsim, hh.shape[0]])]).astype(np.float64) init_state[2, :, 0] = 1. init_state[0, :, 0] = init_state[0, :, 0] - 1. start = time.perf_counter() t, sim = model.sample(init_state, [0., 365.], {'beta': 0.3, 'omega': 1.5, 'nu': 0.25, 'gamma': 0.25}) end = time.perf_counter() print(f"Completed in {end-start} seconds") r = tf.reduce_sum(sim[:, 2, :, :], axis=2) plt.plot(t, r, 'r-', alpha=0.2) plt.show()``````