Add path traversal with infection

network_attack_propagation.py

@@ -25,10 +25,14 @@ class Edge:
     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):
+    def __init__(self, use_weights=False):
         self.edges = []
+        self.use_weights = use_weights
 
     def add_edge(self, edge: Edge):
         self.edges.append(edge)
@@ -43,6 +47,31 @@ class Graph:
             nodes.add(edge.node_b)
         return nodes
 
+    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 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
+
 
 def update(num, layout, g_repr, ax, our_graph: Graph):
     """
@@ -50,9 +79,6 @@ def update(num, layout, g_repr, ax, our_graph: Graph):
     """
     ax.clear()
 
-    for n in our_graph.get_nodes():
-        n.is_infected = bool(random.getrandbits(1))
-
     colors = ['red' if n.is_infected else 'blue' for n in g_repr]
     sizes = [50 if n.is_infected else 1 for n in g_repr]
     nx.draw(
@@ -67,6 +93,8 @@ def update(num, layout, g_repr, ax, our_graph: Graph):
         linewidths=40,
     )
 
+    our_graph.infect_step()
+
 
 def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout):
     g_repr = nx.Graph()
@@ -80,7 +108,7 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int,
     fig.set_figheight(15)
 
     anim = animation.FuncAnimation(
-        fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph)
+        fig, update, frames=frame_count, interval=500, fargs=(layout, g_repr, ax, in_graph)
     )
     anim.save(output_file_name)
 
@@ -88,9 +116,10 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int,
     plt.show()
 
 
-def bus_network(n=30) -> Graph:
+def bus_network(n=30, infected_idx=0) -> Graph:
     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)
@@ -100,6 +129,7 @@ def bus_network(n=30) -> Graph:
 def ring_network(n=30) -> Graph:
     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)
@@ -126,10 +156,10 @@ def main():
     do_graph_animation('test.gif', network, 5, nx.spring_layout)
 
     bus = bus_network()
-    do_graph_animation('bus.gif', bus, 5, nx.spiral_layout)
+    do_graph_animation('bus.gif', bus, 20, nx.spiral_layout)
 
     ring = ring_network()
-    do_graph_animation('ring.gif', ring, 5, nx.circular_layout)
+    do_graph_animation('ring.gif', ring, 20, nx.circular_layout)
 
 
 if __name__ == "__main__":