Updated data generation code

plot_results.py

@@ -10,7 +10,7 @@ from train_and_eval import calculate_f1_score
 from sklearn.metrics import f1_score
 
 if len(sys.argv) < 2:
-    data_dir = "datasets/consensus_dataset"
+    data_dir = "datasets/kuramoto_dataset"
 else:
     data_dir = "datasets/" + sys.argv[1]
 
@@ -22,53 +22,59 @@ for folder in tqdm(os.listdir(data_dir)):
     
     folder_path = os.path.join(data_dir, folder)
 
-    # Load model config from summary json
-    with open(os.path.join(folder_path, "results/NoiseType.NONE", "summary_results.json"), "r") as f:
-        summary_results = json.load(f)
+    for noise_level in os.listdir(os.path.join(folder_path, "results/NoiseType.NONE")):
 
-    
-    for i, graph in enumerate(os.listdir(folder_path)):
 
-        # train_summary_results
-        summ_results = summary_results[i-1]
+        # Load model config from summary json
+        with open(os.path.join(folder_path, "results/NoiseType.NONE", noise_level, "summary_results.json"), "r") as f:
+            summary_results = json.load(f)
 
-        if graph == "results": # ignore the result folder
+        
+        for i, graph in enumerate(os.listdir(folder_path)):
+
+            # train_summary_results
+            summ_results = summary_results[i-1]
+
+            if graph == "results": # ignore the result folder
+                continue
+
+            graph_path = os.path.join(folder_path, graph)
+
+            # Load run data
+            with open(os.path.join(folder_path, graph), "r") as f:
+                run_data = json.load(f)
+
+            true_graph = np.array(run_data["adjacency_matrix"])
+            
+            learned_graph = np.array(summ_results["raw_attention"])
+
+            predicted_graph = (learned_graph > THRESHOLD).astype(int)
+
+            true_flat = true_graph.flatten()
+            pred_flat = predicted_graph.flatten()
+            
+            calc_f1_score = f1_score(true_flat, pred_flat)
+
+
+            datapoints[num_agents] = datapoints.get(num_agents, [])
+            datapoints[num_agents].append(calc_f1_score)
+
+
+    for key in datapoints.keys():
+        try:
+            datapoints[key] = sum(datapoints[key])/len(datapoints[key])
+        except:
             continue
 
-        graph_path = os.path.join(folder_path, graph)
 
-        # Load run data
-        with open(os.path.join(folder_path, graph), "r") as f:
-            run_data = json.load(f)
+    x = []
+    y = []
 
-        true_graph = np.array(run_data["adjacency_matrix"])
-        
-        learned_graph = np.array(summ_results["raw_attention"])
+    for item in datapoints.items():
+        x.append(item[0])
+        y.append(item[1])
 
-        predicted_graph = (learned_graph > THRESHOLD).astype(int)
-
-        true_flat = true_graph.flatten()
-        pred_flat = predicted_graph.flatten()
-        
-        calc_f1_score = f1_score(true_flat, pred_flat)
-
-
-        datapoints[num_agents] = datapoints.get(num_agents, [])
-        datapoints[num_agents].append(calc_f1_score)
-
-
-for key in datapoints.keys():
-    datapoints[key] = sum(datapoints[key])/len(datapoints[key])
-
-
-x = []
-y = []
-
-for item in datapoints.items():
-    x.append(item[0])
-    y.append(item[1])
-
-plt.plot(x, y)
-plt.show()         
+    plt.plot(x, y)
+    plt.show()