Implemented household sim

covid/impl/chainbinom_simulate.py

@@ -31,7 +31,7 @@ def chain_binomial_simulate(hazard_fn, state, start, end, time_step, stoichiomet
     propagate = chain_binomial_propagate(hazard_fn, time_step, stoichiometry)
     times = tf.range(start, end, time_step)
 
-    output = tf.TensorArray(tf.float32, size=times.shape[0])
+    output = tf.TensorArray(tf.float64, size=times.shape[0])
     output = output.write(0, state)
 
     for i in tf.range(1, times.shape[0]):

covid/model.py

@@ -8,6 +8,25 @@ import numpy as np
 
 from covid.impl.chainbinom_simulate import chain_binomial_simulate
 
+def power_iteration(A, tol=1e-3):
+    b_k = tf.random.normal([A.shape[1], 1])
+    epsilon = 1.
+    i = 0
+    while tf.greater(epsilon, tol):
+        b_k1 = tf.matmul(A, b_k)
+        b_k1_norm = tf.linalg.norm(b_k1)
+        b_k_new = b_k1 / b_k1_norm
+        epsilon = tf.reduce_sum(tf.pow(b_k_new-b_k, 2))
+        b_k = b_k_new
+        i += 1
+    return b_k, i
+
+#@tf.function
+def rayleigh_quotient(A, b):
+    b = tf.reshape(b, [b.shape[0], 1])
+    numerator = tf.matmul(tf.transpose(b), tf.matmul(A, b))
+    denominator = tf.matmul(tf.transpose(b), b)
+    return numerator / denominator
 
 class CovidUK:
     def __init__(self, K, T, W):
@@ -19,19 +38,18 @@ class CovidUK:
                               [0, -1, 1, 0],
                               [0, 0, -1, 1]]
 
-    @tf.function
     def h(self, state):
         state = tf.unstack(state, axis=0)
         S, E, I, R = state
 
         hazard_rates = tf.stack([
-            self.param['beta1'] * tf.dot(self.T, tf.dot(self.K, I)),
+            self.param['beta1'] * tf.dot(self.T, tf.dot(self.K, I))/self.K.shape[0],
             self.param['nu'],
             self.param['gamma']
         ])
         return hazard_rates
 
