File size: 4,263 Bytes
4a8b134 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """
Stage 5: Result Processing & Filtering
Sorts results and labels them as 'Known Treatment' or 'Potential Discovery'.
"""
import logging
from typing import List, Dict
logger = logging.getLogger(__name__)
class ResultProcessingPipeline:
@staticmethod
def process_final_results(
results: List[Dict],
known_drugs: List[str],
min_score: float = 0.0
) -> List[Dict]:
"""
Sorts results and labels them as 'Known Treatment' or 'Potential Discovery'.
Args:
results: List of prediction results
known_drugs: List of known drugs for the disease
min_score: Minimum binding affinity score to include (default: 0.0)
Returns:
Processed and sorted results
"""
logger.info(f"Processing {len(results)} results...")
# Filter by minimum score
filtered_results = [r for r in results if r['score'] >= min_score]
logger.info(f"Filtered to {len(filtered_results)} results with score >= {min_score}")
# Sort by score (descending)
sorted_results = sorted(filtered_results, key=lambda x: x['score'], reverse=True)
final_output = []
for res in sorted_results:
# Check if drug is already known for this disease
is_known = any(
known.lower() in res['drug_name'].lower()
for known in known_drugs
)
res['status'] = "✅ Known Treatment" if is_known else "🆕 Potential Discovery"
final_output.append(res)
logger.info(f"✅ Processing complete. Found {len(final_output)} candidates")
return final_output
@staticmethod
def get_top_results(results: List[Dict], top_n: int = 15) -> List[Dict]:
"""
Returns the top N results.
Args:
results: List of processed results
top_n: Number of top results to return
Returns:
Top N results
"""
return results[:top_n]
@staticmethod
def get_potential_discoveries(results: List[Dict]) -> List[Dict]:
"""
Returns only potential discoveries (non-known treatments).
Args:
results: List of processed results
Returns:
Potential discoveries only
"""
return [r for r in results if "Potential Discovery" in r.get('status', '')]
@staticmethod
def get_results_by_target(results: List[Dict], target_symbol: str) -> List[Dict]:
"""
Returns results filtered by target symbol.
Args:
results: List of processed results
target_symbol: Target protein symbol to filter by
Returns:
Filtered results for the specific target
"""
return [r for r in results if r['target_symbol'] == target_symbol]
@staticmethod
def format_results_table(results: List[Dict], top_n: int = 15) -> str:
"""
Formats results as a string table.
Args:
results: List of results to format
top_n: Number of results to display
Returns:
Formatted table string
"""
table = f"{'Drug Name':<20} | {'Target':<10} | {'Score':<8} | {'Status'}\n"
table += "-" * 65 + "\n"
for res in results[:top_n]:
table += (
f"{res['drug_name']:<20} | "
f"{res['target_symbol']:<10} | "
f"{res['score']:<8} | "
f"{res['status']}\n"
)
return table
@staticmethod
def export_results_csv(results: List[Dict], filename: str) -> None:
"""
Exports results to CSV file.
Args:
results: List of results to export
filename: Output CSV filename
"""
try:
import pandas as pd
df = pd.DataFrame(results)
df.to_csv(filename, index=False)
logger.info(f"✅ Results exported to {filename}")
except Exception as e:
logger.error(f"Error exporting results: {str(e)}")
raise
|