Update src/bin/PROBE.py

src/bin/PROBE.py

@@ -17,7 +17,9 @@ def load_representation(multi_col_representation_vector_file_path):
 
 def run_probe(benchmarks, representation_name, representation_file_human, representation_file_affinity, similarity_tasks=["Sparse","200","500"], function_prediction_aspect="All_Aspects", function_prediction_dataset="All_Data_Sets", family_prediction_dataset=["nc","uc50","uc30","mm15"], detailed_output=False):
     print("\n\nPROBE (Protein RepresentatiOn Benchmark) run is started...\n\n")
-    
+    result = {}
+    result['Method'] = representation_name
+
     if any(item in ['similarity', 'function', 'family', 'all'] for item in benchmarks):
         print("\nRepresentation vectors are loading...\n")
         human_representation_dataframe = load_representation(representation_file_human)
@@ -30,9 +32,7 @@ def run_probe(benchmarks, representation_name, representation_file_human, repres
         ssi.similarity_tasks = similarity_tasks
         ssi.detailed_output = detailed_output
         similarity_result = ssi.calculate_all_correlations()
-        print("Similarity Result:")
-        print(similarity_result)
-        
+        result['similarity'] = similarity_result
 
     if "function" in benchmarks:
         print("\n\nOntology-based protein function prediction benchmark is running...\n")
@@ -42,26 +42,26 @@ def run_probe(benchmarks, representation_name, representation_file_human, repres
         fp.representation_name = representation_name
         fp.detailed_output = detailed_output
         function_results = fp.pred_output()
-        print("Function results:")
-        print(function_results)
+        result['function'] = function_results
 
     if "family" in benchmarks:
         print("\n\nDrug target protein family classification benchmark is running...\n")
         tfc.representation_path = representation_file_human
         tfc.representation_name = representation_name
         tfc.detailed_output = detailed_output
+        result['family'] = {}
         for dataset in family_prediction_dataset:
             family_result = tfc.score_protein_rep(dataset)
-            print(f"Family results for {dataset}:")
-            print(family_result)
+            result['family']['dataset'] = family_result
 
     if "affinity" in benchmarks:
         print("\n\nProtein-protein binding affinity estimation benchmark is running...\n")
         bae.skempi_vectors_path = representation_file_affinity
         bae.representation_name = representation_name
         affinity_result = bae.predict_affinities_and_report_results()
-        print("Affinity Results:")
-        print(affinity_result)
+        result['affinity'] = affinity_result
+
+    print(result)
 
     print("\n\nPROBE (Protein RepresentatiOn Benchmark) run is finished...\n")
     return 0