SIRS/test_sirs.py at master · LariDG/SIRS · GitHub

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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
"""
Modelling and Visulaisation in Physics
Checkpoint 2: SIRS
Main for SIRS model.
To run the simulation from the class as well as plotting a phase diagram
based on the changing probabailities of S --> I and R --> S.
Author: Larisa Dorman-Gajic
"""
import numpy as np
import random
import math
from SIRS import SIRS
import matplotlib.pyplot as plt
import sys

def main():

    size_input = sys.argv[1]
    simulate = sys.argv[2]
    size = (50, 50)

    if simulate == "Y":
        p_1 = 0.8
        p_2 = 0.1
        p_3 = 0.01
        game = SIRS(size, p_1, p_2, p_3)
        game.run(10000, 10000)

    elif simulate == "N":
        plot = sys.argv[3]
        p_1_range = np.arange(0, 1, 0.1)
        p_2 = 0.5
        p_3_range = np.arange(0, 1, 0.1)
        i_matrix = []
        i_matrix.append([0.0]*len(p_3_range))
        i_matrix.append([0.0]*len(p_3_range))
        i_matrix.append([0.0]*len(p_3_range))
        i_matrix.append([0.0]*len(p_3_range))
        i_matrix.append([0.0]*len(p_3_range))
        for n in range(5, len(p_1_range)):
            p_1 = p_1_range[n]
            print(p_1)
            i_avg_list = [0.0, 0.0, 0.0, 0.0, 0.0]
            for m in range(5, len(p_3_range)):
                p_3 = p_3_range[m]
                print(p_3)
                game = SIRS(size, p_1, p_2, p_3)
                infected = []
                for i in range(22):
                    for j in range(size[0]*size[1]):
                        game.update()
                    if i > 2:
                        infected_sites = game.infected_sites()
                        infected.append(infected_sites)

                if plot == "PD":
                    infected_avg = np.mean(infected)/(size[0]*size[1])
                    i_avg_list.append(infected_avg)
                elif plot == "waves":
                    infected_variance = np.var(infected)/(size[0]*size[1])
                    i_avg_list.append(infected_variance)
            i_matrix.append(i_avg_list)
        print(i_matrix)

        #plt.imshow(i_matrix, cmap = 'hot', interpolation = 'nearest', extent = [0,1,1,0])
        #plt.show()

        numpy_matrix = np.matrix(i_matrix)
        if plot == "PD":
            np.savetxt("phase.txt", numpy_matrix)
        elif plot == "waves":
            np.savetxt("waves.txt", numpy_matrix)


    elif simulate == "M":
        p_1_range = np.arange(0.2, 0.525, 0.1)
        p_2_range = np.arange(0.2, 0.525, 0.1)
        p_3 = 0.5
        i_matrix = []
        for n in range(len(p_1_range)):
            p_1 = p_1_range[n]
            i_var_list = []
            for m in range(len(p_2_range)):
                p_2 = p_2_range[m]
                print(p_2)
                game = SIRS(size, p_1, p_2, p_3)
                infected = []
                for i in range(10100):
                    for j in range(size[0]*size[1]):
                        game.update()
                    if i > 100:
                        infected_sites = game.infected_sites()
                        infected.append(infected_sites)
                infected_variance = np.var(infected)/(size[0]*size[1])
                i_var_list.append(infected_variance)
            i_matrix.append(i_var_list)

        #plt.imshow(i_matrix, cmap = 'hot', interpolation = 'nearest', extent = [0,1,1,0])
        #plt.show()

        numpy_matrix = np.matrix(i_matrix)
        np.savetxt("p_3_fixed_var.txt", numpy_matrix)


main()