257 lines
7.7 KiB
Python
257 lines
7.7 KiB
Python
import random
|
|
from statistics import mean
|
|
|
|
import matplotlib.pyplot as plt
|
|
import networkx as nx
|
|
from matplotlib import animation
|
|
|
|
|
|
class Node:
|
|
def __init__(self, is_infected=False):
|
|
self.id = random.randint(1, 2000000)
|
|
self.is_infected = is_infected
|
|
self.under_attack = False
|
|
|
|
def as_tuple(self):
|
|
return self.id, self.is_infected
|
|
|
|
def __repr__(self):
|
|
return f'id: {self.id}, infected: {self.is_infected}'
|
|
|
|
|
|
class Edge:
|
|
def __init__(self, node_a: Node, node_b: Node, weight: float):
|
|
self.node_a = node_a
|
|
self.node_b = node_b
|
|
self.weight = weight
|
|
|
|
def as_tuple(self):
|
|
return self.node_a, self.node_b, {'weight': self.weight}
|
|
|
|
def has_node(self, node: Node) -> bool:
|
|
return self.node_a is node or self.node_b is node
|
|
|
|
|
|
class Graph:
|
|
def __init__(self, use_weights=False):
|
|
self.edges = []
|
|
self.use_weights = use_weights
|
|
self.rounds_survived = 0
|
|
|
|
def add_edge(self, edge: Edge):
|
|
self.edges.append(edge)
|
|
|
|
def add_edges(self, edges: [Edge]):
|
|
[self.edges.append(e) for e in edges]
|
|
|
|
def get_nodes(self) -> [Node]:
|
|
nodes = set()
|
|
for edge in self.edges:
|
|
nodes.add(edge.node_a)
|
|
nodes.add(edge.node_b)
|
|
return nodes
|
|
|
|
def clear_attacked(self) -> None:
|
|
for edge in self.edges:
|
|
edge.node_a.under_attack = False
|
|
edge.node_b.under_attack = False
|
|
|
|
def get_adjacent_nodes(self, node: Node) -> [(Node, int)]:
|
|
"""
|
|
:param node: Node to search for
|
|
:return: An array of tuples (node, weight)
|
|
"""
|
|
edges_with_node = filter(lambda ed: ed.has_node(node), self.edges)
|
|
nodes = set()
|
|
for e in edges_with_node:
|
|
if e.node_a is node:
|
|
nodes.add((e.node_b, e.weight))
|
|
else:
|
|
nodes.add((e.node_a, e.weight))
|
|
|
|
return nodes
|
|
|
|
def is_alive(self):
|
|
nodes_alive = list(filter(lambda x: not x.is_infected, self.get_nodes()))
|
|
return len(nodes_alive) > 0
|
|
|
|
def update_survived(self):
|
|
if not self.is_alive():
|
|
return
|
|
|
|
self.rounds_survived += 1
|
|
|
|
def infect_step(self):
|
|
infected_nodes = list(filter(lambda n: n.is_infected, self.get_nodes()))
|
|
for node in infected_nodes:
|
|
adjacent_nodes = self.get_adjacent_nodes(node)
|
|
if self.use_weights:
|
|
to_be_infected = random.choices([n[0] for n in adjacent_nodes], weights=[n[1] for n in adjacent_nodes])[0]
|
|
else:
|
|
to_be_infected = random.choice([n[0] for n in adjacent_nodes])
|
|
to_be_infected.is_infected = True
|
|
to_be_infected.under_attack = True
|
|
self.update_survived()
|
|
|
|
|
|
def update(num, layout, g_repr, ax, our_graph: Graph):
|
|
"""
|
|
This function is called every 'step', so if you wish to update the graph, do it here
|
|
"""
|
|
if not our_graph.is_alive():
|
|
return
|
|
|
|
if num != 0:
|
|
our_graph.infect_step()
|
|
ax.clear()
|
|
|
|
ax.set_title(f'Step: {num}', loc='right', fontsize=30)
|
|
|
|
colors = ['red' if n.is_infected else 'blue' for n in g_repr]
|
|
edgecolors = ['black' if n.under_attack else 'none' for n in g_repr]
|
|
linewidths = [3 if c == 'black' else 0 for c in edgecolors]
|
|
sizes = [300 if n.is_infected else 150 for n in g_repr]
|
|
nx.draw(
|
|
g_repr,
|
|
ax=ax,
|
|
pos=layout,
|
|
node_color=colors,
|
|
linewidths=linewidths,
|
|
edgecolors=edgecolors,
|
|
with_labels=False,
|
|
node_size=sizes,
|
|
alpha=0.7,
|
|
)
|
|
|
|
our_graph.clear_attacked()
|
|
|
|
|
|
def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout):
|
|
g_repr = nx.Graph()
|
|
# Convert our graph class into tuples understood by networkx
|
|
g_repr.add_edges_from([e.as_tuple() for e in in_graph.edges])
|
|
|
|
layout = layout(g_repr)
|
|
fig, ax = plt.subplots()
|
|
|
|
fig.set_figwidth(8)
|
|
fig.set_figheight(8)
|
|
|
|
anim = animation.FuncAnimation(
|
|
fig, update, frames=frame_count, interval=500, fargs=(layout, g_repr, ax, in_graph)
|
|
)
|
|
anim.save(output_file_name)
|
|
|
|
plt.style.use('seaborn')
|
|
plt.show()
|
|
|
|
|
|
def degree_avg(edges, digits=2):
|
|
degrees = {}
|
|
for e in edges:
|
|
degrees[e.node_a] = degrees.get(e.node_a, 0) + 1
|
|
degrees[e.node_b] = degrees.get(e.node_b, 0) + 1
|
|
return round(mean(degrees.values()), digits)
|
|
|
|
|
|
def bus_network(n=30, infected_idx=0) -> tuple[Graph, float]:
|
|
network = Graph()
|
|
nodes = [Node() for _ in range(n)]
|
|
nodes[infected_idx].is_infected = True
|
|
edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)]
|
|
|
|
network.add_edges(edges)
|
|
return network, degree_avg(edges), n
|
|
|
|
|
|
def star_network(cluster_count=5, starsize=6) -> tuple[Graph, float, int]:
|
|
node_count = cluster_count + cluster_count * starsize + 1
|
|
nodes = [Node() for _ in range(node_count)]
|
|
nodes[starsize-1].is_infected = True
|
|
|
|
edges = []
|
|
for x in range(cluster_count):
|
|
center_node = x * starsize + x
|
|
edges += [Edge(nodes[center_node], nodes[i], 1.0) for i in range(center_node + 1, center_node + starsize + 1)]
|
|
edges.append(Edge(nodes[-1], nodes[center_node], 1.0))
|
|
|
|
network = Graph()
|
|
network.add_edges(edges)
|
|
|
|
return network, degree_avg(edges), node_count
|
|
|
|
|
|
def ring_network(n=30) -> tuple[Graph, float]:
|
|
network = Graph()
|
|
nodes = [Node() for _ in range(n)]
|
|
nodes[0].is_infected = True
|
|
edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)]
|
|
end_edge = Edge(nodes[n - 1], nodes[0], 1.0)
|
|
edges.append(end_edge)
|
|
|
|
network.add_edges(edges)
|
|
return network, degree_avg(edges), n
|
|
|
|
|
|
def summary(average_degrees: list[float], propagation_speeds: list[float]) -> None:
|
|
fig, ax = plt.subplots()
|
|
ax.plot(average_degrees, propagation_speeds)
|
|
ax.set(xlabel='Average degree', ylabel='Propagation speed', title='Summary')
|
|
fig.savefig("summary.png")
|
|
plt.show()
|
|
|
|
def bus_experiment():
|
|
degrees = []
|
|
speeds = []
|
|
sizes = [0, 8, 15]
|
|
for i in sizes:
|
|
bus, bus_avg_degree, node_count = bus_network(20 + i)
|
|
do_graph_animation(f'bus{i}.gif', bus, 90, nx.spiral_layout)
|
|
speeds.append(bus.rounds_survived / node_count)
|
|
degrees.append(bus_avg_degree)
|
|
print(f"\n{node_count} NODE BUS")
|
|
print(f"average degree = {bus_avg_degree}")
|
|
print(f"propagation speed = {round(speeds[-1], 2)}")
|
|
print(f"bus_{i} rounds survived = {bus.rounds_survived + 1}")
|
|
|
|
summary(degrees, speeds)
|
|
|
|
def ring_experiment():
|
|
degrees = []
|
|
speeds = []
|
|
sizes = [0, 8, 15]
|
|
for i in sizes:
|
|
ring, ring_avg_degree, node_count = ring_network(20 + i)
|
|
do_graph_animation(f'ring{i}.gif', ring, 90, nx.spiral_layout)
|
|
speeds.append(ring.rounds_survived / node_count)
|
|
degrees.append(ring_avg_degree)
|
|
print(f"\n{node_count} NODE RING")
|
|
print(f"average degree = {ring_avg_degree}")
|
|
print(f"propagation speed = {round(speeds[-1], 2)}")
|
|
print(f"ring_{i} rounds survived = {ring.rounds_survived + 1}")
|
|
|
|
def star_experiment():
|
|
degrees = []
|
|
speeds = []
|
|
sizes = range(0, 8, 2)
|
|
for i in sizes:
|
|
star, star_avg_degree, node_count = star_network(cluster_count=3 + i, starsize=3 + i)
|
|
do_graph_animation(f'star{i}.gif', star, 120, nx.kamada_kawai_layout)
|
|
speeds.append(star.rounds_survived / node_count)
|
|
degrees.append(star_avg_degree)
|
|
print(f"\n{node_count} NODE STAR")
|
|
print(f"average degree = {star_avg_degree}")
|
|
print(f"propagation speed = {round(speeds[-1], 2)}")
|
|
print(f"star_{i} rounds survived = {star.rounds_survived + 1}")
|
|
|
|
summary(degrees, speeds)
|
|
|
|
def main():
|
|
bus_experiment()
|
|
ring_experiment()
|
|
star_experiment()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|