Update analysis.py

analyse.py

@@ -43,7 +43,7 @@ def load_msgs(msg_lens: list[int], file: str | None = None):
 
         while True:
             readable_msg = next(iterator)["text"]
-            try :
+            try:
                 readable_msg[:length].encode("ascii")
                 break
             except Exception as e:
@@ -280,15 +280,20 @@ def process_results(results, save_dir):
         base = r["base"]
         delta = r["delta"]
         msg_rate = r["msg_rate"]
+        msg_len = r["msg_len"]
         perplexity = r["perplexity"]
 
-        if (base, delta) not in data["msg_rates"][msg_type]:
-            data["msg_rates"][msg_type][(base, delta)] = []
-        data["msg_rates"][msg_type][(base, delta)].append(msg_rate)
+        if (base, delta, msg_len) not in data["msg_rates"][msg_type]:
+            data["msg_rates"][msg_type][(base, delta, msg_len)] = []
+        data["msg_rates"][msg_type][(base, delta, msg_len)].append(
+            msg_rate
+        )
 
-        if (base, delta) not in data["perplexities"][msg_type]:
-            data["perplexities"][msg_type][(base, delta)] = []
-        data["perplexities"][msg_type][(base, delta)].append(perplexity)
+        if (base, delta, msg_len) not in data["perplexities"][msg_type]:
+            data["perplexities"][msg_type][(base, delta, msg_len)] = []
+        data["perplexities"][msg_type][(base, delta, msg_len)].append(
+            perplexity
+        )
 
     bases = {
         "perplexities": {
@@ -310,6 +315,16 @@ def process_results(results, save_dir):
             "readable": [],
         },
     }
