diff --git a/network_attack_propagation.py b/network_attack_propagation.py index bcabc56..fa7fc5a 100644 --- a/network_attack_propagation.py +++ b/network_attack_propagation.py @@ -1,6 +1,19 @@ +import random +import networkx as nx +import matplotlib.pyplot as plt +from matplotlib import animation + + class Node: - def __init__(self): - self.is_infected = False + def __init__(self, is_infected=False): + self.id = random.randint(1, 2000000) + self.is_infected = is_infected + + def as_tuple(self): + return self.id, self.is_infected + + def __repr__(self): + return f'id: {self.id}, infected: {self.is_infected}' class Edge: @@ -9,10 +22,57 @@ class Edge: self.node_b = node_b self.weight = weight + def as_tuple(self): + return self.node_a, self.node_b, {'weight': self.weight} + class Graph: def __init__(self): self.edges = [] def add_edge(self, edge: Edge): - self.edges.append(edge) \ No newline at end of file + self.edges.append(edge) + + def add_edges(self, edges: [Edge]): + [self.edges.append(e) for e in edges] + + +def update(num, layout, g_repr, ax, our_graph): + """ + This function is called every 'step', so if you wish to update the graph, do it here + """ + ax.clear() + nx.draw_networkx(g_repr, ax=ax, pos=layout) + + +def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int): + 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 = nx.spring_layout(g_repr) + fig, ax = plt.subplots() + + anim = animation.FuncAnimation(fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph)) + anim.save(output_file_name) + plt.show() + + +def main(): + network = Graph() + nodes = [Node(True), Node(), Node(), Node(True), Node()] + + network.add_edges([ + Edge(nodes[1], nodes[0], 0.02), + Edge(nodes[1], nodes[2], 0.2), + Edge(nodes[2], nodes[0], 0.7), + Edge(nodes[3], nodes[2], 0.2), + Edge(nodes[3], nodes[1], 0.2), + Edge(nodes[4], nodes[3], 0.2) + ]) + + do_graph_animation('test.gif', network, 1) + + +if __name__ == "__main__": + main()