updated edge coloring, modified edge labels to more clear ones

visualize_kg.py

@@ -16,6 +16,23 @@ NODE_TYPE_COLORS = {
  'Protein': '#3aa6a4'
 }
 
+EDGE_LABEL_TRANSLATION = {
+    'Orthology': 'is ortholog to',
+    'Pathway': 'takes part in',
+    'kegg_path_prot': 'takes part in',
+    'protein_domain': 'has',
+    'PPI': 'interacts with',
+    'HPO': 'is associated with',
+    'kegg_dis_prot': 'is related to',
+    'Disease': 'is related to',
+    'Drug': 'targets',
+    'protein_ec': 'catalyzes',
+    'Chembl': 'targets',
+    ('protein_function', 'GO_term_F'): 'enables',
+    ('protein_function', 'GO_term_P'): 'is involved in',
+    ('protein_function', 'GO_term_C'): 'localizes to',    
+}
+
 GO_CATEGORY_MAPPING = {
     'Biological Process': 'GO_term_P',
     'Molecular Function': 'GO_term_F',
@@ -186,7 +203,12 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
     # Add edges and target nodes
     for edge_type, edges in visualized_edges.items():            
         source_type, relation_type, target_type = edge_type
-                
+
+        if relation_type == 'protein_function':
+            relation_type = EDGE_LABEL_TRANSLATION[(relation_type, target_type)]
+        else:
+            relation_type = EDGE_LABEL_TRANSLATION[relation_type]
+
         for edge_info in edges:
             edge, probability, is_ground_truth = edge_info 
             source, target = edge[0], edge[1]
@@ -222,10 +244,11 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
             if probability is not None:
                 if probability == 'no_pred':
                     edge_color = '#219ebc'
-                    edge_label += '(P=Not generated)'
+                    edge_label += ' (P=Not generated)'
                 else:
-                    edge_label += f"(P={probability:.2f})"
-                    edge_color = '#c1121f' if is_ground_truth else '#219ebc'
+                    edge_label += f" (P={probability:.2f})"
+                    edge_color = '#8338ec' if is_ground_truth else '#c1121f'
+                # if validated prediction purple, if non-validated prediction red, if no prediction (directly from database) blue
                 net.add_edge(source_str, target_str, 
                         label=edge_label, 
                         font={'size': 0},