covid_ode.py 1.61 KB
Newer Older
1
import optparse
Chris Jewell's avatar
Chris Jewell committed
2
import numpy as np
3
4
5
import yaml
import time
import matplotlib.pyplot as plt
Chris Jewell's avatar
Chris Jewell committed
6

7
8
from covid.model import CovidUKODE
from covid.rdata import *
Chris Jewell's avatar
Chris Jewell committed
9
10


11
12
13
14
15
def sanitise_parameter(par_dict):
    """Sanitises a dictionary of parameters"""
    par = ['beta','nu','gamma']
    d = {key: np.float64(par_dict[key]) for key in par}
    return d
Chris Jewell's avatar
Chris Jewell committed
16

Chris Jewell's avatar
Chris Jewell committed
17

18
19
20
21
22
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'])}
    return d
Chris Jewell's avatar
Chris Jewell committed
23

Chris Jewell's avatar
Chris Jewell committed
24

25
if __name__ == '__main__':
Chris Jewell's avatar
Chris Jewell committed
26

27
28
29
30
31
32
    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)
Chris Jewell's avatar
Chris Jewell committed
33

34
35
36
37
    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)
    np.fill_diagonal(T, 0.)
    N = load_population(config['data']['population_size'])['n'].to_numpy(dtype=np.float32)
Chris Jewell's avatar
Chris Jewell committed
38

39
40
    param = sanitise_parameter(config['parameter'])
    settings = sanitise_settings(config['settings'])
Chris Jewell's avatar
Chris Jewell committed
41

42
    model = CovidUKODE(K, T, N, T.shape[0])
Chris Jewell's avatar
Chris Jewell committed
43

44
45
46
47
48
49
50
51
52
    state_init = model.create_initial_state()
    start = time.perf_counter()
    t, sim = model.simulate(param, state_init,
                            settings['start'], settings['end'],
                            settings['time_step'])
    end = time.perf_counter()
    print(f"Complete in {end-start} seconds")
    plt.plot(t, sim[:, 2, 0], 'r-')
    plt.show()