-
Notifications
You must be signed in to change notification settings - Fork 0
/
Node.py
105 lines (81 loc) · 2.98 KB
/
Node.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
import random
class Node():
def __init__(self, id, centerValues=0):
self.id = id
self.state = 0 # 1 is infected
self.connections = []
self.connectNumbers = []
self.name = ""
self.visibleToPlayer = False
self.isPatientZero = False
self.timeInfected = None
self.selectedByAI = "False"
self.X = 0
self.Y = 0
self.age = random.randint(1, 80)
def makeConnections(self, nodes, min_connections, max_connections):
numberConnections = min_connections
for i in range(numberConnections):
while True:
target = random.randint(0, len(nodes) - 1)
# If not already connected or target is not self
if nodes[target] not in self.connections and target != self.id:
break
self.connections.append(nodes[target])
self.connectNumbers.append(nodes[target].id)
nodes[target].connections.append(self)
nodes[target].connectNumbers.append(self.id)
def calculateLikelihood(self, enableAlwaysInfection=False):
if 0 <= self.age <= 17:
likelihood = 45.5
elif 18 <= self.age <= 49:
likelihood = 64.9
elif 50 <= self.age <= 64:
likelihood = 53.6
elif self.age >= 65:
likelihood = 42.36
else:
likelihood = 54.6
return 100 if enableAlwaysInfection else likelihood
def infectNeighbors(self, time):
# Infect all uninfected neighbors
for neighbor in self.connections:
if neighbor.state == 0:
# Calculate infection likelihood
likelihood = neighbor.calculateLikelihood()
if random.randint(0, 1000)/10 <= likelihood:
neighbor.state = 1
neighbor.timeInfected = time
def createNodeNetwork(numberOfNodes, nodes, min_connections, max_connections):
# Initialize center values:
centerValues = []
# for x in [i for i in range(30, 180, 30)]:
# for y in [i for i in range(30, 180, 30)]:
# centerValues.append
# Create nodes
for i, newNodeId in enumerate(range(numberOfNodes)):
nodes[i] = (Node(newNodeId, centerValues))
centerValues.append((nodes[newNodeId].X, nodes[newNodeId].Y))
# Make connections
for node in nodes.values():
node.makeConnections(nodes, min_connections, max_connections)
def runInfectionSimulation(numDays, nodes, selected_p_zero=None):
# Select patient zero
p_zero = selected_p_zero if selected_p_zero else random.choice(nodes)
p_zero.state = 1
p_zero.isPatientZero = True
p_zero.timeInfected = 0
# print(f"Patient-Zero is Node #{p_zero.id}, X:{p_zero.X}, Y:{p_zero.Y}\n--------------------------")
# Just to make it easier to read on right
# print("Infected Nodes After Day 0 --------------------")
# print(p_zero.id)
for day in range(1, numDays + 1):
# print(f"\nInfected Nodes --------------------")
infectednodes = []
for node in nodes.values():
if node.state == 1:
infectednodes.append(node)
for node in infectednodes:
if node.state == 1:
node.infectNeighbors(day)
return p_zero.id