refactor analysis and is_around_message_interval

mixnet/v2/sim/analysis.py

@@ -0,0 +1,52 @@
+import pandas as pd
+import seaborn
+from matplotlib import pyplot as plt
+
+from simulation import Simulation
+
+
+class Analysis:
+    def __init__(self, sim: Simulation):
+        self.sim = sim
+
+    def run(self):
+        self.message_size_distribution()
+        self.messages_emitted_around_interval()
+        self.mixed_messages_per_node_over_time()
+
+    def message_size_distribution(self):
+        df = pd.DataFrame(self.sim.p2p.message_sizes, columns=["message_size"])
+        print(df.describe())
+
+    def messages_emitted_around_interval(self):
+        df = pd.DataFrame(
+            [(node.id, cnt, node.id < len(self.sim.config.real_message_prob_weights))
+             for node, cnt in self.sim.p2p.senders_around_interval.items()],
+            columns=["node_id", "msg_count", "expected"]
+        )
+        plt.figure(figsize=(10, 6))
+        seaborn.barplot(data=df, x="node_id", y="msg_count", hue="expected", palette={True: "red", False: "blue"})
+        plt.title("Messages emitted around the promised interval")
+        plt.xlabel("Sender Node ID")
+        plt.ylabel("Msg Count")
+        plt.legend(title="expected")
+        plt.show()
+
+    def mixed_messages_per_node_over_time(self):
+        dataframes = []
+        for mixed_msgs_per_node in self.sim.p2p.mixed_msgs_per_window:
+            df = pd.DataFrame([(node.id, cnt) for node, cnt in mixed_msgs_per_node.items()],
+                              columns=["node_id", "msg_count"])
+            dataframes.append(df)
+        observation_times = range(len(dataframes))
+        df = pd.concat([df.assign(Time=time) for df, time in zip(dataframes, observation_times)], ignore_index=True)
+        df = df.pivot(index="Time", columns="node_id", values="msg_count")
+        plt.figure(figsize=(12, 6))
+        for column in df.columns:
+            plt.plot(df.index, df[column], marker="o", label=column)
+        plt.title("Mixed messages in each mix over time")
+        plt.xlabel("Time")
+        plt.ylabel("Msg Count")
+        plt.legend(title="Node ID")
+        plt.grid(True)
+        plt.show()

mixnet/v2/sim/main.py

@@ -1,10 +1,7 @@
 import argparse
 
-import matplotlib.pyplot as plt
-import pandas as pd
-import seaborn
-
 from config import Config
+from mixnet.v2.sim.analysis import Analysis
 from simulation import Simulation
 
 if __name__ == "__main__":
@@ -16,40 +13,6 @@ if __name__ == "__main__":
     sim = Simulation(config)
     sim.run()
 
-    # Stat the distribution of message sizes
-    df = pd.DataFrame(sim.p2p.message_sizes, columns=["message_size"])
-    print(df.describe())
-
-    # Visualize the nodes emitted messages around the promised interval
-    df = pd.DataFrame(
-        [(node.id, cnt, node.id < len(config.real_message_prob_weights))
-         for node, cnt in sim.p2p.senders_around_interval.items()],
-        columns=["NodeID", "MsgCount", "Expected"]
-    )
-    plt.figure(figsize=(10, 6))
-    seaborn.barplot(data=df, x="NodeID", y="MsgCount", hue="Expected", palette={True: "red", False: "blue"})
-    plt.title("Messages emitted around the promised interval")
-    plt.xlabel("Sender Node ID")
-    plt.ylabel("Msg Count")
-    plt.legend(title="Expected")
-    plt.show()
-
-    # Analyze the number of mixed messages per node per observation window
-    dataframes = []
-    for mixed_msgs_per_node in sim.p2p.mixed_msgs_per_window:
-        df = pd.DataFrame([(node.id, cnt) for node, cnt in mixed_msgs_per_node.items()], columns=["NodeID", "MsgCount"])
-        dataframes.append(df)
-    observation_times = range(len(dataframes))
-    df = pd.concat([df.assign(Time=time) for df, time in zip(dataframes, observation_times)], ignore_index=True)
-    df = df.pivot(index="Time", columns="NodeID", values="MsgCount")
-    plt.figure(figsize=(12, 6))
-    for column in df.columns:
-        plt.plot(df.index, df[column], marker='o', label=column)
-    plt.title('Mixed messages in each mix over time')
-    plt.xlabel('Time')
-    plt.ylabel('Msg Count')
-    plt.legend(title='Node ID')
-    plt.grid(True)
-    plt.show()
+    Analysis(sim).run()
 
     print("Simulation complete!")

mixnet/v2/sim/p2p.py

@@ -3,6 +3,7 @@ import random
 from collections import defaultdict
 
 import simpy
+from simpy.core import SimTime
 
 from config import Config
 from sphinx import SphinxPacket
@@ -30,11 +31,11 @@ class P2p:
     # but we accept SphinxPacket as well because we don't implement Sphinx deserialization.
     def broadcast(self, sender, msg: SphinxPacket | bytes):
         self.log("Broadcasting a msg: %d bytes" % len(msg))
+
+        # Adversary
         self.message_sizes.append(len(msg))
         self.mixed_msgs_per_window[-1][sender] -= 1
-
-        now_frac, now_int = math.modf(self.env.now)
-        if now_int % self.config.message_interval == 0 and now_frac <= self.config.max_message_prep_time:
+        if self.is_around_message_interval(self.env.now):
             self.senders_around_interval[sender] += 1
 
         # Yield 0 to ensure that the broadcast is done in the same time step.
@@ -52,6 +53,11 @@ class P2p:
         self.mixed_msgs_per_window[-1][node] += 1
         self.env.process(node.receive_message(msg))
 
+    # TODO: Move to a separate class `Adversary`.
+    def is_around_message_interval(self, time: SimTime):
+        now_frac, now_int = math.modf(time)
+        return now_int % self.config.message_interval == 0 and now_frac <= self.config.max_message_prep_time
+
     # TODO: Move to a separate class `Adversary`.
     def update_observation_window(self):
         while True: