File size: 4,340 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 | """
Stage 3: PDB Structure & UniProt ID Fetch
Retrieves PDB structure IDs and UniProt accession numbers for target proteins.
Includes file-based caching to avoid redundant API calls.
"""
import requests
import json
import logging
import os
from datetime import datetime
from typing import List, Dict, Optional
logger = logging.getLogger(__name__)
class PdbStructurePipeline:
def __init__(self, api_url: str = "https://rest.uniprot.org/uniprotkb/search",
cache_dir: str = None):
self.api_url = api_url
if cache_dir is None:
cache_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "data")
self.cache_dir = cache_dir
os.makedirs(self.cache_dir, exist_ok=True)
def _get_cache_path(self, symbol: str) -> str:
return os.path.join(self.cache_dir, f"cache_pdb_{symbol.lower()}.json")
def _load_from_cache(self, symbol: str) -> Optional[Dict]:
cache_path = self._get_cache_path(symbol)
if os.path.exists(cache_path):
try:
with open(cache_path, "r") as f:
data = json.load(f)
if data.get("symbol") == symbol:
logger.info(f" 📦 Loaded PDB data from cache for: {symbol}")
return data
except (json.JSONDecodeError, IOError) as e:
logger.warning(f" ⚠️ Failed to read cache for {symbol}: {e}")
return None
def _save_to_cache(self, symbol: str, data: Dict):
cache_path = self._get_cache_path(symbol)
data["fetched_at"] = datetime.utcnow().isoformat()
try:
with open(cache_path, "w") as f:
json.dump(data, f, indent=2)
except IOError as e:
logger.warning(f" ⚠️ Failed to write cache for {symbol}: {e}")
def fetch_pdb_ids(self, targets: List[Dict]) -> List[Dict]:
logger.info(f"Fetching PDB structures for {len(targets)} targets...")
enriched = []
for t in targets:
symbol = t["symbol"]
cached = self._load_from_cache(symbol)
if cached is not None:
t["pdb_ids"] = cached.get("pdb_ids", [])
t["uniprot_id"] = cached.get("uniprot_id", "")
enriched.append(t)
continue
try:
params = {
"query": f"{symbol} AND taxonomy_id:9606",
"format": "json",
"size": 1,
}
response = requests.get(self.api_url, params=params, timeout=30)
response.raise_for_status()
data = response.json()
uniprot_id = ""
pdb_ids = []
if data.get("results") and len(data["results"]) > 0:
result = data["results"][0]
uniprot_id = result.get("primaryAccession", "")
xrefs = result.get("uniProtKBCrossReferences", [])
for xref in xrefs:
if xref.get("database") == "PDB":
pdb_id = xref.get("id")
if pdb_id:
pdb_ids.append(pdb_id)
pdb_ids = pdb_ids[:3]
t["pdb_ids"] = pdb_ids
t["uniprot_id"] = uniprot_id
if pdb_ids:
logger.info(f" ✅ Found {len(pdb_ids)} PDB IDs for: {symbol}")
else:
logger.warning(f" ⚠️ No PDB structures found for: {symbol}")
self._save_to_cache(symbol, {
"symbol": symbol,
"uniprot_id": uniprot_id,
"pdb_ids": pdb_ids,
})
except requests.exceptions.RequestException as e:
logger.warning(f" ⚠️ Failed to fetch PDB for {symbol}: {e}")
t["pdb_ids"] = []
t["uniprot_id"] = ""
except Exception as e:
logger.warning(f" ⚠️ Failed to fetch PDB for {symbol}: {e}")
t["pdb_ids"] = []
t["uniprot_id"] = ""
enriched.append(t)
logger.info(f"✅ PDB fetch complete for {len(enriched)}/{len(targets)} targets")
return enriched
|