mirror of
https://github.com/K-Dense-AI/claude-scientific-skills.git
synced 2026-03-28 07:33:45 +08:00
124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
import pandas as pd
|
|
import os
|
|
import json
|
|
from typing import List, Dict, Optional, Union
|
|
|
|
# Default data path
|
|
DATA_PATH = "/mnt/c/Users/eamon/Documents/Data/PrimeKG/kg.csv"
|
|
|
|
def _load_kg():
|
|
"""Internal helper to load the KG efficiently."""
|
|
if not os.path.exists(DATA_PATH):
|
|
raise FileNotFoundError(f"PrimeKG data not found at {DATA_PATH}. Please ensure the file is downloaded.")
|
|
# For very large files, we might want to use a database or specialized graph library.
|
|
# For now, we'll use pandas for simplicity but with low_memory=True.
|
|
return pd.read_csv(DATA_PATH, low_memory=True)
|
|
|
|
def search_nodes(name_query: str, node_type: Optional[str] = None) -> List[Dict]:
|
|
"""
|
|
Search for nodes in PrimeKG by name and optionally type.
|
|
|
|
Args:
|
|
name_query: String to search for in node names.
|
|
node_type: Optional type of node (e.g., 'gene/protein', 'drug', 'disease').
|
|
|
|
Returns:
|
|
List of matching nodes with their metadata.
|
|
"""
|
|
kg = _load_kg()
|
|
|
|
# Check both x and y columns for unique nodes
|
|
x_nodes = kg[['x_id', 'x_type', 'x_name', 'x_source']].drop_duplicates()
|
|
x_nodes.columns = ['id', 'type', 'name', 'source']
|
|
|
|
y_nodes = kg[['y_id', 'y_type', 'y_name', 'y_source']].drop_duplicates()
|
|
y_nodes.columns = ['id', 'type', 'name', 'source']
|
|
|
|
nodes = pd.concat([x_nodes, y_nodes]).drop_duplicates()
|
|
|
|
mask = nodes['name'].str.contains(name_query, case=False, na=False)
|
|
if node_type:
|
|
mask &= (nodes['type'] == node_type)
|
|
|
|
results = nodes[mask].head(20).to_dict(orient='records')
|
|
return results
|
|
|
|
def get_neighbors(node_id: Union[str, int], relation_type: Optional[str] = None) -> List[Dict]:
|
|
"""
|
|
Get all direct neighbors of a specific node.
|
|
|
|
Args:
|
|
node_id: The ID of the node (e.g., NCBI Gene ID or ChEMBL ID).
|
|
relation_type: Optional filter for specific relationship types.
|
|
|
|
Returns:
|
|
List of neighbors and the relationship metadata.
|
|
"""
|
|
kg = _load_kg()
|
|
node_id = str(node_id)
|
|
|
|
mask_x = (kg['x_id'].astype(str) == node_id)
|
|
mask_y = (kg['y_id'].astype(str) == node_id)
|
|
|
|
if relation_type:
|
|
mask_x &= (kg['relation'] == relation_type)
|
|
mask_y &= (kg['relation'] == relation_type)
|
|
|
|
neighbors_x = kg[mask_x][['relation', 'display_relation', 'y_id', 'y_type', 'y_name', 'y_source']]
|
|
neighbors_x.columns = ['relation', 'display_relation', 'neighbor_id', 'neighbor_type', 'neighbor_name', 'neighbor_source']
|
|
|
|
neighbors_y = kg[mask_y][['relation', 'display_relation', 'x_id', 'x_type', 'x_name', 'x_source']]
|
|
neighbors_y.columns = ['relation', 'display_relation', 'neighbor_id', 'neighbor_type', 'neighbor_name', 'neighbor_source']
|
|
|
|
results = pd.concat([neighbors_x, neighbors_y]).to_dict(orient='records')
|
|
return results
|
|
|
|
def find_paths(start_node_id: str, end_node_id: str, max_depth: int = 2) -> List[List[Dict]]:
|
|
"""
|
|
Find paths between two nodes (e.g., Drug to Disease) up to a certain depth.
|
|
Note: Simple BFS implementation.
|
|
"""
|
|
kg = _load_kg()
|
|
start_node_id = str(start_node_id)
|
|
end_node_id = str(end_node_id)
|
|
|
|
# Simplified path finding for depth 1 and 2
|
|
# Depth 1
|
|
direct = kg[((kg['x_id'].astype(str) == start_node_id) & (kg['y_id'].astype(str) == end_node_id)) |
|
|
((kg['y_id'].astype(str) == start_node_id) & (kg['x_id'].astype(str) == end_node_id))]
|
|
|
|
paths = []
|
|
for _, row in direct.iterrows():
|
|
paths.append([row.to_dict()])
|
|
|
|
if max_depth >= 2:
|
|
# Find neighbors of start
|
|
n1_x = kg[kg['x_id'].astype(str) == start_node_id]
|
|
n1_y = kg[kg['y_id'].astype(str) == start_node_id]
|
|
|
|
# This is computationally expensive in pure pandas for a large KG.
|
|
# Implementation skipped for brevity in this MVP, but suggested for full version.
|
|
pass
|
|
|
|
return paths
|
|
|
|
def get_disease_context(disease_name: str) -> Dict:
|
|
"""
|
|
Analyze the local graph around a disease: associated genes, drugs, and phenotypes.
|
|
"""
|
|
results = search_nodes(disease_name, node_type='disease')
|
|
if not results:
|
|
return {"error": "Disease not found"}
|
|
|
|
disease_id = results[0]['id']
|
|
neighbors = get_neighbors(disease_id)
|
|
|
|
summary = {
|
|
"disease_info": results[0],
|
|
"associated_genes": [n for n in neighbors if n['neighbor_type'] == 'gene/protein'],
|
|
"associated_drugs": [n for n in neighbors if n['neighbor_type'] == 'drug'],
|
|
"phenotypes": [n for n in neighbors if n['neighbor_type'] == 'phenotype'],
|
|
"related_diseases": [n for n in neighbors if n['neighbor_type'] == 'disease']
|
|
}
|
|
return summary
|