added coloring to predicted/ground truth go term edges

visualize_kg.py

@@ -63,7 +63,6 @@ def _gather_protein_edges(data, protein_id):
     return protein_edges
 
 def _filter_edges(protein_id, protein_edges, prediction_df, limit=10):
-
     filtered_edges = {}
     
     prediction_categories = prediction_df['GO_category'].unique()
@@ -75,32 +74,35 @@ def _filter_edges(protein_id, protein_edges, prediction_df, limit=10):
         if edges is None or len(edges) == 0:
             continue
             
-        if edge_type[2] in prediction_categories:
-            category_mask = (prediction_df['GO_category'] == go_category_reverse_mapping[edge_type[2]]) & (prediction_df['UniProt_ID'] == protein_id)
-            category_predictions = prediction_df[category_mask]
-
-            if len(category_predictions) > 0:
-                category_predictions = category_predictions.sort_values(by='Probability', ascending=False)
-                
-                # Convert set to list for easier filtering
-                edges_list = list(edges)
-
-                # Filter valid edges and store with probabilities
-                valid_edges = []
-                for _, row in category_predictions.iterrows():
-                    term = row['GO_ID']
-                    prob = row['Probability']
-                    matching_edges = [(edge, prob) for edge in edges_list if edge[1] == term]
-                    valid_edges.extend(matching_edges)
-                    if len(valid_edges) >= limit:
-                        break
-                filtered_edges[edge_type] = valid_edges  # Remove set conversion to preserve probabilities
+        if edge_type[2].startswith('GO_term'):  # Check if it's any GO term edge
+            if edge_type[2] in prediction_categories:
+                # Handle edges for GO terms that are in prediction_df
+                category_mask = (prediction_df['GO_category'] == go_category_reverse_mapping[edge_type[2]]) & (prediction_df['UniProt_ID'] == protein_id)
+                category_predictions = prediction_df[category_mask]
+
+                if len(category_predictions) > 0:
+                    category_predictions = category_predictions.sort_values(by='Probability', ascending=False)
+                    edges_set = set(edges)  # Convert to set for O(1) lookup
+                    
+                    valid_edges = []
+                    for _, row in category_predictions.iterrows():
+                        term = row['GO_ID']
+                        prob = row['Probability']
+                        edge = (protein_id, term)
+                        is_ground_truth = edge in edges_set
+                        valid_edges.append((edge, prob, is_ground_truth))
+                        if len(valid_edges) >= limit:
+                            break
+                    filtered_edges[edge_type] = valid_edges
+                else:
+                    # If no predictions but it's a GO category in prediction_df
+                    filtered_edges[edge_type] = [(edge, 'no_pred', True) for edge in list(edges)[:limit]]
             else:
-                # If no predictions, include all edges up to limit without probabilities
-                filtered_edges[edge_type] = [(edge, None) for edge in list(edges)[:limit]]
+                # For GO terms not in prediction_df, mark them as ground truth with blue color
+                filtered_edges[edge_type] = [(edge, 'no_pred', True) for edge in list(edges)[:limit]]
         else:
-            # For non-GO edges, include all edges up to limit without probabilities
-            filtered_edges[edge_type] = [(edge, None) for edge in list(edges)[:limit]]
+            # For non-GO edges, include all edges up to limit
+            filtered_edges[edge_type] = [(edge, None, True) for edge in list(edges)[:limit]]
 
     return filtered_edges
 
@@ -186,7 +188,7 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
         source_type, relation_type, target_type = edge_type
                 
         for edge_info in edges:
-            edge, probability = edge_info
+            edge, probability, is_ground_truth = edge_info 
             source, target = edge[0], edge[1]
             source_str = str(source)
             target_str = str(target)
@@ -218,10 +220,16 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
             # Add edge with relationship type and probability as label
             edge_label = f"{relation_type}"
             if probability is not None:
-                edge_label += f"(P={probability:.2f})"
+                if probability == 'no_pred':
+                    edge_color = '#219ebc'
+                    edge_label += '(P=Not generated)'
+                else:
+                    edge_label += f"(P={probability:.2f})"
+                    edge_color = '#c1121f' if is_ground_truth else '#219ebc'
                 net.add_edge(source_str, target_str, 
                         label=edge_label, 
-                        color='#666666',
+                        font={'size': 0},
+                        color=edge_color,
                         title=edge_label,
                         length=200,
                         smooth={'type': 'curvedCW', 'roundness': 0.1})
@@ -229,7 +237,7 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
                 net.add_edge(source_str, target_str, 
                         label=edge_label,
                         font={'size': 0},
-                        color='#666666',
+                        color='#666666',  # Keep default gray for non-GO edges
                         title=edge_label,
                         length=200,
                         smooth={'type': 'curvedCW', 'roundness': 0.1})