Commit after Whole Household Isolation paper.

covid/model.py

@@ -109,6 +109,8 @@ class CovidUKODE:
         N_sum = N_sum[:, None] * tf.ones([1, self.n_ages])
         self.N_sum = tf.reshape(N_sum, [-1])
 
+        self.solver = tode.DormandPrince()
+
     def make_h(self, param):
 
         def h_fn(t, state):
@@ -149,8 +151,8 @@ class CovidUKODE:
         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=t0, initial_state=state_init, solution_times=t,
+
+        results = self.solver.solve(ode_fn=h, initial_time=t0, initial_state=state_init, solution_times=t,
                                previous_solver_internal_state=solver_state)
         return results.times, results.states, results.solver_internal_state
 

covid/pydata.py

+"""Python-based data munging"""
+
+import numpy as np
+import pandas as pd
+import geopandas as gp
+
+def group_ages(df):
+    """
+    Sums age groups
+    :param df: a dataframe with columns 0,1,2,...,90
+    :return: a dataframe with 5-year age groups
+    """
+    ages = np.arange(90).reshape([90//5, 5]).astype(np.str)
+    grouped_ages = pd.DataFrame()
+    for age_group in ages:
+        grouped_ages[f"[{age_group[0]}-{int(age_group[-1])+1})"] = df[age_group].sum(axis=1)
+    grouped_ages['[90,)'] = df[['90']]
+    grouped_ages['[80,inf)'] = grouped_ages[['[80-85)', '[85-90)', '[90,)']].sum(axis=1)
+    grouped_ages = grouped_ages.drop(columns=['[80-85)', '[85-90)', '[90,)'])
+    return grouped_ages
+
+
+def ingest_data(lad_shp, lad_pop):
+    pop = pd.read_csv(lad_pop, skiprows=4, thousands=',')
+    age_pop = group_ages(pop)
+    age_pop.index = pop['Code']
+
+    lad = gp.read_file(lad_shp)
+    lad.index = lad['lad19cd'].rename('Code')
+    lad = lad.iloc[lad.index.str.match('^E0[6-9]'), :]
+    lad = lad.merge(age_pop, on='Code')
+    lad.sort_index(inplace=True)
+    lad.drop(columns=['objectid', 'lad19cd' ,'long', 'lat'])
+
+    N = lad.iloc[:, lad.columns.str.match(pat='^[[0-9]')].stack()
+
+    print(f"Found {lad.shape[0]} LADs")
+
+    return {'geo': lad, 'N': N}
+
+
+if __name__=='__main__':
+    pass

covid/util.py

+"Covid analysis utility functions"
+
+import numpy as np
+
+def sanitise_parameter(par_dict):
+    """Sanitises a dictionary of parameters"""
+    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.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_total_removals(sim):
+    remove = sim[:, 3, :]
+    return remove.sum(axis=1)
+
+
+def doubling_time(t, sim, t1, t2):
+    t1 = np.where(t == np.datetime64(t1))[0]
+    t2 = np.where(t == np.datetime64(t2))[0]
+    delta = t2 - t1
+    r = sum_total_removals(sim)
+    q1 = r[t1]
+    q2 = r[t2]
+    return delta * np.log(2) / np.log(q2 / q1)
+
+
+def final_size(sim):
+    remove = sim[:, 3, :]
+    remove = remove.reshape([remove.shape[0], 152, 17])
+    fs = remove[-1, :, :].sum(axis=0)
+    return fs
+

covid_ode.py

@@ -92,6 +92,7 @@ def plot_total_curve(sim):
     plt.title('UK total cases')
     plt.xlabel('Date')
     plt.ylabel('Num infected or removed')
+    plt.grid()
     plt.legend()
 
 
@@ -220,7 +221,7 @@ if __name__ == '__main__':
         print('R_holiday=', model_holiday.eval_R0(param))
 
         # School holidays and closures
-
+        @tf.function()
         def simulate():
             t0, sim_0, solve0 = model_term.simulate(param, state_init,
                                                     settings['start'], settings['holiday'][0],

ode_config.yaml

@@ -7,8 +7,8 @@ data:
   population_size: data/pop.rds
 
 parameter:
-#  beta1: 0.0347  # R0 2.4
-  beta1: 0.04
+  beta1: 0.035  # R0 2.4
+  #beta1: 0.04
   beta2: 0.33  # Contact with commuters 1/3rd of the time
   nu: 0.25
   gamma: 0.25

prediction.py

+import optparse
+import yaml
+import tensorflow as tf
+import matplotlib.pyplot as plt
+
+from covid.rdata import load_population, load_age_mixing, load_mobility_matrix
+from covid.model import CovidUKODE
+from covid.util import *
+
+
+if __name__ == '__main__':
+
+    parser = optparse.OptionParser()
+    parser.add_option("--config", "-c", dest="config",
+                      help="configuration file")
+    options, args = parser.parse_args()
+    with open(options.config, 'r') as ymlfile:
+        config = yaml.load(ymlfile)
+
+    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, n_names = load_population(config['data']['population_size'])
+
+    param = sanitise_parameter(config['parameter'])
+    settings = sanitise_settings(config['settings'])
+
+    model_term = CovidUKODE(K_tt, T, N)
+    model_holiday = CovidUKODE(K_hh, T, 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)
+
+    @tf.function()
+    def simulate():
+        t0, sim_0, solve0 = model_term.simulate(param, state_init,
+                                                settings['start'], settings['holiday'][0],
+                                                settings['time_step'])
+        t1, sim_1, solve1 = model_holiday.simulate(param, sim_0[-1, :, :],
+                                                   settings['holiday'][0], settings['holiday'][1],
+                                                   settings['time_step'], solver_state=None)
+        t2, sim_2, _ = model_term.simulate(param, sim_1[-1, :, :],
+                                           settings['holiday'][1], settings['end'],
+                                           settings['time_step'], solver_state=None)
+        t = tf.concat([t0, t1 + t0[-1], t2 + t0[-1] + t1[-1]], axis=0)
+        sim = tf.concat([tf.expand_dims(state_init, axis=0), sim_0, sim_1, sim_2], axis=0)
+        return t, sim
+
+    t, sim = simulate()
+    dates = settings['start'] + t.numpy().astype(np.timedelta64)
+    print("Initial R0:", model_term.eval_R0(param))
+    print("Doubling time:", doubling_time(dates, sim.numpy(), '2020-02-27','2020-03-13'))
+
+    fig = plt.figure()
+    removals = tf.reduce_sum(sim[:, 3, :], axis=1)
+    infected = tf.reduce_sum(sim[:, 1:3, :], axis=[1,2])
+    exposed = tf.reduce_sum(sim[:, 1, :], axis=1)
+    date = np.squeeze(np.where(dates == np.datetime64('2020-03-13'))[0])
+    print("Daily incidence 2020-03-13:", exposed[date])
+
+    plt.plot(dates, removals[:-1]*0.10, '-', label='10% reporting')
+    plt.plot(dates, infected[:-1], '-', color='red', label='Total infected')
+    plt.plot(dates, removals[:-1], '-', color='gray', label='Total recovered/detected/died')
+
+    plt.scatter(np.datetime64('2020-03-13'), 600, label='gov.uk cases 13th March 2020')
+    plt.legend()
+    plt.grid(True)
+    fig.autofmt_xdate()
+    plt.show()
+
+

tests/test_data_ingest.py

+import unittest
+import matplotlib.pyplot as plt
+
+
+class TestData(unittest.TestCase):
+
+    def test_ingest_data(self):
+        from covid.pydata import ingest_data
+        x = ingest_data('../data/lad_shp/Local_Authority_Districts_December_2019_Boundaries_UK_BFC.shp',
+                        '../data/ukmidyearestimates20182019ladcodes.csv')
+        x['geo'].plot(column='[80,inf)', legend=True)
+        plt.show()
+
+
+if __name__ == '__main__':
+    unittest.main()

whi.py

+"""Code implementing WHI restrictions"""
+
+import optparse
+import yaml
+import numpy as np
+from scipy import optimize as opt
+import tensorflow as tf
+import matplotlib.pyplot as plt
+
+from covid.rdata import load_population, load_age_mixing, load_mobility_matrix
+from covid.model import CovidUKODE
+from covid.util import final_size
+from covid.util import *
+
+
+def optimise_beta1(R0, model, param):
+
+    def opt_fn(beta, R0):
+        param['beta1'] = beta
+        r0, i = model.eval_R0(param, 1e-9)
+        return (r0 - R0)**2
+
+    res = opt.minimize_scalar(opt_fn, args=R0)
+    return res['x']
+
+
+
+if __name__ == '__main__':
+
+    parser = optparse.OptionParser()
+    parser.add_option("--config", "-c", dest="config",
+                      help="configuration file")
+    options, args = parser.parse_args()
+    with open(options.config, 'r') as ymlfile:
+        config = yaml.load(ymlfile)
+
+    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, n_names = load_population(config['data']['population_size'])
+
+    param = sanitise_parameter(config['parameter'])
+    settings = sanitise_settings(config['settings'])
+
+    model_term = CovidUKODE(K_tt, T, N)
+    model_holiday = CovidUKODE(K_hh, T, 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)
+
+    @tf.function()
+    def simulate(param):
+        t0, sim_0, solve0 = model_term.simulate(param, state_init,
+                                                np.datetime64('2020-02-04'), settings['holiday'][0],
+                                                settings['time_step'])
+        t1, sim_1, solve1 = model_holiday.simulate(param, sim_0[-1, :, :],
+                                                   settings['holiday'][0], settings['holiday'][1],
+                                                   settings['time_step'], solver_state=None)
+        t2, sim_2, _ = model_term.simulate(param, sim_1[-1, :, :],
+                                           settings['holiday'][1], np.datetime64('2021-12-01'),
+                                           settings['time_step'], solver_state=None)
+        t = tf.concat([t0, t1 + t0[-1], t2 + t0[-1] + t1[-1]], axis=0)
+        sim = tf.concat([tf.expand_dims(state_init, axis=0), sim_0, sim_1, sim_2], axis=0)
+        return t, sim
+
+    t, sim = simulate(param)
+    dates = settings['start'] + t.numpy().astype(np.timedelta64)
+    print("Initial R0:", model_term.eval_R0(param))
+    print("Doubling time:", doubling_time(dates, sim.numpy(), '2020-02-27','2020-03-13'))
+    print("Attack rate:", np.sum(final_size(sim.numpy()))/np.sum(N))
+
+    param['beta1'] = optimise_beta1(1.6, model_term, param)
+    t_whi, sim_whi = simulate(param)
+    print("Initial R0:", model_term.eval_R0(param))
+    print("Doubling time:", doubling_time(dates, sim.numpy(), '2020-02-27', '2020-03-13'))
+    print("Attack rate:", np.sum(final_size(sim_whi.numpy()))/np.sum(N))
+
+    param['beta1'] = optimise_beta1(2.2, model_term, param)
+    t_whi, sim_whi_05 = simulate(param)
+    print("Initial R0:", model_term.eval_R0(param))
+    print("Doubling time:", doubling_time(dates, sim.numpy(), '2020-02-27', '2020-03-13'))
+    print("Attack rate:", np.sum(final_size(sim_whi_05.numpy()))/np.sum(N))
+    fig = plt.figure()
+    removals = tf.reduce_sum(sim[:, 3, :], axis=1)
+    infected = tf.reduce_sum(sim[:, 1:3, :], axis=[1, 2])
+    infected_whi = tf.reduce_sum(sim_whi[:, 1:3, :], axis=[1, 2])
+    infected_whi_05 = tf.reduce_sum(sim_whi_05[:, 1:3, :], axis=[1, 2])
+
+    date = np.squeeze(np.where(dates == np.datetime64('2020-03-13'))[0])
+    plt.plot(dates, infected[:-1]/1e6, '-', color='lightblue', label='$\eta=0, R_\star=2.7$')
+    plt.plot(dates, infected_whi[:-1]/1e6, '-', color='pink', label='$\eta=1, R_\star=1.6$')
+    plt.plot(dates, infected_whi_05[:-1]/1e6, '-', color='orange', label='$\eta=0.5, R_\star=2.2$')
+    plt.legend()
+    plt.ylabel("Cases in $10^6$ individuals")
+    plt.xlabel("Date")
+    plt.grid(True)
+    fig.autofmt_xdate()
+    plt.show()
+
+