-
Notifications
You must be signed in to change notification settings - Fork 0
/
dynamic_calculate_r0.py
110 lines (90 loc) · 3.54 KB
/
dynamic_calculate_r0.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
import numpy as np
from graph_tool.generation import price_network, geometric_graph
from graph_tool import infect_vertex_property
import matplotlib.pyplot as plt
SUSCEPTIBLE = 0
INFECTED = 1
RECOVERED = 2
print('Imported all libraries')
def coinflip(p):
"""Return True with probability p and False otherwise."""
return np.random.random() < p
def SIRmodel(g, patient_zero, beta, gamma):
"""
Runs a standard SIR model on graph g with one infected individual,
patient_zero.
g - graph
state - numpy array
"""
# Initialise states
state = g.new_vertex_property('int32_t', False)
state.a[patient_zero] = INFECTED
R0 = 0
# Run a standard SIR model for as long as patient_zero is infected and
# increment R0 value each time patient_zero directly infects another person
while state[patient_zero] == INFECTED:
# Active population with an infected neighbour
g2 = g.copy()
active = g2.new_vertex_property('bool', False)
active.a[np.nonzero(state.a)] = True
infect_vertex_property(g2, active, vals=[True])
g2.set_vertex_filter(active)
# numpy array of all vertices in g
population = g2.get_vertices()
np.random.shuffle(population)
# Iterate through all individuals in population
for person in population:
# Heal each infected person with probability gamma
if state[person] == INFECTED and coinflip(gamma):
state[person] = RECOVERED
if person == patient_zero:
return R0 # we can stop
if state[person] == SUSCEPTIBLE:
for neighbour in g.get_out_neighbors(person):
if state[neighbour] == INFECTED and coinflip(beta):
state[person] = INFECTED
if neighbour == patient_zero:
R0 += 1
break
return R0
def calculate_R0(g, beta, gamma, trials):
vxs = g.get_vertices()
results = np.zeros(trials)
for i in range(trials):
print(f'trial {i}')
# Pick one person uniformly at random
person = np.random.choice(vxs, 1, replace=False)
neighbours = g.get_out_neighbors(person)
if neighbours.size == 0:
results[i] = 0 # trivial edge case
else:
# Pick one neighbour at random to be patient_zero
patient_zero = np.random.choice(neighbours)
# Initialise model with correct state
results[i] = SIRmodel(g, patient_zero, beta, gamma)
"""
I don't think we want to do this? The result would only be an
approximation!
# # Remove all nodes that are not neighbours to patient zero
# new_vxs = state.copy(value_type='bool')
# infect_vertex_property(g, new_vxs, vals=[True])
# g.set_vertex_filter(new_vxs)
"""
return results
if __name__ == "__main__":
n = 10**5 # Population size
# g = price_network(n, directed=False, m=2)
# RANDOM GEOMETRIC NETWORK
points = np.random.random((n, 2)) * 5
g, pos = geometric_graph(points, 0.05, [(0, 4), (0, 4)])
beta = 0.01
gamma = 0.0476
trials = 100
R0_results = calculate_R0(g, beta, gamma, trials)
print(np.mean(R0_results))
print(np.std(R0_results))
plt.hist(R0_results, bins=np.linspace(0, 20, round(trials/10)))
plt.title('mean(R0) = '+str(np.mean(R0_results))+', std(R0) ='+str(round(np.std(R0_results)*100)/100))
plt.xlabel('R0')
plt.ylabel('Frequency')
plt.show()