+    msgs_lens = {
+        "perplexities": {
+            "random": [],
+            "readable": [],
+        },
+        "msg_rates": {
+            "random": [],
+            "readable": [],
+        },
+    }
     values = {
         "perplexities": {
             "random": [],
@@ -322,6 +337,7 @@ def process_results(results, save_dir):
     }
     base_set = set()
     delta_set = set()
+    msgs_lens_set = set()
     for metric in data:
         for msg_type in data[metric]:
             for k in data[metric][msg_type]:
@@ -331,9 +347,12 @@ def process_results(results, save_dir):
 
                 bases[metric][msg_type].append(k[0])
                 deltas[metric][msg_type].append(k[1])
+                msgs_lens[metric][msg_type].append(k[2])
                 values[metric][msg_type].append(s / cnt)
                 base_set.add(k[0])
                 delta_set.add(k[1])
+                msgs_lens_set.add(k[2])
+
     for metric in data:
         for msg_type in data[metric]:
             bases[metric][msg_type] = np.array(
@@ -350,7 +369,7 @@ def process_results(results, save_dir):
     for metric in data:
         for msg_type in data[metric]:
             fig = plt.figure(dpi=300)
-            s = lambda x: 3.0 + x * (3 if metric == "msg_rates" else 0.1)
+            s = lambda x: 3.0 + x * (30 if metric == "msg_rates" else 10)
             plt.scatter(
                 bases[metric][msg_type],
                 deltas[metric][msg_type],
@@ -365,84 +384,92 @@ def process_results(results, save_dir):
     os.makedirs(os.path.join(save_dir, "delta_effect"), exist_ok=True)
     for metric in data:
         for msg_type in data[metric]:
+            fig = plt.figure(dpi=300)
             for base_value in base_set:
-                mask = bases[metric][msg_type] == base_value
-                fig = plt.figure(dpi=300)
-                s = lambda x: x / (1.0 if metric == "msg_rates" else 10.0)
-                plt.plot(
-                    deltas[metric][msg_type][mask],
-                    values[metric][msg_type][mask],
-                )
-                plt.savefig(
-                    os.path.join(
-                        save_dir,
-                        f"delta_effect/{metric}_{msg_type}_base{base_value}.pdf",
-                    ),
-                    bbox_inches="tight",
-                )
-                plt.close(fig)
+                deltas_avg = np.array(list(delta_set))
+                values_avg = np.zeros_like(deltas_avg)
+                for i in range(len(deltas_avg)):
+                    mask = (deltas[metric][msg_type] == deltas_avg[i]) & (
+                        bases[metric][msg_type] == base_value
+                    )
+                    values_avg[i] = np.mean(values[metric][msg_type][mask])
+                plt.plot(deltas_avg, values_avg, label=f"Base {base_value}")
+
+            plt.legend()
+            plt.savefig(
+                os.path.join(
+                    save_dir,
+                    f"delta_effect/{metric}_{msg_type}.pdf",
+                ),
+                bbox_inches="tight",
+            )
+            plt.close(fig)
 
-    os.makedirs(os.path.join(save_dir, "base_effect"), exist_ok=True)
+    os.makedirs(os.path.join(save_dir, "msg_len_effect"), exist_ok=True)
     for metric in data:
         for msg_type in data[metric]:
-            for delta_value in delta_set:
-                mask = deltas[metric][msg_type] == delta_value
-                fig = plt.figure(dpi=300)
-                s = lambda x: x / (1.0 if metric == "msg_rates" else 10.0)
-                plt.plot(
-                    bases[metric][msg_type][mask],
-                    values[metric][msg_type][mask],
-                )
-                plt.savefig(
-                    os.path.join(
-                        save_dir,
-                        f"base_effect/{metric}_{msg_type}_delta{delta_value}.pdf",
-                    ),
-                    bbox_inches="tight",
-                )
-                plt.close(fig)
+            fig = plt.figure(dpi=300)
+            for base_value in base_set:
+                msgs_lens_avg = np.array(list(msgs_lens_set))
+                values_avg = np.zeros_like(msgs_lens_avg)
+                for i in range(len(msgs_lens_avg)):
+                    mask = (msgs_lens[metric][msg_type] == msgs_lens_avg[i]) & (
+                        bases[metric][msg_type] == base_value
+                    )
+                    values_avg[i] = np.mean(values[metric][msg_type][mask])
+
+                plt.plot(msgs_lens_avg, values_avg, label=f"Base {base_value}")
+
+            plt.legend()
+            plt.savefig(
+                os.path.join(
+                    save_dir,
+                    f"msg_len_effect/{metric}_{msg_type}.pdf",
+                ),
+                bbox_inches="tight",
+            )
+            plt.close(fig)
 
 
 def main(args):
-    prompts = load_prompts(
-        args.num_prompts,
-        args.prompt_size,
-        args.prompts_file if not args.overwrite else None,
-    )
+    if not args.results_load_file:
+        prompts = load_prompts(
+            args.num_prompts,
+            args.prompt_size,
+            args.prompts_file if not args.overwrite else None,
+        )
 
-    msgs_lens = []
-    for i in np.linspace(
-        args.msgs_lengths[0],
-        args.msgs_lengths[1],
-        int(args.msgs_lengths[2]),
-        dtype=np.int32,
-    ):
-        for _ in range(args.msgs_per_length):
-            msgs_lens.append(i)
-
-    msgs = load_msgs(
-        msgs_lens,
-        args.msgs_file if not args.overwrite else None,
-    )
+        msgs_lens = []
+        for i in np.linspace(
+            args.msgs_lengths[0],
+            args.msgs_lengths[1],
+            int(args.msgs_lengths[2]),
+            dtype=np.int32,
+        ):
+            for _ in range(args.msgs_per_length):
+                msgs_lens.append(i)
+
+        msgs = load_msgs(
+            msgs_lens,
+            args.msgs_file if not args.overwrite else None,
+        )
 
-    if args.msgs_file:
-        if not os.path.isfile(args.msgs_file) or args.overwrite:
-            os.makedirs(os.path.dirname(args.msgs_file), exist_ok=True)
-            with open(args.msgs_file, "w") as f:
-                json.dump(msgs, f)
-            print(f"Saved messages to {args.msgs_file}")
-    if args.prompts_file:
-        if not os.path.isfile(args.prompts_file) or args.overwrite:
-            os.makedirs(os.path.dirname(args.prompts_file), exist_ok=True)
-            with open(args.prompts_file, "w") as f:
-                json.dump(prompts, f)
-            print(f"Saved prompts to {args.prompts_file}")
-
-    if args.results_load_file:
+        if args.msgs_file:
+            if not os.path.isfile(args.msgs_file) or args.overwrite:
+                os.makedirs(os.path.dirname(args.msgs_file), exist_ok=True)
+                with open(args.msgs_file, "w") as f:
+                    json.dump(msgs, f)
+                print(f"Saved messages to {args.msgs_file}")
+        if args.prompts_file:
+            if not os.path.isfile(args.prompts_file) or args.overwrite:
+                os.makedirs(os.path.dirname(args.prompts_file), exist_ok=True)
+                with open(args.prompts_file, "w") as f:
+                    json.dump(prompts, f)
+                print(f"Saved prompts to {args.prompts_file}")
+        results = get_results(args, prompts, msgs)
+    else:
         with open(args.results_load_file, "r") as f:
             results = json.load(f)
-    else:
-        results = get_results(args, prompts, msgs)
 
     if args.results_save_file:
         os.makedirs(os.path.dirname(args.results_save_file), exist_ok=True)