-
Notifications
You must be signed in to change notification settings - Fork 0
/
SIRSData.py
147 lines (113 loc) · 5.52 KB
/
SIRSData.py
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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 8 22:21:59 2023
@author: djord
"""
from SIRSUpdate import Simulate
import numpy as np
import matplotlib.pyplot as plt
import random
class SIRSfluc(object): #Set up simple class to analyze data dependent on chosen simulation
def __init__(self, pr): # Initially set probability of recovery; only variable constant throughout analyzing all required data
self.pr = pr
def phase_diagram(self):
pi = np.arange(0, 1, 0.05)
ps = np.arange(0, 1, 0.05)
infection_average = []
for i in range (0, len(pi)): # iterate over range of probability for infection and susceptibility
for j in range (0, len(ps)):
self.SIRS = Simulate(50)
infected_data = []
for k in range (0, 500): # chose 500 iterations due to time constraints, found to still be accurate
for l in range (0, 2500):
self.SIRS.update(pi[i], self.pr, ps[j])
if k > 100: # Find infection average after 100 iterations i.e. after equilibrium reached
I = self.SIRS.infected()
infected_data.append(I)
infected_data = np.array(infected_data)
average = np.mean(infected_data)
infection_average.append(average)
file = open("PI VS PS.txt", "a")
file.write("infection average: " + str(average) + ", PI: " + str(pi[i]) + ", PS: "+ str(ps[j])+ ".\n")
file.close()
print(str(i) + " Done")
return infection_average, pi, ps
def waves(self, ps):
pi = np.arange(0.2, 0.51, 0.01)
self.ps = ps
infection_variance = []
error = [0] * len(pi)
for i in range (0, len(pi)):
self.SIRS = Simulate(50)
infected_data = []
for k in range (0, 10000):
for l in range (0, 2500):
self.SIRS.update(pi[i], self.pr, self.ps)
if k > 100:
I = self.SIRS.infected()
infected_data.append(I)
if k%1000 == 0:
print(str(k/100) + "% done.")
infected_data = np.array(infected_data)
variance = np.var(infected_data)
variance = variance/2500
infection_variance.append(variance)
file = open("Waves.txt", "a")
file.write("infection variance: " + str(variance) + ", PI: " + str(pi[i]))
file.close()
print(str(i) + " Done")
error_data = [] # Calculated error using resampling bootstrap method
for j in range(1000):#no. of resamplings
resample_infected_data = random.choices(infected_data, k=len(infected_data))
err = np.var(resample_infected_data)
error_data.append(err)
error[i] = np.std(error_data)
error[i] = error[i]/2500
file = open("Waves.txt", "a")
file.write(" with std. error: " + str(error[i]) + ".\n")
file.close()
return infection_variance, pi, error
def immunity_infection(self):
vax = np.arange(0, 1, 0.025) # Set up possible range of immunity percentage of population
infection_average = []
for i in range (0, len(vax)): # Iterate over different immunities
self.SIRS = Simulate(50)
infected_data = []
self.SIRS.initial_immune(50, vax[i])
for k in range (0, 1000):
for l in range (0, 2500):
self.SIRS.update(0.5, 0.5, 0.5)
if k > 100:
I = self.SIRS.infected()
infected_data.append(I)
infected_data = np.array(infected_data)
average = np.mean(infected_data) # Find average infection
infection_average.append(average)
file = open("Immunity vs Infection.txt", "a")
file.write("infection average: " + str(average) + ", immunity " + str(vax[i]) + ".\n")
file.close()
print(str(i) + " Done")
return infection_average, vax
def main():
# S = SIRSfluc(0.5)
# infection_average, pi, ps = S.phase_diagram()
# infection_average = np.array(infection_average).reshape(len(pi), len(ps))
# fig, ax = plt.subplots()
# im = ax.imshow(infection_average, extent = (pi.min(), pi.max(), ps.min(), ps.max()), cmap='viridis')
# cbar = ax.figure.colorbar(im, ax=ax)
# cbar.ax.set_ylabel("Average infection rate", rotation=-90, va="bottom")
# plt.xlabel("Probability of Susceptibility")
# plt.ylabel("Probability of Infection")
S = SIRSfluc(0.5)
infection_variance, pi, error = S.waves(0.5)
plt.plot(pi, infection_variance)
plt.errorbar(pi, infection_variance, yerr= error, fmt=".k")
plt.title("Infection variance vs Probability of infection at PR = PS = 0.5")
plt.xlabel("Probability of Infection")
plt.ylabel("Infection variance")
# S = SIRSfluc(0.5)
# infection_average, vax = S.immunity_infection()
# plt.plot(vax, infection_average)
# plt.xlabel("Immunity percentage.")
# plt.ylabel("Infection average.")
main()