Skills › Knowledge Work Plugins › nextflow-development › ncbi_utils.py
ncbi_utils.py
808 lines · bio-research/skills/nextflow-development/scripts/utils/ncbi_utils.py at 2574e29
1#!/usr/bin/env python3
2"""
3NCBI Utilities for GEO/SRA Data Access
4======================================
5Shared utilities for fetching metadata and downloading data from NCBI services.
6"""
8import json
9import logging
10import re
11import shutil
12import time
13from pathlib import Path
14from typing import Dict, List, Optional, Tuple
15from urllib.request import Request, urlopen
16from urllib.error import URLError, HTTPError
18# Set up logging
19logging.basicConfig(
20 level=logging.INFO,
21 format='%(asctime)s - %(levelname)s - %(message)s'
22)
23logger = logging.getLogger(__name__)
25# NCBI rate limiting - track last request time
26_last_ncbi_request_time = 0.0
27_NCBI_MIN_DELAY = 0.34 # 3 requests per second max without API key
30def _rate_limit_ncbi():
31 """Enforce NCBI rate limit of 3 requests/second."""
32 global _last_ncbi_request_time
33 current_time = time.time()
34 elapsed = current_time - _last_ncbi_request_time
35 if elapsed < _NCBI_MIN_DELAY:
36 time.sleep(_NCBI_MIN_DELAY - elapsed)
37 _last_ncbi_request_time = time.time()
40# Try to import requests for better HTTP handling
41try:
42 import requests
43 HAS_REQUESTS = True
44except ImportError:
45 HAS_REQUESTS = False
46 logger.debug("requests not installed - using urllib fallback")
49def check_network_access() -> Tuple[bool, str]:
50 """
51 Check if NCBI/ENA servers are accessible.
53 Returns:
54 Tuple of (success, message)
55 """
56 test_urls = [
57 ("NCBI Entrez", "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/einfo.fcgi"),
58 ("NCBI FTP", "https://ftp.ncbi.nlm.nih.gov/"),
59 ("ENA API", "https://www.ebi.ac.uk/ena/portal/api/"),
60 ]
62 results = []
63 for name, url in test_urls:
64 try:
65 if HAS_REQUESTS:
66 # Use GET instead of HEAD - NCBI Entrez returns 405 for HEAD
67 response = requests.get(url, timeout=10)
68 success = response.status_code < 400
69 else:
70 req = Request(url, headers={'User-Agent': 'geo-sra-skill/1.0'})
71 with urlopen(req, timeout=10) as response:
72 success = True
73 results.append((name, success, None))
74 except Exception as e:
75 results.append((name, False, str(e)))
77 all_success = all(r[1] for r in results)
79 msg_parts = []
80 for name, success, error in results:
81 status = "✓" if success else "✗"
82 msg_parts.append(f" {status} {name}: {'OK' if success else error or 'Failed'}")
84 return all_success, "\n".join(msg_parts)
87def fetch_geo_metadata(geo_id: str) -> Optional[Dict]:
88 """
89 Fetch GEO study metadata using NCBI Entrez E-utilities.
91 Args:
92 geo_id: GEO accession (e.g., 'GSE110004')
94 Returns:
95 Dict with study metadata or None if failed
96 """
97 try:
98 # Use esearch to get GEO UID
99 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=gds&term={geo_id}[Accession]&retmode=json"
101 _rate_limit_ncbi()
102 if HAS_REQUESTS:
103 response = requests.get(search_url, timeout=30)
104 data = response.json()
105 else:
106 with urlopen(search_url, timeout=30) as response:
107 data = json.loads(response.read().decode())
109 id_list = data.get('esearchresult', {}).get('idlist', [])
110 if not id_list:
111 logger.warning(f"No GEO entry found for {geo_id}")
112 return None
114 # Use esummary to get metadata
115 uid = id_list[0]
116 summary_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?db=gds&id={uid}&retmode=json"
118 _rate_limit_ncbi()
119 if HAS_REQUESTS:
120 response = requests.get(summary_url, timeout=30)
121 data = response.json()
122 else:
123 with urlopen(summary_url, timeout=30) as response:
124 data = json.loads(response.read().decode())
126 result = data.get('result', {}).get(uid, {})
128 return {
129 'geo_id': geo_id,
130 'title': result.get('title', 'N/A'),
131 'summary': result.get('summary', 'N/A'),
132 'organism': result.get('taxon', 'N/A'),
133 'n_samples': result.get('n_samples', 0),
134 'gpl': result.get('gpl', 'N/A'),
135 'entrytype': result.get('entrytype', 'N/A'),
136 'pubmed_ids': result.get('pubmedids', []),
137 }
139 except Exception as e:
140 logger.error(f"Error fetching GEO metadata for {geo_id}: {e}")
141 return None
144def fetch_sra_study_accession(geo_id: str) -> Optional[str]:
145 """
146 Get the SRA study accession (SRPxxxxxx) for a GEO accession.
148 Args:
149 geo_id: GEO accession (e.g., 'GSE110004')
151 Returns:
152 SRA study accession (e.g., 'SRP126328') or None
153 """
154 try:
155 # Search for SRA study linked to GEO
156 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=sra&term={geo_id}[GEO]&retmode=json"
158 _rate_limit_ncbi()
159 if HAS_REQUESTS:
160 response = requests.get(search_url, timeout=30)
161 data = response.json()
162 else:
163 with urlopen(search_url, timeout=30) as response:
164 data = json.loads(response.read().decode())
166 id_list = data.get('esearchresult', {}).get('idlist', [])
167 if not id_list:
168 return None
170 # Get summary to extract SRP accession
171 uid = id_list[0]
172 summary_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?db=sra&id={uid}&retmode=json"
174 _rate_limit_ncbi()
175 if HAS_REQUESTS:
176 response = requests.get(summary_url, timeout=30)
177 data = response.json()
178 else:
179 with urlopen(summary_url, timeout=30) as response:
180 data = json.loads(response.read().decode())
182 result = data.get('result', {}).get(uid, {})
183 exp_xml = result.get('expxml', '')
185 # Extract SRP from the XML
186 srp_match = re.search(r'<Study acc="(SRP\d+)"', exp_xml)
187 if srp_match:
188 return srp_match.group(1)
190 return None
192 except Exception as e:
193 logger.debug(f"Error fetching SRA study for {geo_id}: {e}")
194 return None
197def fetch_sra_run_info(geo_id: str, bioproject: Optional[str] = None) -> List[Dict]:
198 """
199 Fetch SRA run information for all samples in a GEO study.
201 Args:
202 geo_id: GEO accession (e.g., 'GSE110004')
203 bioproject: Optional BioProject accession for fallback search
205 Returns:
206 List of dicts with run info (srr, gsm, layout, library_strategy, etc.)
207 """
208 runs = []
210 try:
211 # First get the BioProject accession
212 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=sra&term={geo_id}[GEO]&retmax=1000&retmode=json"
214 _rate_limit_ncbi()
215 if HAS_REQUESTS:
216 response = requests.get(search_url, timeout=30)
217 data = response.json()
218 else:
219 with urlopen(search_url, timeout=30) as response:
220 data = json.loads(response.read().decode())
222 id_list = data.get('esearchresult', {}).get('idlist', [])
224 # If no results, try BioProject fallback
225 if not id_list:
226 if not bioproject:
227 bioproject = fetch_bioproject_from_geo(geo_id)
229 if bioproject:
230 logger.info(f"Using BioProject {bioproject} for {geo_id}")
231 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=sra&term={bioproject}&retmax=1000&retmode=json"
233 _rate_limit_ncbi()
234 if HAS_REQUESTS:
235 response = requests.get(search_url, timeout=30)
236 data = response.json()
237 else:
238 with urlopen(search_url, timeout=30) as response:
239 data = json.loads(response.read().decode())
241 id_list = data.get('esearchresult', {}).get('idlist', [])
243 if not id_list:
244 logger.warning(f"No SRA entries found for {geo_id}")
245 return runs
247 # Batch fetch summaries
248 ids_str = ','.join(id_list)
249 summary_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?db=sra&id={ids_str}&retmode=json"
251 _rate_limit_ncbi()
252 if HAS_REQUESTS:
253 response = requests.get(summary_url, timeout=60)
254 data = response.json()
255 else:
256 with urlopen(summary_url, timeout=60) as response:
257 data = json.loads(response.read().decode())
259 result = data.get('result', {})
261 for uid in id_list:
262 entry = result.get(uid, {})
263 if not entry:
264 continue
266 exp_xml = entry.get('expxml', '')
267 runs_xml = entry.get('runs', '')
269 # Extract metadata from XML
270 layout_match = re.search(r'<LIBRARY_LAYOUT>\s*<(\w+)', exp_xml)
271 strategy_match = re.search(r'<LIBRARY_STRATEGY>(\w+)', exp_xml)
272 source_match = re.search(r'<LIBRARY_SOURCE>(\w+)', exp_xml)
273 gsm_match = re.search(r'<Sample acc="(GSM\d+)"', exp_xml)
274 srx_match = re.search(r'<Experiment acc="(SRX\d+)"', exp_xml)
276 # Extract run accessions
277 srr_matches = re.findall(r'<Run acc="(SRR\d+)"[^>]*total_spots="(\d+)"[^>]*total_bases="(\d+)"', runs_xml)
279 for srr, spots, bases in srr_matches:
280 runs.append({
281 'srr': srr,
282 'srx': srx_match.group(1) if srx_match else '',
283 'gsm': gsm_match.group(1) if gsm_match else '',
284 'layout': layout_match.group(1).upper() if layout_match else 'UNKNOWN',
285 'library_strategy': strategy_match.group(1) if strategy_match else 'UNKNOWN',
286 'library_source': source_match.group(1) if source_match else 'UNKNOWN',
287 'spots': int(spots),
288 'bases': int(bases),
289 })
291 return runs
293 except Exception as e:
294 logger.error(f"Error fetching SRA run info for {geo_id}: {e}")
295 return runs
298def fetch_ena_fastq_urls(study_accession: str) -> Dict[str, List[str]]:
299 """
300 Get FASTQ download URLs from ENA for an SRA study.
302 ENA provides faster downloads than SRA with pre-split paired files.
304 Args:
305 study_accession: SRA study accession (e.g., 'SRP126328')
307 Returns:
308 Dict mapping SRR accession to list of FASTQ URLs
309 """
310 fastq_urls = {}
312 try:
313 # Query ENA API
314 ena_url = f"https://www.ebi.ac.uk/ena/portal/api/filereport?accession={study_accession}&result=read_run&fields=run_accession,sample_alias,fastq_ftp&format=tsv"
316 if HAS_REQUESTS:
317 response = requests.get(ena_url, timeout=60)
318 content = response.text
319 else:
320 with urlopen(ena_url, timeout=60) as response:
321 content = response.read().decode()
323 lines = content.strip().split('\n')
324 if len(lines) < 2:
325 logger.warning(f"No FASTQ URLs found in ENA for {study_accession}")
326 return fastq_urls
328 # Parse TSV
329 header = lines[0].split('\t')
330 run_idx = header.index('run_accession') if 'run_accession' in header else 0
331 ftp_idx = header.index('fastq_ftp') if 'fastq_ftp' in header else 2
333 for line in lines[1:]:
334 if not line.strip():
335 continue
336 fields = line.split('\t')
337 if len(fields) > max(run_idx, ftp_idx):
338 srr = fields[run_idx]
339 ftp_urls = fields[ftp_idx]
340 if ftp_urls:
341 # URLs are semicolon-separated, convert to HTTP URLs
342 # ENA supports both FTP and HTTP, HTTP is easier with requests
343 urls = [f"http://{url}" for url in ftp_urls.split(';') if url]
344 fastq_urls[srr] = urls
346 return fastq_urls
348 except Exception as e:
349 logger.error(f"Error fetching ENA URLs for {study_accession}: {e}")
350 return fastq_urls
353def download_file(url: str, output_path: Path, timeout: int = 300, show_progress: bool = True) -> bool:
354 """
355 Download a file with progress indication.
357 Args:
358 url: URL to download
359 output_path: Path to save file
360 timeout: Download timeout in seconds
361 show_progress: Show progress bar
363 Returns:
364 True if successful, False otherwise
365 """
366 try:
367 output_path.parent.mkdir(parents=True, exist_ok=True)
369 if HAS_REQUESTS:
370 response = requests.get(url, stream=True, timeout=timeout)
371 response.raise_for_status()
373 total_size = int(response.headers.get('content-length', 0))
375 with open(output_path, 'wb') as f:
376 downloaded = 0
377 for chunk in response.iter_content(chunk_size=8192):
378 f.write(chunk)
379 downloaded += len(chunk)
380 if show_progress and total_size > 0:
381 pct = (downloaded / total_size) * 100
382 print(f"\r Progress: {pct:.1f}%", end='', flush=True)
383 if show_progress:
384 print() # New line after progress
385 return True
386 else:
387 # Fallback to urllib
388 req = Request(url, headers={'User-Agent': 'geo-sra-skill/1.0'})
389 with urlopen(req, timeout=timeout) as response:
390 with open(output_path, 'wb') as f:
391 shutil.copyfileobj(response, f)
392 return True
394 except Exception as e:
395 logger.error(f"Download error for {url}: {e}")
396 return False
399def fetch_pubmed_metadata(pmid: str, max_retries: int = 3) -> Optional[Dict]:
400 """
401 Fetch paper metadata from PubMed.
403 Args:
404 pmid: PubMed ID
405 max_retries: Number of retries on failure
407 Returns:
408 Dict with 'authors', 'year', 'journal', 'doi' or None
409 """
410 url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?db=pubmed&id={pmid}&retmode=json"
412 for attempt in range(max_retries):
413 try:
414 _rate_limit_ncbi()
415 if HAS_REQUESTS:
416 response = requests.get(url, timeout=30)
417 data = response.json()
418 else:
419 with urlopen(url, timeout=30) as response:
420 data = json.loads(response.read().decode())
422 result = data.get('result', {}).get(pmid, {})
424 if not result or 'error' in result:
425 if attempt < max_retries - 1:
426 time.sleep(1 * (attempt + 1))
427 continue
428 return None
430 # Extract authors
431 authors_list = result.get('authors', [])
432 if not authors_list:
433 if attempt < max_retries - 1:
434 time.sleep(1 * (attempt + 1))
435 continue
436 return None
438 author_names = [f"{a.get('name', '')}" for a in authors_list[:3]]
439 authors = ', '.join(author_names)
440 if len(authors_list) > 3:
441 authors += ', et al.'
443 # Extract year
444 pubdate = result.get('pubdate', '')
445 year_match = re.search(r'\b(20\d{2})\b', pubdate)
446 year = year_match.group(1) if year_match else "Unknown"
448 # Extract journal
449 journal = result.get('source', 'Unknown')
451 # Extract DOI
452 doi = ""
453 for aid in result.get('articleids', []):
454 if aid.get('idtype') == 'doi':
455 doi = aid.get('value', '')
456 break
458 return {
459 'authors': authors,
460 'year': year,
461 'journal': journal,
462 'doi': doi,
463 'title': result.get('title', '')
464 }
466 except Exception as e:
467 logger.debug(f"PubMed fetch error for PMID {pmid} (attempt {attempt + 1}): {e}")
468 if attempt < max_retries - 1:
469 time.sleep(1 * (attempt + 1))
470 continue
472 return None
475def format_file_size(size_bytes: int) -> str:
476 """Format file size in human-readable format."""
477 if size_bytes < 1024:
478 return f"{size_bytes} B"
479 elif size_bytes < 1024 * 1024:
480 return f"{size_bytes / 1024:.1f} KB"
481 elif size_bytes < 1024 * 1024 * 1024:
482 return f"{size_bytes / (1024 * 1024):.1f} MB"
483 else:
484 return f"{size_bytes / (1024 * 1024 * 1024):.1f} GB"
487def estimate_download_size(runs: List[Dict]) -> int:
488 """
489 Estimate total download size from SRA run info.
491 Args:
492 runs: List of run info dicts with 'bases' field
494 Returns:
495 Estimated size in bytes (rough estimate based on bases)
496 """
497 total_bases = sum(r.get('bases', 0) for r in runs)
498 # FASTQ is roughly 1 byte per base when compressed
499 return total_bases // 4 # Rough compression ratio
502def fetch_bioproject_from_geo(geo_id: str) -> Optional[str]:
503 """
504 Fetch BioProject accession linked to a GEO study.
506 Args:
507 geo_id: GEO accession (e.g., 'GSE110004')
509 Returns:
510 BioProject accession (e.g., 'PRJNA432544') or None
511 """
512 try:
513 # First get GDS UID
514 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=gds&term={geo_id}[Accession]&retmode=json"
516 _rate_limit_ncbi()
517 if HAS_REQUESTS:
518 response = requests.get(search_url, timeout=30)
519 data = response.json()
520 else:
521 with urlopen(search_url, timeout=30) as response:
522 data = json.loads(response.read().decode())
524 gds_ids = data.get('esearchresult', {}).get('idlist', [])
525 if not gds_ids:
526 return None
528 # Get linked BioProject
529 elink_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/elink.fcgi?dbfrom=gds&db=bioproject&id={gds_ids[0]}&retmode=json"
531 _rate_limit_ncbi()
532 if HAS_REQUESTS:
533 response = requests.get(elink_url, timeout=30)
534 data = response.json()
535 else:
536 with urlopen(elink_url, timeout=30) as response:
537 data = json.loads(response.read().decode())
539 linksets = data.get('linksets', [])
540 if linksets and linksets[0].get('linksetdbs'):
541 for linksetdb in linksets[0]['linksetdbs']:
542 if linksetdb.get('dbto') == 'bioproject':
543 bp_ids = linksetdb.get('links', [])
544 if bp_ids:
545 # Get BioProject accession
546 summary_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi?db=bioproject&id={bp_ids[0]}&retmode=json"
547 _rate_limit_ncbi()
548 if HAS_REQUESTS:
549 response = requests.get(summary_url, timeout=30)
550 data = response.json()
551 else:
552 with urlopen(summary_url, timeout=30) as response:
553 data = json.loads(response.read().decode())
555 result = data.get('result', {}).get(str(bp_ids[0]), {})
556 return result.get('project_acc')
558 return None
560 except Exception as e:
561 logger.debug(f"Error fetching BioProject for {geo_id}: {e}")
562 return None
565def fetch_sra_run_info_detailed(geo_id: str, bioproject: Optional[str] = None) -> List[Dict]:
566 """
567 Fetch detailed SRA run information using efetch CSV format.
569 This provides richer metadata than esummary, including sample names.
571 Args:
572 geo_id: GEO accession (e.g., 'GSE110004')
573 bioproject: Optional BioProject accession for fallback search
575 Returns:
576 List of dicts with detailed run info
577 """
578 runs = []
580 try:
581 # First get SRA UIDs using GEO search
582 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=sra&term={geo_id}[GEO]&retmax=1000&retmode=json"
584 _rate_limit_ncbi()
585 if HAS_REQUESTS:
586 response = requests.get(search_url, timeout=30)
587 data = response.json()
588 else:
589 with urlopen(search_url, timeout=30) as response:
590 data = json.loads(response.read().decode())
592 id_list = data.get('esearchresult', {}).get('idlist', [])
594 # If no results with GEO search, try BioProject
595 if not id_list:
596 # Try to find BioProject if not provided
597 if not bioproject:
598 logger.info(f"No direct SRA link for {geo_id}, searching for BioProject...")
599 bioproject = fetch_bioproject_from_geo(geo_id)
601 if bioproject:
602 logger.info(f"Found BioProject: {bioproject}")
603 search_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=sra&term={bioproject}&retmax=1000&retmode=json"
605 _rate_limit_ncbi()
606 if HAS_REQUESTS:
607 response = requests.get(search_url, timeout=30)
608 data = response.json()
609 else:
610 with urlopen(search_url, timeout=30) as response:
611 data = json.loads(response.read().decode())
613 id_list = data.get('esearchresult', {}).get('idlist', [])
615 if not id_list:
616 logger.warning(f"No SRA entries found for {geo_id}")
617 return runs
619 # Fetch run info in CSV format using efetch
620 ids_str = ','.join(id_list)
621 efetch_url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi?db=sra&id={ids_str}&rettype=runinfo&retmode=csv"
623 _rate_limit_ncbi()
624 if HAS_REQUESTS:
625 response = requests.get(efetch_url, timeout=60)
626 content = response.text
627 else:
628 with urlopen(efetch_url, timeout=60) as response:
629 content = response.read().decode()
631 lines = content.strip().split('\n')
632 if len(lines) < 1:
633 return runs
635 # NCBI efetch runinfo CSV doesn't include headers
636 # Define the fixed column order for SRA runinfo format
637 header = [
638 'Run', 'ReleaseDate', 'LoadDate', 'spots', 'bases', 'spots_with_mates',
639 'avgLength', 'size_MB', 'AssemblyName', 'download_path', 'Experiment',
640 'LibraryName', 'LibraryStrategy', 'LibrarySelection', 'LibrarySource',
641 'LibraryLayout', 'InsertSize', 'InsertDev', 'Platform', 'Model',
642 'SRAStudy', 'BioProject', 'Study_Pubmed_id', 'ProjectID', 'Sample',
643 'BioSample', 'SampleType', 'TaxID', 'ScientificName', 'SampleName',
644 'g1k_pop_code', 'source', 'g1k_analysis_group', 'Subject_ID', 'Sex',
645 'Disease', 'Tumor', 'Affection_Status', 'Analyte_Type', 'Histological_Type',
646 'Body_Site', 'CenterName', 'Submission', 'dbgap_study_accession', 'Consent',
647 'RunHash', 'ReadHash'
648 ]
650 # Map column names to indices
651 col_map = {col: idx for idx, col in enumerate(header)}
653 for line in lines:
654 if not line.strip():
655 continue
657 # Handle CSV fields (some may contain commas in quotes)
658 fields = _parse_csv_line(line)
659 if len(fields) < len(header):
660 continue
662 def get_field(name, default=''):
663 idx = col_map.get(name, -1)
664 return fields[idx] if idx >= 0 and idx < len(fields) else default
666 run = {
667 'srr': get_field('Run'),
668 'srx': get_field('Experiment'),
669 'gsm': get_field('SampleName'), # Often GSM ID
670 'sample_name': get_field('SampleName'),
671 'library_name': get_field('LibraryName'),
672 'layout': get_field('LibraryLayout', 'UNKNOWN').upper(),
673 'library_strategy': get_field('LibraryStrategy', 'UNKNOWN'),
674 'library_source': get_field('LibrarySource', 'UNKNOWN'),
675 'library_selection': get_field('LibrarySelection', ''),
676 'platform': get_field('Platform'),
677 'model': get_field('Model'),
678 'organism': get_field('ScientificName', ''),
679 'spots': int(get_field('spots', 0) or 0),
680 'bases': int(get_field('bases', 0) or 0),
681 'size_mb': float(get_field('size_MB', 0) or 0),
682 'bioproject': get_field('BioProject'),
683 'biosample': get_field('BioSample'),
684 'sra_study': get_field('SRAStudy'),
685 }
687 # Only add if we have a valid SRR
688 if run['srr'].startswith('SRR'):
689 runs.append(run)
691 return runs
693 except Exception as e:
694 logger.error(f"Error fetching detailed SRA run info for {geo_id}: {e}")
695 return runs
698def _parse_csv_line(line: str) -> List[str]:
699 """Parse a CSV line handling quoted fields."""
700 import csv
701 import io
702 reader = csv.reader(io.StringIO(line))
703 for row in reader:
704 return row
705 return []
708def group_samples_by_type(runs: List[Dict]) -> Dict[str, Dict]:
709 """
710 Group SRA runs by library type and layout.
712 Returns dict with group names as keys and info dicts as values:
713 {
714 'RNA-Seq:PAIRED': {
715 'runs': [...],
716 'count': 18,
717 'gsm_range': 'GSM2879618-GSM2879635',
718 'size_estimate': 50000000000,
719 'description': 'RNA-Seq paired-end'
720 },
721 ...
722 }
723 """
724 groups = {}
726 for run in runs:
727 strategy = run.get('library_strategy', 'UNKNOWN')
728 layout = run.get('layout', 'UNKNOWN')
729 key = f"{strategy}:{layout}"
731 if key not in groups:
732 groups[key] = {
733 'runs': [],
734 'gsm_ids': set(),
735 'total_bases': 0,
736 'strategy': strategy,
737 'layout': layout,
738 }
740 groups[key]['runs'].append(run)
741 gsm = run.get('gsm', '')
742 if gsm.startswith('GSM'):
743 groups[key]['gsm_ids'].add(gsm)
744 groups[key]['total_bases'] += run.get('bases', 0)
746 # Post-process groups
747 result = {}
748 for key, info in groups.items():
749 gsm_list = sorted(info['gsm_ids'])
750 gsm_range = _format_gsm_range(gsm_list) if gsm_list else 'N/A'
752 result[key] = {
753 'runs': info['runs'],
754 'count': len(info['runs']),
755 'gsm_range': gsm_range,
756 'gsm_ids': gsm_list,
757 'size_estimate': info['total_bases'] // 4, # Rough compressed size
758 'strategy': info['strategy'],
759 'layout': info['layout'],
760 'description': f"{info['strategy']} {info['layout'].lower()}",
761 }
763 return result
766def _format_gsm_range(gsm_list: List[str]) -> str:
767 """Format list of GSM IDs as a range if consecutive."""
768 if not gsm_list:
769 return 'N/A'
771 if len(gsm_list) == 1:
772 return gsm_list[0]
774 # Extract numbers and check if consecutive
775 try:
776 numbers = [int(gsm.replace('GSM', '')) for gsm in gsm_list]
777 numbers.sort()
779 if numbers[-1] - numbers[0] == len(numbers) - 1:
780 # Consecutive
781 return f"GSM{numbers[0]}-GSM{numbers[-1]}"
782 else:
783 # Not consecutive, show count
784 return f"{gsm_list[0]}...({len(gsm_list)} samples)"
785 except ValueError:
786 return f"{len(gsm_list)} samples"
789def format_sample_groups_table(groups: Dict[str, Dict]) -> str:
790 """Format sample groups as a readable table."""
791 lines = []
792 lines.append("")
793 lines.append(f"{'Sample Group':<20} {'Count':>6} {'Layout':<10} {'GSM Range':<25} {'Est. Size':>12}")
794 lines.append("-" * 80)
796 for key, info in sorted(groups.items(), key=lambda x: -x[1]['count']):
797 size_str = format_file_size(info['size_estimate'])
798 lines.append(
799 f"{info['strategy']:<20} {info['count']:>6} {info['layout']:<10} "
800 f"{info['gsm_range']:<25} {size_str:>12}"
801 )
803 lines.append("-" * 80)
804 total_runs = sum(g['count'] for g in groups.values())
805 total_size = sum(g['size_estimate'] for g in groups.values())
806 lines.append(f"{'TOTAL':<20} {total_runs:>6} {'':<10} {'':<25} {format_file_size(total_size):>12}")
808 return '\n'.join(lines)