Update calculate_memory_usage.py

calculate_memory_usage.py

@@ -1,8 +1,9 @@
 import os
 import glob
 import argparse
+import json
 from code_efficiency_calculator import run_model_task
-
+data_to_save = {}
 
 def calculate_memory_usage(dat_file_path):
     with open(dat_file_path, 'r') as file:
@@ -51,6 +52,7 @@ def report_max_memory_usage(dat_file_path):
         return max_memory_usage
 
 def report_results(task, model, file):
+    global data_to_save 
     run_model_task(task, model, file)
     dat_directory = f"./results/{task}_{model}"
     canonical_solution_directory = f"./results/{task}_canonical_solution"
@@ -190,6 +192,17 @@ def report_results(task, model, file):
                 max_ntmu_lists[ntmu] = (model, idx)
             max_tmu = max(max_tmu,completion_memory_usage[idx]/canonical_solution_memory_usage[idx])
             total_codes+=1
+            data_to_save[idx] = {
+        "completion_memory_usage": completion_memory_usage[idx],
+        "execution_time": execution_time[idx],
+        "max_memory_usage": max_memory_usage[idx],
+        "canonical_solution_max_memory_usage": canonical_solution_max_memory_usage[idx],
+        "canonical_solution_memory_usage": canonical_solution_memory_usage[idx],
+        "canonical_solution_execution_time": canonical_solution_execution_time[idx],
+        "net": net,
+        "nmu": nmu,
+        "ntmu": ntmu
+    }
 
         if len(normalized_execution_time_list)==0:
             print(model)
@@ -201,12 +214,23 @@ def report_results(task, model, file):
         total_memory_usage = total_memory_usage/len(normalized_execution_time_list)
         total_max_memory_usage = total_max_memory_usage/len(normalized_execution_time_list)
 
-        pass1 = len(completion_memory_usage)/1000*100
+        pass1 = len(completion_memory_usage)/164*100
 
         total_500_net = total_500_net/len(normalized_execution_time_list)*100
         total_500_nmu = total_500_nmu/len(normalized_execution_time_list)*100
         total_500_tmu = total_500_tmu/len(normalized_execution_time_list)*100
-
+        passing_problem_idxs = []
+        for idx in completion_memory_usage.keys():
+            if idx not in canonical_solution_memory_usage.keys():
+                continue
+            else:
+                passing_problem_idxs.append(idx)
+
+        with open("passing_problem_idxs.txt", "w") as file:
+            for problem_idx in passing_problem_idxs:
+                file.write(f"{problem_idx}\n")
+        with open("output_data.json", "w") as json_file:
+            json.dump(data_to_save, json_file, indent=4)
         return f"{model}&{total_execution_time:.2f}&{normalized_execution_time:.2f}&{max_net:.2f}&{total_500_net:.1f}&{total_max_memory_usage:.2f}&{normalized_max_memory_usage:.2f}&{max_nmu:.2f}&{total_500_nmu:.1f}&{total_memory_usage:.2f}&{normalized_memory_usage:.2f}&{max_tmu:.2f}&{total_500_tmu:.1f}&{pass1:.1f}\\\\"