Skill Tracer
Skills › Scientific Agent Skills › gget › enrichment_pipeline.py

enrichment_pipeline.py

246 lines · skills/gget/scripts/enrichment_pipeline.py at 5f2bc03

1#!/usr/bin/env python3
2"""
3Enrichment Analysis Pipeline
4Perform comprehensive enrichment analysis on a gene list
5"""
6
7import argparse
8import sys
9from pathlib import Path
10import gget
11import pandas as pd
14def read_gene_list(file_path):
15 """Read gene list from file (one gene per line or CSV)."""
16 file_path = Path(file_path)
18 if file_path.suffix == ".csv":
19 df = pd.read_csv(file_path)
20 # Assume first column contains gene names
21 genes = df.iloc[:, 0].tolist()
22 else:
23 # Plain text file
24 with open(file_path, "r") as f:
25 genes = [line.strip() for line in f if line.strip()]
27 return genes
30def enrichment_pipeline(
31 gene_list,
32 species="human",
33 background=None,
34 output_prefix="enrichment",
35 plot=True,
36 databases=None,
37):
38 """
39 Perform comprehensive enrichment analysis.
41 Args:
42 gene_list: List of gene symbols
43 species: Species for analysis
44 background: Background gene list (optional)
45 output_prefix: Prefix for output files
46 plot: Whether to generate plots
47 databases: Optional list of full Enrichr library names
48 """
49 print("Enrichment Analysis Pipeline")
50 print("=" * 60)
51 print(f"Analyzing {len(gene_list)} genes")
52 print(f"Species: {species}\n")
54 if not gene_list:
55 raise ValueError("Gene list must not be empty")
56 if species not in {"human", "mouse"} and background is not None:
57 raise ValueError("Custom backgrounds are supported only for human and mouse")
58 if databases is None:
59 if species not in {"human", "mouse"}:
60 raise ValueError("Supply --database with a full species-specific Enrichr library name")
61 databases = {
62 "pathway": "KEGG_2021_Human",
63 "ontology": "GO_Biological_Process_2021",
64 "transcription": "ChEA_2016",
65 "diseases_drugs": "GWAS_Catalog_2019",
66 "celltypes": "PanglaoDB_Augmented_2021",
67 }
68 else:
69 # Stable local output names; library identifiers are remote input, not paths.
70 databases = {f"library_{i+1}": name for i, name in enumerate(databases)}
72 results = {}
74 for db_key, db_name in databases.items():
75 print(f"\nAnalyzing: {db_name}")
76 print("-" * 60)
78 try:
79 enrichment = gget.enrichr(
80 gene_list,
81 database=db_name,
82 species=species,
83 background_list=background,
84 plot=plot,
85 )
87 if enrichment is not None and len(enrichment) > 0:
88 # Save results
89 output_file = f"{output_prefix}_{db_key}.csv"
90 enrichment.to_csv(output_file, index=False)
91 print(f"Results saved to: {output_file}")
93 # Show top 5 results
94 print(f"\nTop 5 enriched terms:")
95 for i, row in enrichment.head(5).iterrows():
96 term = row["path_name"]
97 p_val = row["adj_p_val"]
98 print(f" {i+1}. {term}")
99 print(f" Adjusted P-value: {p_val:.2e}")
101 results[db_key] = enrichment
102 else:
103 print("No enrichment rows returned; significance has not been assessed")
105 except Exception as e:
106 print(f"Error: {e}")
108 # Generate summary report
109 print("\n" + "=" * 60)
110 print("Generating summary report...")
112 summary = []
113 for db_key, db_name in databases.items():
114 if db_key in results and len(results[db_key]) > 0:
115 summary.append(
116 {
117 "Database": db_name,
118 "Total Terms": len(results[db_key]),
119 "Top Term": results[db_key].iloc[0]["path_name"],
120 "Top Adjusted P-value": results[db_key].iloc[0]["adj_p_val"],
121 }
122 )
124 if summary:
125 summary_df = pd.DataFrame(summary)
126 summary_file = f"{output_prefix}_summary.csv"
127 summary_df.to_csv(summary_file, index=False)
128 print(f"\nSummary saved to: {summary_file}")
129 print("\n" + summary_df.to_string(index=False))
130 else:
131 print("\nNo enrichment results to summarize")
133 # Get expression data for genes
134 print("\n" + "=" * 60)
135 print("Getting expression data for input genes...")
137 try:
138 # Get tissue expression for first few genes
139 expr_data = []
140 for gene in gene_list[:5] if species in {"human", "mouse"} else []:
141 print(f" Getting expression for {gene}...")
142 try:
143 tissue_expr = gget.archs4(gene, which="tissue", species=species)
144 top_tissue = tissue_expr.nlargest(1, "median").iloc[0]
145 expr_data.append(
146 {
147 "Gene": gene,
148 "Top Tissue": top_tissue["id"],
149 "Median Expression": top_tissue["median"],
150 }
151 )
152 except Exception as e:
153 print(f" Warning: {e}")
155 if expr_data:
156 expr_df = pd.DataFrame(expr_data)
157 expr_file = f"{output_prefix}_expression.csv"
158 expr_df.to_csv(expr_file, index=False)
159 print(f"\nExpression data saved to: {expr_file}")
161 except Exception as e:
162 print(f"Error getting expression data: {e}")
164 print("\n" + "=" * 60)
165 print("Enrichment analysis complete!")
166 print(f"\nOutput files (prefix: {output_prefix}):")
167 for db_key in databases.keys():
168 if db_key in results:
169 print(f" - {output_prefix}_{db_key}.csv")
170 print(f" - {output_prefix}_summary.csv")
171 print(f" - {output_prefix}_expression.csv")
173 return bool(results)
176def main():
177 parser = argparse.ArgumentParser(
178 description="Perform comprehensive enrichment analysis using gget"
179 )
180 parser.add_argument(
181 "genes",
182 help="Gene list file (one gene per line or CSV with genes in first column)",
183 )
184 parser.add_argument(
185 "-s",
186 "--species",
187 default="human",
188 help="Species (human, mouse, fly, yeast, worm, fish)",
189 )
190 parser.add_argument(
191 "-b", "--background", help="Background gene list file (optional)"
192 )
193 parser.add_argument(
194 "-o", "--output", default="enrichment", help="Output prefix (default: enrichment)"
195 )
196 parser.add_argument(
197 "--no-plot", action="store_true", help="Disable plotting"
198 )
199 parser.add_argument(
200 "--database", action="append", help="Full Enrichr library name; repeat for multiple libraries"
201 )
203 args = parser.parse_args()
205 # Read gene list
206 if not Path(args.genes).exists():
207 print(f"Error: File not found: {args.genes}")
208 sys.exit(1)
210 try:
211 gene_list = read_gene_list(args.genes)
212 print(f"Read {len(gene_list)} genes from {args.genes}")
214 # Read background if provided
215 background = None
216 if args.background:
217 if Path(args.background).exists():
218 background = read_gene_list(args.background)
219 print(f"Read {len(background)} background genes from {args.background}")
220 else:
221 raise FileNotFoundError(f"Background file not found: {args.background}")
223 success = enrichment_pipeline(
224 gene_list,
225 species=args.species,
226 background=background,
227 output_prefix=args.output,
228 plot=not args.no_plot,
229 databases=args.database,
230 )
232 sys.exit(0 if success else 1)
234 except KeyboardInterrupt:
235 print("\n\nAnalysis interrupted by user")
236 sys.exit(1)
237 except Exception as e:
238 print(f"\n\nError: {e}")
239 import traceback
241 traceback.print_exc()
242 sys.exit(1)
245if __name__ == "__main__":
246 main()