-    @tf.function
+
     def sample(self, initial_state, time_lims, param):
         self.param = param
         return chain_binomial_simulate(self.h, initial_state, time_lims[0],
@@ -64,34 +82,42 @@ class Homogeneous:
 
 class CovidUKODE:
 
-    def __init__(self, K, T, N, n_lads):
+    def __init__(self, K, T, N):
         """Represents a CovidUK ODE model
 
-        :param K: a MxM matrix of age group mixing
-        :param T: a n_ladsxn_lads matrix of inter-LAD connectivity
+        :param K: a MxM matrix of age group mixing in term time
+        :param T: a n_ladsxn_lads matrix of inter-LAD commuting
         :param N: a vector of population sizes in each LAD
         :param n_lads: the number of LADS
         """
+        self.n_ages = K.shape[0]
+        self.n_lads = T.shape[0]
+
         self.Kbar = tf.reduce_mean(K)
         self.K = tf.linalg.LinearOperatorFullMatrix(K)
-        eye = tf.linalg.LinearOperatorIdentity(n_lads)
-        self.K = tf.linalg.LinearOperatorKronecker([self.K, eye])
+        eye = tf.linalg.LinearOperatorIdentity(self.n_lads)
+        self.K = tf.linalg.LinearOperatorKronecker([eye, self.K])
 
         self.T = tf.linalg.LinearOperatorFullMatrix(T)
         shp = tf.linalg.LinearOperatorFullMatrix(np.ones_like(K, dtype=np.float32))
         self.T = tf.linalg.LinearOperatorKronecker([self.T, shp])
 
-        self.N = N
-        self.n_lads = n_lads
+        self.N = tf.constant(N, dtype=tf.float32)
+
+        N_matrix = tf.reshape(self.N, [self.n_lads, self.n_ages])
+        N_sum = tf.reduce_sum(N_matrix, axis=1)
+        N_sum = N_sum[:, None] * tf.ones([1, self.n_ages])
+        self.N_sum = tf.reshape(N_sum, [-1])
 
     def make_h(self, param):
+
         def h_fn(t, state):
             state = tf.unstack(state, axis=0)
             S, E, I, R = state
 
-            infec_rate = tf.linalg.matvec(self.K, I)
-            infec_rate += self.Kbar * tf.linalg.matvec(self.T, I / self.N)
-            infec_rate = param['beta'] * S / self.N * infec_rate
+            infec_rate = param['beta1'] * tf.linalg.matvec(self.K, I)
+            infec_rate += param['beta1'] * param['beta2'] * self.Kbar * tf.linalg.matvec(self.T, I / self.N_sum)
+            infec_rate = S / self.N * infec_rate
 
             dS = -infec_rate
             dE = infec_rate - param['nu'] * E
@@ -105,17 +131,37 @@ class CovidUKODE:
 
     def create_initial_state(self, init_matrix=None):
         if init_matrix is None:
-            I = np.ones(self.N.shape, dtype=np.float32)
-            S = self.N - I
-            E = np.zeros(self.N.shape, dtype=np.float32)
-            R = np.zeros(self.N.shape, dtype=np.float32)
-            return np.stack([S, E, I, R])
+            I = np.zeros(self.N.shape, dtype=np.float32)
+            I[149*17+10] = 30. # Middle-aged in Surrey
+        else:
+            np.testing.assert_array_equal(init_matrix.shape, [self.n_lads, self.n_ages],
+                                          err_msg=f"init_matrix does not have shape [<num lads>,<num ages>] \
+                                          ({self.n_lads},{self.n_ages})")
+            I = init_matrix.flatten()
+        S = self.N - I
+        E = np.zeros(self.N.shape, dtype=np.float32)
+        R = np.zeros(self.N.shape, dtype=np.float32)
+        return np.stack([S, E, I, R])
 
     @tf.function
     def simulate(self, param, state_init, t_start, t_end, t_step=1.):
         h = self.make_h(param)
-        t = np.arange(start=t_start, stop=t_end, step=t_step)
-        print(f"Running simulation with times {t_start}-{t_end}:{t_step}")
+        t0 = 0.
+        t1 = (t_end - t_start) / np.timedelta64(1, 'D')
+        t = np.arange(start=t0, stop=t1, step=t_step)[1:]
         solver = tode.DormandPrince()
-        results = solver.solve(ode_fn=h, initial_time=t_start, initial_state=state_init, solution_times=t)
+        results = solver.solve(ode_fn=h, initial_time=t0, initial_state=state_init, solution_times=t)
         return results.times, results.states
+
+    def ngm(self, param):
+        infec_rate = param['beta1'] * self.K.to_dense()
+        infec_rate += param['beta1'] * param['beta2'] * self.Kbar * self.T.to_dense() / self.N_sum[None, :]
+        ngm = infec_rate / param['gamma']
+        return ngm
+
+    def eval_R0(self, param, tol=1e-8):
+        ngm = self.ngm(param)
+        # Dominant eigen value by power iteration
+        dom_eigen_vec, i = power_iteration(ngm, tol=tol)
+        R0 = rayleigh_quotient(ngm, dom_eigen_vec)
+        return tf.squeeze(R0), i

covid/rdata.py

@@ -2,6 +2,7 @@
 
 import numpy
 import pyreadr as pyr
+import numpy as np
 
 
 def load_age_mixing(rds_file: str):
@@ -10,7 +11,9 @@ def load_age_mixing(rds_file: str):
     :param rds_file: a .rds file containing an R data.frame with mixing matrix
     """
     raw = pyr.read_r(rds_file)
-    return list(raw.values())[0]
+    K = list(raw.values())[0]
+    age_groups = K.columns
+    return K.to_numpy(dtype=np.float32), age_groups
 
 
 def load_mobility_matrix(rds_file: str):
@@ -24,7 +27,7 @@ def load_mobility_matrix(rds_file: str):
     colnames = df.columns
     mobility_matrix = df.pivot(index='Workplace', columns='Residence', values=colnames[0])
     mobility_matrix[mobility_matrix.isna()] = 0.
-    return mobility_matrix
+    return mobility_matrix.to_numpy(dtype=np.float32), mobility_matrix.index.to_numpy()
 
 
 def load_population(rds_file: str):
@@ -35,4 +38,4 @@ def load_population(rds_file: str):
     raw = pyr.read_r(rds_file)
     df = list(raw.values())[0]
     df = df.sort_values(by=['LA.code', 'age'])
-    return df
+    return df['n'].to_numpy(dtype=np.float32), df[['name', 'Area.name.2']]

covid_ode.py

 import optparse
-import numpy as np
-import yaml
 import time
+
+import h5py
 import matplotlib.pyplot as plt
+import tensorflow as tf
+import yaml
 
 from covid.model import CovidUKODE
 from covid.rdata import *
@@ -10,18 +12,149 @@ from covid.rdata import *
 
 def sanitise_parameter(par_dict):
     """Sanitises a dictionary of parameters"""
-    par = ['beta','nu','gamma']
+    par = ['beta1', 'beta2', 'nu','gamma']
     d = {key: np.float64(par_dict[key]) for key in par}
     return d
 
 
 def sanitise_settings(par_dict):
-    d = {'start': np.float64(par_dict['start']),
-         'end': np.float64(par_dict['end']),
-         'time_step': np.float64(par_dict['time_step'])}
+    d = {'start': np.datetime64(par_dict['start']),
+         'end': np.datetime64(par_dict['end']),
+         'time_step': float(par_dict['time_step']),
+         'holiday': np.array([np.datetime64(date) for date in par_dict['holiday']])}
     return d
 
 
+def seed_areas(N, names, age_group=8, num_la=152, num_age=17, n_seed=30.):
+    areas = ['Inner London',
+             'Outer London',
+             'West Midlands (Met County)',
+             'Greater Manchester (Met County)']
+
+    names_matrix = names['Area.name.2'].to_numpy().reshape([num_la, num_age])
+
+    seed_areas = np.in1d(names_matrix[:, age_group], areas)
+    N_matrix = N.reshape([num_la, num_age])  # LA x Age
+
+    pop_size_sub = N_matrix[seed_areas, age_group]  # Gather
+    n = np.round(n_seed * pop_size_sub / pop_size_sub.sum())
+
+    seeding = np.zeros_like(N_matrix)
+    seeding[seed_areas, age_group] = n  # Scatter
+    return seeding
+
+
+def sum_age_groups(sim):
+    infec = sim[:, 2, :]
+    infec = infec.reshape([infec.shape[0], 152, 17])
+    infec_uk = infec.sum(axis=2)
+    return infec_uk
+
+
+def sum_la(sim):
+    infec = sim[:, 2, :]
+    infec = infec.reshape([infec.shape[0], 152, 17])
+    infec_uk = infec.sum(axis=1)
+    return infec_uk
+
+
+def sum_total_removals(sim):
+    remove = sim[:, 3, :]
+    return remove.sum(axis=1)
+
+
+def final_size(sim):
+    remove = sim[:, 3, :]
+    remove = remove.reshape([remove.shape[0], 152, 17])
+    fs = remove[-1, :, :].sum(axis=0)
+    return fs
+
+
+def write_hdf5(filename, param, t, sim):
+    with h5py.File(filename, "w") as f:
+        dset_sim = f.create_dataset("simulation", sim.shape, dtype='f')
+        dset_sim[:] = sim
+        dset_t = f.create_dataset("time", t.shape, dtype='f')
+        dset_t[:] = t
+        grp_param = f.create_group("parameter")
+        for k, v in param.items():
+            d_beta = grp_param.create_dataset(k, [1], dtype='f')
+            d_beta[()] = v
+
+
+
+def plot_total_curve(sim):
+    infec_uk = sum_la(sim)
+    infec_uk = infec_uk.sum(axis=1)
+    removals = sum_total_removals(sim)
+    times = np.datetime64('2020-02-20') + np.arange(removals.shape[0])
+    plt.plot(times, infec_uk, 'r-', label='Infected')
+    plt.plot(times, removals, 'b-', label='Removed')
+    plt.title('UK total cases')
+    plt.xlabel('Date')
+    plt.ylabel('Num infected or removed')
+    plt.legend()
+
+
+def plot_by_age(sim, labels, t0=np.datetime64('2020-02-20'), ax=None):
+    if ax is None:
+        ax = plt.figure().gca()
+    infec_uk = sum_la(sim)
+    total_uk = infec_uk.mean(axis=1)
+    t = t0 + np.arange(infec_uk.shape[0])
+    colours = plt.cm.viridis(np.linspace(0., 1., infec_uk.shape[1]))
+    for i in range(infec_uk.shape[1]):
+        ax.plot(t, infec_uk[:, i], 'r-', alpha=0.4, color=colours[i], label=labels[i])
+    ax.plot(t, total_uk, '-', color='black', label='Mean')
+    return ax
+
+
+def plot_by_la(sim, labels,t0=np.datetime64('2020-02-20'), ax=None):
+    if ax is None:
+        ax = plt.figure().gca()
+    infec_uk = sum_age_groups(sim)
+    total_uk = infec_uk.mean(axis=1)
+    t = t0 + np.arange(infec_uk.shape[0])
+    colours = plt.cm.viridis(np.linspace(0., 1., infec_uk.shape[1]))
+    for i in range(infec_uk.shape[1]):
+        ax.plot(t, infec_uk[:, i], 'r-', alpha=0.4, color=colours[i], label=labels[i])
+    ax.plot(t, total_uk, '-', color='black', label='Mean')
+    return ax
+
+
+def draw_figs(sim, N):
+    # Attack rate
+    N = N.reshape([152, 17]).sum(axis=0)
+    fs = final_size(sim)
+    attack_rate = fs / N
+    print("Attack rate:", attack_rate)
+    print("Overall attack rate: ", np.sum(fs)/np.sum(N))
+
+    # Total UK epidemic curve
+    plot_total_curve(sim)
+    plt.xticks(rotation=45, horizontalalignment="right")
+    plt.savefig('total_uk_curve.pdf')
+    plt.show()
+
+    # TotalUK epidemic curve by age-group
+    fig, ax = plt.subplots(1, 2, figsize=[24, 12])
+    plot_by_la(sim, la_names, ax=ax[0])
+    plot_by_age(sim, age_groups, ax=ax[1])
+    ax[1].legend()
+    plt.xticks(rotation=45, horizontalalignment="right")
+    fig.autofmt_xdate()
+    plt.savefig('la_age_infec_curves.pdf')
+    plt.show()
+
+    # Plot attack rate
+    plt.figure(figsize=[4, 2])
+    plt.plot(age_groups, attack_rate, 'o-')
+    plt.xticks(rotation=90)
+    plt.title('Age-specific attack rate')
+    plt.savefig('age_attack_rate.pdf')
+    plt.show()
+
+
 if __name__ == '__main__':
 
     parser = optparse.OptionParser()
@@ -31,22 +164,52 @@ if __name__ == '__main__':
     with open(options.config, 'r') as ymlfile:
         config = yaml.load(ymlfile)
 
-    K = load_age_mixing(config['data']['age_mixing_matrix']).to_numpy(dtype=np.float32)
-    T = load_mobility_matrix(config['data']['mobility_matrix']).to_numpy(dtype=np.float32)
+    K_tt, age_groups = load_age_mixing(config['data']['age_mixing_matrix_term'])
+    K_hh, _ = load_age_mixing(config['data']['age_mixing_matrix_hol'])
+
+    T, la_names = load_mobility_matrix(config['data']['mobility_matrix'])
     np.fill_diagonal(T, 0.)
-    N = load_population(config['data']['population_size'])['n'].to_numpy(dtype=np.float32)
+
+    N, n_names = load_population(config['data']['population_size'])
 
     param = sanitise_parameter(config['parameter'])
     settings = sanitise_settings(config['settings'])
 
-    model = CovidUKODE(K, T, N, T.shape[0])
 
-    state_init = model.create_initial_state()
+# Straight, no school closures
+    model_term = CovidUKODE(K_tt, T, N)
+    model_holiday = CovidUKODE(K_hh/2., T/2., N)
+
+    seeding = seed_areas(N, n_names)  # Seed 40-44 age group, 30 seeds by popn size
+    state_init = model_term.create_initial_state(init_matrix=seeding)
+
+    print('R_term=', model_term.eval_R0(param))
+    print('R_holiday=', model_holiday.eval_R0(param))
+
+    # School holidays and closures
+    @tf.function
+    def simulate():
+        t0, sim_0 = model_term.simulate(param, state_init,
+                                        settings['start'], settings['holiday'][0],
+                                        settings['time_step'])
+        t1, sim_1 = model_holiday.simulate(param, sim_0[-1, :, :],
+                                           settings['holiday'][0], settings['holiday'][1],
+                                           settings['time_step'])
+        t2, sim_2 = model_term.simulate(param, sim_1[-1, :, :],
+                                        settings['holiday'][1], settings['end'],
+                                        settings['time_step'])
+        t = tf.concat([t0, t1, t2], axis=0)
+        sim = tf.concat([tf.expand_dims(state_init, axis=0), sim_0, sim_1, sim_2], axis=0)
+        return t, sim
+
     start = time.perf_counter()
-    t, sim = model.simulate(param, state_init,
-                            settings['start'], settings['end'],
-                            settings['time_step'])
+    t, sim = simulate()
     end = time.perf_counter()
-    print(f"Complete in {end-start} seconds")
-    plt.plot(t, sim[:, 2, 0], 'r-')
-    plt.show()
\ No newline at end of file
+    print(f'Complete in {end-start} seconds')
+
+    draw_figs(sim.numpy(), N)
+
+    if 'simulation' in config['output']:
+        write_hdf5(config['output']['simulation'], param, t, sim)
+
+    print(f"Complete in {end-start} seconds")
\ No newline at end of file

household_sim.py

+"""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()

ode_config.yaml

 # Covid_ODE model configuration
 
 data:
-  age_mixing_matrix: data/polymod_normal_df.rds
+  age_mixing_matrix_term: data/polymod_normal_df.rds
+  age_mixing_matrix_hol: data/polymod_no_school_df.rds
   mobility_matrix: data/movement.rds
   population_size: data/pop.rds
 
 parameter:
-  beta: 0.09
+  beta1: 0.04  # R0 2.4
+  beta2: 0.33  # Contact with commuters 1/3rd of the time
   nu: 0.25
-  gamma: 0.65
+  gamma: 0.25
 
 settings:
-  start: 0.
-  end: 365.
+  start: 2020-02-04
+  end: 2020-12-30
+  holiday:
+    - 2020-04-06
+    - 2020-05-01
   time_step: 1.
 
 output:
-  simulation: covid_ode.sim
+  simulation: covid_ode.hd5

r0_calc.py

+import numpy as np
+from scipy import optimize as opt
+
+from covid.rdata import *
+from covid.model import CovidUKODE
+
+if __name__=='__main__':
+
+    K, _ = load_age_mixing('data/polymod_normal_df.rds')
+    T, _ = load_mobility_matrix('data/movement.rds')
+    np.fill_diagonal(T, 0.)
+    N, _ = load_population('data/pop.rds')
+
+    model = CovidUKODE(K, T, N)
+
+    print("R0 (beta=0.036) = ",  model.eval_R0({'beta1': 0.036, 'beta2': 0.33, 'gamma': 0.25}))
+    R0 = 2.4
+
+    def opt_fn(beta, R0):
+        r0, i = model.eval_R0({'beta1': beta, 'beta2': 0.33, 'gamma': 0.25}, 1e-9)
+        return (r0 - R0)**2
+
+    res = opt.minimize_scalar(opt_fn, args=R0)
+    print(res['x'])

scratch_covid.py

-import tensorflow as tf
 import numpy as np
 
-from covid.model import Homogeneous
+from covid.rdata import *
+from covid.model import CovidUKODE
 
-sw = tf.summary.create_file_writer('tf_log')
+K = load_age_mixing('data/polymod_normal_df.rds').to_numpy(dtype=np.float32)
+T = load_mobility_matrix('data/movement.rds').to_numpy(dtype=np.float32)
+np.fill_diagonal(T, 0.)
+N = load_population('data/pop.rds')['n'].to_numpy(dtype=np.float32)
 
-model = Homogeneous()
-popsize = 2500
-state = np.array([np.full([popsize], 999.),
-                  np.full([popsize], 0.),
-                  np.full([popsize], 1.),
-                  np.full([popsize], 0.)], dtype=np.float32)
+model = CovidUKODE(K, T, N, T.shape[0])
 
-print("Running...", flush=True, sep='')
-tf.summary.trace_on(graph=True, profiler=True)
-
-t, sim_state = model.sample(state,
-                            [0., 10.],
-                            {'beta': 0.2, 'nu': 0.14, 'gamma':0.14})
-print("Done", flush=True)
-print(sim_state)
-
-with sw.as_default():
-    tf.summary.trace_export('profile', step=0, profiler_outdir='tf_log')
-
-sw.flush()
-sw.close()