From e5e0d499c21838d88fa9427945fca36823472741 Mon Sep 17 00:00:00 2001 From: OlgaPushkarev Date: Thu, 8 Oct 2026 17:08:43 -0700 Subject: [PATCH 1/5] Replace study-specific cell type in --cell_type help example Co-Authored-By: Claude Opus 5.5 (1M context) --- .../C_Annotation/0.Prepare/Extract_program_information.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py b/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py index 7847982..228a952 100644 --- a/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py +++ b/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py @@ -284,7 +284,7 @@ def build_parser(): # context info p.add_argument("--programs", type=int, nargs="+", help="Program ids, space separated (e.g. 1 2 3). Default: all.") - p.add_argument("--cell_type", default="", help="Cell type, e.g. 'teloHAEC aortic endothelial cell'.") + p.add_argument("--cell_type", default="", help="Cell type, e.g. 'aortic endothelial cell'.") p.add_argument("--organism", default="human") # select top items of the program From e77572acd6e6c5ff1e02e85572b62d03419eb5ac Mon Sep 17 00:00:00 2001 From: OlgaPushkarev Date: Thu, 8 Oct 2026 17:09:05 -0700 Subject: [PATCH 2/5] Format all Python files with ruff (format-only, no behavior change) Applies 'ruff format' (default black-compatible style, line length 88) to every tracked .py file so later changes have small, reviewable diffs. Notebooks are untouched. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../h5mu-structure/scripts/inspect_h5mu.py | 44 +- .../scripts/add_loadings_to_h5mu.py | 117 +- .../scripts/add_umap_to_h5mu.py | 102 +- .../scripts/build_cnmf_h5mu_from_usages.py | 287 ++- .../scripts/convert_igvf_mudata_to_h5ad.py | 145 +- .../scripts/generate_h5mu_structure.py | 62 +- .../scripts/generate_slurm.py | 414 ++-- .../scripts/prepare_guide_data.py | 43 +- .../scripts/validate_data.py | 210 +- .../scripts/check-pipeline-drift.py | 56 +- .../sk-cNMF_batch_inference_pipeline.py | 398 +++- src/Stage1_Inference/src/__init__.py | 21 +- src/Stage1_Inference/src/plot_diagnostics.py | 58 +- src/Stage1_Inference/src/run_cNMF.py | 414 ++-- .../torch_cnmf_inference_pipeline.py | 410 ++-- .../Slurm_Version/cNMF_evaluation_pipeline.py | 500 +++-- .../A_Metrics/src/Utilities.py | 35 +- .../A_Metrics/src/__init__.py | 16 +- .../A_Metrics/src/association_categorical.py | 572 +++--- .../A_Metrics/src/association_perturbation.py | 244 ++- .../A_Metrics/src/enrichment_geneset.py | 273 +-- .../A_Metrics/src/enrichment_motif.py | 489 +++-- .../A_Metrics/src/enrichment_trait.py | 241 ++- .../A_Metrics/src/explained_variance.py | 102 +- .../B_Calibration/Slurm_version/CRT/CRT.py | 426 ++-- .../U-test_perturbation_calibration.py | 218 ++- .../B_Calibration/src/CRT/ntc_groups.py | 12 +- .../B_Calibration/src/CRT/pipeline.py | 4 +- .../B_Calibration/src/CRT/pipeline_helpers.py | 8 +- .../B_Calibration/src/CRT/qq_plot.py | 36 +- .../B_Calibration/src/U_test/calibration.py | 293 ++- .../B_Calibration/src/check_covariate.py | 49 +- .../Slurm_Version/cNMF_k_selection.py | 311 ++- .../cNMF_perturbed_gene_analysis.py | 420 +++- .../Slurm_Version/cNMF_program_analysis.py | 305 ++- .../A_Plotting/src/Perturbed_gene_QC_plots.py | 1710 +++++++++++------ .../A_Plotting/src/Program_QC_plots.py | 1328 ++++++++----- .../src/Program_expression_weighted_plots.py | 289 ++- .../A_Plotting/src/__init__.py | 88 +- .../A_Plotting/src/efficient_corr.py | 19 +- .../src/html_Perturbed_gene_QC_plots.py | 514 +++-- .../A_Plotting/src/html_Program_QC_plots.py | 628 ++++-- .../A_Plotting/src/k_selection_plots.py | 1197 ++++++++---- .../A_Plotting/src/utilities.py | 143 +- .../Slurm_Version/cNMF_excel_summary.py | 435 +++-- .../src/Compile_excel_sheet.py | 1235 +++++++----- .../B_Summarization/src/__init__.py | 20 +- .../0.Prepare/Extract_program_information.py | 288 ++- .../search_gene/search_MyGene.py | 107 +- .../search_gene/search_NCBI.py | 110 +- .../search_gene/search_UniProt.py | 109 +- .../1.0.Search_database/search_gene/src.py | 61 +- .../search_OmniPath.py | 372 +++- .../1.1.Search_literature/search_arxiv.py | 81 +- .../1.1.Search_literature/search_crossref.py | 119 +- .../1.1.Search_literature/search_generif.py | 215 ++- .../1.1.Search_literature/search_nature.py | 118 +- .../1.1.Search_literature/search_ncbi.py | 178 +- .../1.1.Search_literature/search_openalex.py | 124 +- .../1.1.Search_literature/search_pubmed.py | 67 +- .../1.Search/1.2.Search_web/search_bing.py | 55 +- .../1.2.Search_web/search_duckduckgo.py | 28 +- .../1.Search/1.2.Search_web/search_google.py | 65 +- .../1.2.Search_web/search_uptodate.py | 78 +- .../1.Search/1.2.Search_web/search_wiki.py | 75 +- .../C_Annotation/1.Search/llm_query_agent.py | 263 ++- .../1.Search/llm_summary_agent.py | 181 +- .../C_Annotation/1.Search/pdf_download.py | 38 +- .../1.Search/run_literature_search.py | 393 +++- .../C_Annotation/1.Search/src.py | 27 +- .../2.Evidence_curation/curate_evidence.py | 239 ++- .../2.Evidence_curation/evidence_cache.py | 1 + .../2.Evidence_curation/gpe_export.py | 27 +- .../C_Annotation/2.Evidence_curation/qa.py | 96 +- .../2.Evidence_curation/run_qa.py | 57 +- .../C_Annotation/3.Generate/html_report.py | 330 +++- .../Slurm_Version/run_literature_search.py | 138 +- .../Literature_search/src/__init__.py | 8 +- .../src/literature_mining.py | 154 +- .../Literature_search/src/llm_backend.py | 47 +- .../Literature_search/src/output_writer.py | 92 +- .../Literature_search/src/paper_fetcher.py | 10 +- .../Literature_search/src/preprocess_input.py | 87 +- .../Literature_search/src/query_generator.py | 65 +- .../Literature_search/src/search_engine.py | 9 +- .../Literature_search/src/verification.py | 34 +- .../Slurm_Version/run_annotation.py | 1 + .../src/01_genes_to_string_enrichment.py | 413 ++-- .../ProgramExplorer/src/02_fetch_ncbi_data.py | 1006 ++++++---- .../src/03_submit_and_monitor_batch.py | 510 ++--- .../src/04_parse_and_summarize.py | 104 +- .../src/05_generate_html_report.py | 197 +- .../ProgramExplorer/src/batch_utils.py | 7 +- .../ProgramExplorer/src/column_mapper.py | 329 ++-- .../ProgramExplorer/src/compile_regulators.py | 14 +- .../src/convert_spectra_to_csv.py | 16 +- .../ProgramExplorer/src/gene_extraction.py | 220 ++- .../ProgramExplorer/src/harmonizome_api.py | 96 +- .../ProgramExplorer/src/html_report.py | 143 +- .../ProgramExplorer/src/literature_mining.py | 455 +++-- .../ProgramExplorer/src/llm_submission.py | 126 +- .../ProgramExplorer/src/ncbi_api.py | 100 +- .../ProgramExplorer/src/pipeline_state.py | 4 +- .../ProgramExplorer/src/prompt_builder.py | 60 +- .../ProgramExplorer/src/result_parser.py | 50 +- .../ProgramExplorer/src/run_pipeline.py | 358 ++-- .../ProgramExplorer/src/string_api.py | 263 +-- .../ProgramExplorer/src/string_enrichment.py | 25 +- .../C_Annotation/api/LLM_interface.py | 33 +- tests/Script/Stage1_Inference/conftest.py | 13 +- .../Stage1_Inference/test_edge_cases.py | 50 +- .../test_inference_parallel_sklearn.py | 35 +- .../test_inference_parallel_torch.py | 35 +- .../test_inference_sklearn.py | 32 +- .../Stage1_Inference/test_inference_torch.py | 40 +- tests/Script/Stage2_Evaluation/conftest.py | 29 +- tests/Script/Stage2_Evaluation/test_crt.py | 130 +- .../Script/Stage2_Evaluation/test_metrics.py | 360 +++- tests/Script/Stage2_Evaluation/test_utest.py | 436 ++++- .../A_Plotting/Gene/conftest.py | 45 +- .../Gene/test_perturbed_gene_qc_plots.py | 120 +- .../A_Plotting/KSelection/conftest.py | 90 +- .../KSelection/test_k_selection_plots.py | 146 +- .../A_Plotting/Program/conftest.py | 89 +- .../test_program_expression_weighted_plots.py | 44 +- .../Program/test_program_qc_plots.py | 116 +- .../A_Plotting/gene_html/conftest.py | 28 +- .../A_Plotting/gene_html/test_html_gene.py | 95 +- .../gene_html/test_html_gene_no_perturb.py | 71 +- .../A_Plotting/program_html/conftest.py | 33 +- .../program_html/test_html_program.py | 67 +- .../test_html_program_no_perturb.py | 108 +- .../B_Summarization/conftest.py | 67 +- .../B_Summarization/test_summary_table.py | 120 +- tests/Script/create_mini_dataset.py | 5 +- tools/check_no_lab_specific_content.py | 43 +- 136 files changed, 17913 insertions(+), 8551 deletions(-) diff --git a/.claude/skills/h5mu-structure/scripts/inspect_h5mu.py b/.claude/skills/h5mu-structure/scripts/inspect_h5mu.py index a9d46cd..bdb6631 100644 --- a/.claude/skills/h5mu-structure/scripts/inspect_h5mu.py +++ b/.claude/skills/h5mu-structure/scripts/inspect_h5mu.py @@ -45,7 +45,11 @@ def get_unique_examples(data, n=3): return ", ".join(_fmt_value(x) for x in uniques) elif sp.issparse(data): flat = data.toarray().flatten() - uniques = np.unique(flat[~np.isnan(flat)] if np.issubdtype(flat.dtype, np.floating) else flat)[:n] + uniques = np.unique( + flat[~np.isnan(flat)] + if np.issubdtype(flat.dtype, np.floating) + else flat + )[:n] return ", ".join(_fmt_value(x) for x in uniques) elif isinstance(data, np.ndarray): flat = data.flatten() @@ -182,7 +186,7 @@ def format_mapping_section(name, mapping, indent="", is_last_section=False): lines.append(tree_line(is_last_section, f"{name}/", indent)) child_indent = tree_child_indent(is_last_section, indent) for i, key in enumerate(items): - is_last_item = (i == len(items) - 1) + is_last_item = i == len(items) - 1 val = mapping[key] if name in ("uns",): desc = describe_uns_value(val) @@ -207,9 +211,11 @@ def format_adata(adata, mod_name=None): obs_names_examples = ", ".join(str(x) for x in adata.obs_names[:3]) all_obs_items = ["obs_names"] + obs_cols for i, col in enumerate(all_obs_items): - is_last = (i == len(all_obs_items) - 1) + is_last = i == len(all_obs_items) - 1 if col == "obs_names": - lines.append(tree_line(is_last, f"obs_names: e.g. {obs_names_examples}", obs_indent)) + lines.append( + tree_line(is_last, f"obs_names: e.g. {obs_names_examples}", obs_indent) + ) else: desc = describe_column(adata.obs[col]) lines.append(tree_line(is_last, f"{col} — {desc}", obs_indent)) @@ -221,9 +227,11 @@ def format_adata(adata, mod_name=None): var_names_examples = ", ".join(str(x) for x in adata.var_names[:3]) all_var_items = ["var_names"] + var_cols for i, col in enumerate(all_var_items): - is_last = (i == len(all_var_items) - 1) + is_last = i == len(all_var_items) - 1 if col == "var_names": - lines.append(tree_line(is_last, f"var_names: e.g. {var_names_examples}", var_indent)) + lines.append( + tree_line(is_last, f"var_names: e.g. {var_names_examples}", var_indent) + ) else: desc = describe_column(adata.var[col]) lines.append(tree_line(is_last, f"{col} — {desc}", var_indent)) @@ -238,8 +246,12 @@ def format_adata(adata, mod_name=None): ("varp", adata.varp), ] for i, (sec_name, sec_data) in enumerate(sections): - is_last = (i == len(sections) - 1) - lines.extend(format_mapping_section(sec_name, sec_data, indent="", is_last_section=is_last)) + is_last = i == len(sections) - 1 + lines.extend( + format_mapping_section( + sec_name, sec_data, indent="", is_last_section=is_last + ) + ) return lines @@ -278,7 +290,9 @@ def format_mudata(mdata, filename): obs_indent = tree_child_indent(False) for i, col in enumerate(obs_cols): desc = describe_column(mdata.obs[col]) - lines.append(tree_line(i == len(obs_cols) - 1, f"{col} — {desc}", obs_indent)) + lines.append( + tree_line(i == len(obs_cols) - 1, f"{col} — {desc}", obs_indent) + ) else: lines.append(tree_line(False, "obs columns: (none beyond obs_names)")) @@ -289,7 +303,9 @@ def format_mudata(mdata, filename): var_indent = tree_child_indent(False) for i, col in enumerate(var_cols): desc = describe_column(mdata.var[col]) - lines.append(tree_line(i == len(var_cols) - 1, f"{col} — {desc}", var_indent)) + lines.append( + tree_line(i == len(var_cols) - 1, f"{col} — {desc}", var_indent) + ) else: lines.append(tree_line(False, "var columns: (none beyond var_names)")) @@ -302,8 +318,12 @@ def format_mudata(mdata, filename): ("varp", mdata.varp), ] for i, (sec_name, sec_data) in enumerate(sections): - is_last = (i == len(sections) - 1) - lines.extend(format_mapping_section(sec_name, sec_data, indent="", is_last_section=is_last)) + is_last = i == len(sections) - 1 + lines.extend( + format_mapping_section( + sec_name, sec_data, indent="", is_last_section=is_last + ) + ) return "\n".join(lines) + "\n" diff --git a/.claude/skills/perturbNMF-runner/scripts/add_loadings_to_h5mu.py b/.claude/skills/perturbNMF-runner/scripts/add_loadings_to_h5mu.py index 5a91775..8743ac8 100644 --- a/.claude/skills/perturbNMF-runner/scripts/add_loadings_to_h5mu.py +++ b/.claude/skills/perturbNMF-runner/scripts/add_loadings_to_h5mu.py @@ -36,83 +36,108 @@ def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--h5mu', required=True, help='Existing .h5mu to patch in place') - p.add_argument('--gene_spectra_score', required=True, - help='cNMF gene_spectra_score TSV (K x genes)') - p.add_argument('--prog_key', default='cNMF', help='Program modality to patch') - p.add_argument('--data_key', default='rna', - help='Gene modality to validate the spectra width against ' - '(skipped if absent)') - p.add_argument('--overwrite', action='store_true', - help="Replace varm['loadings'] / uns['var_names'] if already present") - p.add_argument('--dry_run', action='store_true', - help='Validate and report, but write nothing') + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--h5mu", required=True, help="Existing .h5mu to patch in place") + p.add_argument( + "--gene_spectra_score", + required=True, + help="cNMF gene_spectra_score TSV (K x genes)", + ) + p.add_argument("--prog_key", default="cNMF", help="Program modality to patch") + p.add_argument( + "--data_key", + default="rna", + help="Gene modality to validate the spectra width against (skipped if absent)", + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Replace varm['loadings'] / uns['var_names'] if already present", + ) + p.add_argument( + "--dry_run", action="store_true", help="Validate and report, but write nothing" + ) return p.parse_args() def main(): args = parse_args() - print(f'Reading gene spectra score: {args.gene_spectra_score}', flush=True) - loadings = pd.read_csv(args.gene_spectra_score, sep='\t', index_col=0) + print(f"Reading gene spectra score: {args.gene_spectra_score}", flush=True) + loadings = pd.read_csv(args.gene_spectra_score, sep="\t", index_col=0) loadings.index = loadings.index.astype(str) - print(f' loadings: {loadings.shape[0]} programs x {loadings.shape[1]} genes', - flush=True) + print( + f" loadings: {loadings.shape[0]} programs x {loadings.shape[1]} genes", + flush=True, + ) - mode = 'r' if args.dry_run else 'r+' + mode = "r" if args.dry_run else "r+" with h5py.File(args.h5mu, mode) as f: - mod = f'mod/{args.prog_key}' + mod = f"mod/{args.prog_key}" if mod not in f: - raise ValueError(f"modality '{args.prog_key}' not found in {args.h5mu} " - f"(have: {list(f['mod'].keys())})") + raise ValueError( + f"modality '{args.prog_key}' not found in {args.h5mu} " + f"(have: {list(f['mod'].keys())})" + ) - prog_ids = read_elem(f[f'{mod}/var']).index.astype(str) + prog_ids = read_elem(f[f"{mod}/var"]).index.astype(str) if set(prog_ids) != set(loadings.index): raise ValueError( - f'program ids differ between the h5mu ({len(prog_ids)}) and the ' - f'spectra score ({loadings.shape[0]}): e.g. missing from spectra ' - f'{sorted(set(prog_ids) - set(loadings.index))[:5]}') + f"program ids differ between the h5mu ({len(prog_ids)}) and the " + f"spectra score ({loadings.shape[0]}): e.g. missing from spectra " + f"{sorted(set(prog_ids) - set(loadings.index))[:5]}" + ) # Row i of varm must be program var_names[i]; reindex rather than trust order. loadings = loadings.loc[prog_ids] - data_var = f'mod/{args.data_key}/var' + data_var = f"mod/{args.data_key}/var" if data_var in f: n_genes = read_elem(f[data_var]).shape[0] if n_genes != loadings.shape[1]: raise ValueError( f"the spectra score has {loadings.shape[1]} genes but modality " - f"'{args.data_key}' has {n_genes} -- Stage 2a asserts these match") + f"'{args.data_key}' has {n_genes} -- Stage 2a asserts these match" + ) print(f" validated against '{args.data_key}': {n_genes} genes", flush=True) else: - print(f" note: no '{args.data_key}' modality, skipping the gene-count check", - flush=True) + print( + f" note: no '{args.data_key}' modality, skipping the gene-count check", + flush=True, + ) - targets = [(f'{mod}/varm', 'loadings'), (f'{mod}/uns', 'var_names')] - existing = [f'{g}/{k}' for g, k in targets if f'{g}/{k}' in f] + targets = [(f"{mod}/varm", "loadings"), (f"{mod}/uns", "var_names")] + existing = [f"{g}/{k}" for g, k in targets if f"{g}/{k}" in f] if existing and not args.overwrite: - raise ValueError(f'already present: {existing}. Pass --overwrite to replace.') + raise ValueError( + f"already present: {existing}. Pass --overwrite to replace." + ) if args.dry_run: - print('Dry run: validation passed, nothing written.', flush=True) + print("Dry run: validation passed, nothing written.", flush=True) return 0 for group, key in targets: - if f'{group}/{key}' in f: - print(f' overwriting {group}/{key}', flush=True) - del f[f'{group}/{key}'] - - print(f'Writing {mod}/varm/loadings and {mod}/uns/var_names', flush=True) - write_elem(f[f'{mod}/varm'], 'loadings', - np.asarray(loadings.values, dtype=np.float64)) - write_elem(f[f'{mod}/uns'], 'var_names', - loadings.columns.astype(str).to_numpy()) - - print(f'Done. {args.h5mu} now carries loadings for {loadings.shape[0]} programs ' - f'x {loadings.shape[1]} genes.', flush=True) + if f"{group}/{key}" in f: + print(f" overwriting {group}/{key}", flush=True) + del f[f"{group}/{key}"] + + print(f"Writing {mod}/varm/loadings and {mod}/uns/var_names", flush=True) + write_elem( + f[f"{mod}/varm"], "loadings", np.asarray(loadings.values, dtype=np.float64) + ) + write_elem( + f[f"{mod}/uns"], "var_names", loadings.columns.astype(str).to_numpy() + ) + + print( + f"Done. {args.h5mu} now carries loadings for {loadings.shape[0]} programs " + f"x {loadings.shape[1]} genes.", + flush=True, + ) return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/.claude/skills/perturbNMF-runner/scripts/add_umap_to_h5mu.py b/.claude/skills/perturbNMF-runner/scripts/add_umap_to_h5mu.py index 207c0e6..e060984 100755 --- a/.claude/skills/perturbNMF-runner/scripts/add_umap_to_h5mu.py +++ b/.claude/skills/perturbNMF-runner/scripts/add_umap_to_h5mu.py @@ -37,32 +37,48 @@ # Same import convention as the plotting entry points, but resolved relative to this # file rather than hardcoded, since this script sits under .claude/skills/. _HERE = os.path.dirname(os.path.abspath(__file__)) -_PIPELINE_SRC = os.path.abspath(os.path.join(_HERE, '..', '..', '..', '..', 'src')) +_PIPELINE_SRC = os.path.abspath(os.path.join(_HERE, "..", "..", "..", "..", "src")) sys.path.append(_PIPELINE_SRC) from Stage3_Interpretation.A_Plotting.src import ensure_umap # noqa: E402 def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--mdata_path', required=True, help='Input .h5mu') - p.add_argument('--output', default=None, - help='Write the result here. Mutually exclusive with --in_place.') - p.add_argument('--in_place', action='store_true', - help='Overwrite --mdata_path (via a .tmp file + atomic replace)') - p.add_argument('--data_key', default='rna', help='Expression modality') - p.add_argument('--prog_key', default='cNMF', help='Program modality') - p.add_argument('--n_top_genes', type=int, default=2000, - help="Top-variance genes used as the UMAP basis. Must match the " - "plotting scripts' ensure_umap default (2000) or they will " - 'recompute instead of reusing this.') - p.add_argument('--n_comps', type=int, default=50, help='PCA components (default 50)') - p.add_argument('--force', action='store_true', - help='Recompute even if X_umap is already present in both modalities') + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--mdata_path", required=True, help="Input .h5mu") + p.add_argument( + "--output", + default=None, + help="Write the result here. Mutually exclusive with --in_place.", + ) + p.add_argument( + "--in_place", + action="store_true", + help="Overwrite --mdata_path (via a .tmp file + atomic replace)", + ) + p.add_argument("--data_key", default="rna", help="Expression modality") + p.add_argument("--prog_key", default="cNMF", help="Program modality") + p.add_argument( + "--n_top_genes", + type=int, + default=2000, + help="Top-variance genes used as the UMAP basis. Must match the " + "plotting scripts' ensure_umap default (2000) or they will " + "recompute instead of reusing this.", + ) + p.add_argument( + "--n_comps", type=int, default=50, help="PCA components (default 50)" + ) + p.add_argument( + "--force", + action="store_true", + help="Recompute even if X_umap is already present in both modalities", + ) args = p.parse_args() if bool(args.output) == bool(args.in_place): - p.error('pass exactly one of --output or --in_place') + p.error("pass exactly one of --output or --in_place") return args @@ -70,51 +86,63 @@ def main(): args = parse_args() t0 = time.time() - print(f'Reading {args.mdata_path}', flush=True) + print(f"Reading {args.mdata_path}", flush=True) mdata = mu.read_h5mu(args.mdata_path) - print(f' modalities: {list(mdata.mod)}', flush=True) + print(f" modalities: {list(mdata.mod)}", flush=True) for key in (args.data_key, args.prog_key): if key not in mdata.mod: raise KeyError( f"modality '{key}' not in {list(mdata.mod)}. The expression matrix is " - 'required -- a cNMF-only h5mu must be rebuilt with ' - '`build_cnmf_h5mu_from_usages.py --include_rna --tpm_rna` first.') + "required -- a cNMF-only h5mu must be rebuilt with " + "`build_cnmf_h5mu_from_usages.py --include_rna --tpm_rna` first." + ) - present = ['X_umap' in mdata[k].obsm for k in (args.data_key, args.prog_key)] + present = ["X_umap" in mdata[k].obsm for k in (args.data_key, args.prog_key)] if all(present) and not args.force: - print('X_umap already present in both modalities; nothing to do ' - '(pass --force to recompute).', flush=True) + print( + "X_umap already present in both modalities; nothing to do " + "(pass --force to recompute).", + flush=True, + ) return 0 if args.force: for k in (args.data_key, args.prog_key): - for slot in ('X_umap', 'X_pca'): + for slot in ("X_umap", "X_pca"): mdata[k].obsm.pop(slot, None) n_obs = mdata[args.data_key].n_obs - print(f'Computing PCA({args.n_comps}) + neighbors + UMAP on {n_obs} cells from the ' - f'top {args.n_top_genes} variance genes', flush=True) - ensure_umap(mdata, args.data_key, args.prog_key, - n_top_genes=args.n_top_genes, n_comps=args.n_comps) - print(f' done in {time.time() - t0:.0f}s', flush=True) + print( + f"Computing PCA({args.n_comps}) + neighbors + UMAP on {n_obs} cells from the " + f"top {args.n_top_genes} variance genes", + flush=True, + ) + ensure_umap( + mdata, + args.data_key, + args.prog_key, + n_top_genes=args.n_top_genes, + n_comps=args.n_comps, + ) + print(f" done in {time.time() - t0:.0f}s", flush=True) for k in (args.data_key, args.prog_key): - if 'X_umap' not in mdata[k].obsm: + if "X_umap" not in mdata[k].obsm: raise RuntimeError(f"ensure_umap did not populate X_umap on '{k}'") print(f" {k}.obsm['X_umap']: {mdata[k].obsm['X_umap'].shape}", flush=True) target = args.output or args.mdata_path - tmp = f'{target}.tmp' if args.in_place else target + tmp = f"{target}.tmp" if args.in_place else target os.makedirs(os.path.dirname(os.path.abspath(target)), exist_ok=True) - print(f'Writing {tmp}', flush=True) + print(f"Writing {tmp}", flush=True) mdata.write(tmp) if args.in_place: - print(f'Replacing {target}', flush=True) + print(f"Replacing {target}", flush=True) os.replace(tmp, target) - print(f'Done in {time.time() - t0:.0f}s -> {target}', flush=True) + print(f"Done in {time.time() - t0:.0f}s -> {target}", flush=True) return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/.claude/skills/perturbNMF-runner/scripts/build_cnmf_h5mu_from_usages.py b/.claude/skills/perturbNMF-runner/scripts/build_cnmf_h5mu_from_usages.py index 3023f8d..59e66b5 100644 --- a/.claude/skills/perturbNMF-runner/scripts/build_cnmf_h5mu_from_usages.py +++ b/.claude/skills/perturbNMF-runner/scripts/build_cnmf_h5mu_from_usages.py @@ -50,115 +50,175 @@ def _xmax(X): """Max of a dense or sparse matrix, 0 for an all-zero sparse one.""" - data = getattr(X, 'data', X) + data = getattr(X, "data", X) return float(np.asarray(data).max()) if np.size(data) else 0.0 def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--usages', required=True, - help='Consensus usages TSV (cells x K), cell barcodes in the index') - p.add_argument('--gene_spectra_score', - help="cNMF gene_spectra_score TSV (K x genes), the z-scored gene " - "loadings -> varm['loadings'] + uns['var_names']. Stage 2a " - "enrichment and every Stage 3 gene-level plot read these; " - "without them plot_top_gene_per_program raises " - "KeyError: 'loadings'.") - p.add_argument('--raw_mudata', required=True, - help='Raw IGVF .h5mu with gene + guide modalities') - p.add_argument('--output', required=True, - help='Output .h5mu path (e.g. .../adata/cNMF_200_2_0.h5mu)') - p.add_argument('--gene_mod', default='gene', help='Modality holding expression + obs') - p.add_argument('--guide_mod', default='guide', help='Modality holding the guide matrix') - p.add_argument('--guide_layer', default='guide_assignment', - help='Layer in the guide modality holding the BINARY assignment matrix. ' - 'IGVF mudatas keep raw guide UMI counts in X and the called ' - 'assignment in this layer -- using X would treat ambient guide ' - 'reads as perturbations. Pass "X" to use X anyway.') - p.add_argument('--prog_key', default='cNMF', help='Name of the program modality to write') - p.add_argument('--guide_id_col', default='guide_id', - help="Column in guide .var with the guide id -> uns['guide_names']") - p.add_argument('--guide_target_col', default='gene_name', - help="Column in guide .var with the target -> uns['guide_targets']") - p.add_argument('--guide_type_col', default='type', - help='Column in guide .var flagging control guides') - p.add_argument('--control_type', nargs='*', default=['non-targeting'], - help='Value(s) of the type column whose targets get relabelled') - p.add_argument('--control_label', default='non-targeting', - help='Target label written for control guides') - p.add_argument('--condition_key', default='condition', - help='Name of the constant condition column added to obs') - p.add_argument('--condition_value', default='all', - help='Value of the constant condition column') - p.add_argument('--include_rna', action='store_true', - help='Also carry the gene modality over as an "rna" modality') - p.add_argument('--tpm_rna', action='store_true', - help='TPM-normalize the carried-over rna matrix, using the same ' - 'definition cNMF itself uses in compute_tpm() -- ' - 'normalize_total(target_sum=1e6), no log1p. Raw IGVF gene ' - 'matrices hold integer counts, which Stage 3b/3c reject via ' - 'check_normalized().') - p.add_argument('--data_key', default='rna', help='Modality name used with --include_rna') + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument( + "--usages", + required=True, + help="Consensus usages TSV (cells x K), cell barcodes in the index", + ) + p.add_argument( + "--gene_spectra_score", + help="cNMF gene_spectra_score TSV (K x genes), the z-scored gene " + "loadings -> varm['loadings'] + uns['var_names']. Stage 2a " + "enrichment and every Stage 3 gene-level plot read these; " + "without them plot_top_gene_per_program raises " + "KeyError: 'loadings'.", + ) + p.add_argument( + "--raw_mudata", + required=True, + help="Raw IGVF .h5mu with gene + guide modalities", + ) + p.add_argument( + "--output", + required=True, + help="Output .h5mu path (e.g. .../adata/cNMF_200_2_0.h5mu)", + ) + p.add_argument( + "--gene_mod", default="gene", help="Modality holding expression + obs" + ) + p.add_argument( + "--guide_mod", default="guide", help="Modality holding the guide matrix" + ) + p.add_argument( + "--guide_layer", + default="guide_assignment", + help="Layer in the guide modality holding the BINARY assignment matrix. " + "IGVF mudatas keep raw guide UMI counts in X and the called " + "assignment in this layer -- using X would treat ambient guide " + 'reads as perturbations. Pass "X" to use X anyway.', + ) + p.add_argument( + "--prog_key", default="cNMF", help="Name of the program modality to write" + ) + p.add_argument( + "--guide_id_col", + default="guide_id", + help="Column in guide .var with the guide id -> uns['guide_names']", + ) + p.add_argument( + "--guide_target_col", + default="gene_name", + help="Column in guide .var with the target -> uns['guide_targets']", + ) + p.add_argument( + "--guide_type_col", + default="type", + help="Column in guide .var flagging control guides", + ) + p.add_argument( + "--control_type", + nargs="*", + default=["non-targeting"], + help="Value(s) of the type column whose targets get relabelled", + ) + p.add_argument( + "--control_label", + default="non-targeting", + help="Target label written for control guides", + ) + p.add_argument( + "--condition_key", + default="condition", + help="Name of the constant condition column added to obs", + ) + p.add_argument( + "--condition_value", + default="all", + help="Value of the constant condition column", + ) + p.add_argument( + "--include_rna", + action="store_true", + help='Also carry the gene modality over as an "rna" modality', + ) + p.add_argument( + "--tpm_rna", + action="store_true", + help="TPM-normalize the carried-over rna matrix, using the same " + "definition cNMF itself uses in compute_tpm() -- " + "normalize_total(target_sum=1e6), no log1p. Raw IGVF gene " + "matrices hold integer counts, which Stage 3b/3c reject via " + "check_normalized().", + ) + p.add_argument( + "--data_key", default="rna", help="Modality name used with --include_rna" + ) return p.parse_args() def main(): args = parse_args() - print(f'Reading usages: {args.usages}', flush=True) - usages = pd.read_csv(args.usages, sep='\t', index_col=0) + print(f"Reading usages: {args.usages}", flush=True) + usages = pd.read_csv(args.usages, sep="\t", index_col=0) usages.index = usages.index.astype(str) - print(f' usages: {usages.shape[0]} cells x {usages.shape[1]} programs', flush=True) + print(f" usages: {usages.shape[0]} cells x {usages.shape[1]} programs", flush=True) loadings = None if args.gene_spectra_score: - print(f'Reading gene spectra score: {args.gene_spectra_score}', flush=True) - loadings = pd.read_csv(args.gene_spectra_score, sep='\t', index_col=0) + print(f"Reading gene spectra score: {args.gene_spectra_score}", flush=True) + loadings = pd.read_csv(args.gene_spectra_score, sep="\t", index_col=0) loadings.index = loadings.index.astype(str) - print(f' loadings: {loadings.shape[0]} programs x {loadings.shape[1]} genes', - flush=True) + print( + f" loadings: {loadings.shape[0]} programs x {loadings.shape[1]} genes", + flush=True, + ) # The spectra file and the usages file are both keyed by program id, but there # is no guarantee the two are written in the same order -- reindex rather than # trust it, so varm['loadings'] row i really is program var_names[i]. prog_ids = usages.columns.astype(str) if set(prog_ids) != set(loadings.index): raise ValueError( - f'program ids differ between usages ({len(prog_ids)}) and the spectra ' - f'score ({loadings.shape[0]}): e.g. missing from spectra ' - f'{sorted(set(prog_ids) - set(loadings.index))[:5]}') + f"program ids differ between usages ({len(prog_ids)}) and the spectra " + f"score ({loadings.shape[0]}): e.g. missing from spectra " + f"{sorted(set(prog_ids) - set(loadings.index))[:5]}" + ) loadings = loadings.loc[prog_ids] - print(f'Reading raw mudata (metadata only): {args.raw_mudata}', flush=True) - with h5py.File(args.raw_mudata, 'r') as f: - gene_obs = read_elem(f[f'mod/{args.gene_mod}/obs']) - guide_var = read_elem(f[f'mod/{args.guide_mod}/var']) + print(f"Reading raw mudata (metadata only): {args.raw_mudata}", flush=True) + with h5py.File(args.raw_mudata, "r") as f: + gene_obs = read_elem(f[f"mod/{args.gene_mod}/obs"]) + guide_var = read_elem(f[f"mod/{args.guide_mod}/var"]) # Prefer the binary assignment layer over X (raw guide UMI counts). - layer_path = f'mod/{args.guide_mod}/layers/{args.guide_layer}' - if args.guide_layer != 'X' and layer_path in f: + layer_path = f"mod/{args.guide_mod}/layers/{args.guide_layer}" + if args.guide_layer != "X" and layer_path in f: print(f" using guide layer '{args.guide_layer}'", flush=True) guide_X = read_elem(f[layer_path]) else: - print(f" WARNING: layer '{args.guide_layer}' not found, falling back to " - f'guide X -- verify it holds assignments, not raw UMI counts', flush=True) - guide_X = read_elem(f[f'mod/{args.guide_mod}/X']) + print( + f" WARNING: layer '{args.guide_layer}' not found, falling back to " + f"guide X -- verify it holds assignments, not raw UMI counts", + flush=True, + ) + guide_X = read_elem(f[f"mod/{args.guide_mod}/X"]) rna_X = rna_var = None if args.include_rna: - rna_X = read_elem(f[f'mod/{args.gene_mod}/X']) - rna_var = read_elem(f[f'mod/{args.gene_mod}/var']) + rna_X = read_elem(f[f"mod/{args.gene_mod}/X"]) + rna_var = read_elem(f[f"mod/{args.gene_mod}/var"]) gene_obs.index = gene_obs.index.astype(str) - print(f' obs: {gene_obs.shape[0]} cells x {gene_obs.shape[1]} columns', flush=True) - print(f' guides: {guide_X.shape[1]}', flush=True) + print(f" obs: {gene_obs.shape[0]} cells x {gene_obs.shape[1]} columns", flush=True) + print(f" guides: {guide_X.shape[1]}", flush=True) # --- align cells: cNMF usually drops a few cells during consensus --- missing = usages.index.difference(gene_obs.index) if len(missing): raise ValueError( - f'{len(missing)} cells in the usages file are absent from the mudata ' - f'(e.g. {list(missing[:5])}). The usages and mudata do not match.') + f"{len(missing)} cells in the usages file are absent from the mudata " + f"(e.g. {list(missing[:5])}). The usages and mudata do not match." + ) dropped = len(gene_obs.index) - len(usages.index) - print(f' aligning to the usages index ({dropped} mudata cells dropped by cNMF)', - flush=True) + print( + f" aligning to the usages index ({dropped} mudata cells dropped by cNMF)", + flush=True, + ) positions = gene_obs.index.get_indexer(usages.index) obs = gene_obs.iloc[positions].copy() @@ -168,11 +228,17 @@ def main(): # A guide-assignment matrix should be binary; raw UMI counts are a common mix-up. gmax = guide_assignment.data.max() if guide_assignment.nnz else 0 - print(f' guide assignment: {guide_assignment.nnz} nonzeros, max value {gmax}, ' - f'{guide_assignment.nnz / guide_assignment.shape[0]:.2f} guides/cell', flush=True) + print( + f" guide assignment: {guide_assignment.nnz} nonzeros, max value {gmax}, " + f"{guide_assignment.nnz / guide_assignment.shape[0]:.2f} guides/cell", + flush=True, + ) if gmax > 1: - print(' WARNING: assignment matrix is not binary (max > 1) -- this looks like ' - 'raw guide UMI counts, not called assignments', flush=True) + print( + " WARNING: assignment matrix is not binary (max > 1) -- this looks like " + "raw guide UMI counts, not called assignments", + flush=True, + ) # --- guide metadata --- if args.guide_id_col in guide_var.columns: @@ -181,23 +247,33 @@ def main(): guide_names = guide_var.index.astype(str).to_numpy() guide_targets = guide_var[args.guide_target_col].astype(object).to_numpy() - is_control = guide_var[args.guide_type_col].astype(str).isin(args.control_type).to_numpy() + is_control = ( + guide_var[args.guide_type_col].astype(str).isin(args.control_type).to_numpy() + ) guide_targets[is_control] = args.control_label - print(f' relabelled {int(is_control.sum())} control guides as ' - f'"{args.control_label}"', flush=True) + print( + f" relabelled {int(is_control.sum())} control guides as " + f'"{args.control_label}"', + flush=True, + ) # Anything still missing a target (NaN) would silently become the string "nan" # downstream, so surface it here instead. unlabelled = pd.isna(guide_targets) if unlabelled.any(): - print(f' WARNING: {int(unlabelled.sum())} guides have no target and will be ' - f'labelled "unassigned"', flush=True) - guide_targets[unlabelled] = 'unassigned' + print( + f" WARNING: {int(unlabelled.sum())} guides have no target and will be " + f'labelled "unassigned"', + flush=True, + ) + guide_targets[unlabelled] = "unassigned" guide_targets = guide_targets.astype(str) if len(guide_names) != guide_assignment.shape[1]: - raise ValueError(f'guide_names ({len(guide_names)}) does not match the guide ' - f'matrix width ({guide_assignment.shape[1]})') + raise ValueError( + f"guide_names ({len(guide_names)}) does not match the guide " + f"matrix width ({guide_assignment.shape[1]})" + ) # --- program modality --- prog = ad.AnnData( @@ -205,48 +281,53 @@ def main(): obs=obs, var=pd.DataFrame(index=usages.columns.astype(str)), ) - prog.obsm['guide_assignment'] = guide_assignment.tocsr() - prog.uns['guide_names'] = guide_names - prog.uns['guide_targets'] = guide_targets + prog.obsm["guide_assignment"] = guide_assignment.tocsr() + prog.uns["guide_names"] = guide_names + prog.uns["guide_targets"] = guide_targets if loadings is not None: # Stage 2a asserts loadings.shape[1] == mdata[data_key].var.shape[0], so the # spectra columns have to line up with the gene modality one-for-one. if rna_var is not None and loadings.shape[1] != rna_var.shape[0]: raise ValueError( - f'the spectra score has {loadings.shape[1]} genes but the ' - f'{args.gene_mod} modality has {rna_var.shape[0]} -- Stage 2a/3 ' - f'assume they are the same genes in the same order') - prog.varm['loadings'] = np.asarray(loadings.values, dtype=np.float64) - prog.uns['var_names'] = loadings.columns.astype(str).to_numpy() + f"the spectra score has {loadings.shape[1]} genes but the " + f"{args.gene_mod} modality has {rna_var.shape[0]} -- Stage 2a/3 " + f"assume they are the same genes in the same order" + ) + prog.varm["loadings"] = np.asarray(loadings.values, dtype=np.float64) + prog.uns["var_names"] = loadings.columns.astype(str).to_numpy() mods = {args.prog_key: prog} if args.include_rna: rna = ad.AnnData(X=rna_X[positions, :], obs=obs.copy(), var=rna_var) if args.tpm_rna: import scanpy as sc + # Same normalization cNMF applies in compute_tpm() (cnmf/cnmf.py), done # in place rather than via that helper so a ~19GB matrix is not copied. - print(f' TPM-normalizing rna (max before: {_xmax(rna.X):g})', flush=True) + print(f" TPM-normalizing rna (max before: {_xmax(rna.X):g})", flush=True) sc.pp.normalize_total(rna, target_sum=1e6) - print(f' rna TPM-normalized (max after: {_xmax(rna.X):g})', flush=True) - rna.obsm['guide_assignment'] = guide_assignment.tocsr() - rna.uns['guide_names'] = guide_names - rna.uns['guide_targets'] = guide_targets + print(f" rna TPM-normalized (max after: {_xmax(rna.X):g})", flush=True) + rna.obsm["guide_assignment"] = guide_assignment.tocsr() + rna.uns["guide_names"] = guide_names + rna.uns["guide_targets"] = guide_targets mods[args.data_key] = rna mdata = mu.MuData(mods) os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) - print(f'Writing {args.output}', flush=True) + print(f"Writing {args.output}", flush=True) mdata.write(args.output) n_ctrl = int((guide_targets == args.control_label).sum()) - print(f'Done. {args.prog_key}: {prog.shape[0]} cells x {prog.shape[1]} programs; ' - f'{len(np.unique(guide_targets))} unique targets; ' - f'{n_ctrl} control guides labelled "{args.control_label}".', flush=True) + print( + f"Done. {args.prog_key}: {prog.shape[0]} cells x {prog.shape[1]} programs; " + f"{len(np.unique(guide_targets))} unique targets; " + f'{n_ctrl} control guides labelled "{args.control_label}".', + flush=True, + ) return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/.claude/skills/perturbNMF-runner/scripts/convert_igvf_mudata_to_h5ad.py b/.claude/skills/perturbNMF-runner/scripts/convert_igvf_mudata_to_h5ad.py index 6ad97ce..44a1ebf 100644 --- a/.claude/skills/perturbNMF-runner/scripts/convert_igvf_mudata_to_h5ad.py +++ b/.claude/skills/perturbNMF-runner/scripts/convert_igvf_mudata_to_h5ad.py @@ -28,60 +28,92 @@ def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--input', required=True, help='Raw IGVF .h5mu') - p.add_argument('--output', required=True, help='Output .h5ad for --counts_fn') - p.add_argument('--gene_mod', default='gene', help='Expression modality') - p.add_argument('--guide_mod', default='guide', help='Guide modality') - p.add_argument('--guide_layer', default='guide_assignment', - help='Layer with the binary assignment; "X" to use guide X instead') - p.add_argument('--guide_id_col', default='guide_id', help='Guide id column in guide .var') - p.add_argument('--guide_target_col', default='gene_name', - help='Target column in guide .var') - p.add_argument('--guide_type_col', default='type', - help='Column flagging control guides (blank to skip relabelling)') - p.add_argument('--control_type', nargs='*', default=['non-targeting'], - help='Values of the type column treated as controls') - p.add_argument('--control_label', default='non-targeting', - help='Target label written for control guides') - p.add_argument('--condition_key', default='condition', - help='Name of the constant obs column to add (--categorical_key)') - p.add_argument('--condition_value', default='all', help='Value of that column') - p.add_argument('--gene_symbol_col', default='symbol', - help='Column in gene .var with symbols; var_names are set from it ' - 'and the original index is kept in var["gene_ids"]') - p.add_argument('--keep_ensembl_index', action='store_true', - help='Leave var_names as-is instead of switching to symbols') + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--input", required=True, help="Raw IGVF .h5mu") + p.add_argument("--output", required=True, help="Output .h5ad for --counts_fn") + p.add_argument("--gene_mod", default="gene", help="Expression modality") + p.add_argument("--guide_mod", default="guide", help="Guide modality") + p.add_argument( + "--guide_layer", + default="guide_assignment", + help='Layer with the binary assignment; "X" to use guide X instead', + ) + p.add_argument( + "--guide_id_col", default="guide_id", help="Guide id column in guide .var" + ) + p.add_argument( + "--guide_target_col", default="gene_name", help="Target column in guide .var" + ) + p.add_argument( + "--guide_type_col", + default="type", + help="Column flagging control guides (blank to skip relabelling)", + ) + p.add_argument( + "--control_type", + nargs="*", + default=["non-targeting"], + help="Values of the type column treated as controls", + ) + p.add_argument( + "--control_label", + default="non-targeting", + help="Target label written for control guides", + ) + p.add_argument( + "--condition_key", + default="condition", + help="Name of the constant obs column to add (--categorical_key)", + ) + p.add_argument("--condition_value", default="all", help="Value of that column") + p.add_argument( + "--gene_symbol_col", + default="symbol", + help="Column in gene .var with symbols; var_names are set from it " + 'and the original index is kept in var["gene_ids"]', + ) + p.add_argument( + "--keep_ensembl_index", + action="store_true", + help="Leave var_names as-is instead of switching to symbols", + ) return p.parse_args() def main(): args = parse_args() - print(f'Reading {args.input}', flush=True) + print(f"Reading {args.input}", flush=True) mdata = mu.read(args.input) adata = mdata[args.gene_mod].copy() guide = mdata[args.guide_mod] - print(f' gene: {adata.shape[0]} cells x {adata.shape[1]} genes', flush=True) - print(f' guide: {guide.shape[1]} guides', flush=True) + print(f" gene: {adata.shape[0]} cells x {adata.shape[1]} genes", flush=True) + print(f" guide: {guide.shape[1]} guides", flush=True) # --- guide assignment (binary layer, not raw UMI counts in X) --- - if args.guide_layer != 'X' and args.guide_layer in guide.layers: + if args.guide_layer != "X" and args.guide_layer in guide.layers: assignment = guide.layers[args.guide_layer].copy() print(f" using guide layer '{args.guide_layer}'", flush=True) else: assignment = guide.X.copy() - print(f" WARNING: layer '{args.guide_layer}' not found, using guide X -- " - f'verify it holds assignments, not raw UMI counts', flush=True) + print( + f" WARNING: layer '{args.guide_layer}' not found, using guide X -- " + f"verify it holds assignments, not raw UMI counts", + flush=True, + ) gmax = assignment.data.max() if assignment.nnz else 0 - print(f' assignment: {assignment.nnz} nonzeros, max {gmax}, ' - f'{assignment.nnz / assignment.shape[0]:.2f} guides/cell', flush=True) + print( + f" assignment: {assignment.nnz} nonzeros, max {gmax}, " + f"{assignment.nnz / assignment.shape[0]:.2f} guides/cell", + flush=True, + ) if gmax > 1: - print(' WARNING: assignment matrix is not binary (max > 1)', flush=True) + print(" WARNING: assignment matrix is not binary (max > 1)", flush=True) - adata.obsm['guide_assignment'] = assignment + adata.obsm["guide_assignment"] = assignment # --- guide names / targets --- if args.guide_id_col in guide.var.columns: @@ -91,39 +123,52 @@ def main(): guide_targets = guide.var[args.guide_target_col].astype(object).to_numpy() if args.guide_type_col and args.guide_type_col in guide.var.columns: - is_control = guide.var[args.guide_type_col].astype(str).isin( - args.control_type).to_numpy() + is_control = ( + guide.var[args.guide_type_col] + .astype(str) + .isin(args.control_type) + .to_numpy() + ) guide_targets[is_control] = args.control_label - print(f' relabelled {int(is_control.sum())} control guides as ' - f'"{args.control_label}"', flush=True) + print( + f" relabelled {int(is_control.sum())} control guides as " + f'"{args.control_label}"', + flush=True, + ) unlabelled = pd.isna(guide_targets) if unlabelled.any(): - print(f' WARNING: {int(unlabelled.sum())} guides have no target -> "unassigned"', - flush=True) - guide_targets[unlabelled] = 'unassigned' + print( + f' WARNING: {int(unlabelled.sum())} guides have no target -> "unassigned"', + flush=True, + ) + guide_targets[unlabelled] = "unassigned" - adata.uns['guide_names'] = list(np.asarray(guide_names, dtype=str)) - adata.uns['guide_targets'] = list(np.asarray(guide_targets, dtype=str)) + adata.uns["guide_names"] = list(np.asarray(guide_names, dtype=str)) + adata.uns["guide_targets"] = list(np.asarray(guide_targets, dtype=str)) # --- condition column used as --categorical_key --- adata.obs[args.condition_key] = pd.Categorical( - [args.condition_value] * adata.shape[0]) + [args.condition_value] * adata.shape[0] + ) # --- gene names --- if not args.keep_ensembl_index and args.gene_symbol_col in adata.var.columns: - adata.var['gene_ids'] = adata.var_names + adata.var["gene_ids"] = adata.var_names adata.var_names = adata.var[args.gene_symbol_col].astype(str) adata.var_names_make_unique() print(f' var_names set from var["{args.gene_symbol_col}"]', flush=True) os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) - print(f'Writing {args.output}', flush=True) + print(f"Writing {args.output}", flush=True) adata.write(args.output) - print(f'Done. {adata.shape[0]} cells x {adata.shape[1]} genes, ' - f'{len(set(adata.uns["guide_targets"]))} unique targets.', flush=True) + print( + f"Done. {adata.shape[0]} cells x {adata.shape[1]} genes, " + f"{len(set(adata.uns['guide_targets']))} unique targets.", + flush=True, + ) return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/.claude/skills/perturbNMF-runner/scripts/generate_h5mu_structure.py b/.claude/skills/perturbNMF-runner/scripts/generate_h5mu_structure.py index c3fc9e3..4860e6b 100755 --- a/.claude/skills/perturbNMF-runner/scripts/generate_h5mu_structure.py +++ b/.claude/skills/perturbNMF-runner/scripts/generate_h5mu_structure.py @@ -23,7 +23,7 @@ def format_shape_dtype(arr): return f"sparse {fmt}, shape={arr.shape}, dtype={arr.dtype}" elif isinstance(arr, np.ndarray): return f"dense ndarray, shape={arr.shape}, dtype={arr.dtype}" - elif hasattr(arr, 'shape') and hasattr(arr, 'dtype'): + elif hasattr(arr, "shape") and hasattr(arr, "dtype"): return f"{type(arr).__name__}, shape={arr.shape}, dtype={arr.dtype}" else: return f"{type(arr).__name__}" @@ -40,13 +40,13 @@ def write_slot(lines, name, mapping, prefix, is_last): connector = "\u2514" if is_last else "\u251c" continuation = " " if is_last else "\u2502" - items = list(mapping.keys()) if hasattr(mapping, 'keys') else [] + items = list(mapping.keys()) if hasattr(mapping, "keys") else [] if not items: lines.append(f"{prefix}{connector}\u2500\u2500 {name}/ (empty)") else: lines.append(f"{prefix}{connector}\u2500\u2500 {name}/") for i, key in enumerate(items): - is_last_item = (i == len(items) - 1) + is_last_item = i == len(items) - 1 item_connector = "\u2514" if is_last_item else "\u251c" val = mapping[key] desc = format_shape_dtype(val) @@ -55,7 +55,9 @@ def write_slot(lines, name, mapping, prefix, is_last): if isinstance(val, np.ndarray) and val.dtype == object and val.ndim == 1: examples = format_examples(val, n=5) example_str = f", e.g. {examples}" - lines.append(f"{prefix}{continuation} {item_connector}\u2500\u2500 {key} ({desc}{example_str})") + lines.append( + f"{prefix}{continuation} {item_connector}\u2500\u2500 {key} ({desc}{example_str})" + ) def write_modality(lines, mod_key, adata): @@ -79,7 +81,7 @@ def write_modality(lines, mod_key, adata): obs_names_examples = format_examples(adata.obs_names, n=3) lines.append(f"\u2502 \u251c\u2500\u2500 obs_names: e.g. {obs_names_examples}") for i, col in enumerate(obs_cols): - is_last = (i == len(obs_cols) - 1) + is_last = i == len(obs_cols) - 1 connector = "\u2514" if is_last else "\u251c" lines.append(f"\u2502 {connector}\u2500\u2500 {col}") @@ -88,11 +90,15 @@ def write_modality(lines, mod_key, adata): lines.append(f"\u251c\u2500\u2500 var/ ({adata.n_vars} variables)") var_names_examples = format_examples(adata.var_names, n=3) if not var_cols: - lines.append(f"\u2502 \u2514\u2500\u2500 var_names: e.g. {var_names_examples}") + lines.append( + f"\u2502 \u2514\u2500\u2500 var_names: e.g. {var_names_examples}" + ) else: - lines.append(f"\u2502 \u251c\u2500\u2500 var_names: e.g. {var_names_examples}") + lines.append( + f"\u2502 \u251c\u2500\u2500 var_names: e.g. {var_names_examples}" + ) for i, col in enumerate(var_cols): - is_last = (i == len(var_cols) - 1) + is_last = i == len(var_cols) - 1 connector = "\u2514" if is_last else "\u251c" lines.append(f"\u2502 {connector}\u2500\u2500 {col}") @@ -116,7 +122,7 @@ def write_mudata_level(lines, mdata): obs_cols = list(mdata.obs.columns) lines.append(f"\u251c\u2500\u2500 obs columns ({len(obs_cols)} total):") for i, col in enumerate(obs_cols): - is_last = (i == len(obs_cols) - 1) + is_last = i == len(obs_cols) - 1 connector = "\u2514" if is_last else "\u251c" lines.append(f"\u2502 {connector}\u2500\u2500 {col}") @@ -124,7 +130,7 @@ def write_mudata_level(lines, mdata): var_cols = list(mdata.var.columns) lines.append(f"\u251c\u2500\u2500 var columns ({len(var_cols)} total):") for i, col in enumerate(var_cols): - is_last = (i == len(var_cols) - 1) + is_last = i == len(var_cols) - 1 connector = "\u2514" if is_last else "\u251c" lines.append(f"\u2502 {connector}\u2500\u2500 {col}") @@ -169,7 +175,7 @@ def generate_structure(h5mu_path, output_path=None): text = "\n".join(lines) - with open(output_path, 'w') as f: + with open(output_path, "w") as f: f.write(text) print(f"Structure written to: {output_path}") @@ -180,29 +186,33 @@ def main(): parser = argparse.ArgumentParser( description="Generate a structured text summary of an .h5mu file" ) + parser.add_argument("h5mu_path", type=str, help="Path to the .h5mu file") parser.add_argument( - 'h5mu_path', type=str, - help='Path to the .h5mu file' + "--output", + "-o", + type=str, + default=None, + help="Output path for the structure file (default: _structure.txt)", ) parser.add_argument( - '--output', '-o', type=str, default=None, - help='Output path for the structure file (default: _structure.txt)' - ) - parser.add_argument( - '--adata_dir', type=str, default=None, - help='Directory containing .h5mu files. If provided, generates structure ' - 'for all .h5mu files in this directory (h5mu_path is ignored).' + "--adata_dir", + type=str, + default=None, + help="Directory containing .h5mu files. If provided, generates structure " + "for all .h5mu files in this directory (h5mu_path is ignored).", ) args = parser.parse_args() if args.adata_dir: # Batch mode: process all .h5mu files in the directory - h5mu_files = sorted([ - os.path.join(args.adata_dir, f) - for f in os.listdir(args.adata_dir) - if f.endswith('.h5mu') - ]) + h5mu_files = sorted( + [ + os.path.join(args.adata_dir, f) + for f in os.listdir(args.adata_dir) + if f.endswith(".h5mu") + ] + ) if not h5mu_files: print(f"No .h5mu files found in {args.adata_dir}") sys.exit(1) @@ -219,5 +229,5 @@ def main(): generate_structure(args.h5mu_path, args.output) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/.claude/skills/perturbNMF-runner/scripts/generate_slurm.py b/.claude/skills/perturbNMF-runner/scripts/generate_slurm.py index f56bcf6..e7f6b16 100644 --- a/.claude/skills/perturbNMF-runner/scripts/generate_slurm.py +++ b/.claude/skills/perturbNMF-runner/scripts/generate_slurm.py @@ -23,6 +23,7 @@ def get_gpu_mem_constraint(min_mem_gb): eligible = [AVAILABLE_GPU_MEMS_GB[-1]] # fallback to largest return "|".join(f"GPU_MEM:{m}GB" for m in eligible) + # Stage definitions: script path (relative), conda env, whether GPU is needed STAGES = { "inference-sk": { @@ -88,7 +89,9 @@ def generate_script(args, passthrough_args): """Generate the complete SLURM shell script.""" stage_info = STAGES[args.stage] pipeline_root = args.pipeline_root - skill_scripts = os.path.join(pipeline_root, ".claude/skills/perturbNMF-runner/scripts") + skill_scripts = os.path.join( + pipeline_root, ".claude/skills/perturbNMF-runner/scripts" + ) script_path = os.path.join(pipeline_root, stage_info["script"]) conda_env = stage_info["conda_env"] needs_gpu = stage_info["gpu"] or args.gpu @@ -117,18 +120,18 @@ def generate_script(args, passthrough_args): # The sentinel '---COMMENTED---' splits tokens into active flags (before, emitted # inside the python command) and unused flags (after, emitted as commented-out # `# --flag value` lines below the command for easy toggling). - COMMENT_SEP = '---COMMENTED---' + COMMENT_SEP = "---COMMENTED---" passthrough_str = "" commented_lines = [] if passthrough_args: tokens = list(passthrough_args) - if tokens and tokens[0] == '--': + if tokens and tokens[0] == "--": tokens = tokens[1:] if COMMENT_SEP in tokens: idx = tokens.index(COMMENT_SEP) active_tokens = tokens[:idx] - commented_tokens = tokens[idx + 1:] + commented_tokens = tokens[idx + 1 :] else: active_tokens = tokens commented_tokens = [] @@ -137,7 +140,7 @@ def _group_tokens(toks, prefix): grouped = [] current = [] for tok in toks: - if tok.startswith('--') and current: + if tok.startswith("--") and current: grouped.append(prefix + " ".join(current)) current = [tok] else: @@ -154,174 +157,200 @@ def _group_tokens(toks, prefix): # Build script as lines for clean formatting lines = [ - '#!/bin/bash', - '', - '# SLURM job configuration', - f'#SBATCH --job-name={args.job_name}', - f'#SBATCH --output={log_dir}/%j.out', - f'#SBATCH --error={log_dir}/%j.err', + "#!/bin/bash", + "", + "# SLURM job configuration", + f"#SBATCH --job-name={args.job_name}", + f"#SBATCH --output={log_dir}/%j.out", + f"#SBATCH --error={log_dir}/%j.err", ] if args.partition: - lines.append(f'#SBATCH --partition={args.partition}') - lines.extend([ - f'#SBATCH --time={args.time}', - '#SBATCH --nodes=1', - '#SBATCH --ntasks=1', - f'#SBATCH --cpus-per-task={args.cpus}', - f'#SBATCH --mem={args.mem}', - ]) + lines.append(f"#SBATCH --partition={args.partition}") + lines.extend( + [ + f"#SBATCH --time={args.time}", + "#SBATCH --nodes=1", + "#SBATCH --ntasks=1", + f"#SBATCH --cpus-per-task={args.cpus}", + f"#SBATCH --mem={args.mem}", + ] + ) if needs_gpu: - lines.append('#SBATCH --gres=gpu:1') + lines.append("#SBATCH --gres=gpu:1") if args.gpu_sku: # Explicit single-SKU override (cluster-specific GPU_SKU feature) - lines.append(f'#SBATCH -C GPU_SKU:{args.gpu_sku}') + lines.append(f"#SBATCH -C GPU_SKU:{args.gpu_sku}") elif args.gpu_min_mem: # Select all GPUs with sufficient memory (cluster-specific GPU_MEM feature) constraint = get_gpu_mem_constraint(args.gpu_min_mem) lines.append(f'#SBATCH -C "{constraint}"') if args.email: - lines.extend([ - '', - '# Email notifications', - '#SBATCH --mail-type=BEGIN,END,FAIL', - f'#SBATCH --mail-user={args.email}', - ]) - - lines.extend([ - '', - '# Configuration', - f'export PIPELINE_ROOT="{pipeline_root}"', - f'OUT_DIR="{args.output_dir}"', - f'RUN_NAME="{args.run_name}"', - f'LOG_DIR="{log_dir}"', - '', - '# Store start time', - 'START_TIME=$(date +%s)', - '', - '# Print job information', - 'echo "Job started at: $(date)"', - 'echo "Job ID: $SLURM_JOB_ID"', - 'echo "Node: $SLURMD_NODENAME"', - 'echo "Working directory: $(pwd)"', - 'echo "CPUs: $SLURM_CPUS_PER_TASK"', - 'echo "Partition: $SLURM_JOB_PARTITION"', - 'echo "Log directory: $LOG_DIR"', - '', - '# Create log directory', - 'mkdir -p "$LOG_DIR"', - ]) + lines.extend( + [ + "", + "# Email notifications", + "#SBATCH --mail-type=BEGIN,END,FAIL", + f"#SBATCH --mail-user={args.email}", + ] + ) + + lines.extend( + [ + "", + "# Configuration", + f'export PIPELINE_ROOT="{pipeline_root}"', + f'OUT_DIR="{args.output_dir}"', + f'RUN_NAME="{args.run_name}"', + f'LOG_DIR="{log_dir}"', + "", + "# Store start time", + "START_TIME=$(date +%s)", + "", + "# Print job information", + 'echo "Job started at: $(date)"', + 'echo "Job ID: $SLURM_JOB_ID"', + 'echo "Node: $SLURMD_NODENAME"', + 'echo "Working directory: $(pwd)"', + 'echo "CPUs: $SLURM_CPUS_PER_TASK"', + 'echo "Partition: $SLURM_JOB_PARTITION"', + 'echo "Log directory: $LOG_DIR"', + "", + "# Create log directory", + 'mkdir -p "$LOG_DIR"', + ] + ) if needs_gpu: - lines.extend([ - '', - '# Start resource monitoring', - 'MONITOR_LOG="$LOG_DIR/resource_monitor_${SLURM_JOB_ID}.log"', - 'monitor_resources() {', - ' while true; do', - ' echo "$(date \'+%Y-%m-%d %H:%M:%S\')" >> "$MONITOR_LOG"', - ' echo "=== GPU Usage ===" >> "$MONITOR_LOG"', - ' nvidia-smi --query-gpu=timestamp,name,utilization.gpu,utilization.memory,memory.total,memory.used,memory.free,temperature.gpu --format=csv >> "$MONITOR_LOG" 2>/dev/null', - ' echo "=== Memory ===" >> "$MONITOR_LOG"', - ' free -h >> "$MONITOR_LOG"', - ' echo "---" >> "$MONITOR_LOG"', - ' sleep 30', - ' done', - '}', - 'monitor_resources &', - 'MONITOR_PID=$!', - ]) - - lines.extend([ - '', - f'# Activate conda environment', - f'echo "Activating conda environment: {conda_env}"', - 'eval "$(conda shell.bash hook)"', - f'conda activate {conda_env}', - ]) + lines.extend( + [ + "", + "# Start resource monitoring", + 'MONITOR_LOG="$LOG_DIR/resource_monitor_${SLURM_JOB_ID}.log"', + "monitor_resources() {", + " while true; do", + ' echo "$(date \'+%Y-%m-%d %H:%M:%S\')" >> "$MONITOR_LOG"', + ' echo "=== GPU Usage ===" >> "$MONITOR_LOG"', + ' nvidia-smi --query-gpu=timestamp,name,utilization.gpu,utilization.memory,memory.total,memory.used,memory.free,temperature.gpu --format=csv >> "$MONITOR_LOG" 2>/dev/null', + ' echo "=== Memory ===" >> "$MONITOR_LOG"', + ' free -h >> "$MONITOR_LOG"', + ' echo "---" >> "$MONITOR_LOG"', + " sleep 30", + " done", + "}", + "monitor_resources &", + "MONITOR_PID=$!", + ] + ) + + lines.extend( + [ + "", + f"# Activate conda environment", + f'echo "Activating conda environment: {conda_env}"', + 'eval "$(conda shell.bash hook)"', + f"conda activate {conda_env}", + ] + ) # torch-nmf-dl is self-contained (see Tools/PerturbNMF/env/). Pin resolution # to the env so a stale ~/.local/lib/python3.9/site-packages can never shadow # or silently supply its packages. Other envs still rely on user-site. - if conda_env == 'torch-nmf-dl': - lines.extend([ - '', - '# Resolve imports against the env only, never ~/.local (user site)', - 'export PYTHONNOUSERSITE=1', - ]) - - lines.extend([ - '', - 'echo "Active env: $CONDA_DEFAULT_ENV"', - 'echo "Python: $(python --version)"', - ]) - - lines.extend([ - '', - f'# Run pipeline', - f'echo "Running: {args.stage}"', - f'{stage_info.get("interpreter", "python3")} {script_path}{passthrough_str}', - ]) + if conda_env == "torch-nmf-dl": + lines.extend( + [ + "", + "# Resolve imports against the env only, never ~/.local (user site)", + "export PYTHONNOUSERSITE=1", + ] + ) + + lines.extend( + [ + "", + 'echo "Active env: $CONDA_DEFAULT_ENV"', + 'echo "Python: $(python --version)"', + ] + ) + + lines.extend( + [ + "", + f"# Run pipeline", + f'echo "Running: {args.stage}"', + f"{stage_info.get('interpreter', 'python3')} {script_path}{passthrough_str}", + ] + ) if commented_lines: - lines.append('') - lines.append('# Optional / unused flags for this stage (uncomment + add to the command above to enable):') + lines.append("") + lines.append( + "# Optional / unused flags for this stage (uncomment + add to the command above to enable):" + ) lines.extend(commented_lines) - lines.extend([ - '', - '# Capture exit code', - 'EXIT_CODE=$?', - 'if [ $EXIT_CODE -ne 0 ]; then', - ' echo "ERROR: Pipeline exited with code $EXIT_CODE"', - 'fi', - ]) + lines.extend( + [ + "", + "# Capture exit code", + "EXIT_CODE=$?", + "if [ $EXIT_CODE -ne 0 ]; then", + ' echo "ERROR: Pipeline exited with code $EXIT_CODE"', + "fi", + ] + ) if needs_gpu: - lines.extend([ - '', - '# Stop resource monitoring', - 'kill $MONITOR_PID 2>/dev/null', - 'echo "Final GPU status:"', - 'nvidia-smi 2>/dev/null || echo "GPU not available"', - ]) + lines.extend( + [ + "", + "# Stop resource monitoring", + "kill $MONITOR_PID 2>/dev/null", + 'echo "Final GPU status:"', + 'nvidia-smi 2>/dev/null || echo "GPU not available"', + ] + ) # For inference stages, generate h5mu structure files after pipeline completes if args.stage.startswith("inference") and not args.no_structure: h5mu_script = os.path.join(skill_scripts, "generate_h5mu_structure.py") - lines.extend([ - '', - '# Generate h5mu structure summary files', - 'if [ $EXIT_CODE -eq 0 ]; then', - f' ADATA_DIR="{args.output_dir}/{args.run_name}/Inference/adata"', - ' if [ -d "$ADATA_DIR" ]; then', - ' echo "Generating h5mu structure files..."', - f' python3 {h5mu_script} dummy --adata_dir "$ADATA_DIR"', - ' echo "h5mu structure generation complete"', - ' else', - ' echo "WARNING: adata directory not found at $ADATA_DIR, skipping structure generation"', - ' fi', - 'else', - ' echo "Skipping h5mu structure generation due to pipeline error"', - 'fi', - ]) - - lines.extend([ - '', - '# Elapsed time', - 'END_TIME=$(date +%s)', - 'ELAPSED=$((END_TIME - START_TIME))', - 'HOURS=$((ELAPSED / 3600))', - 'MINUTES=$(((ELAPSED % 3600) / 60))', - 'SECONDS=$((ELAPSED % 60))', - '', - 'echo "Job completed at: $(date)"', - 'echo "Total elapsed time: ${HOURS}h ${MINUTES}m ${SECONDS}s (${ELAPSED} seconds)"', - 'exit $EXIT_CODE', - ]) - - script = '\n'.join(lines) + '\n' + lines.extend( + [ + "", + "# Generate h5mu structure summary files", + "if [ $EXIT_CODE -eq 0 ]; then", + f' ADATA_DIR="{args.output_dir}/{args.run_name}/Inference/adata"', + ' if [ -d "$ADATA_DIR" ]; then', + ' echo "Generating h5mu structure files..."', + f' python3 {h5mu_script} dummy --adata_dir "$ADATA_DIR"', + ' echo "h5mu structure generation complete"', + " else", + ' echo "WARNING: adata directory not found at $ADATA_DIR, skipping structure generation"', + " fi", + "else", + ' echo "Skipping h5mu structure generation due to pipeline error"', + "fi", + ] + ) + + lines.extend( + [ + "", + "# Elapsed time", + "END_TIME=$(date +%s)", + "ELAPSED=$((END_TIME - START_TIME))", + "HOURS=$((ELAPSED / 3600))", + "MINUTES=$(((ELAPSED % 3600) / 60))", + "SECONDS=$((ELAPSED % 60))", + "", + 'echo "Job completed at: $(date)"', + 'echo "Total elapsed time: ${HOURS}h ${MINUTES}m ${SECONDS}s (${ELAPSED} seconds)"', + "exit $EXIT_CODE", + ] + ) + + script = "\n".join(lines) + "\n" return script @@ -333,76 +362,89 @@ def main(): # Generator-specific arguments parser.add_argument( - '--stage', type=str, required=True, + "--stage", + type=str, + required=True, choices=list(STAGES.keys()), - help='Pipeline stage to run' - ) - parser.add_argument( - '--job_name', type=str, required=True, - help='SLURM job name' + help="Pipeline stage to run", ) + parser.add_argument("--job_name", type=str, required=True, help="SLURM job name") parser.add_argument( - '--output_dir', type=str, required=True, - help='Output directory for results' + "--output_dir", type=str, required=True, help="Output directory for results" ) parser.add_argument( - '--run_name', type=str, required=True, - help='Run name identifier' + "--run_name", type=str, required=True, help="Run name identifier" ) parser.add_argument( - '--cpus', type=int, default=10, - help='CPUs per task (default: 10)' + "--cpus", type=int, default=10, help="CPUs per task (default: 10)" ) parser.add_argument( - '--mem', type=str, default='64G', - help='Memory allocation (default: 64G)' + "--mem", type=str, default="64G", help="Memory allocation (default: 64G)" ) parser.add_argument( - '--time', type=str, default='10:00:00', - help='Time limit HH:MM:SS (default: 10:00:00)' + "--time", + type=str, + default="10:00:00", + help="Time limit HH:MM:SS (default: 10:00:00)", ) parser.add_argument( - '--partition', type=str, default=os.environ.get('SLURM_PARTITION'), + "--partition", + type=str, + default=os.environ.get("SLURM_PARTITION"), help='SLURM partition(s), e.g. "normal" or "p1,p2" (default: $SLURM_PARTITION; ' - 'if unset, no --partition line is written and the cluster default is used)' + "if unset, no --partition line is written and the cluster default is used)", ) parser.add_argument( - '--gpu', action='store_true', - help='Request GPU resources (auto-set for torch stages)' + "--gpu", + action="store_true", + help="Request GPU resources (auto-set for torch stages)", ) parser.add_argument( - '--gpu_sku', type=str, default=None, - help='Explicit single GPU SKU constraint, written as `-C GPU_SKU:` ' - '(e.g., H100_SXM5; requires a cluster that defines GPU_SKU features). ' - 'Overrides --gpu_min_mem. If neither is set, no GPU constraint is written.' + "--gpu_sku", + type=str, + default=None, + help="Explicit single GPU SKU constraint, written as `-C GPU_SKU:` " + "(e.g., H100_SXM5; requires a cluster that defines GPU_SKU features). " + "Overrides --gpu_min_mem. If neither is set, no GPU constraint is written.", ) parser.add_argument( - '--gpu_min_mem', type=int, default=None, + "--gpu_min_mem", + type=int, + default=None, help='Minimum GPU memory in GB, written as `-C "GPU_MEM:GB|..."` for all ' - 'sizes >= this value (requires a cluster that defines GPU_MEM features). ' - 'Default: no GPU constraint.' + "sizes >= this value (requires a cluster that defines GPU_MEM features). " + "Default: no GPU constraint.", ) parser.add_argument( - '--email', type=str, default=os.environ.get('SLURM_MAIL_USER'), - help='Email for SLURM notifications (default: $SLURM_MAIL_USER; ' - 'if unset, no mail directives are written)' + "--email", + type=str, + default=os.environ.get("SLURM_MAIL_USER"), + help="Email for SLURM notifications (default: $SLURM_MAIL_USER; " + "if unset, no mail directives are written)", ) parser.add_argument( - '--pipeline_root', type=str, required=True, - help='Path to the PerturbNMF repo on the machine that runs the job (required; ' - 'always ask the user — no default is assumed)' + "--pipeline_root", + type=str, + required=True, + help="Path to the PerturbNMF repo on the machine that runs the job (required; " + "always ask the user — no default is assumed)", ) parser.add_argument( - '--log_dir', type=str, default=None, - help='Custom log directory (default: //logs)' + "--log_dir", + type=str, + default=None, + help="Custom log directory (default: //logs)", ) parser.add_argument( - '--script_output_path', type=str, default=None, - help='Where to write the generated .sh file (default: stdout)' + "--script_output_path", + type=str, + default=None, + help="Where to write the generated .sh file (default: stdout)", ) parser.add_argument( - '--no_structure', action='store_true', - help='Skip h5mu structure file generation after inference' + "--no_structure", + action="store_true", + help="Skip h5mu structure file generation after inference", ) # Parse only known args; the rest are passthrough to the pipeline script @@ -412,7 +454,7 @@ def main(): if args.script_output_path: os.makedirs(os.path.dirname(args.script_output_path), exist_ok=True) - with open(args.script_output_path, 'w') as f: + with open(args.script_output_path, "w") as f: f.write(script) os.chmod(args.script_output_path, 0o755) print(f"Script written to: {args.script_output_path}") @@ -420,5 +462,5 @@ def main(): print(script) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/.claude/skills/perturbNMF-runner/scripts/prepare_guide_data.py b/.claude/skills/perturbNMF-runner/scripts/prepare_guide_data.py index 474a2f2..efa5bb3 100644 --- a/.claude/skills/perturbNMF-runner/scripts/prepare_guide_data.py +++ b/.claude/skills/perturbNMF-runner/scripts/prepare_guide_data.py @@ -22,28 +22,37 @@ def main(): description="Prepare guide data h5ad for CRT calibration from h5mu" ) parser.add_argument( - '--h5mu_path', type=str, required=True, - help='Path to input h5mu file (any K value works; use smallest for speed)' + "--h5mu_path", + type=str, + required=True, + help="Path to input h5mu file (any K value works; use smallest for speed)", ) parser.add_argument( - '--output_path', type=str, required=True, - help='Path to save output guide h5ad file' + "--output_path", + type=str, + required=True, + help="Path to save output guide h5ad file", ) parser.add_argument( - '--modality', type=str, default='cNMF', - help='Modality to extract guide info from (default: cNMF)' + "--modality", + type=str, + default="cNMF", + help="Modality to extract guide info from (default: cNMF)", ) parser.add_argument( - '--assignment_key', type=str, default='target_assignment', - help='Source key in obsm for guide assignment matrix (default: target_assignment)' + "--assignment_key", + type=str, + default="target_assignment", + help="Source key in obsm for guide assignment matrix (default: target_assignment)", ) parser.add_argument( - '--names_key', type=str, default='target_names', - help='Source key in uns for guide/target names (default: target_names)' + "--names_key", + type=str, + default="target_names", + help="Source key in uns for guide/target names (default: target_names)", ) parser.add_argument( - '--force', action='store_true', - help='Overwrite output file if it exists' + "--force", action="store_true", help="Overwrite output file if it exists" ) args = parser.parse_args() @@ -75,9 +84,9 @@ def main(): adata_guide.obs[col] = mod.obs[col].values # Map source keys to standard guide_* keys - adata_guide.obsm['guide_assignment'] = np.array(mod.obsm[args.assignment_key]) - adata_guide.uns['guide_names'] = list(mod.uns[args.names_key]) - adata_guide.uns['guide_targets'] = np.array(mod.uns[args.names_key]) + adata_guide.obsm["guide_assignment"] = np.array(mod.obsm[args.assignment_key]) + adata_guide.uns["guide_names"] = list(mod.uns[args.names_key]) + adata_guide.uns["guide_targets"] = np.array(mod.uns[args.names_key]) os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) adata_guide.write(args.output_path) @@ -85,10 +94,10 @@ def main(): print(f"Guide h5ad saved: {args.output_path}") print(f" Cells: {adata_guide.n_obs}") print(f" Targets: {len(adata_guide.uns['guide_names'])}") - names = adata_guide.uns['guide_names'] + names = adata_guide.uns["guide_names"] preview = names[:5] if len(names) > 5 else names print(f" Target names (first 5): {preview}") -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/.claude/skills/perturbNMF-runner/scripts/validate_data.py b/.claude/skills/perturbNMF-runner/scripts/validate_data.py index 4f20824..6d09331 100644 --- a/.claude/skills/perturbNMF-runner/scripts/validate_data.py +++ b/.claude/skills/perturbNMF-runner/scripts/validate_data.py @@ -14,14 +14,16 @@ import sys # Add pipeline root to path -PIPELINE_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +PIPELINE_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) sys.path.insert(0, PIPELINE_ROOT) def get_file_size_str(path): """Return human-readable file size.""" size_bytes = os.path.getsize(path) - for unit in ['B', 'KB', 'MB', 'GB', 'TB']: + for unit in ["B", "KB", "MB", "GB", "TB"]: if size_bytes < 1024: return f"{size_bytes:.1f} {unit}", size_bytes size_bytes /= 1024 @@ -32,6 +34,7 @@ def compute_sparsity(X): """Compute fraction of zeros in the matrix.""" import scipy.sparse as sp import numpy as np + if sp.issparse(X): nnz = X.nnz total = X.shape[0] * X.shape[1] @@ -41,8 +44,13 @@ def compute_sparsity(X): return 1.0 - (nnz / total) if total > 0 else 0.0 -def check_data_format(adata, guide_names_key="guide_names", guide_targets_key="guide_targets", - categorical_key="batch", guide_assignment_key="guide_assignment"): +def check_data_format( + adata, + guide_names_key="guide_names", + guide_targets_key="guide_targets", + categorical_key="batch", + guide_assignment_key="guide_assignment", +): """Validate that an AnnData has the keys cNMF inference relies on.""" is_valid = True if categorical_key not in adata.obs: @@ -83,13 +91,20 @@ def check_data_format(adata, guide_names_key="guide_names", guide_targets_key="g print(f"Found adata.obsm['{guide_assignment_key}']\n") try: import scipy.sparse as sp + if sp.issparse(guide_assignment): - print(f"WARNING: '{guide_assignment_key}' is sparse. Converting to dense array...") + print( + f"WARNING: '{guide_assignment_key}' is sparse. Converting to dense array..." + ) dense_array = guide_assignment.toarray() adata.obsm[guide_assignment_key] = dense_array - print(f"'{guide_assignment_key}' converted to dense array (shape: {dense_array.shape})\n") + print( + f"'{guide_assignment_key}' converted to dense array (shape: {dense_array.shape})\n" + ) else: - print(f"'{guide_assignment_key}' is already dense (shape: {guide_assignment.shape})\n") + print( + f"'{guide_assignment_key}' is already dense (shape: {guide_assignment.shape})\n" + ) except Exception as e: print(f"WARNING: Error checking '{guide_assignment_key}' sparsity: {e}\n") is_valid = False @@ -99,25 +114,43 @@ def check_data_format(adata, guide_names_key="guide_names", guide_targets_key="g def _validate_against_reference_gtf(rna_gene_names, gtf_path): """Compare RNA var_names against gene_name attributes extracted from a GTF.""" import pandas as pd + validation = {"is_valid": True, "gtf_genes": set(), "not_in_gtf": []} try: gtf_df = pd.read_csv( - gtf_path, sep="\t", comment="#", header=None, - dtype={0: str, 1: str, 2: str, 3: int, 4: int, 5: str, 6: str, 7: str, 8: str}, + gtf_path, + sep="\t", + comment="#", + header=None, + dtype={ + 0: str, + 1: str, + 2: str, + 3: int, + 4: int, + 5: str, + 6: str, + 7: str, + 8: str, + }, ) extracted = gtf_df[8].str.extract(r'gene_name "([^"]+)"')[0] if extracted.notna().sum() == 0: print("gene_name not found in any rows\n") else: - print(f"Successfully extracted gene_name from {extracted.notna().sum()} rows") + print( + f"Successfully extracted gene_name from {extracted.notna().sum()} rows" + ) extracted = set(extracted) validation["gtf_genes"] = extracted rna_not_in_gtf = rna_gene_names - extracted print(f"Found {len(rna_gene_names)} gene names in adata.var_name") if rna_not_in_gtf: - print(f"WARNING: {len(rna_not_in_gtf)}/{len(rna_gene_names)} adata.var RNA NOT found in reference GTF." - "\n This might be caused by the adata.var RNA being Ensembl ID instead of gene symbols.") + print( + f"WARNING: {len(rna_not_in_gtf)}/{len(rna_gene_names)} adata.var RNA NOT found in reference GTF." + "\n This might be caused by the adata.var RNA being Ensembl ID instead of gene symbols." + ) print(f"Examples: {list(rna_not_in_gtf)[:5]}\n") validation["not_in_gtf"] = sorted(list(rna_not_in_gtf)) else: @@ -131,16 +164,26 @@ def _validate_against_reference_gtf(rna_gene_names, gtf_path): return validation -def check_guide_names(adata, guide_names_key="guide_names", guide_targets_key="guide_targets", - categorical_key="batch", reference_gtf_path=None, guide_annotation_path=None, - guide_assignment_key="guide_assignment", check_var_names_align_with_guide_targets=True): +def check_guide_names( + adata, + guide_names_key="guide_names", + guide_targets_key="guide_targets", + categorical_key="batch", + reference_gtf_path=None, + guide_annotation_path=None, + guide_assignment_key="guide_assignment", + check_var_names_align_with_guide_targets=True, +): """Validate guide/RNA name consistency, optionally against a GTF and/or guide annotation TSV.""" import pandas as pd results = {"is_valid": True, "missing_from_rna": [], "gtf_validation": None} results["is_valid"] = check_data_format( - adata, guide_names_key=guide_names_key, guide_targets_key=guide_targets_key, - categorical_key=categorical_key, guide_assignment_key=guide_assignment_key, + adata, + guide_names_key=guide_names_key, + guide_targets_key=guide_targets_key, + categorical_key=categorical_key, + guide_assignment_key=guide_assignment_key, ) rna_gene_names = set(adata.var_names) @@ -157,17 +200,24 @@ def check_guide_names(adata, guide_names_key="guide_names", guide_targets_key="g guide_gene_targets = set(guide_targets) guide_gene_names = set(guide_names) - if len(guide_targets) != len(guide_names) and len(guide_gene_targets) != adata.obsm[guide_assignment_key].shape[1]: - print("WARNING: guide_targets and guide_names and guide_assignment col should be equal length but not in here\n") + if ( + len(guide_targets) != len(guide_names) + and len(guide_gene_targets) != adata.obsm[guide_assignment_key].shape[1] + ): + print( + "WARNING: guide_targets and guide_names and guide_assignment col should be equal length but not in here\n" + ) results["is_valid"] = False if check_var_names_align_with_guide_targets: guide_not_in_rna = guide_gene_targets - rna_gene_names print(f"Found {len(guide_gene_targets)} genes in gene_targets") if guide_not_in_rna: - print(f"WARNING: {len(guide_not_in_rna)}/{len(guide_gene_targets)} gene names in guide_targets but NOT in adata.var.\n" - " This might be caused by mismatch between gene symbol vs Ensembl ID or some perturbed genes are not expressed in the dataset." - " Might cause issue in evalutation or plotting if naming conventions are different.") + print( + f"WARNING: {len(guide_not_in_rna)}/{len(guide_gene_targets)} gene names in guide_targets but NOT in adata.var.\n" + " This might be caused by mismatch between gene symbol vs Ensembl ID or some perturbed genes are not expressed in the dataset." + " Might cause issue in evalutation or plotting if naming conventions are different." + ) print(f"Examples: {list(guide_not_in_rna)[:5]}\n") results["missing_from_rna"] = sorted(list(guide_not_in_rna)) else: @@ -177,7 +227,9 @@ def check_guide_names(adata, guide_names_key="guide_names", guide_targets_key="g print("No reference GTF provided (optional check skipped)") results["gtf_validation"] = None else: - gtf_validation = _validate_against_reference_gtf(rna_gene_names, reference_gtf_path) + gtf_validation = _validate_against_reference_gtf( + rna_gene_names, reference_gtf_path + ) results["gtf_validation"] = gtf_validation if not gtf_validation["is_valid"]: results["is_valid"] = False @@ -207,10 +259,14 @@ def check_guide_names(adata, guide_names_key="guide_names", guide_targets_key="g guide_names_not_in_annotation = guide_gene_names - annotation_guide_names if guide_targets_not_in_annotation or guide_names_not_in_annotation: if guide_targets_not_in_annotation: - print(f"\n WARNING: {len(guide_targets_not_in_annotation)}/{len(guide_gene_targets)} guide targets in data but NOT in annotation.") + print( + f"\n WARNING: {len(guide_targets_not_in_annotation)}/{len(guide_gene_targets)} guide targets in data but NOT in annotation." + ) print(f"Examples: {list(guide_targets_not_in_annotation)[:5]}\n") if guide_names_not_in_annotation: - print(f"\n WARNING: {len(guide_names_not_in_annotation)}/{len(guide_gene_names)} guide names in data but NOT in annotation.") + print( + f"\n WARNING: {len(guide_names_not_in_annotation)}/{len(guide_gene_names)} guide names in data but NOT in annotation." + ) print(f"Examples: {list(guide_names_not_in_annotation)[:5]}\n") results["is_valid"] = False else: @@ -244,7 +300,9 @@ def validate_adata(adata, args): # Check categorical key if args.categorical_key not in adata.obs: - issues.append(f"Missing obs['{args.categorical_key}'] (condition key, --categorical_key)") + issues.append( + f"Missing obs['{args.categorical_key}'] (condition key, --categorical_key)" + ) else: n_cats = adata.obs[args.categorical_key].nunique() print(f"\n obs['{args.categorical_key}']: {n_cats} unique values") @@ -252,11 +310,11 @@ def validate_adata(adata, args): # Check obsm print(f"\n=== obsm keys ===") for key in adata.obsm: - shape = adata.obsm[key].shape if hasattr(adata.obsm[key], 'shape') else 'N/A' + shape = adata.obsm[key].shape if hasattr(adata.obsm[key], "shape") else "N/A" print(f" - {key}: {shape}") - has_pca = 'X_pca' in adata.obsm - has_umap = 'X_umap' in adata.obsm + has_pca = "X_pca" in adata.obsm + has_umap = "X_umap" in adata.obsm if not has_pca: issues.append("Missing obsm['X_pca']") if not has_umap: @@ -277,7 +335,7 @@ def validate_adata(adata, args): if has_guide_targets: targets = adata.uns[args.guide_targets_key] - if hasattr(targets, '__len__'): + if hasattr(targets, "__len__"): n_targets = len(set(targets)) print(f" uns['{args.guide_targets_key}']: {n_targets} unique targets") else: @@ -288,7 +346,7 @@ def validate_adata(adata, args): if has_guide_assignment: ga = adata.obsm[args.guide_assignment_key] - ga_shape = ga.shape if hasattr(ga, 'shape') else 'N/A' + ga_shape = ga.shape if hasattr(ga, "shape") else "N/A" ga_sparse = sp.issparse(ga) print(f" obsm['{args.guide_assignment_key}']: {ga_shape} (sparse={ga_sparse})") else: @@ -327,7 +385,7 @@ def validate_mdata(mdata, args): print(f"\n--- Found '{args.prog_key}' modality ---") cnmf_adata = mdata[args.prog_key] print(f" Programs: {cnmf_adata.n_vars}") - if 'loadings' in cnmf_adata.varm: + if "loadings" in cnmf_adata.varm: print(f" varm['loadings']: {cnmf_adata.varm['loadings'].shape}") else: print(f" varm['loadings']: NOT FOUND (expected post-inference)") @@ -340,40 +398,58 @@ def main(): description="Validate input data for cNMF pipeline" ) parser.add_argument( - '--counts_fn', type=str, required=True, - help='Path to input data file (.h5ad or .h5mu)' + "--counts_fn", + type=str, + required=True, + help="Path to input data file (.h5ad or .h5mu)", ) parser.add_argument( - '--categorical_key', type=str, default='sample', - help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample" + "--categorical_key", + type=str, + default="sample", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", ) parser.add_argument( - '--guide_names_key', type=str, default='guide_names', - help='Key in .uns for guide names (default: guide_names)' + "--guide_names_key", + type=str, + default="guide_names", + help="Key in .uns for guide names (default: guide_names)", ) parser.add_argument( - '--guide_targets_key', type=str, default='guide_targets', - help='Key in .uns for guide targets (default: guide_targets)' + "--guide_targets_key", + type=str, + default="guide_targets", + help="Key in .uns for guide targets (default: guide_targets)", ) parser.add_argument( - '--guide_assignment_key', type=str, default='guide_assignment', - help='Key in .obsm for guide assignment matrix (default: guide_assignment)' + "--guide_assignment_key", + type=str, + default="guide_assignment", + help="Key in .obsm for guide assignment matrix (default: guide_assignment)", ) parser.add_argument( - '--data_key', type=str, default='rna', - help='Key for RNA modality in MuData (default: rna)' + "--data_key", + type=str, + default="rna", + help="Key for RNA modality in MuData (default: rna)", ) parser.add_argument( - '--prog_key', type=str, default='cNMF', - help='Key for cNMF modality in MuData (default: cNMF)' + "--prog_key", + type=str, + default="cNMF", + help="Key for cNMF modality in MuData (default: cNMF)", ) parser.add_argument( - '--reference_gtf_path', type=str, default=None, - help='Path to reference GTF for gene name validation (optional)' + "--reference_gtf_path", + type=str, + default=None, + help="Path to reference GTF for gene name validation (optional)", ) parser.add_argument( - '--guide_annotation_path', type=str, default=None, - help='Path to guide annotation TSV file (optional)' + "--guide_annotation_path", + type=str, + default=None, + help="Path to guide annotation TSV file (optional)", ) args = parser.parse_args() @@ -392,13 +468,15 @@ def main(): # Load and validate issues = [] - if ext == '.h5mu': + if ext == ".h5mu": import muon as mu + print("Loading MuData...") mdata = mu.read(args.counts_fn) issues = validate_mdata(mdata, args) - elif ext in ('.h5ad', '.h5'): + elif ext in (".h5ad", ".h5"): import scanpy as sc + print("Loading AnnData...") adata = sc.read_h5ad(args.counts_fn) issues = validate_adata(adata, args) @@ -409,7 +487,7 @@ def main(): # Optional: run full format checking from pipeline if args.reference_gtf_path or args.guide_annotation_path: try: - if ext == '.h5mu': + if ext == ".h5mu": adata_check = mdata[args.data_key] else: adata_check = adata @@ -420,15 +498,21 @@ def main(): categorical_key=args.categorical_key, reference_gtf_path=args.reference_gtf_path, guide_annotation_path=args.guide_annotation_path, - guide_assignment_key=args.guide_assignment_key + guide_assignment_key=args.guide_assignment_key, ) - if not result.get('is_valid', True) if isinstance(result, dict) else not result: - issues.append("Format check reported validation failures (see output above)") + if ( + not result.get("is_valid", True) + if isinstance(result, dict) + else not result + ): + issues.append( + "Format check reported validation failures (see output above)" + ) except Exception as e: print(f"\nWARNING: Could not run extended format checks: {e}") # Resource recommendations - if ext == '.h5mu': + if ext == ".h5mu": n_cells = mdata[args.data_key].n_obs if args.data_key in mdata.mod else 0 n_genes = mdata[args.data_key].n_vars if args.data_key in mdata.mod else 0 else: @@ -447,8 +531,14 @@ def main(): print(f" Cells: {n_cells:,} | Genes: {n_genes:,}") print(f" sk-cNMF: mem={mem_rec}, time={time_rec_sk} (numiter=10, 8 K values)") - print(f" torch-cNMF: mem={'64G' if n_cells < 100000 else '128G' if n_cells < 300000 else '256G'}, time={time_rec_torch}") - partition = "standard CPU partition" if n_cells < 200000 else "large-memory partition (if your cluster has one)" + print( + f" torch-cNMF: mem={'64G' if n_cells < 100000 else '128G' if n_cells < 300000 else '256G'}, time={time_rec_torch}" + ) + partition = ( + "standard CPU partition" + if n_cells < 200000 + else "large-memory partition (if your cluster has one)" + ) print(f" Partition: {partition}") # Report @@ -463,5 +553,5 @@ def main(): sys.exit(0) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/.claude/skills/pipeline-drift-check/scripts/check-pipeline-drift.py b/.claude/skills/pipeline-drift-check/scripts/check-pipeline-drift.py index b3596bf..12ee3ad 100755 --- a/.claude/skills/pipeline-drift-check/scripts/check-pipeline-drift.py +++ b/.claude/skills/pipeline-drift-check/scripts/check-pipeline-drift.py @@ -26,13 +26,21 @@ sys.exit(0) # Primary flag plus one optional alias, e.g. add_argument('--Conditions', '--Sample', ...) -ARG_RE = re.compile(r"add_argument\(\s*['\"]--([a-zA-Z][a-zA-Z0-9_-]+)['\"](?:\s*,\s*['\"]--([a-zA-Z][a-zA-Z0-9_-]+)['\"])?") +ARG_RE = re.compile( + r"add_argument\(\s*['\"]--([a-zA-Z][a-zA-Z0-9_-]+)['\"](?:\s*,\s*['\"]--([a-zA-Z][a-zA-Z0-9_-]+)['\"])?" +) R_ARG_RE = re.compile(r"make_option\(\s*['\"]--([a-zA-Z][a-zA-Z0-9_-]+)['\"]") PYTEST_ARG_RE = re.compile(r"addoption\(\s*['\"]--([a-zA-Z][a-zA-Z0-9_-]+)['\"]") DASH_RE = re.compile(r"--([a-zA-Z][a-zA-Z0-9_-]+)") TABLE_ROW_RE = re.compile(r"^\|\s*`?([a-z][a-z0-9_-]+)`?\s*\|", re.MULTILINE) -SKIP_DIR_PARTS = {"JupterNote_Version", "JupterNote_version", "tests", ".git", "__pycache__"} +SKIP_DIR_PARTS = { + "JupterNote_Version", + "JupterNote_version", + "tests", + ".git", + "__pycache__", +} # Wildcard placeholders in narrative prose ("renamed --foo_* to --bar_*"). WILDCARD_RE = re.compile(r"--[a-zA-Z0-9_]+_\*") # Abstract placeholders used in skill markdown to document conventions, not @@ -42,7 +50,9 @@ # can't tell these apart from real flag references, so we exempt them. SKILL_PLACEHOLDER_ALLOWLIST = {"flag", "flag_name", "COMMENTED---"} # Lines describing a rename — old name on the line is documentation, not drift. -RENAME_RE = re.compile(r"(?:→|->|\bis now\b|\brenamed\b|\bformerly\b|\bwas\s+renamed\b)", re.IGNORECASE) +RENAME_RE = re.compile( + r"(?:→|->|\bis now\b|\brenamed\b|\bformerly\b|\bwas\s+renamed\b)", re.IGNORECASE +) # Lines referencing external CLI tools — their --flags belong to those tools. EXTERNAL_TOOL_RE = re.compile( r"\b(sacct|sbatch|squeue|scancel|sinfo|srun|salloc|jupyter|nbconvert|pytest|conftest)\b" @@ -55,7 +65,12 @@ def skip(p: Path) -> bool: def py_args(py: Path) -> set[str]: try: - return {name for pair in ARG_RE.findall(py.read_text(errors="ignore")) for name in pair if name} + return { + name + for pair in ARG_RE.findall(py.read_text(errors="ignore")) + for name in pair + if name + } except OSError: return set() @@ -129,8 +144,15 @@ def md_param_table_names(md: str) -> set[str]: if in_section and not _skip_line(line): m = TABLE_ROW_RE.match(line) if m and m.group(1).lower() not in { - "parameter", "argument", "name", "type", "default", - "description", "required", "optional", "key", + "parameter", + "argument", + "name", + "type", + "default", + "description", + "required", + "optional", + "key", }: out.add(m.group(1)) return out @@ -165,7 +187,7 @@ def sh_runner_args(sh: Path, target_basename: str | None = None) -> set[str]: if s.startswith("#"): continue if not in_block: - invokes = (("python" in s and ".py" in s) or ("Rscript" in s and ".R" in s)) + invokes = ("python" in s and ".py" in s) or ("Rscript" in s and ".R" in s) if invokes: if target_basename is None or target_basename in s: in_block = True @@ -194,7 +216,9 @@ def find_sibling(py: Path, suffix: str) -> Path | None: return None -def check_per_script(pipeline: dict[Path, set[str]]) -> list[tuple[Path, Path, set[str]]]: +def check_per_script( + pipeline: dict[Path, set[str]], +) -> list[tuple[Path, Path, set[str]]]: """For each pipeline .py, compare its sibling README and .sh to its argparse. Returns list of (consumer_file, pipeline_py, stale_arg_set). @@ -217,7 +241,9 @@ def check_per_script(pipeline: dict[Path, set[str]]) -> list[tuple[Path, Path, s return findings -def check_skill_globally(global_args: set[str], helper_args: set[str]) -> dict[Path, set[str]]: +def check_skill_globally( + global_args: set[str], helper_args: set[str] +) -> dict[Path, set[str]]: """Skill markdown: an arg must be in some pipeline .py OR in a skill helper .py. The latter exempts the skill's own CLI flags (--gpu, --stage, etc.).""" valid = global_args | helper_args @@ -245,14 +271,18 @@ def main() -> int: skill_drift = check_skill_globally(global_args, helper_args) if not per_script and not skill_drift: - print("No drift detected. README/.sh/skill docs are aligned with src/ argparse.") + print( + "No drift detected. README/.sh/skill docs are aligned with src/ argparse." + ) return 0 lines: list[str] = ["", "Param drift detected in PerturbNMF."] if per_script: lines.append("") - lines.append("Per-script drift (sibling README/.sh has args missing from its .py):") + lines.append( + "Per-script drift (sibling README/.sh has args missing from its .py):" + ) # Group by consumer file for readability by_consumer: dict[Path, list[tuple[Path, set[str]]]] = defaultdict(list) for consumer, py, stale in per_script: @@ -267,7 +297,9 @@ def main() -> int: if skill_drift: lines.append("") - lines.append("Skill drift (args mentioned but not defined in ANY pipeline .py):") + lines.append( + "Skill drift (args mentioned but not defined in ANY pipeline .py):" + ) for path, stale in sorted(skill_drift.items()): rel = path.relative_to(ROOT) lines.append(f" {rel}") diff --git a/src/Stage1_Inference/sk-cNMF/Slurm_Version/sk-cNMF_batch_inference_pipeline.py b/src/Stage1_Inference/sk-cNMF/Slurm_Version/sk-cNMF_batch_inference_pipeline.py index fa85475..e83d88e 100644 --- a/src/Stage1_Inference/sk-cNMF/Slurm_Version/sk-cNMF_batch_inference_pipeline.py +++ b/src/Stage1_Inference/sk-cNMF/Slurm_Version/sk-cNMF_batch_inference_pipeline.py @@ -23,70 +23,235 @@ # Change path to wherever you have repo locally from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) + +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) from Stage1_Inference.src import ( - run_cnmf_consensus, get_top_indices_fast, annotate_genes_to_excel, \ - rename_and_move_files_NMF, rename_all_NMF, compile_results, \ - filter_noncoding_genes, filter_noncoding_genes_gtf, \ - add_gene_names_from_gtf + run_cnmf_consensus, + get_top_indices_fast, + annotate_genes_to_excel, + rename_and_move_files_NMF, + rename_all_NMF, + compile_results, + filter_noncoding_genes, + filter_noncoding_genes_gtf, + add_gene_names_from_gtf, ) from Stage1_Inference.src.plot_diagnostics import generate_all_plots def main(): - parser = argparse.ArgumentParser( - description="sk-cNMF inference pipeline" + parser = argparse.ArgumentParser(description="sk-cNMF inference pipeline") + + # IO + parser.add_argument( + "--counts_fn", + type=str, + required=True, + help="Path to input counts file (e.g., .h5mu or .h5ad)", + ) + parser.add_argument( + "--output_directory", + type=str, + required=True, + help="Directory where all outputs will be saved", + ) + parser.add_argument( + "--run_name", + type=str, + required=True, + help="Name for this cNMF run (used for output file naming)", + ) + parser.add_argument( + "--nmf_seeds_path", + type=str, + help="Path to .npy file containing NMF random seeds", + default=None, ) - - #IO - parser.add_argument('--counts_fn', type=str, required=True, help='Path to input counts file (e.g., .h5mu or .h5ad)') - parser.add_argument('--output_directory', type=str, required=True, help='Directory where all outputs will be saved') - parser.add_argument('--run_name', type=str, required=True, help='Name for this cNMF run (used for output file naming)') - parser.add_argument('--nmf_seeds_path', type=str, help='Path to .npy file containing NMF random seeds', default = None) - # cNMF parameters - parser.add_argument('--numiter', type = int, default = 10, help='Number of NMF replicates to run (default: 10)') - parser.add_argument('--numhvgenes', type = int, default = 5451, help='Number of highly variable genes to use (default: 5451)') - parser.add_argument('--seed', type = int, default = 14, help='Random seed for reproducibility (default: 14)') - parser.add_argument('--K', nargs='*', type=int, default=[30, 50, 70, 80, 100, 200, 300], help='List of K values (number of components) to test') - parser.add_argument('--init', type = str, default = 'random', help='Initialization method for NMF (default: random)') - parser.add_argument('--loss', default = 'frobenius', help='Loss function for NMF (default: frobenius)') - parser.add_argument('--algo', type = str, default = 'mu', help='Algorithm for NMF optimization (default: mu - multiplicative update)') - parser.add_argument('--max_NMF_iter', type = int , default = 500, help='Maximum number of iterations for NMF (default: 500)') - parser.add_argument('--tol', type = float , default = 1e4, help='Tolerance for NMF convergence (default: 1e4)') - parser.add_argument('--sel_threshs', nargs='*', type=float, default=[0.2, 2.0], help='Density threshold(s) for consensus matrix filtering') + parser.add_argument( + "--numiter", + type=int, + default=10, + help="Number of NMF replicates to run (default: 10)", + ) + parser.add_argument( + "--numhvgenes", + type=int, + default=5451, + help="Number of highly variable genes to use (default: 5451)", + ) + parser.add_argument( + "--seed", + type=int, + default=14, + help="Random seed for reproducibility (default: 14)", + ) + parser.add_argument( + "--K", + nargs="*", + type=int, + default=[30, 50, 70, 80, 100, 200, 300], + help="List of K values (number of components) to test", + ) + parser.add_argument( + "--init", + type=str, + default="random", + help="Initialization method for NMF (default: random)", + ) + parser.add_argument( + "--loss", default="frobenius", help="Loss function for NMF (default: frobenius)" + ) + parser.add_argument( + "--algo", + type=str, + default="mu", + help="Algorithm for NMF optimization (default: mu - multiplicative update)", + ) + parser.add_argument( + "--max_NMF_iter", + type=int, + default=500, + help="Maximum number of iterations for NMF (default: 500)", + ) + parser.add_argument( + "--tol", + type=float, + default=1e4, + help="Tolerance for NMF convergence (default: 1e4)", + ) + parser.add_argument( + "--sel_threshs", + nargs="*", + type=float, + default=[0.2, 2.0], + help="Density threshold(s) for consensus matrix filtering", + ) # annotation parameters - parser.add_argument('--species', type=str, required=True, help='Species for gene annotation (e.g., human, mouse)') - parser.add_argument('--parallel_running', action="store_true", help='If set, enables parallel processing mode for combining results from multiple K values') - parser.add_argument('--num_gene', type = int, default = 300, help='Number of top genes to use for program annotation (default: 300)') - parser.add_argument('--run_refit', action="store_true", help='If set, run the combine and consensus steps after factorization') - parser.add_argument('--run_complie_annotation', action="store_true", help='If set, compile results and generate gene annotations for all K values') - parser.add_argument('--run_gene_annotation', action="store_true", help='If set, run MyGene annotation (needs network)') - parser.add_argument('--run_factorize', action="store_true", help='If set, run the NMF factorization step') - parser.add_argument('--run_diagnostic_plots', action="store_true", help='Generate diagnostic plots (elbow curves, usage heatmaps, loading violins)') - parser.add_argument('--skip_existing', action="store_true", help='If set, skip NMF replicates already completed on disk (pause/resume mode). Default: re-run all replicates from scratch.') + parser.add_argument( + "--species", + type=str, + required=True, + help="Species for gene annotation (e.g., human, mouse)", + ) + parser.add_argument( + "--parallel_running", + action="store_true", + help="If set, enables parallel processing mode for combining results from multiple K values", + ) + parser.add_argument( + "--num_gene", + type=int, + default=300, + help="Number of top genes to use for program annotation (default: 300)", + ) + parser.add_argument( + "--run_refit", + action="store_true", + help="If set, run the combine and consensus steps after factorization", + ) + parser.add_argument( + "--run_complie_annotation", + action="store_true", + help="If set, compile results and generate gene annotations for all K values", + ) + parser.add_argument( + "--run_gene_annotation", + action="store_true", + help="If set, run MyGene annotation (needs network)", + ) + parser.add_argument( + "--run_factorize", + action="store_true", + help="If set, run the NMF factorization step", + ) + parser.add_argument( + "--run_diagnostic_plots", + action="store_true", + help="Generate diagnostic plots (elbow curves, usage heatmaps, loading violins)", + ) + parser.add_argument( + "--skip_existing", + action="store_true", + help="If set, skip NMF replicates already completed on disk (pause/resume mode). Default: re-run all replicates from scratch.", + ) # keys - parser.add_argument('--data_key', help='Key to access gene expression data in MuData object (default: rna)', type=str, default="rna") - parser.add_argument('--prog_key', help='Key to store cNMF programs in MuData object (default: cNMF)', type=str, default="cNMF") - parser.add_argument('--categorical_key', help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", type=str, default="sample") - parser.add_argument('--guide_names_key', help='Key in .uns to access guide names (default: guide_names)', type=str, default="guide_names") - parser.add_argument('--guide_targets_key', help='Key in .uns to access guide target genes (default: guide_targets)', type=str, default="guide_targets") - parser.add_argument('--guide_assignment_key', help='Key in .obsm to access guide assignment matrix (default: guide_assignment)', type=str, default="guide_assignment") - parser.add_argument('--gene_names_key', type=str, default=None, - help="Column in adata.var with gene names to use in compiled results (e.g. 'symbol'). If None, uses var_names.") + parser.add_argument( + "--data_key", + help="Key to access gene expression data in MuData object (default: rna)", + type=str, + default="rna", + ) + parser.add_argument( + "--prog_key", + help="Key to store cNMF programs in MuData object (default: cNMF)", + type=str, + default="cNMF", + ) + parser.add_argument( + "--categorical_key", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + type=str, + default="sample", + ) + parser.add_argument( + "--guide_names_key", + help="Key in .uns to access guide names (default: guide_names)", + type=str, + default="guide_names", + ) + parser.add_argument( + "--guide_targets_key", + help="Key in .uns to access guide target genes (default: guide_targets)", + type=str, + default="guide_targets", + ) + parser.add_argument( + "--guide_assignment_key", + help="Key in .obsm to access guide assignment matrix (default: guide_assignment)", + type=str, + default="guide_assignment", + ) + parser.add_argument( + "--gene_names_key", + type=str, + default=None, + help="Column in adata.var with gene names to use in compiled results (e.g. 'symbol'). If None, uses var_names.", + ) # preprocessing - parser.add_argument('--remove_noncoding', action='store_true', help='If set, remove non-coding genes before factorization. With --gtf_path, keeps only genes annotated as protein_coding in the GTF (matched by Ensembl ID); otherwise falls back to the Ensembl-prefix heuristic on --gene_names_key.') - parser.add_argument('--ensembl_prefix', type=str, default='ENSG', help='Ensembl ID prefix used by the (fallback) prefix-based non-coding filter (default: ENSG)') - parser.add_argument('--gtf_path', type=str, default=None, help='Path to a GENCODE/Ensembl GTF(.gz). Enables GTF-based --remove_noncoding and --add_gene_names_from_gtf.') - parser.add_argument('--gene_id_key', type=str, default='gene_id', help='Column in adata.var holding Ensembl gene IDs (used for GTF-based filtering / gene-name annotation)') - parser.add_argument('--add_gene_names_from_gtf', action='store_true', help='Populate adata.var[gene_names_key] with gene symbols looked up from --gtf_path by Ensembl ID (from --gene_id_key). Requires --gtf_path.') - - + parser.add_argument( + "--remove_noncoding", + action="store_true", + help="If set, remove non-coding genes before factorization. With --gtf_path, keeps only genes annotated as protein_coding in the GTF (matched by Ensembl ID); otherwise falls back to the Ensembl-prefix heuristic on --gene_names_key.", + ) + parser.add_argument( + "--ensembl_prefix", + type=str, + default="ENSG", + help="Ensembl ID prefix used by the (fallback) prefix-based non-coding filter (default: ENSG)", + ) + parser.add_argument( + "--gtf_path", + type=str, + default=None, + help="Path to a GENCODE/Ensembl GTF(.gz). Enables GTF-based --remove_noncoding and --add_gene_names_from_gtf.", + ) + parser.add_argument( + "--gene_id_key", + type=str, + default="gene_id", + help="Column in adata.var holding Ensembl gene IDs (used for GTF-based filtering / gene-name annotation)", + ) + parser.add_argument( + "--add_gene_names_from_gtf", + action="store_true", + help="Populate adata.var[gene_names_key] with gene symbols looked up from --gtf_path by Ensembl ID (from --gene_id_key). Requires --gtf_path.", + ) args = parser.parse_args() @@ -96,37 +261,36 @@ def main(): else: nmf_seeds = None - # create output directory - run_dir = f'{args.output_directory}/{args.run_name}' - inference_dir = f'{run_dir}/Inference' + run_dir = f"{args.output_directory}/{args.run_name}" + inference_dir = f"{run_dir}/Inference" os.makedirs(inference_dir, exist_ok=True) # --- Save config (incl. SLURM info) --- slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), - 'num_nodes': os.environ.get('SLURM_JOB_NUM_NODES'), - 'ntasks': os.environ.get('SLURM_NTASKS'), - 'cpus_per_task': os.environ.get('SLURM_CPUS_PER_TASK'), - 'mem_per_node': os.environ.get('SLURM_MEM_PER_NODE'), - 'mem_per_cpu': os.environ.get('SLURM_MEM_PER_CPU'), - 'gres': os.environ.get('SLURM_JOB_GRES'), - 'gpu_device_ids': os.environ.get('SLURM_GPUS_ON_NODE'), - 'gpu_type': os.environ.get('SLURM_JOB_GPUS'), - 'time_limit': os.environ.get('SLURM_JOB_TIMELIMIT'), - 'time_remaining': os.environ.get('SLURM_TIMELIMIT'), - 'submit_dir': os.environ.get('SLURM_SUBMIT_DIR'), - 'submit_host': os.environ.get('SLURM_SUBMIT_HOST'), - 'constraint': os.environ.get('SLURM_JOB_CONSTRAINT'), - 'array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), + "num_nodes": os.environ.get("SLURM_JOB_NUM_NODES"), + "ntasks": os.environ.get("SLURM_NTASKS"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + "mem_per_node": os.environ.get("SLURM_MEM_PER_NODE"), + "mem_per_cpu": os.environ.get("SLURM_MEM_PER_CPU"), + "gres": os.environ.get("SLURM_JOB_GRES"), + "gpu_device_ids": os.environ.get("SLURM_GPUS_ON_NODE"), + "gpu_type": os.environ.get("SLURM_JOB_GPUS"), + "time_limit": os.environ.get("SLURM_JOB_TIMELIMIT"), + "time_remaining": os.environ.get("SLURM_TIMELIMIT"), + "submit_dir": os.environ.get("SLURM_SUBMIT_DIR"), + "submit_host": os.environ.get("SLURM_SUBMIT_HOST"), + "constraint": os.environ.get("SLURM_JOB_CONSTRAINT"), + "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" - config_to_save = {'script_args': vars(args), 'slurm_info': slurm_info} - with open(f'{inference_dir}/config_{job_id}.yml', 'w') as f: + config_to_save = {"script_args": vars(args), "slurm_info": slurm_info} + with open(f"{inference_dir}/config_{job_id}.yml", "w") as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) # --- Optionally filter non-coding genes --- @@ -161,59 +325,86 @@ def main(): ) # running cnmf - cnmf_obj = cnmf.cNMF(output_dir=run_dir, name='Inference') - - cnmf_obj.prepare(counts_fn= args.counts_fn, components= args.K, n_iter= args.numiter, densify=False, tpm_fn=None, seed= args.seed, - beta_loss = args.loss,num_highvar_genes=args.numhvgenes, genes_file=None, - alpha_usage=0.0, alpha_spectra=0.0, init=args.init, max_NMF_iter=args.max_NMF_iter, algo = args.algo, tol = args.tol, nmf_seeds = nmf_seeds) + cnmf_obj = cnmf.cNMF(output_dir=run_dir, name="Inference") + + cnmf_obj.prepare( + counts_fn=args.counts_fn, + components=args.K, + n_iter=args.numiter, + densify=False, + tpm_fn=None, + seed=args.seed, + beta_loss=args.loss, + num_highvar_genes=args.numhvgenes, + genes_file=None, + alpha_usage=0.0, + alpha_spectra=0.0, + init=args.init, + max_NMF_iter=args.max_NMF_iter, + algo=args.algo, + tol=args.tol, + nmf_seeds=nmf_seeds, + ) if args.run_factorize: - - cnmf_obj.factorize(total_workers = 1, skip_completed_runs=args.skip_existing) - # the original cNMF parallelism needs additional indexing of worker id and total worker - # here we just use default + cnmf_obj.factorize(total_workers=1, skip_completed_runs=args.skip_existing) + # the original cNMF parallelism needs additional indexing of worker id and total worker + # here we just use default if args.run_refit: - cnmf_obj.combine() cnmf_obj.k_selection_plot() # Consensus plots with all k to choose thresh - run_cnmf_consensus(cnmf_obj, - components=args.K, - density_thresholds=args.sel_threshs) + run_cnmf_consensus( + cnmf_obj, components=args.K, density_thresholds=args.sel_threshs + ) if args.run_complie_annotation: - # Save all cNMF scores in separate mudata objects - compile_results(run_dir, 'Inference', components= args.K, sel_threshs = args.sel_threshs, - guide_names_key = args.guide_names_key, guide_targets_key = args.guide_targets_key, categorical_key= args.categorical_key, - guide_assignment_key = args.guide_assignment_key, gene_names_key = args.gene_names_key ) + compile_results( + run_dir, + "Inference", + components=args.K, + sel_threshs=args.sel_threshs, + guide_names_key=args.guide_names_key, + guide_targets_key=args.guide_targets_key, + categorical_key=args.categorical_key, + guide_assignment_key=args.guide_assignment_key, + gene_names_key=args.gene_names_key, + ) - #---annotate ---- + # ---annotate ---- if args.run_gene_annotation: print("Annotating top genes per program via MyGene...") - os.makedirs(f'{inference_dir}/Annotation', exist_ok=True) + os.makedirs(f"{inference_dir}/Annotation", exist_ok=True) # annotation for all K for i in args.sel_threshs: for k in args.K: - df = pd.read_csv('{inference_dir}/Inference.gene_spectra_score.k_{k}.dt_{sel_thresh}.txt'.format( - inference_dir=inference_dir, - k=k, - sel_thresh = str(i).replace('.','_')), - sep='\t', index_col=0) + df = pd.read_csv( + "{inference_dir}/Inference.gene_spectra_score.k_{k}.dt_{sel_thresh}.txt".format( + inference_dir=inference_dir, + k=k, + sel_thresh=str(i).replace(".", "_"), + ), + sep="\t", + index_col=0, + ) overlap = get_top_indices_fast(df, gene_num=args.num_gene) - annotate_genes_to_excel(overlap, species = args.species, output_file= f"{inference_dir}/Annotation/{k}_{str(i).replace('.','_')}.xlsx") - + annotate_genes_to_excel( + overlap, + species=args.species, + output_file=f"{inference_dir}/Annotation/{k}_{str(i).replace('.', '_')}.xlsx", + ) # --- Diagnostic plots --- if args.run_diagnostic_plots: print("Generating diagnostic plots...") generate_all_plots( run_dir=run_dir, - run_name='Inference', + run_name="Inference", K_list=args.K, sel_thresh_list=args.sel_threshs, categorical_key=args.categorical_key, @@ -226,7 +417,7 @@ def main(): rename_all_NMF( source_folder=f"{args.output_directory}/{args.run_name}/{args.run_name}", destination_folder=f"{args.output_directory}/{args.run_name}/{args.run_name}_all/Inference/cnmf_tmp", - file_name_input='Inference', + file_name_input="Inference", file_name_output="Inference_all", len=args.numiter, components=args.K, @@ -236,10 +427,5 @@ def main(): return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) - - - - - diff --git a/src/Stage1_Inference/src/__init__.py b/src/Stage1_Inference/src/__init__.py index cff635a..aa52152 100644 --- a/src/Stage1_Inference/src/__init__.py +++ b/src/Stage1_Inference/src/__init__.py @@ -1,9 +1,14 @@ -from .run_cNMF import run_cnmf_consensus, get_top_indices_fast, annotate_genes_to_excel, \ - rename_and_move_files_NMF, rename_all_NMF, compile_results, \ - filter_noncoding_genes, filter_noncoding_genes_gtf, \ - load_protein_coding_ids, load_gene_id_to_name, \ - add_gene_names_from_gtf +from .run_cNMF import ( + run_cnmf_consensus, + get_top_indices_fast, + annotate_genes_to_excel, + rename_and_move_files_NMF, + rename_all_NMF, + compile_results, + filter_noncoding_genes, + filter_noncoding_genes_gtf, + load_protein_coding_ids, + load_gene_id_to_name, + add_gene_names_from_gtf, +) from .plot_diagnostics import generate_all_plots, plot_top_gene_overlap - - - diff --git a/src/Stage1_Inference/src/plot_diagnostics.py b/src/Stage1_Inference/src/plot_diagnostics.py index bda7dc3..f637eda 100644 --- a/src/Stage1_Inference/src/plot_diagnostics.py +++ b/src/Stage1_Inference/src/plot_diagnostics.py @@ -16,6 +16,7 @@ import logging import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np @@ -29,6 +30,7 @@ # 1. Elbow curves: ranked gene loadings per program # --------------------------------------------------------------------------- + def plot_elbow_curves(loadings_df, output_path, n_top=2000): """Plot sorted log-loadings per program to visualize the signal/noise boundary. @@ -82,6 +84,7 @@ def plot_elbow_curves(loadings_df, output_path, n_top=2000): # 2. K-selection plot (stability vs error) # --------------------------------------------------------------------------- + def plot_k_selection(run_dir, run_name, output_path): """Plot the k-selection PNG that cNMF generates during combine/consensus. @@ -90,6 +93,7 @@ def plot_k_selection(run_dir, run_name, output_path): cnmf_plot = os.path.join(run_dir, run_name, f"{run_name}.k_selection.png") if os.path.exists(cnmf_plot): import shutil + # copyfile, not copy2: copy2 preserves mtime via os.utime(), which fails # with EPERM when overwriting a file owned by another user (shared group dirs) shutil.copyfile(cnmf_plot, output_path) @@ -102,6 +106,7 @@ def plot_k_selection(run_dir, run_name, output_path): # 3. Program usage heatmap (cells x programs, grouped by batch) # --------------------------------------------------------------------------- + def plot_usage_heatmap(mdata, output_path, categorical_key="batch"): """Heatmap of program usage scores, cells grouped by batch. @@ -165,6 +170,7 @@ def plot_usage_heatmap(mdata, output_path, categorical_key="batch"): # 4. Loading distribution violin plot # --------------------------------------------------------------------------- + def plot_loading_violins(loadings_df, output_path): """Violin plot of gene loading distributions per program. @@ -207,8 +213,10 @@ def plot_loading_violins(loadings_df, output_path): # 5. Program-program top-N gene overlap clustermap # --------------------------------------------------------------------------- -def plot_top_gene_overlap(spectra_df, output_path, top_n=300, title=None, - csv_path=None): + +def plot_top_gene_overlap( + spectra_df, output_path, top_n=300, title=None, csv_path=None +): """Clustermap of shared top-N genes between every pair of programs. For each program take its top-N genes (largest spectra value), then fill @@ -266,11 +274,13 @@ def plot_top_gene_overlap(spectra_df, output_path, top_n=300, title=None, f"Off-diagonal shared top-{n_top} genes: min {offdiag.min()}, " f"mean {offdiag.mean():.1f}, max {offdiag.max()}" ) - pairs = pd.DataFrame({ - "prog_a": [progs[i] for i in iu[0]], - "prog_b": [progs[j] for j in iu[1]], - "shared": offdiag, - }).sort_values("shared", ascending=False) + pairs = pd.DataFrame( + { + "prog_a": [progs[i] for i in iu[0]], + "prog_b": [progs[j] for j in iu[1]], + "shared": offdiag, + } + ).sort_values("shared", ascending=False) for _, r in pairs.head(10).iterrows(): logger.info(f" {r['prog_a']} <-> {r['prog_b']}: {r['shared']} shared") @@ -279,12 +289,17 @@ def plot_top_gene_overlap(spectra_df, output_path, top_n=300, title=None, size = max(6, min(24, 0.35 * n_progs + 4)) g = sns.clustermap( overlap, - row_cluster=True, col_cluster=True, - cmap="RdBu_r", vmin=0, vmax=n_top, + row_cluster=True, + col_cluster=True, + cmap="RdBu_r", + vmin=0, + vmax=n_top, figsize=(size + 1, size), - linewidths=0.1, linecolor="0.9", + linewidths=0.1, + linecolor="0.9", cbar_kws={"label": f"shared genes (top {n_top})"}, - xticklabels=True, yticklabels=True, + xticklabels=True, + yticklabels=True, ) g.ax_row_dendrogram.set_visible(False) g.ax_col_dendrogram.set_visible(False) @@ -296,7 +311,8 @@ def plot_top_gene_overlap(spectra_df, output_path, top_n=300, title=None, g.figure.suptitle( f"{title}\nhierarchically clustered; off-diagonal = redundant / " "co-varying programs", - y=1.02, fontsize=11, + y=1.02, + fontsize=11, ) g.savefig(output_path, dpi=150, bbox_inches="tight") plt.close(g.figure) @@ -308,6 +324,7 @@ def plot_top_gene_overlap(spectra_df, output_path, top_n=300, title=None, # 6. Consensus clustering plot check # --------------------------------------------------------------------------- + def collect_clustering_plots(run_dir, run_name, K_list, sel_thresh_list, output_dir): """Copy consensus clustering PNGs generated by cNMF into the plots dir.""" import shutil @@ -328,8 +345,15 @@ def collect_clustering_plots(run_dir, run_name, K_list, sel_thresh_list, output_ # Main entry point # --------------------------------------------------------------------------- -def generate_all_plots(run_dir, run_name="Inference", K_list=None, sel_thresh_list=None, - categorical_key="batch", top_n_overlap=300): + +def generate_all_plots( + run_dir, + run_name="Inference", + K_list=None, + sel_thresh_list=None, + categorical_key="batch", + top_n_overlap=300, +): """Generate all diagnostic plots for an inference run. Args: @@ -364,7 +388,8 @@ def generate_all_plots(run_dir, run_name="Inference", K_list=None, sel_thresh_li # Load gene spectra scores for elbow + violin plots spectra_path = os.path.join( - run_dir, run_name, + run_dir, + run_name, f"{run_name}.gene_spectra_score.k_{k}.dt_{thresh_str}.txt", ) if os.path.exists(spectra_path): @@ -383,7 +408,8 @@ def generate_all_plots(run_dir, run_name="Inference", K_list=None, sel_thresh_li # Top-N gene overlap clustermap from the median consensus spectra median_path = os.path.join( - run_dir, run_name, + run_dir, + run_name, f"{run_name}.spectra.k_{k}.dt_{thresh_str}.consensus.txt", ) if os.path.exists(median_path): diff --git a/src/Stage1_Inference/src/run_cNMF.py b/src/Stage1_Inference/src/run_cNMF.py index 7b47645..6fdcebb 100644 --- a/src/Stage1_Inference/src/run_cNMF.py +++ b/src/Stage1_Inference/src/run_cNMF.py @@ -6,7 +6,7 @@ import scanpy as sc from pathlib import Path import anndata -import muon +import muon from torch_cnmf import cNMF import os import sys @@ -15,11 +15,17 @@ from pathlib import Path from tqdm.auto import tqdm -def run_cnmf_consensus(cnmf_obj=None, output_dir=None, name=None, - components=[7,8,9,10], density_thresholds=[0.01, 0.05, 2.0]): + +def run_cnmf_consensus( + cnmf_obj=None, + output_dir=None, + name=None, + components=[7, 8, 9, 10], + density_thresholds=[0.01, 0.05, 2.0], +): """ Run cNMF consensus analysis for multiple k values and density thresholds. - + Args: cnmf_obj: Existing cNMF object, or None to create a new one output_dir: Directory to save results @@ -30,12 +36,14 @@ def run_cnmf_consensus(cnmf_obj=None, output_dir=None, name=None, if cnmf_obj is None: cnmf_obj = cNMF(output_dir=output_dir, name=name) - for k in tqdm(components, desc='Running cNMF'): + for k in tqdm(components, desc="Running cNMF"): for thresh in density_thresholds: cnmf_obj.consensus(k=k, density_threshold=thresh, show_clustering=True) -def filter_noncoding_genes(counts_fn, inference_dir, gene_names_key='symbol', ensembl_prefix='ENSG'): +def filter_noncoding_genes( + counts_fn, inference_dir, gene_names_key="symbol", ensembl_prefix="ENSG" +): """Filter out non-coding genes (whose name starts with `ensembl_prefix`) and write a copy. Returns the path to the filtered .h5ad. Caller should rebind whatever variable @@ -53,17 +61,17 @@ def filter_noncoding_genes(counts_fn, inference_dir, gene_names_key='symbol', en adata = sc.read(counts_fn) mask = ~adata.var[gene_names_key].str.startswith(ensembl_prefix) adata = adata[:, mask].copy() - filtered_path = f'{inference_dir}/adata_without_noncoding.h5ad' + filtered_path = f"{inference_dir}/adata_without_noncoding.h5ad" adata.write(filtered_path) return filtered_path def _strip_ensembl_version(ids): """Return Ensembl IDs with any trailing `.` removed (e.g. ENSG0001.12 -> ENSG0001).""" - return pd.Index(ids).astype(str).str.split('.').str[0] + return pd.Index(ids).astype(str).str.split(".").str[0] -def load_protein_coding_ids(gtf_path, gene_type='protein_coding'): +def load_protein_coding_ids(gtf_path, gene_type="protein_coding"): """Parse a (optionally gzipped) GENCODE/Ensembl GTF and return the set of version-stripped gene IDs whose `gene` feature line has the requested `gene_type`. @@ -80,14 +88,14 @@ def load_protein_coding_ids(gtf_path, gene_type='protein_coding'): gene_id_re = re.compile(r'gene_id "([^"]+)"') gene_type_re = re.compile(r'gene_type "([^"]+)"') - opener = gzip.open if str(gtf_path).endswith('.gz') else open + opener = gzip.open if str(gtf_path).endswith(".gz") else open keep = set() - with opener(gtf_path, 'rt') as fh: + with opener(gtf_path, "rt") as fh: for line in fh: - if line.startswith('#'): + if line.startswith("#"): continue - cols = line.split('\t', 8) - if len(cols) < 9 or cols[2] != 'gene': + cols = line.split("\t", 8) + if len(cols) < 9 or cols[2] != "gene": continue attrs = cols[8] gt = gene_type_re.search(attrs) @@ -95,12 +103,17 @@ def load_protein_coding_ids(gtf_path, gene_type='protein_coding'): continue gid = gene_id_re.search(attrs) if gid is not None: - keep.add(gid.group(1).split('.')[0]) + keep.add(gid.group(1).split(".")[0]) return keep -def filter_noncoding_genes_gtf(counts_fn, inference_dir, gtf_path, - gene_id_key='gene_id', gene_type='protein_coding'): +def filter_noncoding_genes_gtf( + counts_fn, + inference_dir, + gtf_path, + gene_id_key="gene_id", + gene_type="protein_coding", +): """Keep only genes whose Ensembl ID is annotated as `gene_type` in `gtf_path`; drop the rest. Unlike `filter_noncoding_genes` (a prefix heuristic on the symbol column), this @@ -133,7 +146,9 @@ def filter_noncoding_genes_gtf(counts_fn, inference_dir, gtf_path, if gene_id_key in adata.var.columns: raw_ids = adata.var[gene_id_key] else: - print(f"WARNING: '{gene_id_key}' not in adata.var columns; falling back to var_names.") + print( + f"WARNING: '{gene_id_key}' not in adata.var columns; falling back to var_names." + ) raw_ids = adata.var_names base_ids = _strip_ensembl_version(raw_ids) mask = np.asarray(base_ids.isin(coding_ids)) @@ -144,11 +159,13 @@ def filter_noncoding_genes_gtf(counts_fn, inference_dir, gtf_path, f"0 of {adata.n_vars} genes matched '{gene_type}' IDs from the GTF. " f"Check that adata.var['{gene_id_key}'] holds Ensembl gene IDs (e.g. ENSG...)." ) - print(f"GTF filter: keeping {n_keep}/{adata.n_vars} genes " - f"({gene_type} in {os.path.basename(str(gtf_path))}).") + print( + f"GTF filter: keeping {n_keep}/{adata.n_vars} genes " + f"({gene_type} in {os.path.basename(str(gtf_path))})." + ) adata = adata[:, mask].copy() - filtered_path = f'{inference_dir}/adata_without_noncoding.h5ad' + filtered_path = f"{inference_dir}/adata_without_noncoding.h5ad" adata.write(filtered_path) return filtered_path @@ -170,26 +187,31 @@ def load_gene_id_to_name(gtf_path): gene_id_re = re.compile(r'gene_id "([^"]+)"') gene_name_re = re.compile(r'gene_name "([^"]+)"') - opener = gzip.open if str(gtf_path).endswith('.gz') else open + opener = gzip.open if str(gtf_path).endswith(".gz") else open id2name = {} - with opener(gtf_path, 'rt') as fh: + with opener(gtf_path, "rt") as fh: for line in fh: - if line.startswith('#'): + if line.startswith("#"): continue - cols = line.split('\t', 8) - if len(cols) < 9 or cols[2] != 'gene': + cols = line.split("\t", 8) + if len(cols) < 9 or cols[2] != "gene": continue attrs = cols[8] gid = gene_id_re.search(attrs) name = gene_name_re.search(attrs) if gid is not None and name is not None: - id2name[gid.group(1).split('.')[0]] = name.group(1) + id2name[gid.group(1).split(".")[0]] = name.group(1) return id2name -def add_gene_names_from_gtf(counts_fn, inference_dir, gtf_path, - gene_id_key='gene_id', gene_names_key='symbol', - keep_unmatched=True): +def add_gene_names_from_gtf( + counts_fn, + inference_dir, + gtf_path, + gene_id_key="gene_id", + gene_names_key="symbol", + keep_unmatched=True, +): """Add/overwrite adata.var[gene_names_key] with gene symbols looked up from the GTF by (version-stripped) Ensembl ID, then write a copy. @@ -221,7 +243,9 @@ def add_gene_names_from_gtf(counts_fn, inference_dir, gtf_path, if gene_id_key in adata.var.columns: raw_ids = adata.var[gene_id_key] else: - print(f"WARNING: '{gene_id_key}' not in adata.var columns; falling back to var_names.") + print( + f"WARNING: '{gene_id_key}' not in adata.var columns; falling back to var_names." + ) raw_ids = adata.var_names base_ids = _strip_ensembl_version(raw_ids) @@ -237,23 +261,33 @@ def add_gene_names_from_gtf(counts_fn, inference_dir, gtf_path, names = mapped.where(mapped.notna(), fallback) adata.var[gene_names_key] = np.asarray(names) - print(f"GTF gene names: matched {n_matched}/{adata.n_vars} genes; " - f"wrote adata.var['{gene_names_key}'] from {os.path.basename(str(gtf_path))}.") + print( + f"GTF gene names: matched {n_matched}/{adata.n_vars} genes; " + f"wrote adata.var['{gene_names_key}'] from {os.path.basename(str(gtf_path))}." + ) - annotated_path = f'{inference_dir}/adata_with_gene_names.h5ad' + annotated_path = f"{inference_dir}/adata_with_gene_names.h5ad" adata.write(annotated_path) return annotated_path -def compile_results(output_directory, run_name, sel_threshs = [2.0], components = [30, 50, 60, 80, 100, 200, 250, 300], - guide_names_key = "guide_names", guide_targets_key = "guide_targets", categorical_key= 'batch', guide_assignment_key ="guide_assignment", - gene_names_key = None ): +def compile_results( + output_directory, + run_name, + sel_threshs=[2.0], + components=[30, 50, 60, 80, 100, 200, 250, 300], + guide_names_key="guide_names", + guide_targets_key="guide_targets", + categorical_key="batch", + guide_assignment_key="guide_assignment", + gene_names_key=None, +): """ Compile cNMF results into correct format for downstream evaluation pipeline. - + Processes cNMF output files and reorganizes them into standardized formats including scores, loadings, and AnnData objects for further analysis. - + Args: output_directory: Base directory containing cNMF results run_name: Name of the cNMF run @@ -261,58 +295,72 @@ def compile_results(output_directory, run_name, sel_threshs = [2.0], components components: List of k values (number of components) to process """ + adata_ = anndata.read_h5ad( + "{output_directory}/{run_name}/cnmf_tmp/{run_name}.tpm.h5ad".format( + output_directory=output_directory, run_name=run_name + ) + ) - adata_ = anndata.read_h5ad('{output_directory}/{run_name}/cnmf_tmp/{run_name}.tpm.h5ad'.format( - output_directory=output_directory, - run_name = run_name)) - for i in sel_threshs: for k in components: - - scores = pd.read_csv('{output_directory}/{run_name}/{run_name}.usages.k_{k}.dt_{sel_thresh}.consensus.txt'.format( - output_directory=output_directory, - run_name = run_name, - k=k, - sel_thresh = str(i).replace('.','_')), - sep='\t', index_col=0) - - loadings = pd.read_csv('{output_directory}/{run_name}/{run_name}.gene_spectra_score.k_{k}.dt_{sel_thresh}.txt'.format( - output_directory=output_directory, - run_name = run_name, - k=k, - sel_thresh = str(i).replace('.','_')), - sep='\t', index_col=0) - - - os.makedirs((f'{output_directory}/{run_name}/loading'), exist_ok=True) - - - scores.to_csv('{output_directory}/{run_name}/loading/cNMF_scores_{k}_{sel_thresh}.txt'.format( - output_directory=output_directory, - run_name = run_name, - k=k, - sel_thresh = str(i).replace('.','_')), sep='\t') - loadings.T.to_csv('{output_directory}/{run_name}/loading/cNMF_loadings_{k}_{sel_thresh}.txt'.format( - output_directory=output_directory, - run_name = run_name, - k=k, - sel_thresh = str(i).replace('.','_')), sep='\t') + scores = pd.read_csv( + "{output_directory}/{run_name}/{run_name}.usages.k_{k}.dt_{sel_thresh}.consensus.txt".format( + output_directory=output_directory, + run_name=run_name, + k=k, + sel_thresh=str(i).replace(".", "_"), + ), + sep="\t", + index_col=0, + ) + + loadings = pd.read_csv( + "{output_directory}/{run_name}/{run_name}.gene_spectra_score.k_{k}.dt_{sel_thresh}.txt".format( + output_directory=output_directory, + run_name=run_name, + k=k, + sel_thresh=str(i).replace(".", "_"), + ), + sep="\t", + index_col=0, + ) + + os.makedirs((f"{output_directory}/{run_name}/loading"), exist_ok=True) + + scores.to_csv( + "{output_directory}/{run_name}/loading/cNMF_scores_{k}_{sel_thresh}.txt".format( + output_directory=output_directory, + run_name=run_name, + k=k, + sel_thresh=str(i).replace(".", "_"), + ), + sep="\t", + ) + loadings.T.to_csv( + "{output_directory}/{run_name}/loading/cNMF_loadings_{k}_{sel_thresh}.txt".format( + output_directory=output_directory, + run_name=run_name, + k=k, + sel_thresh=str(i).replace(".", "_"), + ), + sep="\t", + ) adata_.var_names_make_unique() adata_.obs_names_make_unique() prog_data = anndata.AnnData(X=scores.values, obs=adata_.obs) prog_data.var_names = scores.columns.values - prog_data.varm['loadings'] = loadings.values + prog_data.varm["loadings"] = loadings.values # Use gene_names_key column if provided, otherwise fall back to var_names / loadings columns if gene_names_key is not None and gene_names_key in adata_.var.columns: # Swap var_names to the gene-name column, preserving the original IDs # (e.g. Ensembl) in a 'gene_id' column. Guarded so the original IDs are # only saved on the first loop iteration (var_names is already symbols after). - if 'gene_id' not in adata_.var.columns: - adata_.var['gene_id'] = adata_.var_names + if "gene_id" not in adata_.var.columns: + adata_.var["gene_id"] = adata_.var_names adata_.var_names = adata_.var[gene_names_key].astype(str).values - + adata_.var_names_make_unique() adata_.var[gene_names_key] = adata_.var_names.values @@ -320,224 +368,243 @@ def compile_results(output_directory, run_name, sel_threshs = [2.0], components else: gene_names = adata_.var_names.values - prog_data.uns['var_names'] = gene_names - + prog_data.uns["var_names"] = gene_names # Make adata - os.makedirs((f'{output_directory}/{run_name}/prog_data'), exist_ok=True) - prog_data.write('{output_directory}/{run_name}/prog_data/NMF_{k}_{sel_thresh}.h5ad'.format( - output_directory=output_directory, - run_name = run_name, - k=k, - sel_thresh = str(i).replace('.','_'))) + os.makedirs((f"{output_directory}/{run_name}/prog_data"), exist_ok=True) + prog_data.write( + "{output_directory}/{run_name}/prog_data/NMF_{k}_{sel_thresh}.h5ad".format( + output_directory=output_directory, + run_name=run_name, + k=k, + sel_thresh=str(i).replace(".", "_"), + ) + ) # Make mdata - mdata = muon.MuData({'rna': adata_, 'cNMF': prog_data}) + mdata = muon.MuData({"rna": adata_, "cNMF": prog_data}) if guide_names_key and guide_names_key in adata_.uns: - mdata['cNMF'].uns[guide_names_key] = adata_.uns[guide_names_key] + mdata["cNMF"].uns[guide_names_key] = adata_.uns[guide_names_key] if guide_targets_key and guide_targets_key in adata_.uns: - mdata['cNMF'].uns[guide_targets_key] = adata_.uns[guide_targets_key] + mdata["cNMF"].uns[guide_targets_key] = adata_.uns[guide_targets_key] if categorical_key and categorical_key in adata_.obs.columns: - mdata['cNMF'].obs[categorical_key] = adata_.obs[categorical_key] + mdata["cNMF"].obs[categorical_key] = adata_.obs[categorical_key] if guide_assignment_key and guide_assignment_key in adata_.obsm: - mdata['cNMF'].obsm[guide_assignment_key] = adata_.obsm[guide_assignment_key] - if 'X_pca' in adata_.obsm: - mdata['cNMF'].obsm['X_pca'] = adata_.obsm['X_pca'] - if 'X_umap' in adata_.obsm: - mdata['cNMF'].obsm['X_umap'] = adata_.obsm['X_umap'] + mdata["cNMF"].obsm[guide_assignment_key] = adata_.obsm[ + guide_assignment_key + ] + if "X_pca" in adata_.obsm: + mdata["cNMF"].obsm["X_pca"] = adata_.obsm["X_pca"] + if "X_umap" in adata_.obsm: + mdata["cNMF"].obsm["X_umap"] = adata_.obsm["X_umap"] # var_names is now the de-duplicated gene-name column (or the original IDs # when no gene_names_key was given), so both cases collapse to one line - mdata['rna'].var['var_names'] = adata_.var_names.values - + mdata["rna"].var["var_names"] = adata_.var_names.values - os.makedirs((f'{output_directory}/{run_name}/adata'), exist_ok=True) - mdata.write('{output_directory}/{run_name}/adata/cNMF_{k}_{sel_thresh}.h5mu'.format( - output_directory=output_directory, - run_name = run_name, - k=k, - sel_thresh = str(i).replace('.','_'))) - + os.makedirs((f"{output_directory}/{run_name}/adata"), exist_ok=True) + mdata.write( + "{output_directory}/{run_name}/adata/cNMF_{k}_{sel_thresh}.h5mu".format( + output_directory=output_directory, + run_name=run_name, + k=k, + sel_thresh=str(i).replace(".", "_"), + ) + ) -def get_top_indices_fast(df, gene_num = 300): +def get_top_indices_fast(df, gene_num=300): """ Given a DataFrame from cNMF, return top N genes for each program. - + Args: df: DataFrame with programs as rows and genes as columns gene_num: Number of top genes to return per program - + Returns: DataFrame with programs as rows and top genes as columns """ - + # Get column names col_names = df.columns.values - + # Use argsort to get indices of top 300 values per row # argsort sorts in ascending order, so we use [:, -300:] and reverse top_indices = np.argsort(df.values, axis=1)[:, -gene_num:][:, ::-1] - + # Map indices to column names top_col_names = col_names[top_indices] - + # Create result DataFrame result_df = pd.DataFrame( - top_col_names, - index=df.index, - columns=[f'top_{i+1}' for i in range(gene_num)] + top_col_names, index=df.index, columns=[f"top_{i + 1}" for i in range(gene_num)] ) - result_df.index = [f'Program_{i}' for i in range(1,len(result_df)+1)] - + result_df.index = [f"Program_{i}" for i in range(1, len(result_df) + 1)] + return result_df -def annotate_genes_to_excel(df, species = 'human', output_file='gene_annotations.xlsx'): +def annotate_genes_to_excel(df, species="human", output_file="gene_annotations.xlsx"): """ Annotate genes and export to Excel file. - + Takes a DataFrame with programs as rows and genes as columns, queries MyGene database for annotations, and exports results to Excel. - + Args: df: DataFrame with rows for each program, columns for genes species: Species for gene annotation ('human' or 'mouse') output_file: Path for output Excel file - + Returns: Dictionary of annotations for each program """ - + # Initialize MyGene mg = mygene.MyGeneInfo() - + # Dictionary to store results for each column all_annotations = {} - + # Process each column for row_idx in df.index: # Get unique genes from column (remove NaN) genes = df.loc[row_idx].dropna().unique().tolist() - + if len(genes) == 0: print(f"Column '{row_idx}': No genes found") continue - + print(f"Annotating column '{row_idx}': {len(genes)} genes...") - + # Query MyGene for annotations results = mg.querymany( - genes, - scopes='symbol,alias,ensembl.gene', # Multiple search scopes - fields='symbol,name,entrezgene,summary,type_of_gene', + genes, + scopes="symbol,alias,ensembl.gene", # Multiple search scopes + fields="symbol,name,entrezgene,summary,type_of_gene", species=species, - returnall=True + returnall=True, ) - + # Build lookup: for each query, keep the highest-scoring hit best_hit = {} - for gene_info in results['out']: - query = gene_info.get('query', '') - if 'notfound' in gene_info and gene_info['notfound']: + for gene_info in results["out"]: + query = gene_info.get("query", "") + if "notfound" in gene_info and gene_info["notfound"]: if query not in best_hit: best_hit[query] = gene_info else: - score = gene_info.get('_score', 0) - if query not in best_hit or 'notfound' in best_hit[query] or score > best_hit[query].get('_score', 0): + score = gene_info.get("_score", 0) + if ( + query not in best_hit + or "notfound" in best_hit[query] + or score > best_hit[query].get("_score", 0) + ): best_hit[query] = gene_info # Process results in original gene order annotation_list = [] for query_gene in genes: gene_info = best_hit.get(query_gene) - if gene_info is None or ('notfound' in gene_info and gene_info['notfound']): - annotation_list.append({ - 'Input_Gene': query_gene, - 'Gene_Symbol': 'NOT FOUND', - 'Gene_Name': 'NOT FOUND', - 'Entrez_ID': '', - 'Type': '', - 'Summary': '' - }) + if gene_info is None or ("notfound" in gene_info and gene_info["notfound"]): + annotation_list.append( + { + "Input_Gene": query_gene, + "Gene_Symbol": "NOT FOUND", + "Gene_Name": "NOT FOUND", + "Entrez_ID": "", + "Type": "", + "Summary": "", + } + ) else: - annotation_list.append({ - 'Input_Gene': query_gene, - 'Gene_Symbol': gene_info.get('symbol', query_gene), - 'Gene_Name': gene_info.get('name', ''), - 'Entrez_ID': gene_info.get('entrezgene', ''), - 'Type': gene_info.get('type_of_gene', ''), - 'Summary': gene_info.get('summary', '') - }) - + annotation_list.append( + { + "Input_Gene": query_gene, + "Gene_Symbol": gene_info.get("symbol", query_gene), + "Gene_Name": gene_info.get("name", ""), + "Entrez_ID": gene_info.get("entrezgene", ""), + "Type": gene_info.get("type_of_gene", ""), + "Summary": gene_info.get("summary", ""), + } + ) + # Create DataFrame for this column all_annotations[row_idx] = pd.DataFrame(annotation_list) # Export to Excel with multiple sheets - with pd.ExcelWriter(output_file, engine='openpyxl') as writer: + with pd.ExcelWriter(output_file, engine="openpyxl") as writer: for row_idx, annotations_df in all_annotations.items(): sheet_name = str(row_idx) # Truncate sheet name if too long (Excel limit is 31 chars) sheet_name = sheet_name[:31] if len(sheet_name) > 31 else sheet_name annotations_df.to_excel(writer, sheet_name=sheet_name, index=False) - + return all_annotations -def rename_and_move_files_NMF(file_name_input, file_name_output, source_folder, destination_folder, len = 10): +def rename_and_move_files_NMF( + file_name_input, file_name_output, source_folder, destination_folder, len=10 +): """ Combine parallel inferred cNMF component results into one run. - + Renames and moves NMF result files from parallel runs to a unified location. - + Args: file_name_input: Input filename pattern file_name_output: Output filename pattern source_folder: Source directory containing parallel results destination_folder: Destination directory for combined results len: Number of parallel runs to process - + Returns: List of tuples with (source_filename, new_filename) for processed files """ # Create destination folder if it doesn't exist Path(destination_folder).mkdir(parents=True, exist_ok=True) - + processed_files = [] - - # Process files for i + + # Process files for i for i in range(len): source_filename = f"{file_name_input}_{i}.df.npz" source_path = os.path.join(source_folder, source_filename) - + # Check if source file exists if os.path.exists(source_path): - # Create new filename new_filename = f"{file_name_output}_{i}.df.npz" destination_path = os.path.join(destination_folder, new_filename) - + try: # Copy and rename file to destination shutil.copy2(source_path, destination_path) print(f"Successfully processed: {source_filename} -> {new_filename}") processed_files.append((source_filename, new_filename)) - + except Exception as e: print(f"Error processing {source_filename}: {e}") else: print(f"File not found: {source_filename}") - + return processed_files -def rename_all_NMF(file_name_input, file_name_output, source_folder, destination_folder, len = 10, components = [30, 50, 60, 80, 100, 200, 250, 300]): +def rename_all_NMF( + file_name_input, + file_name_output, + source_folder, + destination_folder, + len=10, + components=[30, 50, 60, 80, 100, 200, 250, 300], +): """ Combine all K values from parallel runs. - + Processes results from multiple parallel cNMF runs across different k values. - + Args: file_name_input: Input filename pattern file_name_output: Output filename pattern @@ -548,10 +615,15 @@ def rename_all_NMF(file_name_input, file_name_output, source_folder, destination """ for k in components: - file_name_input_new = f"{file_name_input}.spectra.k_{k}.iter" file_name_output_new = f"{file_name_output}.spectra.k_{k}.iter" source_folder_new = f"{source_folder}_{k}/Inference/cnmf_tmp" - rename_and_move_files_NMF(file_name_input_new, file_name_output_new, source_folder_new, destination_folder, len=len) + rename_and_move_files_NMF( + file_name_input_new, + file_name_output_new, + source_folder_new, + destination_folder, + len=len, + ) diff --git a/src/Stage1_Inference/torch-cNMF/Slurm_Version/torch_cnmf_inference_pipeline.py b/src/Stage1_Inference/torch-cNMF/Slurm_Version/torch_cnmf_inference_pipeline.py index 0acff46..4107a8d 100644 --- a/src/Stage1_Inference/torch-cNMF/Slurm_Version/torch_cnmf_inference_pipeline.py +++ b/src/Stage1_Inference/torch-cNMF/Slurm_Version/torch_cnmf_inference_pipeline.py @@ -1,4 +1,3 @@ - """ torch-cNMF Inference Pipeline (Updated for torch_cnmf) @@ -23,143 +22,243 @@ # Change path to wherever you have repo locally from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) + +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) from torch_cnmf import cNMF from Stage1_Inference.src import ( - run_cnmf_consensus, get_top_indices_fast, annotate_genes_to_excel, - rename_and_move_files_NMF, rename_all_NMF, compile_results, - filter_noncoding_genes, filter_noncoding_genes_gtf, + run_cnmf_consensus, + get_top_indices_fast, + annotate_genes_to_excel, + rename_and_move_files_NMF, + rename_all_NMF, + compile_results, + filter_noncoding_genes, + filter_noncoding_genes_gtf, ) from Stage1_Inference.src.plot_diagnostics import generate_all_plots - def main(): parser = argparse.ArgumentParser( description="torch-cNMF inference pipeline (torch_cnmf)" ) # --- I/O --- - parser.add_argument('--counts_fn', type=str, required=True, - help="Path to gene expression data (.h5ad)") - parser.add_argument('--output_directory', type=str, required=True, - help="Output directory for results") - parser.add_argument('--run_name', type=str, required=True, - help="Name for this cNMF run") - parser.add_argument('--nmf_seeds_path', type=str, default=None, - help="Path to text file with custom NMF seeds (one integer per line)") + parser.add_argument( + "--counts_fn", + type=str, + required=True, + help="Path to gene expression data (.h5ad)", + ) + parser.add_argument( + "--output_directory", + type=str, + required=True, + help="Output directory for results", + ) + parser.add_argument( + "--run_name", type=str, required=True, help="Name for this cNMF run" + ) + parser.add_argument( + "--nmf_seeds_path", + type=str, + default=None, + help="Path to text file with custom NMF seeds (one integer per line)", + ) # --- cNMF parameters --- - parser.add_argument('--K', nargs='*', type=int, default=[30, 50, 70, 80, 100, 200, 300], - help="K values to test") - parser.add_argument('--sel_threshs', nargs='*', type=float, default=[0.2, 2.0], - help="Density thresholds for consensus") - parser.add_argument('--numiter', type=int, default=10, - help="NMF iterations per K") - parser.add_argument('--densify', action='store_true', - help="Densify sparse matrix before factorization") - parser.add_argument('--tpm_fn', type=str, default=None, - help="Pre-computed TPM file path") - parser.add_argument('--numhvgenes', type=int, default=2000, - help="Number of highly variable genes") - parser.add_argument('--genes_file', type=str, default=None, - help="File with gene list (overrides HVG selection)") - parser.add_argument('--loss', default='frobenius', - help="NMF loss function") - parser.add_argument('--algo', type=str, default='halsvar', - help="NMF algorithm: mu | hals | halsvar | bpp") - parser.add_argument('--mode', type=str, default='batch', - help="Learning mode: batch | minibatch | dataloader") - parser.add_argument('--tol', type=float, default=1e-7, - help="Convergence tolerance") - parser.add_argument('--n_jobs', type=int, default=-1, - help="Number of parallel jobs") - parser.add_argument('--init', type=str, default='random', - help="Initialisation: random | nndsvd") - parser.add_argument('--seed', type=int, default=14, - help="Random seed") - parser.add_argument('--use_gpu', action='store_true', - help="Use GPU acceleration") + parser.add_argument( + "--K", + nargs="*", + type=int, + default=[30, 50, 70, 80, 100, 200, 300], + help="K values to test", + ) + parser.add_argument( + "--sel_threshs", + nargs="*", + type=float, + default=[0.2, 2.0], + help="Density thresholds for consensus", + ) + parser.add_argument("--numiter", type=int, default=10, help="NMF iterations per K") + parser.add_argument( + "--densify", + action="store_true", + help="Densify sparse matrix before factorization", + ) + parser.add_argument( + "--tpm_fn", type=str, default=None, help="Pre-computed TPM file path" + ) + parser.add_argument( + "--numhvgenes", type=int, default=2000, help="Number of highly variable genes" + ) + parser.add_argument( + "--genes_file", + type=str, + default=None, + help="File with gene list (overrides HVG selection)", + ) + parser.add_argument("--loss", default="frobenius", help="NMF loss function") + parser.add_argument( + "--algo", + type=str, + default="halsvar", + help="NMF algorithm: mu | hals | halsvar | bpp", + ) + parser.add_argument( + "--mode", + type=str, + default="batch", + help="Learning mode: batch | minibatch | dataloader", + ) + parser.add_argument("--tol", type=float, default=1e-7, help="Convergence tolerance") + parser.add_argument( + "--n_jobs", type=int, default=-1, help="Number of parallel jobs" + ) + parser.add_argument( + "--init", type=str, default="random", help="Initialisation: random | nndsvd" + ) + parser.add_argument("--seed", type=int, default=14, help="Random seed") + parser.add_argument("--use_gpu", action="store_true", help="Use GPU acceleration") # Regularisation - parser.add_argument('--alpha_usage', type=float, default=0.0) - parser.add_argument('--alpha_spectra', type=float, default=0.0) - parser.add_argument('--l1_ratio_usage', type=float, default=0.0) - parser.add_argument('--l1_ratio_spectra', type=float, default=0.0) + parser.add_argument("--alpha_usage", type=float, default=0.0) + parser.add_argument("--alpha_spectra", type=float, default=0.0) + parser.add_argument("--l1_ratio_usage", type=float, default=0.0) + parser.add_argument("--l1_ratio_spectra", type=float, default=0.0) # Batch-mode NMF - parser.add_argument('--batch_max_epoch', type=int, default=1000) - parser.add_argument('--batch_hals_tol', type=float, default=0.005) - parser.add_argument('--batch_hals_max_iter', type=int, default=1000) + parser.add_argument("--batch_max_epoch", type=int, default=1000) + parser.add_argument("--batch_hals_tol", type=float, default=0.005) + parser.add_argument("--batch_hals_max_iter", type=int, default=1000) # Mini-batch / dataloader mode NMF - parser.add_argument('--minibatch_max_epoch', type=int, default=1000) - parser.add_argument('--minibatch_size', type=int, default=100000) - parser.add_argument('--minibatch_max_iter', type=int, default=1000) - parser.add_argument('--minibatch_usage_tol', type=float, default=0.005) - parser.add_argument('--minibatch_spectra_tol', type=float, default=0.005) - parser.add_argument('--minibatch_shuffle', action='store_true', - help="Shuffle cells in minibatch mode (default: off)") + parser.add_argument("--minibatch_max_epoch", type=int, default=1000) + parser.add_argument("--minibatch_size", type=int, default=100000) + parser.add_argument("--minibatch_max_iter", type=int, default=1000) + parser.add_argument("--minibatch_usage_tol", type=float, default=0.005) + parser.add_argument("--minibatch_spectra_tol", type=float, default=0.005) + parser.add_argument( + "--minibatch_shuffle", + action="store_true", + help="Shuffle cells in minibatch mode (default: off)", + ) # Precision - parser.add_argument('--fp_precision', type=str, default='float', - help="Floating-point precision: float | double") + parser.add_argument( + "--fp_precision", + type=str, + default="float", + help="Floating-point precision: float | double", + ) # Refit - parser.add_argument('--sk_cd_refit', action='store_true', - help="Use sklearn coordinate descent for refitting") + parser.add_argument( + "--sk_cd_refit", + action="store_true", + help="Use sklearn coordinate descent for refitting", + ) # --- Annotation --- - parser.add_argument('--species', type=str, required=True, - help="Species for gene annotation: human | mouse") - parser.add_argument('--num_gene', type=int, default=300, - help="Number of top genes for annotation") - parser.add_argument('--parallel_running', action='store_true', - help="Enable parallel processing mode for multiple K values on when using parallel.sh") + parser.add_argument( + "--species", + type=str, + required=True, + help="Species for gene annotation: human | mouse", + ) + parser.add_argument( + "--num_gene", type=int, default=300, help="Number of top genes for annotation" + ) + parser.add_argument( + "--parallel_running", + action="store_true", + help="Enable parallel processing mode for multiple K values on when using parallel.sh", + ) # --- Preprocessing --- - parser.add_argument('--remove_noncoding', action='store_true', - help="Remove non-coding genes. With --gtf_path, keeps only genes annotated " - "as protein_coding in the GTF; otherwise falls back to the Ensembl-prefix heuristic.") - parser.add_argument('--ensembl_prefix', type=str, default='ENSG', - help="Ensembl ID prefix for the (fallback) prefix-based non-coding filter") - parser.add_argument('--gtf_path', type=str, default=None, - help="Path to a GENCODE/Ensembl GTF(.gz). When set with --remove_noncoding, " - "genes whose Ensembl ID is not a protein_coding 'gene' in the GTF are removed.") - parser.add_argument('--gene_id_key', type=str, default='gene_id', - help="Column in adata.var holding Ensembl gene IDs (used for the GTF-based non-coding filter)") - parser.add_argument('--add_gene_names_from_gtf', action='store_true', - help="Populate adata.var[gene_names_key] with gene symbols looked up from --gtf_path " - "by Ensembl ID (from gene_id_key). Requires --gtf_path.") - parser.add_argument('--gene_names_key', type=str, default='symbol', - help="Column in adata.var with gene names (used for non-coding filter, GTF gene-name " - "annotation, and compiled results)") + parser.add_argument( + "--remove_noncoding", + action="store_true", + help="Remove non-coding genes. With --gtf_path, keeps only genes annotated " + "as protein_coding in the GTF; otherwise falls back to the Ensembl-prefix heuristic.", + ) + parser.add_argument( + "--ensembl_prefix", + type=str, + default="ENSG", + help="Ensembl ID prefix for the (fallback) prefix-based non-coding filter", + ) + parser.add_argument( + "--gtf_path", + type=str, + default=None, + help="Path to a GENCODE/Ensembl GTF(.gz). When set with --remove_noncoding, " + "genes whose Ensembl ID is not a protein_coding 'gene' in the GTF are removed.", + ) + parser.add_argument( + "--gene_id_key", + type=str, + default="gene_id", + help="Column in adata.var holding Ensembl gene IDs (used for the GTF-based non-coding filter)", + ) + parser.add_argument( + "--add_gene_names_from_gtf", + action="store_true", + help="Populate adata.var[gene_names_key] with gene symbols looked up from --gtf_path " + "by Ensembl ID (from gene_id_key). Requires --gtf_path.", + ) + parser.add_argument( + "--gene_names_key", + type=str, + default="symbol", + help="Column in adata.var with gene names (used for non-coding filter, GTF gene-name " + "annotation, and compiled results)", + ) # --- Metadata keys --- - parser.add_argument('--data_key', type=str, default='rna') - parser.add_argument('--prog_key', type=str, default='cNMF') - parser.add_argument('--categorical_key', type=str, default='sample', - help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample") - parser.add_argument('--guide_names_key', type=str, default='guide_names') - parser.add_argument('--guide_targets_key', type=str, default='guide_targets') - parser.add_argument('--guide_assignment_key', type=str, default='guide_assignment') + parser.add_argument("--data_key", type=str, default="rna") + parser.add_argument("--prog_key", type=str, default="cNMF") + parser.add_argument( + "--categorical_key", + type=str, + default="sample", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + ) + parser.add_argument("--guide_names_key", type=str, default="guide_names") + parser.add_argument("--guide_targets_key", type=str, default="guide_targets") + parser.add_argument("--guide_assignment_key", type=str, default="guide_assignment") # --- Step selection --- - parser.add_argument('--run_factorize', action='store_true', - help="Run the factorization step") - parser.add_argument('--run_refit', action='store_true', - help="Run combine + k_selection + consensus") - parser.add_argument('--run_compile_annotation', action='store_true', - help="Run result compilation and gene annotation") - parser.add_argument('--run_gene_annotation', action='store_true', - help="Run MyGene annotation (needs network)") - parser.add_argument('--run_diagnostic_plots', action='store_true', - help="Generate diagnostic plots (elbow curves, usage heatmaps, loading violins)") - parser.add_argument('--skip_existing', action='store_true', - help="If set, skip NMF replicates already completed on disk (pause/resume mode). " - "Default: re-run all replicates from scratch.") + parser.add_argument( + "--run_factorize", action="store_true", help="Run the factorization step" + ) + parser.add_argument( + "--run_refit", action="store_true", help="Run combine + k_selection + consensus" + ) + parser.add_argument( + "--run_compile_annotation", + action="store_true", + help="Run result compilation and gene annotation", + ) + parser.add_argument( + "--run_gene_annotation", + action="store_true", + help="Run MyGene annotation (needs network)", + ) + parser.add_argument( + "--run_diagnostic_plots", + action="store_true", + help="Generate diagnostic plots (elbow curves, usage heatmaps, loading violins)", + ) + parser.add_argument( + "--skip_existing", + action="store_true", + help="If set, skip NMF replicates already completed on disk (pause/resume mode). " + "Default: re-run all replicates from scratch.", + ) args = parser.parse_args() @@ -169,35 +268,35 @@ def main(): nmf_seeds = np.load(args.nmf_seeds_path) # --- Create output directory --- - run_dir = f'{args.output_directory}/{args.run_name}' - inference_dir = f'{run_dir}/Inference' + run_dir = f"{args.output_directory}/{args.run_name}" + inference_dir = f"{run_dir}/Inference" os.makedirs(inference_dir, exist_ok=True) # --- Save config (incl. SLURM info) --- slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), - 'num_nodes': os.environ.get('SLURM_JOB_NUM_NODES'), - 'ntasks': os.environ.get('SLURM_NTASKS'), - 'cpus_per_task': os.environ.get('SLURM_CPUS_PER_TASK'), - 'mem_per_node': os.environ.get('SLURM_MEM_PER_NODE'), - 'mem_per_cpu': os.environ.get('SLURM_MEM_PER_CPU'), - 'gres': os.environ.get('SLURM_JOB_GRES'), - 'gpu_device_ids': os.environ.get('SLURM_GPUS_ON_NODE'), - 'gpu_type': os.environ.get('SLURM_JOB_GPUS'), - 'time_limit': os.environ.get('SLURM_JOB_TIMELIMIT'), - 'time_remaining': os.environ.get('SLURM_TIMELIMIT'), - 'submit_dir': os.environ.get('SLURM_SUBMIT_DIR'), - 'submit_host': os.environ.get('SLURM_SUBMIT_HOST'), - 'constraint': os.environ.get('SLURM_JOB_CONSTRAINT'), - 'array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), + "num_nodes": os.environ.get("SLURM_JOB_NUM_NODES"), + "ntasks": os.environ.get("SLURM_NTASKS"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + "mem_per_node": os.environ.get("SLURM_MEM_PER_NODE"), + "mem_per_cpu": os.environ.get("SLURM_MEM_PER_CPU"), + "gres": os.environ.get("SLURM_JOB_GRES"), + "gpu_device_ids": os.environ.get("SLURM_GPUS_ON_NODE"), + "gpu_type": os.environ.get("SLURM_JOB_GPUS"), + "time_limit": os.environ.get("SLURM_JOB_TIMELIMIT"), + "time_remaining": os.environ.get("SLURM_TIMELIMIT"), + "submit_dir": os.environ.get("SLURM_SUBMIT_DIR"), + "submit_host": os.environ.get("SLURM_SUBMIT_HOST"), + "constraint": os.environ.get("SLURM_JOB_CONSTRAINT"), + "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" - config_to_save = {'script_args': vars(args), 'slurm_info': slurm_info} - with open(f'{inference_dir}/config_{job_id}.yml', 'w') as f: + config_to_save = {"script_args": vars(args), "slurm_info": slurm_info} + with open(f"{inference_dir}/config_{job_id}.yml", "w") as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) # --- Optionally filter non-coding genes --- @@ -234,7 +333,7 @@ def main(): # ====================================================================== # Initialise cNMF (torch_cnmf) # ====================================================================== - cnmf_obj = cNMF(output_dir=run_dir, name='Inference') + cnmf_obj = cNMF(output_dir=run_dir, name="Inference") # --- Prepare --- cnmf_obj.prepare( @@ -279,9 +378,9 @@ def main(): if args.run_refit: cnmf_obj.combine() cnmf_obj.k_selection_plot() - run_cnmf_consensus(cnmf_obj, - components=args.K, - density_thresholds=args.sel_threshs) + run_cnmf_consensus( + cnmf_obj, components=args.K, density_thresholds=args.sel_threshs + ) # One PNG per k: density_filtering_plot() refuses to autodetect k when more # than one is cached, so loop explicitly. Purely diagnostic -- never let a # plotting failure abort compilation/annotation downstream. @@ -293,34 +392,45 @@ def main(): # --- Compile results --- if args.run_compile_annotation: - compile_results(run_dir, 'Inference', - components=args.K, sel_threshs=args.sel_threshs, - guide_names_key=args.guide_names_key, - guide_targets_key=args.guide_targets_key, - categorical_key=args.categorical_key, - guide_assignment_key=args.guide_assignment_key, - gene_names_key=args.gene_names_key) - #---annotate ---- + compile_results( + run_dir, + "Inference", + components=args.K, + sel_threshs=args.sel_threshs, + guide_names_key=args.guide_names_key, + guide_targets_key=args.guide_targets_key, + categorical_key=args.categorical_key, + guide_assignment_key=args.guide_assignment_key, + gene_names_key=args.gene_names_key, + ) + # ---annotate ---- if args.run_gene_annotation: print("Annotating top genes per program via MyGene...") - os.makedirs(f'{inference_dir}/Annotation', exist_ok=True) + os.makedirs(f"{inference_dir}/Annotation", exist_ok=True) for i in args.sel_threshs: for k in args.K: - df = pd.read_csv('{inference_dir}/Inference.gene_spectra_score.k_{k}.dt_{sel_thresh}.txt'.format( - inference_dir=inference_dir, - k=k, - sel_thresh=str(i).replace('.', '_')), - sep='\t', index_col=0) + df = pd.read_csv( + "{inference_dir}/Inference.gene_spectra_score.k_{k}.dt_{sel_thresh}.txt".format( + inference_dir=inference_dir, + k=k, + sel_thresh=str(i).replace(".", "_"), + ), + sep="\t", + index_col=0, + ) overlap = get_top_indices_fast(df, gene_num=args.num_gene) - annotate_genes_to_excel(overlap, species=args.species, - output_file=f"{inference_dir}/Annotation/{k}_{str(i).replace('.','_')}.xlsx") + annotate_genes_to_excel( + overlap, + species=args.species, + output_file=f"{inference_dir}/Annotation/{k}_{str(i).replace('.', '_')}.xlsx", + ) # --- Diagnostic plots --- if args.run_diagnostic_plots: print("Generating diagnostic plots...") generate_all_plots( run_dir=run_dir, - run_name='Inference', + run_name="Inference", K_list=args.K, sel_thresh_list=args.sel_threshs, categorical_key=args.categorical_key, @@ -333,7 +443,7 @@ def main(): rename_all_NMF( source_folder=f"{args.output_directory}/{args.run_name}/{args.run_name}", destination_folder=f"{args.output_directory}/{args.run_name}/{args.run_name}_all/Inference/cnmf_tmp", - file_name_input='Inference', + file_name_input="Inference", file_name_output="Inference_all", len=args.numiter, components=args.K, @@ -343,5 +453,5 @@ def main(): return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py b/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py index bf99350..d41296c 100644 --- a/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py +++ b/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py @@ -36,8 +36,8 @@ def _read_samples_from_h5mu(h5mu_path, data_key, categorical_key): Returns a list of sample names, or None if the column cannot be read (e.g. missing path, decoding error). """ try: - with h5py.File(h5mu_path, 'r') as f: - obs_path = f'mod/{data_key}/obs/{categorical_key}' + with h5py.File(h5mu_path, "r") as f: + obs_path = f"mod/{data_key}/obs/{categorical_key}" if obs_path not in f: return None col = read_elem(f[obs_path]) @@ -49,9 +49,10 @@ def _read_samples_from_h5mu(h5mu_path, data_key, categorical_key): # Make the repo's src/ importable regardless of the current working directory from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) + +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) # Reference data populated by setup_resources.sh -RESOURCES_DIR = Path(__file__).resolve().parents[2] / 'Resources' +RESOURCES_DIR = Path(__file__).resolve().parents[2] / "Resources" from Stage2_Evaluation.A_Metrics.src import ( compute_categorical_association, @@ -59,55 +60,178 @@ def _read_samples_from_h5mu(h5mu_path, data_key, categorical_key): compute_trait_enrichment, compute_perturbation_association, compute_explained_variance, - compute_motif_enrichment + compute_motif_enrichment, ) -def _reassign_name(mdata, gene_names_key='symbol', data_key = 'rna'): +def _reassign_name(mdata, gene_names_key="symbol", data_key="rna"): mdata[data_key].var_names = mdata[data_key].var[gene_names_key].astype(str) def main(): - parser = argparse.ArgumentParser( - description="cNMF evaluation pipeline" + parser = argparse.ArgumentParser(description="cNMF evaluation pipeline") + + # IO info + parser.add_argument( + "--out_dir", + help="Directory containing cNMF output files to evaluate", + type=str, + required=True, + ) + parser.add_argument( + "--run_name", + help="Name of the cNMF run to evaluate (must match name used during inference)", + type=str, + required=True, + ) + parser.add_argument( + "--K", + nargs="*", + type=int, + default=[30, 50, 70, 80, 100, 200, 300], + help="List of K values (number of components) to evaluate", + ) + parser.add_argument( + "--sel_threshs", + nargs="*", + type=float, + default=[0.2, 2.0], + help="List of density thresholds to evaluate", ) - - #IO info - parser.add_argument('--out_dir', help='Directory containing cNMF output files to evaluate', type=str, required=True) - parser.add_argument('--run_name', help='Name of the cNMF run to evaluate (must match name used during inference)', type=str, required=True) - parser.add_argument('--K', nargs='*', type=int, default=[30, 50, 70, 80, 100, 200, 300], help='List of K values (number of components) to evaluate') - parser.add_argument('--sel_threshs', nargs='*', type=float, default=[0.2, 2.0], help='List of density thresholds to evaluate') # running different tests - parser.add_argument('--Perform_categorical', action="store_true", help='If set, compute categorical association between programs and sample labels (Kruskal-Wallis test)') - parser.add_argument('--Perform_perturbation', action="store_true", help='If set, compute perturbation association between programs and guide perturbations') - parser.add_argument('--Perform_geneset', action="store_true", help='If set, perform gene set enrichment analysis (Reactome and GO terms) on program genes') - parser.add_argument('--Perform_trait', action="store_true", help='If set, perform GWAS trait enrichment analysis on program genes') - parser.add_argument('--Perform_explained_variance', action="store_true", help='If set, compute explained variance for each K value') - parser.add_argument('--Perform_motif', action="store_true", help='If set, perform transcription factor motif enrichment analysis') + parser.add_argument( + "--Perform_categorical", + action="store_true", + help="If set, compute categorical association between programs and sample labels (Kruskal-Wallis test)", + ) + parser.add_argument( + "--Perform_perturbation", + action="store_true", + help="If set, compute perturbation association between programs and guide perturbations", + ) + parser.add_argument( + "--Perform_geneset", + action="store_true", + help="If set, perform gene set enrichment analysis (Reactome and GO terms) on program genes", + ) + parser.add_argument( + "--Perform_trait", + action="store_true", + help="If set, perform GWAS trait enrichment analysis on program genes", + ) + parser.add_argument( + "--Perform_explained_variance", + action="store_true", + help="If set, compute explained variance for each K value", + ) + parser.add_argument( + "--Perform_motif", + action="store_true", + help="If set, perform transcription factor motif enrichment analysis", + ) # resources - parser.add_argument('--X_normalized_path', type=str, help='Path to normalized cell x gene matrix (.h5ad) from cNMF pipeline (required for explained variance)', default=None) - parser.add_argument('--guide_annotation_path', type=str, help='Path to tab-separated file with guide annotations including "targeting" column to identify non-targeting controls (optional)') - parser.add_argument('--gwas_data_path', type=str, help='Path to GWAS data file for trait enrichment analysis (required for trait enrichment)', default=None) + parser.add_argument( + "--X_normalized_path", + type=str, + help="Path to normalized cell x gene matrix (.h5ad) from cNMF pipeline (required for explained variance)", + default=None, + ) + parser.add_argument( + "--guide_annotation_path", + type=str, + help='Path to tab-separated file with guide annotations including "targeting" column to identify non-targeting controls (optional)', + ) + parser.add_argument( + "--gwas_data_path", + type=str, + help="Path to GWAS data file for trait enrichment analysis (required for trait enrichment)", + default=None, + ) # keys - parser.add_argument('--data_key', help='Key to access gene expression data in MuData object (default: rna)', type=str, default="rna") - parser.add_argument('--prog_key', help='Key to access cNMF programs in MuData object (default: cNMF)', type=str, default="cNMF") - parser.add_argument('--categorical_key', help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", type=str, default="sample") - parser.add_argument('--gene_names_key', type=str, help='Column in data_guide["rna"].var containing gene names (default: symbol)', default='symbol') - parser.add_argument('--guide_names_key', help='Key in .uns to access guide names (default: guide_names)', type=str, default="guide_names") - parser.add_argument('--guide_targets_key', help='Key in .uns to access guide target genes (default: guide_targets)', type=str, default="guide_targets") - parser.add_argument('--guide_assignment_key', help='Key in .obsm to access guide assignment matrix (default: guide_assignment)', type=str, default="guide_assignment") - parser.add_argument('--guide_annotation_key', nargs='*', type=str, help='name of non-targeting guide targets', default=["non-targeting"]) - - - parser.add_argument('--organism', help='Organism/species for enrichment analysis (default: human)', type=str, default="human") - parser.add_argument('--FDR_method', help='Method for FDR correction in perturbation association (default: StoreyQ)', type=str, default="StoreyQ") - parser.add_argument('--n_top', type = int, help='Number of top loaded genes use to perform enrichment test(default: 300)', default=300) - parser.add_argument('--use_cache', action="store_true", help='If set, load gene set libraries from cached JSON in Resources/ instead of downloading. Falls back to download if cache not found, and saves a cache copy after download.') - parser.add_argument('--skip_existing', action="store_true", help='If set, skip metric computations whose output files already exist. Useful for resuming jobs that were preempted mid-batch.') - parser.add_argument('--reassign_name', action="store_true", help='Reform the data if needed') + parser.add_argument( + "--data_key", + help="Key to access gene expression data in MuData object (default: rna)", + type=str, + default="rna", + ) + parser.add_argument( + "--prog_key", + help="Key to access cNMF programs in MuData object (default: cNMF)", + type=str, + default="cNMF", + ) + parser.add_argument( + "--categorical_key", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + type=str, + default="sample", + ) + parser.add_argument( + "--gene_names_key", + type=str, + help='Column in data_guide["rna"].var containing gene names (default: symbol)', + default="symbol", + ) + parser.add_argument( + "--guide_names_key", + help="Key in .uns to access guide names (default: guide_names)", + type=str, + default="guide_names", + ) + parser.add_argument( + "--guide_targets_key", + help="Key in .uns to access guide target genes (default: guide_targets)", + type=str, + default="guide_targets", + ) + parser.add_argument( + "--guide_assignment_key", + help="Key in .obsm to access guide assignment matrix (default: guide_assignment)", + type=str, + default="guide_assignment", + ) + parser.add_argument( + "--guide_annotation_key", + nargs="*", + type=str, + help="name of non-targeting guide targets", + default=["non-targeting"], + ) + + parser.add_argument( + "--organism", + help="Organism/species for enrichment analysis (default: human)", + type=str, + default="human", + ) + parser.add_argument( + "--FDR_method", + help="Method for FDR correction in perturbation association (default: StoreyQ)", + type=str, + default="StoreyQ", + ) + parser.add_argument( + "--n_top", + type=int, + help="Number of top loaded genes use to perform enrichment test(default: 300)", + default=300, + ) + parser.add_argument( + "--use_cache", + action="store_true", + help="If set, load gene set libraries from cached JSON in Resources/ instead of downloading. Falls back to download if cache not found, and saves a cache copy after download.", + ) + parser.add_argument( + "--skip_existing", + action="store_true", + help="If set, skip metric computations whose output files already exist. Useful for resuming jobs that were preempted mid-batch.", + ) + parser.add_argument( + "--reassign_name", action="store_true", help="Reform the data if needed" + ) args = parser.parse_args() @@ -115,33 +239,36 @@ def main(): parser.error("--gwas_data_path is required when --Perform_trait is set") # create output directory - os.makedirs((f'{args.out_dir}/{args.run_name}/Evaluation'), exist_ok=True) + os.makedirs((f"{args.out_dir}/{args.run_name}/Evaluation"), exist_ok=True) # --- Save config (incl. SLURM info) --- slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), - 'num_nodes': os.environ.get('SLURM_JOB_NUM_NODES'), - 'ntasks': os.environ.get('SLURM_NTASKS'), - 'cpus_per_task': os.environ.get('SLURM_CPUS_PER_TASK'), - 'mem_per_node': os.environ.get('SLURM_MEM_PER_NODE'), - 'mem_per_cpu': os.environ.get('SLURM_MEM_PER_CPU'), - 'time_limit': os.environ.get('SLURM_JOB_TIMELIMIT'), - 'submit_dir': os.environ.get('SLURM_SUBMIT_DIR'), - 'array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), + "num_nodes": os.environ.get("SLURM_JOB_NUM_NODES"), + "ntasks": os.environ.get("SLURM_NTASKS"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + "mem_per_node": os.environ.get("SLURM_MEM_PER_NODE"), + "mem_per_cpu": os.environ.get("SLURM_MEM_PER_CPU"), + "time_limit": os.environ.get("SLURM_JOB_TIMELIMIT"), + "submit_dir": os.environ.get("SLURM_SUBMIT_DIR"), + "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" - config_to_save = {'script_args': vars(args), 'slurm_info': slurm_info} - with open(f'{args.out_dir}/{args.run_name}/Evaluation/config_{job_id}.yml', 'w') as f: + config_to_save = {"script_args": vars(args), "slurm_info": slurm_info} + with open( + f"{args.out_dir}/{args.run_name}/Evaluation/config_{job_id}.yml", "w" + ) as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) - # defind objects used in explained variance if args.Perform_explained_variance: - cnmf_obj = cnmf.cNMF(output_dir=f'{args.out_dir}/{args.run_name}', name='Inference') + cnmf_obj = cnmf.cNMF( + output_dir=f"{args.out_dir}/{args.run_name}", name="Inference" + ) X_norm = sc.read_h5ad(args.X_normalized_path) X = X_norm.X @@ -151,90 +278,121 @@ def main(): # target GROUP names (e.g. "non-targeting") from mdata.uns["guide_targets"]. # Prefer using --guide_annotation_key (default: ["non-targeting"]) for key-based access. if args.guide_annotation_path is not None: - df_target = pd.read_csv(args.guide_annotation_path, sep = "\t", index_col = 0) - df_target_non = df_target[df_target["targeting"]==False] + df_target = pd.read_csv(args.guide_annotation_path, sep="\t", index_col=0) + df_target_non = df_target[df_target["targeting"] == False] reference_targets = df_target_non.index.values.tolist() else: reference_targets = args.guide_annotation_key - for sel_thresh in args.sel_threshs: for k in args.K: - - output_folder = f"{args.out_dir}/{args.run_name}/Evaluation/{k}_{str(sel_thresh).replace('.','_')}" + output_folder = f"{args.out_dir}/{args.run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}" os.makedirs(output_folder, exist_ok=True) # Expected output files for skip checks (perturbation is sample-dependent, checked later) - out_cat = f'{output_folder}/{k}_categorical_association_results.txt' - out_posthoc = f'{output_folder}/{k}_categorical_association_posthoc.txt' - out_geneset = f'{output_folder}/{k}_geneset_enrichment.txt' - out_go = f'{output_folder}/{k}_GO_term_enrichment.txt' - out_trait = f'{output_folder}/{k}_trait_enrichment.txt' - out_expvar = f'{output_folder}/{k}_Explained_Variance.txt' - out_expvar_sum = f'{output_folder}/{k}_Explained_Variance_Summary.txt' + out_cat = f"{output_folder}/{k}_categorical_association_results.txt" + out_posthoc = f"{output_folder}/{k}_categorical_association_posthoc.txt" + out_geneset = f"{output_folder}/{k}_geneset_enrichment.txt" + out_go = f"{output_folder}/{k}_GO_term_enrichment.txt" + out_trait = f"{output_folder}/{k}_trait_enrichment.txt" + out_expvar = f"{output_folder}/{k}_Explained_Variance.txt" + out_expvar_sum = f"{output_folder}/{k}_Explained_Variance_Summary.txt" # Early-skip: if all enabled metrics already have outputs on disk, avoid the mdata load entirely. # For perturbation, sample names come from obs[categorical_key]; read just that column from the # h5mu (cheap) instead of loading the full MuData. if args.skip_existing: early_skip_ok = True - if args.Perform_categorical and not (os.path.exists(out_cat) and os.path.exists(out_posthoc)): + if args.Perform_categorical and not ( + os.path.exists(out_cat) and os.path.exists(out_posthoc) + ): early_skip_ok = False - if args.Perform_geneset and not (os.path.exists(out_geneset) and os.path.exists(out_go)): + if args.Perform_geneset and not ( + os.path.exists(out_geneset) and os.path.exists(out_go) + ): early_skip_ok = False if args.Perform_trait and not os.path.exists(out_trait): early_skip_ok = False - if args.Perform_explained_variance and not (os.path.exists(out_expvar) and os.path.exists(out_expvar_sum)): + if args.Perform_explained_variance and not ( + os.path.exists(out_expvar) and os.path.exists(out_expvar_sum) + ): early_skip_ok = False if args.Perform_perturbation: - h5mu_path = '{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{sel_thresh}.h5mu'.format( - out_dir=args.out_dir, run_name=args.run_name, k=k, - sel_thresh=str(sel_thresh).replace('.', '_')) - samples = _read_samples_from_h5mu(h5mu_path, args.data_key, args.categorical_key) + h5mu_path = "{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{sel_thresh}.h5mu".format( + out_dir=args.out_dir, + run_name=args.run_name, + k=k, + sel_thresh=str(sel_thresh).replace(".", "_"), + ) + samples = _read_samples_from_h5mu( + h5mu_path, args.data_key, args.categorical_key + ) if samples is None: early_skip_ok = False else: - pert_outs = [f'{output_folder}/{k}_perturbation_association_results_{samp}.txt' for samp in samples] + pert_outs = [ + f"{output_folder}/{k}_perturbation_association_results_{samp}.txt" + for samp in samples + ] if not all(os.path.exists(p) for p in pert_outs): early_skip_ok = False if early_skip_ok: - print(f"[K={k}, thresh={sel_thresh}] All requested outputs already exist; skipping iteration.") + print( + f"[K={k}, thresh={sel_thresh}] All requested outputs already exist; skipping iteration." + ) continue # Load mdata - mdata = mu.read('{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{sel_thresh}.h5mu'.format(out_dir = args.out_dir, - run_name =args.run_name, - k=k, - sel_thresh = str(sel_thresh).replace('.','_'))) - - - - - # assign information + mdata = mu.read( + "{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{sel_thresh}.h5mu".format( + out_dir=args.out_dir, + run_name=args.run_name, + k=k, + sel_thresh=str(sel_thresh).replace(".", "_"), + ) + ) + + # assign information if args.reassign_name: - _reassign_name(mdata, gene_names_key=args.gene_names_key, data_key = args.data_key) - + _reassign_name( + mdata, gene_names_key=args.gene_names_key, data_key=args.data_key + ) # Run categorical assocation if args.Perform_categorical: - if args.skip_existing and os.path.exists(out_cat) and os.path.exists(out_posthoc): - print(f"[K={k}, thresh={sel_thresh}] Skipping categorical (existing outputs found).") + if ( + args.skip_existing + and os.path.exists(out_cat) + and os.path.exists(out_posthoc) + ): + print( + f"[K={k}, thresh={sel_thresh}] Skipping categorical (existing outputs found)." + ) else: - results_df, posthoc_df = compute_categorical_association(mdata, prog_key=args.prog_key, categorical_key=args.categorical_key, - pseudobulk_key=None, test='dunn', n_jobs=-1, inplace=False) - - results_df.to_csv(out_cat, sep='\t', index=False) # This was made wide form to insert into .var of the program anndata. - posthoc_df.to_csv(out_posthoc, sep='\t', index=False) + results_df, posthoc_df = compute_categorical_association( + mdata, + prog_key=args.prog_key, + categorical_key=args.categorical_key, + pseudobulk_key=None, + test="dunn", + n_jobs=-1, + inplace=False, + ) + results_df.to_csv( + out_cat, sep="\t", index=False + ) # This was made wide form to insert into .var of the program anndata. + posthoc_df.to_csv(out_posthoc, sep="\t", index=False) # Run perturbation assocation if args.Perform_perturbation: - # Validate reference_targets against mdata guide_targets mdata_targets = set(mdata[args.prog_key].uns[args.guide_targets_key]) matched_ref = mdata_targets.intersection(reference_targets) - print(f"[K={k}] reference_targets overlap with mdata guide_targets: {matched_ref}") + print( + f"[K={k}] reference_targets overlap with mdata guide_targets: {matched_ref}" + ) if len(matched_ref) == 0: raise ValueError( f"No reference_targets found in mdata guide_targets. " @@ -243,61 +401,107 @@ def main(): f"Use --guide_annotation_key instead of --guide_annotation_path." ) - # perform association per condition + # perform association per condition for samp in mdata[args.data_key].obs[args.categorical_key].unique(): - out_pert = '{}/{}_perturbation_association_results_{}.txt'.format(output_folder, k, samp) + out_pert = "{}/{}_perturbation_association_results_{}.txt".format( + output_folder, k, samp + ) if args.skip_existing and os.path.exists(out_pert): - print(f"[K={k}, thresh={sel_thresh}, samp={samp}] Skipping perturbation (existing output found).") + print( + f"[K={k}, thresh={sel_thresh}, samp={samp}] Skipping perturbation (existing output found)." + ) continue - mdata_ = mdata[mdata['rna'].obs[args.categorical_key]==samp] - test_stats_df = compute_perturbation_association(mdata_, prog_key=args.prog_key, - collapse_targets=True, - pseudobulk=False, - reference_targets=reference_targets, - n_jobs=-1, inplace=False, FDR_method = args.FDR_method) - test_stats_df.to_csv(out_pert, sep='\t', index=False) - + mdata_ = mdata[mdata["rna"].obs[args.categorical_key] == samp] + test_stats_df = compute_perturbation_association( + mdata_, + prog_key=args.prog_key, + collapse_targets=True, + pseudobulk=False, + reference_targets=reference_targets, + n_jobs=-1, + inplace=False, + FDR_method=args.FDR_method, + ) + test_stats_df.to_csv(out_pert, sep="\t", index=False) # Gene-set enrichment if args.Perform_geneset: - if args.skip_existing and os.path.exists(out_geneset) and os.path.exists(out_go): - print(f"[K={k}, thresh={sel_thresh}] Skipping geneset/GO enrichment (existing outputs found).") + if ( + args.skip_existing + and os.path.exists(out_geneset) + and os.path.exists(out_go) + ): + print( + f"[K={k}, thresh={sel_thresh}] Skipping geneset/GO enrichment (existing outputs found)." + ) else: - pre_res = compute_geneset_enrichment(mdata, prog_key=args.prog_key, data_key=args.data_key, prog_name=None, - organism=args.organism, library='Reactome_2022', method="fisher", - database='enrichr', n_top=args.n_top, n_jobs=-1, - inplace=False, user_geneset=None, use_loadings_gene=False, - gene_names_key=args.gene_names_key, use_cache=args.use_cache) # use_loadings_gene:use all background genes - pre_res.to_csv(out_geneset, sep='\t', index=False) - + pre_res = compute_geneset_enrichment( + mdata, + prog_key=args.prog_key, + data_key=args.data_key, + prog_name=None, + organism=args.organism, + library="Reactome_2022", + method="fisher", + database="enrichr", + n_top=args.n_top, + n_jobs=-1, + inplace=False, + user_geneset=None, + use_loadings_gene=False, + gene_names_key=args.gene_names_key, + use_cache=args.use_cache, + ) # use_loadings_gene:use all background genes + pre_res.to_csv(out_geneset, sep="\t", index=False) # GO Term enrichment - pre_res = compute_geneset_enrichment(mdata, prog_key=args.prog_key, data_key=args.data_key, prog_name=None, - organism=args.organism, library='GO_Biological_Process_2023', method="fisher", - database='enrichr', n_top=args.n_top, n_jobs=-1, - inplace=False, user_geneset=None, use_loadings_gene=False, - gene_names_key=args.gene_names_key, use_cache=args.use_cache) # use_loadings_gene: use all background genes - pre_res.to_csv(out_go, sep='\t', index=False) - + pre_res = compute_geneset_enrichment( + mdata, + prog_key=args.prog_key, + data_key=args.data_key, + prog_name=None, + organism=args.organism, + library="GO_Biological_Process_2023", + method="fisher", + database="enrichr", + n_top=args.n_top, + n_jobs=-1, + inplace=False, + user_geneset=None, + use_loadings_gene=False, + gene_names_key=args.gene_names_key, + use_cache=args.use_cache, + ) # use_loadings_gene: use all background genes + pre_res.to_csv(out_go, sep="\t", index=False) # Run trait enrichment if args.Perform_trait: if args.skip_existing and os.path.exists(out_trait): - print(f"[K={k}, thresh={sel_thresh}] Skipping trait enrichment (existing output found).") + print( + f"[K={k}, thresh={sel_thresh}] Skipping trait enrichment (existing output found)." + ) else: - pre_res_trait = compute_trait_enrichment(mdata, gwas_data=args.gwas_data_path, - prog_key=args.prog_key, prog_name=None, data_key=args.data_key, - library='OT_GWAS', n_jobs=-1, inplace=False, - key_column='trait_efos', gene_column='gene_name', - method='fisher', n_top=args.n_top, use_loadings_gene=True, - gene_names_key=args.gene_names_key) # use_loadings_gene:use all background genes inter. expressed gene - pre_res_trait.to_csv(out_trait, sep='\t', index=False) - + pre_res_trait = compute_trait_enrichment( + mdata, + gwas_data=args.gwas_data_path, + prog_key=args.prog_key, + prog_name=None, + data_key=args.data_key, + library="OT_GWAS", + n_jobs=-1, + inplace=False, + key_column="trait_efos", + gene_column="gene_name", + method="fisher", + n_top=args.n_top, + use_loadings_gene=True, + gene_names_key=args.gene_names_key, + ) # use_loadings_gene:use all background genes inter. expressed gene + pre_res_trait.to_csv(out_trait, sep="\t", index=False) # Run motif analysis if args.Perform_motif: - - ''' working in progress + """ working in progress fimo_thresh_enhancer = 1e-6 fimo_thresh_promoter = 1e-4 @@ -325,19 +529,31 @@ def main(): motif_match_df.to_csv(os.path.join(args.out_dir, f'cNMF_{class_}_pearson_topn2000_{samp}_motif_match.txt'), sep='\t', index=False) motif_count_df.to_csv(os.path.join(args.out_dir, f'cNMF_{class_}_pearson_topn2000_{samp}_motif_count.txt'), sep='\t', index=False) motif_enrichment_df.to_csv(os.path.join(args.out_dir, f'cNMF_{class_}_pearson_topn2000_{samp}_motif_enrichment.txt'), sep='\t', index=False) - ''' + """ # Run explained variance if args.Perform_explained_variance: - if args.skip_existing and os.path.exists(out_expvar) and os.path.exists(out_expvar_sum): - print(f"[K={k}, thresh={sel_thresh}] Skipping explained variance (existing outputs found).") + if ( + args.skip_existing + and os.path.exists(out_expvar) + and os.path.exists(out_expvar_sum) + ): + print( + f"[K={k}, thresh={sel_thresh}] Skipping explained variance (existing outputs found)." + ) else: - compute_explained_variance(cnmf_obj, X, k, output_folder = output_folder, thre = str(sel_thresh), program_name=mdata[args.prog_key].var_names ) + compute_explained_variance( + cnmf_obj, + X, + k, + output_folder=output_folder, + thre=str(sel_thresh), + program_name=mdata[args.prog_key].var_names, + ) print("Pipeline finished.") return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) - diff --git a/src/Stage2_Evaluation/A_Metrics/src/Utilities.py b/src/Stage2_Evaluation/A_Metrics/src/Utilities.py index 5cb7d1d..f319382 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/Utilities.py +++ b/src/Stage2_Evaluation/A_Metrics/src/Utilities.py @@ -9,13 +9,12 @@ import mygene - def rename_gene_mygene(mdata): """ Convert Ensembl gene IDs to gene symbols using the MyGene web service. - + Rename genes in cNMF modality using MyGene service to convert Ensembl IDs to gene symbols. - + Parameters ---------- mdata : mudata.MuData @@ -23,24 +22,29 @@ def rename_gene_mygene(mdata): """ mg = mygene.MyGeneInfo() - gene_list = mdata['cNMF'].uns['var_names'].tolist() - annotations = mg.querymany(gene_list, scopes='ensembl.gene', fields='symbol', species='human') + gene_list = mdata["cNMF"].uns["var_names"].tolist() + annotations = mg.querymany( + gene_list, scopes="ensembl.gene", fields="symbol", species="human" + ) # Process the results to create mapping gene_dict = {} - + for item in annotations: - if 'symbol' in item: - gene_dict[item['query']] = item['symbol'] + if "symbol" in item: + gene_dict[item["query"]] = item["symbol"] + + mdata["cNMF"].uns["var_names"] = [ + gene_dict.get(x, x) for x in mdata["cNMF"].uns["var_names"] + ] - mdata['cNMF'].uns['var_names'] = [gene_dict.get(x, x) for x in mdata['cNMF'].uns['var_names']] def rename_gene_dictionary(mdata, dictionary_file_path): """ Map Ensembl gene IDs to gene symbols using a provided dictionary file. - + Rename genes in RNA modality using a dictionary file mapping Ensembl IDs to gene symbols. - + Parameters ---------- mdata : mudata.MuData @@ -50,8 +54,7 @@ def rename_gene_dictionary(mdata, dictionary_file_path): """ # Convert mapping result to list before assignment - df = pd.read_csv(dictionary_file_path, sep='\t') - ensemble_to_gene = dict(zip(df['ensembl_id'], df['gene'])) - new_names = [ensemble_to_gene.get(x, x) for x in mdata['rna'].var_names] - mdata['rna'].var_names = (new_names) - + df = pd.read_csv(dictionary_file_path, sep="\t") + ensemble_to_gene = dict(zip(df["ensembl_id"], df["gene"])) + new_names = [ensemble_to_gene.get(x, x) for x in mdata["rna"].var_names] + mdata["rna"].var_names = new_names diff --git a/src/Stage2_Evaluation/A_Metrics/src/__init__.py b/src/Stage2_Evaluation/A_Metrics/src/__init__.py index bf1df49..501b842 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/__init__.py +++ b/src/Stage2_Evaluation/A_Metrics/src/__init__.py @@ -1,10 +1,18 @@ from .association_categorical import compute_categorical_association from .association_perturbation import compute_perturbation_association -from .explained_variance import compute_explained_variance, compute_Var, computeVarianceExplained +from .explained_variance import ( + compute_explained_variance, + compute_Var, + computeVarianceExplained, +) from .enrichment_geneset import compute_geneset_enrichment from .enrichment_motif import compute_motif_enrichment, compute_motif_enrichment_ -from .enrichment_trait import run_opentargets_query, process_json_format_l2g_columns, filter_open_targets_gwas_query, compute_trait_enrichment +from .enrichment_trait import ( + run_opentargets_query, + process_json_format_l2g_columns, + filter_open_targets_gwas_query, + compute_trait_enrichment, +) - -from .Utilities import rename_gene_mygene,rename_gene_dictionary \ No newline at end of file +from .Utilities import rename_gene_mygene, rename_gene_dictionary diff --git a/src/Stage2_Evaluation/A_Metrics/src/association_categorical.py b/src/Stage2_Evaluation/A_Metrics/src/association_categorical.py index 2dad811..cc2a6e0 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/association_categorical.py +++ b/src/Stage2_Evaluation/A_Metrics/src/association_categorical.py @@ -14,11 +14,18 @@ from tqdm.auto import tqdm import logging -logging.basicConfig(level = logging.INFO) + +logging.basicConfig(level=logging.INFO) + # Perform non-parametric test (~one way ANOVA) to compare prog scores across categorical levels -def perform_kruskall_wallis(mdata, prog_key='prog', prog_name=None, - categorical_key='batch', pseudobulk_key='sample'): +def perform_kruskall_wallis( + mdata, + prog_key="prog", + prog_name=None, + categorical_key="batch", + pseudobulk_key="sample", +): """ Performs non-parametric one-way ANOVA using specified categorical levels. If pseudobulk key is provided then data is averaged over those levels. @@ -27,7 +34,7 @@ def perform_kruskall_wallis(mdata, prog_key='prog', prog_name=None, ARGS mdata : MuData mudata object containing anndata of program scores and cell-level metadata. - prog_key: + prog_key: index for the anndata object (mdata[prog_key]) in the mudata object. prog_name: str index of the feature (mdata[prog_key][:, prog_name]) which is the response variable. @@ -35,14 +42,14 @@ def perform_kruskall_wallis(mdata, prog_key='prog', prog_name=None, index of the categorical levels (mdata[prog_key].obs[categorical_key]) being tested. pseudobulk_key: str (optional) index of the feature summarisation (mean) levels (mdata[prog_key].obs[pseudobulk_key]). - + UPDATES if pseudobulk_key is None: store_key = categorical_key else: store_key = '_'.join(categorical_key, pseudobulk_key) - mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_stat'.format(store_key)] - mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_pval'.format(store_key)] + mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_stat'.format(store_key)] + mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_pval'.format(store_key)] mdata[prog_key].uns['{}_association_categories'.format(categorical_key)] """ @@ -50,51 +57,75 @@ def perform_kruskall_wallis(mdata, prog_key='prog', prog_name=None, if pseudobulk_key is None: # Run with single-cells are individual data points for category in mdata[prog_key].obs[categorical_key].astype(str).unique(): - sample_ = mdata[prog_key][mdata[prog_key].obs[categorical_key].astype(str)==category, - prog_name].X[:,0] + sample_ = mdata[prog_key][ + mdata[prog_key].obs[categorical_key].astype(str) == category, prog_name + ].X[:, 0] if sparse.issparse(sample_): sample_ = sample_.toarray().flatten() samples.append(sample_) else: # Pseudobulk the data before performing KW test (ideally these are biological replicates) - if sparse.issparse(mdata[prog_key][:, prog_name].X[:,0]): - prog_data = mdata[prog_key][:, prog_name].X[:,0].toarray() + if sparse.issparse(mdata[prog_key][:, prog_name].X[:, 0]): + prog_data = mdata[prog_key][:, prog_name].X[:, 0].toarray() else: - prog_data = mdata[prog_key][:, prog_name].X[:,0] - prog_data = pd.DataFrame(prog_data, columns=[prog_name], index=mdata[prog_key].obs.index) + prog_data = mdata[prog_key][:, prog_name].X[:, 0] + prog_data = pd.DataFrame( + prog_data, columns=[prog_name], index=mdata[prog_key].obs.index + ) prog_data[categorical_key] = mdata[prog_key].obs[categorical_key] prog_data[pseudobulk_key] = mdata[prog_key].obs[pseudobulk_key] - prog_data = prog_data.groupby([pseudobulk_key, categorical_key]).mean().dropna().reset_index() + prog_data = ( + prog_data.groupby([pseudobulk_key, categorical_key]) + .mean() + .dropna() + .reset_index() + ) for category in mdata[prog_key].obs[categorical_key].astype(str).unique(): - if len(prog_data.loc[prog_data[categorical_key]==category, pseudobulk_key].unique()) < 3: - raise ValueError('Cannot use less than 3 replicates per categorical level') - sample_ = prog_data.loc[prog_data[categorical_key].astype(str)==category, prog_name].values.flatten() + if ( + len( + prog_data.loc[ + prog_data[categorical_key] == category, pseudobulk_key + ].unique() + ) + < 3 + ): + raise ValueError( + "Cannot use less than 3 replicates per categorical level" + ) + sample_ = prog_data.loc[ + prog_data[categorical_key].astype(str) == category, prog_name + ].values.flatten() samples.append(sample_) - stat, pval = stats.kruskal(*samples, nan_policy='propagate') + stat, pval = stats.kruskal(*samples, nan_policy="propagate") # Store results if pseudobulk_key is None: store_key = categorical_key else: - store_key = '_'.join((categorical_key, pseudobulk_key)) + store_key = "_".join((categorical_key, pseudobulk_key)) - mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_stat'.format(store_key)] = stat - mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_pval'.format(store_key)] = pval + mdata[prog_key].var.loc[prog_name, "{}_kruskall_wallis_stat".format(store_key)] = ( + stat + ) + mdata[prog_key].var.loc[prog_name, "{}_kruskall_wallis_pval".format(store_key)] = ( + pval + ) - mdata[prog_key].uns['{}_association_categories'.format(categorical_key)] = \ - mdata[prog_key].obs[categorical_key].astype(str).unique() + mdata[prog_key].uns["{}_association_categories".format(categorical_key)] = ( + mdata[prog_key].obs[categorical_key].astype(str).unique() + ) def perform_correlation( prog_df: pd.DataFrame, group_col: str, val_col: str, - low_threshold: float=0.0, - correlation: Literal['pearsonr', 'spearmanr', 'kendalltau']='pearsonr', - mode: Literal['one_vs_all', 'one_vs_one']='one_vs_all', - df = [], + low_threshold: float = 0.0, + correlation: Literal["pearsonr", "spearmanr", "kendalltau"] = "pearsonr", + mode: Literal["one_vs_all", "one_vs_one"] = "one_vs_all", + df=[], ): """Perform post-hoc analysis with correlation tests @@ -129,83 +160,112 @@ def perform_correlation( # Remove cell memberships that are below the threshold prog_df = prog_df.loc[prog_df[val_col] > low_threshold] - + # Perform correlation tests - if mode=='one_vs_all': - stats_df = pd.DataFrame(index=categories, columns=['stat', 'pval', 'adj_pval', "log2FC"], dtype=float) + if mode == "one_vs_all": + stats_df = pd.DataFrame( + index=categories, + columns=["stat", "pval", "adj_pval", "log2FC"], + dtype=float, + ) for idx in stats_df.index.values: - if correlation=='pearsonr': - prog_df['binarized'] = prog_df[group_col].apply(lambda x: 1 if x==idx else 0) - stat, pval = stats.pearsonr(prog_df[val_col], prog_df['binarized']) + if correlation == "pearsonr": + prog_df["binarized"] = prog_df[group_col].apply( + lambda x: 1 if x == idx else 0 + ) + stat, pval = stats.pearsonr(prog_df[val_col], prog_df["binarized"]) _, adj_pval = fdrcorrection([pval]) - log2FC = np.log2(prog_df.loc[prog_df[group_col]==idx, val_col].mean() / prog_df.loc[prog_df[group_col]!=idx, val_col].mean()) - stats_df.loc[idx, 'stat'] = stat - stats_df.loc[idx, 'pval'] = pval - stats_df.loc[idx, 'adj_pval'] = adj_pval - stats_df.loc[idx, 'log2FC'] = log2FC - elif correlation=='spearmanr': - stat, pval = stats.spearmanr(prog_df[val_col], prog_df[group_col].astype('category').cat.codes) + log2FC = np.log2( + prog_df.loc[prog_df[group_col] == idx, val_col].mean() + / prog_df.loc[prog_df[group_col] != idx, val_col].mean() + ) + stats_df.loc[idx, "stat"] = stat + stats_df.loc[idx, "pval"] = pval + stats_df.loc[idx, "adj_pval"] = adj_pval + stats_df.loc[idx, "log2FC"] = log2FC + elif correlation == "spearmanr": + stat, pval = stats.spearmanr( + prog_df[val_col], prog_df[group_col].astype("category").cat.codes + ) _, adj_pval = fdrcorrection([pval]) - log2FC = np.log2(prog_df.loc[prog_df[group_col]==idx, val_col].mean() / prog_df.loc[prog_df[group_col]!=idx, val_col].mean()) - stats_df.loc[idx, 'stat'] = stat - stats_df.loc[idx, 'pval'] = pval - stats_df.loc[idx, 'adj_pval'] = adj_pval - stats_df.loc[idx, 'log2FC'] = log2FC - elif correlation=='kendalltau': - stat, pval = stats.kendalltau(prog_df[val_col], prog_df[group_col].astype('category').cat.codes) + log2FC = np.log2( + prog_df.loc[prog_df[group_col] == idx, val_col].mean() + / prog_df.loc[prog_df[group_col] != idx, val_col].mean() + ) + stats_df.loc[idx, "stat"] = stat + stats_df.loc[idx, "pval"] = pval + stats_df.loc[idx, "adj_pval"] = adj_pval + stats_df.loc[idx, "log2FC"] = log2FC + elif correlation == "kendalltau": + stat, pval = stats.kendalltau( + prog_df[val_col], prog_df[group_col].astype("category").cat.codes + ) _, adj_pval = fdrcorrection([pval]) - log2FC = np.log2(prog_df.loc[prog_df[group_col]==idx, val_col].mean() / prog_df.loc[prog_df[group_col]!=idx, val_col].mean()) - stats_df.loc[idx, 'stat'] = stat - stats_df.loc[idx, 'pval'] = pval - stats_df.loc[idx, 'adj_pval'] = adj_pval - stats_df.loc[idx, 'log2FC'] = log2FC - + log2FC = np.log2( + prog_df.loc[prog_df[group_col] == idx, val_col].mean() + / prog_df.loc[prog_df[group_col] != idx, val_col].mean() + ) + stats_df.loc[idx, "stat"] = stat + stats_df.loc[idx, "pval"] = pval + stats_df.loc[idx, "adj_pval"] = adj_pval + stats_df.loc[idx, "log2FC"] = log2FC + # Format wide_df = pd.DataFrame() - columns = ['stat', 'pval', 'adj_pval', 'log2FC'] + columns = ["stat", "pval", "adj_pval", "log2FC"] for idx, row in stats_df.iterrows(): for col in columns: - wide_df[f'{group_col}_{idx}_{correlation}_{col}'] = [row[col]] + wide_df[f"{group_col}_{idx}_{correlation}_{col}"] = [row[col]] wide_df.index = [val_col] - wide_df.index.name = 'program_name' + wide_df.index.name = "program_name" df.append(wide_df) - - elif mode=='one_vs_one': + + elif mode == "one_vs_one": pvals = pd.DataFrame(index=categories, columns=categories) for idx in pvals.index.values: for col in pvals.columns.values: - if idx==col: + if idx == col: continue else: test_df = prog_df.loc[prog_df[group_col].isin([idx, col])] - if correlation=='pearsonr': - stat, pval = stats.pearsonr(test_df[val_col], - test_df[group_col].astype('category').cat.codes) + if correlation == "pearsonr": + stat, pval = stats.pearsonr( + test_df[val_col], + test_df[group_col].astype("category").cat.codes, + ) pvals.loc[idx, col] = pval - elif correlation=='kendalltau': - stat, pval = stats.kendalltau(test_df[val_col], - test_df[group_col].astype('category').cat.codes) + elif correlation == "kendalltau": + stat, pval = stats.kendalltau( + test_df[val_col], + test_df[group_col].astype("category").cat.codes, + ) pvals.loc[idx, col] = pval return pvals # Perfom posthoc test and compute categorical-program score -def perform_posthoc(mdata, prog_key='prog', prog_name=None, - categorical_key='batch', pseudobulk_key='sample', - test='dunn', mode='one_vs_one', df=[]): - +def perform_posthoc( + mdata, + prog_key="prog", + prog_name=None, + categorical_key="batch", + pseudobulk_key="sample", + test="dunn", + mode="one_vs_one", + df=[], +): """ - Performs post-hoc tests for Kruskall-Wallis test and + Performs post-hoc tests for Kruskall-Wallis test and pairwise p-vals b/w categorical levels are reduced via mean and min. If pseudobulk key is provided then data is averaged over those levels. - Mudata object is updated in place with KW test results. + Mudata object is updated in place with KW test results. ARGS mdata : MuData mudata object containing anndata of program scores and cell-level metadata. - prog_key: + prog_key: index for the anndata object (mdata[prog_key]) in the mudata object. prog_name: str index of the feature (mdata[prog_key][:, prog_name]) which is the response variable. @@ -215,76 +275,95 @@ def perform_posthoc(mdata, prog_key='prog', prog_name=None, index of the feature summarisation (mean) levels (mdata[prog_key].obs[pseudobulk_key]). test: {'conover','dunn', 'dscf', 'pearsonr', 'kendalltau'} posthoc test to use to test categorical levels. - + UPDATES if pseudobulk_key is None: store_key = categorical_key else: store_key = '_'.join(categorical_key, pseudobulk_key) mdata[prog_key].varm['{}_association_{}_min_pval'.format(store_key, test)] - mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)] + mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)] mdata[prog_key].uns['{}_association_{}_pvals'.format(store_key, test)] """ - + store_key = categorical_key prog_data_ = mdata[prog_key][:, prog_name].X if sparse.issparse(prog_data_): prog_data_ = prog_data_.toarray() - prog_df = pd.DataFrame(prog_data_, - columns=[prog_name], - index=mdata[prog_key].obs.index) + prog_df = pd.DataFrame( + prog_data_, columns=[prog_name], index=mdata[prog_key].obs.index + ) prog_df[categorical_key] = mdata[prog_key].obs[categorical_key].astype(str).values if pseudobulk_key is not None: - store_key = '_'.join((categorical_key, pseudobulk_key)) + store_key = "_".join((categorical_key, pseudobulk_key)) prog_df[pseudobulk_key] = pseudobulk_key - prog_df = prog_df.groupby([pseudobulk_key, categorical_key]).mean().reset_index() + prog_df = ( + prog_df.groupby([pseudobulk_key, categorical_key]).mean().reset_index() + ) # Peform post hoc tests - if test=='dunn': + if test == "dunn": p_vals = posthoc_dunn(prog_df, val_col=prog_name, group_col=categorical_key) - elif test=='conover': + elif test == "conover": p_vals = posthoc_conover(prog_df, val_col=prog_name, group_col=categorical_key) - elif test=='dscf': + elif test == "dscf": p_vals = posthoc_dscf(prog_df, val_col=prog_name, group_col=categorical_key) - elif test=='pearsonr': - p_vals = perform_correlation(prog_df, val_col=prog_name, group_col=categorical_key, - correlation='pearsonr', mode=mode, df=df) - if mode == 'one_vs_all': + elif test == "pearsonr": + p_vals = perform_correlation( + prog_df, + val_col=prog_name, + group_col=categorical_key, + correlation="pearsonr", + mode=mode, + df=df, + ) + if mode == "one_vs_all": return - elif test=='kendalltau': - p_vals = perform_correlation(prog_df, val_col=prog_name, group_col=categorical_key, - correlation='kendalltau') + elif test == "kendalltau": + p_vals = perform_correlation( + prog_df, + val_col=prog_name, + group_col=categorical_key, + correlation="kendalltau", + ) # Create combined statistic with p-vals for each categorical level - for i, category in enumerate(mdata[prog_key].obs[categorical_key].astype(str).unique()): - min_pval = np.min(p_vals.loc[p_vals.index!=category, category]) - mean_pval = np.mean(p_vals.loc[p_vals.index!=category, category]) + for i, category in enumerate( + mdata[prog_key].obs[categorical_key].astype(str).unique() + ): + min_pval = np.min(p_vals.loc[p_vals.index != category, category]) + mean_pval = np.mean(p_vals.loc[p_vals.index != category, category]) prog_idx = mdata[prog_key].var.index.get_loc(prog_name) # Store results - mdata[prog_key].varm['{}_association_{}_min_pval'.format(store_key, test)][prog_idx,i] = min_pval - mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)][prog_idx,i] = mean_pval - - mdata[prog_key].uns['{}_association_{}_pvals'.format(store_key, test)][prog_name] = p_vals + mdata[prog_key].varm["{}_association_{}_min_pval".format(store_key, test)][ + prog_idx, i + ] = min_pval + mdata[prog_key].varm["{}_association_{}_mean_pval".format(store_key, test)][ + prog_idx, i + ] = mean_pval + + mdata[prog_key].uns["{}_association_{}_pvals".format(store_key, test)][ + prog_name + ] = p_vals # TODO: Add one way ANOVA, MANOVA, multi-variate KW, logistic-regression def compute_categorical_association( - mdata, - prog_key='prog', - categorical_key='batch', - pseudobulk_key=None, - test='dunn', - mode='one_vs_one', - n_jobs=1, - inplace=True, - **kwargs + mdata, + prog_key="prog", + categorical_key="batch", + pseudobulk_key=None, + test="dunn", + mode="one_vs_one", + n_jobs=1, + inplace=True, + **kwargs, ): - """ Compute association of continous features with categorical levels. Currently only the Kruskall-Wallis test is implemented to determine if @@ -294,7 +373,7 @@ def compute_categorical_association( ARGS mdata : MuData mudata object containing anndata of program scores and cell-level metadata. - prog_key: + prog_key: index for the anndata object (mdata[prog_key]) in the mudata object. categorical_key: str index of the categorical levels (mdata[prog_key].obs[categorical_key]) being tested. @@ -306,161 +385,226 @@ def compute_categorical_association( number of threads to run processes on. inplace: Bool (default: True) update the mudata object inplace or return a copy - + --- if inplace: UPDATES - mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_stat'.format(store_key)] - mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_pval'.format(store_key)] + mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_stat'.format(store_key)] + mdata[prog_key].var.loc[prog_name, '{}_kruskall_wallis_pval'.format(store_key)] mdata[prog_key].varm['{}_association_{}_min_pval'.format(store_key, test)] - mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)] + mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)] mdata[prog_key].uns['{}_association_{}_pvals'.format(store_key, test)] - mdata[prog_key].uns['{}_association_categories'.format(categorical_key)] + mdata[prog_key].uns['{}_association_categories'.format(categorical_key)] else: RETURNS - results_df, posthoc_df + results_df, posthoc_df """ # Read in mudata if it is provided as a path - frompath=False + frompath = False if isinstance(mdata, str): if os.path.exists(mdata): mdata = mudata.read(mdata) if inplace: - logging.warning('Changed to inplace=False since path was provided') - inplace=False - frompath=True - else: raise ValueError('Incorrect mudata specification.') + logging.warning("Changed to inplace=False since path was provided") + inplace = False + frompath = True + else: + raise ValueError("Incorrect mudata specification.") if not inplace and not frompath: mdata = mudata.MuData({prog_key: mdata[prog_key].copy()}) - + if pseudobulk_key is None: - logging.info('Performing tests at single-cell level. Significance will likely be inflated') + logging.info( + "Performing tests at single-cell level. Significance will likely be inflated" + ) store_key = categorical_key else: - logging.info('Perform testing by averaging over {}'.format(pseudobulk_key)) - store_key = '_'.join((categorical_key, pseudobulk_key)) + logging.info("Perform testing by averaging over {}".format(pseudobulk_key)) + store_key = "_".join((categorical_key, pseudobulk_key)) + + mdata[prog_key].var["{}_kruskall_wallis_stat".format(store_key)] = None + mdata[prog_key].var["{}_kruskall_wallis_pval".format(store_key)] = None - mdata[prog_key].var['{}_kruskall_wallis_stat'.format(store_key)] = None - mdata[prog_key].var['{}_kruskall_wallis_pval'.format(store_key)] = None - # Run KW test - Parallel(n_jobs=n_jobs, - backend='threading')(delayed(perform_kruskall_wallis)(mdata, - prog_key=prog_key, - prog_name=prog_name, - categorical_key=categorical_key, - pseudobulk_key=pseudobulk_key) \ - for prog_name in tqdm(mdata[prog_key].var_names, - desc='Testing {} association'.format(categorical_key), - unit='programs')) + Parallel(n_jobs=n_jobs, backend="threading")( + delayed(perform_kruskall_wallis)( + mdata, + prog_key=prog_key, + prog_name=prog_name, + categorical_key=categorical_key, + pseudobulk_key=pseudobulk_key, + ) + for prog_name in tqdm( + mdata[prog_key].var_names, + desc="Testing {} association".format(categorical_key), + unit="programs", + ) + ) # Convert to float to prevent error when saving mudata - mdata[prog_key].var['{}_kruskall_wallis_stat'.format(store_key)] = \ - mdata[prog_key].var['{}_kruskall_wallis_stat'.format(store_key)].astype(float) - mdata[prog_key].var['{}_kruskall_wallis_pval'.format(store_key)] = \ - mdata[prog_key].var['{}_kruskall_wallis_pval'.format(store_key)].astype(float) - + mdata[prog_key].var["{}_kruskall_wallis_stat".format(store_key)] = ( + mdata[prog_key].var["{}_kruskall_wallis_stat".format(store_key)].astype(float) + ) + mdata[prog_key].var["{}_kruskall_wallis_pval".format(store_key)] = ( + mdata[prog_key].var["{}_kruskall_wallis_pval".format(store_key)].astype(float) + ) # If test = 'pearsonr' with mode = 'one_vs_all' we are going to do something special until we can more properly integrate this - if test=='pearsonr' and mode=='one_vs_all': - logging.info('Running jamboree specific version of posthoc with pearsonr, this is not yet integrated into the main pipeline') - res = [] + if test == "pearsonr" and mode == "one_vs_all": + logging.info( + "Running jamboree specific version of posthoc with pearsonr, this is not yet integrated into the main pipeline" + ) + res = [] # Run posthoc-tests and append results - Parallel(n_jobs=n_jobs, - backend='threading')(delayed(perform_posthoc)(mdata, - prog_key=prog_key, - prog_name=prog_name, - categorical_key=categorical_key, - pseudobulk_key=pseudobulk_key, - test=test, - mode=mode, - df=res) \ - for prog_name in tqdm(mdata[prog_key].var_names, - desc='Identifying differential {}'.format(categorical_key), - unit='programs')) - + Parallel(n_jobs=n_jobs, backend="threading")( + delayed(perform_posthoc)( + mdata, + prog_key=prog_key, + prog_name=prog_name, + categorical_key=categorical_key, + pseudobulk_key=pseudobulk_key, + test=test, + mode=mode, + df=res, + ) + for prog_name in tqdm( + mdata[prog_key].var_names, + desc="Identifying differential {}".format(categorical_key), + unit="programs", + ) + ) + posthoc_df = pd.concat(res, axis=0) - results_df = mdata[prog_key].var.loc[:, ['{}_kruskall_wallis_stat'.format(store_key), - '{}_kruskall_wallis_pval'.format(store_key)]] + results_df = mdata[prog_key].var.loc[ + :, + [ + "{}_kruskall_wallis_stat".format(store_key), + "{}_kruskall_wallis_pval".format(store_key), + ], + ] posthoc_df["program_name"] = mdata[prog_key].var_names results_df["program_name"] = mdata[prog_key].var_names return (results_df, posthoc_df) else: + mdata[prog_key].varm["{}_association_{}_min_pval".format(store_key, test)] = ( + np.zeros( + ( + mdata[prog_key].shape[1], + mdata[prog_key].obs[categorical_key].unique().shape[0], + ) + ) + ) + mdata[prog_key].varm["{}_association_{}_mean_pval".format(store_key, test)] = ( + np.ones( + ( + mdata[prog_key].shape[1], + mdata[prog_key].obs[categorical_key].unique().shape[0], + ) + ) + ) + mdata[prog_key].uns["{}_association_{}_pvals".format(store_key, test)] = {} - mdata[prog_key].varm['{}_association_{}_min_pval'.format(store_key, test)] = \ - np.zeros((mdata[prog_key].shape[1], mdata[prog_key].obs[categorical_key].unique().shape[0])) - mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)] = \ - np.ones((mdata[prog_key].shape[1], mdata[prog_key].obs[categorical_key].unique().shape[0])) - mdata[prog_key].uns['{}_association_{}_pvals'.format(store_key, test)] = {} - # Run posthoc-tests - Parallel(n_jobs=n_jobs, - backend='threading')(delayed(perform_posthoc)(mdata, - prog_key=prog_key, - prog_name=prog_name, - categorical_key=categorical_key, - pseudobulk_key=pseudobulk_key, - test=test) \ - for prog_name in tqdm(mdata[prog_key].var_names, - desc='Identifying differential {}'.format(categorical_key), - unit='programs')) - + Parallel(n_jobs=n_jobs, backend="threading")( + delayed(perform_posthoc)( + mdata, + prog_key=prog_key, + prog_name=prog_name, + categorical_key=categorical_key, + pseudobulk_key=pseudobulk_key, + test=test, + ) + for prog_name in tqdm( + mdata[prog_key].var_names, + desc="Identifying differential {}".format(categorical_key), + unit="programs", + ) + ) + # Returning test results only - if not inplace: - - results_df = mdata[prog_key].var.loc[:, ['{}_kruskall_wallis_stat'.format(store_key), - '{}_kruskall_wallis_pval'.format(store_key)]] - min_pval_df = pd.DataFrame(mdata[prog_key].varm['{}_association_{}_min_pval'.format(store_key, test)], - index=mdata[prog_key].var.index, - columns=['{}_{}_association_{}_min_pval'.format(col, store_key, test) \ - for col in mdata[prog_key].uns['{}_association_categories'.format(categorical_key)]] - ) - results_df = results_df.merge(min_pval_df, left_index=True, right_index=True) - - mean_pval_df = pd.DataFrame(mdata[prog_key].varm['{}_association_{}_mean_pval'.format(store_key, test)], - index=mdata[prog_key].var.index, - columns=['{}_{}_association_{}_mean_pval'.format(col, store_key, test) \ - for col in mdata[prog_key].uns['{}_association_categories'.format(categorical_key)]] - ) - results_df = results_df.merge(mean_pval_df, left_index=True, right_index=True) + if not inplace: + results_df = mdata[prog_key].var.loc[ + :, + [ + "{}_kruskall_wallis_stat".format(store_key), + "{}_kruskall_wallis_pval".format(store_key), + ], + ] + min_pval_df = pd.DataFrame( + mdata[prog_key].varm[ + "{}_association_{}_min_pval".format(store_key, test) + ], + index=mdata[prog_key].var.index, + columns=[ + "{}_{}_association_{}_min_pval".format(col, store_key, test) + for col in mdata[prog_key].uns[ + "{}_association_categories".format(categorical_key) + ] + ], + ) + results_df = results_df.merge( + min_pval_df, left_index=True, right_index=True + ) + + mean_pval_df = pd.DataFrame( + mdata[prog_key].varm[ + "{}_association_{}_mean_pval".format(store_key, test) + ], + index=mdata[prog_key].var.index, + columns=[ + "{}_{}_association_{}_mean_pval".format(col, store_key, test) + for col in mdata[prog_key].uns[ + "{}_association_categories".format(categorical_key) + ] + ], + ) + results_df = results_df.merge( + mean_pval_df, left_index=True, right_index=True + ) posthoc_df = [] - dict_ = mdata[prog_key].uns['{}_association_{}_pvals'.format(store_key, test)] + dict_ = mdata[prog_key].uns[ + "{}_association_{}_pvals".format(store_key, test) + ] for key, values in dict_.items(): - tmp_df = values.where(np.triu(np.ones(values.shape), k=1).astype(bool)) tmp_df = tmp_df.stack().reset_index() - tmp_df.columns = ['{}_a'.format(store_key), - '{}_b'.format(store_key), - 'p_value'] - tmp_df['program_name'] = key + tmp_df.columns = [ + "{}_a".format(store_key), + "{}_b".format(store_key), + "p_value", + ] + tmp_df["program_name"] = key posthoc_df.append(tmp_df) posthoc_df = pd.concat(posthoc_df, ignore_index=True) return (results_df, posthoc_df) - - -if __name__=='__main__': + + +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('mudataObj_path') - parser.add_argument('categorical_key', type=str) - parser.add_argument('--pseudobulk_key', default=None, type=str) - parser.add_argument('--test', default='dunn', type=str) - parser.add_argument('-pk', '--prog_key', default='prog', type=str) - parser.add_argument('-n', '--n_jobs', default=1, type=int) - parser.add_argument('--output', action='store_false') + parser.add_argument("mudataObj_path") + parser.add_argument("categorical_key", type=str) + parser.add_argument("--pseudobulk_key", default=None, type=str) + parser.add_argument("--test", default="dunn", type=str) + parser.add_argument("-pk", "--prog_key", default="prog", type=str) + parser.add_argument("-n", "--n_jobs", default=1, type=int) + parser.add_argument("--output", action="store_false") args = parser.parse_args() mdata = mudata.read(args.mudataObj_path) - compute_categorical_association(mdata, prog_key=args.prog_key, categorical_key=args.categorical_key, - pseudobulk_key=args.pseudobulk_key, test=args.test, n_jobs=args.n_jobs, - inplace=args.output) - - - + compute_categorical_association( + mdata, + prog_key=args.prog_key, + categorical_key=args.categorical_key, + pseudobulk_key=args.pseudobulk_key, + test=args.test, + n_jobs=args.n_jobs, + inplace=args.output, + ) diff --git a/src/Stage2_Evaluation/A_Metrics/src/association_perturbation.py b/src/Stage2_Evaluation/A_Metrics/src/association_perturbation.py index 1f96199..ee64963 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/association_perturbation.py +++ b/src/Stage2_Evaluation/A_Metrics/src/association_perturbation.py @@ -14,14 +14,15 @@ import mudata import logging -logging.basicConfig(level = logging.INFO) + +logging.basicConfig(level=logging.INFO) def get_guide_metadata( mdata: mudata.MuData, prog_key: str, - guide_names_key: str = 'guide_names', - guide_targets_key: str = 'guide_targets' + guide_names_key: str = "guide_names", + guide_targets_key: str = "guide_targets", ) -> pd.DataFrame: """ Get guide metadata from the mudata object. @@ -38,8 +39,10 @@ def get_guide_metadata( pd.DataFrame The guide metadata. """ - guide_metadata = pd.DataFrame(index=mdata[prog_key].uns[guide_names_key], columns=['Target']) - guide_metadata['Target'] = mdata[prog_key].uns[guide_targets_key] + guide_metadata = pd.DataFrame( + index=mdata[prog_key].uns[guide_names_key], columns=["Target"] + ) + guide_metadata["Target"] = mdata[prog_key].uns[guide_targets_key] return guide_metadata @@ -53,7 +56,7 @@ def compute_perturbation_association_( program: str, level_name: str, test_stats_df: List[List], - balanced: bool = False + balanced: bool = False, ): # Get perturbation data test_data_ = test_data[:, program].X.toarray() @@ -61,34 +64,46 @@ def compute_perturbation_association_( # TODO: Resample reference data to match number of obs reference_data_ = reference_data[:, program].X.toarray() if balanced: - reference_data_ = reference_data_[np.random.choice(reference_data_.shape[0], test_data_.shape[0], replace=False)] + reference_data_ = reference_data_[ + np.random.choice( + reference_data_.shape[0], test_data_.shape[0], replace=False + ) + ] # Calculate log2FC ref_mean = np.mean(reference_data_) test_mean = np.mean(test_data_) log2fc = np.log2((test_mean + 1e-9) / (ref_mean + 1e-9)) - + # Compute Mann-Whitney U test - results = stats.mannwhitneyu(test_data_, reference_data_, alternative='two-sided', nan_policy='omit', method='asymptotic') + results = stats.mannwhitneyu( + test_data_, + reference_data_, + alternative="two-sided", + nan_policy="omit", + method="asymptotic", + ) # Append to test stats df - test_stats_df.append([level_name, program, ref_mean, test_mean, log2fc, results[0][0], results[1][0]]) + test_stats_df.append( + [level_name, program, ref_mean, test_mean, log2fc, results[0][0], results[1][0]] + ) # TODO: Add support for stratification by categorical levels def compute_perturbation_association( mdata: Union[os.PathLike, mudata.MuData], prog_key: str, - guide_names_key: str = 'guide_names', - guide_targets_key: str = 'guide_targets', - guide_assignments_key: str = 'guide_assignment', + guide_names_key: str = "guide_names", + guide_targets_key: str = "guide_targets", + guide_assignments_key: str = "guide_assignment", collapse_targets: bool = True, pseudobulk: bool = False, - reference_targets: Union[str, List[str]] = ['non-targeting'], + reference_targets: Union[str, List[str]] = ["non-targeting"], balanced: bool = False, n_jobs: int = 1, inplace: bool = True, - FDR_method: Literal['BH', 'StoreyQ'] = 'StoreyQ' + FDR_method: Literal["BH", "StoreyQ"] = "StoreyQ", ): """Compute Mann-Whitney U test for cells targeted with a perturbation against a reference set of cells. @@ -105,7 +120,7 @@ def compute_perturbation_association( guide_assignments_key : str (default: 'guide_assignment') key in mdata[prog_key].obsm for 2D cell x guide assignment matrix. Should correspond to guide_names and guide_targets. collapse_targets : bool (default:True) - If target gene per guide is provided, perform tests on target levels. + If target gene per guide is provided, perform tests on target levels. Mutually exclusive with pseudobulk. pseudobulk : bool (default: False) If multiple non-targeting guides are available - optionally test at pseudobulk level. @@ -126,50 +141,54 @@ def compute_perturbation_association( if inplace: UPDATES mdata[prog_key].varm['perturbation_association_pval'] - mdata[prog_key].varm['perturbation_association_stat'] + mdata[prog_key].varm['perturbation_association_stat'] else: RETURNS returns test_stats_df """ # Read in mudata if it is provided as a path - frompath=False + frompath = False if isinstance(mdata, str): if os.path.exists(mdata): mdata = mudata.read(mdata) if inplace: - logging.warning('Changed to inplace=False since path was provided') - inplace=False - frompath=True - else: raise ValueError('Incorrect mudata specification.') + logging.warning("Changed to inplace=False since path was provided") + inplace = False + frompath = True + else: + raise ValueError("Incorrect mudata specification.") if not inplace and not frompath: - mdata = mudata.MuData({ - prog_key: mdata[prog_key].copy() - }) + mdata = mudata.MuData({prog_key: mdata[prog_key].copy()}) # Guide metadata guide_metadata = get_guide_metadata( - mdata, + mdata, prog_key=prog_key, - guide_names_key=guide_names_key, - guide_targets_key=guide_targets_key + guide_names_key=guide_names_key, + guide_targets_key=guide_targets_key, ) # Run tests on pseudobulk if multiple non-targeting guides otherwise at single-cell level if collapse_targets: if pseudobulk: - pseudobulk=False - logging.info('Setting pseudobulk to False since collapse_targets is True') + pseudobulk = False + logging.info("Setting pseudobulk to False since collapse_targets is True") if pseudobulk: - pseudobulked_scores = pd.DataFrame(index=mdata[prog_key].uns[guide_names_key], columns=mdata[prog_key].obs_names) + pseudobulked_scores = pd.DataFrame( + index=mdata[prog_key].uns[guide_names_key], + columns=mdata[prog_key].obs_names, + ) - #TODO: Parallelize - for i, guide in enumerate(mdata[prog_key].uns['guide_names']): + # TODO: Parallelize + for i, guide in enumerate(mdata[prog_key].uns["guide_names"]): for program in mdata[prog_key].obs_names: - pseudobulked_scores.loc[guide, program] = \ - mdata[prog_key][mdata[prog_key].obsm[guide_assignments_key][:,i].astype(bool), program].X.mean() - - pseudobulked_scores['Target'] = mdata[prog_key].uns['guide_targets'] + pseudobulked_scores.loc[guide, program] = mdata[prog_key][ + mdata[prog_key].obsm[guide_assignments_key][:, i].astype(bool), + program, + ].X.mean() + + pseudobulked_scores["Target"] = mdata[prog_key].uns["guide_targets"] # Create reference data if type(reference_targets) is str: @@ -177,53 +196,70 @@ def compute_perturbation_association( # If pseudobulk, grab a DataFrame of all non-targeting guides if pseudobulk: - reference_data = pseudobulked_scores.loc[pseudobulked_scores.Target.isin(reference_targets)] + reference_data = pseudobulked_scores.loc[ + pseudobulked_scores.Target.isin(reference_targets) + ] # Otherwise, grab a MuData for all those guides assigned to the reference targets else: - reference_guide_idx = guide_metadata.index.get_indexer(guide_metadata.loc[guide_metadata.Target.isin(reference_targets)].index.values) + reference_guide_idx = guide_metadata.index.get_indexer( + guide_metadata.loc[ + guide_metadata.Target.isin(reference_targets) + ].index.values + ) assign_mat = mdata[prog_key].obsm[guide_assignments_key] if sparse.issparse(assign_mat): assign_mat = assign_mat.toarray() - reference_data = mdata[prog_key][np.any(assign_mat[:,reference_guide_idx].astype(bool), axis=1)] + reference_data = mdata[prog_key][ + np.any(assign_mat[:, reference_guide_idx].astype(bool), axis=1) + ] # Get the rest of the guides to test against the reference - test_guides = guide_metadata.loc[~guide_metadata.Target.isin(reference_targets)].index.values + test_guides = guide_metadata.loc[ + ~guide_metadata.Target.isin(reference_targets) + ].index.values # Run tests test_stats_df = [] # If collapse_targets, run tests for each target if collapse_targets: - # Set level key - level_key = 'target' + level_key = "target" # Run tests for each target (with multiple guides) - targets_no_ref = [targ for targ in guide_metadata['Target'].unique() if targ not in reference_targets] - for target in tqdm(targets_no_ref, desc='Testing perturbation association', unit='targets'): - + targets_no_ref = [ + targ + for targ in guide_metadata["Target"].unique() + if targ not in reference_targets + ] + for target in tqdm( + targets_no_ref, desc="Testing perturbation association", unit="targets" + ): # Grab data for the current target - test_guide_idx = guide_metadata.index.get_indexer(guide_metadata.index[guide_metadata['Target']==target]) - test_data = mdata[prog_key][assign_mat[:,test_guide_idx].any(-1).astype(bool)] + test_guide_idx = guide_metadata.index.get_indexer( + guide_metadata.index[guide_metadata["Target"] == target] + ) + test_data = mdata[prog_key][ + assign_mat[:, test_guide_idx].any(-1).astype(bool) + ] # Run test for every program - Parallel(n_jobs=n_jobs, backend='threading')(delayed(compute_perturbation_association_)( - test_data, - reference_data, - program, - target, - test_stats_df, - balanced) for program in mdata[prog_key].var_names) + Parallel(n_jobs=n_jobs, backend="threading")( + delayed(compute_perturbation_association_)( + test_data, reference_data, program, target, test_stats_df, balanced + ) + for program in mdata[prog_key].var_names + ) else: - # Set level key - level_key='guide' + level_key = "guide" # Run tests for each guide - for guide in tqdm(test_guides, desc='Testing perturbation association', unit='guides'): - + for guide in tqdm( + test_guides, desc="Testing perturbation association", unit="guides" + ): # Run test at using pseudobulks at guide level -> requires multiple non-targeting guides as reference if pseudobulk: # TODO: Implement per guide probability under reference approach @@ -231,50 +267,84 @@ def compute_perturbation_association( else: # Grab data for the current guide test_guide_idx = guide_metadata.index.get_loc(guide) - test_data = mdata[prog_key][assign_mat[:,test_guide_idx].astype(bool)] + test_data = mdata[prog_key][assign_mat[:, test_guide_idx].astype(bool)] # Run test for every program - Parallel(n_jobs=n_jobs, backend='threading')(delayed(compute_perturbation_association_)( - test_data, - reference_data, - program, - guide, - test_stats_df, - balanced) for program in mdata[prog_key].var_names) - + Parallel(n_jobs=n_jobs, backend="threading")( + delayed(compute_perturbation_association_)( + test_data, + reference_data, + program, + guide, + test_stats_df, + balanced, + ) + for program in mdata[prog_key].var_names + ) + # Create DataFrame - test_stats_df = pd.DataFrame(test_stats_df, columns=['{}_name'.format(level_key), 'program_name', 'ref_mean', 'test_mean', 'log2FC', 'stat', 'pval']) + test_stats_df = pd.DataFrame( + test_stats_df, + columns=[ + "{}_name".format(level_key), + "program_name", + "ref_mean", + "test_mean", + "log2FC", + "stat", + "pval", + ], + ) # correct invalid p-value by resetting to 1 - p = test_stats_df['pval'].copy() + p = test_stats_df["pval"].copy() invalid = (~np.isfinite(p)) | (p < 0) | (p > 1) p[invalid] = 1.0 - test_stats_df['pval'] = p + test_stats_df["pval"] = p # Correct for multiple testing - if FDR_method == 'BH': - test_stats_df['adj_pval'] = multipletests(test_stats_df['pval'], method='fdr_bh')[1] - elif FDR_method == 'StoreyQ': + if FDR_method == "BH": + test_stats_df["adj_pval"] = multipletests( + test_stats_df["pval"], method="fdr_bh" + )[1] + elif FDR_method == "StoreyQ": from multipy.fdr import qvalue import builtins - if not hasattr(builtins, 'xrange'): + + if not hasattr(builtins, "xrange"): builtins.xrange = range - test_stats_df['adj_pval'] = qvalue(test_stats_df['pval'].values, threshold=0.05, verbose=False)[1] + test_stats_df["adj_pval"] = qvalue( + test_stats_df["pval"].values, threshold=0.05, verbose=False + )[1] # Return only the evaluations if not in inplace mode - if not inplace: return test_stats_df + if not inplace: + return test_stats_df else: - init_array = np.empty((mdata[prog_key].shape[1], len(test_stats_df['{}_name'.format(level_key)].unique()))) + init_array = np.empty( + ( + mdata[prog_key].shape[1], + len(test_stats_df["{}_name".format(level_key)].unique()), + ) + ) init_array[:] = np.nan stats, pvals = init_array.copy(), init_array.copy() - for level_idx, level_name in enumerate(test_stats_df['{}_name'.format(level_key)].unique()): - level_df = test_stats_df[test_stats_df['{}_name'.format(level_key)]==level_name].set_index('program_name') - stats[:, level_idx] = level_df.loc[mdata[prog_key].var_names, 'stat'].values - pvals[:, level_idx] = level_df.loc[mdata[prog_key].var_names, 'pval'].values - mdata[prog_key].varm['perturbation_association_{}_stat'.format(level_key)] = stats - mdata[prog_key].varm['perturbation_association_{}_pval'.format(level_key)] = pvals - mdata[prog_key].uns['perturbation_association_{}_names'.format(level_key)] = test_stats_df['{}_name'.format(level_key)].unique() - - \ No newline at end of file + for level_idx, level_name in enumerate( + test_stats_df["{}_name".format(level_key)].unique() + ): + level_df = test_stats_df[ + test_stats_df["{}_name".format(level_key)] == level_name + ].set_index("program_name") + stats[:, level_idx] = level_df.loc[mdata[prog_key].var_names, "stat"].values + pvals[:, level_idx] = level_df.loc[mdata[prog_key].var_names, "pval"].values + mdata[prog_key].varm["perturbation_association_{}_stat".format(level_key)] = ( + stats + ) + mdata[prog_key].varm["perturbation_association_{}_pval".format(level_key)] = ( + pvals + ) + mdata[prog_key].uns["perturbation_association_{}_names".format(level_key)] = ( + test_stats_df["{}_name".format(level_key)].unique() + ) diff --git a/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py b/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py index 6fb9bd0..209dc5c 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py +++ b/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py @@ -14,13 +14,14 @@ logging.basicConfig(level=logging.INFO) -_RESOURCES_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "Resources") +_RESOURCES_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), + "Resources", +) def create_geneset_dict( - dataframe: pd.DataFrame, - key_column='trait_efos', - gene_column='gene_name' + dataframe: pd.DataFrame, key_column="trait_efos", gene_column="gene_name" ): """Create custom geneset""" @@ -31,7 +32,8 @@ def create_geneset_dict( geneset_dict.setdefault(key, []).append(gene) return geneset_dict -def get_idconversion(var_names, organism='human', chunk_size=100): + +def get_idconversion(var_names, organism="human", chunk_size=100): """Match gene IDs with database, chunking large queries to avoid HTML errors.""" # Ensure var_names is a simple list var_names = list(var_names) @@ -40,55 +42,74 @@ def get_idconversion(var_names, organism='human', chunk_size=100): gene_names = [v.upper() for v in var_names] # Detect Ensembl IDs by looking at the first element safely - is_ensembl = (len(var_names) > 0 and var_names[0].lower().startswith('ens')) + is_ensembl = len(var_names) > 0 and var_names[0].lower().startswith("ens") if is_ensembl: bm = Biomart() dataset = ( - 'hsapiens_gene_ensembl' - if organism=='human' - else 'mmusculus_gene_ensembl' + "hsapiens_gene_ensembl" if organism == "human" else "mmusculus_gene_ensembl" ) lookup = {} for i in range(0, len(var_names), chunk_size): - chunk = var_names[i:i+chunk_size] + chunk = var_names[i : i + chunk_size] try: res = bm.query( dataset=dataset, - attributes=['ensembl_gene_id', 'external_gene_name'], - filters={'ensembl_gene_id': chunk} + attributes=["ensembl_gene_id", "external_gene_name"], + filters={"ensembl_gene_id": chunk}, + ) + if not hasattr(res, "columns"): + raise RuntimeError("BioMart returned HTML or invalid format") + df = res[["ensembl_gene_id", "external_gene_name"]].astype(str) + df["external_gene_name"] = df["external_gene_name"].str.upper() + lookup.update( + dict(zip(df["ensembl_gene_id"], df["external_gene_name"])) ) - if not hasattr(res, 'columns'): - raise RuntimeError('BioMart returned HTML or invalid format') - df = res[['ensembl_gene_id', 'external_gene_name']].astype(str) - df['external_gene_name'] = df['external_gene_name'].str.upper() - lookup.update(dict(zip(df['ensembl_gene_id'], df['external_gene_name']))) except Exception as e: - logging.info(f"BioMart chunk lookup failed for IDs {i}-{i+len(chunk)} ({e}); skipping these.") + logging.info( + f"BioMart chunk lookup failed for IDs {i}-{i + len(chunk)} ({e}); skipping these." + ) gene_names = [lookup.get(v, v.upper()) for v in var_names] - elif any(':ens' in v.lower() for v in var_names): - gene_names = [v.split(':', 1)[0].upper() for v in var_names] + elif any(":ens" in v.lower() for v in var_names): + gene_names = [v.split(":", 1)[0].upper() for v in var_names] return gene_names - -def get_program_gene_loadings(mdata, prog_key='prog', prog_nam=None, data_key='rna', organism='human', gene_names_key=None): + + +def get_program_gene_loadings( + mdata, + prog_key="prog", + prog_nam=None, + data_key="rna", + organism="human", + gene_names_key=None, +): """Get gene loadings for each program in the mudata object.""" if gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: - gene_names = mdata[data_key].var[gene_names_key].astype(str).str.upper().tolist() - elif 'var_names' in mdata[prog_key].uns.keys(): - gene_names = get_idconversion(mdata[prog_key].uns['var_names'], organism=organism) + gene_names = ( + mdata[data_key].var[gene_names_key].astype(str).str.upper().tolist() + ) + elif "var_names" in mdata[prog_key].uns.keys(): + gene_names = get_idconversion( + mdata[prog_key].uns["var_names"], organism=organism + ) else: - assert mdata[prog_key].varm['loadings'].shape[1] == mdata[data_key].var.shape[0] + assert mdata[prog_key].varm["loadings"].shape[1] == mdata[data_key].var.shape[0] gene_names = get_idconversion(mdata[data_key].var_names, organism=organism) if prog_nam: - loadings = pd.DataFrame(data=mdata[prog_key][:, prog_nam].varm['loadings'].flatten(), index=gene_names) + loadings = pd.DataFrame( + data=mdata[prog_key][:, prog_nam].varm["loadings"].flatten(), + index=gene_names, + ) loadings.columns = [prog_nam] else: - loadings = pd.DataFrame(data=mdata[prog_key].varm['loadings'], index=mdata[prog_key].var.index).T + loadings = pd.DataFrame( + data=mdata[prog_key].varm["loadings"], index=mdata[prog_key].var.index + ).T loadings["gene_names"] = gene_names loadings.set_index("gene_names", inplace=True) @@ -99,7 +120,14 @@ def get_program_gene_loadings(mdata, prog_key='prog', prog_nam=None, data_key='r return loadings -def get_geneset(organism='human', library='h.all', database='msigdb', min_size: int = 0, max_size: int = 2000, use_cache: bool = False): +def get_geneset( + organism="human", + library="h.all", + database="msigdb", + min_size: int = 0, + max_size: int = 2000, + use_cache: bool = False, +): """Download gene set from MsigDB or Enrichr. Parameters @@ -116,21 +144,28 @@ def get_geneset(organism='human', library='h.all', database='msigdb', min_size: with open(cache_file) as f: return json.load(f) else: - logging.info(f"Cache file not found ({cache_file}), downloading from {database}") + logging.info( + f"Cache file not found ({cache_file}), downloading from {database}" + ) - if database == 'msigdb': + if database == "msigdb": msig = Msigdb() - dbver = '2023.2.Hs' if organism == 'human' else '2023.1.Mm' + dbver = "2023.2.Hs" if organism == "human" else "2023.1.Mm" gmt = msig.get_gmt(category=library, dbver=dbver) if gmt is None: - raise ValueError('Library does not exist') - elif database == 'enrichr': - gmt = gp.get_library(name=library, organism=organism.capitalize(), min_size=min_size, max_size=max_size) + raise ValueError("Library does not exist") + elif database == "enrichr": + gmt = gp.get_library( + name=library, + organism=organism.capitalize(), + min_size=min_size, + max_size=max_size, + ) # Cache for future use try: os.makedirs(os.path.dirname(cache_file), exist_ok=True) - with open(cache_file, 'w') as f: + with open(cache_file, "w") as f: json.dump(gmt, f) logging.info(f"Saved gene set cache to {cache_file}") except Exception as e: @@ -145,10 +180,10 @@ def perform_prerank( n_jobs: int = 1, low_cutoff: float = -np.inf, n_top: int = None, - **kwargs + **kwargs, ) -> pd.DataFrame: """Run GSEA prerank on each gene program in the loadings matrix. - + Parameters ---------- loadings : pd.DataFrame @@ -161,7 +196,7 @@ def perform_prerank( Remove features with loadings at or below this value. n_top : int Take the top n features with the highest loadings. - + Returns ------- pd.DataFrame @@ -181,8 +216,7 @@ def perform_prerank( # Run GSEA prerank for each column of loadings (each cell program) pre_res = pd.DataFrame() - for i in tqdm(loadings.columns, desc='Running GSEA', unit='programs'): - + for i in tqdm(loadings.columns, desc="Running GSEA", unit="programs"): # Filter out low loadings temp_loadings = loadings[i][(loadings[i] > low_cutoff)] @@ -190,27 +224,26 @@ def perform_prerank( if n_top is not None: temp_loadings = temp_loadings.sort_values(ascending=False).head(n_top) if len(temp_loadings) < n_top: - logging.warning(f"Program {i} has less than {n_top} features after filtering. Only {len(temp_loadings)} features will be used.") + logging.warning( + f"Program {i} has less than {n_top} features after filtering. Only {len(temp_loadings)} features will be used." + ) # Run GSEA prerank temp_res = gp.prerank( - rnk=temp_loadings, - gene_sets=geneset, - threads=n_jobs, - **kwargs + rnk=temp_loadings, gene_sets=geneset, threads=n_jobs, **kwargs ).res2d # Post-process results - temp_res['Gene %'] = temp_res['Gene %'].apply(lambda x: float(x[:-1])) - temp_res['tag_before'] = temp_res['Tag %'].apply(lambda x: int(x.split('/')[0])) - temp_res['tag_after'] = temp_res['Tag %'].apply(lambda x: int(x.split('/')[1])) - temp_res.drop(columns=['Tag %'], inplace=True) - if 'Name' in temp_res.columns and temp_res['Name'][0] == "prerank": - temp_res['Name'] = i - temp_res.rename(columns={'Name': 'program_name'}, inplace=True) - temp_res = temp_res.sort_values(['program_name', 'FDR q-val']) + temp_res["Gene %"] = temp_res["Gene %"].apply(lambda x: float(x[:-1])) + temp_res["tag_before"] = temp_res["Tag %"].apply(lambda x: int(x.split("/")[0])) + temp_res["tag_after"] = temp_res["Tag %"].apply(lambda x: int(x.split("/")[1])) + temp_res.drop(columns=["Tag %"], inplace=True) + if "Name" in temp_res.columns and temp_res["Name"][0] == "prerank": + temp_res["Name"] = i + temp_res.rename(columns={"Name": "program_name"}, inplace=True) + temp_res = temp_res.sort_values(["program_name", "FDR q-val"]) pre_res = pd.concat([pre_res, temp_res], ignore_index=True) - + return pre_res @@ -220,11 +253,7 @@ def perform_ssGSEA(): def perform_fisher_enrich( - loadings, - geneset, - n_top=500, - use_loadings_gene = False, - **kwargs + loadings, geneset, n_top=500, use_loadings_gene=False, **kwargs ): """Run Fisher enrichment on each gene program in the loadings matrix. @@ -236,7 +265,7 @@ def perform_fisher_enrich( Dictionary of gene sets. n_top : int Number of top features to take. - + Returns ------- ['Gene_set', 'Term', 'P-value', 'Adjusted P-value', 'Odds Ratio', @@ -261,41 +290,44 @@ def perform_fisher_enrich( # Find the intersection of genes present in the mudata object and in the library background_genes = set(value for sublist in geneset.values() for value in sublist) - if use_loadings_gene: # if True, only use expressed genes as background genes + if use_loadings_gene: # if True, only use expressed genes as background genes expressed_gene_list = set(loadings.index) background_genes = background_genes.intersection(expressed_gene_list) - enr_res = pd.DataFrame() - for i in tqdm(loadings.columns, desc='Running Fisher enrichment', unit='programs'): + for i in tqdm(loadings.columns, desc="Running Fisher enrichment", unit="programs"): gene_list = list(loadings[i].sort_values(ascending=False).head(n_top).index) temp_res = gp.enrich( - gene_list=gene_list, - gene_sets=geneset, - background=background_genes + gene_list=gene_list, gene_sets=geneset, background=background_genes ).res2d temp_res["program_name"] = i enr_res = pd.concat([enr_res, temp_res], ignore_index=True) - enr_res['overlap_numerator'] = enr_res['Overlap'].apply(lambda x: int(x.split('/')[0])) - enr_res['overlap_denominator'] = enr_res['Overlap'].apply(lambda x: int(x.split('/')[1])) - enr_res.drop(columns=['Overlap', 'Gene_set'], inplace=True) - enr_res = enr_res[["program_name"] + [col for col in enr_res.columns if col != "program_name"]] - enr_res = enr_res.sort_values(['program_name', 'Adjusted P-value']) - + enr_res["overlap_numerator"] = enr_res["Overlap"].apply( + lambda x: int(x.split("/")[0]) + ) + enr_res["overlap_denominator"] = enr_res["Overlap"].apply( + lambda x: int(x.split("/")[1]) + ) + enr_res.drop(columns=["Overlap", "Gene_set"], inplace=True) + enr_res = enr_res[ + ["program_name"] + [col for col in enr_res.columns if col != "program_name"] + ] + enr_res = enr_res.sort_values(["program_name", "Adjusted P-value"]) + return enr_res def insert_enrichment( mdata: mudata.MuData, df: pd.DataFrame, - library="GSEA", + library="GSEA", prog_key="prog", - geneset_index="Term", + geneset_index="Term", program_index="program_name", - varmap_name_prefix="gsea_varmap" + varmap_name_prefix="gsea_varmap", ) -> None: """Insert geneset enrichment into mudata - + Parameters ---------- mdata : mudata.MuData @@ -312,40 +344,41 @@ def insert_enrichment( index for the program in the DataFrame. varmap_name_prefix : str prefix for the varmap name. - + Returns ------- None """ - + # Create a mudata key to column name mapping dictionary mudata_keys_dict = {} for col in df.columns: if col not in [geneset_index, program_index]: key = f"{col}_{library}" - key = key.replace(' ', '_').replace('%', 'percent') + key = key.replace(" ", "_").replace("%", "percent") mudata_keys_dict[key] = col # Insert the values from the dataframe into the array for each key for key, colname in mudata_keys_dict.items(): # Create an empty dataframe with the right dimensions - all_progs_df = pd.DataFrame(index=df[geneset_index].unique(), - columns=mdata[prog_key].var.index) - + all_progs_df = pd.DataFrame( + index=df[geneset_index].unique(), columns=mdata[prog_key].var.index + ) + # Pivot the dataframe for gene sets and programs - pivot_df = df[[geneset_index, program_index, colname]].pivot(index=geneset_index, - columns=program_index, - values=colname) - + pivot_df = df[[geneset_index, program_index, colname]].pivot( + index=geneset_index, columns=program_index, values=colname + ) + # Update the empty dataframe with new values all_progs_df[pivot_df.columns] = pivot_df - + # Convert dataframe to a numpy array all_progs_array = all_progs_df.T.to_numpy() - + # Add the array into the MuData object mdata[prog_key].varm[key] = all_progs_array - + # Add the varmap to the mudata object varmap_name = f"{varmap_name_prefix}_{library}" mdata[prog_key].uns[varmap_name] = all_progs_df.index @@ -353,13 +386,13 @@ def insert_enrichment( def compute_geneset_enrichment( mdata: Union[str, mudata.MuData], - prog_key: str = 'prog', - data_key: str = 'rna', + prog_key: str = "prog", + data_key: str = "rna", prog_name: Optional[str] = None, - method: Literal['gsea', 'fisher'] = 'gsea', - organism: Literal['human', 'mouse'] = 'human', - library: str = 'Reactome_2022', - database: Literal['msigdb', 'enrichr'] = 'enrichr', + method: Literal["gsea", "fisher"] = "gsea", + organism: Literal["human", "mouse"] = "human", + library: str = "Reactome_2022", + database: Literal["msigdb", "enrichr"] = "enrichr", user_geneset: Optional[Dict[str, List[str]]] = None, min_size: int = 0, max_size: int = 2000, @@ -367,12 +400,11 @@ def compute_geneset_enrichment( n_top: int = 2000, n_jobs: int = 1, inplace: bool = True, - use_loadings_gene = False, # False = use all oncology gene, True = use interact expresseed gene + use_loadings_gene=False, # False = use all oncology gene, True = use interact expresseed gene gene_names_key: Optional[str] = None, use_cache: bool = False, - **kwargs + **kwargs, ) -> Optional[pd.DataFrame]: - """ Wrapper function to compute gene set enrichment for each program in the MuData object. @@ -410,7 +442,7 @@ def compute_geneset_enrichment( number of threads to run processes on. inplace : bool (default: True) whether to insert the results back into the mudata object. - + Returns ------- if not inplace: @@ -420,16 +452,17 @@ def compute_geneset_enrichment( """ # Read in mudata if it is provided as a path - frompath=False + frompath = False if isinstance(mdata, str): if os.path.exists(mdata): mdata = mudata.read(mdata) if inplace: - logging.warning('Changed to inplace=False since path was provided') - inplace=False - frompath=True - else: raise ValueError('Incorrect mudata specification.') - + logging.warning("Changed to inplace=False since path was provided") + inplace = False + frompath = True + else: + raise ValueError("Incorrect mudata specification.") + # get the geneset if user_geneset is not None: geneset = user_geneset @@ -440,9 +473,9 @@ def compute_geneset_enrichment( database=database, max_size=max_size, min_size=min_size, - use_cache=use_cache + use_cache=use_cache, ) - + # get the gene loadings for each program loadings = get_program_gene_loadings( mdata, @@ -452,16 +485,16 @@ def compute_geneset_enrichment( organism=organism, gene_names_key=gene_names_key, ) - + # run enrichment if method == "gsea": pre_res = perform_prerank( - loadings=loadings, + loadings=loadings, geneset=geneset, n_jobs=n_jobs, low_cutoff=low_cutoff, n_top=n_top, - **kwargs + **kwargs, ) elif method == "fisher": pre_res = perform_fisher_enrich( @@ -469,19 +502,19 @@ def compute_geneset_enrichment( geneset=geneset, n_top=n_top, use_loadings_gene=use_loadings_gene, - **kwargs + **kwargs, ) - + # insert results back into the mudata object if inplace if inplace: - mdata=insert_enrichment( - mdata, df=pre_res, - library=library, + mdata = insert_enrichment( + mdata, + df=pre_res, + library=library, prog_key=prog_key, - geneset_index="Term", + geneset_index="Term", program_index="program_name", - varmap_name_prefix="gsea_varmap" + varmap_name_prefix="gsea_varmap", ) else: - return(pre_res) - \ No newline at end of file + return pre_res diff --git a/src/Stage2_Evaluation/A_Metrics/src/enrichment_motif.py b/src/Stage2_Evaluation/A_Metrics/src/enrichment_motif.py index f6d28f3..0b58abb 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/enrichment_motif.py +++ b/src/Stage2_Evaluation/A_Metrics/src/enrichment_motif.py @@ -17,9 +17,7 @@ from typing import List, Dict, Tuple, Union, Optional, Literal, Mapping -def read_loci( - path_loci: os.PathLike -) -> pd.DataFrame: +def read_loci(path_loci: os.PathLike) -> pd.DataFrame: """Read loci to gene links file Read promoter/enhancer to gene links tab delimited file with the following column headers: @@ -43,27 +41,31 @@ def read_loci( """ # Read formatted coords file - loci = pd.read_csv(path_loci, sep='\t') + loci = pd.read_csv(path_loci, sep="\t") # Enforce header format expected_headers = [ - 'chromosome', - 'start', - 'end', - 'seq_name', - 'seq_class', - 'seq_score', - 'gene_name' + "chromosome", + "start", + "end", + "seq_name", + "seq_class", + "seq_score", + "gene_name", ] # Check for expected headers for col in expected_headers: - try: assert col in loci.columns - except: raise ValueError('Coordinate file is not formatted correctly') + try: + assert col in loci.columns + except: + raise ValueError("Coordinate file is not formatted correctly") # TODO: Support multiple classes at once - try: assert len(loci['seq_class']==1) - except: raise ValueError('Coordinate file contains multiple sequence classes') + try: + assert len(loci["seq_class"] == 1) + except: + raise ValueError("Coordinate file contains multiple sequence classes") return loci @@ -72,56 +74,81 @@ def perform_motif_match( loci: pd.DataFrame, sequences: os.PathLike, pwms: Mapping[str, np.ndarray], - in_window: int=1000, - threshold: float=1e-4, - eps: float=1e-4, - reverse_complement: bool=True, - output_loc: os.PathLike=None + in_window: int = 1000, + threshold: float = 1e-4, + eps: float = 1e-4, + reverse_complement: bool = True, + output_loc: os.PathLike = None, ): - """Score motif matches - - Perform motif matching on sequences linked to genes - via enhancer/promoter coordinates. - - Parameters - ---------- - loci : pd.DataFrame - DataFrame containing sequence coordinates. - sequences : os.PathLike - Path to FASTA formatted genomic sequence. - pwms : Mapping[str, np.ndarray] - Dictionary of PWMs where keys are motif names and values are PWMs. - in_window : int - Window size to extract sequences around center of loci. - threshold : float - Threshold for motif matching. - output_loc : os.PathLike - Path to directory to store motif matches for individual motifs. - """ - # Perform motif matching - X = extract_loci(loci, sequences, in_window=in_window).float() - hits = fimo(pwms, X, threshold=threshold, eps=eps, reverse_complement=reverse_complement) - - # Create motif match dataframe - motif_match_df = pd.DataFrame() - for i, hit in enumerate(hits): - annotated_hit = hit.merge(loci[["chromosome", "seq_name", "seq_class", "gene_name"]], left_on="sequence_name", right_index=True).drop(columns=["sequence_name"]) - annotated_hit["motif_name"] = list(pwms.keys())[i] - annotated_hit = annotated_hit[["chromosome", "start", "end", "strand", "motif_name", "score", "p-value", "seq_name", "seq_class", "gene_name"]] - if output_loc is not None: - annotated_hit.to_csv(os.path.join(output_loc, f"motif_match_{list(pwms.keys())[i]}.txt"), sep='\t', index=False) - motif_match_df = pd.concat([motif_match_df, annotated_hit]) - motif_match_df["adj_pval"] = multipletests(motif_match_df["p-value"], method="fdr_bh")[1] - motif_match_df.reset_index(drop=True, inplace=True) - return motif_match_df + """Score motif matches + + Perform motif matching on sequences linked to genes + via enhancer/promoter coordinates. + + Parameters + ---------- + loci : pd.DataFrame + DataFrame containing sequence coordinates. + sequences : os.PathLike + Path to FASTA formatted genomic sequence. + pwms : Mapping[str, np.ndarray] + Dictionary of PWMs where keys are motif names and values are PWMs. + in_window : int + Window size to extract sequences around center of loci. + threshold : float + Threshold for motif matching. + output_loc : os.PathLike + Path to directory to store motif matches for individual motifs. + """ + # Perform motif matching + X = extract_loci(loci, sequences, in_window=in_window).float() + hits = fimo( + pwms, X, threshold=threshold, eps=eps, reverse_complement=reverse_complement + ) + + # Create motif match dataframe + motif_match_df = pd.DataFrame() + for i, hit in enumerate(hits): + annotated_hit = hit.merge( + loci[["chromosome", "seq_name", "seq_class", "gene_name"]], + left_on="sequence_name", + right_index=True, + ).drop(columns=["sequence_name"]) + annotated_hit["motif_name"] = list(pwms.keys())[i] + annotated_hit = annotated_hit[ + [ + "chromosome", + "start", + "end", + "strand", + "motif_name", + "score", + "p-value", + "seq_name", + "seq_class", + "gene_name", + ] + ] + if output_loc is not None: + annotated_hit.to_csv( + os.path.join(output_loc, f"motif_match_{list(pwms.keys())[i]}.txt"), + sep="\t", + index=False, + ) + motif_match_df = pd.concat([motif_match_df, annotated_hit]) + motif_match_df["adj_pval"] = multipletests( + motif_match_df["p-value"], method="fdr_bh" + )[1] + motif_match_df.reset_index(drop=True, inplace=True) + return motif_match_df def compute_motif_instances( motif_match_df: pd.DataFrame, - motif_var: str = 'motif_name', - sig: float=0.05, - sig_var: str='adj_pval', - gene_names: Optional[np.ndarray]=None + motif_var: str = "motif_name", + sig: float = 0.05, + sig_var: str = "adj_pval", + gene_names: Optional[np.ndarray] = None, ): """Count motif instances per gene (via enahncer/promoter linking) @@ -137,7 +164,7 @@ def compute_motif_instances( Column name for significance values. Default is 'adj_pval'. gene_names : np.ndarray Array of gene names. Default is None. - + Returns ------- motif_count_df : pd.DataFrame @@ -147,11 +174,17 @@ def compute_motif_instances( # Count up significant occurences of motif motif_match_df_ = motif_match_df.loc[motif_match_df[sig_var] < sig] print(f"There are {motif_match_df_.shape[0]} significant motif matches.") - motif_match_df_ = motif_match_df.value_counts(subset=['gene_name', motif_var]).reset_index() - motif_match_df_.columns = ['gene_name', motif_var, 'motif_count'] - motif_match_df_ = motif_match_df_.pivot(index='gene_name', columns=motif_var, values='motif_count') + motif_match_df_ = motif_match_df.value_counts( + subset=["gene_name", motif_var] + ).reset_index() + motif_match_df_.columns = ["gene_name", motif_var, "motif_count"] + motif_match_df_ = motif_match_df_.pivot( + index="gene_name", columns=motif_var, values="motif_count" + ) motif_count_df = pd.DataFrame(index=gene_names, columns=motif_match_df_.columns) - motif_count_df.loc[motif_match_df_.index.values] = motif_match_df_ # Gene names should match as this point + motif_count_df.loc[motif_match_df_.index.values] = ( + motif_match_df_ # Gene names should match as this point + ) return motif_count_df @@ -162,10 +195,10 @@ def perform_correlation( motif_enrich_pval_df: pd.DataFrame, motif_idx: int, prog_idx: int, - correlation: str='pearsonr' + correlation: str = "pearsonr", ): """Compute motif enrichment as correlation b/w gene weights/ranks and motif counts - + Perform pearson correlation test for motif count enrichment vs gene loadings If loadings are dichotomized then this is equivalent to a point biserial correlation test. @@ -189,34 +222,34 @@ def perform_correlation( loadings = prog_genes.iloc[prog_idx].values.flatten() counts = motif_count_df.T.iloc[motif_idx].fillna(0).values.flatten() - - if correlation=='pearsonr': + + if correlation == "pearsonr": stat, pval = pearsonr(loadings, counts) - elif correlation=='spearmanr': + elif correlation == "spearmanr": stat, pval = spearmanr(loadings, counts) - elif correlation=='kendalltau': + elif correlation == "kendalltau": stat, pval = kendalltau(loadings, counts) - motif_enrich_stat_df.iloc[prog_idx, motif_idx] = stat - motif_enrich_pval_df.iloc[prog_idx, motif_idx] = pval + motif_enrich_stat_df.iloc[prog_idx, motif_idx] = stat + motif_enrich_pval_df.iloc[prog_idx, motif_idx] = pval def compute_motif_enrichment_( mdata: mudata.MuData, motif_count_df: pd.DataFrame, - prog_key: str='prog', - gene_names: Optional[np.ndarray]=None, - weighted: bool=True, - num_genes: Optional[int]=None, - correlation: str='pearsonr', - n_jobs: int=1 + prog_key: str = "prog", + gene_names: Optional[np.ndarray] = None, + weighted: bool = True, + num_genes: Optional[int] = None, + correlation: str = "pearsonr", + n_jobs: int = 1, ): """Count up motif ocurrences and perform diff. test Perform motif enrichment using gene program loadings and motif counts linked to genes via motif scanning of linked enhancer/promoter sequences. - + Parameters ---------- mdata : MuData @@ -239,74 +272,82 @@ def compute_motif_enrichment_( # Both weighted and num_genes cannot be set if num_genes is not None and weighted: - raise ValueError('Will not use weighted when num_genes specified.') + raise ValueError("Will not use weighted when num_genes specified.") - loadings = pd.DataFrame(data=mdata[prog_key].varm['loadings'], - index=mdata[prog_key].var_names, - columns=gene_names) + loadings = pd.DataFrame( + data=mdata[prog_key].varm["loadings"], + index=mdata[prog_key].var_names, + columns=gene_names, + ) # FIXME: Causes expansion due to duplications in gene_names # Ensure index matches b/w loadings and counts # loadings = loadings.loc[:, motif_count_df.index.values] - # Binary matrix + # Binary matrix if not weighted: - prog_genes = (loadings.rank(axis=1)1 then parallel processes dont seem to terminate. # Perform test in parallel across motifs and programs - Parallel(n_jobs=n_jobs, - backend='threading')(delayed(perform_correlation)(motif_count_df, - prog_genes, - motif_enrich_stat_df, - motif_enrich_pval_df, - motif_idx, - prog_idx, - correlation=correlation) \ - for motif_idx in tqdm(range(motif_count_df.columns.values.shape[-1]), - desc='Computing motif enrichment', - unit='motifs') \ - for prog_idx in range(mdata[prog_key].var_names.shape[0])) + Parallel(n_jobs=n_jobs, backend="threading")( + delayed(perform_correlation)( + motif_count_df, + prog_genes, + motif_enrich_stat_df, + motif_enrich_pval_df, + motif_idx, + prog_idx, + correlation=correlation, + ) + for motif_idx in tqdm( + range(motif_count_df.columns.values.shape[-1]), + desc="Computing motif enrichment", + unit="motifs", + ) + for prog_idx in range(mdata[prog_key].var_names.shape[0]) + ) return motif_enrich_stat_df, motif_enrich_pval_df def compute_motif_enrichment( mdata: Union[mudata.MuData, os.PathLike], - prog_key: str='prog', - data_key: str='data', - motif_file: Optional[os.PathLike]=None, - seq_file: Optional[os.PathLike]=None, - loci_file: Optional[os.PathLike]=None, - output_loc: Optional[os.PathLike]=None, - window: int=1000, - threshold: float=1e-4, - eps: float=1e-4, - reverse_complement: bool=True, - sig: float=0.05, - num_genes: Optional[int]=None, - correlation: str='pearsonr', - n_jobs: int=1, - inplace: bool=True, - gene_names_key: Optional[str]=None, - **kwargs + prog_key: str = "prog", + data_key: str = "data", + motif_file: Optional[os.PathLike] = None, + seq_file: Optional[os.PathLike] = None, + loci_file: Optional[os.PathLike] = None, + output_loc: Optional[os.PathLike] = None, + window: int = 1000, + threshold: float = 1e-4, + eps: float = 1e-4, + reverse_complement: bool = True, + sig: float = 0.05, + num_genes: Optional[int] = None, + correlation: str = "pearsonr", + n_jobs: int = 1, + inplace: bool = True, + gene_names_key: Optional[str] = None, + **kwargs, ): - """Compute motif enrichment in enhancers or promoters associated with a gene - + Perform motif enrichment using gene program loadings and motif counts linked to genes via motif scanning of linked enhancer/promoter sequences. @@ -315,7 +356,7 @@ def compute_motif_enrichment( ---------- mdata : MuData mudata object containing anndata of program scores and cell-level metadata. - prog_key: + prog_key: index for the anndata object (mdata[prog_key]) in the mudata object. data_key: str index of the genomic data anndata object (mdata[data_key]) in the mudata object. @@ -326,7 +367,7 @@ def compute_motif_enrichment( loci_file: str path to enhancer/promoter gene links file with sequence coordinates. Tab delimited file with the following column headers: - chr, start, end, seq_name, seq_class {promoter, enhancer}, + chr, start, end, seq_name, seq_class {promoter, enhancer}, seq_score, gene_name. output_loc: str path to directory to store motif - gene counts. @@ -343,7 +384,7 @@ def compute_motif_enrichment( number of threads to run processes on. inplace: Bool (default: True) update the mudata object inplace or return a copy - + Returns ------- if not inplace: @@ -355,65 +396,78 @@ def compute_motif_enrichment( """ # Read in mudata if it is provided as a path - frompath=False + frompath = False if isinstance(mdata, str): if os.path.exists(mdata): mdata = mudata.read(mdata) if inplace: - logging.warning('Changed to inplace=False since path was provided') - inplace=False - frompath=True - else: raise ValueError('Incorrect mudata specification.') - + logging.warning("Changed to inplace=False since path was provided") + inplace = False + frompath = True + else: + raise ValueError("Incorrect mudata specification.") + if not inplace and not frompath: - mdata = mudata.MuData({prog_key: mdata[prog_key].copy(), - data_key: mdata[data_key].copy()}) + mdata = mudata.MuData( + {prog_key: mdata[prog_key].copy(), data_key: mdata[data_key].copy()} + ) # Get gene names in MuData if gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: gene_names = mdata[data_key].var[gene_names_key].astype(str).tolist() - elif 'var_names' in mdata[prog_key].uns.keys(): - gene_names = mdata[prog_key].uns['var_names'] + elif "var_names" in mdata[prog_key].uns.keys(): + gene_names = mdata[prog_key].uns["var_names"] else: - try: assert mdata[prog_key].varm['loadings'].shape[1]==mdata[data_key].var.shape[0] - except: raise ValueError('Different number of genes present in data and program loadings') + try: + assert ( + mdata[prog_key].varm["loadings"].shape[1] + == mdata[data_key].var.shape[0] + ) + except: + raise ValueError( + "Different number of genes present in data and program loadings" + ) gene_names = mdata[data_key].var_names - if ':ens' in gene_names[0].lower(): - gene_names = [name.split(':')[0] for name in gene_names] + if ":ens" in gene_names[0].lower(): + gene_names = [name.split(":")[0] for name in gene_names] # Check if output loc exists if output_loc is not None: - try: os.makedirs(output_loc, exist_ok=True) - except: raise ValueError('Output location does not exist.') + try: + os.makedirs(output_loc, exist_ok=True) + except: + raise ValueError("Output location does not exist.") # If num_genes specified then cannot be weighted if num_genes is None: - weighted=True + weighted = True elif isinstance(num_genes, int): - weighted=False + weighted = False # Intake motif file path or in memory if isinstance(motif_file, str) and os.path.exists(motif_file): pwms = read_meme(motif_file) else: - raise ValueError('Motif file not found.') - + raise ValueError("Motif file not found.") + # Intake coord file path or in memory if isinstance(loci_file, str) and os.path.exists(loci_file): loci = read_loci(loci_file) else: - raise ValueError('Coordinate file not found.') + raise ValueError("Coordinate file not found.") # Valid genes - matching_gene_names = np.intersect1d(loci['gene_name'].unique(), gene_names) - print(f'Number of matching genes: {len(matching_gene_names)}') - try: assert len(matching_gene_names) > 0 - except: raise ValueError('No matching genes b/w data and coordinate files') + matching_gene_names = np.intersect1d(loci["gene_name"].unique(), gene_names) + print(f"Number of matching genes: {len(matching_gene_names)}") + try: + assert len(matching_gene_names) > 0 + except: + raise ValueError("No matching genes b/w data and coordinate files") # Compute motif matching - loci_ = loci[loci['gene_name'].isin(matching_gene_names)] - print(f'Number of loci: {loci_.shape[0]}') + loci_ = loci[loci["gene_name"].isin(matching_gene_names)] + print(f"Number of loci: {loci_.shape[0]}") motif_match_df = perform_motif_match( loci=loci_, sequences=seq_file, @@ -422,16 +476,16 @@ def compute_motif_enrichment( threshold=threshold, eps=eps, reverse_complement=reverse_complement, - output_loc=output_loc + output_loc=output_loc, ) # Count motif enrichment motif_count_df = compute_motif_instances( motif_match_df, - motif_var='motif_name', + motif_var="motif_name", sig=sig, - sig_var='adj_pval', - gene_names=gene_names + sig_var="adj_pval", + gene_names=gene_names, ) motif_enrich_stat_df, motif_enrich_pval_df = compute_motif_enrichment_( mdata, @@ -440,74 +494,85 @@ def compute_motif_enrichment( gene_names=gene_names, weighted=weighted, num_genes=num_genes, - n_jobs=n_jobs) - + n_jobs=n_jobs, + ) + # Store motif counts if inplace: - mdata[prog_key].uns['motif_counts'] = motif_count_df.loc[gene_names].values - mdata[prog_key].uns['motif_names'] = motif_count_df.columns.values + mdata[prog_key].uns["motif_counts"] = motif_count_df.loc[gene_names].values + mdata[prog_key].uns["motif_names"] = motif_count_df.columns.values - mdata[prog_key].varm['motif_enrich_{}_stat'.format(correlation)] = motif_enrich_stat_df.values - mdata[prog_key].varm['motif_enrich_{}_pval'.format(correlation)] = motif_enrich_pval_df.values - mdata[prog_key].uns['motif_names'] = motif_count_df.columns.values + mdata[prog_key].varm["motif_enrich_{}_stat".format(correlation)] = ( + motif_enrich_stat_df.values + ) + mdata[prog_key].varm["motif_enrich_{}_pval".format(correlation)] = ( + motif_enrich_pval_df.values + ) + mdata[prog_key].uns["motif_names"] = motif_count_df.columns.values else: - - motif_enrich_stat_df = motif_enrich_stat_df.reset_index().melt(id_vars='index', - var_name='motif', - value_name='stat') - motif_enrich_stat_df = motif_enrich_stat_df.set_index(['index', 'motif']) - - motif_enrich_pval_df = motif_enrich_pval_df.reset_index().melt(id_vars='index', - var_name='motif', - value_name='pval') - motif_enrich_pval_df = motif_enrich_pval_df.set_index(['index', 'motif']) - - motif_enrichment_df = motif_enrich_stat_df.merge(motif_enrich_pval_df, - left_index=True, - right_index=True) + motif_enrich_stat_df = motif_enrich_stat_df.reset_index().melt( + id_vars="index", var_name="motif", value_name="stat" + ) + motif_enrich_stat_df = motif_enrich_stat_df.set_index(["index", "motif"]) + + motif_enrich_pval_df = motif_enrich_pval_df.reset_index().melt( + id_vars="index", var_name="motif", value_name="pval" + ) + motif_enrich_pval_df = motif_enrich_pval_df.set_index(["index", "motif"]) + + motif_enrichment_df = motif_enrich_stat_df.merge( + motif_enrich_pval_df, left_index=True, right_index=True + ) motif_enrichment_df = motif_enrichment_df.reset_index() - motif_enrichment_df['program_name'] = motif_enrichment_df['index'] - motif_enrichment_df.drop('index', axis=1, inplace=True) - motif_enrichment_df = motif_enrichment_df.sort_values(['program_name', 'pval']) - motif_enrichment_df["adj_pval"] = multipletests(motif_enrichment_df["pval"], method="fdr_bh")[1] + motif_enrichment_df["program_name"] = motif_enrichment_df["index"] + motif_enrichment_df.drop("index", axis=1, inplace=True) + motif_enrichment_df = motif_enrichment_df.sort_values(["program_name", "pval"]) + motif_enrichment_df["adj_pval"] = multipletests( + motif_enrichment_df["pval"], method="fdr_bh" + )[1] - motif_count_df.columns.name='' + motif_count_df.columns.name = "" - return (motif_match_df, - motif_count_df.loc[gene_names].fillna(0), - motif_enrichment_df) + return ( + motif_match_df, + motif_count_df.loc[gene_names].fillna(0), + motif_enrichment_df, + ) -if __name__=='__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('mudataObj_path') - parser.add_argument('-mf', '--motif_file', default=None, type=str, required=True) - parser.add_argument('-sf', '--seq_file', default=None, type=str, required=True) - parser.add_argument('-cf', '--coords_file', default=None, type=str, required=True) - parser.add_argument('--store_files', default=None, type=str) - parser.add_argument('--significance', default=0.05, choices=range(0,1), - metavar="(0,1)", type=float) - parser.add_argument('--num_genes', default=None, type=int) - parser.add_argument('--correlation', default='pearsonr', type=str) - parser.add_argument('-pk', '--prog_key', default='prog', type=str) - parser.add_argument('-dk', '--data_key', default='rna', type=str) - parser.add_argument('-n', '--n_jobs', default=1, type=int) - parser.add_argument('--output', action='store_false') + parser.add_argument("mudataObj_path") + parser.add_argument("-mf", "--motif_file", default=None, type=str, required=True) + parser.add_argument("-sf", "--seq_file", default=None, type=str, required=True) + parser.add_argument("-cf", "--coords_file", default=None, type=str, required=True) + parser.add_argument("--store_files", default=None, type=str) + parser.add_argument( + "--significance", default=0.05, choices=range(0, 1), metavar="(0,1)", type=float + ) + parser.add_argument("--num_genes", default=None, type=int) + parser.add_argument("--correlation", default="pearsonr", type=str) + parser.add_argument("-pk", "--prog_key", default="prog", type=str) + parser.add_argument("-dk", "--data_key", default="rna", type=str) + parser.add_argument("-n", "--n_jobs", default=1, type=int) + parser.add_argument("--output", action="store_false") args = parser.parse_args() mdata = mudata.read(args.mudataObj_path) - compute_motif_enrichment(mdata, prog_key=args.prog_key, data_key=args.data_key, - motif_file=args.motif_file, seq_file=args.seq_file, - coords_file=args.coords_file, output_loc=args.store_files, - sig=args.significance, num_genes=args.num_genes, - n_jobs=args.n_jobs, correlation=args.correlation, - inplace=args.output) - - - - - - + compute_motif_enrichment( + mdata, + prog_key=args.prog_key, + data_key=args.data_key, + motif_file=args.motif_file, + seq_file=args.seq_file, + coords_file=args.coords_file, + output_loc=args.store_files, + sig=args.significance, + num_genes=args.num_genes, + n_jobs=args.n_jobs, + correlation=args.correlation, + inplace=args.output, + ) diff --git a/src/Stage2_Evaluation/A_Metrics/src/enrichment_trait.py b/src/Stage2_Evaluation/A_Metrics/src/enrichment_trait.py index 6e48f35..43e4cbf 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/enrichment_trait.py +++ b/src/Stage2_Evaluation/A_Metrics/src/enrichment_trait.py @@ -8,11 +8,18 @@ from .enrichment_geneset import compute_geneset_enrichment, create_geneset_dict import logging -logging.basicConfig(level = logging.INFO) +logging.basicConfig(level=logging.INFO) -def run_opentargets_query(credentials_path, output_file, min_assoc_loci=10, - min_n_cases=1000, min_l2g_score=0.2, study_ids_to_keep=None): + +def run_opentargets_query( + credentials_path, + output_file, + min_assoc_loci=10, + min_n_cases=1000, + min_l2g_score=0.2, + study_ids_to_keep=None, +): """Query Open Targets using BigQuery to obtain GWAS to Gene mappings across many traits""" from google.cloud import bigquery @@ -24,7 +31,7 @@ def run_opentargets_query(credentials_path, output_file, min_assoc_loci=10, client = bigquery.Client.from_service_account_json(credentials_path) # Construct parameterized query - query = ''' + query = """ WITH ranked_genes AS ( SELECT locus2gene.study_id, locus2gene.chrom, locus2gene.pos, locus2gene.ref, locus2gene.alt, @@ -48,10 +55,10 @@ def run_opentargets_query(credentials_path, output_file, min_assoc_loci=10, ON locus2gene.pos = lead_variants.lead_pos AND locus2gene.chrom = lead_variants.lead_chrom AND locus2gene.study_id = lead_variants.study_id - ''' + """ # Add filter conditions - query += f''' + query += f""" WHERE -- Remove the "raw" Neale lab results -- I'm not sure what this is locus2gene.study_id NOT LIKE '%raw%' @@ -63,31 +70,39 @@ def run_opentargets_query(credentials_path, output_file, min_assoc_loci=10, AND study_metadata.num_assoc_loci >= {min_assoc_loci} -- Filter to n_cases of at least {min_n_cases} - ''' + """ # Optionally filter by study_ids_to_keep if provided - if study_ids_to_keep and isinstance(study_ids_to_keep, list) and any(isinstance(x, str) for x in study_ids_to_keep): - query += f''' + if ( + study_ids_to_keep + and isinstance(study_ids_to_keep, list) + and any(isinstance(x, str) for x in study_ids_to_keep) + ): + query += f""" AND locus2gene.study_id IN UNNEST(@study_ids_to_keep) - ''' + """ - query += ''' + query += """ ) SELECT * FROM ranked_genes WHERE rn = 1; - ''' + """ # Set query parameters job_config = bigquery.QueryJobConfig() if study_ids_to_keep: - job_config.query_parameters = [bigquery.ArrayQueryParameter("study_ids_to_keep", "STRING", study_ids_to_keep)] + job_config.query_parameters = [ + bigquery.ArrayQueryParameter( + "study_ids_to_keep", "STRING", study_ids_to_keep + ) + ] # Run the query query_job = client.query(query, job_config=job_config) # Convert the query results to a Pandas dataframe l2g = query_job.to_dataframe() - l2g.sort_values(by=['study_id', 'pval']) + l2g.sort_values(by=["study_id", "pval"]) # Save dataframe to output file l2g.to_csv(output_file, compression="gzip", index=False) @@ -99,7 +114,7 @@ def run_opentargets_query(credentials_path, output_file, min_assoc_loci=10, def process_json_format_l2g_columns(row, column_name): """Filter Open Targets query results to a set of high-quality GWAS and L2G scores - + Extracts a comma-separated list of IDs from a string representation of a list of dictionaries in a DataFrame row. Parameters: @@ -113,17 +128,23 @@ def process_json_format_l2g_columns(row, column_name): try: elements_str = row[column_name] start_index = elements_str.find("[{") # Find the starting index of the list - end_index = elements_str.find("}]") + 2 # Find the ending index of the list and include the closing bracket - - if start_index != -1 and end_index != -1: # Ensure both start and end indices are found + end_index = ( + elements_str.find("}]") + 2 + ) # Find the ending index of the list and include the closing bracket + + if ( + start_index != -1 and end_index != -1 + ): # Ensure both start and end indices are found elements_list_str = elements_str[start_index:end_index] - + # Remove newline characters and convert to a list of dictionaries - elements_list = eval(elements_list_str.replace('\n', '').replace('array', 'list')) + elements_list = eval( + elements_list_str.replace("\n", "").replace("array", "list") + ) - ids = [elem['element'] for elem in elements_list] + ids = [elem["element"] for elem in elements_list] ids.sort() - return ', '.join(ids) + return ", ".join(ids) else: return None # Return None if start or end index not found except Exception as e: @@ -131,17 +152,23 @@ def process_json_format_l2g_columns(row, column_name): return None -def filter_open_targets_gwas_query(input_file, output_file, min_l2g_score=None, remove_mhc_region=True): - l2g = pd.read_csv(input_file, compression='gzip', low_memory=False) +def filter_open_targets_gwas_query( + input_file, output_file, min_l2g_score=None, remove_mhc_region=True +): + l2g = pd.read_csv(input_file, compression="gzip", low_memory=False) # Convert the trait_efos into just a flat list of EFO IDs rather than json formatted - l2g['trait_efos'] = l2g.apply(lambda row: process_json_format_l2g_columns(row, "trait_efos"), axis=1) + l2g["trait_efos"] = l2g.apply( + lambda row: process_json_format_l2g_columns(row, "trait_efos"), axis=1 + ) # Remove GWAS with lots of EFO IDs - filtered_l2g = l2g[l2g['trait_efos'].fillna('').apply(lambda x: x.count(',')) <= 2] + filtered_l2g = l2g[l2g["trait_efos"].fillna("").apply(lambda x: x.count(",")) <= 2] # Remove double quotes present in some trait_reported rows, but not others - filtered_l2g.loc[:, 'trait_reported'] = filtered_l2g['trait_reported'].str.replace('"', '') + filtered_l2g.loc[:, "trait_reported"] = filtered_l2g["trait_reported"].str.replace( + '"', "" + ) # Remove unusual traits filtered_l2g = filtered_l2g.query( @@ -156,42 +183,68 @@ def filter_open_targets_gwas_query(input_file, output_file, min_l2g_score=None, ) # Retain the GWAS with the largest sample size by number of cases by EFO group - filtered_l2g.loc[:, 'n_cases'] = filtered_l2g['n_cases'].fillna(filtered_l2g['n_initial']) - filtered_l2g = (filtered_l2g.assign(rank=filtered_l2g.groupby(["trait_efos"])["n_cases"].rank(method="min", ascending=False)) - .query("rank == 1") - .drop(columns=["rank"]) - .reset_index(drop=True)) + filtered_l2g.loc[:, "n_cases"] = filtered_l2g["n_cases"].fillna( + filtered_l2g["n_initial"] + ) + filtered_l2g = ( + filtered_l2g.assign( + rank=filtered_l2g.groupby(["trait_efos"])["n_cases"].rank( + method="min", ascending=False + ) + ) + .query("rank == 1") + .drop(columns=["rank"]) + .reset_index(drop=True) + ) # Rename the y_proba_full_model to L2G - filtered_l2g.rename(columns={'y_proba_full_model': 'L2G'}, inplace=True) + filtered_l2g.rename(columns={"y_proba_full_model": "L2G"}, inplace=True) # Filter to positions outside the extended MHC region, which can have poor L2G mapping if remove_mhc_region: mhc_start = 28510120 mhc_end = 33480577 - filtered_l2g = filtered_l2g[~((filtered_l2g['chrom'] == '6') & - (filtered_l2g['pos'] >= mhc_start) & - (filtered_l2g['pos'] <= mhc_end))] + filtered_l2g = filtered_l2g[ + ~( + (filtered_l2g["chrom"] == "6") + & (filtered_l2g["pos"] >= mhc_start) + & (filtered_l2g["pos"] <= mhc_end) + ) + ] # Filter rows based on min L2G score if min_l2g_score is not None: - filtered_l2g = filtered_l2g[filtered_l2g['L2G'] >= min_l2g_score] + filtered_l2g = filtered_l2g[filtered_l2g["L2G"] >= min_l2g_score] # Filter down the columns to make this DataFrame a bit easier to work with - desired_columns = ['trait_category', 'trait_efos', 'trait_reported', - 'gene_name', 'L2G', - 'chrom', 'pos', 'ref', 'alt', 'pval', - 'source', 'study_id', 'pmid', - 'pub_date', 'num_assoc_loci'] + desired_columns = [ + "trait_category", + "trait_efos", + "trait_reported", + "gene_name", + "L2G", + "chrom", + "pos", + "ref", + "alt", + "pval", + "source", + "study_id", + "pmid", + "pub_date", + "num_assoc_loci", + ] filtered_l2g = filtered_l2g[desired_columns] # Sort the DataFrame - filtered_l2g = filtered_l2g.sort_values(by=['trait_category', 'study_id', 'pval', 'L2G'], - ascending=[True, True, True, False]) + filtered_l2g = filtered_l2g.sort_values( + by=["trait_category", "study_id", "pval", "L2G"], + ascending=[True, True, True, False], + ) # Save the filtered DataFrame to a CSV file filtered_l2g.to_csv(output_file, index=False) - + def process_enrichment_data( enrich_res, @@ -201,7 +254,7 @@ def process_enrichment_data( metadata_geneset_id_col="trait_efos", color_category_col="trait_category", program_name_col="program_name", - annotation_cols=["trait_reported", "Genes", "study_id", "pmid"] + annotation_cols=["trait_reported", "Genes", "study_id", "pmid"], ): # Read in enrichment results @@ -210,31 +263,53 @@ def process_enrichment_data( elif isinstance(enrich_res, pd.DataFrame): enrich_df = enrich_res else: - raise ValueError("enrich_res must be either a pandas DataFrame or a file path to a CSV file.") + raise ValueError( + "enrich_res must be either a pandas DataFrame or a file path to a CSV file." + ) if isinstance(metadata, str): - metadata_df = pd.read_csv(metadata, compression='gzip', low_memory=False) + metadata_df = pd.read_csv(metadata, compression="gzip", low_memory=False) elif isinstance(metadata, pd.DataFrame): metadata_df = metadata else: - raise ValueError("metadata must be either a pandas DataFrame or a file path to a CSV file.") + raise ValueError( + "metadata must be either a pandas DataFrame or a file path to a CSV file." + ) # Join the enrichment results and the metadata - enrich_ps = enrich_df.merge(metadata_df, left_on=enrich_geneset_id_col, right_on=metadata_geneset_id_col, how="left") + enrich_ps = enrich_df.merge( + metadata_df, + left_on=enrich_geneset_id_col, + right_on=metadata_geneset_id_col, + how="left", + ) # Only keep the relevant columns - keep_cols = list([enrich_geneset_id_col, pval_col, metadata_geneset_id_col, color_category_col, program_name_col] + annotation_cols) + keep_cols = list( + [ + enrich_geneset_id_col, + pval_col, + metadata_geneset_id_col, + color_category_col, + program_name_col, + ] + + annotation_cols + ) enrich_ps = enrich_ps[keep_cols] # Sort by P-value - enrich_ps = enrich_ps.drop_duplicates().sort_values(by=[color_category_col, pval_col]) + enrich_ps = enrich_ps.drop_duplicates().sort_values( + by=[color_category_col, pval_col] + ) # If the input P-value == 0, then replace it with the lowest non-zero P-value in the dataframe min_value = enrich_ps.query(f"`{pval_col}` > 0")[pval_col].min() # Compute the -log(10) P-value and deal with edge-cases (e.g. P=0, P=1) - enrich_ps.loc[enrich_ps[pval_col] == 0, pval_col] = min_value # Replace P=0 with min non-0 p-value - enrich_ps[f'-log10({pval_col})'] = abs(-1 * np.log10(enrich_ps[pval_col])) + enrich_ps.loc[enrich_ps[pval_col] == 0, pval_col] = ( + min_value # Replace P=0 with min non-0 p-value + ) + enrich_ps[f"-log10({pval_col})"] = abs(-1 * np.log10(enrich_ps[pval_col])) enrich_ps.reset_index(drop=True, inplace=True) @@ -244,35 +319,37 @@ def process_enrichment_data( def compute_trait_enrichment( mdata, gwas_data, - prog_key='prog', + prog_key="prog", prog_name=None, - data_key='rna', - library='OT_GWAS', + data_key="rna", + library="OT_GWAS", n_jobs=1, inplace=False, - key_column='trait_efos', - gene_column='gene_name', - method='fisher', + key_column="trait_efos", + gene_column="gene_name", + method="fisher", min_size=0, max_size=2000, n_top=500, low_cutoff=-np.inf, use_loadings_gene=False, gene_names_key=None, - **kwargs + **kwargs, ): """Compute Trait enrichment using open-targets GWAS""" - #read in gwas data + # read in gwas data if isinstance(gwas_data, str): - df = pd.read_csv(gwas_data, compression='gzip', low_memory=False) + df = pd.read_csv(gwas_data, compression="gzip", low_memory=False) elif isinstance(gwas_data, pd.DataFrame): df = gwas_data else: - raise ValueError("gwas_data must be either a pandas DataFrame or a file path to a CSV file.") - + raise ValueError( + "gwas_data must be either a pandas DataFrame or a file path to a CSV file." + ) + gmt = create_geneset_dict(df, key_column=key_column, gene_column=gene_column) - return (compute_geneset_enrichment( + return compute_geneset_enrichment( mdata=mdata, prog_key=prog_key, data_key=data_key, @@ -290,26 +367,28 @@ def compute_trait_enrichment( inplace=inplace, use_loadings_gene=use_loadings_gene, gene_names_key=gene_names_key, - **kwargs - )) + **kwargs, + ) -if __name__=='__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('mudataObj_path') - parser.add_argument('GWAS_data_path') - parser.add_argument('-pk', '--prog_key', default='prog', type=str) - parser.add_argument('-dk', '--data_key', default='rna', type=str) - parser.add_argument('-n', '--n_jobs', default=1, type=int) - parser.add_argument('--output', action='store_false') + parser.add_argument("mudataObj_path") + parser.add_argument("GWAS_data_path") + parser.add_argument("-pk", "--prog_key", default="prog", type=str) + parser.add_argument("-dk", "--data_key", default="rna", type=str) + parser.add_argument("-n", "--n_jobs", default=1, type=int) + parser.add_argument("--output", action="store_false") args = parser.parse_args() mdata = mudata.read(args.mudataObj_path) - compute_trait_enrichment(mdata, args.GWAS_data_path, - prog_key=args.prog_key, - data_key=args.data_key, - n_jobs=args.n_jobs, - inplace=args.output) - \ No newline at end of file + compute_trait_enrichment( + mdata, + args.GWAS_data_path, + prog_key=args.prog_key, + data_key=args.data_key, + n_jobs=args.n_jobs, + inplace=args.output, + ) diff --git a/src/Stage2_Evaluation/A_Metrics/src/explained_variance.py b/src/Stage2_Evaluation/A_Metrics/src/explained_variance.py index b9961ae..c942bb7 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/explained_variance.py +++ b/src/Stage2_Evaluation/A_Metrics/src/explained_variance.py @@ -8,15 +8,16 @@ from torch_cnmf import cNMF from sklearn.decomposition import PCA + def compute_Var(X): """ Calculate total variance across all features in a matrix. - + Parameters ---------- X : array-like Input matrix where variance is computed across rows for each column - + Returns ------- float @@ -24,10 +25,11 @@ def compute_Var(X): """ return np.sum(np.var(X, axis=0, ddof=1)) + def computeVarianceExplained(X, H, Var_X, i): """ Calculate variance explained by a specific gene program component. - + Parameters ---------- X : array-like @@ -38,25 +40,27 @@ def computeVarianceExplained(X, H, Var_X, i): Total variance of original data i : int Index of the specific component to analyze - + Returns ------- float Fraction of variance explained by component i """ if not isinstance(H, (pd.DataFrame)): - B_k = X @ H[i,:].T / np.sqrt((np.asarray(H[i,:])**2).sum()) - numerator = compute_Var(X - np.outer(B_k, H[i,:])) + B_k = X @ H[i, :].T / np.sqrt((np.asarray(H[i, :]) ** 2).sum()) + numerator = compute_Var(X - np.outer(B_k, H[i, :])) else: - B_k = X @ H.iloc[i,:] / np.sqrt((H.iloc[i,:]**2).sum()) - numerator = compute_Var(X - np.outer(B_k, H.iloc[i,:])) - return (1 - numerator / Var_X) + B_k = X @ H.iloc[i, :] / np.sqrt((H.iloc[i, :] ** 2).sum()) + numerator = compute_Var(X - np.outer(B_k, H.iloc[i, :])) + return 1 - numerator / Var_X -def compute_explained_variance(cnmf_obj, X, k, output_folder, program_name = range(1, 51), thre = '2.0'): +def compute_explained_variance( + cnmf_obj, X, k, output_folder, program_name=range(1, 51), thre="2.0" +): """ Compute explained variance for all gene programs in a cNMF factorization. - + Parameters ---------- cnmf_obj : cnmf.cNMF @@ -70,58 +74,74 @@ def compute_explained_variance(cnmf_obj, X, k, output_folder, program_name = ran thre : str, default '2.0' Threshold parameter for consensus matrices program_name: list - List of program names - + List of program names + Returns ------- None Saves variance metrics to TSV files in output_folder """ - # Convert sparse matrix to dense if needed - if hasattr(X, 'toarray'): + # Convert sparse matrix to dense if needed + if hasattr(X, "toarray"): X = X.toarray() # Load consensus gene program matrices (W: usage, H: spectra) - thre_name = str(thre).replace('.', '_') - H_path = cnmf_obj.paths['consensus_spectra__txt'] % (k, thre_name) ## median_spectra_file - H_df = pd.read_csv(H_path, sep='\t', index_col=0).T + thre_name = str(thre).replace(".", "_") + H_path = cnmf_obj.paths["consensus_spectra__txt"] % ( + k, + thre_name, + ) ## median_spectra_file + H_df = pd.read_csv(H_path, sep="\t", index_col=0).T H = H_df.to_numpy() - H = (H/H.sum(0)) + H = H / H.sum(0) - W_path = cnmf_obj.paths['consensus_usages__txt'] % (k, thre_name) ## median_spectra_file - W_df = pd.read_csv(W_path, sep='\t', index_col=0) + W_path = cnmf_obj.paths["consensus_usages__txt"] % ( + k, + thre_name, + ) ## median_spectra_file + W_df = pd.read_csv(W_path, sep="\t", index_col=0) W = W_df.to_numpy() - W = (W/W.sum(0)) - + W = W / W.sum(0) WH = W @ H.T diff = X - WH - diff_sumOfSquaresError = (np.asarray(diff)**2).sum() + diff_sumOfSquaresError = (np.asarray(diff) ** 2).sum() Var_diff = compute_Var(diff) Var_X = compute_Var(X) TotalVarianceExplained = 1 - Var_diff / Var_X V_k = np.empty([k]) - - # Calculate variance explained by each individual program + # Calculate variance explained by each individual program for i in range(k): V_k[i] = computeVarianceExplained(X, H.T, Var_X, i) - ProgramID = ['K' + str(k) + '_' + str(i+1) for i in range(k)] + ProgramID = ["K" + str(k) + "_" + str(i + 1) for i in range(k)] # Save variance metrics and summary statistics to files - metrics_df = pd.DataFrame({'VarianceExplained': V_k, - 'ProgramID': ProgramID, - 'program_name': list(program_name) }) - - metrics_summary = pd.DataFrame({'Sum' : metrics_df['VarianceExplained'].sum(), - 'Median' : metrics_df['VarianceExplained'].median(), - 'Max' : metrics_df['VarianceExplained'].max(), - 'Total' : TotalVarianceExplained}, - index = [0]) - - - metrics_df.to_csv(os.path.join(output_folder, f"{k}_Explained_Variance.txt"), index = None, sep="\t") - metrics_summary.to_csv(os.path.join(output_folder, f"{k}_Explained_Variance_Summary.txt"), index = None, sep="\t") - \ No newline at end of file + metrics_df = pd.DataFrame( + { + "VarianceExplained": V_k, + "ProgramID": ProgramID, + "program_name": list(program_name), + } + ) + + metrics_summary = pd.DataFrame( + { + "Sum": metrics_df["VarianceExplained"].sum(), + "Median": metrics_df["VarianceExplained"].median(), + "Max": metrics_df["VarianceExplained"].max(), + "Total": TotalVarianceExplained, + }, + index=[0], + ) + + metrics_df.to_csv( + os.path.join(output_folder, f"{k}_Explained_Variance.txt"), index=None, sep="\t" + ) + metrics_summary.to_csv( + os.path.join(output_folder, f"{k}_Explained_Variance_Summary.txt"), + index=None, + sep="\t", + ) diff --git a/src/Stage2_Evaluation/B_Calibration/Slurm_version/CRT/CRT.py b/src/Stage2_Evaluation/B_Calibration/Slurm_version/CRT/CRT.py index 86bfe69..acd396c 100644 --- a/src/Stage2_Evaluation/B_Calibration/Slurm_version/CRT/CRT.py +++ b/src/Stage2_Evaluation/B_Calibration/Slurm_version/CRT/CRT.py @@ -1,4 +1,3 @@ - """ Conditional Randomization Test (CRT) Calibration Pipeline @@ -26,7 +25,13 @@ # (this file is CRT.py; the interpreter puts its own dir on sys.path[0], which would # otherwise shadow the `CRT` package and cause a circular self-import). from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[5] / 'src/Stage2_Evaluation/B_Calibration/src')) + +sys.path.insert( + 0, + str( + Path(__file__).resolve().parents[5] / "src/Stage2_Evaluation/B_Calibration/src" + ), +) from CRT import ( prepare_crt_inputs, @@ -42,7 +47,6 @@ from check_covariate import run_covariate_check - # reformat adata info for CRT def reformat_data_for_CRT(mdata, covariates=None, log_covariates=None): @@ -69,7 +73,7 @@ def reformat_data_for_CRT(mdata, covariates=None, log_covariates=None): covar_dict[key] = adata.obs[key] if log_covariates: for key in log_covariates: - covar_dict[f'log_{key}'] = np.log1p(adata.obs[key]) + covar_dict[f"log_{key}"] = np.log1p(adata.obs[key]) adata.obsm["covar"] = pd.DataFrame(covar_dict, index=adata.obs_names) @@ -90,11 +94,14 @@ def run_CRT(adata, k, sel_thresh, output_folder, args): # Cache path: if both real and fake results already exist, skip the (expensive) # CRT recompute and just (re)generate the QQ plot from the cached raw p-values. if args.skip_existing and os.path.exists(real_txt) and os.path.exists(fake_txt): - print(f" Cached K={k}, sel_thresh={sel_thresh}, condition={condition}: " - f"skipping compute, plotting from cache") + print( + f" Cached K={k}, sel_thresh={sel_thresh}, condition={condition}: " + f"skipping compute, plotting from cache" + ) plot_qq_from_cache(real_txt, fake_txt, png, condition) summarize_ntc_significance_from_cache( - real_txt, args, k, output_folder, condition, covar_tag) + real_txt, args, k, output_folder, condition, covar_tag + ) continue adata_con = adata[adata.obs[args.categorical_key] == condition].copy() @@ -104,10 +111,9 @@ def run_CRT(adata, k, sel_thresh, output_folder, args): usage_key="cnmf_usage", covar_key="covar", guide_assignment_key="guide_assignment", - guide2gene_key="guide2gene" + guide2gene_key="guide2gene", ) - out = run_all_genes_union_crt( inputs=inputs, B=args.number_permutations, @@ -117,7 +123,9 @@ def run_CRT(adata, k, sel_thresh, output_folder, args): return_skew_normal=True, ) - ntc_labels = args.guide_annotation_key # Identify NTC guides and build guide-frequency bins / real-gene signatures + ntc_labels = ( + args.guide_annotation_key + ) # Identify NTC guides and build guide-frequency bins / real-gene signatures ntc_guides, guide_freq, guide_to_bin, real_sigs = build_ntc_group_inputs( inputs=inputs, ntc_label=ntc_labels, @@ -147,7 +155,9 @@ def run_CRT(adata, k, sel_thresh, output_folder, args): # Save real + fake (NTC null) results (raw p-values for both, so they are comparable) save_result(out, k, thresh_tag, output_folder, condition, args, covar_tag) - save_fake_result(ntc_group_pvals_ens, k, output_folder, condition, args, covar_tag) + save_fake_result( + ntc_group_pvals_ens, k, output_folder, condition, args, covar_tag + ) # Empirical-FDR calibration check: treat non-targeting controls as genes and # count NTC regulator-program pairs passing FDR < 0.05 (per condition). @@ -168,22 +178,24 @@ def run_CRT(adata, k, sel_thresh, output_folder, args): def _covariate_tag(args): """Build filename token from covariate args, e.g. percent_mito_log_n_counts.""" - parts = list(args.covariates or []) + [f"log_{c}" for c in (args.log_covariates or [])] + parts = list(args.covariates or []) + [ + f"log_{c}" for c in (args.log_covariates or []) + ] return "_".join(parts) if parts else "no_covariates" -def _add_adj_pval(df, pcol, args, out_col='adj_pval'): +def _add_adj_pval(df, pcol, args, out_col="adj_pval"): """Add column `out_col` = FDR-corrected p-values from column `pcol` (BH or StoreyQ).""" - if args.FDR_method == 'BH': - + if args.FDR_method == "BH": from statsmodels.stats.multitest import multipletests - df[out_col] = multipletests(df[pcol], method='fdr_bh')[1] - elif args.FDR_method == 'StoreyQ': + df[out_col] = multipletests(df[pcol], method="fdr_bh")[1] + elif args.FDR_method == "StoreyQ": from multipy.fdr import qvalue import builtins - if not hasattr(builtins, 'xrange'): + + if not hasattr(builtins, "xrange"): builtins.xrange = range df[out_col] = qvalue(df[pcol].values, threshold=0.05, verbose=False)[1] @@ -192,32 +204,36 @@ def _add_adj_pval(df, pcol, args, out_col='adj_pval'): def _melt_programs(df, value_name): """Melt a genes(rows) x programs(cols) DataFrame to long (target_name, program_name, value).""" - return df.reset_index().melt( - id_vars='index', - var_name='program_name', - value_name=value_name, - ).rename(columns={'index': 'target_name'}) + return ( + df.reset_index() + .melt( + id_vars="index", + var_name="program_name", + value_name=value_name, + ) + .rename(columns={"index": "target_name"}) + ) # save real perturbation results def save_result(out, k, thresh_tag, output_folder, condition, args, covar_tag=None): - pval_long = _melt_programs(out['pvals_skew_df'], 'p-value') - beta_long = _melt_programs(out['betas_df'], 'beta_clr') + pval_long = _melt_programs(out["pvals_skew_df"], "p-value") + beta_long = _melt_programs(out["betas_df"], "beta_clr") # Merge on program and target_gene result_df = pval_long.merge( - beta_long, - on=['program_name', 'target_name'], - how='inner' + beta_long, on=["program_name", "target_name"], how="inner" ) # Also carry the raw (uncalibrated) CRT p-value so real and NTC-null are on the # same scale (the null / QQ plot use raw p-values). - has_raw = 'pvals_raw_df' in out + has_raw = "pvals_raw_df" in out if has_raw: - raw_long = _melt_programs(out['pvals_raw_df'], 'p-value_raw') - result_df = result_df.merge(raw_long, on=['program_name', 'target_name'], how='left') + raw_long = _melt_programs(out["pvals_raw_df"], "p-value_raw") + result_df = result_df.merge( + raw_long, on=["program_name", "target_name"], how="left" + ) # The CRT effect size is beta_hat on the CLR (natural-log) scale, NOT a log2 fold # change. A shift of beta in CLR coordinate k corresponds to a natural-log fold @@ -225,28 +241,34 @@ def save_result(out, k, thresh_tag, output_folder, condition, args, covar_tag=No # shift across the K programs, so ln-FC = beta * K/(K-1). Converting to base 2: # approx_log2FC = [K / (K - 1)] * beta_hat / ln(2) # K = k = number of programs (CLR is centered over all program columns). - result_df['log2FC'] = result_df['beta_clr'] * (k / (k - 1)) / np.log(2) + result_df["log2FC"] = result_df["beta_clr"] * (k / (k - 1)) / np.log(2) # Multiple-testing correction: FDR for the skew-calibrated AND the raw p-values. - result_df = _add_adj_pval(result_df, 'p-value', args, out_col='adj_pval') + result_df = _add_adj_pval(result_df, "p-value", args, out_col="adj_pval") if has_raw: - result_df = _add_adj_pval(result_df, 'p-value_raw', args, out_col='adj_pval_raw') + result_df = _add_adj_pval( + result_df, "p-value_raw", args, out_col="adj_pval_raw" + ) # Reorder columns (each p-value followed by its adjusted p-value) - cols = ['target_name', 'program_name', 'log2FC', 'p-value', 'adj_pval'] + cols = ["target_name", "program_name", "log2FC", "p-value", "adj_pval"] if has_raw: - cols += ['p-value_raw', 'adj_pval_raw'] + cols += ["p-value_raw", "adj_pval_raw"] result_df = result_df[cols] if covar_tag is None: covar_tag = _covariate_tag(args) - result_df.to_csv(f'{output_folder}/{k}_CRT_{covar_tag}_{condition}.txt', sep='\t', index=False) + result_df.to_csv( + f"{output_folder}/{k}_CRT_{covar_tag}_{condition}.txt", sep="\t", index=False + ) return result_df # save fake / null (NTC group ensemble) raw p-values -def save_fake_result(ntc_group_pvals_ens, k, output_folder, condition, args, covar_tag=None): +def save_fake_result( + ntc_group_pvals_ens, k, output_folder, condition, args, covar_tag=None +): """ Write the NTC null distribution (raw p-values only) so it can be re-analyzed / re-plotted offline. `ntc_group_pvals_ens` is a dict {ensemble -> DataFrame(rows=NTC @@ -255,24 +277,31 @@ def save_fake_result(ntc_group_pvals_ens, k, output_folder, condition, args, cov """ frames = [] for e in sorted(ntc_group_pvals_ens.keys()): - long_e = _melt_programs(ntc_group_pvals_ens[e], 'p-value_raw') - long_e['ensemble'] = e + long_e = _melt_programs(ntc_group_pvals_ens[e], "p-value_raw") + long_e["ensemble"] = e frames.append(long_e) fake_df = pd.concat(frames, ignore_index=True) - fake_df = _add_adj_pval(fake_df, 'p-value_raw', args, out_col='adj_pval_raw') - fake_df = fake_df[['ensemble', 'target_name', 'program_name', 'p-value_raw', 'adj_pval_raw']] + fake_df = _add_adj_pval(fake_df, "p-value_raw", args, out_col="adj_pval_raw") + fake_df = fake_df[ + ["ensemble", "target_name", "program_name", "p-value_raw", "adj_pval_raw"] + ] if covar_tag is None: covar_tag = _covariate_tag(args) - fake_df.to_csv(f'{output_folder}/{k}_CRT_fake_{covar_tag}_{condition}.txt', sep='\t', index=False) + fake_df.to_csv( + f"{output_folder}/{k}_CRT_fake_{covar_tag}_{condition}.txt", + sep="\t", + index=False, + ) return fake_df # negative-control calibration: empirical FDR from non-targeting "genes" -def _ntc_significance_core(raw_long, args, k, output_folder, condition, - covar_tag=None, p_thresh=0.05): +def _ntc_significance_core( + raw_long, args, k, output_folder, condition, covar_tag=None, p_thresh=0.05 +): """ Shared core for the NTC empirical-FDR check. @@ -301,66 +330,80 @@ def _ntc_significance_core(raw_long, args, k, output_folder, condition, ntc_labels = [ntc_labels] ntc_labels = set(ntc_labels) - ntc_present = sorted(set(raw_long['target_name']) & ntc_labels) + ntc_present = sorted(set(raw_long["target_name"]) & ntc_labels) if not ntc_present: - print(f" [NTC significance] none of {sorted(ntc_labels)} found in tested " - f"targets (condition={condition}); skipping.") + print( + f" [NTC significance] none of {sorted(ntc_labels)} found in tested " + f"targets (condition={condition}); skipping." + ) return None raw_long = raw_long.copy() - raw_long['significant'] = raw_long['adj_pval_raw'] < p_thresh + raw_long["significant"] = raw_long["adj_pval_raw"] < p_thresh # total discoveries across every regulator-program pair (all genes) - n_sig_total_all = int(raw_long['significant'].sum()) + n_sig_total_all = int(raw_long["significant"].sum()) - ntc_long = raw_long[raw_long['target_name'].isin(ntc_labels)] + ntc_long = raw_long[raw_long["target_name"].isin(ntc_labels)] # per-control-gene breakdown per_gene = ( - ntc_long.groupby('target_name') - .agg(n_significant=('significant', 'sum'), n_pairs=('significant', 'size')) + ntc_long.groupby("target_name") + .agg(n_significant=("significant", "sum"), n_pairs=("significant", "size")) .reset_index() ) - per_gene['n_significant'] = per_gene['n_significant'].astype(int) - per_gene['fraction_of_ntc_pairs'] = per_gene['n_significant'] / per_gene['n_pairs'] + per_gene["n_significant"] = per_gene["n_significant"].astype(int) + per_gene["fraction_of_ntc_pairs"] = per_gene["n_significant"] / per_gene["n_pairs"] - n_sig_ntc = int(ntc_long['significant'].sum()) + n_sig_ntc = int(ntc_long["significant"].sum()) n_ntc_pairs = int(len(ntc_long)) - frac_ntc = n_sig_ntc / n_ntc_pairs if n_ntc_pairs else float('nan') - empirical_fdr = n_sig_ntc / n_sig_total_all if n_sig_total_all else float('nan') - - print(f" [NTC significance] condition={condition}: {n_sig_ntc}/{n_ntc_pairs} " - f"NTC regulator-program pairs pass FDR < {p_thresh} " - f"(fraction_of_ntc_pairs={frac_ntc:.4f}); " - f"empirical_fdr={n_sig_ntc}/{n_sig_total_all}={empirical_fdr:.4f} " - f"across {len(ntc_present)} control target(s): {ntc_present}") + frac_ntc = n_sig_ntc / n_ntc_pairs if n_ntc_pairs else float("nan") + empirical_fdr = n_sig_ntc / n_sig_total_all if n_sig_total_all else float("nan") + + print( + f" [NTC significance] condition={condition}: {n_sig_ntc}/{n_ntc_pairs} " + f"NTC regulator-program pairs pass FDR < {p_thresh} " + f"(fraction_of_ntc_pairs={frac_ntc:.4f}); " + f"empirical_fdr={n_sig_ntc}/{n_sig_total_all}={empirical_fdr:.4f} " + f"across {len(ntc_present)} control target(s): {ntc_present}" + ) if covar_tag is None: covar_tag = _covariate_tag(args) - overall = pd.DataFrame([{ - 'target_name': 'ALL_NTC', - 'n_significant': n_sig_ntc, - 'n_pairs': n_ntc_pairs, - 'fraction_of_ntc_pairs': frac_ntc, - 'n_significant_all_genes': n_sig_total_all, - 'empirical_fdr': empirical_fdr, - }]) + overall = pd.DataFrame( + [ + { + "target_name": "ALL_NTC", + "n_significant": n_sig_ntc, + "n_pairs": n_ntc_pairs, + "fraction_of_ntc_pairs": frac_ntc, + "n_significant_all_genes": n_sig_total_all, + "empirical_fdr": empirical_fdr, + } + ] + ) summary = pd.concat([per_gene, overall], ignore_index=True) summary.to_csv( - f'{output_folder}/{k}_CRT_ntc_significance_{covar_tag}_{condition}.txt', - sep='\t', index=False, + f"{output_folder}/{k}_CRT_ntc_significance_{covar_tag}_{condition}.txt", + sep="\t", + index=False, ) - return {'n_significant': n_sig_ntc, 'n_ntc_pairs': n_ntc_pairs, - 'fraction_of_ntc_pairs': frac_ntc, - 'n_significant_all_genes': n_sig_total_all, - 'empirical_fdr': empirical_fdr, - 'per_gene': per_gene, 'p_thresh': p_thresh} + return { + "n_significant": n_sig_ntc, + "n_ntc_pairs": n_ntc_pairs, + "fraction_of_ntc_pairs": frac_ntc, + "n_significant_all_genes": n_sig_total_all, + "empirical_fdr": empirical_fdr, + "per_gene": per_gene, + "p_thresh": p_thresh, + } -def summarize_ntc_significance(out, args, k, output_folder, condition, - covar_tag=None, p_thresh=0.05): +def summarize_ntc_significance( + out, args, k, output_folder, condition, covar_tag=None, p_thresh=0.05 +): """ Empirical-FDR calibration check from the freshly-computed CRT results. @@ -370,47 +413,63 @@ def summarize_ntc_significance(out, args, k, output_folder, condition, ``_ntc_significance_core`` to score the non-targeting controls. See that function for the reported metrics and output file. """ - pvals_raw = out.get('pvals_raw_df') + pvals_raw = out.get("pvals_raw_df") if pvals_raw is None: raise KeyError( "out['pvals_raw_df'] not found — run_all_genes_union_crt must be called " "with return_raw_pvals=True." ) - raw_long = _melt_programs(pvals_raw, 'p-value_raw') - raw_long = _add_adj_pval(raw_long, 'p-value_raw', args, out_col='adj_pval_raw') - return _ntc_significance_core(raw_long, args, k, output_folder, condition, - covar_tag=covar_tag, p_thresh=p_thresh) + raw_long = _melt_programs(pvals_raw, "p-value_raw") + raw_long = _add_adj_pval(raw_long, "p-value_raw", args, out_col="adj_pval_raw") + return _ntc_significance_core( + raw_long, + args, + k, + output_folder, + condition, + covar_tag=covar_tag, + p_thresh=p_thresh, + ) -def summarize_ntc_significance_from_cache(real_txt, args, k, output_folder, condition, - covar_tag=None, p_thresh=0.05): +def summarize_ntc_significance_from_cache( + real_txt, args, k, output_folder, condition, covar_tag=None, p_thresh=0.05 +): """ Same empirical-FDR check as summarize_ntc_significance, but reusing the cached real result file (which already carries adj_pval_raw FDR-corrected over the full family in save_result). Returns None if the cache predates the raw-p-value columns. """ - real_df = pd.read_csv(real_txt, sep='\t') - if 'adj_pval_raw' not in real_df.columns: - print(f" [NTC significance] cached {os.path.basename(real_txt)} lacks " - f"'adj_pval_raw' (condition={condition}); skipping.") + real_df = pd.read_csv(real_txt, sep="\t") + if "adj_pval_raw" not in real_df.columns: + print( + f" [NTC significance] cached {os.path.basename(real_txt)} lacks " + f"'adj_pval_raw' (condition={condition}); skipping." + ) return None - return _ntc_significance_core(real_df[['target_name', 'adj_pval_raw']], - args, k, output_folder, condition, - covar_tag=covar_tag, p_thresh=p_thresh) + return _ntc_significance_core( + real_df[["target_name", "adj_pval_raw"]], + args, + k, + output_folder, + condition, + covar_tag=covar_tag, + p_thresh=p_thresh, + ) # regenerate the QQ plot from cached result files (no CRT recompute) def plot_qq_from_cache(real_txt, fake_txt, png, condition): import matplotlib.pyplot as plt - real_df = pd.read_csv(real_txt, sep='\t') - fake_df = pd.read_csv(fake_txt, sep='\t') + real_df = pd.read_csv(real_txt, sep="\t") + fake_df = pd.read_csv(fake_txt, sep="\t") qq_plot_real_vs_null( - real_pvals=real_df['p-value_raw'], - null_pvals=fake_df['p-value_raw'], + real_pvals=real_df["p-value_raw"], + null_pvals=fake_df["p-value_raw"], title=f"CRT QQ: real vs null (NTC) p-values for {condition}", ) plt.tight_layout() @@ -423,52 +482,131 @@ def main(): description="Conditional Randomization Test (CRT) calibration pipeline" ) - #IO info - parser.add_argument('--out_dir', help='Directory containing cNMF output files for calibration analysis', type=str, required=True) - parser.add_argument('--run_name', help='Name of the cNMF run to perform calibration on (must match name used during inference)', type=str, required=True) - parser.add_argument('--K', nargs='*', type=int, help="list of K values (number of components) to test", default=[30, 50, 70, 80, 100, 200, 300]) - parser.add_argument('--sel_threshs', nargs='*', type=float, help="list of density threshold values for consensus selection", default=[0.2, 2.0]) + # IO info + parser.add_argument( + "--out_dir", + help="Directory containing cNMF output files for calibration analysis", + type=str, + required=True, + ) + parser.add_argument( + "--run_name", + help="Name of the cNMF run to perform calibration on (must match name used during inference)", + type=str, + required=True, + ) + parser.add_argument( + "--K", + nargs="*", + type=int, + help="list of K values (number of components) to test", + default=[30, 50, 70, 80, 100, 200, 300], + ) + parser.add_argument( + "--sel_threshs", + nargs="*", + type=float, + help="list of density threshold values for consensus selection", + default=[0.2, 2.0], + ) # keys - parser.add_argument('--categorical_key', help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", type=str, default="sample") + parser.add_argument( + "--categorical_key", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + type=str, + default="sample", + ) # Covariates - parser.add_argument('--covariates', nargs='*', type=str, help='Covariate keys in .obs to include as-is (e.g., any continues value, categorical values)', default=None) - parser.add_argument('--log_covariates', nargs='*', type=str, help='Covariate keys in .obs to log1p-transform before inclusion (e.g., any count based values)', default=None) + parser.add_argument( + "--covariates", + nargs="*", + type=str, + help="Covariate keys in .obs to include as-is (e.g., any continues value, categorical values)", + default=None, + ) + parser.add_argument( + "--log_covariates", + nargs="*", + type=str, + help="Covariate keys in .obs to log1p-transform before inclusion (e.g., any count based values)", + default=None, + ) # Covariate independence diagnostic (run before CRT) - parser.add_argument('--check_covariate', action='store_true', help='Before running CRT, run the covariate independence diagnostic (Pearson correlation heatmap + VIF) on the encoded design matrix and write it to /covariate_check/. Useful for catching collinear covariates that would make CRT p-values unreliable.') - parser.add_argument('--numeric_as_category_threshold', type=int, default=20, help='For the covariate check: numeric covariates with <= this many unique values are one-hot encoded (matches CRT get_covar_matrix default: 20).') + parser.add_argument( + "--check_covariate", + action="store_true", + help="Before running CRT, run the covariate independence diagnostic (Pearson correlation heatmap + VIF) on the encoded design matrix and write it to /covariate_check/. Useful for catching collinear covariates that would make CRT p-values unreliable.", + ) + parser.add_argument( + "--numeric_as_category_threshold", + type=int, + default=20, + help="For the covariate check: numeric covariates with <= this many unique values are one-hot encoded (matches CRT get_covar_matrix default: 20).", + ) # Calibration parameters - parser.add_argument('--number_guide', help='Number of non-targeting guides to randomly designate as "targeting" in each calibration iteration (default: 6)', type=int, default=6) - parser.add_argument('--number_permutations', help='Number of calibration iterations to run with (default: 1024)', type=int, default=1024) - parser.add_argument('--guide_annotation_key', nargs='*', type=str, help='Name of target for non-targeting/safe-targeting guides,default="non-targeting"', default='non-targeting') - parser.add_argument('--FDR_method', type=str, choices=['BH', 'StoreyQ'], default='BH', help='FDR correction method: BH (Benjamini-Hochberg) or StoreyQ (Storey Q-value) (default: BH)') - parser.add_argument('--save_dir', type=str, default=None, help='Base directory under which {K}_{thresh} subdirs are created. Default: //Evaluation/') - parser.add_argument('--skip_existing', help='If set, skip per-(K, sel_thresh, condition) computations whose output .txt files already exist. Useful for resuming preempted jobs.', action='store_true') + parser.add_argument( + "--number_guide", + help='Number of non-targeting guides to randomly designate as "targeting" in each calibration iteration (default: 6)', + type=int, + default=6, + ) + parser.add_argument( + "--number_permutations", + help="Number of calibration iterations to run with (default: 1024)", + type=int, + default=1024, + ) + parser.add_argument( + "--guide_annotation_key", + nargs="*", + type=str, + help='Name of target for non-targeting/safe-targeting guides,default="non-targeting"', + default="non-targeting", + ) + parser.add_argument( + "--FDR_method", + type=str, + choices=["BH", "StoreyQ"], + default="BH", + help="FDR correction method: BH (Benjamini-Hochberg) or StoreyQ (Storey Q-value) (default: BH)", + ) + parser.add_argument( + "--save_dir", + type=str, + default=None, + help="Base directory under which {K}_{thresh} subdirs are created. Default: //Evaluation/", + ) + parser.add_argument( + "--skip_existing", + help="If set, skip per-(K, sel_thresh, condition) computations whose output .txt files already exist. Useful for resuming preempted jobs.", + action="store_true", + ) args = parser.parse_args() # --- Save config (incl. SLURM info) --- slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), - 'cpus_per_task': os.environ.get('SLURM_CPUS_PER_TASK'), - 'mem_per_node': os.environ.get('SLURM_MEM_PER_NODE'), - 'time_limit': os.environ.get('SLURM_JOB_TIMELIMIT'), - 'submit_dir': os.environ.get('SLURM_SUBMIT_DIR'), - 'array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + "mem_per_node": os.environ.get("SLURM_MEM_PER_NODE"), + "time_limit": os.environ.get("SLURM_JOB_TIMELIMIT"), + "submit_dir": os.environ.get("SLURM_SUBMIT_DIR"), + "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" - save_config_dir = args.save_dir or f'{args.out_dir}/{args.run_name}/Evaluation' + save_config_dir = args.save_dir or f"{args.out_dir}/{args.run_name}/Evaluation" os.makedirs(save_config_dir, exist_ok=True) - config_to_save = {'script_args': vars(args), 'slurm_info': slurm_info} - with open(f'{save_config_dir}/config_{job_id}.yml', 'w') as f: + config_to_save = {"script_args": vars(args), "slurm_info": slurm_info} + with open(f"{save_config_dir}/config_{job_id}.yml", "w") as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) # Optional covariate independence diagnostic. Covariates are K-independent, so run @@ -477,12 +615,14 @@ def main(): if args.check_covariate: first_k = args.K[0] first_thresh = args.sel_threshs[0] - first_tag = str(first_thresh).replace('.', '_') - check_save = f'{save_config_dir}/covariate_check' - print(f"Running covariate independence check (K={first_k}, " - f"sel_thresh={first_thresh}) -> {check_save}") + first_tag = str(first_thresh).replace(".", "_") + check_save = f"{save_config_dir}/covariate_check" + print( + f"Running covariate independence check (K={first_k}, " + f"sel_thresh={first_thresh}) -> {check_save}" + ) check_mdata = mu.read( - f'{args.out_dir}/{args.run_name}/Inference/adata/cNMF_{first_k}_{first_tag}.h5mu' + f"{args.out_dir}/{args.run_name}/Inference/adata/cNMF_{first_k}_{first_tag}.h5mu" ) run_covariate_check( check_mdata["cNMF"], @@ -499,20 +639,24 @@ def main(): for k in args.K: print(f"Processing K={k}, sel_thresh={sel_thresh}") - thresh_tag = str(sel_thresh).replace('.', '_') + thresh_tag = str(sel_thresh).replace(".", "_") if args.save_dir: output_folder = f"{args.save_dir}/{k}_{thresh_tag}" else: - output_folder = f"{args.out_dir}/{args.run_name}/Evaluation/{k}_{thresh_tag}" + output_folder = ( + f"{args.out_dir}/{args.run_name}/Evaluation/{k}_{thresh_tag}" + ) os.makedirs(output_folder, exist_ok=True) # Load mdata - mdata = mu.read(f'{args.out_dir}/{args.run_name}/Inference/adata/cNMF_{k}_{str(sel_thresh).replace(".","_")}.h5mu') + mdata = mu.read( + f"{args.out_dir}/{args.run_name}/Inference/adata/cNMF_{k}_{str(sel_thresh).replace('.', '_')}.h5mu" + ) # Assign guide - adata = reformat_data_for_CRT(mdata, - covariates=args.covariates, - log_covariates=args.log_covariates) + adata = reformat_data_for_CRT( + mdata, covariates=args.covariates, log_covariates=args.log_covariates + ) # add flooring for program matrix with exceesive zeros U = adata.obsm["cnmf_usage"].copy() @@ -527,5 +671,5 @@ def main(): return 0 -if __name__ == '__main__': - sys.exit(main()) \ No newline at end of file +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/Stage2_Evaluation/B_Calibration/Slurm_version/U-test_perturbation_calibration/U-test_perturbation_calibration.py b/src/Stage2_Evaluation/B_Calibration/Slurm_version/U-test_perturbation_calibration/U-test_perturbation_calibration.py index 232d56d..87518f7 100644 --- a/src/Stage2_Evaluation/B_Calibration/Slurm_version/U-test_perturbation_calibration/U-test_perturbation_calibration.py +++ b/src/Stage2_Evaluation/B_Calibration/Slurm_version/U-test_perturbation_calibration/U-test_perturbation_calibration.py @@ -16,10 +16,16 @@ # Change path to wherever you have repo locally from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[5] / 'src')) + +sys.path.append(str(Path(__file__).resolve().parents[5] / "src")) # Point at B_Calibration/src so the U_test package (src/U_test/) is importable. -sys.path.insert(0, str(Path(__file__).resolve().parents[5] / 'src/Stage2_Evaluation/B_Calibration/src')) +sys.path.insert( + 0, + str( + Path(__file__).resolve().parents[5] / "src/Stage2_Evaluation/B_Calibration/src" + ), +) from U_test import ( compute_real_perturbation_tests, @@ -38,59 +44,166 @@ def main(): description="U-test perturbation calibration analysis" ) - #IO info - parser.add_argument('--out_dir', help='Directory containing cNMF output files for calibration analysis', type=str, required=True) - parser.add_argument('--run_name', help='Name of the cNMF run to perform calibration on (must match name used during inference)', type=str, required=True) - parser.add_argument('--K', nargs='*', type=int, help="list of K values (number of components) to test", default=[30, 50, 70, 80, 100, 200, 300]) - parser.add_argument('--sel_threshs', nargs='*', type=float, help="list of density threshold values for consensus selection", default=[0.2, 2.0]) + # IO info + parser.add_argument( + "--out_dir", + help="Directory containing cNMF output files for calibration analysis", + type=str, + required=True, + ) + parser.add_argument( + "--run_name", + help="Name of the cNMF run to perform calibration on (must match name used during inference)", + type=str, + required=True, + ) + parser.add_argument( + "--K", + nargs="*", + type=int, + help="list of K values (number of components) to test", + default=[30, 50, 70, 80, 100, 200, 300], + ) + parser.add_argument( + "--sel_threshs", + nargs="*", + type=float, + help="list of density threshold values for consensus selection", + default=[0.2, 2.0], + ) # resources - parser.add_argument('--guide_annotation_path', type=str, help='Path to tab-separated file with guide annotations including "targeting" column (True/False) to identify non-targeting guides for calibration') - parser.add_argument('--guide_annotation_key', nargs='+', type=str, help='Name of target for non-targeting/safe-targeting guides, default="non-targeting"', default=['non-targeting']) - parser.add_argument('--reference_gtf_path', type=str, help='Path to reference GTF file for validating gene names during format checking (optional)') + parser.add_argument( + "--guide_annotation_path", + type=str, + help='Path to tab-separated file with guide annotations including "targeting" column (True/False) to identify non-targeting guides for calibration', + ) + parser.add_argument( + "--guide_annotation_key", + nargs="+", + type=str, + help='Name of target for non-targeting/safe-targeting guides, default="non-targeting"', + default=["non-targeting"], + ) + parser.add_argument( + "--reference_gtf_path", + type=str, + help="Path to reference GTF file for validating gene names during format checking (optional)", + ) # keys - parser.add_argument('--data_key', help='Key to access gene expression data in MuData object (default: rna)', type=str, default="rna") - parser.add_argument('--prog_key', help='Key to access cNMF programs in MuData object (default: cNMF)', type=str, default="cNMF") - parser.add_argument('--categorical_key', help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", type=str, default="sample") - parser.add_argument('--guide_names_key', help='Key in .uns to access guide names (default: guide_names)', type=str, default="guide_names") - parser.add_argument('--guide_targets_key', help='Key in .uns to access guide target genes (default: guide_targets)', type=str, default="guide_targets") - parser.add_argument('--guide_assignment_key', help='Key in .obsm to access guide assignment matrix (default: guide_assignment)', type=str, default="guide_assignment") - parser.add_argument('--organism', help='Organism/species for analysis (default: human)', type=str, default="human") - parser.add_argument('--FDR_method', help='Method for FDR correction in real perturbation tests (default: StoreyQ)', type=str, default="StoreyQ") + parser.add_argument( + "--data_key", + help="Key to access gene expression data in MuData object (default: rna)", + type=str, + default="rna", + ) + parser.add_argument( + "--prog_key", + help="Key to access cNMF programs in MuData object (default: cNMF)", + type=str, + default="cNMF", + ) + parser.add_argument( + "--categorical_key", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + type=str, + default="sample", + ) + parser.add_argument( + "--guide_names_key", + help="Key in .uns to access guide names (default: guide_names)", + type=str, + default="guide_names", + ) + parser.add_argument( + "--guide_targets_key", + help="Key in .uns to access guide target genes (default: guide_targets)", + type=str, + default="guide_targets", + ) + parser.add_argument( + "--guide_assignment_key", + help="Key in .obsm to access guide assignment matrix (default: guide_assignment)", + type=str, + default="guide_assignment", + ) + parser.add_argument( + "--organism", + help="Organism/species for analysis (default: human)", + type=str, + default="human", + ) + parser.add_argument( + "--FDR_method", + help="Method for FDR correction in real perturbation tests (default: StoreyQ)", + type=str, + default="StoreyQ", + ) # check format - parser.add_argument('--check_format', help='If set, validate MuData format and check for all necessary keys before running calibration', action="store_true") + parser.add_argument( + "--check_format", + help="If set, validate MuData format and check for all necessary keys before running calibration", + action="store_true", + ) # Calibration parameters - parser.add_argument('--number_run', help='Number of calibration iterations to run with randomly selected fake targeting guides (default: 300)', type=int, default=300) - parser.add_argument('--number_guide', help='Number of non-targeting guides to randomly designate as "targeting" in each calibration iteration (default: 6)', type=int, default=6) - parser.add_argument('--compute_real_perturbation_tests', help='If set, compute perturbation association tests on real targeting guides', action="store_true") - parser.add_argument('--compute_fake_perturbation_tests', help='If set, compute perturbation association tests on fake targeting guides (calibration null distribution)', action="store_true") - parser.add_argument('--visualizations', help='If set, generate and save QQ plots and violin plots comparing real vs null distributions', action="store_true") - parser.add_argument('--skip_existing', help='If set, skip per-(K, sel_thresh, sample) computations whose output files already exist. Existing outputs are loaded into the accumulator so visualizations still work. Useful for resuming preempted jobs.', action="store_true") + parser.add_argument( + "--number_run", + help="Number of calibration iterations to run with randomly selected fake targeting guides (default: 300)", + type=int, + default=300, + ) + parser.add_argument( + "--number_guide", + help='Number of non-targeting guides to randomly designate as "targeting" in each calibration iteration (default: 6)', + type=int, + default=6, + ) + parser.add_argument( + "--compute_real_perturbation_tests", + help="If set, compute perturbation association tests on real targeting guides", + action="store_true", + ) + parser.add_argument( + "--compute_fake_perturbation_tests", + help="If set, compute perturbation association tests on fake targeting guides (calibration null distribution)", + action="store_true", + ) + parser.add_argument( + "--visualizations", + help="If set, generate and save QQ plots and violin plots comparing real vs null distributions", + action="store_true", + ) + parser.add_argument( + "--skip_existing", + help="If set, skip per-(K, sel_thresh, sample) computations whose output files already exist. Existing outputs are loaded into the accumulator so visualizations still work. Useful for resuming preempted jobs.", + action="store_true", + ) args = parser.parse_args() # --- Save config (incl. SLURM info) --- slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), - 'cpus_per_task': os.environ.get('SLURM_CPUS_PER_TASK'), - 'mem_per_node': os.environ.get('SLURM_MEM_PER_NODE'), - 'time_limit': os.environ.get('SLURM_JOB_TIMELIMIT'), - 'submit_dir': os.environ.get('SLURM_SUBMIT_DIR'), - 'array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + "mem_per_node": os.environ.get("SLURM_MEM_PER_NODE"), + "time_limit": os.environ.get("SLURM_JOB_TIMELIMIT"), + "submit_dir": os.environ.get("SLURM_SUBMIT_DIR"), + "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" - config_to_save = {'script_args': vars(args), 'slurm_info': slurm_info} - with open(f'{args.out_dir}/{args.run_name}/Evaluation/config_{job_id}.yml', 'w') as f: + config_to_save = {"script_args": vars(args), "slurm_info": slurm_info} + with open( + f"{args.out_dir}/{args.run_name}/Evaluation/config_{job_id}.yml", "w" + ) as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) - print("=" * 80) print("U-test Perturbation Calibration Analysis") print("=" * 80) @@ -136,36 +249,51 @@ def main(): # Load Merge datasets for visualizations if args.visualizations: - # Load results if not computed in this run if not args.compute_real_perturbation_tests: print("\nLoading pre-computed real perturbation results...") test_stats_real_df = load_real_perturbation_tests( - out_dir=args.out_dir, run_name=args.run_name, K=args.K, sel_threshs=args.sel_threshs, + out_dir=args.out_dir, + run_name=args.run_name, + K=args.K, + sel_threshs=args.sel_threshs, ) if not args.compute_fake_perturbation_tests: print("\nLoading pre-computed fake perturbation results...") test_stats_fake_df = load_fake_perturbation_tests( - out_dir=args.out_dir, run_name=args.run_name, K=args.K, sel_threshs=args.sel_threshs, + out_dir=args.out_dir, + run_name=args.run_name, + K=args.K, + sel_threshs=args.sel_threshs, ) print("\nMerging datasets...") - test_stats_dfs = pd.concat([test_stats_real_df, test_stats_fake_df], ignore_index=True) + test_stats_dfs = pd.concat( + [test_stats_real_df, test_stats_fake_df], ignore_index=True + ) # Create visualizations print("\nCreating visualizations...") plot_calibration_comparison( - test_stats_dfs, out_dir=args.out_dir, run_name=args.run_name, K=args.K, sel_threshs=args.sel_threshs, + test_stats_dfs, + out_dir=args.out_dir, + run_name=args.run_name, + K=args.K, + sel_threshs=args.sel_threshs, ) plot_qq_comparison( - test_stats_dfs, out_dir=args.out_dir, run_name=args.run_name, K=args.K, sel_threshs=args.sel_threshs, + test_stats_dfs, + out_dir=args.out_dir, + run_name=args.run_name, + K=args.K, + sel_threshs=args.sel_threshs, ) print("\nPipeline finished.") return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage2_Evaluation/B_Calibration/src/CRT/ntc_groups.py b/src/Stage2_Evaluation/B_Calibration/src/CRT/ntc_groups.py index 982315c..573c9a3 100644 --- a/src/Stage2_Evaluation/B_Calibration/src/CRT/ntc_groups.py +++ b/src/Stage2_Evaluation/B_Calibration/src/CRT/ntc_groups.py @@ -21,9 +21,7 @@ logger = logging.getLogger(__name__) -def guide_frequency( - G: sp.spmatrix, guide_names: Sequence[str] -) -> Dict[str, float]: +def guide_frequency(G: sp.spmatrix, guide_names: Sequence[str]) -> Dict[str, float]: """ Returns prevalence per guide: freq[g] = mean(G[:,g] > 0) across cells. """ @@ -209,7 +207,9 @@ def make_ntc_groups_ensemble( def _validate_group_sizes( - groups: Mapping[str, Sequence[str]], # CHANGED: dropped `expected_size` param — size now follows --number_guide, not a hardcoded value + groups: Mapping[ + str, Sequence[str] + ], # CHANGED: dropped `expected_size` param — size now follows --number_guide, not a hardcoded value ) -> Dict[str, float]: sizes = np.array([len(guides) for guides in groups.values()], dtype=np.int32) if sizes.size == 0: @@ -305,7 +305,7 @@ def crt_pvals_for_ntc_groups_ensemble( cols = [guide_to_col[g] for g in guides if g in guide_to_col] if not cols: continue - seed = (hash((seed0, e, group_id)) & 0xFFFFFFFF) + seed = hash((seed0, e, group_id)) & 0xFFFFFFFF pvals = crt_pvals_for_guide_set( inputs=inputs, guide_idx=np.asarray(cols, dtype=np.int32), @@ -347,7 +347,7 @@ def crt_pvals_for_ntc_groups_ensemble_skew( cols = [guide_to_col[g] for g in guides if g in guide_to_col] if not cols: continue - seed = (hash((seed0, e, group_id)) & 0xFFFFFFFF) + seed = hash((seed0, e, group_id)) & 0xFFFFFFFF pvals_skew, _ = crt_pvals_for_guide_set_skew( inputs=inputs, guide_idx=np.asarray(cols, dtype=np.int32), diff --git a/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline.py b/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline.py index a9f4d23..afcac7b 100644 --- a/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline.py +++ b/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline.py @@ -299,7 +299,9 @@ def run_all_genes_union_crt( skew_params = None if return_skew_normal: if not calibrate_skew_normal: - raise ValueError("return_skew_normal=True requires calibrate_skew_normal=True.") + raise ValueError( + "return_skew_normal=True requires calibrate_skew_normal=True." + ) pvals_skew_df, skew_params = _stack_skew_outputs( results, gene_list, inputs.program_names ) diff --git a/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline_helpers.py b/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline_helpers.py index 3e59ee4..14149ba 100644 --- a/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline_helpers.py +++ b/src/Stage2_Evaluation/B_Calibration/src/CRT/pipeline_helpers.py @@ -46,7 +46,9 @@ def _fit_propensity( return _extract_probabilities(propensity_model(inputs.C, y01)) -def _sample_crt_indices(p: np.ndarray, B: int, seed: int) -> Tuple[np.ndarray, np.ndarray]: +def _sample_crt_indices( + p: np.ndarray, B: int, seed: int +) -> Tuple[np.ndarray, np.ndarray]: return crt_index_sampler_fast_numba(p, B, seed) @@ -161,9 +163,7 @@ def _stack_skew_outputs( ) -> Tuple[pd.DataFrame, np.ndarray]: pvals_sn_mat = np.vstack([r.pvals_sn for r in results]) skew_params = np.stack([r.skew_params for r in results]) - pvals_skew_df = pd.DataFrame( - pvals_sn_mat, index=gene_list, columns=program_names - ) + pvals_skew_df = pd.DataFrame(pvals_sn_mat, index=gene_list, columns=program_names) return pvals_skew_df, skew_params diff --git a/src/Stage2_Evaluation/B_Calibration/src/CRT/qq_plot.py b/src/Stage2_Evaluation/B_Calibration/src/CRT/qq_plot.py index b5daf63..4b71dd6 100644 --- a/src/Stage2_Evaluation/B_Calibration/src/CRT/qq_plot.py +++ b/src/Stage2_Evaluation/B_Calibration/src/CRT/qq_plot.py @@ -19,7 +19,11 @@ def _flatten_pvals_any(obj) -> np.ndarray: elif isinstance(obj, Mapping): parts = [] for df in obj.values(): - a = df.to_numpy() if isinstance(df, (pd.DataFrame, pd.Series)) else np.asarray(df) + a = ( + df.to_numpy() + if isinstance(df, (pd.DataFrame, pd.Series)) + else np.asarray(df) + ) parts.append(np.asarray(a, dtype=np.float64).ravel()) arr = np.concatenate(parts) if parts else np.empty(0, dtype=np.float64) else: @@ -65,14 +69,32 @@ def qq_plot_real_vs_null(real_pvals, null_pvals, ax=None, title=None): # Fixed CRT style/colors: real = purple ("All observed"), null = blue (NTC). xmax = max(float(x_real.max()), float(x_null.max())) - ax.plot([0.0, xmax], [0.0, xmax], color="#333333", linewidth=1.0, label="y = x (null)") - ax.scatter(x_null, y_null, label="Null (NTC / fake)", color="#1f77b4", - marker=".", s=14.0, alpha=0.6) - ax.scatter(x_real, y_real, label="Real (perturbations)", color="#9467bd", - marker=".", s=14.0, alpha=0.6) + ax.plot( + [0.0, xmax], [0.0, xmax], color="#333333", linewidth=1.0, label="y = x (null)" + ) + ax.scatter( + x_null, + y_null, + label="Null (NTC / fake)", + color="#1f77b4", + marker=".", + s=14.0, + alpha=0.6, + ) + ax.scatter( + x_real, + y_real, + label="Real (perturbations)", + color="#9467bd", + marker=".", + s=14.0, + alpha=0.6, + ) ax.set_xlabel("Expected -log10(p)") ax.set_ylabel("Observed -log10(p)") - ax.set_title(title if title is not None else "CRT QQ plot: real vs null (NTC) p-values") + ax.set_title( + title if title is not None else "CRT QQ plot: real vs null (NTC) p-values" + ) ax.legend() return ax diff --git a/src/Stage2_Evaluation/B_Calibration/src/U_test/calibration.py b/src/Stage2_Evaluation/B_Calibration/src/U_test/calibration.py index 240bec7..9a53113 100644 --- a/src/Stage2_Evaluation/B_Calibration/src/U_test/calibration.py +++ b/src/Stage2_Evaluation/B_Calibration/src/U_test/calibration.py @@ -31,7 +31,9 @@ # Make the repo `src` root importable so we can reach the Stage2 metrics package # regardless of the caller's own sys.path setup. -_SRC_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..')) +_SRC_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..") +) if _SRC_ROOT not in sys.path: sys.path.append(_SRC_ROOT) @@ -64,12 +66,12 @@ def resolve_reference_targets(patterns, available_targets): available = set(available_targets) matched = set() for pat in patterns: - if pat in available: # fast path / exact label + if pat in available: # fast path / exact label matched.add(pat) continue try: rx = re.compile(pat) - except re.error: # not a valid regex and not an exact label + except re.error: # not a valid regex and not an exact label continue matched.update(t for t in available if rx.search(t)) return sorted(matched) if matched else list(patterns) @@ -120,8 +122,10 @@ def compute_real_perturbation_tests( # Validate once using the first K and sel_thresh first_k = K[0] first_thresh = sel_threshs[0] - thresh_str = str(first_thresh).replace('.', '_') - mdata_check = mu.read(f'{out_dir}/{run_name}/Inference/adata/cNMF_{first_k}_{thresh_str}.h5mu') + thresh_str = str(first_thresh).replace(".", "_") + mdata_check = mu.read( + f"{out_dir}/{run_name}/Inference/adata/cNMF_{first_k}_{thresh_str}.h5mu" + ) mdata_targets = set(mdata_check[prog_key].uns[guide_targets_key]) @@ -131,7 +135,9 @@ def compute_real_perturbation_tests( # rather than a single "non-targeting" label. Path-based reference_targets (guide # names) are left untouched. if guide_annotation_path is None: - reference_targets = resolve_reference_targets(guide_annotation_key, mdata_targets) + reference_targets = resolve_reference_targets( + guide_annotation_key, mdata_targets + ) matched_ref = mdata_targets.intersection(reference_targets) print(f"reference_targets overlap with mdata guide_targets: {matched_ref}") @@ -159,19 +165,25 @@ def compute_real_perturbation_tests( for k in K: print(f"Processing K={k}, sel_thresh={sel_thresh}") - output_folder = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.','_')}" + output_folder = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}" os.makedirs(output_folder, exist_ok=True) # Load mdata - mdata = mu.read(f'{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{str(sel_thresh).replace(".","_")}.h5mu') + mdata = mu.read( + f"{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{str(sel_thresh).replace('.', '_')}.h5mu" + ) # Run perturbation association for each sample for samp in mdata[data_key].obs[categorical_key].unique(): - out_path = f'{output_folder}/{k}_perturbation_association_results_{samp}.txt' + out_path = ( + f"{output_folder}/{k}_perturbation_association_results_{samp}.txt" + ) if skip_existing and os.path.exists(out_path): - print(f" Skipping K={k}, sel_thresh={sel_thresh}, samp={samp}: output exists") - test_stats_df = pd.read_csv(out_path, sep='\t') + print( + f" Skipping K={k}, sel_thresh={sel_thresh}, samp={samp}: output exists" + ) + test_stats_df = pd.read_csv(out_path, sep="\t") else: mdata_ = mdata[mdata[data_key].obs[categorical_key] == samp] @@ -183,21 +195,21 @@ def compute_real_perturbation_tests( reference_targets=reference_targets, FDR_method=FDR_method, n_jobs=-1, - inplace=False + inplace=False, ) # Save results - test_stats_df.to_csv(out_path, sep='\t', index=False) + test_stats_df.to_csv(out_path, sep="\t", index=False) # Add metadata - test_stats_df['sample'] = samp - test_stats_df['K'] = k - test_stats_df['sel_thresh'] = sel_thresh + test_stats_df["sample"] = samp + test_stats_df["K"] = k + test_stats_df["sel_thresh"] = sel_thresh test_stats_real_df.append(test_stats_df) # Concatenate all results test_stats_real_df = pd.concat(test_stats_real_df, ignore_index=True) - test_stats_real_df['real'] = True + test_stats_real_df["real"] = True return test_stats_real_df @@ -232,10 +244,16 @@ def compute_fake_perturbation_tests( # read guide annotation file to find non-targeting guide names and targets if guide_annotation_path is not None: - guide_target = pd.read_csv(guide_annotation_path, sep='\t') - non_targeting_idx = guide_target.index[guide_target.targeting == False] # targeting is col with True / False - guide_target = guide_target.loc[non_targeting_idx] # subset only non-targeting/safe-targeting guides - guide_target['type'] = "non-targeting" # set both safe-targeting and non-targeting to be non-targeting + guide_target = pd.read_csv(guide_annotation_path, sep="\t") + non_targeting_idx = guide_target.index[ + guide_target.targeting == False + ] # targeting is col with True / False + guide_target = guide_target.loc[ + non_targeting_idx + ] # subset only non-targeting/safe-targeting guides + guide_target["type"] = ( + "non-targeting" # set both safe-targeting and non-targeting to be non-targeting + ) else: guide_target = None non_targeting_idx = None @@ -245,16 +263,22 @@ def compute_fake_perturbation_tests( # Extract guide info from mdata if no annotation file provided if guide_target is None: first_k = K[0] - thresh_str = str(sel_threshs[0]).replace('.', '_') - mdata_check = mu.read(f'{out_dir}/{run_name}/Inference/adata/cNMF_{first_k}_{thresh_str}.h5mu') + thresh_str = str(sel_threshs[0]).replace(".", "_") + mdata_check = mu.read( + f"{out_dir}/{run_name}/Inference/adata/cNMF_{first_k}_{thresh_str}.h5mu" + ) guide_targets_arr = mdata_check[prog_key].uns[guide_targets_key] guide_names_arr = mdata_check[prog_key].uns[guide_names_key] # Expand guide_annotation_key (exact label OR regex) against the dataset's # actual target labels so suffixed control variants are all captured. - reference_targets = resolve_reference_targets(guide_annotation_key, set(guide_targets_arr)) - non_targeting_idx = np.where(np.isin(guide_targets_arr, list(reference_targets)))[0] + reference_targets = resolve_reference_targets( + guide_annotation_key, set(guide_targets_arr) + ) + non_targeting_idx = np.where( + np.isin(guide_targets_arr, list(reference_targets)) + )[0] if len(non_targeting_idx) == 0: raise ValueError( @@ -264,22 +288,26 @@ def compute_fake_perturbation_tests( f"Example mdata guide_targets values: {list(set(guide_targets_arr))[:5]}" ) - print(f' Non-targeting guides: {len(non_targeting_idx)}') + print(f" Non-targeting guides: {len(non_targeting_idx)}") # make a new annotation df when we don't have annotation file - guide_target = pd.DataFrame({ - guide_names_key: guide_names_arr[non_targeting_idx], - 'type': 'non-targeting' - }) + guide_target = pd.DataFrame( + { + guide_names_key: guide_names_arr[non_targeting_idx], + "type": "non-targeting", + } + ) del mdata_check # Validate guide count between annotation file and mdata if guide_annotation_path is not None: first_k = K[0] - thresh_str = str(sel_threshs[0]).replace('.', '_') - mdata_check = mu.read(f'{out_dir}/{run_name}/Inference/adata/cNMF_{first_k}_{thresh_str}.h5mu') + thresh_str = str(sel_threshs[0]).replace(".", "_") + mdata_check = mu.read( + f"{out_dir}/{run_name}/Inference/adata/cNMF_{first_k}_{thresh_str}.h5mu" + ) - n_file_guides = len(pd.read_csv(guide_annotation_path, sep='\t')) + n_file_guides = len(pd.read_csv(guide_annotation_path, sep="\t")) n_mdata_guides = len(mdata_check[prog_key].uns[guide_names_key]) if n_file_guides != n_mdata_guides: @@ -306,46 +334,70 @@ def compute_fake_perturbation_tests( print(f"Processing K={k}, sel_thresh={sel_thresh}") # Load mdata - output_folder = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.','_')}" + output_folder = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}" os.makedirs(output_folder, exist_ok=True) mdata = mu.read( - f'{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{str(sel_thresh).replace(".","_")}.h5mu' + f"{out_dir}/{run_name}/Inference/adata/cNMF_{k}_{str(sel_thresh).replace('.', '_')}.h5mu" ) # Skip this (K, sel_thresh) if all fake outputs already exist; load them into the accumulator. if skip_existing: samples_for_iter = list(mdata[data_key].obs[categorical_key].unique()) - expected = [f'{output_folder}/{k}_fake_perturbation_association_results_{samp}.txt' for samp in samples_for_iter] + expected = [ + f"{output_folder}/{k}_fake_perturbation_association_results_{samp}.txt" + for samp in samples_for_iter + ] if all(os.path.exists(p) for p in expected): - print(f" Skipping K={k}, sel_thresh={sel_thresh}: all fake outputs exist; loading from disk") + print( + f" Skipping K={k}, sel_thresh={sel_thresh}: all fake outputs exist; loading from disk" + ) for samp, p in zip(samples_for_iter, expected): - samp_df = pd.read_csv(p, sep='\t') + samp_df = pd.read_csv(p, sep="\t") test_stats_fake_dfs.append(samp_df) continue test_stats_fake_dfs_temp = [] for i in range(number_run): - print(f" Running iteration {i+1}/{number_run}") + print(f" Running iteration {i + 1}/{number_run}") # Randomly make number_guide non-targeting guides "targeting" guide_target_ = guide_target.copy() - selected_guides = np.random.choice(guide_target_[guide_names_key], number_guide, replace=False) - guide_target_.loc[guide_target_.guide_names.isin(selected_guides),'type'] = 'targeting' # type is a col with targeting / non-targeting + selected_guides = np.random.choice( + guide_target_[guide_names_key], number_guide, replace=False + ) + guide_target_.loc[ + guide_target_.guide_names.isin(selected_guides), "type" + ] = "targeting" # type is a col with targeting / non-targeting # Filter to only non-targeting guides that exist in both datasets - valid_guide_mask = np.isin(mdata[prog_key].uns[guide_names_key],guide_target_[guide_names_key].values) + valid_guide_mask = np.isin( + mdata[prog_key].uns[guide_names_key], + guide_target_[guide_names_key].values, + ) valid_indices = np.where(valid_guide_mask)[0] - print(f" Found {len(valid_indices)} valid non-targeting guides out of " f"{len(mdata[prog_key].uns[guide_names_key])} total") + print( + f" Found {len(valid_indices)} valid non-targeting guides out of " + f"{len(mdata[prog_key].uns[guide_names_key])} total" + ) if len(valid_indices) == 0: raise ValueError("No valid guides found") _mdata = mdata.copy() - _mdata[prog_key].obsm[guide_assignment_key] = mdata[prog_key].obsm[guide_assignment_key][:, non_targeting_idx] - _mdata[prog_key].uns[guide_names_key] = mdata[prog_key].uns[guide_names_key][non_targeting_idx] - _mdata[prog_key].uns[guide_targets_key] = guide_target_.loc[guide_target_.guide_names.isin(mdata[prog_key].uns[guide_names_key]),'type'].values + _mdata[prog_key].obsm[guide_assignment_key] = mdata[prog_key].obsm[ + guide_assignment_key + ][:, non_targeting_idx] + _mdata[prog_key].uns[guide_names_key] = mdata[prog_key].uns[ + guide_names_key + ][non_targeting_idx] + _mdata[prog_key].uns[guide_targets_key] = guide_target_.loc[ + guide_target_.guide_names.isin( + mdata[prog_key].uns[guide_names_key] + ), + "type", + ].values # Run perturbation association for each sample for samp in _mdata[data_key].obs[categorical_key].unique(): @@ -356,32 +408,36 @@ def compute_fake_perturbation_tests( prog_key=prog_key, collapse_targets=True, pseudobulk=False, - reference_targets=['non-targeting'], + reference_targets=["non-targeting"], FDR_method=FDR_method, n_jobs=-1, - inplace=False + inplace=False, ) test_stats_df[categorical_key] = samp - test_stats_df['K'] = k - test_stats_df['run'] = i - test_stats_df['sel_thresh'] = sel_thresh - test_stats_fake_dfs.append(test_stats_df) # combine all - test_stats_fake_dfs_temp.append(test_stats_df) # combine for each k and sel_thresh + test_stats_df["K"] = k + test_stats_df["run"] = i + test_stats_df["sel_thresh"] = sel_thresh + test_stats_fake_dfs.append(test_stats_df) # combine all + test_stats_fake_dfs_temp.append( + test_stats_df + ) # combine for each k and sel_thresh # Save results per condition - test_stats_fake_dfs_temp = pd.concat(test_stats_fake_dfs_temp, ignore_index=True) + test_stats_fake_dfs_temp = pd.concat( + test_stats_fake_dfs_temp, ignore_index=True + ) for samp, samp_df in test_stats_fake_dfs_temp.groupby(categorical_key): samp_df.to_csv( - f'{output_folder}/{k}_fake_perturbation_association_results_{samp}.txt', - sep='\t', - index=False + f"{output_folder}/{k}_fake_perturbation_association_results_{samp}.txt", + sep="\t", + index=False, ) # Concatenate all results test_stats_fake_dfs = pd.concat(test_stats_fake_dfs, ignore_index=True) - test_stats_fake_dfs['real'] = False + test_stats_fake_dfs["real"] = False return test_stats_fake_dfs @@ -394,10 +450,20 @@ def load_real_perturbation_tests(out_dir, run_name, K, sel_threshs): # Discover sample names from the first K directory first_k = K[0] - thresh_str = str(sel_threshs[0]).replace('.', '_') - first_dir = f'{out_dir}/{run_name}/Evaluation/{first_k}_{thresh_str}' - sample_files = [f for f in os.listdir(first_dir) if f.startswith(f'{first_k}_perturbation_association_results_') and f.endswith('.txt')] - samples = [f.replace(f'{first_k}_perturbation_association_results_', '').replace('.txt', '') for f in sample_files] + thresh_str = str(sel_threshs[0]).replace(".", "_") + first_dir = f"{out_dir}/{run_name}/Evaluation/{first_k}_{thresh_str}" + sample_files = [ + f + for f in os.listdir(first_dir) + if f.startswith(f"{first_k}_perturbation_association_results_") + and f.endswith(".txt") + ] + samples = [ + f.replace(f"{first_k}_perturbation_association_results_", "").replace( + ".txt", "" + ) + for f in sample_files + ] print(f"Discovered samples: {samples}") test_stats_real_df = [] @@ -405,18 +471,18 @@ def load_real_perturbation_tests(out_dir, run_name, K, sel_threshs): for sel_thresh in sel_threshs: for k in K: for samp in samples: - thresh_str = str(sel_thresh).replace('.', '_') + thresh_str = str(sel_thresh).replace(".", "_") test_stats_df_ = pd.read_csv( - f'{out_dir}/{run_name}/Evaluation/{k}_{thresh_str}/{k}_perturbation_association_results_{samp}.txt', - sep='\t' + f"{out_dir}/{run_name}/Evaluation/{k}_{thresh_str}/{k}_perturbation_association_results_{samp}.txt", + sep="\t", ) - test_stats_df_['sample'] = samp - test_stats_df_['K'] = k - test_stats_df_['sel_thresh'] = sel_thresh + test_stats_df_["sample"] = samp + test_stats_df_["K"] = k + test_stats_df_["sel_thresh"] = sel_thresh test_stats_real_df.append(test_stats_df_) test_stats_real_df = pd.concat(test_stats_real_df, ignore_index=True) - test_stats_real_df['real'] = True + test_stats_real_df["real"] = True return test_stats_real_df @@ -426,10 +492,20 @@ def load_fake_perturbation_tests(out_dir, run_name, K, sel_threshs): # Discover sample names from the first K directory first_k = K[0] - thresh_str = str(sel_threshs[0]).replace('.', '_') - first_dir = f'{out_dir}/{run_name}/Evaluation/{first_k}_{thresh_str}' - sample_files = [f for f in os.listdir(first_dir) if f.startswith(f'{first_k}_fake_perturbation_association_results_') and f.endswith('.txt')] - samples = [f.replace(f'{first_k}_fake_perturbation_association_results_', '').replace('.txt', '') for f in sample_files] + thresh_str = str(sel_threshs[0]).replace(".", "_") + first_dir = f"{out_dir}/{run_name}/Evaluation/{first_k}_{thresh_str}" + sample_files = [ + f + for f in os.listdir(first_dir) + if f.startswith(f"{first_k}_fake_perturbation_association_results_") + and f.endswith(".txt") + ] + samples = [ + f.replace(f"{first_k}_fake_perturbation_association_results_", "").replace( + ".txt", "" + ) + for f in sample_files + ] print(f"Discovered samples for fake tests: {samples}") test_stats_fake_df = [] @@ -437,15 +513,15 @@ def load_fake_perturbation_tests(out_dir, run_name, K, sel_threshs): for sel_thresh in sel_threshs: for k in K: for samp in samples: - thresh_str = str(sel_thresh).replace('.', '_') + thresh_str = str(sel_thresh).replace(".", "_") test_stats_df_ = pd.read_csv( - f'{out_dir}/{run_name}/Evaluation/{k}_{thresh_str}/{k}_fake_perturbation_association_results_{samp}.txt', - sep='\t' + f"{out_dir}/{run_name}/Evaluation/{k}_{thresh_str}/{k}_fake_perturbation_association_results_{samp}.txt", + sep="\t", ) test_stats_fake_df.append(test_stats_df_) test_stats_fake_df = pd.concat(test_stats_fake_df, ignore_index=True) - test_stats_fake_df['real'] = False + test_stats_fake_df["real"] = False return test_stats_fake_df @@ -458,7 +534,7 @@ def plot_calibration_comparison(test_stats_dfs, out_dir, run_name, K, sel_thresh for sel_thresh in sel_threshs: for k in K: - thresh_str = str(sel_thresh).replace('.', '_') + thresh_str = str(sel_thresh).replace(".", "_") output_folder = f"{out_dir}/{run_name}/Evaluation/{k}_{thresh_str}" os.makedirs(output_folder, exist_ok=True) @@ -470,26 +546,27 @@ def plot_calibration_comparison(test_stats_dfs, out_dir, run_name, K, sel_thresh print(f" Skipping K={k}, sel_thresh={sel_thresh}: no data") continue - test_stats_k['neg_log_pval'] = test_stats_k['pval'].apply(lambda x: -np.log(x)) + test_stats_k["neg_log_pval"] = test_stats_k["pval"].apply( + lambda x: -np.log(x) + ) fig, ax = plt.subplots(figsize=(6, 5)) sns.violinplot( - x='sample', - y='neg_log_pval', - hue='real', - data=test_stats_k, - ax=ax + x="sample", y="neg_log_pval", hue="real", data=test_stats_k, ax=ax ) - ax.set_title(f'K={k}, sel_thresh={sel_thresh}') - ax.set_xlabel('Sample ID') - ax.set_ylabel('-ln(p-value)') + ax.set_title(f"K={k}, sel_thresh={sel_thresh}") + ax.set_xlabel("Sample ID") + ax.set_ylabel("-ln(p-value)") ax.set_ylim(0, 50) - ax.axhline(8, color='grey', linestyle='dashed') + ax.axhline(8, color="grey", linestyle="dashed") plt.tight_layout() - plt.savefig(f'{output_folder}/U_test_perturbation_association_calibration.png', dpi=100) + plt.savefig( + f"{output_folder}/U_test_perturbation_association_calibration.png", + dpi=100, + ) plt.close(fig) @@ -498,23 +575,23 @@ def plot_qq_comparison(test_stats_dfs, out_dir, run_name, K, sel_threshs): for sel_thresh in sel_threshs: for k in K: - thresh_str = str(sel_thresh).replace('.', '_') + thresh_str = str(sel_thresh).replace(".", "_") output_folder = f"{out_dir}/{run_name}/Evaluation/{k}_{thresh_str}" os.makedirs(output_folder, exist_ok=True) real_pvals = test_stats_dfs.loc[ - (test_stats_dfs.K == k) & - (test_stats_dfs.sel_thresh == sel_thresh) & - (test_stats_dfs.real == True), - 'pval' + (test_stats_dfs.K == k) + & (test_stats_dfs.sel_thresh == sel_thresh) + & (test_stats_dfs.real == True), + "pval", ] null_pvals = test_stats_dfs.loc[ - (test_stats_dfs.K == k) & - (test_stats_dfs.sel_thresh == sel_thresh) & - (test_stats_dfs.real == False) & - (test_stats_dfs.target_name == 'targeting'), - 'pval' + (test_stats_dfs.K == k) + & (test_stats_dfs.sel_thresh == sel_thresh) + & (test_stats_dfs.real == False) + & (test_stats_dfs.target_name == "targeting"), + "pval", ] if len(real_pvals) == 0 and len(null_pvals) == 0: @@ -523,8 +600,8 @@ def plot_qq_comparison(test_stats_dfs, out_dir, run_name, K, sel_threshs): fig, ax = plt.subplots(figsize=(6, 5)) - qqplot(data=real_pvals, ax=ax, color='blue', label='Real') - qqplot(data=null_pvals, ax=ax, color='red', label='Null') + qqplot(data=real_pvals, ax=ax, color="blue", label="Real") + qqplot(data=null_pvals, ax=ax, color="red", label="Null") lines = ax.get_lines() all_x = np.concatenate([line.get_xdata() for line in lines]) @@ -532,12 +609,18 @@ def plot_qq_comparison(test_stats_dfs, out_dir, run_name, K, sel_threshs): padding = 0.05 x_range = all_x.max() - all_x.min() y_range = all_y.max() - all_y.min() - ax.set_xlim(all_x.min() - padding * x_range, all_x.max() + padding * x_range) - ax.set_ylim(all_y.min() - padding * y_range, all_y.max() + padding * y_range) + ax.set_xlim( + all_x.min() - padding * x_range, all_x.max() + padding * x_range + ) + ax.set_ylim( + all_y.min() - padding * y_range, all_y.max() + padding * y_range + ) - ax.set_title(f'K={k}, sel_thresh={sel_thresh}') + ax.set_title(f"K={k}, sel_thresh={sel_thresh}") ax.legend() plt.tight_layout() - plt.savefig(f'{output_folder}/U_test_perturbation_association_qqplot.png', dpi=100) + plt.savefig( + f"{output_folder}/U_test_perturbation_association_qqplot.png", dpi=100 + ) plt.close(fig) diff --git a/src/Stage2_Evaluation/B_Calibration/src/check_covariate.py b/src/Stage2_Evaluation/B_Calibration/src/check_covariate.py index e7dd5ef..6af9bfc 100644 --- a/src/Stage2_Evaluation/B_Calibration/src/check_covariate.py +++ b/src/Stage2_Evaluation/B_Calibration/src/check_covariate.py @@ -28,6 +28,7 @@ import os import matplotlib + matplotlib.use("Agg") # headless / SLURM-safe import matplotlib.pyplot as plt import numpy as np @@ -58,7 +59,9 @@ def build_covar_df(adata, covariates=None, log_covariates=None): covar_dict[f"log_{key}"] = np.log1p(adata.obs[key]) if not covar_dict: - raise ValueError("No covariates provided; pass covariates and/or log_covariates.") + raise ValueError( + "No covariates provided; pass covariates and/or log_covariates." + ) return pd.DataFrame(covar_dict, index=adata.obs_names) @@ -86,7 +89,9 @@ def compute_vif(encoded_df): vifs.append((col, v)) vif_df = pd.DataFrame(vifs, columns=["covariate", "VIF"]) - return vif_df.sort_values("VIF", ascending=False, na_position="first").reset_index(drop=True) + return vif_df.sort_values("VIF", ascending=False, na_position="first").reset_index( + drop=True + ) def plot_correlation_heatmap(corr, png_path, title): @@ -105,8 +110,12 @@ def plot_correlation_heatmap(corr, png_path, title): for j in range(n): val = corr.iloc[i, j] ax.text( - j, i, f"{val:.2f}", - ha="center", va="center", fontsize=7, + j, + i, + f"{val:.2f}", + ha="center", + va="center", + fontsize=7, color="white" if abs(val) > 0.5 else "black", ) @@ -143,7 +152,9 @@ def run_check(covar_df, save_path, prefix, threshold, corr_flag=0.8, vif_flag=5. # --- stdout summary --- label = f"[{prefix}] " if prefix else "" - print(f"\n{label}Encoded covariate columns ({encoded.shape[1]}): {list(encoded.columns)}") + print( + f"\n{label}Encoded covariate columns ({encoded.shape[1]}): {list(encoded.columns)}" + ) # high-correlation off-diagonal pairs high_pairs = [] @@ -154,7 +165,9 @@ def run_check(covar_df, save_path, prefix, threshold, corr_flag=0.8, vif_flag=5. if abs(r) > corr_flag: high_pairs.append((cols[i], cols[j], r)) if high_pairs: - print(f"{label}WARNING: |Pearson r| > {corr_flag} pairs (potential redundancy):") + print( + f"{label}WARNING: |Pearson r| > {corr_flag} pairs (potential redundancy):" + ) for a, b, r in sorted(high_pairs, key=lambda t: -abs(t[2])): print(f" {a} <-> {b} : r = {r:+.2f}") else: @@ -171,9 +184,15 @@ def run_check(covar_df, save_path, prefix, threshold, corr_flag=0.8, vif_flag=5. print(f"{label}Wrote:\n {png}\n {corr_tsv}\n {vif_tsv}") -def run_covariate_check(adata, save_path, covariates=None, log_covariates=None, - categorical_key=None, per_condition=False, - numeric_as_category_threshold=20): +def run_covariate_check( + adata, + save_path, + covariates=None, + log_covariates=None, + categorical_key=None, + per_condition=False, + numeric_as_category_threshold=20, +): """High-level entry point: build the covariate design and run correlation + VIF. Intended for programmatic use (e.g. CRT.py's ``--check_covariate``). Takes an @@ -199,7 +218,9 @@ def run_covariate_check(adata, save_path, covariates=None, log_covariates=None, Numeric columns with <= this many unique values are one-hot encoded (matches CRT ``get_covar_matrix`` default: 20). """ - covar_df = build_covar_df(adata, covariates=covariates, log_covariates=log_covariates) + covar_df = build_covar_df( + adata, covariates=covariates, log_covariates=log_covariates + ) if per_condition: if not categorical_key: @@ -211,11 +232,15 @@ def run_covariate_check(adata, save_path, covariates=None, log_covariates=None, for cond in conditions: mask = (adata.obs[categorical_key] == cond).to_numpy() run_check( - covar_df.loc[mask], save_path, prefix=str(cond), + covar_df.loc[mask], + save_path, + prefix=str(cond), threshold=numeric_as_category_threshold, ) else: run_check( - covar_df, save_path, prefix="", + covar_df, + save_path, + prefix="", threshold=numeric_as_category_threshold, ) diff --git a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py index b87e853..dc2392a 100644 --- a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py +++ b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py @@ -6,90 +6,178 @@ # Change path to wherever you have repo locally from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) -from Stage3_Interpretation.A_Plotting.src import (load_stablity_error_data, plot_stablity_error,\ - load_enrichment_data, plot_enrichment,\ - load_perturbation_data, plot_perturbation,\ - load_explained_variance_data,plot_explained_variance, programs_dotplots,plot_k_selection_panel,\ - read_condition_labels_from_h5mu - ) +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) + +from Stage3_Interpretation.A_Plotting.src import ( + load_stablity_error_data, + plot_stablity_error, + load_enrichment_data, + plot_enrichment, + load_perturbation_data, + plot_perturbation, + load_explained_variance_data, + plot_explained_variance, + programs_dotplots, + plot_k_selection_panel, + read_condition_labels_from_h5mu, +) def main(): parser = argparse.ArgumentParser() - parser.add_argument('--output_directory', type=str, required=True) - parser.add_argument('--run_name', type=str, required=True) - parser.add_argument('--groupby', type=str, default="sample", - help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample") - parser.add_argument('--K', nargs='*', type=int, default=[30, 50, 70, 80, 100, 200, 300], help='list of K values (number of components)') - parser.add_argument('--save_folder_name', type=str, required=True) - parser.add_argument('--pval', type=float, default=0.05) - parser.add_argument('--eval_folder_name', type=str, required=True) - parser.add_argument('--sel_threshs', nargs='*', type=float, default=[0.2, 2.0], help='list of density thresholds') - parser.add_argument('--Conditions', nargs='*', type=str, default=None, help='Condition labels (values of obs[groupby]); default: all labels found in the data (first cNMF_{K}_{thresh}.h5mu)') - parser.add_argument('--selected_k', type=int, default=None) + parser.add_argument("--output_directory", type=str, required=True) + parser.add_argument("--run_name", type=str, required=True) + parser.add_argument( + "--groupby", + type=str, + default="sample", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + ) + parser.add_argument( + "--K", + nargs="*", + type=int, + default=[30, 50, 70, 80, 100, 200, 300], + help="list of K values (number of components)", + ) + parser.add_argument("--save_folder_name", type=str, required=True) + parser.add_argument("--pval", type=float, default=0.05) + parser.add_argument("--eval_folder_name", type=str, required=True) + parser.add_argument( + "--sel_threshs", + nargs="*", + type=float, + default=[0.2, 2.0], + help="list of density thresholds", + ) + parser.add_argument( + "--Conditions", + nargs="*", + type=str, + default=None, + help="Condition labels (values of obs[groupby]); default: all labels found in the data (first cNMF_{K}_{thresh}.h5mu)", + ) + parser.add_argument("--selected_k", type=int, default=None) # Enrichment file name and column name arguments - parser.add_argument('--go_file', type=str, default=None, - help='GO enrichment file name pattern. Use {k} as placeholder for K value. Default: {k}_GO_term_enrichment.txt') - parser.add_argument('--geneset_file', type=str, default=None, - help='Geneset enrichment file name pattern. Use {k} as placeholder for K value. Default: {k}_geneset_enrichment.txt') - parser.add_argument('--trait_file', type=str, default=None, - help='Trait enrichment file name pattern. Use {k} as placeholder for K value. Default: {k}_trait_enrichment.txt') - parser.add_argument('--term_col', type=str, default='Term', - help='Column name for the term/pathway name in enrichment files. Default: Term') - parser.add_argument('--adjpval_col', type=str, default='Adjusted P-value', - help='Column name for adjusted p-value in enrichment files. Default: Adjusted P-value') - parser.add_argument('--perturbation_file', type=str, default=None, - help='Perturbation file name pattern. Use {k} and {sample} as placeholders. Default: {k}_perturbation_association_results_{sample}.txt') - parser.add_argument('--perturb_adjpval_col', type=str, default='adj_pval', - help='Column name for adjusted p-value in perturbation files. Default: adj_pval') - parser.add_argument('--perturb_target_col', type=str, default='target_name', - help='Column name for target/regulator name in perturbation files. Default: target_name') - parser.add_argument('--perturb_log2fc_col', type=str, default='log2FC', - help='Column name for log2 fold change in perturbation files. Default: log2FC') - parser.add_argument('--variance_file', type=str, default=None, - help='Explained variance file name pattern. Use {k} as placeholder. Default: {k}_Explained_Variance_Summary.txt') - parser.add_argument('--variance_col', type=str, default='Total', - help='Column name for variance values. Use "Total" for summary files or column name for per-program files (will be summed). Default: Total') - parser.add_argument('--stability_file', type=str, default=None, - help='Path to a pre-computed stability/error file (TSV or NPZ). Bypasses cnmf.consensus(). ' - 'Useful for torch-cNMF runs where the cnmf package is not installed.') - parser.add_argument('--run_program_dotplot', action='store_true', - help='If set, also generate per-(K, sel_thresh) program dotplots. Requires the inference ' - 'h5mu (cNMF_{K}_{thresh}.h5mu) to exist for each (K, sel_thresh). Default off.') - + parser.add_argument( + "--go_file", + type=str, + default=None, + help="GO enrichment file name pattern. Use {k} as placeholder for K value. Default: {k}_GO_term_enrichment.txt", + ) + parser.add_argument( + "--geneset_file", + type=str, + default=None, + help="Geneset enrichment file name pattern. Use {k} as placeholder for K value. Default: {k}_geneset_enrichment.txt", + ) + parser.add_argument( + "--trait_file", + type=str, + default=None, + help="Trait enrichment file name pattern. Use {k} as placeholder for K value. Default: {k}_trait_enrichment.txt", + ) + parser.add_argument( + "--term_col", + type=str, + default="Term", + help="Column name for the term/pathway name in enrichment files. Default: Term", + ) + parser.add_argument( + "--adjpval_col", + type=str, + default="Adjusted P-value", + help="Column name for adjusted p-value in enrichment files. Default: Adjusted P-value", + ) + parser.add_argument( + "--perturbation_file", + type=str, + default=None, + help="Perturbation file name pattern. Use {k} and {sample} as placeholders. Default: {k}_perturbation_association_results_{sample}.txt", + ) + parser.add_argument( + "--perturb_adjpval_col", + type=str, + default="adj_pval", + help="Column name for adjusted p-value in perturbation files. Default: adj_pval", + ) + parser.add_argument( + "--perturb_target_col", + type=str, + default="target_name", + help="Column name for target/regulator name in perturbation files. Default: target_name", + ) + parser.add_argument( + "--perturb_log2fc_col", + type=str, + default="log2FC", + help="Column name for log2 fold change in perturbation files. Default: log2FC", + ) + parser.add_argument( + "--variance_file", + type=str, + default=None, + help="Explained variance file name pattern. Use {k} as placeholder. Default: {k}_Explained_Variance_Summary.txt", + ) + parser.add_argument( + "--variance_col", + type=str, + default="Total", + help='Column name for variance values. Use "Total" for summary files or column name for per-program files (will be summed). Default: Total', + ) + parser.add_argument( + "--stability_file", + type=str, + default=None, + help="Path to a pre-computed stability/error file (TSV or NPZ). Bypasses cnmf.consensus(). " + "Useful for torch-cNMF runs where the cnmf package is not installed.", + ) + parser.add_argument( + "--run_program_dotplot", + action="store_true", + help="If set, also generate per-(K, sel_thresh) program dotplots. Requires the inference " + "h5mu (cNMF_{K}_{thresh}.h5mu) to exist for each (K, sel_thresh). Default off.", + ) args = parser.parse_args() # save comfigs used args_dict = vars(args) - job_id = os.environ.get('SLURM_JOB_ID') + job_id = os.environ.get("SLURM_JOB_ID") - os.makedirs(f'{args.save_folder_name}', exist_ok=True) - with open(f'{args.save_folder_name}/config_{job_id}.yml', 'w') as f: + os.makedirs(f"{args.save_folder_name}", exist_ok=True) + with open(f"{args.save_folder_name}/config_{job_id}.yml", "w") as f: yaml.dump(args_dict, f, default_flow_style=False, width=1000) # ---- Condition labels: default to the values of --groupby in the inference h5mu ---- if args.Conditions is None: - first_h5mu = (f"{args.output_directory}/{args.run_name}/Inference/adata/" - f"cNMF_{args.K[0]}_{str(args.sel_threshs[0]).replace('.', '_')}.h5mu") - args.Conditions = read_condition_labels_from_h5mu(first_h5mu, 'rna', args.groupby) - print(f"--Conditions not given; using {args.groupby} values from {first_h5mu}: {args.Conditions}") + first_h5mu = ( + f"{args.output_directory}/{args.run_name}/Inference/adata/" + f"cNMF_{args.K[0]}_{str(args.sel_threshs[0]).replace('.', '_')}.h5mu" + ) + args.Conditions = read_condition_labels_from_h5mu( + first_h5mu, "rna", args.groupby + ) + print( + f"--Conditions not given; using {args.groupby} values from {first_h5mu}: {args.Conditions}" + ) # ---- Pre-flight check: verify all required evaluation files exist ---- - _go_pat = args.go_file or '{k}_GO_term_enrichment.txt' - _gs_pat = args.geneset_file or '{k}_geneset_enrichment.txt' - _tr_pat = args.trait_file or '{k}_trait_enrichment.txt' - _pt_pat = args.perturbation_file or '{k}_perturbation_association_results_{sample}.txt' - _ev_pat = args.variance_file or '{k}_Explained_Variance_Summary.txt' + _go_pat = args.go_file or "{k}_GO_term_enrichment.txt" + _gs_pat = args.geneset_file or "{k}_geneset_enrichment.txt" + _tr_pat = args.trait_file or "{k}_trait_enrichment.txt" + _pt_pat = ( + args.perturbation_file or "{k}_perturbation_association_results_{sample}.txt" + ) + _ev_pat = args.variance_file or "{k}_Explained_Variance_Summary.txt" missing = [] for sel_thresh in args.sel_threshs: - sel_str = str(sel_thresh).replace('.', '_') + sel_str = str(sel_thresh).replace(".", "_") for k in args.K: k_folder = os.path.join(args.eval_folder_name, f"{k}_{sel_str}") for pat in [_go_pat, _gs_pat, _tr_pat, _ev_pat]: @@ -104,48 +192,105 @@ def main(): print(f"ERROR: {len(missing)} required evaluation file(s) not found:") for m in missing: print(f" {m}") - raise FileNotFoundError(f"{len(missing)} required evaluation file(s) missing. See list above.") + raise FileNotFoundError( + f"{len(missing)} required evaluation file(s) missing. See list above." + ) # Stability & Error - stats_SE = load_stablity_error_data(output_directory = f'{args.output_directory}/{args.run_name}', run_name = 'Inference', components = args.K, - stability_file = args.stability_file) - plot_stablity_error(stats = stats_SE,folder_name = args.save_folder_name, file_name = "Stability_Error") + stats_SE = load_stablity_error_data( + output_directory=f"{args.output_directory}/{args.run_name}", + run_name="Inference", + components=args.K, + stability_file=args.stability_file, + ) + plot_stablity_error( + stats=stats_SE, folder_name=args.save_folder_name, file_name="Stability_Error" + ) for sel_thresh in args.sel_threshs: - # Enrichement - count_df = load_enrichment_data(folder = args.eval_folder_name, components = args.K, sel_thresh = sel_thresh, - go_file=args.go_file, geneset_file=args.geneset_file, trait_file=args.trait_file, - term_col=args.term_col, adjpval_col=args.adjpval_col) - plot_enrichment(count_df,folder_name = args.save_folder_name, file_name = f"Enrichment_{sel_thresh}") + count_df = load_enrichment_data( + folder=args.eval_folder_name, + components=args.K, + sel_thresh=sel_thresh, + go_file=args.go_file, + geneset_file=args.geneset_file, + trait_file=args.trait_file, + term_col=args.term_col, + adjpval_col=args.adjpval_col, + ) + plot_enrichment( + count_df, + folder_name=args.save_folder_name, + file_name=f"Enrichment_{sel_thresh}", + ) # Perturbation - test_stats_df = load_perturbation_data(folder = args.eval_folder_name, components = args.K, sel_thresh = sel_thresh, - conditions = args.Conditions, pval = args.pval, perturbation_file=args.perturbation_file, - perturb_adjpval_col=args.perturb_adjpval_col, perturb_target_col=args.perturb_target_col, - perturb_log2fc_col=args.perturb_log2fc_col) - plot_perturbation(test_stats_df, folder_name = args.save_folder_name, pval=args.pval,file_name = f"Perturbation_{sel_thresh}") + test_stats_df = load_perturbation_data( + folder=args.eval_folder_name, + components=args.K, + sel_thresh=sel_thresh, + conditions=args.Conditions, + pval=args.pval, + perturbation_file=args.perturbation_file, + perturb_adjpval_col=args.perturb_adjpval_col, + perturb_target_col=args.perturb_target_col, + perturb_log2fc_col=args.perturb_log2fc_col, + ) + plot_perturbation( + test_stats_df, + folder_name=args.save_folder_name, + pval=args.pval, + file_name=f"Perturbation_{sel_thresh}", + ) # Explained Variance - stats_EV = load_explained_variance_data(folder = args.eval_folder_name, components=args.K, sel_thresh = sel_thresh, - variance_file=args.variance_file, variance_col=args.variance_col) - plot_explained_variance(stats_EV, folder_name = args.save_folder_name, file_name = f"Explained_Variance_{sel_thresh}") + stats_EV = load_explained_variance_data( + folder=args.eval_folder_name, + components=args.K, + sel_thresh=sel_thresh, + variance_file=args.variance_file, + variance_col=args.variance_col, + ) + plot_explained_variance( + stats_EV, + folder_name=args.save_folder_name, + file_name=f"Explained_Variance_{sel_thresh}", + ) # Motif (working in progress) - # K-selection panel - plot_k_selection_panel(stats_SE, count_df, test_stats_df, stats_EV, - pval=args.pval, folder_name= args.save_folder_name, file_name=f'K-selection_panel_{sel_thresh}', selected_k=args.selected_k) + # K-selection panel + plot_k_selection_panel( + stats_SE, + count_df, + test_stats_df, + stats_EV, + pval=args.pval, + folder_name=args.save_folder_name, + file_name=f"K-selection_panel_{sel_thresh}", + selected_k=args.selected_k, + ) # program dotplots (optional — requires inference h5mu per (K, sel_thresh)) if args.run_program_dotplot: for sel_thresh in args.sel_threshs: for k in args.K: - fig = programs_dotplots(k, args.output_directory, args.run_name, sel_thresh = sel_thresh, groupby=args.groupby, figsize=(4, 30), - show = False, save_name=f"Program_dotplot_{k}_{sel_thresh}", save_path = args.save_folder_name, ax = None) + fig = programs_dotplots( + k, + args.output_directory, + args.run_name, + sel_thresh=sel_thresh, + groupby=args.groupby, + figsize=(4, 30), + show=False, + save_name=f"Program_dotplot_{k}_{sel_thresh}", + save_path=args.save_folder_name, + ax=None, + ) return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_perturbed_gene_analysis.py b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_perturbed_gene_analysis.py index 7fc13e1..2bc2a97 100644 --- a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_perturbed_gene_analysis.py +++ b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_perturbed_gene_analysis.py @@ -10,87 +10,271 @@ from pathlib import Path # Change path to wherever you have repo locally -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) - -from Stage3_Interpretation.A_Plotting.src import plot_umap_per_gene, plot_top_program_per_gene, perturbed_gene_dotplot,\ - plot_log2FC, plot_volcano, programs_dotplot, analyze_correlations, \ - create_gene_correlation_waterfall, \ - convert_with_mygene, convert_adata_with_mygene, read_npz, \ - merge_pdfs_in_folder, merge_svgs_to_pdf, create_comprehensive_plot, rename_adata_gene_dictionary, \ - rename_list_gene_dictionary, plot_umap_per_gene_guide, process_single_gene, parallel_gene_processing,_process_gene_worker,\ - compute_gene_correlation_matrix, compute_gene_waterfall_cor,perturbed_program_dotplot, \ - plot_perturbation_vs_control, \ - export_gene_html, write_gene_share_index, \ - check_normalized, ensure_umap, condition_labels_from_obs +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) + +from Stage3_Interpretation.A_Plotting.src import ( + plot_umap_per_gene, + plot_top_program_per_gene, + perturbed_gene_dotplot, + plot_log2FC, + plot_volcano, + programs_dotplot, + analyze_correlations, + create_gene_correlation_waterfall, + convert_with_mygene, + convert_adata_with_mygene, + read_npz, + merge_pdfs_in_folder, + merge_svgs_to_pdf, + create_comprehensive_plot, + rename_adata_gene_dictionary, + rename_list_gene_dictionary, + plot_umap_per_gene_guide, + process_single_gene, + parallel_gene_processing, + _process_gene_worker, + compute_gene_correlation_matrix, + compute_gene_waterfall_cor, + perturbed_program_dotplot, + plot_perturbation_vs_control, + export_gene_html, + write_gene_share_index, + check_normalized, + ensure_umap, + condition_labels_from_obs, +) + def main(): parser = argparse.ArgumentParser() - #io path - parser.add_argument('--mdata_path', type=str, required=True, help='path to the MuData (.h5mu) file') - parser.add_argument('--perturb_path_base', type=str, default=None, help='base path for perturbation result files (sample suffix appended automatically). Omit to skip every per-condition perturbation panel (log2FC, volcano, program dotplot, waterfall) and plot only the h5mu-derived rows. Required for --output_format HTML.') - parser.add_argument('--ensembl_to_symbol_file', type=str, default=None, help='path to gene name mapping dictionary file for ID-to-name conversion') - parser.add_argument('--reference_gtf_path', type=str, default=None, help='path to reference GTF file for checking gene names') + # io path + parser.add_argument( + "--mdata_path", type=str, required=True, help="path to the MuData (.h5mu) file" + ) + parser.add_argument( + "--perturb_path_base", + type=str, + default=None, + help="base path for perturbation result files (sample suffix appended automatically). Omit to skip every per-condition perturbation panel (log2FC, volcano, program dotplot, waterfall) and plot only the h5mu-derived rows. Required for --output_format HTML.", + ) + parser.add_argument( + "--ensembl_to_symbol_file", + type=str, + default=None, + help="path to gene name mapping dictionary file for ID-to-name conversion", + ) + parser.add_argument( + "--reference_gtf_path", + type=str, + default=None, + help="path to reference GTF file for checking gene names", + ) # plotting variables - parser.add_argument('--perturb_target_col', type=str, default="target_name", help='column name for target genes in perturbation results') - parser.add_argument('--perturb_program_col', type=str, default="program_name", help='column name for programs in perturbation results') - parser.add_argument('--perturb_log2fc_col', type=str, default="log2FC", help='column name for log2 fold change values') - parser.add_argument('--top_corr_genes', type=int, default=5, help='number of top correlated genes to display per program') - parser.add_argument('--top_n_programs', type=int, default=10, help='number of top programs to display per gene') - parser.add_argument('--significance_threshold', type=float, default=0.05, help='p-value threshold for significance') - parser.add_argument('--volcano_log2fc_min', type=float, default=-0.00, help='lower log2FC threshold for volcano plot') - parser.add_argument('--volcano_log2fc_max', type=float, default=0.00, help='upper log2FC threshold for volcano plot') - parser.add_argument('--save_path', type=str, required=True, help='directory path to save output (PDF/SVG files or HTML share tree)') - parser.add_argument('--square_plots', action="store_true", help='use square aspect ratio for plots') - parser.add_argument('--figsize', type=float, nargs=2, default=(35, 35), help='figure size as width height') - parser.add_argument('--show', action="store_true", help='display plots interactively') - parser.add_argument('--output_format', type=str, default='SVG', choices=['PDF', 'SVG', 'HTML'], help='output format: PDF (matplotlib + PyPDF2 merge), SVG (matplotlib + svglib merge), HTML (interactive Plotly share folder)') - parser.add_argument('--n_processes', type=int, default=4, help='number of parallel processes for --parallel mode. Each worker holds a copy-on-write fork of mdata, so RAM scales with worker count: on large data (e.g. 1M cells x 30K genes), -1 (all cores) on a high-CPU node can hit 30-70 GB RSS. Default 4 keeps it bounded. Set -1 only when you know the per-worker RAM cost.') - parser.add_argument('--Conditions', nargs='*', type=str, default=None, help='Condition labels (values of obs[categorical_key]); default: all labels found in the data') - parser.add_argument('--umap_dot_size', type=int, default=10, help='dot size for UMAP plots') - parser.add_argument('--expressed_only', action="store_true", help='only plot perturbed genes found in the gene expression matrix (default: plot all perturbed genes)') - parser.add_argument('--gene_list_file', type=str, default=None, help='path to a file with one gene name per line to process (overrides automatic perturbed gene detection)') - parser.add_argument('--subsample_frac', type=float, default=None, help='fraction of cells to subsample for UMAP plots (e.g. 0.1 for 10%%). Default: None (plot all cells)') - parser.add_argument('--parallel', action="store_true", help='use fork-based multiprocessing to plot genes in parallel (Linux only)') - parser.add_argument('--corr_matrix_path', type=str, default=None, help='directory for the full correlation matrices written every run: regulator x regulator /regulator_corr_.npz and gene x gene /gene_loading_corr.npz. Defaults to --save_path.') - parser.add_argument('--no_save_regulator_corr', dest='save_regulator_corr', action='store_false', help='do not write /regulator_corr_.npz (~1.6 GB per sample at 20k regulators; the full matrix is built in RAM once to save it). Plots are unaffected.') - parser.add_argument('--no_save_gene_corr', dest='save_gene_corr', action='store_false', help='do not write /gene_loading_corr.npz (~0.4 GB at 10k genes, ~3.6 GB at 30k; the full matrix is built in RAM once to save it). Plots are unaffected.') - parser.add_argument('--skip_existing', action='store_false', help='[default on] skip genes whose output already exists. Pass --skip_existing to force re-process all.') + parser.add_argument( + "--perturb_target_col", + type=str, + default="target_name", + help="column name for target genes in perturbation results", + ) + parser.add_argument( + "--perturb_program_col", + type=str, + default="program_name", + help="column name for programs in perturbation results", + ) + parser.add_argument( + "--perturb_log2fc_col", + type=str, + default="log2FC", + help="column name for log2 fold change values", + ) + parser.add_argument( + "--top_corr_genes", + type=int, + default=5, + help="number of top correlated genes to display per program", + ) + parser.add_argument( + "--top_n_programs", + type=int, + default=10, + help="number of top programs to display per gene", + ) + parser.add_argument( + "--significance_threshold", + type=float, + default=0.05, + help="p-value threshold for significance", + ) + parser.add_argument( + "--volcano_log2fc_min", + type=float, + default=-0.00, + help="lower log2FC threshold for volcano plot", + ) + parser.add_argument( + "--volcano_log2fc_max", + type=float, + default=0.00, + help="upper log2FC threshold for volcano plot", + ) + parser.add_argument( + "--save_path", + type=str, + required=True, + help="directory path to save output (PDF/SVG files or HTML share tree)", + ) + parser.add_argument( + "--square_plots", action="store_true", help="use square aspect ratio for plots" + ) + parser.add_argument( + "--figsize", + type=float, + nargs=2, + default=(35, 35), + help="figure size as width height", + ) + parser.add_argument( + "--show", action="store_true", help="display plots interactively" + ) + parser.add_argument( + "--output_format", + type=str, + default="SVG", + choices=["PDF", "SVG", "HTML"], + help="output format: PDF (matplotlib + PyPDF2 merge), SVG (matplotlib + svglib merge), HTML (interactive Plotly share folder)", + ) + parser.add_argument( + "--n_processes", + type=int, + default=4, + help="number of parallel processes for --parallel mode. Each worker holds a copy-on-write fork of mdata, so RAM scales with worker count: on large data (e.g. 1M cells x 30K genes), -1 (all cores) on a high-CPU node can hit 30-70 GB RSS. Default 4 keeps it bounded. Set -1 only when you know the per-worker RAM cost.", + ) + parser.add_argument( + "--Conditions", + nargs="*", + type=str, + default=None, + help="Condition labels (values of obs[categorical_key]); default: all labels found in the data", + ) + parser.add_argument( + "--umap_dot_size", type=int, default=10, help="dot size for UMAP plots" + ) + parser.add_argument( + "--expressed_only", + action="store_true", + help="only plot perturbed genes found in the gene expression matrix (default: plot all perturbed genes)", + ) + parser.add_argument( + "--gene_list_file", + type=str, + default=None, + help="path to a file with one gene name per line to process (overrides automatic perturbed gene detection)", + ) + parser.add_argument( + "--subsample_frac", + type=float, + default=None, + help="fraction of cells to subsample for UMAP plots (e.g. 0.1 for 10%%). Default: None (plot all cells)", + ) + parser.add_argument( + "--parallel", + action="store_true", + help="use fork-based multiprocessing to plot genes in parallel (Linux only)", + ) + parser.add_argument( + "--corr_matrix_path", + type=str, + default=None, + help="directory for the full correlation matrices written every run: regulator x regulator /regulator_corr_.npz and gene x gene /gene_loading_corr.npz. Defaults to --save_path.", + ) + parser.add_argument( + "--no_save_regulator_corr", + dest="save_regulator_corr", + action="store_false", + help="do not write /regulator_corr_.npz (~1.6 GB per sample at 20k regulators; the full matrix is built in RAM once to save it). Plots are unaffected.", + ) + parser.add_argument( + "--no_save_gene_corr", + dest="save_gene_corr", + action="store_false", + help="do not write /gene_loading_corr.npz (~0.4 GB at 10k genes, ~3.6 GB at 30k; the full matrix is built in RAM once to save it). Plots are unaffected.", + ) + parser.add_argument( + "--skip_existing", + action="store_false", + help="[default on] skip genes whose output already exists. Pass --skip_existing to force re-process all.", + ) # keys - parser.add_argument('--data_key', type=str, default="rna", help='key to access gene expression data in MuData') - parser.add_argument('--prog_key', type=str, default="cNMF", help='key to access cNMF programs in MuData') - parser.add_argument('--gene_name_key', type=str, default="gene_names", help='key to access gene names in var') - parser.add_argument('--categorical_key', type=str, default="sample", help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample") - parser.add_argument('--guide_targets_key', type=str, default="guide_targets", help='key in .uns to access guide target genes (default: guide_targets)') - parser.add_argument('--guide_assignment_key', type=str, default="guide_assignment", help='key in .obsm to access the guide-assignment matrix (default: guide_assignment)') - parser.add_argument('--control_target_name', type=str, nargs='+', default=["non-targeting"], help='one or more control labels in guide_targets (e.g. non-targeting, or WT WT111 WT4). A cell is a control if its guide target matches any of these.') - - + parser.add_argument( + "--data_key", + type=str, + default="rna", + help="key to access gene expression data in MuData", + ) + parser.add_argument( + "--prog_key", + type=str, + default="cNMF", + help="key to access cNMF programs in MuData", + ) + parser.add_argument( + "--gene_name_key", + type=str, + default="gene_names", + help="key to access gene names in var", + ) + parser.add_argument( + "--categorical_key", + type=str, + default="sample", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + ) + parser.add_argument( + "--guide_targets_key", + type=str, + default="guide_targets", + help="key in .uns to access guide target genes (default: guide_targets)", + ) + parser.add_argument( + "--guide_assignment_key", + type=str, + default="guide_assignment", + help="key in .obsm to access the guide-assignment matrix (default: guide_assignment)", + ) + parser.add_argument( + "--control_target_name", + type=str, + nargs="+", + default=["non-targeting"], + help="one or more control labels in guide_targets (e.g. non-targeting, or WT WT111 WT4). A cell is a control if its guide target matches any of these.", + ) + args = parser.parse_args() # export_gene_html renders the perturbation sections unconditionally, so the # HTML report cannot be produced without the per-sample association files. - if args.output_format == 'HTML' and args.perturb_path_base is None: - parser.error("--output_format HTML requires --perturb_path_base " - "(the HTML export always renders the perturbation sections). " - "Use --output_format PDF or SVG to plot without perturbation results.") + if args.output_format == "HTML" and args.perturb_path_base is None: + parser.error( + "--output_format HTML requires --perturb_path_base " + "(the HTML export always renders the perturbation sections). " + "Use --output_format PDF or SVG to plot without perturbation results." + ) # save comfigs used args_dict = vars(args) - job_id = os.environ.get('SLURM_JOB_ID') - os.makedirs(f'{args.save_path}', exist_ok=True) - with open(f'{args.save_path}/config_{job_id}.yml', 'w') as f: + job_id = os.environ.get("SLURM_JOB_ID") + os.makedirs(f"{args.save_path}", exist_ok=True) + with open(f"{args.save_path}/config_{job_id}.yml", "w") as f: yaml.dump(args_dict, f, default_flow_style=False, width=1000) - - #read mdata + # read mdata mdata = mu.read_h5mu(args.mdata_path) - - # validate that expression matrix is normalized (TPM/CPM/log), not raw counts. # Downstream UMAP/PCA + per-gene expression plots assume normalized values. check_normalized(mdata[args.data_key], args.data_key) @@ -100,9 +284,12 @@ def main(): # condition labels default to the values of --categorical_key in the data if args.Conditions is None: - args.Conditions = condition_labels_from_obs(mdata[args.data_key].obs, args.categorical_key) - print(f"--Conditions not given; using {args.categorical_key} values from the h5mu: {args.Conditions}") - + args.Conditions = condition_labels_from_obs( + mdata[args.data_key].obs, args.categorical_key + ) + print( + f"--Conditions not given; using {args.categorical_key} values from the h5mu: {args.Conditions}" + ) # found detected perturbed gene (use gene symbols from var column when var_names are Ensembl IDs) perturbed_gene = np.unique(mdata[args.prog_key].uns[args.guide_targets_key]) @@ -111,25 +298,33 @@ def main(): else: gene_symbols = mdata[args.data_key].var_names.tolist() - # print out how many perturbation is not found expressed + # print out how many perturbation is not found expressed perturbed_gene_found = sorted(set(gene_symbols) & set(perturbed_gene.tolist())) perturbed_gene_not_found = sorted(set(perturbed_gene.tolist()) - set(gene_symbols)) - print(f"there are {len(perturbed_gene_found)} perturbed genes found in expression matrix") - print(f"there are {len(perturbed_gene_not_found)} perturbed genes NOT found in expression matrix: {perturbed_gene_not_found}") + print( + f"there are {len(perturbed_gene_found)} perturbed genes found in expression matrix" + ) + print( + f"there are {len(perturbed_gene_not_found)} perturbed genes NOT found in expression matrix: {perturbed_gene_not_found}" + ) - # decide which genes are being ploted - if args.gene_list_file is not None: # a given list of genes - with open(args.gene_list_file, 'r') as f: + # decide which genes are being ploted + if args.gene_list_file is not None: # a given list of genes + with open(args.gene_list_file, "r") as f: genes_requested = sorted([line.strip() for line in f if line.strip()]) genes_valid = sorted(set(genes_requested) & set(gene_symbols)) genes_missing = sorted(set(genes_requested) - set(gene_symbols)) if genes_missing: - print(f"WARNING: {len(genes_missing)} genes from {args.gene_list_file} not found in expression matrix: {genes_missing}") + print( + f"WARNING: {len(genes_missing)} genes from {args.gene_list_file} not found in expression matrix: {genes_missing}" + ) genes_to_plot = genes_valid - print(f"Using {len(genes_to_plot)}/{len(genes_requested)} genes from {args.gene_list_file}") - elif args.expressed_only: # only expressed + print( + f"Using {len(genes_to_plot)}/{len(genes_requested)} genes from {args.gene_list_file}" + ) + elif args.expressed_only: # only expressed genes_to_plot = perturbed_gene_found - else: # all genes + else: # all genes genes_to_plot = sorted(perturbed_gene.tolist()) # The control target is never tested in the perturbation-association results, @@ -138,46 +333,58 @@ def main(): present = control_set & set(genes_to_plot) if present: genes_to_plot = [g for g in genes_to_plot if g not in control_set] - print(f"Excluding control target(s) {sorted(present)} from perturbed-gene plots") + print( + f"Excluding control target(s) {sorted(present)} from perturbed-gene plots" + ) # Skip-existing support: build the set of genes that actually need processing, # but keep `genes_to_plot` as the full ordered list so HTML nav/index stay correct. if args.skip_existing: - if args.output_format == 'HTML': - done = {p.parent.name[len('gene_'):] - for p in Path(args.save_path).glob('gene_*/metadata.json')} + if args.output_format == "HTML": + done = { + p.parent.name[len("gene_") :] + for p in Path(args.save_path).glob("gene_*/metadata.json") + } else: - ext = '.pdf' if args.output_format == 'PDF' else '.svg' - done = {p.stem for p in Path(args.save_path).glob(f'*{ext}')} - + ext = ".pdf" if args.output_format == "PDF" else ".svg" + done = {p.stem for p in Path(args.save_path).glob(f"*{ext}")} + process_set = {g for g in genes_to_plot if g not in done} skipped = len(genes_to_plot) - len(process_set) if skipped: - print(f"Skip-existing: skipping {skipped} already-produced gene(s); {len(process_set)} remaining.") + print( + f"Skip-existing: skipping {skipped} already-produced gene(s); {len(process_set)} remaining." + ) else: process_set = set(genes_to_plot) - - corr_dir = args.corr_matrix_path or args.save_path # The waterfall correlation is built from the perturbation files; without them # the per-condition rows are skipped and nothing needs to be computed. if args.perturb_path_base is None: waterfall_correlation = None - print("No --perturb_path_base given: skipping per-condition perturbation panels " - "(log2FC, volcano, program dotplot, waterfall).") + print( + "No --perturb_path_base given: skipping per-condition perturbation panels " + "(log2FC, volcano, program dotplot, waterfall)." + ) else: waterfall_correlation = {} for samp in args.Conditions: - save = f"{corr_dir}/regulator_corr_{samp}.npz" if args.save_regulator_corr else None + save = ( + f"{corr_dir}/regulator_corr_{samp}.npz" + if args.save_regulator_corr + else None + ) waterfall_correlation[samp] = compute_gene_waterfall_cor( f"{args.perturb_path_base}_{samp}.txt", perturb_log2fc_col=args.perturb_log2fc_col, save_path=save, ) if save is not None: - print(f"Wrote regulator x regulator correlation ({len(waterfall_correlation[samp].index)} regulators): {save}") + print( + f"Wrote regulator x regulator correlation ({len(waterfall_correlation[samp].index)} regulators): {save}" + ) loading_corr = f"{corr_dir}/gene_loading_corr.npz" if args.save_gene_corr else None correlation_matrix = compute_gene_correlation_matrix( @@ -189,21 +396,22 @@ def main(): prog_key=args.prog_key, ) if loading_corr is not None: - print(f"Wrote gene x gene loading correlation ({len(correlation_matrix.columns)} genes): {loading_corr}") - + print( + f"Wrote gene x gene loading correlation ({len(correlation_matrix.columns)} genes): {loading_corr}" + ) # Graph all genes - if args.output_format == 'HTML': + if args.output_format == "HTML": n_genes = len(genes_to_plot) for i, gene in enumerate(genes_to_plot): if gene not in process_set: continue - # enable pre/after browsing - prev_g = genes_to_plot[i - 1] if i > 0 else None + # enable pre/after browsing + prev_g = genes_to_plot[i - 1] if i > 0 else None next_g = genes_to_plot[i + 1] if i + 1 < n_genes else None - # make html per gene + # make html per gene export_gene_html( mdata=mdata, perturb_path_base=args.perturb_path_base, @@ -235,7 +443,7 @@ def main(): guide_targets_key=args.guide_targets_key, guide_assignment_key=args.guide_assignment_key, ) - plt.close('all') + plt.close("all") if (i + 1) % 20 == 0: gc.collect() @@ -265,7 +473,7 @@ def main(): sample=args.Conditions, square_plots=args.square_plots, show=args.show, - PDF=(args.output_format == 'PDF'), + PDF=(args.output_format == "PDF"), n_processes=args.n_processes, gene_name_key=args.gene_name_key, umap_dot_size=args.umap_dot_size, @@ -274,16 +482,17 @@ def main(): data_key=args.data_key, prog_key=args.prog_key, guide_targets_key=args.guide_targets_key, - guide_assignment_key=args.guide_assignment_key + guide_assignment_key=args.guide_assignment_key, + ) + print( + f"Parallel processing completed. Results: {len(result) if result else 'None'}" ) - print(f"Parallel processing completed. Results: {len(result) if result else 'None'}") except Exception as e: print(f"ERROR in parallel_gene_processing: {e}") else: for i, gene in enumerate(genes_to_plot): - if gene not in process_set: continue @@ -309,7 +518,7 @@ def main(): sample=args.Conditions, square_plots=args.square_plots, show=args.show, - PDF=(args.output_format == 'PDF'), + PDF=(args.output_format == "PDF"), umap_dot_size=args.umap_dot_size, umap_subsample_frac=args.subsample_frac, gene_name_key=args.gene_name_key, @@ -317,22 +526,21 @@ def main(): data_key=args.data_key, prog_key=args.prog_key, guide_targets_key=args.guide_targets_key, - guide_assignment_key=args.guide_assignment_key + guide_assignment_key=args.guide_assignment_key, ) - plt.close('all') + plt.close("all") if (i + 1) % 20 == 0: gc.collect() # post-loop assembly - if args.output_format == 'PDF': + if args.output_format == "PDF": merge_pdfs_in_folder(args.save_path, output_filename="gene.pdf") - elif args.output_format == 'SVG': + elif args.output_format == "SVG": merge_svgs_to_pdf(args.save_path, output_filename="gene.pdf") # no HTML, index already written inside the HTML branch above return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) - diff --git a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py index 4d950c9..5324231 100644 --- a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py +++ b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py @@ -8,76 +8,215 @@ from pathlib import Path - # Change path to wherever you have repo locally -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) - -from Stage3_Interpretation.A_Plotting.src import merge_pdfs_in_folder, merge_svgs_to_pdf, ensure_umap, condition_labels_from_obs - -from Stage3_Interpretation.A_Plotting.src import plot_umap_per_program, plot_top_gene_per_program, top_GO_per_program, compute_program_correlation_matrix,\ - analyze_program_correlations, plot_violin, plot_program_log2FC, plot_program_heatmap, plot_program_volcano, \ - perturbed_program_dotplot, compute_program_waterfall_cor, create_program_correlation_waterfall, create_comprehensive_program_plot, \ - export_program_html, write_share_index +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) + +from Stage3_Interpretation.A_Plotting.src import ( + merge_pdfs_in_folder, + merge_svgs_to_pdf, + ensure_umap, + condition_labels_from_obs, +) + +from Stage3_Interpretation.A_Plotting.src import ( + plot_umap_per_program, + plot_top_gene_per_program, + top_GO_per_program, + compute_program_correlation_matrix, + analyze_program_correlations, + plot_violin, + plot_program_log2FC, + plot_program_heatmap, + plot_program_volcano, + perturbed_program_dotplot, + compute_program_waterfall_cor, + create_program_correlation_waterfall, + create_comprehensive_program_plot, + export_program_html, + write_share_index, +) def main(): parser = argparse.ArgumentParser() - #io path - parser.add_argument('--mdata_path', type=str, required=True, help='path to the MuData (.h5mu) file') - parser.add_argument('--perturb_path_base', type=str, default=None, help='base path for perturbation result files (sample suffix appended automatically). Omit to skip every per-condition perturbation panel (log2FC, volcano, regulator dotplot, waterfall) plus the regulator heatmap, and plot only the h5mu/GO-derived header row. Required for --output_format HTML.') - parser.add_argument('--file_to_dictionary', type=str, default=None, help='path to gene name mapping dictionary file for ID-to-name conversion') - parser.add_argument('--reference_gtf_path', type=str, default=None, help='path to reference GTF file for checking gene names') - parser.add_argument('--GO_path', type=str, required=True, help='path to Gene Ontology enrichment results directory') + # io path + parser.add_argument( + "--mdata_path", type=str, required=True, help="path to the MuData (.h5mu) file" + ) + parser.add_argument( + "--perturb_path_base", + type=str, + default=None, + help="base path for perturbation result files (sample suffix appended automatically). Omit to skip every per-condition perturbation panel (log2FC, volcano, regulator dotplot, waterfall) plus the regulator heatmap, and plot only the h5mu/GO-derived header row. Required for --output_format HTML.", + ) + parser.add_argument( + "--file_to_dictionary", + type=str, + default=None, + help="path to gene name mapping dictionary file for ID-to-name conversion", + ) + parser.add_argument( + "--reference_gtf_path", + type=str, + default=None, + help="path to reference GTF file for checking gene names", + ) + parser.add_argument( + "--GO_path", + type=str, + required=True, + help="path to Gene Ontology enrichment results directory", + ) # plotting variables - parser.add_argument('--tagert_col_name', type=str, default="program_name", help='column name for target programs in perturbation results') - parser.add_argument('--plot_col_name', type=str, default="target_name", help='column name for genes in perturbation results') - parser.add_argument('--log2fc_col', type=str, default="log2FC", help='column name for log2 fold change values') - parser.add_argument('--top_program', type=int, default=10, help='number of top programs to display') - parser.add_argument('--top_enrichned_term', type=int, default=10, help='number of top GO enrichment terms to display per program') - parser.add_argument('--p_value', type=float, default=0.05, help='p-value threshold for significance') - parser.add_argument('--down_thred_log', type=float, default=-0.00, help='lower log2FC threshold for volcano plot') - parser.add_argument('--up_thred_log', type=float, default=0.00, help='upper log2FC threshold for volcano plot') - parser.add_argument('--save_path', type=str, required=True, help='directory path to save output (PDF/SVG files or HTML share tree)') - parser.add_argument('--square_plots', action="store_true", help='use square aspect ratio for plots') - parser.add_argument('--figsize', type=float, nargs=2, default=(35, 35), help='figure size as width height') - parser.add_argument('--show', action="store_true", help='display plots interactively') - parser.add_argument('--output_format', type=str, default='SVG', choices=['PDF', 'SVG', 'HTML'], help='output format: PDF (matplotlib + PyPDF2 merge), SVG (matplotlib + svglib merge), HTML (interactive Plotly share folder)') - parser.add_argument('--Conditions', nargs='*', type=str, default=None, help='Condition labels (values of obs[categorical_key]); default: all labels found in the data') - parser.add_argument('--programs', nargs='+', type=int, default=None, help='specific program numbers to plot (e.g. 4 5 6 ... 100). If omitted, all programs are plotted.') - parser.add_argument('--subsample_frac', type=float, default=None, help='fraction of cells to subsample for UMAP plots (e.g. 0.1 for 10%%). Default: None (plot all cells)') - parser.add_argument('--corr_matrix_path', type=str, default=None, help='directory for the full program x program correlation matrices written every run: /program_corr.npz (usage correlation) and /program_waterfall_corr_.npz (perturbation-effect correlation). Defaults to --save_path.') - parser.add_argument('--skip_existing', action='store_false', help='[default on] skip programs whose output already exists. Pass --skip_existing to force re-process all.') + parser.add_argument( + "--tagert_col_name", + type=str, + default="program_name", + help="column name for target programs in perturbation results", + ) + parser.add_argument( + "--plot_col_name", + type=str, + default="target_name", + help="column name for genes in perturbation results", + ) + parser.add_argument( + "--log2fc_col", + type=str, + default="log2FC", + help="column name for log2 fold change values", + ) + parser.add_argument( + "--top_program", type=int, default=10, help="number of top programs to display" + ) + parser.add_argument( + "--top_enrichned_term", + type=int, + default=10, + help="number of top GO enrichment terms to display per program", + ) + parser.add_argument( + "--p_value", type=float, default=0.05, help="p-value threshold for significance" + ) + parser.add_argument( + "--down_thred_log", + type=float, + default=-0.00, + help="lower log2FC threshold for volcano plot", + ) + parser.add_argument( + "--up_thred_log", + type=float, + default=0.00, + help="upper log2FC threshold for volcano plot", + ) + parser.add_argument( + "--save_path", + type=str, + required=True, + help="directory path to save output (PDF/SVG files or HTML share tree)", + ) + parser.add_argument( + "--square_plots", action="store_true", help="use square aspect ratio for plots" + ) + parser.add_argument( + "--figsize", + type=float, + nargs=2, + default=(35, 35), + help="figure size as width height", + ) + parser.add_argument( + "--show", action="store_true", help="display plots interactively" + ) + parser.add_argument( + "--output_format", + type=str, + default="SVG", + choices=["PDF", "SVG", "HTML"], + help="output format: PDF (matplotlib + PyPDF2 merge), SVG (matplotlib + svglib merge), HTML (interactive Plotly share folder)", + ) + parser.add_argument( + "--Conditions", + nargs="*", + type=str, + default=None, + help="Condition labels (values of obs[categorical_key]); default: all labels found in the data", + ) + parser.add_argument( + "--programs", + nargs="+", + type=int, + default=None, + help="specific program numbers to plot (e.g. 4 5 6 ... 100). If omitted, all programs are plotted.", + ) + parser.add_argument( + "--subsample_frac", + type=float, + default=None, + help="fraction of cells to subsample for UMAP plots (e.g. 0.1 for 10%%). Default: None (plot all cells)", + ) + parser.add_argument( + "--corr_matrix_path", + type=str, + default=None, + help="directory for the full program x program correlation matrices written every run: /program_corr.npz (usage correlation) and /program_waterfall_corr_.npz (perturbation-effect correlation). Defaults to --save_path.", + ) + parser.add_argument( + "--skip_existing", + action="store_false", + help="[default on] skip programs whose output already exists. Pass --skip_existing to force re-process all.", + ) # keys - parser.add_argument('--data_key', type=str, default="rna", help='key to access gene expression data in MuData') - parser.add_argument('--prog_key', type=str, default="cNMF", help='key to access cNMF programs in MuData') - parser.add_argument('--gene_name_key', type=str, default="gene_names", help='key to access gene names in var') - parser.add_argument('--categorical_key', type=str, default="sample", help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample") - - + parser.add_argument( + "--data_key", + type=str, + default="rna", + help="key to access gene expression data in MuData", + ) + parser.add_argument( + "--prog_key", + type=str, + default="cNMF", + help="key to access cNMF programs in MuData", + ) + parser.add_argument( + "--gene_name_key", + type=str, + default="gene_names", + help="key to access gene names in var", + ) + parser.add_argument( + "--categorical_key", + type=str, + default="sample", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + ) args = parser.parse_args() # export_program_html renders the per-sample perturbation sections and the # heatmap unconditionally, so the HTML report cannot be produced without the # per-sample association files. - if args.output_format == 'HTML' and args.perturb_path_base is None: - parser.error("--output_format HTML requires --perturb_path_base " - "(the HTML export always renders the perturbation sections). " - "Use --output_format PDF or SVG to plot without perturbation results.") + if args.output_format == "HTML" and args.perturb_path_base is None: + parser.error( + "--output_format HTML requires --perturb_path_base " + "(the HTML export always renders the perturbation sections). " + "Use --output_format PDF or SVG to plot without perturbation results." + ) # save comfigs used args_dict = vars(args) - job_id = os.environ.get('SLURM_JOB_ID') - os.makedirs(f'{args.save_path}', exist_ok=True) - with open(f'{args.save_path}/config_{job_id}.yml', 'w') as f: + job_id = os.environ.get("SLURM_JOB_ID") + os.makedirs(f"{args.save_path}", exist_ok=True) + with open(f"{args.save_path}/config_{job_id}.yml", "w") as f: yaml.dump(args_dict, f, default_flow_style=False, width=1000) - - #read mdata + # read mdata mdata = mu.read_h5mu(args.mdata_path) # compute UMAP/PCA from top-variance genes if missing @@ -85,13 +224,16 @@ def main(): # condition labels default to the values of --categorical_key in the data if args.Conditions is None: - args.Conditions = condition_labels_from_obs(mdata[args.data_key].obs, args.categorical_key) - print(f"--Conditions not given; using {args.categorical_key} values from the h5mu: {args.Conditions}") - - program_len = len(mdata[args.prog_key].var) # find out list of programs to process + args.Conditions = condition_labels_from_obs( + mdata[args.data_key].obs, args.categorical_key + ) + print( + f"--Conditions not given; using {args.categorical_key} values from the h5mu: {args.Conditions}" + ) + + program_len = len(mdata[args.prog_key].var) # find out list of programs to process print(f"there are {program_len} Program found") - # found detected perturbed gene (use gene symbols from var column when var_names are Ensembl IDs) perturbed_gene = np.unique(mdata[args.prog_key].uns["guide_targets"]) if args.gene_name_key in mdata[args.data_key].var.columns: @@ -100,9 +242,6 @@ def main(): gene_symbols = mdata[args.data_key].var_names.tolist() perturbed_gene_found = sorted(set(gene_symbols) & set(perturbed_gene.tolist())) - - - # compute correlations (both K x K; full matrices written to corr_dir every run) corr_dir = args.corr_matrix_path or args.save_path @@ -110,48 +249,60 @@ def main(): # the per-condition rows and the heatmap are skipped and nothing needs computing. if args.perturb_path_base is None: waterfall_correlation = None - print("No --perturb_path_base given: skipping per-condition perturbation panels " - "(log2FC, volcano, regulator dotplot, waterfall) and the regulator heatmap.") + print( + "No --perturb_path_base given: skipping per-condition perturbation panels " + "(log2FC, volcano, regulator dotplot, waterfall) and the regulator heatmap." + ) else: waterfall_correlation = {} for samp in args.Conditions: save = f"{corr_dir}/program_waterfall_corr_{samp}.npz" waterfall_correlation[samp] = compute_program_waterfall_cor( - f"{args.perturb_path_base}_{samp}.txt", save_path=save, log2fc_col=args.log2fc_col) + f"{args.perturb_path_base}_{samp}.txt", + save_path=save, + log2fc_col=args.log2fc_col, + ) print(f"Wrote program x program perturbation correlation: {save}") program_corr_path = f"{corr_dir}/program_corr.npz" program_correlation = compute_program_correlation_matrix( - mdata, prog_key=args.prog_key, save_path=program_corr_path) + mdata, prog_key=args.prog_key, save_path=program_corr_path + ) print(f"Wrote program x program usage correlation: {program_corr_path}") - - - programs_to_plot = args.programs if args.programs is not None else list(mdata[args.prog_key].var_names) + + programs_to_plot = ( + args.programs + if args.programs is not None + else list(mdata[args.prog_key].var_names) + ) # Skip-existing support: build the set of programs that actually need processing, # but keep `programs_to_plot` as the full ordered list so HTML nav/index stay correct. if args.skip_existing: - if args.output_format == 'HTML': - done = {p.parent.name[len('program_'):] - for p in Path(args.save_path).glob('program_*/metadata.json')} + if args.output_format == "HTML": + done = { + p.parent.name[len("program_") :] + for p in Path(args.save_path).glob("program_*/metadata.json") + } else: - ext = '.pdf' if args.output_format == 'PDF' else '.svg' - done = {p.stem for p in Path(args.save_path).glob(f'*{ext}')} + ext = ".pdf" if args.output_format == "PDF" else ".svg" + done = {p.stem for p in Path(args.save_path).glob(f"*{ext}")} process_set = {str(p) for p in programs_to_plot if str(p) not in done} skipped = len(programs_to_plot) - len(process_set) if skipped: - print(f"Skip-existing: skipping {skipped} already-produced program(s); {len(process_set)} remaining.") + print( + f"Skip-existing: skipping {skipped} already-produced program(s); {len(process_set)} remaining." + ) else: process_set = {str(p) for p in programs_to_plot} programs_str = [str(p) for p in programs_to_plot] n_progs = len(programs_str) for i, program in enumerate(programs_to_plot): - if str(program) not in process_set: continue - if args.output_format == 'HTML': + if args.output_format == "HTML": prev_pid = programs_str[i - 1] if i > 0 else None next_pid = programs_str[i + 1] if i + 1 < n_progs else None export_program_html( @@ -207,19 +358,18 @@ def main(): sample=args.Conditions, square_plots=args.square_plots, show=args.show, - PDF=(args.output_format == 'PDF'), + PDF=(args.output_format == "PDF"), gene_list=perturbed_gene_found, subsample_frac=args.subsample_frac, gene_name_key=args.gene_name_key, data_key=args.data_key, - prog_key=args.prog_key + prog_key=args.prog_key, ) - # post-loop assembly - if args.output_format == 'PDF': + if args.output_format == "PDF": merge_pdfs_in_folder(args.save_path, output_filename="program.pdf") - elif args.output_format == 'SVG': + elif args.output_format == "SVG": merge_svgs_to_pdf(args.save_path) else: # HTML write_share_index(args.save_path, programs_to_plot, args_dict) @@ -227,6 +377,5 @@ def main(): return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) - diff --git a/src/Stage3_Interpretation/A_Plotting/src/Perturbed_gene_QC_plots.py b/src/Stage3_Interpretation/A_Plotting/src/Perturbed_gene_QC_plots.py index 41d5d0b..f1c4cfd 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/Perturbed_gene_QC_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/Perturbed_gene_QC_plots.py @@ -1,5 +1,4 @@ - -import muon as mu +import muon as mu import scanpy as sc import numpy as np from scipy.stats import pearsonr @@ -29,30 +28,39 @@ plt.rcParams["axes.spines.top"] = False plt.rcParams["axes.spines.right"] = False -#plt.rcParams["axes.spines.bottom"] = False -#plt.rcParams["axes.spines.left"] = False -plt.rcParams['pdf.fonttype'] = 42 # TrueType fonts (editable text) -plt.rcParams['ps.fonttype'] = 42 # For EPS as well -plt.rcParams['svg.fonttype'] = 'none' # SVG: write real (editable in Illustrator) - - - -plt.rcParams.update({ - 'font.size': 16, - 'axes.titlesize': 16, - 'axes.labelsize': 16, - 'legend.fontsize': 16, # Legend text size - 'legend.title_fontsize': 16, # Legend title size - 'xtick.labelsize': 16, - 'ytick.labelsize': 16, - 'figure.titlesize': 20 -}) +# plt.rcParams["axes.spines.bottom"] = False +# plt.rcParams["axes.spines.left"] = False +plt.rcParams["pdf.fonttype"] = 42 # TrueType fonts (editable text) +plt.rcParams["ps.fonttype"] = 42 # For EPS as well +plt.rcParams["svg.fonttype"] = ( + "none" # SVG: write real (editable in Illustrator) +) + + +plt.rcParams.update( + { + "font.size": 16, + "axes.titlesize": 16, + "axes.labelsize": 16, + "legend.fontsize": 16, # Legend text size + "legend.title_fontsize": 16, # Legend title size + "xtick.labelsize": 16, + "ytick.labelsize": 16, + "figure.titlesize": 20, + } +) import sys -# Change path to wherever you have repo locally -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) -from .utilities import convert_adata_with_mygene, convert_with_mygene, rename_list_gene_dictionary, rename_adata_gene_dictionary +# Change path to wherever you have repo locally +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) + +from .utilities import ( + convert_adata_with_mygene, + convert_with_mygene, + rename_list_gene_dictionary, + rename_adata_gene_dictionary, +) from .efficient_corr import RowCorr @@ -62,16 +70,23 @@ def _blank_ax(ax, gene, message=None): Hides spines, ticks, and tick labels, then centres a message. """ if message is None: - message = f'No Guide Found\nfor {gene}' + message = f"No Guide Found\nfor {gene}" for spine in ax.spines.values(): spine.set_visible(False) ax.set_xticks([]) ax.set_yticks([]) ax.set_xticklabels([]) ax.set_yticklabels([]) - ax.text(0.5, 0.5, message, - ha='center', va='center', transform=ax.transAxes, - fontsize=12, color='grey') + ax.text( + 0.5, + 0.5, + message, + ha="center", + va="center", + transform=ax.transAxes, + fontsize=12, + color="grey", + ) return ax @@ -86,10 +101,22 @@ def _safe_name(name): return "".join(c if c.isalnum() or c in "_-" else "_" for c in str(name)) - -def plot_umap_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None, ax=None, - color='purple', save_path=None, save_name=None, figsize=(8,6), show=False, size = None, - umap_subsample_frac=None, random_state=42, gene_name_key='gene_names', data_key='rna'): +def plot_umap_per_gene( + mdata, + Target_Gene, + ensembl_to_symbol_file=None, + ax=None, + color="purple", + save_path=None, + save_name=None, + figsize=(8, 6), + show=False, + size=None, + umap_subsample_frac=None, + random_state=42, + gene_name_key="gene_names", + data_key="rna", +): """Plot gene expression on UMAP embedding. Colors cells by expression level of a single gene on a precomputed UMAP. @@ -164,24 +191,29 @@ def plot_umap_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None, ax=Non img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No UMAP to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No UMAP to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax - - + # Check if gene exists gene_name_list = renamed.var_names.tolist() if Target_Gene not in gene_name_list: print("gene name is not found in adata") return None - + # Set color - colors = ['lightgrey', color] + colors = ["lightgrey", color] n_bins = 100 - cmap = mcolors.LinearSegmentedColormap.from_list('custom', colors, N=n_bins) - - title = f'{Target_Gene} Expression' - + cmap = mcolors.LinearSegmentedColormap.from_list("custom", colors, N=n_bins) + + title = f"{Target_Gene} Expression" + # If no axis provided, create new figure if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -189,44 +221,71 @@ def plot_umap_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None, ax=Non else: fig = ax.get_figure() standalone = False - + # Plot on the provided/created axis plt.sca(ax) # Set current axis to ensure scanpy uses correct figure -CC change # use_raw=False: var_names were renamed to symbols on .X only, not on .raw # (which keeps Ensembl IDs). scanpy defaults use_raw=True when .raw exists, so # it would look Target_Gene up in .raw and raise KeyError. Force it to read .X, # matching the existence check against renamed.var_names above. - sc.pl.umap(renamed, color=Target_Gene, title=title, cmap=cmap, ax=ax, show=False, size = size, use_raw=False) - + sc.pl.umap( + renamed, + color=Target_Gene, + title=title, + cmap=cmap, + ax=ax, + show=False, + size=size, + use_raw=False, + ) + # Rasterize ONLY the scatter points (collections) in this axis for collection in ax.collections: collection.set_rasterized(True) - - ax.set_title(title, fontsize=18, fontweight='bold', loc='center') - ax.set_xlabel('UMAP 1', fontsize=10, fontweight = 'bold') - ax.set_ylabel('UMAP 2', fontsize=10, fontweight = 'bold') + + ax.set_title(title, fontsize=18, fontweight="bold", loc="center") + ax.set_xlabel("UMAP 1", fontsize=10, fontweight="bold") + ax.set_ylabel("UMAP 2", fontsize=10, fontweight="bold") if ax is None: plt.tight_layout() # Only save if standalone and save_path provided if standalone and save_path and save_name: - fig.savefig(f"{save_path}/{_safe_name(save_name)}.svg", format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) + # Only show/close if standalone if standalone: if show: plt.show() else: plt.close(fig) - + return ax -def plot_umap_per_gene_guide(mdata, Target_Gene, ax=None, color='red', save_path=None, -save_name=None, figsize=(8,6), show=False, size = None, -umap_subsample_frac=None, random_state=42, data_key='rna', prog_key='cNMF', -guide_targets_key='guide_targets', guide_assignment_key='guide_assignment'): +def plot_umap_per_gene_guide( + mdata, + Target_Gene, + ax=None, + color="red", + save_path=None, + save_name=None, + figsize=(8, 6), + show=False, + size=None, + umap_subsample_frac=None, + random_state=42, + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): """Plot guide perturbation signal on UMAP embedding. Colors cells by the sum of guide-assignment counts targeting a specific gene, @@ -279,20 +338,30 @@ def plot_umap_per_gene_guide(mdata, Target_Gene, ax=None, color='red', save_path img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No UMAP to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No UMAP to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax guide_mat = mdata[prog_key].obsm[guide_assignment_key] # (cells x guides), sparse - target_counts = np.asarray(guide_mat[:, target_mask].sum(axis=1)).flatten() # collapse guide all guides for target gene only + target_counts = np.asarray( + guide_mat[:, target_mask].sum(axis=1) + ).flatten() # collapse guide all guides for target gene only - adata = ad.AnnData(X=target_counts.reshape(-1, 1), obs=mdata[data_key].obs) # make adata with only one column cell x target_gene_guide + adata = ad.AnnData( + X=target_counts.reshape(-1, 1), obs=mdata[data_key].obs + ) # make adata with only one column cell x target_gene_guide adata.var_names = pd.Index([Target_Gene]) # X_pca is not used for the UMAP scatter below; copy only if present so h5mu # files that ship X_umap without a matching X_pca don't crash here. - if 'X_pca' in mdata[data_key].obsm: - adata.obsm['X_pca'] = mdata[data_key].obsm['X_pca'] - adata.obsm['X_umap'] = mdata[data_key].obsm['X_umap'] + if "X_pca" in mdata[data_key].obsm: + adata.obsm["X_pca"] = mdata[data_key].obsm["X_pca"] + adata.obsm["X_umap"] = mdata[data_key].obsm["X_umap"] # Optionally subsample cells for faster plotting if umap_subsample_frac is not None and 0 < umap_subsample_frac < 1.0: @@ -303,12 +372,12 @@ def plot_umap_per_gene_guide(mdata, Target_Gene, ax=None, color='red', save_path adata = adata[sorted(indices)] # Set color - colors = ['lightgrey', color] + colors = ["lightgrey", color] n_bins = 100 - cmap = mcolors.LinearSegmentedColormap.from_list('custom', colors, N=n_bins) - - title = f'{Target_Gene} Perturbation' - + cmap = mcolors.LinearSegmentedColormap.from_list("custom", colors, N=n_bins) + + title = f"{Target_Gene} Perturbation" + # If no axis provided, create new figure if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -316,41 +385,58 @@ def plot_umap_per_gene_guide(mdata, Target_Gene, ax=None, color='red', save_path else: fig = ax.get_figure() standalone = False - + # Plot on the provided/created axis plt.sca(ax) # Set current axis to ensure scanpy uses correct figure -CC change - sc.pl.umap(adata, color=Target_Gene, title=title, cmap=cmap, ax=ax, show=False, size=size) - + sc.pl.umap( + adata, color=Target_Gene, title=title, cmap=cmap, ax=ax, show=False, size=size + ) + # Rasterize ONLY the scatter points (collections) in this axis for collection in ax.collections: collection.set_rasterized(True) - - ax.set_title(title, fontsize=18, fontweight='bold', loc='center') - ax.set_xlabel('UMAP 1', fontsize=10, fontweight = 'bold') - ax.set_ylabel('UMAP 2', fontsize=10, fontweight = 'bold') + + ax.set_title(title, fontsize=18, fontweight="bold", loc="center") + ax.set_xlabel("UMAP 1", fontsize=10, fontweight="bold") + ax.set_ylabel("UMAP 2", fontsize=10, fontweight="bold") if ax is None: plt.tight_layout() # Only save if standalone and save_path provided if standalone and save_path and save_name: - fig.savefig(f"{save_path}/{_safe_name(save_name)}.svg", format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) + # Only show/close if standalone if standalone: if show: plt.show() else: plt.close(fig) - + return ax -def plot_top_program_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None,top_n_programs=10, - ax=None, save_path=None, save_name=None, - figsize=(5,8), show=False, gene_name_key='gene_names', - data_key='rna', prog_key='cNMF'): +def plot_top_program_per_gene( + mdata, + Target_Gene, + ensembl_to_symbol_file=None, + top_n_programs=10, + ax=None, + save_path=None, + save_name=None, + figsize=(5, 8), + show=False, + gene_name_key="gene_names", + data_key="rna", + prog_key="cNMF", +): """Horizontal bar plot of the top cNMF programs by gene-loading score for a gene. Reads the cNMF loadings matrix from ``mdata['cNMF'].varm['loadings']`` and @@ -386,7 +472,7 @@ def plot_top_program_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None """ # read cNMF gene program matrix - X = mdata[prog_key].varm["loadings"] + X = mdata[prog_key].varm["loadings"] # rename gene if ensembl_to_symbol_file is None: @@ -394,12 +480,13 @@ def plot_top_program_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None col_names = mdata[data_key].var[gene_name_key].astype(str) else: col_names = mdata[data_key].var_names - df_renamed = pd.DataFrame(data=X, columns = col_names) + df_renamed = pd.DataFrame(data=X, columns=col_names) else: - renamed_gene_list = rename_list_gene_dictionary( mdata[data_key].var_names, ensembl_to_symbol_file) - df_renamed = pd.DataFrame(data=X, columns = renamed_gene_list) + renamed_gene_list = rename_list_gene_dictionary( + mdata[data_key].var_names, ensembl_to_symbol_file + ) + df_renamed = pd.DataFrame(data=X, columns=renamed_gene_list) - # Check if gene exists if Target_Gene not in df_renamed.columns: print(f"Gene {Target_Gene} not found in mdata") @@ -412,12 +499,16 @@ def plot_top_program_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No program to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No program to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax - - - + # Multiple Ensembl IDs collapsing to the same symbol would make # df_renamed[Target_Gene] return a DataFrame and break nlargest(). # Fail explicitly rather than silently averaging or picking one column. @@ -441,7 +532,6 @@ def plot_top_program_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None df_sorted = df_renamed[Target_Gene].nlargest(top_n_programs) df_sorted = df_sorted.sort_values(ascending=True) - # If no axis provided, create new figure if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -449,49 +539,71 @@ def plot_top_program_per_gene(mdata, Target_Gene, ensembl_to_symbol_file = None else: fig = ax.get_figure() standalone = False - + # Create horizontal bar plot - bars = ax.barh(range(len(df_sorted)), df_sorted.values, color='#808080', alpha=0.8) - + bars = ax.barh(range(len(df_sorted)), df_sorted.values, color="#808080", alpha=0.8) + # Customize the plot - ax.set_title(f"Top Loading Program for {Target_Gene}",fontsize=12, fontweight='bold', loc='center') + ax.set_title( + f"Top Loading Program for {Target_Gene}", + fontsize=12, + fontweight="bold", + loc="center", + ) ax.set_yticks(range(len(df_sorted))) ax.set_yticklabels(df_sorted.index, fontsize=10) - ax.set_xlabel('Gene Loading Score (z-score)',fontsize=10, fontweight='bold') #, loc='center') - ax.set_ylabel(f'Program Name',fontsize=10, fontweight='bold', loc='center') - + ax.set_xlabel( + "Gene Loading Score (z-score)", fontsize=10, fontweight="bold" + ) # , loc='center') + ax.set_ylabel(f"Program Name", fontsize=10, fontweight="bold", loc="center") + # Format x-axis ax.set_xlim(0, max(df_sorted.values) * 1.1) - ax.ticklabel_format(style='scientific', axis='x' )#scilimits=(0,0)) - + ax.ticklabel_format(style="scientific", axis="x") # scilimits=(0,0)) + # Add grid - ax.grid(axis='x', alpha=0.3, linestyle='-', linewidth=0.5) + ax.grid(axis="x", alpha=0.3, linestyle="-", linewidth=0.5) ax.set_axisbelow(True) - + # Remove top and right spines - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + if standalone: plt.tight_layout() # Only save if standalone and save_path provided if standalone and save_path and save_name: - fig.savefig(f"{save_path}/{_safe_name(save_name)}.svg", format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) + # Only show/close if standalone if standalone: if show: plt.show() else: plt.close(fig) - - return ax + return ax -def perturbed_gene_dotplot(mdata, Target_Gene, ensembl_to_symbol_file=None, dotplot_groupby='sample', save_name=None, - save_path=None, figsize=(3, 2), show=False, ax=None, gene_name_key='gene_names', data_key='rna'): +def perturbed_gene_dotplot( + mdata, + Target_Gene, + ensembl_to_symbol_file=None, + dotplot_groupby="sample", + save_name=None, + save_path=None, + figsize=(3, 2), + show=False, + ax=None, + gene_name_key="gene_names", + data_key="rna", +): """Scanpy dot plot of a single gene's expression grouped by a categorical variable. Wraps ``sc.pl.dotplot`` for a single perturbed gene, split by e.g. sample/day. @@ -529,11 +641,10 @@ def perturbed_gene_dotplot(mdata, Target_Gene, ensembl_to_symbol_file=None, dot if gene_name_key is not None and gene_name_key in renamed.var.columns: renamed.var_names = renamed.var[gene_name_key].astype(str) else: - renamed = rename_adata_gene_dictionary(mdata[data_key],ensembl_to_symbol_file) + renamed = rename_adata_gene_dictionary(mdata[data_key], ensembl_to_symbol_file) renamed.var_names_make_unique() - # Check if query gene exists if Target_Gene not in renamed.var_names.values: print(f"Gene {Target_Gene} not found in mdata") @@ -546,40 +657,65 @@ def perturbed_gene_dotplot(mdata, Target_Gene, ensembl_to_symbol_file=None, dot img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No gene to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No gene to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax - - + if save_name is None: save_name = f"{Target_Gene} Expression by days" - + # Create the dotplot if ax is None: # Standalone mode - let scanpy create its own figure standalone = True - dp = sc.pl.dotplot(renamed, Target_Gene, groupby=dotplot_groupby, - figsize=figsize, swap_axes=True, dendrogram=False, - show=False, return_fig=True, use_raw=False) + dp = sc.pl.dotplot( + renamed, + Target_Gene, + groupby=dotplot_groupby, + figsize=figsize, + swap_axes=True, + dendrogram=False, + show=False, + return_fig=True, + use_raw=False, + ) dp.make_figure() fig = dp.fig - ax = dp.ax_dict['mainplot_ax'] + ax = dp.ax_dict["mainplot_ax"] else: # Gridspec mode - use provided ax standalone = False fig = ax.get_figure() - dp = sc.pl.dotplot(renamed, Target_Gene, groupby=dotplot_groupby, - swap_axes=True, dendrogram=False, show=False, - return_fig=True, ax=ax, use_raw=False) + dp = sc.pl.dotplot( + renamed, + Target_Gene, + groupby=dotplot_groupby, + swap_axes=True, + dendrogram=False, + show=False, + return_fig=True, + ax=ax, + use_raw=False, + ) dp.make_figure() - - ax.set_title(f"{Target_Gene} Expression", fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel('Gene', fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel('Day', fontsize=10, fontweight='bold', loc='center') - + + ax.set_title( + f"{Target_Gene} Expression", fontsize=14, fontweight="bold", loc="center" + ) + ax.set_ylabel("Gene", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel("Day", fontsize=10, fontweight="bold", loc="center") + # Get labels and set ticks properly - label = list(mdata[data_key].obs[dotplot_groupby].cat.categories) # Use categories instead - + label = list( + mdata[data_key].obs[dotplot_groupby].cat.categories + ) # Use categories instead + # Set both ticks and labels together tick_positions = range(len(label)) ax.set_xticks(tick_positions) @@ -590,7 +726,12 @@ def perturbed_gene_dotplot(mdata, Target_Gene, ensembl_to_symbol_file=None, dot # save fig (only in standalone mode) if standalone and save_name and save_path: - fig.savefig(f"{save_path}/{_safe_name(save_name)}.png", format='png', bbox_inches='tight', dpi=300) # Changed to png + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.png", + format="png", + bbox_inches="tight", + dpi=300, + ) # Changed to png # Control whether to display the plot (only in standalone mode) if standalone: @@ -602,10 +743,22 @@ def perturbed_gene_dotplot(mdata, Target_Gene, ensembl_to_symbol_file=None, dot return ax - -def plot_log2FC(perturb_path, Target, perturb_target_col="target_name", perturb_program_col="program_name", - perturb_log2fc_col='log2FC', num_item=5, significance_threshold=0.05, save_path=None, save_name=None, - figsize=(5, 4), show=False, ax=None, Day ="", perturb_df=None): +def plot_log2FC( + perturb_path, + Target, + perturb_target_col="target_name", + perturb_program_col="program_name", + perturb_log2fc_col="log2FC", + num_item=5, + significance_threshold=0.05, + save_path=None, + save_name=None, + figsize=(5, 4), + show=False, + ax=None, + Day="", + perturb_df=None, +): """Horizontal bar plot of top up- and down-regulated programs by log2 fold-change. Reads perturbation-association results (one row per target-gene × program), @@ -656,18 +809,24 @@ def plot_log2FC(perturb_path, Target, perturb_target_col="target_name", perturb_ df = perturb_df else: df = pd.read_csv(perturb_path, sep="\t") - + # Check if query gene exists if Target not in df[perturb_target_col].values: print(f"Gene {Target} not found in mdata") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax, pd.DataFrame() # Sort by log2FC df_sorted = df.loc[df[perturb_target_col] == Target] - df_sorted = df_sorted[df_sorted['adj_pval'] < significance_threshold] + df_sorted = df_sorted[df_sorted["adj_pval"] < significance_threshold] df_sorted = df_sorted.sort_values(by=perturb_log2fc_col, ascending=False) # Get top and bottom gene top_df = df_sorted.head(num_item) @@ -676,9 +835,9 @@ def plot_log2FC(perturb_path, Target, perturb_target_col="target_name", perturb_ top = top_df.copy() bottom = bottom_df.copy() # Combine data - plot_data = pd.merge(top, bottom, how='outer') + plot_data = pd.merge(top, bottom, how="outer") plot_data = plot_data.sort_values(by=perturb_log2fc_col, ascending=False) - + # Create the plot if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -688,38 +847,52 @@ def plot_log2FC(perturb_path, Target, perturb_target_col="target_name", perturb_ standalone = False # Create horizontal bar plot - colors = ['red' if x > 0 else 'blue' for x in plot_data[perturb_log2fc_col]] - bars = ax.barh(range(len(plot_data)), plot_data[perturb_log2fc_col], color=colors, alpha=0.7) - + colors = ["red" if x > 0 else "blue" for x in plot_data[perturb_log2fc_col]] + bars = ax.barh( + range(len(plot_data)), plot_data[perturb_log2fc_col], color=colors, alpha=0.7 + ) + # Customize the plot ax.set_yticks(range(len(plot_data))) - ax.set_yticklabels(plot_data[perturb_program_col],fontsize=10) + ax.set_yticklabels(plot_data[perturb_program_col], fontsize=10) - - ax.set_xlabel('Effect on Program Expression',fontsize=10, fontweight='bold', loc='center') - ax.set_ylabel( "Name",fontsize=10, fontweight='bold', loc='center') - ax.set_title(f"Log2 Fold-Change with {Target} {Day}",fontsize=14, fontweight='bold', loc='center') + ax.set_xlabel( + "Effect on Program Expression", fontsize=10, fontweight="bold", loc="center" + ) + ax.set_ylabel("Name", fontsize=10, fontweight="bold", loc="center") + ax.set_title( + f"Log2 Fold-Change with {Target} {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) # Add a vertical line at x=0 - ax.axvline(x=0, color='black', linestyle='-', linewidth=0.5) - + ax.axvline(x=0, color="black", linestyle="-", linewidth=0.5) - # Add legend from matplotlib.patches import Patch - legend_elements = [Patch(facecolor='red', alpha=0.7, label='Upregulated'), - Patch(facecolor='blue', alpha=0.7, label='Downregulated')] - ax.legend(handles=legend_elements) #loc='lower right') - + + legend_elements = [ + Patch(facecolor="red", alpha=0.7, label="Upregulated"), + Patch(facecolor="blue", alpha=0.7, label="Downregulated"), + ] + ax.legend(handles=legend_elements) # loc='lower right') + # Adjust layout - ax.grid(axis='x', alpha=0.3) + ax.grid(axis="x", alpha=0.3) if standalone: plt.tight_layout() # Save if path provided (only when standalone) if standalone and save_path: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only when standalone) if standalone: @@ -729,14 +902,26 @@ def plot_log2FC(perturb_path, Target, perturb_target_col="target_name", perturb_ plt.close(fig) return ax, plot_data - - -def plot_volcano(perturb_path, Target, perturb_target_col="target_name", perturb_program_col="program_name", perturb_log2fc_col = 'log2FC', - volcano_log2fc_min=-0.05, volcano_log2fc_max=0.05, significance_threshold=0.05, save_path=None, - save_name=None, figsize=(5, 4), show=False, ax=None, Day ="", - perturb_df=None, run_adjust_text=True): +def plot_volcano( + perturb_path, + Target, + perturb_target_col="target_name", + perturb_program_col="program_name", + perturb_log2fc_col="log2FC", + volcano_log2fc_min=-0.05, + volcano_log2fc_max=0.05, + significance_threshold=0.05, + save_path=None, + save_name=None, + figsize=(5, 4), + show=False, + ax=None, + Day="", + perturb_df=None, + run_adjust_text=True, +): """Volcano plot of perturbation log2FC vs adjusted p-value for one target gene. Programs exceeding the fold-change and significance thresholds are @@ -791,17 +976,23 @@ def plot_volcano(perturb_path, Target, perturb_target_col="target_name", perturb df = perturb_df else: df = pd.read_csv(perturb_path, sep="\t") - + # Check if query gene exists if Target not in df[perturb_target_col].values: print(f"Gene {Target} not found in mdata") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax, pd.DataFrame(), [] df_program = df.loc[df[perturb_target_col] == Target] - + # Create figure/axes if not provided if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -811,53 +1002,97 @@ def plot_volcano(perturb_path, Target, perturb_target_col="target_name", perturb standalone = False # Plot all points - ax.scatter(x=df_program[perturb_log2fc_col], - y=df_program['adj_pval'].apply(lambda x: -np.log10(x)), - s=10, label="Not significant", color="grey") - + ax.scatter( + x=df_program[perturb_log2fc_col], + y=df_program["adj_pval"].apply(lambda x: -np.log10(x)), + s=10, + label="Not significant", + color="grey", + ) + # Highlight down- or up-regulated genes - down = df_program[(df_program[perturb_log2fc_col] <= volcano_log2fc_min) & - (df_program['adj_pval'] <= significance_threshold)] - up = df_program[(df_program[perturb_log2fc_col] >= volcano_log2fc_max) & - (df_program['adj_pval'] <= significance_threshold)] - - ax.scatter(x=down[perturb_log2fc_col], - y=down['adj_pval'].apply(lambda x: -np.log10(x)), - s=10, label="Down-regulated", color="blue") - ax.scatter(x=up[perturb_log2fc_col], - y=up['adj_pval'].apply(lambda x: -np.log10(x)), - s=10, label="Up-regulated", color="red") - + down = df_program[ + (df_program[perturb_log2fc_col] <= volcano_log2fc_min) + & (df_program["adj_pval"] <= significance_threshold) + ] + up = df_program[ + (df_program[perturb_log2fc_col] >= volcano_log2fc_max) + & (df_program["adj_pval"] <= significance_threshold) + ] + + ax.scatter( + x=down[perturb_log2fc_col], + y=down["adj_pval"].apply(lambda x: -np.log10(x)), + s=10, + label="Down-regulated", + color="blue", + ) + ax.scatter( + x=up[perturb_log2fc_col], + y=up["adj_pval"].apply(lambda x: -np.log10(x)), + s=10, + label="Up-regulated", + color="red", + ) + # Add text labels texts = [] for i, r in up.iterrows(): - texts.append(ax.text(x=r[perturb_log2fc_col], y=-np.log10(r['adj_pval']), - s=r[perturb_program_col], fontsize=10, zorder=10)) + texts.append( + ax.text( + x=r[perturb_log2fc_col], + y=-np.log10(r["adj_pval"]), + s=r[perturb_program_col], + fontsize=10, + zorder=10, + ) + ) for i, r in down.iterrows(): - texts.append(ax.text(x=r[perturb_log2fc_col], y=-np.log10(r['adj_pval']), - s=r[perturb_program_col], fontsize=10, zorder=10)) - + texts.append( + ax.text( + x=r[perturb_log2fc_col], + y=-np.log10(r["adj_pval"]), + s=r[perturb_program_col], + fontsize=10, + zorder=10, + ) + ) + # Adjust text to avoid overlaps if run_adjust_text: - adjust_text(texts, arrowprops=dict(arrowstyle='->', color='black', lw=0.5), ax=ax) - - + adjust_text( + texts, arrowprops=dict(arrowstyle="->", color="black", lw=0.5), ax=ax + ) + # Set labels and lines - ax.set_title(f"Volcano Plot for {Target} {Day}",fontsize=14, fontweight='bold', loc='center') - ax.set_xlabel('Effect on Program Expression',fontsize=10, fontweight='bold', loc='center') - ax.set_ylabel("Adjusted p-value, -log10",fontsize=10, fontweight='bold', loc='center') + ax.set_title( + f"Volcano Plot for {Target} {Day}", fontsize=14, fontweight="bold", loc="center" + ) + ax.set_xlabel( + "Effect on Program Expression", fontsize=10, fontweight="bold", loc="center" + ) + ax.set_ylabel( + "Adjusted p-value, -log10", fontsize=10, fontweight="bold", loc="center" + ) ax.axvline(volcano_log2fc_min, color="grey", linestyle="--") ax.axvline(volcano_log2fc_max, color="grey", linestyle="--") ax.axhline(-np.log10(significance_threshold), color="grey", linestyle="--") ax.legend() - ax.set_title(f"Volcano Plot for {Target} {Day}",fontsize=14, fontweight='bold', loc='center') + ax.set_title( + f"Volcano Plot for {Target} {Day}", fontsize=14, fontweight="bold", loc="center" + ) if standalone: plt.tight_layout() # Save if path provided (only when standalone) if standalone and save_path: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only when standalone) if standalone: @@ -866,13 +1101,22 @@ def plot_volcano(perturb_path, Target, perturb_target_col="target_name", perturb else: plt.close(fig) - return ax, pd.merge(up, down, how='outer'), texts - + return ax, pd.merge(up, down, how="outer"), texts - -def programs_dotplot(mdata, Target, dotplot_groupby="sample", program_list=None, - save_path=None, save_name=None, figsize=(5, 4), show=False, ax=None, Day="", prog_key='cNMF'): +def programs_dotplot( + mdata, + Target, + dotplot_groupby="sample", + program_list=None, + save_path=None, + save_name=None, + figsize=(5, 4), + show=False, + ax=None, + Day="", + prog_key="cNMF", +): """Scanpy dot plot of cNMF program loading scores grouped by a categorical variable. Shows selected programs (e.g. significant ones from ``plot_log2FC``) as a @@ -911,13 +1155,16 @@ def programs_dotplot(mdata, Target, dotplot_groupby="sample", program_list=None, # make anndata from program loadings adata_new = mdata[prog_key] - program_name_list = adata_new.var_names.tolist() # get all programs - - program_names = list(map(str, program_list)) if program_list is not None else [] # make program str if given - program_names = [program for program in program_names if program in program_name_list] - - if not program_names: + program_name_list = adata_new.var_names.tolist() # get all programs + + program_names = ( + list(map(str, program_list)) if program_list is not None else [] + ) # make program str if given + program_names = [ + program for program in program_names if program in program_name_list + ] + if not program_names: print("No significant Programs") # Handle empty gene list if ax is None: @@ -928,43 +1175,75 @@ def programs_dotplot(mdata, Target, dotplot_groupby="sample", program_list=None, img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No programs to display', - ha='center', va='center', transform=ax.transAxes) - ax.set_title(f"{Target} Perturbed Program Loading Scores {Day}", - fontweight='bold', loc='center') + ax.text( + 0.5, + 0.5, + "No programs to display", + ha="center", + va="center", + transform=ax.transAxes, + ) + ax.set_title( + f"{Target} Perturbed Program Loading Scores {Day}", + fontweight="bold", + loc="center", + ) return ax program_names = program_names[::-1] - + # Create dotplot if ax is None: # Standalone mode - let scanpy create its own figure standalone = True - dp = sc.pl.dotplot(adata_new, program_names, groupby=dotplot_groupby, swap_axes=True, figsize=figsize, - dendrogram=False, show=False, return_fig=True, use_raw=False) + dp = sc.pl.dotplot( + adata_new, + program_names, + groupby=dotplot_groupby, + swap_axes=True, + figsize=figsize, + dendrogram=False, + show=False, + return_fig=True, + use_raw=False, + ) dp.make_figure() fig = dp.fig - ax = dp.ax_dict['mainplot_ax'] + ax = dp.ax_dict["mainplot_ax"] else: # Gridspec mode - use provided ax standalone = False fig = ax.get_figure() - dp = sc.pl.dotplot(adata_new, program_names, groupby=dotplot_groupby, swap_axes=True, figsize=figsize, - dendrogram=False, show=False, return_fig=True, ax=ax, use_raw=False) + dp = sc.pl.dotplot( + adata_new, + program_names, + groupby=dotplot_groupby, + swap_axes=True, + figsize=figsize, + dendrogram=False, + show=False, + return_fig=True, + ax=ax, + use_raw=False, + ) dp.make_figure() - - ax.set_title(f"{Target} Perturbed Program Loading Scores {Day}", - fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel('Program Name', fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel("Day", fontsize=10, fontweight='bold', loc='center') - + + ax.set_title( + f"{Target} Perturbed Program Loading Scores {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) + ax.set_ylabel("Program Name", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel("Day", fontsize=10, fontweight="bold", loc="center") + # Get labels and set ticks properly # Check if it's already categorical - if hasattr(adata_new.obs[dotplot_groupby], 'cat'): + if hasattr(adata_new.obs[dotplot_groupby], "cat"): label = list(adata_new.obs[dotplot_groupby].cat.categories) else: label = list(adata_new.obs[dotplot_groupby].unique()) - + # Set both ticks and labels together tick_positions = range(len(label)) ax.set_xticks(tick_positions) @@ -975,7 +1254,12 @@ def programs_dotplot(mdata, Target, dotplot_groupby="sample", program_list=None, # Save if path provided (only when standalone) if standalone and save_path and save_name: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only in standalone mode) if standalone: @@ -985,10 +1269,16 @@ def programs_dotplot(mdata, Target, dotplot_groupby="sample", program_list=None, plt.close(fig) return ax - -def compute_gene_correlation_matrix(mdata, ensembl_to_symbol_file=None, gene_name_key='gene_names', - save_path=None, data_key='rna', prog_key='cNMF'): + +def compute_gene_correlation_matrix( + mdata, + ensembl_to_symbol_file=None, + gene_name_key="gene_names", + save_path=None, + data_key="rna", + prog_key="cNMF", +): """Gene x gene Pearson correlation of cNMF gene loadings across programs. Same values as ``pd.DataFrame(loadings).corr()``, but @@ -1025,7 +1315,11 @@ def compute_gene_correlation_matrix(mdata, ensembl_to_symbol_file=None, gene_nam else: col_names = list(mdata[data_key].var_names) else: - col_names = list(rename_list_gene_dictionary(mdata[data_key].var_names, ensembl_to_symbol_file)) + col_names = list( + rename_list_gene_dictionary( + mdata[data_key].var_names, ensembl_to_symbol_file + ) + ) # genes as rows, correlated across programs; plots expect NaN filled with 0 corr = RowCorr(X.T, col_names, row_label="genes", col_label="programs") @@ -1036,10 +1330,17 @@ def compute_gene_correlation_matrix(mdata, ensembl_to_symbol_file=None, gene_nam return corr - -def analyze_correlations(gene_loading_corr_matrix, Target, top_corr_genes=5, save_path=None, - save_name=None, figsize=(10, 8), show=False, ax=None): +def analyze_correlations( + gene_loading_corr_matrix, + Target, + top_corr_genes=5, + save_path=None, + save_name=None, + figsize=(10, 8), + show=False, + ax=None, +): """Horizontal bar plot of genes most/least correlated with a target gene in cNMF loadings. Uses the precomputed gene × gene correlation matrix from @@ -1084,26 +1385,33 @@ def analyze_correlations(gene_loading_corr_matrix, Target, top_corr_genes=5, sav img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No correlation to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No correlation to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax - # Get correlations with the target program target_correlations = gene_loading_corr_matrix[Target] - target_correlations = target_correlations.drop(Target, errors="ignore") # Remove self-correlation - + target_correlations = target_correlations.drop( + Target, errors="ignore" + ) # Remove self-correlation + # Sort correlations sorted_correlations = target_correlations.sort_values(ascending=False) - + # Get top and bottom gene top = sorted_correlations[sorted_correlations > 0].head(top_corr_genes) bottom = sorted_correlations[sorted_correlations < 0].tail(top_corr_genes) - + # Combine for plotting combined_correlations = pd.concat([bottom, top]) combined_correlations = combined_correlations.sort_values(ascending=False) - + # Create figure/axes if not provided if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -1113,38 +1421,57 @@ def analyze_correlations(gene_loading_corr_matrix, Target, top_corr_genes=5, sav standalone = False # Create horizontal bar plot - colors = ['blue' if x < 0 else 'red' for x in combined_correlations] - bars = ax.barh(range(len(combined_correlations)), combined_correlations.values, - color=colors, alpha=0.7) - + colors = ["blue" if x < 0 else "red" for x in combined_correlations] + bars = ax.barh( + range(len(combined_correlations)), + combined_correlations.values, + color=colors, + alpha=0.7, + ) + # Customize the plot - ax.set_title(f"Gene Loading Correlation for {Target}",fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel('Gene Name',fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel('Correlation Coefficient',fontsize=10, fontweight='bold', loc='center') - + ax.set_title( + f"Gene Loading Correlation for {Target}", + fontsize=14, + fontweight="bold", + loc="center", + ) + ax.set_ylabel("Gene Name", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel( + "Correlation Coefficient", fontsize=10, fontweight="bold", loc="center" + ) + # Set y-axis labels ax.set_yticks(range(len(combined_correlations))) - ax.set_yticklabels(combined_correlations.index, rotation=45, ha='right') - + ax.set_yticklabels(combined_correlations.index, rotation=45, ha="right") + # Add a vertical line at x=0 - ax.axvline(x=0, color='black', linestyle='-', linewidth=0.5) - + ax.axvline(x=0, color="black", linestyle="-", linewidth=0.5) + # Add legend from matplotlib.patches import Patch - legend_elements = [Patch(facecolor='red', alpha=0.7, label='Positive correlation'), - Patch(facecolor='blue', alpha=0.7, label='Negative correlation')] - ax.legend(handles=legend_elements )# loc='lower right') - + + legend_elements = [ + Patch(facecolor="red", alpha=0.7, label="Positive correlation"), + Patch(facecolor="blue", alpha=0.7, label="Negative correlation"), + ] + ax.legend(handles=legend_elements) # loc='lower right') + # Add grid - ax.grid(axis='x', alpha=0.3) - + ax.grid(axis="x", alpha=0.3) + # Adjust layout (only in standalone mode) if standalone: plt.tight_layout() # Save if path provided (only in standalone mode) if standalone and save_path: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only in standalone mode) if standalone: @@ -1156,8 +1483,9 @@ def analyze_correlations(gene_loading_corr_matrix, Target, top_corr_genes=5, sav return ax - -def compute_gene_waterfall_cor(perturb_path, perturb_log2fc_col='log2FC', save_path=None): +def compute_gene_waterfall_cor( + perturb_path, perturb_log2fc_col="log2FC", save_path=None +): """Regulator x regulator (target x target) correlation of per-program log2FC. Pivots the perturbation results to a (target x program) table and @@ -1182,25 +1510,40 @@ def compute_gene_waterfall_cor(perturb_path, perturb_log2fc_col='log2FC', save_p (self = 1; the waterfall plots drop it). """ - df = pd.read_csv(perturb_path, sep='\t', index_col=0) - pivot_df = df.pivot_table(index='target_name', columns='program_name', - values=perturb_log2fc_col) + df = pd.read_csv(perturb_path, sep="\t", index_col=0) + pivot_df = df.pivot_table( + index="target_name", columns="program_name", values=perturb_log2fc_col + ) # targets as rows, correlated across programs; the waterfall plots drop self - corr = RowCorr(pivot_df.to_numpy(dtype=np.float64), pivot_df.index, - row_label="targets", col_label="programs") + corr = RowCorr( + pivot_df.to_numpy(dtype=np.float64), + pivot_df.index, + row_label="targets", + col_label="programs", + ) if save_path is not None: Path(save_path).parent.mkdir(parents=True, exist_ok=True) - np.savez(save_path, corr=corr.full(), regulator=np.asarray(corr.index, dtype=str)) + np.savez( + save_path, corr=corr.full(), regulator=np.asarray(corr.index, dtype=str) + ) return corr - -def create_gene_correlation_waterfall(corr_matrix, Target_Gene, top_corr_genes=5, save_path=None, - save_name=None, figsize=(3, 5), show=False, ax=None, Day ="", - run_adjust_text=True): +def create_gene_correlation_waterfall( + corr_matrix, + Target_Gene, + top_corr_genes=5, + save_path=None, + save_name=None, + figsize=(3, 5), + show=False, + ax=None, + Day="", + run_adjust_text=True, +): """Waterfall scatter plot of gene-gene perturbation correlation for one target. Plots all genes' perturbation-profile correlations with ``Target_Gene`` @@ -1254,18 +1597,34 @@ def create_gene_correlation_waterfall(corr_matrix, Target_Gene, top_corr_genes=5 # empty panel rather than raising KeyError, so the rest of the # comprehensive plot for this gene is still produced. if Target_Gene not in corr_matrix: - ax.text(0.5, 0.5, f"No perturbation\ncorrelation data\nfor {Target_Gene}", - ha='center', va='center', transform=ax.transAxes, - fontsize=10, style='italic', color='gray') + ax.text( + 0.5, + 0.5, + f"No perturbation\ncorrelation data\nfor {Target_Gene}", + ha="center", + va="center", + transform=ax.transAxes, + fontsize=10, + style="italic", + color="gray", + ) ax.set_xticks([]) ax.set_yticks([]) - ax.set_title(f"Gene Perturbation Correlation for {Target_Gene} {Day}", - fontsize=14, fontweight='bold', loc='center') + ax.set_title( + f"Gene Perturbation Correlation for {Target_Gene} {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) if standalone: plt.tight_layout() if save_path: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', - bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) if show: plt.show() else: @@ -1274,46 +1633,60 @@ def create_gene_correlation_waterfall(corr_matrix, Target_Gene, top_corr_genes=5 # Convert to DataFrame and sort gene_corrs = corr_matrix.loc[Target_Gene].dropna() - corr_df = (gene_corrs).sort_values(ascending = False) - corr_df = corr_df.drop(str(Target_Gene), errors="ignore") # drop self corr (already gone if NaN) + corr_df = (gene_corrs).sort_values(ascending=False) + corr_df = corr_df.drop( + str(Target_Gene), errors="ignore" + ) # drop self corr (already gone if NaN) # Get top N positive and bottom N negative correlations for labeling top_positive = corr_df.head(top_corr_genes) top_negative = corr_df.tail(top_corr_genes) genes_to_label = set(top_positive.index.tolist() + top_negative.index.tolist()) - + # Use ALL correlations for plotting plot_data = corr_df # Plot only the markers (no line) - ax.scatter(range(len(plot_data)), plot_data.values, - color='gray', s=16) - + ax.scatter(range(len(plot_data)), plot_data.values, color="gray", s=16) + # Add horizontal reference line at y=0 - ax.axhline(y=0, color='gray', linestyle='--', linewidth=1, alpha=0.5) - + ax.axhline(y=0, color="gray", linestyle="--", linewidth=1, alpha=0.5) + # Remove x-axis labels ax.set_xticks([]) - # Highlight the labeled genes with larger markers and add annotations texts = [] for i, (gene, values) in enumerate(plot_data.items()): if gene in genes_to_label: - ax.plot(i, values, 'o', color='darkgray', markersize=5, zorder=5) - text = ax.text(i, values, gene, fontsize=10, style='italic', ha='center') + ax.plot(i, values, "o", color="darkgray", markersize=5, zorder=5) + text = ax.text(i, values, gene, fontsize=10, style="italic", ha="center") texts.append(text) - + if run_adjust_text: - adjust_text(texts, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5), ax=ax, fontsize=12) + adjust_text( + texts, + arrowprops=dict(arrowstyle="-", color="gray", lw=0.5), + ax=ax, + fontsize=12, + ) - ax.set_title(f"Gene Perturbation Correlation for {Target_Gene} {Day}",fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel(f'Correlation of Program\nExpression with\n{Target_Gene} Perturbation',fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel('Perturbed Genes',fontsize=10, fontweight='bold', loc='center') + ax.set_title( + f"Gene Perturbation Correlation for {Target_Gene} {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) + ax.set_ylabel( + f"Correlation of Program\nExpression with\n{Target_Gene} Perturbation", + fontsize=10, + fontweight="bold", + loc="center", + ) + ax.set_xlabel("Perturbed Genes", fontsize=10, fontweight="bold", loc="center") - # Add grid - ax.grid(True, alpha=0.3, linestyle='--') + ax.grid(True, alpha=0.3, linestyle="--") ax.set_axisbelow(True) # Adjust layout (only in standalone mode) @@ -1322,7 +1695,12 @@ def create_gene_correlation_waterfall(corr_matrix, Target_Gene, top_corr_genes=5 # Save if path provided (only in standalone mode) if standalone and save_path: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only in standalone mode) if standalone: @@ -1330,17 +1708,25 @@ def create_gene_correlation_waterfall(corr_matrix, Target_Gene, top_corr_genes=5 plt.show() else: plt.close(fig) - - - return ax, texts + return ax, texts -def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', ax=None, figsize=(3.2, 5.0), - control_target_name='non-targeting', - save_path=None, save_name=None, show=True, - data_key='rna', prog_key='cNMF', - guide_targets_key='guide_targets', guide_assignment_key='guide_assignment'): +def plot_perturbation_vs_control( + mdata, + target_gene, + gene_name_key="gene_names", + ax=None, + figsize=(3.2, 5.0), + control_target_name="non-targeting", + save_path=None, + save_name=None, + show=True, + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): """ Bar plot comparing expression of a targeted gene in perturbed vs non-targeting control cells, normalized to control mean (displayed as %). @@ -1370,8 +1756,9 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', """ from scipy import sparse - assert mdata[data_key].n_obs == mdata[prog_key].n_obs, \ + assert mdata[data_key].n_obs == mdata[prog_key].n_obs, ( "RNA and cNMF have different number of cells" + ) guide_targets = mdata[prog_key].uns[guide_targets_key] ga = mdata[prog_key].obsm[guide_assignment_key] @@ -1386,8 +1773,14 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', if gene_mask.sum() == 0: print(f"Gene {target_gene} not found in expression matrix") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax return None gene_idx = np.where(gene_mask)[0][0] @@ -1416,8 +1809,14 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', if len(target_idx) == 0: print(f"Gene {target_gene} not found in guide_targets") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax return None perturbed_mask = np.array(ga[:, target_idx].sum(axis=1)).flatten() > 0 @@ -1428,20 +1827,27 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', # Normalize both to control mean → values are now fractions (1.0 = 100%) ctrl_mean = np.mean(ctrl_expr) if np.isnan(ctrl_mean): - print(f"WARNING: ctrl_mean is NaN (n_ctrl={len(ctrl_expr)}); bar plot will be empty.") + print( + f"WARNING: ctrl_mean is NaN (n_ctrl={len(ctrl_expr)}); bar plot will be empty." + ) norm_means = np.array([np.mean(pert_expr), ctrl_mean]) / ctrl_mean - norm_sems = np.array([ - np.std(pert_expr) / np.sqrt(len(pert_expr)), - np.std(ctrl_expr) / np.sqrt(len(ctrl_expr)), - ]) / ctrl_mean + norm_sems = ( + np.array( + [ + np.std(pert_expr) / np.sqrt(len(pert_expr)), + np.std(ctrl_expr) / np.sqrt(len(ctrl_expr)), + ] + ) + / ctrl_mean + ) n_pert = len(pert_expr) n_ctrl = len(ctrl_expr) # ── Style ──────────────────────────────────────────────────────────────── - COLOR_PERT = '#D62728' - COLOR_CTRL = '#999999' - ECAP = dict(linewidth=1.5, ecolor='black') + COLOR_PERT = "#D62728" + COLOR_CTRL = "#999999" + ECAP = dict(linewidth=1.5, ecolor="black") # If no axis provided, create new figure if ax is None: @@ -1452,39 +1858,49 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', standalone = False bars = ax.bar( - [0, 1], norm_means, yerr=norm_sems, - capsize=6, color=[COLOR_PERT, COLOR_CTRL], edgecolor='none', - width=0.55, error_kw=ECAP, + [0, 1], + norm_means, + yerr=norm_sems, + capsize=6, + color=[COLOR_PERT, COLOR_CTRL], + edgecolor="none", + width=0.55, + error_kw=ECAP, ) # Dashed reference line at 100 % - ax.axhline(1.0, color='gray', linestyle='--', linewidth=1.2, zorder=0) + ax.axhline(1.0, color="gray", linestyle="--", linewidth=1.2, zorder=0) # Y-axis: percentage ticks ax.set_ylim(0, 1.18) - ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f'{int(round(v * 100))}%')) + ax.yaxis.set_major_formatter( + plt.FuncFormatter(lambda v, _: f"{int(round(v * 100))}%") + ) ax.set_yticks([0, 0.25, 0.50, 0.75, 1.00]) # X-axis tick labels (gene name + n on separate line, matching reference) ax.set_xticks([0, 1]) - ax.set_xticklabels([ - f'{target_gene}\n[n = {n_pert}]', - f'Non\nTargeting\n[n = {n_ctrl}]', - ], fontsize=10) + ax.set_xticklabels( + [ + f"{target_gene}\n[n = {n_pert}]", + f"Non\nTargeting\n[n = {n_ctrl}]", + ], + fontsize=10, + ) - ax.set_ylabel(f'{target_gene} Expression', fontsize=11, labelpad=8) + ax.set_ylabel(f"{target_gene} Expression", fontsize=11, labelpad=8) effect_pct = (1 - norm_means[0]) * 100 if np.isnan(effect_pct): - title_str = 'CRISPRi knockdown:\neffect N/A' + title_str = "CRISPRi knockdown:\neffect N/A" else: - title_str = f'CRISPRi knockdown:\n>{int(round(effect_pct))}% effect' - ax.set_title(title_str, fontsize=12, fontweight='bold', pad=12) + title_str = f"CRISPRi knockdown:\n>{int(round(effect_pct))}% effect" + ax.set_title(title_str, fontsize=12, fontweight="bold", pad=12) - ax.tick_params(axis='y', labelsize=10) - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - ax.spines['bottom'].set_linewidth(0.8) - ax.spines['left'].set_linewidth(0.8) + ax.tick_params(axis="y", labelsize=10) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + ax.spines["bottom"].set_linewidth(0.8) + ax.spines["left"].set_linewidth(0.8) # Adjust layout (only in standalone mode) if standalone: @@ -1492,7 +1908,12 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', # Save if path provided (only in standalone mode) if standalone and save_path: - fig.savefig(f'{save_path}/{_safe_name(save_name)}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{_safe_name(save_name)}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only in standalone mode) if standalone: @@ -1504,13 +1925,23 @@ def plot_perturbation_vs_control(mdata, target_gene, gene_name_key='gene_names', return ax -def plot_perturbation_vs_control_by_condition(mdata, target_gene, condition_key='sample', - gene_name_key='gene_names', figsize=(8, 5), - control_target_name='non-targeting', - ax=None, - save_path=None, save_name=None, show=True, PDF=True, - data_key='rna', prog_key='cNMF', - guide_targets_key='guide_targets', guide_assignment_key='guide_assignment'): +def plot_perturbation_vs_control_by_condition( + mdata, + target_gene, + condition_key="sample", + gene_name_key="gene_names", + figsize=(8, 5), + control_target_name="non-targeting", + ax=None, + save_path=None, + save_name=None, + show=True, + PDF=True, + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): """ Grouped bar chart comparing perturbed vs control expression across conditions. @@ -1556,8 +1987,9 @@ def plot_perturbation_vs_control_by_condition(mdata, target_gene, condition_key= """ from scipy import sparse - assert mdata[data_key].n_obs == mdata[prog_key].n_obs, \ + assert mdata[data_key].n_obs == mdata[prog_key].n_obs, ( "RNA and cNMF have different number of cells" + ) guide_targets = mdata[prog_key].uns[guide_targets_key] ga = mdata[prog_key].obsm[guide_assignment_key] @@ -1572,8 +2004,14 @@ def plot_perturbation_vs_control_by_condition(mdata, target_gene, condition_key= if gene_mask.sum() == 0: print(f"Gene {target_gene} not found in expression matrix") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax return None @@ -1593,29 +2031,37 @@ def plot_perturbation_vs_control_by_condition(mdata, target_gene, condition_key= nt_idx = np.where(np.isin(guide_targets, control_target_name))[0] control_mask = np.array(ga[:, nt_idx].sum(axis=1)).flatten() > 0 if control_mask.sum() == 0: - print(f"WARNING: 0 control cells for control_target_name='{control_target_name}'. " - f"Check np.unique(mdata['{prog_key}'].uns['{guide_targets_key}']) for the correct label.") + print( + f"WARNING: 0 control cells for control_target_name='{control_target_name}'. " + f"Check np.unique(mdata['{prog_key}'].uns['{guide_targets_key}']) for the correct label." + ) # Perturbed cells target_idx = np.where(guide_targets == target_gene)[0] if len(target_idx) == 0: print(f"Gene {target_gene} not found in guide_targets") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax return None perturbed_mask = np.array(ga[:, target_idx].sum(axis=1)).flatten() > 0 # Get conditions conditions_series = mdata[data_key].obs[condition_key] - if hasattr(conditions_series, 'cat'): + if hasattr(conditions_series, "cat"): conditions = list(conditions_series.cat.categories) else: conditions = sorted(conditions_series.unique()) # Group labels: "All" + each condition - group_labels = ['All'] + [str(c) for c in conditions] + group_labels = ["All"] + [str(c) for c in conditions] n_groups = len(group_labels) # Compute means and SEMs for each group @@ -1624,7 +2070,7 @@ def plot_perturbation_vs_control_by_condition(mdata, target_gene, condition_key= pert_ns, ctrl_ns = [], [] for label in group_labels: - if label == 'All': + if label == "All": cond_mask = np.ones(len(gene_expr_norm), dtype=bool) else: cond_mask = np.array(conditions_series == label) @@ -1667,57 +2113,75 @@ def plot_perturbation_vs_control_by_condition(mdata, target_gene, condition_key= x = np.arange(n_groups) bar_width = 0.35 - COLOR_PERT = '#D62728' - COLOR_CTRL = '#999999' - ECAP = dict(linewidth=1.5, ecolor='black') - - ax.bar(x - bar_width / 2, pert_means, bar_width, yerr=pert_sems, - capsize=4, color=COLOR_PERT, edgecolor='none', error_kw=ECAP, - label=target_gene) - ax.bar(x + bar_width / 2, ctrl_means, bar_width, yerr=ctrl_sems, - capsize=4, color=COLOR_CTRL, edgecolor='none', error_kw=ECAP, - label='Non-Targeting') + COLOR_PERT = "#D62728" + COLOR_CTRL = "#999999" + ECAP = dict(linewidth=1.5, ecolor="black") + + ax.bar( + x - bar_width / 2, + pert_means, + bar_width, + yerr=pert_sems, + capsize=4, + color=COLOR_PERT, + edgecolor="none", + error_kw=ECAP, + label=target_gene, + ) + ax.bar( + x + bar_width / 2, + ctrl_means, + bar_width, + yerr=ctrl_sems, + capsize=4, + color=COLOR_CTRL, + edgecolor="none", + error_kw=ECAP, + label="Non-Targeting", + ) # Reference line at 100% - ax.axhline(1.0, color='gray', linestyle='--', linewidth=1.2, zorder=0) + ax.axhline(1.0, color="gray", linestyle="--", linewidth=1.2, zorder=0) # X-axis: condition labels with cell counts tick_labels = [] for i, label in enumerate(group_labels): - tick_labels.append(f'{label}\n[{pert_ns[i]}|{ctrl_ns[i]}]') + tick_labels.append(f"{label}\n[{pert_ns[i]}|{ctrl_ns[i]}]") ax.set_xticks(x) ax.set_xticklabels(tick_labels, fontsize=10) # Y-axis: percentage ax.set_ylim(0, 1.18) - ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f'{int(round(v * 100))}%')) + ax.yaxis.set_major_formatter( + plt.FuncFormatter(lambda v, _: f"{int(round(v * 100))}%") + ) ax.set_yticks([0, 0.25, 0.50, 0.75, 1.00]) - ax.set_ylabel(f'{target_gene} Expression', fontsize=11, labelpad=8) + ax.set_ylabel(f"{target_gene} Expression", fontsize=11, labelpad=8) # Title with overall effect all_effect = (1 - pert_means[0]) * 100 if not np.isnan(pert_means[0]) else np.nan if np.isnan(all_effect): - title_str = f'CRISPRi knockdown: {target_gene}\noverall effect N/A' + title_str = f"CRISPRi knockdown: {target_gene}\noverall effect N/A" else: - title_str = f'CRISPRi knockdown: {target_gene}\n>{int(round(all_effect))}% overall effect' - ax.set_title(title_str, fontsize=12, fontweight='bold', pad=12) + title_str = f"CRISPRi knockdown: {target_gene}\n>{int(round(all_effect))}% overall effect" + ax.set_title(title_str, fontsize=12, fontweight="bold", pad=12) - ax.tick_params(axis='y', labelsize=10) - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - ax.spines['bottom'].set_linewidth(0.8) - ax.spines['left'].set_linewidth(0.8) + ax.tick_params(axis="y", labelsize=10) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + ax.spines["bottom"].set_linewidth(0.8) + ax.spines["left"].set_linewidth(0.8) - ax.legend(fontsize=9, loc='upper right', framealpha=0.8) + ax.legend(fontsize=9, loc="upper right", framealpha=0.8) if standalone: plt.tight_layout() if save_path and save_name: - ext = 'pdf' if PDF else 'svg' + ext = "pdf" if PDF else "svg" full_path = f"{save_path}/{_safe_name(save_name)}.{ext}" - fig.savefig(full_path, format=ext, bbox_inches='tight', dpi=300) + fig.savefig(full_path, format=ext, bbox_inches="tight", dpi=300) if show: plt.show() @@ -1737,10 +2201,10 @@ def create_comprehensive_plot( gene_loading_corr_matrix=None, perturb_corr_by_sample=None, top_n_programs=10, - dotplot_groupby='sample', + dotplot_groupby="sample", perturb_target_col="target_name", perturb_program_col="program_name", - perturb_log2fc_col='log2FC', + perturb_log2fc_col="log2FC", top_corr_genes=5, volcano_log2fc_min=-0.05, volcano_log2fc_max=0.05, @@ -1752,14 +2216,14 @@ def create_comprehensive_plot( square_plots=True, show=True, PDF=True, - umap_dot_size = None, + umap_dot_size=None, umap_subsample_frac=None, - gene_name_key='gene_names', - control_target_name='non-targeting', - data_key='rna', - prog_key='cNMF', - guide_targets_key='guide_targets', - guide_assignment_key='guide_assignment' + gene_name_key="gene_names", + control_target_name="non-targeting", + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", ): """Create a comprehensive multi-panel figure for one perturbed gene. @@ -1851,7 +2315,7 @@ def create_comprehensive_plot( Name of the non-targeting control in guide_targets (e.g. 'non-targeting', 'CTRL'). Passed to ``plot_perturbation_vs_control``. """ - + # These carry defaults only so perturb_path_base can be optional (a # non-default parameter cannot follow a defaulted one); they are still required. if Target_Gene is None: @@ -1872,7 +2336,7 @@ def create_comprehensive_plot( # Set default samples if not provided if sample is None: - sample = ['D0', 'sample_D1', 'sample_D2', 'sample_D3'] + sample = ["D0", "sample_D1", "sample_D2", "sample_D3"] # Calculate figure dimensions num_samples = len(sample) if plot_perturbation else 0 @@ -1881,7 +2345,7 @@ def create_comprehensive_plot( plt.ioff() # Close any existing figures to prevent memory leaks - plt.close('all') + plt.close("all") # Create figure with custom gridspec # 2 header rows + n sample rows (0 sample rows when no perturbation files) @@ -1901,21 +2365,25 @@ def create_comprehensive_plot( fig = plt.figure(figsize=figsize) # Create gridspec with dynamic number of rows using 27 columns for flexibility - gs = gridspec.GridSpec(num_rows, 27, figure=fig, - hspace=0.4, # Vertical spacing between rows - wspace=0.5, # Horizontal spacing between columns - left=0.10, # Left margin - right=0.90) # Right margin + gs = gridspec.GridSpec( + num_rows, + 27, + figure=fig, + hspace=0.4, # Vertical spacing between rows + wspace=0.5, # Horizontal spacing between columns + left=0.10, # Left margin + right=0.90, + ) # Right margin # Row 0: 4 square plots — UMAP Expression | UMAP Perturbation | Gene Dotplot | Top Program - ax_umap_expr = fig.add_subplot(gs[0, 1:6]) # UMAP Expression - ax_umap_pert = fig.add_subplot(gs[0, 7:12]) # UMAP Perturbation - ax_dotplot = fig.add_subplot(gs[0, 13:18]) # Gene Expression Dotplot - ax_top_prog = fig.add_subplot(gs[0, 19:24]) # Top Program + ax_umap_expr = fig.add_subplot(gs[0, 1:6]) # UMAP Expression + ax_umap_pert = fig.add_subplot(gs[0, 7:12]) # UMAP Perturbation + ax_dotplot = fig.add_subplot(gs[0, 13:18]) # Gene Expression Dotplot + ax_top_prog = fig.add_subplot(gs[0, 19:24]) # Top Program # Row 1: KD grouped bar chart | Correlation analysis - ax_kd_grouped = fig.add_subplot(gs[1, 1:18]) # KD grouped bar chart - ax_corr = fig.add_subplot(gs[1, 19:24]) # Correlation analysis + ax_kd_grouped = fig.add_subplot(gs[1, 1:18]) # KD grouped bar chart + ax_corr = fig.add_subplot(gs[1, 19:24]) # Correlation analysis # Initialize list with header axes axes = [ax_umap_expr, ax_umap_pert, ax_top_prog, ax_dotplot, ax_corr] @@ -1925,10 +2393,10 @@ def create_comprehensive_plot( for sample_idx in range(num_samples): row_idx = sample_idx + 2 # Start from row 2 (after 2 header rows) - ax_log2fc = fig.add_subplot(gs[row_idx, 1:6]) # Log2FC + ax_log2fc = fig.add_subplot(gs[row_idx, 1:6]) # Log2FC ax_volcano = fig.add_subplot(gs[row_idx, 7:12]) # Volcano - ax_s_dotplot = fig.add_subplot(gs[row_idx, 13:18]) # Dotplot - ax_waterfall = fig.add_subplot(gs[row_idx, 20:25]) # Waterfall + ax_s_dotplot = fig.add_subplot(gs[row_idx, 13:18]) # Dotplot + ax_waterfall = fig.add_subplot(gs[row_idx, 20:25]) # Waterfall axes.extend([ax_log2fc, ax_volcano, ax_s_dotplot, ax_waterfall]) @@ -1942,11 +2410,13 @@ def create_comprehensive_plot( ax=ax_umap_expr, umap_subsample_frac=umap_subsample_frac, gene_name_key=gene_name_key, - data_key=data_key + data_key=data_key, ) - ax_umap_expr.set_title(f"{Target_Gene} Expression", fontsize=18, fontweight='bold', loc='center') - ax_umap_expr.set_xlabel('UMAP 1', fontsize=14, fontweight='bold') - ax_umap_expr.set_ylabel('UMAP 2', fontsize=14, fontweight='bold') + ax_umap_expr.set_title( + f"{Target_Gene} Expression", fontsize=18, fontweight="bold", loc="center" + ) + ax_umap_expr.set_xlabel("UMAP 1", fontsize=14, fontweight="bold") + ax_umap_expr.set_ylabel("UMAP 2", fontsize=14, fontweight="bold") # Row 0, Plot 2: UMAP Perturbation ax_umap_pert = plot_umap_per_gene_guide( @@ -1959,11 +2429,13 @@ def create_comprehensive_plot( data_key=data_key, prog_key=prog_key, guide_targets_key=guide_targets_key, - guide_assignment_key=guide_assignment_key + guide_assignment_key=guide_assignment_key, + ) + ax_umap_pert.set_title( + f"{Target_Gene} Perturbation", fontsize=18, fontweight="bold", loc="center" ) - ax_umap_pert.set_title(f"{Target_Gene} Perturbation", fontsize=18, fontweight='bold', loc='center') - ax_umap_pert.set_xlabel('UMAP 1', fontsize=14, fontweight='bold') - ax_umap_pert.set_ylabel('UMAP 2', fontsize=14, fontweight='bold') + ax_umap_pert.set_xlabel("UMAP 1", fontsize=14, fontweight="bold") + ax_umap_pert.set_ylabel("UMAP 2", fontsize=14, fontweight="bold") # Row 0, Plot 3: Gene expression dotplot ax_dotplot = perturbed_gene_dotplot( @@ -1974,11 +2446,13 @@ def create_comprehensive_plot( figsize=(3, 2), ax=ax_dotplot, gene_name_key=gene_name_key, - data_key=data_key + data_key=data_key, + ) + ax_dotplot.set_title( + f"{Target_Gene} Expression", fontsize=20, fontweight="bold", loc="center" ) - ax_dotplot.set_title(f"{Target_Gene} Expression", fontsize=20, fontweight='bold', loc='center') - ax_dotplot.set_ylabel('Gene', fontsize=14, fontweight='bold', loc='center') - ax_dotplot.set_xlabel('Day', fontsize=14, fontweight='bold', loc='center') + ax_dotplot.set_ylabel("Gene", fontsize=14, fontweight="bold", loc="center") + ax_dotplot.set_xlabel("Day", fontsize=14, fontweight="bold", loc="center") # Row 0, Plot 4: Top program ax_top_prog = plot_top_program_per_gene( @@ -1990,11 +2464,20 @@ def create_comprehensive_plot( figsize=(2, 4), gene_name_key=gene_name_key, data_key=data_key, - prog_key=prog_key + prog_key=prog_key, + ) + ax_top_prog.set_title( + f"Top Loading Program for {Target_Gene}", + fontsize=20, + fontweight="bold", + loc="center", + ) + ax_top_prog.set_xlabel( + "Gene Loading Score (z-score)", fontsize=14, fontweight="bold" + ) + ax_top_prog.set_ylabel( + f"Program Name", fontsize=14, fontweight="bold", loc="center" ) - ax_top_prog.set_title(f"Top Loading Program for {Target_Gene}", fontsize=20, fontweight='bold', loc='center') - ax_top_prog.set_xlabel('Gene Loading Score (z-score)', fontsize=14, fontweight='bold') - ax_top_prog.set_ylabel(f'Program Name', fontsize=14, fontweight='bold', loc='center') # Row 1, Plot 1: KD grouped bar chart (All + per-condition on one axis) plot_perturbation_vs_control_by_condition( @@ -2016,11 +2499,18 @@ def create_comprehensive_plot( Target=Target_Gene, top_corr_genes=top_corr_genes, figsize=(4, 3), - ax=ax_corr + ax=ax_corr, + ) + ax_corr.set_title( + f"Gene Loading Correlation for {Target_Gene}", + fontsize=20, + fontweight="bold", + loc="center", + ) + ax_corr.set_ylabel("Gene Name", fontsize=14, fontweight="bold", loc="center") + ax_corr.set_xlabel( + "Correlation Coefficient", fontsize=14, fontweight="bold", loc="center" ) - ax_corr.set_title(f"Gene Loading Correlation for {Target_Gene}", fontsize=20, fontweight='bold', loc='center') - ax_corr.set_ylabel('Gene Name', fontsize=14, fontweight='bold', loc='center') - ax_corr.set_xlabel('Correlation Coefficient', fontsize=14, fontweight='bold', loc='center') # Loop through samples and create 4 plots per sample (Log2FC, Volcano, Dotplot, Waterfall) ax_index = 5 # Start after header axes (0: umap_expr, 1: umap_pert, 2: top_prog, 3: dotplot, 4: corr) @@ -2041,12 +2531,21 @@ def create_comprehensive_plot( figsize=(4, 3), ax=current_ax, Day=samp, - perturb_df=perturb_df + perturb_df=perturb_df, ) ax_index += 1 - current_ax.set_xlabel('Effect on Program Expression', fontsize=14, fontweight='bold', loc='center') - current_ax.set_ylabel("Program Name", fontsize=14, fontweight='bold', loc='center') - current_ax.set_title(f"Program Log2 Fold-Change, {samp} \n {Target_Gene}", fontsize=20, fontweight='bold', loc='center') + current_ax.set_xlabel( + "Effect on Program Expression", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_ylabel( + "Program Name", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_title( + f"Program Log2 Fold-Change, {samp} \n {Target_Gene}", + fontsize=20, + fontweight="bold", + loc="center", + ) # Plot 2: Volcano plot current_ax = axes[ax_index] @@ -2061,35 +2560,59 @@ def create_comprehensive_plot( ax=current_ax, Day=samp, perturb_df=perturb_df, - run_adjust_text=False + run_adjust_text=False, ) ax_index += 1 - current_ax.set_title(f"Volcano Plot, {samp} \n {Target_Gene}", fontsize=20, fontweight='bold', loc='center') - current_ax.set_xlabel('Effect on Program Expression', fontsize=14, fontweight='bold', loc='center') - current_ax.set_ylabel("Adjusted p-value, -log10", fontsize=14, fontweight='bold', loc='center') + current_ax.set_title( + f"Volcano Plot, {samp} \n {Target_Gene}", + fontsize=20, + fontweight="bold", + loc="center", + ) + current_ax.set_xlabel( + "Effect on Program Expression", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_ylabel( + "Adjusted p-value, -log10", fontsize=14, fontweight="bold", loc="center" + ) for t in text: t.set_fontsize(14) - adjust_text(text, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5), ax=current_ax, fontsize=14) + adjust_text( + text, + arrowprops=dict(arrowstyle="-", color="gray", lw=0.5), + ax=current_ax, + fontsize=14, + ) # Plot 3: Programs dotplot current_ax = axes[ax_index] current_ax = programs_dotplot( mdata=mdata, - program_list=(df[perturb_program_col].tolist() - if perturb_program_col in df.columns else []), + program_list=( + df[perturb_program_col].tolist() + if perturb_program_col in df.columns + else [] + ), Target=Target_Gene, dotplot_groupby=dotplot_groupby, figsize=(5, 3), ax=current_ax, Day=samp, - prog_key=prog_key + prog_key=prog_key, ) ax_index += 1 - current_ax.set_title(f"Perturbed Program Loading Scores, {samp} \n {Target_Gene}", fontsize=18, fontweight='bold', loc='center') - current_ax.set_ylabel('Program Name', fontsize=14, fontweight='bold', loc='center') - current_ax.set_xlabel("Day", fontsize=14, fontweight='bold', loc='center') - - # Plot 4: Waterfall plot + current_ax.set_title( + f"Perturbed Program Loading Scores, {samp} \n {Target_Gene}", + fontsize=18, + fontweight="bold", + loc="center", + ) + current_ax.set_ylabel( + "Program Name", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_xlabel("Day", fontsize=14, fontweight="bold", loc="center") + + # Plot 4: Waterfall plot current_ax = axes[ax_index] current_ax, text = create_gene_correlation_waterfall( corr_matrix=perturb_corr_by_sample[samp], @@ -2097,72 +2620,95 @@ def create_comprehensive_plot( ax=current_ax, Day=samp, figsize=(3, 4), - run_adjust_text=False + run_adjust_text=False, ) ax_index += 1 - current_ax.set_title(f"Gene Perturbation Correlation, {samp} \n {Target_Gene}", fontsize=18, fontweight='bold', loc='center') - current_ax.set_ylabel(f'Correlation of Program\nExpression with\n{Target_Gene} Perturbation', fontsize=14, fontweight='bold', loc='center') - current_ax.set_xlabel('Perturbed Genes', fontsize=14, fontweight='bold', loc='center') + current_ax.set_title( + f"Gene Perturbation Correlation, {samp} \n {Target_Gene}", + fontsize=18, + fontweight="bold", + loc="center", + ) + current_ax.set_ylabel( + f"Correlation of Program\nExpression with\n{Target_Gene} Perturbation", + fontsize=14, + fontweight="bold", + loc="center", + ) + current_ax.set_xlabel( + "Perturbed Genes", fontsize=14, fontweight="bold", loc="center" + ) for t in text: t.set_fontsize(14) - adjust_text(text, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5), ax=current_ax, fontsize=14) + adjust_text( + text, + arrowprops=dict(arrowstyle="-", color="gray", lw=0.5), + ax=current_ax, + fontsize=14, + ) # Add main title - fig.suptitle(f'Comprehensive Analysis: {Target_Gene}', fontweight='bold', y=0.995, fontsize=30) - + fig.suptitle( + f"Comprehensive Analysis: {Target_Gene}", + fontweight="bold", + y=0.995, + fontsize=30, + ) + # Adjust layout plt.tight_layout() - + # Save if path provided if save_path and save_name and PDF: full_path = f"{save_path}/{_safe_name(save_name)}.pdf" print(f"Saving figure to {full_path}...") - fig.savefig(full_path, format='pdf', bbox_inches='tight', dpi=300) + fig.savefig(full_path, format="pdf", bbox_inches="tight", dpi=300) if save_path and save_name and not PDF: full_path = f"{save_path}/{_safe_name(save_name)}.svg" print(f"Saving figure to {full_path}...") - fig.savefig(full_path, format='svg', bbox_inches='tight', dpi=300) - + fig.savefig(full_path, format="svg", bbox_inches="tight", dpi=300) + # Control whether to display the plot if show: plt.show() - + # Always close the figure to prevent memory leak plt.close(fig) - -def process_single_gene(Target_Gene, - mdata, - perturb_path_base, - ensembl_to_symbol_file, - gene_loading_corr_matrix, - perturb_corr_by_sample, - top_n_programs=10, - dotplot_groupby='sample', - perturb_target_col="target_name", - perturb_program_col="program_name", - perturb_log2fc_col='log2FC', - top_corr_genes=5, - volcano_log2fc_min=-0.05, - volcano_log2fc_max=0.05, - significance_threshold=0.05, - save_path=None, - figsize=None, - sample=None, - square_plots=True, - show=True, - PDF=True, - gene_name_key='gene_names', - umap_dot_size=None, - umap_subsample_frac=None, - control_target_name='non-targeting', - data_key='rna', - prog_key='cNMF', - guide_targets_key='guide_targets', - guide_assignment_key='guide_assignment'): +def process_single_gene( + Target_Gene, + mdata, + perturb_path_base, + ensembl_to_symbol_file, + gene_loading_corr_matrix, + perturb_corr_by_sample, + top_n_programs=10, + dotplot_groupby="sample", + perturb_target_col="target_name", + perturb_program_col="program_name", + perturb_log2fc_col="log2FC", + top_corr_genes=5, + volcano_log2fc_min=-0.05, + volcano_log2fc_max=0.05, + significance_threshold=0.05, + save_path=None, + figsize=None, + sample=None, + square_plots=True, + show=True, + PDF=True, + gene_name_key="gene_names", + umap_dot_size=None, + umap_subsample_frac=None, + control_target_name="non-targeting", + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): """Wrapper around ``create_comprehensive_plot`` for a single gene with error handling. Catches exceptions so that one failing gene does not abort a batch run. @@ -2214,7 +2760,7 @@ def process_single_gene(Target_Gene, data_key=data_key, prog_key=prog_key, guide_targets_key=guide_targets_key, - guide_assignment_key=guide_assignment_key + guide_assignment_key=guide_assignment_key, ) return f"Success: {Target_Gene}" @@ -2229,6 +2775,7 @@ def process_single_gene(Target_Gene, # child processes inherit via copy-on-write (fork) without pickling. _shared_pool_data = {} + def _process_gene_worker(Target_Gene): """Worker function for ``multiprocessing.Pool.map``. @@ -2250,73 +2797,76 @@ def _process_gene_worker(Target_Gene): try: return process_single_gene( Target_Gene, - mdata=d['mdata'], - perturb_path_base=d['perturb_path_base'], - ensembl_to_symbol_file=d['ensembl_to_symbol_file'], - gene_loading_corr_matrix=d['gene_loading_corr_matrix'], - perturb_corr_by_sample=d['perturb_corr_by_sample'], - top_n_programs=d['top_n_programs'], - dotplot_groupby=d['dotplot_groupby'], - perturb_target_col=d['perturb_target_col'], - perturb_program_col=d['perturb_program_col'], - perturb_log2fc_col=d['perturb_log2fc_col'], - top_corr_genes=d['top_corr_genes'], - volcano_log2fc_min=d['volcano_log2fc_min'], - volcano_log2fc_max=d['volcano_log2fc_max'], - significance_threshold=d['significance_threshold'], - save_path=d['save_path'], - figsize=d['figsize'], - sample=d['sample'], - square_plots=d['square_plots'], - show=d['show'], - PDF=d['PDF'], - gene_name_key=d['gene_name_key'], - umap_dot_size=d['umap_dot_size'], - umap_subsample_frac=d['umap_subsample_frac'], - control_target_name=d['control_target_name'], - data_key=d['data_key'], - prog_key=d['prog_key'], - guide_targets_key=d['guide_targets_key'], - guide_assignment_key=d['guide_assignment_key'], + mdata=d["mdata"], + perturb_path_base=d["perturb_path_base"], + ensembl_to_symbol_file=d["ensembl_to_symbol_file"], + gene_loading_corr_matrix=d["gene_loading_corr_matrix"], + perturb_corr_by_sample=d["perturb_corr_by_sample"], + top_n_programs=d["top_n_programs"], + dotplot_groupby=d["dotplot_groupby"], + perturb_target_col=d["perturb_target_col"], + perturb_program_col=d["perturb_program_col"], + perturb_log2fc_col=d["perturb_log2fc_col"], + top_corr_genes=d["top_corr_genes"], + volcano_log2fc_min=d["volcano_log2fc_min"], + volcano_log2fc_max=d["volcano_log2fc_max"], + significance_threshold=d["significance_threshold"], + save_path=d["save_path"], + figsize=d["figsize"], + sample=d["sample"], + square_plots=d["square_plots"], + show=d["show"], + PDF=d["PDF"], + gene_name_key=d["gene_name_key"], + umap_dot_size=d["umap_dot_size"], + umap_subsample_frac=d["umap_subsample_frac"], + control_target_name=d["control_target_name"], + data_key=d["data_key"], + prog_key=d["prog_key"], + guide_targets_key=d["guide_targets_key"], + guide_assignment_key=d["guide_assignment_key"], ) finally: # Each worker is reused for many genes (pool.map default chunking), # so explicit cleanup prevents matplotlib/pandas state from accumulating per worker. import gc - plt.close('all') + + plt.close("all") gc.collect() -def parallel_gene_processing(perturbed_gene_list, - mdata, - perturb_path_base, - ensembl_to_symbol_file, - gene_loading_corr_matrix, - perturb_corr_by_sample, - top_n_programs=10, - dotplot_groupby='sample', - perturb_target_col="target_name", - perturb_program_col="program_name", - perturb_log2fc_col='log2FC', - top_corr_genes=5, - volcano_log2fc_min=-0.05, - volcano_log2fc_max=0.05, - significance_threshold=0.05, - save_path=None, - figsize=None, - sample=None, - square_plots=True, - show=True, - PDF=True, - n_processes=-1, - gene_name_key='gene_names', - umap_dot_size=None, - umap_subsample_frac=None, - control_target_name='non-targeting', - data_key='rna', - prog_key='cNMF', - guide_targets_key='guide_targets', - guide_assignment_key='guide_assignment'): +def parallel_gene_processing( + perturbed_gene_list, + mdata, + perturb_path_base, + ensembl_to_symbol_file, + gene_loading_corr_matrix, + perturb_corr_by_sample, + top_n_programs=10, + dotplot_groupby="sample", + perturb_target_col="target_name", + perturb_program_col="program_name", + perturb_log2fc_col="log2FC", + top_corr_genes=5, + volcano_log2fc_min=-0.05, + volcano_log2fc_max=0.05, + significance_threshold=0.05, + save_path=None, + figsize=None, + sample=None, + square_plots=True, + show=True, + PDF=True, + n_processes=-1, + gene_name_key="gene_names", + umap_dot_size=None, + umap_subsample_frac=None, + control_target_name="non-targeting", + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): """Process multiple genes in parallel using fork-based multiprocessing. Stores all shared data (mdata, correlation matrices, etc.) in the @@ -2350,37 +2900,39 @@ def parallel_gene_processing(perturbed_gene_list, # Populate module globals BEFORE fork so workers inherit via copy-on-write _shared_pool_data = { - 'mdata': mdata, - 'perturb_path_base': perturb_path_base, - 'ensembl_to_symbol_file': ensembl_to_symbol_file, - 'gene_loading_corr_matrix': gene_loading_corr_matrix, - 'perturb_corr_by_sample': perturb_corr_by_sample, - 'top_n_programs': top_n_programs, - 'dotplot_groupby': dotplot_groupby, - 'perturb_target_col': perturb_target_col, - 'perturb_program_col': perturb_program_col, - 'perturb_log2fc_col': perturb_log2fc_col, - 'top_corr_genes': top_corr_genes, - 'volcano_log2fc_min': volcano_log2fc_min, - 'volcano_log2fc_max': volcano_log2fc_max, - 'significance_threshold': significance_threshold, - 'save_path': save_path, - 'figsize': figsize, - 'sample': sample, - 'square_plots': square_plots, - 'show': show, - 'PDF': PDF, - 'gene_name_key': gene_name_key, - 'umap_dot_size': umap_dot_size, - 'umap_subsample_frac': umap_subsample_frac, - 'control_target_name': control_target_name, - 'data_key': data_key, - 'prog_key': prog_key, - 'guide_targets_key': guide_targets_key, - 'guide_assignment_key': guide_assignment_key, + "mdata": mdata, + "perturb_path_base": perturb_path_base, + "ensembl_to_symbol_file": ensembl_to_symbol_file, + "gene_loading_corr_matrix": gene_loading_corr_matrix, + "perturb_corr_by_sample": perturb_corr_by_sample, + "top_n_programs": top_n_programs, + "dotplot_groupby": dotplot_groupby, + "perturb_target_col": perturb_target_col, + "perturb_program_col": perturb_program_col, + "perturb_log2fc_col": perturb_log2fc_col, + "top_corr_genes": top_corr_genes, + "volcano_log2fc_min": volcano_log2fc_min, + "volcano_log2fc_max": volcano_log2fc_max, + "significance_threshold": significance_threshold, + "save_path": save_path, + "figsize": figsize, + "sample": sample, + "square_plots": square_plots, + "show": show, + "PDF": PDF, + "gene_name_key": gene_name_key, + "umap_dot_size": umap_dot_size, + "umap_subsample_frac": umap_subsample_frac, + "control_target_name": control_target_name, + "data_key": data_key, + "prog_key": prog_key, + "guide_targets_key": guide_targets_key, + "guide_assignment_key": guide_assignment_key, } - print(f"Starting parallel processing of {len(perturbed_gene_list)} genes using {n_processes} processes...") + print( + f"Starting parallel processing of {len(perturbed_gene_list)} genes using {n_processes} processes..." + ) # Pool uses fork on Linux — globals inherited without pickling with Pool(processes=n_processes) as pool: @@ -2390,4 +2942,4 @@ def parallel_gene_processing(perturbed_gene_list, print("Parallel processing completed!") - return results \ No newline at end of file + return results diff --git a/src/Stage3_Interpretation/A_Plotting/src/Program_QC_plots.py b/src/Stage3_Interpretation/A_Plotting/src/Program_QC_plots.py index ca31dda..0ec55cc 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/Program_QC_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/Program_QC_plots.py @@ -1,4 +1,4 @@ -import muon as mu +import muon as mu import scanpy as sc import numpy as np import pandas as pd @@ -19,40 +19,58 @@ plt.rcParams["axes.spines.top"] = False plt.rcParams["axes.spines.right"] = False -#plt.rcParams["axes.spines.bottom"] = False -#plt.rcParams["axes.spines.left"] = False -plt.rcParams['pdf.fonttype'] = 42 # TrueType fonts (editable text) -plt.rcParams['ps.fonttype'] = 42 # For EPS as well -plt.rcParams['svg.fonttype'] = 'none' # SVG: write real (editable in Illustrator) - - -plt.rcParams.update({ - 'font.size': 16, - 'axes.titlesize': 16, - 'axes.labelsize': 16, - 'legend.fontsize': 16, # Legend text size - 'legend.title_fontsize': 14, # Legend title size - 'xtick.labelsize': 16, - 'ytick.labelsize': 16, - 'figure.titlesize': 20 -}) +# plt.rcParams["axes.spines.bottom"] = False +# plt.rcParams["axes.spines.left"] = False +plt.rcParams["pdf.fonttype"] = 42 # TrueType fonts (editable text) +plt.rcParams["ps.fonttype"] = 42 # For EPS as well +plt.rcParams["svg.fonttype"] = ( + "none" # SVG: write real (editable in Illustrator) +) + + +plt.rcParams.update( + { + "font.size": 16, + "axes.titlesize": 16, + "axes.labelsize": 16, + "legend.fontsize": 16, # Legend text size + "legend.title_fontsize": 14, # Legend title size + "xtick.labelsize": 16, + "ytick.labelsize": 16, + "figure.titlesize": 20, + } +) import sys -# Change path to wherever you have repo locally -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) -from .utilities import convert_adata_with_mygene, convert_with_mygene, rename_list_gene_dictionary, rename_adata_gene_dictionary +# Change path to wherever you have repo locally +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) + +from .utilities import ( + convert_adata_with_mygene, + convert_with_mygene, + rename_list_gene_dictionary, + rename_adata_gene_dictionary, +) from .Program_expression_weighted_plots import plot_program_heatmap_weighted from .efficient_corr import column_pearson_streaming - - # make UMAP of program score -def plot_umap_per_program(mdata, Target_Program, ax=None, - color='purple', save_path=None, save_name=None, figsize=(8,6), show=False, - subsample_frac=None, random_state=42, prog_key='cNMF'): +def plot_umap_per_program( + mdata, + Target_Program, + ax=None, + color="purple", + save_path=None, + save_name=None, + figsize=(8, 6), + show=False, + subsample_frac=None, + random_state=42, + prog_key="cNMF", +): # read cell and program data adata_plot = mdata[prog_key].copy() @@ -64,10 +82,10 @@ def plot_umap_per_program(mdata, Target_Program, ax=None, n_sample = int(n_cells * subsample_frac) indices = np.random.choice(n_cells, size=n_sample, replace=False) adata_plot = adata_plot[sorted(indices)].copy() - - colors = ['lightgrey', color] + + colors = ["lightgrey", color] n_bins = 100 - cmap = mcolors.LinearSegmentedColormap.from_list('custom', colors, N=n_bins) + cmap = mcolors.LinearSegmentedColormap.from_list("custom", colors, N=n_bins) title = f"Program {Target_Program} Expression" @@ -78,18 +96,20 @@ def plot_umap_per_program(mdata, Target_Program, ax=None, else: fig = ax.get_figure() standalone = False - + # Plot on the provided/created axis plt.sca(ax) # Set current axis to ensure scanpy uses correct figure -CC change - sc.pl.umap(adata_plot, color=str(Target_Program), title=title, cmap=cmap, ax=ax, show=False) - + sc.pl.umap( + adata_plot, color=str(Target_Program), title=title, cmap=cmap, ax=ax, show=False + ) + # Rasterize ONLY the scatter points (collections) in this axis for collection in ax.collections: collection.set_rasterized(True) - - ax.set_title(title, fontsize=18, fontweight='bold', loc='center') - ax.set_xlabel('UMAP 1', fontsize=10, fontweight = 'bold') - ax.set_ylabel('UMAP 2', fontsize=10, fontweight = 'bold') + + ax.set_title(title, fontsize=18, fontweight="bold", loc="center") + ax.set_xlabel("UMAP 1", fontsize=10, fontweight="bold") + ax.set_ylabel("UMAP 2", fontsize=10, fontweight="bold") # Adjust layout (only in standalone mode) if standalone: @@ -97,41 +117,62 @@ def plot_umap_per_program(mdata, Target_Program, ax=None, # Save if path provided (only in standalone mode) if save_path and save_name: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{save_name}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only in standalone mode) if show: plt.show() else: plt.close(fig) - - return ax - + return ax -# plot x top loading genes for target program -def plot_top_gene_per_program(mdata, Target_Program, - num_gene=5, ax=None, - save_path=None, save_name=None, - figsize=(5,8), show=False, file_to_dictionary=None, - gene_name_key='gene_names', data_key='rna', prog_key='cNMF'): +# plot x top loading genes for target program +def plot_top_gene_per_program( + mdata, + Target_Program, + num_gene=5, + ax=None, + save_path=None, + save_name=None, + figsize=(5, 8), + show=False, + file_to_dictionary=None, + gene_name_key="gene_names", + data_key="rna", + prog_key="cNMF", +): # Read cNMF gene program matrix X = mdata[prog_key].varm["loadings"] - gene_names = mdata[prog_key].uns['var_names'] + gene_names = mdata[prog_key].uns["var_names"] # Map Ensembl IDs to gene symbols using gene_name_key column when available if gene_name_key is not None and gene_name_key in mdata[data_key].var.columns: - ensembl_to_symbol = dict(zip(mdata[data_key].var_names, mdata[data_key].var[gene_name_key].astype(str))) + ensembl_to_symbol = dict( + zip( + mdata[data_key].var_names, + mdata[data_key].var[gene_name_key].astype(str), + ) + ) gene_names = np.array([ensembl_to_symbol.get(g, g) for g in gene_names]) # Rename genes if file_to_dictionary is None: - df_renamed = pd.DataFrame(data=X, columns=gene_names, index=mdata[prog_key].var_names) + df_renamed = pd.DataFrame( + data=X, columns=gene_names, index=mdata[prog_key].var_names + ) else: renamed_gene_list = rename_list_gene_dictionary(gene_names, file_to_dictionary) - df_renamed = pd.DataFrame(data=X, columns=renamed_gene_list, index=mdata[prog_key].var_names) + df_renamed = pd.DataFrame( + data=X, columns=renamed_gene_list, index=mdata[prog_key].var_names + ) # Validate the program exists, and is unique, in the loading matrix. matching_rows = (df_renamed.index == str(Target_Program)).sum() @@ -157,7 +198,7 @@ def plot_top_gene_per_program(mdata, Target_Program, # Sort top x genes for the program df_sorted = df_renamed.loc[str(Target_Program)].nlargest(num_gene) df_sorted = df_sorted.sort_values(ascending=True) - + # If no axis provided, create new figure if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -165,58 +206,74 @@ def plot_top_gene_per_program(mdata, Target_Program, else: fig = ax.get_figure() standalone = False - + # Create horizontal bar plot - bars = ax.barh(range(len(df_sorted)), df_sorted.values, color='#808080', alpha=0.8) - + bars = ax.barh(range(len(df_sorted)), df_sorted.values, color="#808080", alpha=0.8) + # Customize the plot - ax.set_title(f"Top Loading Genes for Program {Target_Program}", - fontsize=12, fontweight='bold', loc='center') + ax.set_title( + f"Top Loading Genes for Program {Target_Program}", + fontsize=12, + fontweight="bold", + loc="center", + ) ax.set_yticks(range(len(df_sorted))) ax.set_yticklabels(df_sorted.index, fontsize=10) # Keep gene names - ax.set_xlabel('Gene Loading Score (z-score)', fontsize=10, fontweight='bold', loc='center') - ax.set_ylabel('Gene Name', fontsize=10, fontweight='bold', loc='center') - ax.set_yticklabels(df_sorted.index, fontsize=10, rotation=45, ha='right') + ax.set_xlabel( + "Gene Loading Score (z-score)", fontsize=10, fontweight="bold", loc="center" + ) + ax.set_ylabel("Gene Name", fontsize=10, fontweight="bold", loc="center") + ax.set_yticklabels(df_sorted.index, fontsize=10, rotation=45, ha="right") - # Format x-axis ax.set_xlim(0, max(df_sorted.values) * 1.1) - ax.ticklabel_format(style='scientific', axis='x') - + ax.ticklabel_format(style="scientific", axis="x") + # Add grid - ax.grid(axis='x', alpha=0.3, linestyle='-', linewidth=0.5) + ax.grid(axis="x", alpha=0.3, linestyle="-", linewidth=0.5) ax.set_axisbelow(True) - + # Remove top and right spines - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + if standalone: plt.tight_layout() - + # Save if path and name provided if save_path and save_name: - fig.savefig(f"{save_path}/{save_name}.svg", format='svg', - bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) + # Show or close if show: plt.show() else: plt.close(fig) - - return ax - + return ax # plot x top enriched GO term for target program -def top_GO_per_program(GO_path, Target_Program, num_term = 5, p_value_name = "Adjusted P-value", - term_col = "Term", - save_path = None, save_name = None, ax=None, figsize=(8,6), show=False): +def top_GO_per_program( + GO_path, + Target_Program, + num_term=5, + p_value_name="Adjusted P-value", + term_col="Term", + save_path=None, + save_name=None, + ax=None, + figsize=(8, 6), + show=False, +): # read txt file - df = pd.read_csv(GO_path, sep='\t', index_col=0) + df = pd.read_csv(GO_path, sep="\t", index_col=0) df.index = df.index.astype(str) # local to a program and isolate Term @@ -224,8 +281,14 @@ def top_GO_per_program(GO_path, Target_Program, num_term = 5, p_value_name = "Ad print(f"Program {Target_Program} not found in GO data") if ax is None: fig, ax = plt.subplots(figsize=figsize) - ax.text(0.5, 0.5, 'No GO data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No GO data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax, [] df_program = df.loc[str(Target_Program)] @@ -249,66 +312,76 @@ def top_GO_per_program(GO_path, Target_Program, num_term = 5, p_value_name = "Ad # Wrap long labels - break at natural points import textwrap + wrapped_labels = [] for label in df_sort_log.index: # Try to wrap at 40 characters, preferring breaks at spaces - wrapped = textwrap.fill(label, width=30, break_long_words=False, break_on_hyphens=False) + wrapped = textwrap.fill( + label, width=30, break_long_words=False, break_on_hyphens=False + ) wrapped_labels.append(wrapped) - # Compute bar spacing based on number of wrapped lines to avoid label overlap - line_counts = [wrapped.count('\n') + 1 for wrapped in wrapped_labels] - max_lines = max(line_counts) + line_counts = [wrapped.count("\n") + 1 for wrapped in wrapped_labels] + max_lines = max(line_counts) spacing = max(1.0, max_lines * 0.7) positions = [i * spacing for i in range(len(wrapped_labels))] # Create horizontal bar plot - #bars = ax.barh(range(len(df_sort_log)), df_sort_log.values, color='#808080', alpha=0.8) - bars = ax.barh(positions, df_sort_log.values, color='#808080', alpha=0.8, height=0.8 * spacing) + # bars = ax.barh(range(len(df_sort_log)), df_sort_log.values, color='#808080', alpha=0.8) + bars = ax.barh( + positions, df_sort_log.values, color="#808080", alpha=0.8, height=0.8 * spacing + ) # Customize the plot - ax.set_title(f"GO enrichment for Program {Target_Program}",fontsize=14, fontweight='bold', loc='center') - #ax.set_yticks(range(len(df_sort_log))) + ax.set_title( + f"GO enrichment for Program {Target_Program}", + fontsize=14, + fontweight="bold", + loc="center", + ) + # ax.set_yticks(range(len(df_sort_log))) ax.set_yticks(positions) - ax.set_yticklabels(wrapped_labels, fontsize=6, fontweight='bold') - - ax.set_xlabel('-log10 Adjusted P-value',fontsize=11, fontweight='bold', loc='center') + ax.set_yticklabels(wrapped_labels, fontsize=6, fontweight="bold") + + ax.set_xlabel( + "-log10 Adjusted P-value", fontsize=11, fontweight="bold", loc="center" + ) # Format x-axis to match your reference - ax.set_xlim(0, max(df_sort_log.values)* 1.05) # keep the plotting part smaller - ax.ticklabel_format(style='scientific', axis='x', scilimits=(0,0)) + ax.set_xlim(0, max(df_sort_log.values) * 1.05) # keep the plotting part smaller + ax.ticklabel_format(style="scientific", axis="x", scilimits=(0, 0)) plt.subplots_adjust(left=0.3, right=0.7, top=0.93, bottom=0.08) # Add grid - ax.grid(axis='x', alpha=0.3, linestyle='-', linewidth=0.5) + ax.grid(axis="x", alpha=0.3, linestyle="-", linewidth=0.5) ax.set_axisbelow(True) # Remove top and right spines - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) # Adjust layout plt.tight_layout() - + # Only save if standalone and save_path provided if standalone and save_path and save_name: - fig.savefig(f"{save_path}/{save_name}.svg", format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) + # Only show/close if standalone if standalone: if show: plt.show() else: plt.close(fig) - - return ax,wrapped_labels - + return ax, wrapped_labels # helper method to compute corr for analyze_correlations -def compute_program_correlation_matrix(mdata, prog_key='cNMF', save_path=None): +def compute_program_correlation_matrix(mdata, prog_key="cNMF", save_path=None): """Program x program Pearson correlation of cNMF usages across cells. Same values as ``pd.DataFrame(mdata[prog_key].X).corr().fillna(0)``, but @@ -335,13 +408,24 @@ def compute_program_correlation_matrix(mdata, prog_key='cNMF', save_path=None): if save_path is not None: Path(save_path).parent.mkdir(parents=True, exist_ok=True) - np.savez(save_path, corr=R.astype(np.float32), program=np.asarray(names, dtype=str)) + np.savez( + save_path, corr=R.astype(np.float32), program=np.asarray(names, dtype=str) + ) return pd.DataFrame(R, index=names, columns=names).fillna(0) + # find most and least simliar programs -def analyze_program_correlations(program_correlation, Target_Program , num_program = 5, -save_path=None, save_name = None, figsize = (5, 4),show=False, ax=None): +def analyze_program_correlations( + program_correlation, + Target_Program, + num_program=5, + save_path=None, + save_name=None, + figsize=(5, 4), + show=False, + ax=None, +): # check if program exists if str(Target_Program) not in program_correlation.columns: @@ -354,14 +438,22 @@ def analyze_program_correlations(program_correlation, Target_Program , num_progr img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No correlation to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No correlation to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax # Get correlations with the target program target_correlations = program_correlation.loc[str(Target_Program)] - target_correlations = target_correlations.drop(str(Target_Program), errors="ignore") # Remove self-correlation - + target_correlations = target_correlations.drop( + str(Target_Program), errors="ignore" + ) # Remove self-correlation + # Sort correlations sorted_correlations = target_correlations.sort_values(ascending=False) @@ -377,83 +469,105 @@ def analyze_program_correlations(program_correlation, Target_Program , num_progr fig, ax = plt.subplots(figsize=figsize) else: fig = ax.get_figure() - - + # Create horizontal bar plot - colors = ['blue' if x < 0 else 'red' for x in combined_correlations] - - bars = ax.barh(range(len(combined_correlations)), combined_correlations.values, - color=colors, alpha=0.7) + colors = ["blue" if x < 0 else "red" for x in combined_correlations] + + bars = ax.barh( + range(len(combined_correlations)), + combined_correlations.values, + color=colors, + alpha=0.7, + ) # Customize the plot - ax.set_title( f"Programs with Correlated Loading Across Cells, \n for Program {Target_Program}", fontsize=14, fontweight='bold', pad=20) - ax.set_ylabel('Program Name', fontsize=12, fontweight='bold') - ax.set_xlabel('Correlation Coefficient', fontsize=12, fontweight='bold') - + ax.set_title( + f"Programs with Correlated Loading Across Cells, \n for Program {Target_Program}", + fontsize=14, + fontweight="bold", + pad=20, + ) + ax.set_ylabel("Program Name", fontsize=12, fontweight="bold") + ax.set_xlabel("Correlation Coefficient", fontsize=12, fontweight="bold") + # Set x-axis labels - ax.set_yticks(range(len(combined_correlations)), - combined_correlations.index, - #rotation=45, - ha='right') + ax.set_yticks( + range(len(combined_correlations)), + combined_correlations.index, + # rotation=45, + ha="right", + ) - # Add a vertical line at x=0 - ax.axvline(x=0, color='black', linestyle='-', linewidth=0.5) - - + ax.axvline(x=0, color="black", linestyle="-", linewidth=0.5) + # Add legend from matplotlib.patches import Patch - #legend_elements = [Patch(facecolor='red', alpha=0.7, label='Positive correlation'), + # legend_elements = [Patch(facecolor='red', alpha=0.7, label='Positive correlation'), # Patch(facecolor='blue', alpha=0.7, label='Negative correlation')] - #ax.legend(handles=legend_elements)# loc='lower right') - - # Adjust layout (only in standalone mode) + # ax.legend(handles=legend_elements)# loc='lower right') + + # Adjust layout (only in standalone mode) # Add grid - ax.grid(axis='x', alpha=0.3) + ax.grid(axis="x", alpha=0.3) - # Save if path provided (only in standalone mode) if save_path and ax is None: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) + # Control whether to display the plot (only in standalone mode) if ax is None: if show: plt.show() else: plt.close(fig) - - return ax - + return ax -# plot violin plots for program expression on target program -def plot_violin(mdata, Target_Program, save_path=None, save_name=None, groupby = 'sample', figsize=(3, 5), show=False, ax=None, data_key='rna', prog_key='cNMF'): +# plot violin plots for program expression on target program +def plot_violin( + mdata, + Target_Program, + save_path=None, + save_name=None, + groupby="sample", + figsize=(3, 5), + show=False, + ax=None, + data_key="rna", + prog_key="cNMF", +): # Build dataframe from cNMF loadings X = mdata[prog_key].X - if hasattr(X, 'toarray'): + if hasattr(X, "toarray"): X = X.toarray() - df = pd.DataFrame(data=X, index=mdata[prog_key].obs_names, columns=mdata[prog_key].var_names) + df = pd.DataFrame( + data=X, index=mdata[prog_key].obs_names, columns=mdata[prog_key].var_names + ) # Create dataframe with expression and cell type - df = pd.DataFrame({ - "expression": df[str(Target_Program)], - "cell_type": mdata[data_key].obs[groupby].values - }) - + df = pd.DataFrame( + { + "expression": df[str(Target_Program)], + "cell_type": mdata[data_key].obs[groupby].values, + } + ) + # Compute summary stats per cell_type summary = ( df.groupby("cell_type") .agg( mean_expr=("expression", "mean"), - frac_cells=("expression", lambda x: (x > x.mean()).mean()) + frac_cells=("expression", lambda x: (x > x.mean()).mean()), ) .reset_index() ) - + # Create figure/axes if not provided if ax is None: fig, ax = plt.subplots(figsize=figsize) @@ -468,28 +582,40 @@ def plot_violin(mdata, Target_Program, save_path=None, save_name=None, groupby = x="cell_type", y="expression", inner="quartile", - density_norm='width', + density_norm="width", cut=0, - ax=ax + ax=ax, ) # Annotate mean & fraction above violins - cell_types_order = df['cell_type'].cat.categories if hasattr(df['cell_type'], 'cat') else sorted(df['cell_type'].unique()) + cell_types_order = ( + df["cell_type"].cat.categories + if hasattr(df["cell_type"], "cat") + else sorted(df["cell_type"].unique()) + ) for i, row in summary.iterrows(): # Get x position based on cell type order - x_pos = list(cell_types_order).index(row['cell_type']) + x_pos = list(cell_types_order).index(row["cell_type"]) ax.text( x_pos, # Use correct position - df["expression"].max()*0.90, # Just above max y + df["expression"].max() * 0.90, # Just above max y f"Mean={row['mean_expr']:.2f}\nFrac={row['frac_cells']:.2f}", - ha="center", va="bottom", fontsize=9, color="black" + ha="center", + va="bottom", + fontsize=9, + color="black", ) # Set labels and title - ax.set_title(f"Program {Target_Program} Expression by Conditions", fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel('Program Expression', fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel('Conditions', fontsize=10, fontweight='bold', loc='center') + ax.set_title( + f"Program {Target_Program} Expression by Conditions", + fontsize=14, + fontweight="bold", + loc="center", + ) + ax.set_ylabel("Program Expression", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel("Conditions", fontsize=10, fontweight="bold", loc="center") # Adjust layout (only in standalone mode) if standalone: @@ -497,7 +623,12 @@ def plot_violin(mdata, Target_Program, save_path=None, save_name=None, groupby = # Save if path provided (only in standalone mode) if save_path and save_name: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{save_name}.svg", + format="svg", + bbox_inches="tight", + dpi=300, + ) # Control whether to display the plot (only in standalone mode) if show: @@ -506,37 +637,52 @@ def plot_violin(mdata, Target_Program, save_path=None, save_name=None, groupby = plt.close(fig) return ax - - # plot barplot for up/down regulated genes log2FC given results from perturbation analysis, return genes in df -def plot_program_log2FC(perturb_path, Target, tagert_col_name="target_name", plot_col_name="program_name", - log2fc_col='log2FC', num_item=5, p_value=0.05, save_path=None, save_name=None, gene_list = None, - figsize=(5, 4), show=False, ax=None, Day =""): +def plot_program_log2FC( + perturb_path, + Target, + tagert_col_name="target_name", + plot_col_name="program_name", + log2fc_col="log2FC", + num_item=5, + p_value=0.05, + save_path=None, + save_name=None, + gene_list=None, + figsize=(5, 4), + show=False, + ax=None, + Day="", +): # read df df = pd.read_csv(perturb_path, sep="\t") - df['program_name'] = df['program_name'].astype(str) + df["program_name"] = df["program_name"].astype(str) Target = str(Target) - # Check if query gene exists if Target not in df[tagert_col_name].values: print(f"{Target} not found in mdata") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax, pd.DataFrame(columns=df.columns) if gene_list is not None: - expressed_gene = set(df[plot_col_name].values).intersection(set(gene_list)) + expressed_gene = set(df[plot_col_name].values).intersection(set(gene_list)) df = df[df[plot_col_name].isin(expressed_gene)] - # Sort by log2FC df_sorted = df.loc[df[tagert_col_name] == Target] - df_sorted = df_sorted[df_sorted['adj_pval'] < p_value] + df_sorted = df_sorted[df_sorted["adj_pval"] < p_value] df_sorted = df_sorted.sort_values(by=log2fc_col, ascending=False) # Get top and bottom gene top_df = df_sorted.head(num_item) @@ -545,91 +691,110 @@ def plot_program_log2FC(perturb_path, Target, tagert_col_name="target_name", plo top = top_df.copy() bottom = bottom_df.copy() # Combine data - plot_data = pd.merge(top, bottom, how='outer') + plot_data = pd.merge(top, bottom, how="outer") plot_data = plot_data.sort_values(by=log2fc_col, ascending=False) - + # Create the plot if ax is None: fig, ax = plt.subplots(figsize=figsize) else: fig = ax.get_figure() - + # Create horizontal bar plot - colors = ['red' if x > 0 else 'blue' for x in plot_data[log2fc_col]] - bars = ax.barh(range(len(plot_data)), plot_data[log2fc_col], color=colors, alpha=0.7) - + colors = ["red" if x > 0 else "blue" for x in plot_data[log2fc_col]] + bars = ax.barh( + range(len(plot_data)), plot_data[log2fc_col], color=colors, alpha=0.7 + ) + # Customize the plot ax.set_yticks(range(len(plot_data))) - ax.set_yticklabels(plot_data[plot_col_name],fontsize=10) - + ax.set_yticklabels(plot_data[plot_col_name], fontsize=10) - ax.set_xlabel('Effect on Program Loadings',fontsize=10, fontweight='bold', loc='center') - ax.set_ylabel( "Regulator Name",fontsize=10, fontweight='bold', loc='center') - ax.set_title(f"Regulator Effect on Program Loadings, \n Program {Target}, {Day}",fontsize=14, fontweight='bold', loc='center') + ax.set_xlabel( + "Effect on Program Loadings", fontsize=10, fontweight="bold", loc="center" + ) + ax.set_ylabel("Regulator Name", fontsize=10, fontweight="bold", loc="center") + ax.set_title( + f"Regulator Effect on Program Loadings, \n Program {Target}, {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) # Add a vertical line at x=0 - ax.axvline(x=0, color='black', linestyle='-', linewidth=0.5) + ax.axvline(x=0, color="black", linestyle="-", linewidth=0.5) - - # Add legend - #from matplotlib.patches import Patch - #legend_elements = [Patch(facecolor='red', alpha=0.7, label='Upregulated'), + # from matplotlib.patches import Patch + # legend_elements = [Patch(facecolor='red', alpha=0.7, label='Upregulated'), # Patch(facecolor='blue', alpha=0.7, label='Downregulated')] - #ax.legend(handles=legend_elements) #loc='lower right') - + # ax.legend(handles=legend_elements) #loc='lower right') + # Adjust layout - ax.grid(axis='x', alpha=0.3) + ax.grid(axis="x", alpha=0.3) if ax is None: plt.tight_layout() - + # Save if path provided (only when ax is not provided) if save_path and ax is None: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) + # Control whether to display the plot (only when ax is not provided) if ax is None: if show: plt.show() else: plt.close(fig) - - return ax, plot_data - - - - + return ax, plot_data # plot heatmap for program perturbation on different conditions -def plot_program_heatmap(perturb_path_base, Target_Program, tagert_col_name="target_name", plot_col_name="program_name", sample= ['D0', 'sample_D1', 'sample_D2', 'sample_D3'], - log2fc_col='log2FC', p_value=0.05, save_path=None, save_name=None, groupby = 'sample', - figsize=(5, 4), show=False, ax=None): - +def plot_program_heatmap( + perturb_path_base, + Target_Program, + tagert_col_name="target_name", + plot_col_name="program_name", + sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + log2fc_col="log2FC", + p_value=0.05, + save_path=None, + save_name=None, + groupby="sample", + figsize=(5, 4), + show=False, + ax=None, +): - perturbed_df_list = [pd.read_csv(f"{perturb_path_base}_{samp}.txt", sep="\t").assign(sample=samp) for samp in sample] + perturbed_df_list = [ + pd.read_csv(f"{perturb_path_base}_{samp}.txt", sep="\t").assign(sample=samp) + for samp in sample + ] perturbed_df = pd.concat(perturbed_df_list, ignore_index=True) - perturbed_df['program_name'] = perturbed_df['program_name'].astype(str) + perturbed_df["program_name"] = perturbed_df["program_name"].astype(str) - #sort by target programs + # sort by target programs df_sorted = perturbed_df.loc[perturbed_df[tagert_col_name] == str(Target_Program)] # Get unique genes that have adj_pval > 0.05 in any sample - genes_to_keep = df_sorted[df_sorted['adj_pval'] < p_value][plot_col_name].unique() + genes_to_keep = df_sorted[df_sorted["adj_pval"] < p_value][plot_col_name].unique() # Filter df to keep only those genes - df_filtered = df_sorted[df_sorted['target_name'].isin(genes_to_keep)] - - df_pivot = df_filtered.pivot(columns = "target_name", index = groupby, values = log2fc_col) - df_pivot_pval = df_filtered.pivot(columns = "target_name", index = groupby, values = "adj_pval") + df_filtered = df_sorted[df_sorted["target_name"].isin(genes_to_keep)] + df_pivot = df_filtered.pivot( + columns="target_name", index=groupby, values=log2fc_col + ) + df_pivot_pval = df_filtered.pivot( + columns="target_name", index=groupby, values="adj_pval" + ) # Assuming your df has 'sample' column and gene columns with log2FC values - heatmap_data = df_pivot #.rename(index={'sample_D1': 'D1', 'sample_D2': 'D2', 'sample_D3': 'D3'}) - + heatmap_data = df_pivot # .rename(index={'sample_D1': 'D1', 'sample_D2': 'D2', 'sample_D3': 'D3'}) # Create the plot if ax is None: @@ -637,148 +802,236 @@ def plot_program_heatmap(perturb_path_base, Target_Program, tagert_col_name="tar else: fig = ax.get_figure() - # Create heatmap - sns.heatmap(heatmap_data, cmap='RdBu_r', center=0, - cbar_kws={'label': 'log2FC','shrink': 0.9}, ax=ax, square=True, - vmin=-1, vmax=1) + sns.heatmap( + heatmap_data, + cmap="RdBu_r", + center=0, + cbar_kws={"label": "log2FC", "shrink": 0.9}, + ax=ax, + square=True, + vmin=-1, + vmax=1, + ) - ax.set_xticklabels(ax.get_xticklabels(), fontsize=8, rotation=45, ha='right') + ax.set_xticklabels(ax.get_xticklabels(), fontsize=8, rotation=45, ha="right") ax.set_yticklabels(ax.get_yticklabels(), fontsize=8, rotation=0) - ax.set_ylabel('Conditions',fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel( "Regulator Name",fontsize=10, fontweight='bold', loc='center') - ax.set_title(f"Heatmap of Regulator Effects on Program {Target_Program} Across Conditions",fontsize=14, fontweight='bold', loc='center') + ax.set_ylabel("Conditions", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel("Regulator Name", fontsize=10, fontweight="bold", loc="center") + ax.set_title( + f"Heatmap of Regulator Effects on Program {Target_Program} Across Conditions", + fontsize=14, + fontweight="bold", + loc="center", + ) # Add X marks for values with adj_pval < 0.05 # You'll need to have a separate df with p-values in the same format for i in range(len(heatmap_data)): for j in range(len(heatmap_data.columns)): if df_pivot_pval.iloc[i, j] < 0.05: # pval_data should have same structure - ax.text(j+0.5, i+0.5, '*', ha='center', va='center', - color='black', fontsize=14, weight='bold') - + ax.text( + j + 0.5, + i + 0.5, + "*", + ha="center", + va="center", + color="black", + fontsize=14, + weight="bold", + ) if ax is None: plt.tight_layout() - + # Save if path provided (only when ax is not provided) if save_path and ax is None: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) + # Control whether to display the plot (only when ax is not provided) if ax is None: if show: plt.show() else: plt.close(fig) - - return ax + return ax # plot one volcone plot given perturbation analysis, return genes in df -def plot_program_volcano(perturb_path, Target, tagert_col_name="target_name", plot_col_name="program_name", - log2fc_col='log2FC', down_thred_log=-0.05, up_thred_log=0.05, p_value=0.05, save_path=None, - save_name=None, figsize=(5, 4), show=False, ax=None, Day ="", gene_list = None): - +def plot_program_volcano( + perturb_path, + Target, + tagert_col_name="target_name", + plot_col_name="program_name", + log2fc_col="log2FC", + down_thred_log=-0.05, + up_thred_log=0.05, + p_value=0.05, + save_path=None, + save_name=None, + figsize=(5, 4), + show=False, + ax=None, + Day="", + gene_list=None, +): + df = pd.read_csv(perturb_path, sep="\t") - df['program_name'] = df['program_name'].astype(str) + df["program_name"] = df["program_name"].astype(str) Target = str(Target) - # Check if query gene exists if Target not in df[tagert_col_name].values: print(f"Program {Target} not found in data") if ax is not None: - ax.text(0.5, 0.5, 'No data to display', - ha='center', va='center', transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + "No data to display", + ha="center", + va="center", + transform=ax.transAxes, + ) return ax, pd.DataFrame(), [] df_program = df.loc[df[tagert_col_name] == Target] if gene_list is not None: - expressed_gene = set(df_program[plot_col_name].values).intersection(set(gene_list)) + expressed_gene = set(df_program[plot_col_name].values).intersection( + set(gene_list) + ) df_program = df_program[df_program[plot_col_name].isin(expressed_gene)] - - + # Create figure/axes if not provided if ax is None: fig, ax = plt.subplots(figsize=figsize) else: fig = ax.get_figure() - + # Plot all points - ax.scatter(x=df_program[log2fc_col], - y=df_program['adj_pval'].apply(lambda x: -np.log10(x)), - s=10, label="Not significant", color="grey") + ax.scatter( + x=df_program[log2fc_col], + y=df_program["adj_pval"].apply(lambda x: -np.log10(x)), + s=10, + label="Not significant", + color="grey", + ) # Highlight down- or up-regulated genes - down = df_program[(df_program[log2fc_col] <= down_thred_log) & - (df_program['adj_pval'] <= p_value)] - up = df_program[(df_program[log2fc_col] >= up_thred_log) & - (df_program['adj_pval'] <= p_value)] - - ax.scatter(x=down[log2fc_col], - y=down['adj_pval'].apply(lambda x: -np.log10(x)), - s=10, label="Down-regulated", color="blue") - ax.scatter(x=up[log2fc_col], - y=up['adj_pval'].apply(lambda x: -np.log10(x)), - s=10, label="Up-regulated", color="red") + down = df_program[ + (df_program[log2fc_col] <= down_thred_log) & (df_program["adj_pval"] <= p_value) + ] + up = df_program[ + (df_program[log2fc_col] >= up_thred_log) & (df_program["adj_pval"] <= p_value) + ] + + ax.scatter( + x=down[log2fc_col], + y=down["adj_pval"].apply(lambda x: -np.log10(x)), + s=10, + label="Down-regulated", + color="blue", + ) + ax.scatter( + x=up[log2fc_col], + y=up["adj_pval"].apply(lambda x: -np.log10(x)), + s=10, + label="Up-regulated", + color="red", + ) # Add text labels texts = [] for i, r in up.iterrows(): - texts.append(ax.text(x=r[log2fc_col], y=-np.log10(r['adj_pval']), - s=r[plot_col_name], fontsize=10, zorder=10)) + texts.append( + ax.text( + x=r[log2fc_col], + y=-np.log10(r["adj_pval"]), + s=r[plot_col_name], + fontsize=10, + zorder=10, + ) + ) for i, r in down.iterrows(): - texts.append(ax.text(x=r[log2fc_col], y=-np.log10(r['adj_pval']), - s=r[plot_col_name], fontsize=10, zorder=10)) - + texts.append( + ax.text( + x=r[log2fc_col], + y=-np.log10(r["adj_pval"]), + s=r[plot_col_name], + fontsize=10, + zorder=10, + ) + ) + # Adjust text to avoid overlaps - adjust_text(texts, arrowprops=dict(arrowstyle='->', color='black', lw=0.5), ax=ax) - - + adjust_text(texts, arrowprops=dict(arrowstyle="->", color="black", lw=0.5), ax=ax) + # Set labels and lines - ax.set_xlabel('Effect on Program Expression',fontsize=10, fontweight='bold', loc='center') - ax.set_ylabel("-log10 Adjusted p-value ",fontsize=10, fontweight='bold', loc='center') + ax.set_xlabel( + "Effect on Program Expression", fontsize=10, fontweight="bold", loc="center" + ) + ax.set_ylabel( + "-log10 Adjusted p-value ", fontsize=10, fontweight="bold", loc="center" + ) ax.axvline(down_thred_log, color="grey", linestyle="--") ax.axvline(up_thred_log, color="grey", linestyle="--") ax.axhline(-np.log10(p_value), color="grey", linestyle="--") ax.legend() - ax.set_title(f"Volcano Plot for Program Expression, \n Program {Target}, {Day}",fontsize=14, fontweight='bold', loc='center') + ax.set_title( + f"Volcano Plot for Program Expression, \n Program {Target}, {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) if ax is None: plt.tight_layout() - + # Save if path provided (only when ax is not provided) if save_path and ax is None: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) + # Control whether to display the plot (only when ax is not provided) if ax is None: if show: plt.show() else: plt.close(fig) - - return ax, pd.merge(up, down, how='outer'), texts + return ax, pd.merge(up, down, how="outer"), texts # given mdata, list of programs to plot, plot dotplot for programs, split by days -def perturbed_program_dotplot(mdata, Target_Program, groupby="sample", gene_list=None, - save_path=None, save_name=None, figsize=(5, 4), show=False, ax=None, Day="", - gene_name_key='gene_names', data_key='rna'): +def perturbed_program_dotplot( + mdata, + Target_Program, + groupby="sample", + gene_list=None, + save_path=None, + save_name=None, + figsize=(5, 4), + show=False, + ax=None, + Day="", + gene_name_key="gene_names", + data_key="rna", +): if gene_name_key is not None and gene_name_key in mdata[data_key].var.columns: adata_gene_list = mdata[data_key].var[gene_name_key].astype(str).tolist() else: adata_gene_list = mdata[data_key].var_names.tolist() - gene_list = [gene for gene in gene_list if gene in adata_gene_list] if gene_list else [] + gene_list = ( + [gene for gene in gene_list if gene in adata_gene_list] if gene_list else [] + ) # check when there is gene that perturbed this program if not gene_list: - print("No significant Gene or Gene not Found") # Handle empty gene list @@ -790,43 +1043,76 @@ def perturbed_program_dotplot(mdata, Target_Program, groupby="sample", gene_list img.save(full_path) return None else: - ax.text(0.5, 0.5, 'No Genes to display', - ha='center', va='center', transform=ax.transAxes) - ax.set_title(f"{Target_Program} Perturbed Program Loading Scores {Day}", - fontweight='bold', loc='center') + ax.text( + 0.5, + 0.5, + "No Genes to display", + ha="center", + va="center", + transform=ax.transAxes, + ) + ax.set_title( + f"{Target_Program} Perturbed Program Loading Scores {Day}", + fontweight="bold", + loc="center", + ) return ax # check if the regulators are in the adata to plot if not set(gene_list).issubset(set(adata_gene_list)): missing = set(gene_list) - set(adata_gene_list) print(f"The following regulators are not in adata: {missing}") - - + gene_list = gene_list[::-1] - + # Create dotplot if ax is None: # Standalone mode - let scanpy create its own figure - dp = sc.pl.dotplot(mdata[data_key] , gene_list, groupby=groupby, swap_axes=True, figsize=figsize, - dendrogram=False, show=False, return_fig=True, gene_symbols=gene_name_key, use_raw=False) + dp = sc.pl.dotplot( + mdata[data_key], + gene_list, + groupby=groupby, + swap_axes=True, + figsize=figsize, + dendrogram=False, + show=False, + return_fig=True, + gene_symbols=gene_name_key, + use_raw=False, + ) dp.make_figure() fig = dp.fig - ax = dp.ax_dict['mainplot_ax'] + ax = dp.ax_dict["mainplot_ax"] else: # Gridspec mode - use provided ax fig = ax.get_figure() - dp = sc.pl.dotplot(mdata[data_key] , gene_list, groupby=groupby, swap_axes=True, figsize=figsize, - dendrogram=False, show=False, return_fig=True, ax=ax, gene_symbols=gene_name_key, use_raw=False) + dp = sc.pl.dotplot( + mdata[data_key], + gene_list, + groupby=groupby, + swap_axes=True, + figsize=figsize, + dendrogram=False, + show=False, + return_fig=True, + ax=ax, + gene_symbols=gene_name_key, + use_raw=False, + ) dp.make_figure() - - ax.set_title(f"Regulator Expression \n Program {Target_Program}, {Day}", - fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel('Regulator Name', fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel("Conditions", fontsize=10, fontweight='bold', loc='center') - + + ax.set_title( + f"Regulator Expression \n Program {Target_Program}, {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) + ax.set_ylabel("Regulator Name", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel("Conditions", fontsize=10, fontweight="bold", loc="center") + # Get labels and set ticks properly label = list(mdata[data_key].obs[groupby].cat.categories) - + # Set both ticks and labels together tick_positions = range(len(label)) ax.set_xticks(tick_positions) @@ -834,26 +1120,25 @@ def perturbed_program_dotplot(mdata, Target_Program, groupby="sample", gene_list if ax is None: plt.tight_layout() - + # Save if path provided (only when ax is not provided) if save_path and save_name and ax is None: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) + # Control whether to display the plot (only in standalone mode) if ax is None: if show: plt.show() else: plt.close(fig) - - return ax - - + return ax # helper method for computing waterfall corr -def compute_program_waterfall_cor(perturb_path, save_path=None, log2fc_col='log2FC'): +def compute_program_waterfall_cor(perturb_path, save_path=None, log2fc_col="log2FC"): """Program x program correlation of perturbation effects across regulators. Pivots the perturbation results to a (program x target) log2FC table and @@ -877,88 +1162,109 @@ def compute_program_waterfall_cor(perturb_path, save_path=None, log2fc_col='log2 K x K correlation indexed by program name (diagonal 1; the waterfall plots drop the program itself). """ - df = pd.read_csv(perturb_path, sep='\t', index_col=0) - df['program_name'] = df['program_name'].astype(str) + df = pd.read_csv(perturb_path, sep="\t", index_col=0) + df["program_name"] = df["program_name"].astype(str) # programs as rows, perturbed targets as columns: (K, targets) - pivot_df = df.pivot_table(index='program_name', columns='target_name', values=log2fc_col) + pivot_df = df.pivot_table( + index="program_name", columns="target_name", values=log2fc_col + ) names = pivot_df.index.astype(str).tolist() R = pivot_df.T.corr().to_numpy() # column-wise corr on (targets, K) -> (K, K) if save_path is not None: Path(save_path).parent.mkdir(parents=True, exist_ok=True) - np.savez(save_path, corr=R.astype(np.float32), program=np.asarray(names, dtype=str)) + np.savez( + save_path, corr=R.astype(np.float32), program=np.asarray(names, dtype=str) + ) corr_matrix = pd.DataFrame(R, index=names, columns=names) return corr_matrix +# plot the waterfall plot for genes that have simliar program loading scores when perturbed +def create_program_correlation_waterfall( + corr_matrix, + Target_Program, + top_num=5, + save_path=None, + save_name=None, + figsize=(3, 5), + show=False, + ax=None, + Day="", +): - -# plot the waterfall plot for genes that have simliar program loading scores when perturbed -def create_program_correlation_waterfall(corr_matrix, Target_Program, top_num=5, save_path=None, - save_name=None, figsize=(3, 5), show=False, ax=None, Day =""): - - # Convert to DataFrame and sort - corr_matrix.index = corr_matrix.index.astype(str) # convert type + corr_matrix.index = corr_matrix.index.astype(str) # convert type gene_correlations = corr_matrix.loc[str(Target_Program)].dropna() - corr_df = (gene_correlations).sort_values(ascending = False) - corr_df = corr_df.drop(str(Target_Program), errors="ignore") # drop self corr (already gone if NaN) - + corr_df = (gene_correlations).sort_values(ascending=False) + corr_df = corr_df.drop( + str(Target_Program), errors="ignore" + ) # drop self corr (already gone if NaN) + # Get top N positive and bottom N negative correlations for labeling top_positive = corr_df.head(top_num) top_negative = corr_df.tail(top_num) genes_to_label = set(top_positive.index.tolist() + top_negative.index.tolist()) - + # Use ALL correlations for plotting - plot_data = corr_df + plot_data = corr_df # Create figure/axes if not provided if ax is None: fig, ax = plt.subplots(figsize=figsize) else: fig = ax.get_figure() - - + # Plot only the markers (no line) - ax.scatter(range(len(plot_data)), plot_data.values, - color='gray', s=16) - + ax.scatter(range(len(plot_data)), plot_data.values, color="gray", s=16) + # Add horizontal reference line at y=0 - ax.axhline(y=0, color='gray', linestyle='--', linewidth=1, alpha=0.5) - + ax.axhline(y=0, color="gray", linestyle="--", linewidth=1, alpha=0.5) + # Remove x-axis labels ax.set_xticks([]) - # Highlight the labeled genes with larger markers and add annotations texts = [] for i, (gene, values) in enumerate(plot_data.items()): if gene in genes_to_label: - ax.plot(i, values, 'o', color='darkgray', markersize=5, zorder=5) - text = ax.text(i, values, gene, fontsize=10, style='italic', ha='center') + ax.plot(i, values, "o", color="darkgray", markersize=5, zorder=5) + text = ax.text(i, values, gene, fontsize=10, style="italic", ha="center") texts.append(text) - - adjust_text(texts, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5), ax=ax, fontsize=12) - - ax.set_title(f"Programs Regulated by Similiar Regulators \n Program {Target_Program}, {Day}",fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel(f'Correlation of Perturbed Genes\n Effect on Program {Target_Program}',fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel('Program Name',fontsize=10, fontweight='bold', loc='center') - - + + adjust_text( + texts, arrowprops=dict(arrowstyle="-", color="gray", lw=0.5), ax=ax, fontsize=12 + ) + + ax.set_title( + f"Programs Regulated by Similiar Regulators \n Program {Target_Program}, {Day}", + fontsize=14, + fontweight="bold", + loc="center", + ) + ax.set_ylabel( + f"Correlation of Perturbed Genes\n Effect on Program {Target_Program}", + fontsize=10, + fontweight="bold", + loc="center", + ) + ax.set_xlabel("Program Name", fontsize=10, fontweight="bold", loc="center") + # Add grid - ax.grid(True, alpha=0.3, linestyle='--') + ax.grid(True, alpha=0.3, linestyle="--") ax.set_axisbelow(True) # Adjust layout (only in standalone mode) if ax is None: plt.tight_layout() - + # Save if path provided (only in standalone mode) if save_path and ax is None: - fig.savefig(f'{save_path}/{save_name}.svg', format='svg', bbox_inches='tight', dpi=300) - + fig.savefig( + f"{save_path}/{save_name}.svg", format="svg", bbox_inches="tight", dpi=300 + ) # Control whether to display the plot (only in standalone mode) if ax is None: @@ -966,13 +1272,11 @@ def create_program_correlation_waterfall(corr_matrix, Target_Program, top_num=5, plt.show() else: plt.close(fig) - - - return ax, texts + return ax, texts -# plot program PDF +# plot program PDF def create_comprehensive_program_plot( mdata, perturb_path_base=None, @@ -982,10 +1286,10 @@ def create_comprehensive_program_plot( program_correlation=None, waterfall_correlation=None, top_program=5, - groupby='sample', + groupby="sample", tagert_col_name="program_name", plot_col_name="target_name", - log2fc_col='log2FC', + log2fc_col="log2FC", top_enrichned_term=10, down_thred_log=-0.05, up_thred_log=0.05, @@ -997,13 +1301,13 @@ def create_comprehensive_program_plot( square_plots=True, show=True, PDF=True, - gene_list = None, + gene_list=None, subsample_frac=None, go_p_value_name="Adjusted P-value", go_term_col="Term", - gene_name_key='gene_names', - data_key='rna', - prog_key='cNMF' + gene_name_key="gene_names", + data_key="rna", + prog_key="cNMF", ): """Create a comprehensive multi-panel figure for one cNMF program. @@ -1075,7 +1379,7 @@ def create_comprehensive_program_plot( # Set default samples if not provided if sample is None: - sample = ['D0', 'sample_D1', 'sample_D2', 'sample_D3'] + sample = ["D0", "sample_D1", "sample_D2", "sample_D3"] # Calculate figure dimensions num_samples = len(sample) if plot_perturbation else 0 @@ -1084,7 +1388,7 @@ def create_comprehensive_program_plot( plt.ioff() # Close any existing figures to prevent memory leaks - plt.close('all') + plt.close("all") # Create figure with custom gridspec # 1 header row + n sample rows + heatmap @@ -1103,110 +1407,143 @@ def create_comprehensive_program_plot( figsize = (35, num_rows * PER_ROW_HEIGHT_IN) fig = plt.figure(figsize=figsize) - + # Create gridspec with dynamic number of rows using 27 columns for flexibility - gs = gridspec.GridSpec(num_rows, 27, figure=fig, - hspace=0.4, # Vertical spacing between rows - wspace=0.5, # Horizontal spacing between columns - left=0.10, # Left margin - right=0.90) # Right margin - + gs = gridspec.GridSpec( + num_rows, + 27, + figure=fig, + hspace=0.4, # Vertical spacing between rows + wspace=0.5, # Horizontal spacing between columns + left=0.10, # Left margin + right=0.90, + ) # Right margin + # First row: 5 header plots - ax1 = fig.add_subplot(gs[0, 0:5]) # UMAP Expression - ax21 = fig.add_subplot(gs[0, 6:11]) # Top program + ax1 = fig.add_subplot(gs[0, 0:5]) # UMAP Expression + ax21 = fig.add_subplot(gs[0, 6:11]) # Top program ax3 = fig.add_subplot(gs[0, 12:16]) # Perturbed gene dotplot ax4 = fig.add_subplot(gs[0, 18:21]) # Correlation analysis ax2 = fig.add_subplot(gs[0, 22:27]) # UMAP Perturbation # Heatmap occupies the last row and is fed entirely by the perturbation files, # so it is only allocated when those files are available. - ax22 = fig.add_subplot(gs[num_rows-1, 2:26]) if plot_perturbation else None + ax22 = fig.add_subplot(gs[num_rows - 1, 2:26]) if plot_perturbation else None # Initialize list with header axes axes = [ax1, ax2, ax3, ax4, ax21] - + # Dynamically create axes for each sample # Each sample row has 4 plots: Log2FC, Volcano, Dotplot, Waterfall for sample_idx in range(num_samples): row_idx = sample_idx + 1 # Start from row 1 (after header) - - ax_log2fc = fig.add_subplot(gs[row_idx, 1:6]) # Log2FC + + ax_log2fc = fig.add_subplot(gs[row_idx, 1:6]) # Log2FC ax_volcano = fig.add_subplot(gs[row_idx, 7:12]) # Volcano - ax_dotplot = fig.add_subplot(gs[row_idx, 13:18]) # Dotplot - ax_waterfall = fig.add_subplot(gs[row_idx, 20:25]) # Waterfall - + ax_dotplot = fig.add_subplot(gs[row_idx, 13:18]) # Dotplot + ax_waterfall = fig.add_subplot(gs[row_idx, 20:25]) # Waterfall + axes.extend([ax_log2fc, ax_volcano, ax_dotplot, ax_waterfall]) # Plot 1: UMAP ax1 = plot_umap_per_program( mdata=mdata, Target_Program=Target_Program, - color='purple', + color="purple", figsize=(4, 3), ax=ax1, subsample_frac=subsample_frac, - prog_key=prog_key + prog_key=prog_key, + ) + ax1.set_title( + f"Program {Target_Program} Expression", + fontsize=18, + fontweight="bold", + loc="center", ) - ax1.set_title(f"Program {Target_Program} Expression", fontsize=18, fontweight='bold', loc='center') - ax1.set_xlabel('UMAP 1', fontsize=14, fontweight='bold') - ax1.set_ylabel('UMAP 2', fontsize=14, fontweight='bold') - - # Plot 21: violin + ax1.set_xlabel("UMAP 1", fontsize=14, fontweight="bold") + ax1.set_ylabel("UMAP 2", fontsize=14, fontweight="bold") + + # Plot 21: violin ax21 = plot_violin( - mdata = mdata, - Target_Program = Target_Program, - groupby = groupby, - figsize=(5,3), + mdata=mdata, + Target_Program=Target_Program, + groupby=groupby, + figsize=(5, 3), ax=ax21, data_key=data_key, - prog_key=prog_key - ) - ax21.set_title(f"Program {Target_Program} Expression by Conditions", fontsize=18, fontweight='bold', loc='center') - ax21.set_ylabel('Program Expression', fontsize=14, fontweight='bold', loc='center') - ax21.set_xlabel('Conditions', fontsize=14, fontweight='bold', loc='center') - + prog_key=prog_key, + ) + ax21.set_title( + f"Program {Target_Program} Expression by Conditions", + fontsize=18, + fontweight="bold", + loc="center", + ) + ax21.set_ylabel("Program Expression", fontsize=14, fontweight="bold", loc="center") + ax21.set_xlabel("Conditions", fontsize=14, fontweight="bold", loc="center") + # Plot 2: Correlation analysis - ax2 =analyze_program_correlations( - program_correlation = program_correlation, - Target_Program = Target_Program, - num_program = top_program, - figsize = (4, 4), - ax=ax2 - ) - ax2.set_title(f"Programs with Correlated Loading \nAcross Cells for Program {Target_Program}", fontsize=20, fontweight='bold', loc='center') - ax2.set_ylabel('Program Name', fontsize=14, fontweight='bold') - ax2.set_xlabel(f'Correlation Coefficient', fontsize=14, fontweight='bold', loc='center') + ax2 = analyze_program_correlations( + program_correlation=program_correlation, + Target_Program=Target_Program, + num_program=top_program, + figsize=(4, 4), + ax=ax2, + ) + ax2.set_title( + f"Programs with Correlated Loading \nAcross Cells for Program {Target_Program}", + fontsize=20, + fontweight="bold", + loc="center", + ) + ax2.set_ylabel("Program Name", fontsize=14, fontweight="bold") + ax2.set_xlabel( + f"Correlation Coefficient", fontsize=14, fontweight="bold", loc="center" + ) - # Plot 3: Top Gene for program ax3 = plot_top_gene_per_program( mdata=mdata, - Target_Program = Target_Program, - num_gene = top_enrichned_term, + Target_Program=Target_Program, + num_gene=top_enrichned_term, figsize=(4, 4), file_to_dictionary=file_to_dictionary, ax=ax3, gene_name_key=gene_name_key, data_key=data_key, - prog_key=prog_key) - ax3.set_title(f"Top Loading Genes for Program {Target_Program}", fontsize=20, fontweight='bold', loc='center') - ax3.set_xlabel('Gene Loading Score (z-score)', fontsize=14, fontweight='bold', loc='center') - ax3.set_ylabel('Gene Name', fontsize=14, fontweight='bold', loc='center') - - - + prog_key=prog_key, + ) + ax3.set_title( + f"Top Loading Genes for Program {Target_Program}", + fontsize=20, + fontweight="bold", + loc="center", + ) + ax3.set_xlabel( + "Gene Loading Score (z-score)", fontsize=14, fontweight="bold", loc="center" + ) + ax3.set_ylabel("Gene Name", fontsize=14, fontweight="bold", loc="center") + # Plot 4: GO enrichment - ax4,wrapped_labels=top_GO_per_program( + ax4, wrapped_labels = top_GO_per_program( GO_path=GO_path, - Target_Program = Target_Program, - num_term = top_enrichned_term, - p_value_name = go_p_value_name, - term_col = go_term_col, - figsize=(6,6), - ax=ax4) - ax4.set_title(f"GO enrichment for Program {Target_Program}",fontsize=18, fontweight='bold', loc='center') - ax4.set_yticklabels(wrapped_labels, fontsize=8, fontweight='bold') - ax4.set_xlabel('-log10 Adjusted P-value',fontsize=11, fontweight='bold', loc='center') - + Target_Program=Target_Program, + num_term=top_enrichned_term, + p_value_name=go_p_value_name, + term_col=go_term_col, + figsize=(6, 6), + ax=ax4, + ) + ax4.set_title( + f"GO enrichment for Program {Target_Program}", + fontsize=18, + fontweight="bold", + loc="center", + ) + ax4.set_yticklabels(wrapped_labels, fontsize=8, fontweight="bold") + ax4.set_xlabel( + "-log10 Adjusted P-value", fontsize=11, fontweight="bold", loc="center" + ) # Loop through samples and create 4 plots per sample (Log2FC, Volcano, Dotplot, Waterfall) ax_index = 5 # Start after header axes (0-5) @@ -1226,13 +1563,21 @@ def create_comprehensive_program_plot( figsize=(4, 3), gene_list=gene_list, ax=current_ax, - Day=samp + Day=samp, ) ax_index += 1 - current_ax.set_xlabel('Effect on Program Expression', fontsize=14, fontweight='bold', loc='center') - current_ax.set_ylabel("Regulator Name", fontsize=14, fontweight='bold', loc='center') - current_ax.set_title(f"Regulator Effect on Program Expression, \n Program {Target_Program}, {samp}", fontsize=20, fontweight='bold', loc='center') - + current_ax.set_xlabel( + "Effect on Program Expression", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_ylabel( + "Regulator Name", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_title( + f"Regulator Effect on Program Expression, \n Program {Target_Program}, {samp}", + fontsize=20, + fontweight="bold", + loc="center", + ) # Plot 2: Volcano plot current_ax = axes[ax_index] @@ -1248,17 +1593,31 @@ def create_comprehensive_program_plot( figsize=(5, 3), gene_list=gene_list, ax=current_ax, - Day=samp + Day=samp, ) ax_index += 1 - current_ax.set_title(f"Volcano Plot for Program Expression, \n Program {Target_Program}, {samp}", fontsize=20, fontweight='bold', loc='center') - current_ax.set_xlabel('Effect on Program Expression', fontsize=14, fontweight='bold', loc='center') - current_ax.set_ylabel(" -log10 Adjusted p-value", fontsize=14, fontweight='bold', loc='center') + current_ax.set_title( + f"Volcano Plot for Program Expression, \n Program {Target_Program}, {samp}", + fontsize=20, + fontweight="bold", + loc="center", + ) + current_ax.set_xlabel( + "Effect on Program Expression", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_ylabel( + " -log10 Adjusted p-value", fontsize=14, fontweight="bold", loc="center" + ) for t in text: t.set_fontsize(14) - adjust_text(text, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5), ax=current_ax, fontsize=14) - - # Plot 3: Programs dotplot + adjust_text( + text, + arrowprops=dict(arrowstyle="-", color="gray", lw=0.5), + ax=current_ax, + fontsize=14, + ) + + # Plot 3: Programs dotplot current_ax = axes[ax_index] current_ax = perturbed_program_dotplot( mdata=mdata, @@ -1269,17 +1628,26 @@ def create_comprehensive_program_plot( ax=current_ax, Day=samp, gene_name_key=gene_name_key, - data_key=data_key + data_key=data_key, + ) + current_ax.set_title( + f"Regulator Expression \n Program {Target_Program}, {samp} ", + fontsize=18, + fontweight="bold", + loc="center", + ) + current_ax.set_ylabel( + "Regulator Name", fontsize=14, fontweight="bold", loc="center" + ) + current_ax.set_xlabel( + "Conditions", fontsize=14, fontweight="bold", loc="center" ) - current_ax.set_title(f"Regulator Expression \n Program {Target_Program}, {samp} ", fontsize=18, fontweight='bold', loc='center') - current_ax.set_ylabel('Regulator Name', fontsize=14, fontweight='bold', loc='center') - current_ax.set_xlabel("Conditions", fontsize=14, fontweight='bold', loc='center') ax_index += 1 - # Plot 4: Waterfall plot + # Plot 4: Waterfall plot current_ax = axes[ax_index] current_ax, text = create_program_correlation_waterfall( - corr_matrix= waterfall_correlation[samp], + corr_matrix=waterfall_correlation[samp], Target_Program=Target_Program, top_num=top_enrichned_term, ax=current_ax, @@ -1287,13 +1655,29 @@ def create_comprehensive_program_plot( figsize=(3, 4), ) ax_index += 1 - current_ax.set_title(f"Programs Regulated by Similiar Regulators \n Program {Target_Program}, {samp}", fontsize=18, fontweight='bold', loc='center') - current_ax.set_ylabel(f'Correlation of Perturbed Genes\n Effect on Program {Target_Program}', fontsize=14, fontweight='bold', loc='center') - current_ax.set_xlabel('Program Name', fontsize=14, fontweight='bold', loc='center') + current_ax.set_title( + f"Programs Regulated by Similiar Regulators \n Program {Target_Program}, {samp}", + fontsize=18, + fontweight="bold", + loc="center", + ) + current_ax.set_ylabel( + f"Correlation of Perturbed Genes\n Effect on Program {Target_Program}", + fontsize=14, + fontweight="bold", + loc="center", + ) + current_ax.set_xlabel( + "Program Name", fontsize=14, fontweight="bold", loc="center" + ) for t in text: - t.set_fontsize(14) - adjust_text(text, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5), ax=current_ax, fontsize=14) - + t.set_fontsize(14) + adjust_text( + text, + arrowprops=dict(arrowstyle="-", color="gray", lw=0.5), + ax=current_ax, + fontsize=14, + ) # Plot 5: Heatmap with expression sidebar (perturbation files only) if plot_perturbation: @@ -1310,51 +1694,69 @@ def create_comprehensive_program_plot( data_key=data_key, prog_key=prog_key, figsize=(15, 5), - ax=ax22) + ax=ax22, + ) # plot_program_heatmap_weighted returns None when a program has no # significant regulators; keep the pre-allocated gridspec axes (ax22) and # mark it empty instead of crashing on the styling calls below. if heatmap_ax is not None: ax22 = heatmap_ax - ax22.set_xticklabels(ax22.get_xticklabels(), fontsize=13, rotation=45, ha='right') + ax22.set_xticklabels( + ax22.get_xticklabels(), fontsize=13, rotation=45, ha="right" + ) ax22.set_yticklabels(ax22.get_yticklabels(), fontsize=13, rotation=0) - ax22.set_ylabel('',fontsize=14, fontweight='bold', loc='center') - ax22.set_xlabel('',fontsize=14, fontweight='bold', loc='center') + ax22.set_ylabel("", fontsize=14, fontweight="bold", loc="center") + ax22.set_xlabel("", fontsize=14, fontweight="bold", loc="center") else: - ax22.axis('off') - ax22.text(0.5, 0.5, f"No significant regulators for Program {Target_Program}", - ha='center', va='center', transform=ax22.transAxes, fontsize=14) - ax22.set_title(f"Regulator effects on Program {Target_Program}",fontsize=18, fontweight='bold', loc='center') - - + ax22.axis("off") + ax22.text( + 0.5, + 0.5, + f"No significant regulators for Program {Target_Program}", + ha="center", + va="center", + transform=ax22.transAxes, + fontsize=14, + ) + ax22.set_title( + f"Regulator effects on Program {Target_Program}", + fontsize=18, + fontweight="bold", + loc="center", + ) # Add main title - fig.suptitle(f'Comprehensive Analysis: Program {Target_Program}', fontweight='bold', y=0.995, fontsize=30) - + fig.suptitle( + f"Comprehensive Analysis: Program {Target_Program}", + fontweight="bold", + y=0.995, + fontsize=30, + ) + # Adjust layout plt.tight_layout() - + # Save if path provided if save_path and save_name and PDF: full_path = f"{save_path}/{save_name}.pdf" print(f"Saving figure to {full_path}...") - fig.savefig(full_path, format='pdf', bbox_inches='tight', dpi=300) - + fig.savefig(full_path, format="pdf", bbox_inches="tight", dpi=300) + if save_path and save_name and not PDF: full_path = f"{save_path}/{save_name}.svg" print(f"Saving figure to {full_path}...") - fig.savefig(full_path, format='svg', bbox_inches='tight', dpi=300) - + fig.savefig(full_path, format="svg", bbox_inches="tight", dpi=300) + # Control whether to display the plot if show: plt.show() - + # Always close the figure to prevent memory leak plt.close(fig) -'''working on progress +"""working on progress # plot top enriched enhancer or promoter def plot_motif_per_program(path, num_term = 10, program_num = 0, day = None, title = "Enhancer", p_value_name = "Adjusted P-value", folder_name = None, file_name = None): @@ -1419,4 +1821,4 @@ def plot_motif_per_program(path, num_term = 10, program_num = 0, day = None, ti def plot_all_days_motif(path, num_term = 10, program_num = 0, title = "Enhancer", p_value_name = "Adjusted P-value", folder_name = None, file_name = None): for samp in ['D0', 'D1', 'D2', 'D3']: plot_motif_per_program(path,num_term=num_term, program_num=program_num,day = samp, p_value_name ="adj_pval", title = title,folder_name=folder_name,file_name=file_name) -''' \ No newline at end of file +""" diff --git a/src/Stage3_Interpretation/A_Plotting/src/Program_expression_weighted_plots.py b/src/Stage3_Interpretation/A_Plotting/src/Program_expression_weighted_plots.py index 65a667b..3fa5f02 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/Program_expression_weighted_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/Program_expression_weighted_plots.py @@ -7,26 +7,32 @@ plt.rcParams["axes.spines.top"] = False plt.rcParams["axes.spines.right"] = False -plt.rcParams['pdf.fonttype'] = 42 -plt.rcParams['ps.fonttype'] = 42 -plt.rcParams['svg.fonttype'] = 'none' # SVG: write real (editable in Illustrator) - -plt.rcParams.update({ - 'font.family': 'sans-serif', - 'font.size': 12, - 'axes.titlesize': 13, - 'axes.labelsize': 11, - 'legend.fontsize': 10, - 'legend.title_fontsize': 10, - 'xtick.labelsize': 9, - 'ytick.labelsize': 9, - 'figure.titlesize': 14, - 'figure.facecolor': 'white', - 'axes.facecolor': 'white', -}) - - -def compute_program_expression_by_condition(mdata, Target_Program, groupby='sample', data_key='rna', prog_key='cNMF'): +plt.rcParams["pdf.fonttype"] = 42 +plt.rcParams["ps.fonttype"] = 42 +plt.rcParams["svg.fonttype"] = ( + "none" # SVG: write real (editable in Illustrator) +) + +plt.rcParams.update( + { + "font.family": "sans-serif", + "font.size": 12, + "axes.titlesize": 13, + "axes.labelsize": 11, + "legend.fontsize": 10, + "legend.title_fontsize": 10, + "xtick.labelsize": 9, + "ytick.labelsize": 9, + "figure.titlesize": 14, + "figure.facecolor": "white", + "axes.facecolor": "white", + } +) + + +def compute_program_expression_by_condition( + mdata, Target_Program, groupby="sample", data_key="rna", prog_key="cNMF" +): """Compute mean program expression per condition. Parameters @@ -44,7 +50,7 @@ def compute_program_expression_by_condition(mdata, Target_Program, groupby='samp Mean program expression indexed by condition name. """ X = mdata[prog_key].X - if hasattr(X, 'toarray'): + if hasattr(X, "toarray"): X = X.toarray() prog_col = list(mdata[prog_key].var_names).index(str(Target_Program)) @@ -56,19 +62,29 @@ def compute_program_expression_by_condition(mdata, Target_Program, groupby='samp else: groups = mdata[data_key].obs[groupby].values - df = pd.DataFrame({'score': scores, 'group': groups}) - return df.groupby('group')['score'].mean() + df = pd.DataFrame({"score": scores, "group": groups}) + return df.groupby("group")["score"].mean() def plot_program_heatmap_weighted( - perturb_path_base, mdata, Target_Program, - tagert_col_name="target_name", plot_col_name="program_name", - sample=None, groupby='sample', - log2fc_col='log2FC', p_value=0.05, - save_path=None, save_name=None, - figsize=(12, 5), show=False, ax=None, - vmin=-1, vmax=1, - data_key='rna', prog_key='cNMF', + perturb_path_base, + mdata, + Target_Program, + tagert_col_name="target_name", + plot_col_name="program_name", + sample=None, + groupby="sample", + log2fc_col="log2FC", + p_value=0.05, + save_path=None, + save_name=None, + figsize=(12, 5), + show=False, + ax=None, + vmin=-1, + vmax=1, + data_key="rna", + prog_key="cNMF", ): """Heatmap of regulator effects with a side bar for program expression. @@ -119,7 +135,7 @@ def plot_program_heatmap_weighted( matplotlib.axes.Axes """ if sample is None: - sample = ['D0', 'sample_D1', 'sample_D2', 'sample_D3'] + sample = ["D0", "sample_D1", "sample_D2", "sample_D3"] # --- load perturbation data across conditions --- perturbed_df_list = [ @@ -127,17 +143,18 @@ def plot_program_heatmap_weighted( for samp in sample ] perturbed_df = pd.concat(perturbed_df_list, ignore_index=True) - perturbed_df['program_name'] = perturbed_df['program_name'].astype(str) + perturbed_df["program_name"] = perturbed_df["program_name"].astype(str) # filter for target program df_sorted = perturbed_df.loc[perturbed_df[tagert_col_name] == str(Target_Program)] # keep regulators significant in at least one condition - genes_to_keep = df_sorted[df_sorted['adj_pval'] < p_value][plot_col_name].unique() + genes_to_keep = df_sorted[df_sorted["adj_pval"] < p_value][plot_col_name].unique() df_filtered = df_sorted[df_sorted[plot_col_name].isin(genes_to_keep)] if df_filtered.empty: import warnings + warnings.warn( f"No regulators with adj_pval < {p_value} for program {Target_Program}. " "Nothing to plot." @@ -145,11 +162,17 @@ def plot_program_heatmap_weighted( return None # pivot tables - df_pivot = df_filtered.pivot(columns=plot_col_name, index='sample', values=log2fc_col) - df_pivot_pval = df_filtered.pivot(columns=plot_col_name, index='sample', values='adj_pval') + df_pivot = df_filtered.pivot( + columns=plot_col_name, index="sample", values=log2fc_col + ) + df_pivot_pval = df_filtered.pivot( + columns=plot_col_name, index="sample", values="adj_pval" + ) # --- compute mean program expression per condition --- - expr_series = compute_program_expression_by_condition(mdata, Target_Program, groupby=groupby, data_key=data_key, prog_key=prog_key) + expr_series = compute_program_expression_by_condition( + mdata, Target_Program, groupby=groupby, data_key=data_key, prog_key=prog_key + ) # --- plot --- if ax is None: @@ -163,27 +186,42 @@ def plot_program_heatmap_weighted( # heatmap (no automatic colorbar — we place it manually) sns.heatmap( - df_pivot, cmap='RdBu_r', center=0, + df_pivot, + cmap="RdBu_r", + center=0, cbar=False, - ax=ax_heat, square=True, - vmin=vmin, vmax=vmax, - linewidths=0.5, linecolor='white', + ax=ax_heat, + square=True, + vmin=vmin, + vmax=vmax, + linewidths=0.5, + linecolor="white", ) # significance asterisks for i in range(len(df_pivot)): for j in range(len(df_pivot.columns)): if df_pivot_pval.iloc[i, j] < p_value: - ax_heat.text(j + 0.5, i + 0.5, '*', ha='center', va='center', - color='black', fontsize=12, weight='bold') - - ax_heat.set_xticklabels(ax_heat.get_xticklabels(), rotation=45, ha='right') + ax_heat.text( + j + 0.5, + i + 0.5, + "*", + ha="center", + va="center", + color="black", + fontsize=12, + weight="bold", + ) + + ax_heat.set_xticklabels(ax_heat.get_xticklabels(), rotation=45, ha="right") ax_heat.set_yticklabels(ax_heat.get_yticklabels(), rotation=0) - ax_heat.set_ylabel('') - ax_heat.set_xlabel('') + ax_heat.set_ylabel("") + ax_heat.set_xlabel("") ax_heat.set_title( f"Regulator effects on Program {Target_Program}", - fontsize=13, fontweight='bold', pad=12, + fontsize=13, + fontweight="bold", + pad=12, ) # Draw the figure so heatmap position is finalized, then add bar + colorbar @@ -193,10 +231,14 @@ def plot_program_heatmap_weighted( # Expression bar: same top/bottom as heatmap, thin strip to its right bar_width = 0.03 bar_gap = 0.005 - ax_bar = fig.add_axes([ - heat_bbox.x1 + bar_gap, heat_bbox.y0, - bar_width, heat_bbox.height, - ]) + ax_bar = fig.add_axes( + [ + heat_bbox.x1 + bar_gap, + heat_bbox.y0, + bar_width, + heat_bbox.height, + ] + ) conditions = list(df_pivot.index) expr_vals = [expr_series.get(c, 0.0) for c in conditions] @@ -204,39 +246,49 @@ def plot_program_heatmap_weighted( norm_vals = [v / max_e for v in expr_vals] y_positions = np.arange(len(conditions)) + 0.5 - ax_bar.barh(y_positions, norm_vals, height=0.8, - color='#2CA02C', edgecolor='none', alpha=0.8) + ax_bar.barh( + y_positions, norm_vals, height=0.8, color="#2CA02C", edgecolor="none", alpha=0.8 + ) ax_bar.set_ylim(ax_heat.get_ylim()) ax_bar.set_xlim(0, 1.05) ax_bar.set_yticks([]) ax_bar.set_xticks([0, 1]) - ax_bar.set_xticklabels(['0', '1'], fontsize=7) - ax_bar.set_xlabel('Expr.', fontsize=9) - ax_bar.spines['top'].set_visible(False) - ax_bar.spines['right'].set_visible(False) - ax_bar.spines['left'].set_visible(False) - ax_bar.tick_params(axis='y', length=0) - ax_bar.tick_params(axis='x', length=2, pad=2) + ax_bar.set_xticklabels(["0", "1"], fontsize=7) + ax_bar.set_xlabel("Expr.", fontsize=9) + ax_bar.spines["top"].set_visible(False) + ax_bar.spines["right"].set_visible(False) + ax_bar.spines["left"].set_visible(False) + ax_bar.tick_params(axis="y", length=0) + ax_bar.tick_params(axis="x", length=2, pad=2) # Colorbar: same top/bottom, to the right of expression bar import matplotlib.colors as mcolors + cbar_gap = 0.015 cbar_width = 0.012 - ax_cbar = fig.add_axes([ - heat_bbox.x1 + bar_gap + bar_width + cbar_gap, heat_bbox.y0, - cbar_width, heat_bbox.height, - ]) + ax_cbar = fig.add_axes( + [ + heat_bbox.x1 + bar_gap + bar_width + cbar_gap, + heat_bbox.y0, + cbar_width, + heat_bbox.height, + ] + ) norm = mcolors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) - sm = plt.cm.ScalarMappable(cmap='RdBu_r', norm=norm) + sm = plt.cm.ScalarMappable(cmap="RdBu_r", norm=norm) sm.set_array([]) cb = fig.colorbar(sm, cax=ax_cbar) - cb.set_label('log2FC', fontsize=9) + cb.set_label("log2FC", fontsize=9) cb.ax.tick_params(labelsize=7) if standalone: if save_path and save_name: - fig.savefig(f'{save_path}/{save_name}.pdf', format='pdf', - bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{save_name}.pdf", + format="pdf", + bbox_inches="tight", + dpi=300, + ) if show: plt.show() else: @@ -246,13 +298,22 @@ def plot_program_heatmap_weighted( def plot_program_heatmap_expression_scaled( - perturb_path_base, mdata, Target_Program, - tagert_col_name="target_name", plot_col_name="program_name", - sample=None, groupby='sample', - log2fc_col='log2FC', p_value=0.05, - save_path=None, save_name=None, - figsize=(12, 5), show=False, ax=None, - data_key='rna', prog_key='cNMF', + perturb_path_base, + mdata, + Target_Program, + tagert_col_name="target_name", + plot_col_name="program_name", + sample=None, + groupby="sample", + log2fc_col="log2FC", + p_value=0.05, + save_path=None, + save_name=None, + figsize=(12, 5), + show=False, + ax=None, + data_key="rna", + prog_key="cNMF", ): """Heatmap of log2FC scaled by relative program expression per condition. @@ -269,7 +330,7 @@ def plot_program_heatmap_expression_scaled( matplotlib.axes.Axes """ if sample is None: - sample = ['D0', 'sample_D1', 'sample_D2', 'sample_D3'] + sample = ["D0", "sample_D1", "sample_D2", "sample_D3"] # --- load perturbation data --- perturbed_df_list = [ @@ -277,17 +338,23 @@ def plot_program_heatmap_expression_scaled( for samp in sample ] perturbed_df = pd.concat(perturbed_df_list, ignore_index=True) - perturbed_df['program_name'] = perturbed_df['program_name'].astype(str) + perturbed_df["program_name"] = perturbed_df["program_name"].astype(str) df_sorted = perturbed_df.loc[perturbed_df[tagert_col_name] == str(Target_Program)] - genes_to_keep = df_sorted[df_sorted['adj_pval'] < p_value][plot_col_name].unique() + genes_to_keep = df_sorted[df_sorted["adj_pval"] < p_value][plot_col_name].unique() df_filtered = df_sorted[df_sorted[plot_col_name].isin(genes_to_keep)] - df_pivot = df_filtered.pivot(columns=plot_col_name, index='sample', values=log2fc_col) - df_pivot_pval = df_filtered.pivot(columns=plot_col_name, index='sample', values='adj_pval') + df_pivot = df_filtered.pivot( + columns=plot_col_name, index="sample", values=log2fc_col + ) + df_pivot_pval = df_filtered.pivot( + columns=plot_col_name, index="sample", values="adj_pval" + ) # --- compute expression weights --- - expr_series = compute_program_expression_by_condition(mdata, Target_Program, groupby=groupby, data_key=data_key, prog_key=prog_key) + expr_series = compute_program_expression_by_condition( + mdata, Target_Program, groupby=groupby, data_key=data_key, prog_key=prog_key + ) max_expr = expr_series.max() if max_expr > 0: weights = expr_series / max_expr @@ -312,40 +379,66 @@ def plot_program_heatmap_expression_scaled( standalone = False sns.heatmap( - scaled_pivot, cmap='RdBu_r', center=0, - cbar_kws={'label': 'Expr-scaled log2FC', 'shrink': 0.6, 'aspect': 20}, - ax=ax, square=True, - vmin=-abs_max, vmax=abs_max, - linewidths=0.5, linecolor='white', + scaled_pivot, + cmap="RdBu_r", + center=0, + cbar_kws={"label": "Expr-scaled log2FC", "shrink": 0.6, "aspect": 20}, + ax=ax, + square=True, + vmin=-abs_max, + vmax=abs_max, + linewidths=0.5, + linecolor="white", ) # significance markers for i in range(len(scaled_pivot)): for j in range(len(scaled_pivot.columns)): if df_pivot_pval.iloc[i, j] < p_value: - ax.text(j + 0.5, i + 0.5, '*', ha='center', va='center', - color='black', fontsize=12, weight='bold') - - ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right') + ax.text( + j + 0.5, + i + 0.5, + "*", + ha="center", + va="center", + color="black", + fontsize=12, + weight="bold", + ) + + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") ax.set_yticklabels(ax.get_yticklabels(), rotation=0) - ax.set_ylabel('') - ax.set_xlabel('') + ax.set_ylabel("") + ax.set_xlabel("") ax.set_title( f"Expression-scaled regulator effects on Program {Target_Program}", - fontsize=13, fontweight='bold', pad=12, + fontsize=13, + fontweight="bold", + pad=12, ) # annotate condition weights on the right side for i, cond in enumerate(scaled_pivot.index): w = weights.get(cond, 0.0) - ax.text(len(scaled_pivot.columns) + 0.3, i + 0.5, - f'w={w:.2f}', ha='left', va='center', fontsize=8, color='#666666') + ax.text( + len(scaled_pivot.columns) + 0.3, + i + 0.5, + f"w={w:.2f}", + ha="left", + va="center", + fontsize=8, + color="#666666", + ) if standalone: fig.subplots_adjust(bottom=0.2) if save_path and save_name: - fig.savefig(f'{save_path}/{save_name}.pdf', format='pdf', - bbox_inches='tight', dpi=300) + fig.savefig( + f"{save_path}/{save_name}.pdf", + format="pdf", + bbox_inches="tight", + dpi=300, + ) if show: plt.show() else: diff --git a/src/Stage3_Interpretation/A_Plotting/src/__init__.py b/src/Stage3_Interpretation/A_Plotting/src/__init__.py index 4dfcb9d..104f6da 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/__init__.py +++ b/src/Stage3_Interpretation/A_Plotting/src/__init__.py @@ -1,38 +1,80 @@ # helper functions -from .utilities import convert_with_mygene, convert_adata_with_mygene, rename_adata_gene_dictionary, \ - rename_list_gene_dictionary, read_npz, merge_pdfs_in_folder, merge_svgs_to_pdf, \ - check_normalized, ensure_umap, condition_labels_from_obs, \ - read_condition_labels_from_h5mu +from .utilities import ( + convert_with_mygene, + convert_adata_with_mygene, + rename_adata_gene_dictionary, + rename_list_gene_dictionary, + read_npz, + merge_pdfs_in_folder, + merge_svgs_to_pdf, + check_normalized, + ensure_umap, + condition_labels_from_obs, + read_condition_labels_from_h5mu, +) # K selection plots -from .k_selection_plots import load_stablity_error_data, plot_stablity_error,\ - load_enrichment_data, plot_enrichment,\ - load_perturbation_data, plot_perturbation,\ - load_explained_variance_data,plot_explained_variance, plot_k_selection_panel, \ - plot_k_selection_panel_no_traits, programs_dotplots +from .k_selection_plots import ( + load_stablity_error_data, + plot_stablity_error, + load_enrichment_data, + plot_enrichment, + load_perturbation_data, + plot_perturbation, + load_explained_variance_data, + plot_explained_variance, + plot_k_selection_panel, + plot_k_selection_panel_no_traits, + programs_dotplots, +) # gene QC plots -from .Perturbed_gene_QC_plots import plot_umap_per_gene, plot_top_program_per_gene, perturbed_gene_dotplot,\ - plot_log2FC, plot_volcano, programs_dotplot, analyze_correlations, \ - create_comprehensive_plot,create_gene_correlation_waterfall,\ - plot_umap_per_gene_guide, process_single_gene, _process_gene_worker, parallel_gene_processing,\ - compute_gene_correlation_matrix, compute_gene_waterfall_cor, \ - plot_perturbation_vs_control, \ - plot_perturbation_vs_control_by_condition - +from .Perturbed_gene_QC_plots import ( + plot_umap_per_gene, + plot_top_program_per_gene, + perturbed_gene_dotplot, + plot_log2FC, + plot_volcano, + programs_dotplot, + analyze_correlations, + create_comprehensive_plot, + create_gene_correlation_waterfall, + plot_umap_per_gene_guide, + process_single_gene, + _process_gene_worker, + parallel_gene_processing, + compute_gene_correlation_matrix, + compute_gene_waterfall_cor, + plot_perturbation_vs_control, + plot_perturbation_vs_control_by_condition, +) # program QC plots -from .Program_QC_plots import plot_umap_per_program, plot_top_gene_per_program, top_GO_per_program, compute_program_correlation_matrix,\ - analyze_program_correlations, plot_violin, plot_program_log2FC, plot_program_heatmap, plot_program_volcano, \ - perturbed_program_dotplot, compute_program_waterfall_cor, create_program_correlation_waterfall, create_comprehensive_program_plot +from .Program_QC_plots import ( + plot_umap_per_program, + plot_top_gene_per_program, + top_GO_per_program, + compute_program_correlation_matrix, + analyze_program_correlations, + plot_violin, + plot_program_log2FC, + plot_program_heatmap, + plot_program_volcano, + perturbed_program_dotplot, + compute_program_waterfall_cor, + create_program_correlation_waterfall, + create_comprehensive_program_plot, +) # expression-weighted program plots -from .Program_expression_weighted_plots import compute_program_expression_by_condition, \ - plot_program_heatmap_weighted, \ - plot_program_heatmap_expression_scaled +from .Program_expression_weighted_plots import ( + compute_program_expression_by_condition, + plot_program_heatmap_weighted, + plot_program_heatmap_expression_scaled, +) # interactive HTML export from .html_Program_QC_plots import export_program_html, write_share_index diff --git a/src/Stage3_Interpretation/A_Plotting/src/efficient_corr.py b/src/Stage3_Interpretation/A_Plotting/src/efficient_corr.py index c0489a6..aac117b 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/efficient_corr.py +++ b/src/Stage3_Interpretation/A_Plotting/src/efficient_corr.py @@ -37,9 +37,12 @@ class RowCorr: """ def __init__(self, X, names, row_label="rows", col_label="columns"): - self.Z = np.ascontiguousarray(_zscore_rows(X, names, row_label, col_label), - dtype=np.float32) - self.index = self.columns = pd.Index(np.asarray(names).astype(str), dtype=object) + self.Z = np.ascontiguousarray( + _zscore_rows(X, names, row_label, col_label), dtype=np.float32 + ) + self.index = self.columns = pd.Index( + np.asarray(names).astype(str), dtype=object + ) self._pos = {g: i for i, g in enumerate(self.index)} @property @@ -89,7 +92,9 @@ def _zscore_rows(X, names, row_label="rows", col_label="columns"): ) with np.errstate(invalid="ignore", divide="ignore"): Z = (X - X.mean(axis=1, keepdims=True)) / X.std(axis=1, keepdims=True) - Z[np.ptp(X, axis=1) == 0] = np.nan # exact test; round-off can leave std tiny but nonzero + Z[np.ptp(X, axis=1) == 0] = ( + np.nan + ) # exact test; round-off can leave std tiny but nonzero return Z @@ -105,7 +110,7 @@ def column_pearson_streaming(X, chunk_rows=100_000): def chunks(): for s in range(0, N, chunk_rows): - C = X[s:s + chunk_rows] + C = X[s : s + chunk_rows] C = C.toarray() if hasattr(C, "toarray") else np.asarray(C) yield C.astype(np.float64, copy=False) @@ -120,6 +125,8 @@ def chunks(): C = C - mean S += C.T @ C sd = np.sqrt(np.diag(S)) - sd[hi == lo] = np.nan # exact constant test; round-off can leave sd tiny but nonzero + sd[hi == lo] = ( + np.nan + ) # exact constant test; round-off can leave sd tiny but nonzero with np.errstate(invalid="ignore", divide="ignore"): return S / np.outer(sd, sd) diff --git a/src/Stage3_Interpretation/A_Plotting/src/html_Perturbed_gene_QC_plots.py b/src/Stage3_Interpretation/A_Plotting/src/html_Perturbed_gene_QC_plots.py index d3e7429..83e284a 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/html_Perturbed_gene_QC_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/html_Perturbed_gene_QC_plots.py @@ -14,6 +14,7 @@ Re-uses the styling, Plotly template, and most figure makers from html_Program_QC_plots.py so the two share folders look identical. """ + from __future__ import annotations import os @@ -34,12 +35,26 @@ perturbed_gene_dotplot, ) from .html_Program_QC_plots import ( - PLOTLY_CDN, GOOGLE_FONTS, - COLOR_UP, COLOR_DOWN, COLOR_NS, COLOR_NEUTRAL, COLOR_ACCENT, - _PLOTLY_TEMPLATE, _STYLE_CSS, - _sanitize, _write_json, _safe_sample_key, - _fmt_layout, _empty_fig, _fig_to_div, _write_plotly_js, - _make_correlations_fig, _make_log2fc_fig, _make_volcano_fig, _make_waterfall_fig, + PLOTLY_CDN, + GOOGLE_FONTS, + COLOR_UP, + COLOR_DOWN, + COLOR_NS, + COLOR_NEUTRAL, + COLOR_ACCENT, + _PLOTLY_TEMPLATE, + _STYLE_CSS, + _sanitize, + _write_json, + _safe_sample_key, + _fmt_layout, + _empty_fig, + _fig_to_div, + _write_plotly_js, + _make_correlations_fig, + _make_log2fc_fig, + _make_volcano_fig, + _make_waterfall_fig, ) @@ -64,13 +79,22 @@ def _get_dense_cnmf_X(mdata, prog_key="cNMF"): if key in _cnmf_X_cache: return _cnmf_X_cache[key] _cnmf_X_cache.clear() # drop stale entry from any prior mdata - dense = np.asarray(X_raw.toarray()) if hasattr(X_raw, "toarray") else np.asarray(X_raw) + dense = ( + np.asarray(X_raw.toarray()) if hasattr(X_raw, "toarray") else np.asarray(X_raw) + ) _cnmf_X_cache[key] = dense return dense -def _build_top_programs(mdata, target_gene, top_n, ensembl_to_symbol_file, gene_name_key, - data_key="rna", prog_key="cNMF"): +def _build_top_programs( + mdata, + target_gene, + top_n, + ensembl_to_symbol_file, + gene_name_key, + data_key="rna", + prog_key="cNMF", +): """Top N programs by gene-loading score for one gene.""" X = mdata[prog_key].varm["loadings"] # (programs x genes) if ensembl_to_symbol_file is None: @@ -80,7 +104,10 @@ def _build_top_programs(mdata, target_gene, top_n, ensembl_to_symbol_file, gene_ col_names = list(mdata[data_key].var_names) else: from .utilities import rename_list_gene_dictionary - col_names = rename_list_gene_dictionary(list(mdata[data_key].var_names), ensembl_to_symbol_file) + + col_names = rename_list_gene_dictionary( + list(mdata[data_key].var_names), ensembl_to_symbol_file + ) df = pd.DataFrame(data=X, columns=col_names, index=mdata[prog_key].var_names) if target_gene not in df.columns: return {"programs": [], "loadings": []} @@ -107,7 +134,11 @@ def _build_gene_correlations(corr_matrix, target_gene, top_n): """Top ± correlated genes from the gene × gene correlation matrix.""" if target_gene not in corr_matrix.columns: return {"programs": [], "r": [], "direction": []} - s = corr_matrix.loc[target_gene].drop(target_gene, errors="ignore").sort_values(ascending=False) + s = ( + corr_matrix.loc[target_gene] + .drop(target_gene, errors="ignore") + .sort_values(ascending=False) + ) top = s[s > 0].head(top_n) bottom = s[s < 0].tail(top_n) combined = pd.concat([bottom, top]) @@ -118,8 +149,9 @@ def _build_gene_correlations(corr_matrix, target_gene, top_n): } -def _build_log2fc_gene(perturb_path, target_gene, target_col, program_col, - log2fc_col, num_item, p_value): +def _build_log2fc_gene( + perturb_path, target_gene, target_col, program_col, log2fc_col, num_item, p_value +): """log2FC bars: programs most affected by perturbation of `target_gene`.""" df = pd.read_csv(perturb_path, sep="\t") df[target_col] = df[target_col].astype(str) @@ -129,7 +161,9 @@ def _build_log2fc_gene(perturb_path, target_gene, target_col, program_col, sub = sub[sub["adj_pval"] < p_value].sort_values(by=log2fc_col, ascending=False) top = sub.head(num_item) bottom = sub.tail(num_item) - plot_data = pd.merge(top, bottom, how="outer").sort_values(by=log2fc_col, ascending=False) + plot_data = pd.merge(top, bottom, how="outer").sort_values( + by=log2fc_col, ascending=False + ) return { "regulators": [str(p) for p in plot_data[program_col].values], "log2fc": [float(x) for x in plot_data[log2fc_col].values], @@ -137,14 +171,27 @@ def _build_log2fc_gene(perturb_path, target_gene, target_col, program_col, } -def _build_volcano_gene(perturb_path, target_gene, target_col, program_col, - log2fc_col, down_thred_log, up_thred_log, p_value): +def _build_volcano_gene( + perturb_path, + target_gene, + target_col, + program_col, + log2fc_col, + down_thred_log, + up_thred_log, + p_value, +): """Volcano: all programs as effect of perturbing `target_gene`.""" df = pd.read_csv(perturb_path, sep="\t") df[target_col] = df[target_col].astype(str) if target_gene not in df[target_col].values: - return {"regulators": [], "log2fc": [], "neglog10p": [], "category": [], - "thresholds": {"down": down_thred_log, "up": up_thred_log, "p": p_value}} + return { + "regulators": [], + "log2fc": [], + "neglog10p": [], + "category": [], + "thresholds": {"down": down_thred_log, "up": up_thred_log, "p": p_value}, + } sub = df.loc[df[target_col] == target_gene].copy() sub["neglog10p"] = -np.log10(sub["adj_pval"].clip(lower=1e-300)) cats = [] @@ -162,11 +209,17 @@ def _build_volcano_gene(perturb_path, target_gene, target_col, program_col, "neglog10p": [float(x) for x in sub["neglog10p"].values], "adj_pval": [float(x) for x in sub["adj_pval"].values], "category": cats, - "thresholds": {"down": float(down_thred_log), "up": float(up_thred_log), "p": float(p_value)}, + "thresholds": { + "down": float(down_thred_log), + "up": float(up_thred_log), + "p": float(p_value), + }, } -def _build_programs_dotplot(mdata, program_list, groupby, data_key="rna", prog_key="cNMF"): +def _build_programs_dotplot( + mdata, program_list, groupby, data_key="rna", prog_key="cNMF" +): """Program loading dotplot data: program × condition (mean + frac>0).""" if not program_list: return {"genes": [], "conditions": [], "mean": [], "frac": []} @@ -223,9 +276,17 @@ def _build_gene_waterfall(corr_matrix, target_gene, top_num): } -def _build_kd_vs_control(mdata, target_gene, condition_key, control_target_name, gene_name_key, - data_key="rna", prog_key="cNMF", - guide_targets_key="guide_targets", guide_assignment_key="guide_assignment"): +def _build_kd_vs_control( + mdata, + target_gene, + condition_key, + control_target_name, + gene_name_key, + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): """KD vs control bar data per condition. Mirrors plot_perturbation_vs_control_by_condition logic.""" from scipy import sparse @@ -237,7 +298,9 @@ def _build_kd_vs_control(mdata, target_gene, condition_key, control_target_name, X = mdata[data_key].X if gene_name_key is not None and gene_name_key in mdata[data_key].var.columns: - gene_mask = (mdata[data_key].var[gene_name_key].astype(str) == target_gene).values + gene_mask = ( + mdata[data_key].var[gene_name_key].astype(str) == target_gene + ).values else: gene_mask = np.array([v == target_gene for v in mdata[data_key].var_names]) if gene_mask.sum() == 0: @@ -255,7 +318,10 @@ def _build_kd_vs_control(mdata, target_gene, condition_key, control_target_name, nt_idx = np.where(np.isin(guide_targets, control_target_name))[0] if len(nt_idx) == 0: - return {"groups": [], "error": f"control target '{control_target_name}' not in guide_targets"} + return { + "groups": [], + "error": f"control target '{control_target_name}' not in guide_targets", + } control_mask = np.asarray(ga[:, nt_idx].sum(axis=1)).flatten() > 0 target_idx = np.where(guide_targets == target_gene)[0] @@ -270,7 +336,14 @@ def _build_kd_vs_control(mdata, target_gene, condition_key, control_target_name, conditions = sorted(conditions_series.unique()) group_labels = ["All"] + [str(c) for c in conditions] - pert_means, ctrl_means, pert_sems, ctrl_sems, pert_ns, ctrl_ns = [], [], [], [], [], [] + pert_means, ctrl_means, pert_sems, ctrl_sems, pert_ns, ctrl_ns = ( + [], + [], + [], + [], + [], + [], + ) for label in group_labels: if label == "All": cond_mask = np.ones(len(gene_expr_norm), dtype=bool) @@ -316,17 +389,25 @@ def _build_kd_vs_control(mdata, target_gene, condition_key, control_target_name, # Plotly figure makers for the gene-specific panels # --------------------------------------------------------------------------- + def _make_top_programs_fig(data): if not data["programs"]: return _empty_fig("No program loadings") - fig = go.Figure(go.Bar( - x=data["loadings"], y=data["programs"], orientation="h", - marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), - hovertemplate="Program %{y}
Loading: %{x:.3g}", - )) - fig.update_layout(xaxis_title="Gene loading", yaxis_title="Program", - xaxis=dict(exponentformat="none"), - yaxis=dict(type="category")) + fig = go.Figure( + go.Bar( + x=data["loadings"], + y=data["programs"], + orientation="h", + marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), + hovertemplate="Program %{y}
Loading: %{x:.3g}", + ) + ) + fig.update_layout( + xaxis_title="Gene loading", + yaxis_title="Program", + xaxis=dict(exponentformat="none"), + yaxis=dict(type="category"), + ) return _fmt_layout(fig) @@ -344,15 +425,33 @@ def _make_programs_dotplot_fig(data): sizes.append(4 + 20 * (f / max_frac)) colors.append(m) hovers.append(f"Program {prog}
{cond}
mean={m:.3f}
frac={f:.2f}") - fig = go.Figure(go.Scatter( - x=xs, y=ys, mode="markers", - marker=dict(size=sizes, color=colors, colorscale="Blues", showscale=True, - colorbar=dict(title=dict(text="mean
loading", side="right"), thickness=8, len=0.7), - line=dict(width=0)), - text=hovers, hovertemplate="%{text}", - )) - fig.update_layout(xaxis_title="Condition", yaxis_title="Program", - xaxis=dict(type="category"), yaxis=dict(type="category")) + fig = go.Figure( + go.Scatter( + x=xs, + y=ys, + mode="markers", + marker=dict( + size=sizes, + color=colors, + colorscale="Blues", + showscale=True, + colorbar=dict( + title=dict(text="mean
loading", side="right"), + thickness=8, + len=0.7, + ), + line=dict(width=0), + ), + text=hovers, + hovertemplate="%{text}", + ) + ) + fig.update_layout( + xaxis_title="Condition", + yaxis_title="Program", + xaxis=dict(type="category"), + yaxis=dict(type="category"), + ) return _fmt_layout(fig) @@ -365,37 +464,58 @@ def _make_kd_vs_control_fig(data): ctrl = [v if v is not None else 0.0 for v in data["ctrl_means"]] pert_sem = data["pert_sems"] ctrl_sem = data["ctrl_sems"] - tick_labels = [f"{g}
[{p}|{c}]" for g, p, c in zip(groups, data["pert_n"], data["ctrl_n"])] + tick_labels = [ + f"{g}
[{p}|{c}]" for g, p, c in zip(groups, data["pert_n"], data["ctrl_n"]) + ] fig = go.Figure() - fig.add_trace(go.Bar( - name=data["target_gene"], - x=groups, y=pert, - error_y=dict(type="data", array=pert_sem, color="#222", thickness=1.2, width=4), - marker=dict(color=COLOR_UP, line=dict(width=0)), - customdata=list(zip(data["pert_n"], pert_sem)), - hovertemplate="%{x}
" - f"{data['target_gene']}: " "%{y:.1%} of control
" - "n=%{customdata[0]}
SEM=%{customdata[1]:.3f}", - )) - fig.add_trace(go.Bar( - name=f"control ({data['control_name']})", - x=groups, y=ctrl, - error_y=dict(type="data", array=ctrl_sem, color="#222", thickness=1.2, width=4), - marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), - customdata=list(zip(data["ctrl_n"], ctrl_sem)), - hovertemplate="%{x}
control: %{y:.1%}
" - "n=%{customdata[0]}
SEM=%{customdata[1]:.3f}", - )) + fig.add_trace( + go.Bar( + name=data["target_gene"], + x=groups, + y=pert, + error_y=dict( + type="data", array=pert_sem, color="#222", thickness=1.2, width=4 + ), + marker=dict(color=COLOR_UP, line=dict(width=0)), + customdata=list(zip(data["pert_n"], pert_sem)), + hovertemplate="%{x}
" + f"{data['target_gene']}: " + "%{y:.1%} of control
" + "n=%{customdata[0]}
SEM=%{customdata[1]:.3f}", + ) + ) + fig.add_trace( + go.Bar( + name=f"control ({data['control_name']})", + x=groups, + y=ctrl, + error_y=dict( + type="data", array=ctrl_sem, color="#222", thickness=1.2, width=4 + ), + marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), + customdata=list(zip(data["ctrl_n"], ctrl_sem)), + hovertemplate="%{x}
control: %{y:.1%}
" + "n=%{customdata[0]}
SEM=%{customdata[1]:.3f}", + ) + ) fig.add_hline(y=1.0, line_dash="dash", line_color="#888", line_width=1) fig.update_layout( barmode="group", - xaxis=dict(tickmode="array", tickvals=groups, ticktext=tick_labels, type="category"), + xaxis=dict( + tickmode="array", tickvals=groups, ticktext=tick_labels, type="category" + ), yaxis=dict(tickformat=".0%", range=[0, 1.20]), xaxis_title="Condition [n perturbed | n control]", yaxis_title=f"{data['target_gene']} expression (% of control)", - legend=dict(orientation="h", x=1, xanchor="right", y=1.08, yanchor="bottom", - bgcolor="rgba(0,0,0,0)"), + legend=dict( + orientation="h", + x=1, + xanchor="right", + y=1.08, + yanchor="bottom", + bgcolor="rgba(0,0,0,0)", + ), ) return _fmt_layout(fig, height=300) @@ -404,14 +524,27 @@ def _make_kd_vs_control_fig(data): # Matplotlib → PNG renderers # --------------------------------------------------------------------------- -def _render_umap_expression_png(mdata, target_gene, out_path, ensembl_to_symbol_file, - gene_name_key, umap_dot_size, subsample_frac, data_key="rna"): + +def _render_umap_expression_png( + mdata, + target_gene, + out_path, + ensembl_to_symbol_file, + gene_name_key, + umap_dot_size, + subsample_frac, + data_key="rna", +): fig, ax = plt.subplots(figsize=(5, 4)) plot_umap_per_gene( - mdata=mdata, Target_Gene=target_gene, ax=ax, + mdata=mdata, + Target_Gene=target_gene, + ax=ax, ensembl_to_symbol_file=ensembl_to_symbol_file, - gene_name_key=gene_name_key, size=umap_dot_size, - umap_subsample_frac=subsample_frac, data_key=data_key, + gene_name_key=gene_name_key, + size=umap_dot_size, + umap_subsample_frac=subsample_frac, + data_key=data_key, ) fig.tight_layout() out_path.parent.mkdir(parents=True, exist_ok=True) @@ -419,15 +552,28 @@ def _render_umap_expression_png(mdata, target_gene, out_path, ensembl_to_symbol_ plt.close(fig) -def _render_umap_guide_png(mdata, target_gene, out_path, umap_dot_size, subsample_frac, - data_key="rna", prog_key="cNMF", - guide_targets_key="guide_targets", guide_assignment_key="guide_assignment"): +def _render_umap_guide_png( + mdata, + target_gene, + out_path, + umap_dot_size, + subsample_frac, + data_key="rna", + prog_key="cNMF", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", +): fig, ax = plt.subplots(figsize=(5, 4)) plot_umap_per_gene_guide( - mdata=mdata, Target_Gene=target_gene, ax=ax, - size=umap_dot_size, umap_subsample_frac=subsample_frac, - data_key=data_key, prog_key=prog_key, - guide_targets_key=guide_targets_key, guide_assignment_key=guide_assignment_key, + mdata=mdata, + Target_Gene=target_gene, + ax=ax, + size=umap_dot_size, + umap_subsample_frac=subsample_frac, + data_key=data_key, + prog_key=prog_key, + guide_targets_key=guide_targets_key, + guide_assignment_key=guide_assignment_key, ) fig.tight_layout() out_path.parent.mkdir(parents=True, exist_ok=True) @@ -435,8 +581,15 @@ def _render_umap_guide_png(mdata, target_gene, out_path, umap_dot_size, subsampl plt.close(fig) -def _render_gene_dotplot_png(mdata, target_gene, out_path, - ensembl_to_symbol_file, groupby, gene_name_key, data_key="rna"): +def _render_gene_dotplot_png( + mdata, + target_gene, + out_path, + ensembl_to_symbol_file, + groupby, + gene_name_key, + data_key="rna", +): """Render the scanpy dotplot for the target gene to a PNG (standalone mode). Note: perturbed_gene_dotplot calls plt.close() internally when show=False, @@ -446,11 +599,16 @@ def _render_gene_dotplot_png(mdata, target_gene, out_path, out_path.parent.mkdir(parents=True, exist_ok=True) try: ax = perturbed_gene_dotplot( - mdata=mdata, Target_Gene=target_gene, + mdata=mdata, + Target_Gene=target_gene, ensembl_to_symbol_file=ensembl_to_symbol_file, - dotplot_groupby=groupby, gene_name_key=gene_name_key, + dotplot_groupby=groupby, + gene_name_key=gene_name_key, figsize=(4, 3), - save_path=None, save_name=None, show=True, ax=None, + save_path=None, + save_name=None, + show=True, + ax=None, data_key=data_key, ) if ax is None: @@ -460,8 +618,15 @@ def _render_gene_dotplot_png(mdata, target_gene, out_path, plt.close(fig) except Exception as e: fig, ax = plt.subplots(figsize=(4, 3)) - ax.text(0.5, 0.5, f"dotplot unavailable\n({e})", - ha="center", va="center", fontsize=10, color="#888") + ax.text( + 0.5, + 0.5, + f"dotplot unavailable\n({e})", + ha="center", + va="center", + fontsize=10, + color="#888", + ) ax.set_axis_off() fig.savefig(str(out_path), format="png", bbox_inches="tight", dpi=150) plt.close(fig) @@ -561,6 +726,7 @@ def _render_gene_dotplot_png(mdata, target_gene, out_path, # Public API # --------------------------------------------------------------------------- + def export_gene_html( mdata, perturb_path_base, @@ -621,21 +787,50 @@ def export_gene_html( (gene_dir / "images").mkdir(parents=True, exist_ok=True) # ---- images (matplotlib → PNG) ---- - _render_umap_expression_png(mdata, Target_Gene, gene_dir / "images" / "umap_expression.png", - ensembl_to_symbol_file, gene_name_key, umap_dot_size, subsample_frac, - data_key=data_key) - _render_umap_guide_png(mdata, Target_Gene, gene_dir / "images" / "umap_guide.png", - umap_dot_size, subsample_frac, - data_key=data_key, prog_key=prog_key, - guide_targets_key=guide_targets_key, guide_assignment_key=guide_assignment_key) - _render_gene_dotplot_png(mdata, Target_Gene, gene_dir / "images" / "gene_dotplot.png", - ensembl_to_symbol_file, groupby, gene_name_key, data_key=data_key) + _render_umap_expression_png( + mdata, + Target_Gene, + gene_dir / "images" / "umap_expression.png", + ensembl_to_symbol_file, + gene_name_key, + umap_dot_size, + subsample_frac, + data_key=data_key, + ) + _render_umap_guide_png( + mdata, + Target_Gene, + gene_dir / "images" / "umap_guide.png", + umap_dot_size, + subsample_frac, + data_key=data_key, + prog_key=prog_key, + guide_targets_key=guide_targets_key, + guide_assignment_key=guide_assignment_key, + ) + _render_gene_dotplot_png( + mdata, + Target_Gene, + gene_dir / "images" / "gene_dotplot.png", + ensembl_to_symbol_file, + groupby, + gene_name_key, + data_key=data_key, + ) # ---- header-row interactive panels ---- - top_programs_d = _build_top_programs(mdata, Target_Gene, top_n_programs, - ensembl_to_symbol_file, gene_name_key, - data_key=data_key, prog_key=prog_key) - corr_d = _build_gene_correlations(gene_loading_corr_matrix, Target_Gene, top_corr_genes) + top_programs_d = _build_top_programs( + mdata, + Target_Gene, + top_n_programs, + ensembl_to_symbol_file, + gene_name_key, + data_key=data_key, + prog_key=prog_key, + ) + corr_d = _build_gene_correlations( + gene_loading_corr_matrix, Target_Gene, top_corr_genes + ) _write_json(gene_dir / "data" / "top_programs.json", top_programs_d) _write_json(gene_dir / "data" / "correlations.json", corr_d) @@ -643,29 +838,52 @@ def export_gene_html( corr_fig = _make_correlations_fig(corr_d) # ---- wide KD vs control bar ---- - kd_d = _build_kd_vs_control(mdata, Target_Gene, groupby, control_target_name, gene_name_key, - data_key=data_key, prog_key=prog_key, - guide_targets_key=guide_targets_key, guide_assignment_key=guide_assignment_key) + kd_d = _build_kd_vs_control( + mdata, + Target_Gene, + groupby, + control_target_name, + gene_name_key, + data_key=data_key, + prog_key=prog_key, + guide_targets_key=guide_targets_key, + guide_assignment_key=guide_assignment_key, + ) _write_json(gene_dir / "data" / "kd_vs_control.json", kd_d) kd_fig = _make_kd_vs_control_fig(kd_d) # ---- per-sample panels ---- sample_blocks = [] sig_per_sample = {} - for samp in (sample if plot_perturbation else []): + for samp in sample if plot_perturbation else []: skey = _safe_sample_key(samp) perturb_path = f"{perturb_path_base}_{samp}.txt" - log2fc_d = _build_log2fc_gene(perturb_path, Target_Gene, perturb_target_col, - perturb_program_col, perturb_log2fc_col, - top_n_programs, significance_threshold) - volcano_d = _build_volcano_gene(perturb_path, Target_Gene, perturb_target_col, - perturb_program_col, perturb_log2fc_col, - volcano_log2fc_min, volcano_log2fc_max, - significance_threshold) - programs_dot_d = _build_programs_dotplot(mdata, log2fc_d["regulators"], groupby, - data_key=data_key, prog_key=prog_key) - waterfall_d = _build_gene_waterfall(perturb_corr_by_sample[samp], Target_Gene, top_corr_genes) + log2fc_d = _build_log2fc_gene( + perturb_path, + Target_Gene, + perturb_target_col, + perturb_program_col, + perturb_log2fc_col, + top_n_programs, + significance_threshold, + ) + volcano_d = _build_volcano_gene( + perturb_path, + Target_Gene, + perturb_target_col, + perturb_program_col, + perturb_log2fc_col, + volcano_log2fc_min, + volcano_log2fc_max, + significance_threshold, + ) + programs_dot_d = _build_programs_dotplot( + mdata, log2fc_d["regulators"], groupby, data_key=data_key, prog_key=prog_key + ) + waterfall_d = _build_gene_waterfall( + perturb_corr_by_sample[samp], Target_Gene, top_corr_genes + ) _write_json(gene_dir / "data" / f"log2fc_{skey}.json", log2fc_d) _write_json(gene_dir / "data" / f"volcano_{skey}.json", volcano_d) @@ -674,13 +892,24 @@ def export_gene_html( sig_per_sample[samp] = sum(1 for c in volcano_d["category"] if c != "ns") - sample_blocks.append(_SAMPLE_BLOCK.format( - sample=samp, - log2fc_div=_fig_to_div(_make_log2fc_fig(log2fc_d), f"div-log2fc-{skey}"), - volcano_div=_fig_to_div(_make_volcano_fig(volcano_d), f"div-volcano-{skey}"), - dotplot_div=_fig_to_div(_make_programs_dotplot_fig(programs_dot_d), f"div-dotplot-{skey}"), - waterfall_div=_fig_to_div(_make_waterfall_fig(waterfall_d, Target_Gene), f"div-waterfall-{skey}"), - )) + sample_blocks.append( + _SAMPLE_BLOCK.format( + sample=samp, + log2fc_div=_fig_to_div( + _make_log2fc_fig(log2fc_d), f"div-log2fc-{skey}" + ), + volcano_div=_fig_to_div( + _make_volcano_fig(volcano_d), f"div-volcano-{skey}" + ), + dotplot_div=_fig_to_div( + _make_programs_dotplot_fig(programs_dot_d), f"div-dotplot-{skey}" + ), + waterfall_div=_fig_to_div( + _make_waterfall_fig(waterfall_d, Target_Gene), + f"div-waterfall-{skey}", + ), + ) + ) # ---- nav strings ---- position = "" @@ -695,10 +924,14 @@ def export_gene_html( # ---- assemble page ---- page = _PAGE_TEMPLATE.format( - gene=Target_Gene, cdn=PLOTLY_CDN, fonts=GOOGLE_FONTS, + gene=Target_Gene, + cdn=PLOTLY_CDN, + fonts=GOOGLE_FONTS, position=position, - prev_href=prev_href, prev_disabled=prev_disabled, - next_href=next_href, next_disabled=next_disabled, + prev_href=prev_href, + prev_disabled=prev_disabled, + next_href=next_href, + next_disabled=next_disabled, top_programs_div=_fig_to_div(top_programs_fig, "div-top-programs"), corr_div=_fig_to_div(corr_fig, "div-gene-corr"), kd_div=_fig_to_div(kd_fig, "div-kd"), @@ -719,11 +952,21 @@ def export_gene_html( }, "top_loaded_programs": top_programs_d["programs"][::-1], "top_correlated_genes": { - "positive": [p for p, d in zip(corr_d["programs"], corr_d["direction"]) if d == "positive"], - "negative": [p for p, d in zip(corr_d["programs"], corr_d["direction"]) if d == "negative"], + "positive": [ + p + for p, d in zip(corr_d["programs"], corr_d["direction"]) + if d == "positive" + ], + "negative": [ + p + for p, d in zip(corr_d["programs"], corr_d["direction"]) + if d == "negative" + ], }, "n_significant_program_perturbations_total": int(sum(sig_per_sample.values())), - "n_significant_program_perturbations_per_sample": {str(k): int(v) for k, v in sig_per_sample.items()}, + "n_significant_program_perturbations_per_sample": { + str(k): int(v) for k, v in sig_per_sample.items() + }, "kd_effect_pct_of_control": kd_d.get("all_effect_pct"), "control_target_name": control_target_name, "generated_at": datetime.datetime.utcnow().isoformat() + "Z", @@ -743,8 +986,9 @@ def write_gene_share_index(html_share_path, gene_list, config_snapshot): "generated_at": datetime.datetime.utcnow().isoformat() + "Z", "kind": "perturbed_gene", "gene_list": list(map(str, gene_list)), - "config": {k: v for k, v in config_snapshot.items() - if not isinstance(v, (np.ndarray,))}, + "config": { + k: v for k, v in config_snapshot.items() if not isinstance(v, (np.ndarray,)) + }, } _write_json(share_root / "shared" / "manifest.json", manifest) @@ -761,8 +1005,10 @@ def write_gene_share_index(html_share_path, gene_list, config_snapshot): n_sig = int(m.get("n_significant_program_perturbations_total", 0)) kd_pct = m.get("kd_effect_pct_of_control") - chips = "".join(f'P{p}' for p in top_progs) \ - or '—' + chips = ( + "".join(f'P{p}' for p in top_progs) + or '—' + ) if kd_pct is None: kd_label = "n/a" kd_cls = "zero" @@ -776,20 +1022,20 @@ def write_gene_share_index(html_share_path, gene_list, config_snapshot): f'
' f' Gene' f' {gene}' - f'
' + f" " f'
' f'
Top loaded programs
' f'
{chips}
' - f'
' + f" " f'
' f' {kd_label}' f' vs control' - f'
' + f" " f'
' f' {n_sig}' f' sig programs' - f'
' - f'' + f" " + f"" ) n_genes = len(cards) diff --git a/src/Stage3_Interpretation/A_Plotting/src/html_Program_QC_plots.py b/src/Stage3_Interpretation/A_Plotting/src/html_Program_QC_plots.py index f71873d..c76d8fb 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/html_Program_QC_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/html_Program_QC_plots.py @@ -13,6 +13,7 @@ Plotly is loaded from CDN; the folder is movable. """ + from __future__ import annotations import os @@ -37,11 +38,11 @@ GOOGLE_FONTS = "https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@500&display=swap" # Unified palette — used in every Plotly figure -COLOR_UP = "#c0392b" # positive effect / upregulated -COLOR_DOWN = "#1c5f8f" # negative effect / downregulated -COLOR_NS = "#cbd5db" # not significant +COLOR_UP = "#c0392b" # positive effect / upregulated +COLOR_DOWN = "#1c5f8f" # negative effect / downregulated +COLOR_NS = "#cbd5db" # not significant COLOR_NEUTRAL = "#7f8c8d" # neutral grey bar -COLOR_ACCENT = "#1769aa" # link / accent +COLOR_ACCENT = "#1769aa" # link / accent # Plotly template applied to every figure for consistency _PLOTLY_TEMPLATE = go.layout.Template( @@ -49,10 +50,22 @@ font=dict(family="Inter, system-ui, sans-serif", size=11, color="#222"), plot_bgcolor="#ffffff", paper_bgcolor="rgba(0,0,0,0)", - xaxis=dict(gridcolor="#eef0f2", linecolor="#cfd4d8", zerolinecolor="#dfe3e6", - ticks="outside", tickcolor="#cfd4d8", tickfont=dict(size=10)), - yaxis=dict(gridcolor="#eef0f2", linecolor="#cfd4d8", zerolinecolor="#dfe3e6", - ticks="outside", tickcolor="#cfd4d8", tickfont=dict(size=10)), + xaxis=dict( + gridcolor="#eef0f2", + linecolor="#cfd4d8", + zerolinecolor="#dfe3e6", + ticks="outside", + tickcolor="#cfd4d8", + tickfont=dict(size=10), + ), + yaxis=dict( + gridcolor="#eef0f2", + linecolor="#cfd4d8", + zerolinecolor="#dfe3e6", + ticks="outside", + tickcolor="#cfd4d8", + tickfont=dict(size=10), + ), colorway=[COLOR_ACCENT, COLOR_UP, COLOR_DOWN, COLOR_NEUTRAL], margin=dict(l=50, r=20, t=20, b=50), hoverlabel=dict(bgcolor="#222", font=dict(color="#fff", size=11)), @@ -64,6 +77,7 @@ # JSON / file helpers # --------------------------------------------------------------------------- + def _sanitize(o): if isinstance(o, float): return None if (math.isnan(o) or math.isinf(o)) else o @@ -103,13 +117,26 @@ def _safe_sample_key(samp: str) -> str: # _make_*_fig functions below. # --------------------------------------------------------------------------- -def _build_top_genes(mdata, target_program, num_gene, file_to_dictionary, gene_name_key, - data_key="rna", prog_key="cNMF"): + +def _build_top_genes( + mdata, + target_program, + num_gene, + file_to_dictionary, + gene_name_key, + data_key="rna", + prog_key="cNMF", +): # Mirror of plot_top_gene_per_program data prep. X = mdata[prog_key].varm["loadings"] gene_names = mdata[prog_key].uns["var_names"] if gene_name_key is not None and gene_name_key in mdata[data_key].var.columns: - ens_to_sym = dict(zip(mdata[data_key].var_names, mdata[data_key].var[gene_name_key].astype(str))) + ens_to_sym = dict( + zip( + mdata[data_key].var_names, + mdata[data_key].var[gene_name_key].astype(str), + ) + ) gene_names = np.array([ens_to_sym.get(g, g) for g in gene_names]) if file_to_dictionary is not None: gene_names = rename_list_gene_dictionary(list(gene_names), file_to_dictionary) @@ -154,7 +181,11 @@ def _build_go_terms(GO_path, target_program, num_term, p_value_name, term_col): def _build_correlations(program_correlation, target_program, num_program): if str(target_program) not in program_correlation.columns: return {"programs": [], "r": [], "direction": []} - s = program_correlation.loc[str(target_program)].drop(str(target_program), errors="ignore").sort_values(ascending=False) + s = ( + program_correlation.loc[str(target_program)] + .drop(str(target_program), errors="ignore") + .sort_values(ascending=False) + ) top = s[s > 0].head(num_program) bottom = s[s < 0].tail(num_program) combined = pd.concat([bottom, top]) @@ -176,19 +207,40 @@ def _build_violin(mdata, target_program, groupby, data_key="rna", prog_key="cNMF else: groups = mdata[data_key].obs[groupby].values df = pd.DataFrame({"expression": scores, "group": list(groups)}) - order = list(df["group"].astype("category").cat.categories) if hasattr(df["group"], "cat") else sorted(df["group"].unique()) - summary = df.groupby("group")["expression"].agg(["mean", lambda x: float((x > x.mean()).mean())]) + order = ( + list(df["group"].astype("category").cat.categories) + if hasattr(df["group"], "cat") + else sorted(df["group"].unique()) + ) + summary = df.groupby("group")["expression"].agg( + ["mean", lambda x: float((x > x.mean()).mean())] + ) summary.columns = ["mean", "frac_above"] return { "groups": list(map(str, order)), - "per_group_expression": {str(g): list(df.loc[df["group"] == g, "expression"].values) for g in order}, - "summary": {str(g): {"mean": float(summary.loc[g, "mean"]), - "frac_above": float(summary.loc[g, "frac_above"])} for g in order}, + "per_group_expression": { + str(g): list(df.loc[df["group"] == g, "expression"].values) for g in order + }, + "summary": { + str(g): { + "mean": float(summary.loc[g, "mean"]), + "frac_above": float(summary.loc[g, "frac_above"]), + } + for g in order + }, } -def _build_log2fc(perturb_path, target_program, tagert_col_name, plot_col_name, - log2fc_col, num_item, p_value, gene_list): +def _build_log2fc( + perturb_path, + target_program, + tagert_col_name, + plot_col_name, + log2fc_col, + num_item, + p_value, + gene_list, +): df = pd.read_csv(perturb_path, sep="\t") df["program_name"] = df["program_name"].astype(str) if str(target_program) not in df[tagert_col_name].astype(str).values: @@ -199,7 +251,9 @@ def _build_log2fc(perturb_path, target_program, tagert_col_name, plot_col_name, sub = sub[sub["adj_pval"] < p_value].sort_values(by=log2fc_col, ascending=False) top = sub.head(num_item) bottom = sub.tail(num_item) - plot_data = pd.merge(top, bottom, how="outer").sort_values(by=log2fc_col, ascending=False) + plot_data = pd.merge(top, bottom, how="outer").sort_values( + by=log2fc_col, ascending=False + ) return { "regulators": list(plot_data[plot_col_name].astype(str).values), "log2fc": [float(x) for x in plot_data[log2fc_col].values], @@ -207,13 +261,27 @@ def _build_log2fc(perturb_path, target_program, tagert_col_name, plot_col_name, } -def _build_volcano(perturb_path, target_program, tagert_col_name, plot_col_name, - log2fc_col, down_thred_log, up_thred_log, p_value, gene_list): +def _build_volcano( + perturb_path, + target_program, + tagert_col_name, + plot_col_name, + log2fc_col, + down_thred_log, + up_thred_log, + p_value, + gene_list, +): df = pd.read_csv(perturb_path, sep="\t") df["program_name"] = df["program_name"].astype(str) if str(target_program) not in df[tagert_col_name].astype(str).values: - return {"regulators": [], "log2fc": [], "neglog10p": [], "category": [], - "thresholds": {"down": down_thred_log, "up": up_thred_log, "p": p_value}} + return { + "regulators": [], + "log2fc": [], + "neglog10p": [], + "category": [], + "thresholds": {"down": down_thred_log, "up": up_thred_log, "p": p_value}, + } sub = df.loc[df[tagert_col_name].astype(str) == str(target_program)].copy() if gene_list: sub = sub[sub[plot_col_name].isin(set(gene_list))] @@ -233,7 +301,11 @@ def _build_volcano(perturb_path, target_program, tagert_col_name, plot_col_name, "neglog10p": [float(x) for x in sub["neglog10p"].values], "adj_pval": [float(x) for x in sub["adj_pval"].values], "category": cats, - "thresholds": {"down": float(down_thred_log), "up": float(up_thred_log), "p": float(p_value)}, + "thresholds": { + "down": float(down_thred_log), + "up": float(up_thred_log), + "p": float(p_value), + }, } @@ -242,13 +314,17 @@ def _build_dotplot(mdata, gene_list, groupby, gene_name_key, data_key="rna"): return {"genes": [], "conditions": [], "mean": [], "frac": []} if gene_name_key is not None and gene_name_key in mdata[data_key].var.columns: sym_to_ens = {} - for ens, sym in zip(mdata[data_key].var_names, mdata[data_key].var[gene_name_key].astype(str)): + for ens, sym in zip( + mdata[data_key].var_names, mdata[data_key].var[gene_name_key].astype(str) + ): sym_to_ens.setdefault(sym, ens) gene_indices = [] kept_genes = [] for g in gene_list: if g in sym_to_ens: - gene_indices.append(list(mdata[data_key].var_names).index(sym_to_ens[g])) + gene_indices.append( + list(mdata[data_key].var_names).index(sym_to_ens[g]) + ) kept_genes.append(g) gene_list = kept_genes else: @@ -264,7 +340,11 @@ def _build_dotplot(mdata, gene_list, groupby, gene_name_key, data_key="rna"): X = mdata[data_key].X groups = mdata[data_key].obs[groupby].values - conds_cat = mdata[data_key].obs[groupby].astype("category").cat.categories if hasattr(mdata[data_key].obs[groupby], "cat") else sorted(np.unique(groups)) + conds_cat = ( + mdata[data_key].obs[groupby].astype("category").cat.categories + if hasattr(mdata[data_key].obs[groupby], "cat") + else sorted(np.unique(groups)) + ) conditions = [str(c) for c in conds_cat] mean_mat = [] frac_mat = [] @@ -298,30 +378,65 @@ def _build_waterfall(corr_matrix, target_program, top_num): } -def _build_heatmap(perturb_path_base, mdata, target_program, sample, - tagert_col_name, plot_col_name, log2fc_col, p_value, groupby, - data_key="rna", prog_key="cNMF"): - dfs = [pd.read_csv(f"{perturb_path_base}_{samp}.txt", sep="\t").assign(sample=samp) for samp in sample] +def _build_heatmap( + perturb_path_base, + mdata, + target_program, + sample, + tagert_col_name, + plot_col_name, + log2fc_col, + p_value, + groupby, + data_key="rna", + prog_key="cNMF", +): + dfs = [ + pd.read_csv(f"{perturb_path_base}_{samp}.txt", sep="\t").assign(sample=samp) + for samp in sample + ] combined = pd.concat(dfs, ignore_index=True) combined["program_name"] = combined["program_name"].astype(str) sub = combined.loc[combined[tagert_col_name].astype(str) == str(target_program)] keep = sub[sub["adj_pval"] < p_value][plot_col_name].unique() sub = sub[sub[plot_col_name].isin(keep)] if sub.empty: - return {"conditions": [], "regulators": [], "log2fc": [], "sig": [], "expression": []} + return { + "conditions": [], + "regulators": [], + "log2fc": [], + "sig": [], + "expression": [], + } pv = sub.pivot(columns=plot_col_name, index="sample", values=log2fc_col) pp = sub.pivot(columns=plot_col_name, index="sample", values="adj_pval") pv = pv.reindex(sample) pp = pp.reindex(sample) - expr = compute_program_expression_by_condition(mdata, target_program, groupby=groupby, data_key=data_key, prog_key=prog_key) + expr = compute_program_expression_by_condition( + mdata, target_program, groupby=groupby, data_key=data_key, prog_key=prog_key + ) expr_vals = [float(expr.get(c, 0.0)) for c in pv.index] return { "conditions": list(map(str, pv.index)), "regulators": list(map(str, pv.columns)), - "log2fc": [[None if (isinstance(v, float) and math.isnan(v)) else float(v) for v in row] - for row in pv.values], - "sig": [[(False if (isinstance(v, float) and math.isnan(v)) else bool(v < p_value)) for v in row] - for row in pp.values], + "log2fc": [ + [ + None if (isinstance(v, float) and math.isnan(v)) else float(v) + for v in row + ] + for row in pv.values + ], + "sig": [ + [ + ( + False + if (isinstance(v, float) and math.isnan(v)) + else bool(v < p_value) + ) + for v in row + ] + for row in pp.values + ], "expression": expr_vals, } @@ -330,6 +445,7 @@ def _build_heatmap(perturb_path_base, mdata, target_program, sample, # Plotly figure builders. No internal titles — titles are emitted in HTML. # --------------------------------------------------------------------------- + def _fmt_layout(fig: go.Figure, height: int = 320) -> go.Figure: fig.update_layout(template=_PLOTLY_TEMPLATE, height=height) return fig @@ -338,13 +454,19 @@ def _fmt_layout(fig: go.Figure, height: int = 320) -> go.Figure: def _empty_fig(message: str) -> go.Figure: fig = go.Figure() fig.add_annotation( - text=f"{message}", showarrow=False, - xref="paper", yref="paper", x=0.5, y=0.5, + text=f"{message}", + showarrow=False, + xref="paper", + yref="paper", + x=0.5, + y=0.5, font=dict(color="#999", size=12), ) fig.update_layout( - template=_PLOTLY_TEMPLATE, height=160, - xaxis=dict(visible=False), yaxis=dict(visible=False), + template=_PLOTLY_TEMPLATE, + height=160, + xaxis=dict(visible=False), + yaxis=dict(visible=False), margin=dict(l=10, r=10, t=10, b=10), ) return fig @@ -353,11 +475,15 @@ def _empty_fig(message: str) -> go.Figure: def _make_top_genes_fig(data): if not data["genes"]: return _empty_fig("No gene loadings") - fig = go.Figure(go.Bar( - x=data["loadings"], y=data["genes"], orientation="h", - marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), - hovertemplate="%{y}
Loading: %{x:.3g}", - )) + fig = go.Figure( + go.Bar( + x=data["loadings"], + y=data["genes"], + orientation="h", + marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), + hovertemplate="%{y}
Loading: %{x:.3g}", + ) + ) fig.update_layout(xaxis_title="Gene loading (z-score)", yaxis_title="") return _fmt_layout(fig) @@ -365,14 +491,19 @@ def _make_top_genes_fig(data): def _make_go_fig(data): if not data["terms"]: return _empty_fig("No GO enrichment") - fig = go.Figure(go.Bar( - x=data["neglog10p"], y=data["terms"], orientation="h", - marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), - customdata=data["adj_pval"], - hovertemplate="%{y}
-log10(adj_p): %{x:.2f}
adj_p: %{customdata:.2e}", - )) - fig.update_layout(xaxis_title="-log10 adj. p-value", yaxis_title="", - yaxis=dict(automargin=True)) + fig = go.Figure( + go.Bar( + x=data["neglog10p"], + y=data["terms"], + orientation="h", + marker=dict(color=COLOR_NEUTRAL, line=dict(width=0)), + customdata=data["adj_pval"], + hovertemplate="%{y}
-log10(adj_p): %{x:.2f}
adj_p: %{customdata:.2e}", + ) + ) + fig.update_layout( + xaxis_title="-log10 adj. p-value", yaxis_title="", yaxis=dict(automargin=True) + ) return _fmt_layout(fig) @@ -380,14 +511,19 @@ def _make_correlations_fig(data): if not data["programs"]: return _empty_fig("No correlations") colors = [COLOR_DOWN if d == "negative" else COLOR_UP for d in data["direction"]] - fig = go.Figure(go.Bar( - x=data["r"], y=data["programs"], orientation="h", - marker=dict(color=colors, line=dict(width=0)), - hovertemplate="Program %{y}
r = %{x:.3f}", - )) + fig = go.Figure( + go.Bar( + x=data["r"], + y=data["programs"], + orientation="h", + marker=dict(color=colors, line=dict(width=0)), + hovertemplate="Program %{y}
r = %{x:.3f}", + ) + ) fig.add_vline(x=0, line_width=0.5, line_color="#555") - fig.update_layout(xaxis_title="Pearson r", yaxis_title="Program", - yaxis=dict(type="category")) + fig.update_layout( + xaxis_title="Pearson r", yaxis_title="Program", yaxis=dict(type="category") + ) return _fmt_layout(fig) @@ -397,15 +533,21 @@ def _make_violin_fig(data): fig = go.Figure() for g in data["groups"]: s = data["summary"][g] - fig.add_trace(go.Violin( - y=data["per_group_expression"][g], name=g, - box_visible=True, meanline_visible=True, points=False, - line=dict(color=COLOR_ACCENT, width=1), - fillcolor="rgba(23,105,170,0.20)", - hovertemplate=f"{g}
mean={s['mean']:.3f}
frac>mean={s['frac_above']:.2f}", - )) - fig.update_layout(xaxis_title="Condition", yaxis_title="Program expression", - showlegend=False) + fig.add_trace( + go.Violin( + y=data["per_group_expression"][g], + name=g, + box_visible=True, + meanline_visible=True, + points=False, + line=dict(color=COLOR_ACCENT, width=1), + fillcolor="rgba(23,105,170,0.20)", + hovertemplate=f"{g}
mean={s['mean']:.3f}
frac>mean={s['frac_above']:.2f}", + ) + ) + fig.update_layout( + xaxis_title="Condition", yaxis_title="Program expression", showlegend=False + ) return _fmt_layout(fig) @@ -413,17 +555,22 @@ def _make_log2fc_fig(data): if not data["regulators"]: return _empty_fig("No significant regulators") colors = [COLOR_UP if v > 0 else COLOR_DOWN for v in data["log2fc"]] - fig = go.Figure(go.Bar( - x=data["log2fc"], y=data["regulators"], orientation="h", - marker=dict(color=colors, line=dict(width=0)), - customdata=data["adj_pval"], - hovertemplate="%{y}
log2FC: %{x:.3f}
adj_p: %{customdata:.2e}", - )) + fig = go.Figure( + go.Bar( + x=data["log2fc"], + y=data["regulators"], + orientation="h", + marker=dict(color=colors, line=dict(width=0)), + customdata=data["adj_pval"], + hovertemplate="%{y}
log2FC: %{x:.3f}
adj_p: %{customdata:.2e}", + ) + ) fig.add_vline(x=0, line_width=0.5, line_color="#555") # Program/regulator labels are numeric-looking strings ("3", "28", ...); without # type="category" Plotly renders a linear y-axis and scatters the bars by value. - fig.update_layout(xaxis_title="log2FC", yaxis_title="Regulator", - yaxis=dict(type="category")) + fig.update_layout( + xaxis_title="log2FC", yaxis_title="Regulator", yaxis=dict(type="category") + ) return _fmt_layout(fig) @@ -437,26 +584,41 @@ def _make_volcano_fig(data): idx = [i for i, c in enumerate(data["category"]) if c == cat] if not idx: continue - fig.add_trace(go.Scattergl( - x=[data["log2fc"][i] for i in idx], - y=[data["neglog10p"][i] for i in idx], - mode="markers", name=cat_names[cat], - text=[data["regulators"][i] for i in idx], - customdata=[data["adj_pval"][i] for i in idx], - marker=dict(color=cat_colors[cat], size=6, - opacity=0.85 if cat != "ns" else 0.45, - line=dict(width=0)), - hovertemplate="%{text}
log2FC: %{x:.3f}
-log10 adj_p: %{y:.2f}
adj_p: %{customdata:.2e}", - )) + fig.add_trace( + go.Scattergl( + x=[data["log2fc"][i] for i in idx], + y=[data["neglog10p"][i] for i in idx], + mode="markers", + name=cat_names[cat], + text=[data["regulators"][i] for i in idx], + customdata=[data["adj_pval"][i] for i in idx], + marker=dict( + color=cat_colors[cat], + size=6, + opacity=0.85 if cat != "ns" else 0.45, + line=dict(width=0), + ), + hovertemplate="%{text}
log2FC: %{x:.3f}
-log10 adj_p: %{y:.2f}
adj_p: %{customdata:.2e}", + ) + ) th = data["thresholds"] fig.add_vline(x=th["down"], line_dash="dash", line_color="#999", line_width=1) fig.add_vline(x=th["up"], line_dash="dash", line_color="#999", line_width=1) if th["p"] > 0: - fig.add_hline(y=-math.log10(th["p"]), line_dash="dash", line_color="#999", line_width=1) + fig.add_hline( + y=-math.log10(th["p"]), line_dash="dash", line_color="#999", line_width=1 + ) fig.update_layout( - xaxis_title="log2FC", yaxis_title="-log10 adj. p-value", - legend=dict(orientation="h", x=1, xanchor="right", y=1.08, yanchor="bottom", - bgcolor="rgba(0,0,0,0)"), + xaxis_title="log2FC", + yaxis_title="-log10 adj. p-value", + legend=dict( + orientation="h", + x=1, + xanchor="right", + y=1.08, + yanchor="bottom", + bgcolor="rgba(0,0,0,0)", + ), ) return _fmt_layout(fig) @@ -475,15 +637,31 @@ def _make_dotplot_fig(data): sizes.append(4 + 20 * (f / max_frac)) colors.append(m) hovers.append(f"{gene}
{cond}
mean={m:.3f}
frac={f:.2f}") - fig = go.Figure(go.Scatter( - x=xs, y=ys, mode="markers", - marker=dict(size=sizes, color=colors, colorscale="Reds", showscale=True, - colorbar=dict(title=dict(text="mean", side="right"), thickness=8, len=0.7), - line=dict(width=0)), - text=hovers, hovertemplate="%{text}", - )) - fig.update_layout(xaxis_title="Condition", yaxis_title="Regulator", - xaxis=dict(type="category"), yaxis=dict(type="category")) + fig = go.Figure( + go.Scatter( + x=xs, + y=ys, + mode="markers", + marker=dict( + size=sizes, + color=colors, + colorscale="Reds", + showscale=True, + colorbar=dict( + title=dict(text="mean", side="right"), thickness=8, len=0.7 + ), + line=dict(width=0), + ), + text=hovers, + hovertemplate="%{text}", + ) + ) + fig.update_layout( + xaxis_title="Condition", + yaxis_title="Regulator", + xaxis=dict(type="category"), + yaxis=dict(type="category"), + ) return _fmt_layout(fig) @@ -491,14 +669,20 @@ def _make_waterfall_fig(data, target_program): if not data["programs"]: return _empty_fig("No correlation data") xs = list(range(len(data["programs"]))) - fig = go.Figure(go.Scattergl( - x=xs, y=data["r"], mode="markers", - text=data["programs"], - marker=dict(size=[8 if lab else 5 for lab in data["labeled"]], - color=[COLOR_ACCENT if lab else COLOR_NS for lab in data["labeled"]], - line=dict(width=0)), - hovertemplate="Program %{text}
r = %{y:.3f}", - )) + fig = go.Figure( + go.Scattergl( + x=xs, + y=data["r"], + mode="markers", + text=data["programs"], + marker=dict( + size=[8 if lab else 5 for lab in data["labeled"]], + color=[COLOR_ACCENT if lab else COLOR_NS for lab in data["labeled"]], + line=dict(width=0), + ), + hovertemplate="Program %{text}
r = %{y:.3f}", + ) + ) fig.add_hline(y=0, line_dash="dash", line_color="#999", line_width=1) fig.update_layout(xaxis_title="Rank", yaxis_title=f"r vs Program {target_program}") return _fmt_layout(fig) @@ -512,22 +696,42 @@ def _make_heatmap_fig(data): line = [] for j, v in enumerate(_row): s = " *" if data["sig"][i][j] else "" - line.append(f"{data['regulators'][j]}
{data['conditions'][i]}
log2FC={v}{s}") + line.append( + f"{data['regulators'][j]}
{data['conditions'][i]}
log2FC={v}{s}" + ) htext.append(line) - fig = go.Figure(go.Heatmap( - z=data["log2fc"], x=data["regulators"], y=data["conditions"], - zmin=-1, zmax=1, zmid=0, colorscale="RdBu_r", reversescale=False, - text=htext, hovertemplate="%{text}", - colorbar=dict(title=dict(text="log2FC", side="right"), thickness=8, len=0.7), - )) + fig = go.Figure( + go.Heatmap( + z=data["log2fc"], + x=data["regulators"], + y=data["conditions"], + zmin=-1, + zmax=1, + zmid=0, + colorscale="RdBu_r", + reversescale=False, + text=htext, + hovertemplate="%{text}", + colorbar=dict( + title=dict(text="log2FC", side="right"), thickness=8, len=0.7 + ), + ) + ) annotations = [] for i, row in enumerate(data["sig"]): for j, sig in enumerate(row): if sig: - annotations.append(dict( - x=data["regulators"][j], y=data["conditions"][i], - text="*", showarrow=False, - font=dict(color="#111", size=14, family="JetBrains Mono, monospace"))) + annotations.append( + dict( + x=data["regulators"][j], + y=data["conditions"][i], + text="*", + showarrow=False, + font=dict( + color="#111", size=14, family="JetBrains Mono, monospace" + ), + ) + ) height = max(280, 80 + 36 * len(data["conditions"])) fig.update_layout(annotations=annotations, xaxis=dict(tickangle=-45), height=height) return _fmt_layout(fig, height=height) @@ -537,11 +741,15 @@ def _make_heatmap_fig(data): # UMAP PNG rendering — reuses the existing matplotlib function # --------------------------------------------------------------------------- + def _render_umap_png(mdata, target_program, out_path, subsample_frac, prog_key="cNMF"): fig, ax = plt.subplots(figsize=(5, 4)) plot_umap_per_program( - mdata=mdata, Target_Program=target_program, ax=ax, - subsample_frac=subsample_frac, prog_key=prog_key, + mdata=mdata, + Target_Program=target_program, + ax=ax, + subsample_frac=subsample_frac, + prog_key=prog_key, ) fig.tight_layout() out_path.parent.mkdir(parents=True, exist_ok=True) @@ -831,8 +1039,12 @@ def _render_umap_png(mdata, target_program, out_path, subsample_frac, prog_key=" def _fig_to_div(fig: go.Figure, div_id: str) -> str: - return fig.to_html(include_plotlyjs=False, full_html=False, div_id=div_id, - config={"displaylogo": False, "responsive": True}) + return fig.to_html( + include_plotlyjs=False, + full_html=False, + div_id=div_id, + config={"displaylogo": False, "responsive": True}, + ) # Pages load Plotly from this local, relative path (every page lives one level @@ -851,6 +1063,7 @@ def _write_plotly_js(share_root): so the renderer always matches the version that produced the figure divs. """ from plotly.offline import get_plotlyjs + shared = Path(share_root) / "shared" shared.mkdir(parents=True, exist_ok=True) js_path = shared / "plotly.min.js" @@ -862,6 +1075,7 @@ def _write_plotly_js(share_root): # Public API # --------------------------------------------------------------------------- + def export_program_html( mdata, perturb_path_base, @@ -919,14 +1133,31 @@ def export_program_html( (prog_dir / "images").mkdir(parents=True, exist_ok=True) # UMAP PNG - _render_umap_png(mdata, Target_Program, prog_dir / "images" / "umap.png", subsample_frac, prog_key=prog_key) + _render_umap_png( + mdata, + Target_Program, + prog_dir / "images" / "umap.png", + subsample_frac, + prog_key=prog_key, + ) # ---- header panels ---- - top_genes_d = _build_top_genes(mdata, Target_Program, top_enrichned_term, file_to_dictionary, gene_name_key, - data_key=data_key, prog_key=prog_key) - go_d = _build_go_terms(GO_path, Target_Program, top_enrichned_term, go_p_value_name, go_term_col) + top_genes_d = _build_top_genes( + mdata, + Target_Program, + top_enrichned_term, + file_to_dictionary, + gene_name_key, + data_key=data_key, + prog_key=prog_key, + ) + go_d = _build_go_terms( + GO_path, Target_Program, top_enrichned_term, go_p_value_name, go_term_col + ) corr_d = _build_correlations(program_correlation, Target_Program, top_program) - violin_d = _build_violin(mdata, Target_Program, groupby, data_key=data_key, prog_key=prog_key) + violin_d = _build_violin( + mdata, Target_Program, groupby, data_key=data_key, prog_key=prog_key + ) _write_json(prog_dir / "data" / "top_genes.json", top_genes_d) _write_json(prog_dir / "data" / "go_terms.json", go_d) @@ -940,35 +1171,77 @@ def export_program_html( # ---- per-sample panels ---- sample_blocks = [] - for samp in (sample if plot_perturbation else []): + for samp in sample if plot_perturbation else []: skey = _safe_sample_key(samp) perturb_path = f"{perturb_path_base}_{samp}.txt" - log2fc_d = _build_log2fc(perturb_path, Target_Program, tagert_col_name, plot_col_name, - log2fc_col, top_enrichned_term, p_value, perturbed_gene_found) - volcano_d = _build_volcano(perturb_path, Target_Program, tagert_col_name, plot_col_name, - log2fc_col, down_thred_log, up_thred_log, p_value, perturbed_gene_found) - dotplot_d = _build_dotplot(mdata, log2fc_d["regulators"], groupby, gene_name_key, data_key=data_key) - waterfall_d = _build_waterfall(waterfall_correlation[samp], Target_Program, top_enrichned_term) + log2fc_d = _build_log2fc( + perturb_path, + Target_Program, + tagert_col_name, + plot_col_name, + log2fc_col, + top_enrichned_term, + p_value, + perturbed_gene_found, + ) + volcano_d = _build_volcano( + perturb_path, + Target_Program, + tagert_col_name, + plot_col_name, + log2fc_col, + down_thred_log, + up_thred_log, + p_value, + perturbed_gene_found, + ) + dotplot_d = _build_dotplot( + mdata, log2fc_d["regulators"], groupby, gene_name_key, data_key=data_key + ) + waterfall_d = _build_waterfall( + waterfall_correlation[samp], Target_Program, top_enrichned_term + ) _write_json(prog_dir / "data" / f"log2fc_{skey}.json", log2fc_d) _write_json(prog_dir / "data" / f"volcano_{skey}.json", volcano_d) _write_json(prog_dir / "data" / f"dotplot_{skey}.json", dotplot_d) _write_json(prog_dir / "data" / f"waterfall_{skey}.json", waterfall_d) - sample_blocks.append(_SAMPLE_BLOCK.format( - sample=samp, - log2fc_div=_fig_to_div(_make_log2fc_fig(log2fc_d), f"div-log2fc-{skey}"), - volcano_div=_fig_to_div(_make_volcano_fig(volcano_d), f"div-volcano-{skey}"), - dotplot_div=_fig_to_div(_make_dotplot_fig(dotplot_d), f"div-dotplot-{skey}"), - waterfall_div=_fig_to_div(_make_waterfall_fig(waterfall_d, Target_Program), f"div-waterfall-{skey}"), - )) + sample_blocks.append( + _SAMPLE_BLOCK.format( + sample=samp, + log2fc_div=_fig_to_div( + _make_log2fc_fig(log2fc_d), f"div-log2fc-{skey}" + ), + volcano_div=_fig_to_div( + _make_volcano_fig(volcano_d), f"div-volcano-{skey}" + ), + dotplot_div=_fig_to_div( + _make_dotplot_fig(dotplot_d), f"div-dotplot-{skey}" + ), + waterfall_div=_fig_to_div( + _make_waterfall_fig(waterfall_d, Target_Program), + f"div-waterfall-{skey}", + ), + ) + ) # ---- heatmap ---- if plot_perturbation: - heatmap_d = _build_heatmap(perturb_path_base, mdata, Target_Program, sample, - tagert_col_name, plot_col_name, log2fc_col, p_value, groupby, - data_key=data_key, prog_key=prog_key) + heatmap_d = _build_heatmap( + perturb_path_base, + mdata, + Target_Program, + sample, + tagert_col_name, + plot_col_name, + log2fc_col, + p_value, + groupby, + data_key=data_key, + prog_key=prog_key, + ) _write_json(prog_dir / "data" / "heatmap.json", heatmap_d) heatmap_fig = _make_heatmap_fig(heatmap_d) else: @@ -981,11 +1254,17 @@ def export_program_html( if prev_program_id is None: prev_href, prev_disabled = "#", " disabled" else: - prev_href, prev_disabled = f"../program_{prev_program_id}/program_{prev_program_id}.html", "" + prev_href, prev_disabled = ( + f"../program_{prev_program_id}/program_{prev_program_id}.html", + "", + ) if next_program_id is None: next_href, next_disabled = "#", " disabled" else: - next_href, next_disabled = f"../program_{next_program_id}/program_{next_program_id}.html", "" + next_href, next_disabled = ( + f"../program_{next_program_id}/program_{next_program_id}.html", + "", + ) # ---- assemble page ---- page = _PAGE_TEMPLATE.format( @@ -993,8 +1272,10 @@ def export_program_html( cdn=PLOTLY_CDN, fonts=GOOGLE_FONTS, position=position, - prev_href=prev_href, prev_disabled=prev_disabled, - next_href=next_href, next_disabled=next_disabled, + prev_href=prev_href, + prev_disabled=prev_disabled, + next_href=next_href, + next_disabled=next_disabled, violin_div=_fig_to_div(violin_fig, "div-violin"), go_div=_fig_to_div(go_fig, "div-go"), corr_div=_fig_to_div(corr_fig, "div-corr"), @@ -1009,7 +1290,7 @@ def export_program_html( # ---- metadata.json ---- n_sig_total = 0 top_go_term = go_d["terms"][0] if go_d["terms"] else None - for samp in (sample if plot_perturbation else []): + for samp in sample if plot_perturbation else []: skey = _safe_sample_key(samp) with open(prog_dir / "data" / f"volcano_{skey}.json") as f: v = json.load(f) @@ -1017,14 +1298,27 @@ def export_program_html( metadata = { "program_id": pid, "samples": list(map(str, sample)), - "thresholds": {"p_value": float(p_value), "down_thred_log": float(down_thred_log), - "up_thred_log": float(up_thred_log)}, + "thresholds": { + "p_value": float(p_value), + "down_thred_log": float(down_thred_log), + "up_thred_log": float(up_thred_log), + }, "top_loaded_genes": top_genes_d["genes"][::-1], "top_correlated_programs": { - "positive": [p for p, d in zip(corr_d["programs"], corr_d["direction"]) if d == "positive"], - "negative": [p for p, d in zip(corr_d["programs"], corr_d["direction"]) if d == "negative"], + "positive": [ + p + for p, d in zip(corr_d["programs"], corr_d["direction"]) + if d == "positive" + ], + "negative": [ + p + for p, d in zip(corr_d["programs"], corr_d["direction"]) + if d == "negative" + ], }, - "top_GO_terms": [{"term": t, "adj_pval": p} for t, p in zip(go_d["terms"], go_d["adj_pval"])], + "top_GO_terms": [ + {"term": t, "adj_pval": p} for t, p in zip(go_d["terms"], go_d["adj_pval"]) + ], "top_GO_term": top_go_term, "n_significant_regulators_total": n_sig_total, "generated_at": datetime.datetime.utcnow().isoformat() + "Z", @@ -1044,8 +1338,9 @@ def write_share_index(html_share_path, program_ids, config_snapshot): manifest = { "generated_at": datetime.datetime.utcnow().isoformat() + "Z", "program_ids": [str(p) for p in program_ids], - "config": {k: v for k, v in config_snapshot.items() - if not isinstance(v, (np.ndarray,))}, + "config": { + k: v for k, v in config_snapshot.items() if not isinstance(v, (np.ndarray,)) + }, } _write_json(share_root / "shared" / "manifest.json", manifest) @@ -1061,7 +1356,10 @@ def write_share_index(html_share_path, program_ids, config_snapshot): top_go = m.get("top_GO_term") or "—" n_sig = int(m.get("n_significant_regulators_total", 0)) - chips = "".join(f'{g}' for g in top_genes) or '—' + chips = ( + "".join(f'{g}' for g in top_genes) + or '—' + ) badge_cls = "sig" if n_sig > 0 else "zero" cards.append( @@ -1069,16 +1367,16 @@ def write_share_index(html_share_path, program_ids, config_snapshot): f'
' f' Program' f' {pid}' - f'
' + f" " f'
' f'
{top_go}
' f'
{chips}
' - f'
' + f" " f'
' f' {n_sig}' f' sig regulators' - f'
' - f'' + f" " + f"" ) n_progs = len(cards) diff --git a/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py b/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py index 150b033..1b31c61 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py @@ -21,69 +21,96 @@ # Keep text as editable text (not outlined paths) in every saved SVG/PDF/EPS so # the figures can be edited in Illustrator. Set at module import so all plotting # functions below inherit it regardless of call order. -plt.rcParams["svg.fonttype"] = "none" # SVG: write real elements -plt.rcParams["pdf.fonttype"] = 42 # PDF: embed TrueType (editable text) -plt.rcParams["ps.fonttype"] = 42 # EPS/PS: same +plt.rcParams["svg.fonttype"] = "none" # SVG: write real elements +plt.rcParams["pdf.fonttype"] = 42 # PDF: embed TrueType (editable text) +plt.rcParams["ps.fonttype"] = 42 # EPS/PS: same # Publication-quality color palette — one distinct color per plot _COLORS = { - 'primary': '#2c3e50', # dark slate for line strokes - 'stability': '#3498db', # blue - 'error': '#e74c3c', # red - 'go_terms': '#2ecc71', # green - 'genesets': '#d35400', # burnt orange - 'traits': '#9b59b6', # purple - 'perturbation': '#1abc9c', # teal - 'explained_var': '#f1c40f', # yellow + "primary": "#2c3e50", # dark slate for line strokes + "stability": "#3498db", # blue + "error": "#e74c3c", # red + "go_terms": "#2ecc71", # green + "genesets": "#d35400", # burnt orange + "traits": "#9b59b6", # purple + "perturbation": "#1abc9c", # teal + "explained_var": "#f1c40f", # yellow } # Qualitative palette for multi-line plots -_PALETTE = ['#3498db', '#e74c3c', '#2ecc71', '#9b59b6', '#f39c12', '#1abc9c', '#e67e22', '#34495e'] +_PALETTE = [ + "#3498db", + "#e74c3c", + "#2ecc71", + "#9b59b6", + "#f39c12", + "#1abc9c", + "#e67e22", + "#34495e", +] def _style_ax(ax, xlabel=None, ylabel=None, title=None, hide_top_right=True): """Apply publication-quality styling to an axis.""" if hide_top_right: - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - ax.spines['left'].set_linewidth(0.8) - ax.spines['bottom'].set_linewidth(0.8) - ax.tick_params(axis='both', which='major', labelsize=9, width=0.8, length=4, - direction='out', color='#333333') - ax.tick_params(axis='both', which='minor', width=0.5, length=2, - direction='out', color='#333333') + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + ax.spines["left"].set_linewidth(0.8) + ax.spines["bottom"].set_linewidth(0.8) + ax.tick_params( + axis="both", + which="major", + labelsize=9, + width=0.8, + length=4, + direction="out", + color="#333333", + ) + ax.tick_params( + axis="both", + which="minor", + width=0.5, + length=2, + direction="out", + color="#333333", + ) if xlabel: - ax.set_xlabel(xlabel, fontsize=11, fontweight='medium', labelpad=6) + ax.set_xlabel(xlabel, fontsize=11, fontweight="medium", labelpad=6) if ylabel: - ax.set_ylabel(ylabel, fontsize=11, fontweight='medium', labelpad=6) + ax.set_ylabel(ylabel, fontsize=11, fontweight="medium", labelpad=6) if title: - ax.set_title(title, fontsize=12, fontweight='semibold', pad=8) + ax.set_title(title, fontsize=12, fontweight="semibold", pad=8) # --------------------------------------------------------------------------- # Helpers for load_stablity_error_data # --------------------------------------------------------------------------- + def _load_stats_file(path): """Load a stability/error DataFrame from TSV or NPZ.""" - if path.endswith('.npz'): + if path.endswith(".npz"): npz = np.load(path, allow_pickle=True) - return pd.DataFrame(data=npz['data'], index=npz['index'], columns=npz['columns']) + return pd.DataFrame( + data=npz["data"], index=npz["index"], columns=npz["columns"] + ) else: - return pd.read_csv(path, sep='\t') + return pd.read_csv(path, sep="\t") def _filter_components(stats, components): """Keep only the requested K values, if present.""" stats = stats.copy() - stats['k'] = stats['k'].astype(int) - available = set(stats['k']) + stats["k"] = stats["k"].astype(int) + available = set(stats["k"]) requested = set(components) missing = requested - available if missing: - print(f" Warning: requested K values not in stats file and will be skipped: {sorted(missing)}") - return stats.loc[stats['k'].isin(requested)].sort_values('k').reset_index(drop=True) + print( + f" Warning: requested K values not in stats file and will be skipped: {sorted(missing)}" + ) + return stats.loc[stats["k"].isin(requested)].sort_values("k").reset_index(drop=True) def _cache_has_all_components(stats, components): @@ -95,35 +122,47 @@ def _cache_has_all_components(stats, components): """ if stats is None or len(stats) == 0: return False - available_k = set(stats['k'].astype(int)) + available_k = set(stats["k"].astype(int)) missing_k = set(components) - available_k if missing_k: - print(f" Cache missing K values: {sorted(missing_k)} — falling through to next tier") + print( + f" Cache missing K values: {sorted(missing_k)} — falling through to next tier" + ) return False return True def _print_stats_summary(stats): """Print min/max summary for stability and error.""" - print("min stablity is", stats['silhouette'].min(), flush=True) - print("max stablity is", stats['silhouette'].max(), flush=True) - print("min error is", stats['prediction_error'].min(), flush=True) - print("max error is", stats['prediction_error'].max(), flush=True) - - -def _compute_stability_error_manual(output_directory, run_name, components, - density_threshold=2.0, local_neighborhood_size=0.30): + print("min stablity is", stats["silhouette"].min(), flush=True) + print("max stablity is", stats["silhouette"].max(), flush=True) + print("min error is", stats["prediction_error"].min(), flush=True) + print("max error is", stats["prediction_error"].max(), flush=True) + + +def _compute_stability_error_manual( + output_directory, + run_name, + components, + density_threshold=2.0, + local_neighborhood_size=0.30, +): """Tier 5: compute stability/error from merged_spectra + usages without cnmf.consensus().""" run_dir = os.path.join(output_directory, run_name) - cnmf_tmp = os.path.join(run_dir, 'cnmf_tmp') + cnmf_tmp = os.path.join(run_dir, "cnmf_tmp") # Load normalized counts — fall back to any per-K run dir if shared path missing - norm_path = os.path.join(run_dir, run_name + '.norm_counts.h5ad') + norm_path = os.path.join(run_dir, run_name + ".norm_counts.h5ad") if not os.path.exists(norm_path): - norm_path = os.path.join(cnmf_tmp, run_name + '.norm_counts.h5ad') + norm_path = os.path.join(cnmf_tmp, run_name + ".norm_counts.h5ad") if not os.path.exists(norm_path): for k in components: - fallback = os.path.join(f'{output_directory}_K{k}', run_name, 'cnmf_tmp', run_name + '.norm_counts.h5ad') + fallback = os.path.join( + f"{output_directory}_K{k}", + run_name, + "cnmf_tmp", + run_name + ".norm_counts.h5ad", + ) if os.path.exists(fallback): norm_path = fallback print(f" Using per-K norm_counts from K={k}: {fallback}") @@ -133,41 +172,57 @@ def _compute_stability_error_manual(output_directory, run_name, components, rows = [] for k in components: # Load merged spectra — try shared run dir first, then per-K fallback - spectra_path = os.path.join(cnmf_tmp, f'{run_name}.spectra.k_{k}.merged.df.npz') + spectra_path = os.path.join(cnmf_tmp, f"{run_name}.spectra.k_{k}.merged.df.npz") if not os.path.exists(spectra_path): - per_k_spectra = os.path.join(f'{output_directory}_K{k}', run_name, 'cnmf_tmp', f'{run_name}.spectra.k_{k}.merged.df.npz') + per_k_spectra = os.path.join( + f"{output_directory}_K{k}", + run_name, + "cnmf_tmp", + f"{run_name}.spectra.k_{k}.merged.df.npz", + ) if os.path.exists(per_k_spectra): spectra_path = per_k_spectra print(f" Using per-K spectra fallback for k={k}: {per_k_spectra}") else: - print(f" Warning: merged spectra not found for k={k}, skipping: {spectra_path}") + print( + f" Warning: merged spectra not found for k={k}, skipping: {spectra_path}" + ) continue npz = np.load(spectra_path, allow_pickle=True) - merged_spectra = pd.DataFrame(data=npz['data'], index=npz['index'], columns=npz['columns']) + merged_spectra = pd.DataFrame( + data=npz["data"], index=npz["index"], columns=npz["columns"] + ) # Stability via KMeans + silhouette l2_spectra = (merged_spectra.T / np.sqrt((merged_spectra**2).sum(axis=1))).T kmeans = KMeans(n_clusters=k, n_init=10, random_state=1) kmeans.fit(l2_spectra) labels = pd.Series(kmeans.labels_ + 1, index=l2_spectra.index) - sil = silhouette_score(l2_spectra.values, labels, metric='euclidean') + sil = silhouette_score(l2_spectra.values, labels, metric="euclidean") # Median spectra (normalized to probability distributions) median_spectra = l2_spectra.groupby(labels).median() median_spectra = (median_spectra.T / median_spectra.sum(axis=1)).T # Load refit usages — try shared run dir first, then per-K fallback - dt_str = str(density_threshold).replace('.', '_') + dt_str = str(density_threshold).replace(".", "_") usage_path = os.path.join( - run_dir, f'{run_name}.usages.k_{k}.dt_{dt_str}.consensus.txt') + run_dir, f"{run_name}.usages.k_{k}.dt_{dt_str}.consensus.txt" + ) if not os.path.exists(usage_path): - per_k_usage = os.path.join(f'{output_directory}_K{k}', run_name, f'{run_name}.usages.k_{k}.dt_{dt_str}.consensus.txt') + per_k_usage = os.path.join( + f"{output_directory}_K{k}", + run_name, + f"{run_name}.usages.k_{k}.dt_{dt_str}.consensus.txt", + ) if os.path.exists(per_k_usage): usage_path = per_k_usage else: - print(f" Warning: usage file not found for k={k}, skipping: {usage_path}") + print( + f" Warning: usage file not found for k={k}, skipping: {usage_path}" + ) continue - rf_usages = pd.read_csv(usage_path, sep='\t', index_col=0) + rf_usages = pd.read_csv(usage_path, sep="\t", index_col=0) # Align dimensions median_spectra.columns = range(median_spectra.shape[1]) @@ -182,17 +237,20 @@ def _compute_stability_error_manual(output_directory, run_name, components, X = X.todense() prediction_error = float(((np.asarray(X) - pred) ** 2).sum()) - rows.append({ - 'k': k, - 'local_density_threshold': density_threshold, - 'silhouette': sil, - 'prediction_error': prediction_error, - }) + rows.append( + { + "k": k, + "local_density_threshold": density_threshold, + "silhouette": sil, + "prediction_error": prediction_error, + } + ) if not rows: raise FileNotFoundError( f"Could not compute stability/error for any K in {components}. " - f"Check that merged spectra and usage files exist under {run_dir}.") + f"Check that merged spectra and usage files exist under {run_dir}." + ) return pd.DataFrame(rows) @@ -200,8 +258,13 @@ def _compute_stability_error_manual(output_directory, run_name, components, # Main loader — 3-tier fallback # --------------------------------------------------------------------------- -def load_stablity_error_data(output_directory, run_name, components=[30, 50, 60, 80, 100, 200, 250, 300], - stability_file=None): + +def load_stablity_error_data( + output_directory, + run_name, + components=[30, 50, 60, 80, 100, 200, 250, 300], + stability_file=None, +): """Load stability and prediction-error stats with a 3-tier fallback. Tier 1: Load from file — user-supplied path, or auto-detect @@ -210,7 +273,7 @@ def load_stablity_error_data(output_directory, run_name, components=[30, 50, 60, Tier 3: Manual computation from merged_spectra + usages """ run_dir = os.path.join(output_directory, run_name) - npz_path = os.path.join(run_dir, f'{run_name}.k_selection_stats.df.npz') + npz_path = os.path.join(run_dir, f"{run_name}.k_selection_stats.df.npz") stats = None @@ -219,7 +282,9 @@ def load_stablity_error_data(output_directory, run_name, components=[30, 50, 60, # if any K is missing, drop stats and fall through to Tier 2 / Tier 3. if stability_file is not None: if os.path.exists(stability_file): - print(f"Loading stability/error data from user-supplied file: {stability_file}") + print( + f"Loading stability/error data from user-supplied file: {stability_file}" + ) cached = _load_stats_file(stability_file) if _cache_has_all_components(cached, components): stats = _filter_components(cached, components) @@ -236,16 +301,20 @@ def load_stablity_error_data(output_directory, run_name, components=[30, 50, 60, if stats is None: try: from torch_cnmf import cNMF + print("Running torch_cnmf.consensus() for each K ...") cnmf_obj = cNMF(output_dir=output_directory, name=run_name) rows = [] - norm_counts = sc.read(cnmf_obj.paths['normalized_counts']) + norm_counts = sc.read(cnmf_obj.paths["normalized_counts"]) for k in components: rows.append( cnmf_obj.consensus( - k, skip_density_and_return_after_stats=True, - show_clustering=False, close_clustergram_fig=True, - norm_counts=norm_counts, density_threshold=2.0, + k, + skip_density_and_return_after_stats=True, + show_clustering=False, + close_clustergram_fig=True, + norm_counts=norm_counts, + density_threshold=2.0, local_neighborhood_size=0.3, ).stats ) @@ -266,110 +335,179 @@ def load_stablity_error_data(output_directory, run_name, components=[30, 50, 60, def plot_stablity_error(stats, folder_name=None, file_name=None, selected_k=None): # --- Stability --- fig1, ax1 = plt.subplots(figsize=(4, 3.5)) - ax1.plot(stats['k'], stats['silhouette'], color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['stability'], markeredgecolor=_COLORS['primary'], - markeredgewidth=0.8, zorder=3) - ax1.fill_between(stats['k'], stats['silhouette'], alpha=0.08, color=_COLORS['stability']) + ax1.plot( + stats["k"], + stats["silhouette"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["stability"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax1.fill_between( + stats["k"], stats["silhouette"], alpha=0.08, color=_COLORS["stability"] + ) if selected_k is not None: - ax1.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) - _style_ax(ax1, xlabel='Number of components (k)', ylabel='Stability (silhouette)', - title='Program stability') + ax1.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) + _style_ax( + ax1, + xlabel="Number of components (k)", + ylabel="Stability (silhouette)", + title="Program stability", + ) if folder_name and file_name: fig1.savefig(f"{folder_name}/{file_name}_stability.svg", bbox_inches="tight") - fig1.savefig(f"{folder_name}/{file_name}_stability.png", dpi=300, bbox_inches="tight") + fig1.savefig( + f"{folder_name}/{file_name}_stability.png", dpi=300, bbox_inches="tight" + ) plt.show() # --- Error --- fig2, ax2 = plt.subplots(figsize=(4, 3.5)) - ax2.plot(stats['k'], stats['prediction_error'], color=_COLORS['primary'], linewidth=1.5, - marker='s', markersize=4, markerfacecolor=_COLORS['error'], markeredgecolor=_COLORS['primary'], - markeredgewidth=0.8, zorder=3) - ax2.fill_between(stats['k'], stats['prediction_error'], alpha=0.08, color=_COLORS['error']) + ax2.plot( + stats["k"], + stats["prediction_error"], + color=_COLORS["primary"], + linewidth=1.5, + marker="s", + markersize=4, + markerfacecolor=_COLORS["error"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax2.fill_between( + stats["k"], stats["prediction_error"], alpha=0.08, color=_COLORS["error"] + ) if selected_k is not None: - ax2.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) - _style_ax(ax2, xlabel='Number of components (k)', ylabel='Prediction error', - title='Reconstruction error') + ax2.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) + _style_ax( + ax2, + xlabel="Number of components (k)", + ylabel="Prediction error", + title="Reconstruction error", + ) if folder_name and file_name: fig2.savefig(f"{folder_name}/{file_name}_error.svg", bbox_inches="tight") - fig2.savefig(f"{folder_name}/{file_name}_error.png", dpi=300, bbox_inches="tight") + fig2.savefig( + f"{folder_name}/{file_name}_error.png", dpi=300, bbox_inches="tight" + ) plt.show() - - - # Load data for differeent enrichment test -def load_enrichment_data(folder, components = [30, 50, 60, 80, 100, 200, 250, 300], sel_thresh = 2.0, pval = 0.05, - go_file=None, geneset_file=None, trait_file=None, - term_col='Term', adjpval_col='Adjusted P-value'): +def load_enrichment_data( + folder, + components=[30, 50, 60, 80, 100, 200, 250, 300], + sel_thresh=2.0, + pval=0.05, + go_file=None, + geneset_file=None, + trait_file=None, + term_col="Term", + adjpval_col="Adjusted P-value", +): # loading function def load(k, term, folder, term_col, adjpval_col): # read evaluation results if os.path.getsize(folder) == 0: df = pd.DataFrame(columns=[term_col, adjpval_col]) - df['num_programs'] = k - df['test_term'] = term + df["num_programs"] = k + df["test_term"] = term return df - df = pd.read_csv(folder, sep='\t') - df = df.loc[df[adjpval_col]<=pval] - df['num_programs'] = k - df['test_term'] = term + df = pd.read_csv(folder, sep="\t") + df = df.loc[df[adjpval_col] <= pval] + df["num_programs"] = k + df["test_term"] = term # normalize the term column name for downstream counting - if term_col != 'Term': - df = df.rename(columns={term_col: 'Term'}) + if term_col != "Term": + df = df.rename(columns={term_col: "Term"}) return df # Default file name patterns (legacy convention) if go_file is None: - go_file = '{k}_GO_term_enrichment.txt' + go_file = "{k}_GO_term_enrichment.txt" if geneset_file is None: - geneset_file = '{k}_geneset_enrichment.txt' + geneset_file = "{k}_geneset_enrichment.txt" if trait_file is None: - trait_file = '{k}_trait_enrichment.txt' + trait_file = "{k}_trait_enrichment.txt" # collect all the results for each k term_df = [] for k in components: - k_folder = '{}/{}_{}'.format(folder, k, str(sel_thresh).replace('.','_')) - - term_df.append(load(k, 'go_terms', '{}/{}'.format(k_folder, go_file.format(k=k)), term_col, adjpval_col)) - term_df.append(load(k, 'genesets', '{}/{}'.format(k_folder, geneset_file.format(k=k)), term_col, adjpval_col)) - term_df.append(load(k, 'traits', '{}/{}'.format(k_folder, trait_file.format(k=k)), term_col, adjpval_col)) + k_folder = "{}/{}_{}".format(folder, k, str(sel_thresh).replace(".", "_")) + + term_df.append( + load( + k, + "go_terms", + "{}/{}".format(k_folder, go_file.format(k=k)), + term_col, + adjpval_col, + ) + ) + term_df.append( + load( + k, + "genesets", + "{}/{}".format(k_folder, geneset_file.format(k=k)), + term_col, + adjpval_col, + ) + ) + term_df.append( + load( + k, + "traits", + "{}/{}".format(k_folder, trait_file.format(k=k)), + term_col, + adjpval_col, + ) + ) term_df = pd.concat(term_df, ignore_index=True) # Count unique terms per k - count_df = pd.DataFrame(index=components, columns=term_df['test_term'].unique()) + count_df = pd.DataFrame(index=components, columns=term_df["test_term"].unique()) for k in components: for col in count_df.columns: - count_df.loc[k, col] = term_df.loc[(term_df['num_programs']==k) & (term_df['test_term']==col), 'Term'].unique().shape[0] + count_df.loc[k, col] = ( + term_df.loc[ + (term_df["num_programs"] == k) & (term_df["test_term"] == col), + "Term", + ] + .unique() + .shape[0] + ) # ensure expected columns exist even when no significant results (e.g. no GWAS hits for mouse data) - for col in ['go_terms', 'genesets', 'traits']: + for col in ["go_terms", "genesets", "traits"]: if col not in count_df.columns: count_df[col] = 0 - #print out some stats - print(f"min go_terms for {sel_thresh} is", count_df['go_terms'].min(), flush=True) - print(f"max go_terms for {sel_thresh} is", count_df['go_terms'].max(), flush=True) - - print(f"min genesets for {sel_thresh} is", count_df['genesets'].min(), flush=True) - print(f"max genesets for {sel_thresh} is", count_df['genesets'].max(), flush=True) + # print out some stats + print(f"min go_terms for {sel_thresh} is", count_df["go_terms"].min(), flush=True) + print(f"max go_terms for {sel_thresh} is", count_df["go_terms"].max(), flush=True) - print(f"min traits for {sel_thresh} is", count_df['traits'].min(), flush=True) - print(f"max traits for {sel_thresh} is", count_df['traits'].max(), flush=True) + print(f"min genesets for {sel_thresh} is", count_df["genesets"].min(), flush=True) + print(f"max genesets for {sel_thresh} is", count_df["genesets"].max(), flush=True) - #print("min motif is", count_df['motifs'].min()) - #print("max motif is", count_df['motifs'].max()) + print(f"min traits for {sel_thresh} is", count_df["traits"].min(), flush=True) + print(f"max traits for {sel_thresh} is", count_df["traits"].max(), flush=True) + # print("min motif is", count_df['motifs'].min()) + # print("max motif is", count_df['motifs'].max()) return count_df @@ -377,40 +515,65 @@ def load(k, term, folder, term_col, adjpval_col): # plot loaded df — 3 independent square figures def plot_enrichment(count_df, folder_name=None, file_name=None, selected_k=None): enrichment_config = [ - ('go_terms', 'Unique GO terms', 'GO term enrichment', _COLORS['go_terms'], 'o'), - ('genesets', 'Unique gene sets', 'Gene set enrichment', _COLORS['genesets'], 's'), - ('traits', 'Unique traits', 'Trait enrichment', _COLORS['traits'], 'D'), + ("go_terms", "Unique GO terms", "GO term enrichment", _COLORS["go_terms"], "o"), + ( + "genesets", + "Unique gene sets", + "Gene set enrichment", + _COLORS["genesets"], + "s", + ), + ("traits", "Unique traits", "Trait enrichment", _COLORS["traits"], "D"), ] for col, ylabel, title, color, marker in enrichment_config: fig, ax = plt.subplots(figsize=(4, 3.5)) vals = count_df[col].astype(float) - ax.plot(count_df.index, vals, color=_COLORS['primary'], linewidth=1.5, - marker=marker, markersize=4, markerfacecolor=color, - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) + ax.plot( + count_df.index, + vals, + color=_COLORS["primary"], + linewidth=1.5, + marker=marker, + markersize=4, + markerfacecolor=color, + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) ax.fill_between(count_df.index, vals, alpha=0.08, color=color) if selected_k is not None: - ax.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) - _style_ax(ax, xlabel='Number of components (k)', ylabel=ylabel, title=title) + ax.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) + _style_ax(ax, xlabel="Number of components (k)", ylabel=ylabel, title=title) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) if folder_name and file_name: fig.savefig(f"{folder_name}/{file_name}_{col}.svg", bbox_inches="tight") - fig.savefig(f"{folder_name}/{file_name}_{col}.png", dpi=300, bbox_inches="tight") + fig.savefig( + f"{folder_name}/{file_name}_{col}.png", dpi=300, bbox_inches="tight" + ) plt.show() # load perturbation data -def load_perturbation_data(folder, pval = 0.000335, components = [30, 50, 60, 80, 100, 200, 250, 300], sel_thresh = 2.0, - conditions = ['D0', 'sample_D1', 'sample_D2', 'sample_D3'], perturbation_file=None, - perturb_adjpval_col='adj_pval', perturb_target_col='target_name', perturb_log2fc_col='log2FC'): +def load_perturbation_data( + folder, + pval=0.000335, + components=[30, 50, 60, 80, 100, 200, 250, 300], + sel_thresh=2.0, + conditions=["D0", "sample_D1", "sample_D2", "sample_D3"], + perturbation_file=None, + perturb_adjpval_col="adj_pval", + perturb_target_col="target_name", + perturb_log2fc_col="log2FC", +): # Default file name pattern (legacy convention): {k}_perturbation_association_results_{sample}.txt # Use {k} and {sample} as placeholders in custom patterns. if perturbation_file is None: - perturbation_file = '{k}_perturbation_association_results_{sample}.txt' + perturbation_file = "{k}_perturbation_association_results_{sample}.txt" # Compute no. of unique regulators test_stats_df = [] @@ -418,29 +581,38 @@ def load_perturbation_data(folder, pval = 0.000335, components = [30, 50, 60, 80 for k in components: # Run perturbation assocation for cond in conditions: - file_path = '{}/{}_{}/{}'.format( - folder, k, str(sel_thresh).replace('.','_'), - perturbation_file.format(k=k, sample=cond)) - test_stats_df_ = pd.read_csv(file_path, sep='\t') - test_stats_df_['condition'] = cond - test_stats_df_['K'] = k + file_path = "{}/{}_{}/{}".format( + folder, + k, + str(sel_thresh).replace(".", "_"), + perturbation_file.format(k=k, sample=cond), + ) + test_stats_df_ = pd.read_csv(file_path, sep="\t") + test_stats_df_["condition"] = cond + test_stats_df_["K"] = k test_stats_df.append(test_stats_df_) test_stats_df = pd.concat(test_stats_df, ignore_index=True) # Normalize column names so downstream functions work with standard names rename_map = {} - if perturb_adjpval_col != 'adj_pval': - rename_map[perturb_adjpval_col] = 'adj_pval' - if perturb_target_col != 'target_name': - rename_map[perturb_target_col] = 'target_name' - if perturb_log2fc_col != 'log2FC': - rename_map[perturb_log2fc_col] = 'log2FC' + if perturb_adjpval_col != "adj_pval": + rename_map[perturb_adjpval_col] = "adj_pval" + if perturb_target_col != "target_name": + rename_map[perturb_target_col] = "target_name" + if perturb_log2fc_col != "log2FC": + rename_map[perturb_log2fc_col] = "log2FC" if rename_map: test_stats_df = test_stats_df.rename(columns=rename_map) # pring some stats - plotting_df = test_stats_df.loc[test_stats_df.adj_pval < pval, ['K','target_name']].drop_duplicates().groupby(['K']).count().reset_index() + plotting_df = ( + test_stats_df.loc[test_stats_df.adj_pval < pval, ["K", "target_name"]] + .drop_duplicates() + .groupby(["K"]) + .count() + .reset_index() + ) print("min regulators is", plotting_df["target_name"].min(), flush=True) print("max regulators is", plotting_df["target_name"].max(), flush=True) @@ -449,84 +621,143 @@ def load_perturbation_data(folder, pval = 0.000335, components = [30, 50, 60, 80 # plot perturbation data — 2 independent square figures -def plot_perturbation(test_stats_df, pval=0.000335, folder_name=None, file_name=None, selected_k=None): - pval_str = f'{pval:.1e}' if pval < 0.01 else str(pval) +def plot_perturbation( + test_stats_df, pval=0.000335, folder_name=None, file_name=None, selected_k=None +): + pval_str = f"{pval:.1e}" if pval < 0.01 else str(pval) # --- Per-sample plot --- fig1, ax1 = plt.subplots(figsize=(4, 3.5)) - plotting_df_sample = (test_stats_df - .loc[test_stats_df.adj_pval <= pval, ['K', 'condition', 'target_name']] + plotting_df_sample = ( + test_stats_df.loc[ + test_stats_df.adj_pval <= pval, ["K", "condition", "target_name"] + ] .drop_duplicates() - .groupby(['K', 'condition']).count().reset_index()) - sns.lineplot(x='K', y='target_name', hue='condition', data=plotting_df_sample, - palette=_PALETTE[:plotting_df_sample['condition'].nunique()], linewidth=1.5, marker='o', markersize=4, - ax=ax1, legend='brief') + .groupby(["K", "condition"]) + .count() + .reset_index() + ) + sns.lineplot( + x="K", + y="target_name", + hue="condition", + data=plotting_df_sample, + palette=_PALETTE[: plotting_df_sample["condition"].nunique()], + linewidth=1.5, + marker="o", + markersize=4, + ax=ax1, + legend="brief", + ) if selected_k is not None: - ax1.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) - _style_ax(ax1, xlabel='Number of components (k)', ylabel='No. unique regulators', - title=f'Unique regulators per condition') # (adj. p-value \u2264 {pval_str})') - ax1.legend(title='Condition', fontsize=8, title_fontsize=9, - frameon=True, fancybox=False, edgecolor='#cccccc', - bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0) + ax1.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) + _style_ax( + ax1, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title=f"Unique regulators per condition", + ) # (adj. p-value \u2264 {pval_str})') + ax1.legend( + title="Condition", + fontsize=8, + title_fontsize=9, + frameon=True, + fancybox=False, + edgecolor="#cccccc", + bbox_to_anchor=(1.02, 1), + loc="upper left", + borderaxespad=0, + ) ax1.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) if folder_name and file_name: - fig1.savefig(f"{folder_name}/{file_name}_per_condition.svg", bbox_inches="tight") - fig1.savefig(f"{folder_name}/{file_name}_per_condition.png", dpi=300, bbox_inches="tight") + fig1.savefig( + f"{folder_name}/{file_name}_per_condition.svg", bbox_inches="tight" + ) + fig1.savefig( + f"{folder_name}/{file_name}_per_condition.png", dpi=300, bbox_inches="tight" + ) plt.show() # --- Aggregated plot --- fig2, ax2 = plt.subplots(figsize=(4, 3.5)) - plotting_df = (test_stats_df - .loc[test_stats_df.adj_pval <= pval, ['K', 'target_name']] + plotting_df = ( + test_stats_df.loc[test_stats_df.adj_pval <= pval, ["K", "target_name"]] .drop_duplicates() - .groupby(['K']).count().reset_index()) - ax2.plot(plotting_df['K'], plotting_df['target_name'], - color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['perturbation'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax2.fill_between(plotting_df['K'], plotting_df['target_name'], - alpha=0.08, color=_COLORS['perturbation']) + .groupby(["K"]) + .count() + .reset_index() + ) + ax2.plot( + plotting_df["K"], + plotting_df["target_name"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["perturbation"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax2.fill_between( + plotting_df["K"], + plotting_df["target_name"], + alpha=0.08, + color=_COLORS["perturbation"], + ) if selected_k is not None: - ax2.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) - _style_ax(ax2, xlabel='Number of components (k)', ylabel='No. unique regulators', - title=f'Unique regulators for all conditions') #(adj. p-value \u2264 {pval_str})') + ax2.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) + _style_ax( + ax2, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title=f"Unique regulators for all conditions", + ) # (adj. p-value \u2264 {pval_str})') ax2.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) if folder_name and file_name: - fig2.savefig(f"{folder_name}/{file_name}_all_conditions.svg", bbox_inches="tight") - fig2.savefig(f"{folder_name}/{file_name}_all_conditions.png", dpi=300, bbox_inches="tight") + fig2.savefig( + f"{folder_name}/{file_name}_all_conditions.svg", bbox_inches="tight" + ) + fig2.savefig( + f"{folder_name}/{file_name}_all_conditions.png", + dpi=300, + bbox_inches="tight", + ) plt.show() return plotting_df - # load total explained variance -def load_explained_variance_data(folder, components = [30, 50, 60, 80, 100, 200, 250, 300], sel_thresh = 2.0, - variance_file=None, variance_col='Total'): +def load_explained_variance_data( + folder, + components=[30, 50, 60, 80, 100, 200, 250, 300], + sel_thresh=2.0, + variance_file=None, + variance_col="Total", +): # Default file name pattern if variance_file is None: - variance_file = '{k}_Explained_Variance_Summary.txt' + variance_file = "{k}_Explained_Variance_Summary.txt" stats = {} for k in components: + input_path = f"{folder}/{k}_{str(sel_thresh).replace('.', '_')}/{variance_file.format(k=k)}" + df = pd.read_csv(input_path, sep="\t", index_col=0) - input_path = f"{folder}/{k}_{str(sel_thresh).replace('.','_')}/{variance_file.format(k=k)}" - df = pd.read_csv(input_path, sep = '\t', index_col = 0) - - if variance_col == 'Total': - stats[k] = df['Total'].values[0] + if variance_col == "Total": + stats[k] = df["Total"].values[0] else: # For per-program variance (e.g. Morphic: variance_explained_ratio), sum all programs stats[k] = df[variance_col].sum() - print("min Explained_variance is", min(stats.values()), flush=True) print("max Explained_variance is", max(stats.values()), flush=True) @@ -540,14 +771,27 @@ def plot_explained_variance(stats, folder_name=None, file_name=None, selected_k= fig, ax = plt.subplots(figsize=(4, 3.5)) - ax.plot(ks, vals, color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['explained_var'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(ks, vals, alpha=0.08, color=_COLORS['explained_var']) + ax.plot( + ks, + vals, + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["explained_var"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between(ks, vals, alpha=0.08, color=_COLORS["explained_var"]) if selected_k is not None: - ax.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Total explained variance', - title='Explained variance') + ax.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Total explained variance", + title="Explained variance", + ) if folder_name and file_name: fig.savefig(f"{folder_name}/{file_name}.svg", bbox_inches="tight") @@ -556,110 +800,249 @@ def plot_explained_variance(stats, folder_name=None, file_name=None, selected_k= plt.show() - # Combined panel figure: 3 rows x 3 columns # Row 1: Stability, Error, Explained variance # Row 2: GO terms, Gene sets, Traits # Row 3: Regulators (all conditions), Regulators (per condition) -def plot_k_selection_panel(stability_stats, count_df, test_stats_df, explained_var_stats, - pval=0.05, folder_name=None, file_name=None, selected_k=None): +def plot_k_selection_panel( + stability_stats, + count_df, + test_stats_df, + explained_var_stats, + pval=0.05, + folder_name=None, + file_name=None, + selected_k=None, +): fig, axes = plt.subplots(3, 3, figsize=(13, 10.5)) fig.subplots_adjust(hspace=0.45, wspace=0.35) - pval_str = f'{pval:.1e}' if pval < 0.01 else str(pval) + pval_str = f"{pval:.1e}" if pval < 0.01 else str(pval) def _add_vline(ax): if selected_k is not None: - ax.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) + ax.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) # --- Row 1, Col 0: Stability --- ax = axes[0, 0] - ax.plot(stability_stats['k'], stability_stats['silhouette'], color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['stability'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(stability_stats['k'], stability_stats['silhouette'], alpha=0.08, color=_COLORS['stability']) + ax.plot( + stability_stats["k"], + stability_stats["silhouette"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["stability"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + stability_stats["k"], + stability_stats["silhouette"], + alpha=0.08, + color=_COLORS["stability"], + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Stability (silhouette)', title='Program stability') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Stability (silhouette)", + title="Program stability", + ) # --- Row 1, Col 1: Error --- ax = axes[0, 1] - ax.plot(stability_stats['k'], stability_stats['prediction_error'], color=_COLORS['primary'], linewidth=1.5, - marker='s', markersize=4, markerfacecolor=_COLORS['error'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(stability_stats['k'], stability_stats['prediction_error'], alpha=0.08, color=_COLORS['error']) + ax.plot( + stability_stats["k"], + stability_stats["prediction_error"], + color=_COLORS["primary"], + linewidth=1.5, + marker="s", + markersize=4, + markerfacecolor=_COLORS["error"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + stability_stats["k"], + stability_stats["prediction_error"], + alpha=0.08, + color=_COLORS["error"], + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Prediction error', title='Reconstruction error') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Prediction error", + title="Reconstruction error", + ) # --- Row 1, Col 2: Explained variance --- ax = axes[0, 2] ev_ks = list(explained_var_stats.keys()) ev_vals = list(explained_var_stats.values()) - ax.plot(ev_ks, ev_vals, color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['explained_var'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(ev_ks, ev_vals, alpha=0.08, color=_COLORS['explained_var']) + ax.plot( + ev_ks, + ev_vals, + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["explained_var"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between(ev_ks, ev_vals, alpha=0.08, color=_COLORS["explained_var"]) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Total explained variance', title='Explained variance') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Total explained variance", + title="Explained variance", + ) # --- Row 2: Enrichment (GO, genesets, traits) --- enrichment_config = [ - (0, 'go_terms', 'Unique GO terms', 'GO term enrichment', _COLORS['go_terms'], 'o'), - (1, 'genesets', 'Unique gene sets', 'Gene set enrichment', _COLORS['genesets'], 's'), - (2, 'traits', 'Unique traits', 'Trait enrichment', _COLORS['traits'], 'D'), + ( + 0, + "go_terms", + "Unique GO terms", + "GO term enrichment", + _COLORS["go_terms"], + "o", + ), + ( + 1, + "genesets", + "Unique gene sets", + "Gene set enrichment", + _COLORS["genesets"], + "s", + ), + (2, "traits", "Unique traits", "Trait enrichment", _COLORS["traits"], "D"), ] for col_idx, col, ylabel, title, color, marker in enrichment_config: ax = axes[1, col_idx] vals = count_df[col].astype(float) - ax.plot(count_df.index, vals, color=_COLORS['primary'], linewidth=1.5, - marker=marker, markersize=4, markerfacecolor=color, - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) + ax.plot( + count_df.index, + vals, + color=_COLORS["primary"], + linewidth=1.5, + marker=marker, + markersize=4, + markerfacecolor=color, + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) ax.fill_between(count_df.index, vals, alpha=0.08, color=color) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel=ylabel, title=title) + _style_ax(ax, xlabel="Number of components (k)", ylabel=ylabel, title=title) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) # --- Row 3, Col 0: Regulators (all conditions) --- ax = axes[2, 0] - plotting_df = (test_stats_df - .loc[test_stats_df.adj_pval <= pval, ['K', 'target_name']] + plotting_df = ( + test_stats_df.loc[test_stats_df.adj_pval <= pval, ["K", "target_name"]] .drop_duplicates() - .groupby(['K']).count().reset_index()) - ax.plot(plotting_df['K'], plotting_df['target_name'], color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['perturbation'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(plotting_df['K'], plotting_df['target_name'], alpha=0.08, color=_COLORS['perturbation']) + .groupby(["K"]) + .count() + .reset_index() + ) + ax.plot( + plotting_df["K"], + plotting_df["target_name"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["perturbation"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + plotting_df["K"], + plotting_df["target_name"], + alpha=0.08, + color=_COLORS["perturbation"], + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='No. unique regulators', - title='Unique regulators for all conditions') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title="Unique regulators for all conditions", + ) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) # --- Row 3, Col 1: Regulators (per sample) --- ax = axes[2, 1] - plotting_df_sample = (test_stats_df - .loc[test_stats_df.adj_pval <= pval, ['K', 'condition', 'target_name']] + plotting_df_sample = ( + test_stats_df.loc[ + test_stats_df.adj_pval <= pval, ["K", "condition", "target_name"] + ] .drop_duplicates() - .groupby(['K', 'condition']).count().reset_index()) - sns.lineplot(x='K', y='target_name', hue='condition', data=plotting_df_sample, - palette=_PALETTE[:plotting_df_sample['condition'].nunique()], linewidth=1.5, marker='o', markersize=4, - ax=ax, legend='brief') + .groupby(["K", "condition"]) + .count() + .reset_index() + ) + sns.lineplot( + x="K", + y="target_name", + hue="condition", + data=plotting_df_sample, + palette=_PALETTE[: plotting_df_sample["condition"].nunique()], + linewidth=1.5, + marker="o", + markersize=4, + ax=ax, + legend="brief", + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='No. unique regulators', - title='Unique regulators per condition') - ax.legend(title='Condition', fontsize=7, title_fontsize=8, - frameon=True, fancybox=False, edgecolor='#cccccc', - bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0) + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title="Unique regulators per condition", + ) + ax.legend( + title="Condition", + fontsize=7, + title_fontsize=8, + frameon=True, + fancybox=False, + edgecolor="#cccccc", + bbox_to_anchor=(1.02, 1), + loc="upper left", + borderaxespad=0, + ) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) # --- Row 3, Col 2: hide empty panel --- axes[2, 2].set_visible(False) # Add panel labels (A, B, C, ...) - panel_labels = 'ABCDEFGH' - visible_axes = [axes[r, c] for r in range(3) for c in range(3) if axes[r, c].get_visible()] + panel_labels = "ABCDEFGH" + visible_axes = [ + axes[r, c] for r in range(3) for c in range(3) if axes[r, c].get_visible() + ] for label, ax in zip(panel_labels, visible_axes): - ax.text(-0.15, 1.12, label, transform=ax.transAxes, - fontsize=14, fontweight='bold', va='top', ha='left') + ax.text( + -0.15, + 1.12, + label, + transform=ax.transAxes, + fontsize=14, + fontweight="bold", + va="top", + ha="left", + ) if folder_name and file_name: fig.savefig(f"{folder_name}/{file_name}.svg", bbox_inches="tight") @@ -672,61 +1055,144 @@ def _add_vline(ax): # Row 1: Stability, Error, Explained variance # Row 2: GO terms, Gene sets # Row 3: Regulators (all conditions), Regulators (per condition) -def plot_k_selection_panel_no_traits(stability_stats, count_df, test_stats_df, explained_var_stats, - pval=0.05, folder_name=None, file_name=None, selected_k=None): +def plot_k_selection_panel_no_traits( + stability_stats, + count_df, + test_stats_df, + explained_var_stats, + pval=0.05, + folder_name=None, + file_name=None, + selected_k=None, +): fig, axes = plt.subplots(3, 3, figsize=(13, 10.5)) fig.subplots_adjust(hspace=0.45, wspace=0.35) - pval_str = f'{pval:.1e}' if pval < 0.01 else str(pval) + pval_str = f"{pval:.1e}" if pval < 0.01 else str(pval) def _add_vline(ax): if selected_k is not None: - ax.axvline(x=selected_k, color='red', linestyle='--', linewidth=1, zorder=2) + ax.axvline(x=selected_k, color="red", linestyle="--", linewidth=1, zorder=2) # --- Row 1, Col 0: Stability --- ax = axes[0, 0] - ax.plot(stability_stats['k'], stability_stats['silhouette'], color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['stability'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(stability_stats['k'], stability_stats['silhouette'], alpha=0.08, color=_COLORS['stability']) + ax.plot( + stability_stats["k"], + stability_stats["silhouette"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["stability"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + stability_stats["k"], + stability_stats["silhouette"], + alpha=0.08, + color=_COLORS["stability"], + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Stability (silhouette)', title='Program stability') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Stability (silhouette)", + title="Program stability", + ) # --- Row 1, Col 1: Error --- ax = axes[0, 1] - ax.plot(stability_stats['k'], stability_stats['prediction_error'], color=_COLORS['primary'], linewidth=1.5, - marker='s', markersize=4, markerfacecolor=_COLORS['error'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(stability_stats['k'], stability_stats['prediction_error'], alpha=0.08, color=_COLORS['error']) + ax.plot( + stability_stats["k"], + stability_stats["prediction_error"], + color=_COLORS["primary"], + linewidth=1.5, + marker="s", + markersize=4, + markerfacecolor=_COLORS["error"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + stability_stats["k"], + stability_stats["prediction_error"], + alpha=0.08, + color=_COLORS["error"], + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Prediction error', title='Reconstruction error') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Prediction error", + title="Reconstruction error", + ) # --- Row 1, Col 2: Explained variance --- ax = axes[0, 2] ev_ks = list(explained_var_stats.keys()) ev_vals = list(explained_var_stats.values()) - ax.plot(ev_ks, ev_vals, color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['explained_var'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(ev_ks, ev_vals, alpha=0.08, color=_COLORS['explained_var']) + ax.plot( + ev_ks, + ev_vals, + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["explained_var"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between(ev_ks, ev_vals, alpha=0.08, color=_COLORS["explained_var"]) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='Total explained variance', title='Explained variance') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="Total explained variance", + title="Explained variance", + ) # --- Row 2: Enrichment (GO, genesets) — no traits --- enrichment_config = [ - (0, 'go_terms', 'Unique GO terms', 'GO term enrichment', _COLORS['go_terms'], 'o'), - (1, 'genesets', 'Unique gene sets', 'Gene set enrichment', _COLORS['genesets'], 's'), + ( + 0, + "go_terms", + "Unique GO terms", + "GO term enrichment", + _COLORS["go_terms"], + "o", + ), + ( + 1, + "genesets", + "Unique gene sets", + "Gene set enrichment", + _COLORS["genesets"], + "s", + ), ] for col_idx, col, ylabel, title, color, marker in enrichment_config: ax = axes[1, col_idx] vals = count_df[col].astype(float) - ax.plot(count_df.index, vals, color=_COLORS['primary'], linewidth=1.5, - marker=marker, markersize=4, markerfacecolor=color, - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) + ax.plot( + count_df.index, + vals, + color=_COLORS["primary"], + linewidth=1.5, + marker=marker, + markersize=4, + markerfacecolor=color, + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) ax.fill_between(count_df.index, vals, alpha=0.08, color=color) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel=ylabel, title=title) + _style_ax(ax, xlabel="Number of components (k)", ylabel=ylabel, title=title) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) # --- Row 2, Col 2: hide empty panel --- @@ -734,45 +1200,102 @@ def _add_vline(ax): # --- Row 3, Col 0: Regulators (all conditions) --- ax = axes[2, 0] - plotting_df = (test_stats_df - .loc[test_stats_df.adj_pval <= pval, ['K', 'target_name']] + plotting_df = ( + test_stats_df.loc[test_stats_df.adj_pval <= pval, ["K", "target_name"]] .drop_duplicates() - .groupby(['K']).count().reset_index()) - ax.plot(plotting_df['K'], plotting_df['target_name'], color=_COLORS['primary'], linewidth=1.5, - marker='o', markersize=4, markerfacecolor=_COLORS['perturbation'], - markeredgecolor=_COLORS['primary'], markeredgewidth=0.8, zorder=3) - ax.fill_between(plotting_df['K'], plotting_df['target_name'], alpha=0.08, color=_COLORS['perturbation']) + .groupby(["K"]) + .count() + .reset_index() + ) + ax.plot( + plotting_df["K"], + plotting_df["target_name"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["perturbation"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + plotting_df["K"], + plotting_df["target_name"], + alpha=0.08, + color=_COLORS["perturbation"], + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='No. unique regulators', - title='Unique regulators for all conditions') + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title="Unique regulators for all conditions", + ) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) # --- Row 3, Col 1: Regulators (per sample) --- ax = axes[2, 1] - plotting_df_sample = (test_stats_df - .loc[test_stats_df.adj_pval <= pval, ['K', 'condition', 'target_name']] + plotting_df_sample = ( + test_stats_df.loc[ + test_stats_df.adj_pval <= pval, ["K", "condition", "target_name"] + ] .drop_duplicates() - .groupby(['K', 'condition']).count().reset_index()) - sns.lineplot(x='K', y='target_name', hue='condition', data=plotting_df_sample, - palette=_PALETTE[:plotting_df_sample['condition'].nunique()], linewidth=1.5, marker='o', markersize=4, - ax=ax, legend='brief') + .groupby(["K", "condition"]) + .count() + .reset_index() + ) + sns.lineplot( + x="K", + y="target_name", + hue="condition", + data=plotting_df_sample, + palette=_PALETTE[: plotting_df_sample["condition"].nunique()], + linewidth=1.5, + marker="o", + markersize=4, + ax=ax, + legend="brief", + ) _add_vline(ax) - _style_ax(ax, xlabel='Number of components (k)', ylabel='No. unique regulators', - title='Unique regulators per condition') - ax.legend(title='Condition', fontsize=7, title_fontsize=8, - frameon=True, fancybox=False, edgecolor='#cccccc', - bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0) + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title="Unique regulators per condition", + ) + ax.legend( + title="Condition", + fontsize=7, + title_fontsize=8, + frameon=True, + fancybox=False, + edgecolor="#cccccc", + bbox_to_anchor=(1.02, 1), + loc="upper left", + borderaxespad=0, + ) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) # --- Row 3, Col 2: hide empty panel --- axes[2, 2].set_visible(False) # Add panel labels (A, B, C, ...) - panel_labels = 'ABCDEFG' - visible_axes = [axes[r, c] for r in range(3) for c in range(3) if axes[r, c].get_visible()] + panel_labels = "ABCDEFG" + visible_axes = [ + axes[r, c] for r in range(3) for c in range(3) if axes[r, c].get_visible() + ] for label, ax in zip(panel_labels, visible_axes): - ax.text(-0.15, 1.12, label, transform=ax.transAxes, - fontsize=14, fontweight='bold', va='top', ha='left') + ax.text( + -0.15, + 1.12, + label, + transform=ax.transAxes, + fontsize=14, + fontweight="bold", + va="top", + ha="left", + ) if folder_name and file_name: fig.savefig(f"{folder_name}/{file_name}.svg", bbox_inches="tight") @@ -782,86 +1305,110 @@ def _add_vline(ax): # --- per-(K, sel_thresh) program dotplot (moved here from k_quality_plots.py) --- -def programs_dotplots(k, output_dir, run_name, sel_thresh = 2.0, groupby='sample', save_name=None, save_path=None, -figsize=(4, 30), show=False, ax=None): +def programs_dotplots( + k, + output_dir, + run_name, + sel_thresh=2.0, + groupby="sample", + save_name=None, + save_path=None, + figsize=(4, 30), + show=False, + ax=None, +): standalone_mode = ax is None - def get_gene_path(output_dir, run_name, k, sel_thresh): """Helper to build path consistently. Tries standard PerturbNMF layout, then flat shared-run, then per-K run layout.""" sel_str = str(sel_thresh).replace(".", "_") - standard = f'{output_dir}/{run_name}/Inference/adata/cNMF_{k}_{sel_str}.h5mu' + standard = f"{output_dir}/{run_name}/Inference/adata/cNMF_{k}_{sel_str}.h5mu" if os.path.isfile(standard): return standard - flat = f'{output_dir}/{run_name}/adata/cNMF_{k}_{sel_str}.h5mu' + flat = f"{output_dir}/{run_name}/adata/cNMF_{k}_{sel_str}.h5mu" if os.path.isfile(flat): return flat - per_k = f'{output_dir}/{run_name}_K{k}/Inference/adata/cNMF_{k}_{sel_str}.h5mu' + per_k = f"{output_dir}/{run_name}_K{k}/Inference/adata/cNMF_{k}_{sel_str}.h5mu" if os.path.isfile(per_k): return per_k return standard - # read in adata path = get_gene_path(output_dir, run_name, k, sel_thresh) if not os.path.isfile(path): - print(f"[programs_dotplots] WARNING: h5mu not found for K={k}, sel_thresh={sel_thresh}; " - f"tried standard / flat / per-K layouts. Skipping this combination.") + print( + f"[programs_dotplots] WARNING: h5mu not found for K={k}, sel_thresh={sel_thresh}; " + f"tried standard / flat / per-K layouts. Skipping this combination." + ) return None mdata = mu.read_h5mu(path) - adata_new = mdata['cNMF'].copy() - + adata_new = mdata["cNMF"].copy() + if save_name is None: save_name = "Program Loadings by Days" - + # Create the dotplot if ax is None: # Standalone mode - let scanpy create its own figure - + grogram_list = adata_new.var_names.tolist() - dp = sc.pl.dotplot(adata_new, grogram_list, groupby=groupby, - figsize=figsize, swap_axes=True, dendrogram=False, - show=False, return_fig=True) + dp = sc.pl.dotplot( + adata_new, + grogram_list, + groupby=groupby, + figsize=figsize, + swap_axes=True, + dendrogram=False, + show=False, + return_fig=True, + ) dp.make_figure() fig = dp.fig - ax = dp.ax_dict['mainplot_ax'] + ax = dp.ax_dict["mainplot_ax"] else: # Gridspec mode - use provided ax fig = ax.get_figure() grogram_list = adata_new.var_names.tolist() - dp = sc.pl.dotplot(adata_new, grogram_list, groupby=groupby, - swap_axes=True, dendrogram=False, show=False, - return_fig=True, ax=ax) + dp = sc.pl.dotplot( + adata_new, + grogram_list, + groupby=groupby, + swap_axes=True, + dendrogram=False, + show=False, + return_fig=True, + ax=ax, + ) dp.make_figure() - - ax.set_title(save_name, fontsize=14, fontweight='bold', loc='center') - ax.set_ylabel('Program', fontsize=10, fontweight='bold', loc='center') - ax.set_xlabel(groupby, fontsize=10, fontweight='bold', loc='center') - + + ax.set_title(save_name, fontsize=14, fontweight="bold", loc="center") + ax.set_ylabel("Program", fontsize=10, fontweight="bold", loc="center") + ax.set_xlabel(groupby, fontsize=10, fontweight="bold", loc="center") + # Get labels and set ticks properly - label = list(mdata['rna'].obs[groupby].cat.categories) # Use categories instead - + label = list(mdata["rna"].obs[groupby].cat.categories) # Use categories instead + # Set both ticks and labels together tick_positions = range(len(label)) ax.set_xticks(tick_positions) ax.set_xticklabels(label, fontsize=8) - # save fig (only in standalone mode) if save_name and save_path and standalone_mode: - fig.savefig(f"{save_path}/{save_name}.png", format='png', bbox_inches='tight', dpi=300) # Changed to png - + fig.savefig( + f"{save_path}/{save_name}.png", format="png", bbox_inches="tight", dpi=300 + ) # Changed to png + # Control whether to display the plot (only in standalone mode) if ax is None: if show: plt.show() else: plt.close(fig) - - return ax + return ax # Plot clustermap given paths to program by gene matrix diff --git a/src/Stage3_Interpretation/A_Plotting/src/utilities.py b/src/Stage3_Interpretation/A_Plotting/src/utilities.py index bf9b0a3..9fa8bf3 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/utilities.py +++ b/src/Stage3_Interpretation/A_Plotting/src/utilities.py @@ -1,4 +1,4 @@ -import muon as mu +import muon as mu import scanpy as sc import numpy as np import pandas as pd @@ -14,100 +14,100 @@ import subprocess from PyPDF2 import PdfMerger import glob -import re -from svglib.svglib import svg2rlg +import re +from svglib.svglib import svg2rlg from reportlab.graphics import renderPDF + # replace EnsemblID by gene name given the dataframe with EnsemblID as the index -def convert_with_mygene(dataframe, species='human', index = True): +def convert_with_mygene(dataframe, species="human", index=True): mg = mygene.MyGeneInfo() - + if index: # Query multiple genes at once - result = mg.querymany(dataframe.index, - scopes='ensembl.gene', - fields='symbol,name', - species='human') + result = mg.querymany( + dataframe.index, + scopes="ensembl.gene", + fields="symbol,name", + species="human", + ) else: - # Query multiple genes at once - result = mg.querymany(dataframe.columns, - scopes='ensembl.gene', - fields='symbol,name', - species='human') + # Query multiple genes at once + result = mg.querymany( + dataframe.columns, + scopes="ensembl.gene", + fields="symbol,name", + species="human", + ) - # Create mapping dictionary mapping = {} for item in result: - if 'symbol' in item and 'query' in item: - mapping[item['query']] = item['symbol'] + if "symbol" in item and "query" in item: + mapping[item["query"]] = item["symbol"] - elif 'query' in item: - mapping[item['query']] = item['query'] # Keep original if no symbol + elif "query" in item: + mapping[item["query"]] = item["query"] # Keep original if no symbol - if index: + if index: new_dataframe = dataframe.rename(index=mapping) else: new_dataframe = dataframe.rename(columns=mapping) - - return new_dataframe + return new_dataframe # same, input is adata -def convert_adata_with_mygene(adata, species='human'): +def convert_adata_with_mygene(adata, species="human"): mg = mygene.MyGeneInfo() gene_list = adata.var_names.tolist() - annotations = mg.querymany(gene_list, scopes='ensembl.gene', - fields='symbol', species='human') + annotations = mg.querymany( + gene_list, scopes="ensembl.gene", fields="symbol", species="human" + ) # Process the results to create mapping gene_dict = {} for item in annotations: - if 'symbol' in item and 'query' in item: - gene_dict[item['query']] = item['symbol'] + if "symbol" in item and "query" in item: + gene_dict[item["query"]] = item["symbol"] - elif 'query' in item: - gene_dict[item['query']] = item['query'] # Keep original if no symbol + elif "query" in item: + gene_dict[item["query"]] = item["query"] # Keep original if no symbol adata_new = adata.copy() - - adata_new.var['gene_name'] = [gene_dict.get(x, x) for x in adata_new.var_names] - adata_new.var_names = adata_new.var['gene_name'] + adata_new.var["gene_name"] = [gene_dict.get(x, x) for x in adata_new.var_names] + adata_new.var_names = adata_new.var["gene_name"] return adata_new - -# given a tsv dictionary, convert EnsemblID to gene name for adata +# given a tsv dictionary, convert EnsemblID to gene name for adata def rename_adata_gene_dictionary(adata, dictionary_file_path): adata_new = adata.copy() # Convert mapping result to list before assignment - df = pd.read_csv(dictionary_file_path, sep='\t', low_memory=False) - ensemble_to_gene = dict(zip(df['ensembl_id'], df['gene'])) + df = pd.read_csv(dictionary_file_path, sep="\t", low_memory=False) + ensemble_to_gene = dict(zip(df["ensembl_id"], df["gene"])) new_names = [ensemble_to_gene.get(x, x) for x in adata.var_names] - - adata_new.var_names = (new_names) - return adata_new + adata_new.var_names = new_names + return adata_new # given a tsv dictionary, convert EnsemblID to gene name for list def rename_list_gene_dictionary(list_input, dictionary_file_path): # Convert mapping result to list before assignment - df = pd.read_csv(dictionary_file_path, sep='\t', low_memory=False) - ensemble_to_gene = dict(zip(df['ensembl_id'], df['gene'])) + df = pd.read_csv(dictionary_file_path, sep="\t", low_memory=False) + ensemble_to_gene = dict(zip(df["ensembl_id"], df["gene"])) new_names = [ensemble_to_gene.get(x, x) for x in list_input] - - return new_names + return new_names # read cNMF programs @@ -118,13 +118,10 @@ def read_npz(path): # Reconstruct the DataFrame df = pd.DataFrame( - data=npz_data['data'], - index=npz_data['index'], - columns=npz_data['columns'] + data=npz_data["data"], index=npz_data["index"], columns=npz_data["columns"] ) - - return df + return df # Filter out the merged output itself and any 0-byte PDFs (silent failures from upstream workers). @@ -139,7 +136,9 @@ def _filter_mergeable_pdfs(pdf_files, output_path): continue kept.append(p) if skipped_empty: - print(f"Skipping {len(skipped_empty)} empty (0-byte) PDF(s): {skipped_empty[:5]}{'...' if len(skipped_empty) > 5 else ''}") + print( + f"Skipping {len(skipped_empty)} empty (0-byte) PDF(s): {skipped_empty[:5]}{'...' if len(skipped_empty) > 5 else ''}" + ) return kept @@ -166,7 +165,7 @@ def _merge_pdfs(pdf_files, output_path): print(f"Error processing {pdf_file}: {str(e)}") continue try: - with open(output_path, 'wb') as output_file: + with open(output_path, "wb") as output_file: merger.write(output_file) except Exception as e: print(f"Error saving merged PDF: {str(e)}") @@ -186,7 +185,6 @@ def merge_pdfs_in_folder(folder_path, output_filename="merged_perturbed_gene_QC. _merge_pdfs(pdf_files, output_path) - # merge all svgs to pdf def merge_svgs_to_pdf(folder_path, output_filename="merged_perturbed_gene_QC.pdf"): @@ -200,7 +198,7 @@ def merge_svgs_to_pdf(folder_path, output_filename="merged_perturbed_gene_QC.pdf for svg_file in svg_files: try: drawing = svg2rlg(svg_file) - temp_pdf = svg_file.replace('.svg', '_temp.pdf') + temp_pdf = svg_file.replace(".svg", "_temp.pdf") renderPDF.drawToFile(drawing, temp_pdf) temp_pdfs.append(temp_pdf) except Exception as e: @@ -218,18 +216,19 @@ def merge_svgs_to_pdf(folder_path, output_filename="merged_perturbed_gene_QC.pdf print(f"PDF created with {len(temp_pdfs)} pages: {output_path}") - def _natural_sort_key(filepath): """Sort key that orders numbers numerically (2 before 100) instead of lexicographically.""" basename = os.path.basename(filepath) - return [int(part) if part.isdigit() else part.lower() - for part in re.split(r'(\d+)', basename)] + return [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", basename) + ] def ensure_umap(mdata, data_key, prog_key, n_top_genes=2000, n_comps=50): """Compute PCA/UMAP from top-variance genes and write to both modalities; no-op if X_umap already exists in both.""" - has_prog = 'X_umap' in mdata[prog_key].obsm - has_data = 'X_umap' in mdata[data_key].obsm + has_prog = "X_umap" in mdata[prog_key].obsm + has_data = "X_umap" in mdata[data_key].obsm if has_prog and has_data: return # If exactly one modality already has UMAP and the modalities share n_obs, mirror it @@ -238,29 +237,33 @@ def ensure_umap(mdata, data_key, prog_key, n_top_genes=2000, n_comps=50): if has_prog ^ has_data: src_key, dst_key = (prog_key, data_key) if has_prog else (data_key, prog_key) if mdata[src_key].n_obs == mdata[dst_key].n_obs: - mdata[dst_key].obsm['X_umap'] = mdata[src_key].obsm['X_umap'] - if 'X_pca' in mdata[src_key].obsm: - mdata[dst_key].obsm['X_pca'] = mdata[src_key].obsm['X_pca'] + mdata[dst_key].obsm["X_umap"] = mdata[src_key].obsm["X_umap"] + if "X_pca" in mdata[src_key].obsm: + mdata[dst_key].obsm["X_pca"] = mdata[src_key].obsm["X_pca"] return # Compute variances on the live matrix (no copy) so a 1M×30K matrix doesn't get duplicated. X = mdata[data_key].X - variances = np.array(X.power(2).mean(axis=0) - np.power(X.mean(axis=0), 2)).flatten() \ - if hasattr(X, 'power') else X.var(axis=0) + variances = ( + np.array(X.power(2).mean(axis=0) - np.power(X.mean(axis=0), 2)).flatten() + if hasattr(X, "power") + else X.var(axis=0) + ) top_idx = np.argsort(variances)[-n_top_genes:] # Slice first, copy second: the copy is now ~n_top_genes columns instead of the full matrix. adata_tmp = mdata[data_key][:, top_idx].copy() sc.tl.pca(adata_tmp, n_comps=n_comps) sc.pp.neighbors(adata_tmp) sc.tl.umap(adata_tmp) - mdata[prog_key].obsm['X_pca'] = adata_tmp.obsm['X_pca'] - mdata[prog_key].obsm['X_umap'] = adata_tmp.obsm['X_umap'] - mdata[data_key].obsm['X_pca'] = adata_tmp.obsm['X_pca'] - mdata[data_key].obsm['X_umap'] = adata_tmp.obsm['X_umap'] + mdata[prog_key].obsm["X_pca"] = adata_tmp.obsm["X_pca"] + mdata[prog_key].obsm["X_umap"] = adata_tmp.obsm["X_umap"] + mdata[data_key].obsm["X_pca"] = adata_tmp.obsm["X_pca"] + mdata[data_key].obsm["X_umap"] = adata_tmp.obsm["X_umap"] def check_normalized(adata, key): """Raise if adata.X looks like raw integer counts instead of normalized expression.""" import scipy.sparse as sp + X = adata.X vals = X.data if sp.issparse(X) else np.asarray(X).ravel() if vals.size == 0: @@ -299,14 +302,16 @@ def read_condition_labels_from_h5mu(h5mu_path, data_key, categorical_key): """Like condition_labels_from_obs, but reads only mod//obs/ from the .h5mu instead of loading the full MuData.""" import h5py + try: from anndata.io import read_elem except ImportError: from anndata.experimental import read_elem - with h5py.File(h5mu_path, 'r') as f: - obs_path = f'mod/{data_key}/obs/{categorical_key}' + with h5py.File(h5mu_path, "r") as f: + obs_path = f"mod/{data_key}/obs/{categorical_key}" if obs_path not in f: raise KeyError( - f"'{obs_path}' not found in {h5mu_path}; pass the condition labels explicitly") + f"'{obs_path}' not found in {h5mu_path}; pass the condition labels explicitly" + ) column = read_elem(f[obs_path]) return [str(v) for v in pd.Series(column).dropna().unique()] diff --git a/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py b/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py index 3cce7ad..ef48ea4 100644 --- a/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py +++ b/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py @@ -27,7 +27,8 @@ # Change path to wherever you have repo locally from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) + +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) from Stage3_Interpretation.B_Summarization.src import ( compile_Program_loading_score_sheet_long, @@ -51,96 +52,207 @@ def main(): parser = argparse.ArgumentParser( description="Compile Stage 1 (.h5mu) + Stage 2 (Evaluation) outputs into a " - "single multi-sheet Excel summary workbook for one (K, sel_thresh).") + "single multi-sheet Excel summary workbook for one (K, sel_thresh)." + ) # IO info - parser.add_argument('--out_dir', help='Output root directory (contains {run_name}/).', type=str, required=True) - parser.add_argument('--run_name', help='Run name (subdirectory under out_dir).', type=str, required=True) - parser.add_argument('--K', help='Number of components (K) to summarize.', type=int, required=True) - parser.add_argument('--sel_thresh', help='Density selection threshold (dot->underscore handled internally). Default: 0.2', type=float, default=0.2) + parser.add_argument( + "--out_dir", + help="Output root directory (contains {run_name}/).", + type=str, + required=True, + ) + parser.add_argument( + "--run_name", + help="Run name (subdirectory under out_dir).", + type=str, + required=True, + ) + parser.add_argument( + "--K", help="Number of components (K) to summarize.", type=int, required=True + ) + parser.add_argument( + "--sel_thresh", + help="Density selection threshold (dot->underscore handled internally). Default: 0.2", + type=float, + default=0.2, + ) # Optional paths (derived from out_dir/run_name/K/thresh if not given) - parser.add_argument('--save_path', type=str, default=None, - help='Output directory for the .xlsx + TSV + sidecar files. ' - 'Default: {out_dir}/{run_name}/Interpretation/Summary_table/{K}_{thresh}') - parser.add_argument('--mdata_path', type=str, default=None, - help='Path to the Stage 1 .h5mu. ' - 'Default: {out_dir}/{run_name}/Inference/adata/cNMF_{K}_{thresh}.h5mu') + parser.add_argument( + "--save_path", + type=str, + default=None, + help="Output directory for the .xlsx + TSV + sidecar files. " + "Default: {out_dir}/{run_name}/Interpretation/Summary_table/{K}_{thresh}", + ) + parser.add_argument( + "--mdata_path", + type=str, + default=None, + help="Path to the Stage 1 .h5mu. " + "Default: {out_dir}/{run_name}/Inference/adata/cNMF_{K}_{thresh}.h5mu", + ) # Compile options - parser.add_argument('--num_gene', type=int, default=300, - help='Number of top genes per program / per enriched term to keep. Default: 300') - parser.add_argument('--perturbation_file_name', type=str, default='perturbation_association_results', + parser.add_argument( + "--num_gene", + type=int, + default=300, + help="Number of top genes per program / per enriched term to keep. Default: 300", + ) + parser.add_argument( + "--perturbation_file_name", + type=str, + default="perturbation_association_results", help='Perturbation result file stem (between "{K}_" and "_{Condition}.txt"). ' - 'Default: perturbation_association_results') - parser.add_argument('--Conditions', '--Sample', dest='Conditions', nargs='*', type=str, default=None, - help='Condition labels (values of obs[categorical_key]); default: all labels found in the data. (--Sample is a deprecated alias.)') - parser.add_argument('--effect_size', type=str, default='log2FC', - help='Effect-size column in perturbation files. Default: log2FC') - parser.add_argument('--control_target_name', type=str, default='non-targeting', - help='Control target name for KD efficiency. Default: non-targeting') - + "Default: perturbation_association_results", + ) + parser.add_argument( + "--Conditions", + "--Sample", + dest="Conditions", + nargs="*", + type=str, + default=None, + help="Condition labels (values of obs[categorical_key]); default: all labels found in the data. (--Sample is a deprecated alias.)", + ) + parser.add_argument( + "--effect_size", + type=str, + default="log2FC", + help="Effect-size column in perturbation files. Default: log2FC", + ) + parser.add_argument( + "--control_target_name", + type=str, + default="non-targeting", + help="Control target name for KD efficiency. Default: non-targeting", + ) # keys - parser.add_argument('--categorical_key', type=str, default='sample', - help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample") - parser.add_argument('--prog_key', type=str, default='cNMF', - help='Modality key for cNMF programs in the MuData. Default: cNMF') - parser.add_argument('--data_key', type=str, default='rna', - help='Modality key for RNA expression in the MuData. Default: rna') - parser.add_argument('--guide_targets_key', type=str, default='guide_targets', - help='uns key for guide target names. Default: guide_targets') - parser.add_argument('--gene_names_key', type=str, default='symbol', - help='var column with gene symbols (for loadings / KD efficiency). Default: symbol') - parser.add_argument('--adjusted_pval_key', type=str, default='Adjusted P-value', - help='Adjusted p-value column in enrichment files. Default: "Adjusted P-value"') - parser.add_argument('--non_targeting_key', nargs='*', type=str, default=['non-targeting'], - help='Control target name(s) for the program summary sheet. Default: non-targeting') + parser.add_argument( + "--categorical_key", + type=str, + default="sample", + help="Key in .obs holding each cell's condition label (e.g. timepoint, stimulus, donor); default: sample", + ) + parser.add_argument( + "--prog_key", + type=str, + default="cNMF", + help="Modality key for cNMF programs in the MuData. Default: cNMF", + ) + parser.add_argument( + "--data_key", + type=str, + default="rna", + help="Modality key for RNA expression in the MuData. Default: rna", + ) + parser.add_argument( + "--guide_targets_key", + type=str, + default="guide_targets", + help="uns key for guide target names. Default: guide_targets", + ) + parser.add_argument( + "--gene_names_key", + type=str, + default="symbol", + help="var column with gene symbols (for loadings / KD efficiency). Default: symbol", + ) + parser.add_argument( + "--adjusted_pval_key", + type=str, + default="Adjusted P-value", + help='Adjusted p-value column in enrichment files. Default: "Adjusted P-value"', + ) + parser.add_argument( + "--non_targeting_key", + nargs="*", + type=str, + default=["non-targeting"], + help="Control target name(s) for the program summary sheet. Default: non-targeting", + ) # enrichment / perturbation column-name keys - parser.add_argument('--GO_Term_key', type=str, default='Term', help='Index column in GO enrichment file. Default: Term') - parser.add_argument('--GO_Genes_key', type=str, default='Genes', help='Gene-list column in GO enrichment file. Default: Genes') - parser.add_argument('--Geneset_Term_key', type=str, default='Term', help='Index column in geneset enrichment file. Default: Term') - parser.add_argument('--Geneset_Genes_key', type=str, default='Genes', help='Gene-list column in geneset enrichment file. Default: Genes') - parser.add_argument('--Trait_Term_key', type=str, default='Term', help='Index column in trait enrichment file. Default: Term') - parser.add_argument('--Trait_Genes_key', type=str, default='Genes', help='Gene-list column in trait enrichment file. Default: Genes') - parser.add_argument('--Perturbation_Sample_key', type=str, default='Sample', - help='Sample-label column in perturbation results. Default: Sample') + parser.add_argument( + "--GO_Term_key", + type=str, + default="Term", + help="Index column in GO enrichment file. Default: Term", + ) + parser.add_argument( + "--GO_Genes_key", + type=str, + default="Genes", + help="Gene-list column in GO enrichment file. Default: Genes", + ) + parser.add_argument( + "--Geneset_Term_key", + type=str, + default="Term", + help="Index column in geneset enrichment file. Default: Term", + ) + parser.add_argument( + "--Geneset_Genes_key", + type=str, + default="Genes", + help="Gene-list column in geneset enrichment file. Default: Genes", + ) + parser.add_argument( + "--Trait_Term_key", + type=str, + default="Term", + help="Index column in trait enrichment file. Default: Term", + ) + parser.add_argument( + "--Trait_Genes_key", + type=str, + default="Genes", + help="Gene-list column in trait enrichment file. Default: Genes", + ) + parser.add_argument( + "--Perturbation_Sample_key", + type=str, + default="Sample", + help="Sample-label column in perturbation results. Default: Sample", + ) args = parser.parse_args() - if any(a == '--Sample' or a.startswith('--Sample=') for a in sys.argv[1:]): - print('Note: --Sample is deprecated; use --Conditions (same meaning).') + if any(a == "--Sample" or a.startswith("--Sample=") for a in sys.argv[1:]): + print("Note: --Sample is deprecated; use --Conditions (same meaning).") # Resolve derived defaults - thresh_str = str(args.sel_thresh).replace('.', '_') - eval_base = f'{args.out_dir}/{args.run_name}/Evaluation/{args.K}_{thresh_str}' + thresh_str = str(args.sel_thresh).replace(".", "_") + eval_base = f"{args.out_dir}/{args.run_name}/Evaluation/{args.K}_{thresh_str}" if args.save_path is None: - args.save_path = f'{args.out_dir}/{args.run_name}/Interpretation/Summary_table/{args.K}_{thresh_str}' + args.save_path = f"{args.out_dir}/{args.run_name}/Interpretation/Summary_table/{args.K}_{thresh_str}" if args.mdata_path is None: - args.mdata_path = f'{args.out_dir}/{args.run_name}/Inference/adata/cNMF_{args.K}_{thresh_str}.h5mu' + args.mdata_path = f"{args.out_dir}/{args.run_name}/Inference/adata/cNMF_{args.K}_{thresh_str}.h5mu" # create output directory - os.makedirs(f'{args.save_path}', exist_ok=True) + os.makedirs(f"{args.save_path}", exist_ok=True) # --- Save config (incl. SLURM info) --- slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), - 'num_nodes': os.environ.get('SLURM_JOB_NUM_NODES'), - 'ntasks': os.environ.get('SLURM_NTASKS'), - 'cpus_per_task': os.environ.get('SLURM_CPUS_PER_TASK'), - 'mem_per_node': os.environ.get('SLURM_MEM_PER_NODE'), - 'mem_per_cpu': os.environ.get('SLURM_MEM_PER_CPU'), - 'time_limit': os.environ.get('SLURM_JOB_TIMELIMIT'), - 'submit_dir': os.environ.get('SLURM_SUBMIT_DIR'), - 'array_task_id': os.environ.get('SLURM_ARRAY_TASK_ID'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), + "num_nodes": os.environ.get("SLURM_JOB_NUM_NODES"), + "ntasks": os.environ.get("SLURM_NTASKS"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + "mem_per_node": os.environ.get("SLURM_MEM_PER_NODE"), + "mem_per_cpu": os.environ.get("SLURM_MEM_PER_CPU"), + "time_limit": os.environ.get("SLURM_JOB_TIMELIMIT"), + "submit_dir": os.environ.get("SLURM_SUBMIT_DIR"), + "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" - config_to_save = {'script_args': vars(args), 'slurm_info': slurm_info} - with open(f'{args.save_path}/config_{job_id}.yml', 'w') as f: + config_to_save = {"script_args": vars(args), "slurm_info": slurm_info} + with open(f"{args.save_path}/config_{job_id}.yml", "w") as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) # ---- Pre-flight check: verify the input h5mu exists ---- @@ -148,89 +260,168 @@ def main(): raise FileNotFoundError(f"Input .h5mu not found: {args.mdata_path}") # ── Load MuData ── - print(f'Loading MuData: {args.mdata_path}') + print(f"Loading MuData: {args.mdata_path}") mdata = mu.read(args.mdata_path) # condition labels default to the values of --categorical_key in the data if args.Conditions is None: obs = mdata[args.data_key].obs if args.categorical_key not in obs.columns: - raise KeyError(f"obs column '{args.categorical_key}' not found; pass --Conditions explicitly") + raise KeyError( + f"obs column '{args.categorical_key}' not found; pass --Conditions explicitly" + ) args.Conditions = [str(v) for v in obs[args.categorical_key].dropna().unique()] - print(f"--Conditions not given; using {args.categorical_key} values from the h5mu: {args.Conditions}") + print( + f"--Conditions not given; using {args.categorical_key} values from the h5mu: {args.Conditions}" + ) # ── Evaluation file paths ── - GO_path = f'{eval_base}/{args.K}_GO_term_enrichment.txt' - Geneset_path = f'{eval_base}/{args.K}_geneset_enrichment.txt' - Trait_path = f'{eval_base}/{args.K}_trait_enrichment.txt' - Perturbation_path_base = f'{eval_base}/{args.K}_{args.perturbation_file_name}' - Association_path = f'{eval_base}/{args.K}_categorical_association_results.txt' - Explained_Variance_path = f'{eval_base}/{args.K}_Explained_Variance.txt' + GO_path = f"{eval_base}/{args.K}_GO_term_enrichment.txt" + Geneset_path = f"{eval_base}/{args.K}_geneset_enrichment.txt" + Trait_path = f"{eval_base}/{args.K}_trait_enrichment.txt" + Perturbation_path_base = f"{eval_base}/{args.K}_{args.perturbation_file_name}" + Association_path = f"{eval_base}/{args.K}_categorical_association_results.txt" + Explained_Variance_path = f"{eval_base}/{args.K}_Explained_Variance.txt" # ── Program loadings ── df_Program_loading_long = compile_Program_loading_score_sheet_long( - mdata, num_gene=args.num_gene, data_key=args.data_key, gene_names_key=args.gene_names_key) + mdata, + num_gene=args.num_gene, + data_key=args.data_key, + gene_names_key=args.gene_names_key, + ) df_Program_loading_flat = compile_Program_loading_score_sheet_flat( - mdata, num_gene=args.num_gene, data_key=args.data_key, gene_names_key=args.gene_names_key) + mdata, + num_gene=args.num_gene, + data_key=args.data_key, + gene_names_key=args.gene_names_key, + ) # ── Enrichment sheets ── - df_GO = Compile_GO_sheet(GO_path, gene_num=args.num_gene, term_key=args.GO_Term_key, genes_key=args.GO_Genes_key) if os.path.exists(GO_path) else None - df_Geneset = Compile_Geneset_sheet(Geneset_path, gene_num=args.num_gene, term_key=args.Geneset_Term_key, genes_key=args.Geneset_Genes_key) if os.path.exists(Geneset_path) else None - df_Trait = Compile_Trait_sheet(Trait_path, gene_num=args.num_gene, term_key=args.Trait_Term_key, genes_key=args.Trait_Genes_key) if os.path.exists(Trait_path) else None + df_GO = ( + Compile_GO_sheet( + GO_path, + gene_num=args.num_gene, + term_key=args.GO_Term_key, + genes_key=args.GO_Genes_key, + ) + if os.path.exists(GO_path) + else None + ) + df_Geneset = ( + Compile_Geneset_sheet( + Geneset_path, + gene_num=args.num_gene, + term_key=args.Geneset_Term_key, + genes_key=args.Geneset_Genes_key, + ) + if os.path.exists(Geneset_path) + else None + ) + df_Trait = ( + Compile_Trait_sheet( + Trait_path, + gene_num=args.num_gene, + term_key=args.Trait_Term_key, + genes_key=args.Trait_Genes_key, + ) + if os.path.exists(Trait_path) + else None + ) # ── Perturbation ── - perturbation_files = [f'{Perturbation_path_base}_{samp}.txt' for samp in args.Conditions] + perturbation_files = [ + f"{Perturbation_path_base}_{samp}.txt" for samp in args.Conditions + ] if any(os.path.exists(f) for f in perturbation_files): - df_Perturbation = Compile_Perturbation_sheet(Perturbation_path_base, Sample=args.Conditions, sample_key=args.Perturbation_Sample_key) + df_Perturbation = Compile_Perturbation_sheet( + Perturbation_path_base, + Sample=args.Conditions, + sample_key=args.Perturbation_Sample_key, + ) else: - print(f'No perturbation files found for: {Perturbation_path_base}') + print(f"No perturbation files found for: {Perturbation_path_base}") df_Perturbation = None # ── Categorical association & explained variance ── - df_Association = Compile_Association_sheet(Association_path, gene_num=args.num_gene) if os.path.exists(Association_path) else None - df_Explained_Variance = Compile_Explained_variance(Explained_Variance_path) if os.path.exists(Explained_Variance_path) else None + df_Association = ( + Compile_Association_sheet(Association_path, gene_num=args.num_gene) + if os.path.exists(Association_path) + else None + ) + df_Explained_Variance = ( + Compile_Explained_variance(Explained_Variance_path) + if os.path.exists(Explained_Variance_path) + else None + ) # ── Validate program names align across loaded DataFrames ── - check_program_name_match(mdata, prog_key=args.prog_key, dataframes=[ - df_GO, df_Geneset, df_Trait, df_Perturbation, - df_Association, df_Explained_Variance - ]) + check_program_name_match( + mdata, + prog_key=args.prog_key, + dataframes=[ + df_GO, + df_Geneset, + df_Trait, + df_Perturbation, + df_Association, + df_Explained_Variance, + ], + ) # ── Target Summary (writes specificity / correlation / KD-efficiency sidecars to save_path) ── df_Target_Summary = Compile_Target_Summary_sheet( - mdata, Perturbation_path_base, - Sample=args.Conditions, categorical_key=args.categorical_key, - prog_key=args.prog_key, data_key=args.data_key, + mdata, + Perturbation_path_base, + Sample=args.Conditions, + categorical_key=args.categorical_key, + prog_key=args.prog_key, + data_key=args.data_key, guide_targets_key=args.guide_targets_key, - save_path=args.save_path, effect_size=args.effect_size, + save_path=args.save_path, + effect_size=args.effect_size, control_target_name=args.control_target_name, gene_names_key=args.gene_names_key, ) # ── Program Summary ── df_Summary = Compile_Summary_sheet( - mdata, df_GO, df_Geneset, df_Perturbation, df_Program_loading_flat, df_Explained_Variance, - Sample=args.Conditions, specicicity_path=args.save_path, - categorical_key=args.categorical_key, non_tagerting_key=args.non_targeting_key, - effect_size=args.effect_size, adjusted_pval_key=args.adjusted_pval_key, + mdata, + df_GO, + df_Geneset, + df_Perturbation, + df_Program_loading_flat, + df_Explained_Variance, + Sample=args.Conditions, + specicicity_path=args.save_path, + categorical_key=args.categorical_key, + non_tagerting_key=args.non_targeting_key, + effect_size=args.effect_size, + adjusted_pval_key=args.adjusted_pval_key, ) # ── Save key DataFrames as separate TSV files ── - df_Summary.to_csv(f'{args.save_path}/Summary_{args.K}_{thresh_str}.tsv', sep='\t') - df_Program_loading_long.to_csv(f'{args.save_path}/Program_Loadings_{args.K}_{thresh_str}.tsv', sep='\t') - df_Target_Summary.to_csv(f'{args.save_path}/Targets_Summary_{args.K}_{thresh_str}.tsv', sep='\t') - print(f'Saved separate TSV files to {args.save_path}') + df_Summary.to_csv(f"{args.save_path}/Summary_{args.K}_{thresh_str}.tsv", sep="\t") + df_Program_loading_long.to_csv( + f"{args.save_path}/Program_Loadings_{args.K}_{thresh_str}.tsv", sep="\t" + ) + df_Target_Summary.to_csv( + f"{args.save_path}/Targets_Summary_{args.K}_{thresh_str}.tsv", sep="\t" + ) + print(f"Saved separate TSV files to {args.save_path}") # ── Write Excel ── - excel_path = f'{args.save_path}/cNMF_{args.K}_{thresh_str}.xlsx' - print(f'Compiling Excel for K={args.K}, threshold={args.sel_thresh}') + excel_path = f"{args.save_path}/cNMF_{args.K}_{thresh_str}.xlsx" + print(f"Compiling Excel for K={args.K}, threshold={args.sel_thresh}") with pd.ExcelWriter(excel_path) as writer: - df_Summary.to_excel(writer, sheet_name='Summary', index=True) - df_Program_loading_long.to_excel(writer, sheet_name='Program Loadings', index=True) - df_Target_Summary.to_excel(writer, sheet_name='Targets Summary', index=True) + df_Summary.to_excel(writer, sheet_name="Summary", index=True) + df_Program_loading_long.to_excel( + writer, sheet_name="Program Loadings", index=True + ) + df_Target_Summary.to_excel(writer, sheet_name="Targets Summary", index=True) if df_Association is not None: - df_Association.to_excel(writer, sheet_name='Sample Association', index=True) + df_Association.to_excel(writer, sheet_name="Sample Association", index=True) if df_Perturbation is not None: # Re-load perturbation per sample with specificity scores merged in. @@ -238,7 +429,9 @@ def main(): # both carry the specificity_scores column. combined_full, combined_sig = [], [] for samp in args.Conditions: - df_full_, df_sig_ = add_specificity_scores_file(args.save_path, Perturbation_path_base, samp) + df_full_, df_sig_ = add_specificity_scores_file( + args.save_path, Perturbation_path_base, samp + ) df_full_[args.Perturbation_Sample_key] = samp df_sig_[args.Perturbation_Sample_key] = samp combined_full.append(df_full_) @@ -248,25 +441,31 @@ def main(): for i in range(0, len(df_Perturbation_with_spec), MAX_ROWS): sheet_num = i // MAX_ROWS + 1 - df_Perturbation_with_spec.iloc[i:i + MAX_ROWS].to_excel( - writer, sheet_name=f'Perturbation Association {sheet_num}', index=True) + df_Perturbation_with_spec.iloc[i : i + MAX_ROWS].to_excel( + writer, + sheet_name=f"Perturbation Association {sheet_num}", + index=True, + ) for i in range(0, len(df_Perturbation_sig_with_spec), MAX_ROWS): sheet_num = i // MAX_ROWS + 1 - df_Perturbation_sig_with_spec.iloc[i:i + MAX_ROWS].to_excel( - writer, sheet_name=f'significant regulators only {sheet_num}', index=True) + df_Perturbation_sig_with_spec.iloc[i : i + MAX_ROWS].to_excel( + writer, + sheet_name=f"significant regulators only {sheet_num}", + index=True, + ) if df_Trait is not None: - df_Trait.to_excel(writer, sheet_name='Trait Enrichment', index=True) + df_Trait.to_excel(writer, sheet_name="Trait Enrichment", index=True) if df_GO is not None: - df_GO.to_excel(writer, sheet_name='GO Term Enrichment', index=True) + df_GO.to_excel(writer, sheet_name="GO Term Enrichment", index=True) if df_Geneset is not None: - df_Geneset.to_excel(writer, sheet_name='Geneset Enrichment', index=True) + df_Geneset.to_excel(writer, sheet_name="Geneset Enrichment", index=True) - print(f'Done. Saved to {excel_path}') + print(f"Done. Saved to {excel_path}") print("Pipeline finished.") return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py b/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py index e453426..924b338 100644 --- a/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py +++ b/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py @@ -13,18 +13,22 @@ from functools import reduce - # Change path to wherever you have repo locally from pathlib import Path -sys.path.append(str(Path(__file__).resolve().parents[4] / 'src')) +sys.path.append(str(Path(__file__).resolve().parents[4] / "src")) + + +from Stage3_Interpretation.A_Plotting.src import ( + rename_adata_gene_dictionary, + compute_gene_waterfall_cor, +) -from Stage3_Interpretation.A_Plotting.src import rename_adata_gene_dictionary ,compute_gene_waterfall_cor +# -------------- helper methods compile into excel sheets, not used here -------------- -#-------------- helper methods compile into excel sheets, not used here -------------- # check matching names in program -def check_program_name_match(mdata, dataframes, prog_key='cNMF'): +def check_program_name_match(mdata, dataframes, prog_key="cNMF"): """Check that program names in all loaded DataFrames match mdata var_names. Parameters @@ -47,9 +51,9 @@ def check_program_name_match(mdata, dataframes, prog_key='cNMF'): for df in dataframes: if df is None: continue - if 'program_name' not in df.columns: + if "program_name" not in df.columns: continue - file_programs = set(str(v) for v in df['program_name'].unique()) + file_programs = set(str(v) for v in df["program_name"].unique()) if file_programs != mdata_programs: try: file_sorted = sorted(file_programs, key=int) @@ -63,55 +67,78 @@ def check_program_name_match(mdata, dataframes, prog_key='cNMF'): for file_sorted in mismatched: print(f" file programs: {file_sorted}") -# merge to have specificity score added for perturbation results -def add_specificity_scores_file(save_path, Perturbation_path_base, samp, adj_pval_threshold=0.05): - - PMI = pd.read_csv(f'{save_path}/specificity_score_{samp}.txt', sep='\t', index_col=0) - df_perturbation = pd.read_csv(f'{Perturbation_path_base}_{samp}.txt', sep='\t') - # convert PMI to long form - df_PMI = PMI.reset_index().melt(id_vars="target_name", var_name="program_name", value_name="specificity_scores") +# merge to have specificity score added for perturbation results +def add_specificity_scores_file( + save_path, Perturbation_path_base, samp, adj_pval_threshold=0.05 +): - # make sure program name is in str form - df_PMI["program_name"] = df_PMI["program_name"].astype(str) - df_perturbation["program_name"] = df_perturbation["program_name"].astype(str) + PMI = pd.read_csv( + f"{save_path}/specificity_score_{samp}.txt", sep="\t", index_col=0 + ) + df_perturbation = pd.read_csv(f"{Perturbation_path_base}_{samp}.txt", sep="\t") - # merge to have specificity score added for perturbation results - df_perturbation_merged = df_PMI.merge(df_perturbation, on=["target_name", "program_name"]) + # convert PMI to long form + df_PMI = PMI.reset_index().melt( + id_vars="target_name", var_name="program_name", value_name="specificity_scores" + ) - # significant-only subset, derived from the merged frame so it keeps specificity_scores - df_sig_df_perturbation_merged = df_perturbation_merged[df_perturbation_merged['adj_pval'] < adj_pval_threshold] + # make sure program name is in str form + df_PMI["program_name"] = df_PMI["program_name"].astype(str) + df_perturbation["program_name"] = df_perturbation["program_name"].astype(str) - df_perturbation_merged.to_csv(f'{save_path}/perturbation_merged_{samp}.txt', sep='\t') - df_sig_df_perturbation_merged.to_csv(f'{save_path}/perturbation_merged_{samp}_significant.txt', sep='\t') + # merge to have specificity score added for perturbation results + df_perturbation_merged = df_PMI.merge( + df_perturbation, on=["target_name", "program_name"] + ) - return df_perturbation_merged, df_sig_df_perturbation_merged + # significant-only subset, derived from the merged frame so it keeps specificity_scores + df_sig_df_perturbation_merged = df_perturbation_merged[ + df_perturbation_merged["adj_pval"] < adj_pval_threshold + ] + df_perturbation_merged.to_csv( + f"{save_path}/perturbation_merged_{samp}.txt", sep="\t" + ) + df_sig_df_perturbation_merged.to_csv( + f"{save_path}/perturbation_merged_{samp}_significant.txt", sep="\t" + ) + return df_perturbation_merged, df_sig_df_perturbation_merged # Compile simple sheets -#-------------- helper methods for loading sheets-------------- -def compile_Program_loading_score_sheet_long(mdata, num_gene = 300, data_key='rna', gene_names_key=None): - print('Load program loadings data in long form') +# -------------- helper methods for loading sheets-------------- +def compile_Program_loading_score_sheet_long( + mdata, num_gene=300, data_key="rna", gene_names_key=None +): - program_loading_df = pd.DataFrame(data=mdata['cNMF'].varm["loadings"], columns=mdata['cNMF'].uns['var_names'], index=mdata['cNMF'].var_names) + print("Load program loadings data in long form") + + program_loading_df = pd.DataFrame( + data=mdata["cNMF"].varm["loadings"], + columns=mdata["cNMF"].uns["var_names"], + index=mdata["cNMF"].var_names, + ) # Rename columns to curated gene names if gene_names_key is not None: - rename_map = dict(zip(mdata[data_key].var_names, mdata[data_key].var[gene_names_key])) + rename_map = dict( + zip(mdata[data_key].var_names, mdata[data_key].var[gene_names_key]) + ) else: - rename_map = dict(zip(mdata['cNMF'].uns['var_names'], mdata[data_key].var_names)) + rename_map = dict( + zip(mdata["cNMF"].uns["var_names"], mdata[data_key].var_names) + ) program_loading_df.rename(columns=rename_map, inplace=True) top_df = program_loading_df.apply( - lambda row: row.nlargest(num_gene).index.tolist(), - axis=1 + lambda row: row.nlargest(num_gene).index.tolist(), axis=1 ) - result_df = pd.DataFrame(top_df.tolist(), columns=range(1, num_gene+1)) + result_df = pd.DataFrame(top_df.tolist(), columns=range(1, num_gene + 1)) result_df.index = program_loading_df.index result_df.index.name = "Program" @@ -120,371 +147,459 @@ def compile_Program_loading_score_sheet_long(mdata, num_gene = 300, data_key='rn long_data = [] for program_idx, row in result_df.iterrows(): - # Query all genes for this program at once genes_list = row.dropna().astype(str).str.strip().tolist() genes_list = [g for g in genes_list if g] # Remove empty strings - + if not genes_list: continue - + # Get results for all genes in this program - results = mg.querymany(genes_list, scopes='symbol', fields='summary', species='human', verbose=False) - + results = mg.querymany( + genes_list, + scopes="symbol", + fields="summary", + species="human", + verbose=False, + ) + # Create a lookup dictionary for faster access summary_dict = {} for r in results: - q = r.get('query') + q = r.get("query") if q not in summary_dict: # keep first (best) match - summary_dict[q] = r.get('summary', 'N/A') - + summary_dict[q] = r.get("summary", "N/A") + # Now iterate through the ranked genes for rank, gene in enumerate(row, 1): - # Skip NaN values - if pd.notna(gene) and gene != '': + if pd.notna(gene) and gene != "": gene_clean = str(gene).strip() - annotation = summary_dict.get(gene_clean, 'N/A') - - long_data.append({ - 'Program': program_idx, - 'Rank': rank, - 'Gene': gene_clean, - 'Annotation': annotation - }) + annotation = summary_dict.get(gene_clean, "N/A") + + long_data.append( + { + "Program": program_idx, + "Rank": rank, + "Gene": gene_clean, + "Annotation": annotation, + } + ) # Convert to DataFrame annotation_df = pd.DataFrame(long_data) return annotation_df -def compile_Program_loading_score_sheet_flat(mdata, num_gene = 300, data_key='rna', gene_names_key=None): - print('Load program loadings data in flat form') +def compile_Program_loading_score_sheet_flat( + mdata, num_gene=300, data_key="rna", gene_names_key=None +): + print("Load program loadings data in flat form") - program_loading_df = pd.DataFrame(data=mdata['cNMF'].varm["loadings"], columns=mdata['cNMF'].uns['var_names'], index=mdata['cNMF'].var_names) + program_loading_df = pd.DataFrame( + data=mdata["cNMF"].varm["loadings"], + columns=mdata["cNMF"].uns["var_names"], + index=mdata["cNMF"].var_names, + ) # Rename columns to curated gene names if gene_names_key is not None: - rename_map = dict(zip(mdata[data_key].var_names, mdata[data_key].var[gene_names_key])) + rename_map = dict( + zip(mdata[data_key].var_names, mdata[data_key].var[gene_names_key]) + ) else: - rename_map = dict(zip(mdata['cNMF'].uns['var_names'], mdata[data_key].var_names)) + rename_map = dict( + zip(mdata["cNMF"].uns["var_names"], mdata[data_key].var_names) + ) program_loading_df.rename(columns=rename_map, inplace=True) top_df = program_loading_df.apply( - lambda row: row.nlargest(num_gene).index.tolist(), - axis=1 + lambda row: row.nlargest(num_gene).index.tolist(), axis=1 ) - result_df = pd.DataFrame(top_df.tolist(), columns=range(1, num_gene+1)) + result_df = pd.DataFrame(top_df.tolist(), columns=range(1, num_gene + 1)) result_df.index = program_loading_df.index result_df.index.name = "Program" return result_df -def Compile_GO_sheet(GO_path, gene_num = 5, term_key = "Term", genes_key = "Genes"): - print('Load GO data') +def Compile_GO_sheet(GO_path, gene_num=5, term_key="Term", genes_key="Genes"): + + print("Load GO data") - df = pd.read_csv(GO_path, sep = "\t", index_col = 0) + df = pd.read_csv(GO_path, sep="\t", index_col=0) df = df.reset_index().set_index(term_key) - df[genes_key] = df[genes_key].str.split(';').str[:gene_num].str.join(";") + df[genes_key] = df[genes_key].str.split(";").str[:gene_num].str.join(";") return df -def Compile_Geneset_sheet(Geneset_path, gene_num = 5, term_key = "Term", genes_key = "Genes"): - print('Load geneset data') +def Compile_Geneset_sheet(Geneset_path, gene_num=5, term_key="Term", genes_key="Genes"): + print("Load geneset data") - df = pd.read_csv(Geneset_path, sep = "\t", index_col = 0) + df = pd.read_csv(Geneset_path, sep="\t", index_col=0) df = df.reset_index().set_index(term_key) - df[genes_key] = df[genes_key].str.split(';').str[:gene_num].str.join(";") + df[genes_key] = df[genes_key].str.split(";").str[:gene_num].str.join(";") return df -def Compile_Trait_sheet(Trait_path, gene_num = 5, term_key = "Term", genes_key = "Genes"): - print('Load trait data') +def Compile_Trait_sheet(Trait_path, gene_num=5, term_key="Term", genes_key="Genes"): - df = pd.read_csv(Trait_path, sep = "\t", index_col = 0) + print("Load trait data") + + df = pd.read_csv(Trait_path, sep="\t", index_col=0) df = df.reset_index().set_index(term_key) - df[genes_key] = df[genes_key].str.split(';').str[:gene_num].str.join(";") + df[genes_key] = df[genes_key].str.split(";").str[:gene_num].str.join(";") return df -def Compile_Perturbation_sheet(Perturbation_path, Sample = ["D0", "sample_D1","sample_D2","sample_D3" ], sample_key = "Sample"): - print('Load perturbation data') +def Compile_Perturbation_sheet( + Perturbation_path, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + sample_key="Sample", +): + print("Load perturbation data") combined_conditions = [] for samp in Sample: - df = pd.read_csv(f"{Perturbation_path}_{samp}.txt", sep = "\t") + df = pd.read_csv(f"{Perturbation_path}_{samp}.txt", sep="\t") df[sample_key] = samp combined_conditions.append(df) df = pd.concat(combined_conditions) - - return df -def Compile_Association_sheet(Association_path, gene_num = 5): + return df + - print('Load categorical associa data') +def Compile_Association_sheet(Association_path, gene_num=5): - df = pd.read_csv(Association_path, sep = "\t", index_col = 0) + print("Load categorical associa data") + + df = pd.read_csv(Association_path, sep="\t", index_col=0) df.index.name = "program_name" df.sort_index(inplace=True, ascending=True) return df + def Compile_Explained_variance(Explained_Variance_path): - print('Load explained variance data') + print("Load explained variance data") - df = pd.read_csv(Explained_Variance_path, sep = "\t") - df = df.set_index('program_name') + df = pd.read_csv(Explained_Variance_path, sep="\t") + df = df.set_index("program_name") return df -#-------------- helper methods for loading sheets-------------- -# combine methods for simple sheets -def load_simple_sheets(mdata, out_dir, run_name, k, sel_thresh, num_gene = 300, perturbation_file_name = "perturbation_association_results", Sample = ['1', '2', '3'], - GO_Term_key = "Term", GO_Genes_key = "Genes", - Geneset_Term_key = "Term", Geneset_Genes_key = "Genes", - Trait_Term_key = "Term", Trait_Genes_key = "Genes", - Perturbation_Sample_key = "Sample"): +# -------------- helper methods for loading sheets-------------- - print('Load simple sheet') - GO_path = f'{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace(".","_")}/{k}_GO_term_enrichment.txt' - Geneset_path = f'{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace(".","_")}/{k}_geneset_enrichment.txt' - Trait_path = f'{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace(".","_")}/{k}_trait_enrichment.txt' - Perturbation_path_base = f'{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace(".","_")}/{k}_{perturbation_file_name}' - Association_path = f'{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace(".","_")}/{k}_categorical_association_results.txt' - Explained_Variance_path = f'{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace(".","_")}/{k}_Explained_Variance.txt' +# combine methods for simple sheets +def load_simple_sheets( + mdata, + out_dir, + run_name, + k, + sel_thresh, + num_gene=300, + perturbation_file_name="perturbation_association_results", + Sample=["1", "2", "3"], + GO_Term_key="Term", + GO_Genes_key="Genes", + Geneset_Term_key="Term", + Geneset_Genes_key="Genes", + Trait_Term_key="Term", + Trait_Genes_key="Genes", + Perturbation_Sample_key="Sample", +): + + print("Load simple sheet") + + GO_path = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}/{k}_GO_term_enrichment.txt" + Geneset_path = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}/{k}_geneset_enrichment.txt" + Trait_path = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}/{k}_trait_enrichment.txt" + Perturbation_path_base = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}/{k}_{perturbation_file_name}" + Association_path = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}/{k}_categorical_association_results.txt" + Explained_Variance_path = f"{out_dir}/{run_name}/Evaluation/{k}_{str(sel_thresh).replace('.', '_')}/{k}_Explained_Variance.txt" # compile program loadings - df_Program_loading_long = compile_Program_loading_score_sheet_long(mdata, num_gene = num_gene) - df_Program_loading_flat = compile_Program_loading_score_sheet_flat(mdata, num_gene = num_gene) + df_Program_loading_long = compile_Program_loading_score_sheet_long( + mdata, num_gene=num_gene + ) + df_Program_loading_flat = compile_Program_loading_score_sheet_flat( + mdata, num_gene=num_gene + ) # compile GO if os.path.exists(GO_path): - df_GO = Compile_GO_sheet(GO_path, gene_num = num_gene, term_key = GO_Term_key, genes_key = GO_Genes_key) + df_GO = Compile_GO_sheet( + GO_path, gene_num=num_gene, term_key=GO_Term_key, genes_key=GO_Genes_key + ) else: - print(f'GO file not found: {GO_path}') + print(f"GO file not found: {GO_path}") df_GO = None # compile Genesets if os.path.exists(Geneset_path): - df_Geneset = Compile_Geneset_sheet(Geneset_path, gene_num = num_gene, term_key = Geneset_Term_key, genes_key = Geneset_Genes_key) + df_Geneset = Compile_Geneset_sheet( + Geneset_path, + gene_num=num_gene, + term_key=Geneset_Term_key, + genes_key=Geneset_Genes_key, + ) else: - print(f'Geneset file not found: {Geneset_path}') + print(f"Geneset file not found: {Geneset_path}") df_Geneset = None # compile Trait if os.path.exists(Trait_path): - df_Trait = Compile_Trait_sheet(Trait_path, gene_num = num_gene, term_key = Trait_Term_key, genes_key = Trait_Genes_key) + df_Trait = Compile_Trait_sheet( + Trait_path, + gene_num=num_gene, + term_key=Trait_Term_key, + genes_key=Trait_Genes_key, + ) else: - print(f'Trait file not found: {Trait_path}') + print(f"Trait file not found: {Trait_path}") df_Trait = None # compile perturbation perturbation_files = [f"{Perturbation_path_base}_{samp}.txt" for samp in Sample] if any(os.path.exists(f) for f in perturbation_files): - df_Perturbation = Compile_Perturbation_sheet(Perturbation_path_base, Sample = Sample, sample_key = Perturbation_Sample_key) + df_Perturbation = Compile_Perturbation_sheet( + Perturbation_path_base, Sample=Sample, sample_key=Perturbation_Sample_key + ) else: - print(f'No perturbation files found for base: {Perturbation_path_base}') + print(f"No perturbation files found for base: {Perturbation_path_base}") df_Perturbation = None # compile association if os.path.exists(Association_path): - df_Association = Compile_Association_sheet(Association_path, gene_num = num_gene) + df_Association = Compile_Association_sheet(Association_path, gene_num=num_gene) else: - print(f'Association file not found: {Association_path}') + print(f"Association file not found: {Association_path}") df_Association = None # compile explained variance if os.path.exists(Explained_Variance_path): df_Explained_Variance = Compile_Explained_variance(Explained_Variance_path) else: - print(f'Explained variance file not found: {Explained_Variance_path}') + print(f"Explained variance file not found: {Explained_Variance_path}") df_Explained_Variance = None - return df_Program_loading_long, df_Program_loading_flat, df_GO, df_Geneset, df_Trait, df_Perturbation, df_Association, df_Explained_Variance + return ( + df_Program_loading_long, + df_Program_loading_flat, + df_GO, + df_Geneset, + df_Trait, + df_Perturbation, + df_Association, + df_Explained_Variance, + ) +# calcaulte specicicity scores +# -------------- helper methods for specicicity scores-------------- -# calcaulte specicicity scores +# helper methods for specicicity score calcualtion +def Compute_regulator_zscore(perturb_df, effect_size="log2FC"): -#-------------- helper methods for specicicity scores-------------- + # Pivot to get program × gene log2FC table + pivot_table = perturb_df.pivot_table( + index="target_name", columns="program_name", values=effect_size, aggfunc="first" + ) -# helper methods for specicicity score calcualtion -def Compute_regulator_zscore(perturb_df, effect_size='log2FC'): + # convert to abs values + # pivot_table_abs = pivot_table.abs() - # Pivot to get program × gene log2FC table - pivot_table = perturb_df.pivot_table( - index='target_name', - columns='program_name', - values=effect_size, - aggfunc='first' - ) + # Compute z-scores row-wise + median = pivot_table.median(axis=1) + mad = pivot_table.sub(median, axis=0).abs().median(axis=1) + epsilon = 1e-10 - # convert to abs values - # pivot_table_abs = pivot_table.abs() + z_score_table = pivot_table.sub(median, axis=0).div(1.4826 * mad + epsilon, axis=0) - # Compute z-scores row-wise - median = pivot_table.median(axis=1) - mad = pivot_table.sub(median, axis=0).abs().median(axis=1) - epsilon = 1e-10 + return z_score_table - z_score_table = pivot_table.sub(median, axis=0).div(1.4826 * mad + epsilon, axis=0) - - return z_score_table def softmax_with_temperature(effects_df, T=1.0): # eq = e^{z/T}/ sum[ e^{z/T}] - + # Convert pandas to numpy and abs effects = effects_df.abs().values - + effects_scaled = effects / T - effects_scaled = effects_scaled - effects_scaled.max(axis=1, keepdims=True) # for overflow prevention - + effects_scaled = effects_scaled - effects_scaled.max( + axis=1, keepdims=True + ) # for overflow prevention + exp_effects = np.exp(effects_scaled) weights = exp_effects / exp_effects.sum() - + # Convert back to DataFrame - return pd.DataFrame(weights, - index=effects_df.index, - columns=effects_df.columns) + return pd.DataFrame(weights, index=effects_df.index, columns=effects_df.columns) + def compute_joint_distribution(weights): weights_array = weights.values - + # Joint: divide entire matrix by total sum P_tp = weights_array / weights_array.sum() - + # Marginals: sum by rows/columns P_t = P_tp.sum(axis=1) # Row sums P_p = P_tp.sum(axis=0) # Column sums # Convert back to DataFrame - P_tp = pd.DataFrame(P_tp, - index=weights.index, - columns=weights.columns) + P_tp = pd.DataFrame(P_tp, index=weights.index, columns=weights.columns) - return P_tp, P_t, P_p + def compute_pointwise_mutual_information(P_tp, P_t, P_p): # Reshape marginals for broadcasting P_t_col = P_t.reshape(-1, 1) # Column vector P_p_row = P_p.reshape(1, -1) # Row vector - + # P(t) * P(p) for all combinations P_independent = P_t_col @ P_p_row # Outer product - + # Avoid log(0) epsilon = 1e-10 - + # PMI = log(P_tp / P_independent) PMI = np.log((P_tp.values + epsilon) / (P_independent + epsilon)) - - # Convert back to DataFrame for readability - PMI_df = pd.DataFrame(PMI, - index=P_tp.index, - columns=P_tp.columns) - - return PMI_df -#-------------- helper methods specicicity scores-------------- + # Convert back to DataFrame for readability + PMI_df = pd.DataFrame(PMI, index=P_tp.index, columns=P_tp.columns) -# get specific programs and scores -def get_specificity_program(Perturbation_path, Sample = ["D0", "sample_D1", "sample_D2", "sample_D3"], T=0.5, save_path = - None, effect_size='log2FC'): - - print("Calculate specificity scores for program") - - all_specificity = [] + return PMI_df - for samp in Sample: - df = pd.read_csv(f"{Perturbation_path}_{samp}.txt", sep="\t") - df['program_name'] = df['program_name'].astype(str) - z_score_table = Compute_regulator_zscore(df, effect_size=effect_size) - weights = softmax_with_temperature(z_score_table, T=T) - P_tp, P_t, P_p = compute_joint_distribution(weights) - PMI = compute_pointwise_mutual_information(P_tp, P_t, P_p) +# -------------- helper methods specicicity scores-------------- - if save_path is not None: - PMI.to_csv(f'{save_path}/specificity_score_{samp}.txt', sep = '\t') - # Fix #3: Ensure PMI columns are strings to match sig_mask columns - PMI.columns = PMI.columns.astype(str) +# get specific programs and scores +def get_specificity_program( + Perturbation_path, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + T=0.5, + save_path=None, + effect_size="log2FC", +): - # Fix #4: Deduplicate PMI index to prevent .loc returning a DataFrame - PMI = PMI[~PMI.index.duplicated(keep='first')] + print("Calculate specificity scores for program") - # Build a mask of significant gene-program pairs (adj_pval < 0.1) - sig_mask = df[df['adj_pval'] < 0.1].pivot_table( - index='target_name', columns='program_name', values='adj_pval', aggfunc='first' - ).notna() + all_specificity = [] - # Mask PMI: keep only significant target-program pairs, NaN the rest - sig_mask_aligned = sig_mask.reindex(index=PMI.index, columns=PMI.columns, fill_value=False) - PMI_filtered = PMI.where(sig_mask_aligned) + for samp in Sample: + df = pd.read_csv(f"{Perturbation_path}_{samp}.txt", sep="\t") + df["program_name"] = df["program_name"].astype(str) - # For each target, get top 5 programs by PMI score - program_specificity = [] - for gene_name in PMI_filtered.index: - top_5 = PMI_filtered.loc[gene_name].dropna().nlargest(5) - top_5_programs = ', '.join(top_5.index.astype(str)) if len(top_5) > 0 else '' - top_5_scores = ', '.join([f"{score:.4f}" for score in top_5.values]) if len(top_5) > 0 else '' + z_score_table = Compute_regulator_zscore(df, effect_size=effect_size) + weights = softmax_with_temperature(z_score_table, T=T) + P_tp, P_t, P_p = compute_joint_distribution(weights) + PMI = compute_pointwise_mutual_information(P_tp, P_t, P_p) - program_specificity.append({ - 'target_name': gene_name, - f'top 5 specific programs (FDR < 0.1) {samp}': top_5_programs, - f'top 5 specificity scores (FDR < 0.1) {samp}': top_5_scores, - }) + if save_path is not None: + PMI.to_csv(f"{save_path}/specificity_score_{samp}.txt", sep="\t") + + # Fix #3: Ensure PMI columns are strings to match sig_mask columns + PMI.columns = PMI.columns.astype(str) + + # Fix #4: Deduplicate PMI index to prevent .loc returning a DataFrame + PMI = PMI[~PMI.index.duplicated(keep="first")] + + # Build a mask of significant gene-program pairs (adj_pval < 0.1) + sig_mask = ( + df[df["adj_pval"] < 0.1] + .pivot_table( + index="target_name", + columns="program_name", + values="adj_pval", + aggfunc="first", + ) + .notna() + ) - df_specificity = pd.DataFrame(program_specificity).set_index('target_name') - all_specificity.append(df_specificity) + # Mask PMI: keep only significant target-program pairs, NaN the rest + sig_mask_aligned = sig_mask.reindex( + index=PMI.index, columns=PMI.columns, fill_value=False + ) + PMI_filtered = PMI.where(sig_mask_aligned) + + # For each target, get top 5 programs by PMI score + program_specificity = [] + for gene_name in PMI_filtered.index: + top_5 = PMI_filtered.loc[gene_name].dropna().nlargest(5) + top_5_programs = ( + ", ".join(top_5.index.astype(str)) if len(top_5) > 0 else "" + ) + top_5_scores = ( + ", ".join([f"{score:.4f}" for score in top_5.values]) + if len(top_5) > 0 + else "" + ) - # Fix #5: Guard against empty Sample list - if not all_specificity: - return pd.DataFrame() + program_specificity.append( + { + "target_name": gene_name, + f"top 5 specific programs (FDR < 0.1) {samp}": top_5_programs, + f"top 5 specificity scores (FDR < 0.1) {samp}": top_5_scores, + } + ) - # Merge all results by index (gene) - merged_df = pd.concat(all_specificity, axis=1) + df_specificity = pd.DataFrame(program_specificity).set_index("target_name") + all_specificity.append(df_specificity) - return merged_df + # Fix #5: Guard against empty Sample list + if not all_specificity: + return pd.DataFrame() + # Merge all results by index (gene) + merged_df = pd.concat(all_specificity, axis=1) + return merged_df # load target summary -#-------------- helper methods target summary-------------- +# -------------- helper methods target summary-------------- + # Get cells with the guide per day -def get_guide_cells_per_days(mdata, categorical_key="sample", data_key='rna', prog_key='cNMF', - guide_assignment_key="guide_assignment", guide_targets_key="guide_targets"): - - print('Compute guides cells per days') - - guide_assignment = mdata[prog_key].obsm[guide_assignment_key] # (n_cells × n_targets), no copy +def get_guide_cells_per_days( + mdata, + categorical_key="sample", + data_key="rna", + prog_key="cNMF", + guide_assignment_key="guide_assignment", + guide_targets_key="guide_targets", +): + + print("Compute guides cells per days") + + guide_assignment = mdata[prog_key].obsm[ + guide_assignment_key + ] # (n_cells × n_targets), no copy guide_targets = mdata[prog_key].uns[guide_targets_key] if not sparse.issparse(guide_assignment): @@ -498,7 +613,7 @@ def get_guide_cells_per_days(mdata, categorical_key="sample", data_key='rna', pr indicator = sparse.csc_matrix( (np.ones(len(groups), dtype=np.float32), (np.arange(len(groups)), col_idx)), - shape=(len(groups), len(unique_groups)) + shape=(len(groups), len(unique_groups)), ) # Sparse matmul: (n_groups × n_cells) @ (n_cells × n_targets) = (n_groups × n_targets) @@ -508,7 +623,7 @@ def get_guide_cells_per_days(mdata, categorical_key="sample", data_key='rna', pr df_merge = pd.DataFrame( np.asarray(group_sums.todense()).T, index=guide_targets, - columns=[f'# Cells {g}' for g in unique_groups] + columns=[f"# Cells {g}" for g in unique_groups], ) # Collapse duplicate target names @@ -517,12 +632,19 @@ def get_guide_cells_per_days(mdata, categorical_key="sample", data_key='rna', pr return df_merge -# Get the gene expression average per day -def get_guide_mean_expr_per_day(mdata, categorical_key = "sample", prog_key = 'cNMF', data_key = 'rna', guide_targets_key = "guide_targets"): - - print('Get targeted gene mean expressuion per day') - - adata = mdata[data_key] # no .copy() — avoid duplicating the whole object + +# Get the gene expression average per day +def get_guide_mean_expr_per_day( + mdata, + categorical_key="sample", + prog_key="cNMF", + data_key="rna", + guide_targets_key="guide_targets", +): + + print("Get targeted gene mean expressuion per day") + + adata = mdata[data_key] # no .copy() — avoid duplicating the whole object perturbed_genes = mdata[prog_key].uns[guide_targets_key] @@ -538,24 +660,34 @@ def get_guide_mean_expr_per_day(mdata, categorical_key = "sample", prog_key = 'c mean_per_group = {} for group_name in groups.unique(): row_mask = (groups == group_name).values - mean_per_group[f'mean_expression_{group_name}'] = np.asarray(X_sub[row_mask].mean(axis=0)).ravel() + mean_per_group[f"mean_expression_{group_name}"] = np.asarray( + X_sub[row_mask].mean(axis=0) + ).ravel() mean_expr_per_day = pd.DataFrame(mean_per_group, index=gene_names) mean_expr_per_day.sort_index(inplace=True) - mean_expr_per_day.index.name = 'target_name' + mean_expr_per_day.index.name = "target_name" return mean_expr_per_day -def compute_kd_efficiency(mdata, categorical_key='sample', prog_key='cNMF', data_key='rna', - gene_names_key=None, guide_targets_key='guide_targets', - guide_assignment_key='guide_assignment', - control_target_name='non-targeting', save_path=None): + +def compute_kd_efficiency( + mdata, + categorical_key="sample", + prog_key="cNMF", + data_key="rna", + gene_names_key=None, + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", + control_target_name="non-targeting", + save_path=None, +): """Compute knockdown efficiency for each perturbation target per condition. Returns a DataFrame with index=target_name and columns=kd_efficiency_{condition}. KD efficiency = (1 - mean_perturbed / mean_control) * 100, using CP10K-normalized expression. """ - print('Compute KD efficiency per target per condition') + print("Compute KD efficiency per target per condition") guide_targets = mdata[prog_key].uns[guide_targets_key] ga = mdata[prog_key].obsm[guide_assignment_key] @@ -607,7 +739,9 @@ def compute_kd_efficiency(mdata, categorical_key='sample', prog_key='cNMF', data # Perturbed cells in this condition target_idx = np.where(guide_targets == target)[0] - pert_mask = (np.array(ga[:, target_idx].sum(axis=1)).flatten() > 0) & cond_mask + pert_mask = ( + np.array(ga[:, target_idx].sum(axis=1)).flatten() > 0 + ) & cond_mask pert_expr = gene_expr_norm[pert_mask] pert_mean = np.mean(pert_expr) if pert_mask.sum() > 0 else np.nan @@ -616,76 +750,98 @@ def compute_kd_efficiency(mdata, categorical_key='sample', prog_key='cNMF', data else: kd_vals[target] = np.nan - results[f'kd_efficiency_{condition}'] = kd_vals + results[f"kd_efficiency_{condition}"] = kd_vals df_kd = pd.DataFrame(results) - df_kd.index.name = 'target_name' - df_kd['Mean KD efficiency (%) across conditions'] = df_kd.mean(axis=1) + df_kd.index.name = "target_name" + df_kd["Mean KD efficiency (%) across conditions"] = df_kd.mean(axis=1) df_kd.sort_index(inplace=True) if save_path is not None: os.makedirs(save_path, exist_ok=True) - df_kd.to_csv(os.path.join(save_path, 'kd_efficiency.txt'), sep='\t') + df_kd.to_csv(os.path.join(save_path, "kd_efficiency.txt"), sep="\t") return df_kd + # Function to get significant programs for each gene across all days -def get_significant_programs(Perturbation_path, Sample = ["D0", "sample_D1","sample_D2","sample_D3"], adj_pval_threshold=0.05, effect_size='log2FC'): +def get_significant_programs( + Perturbation_path, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + adj_pval_threshold=0.05, + effect_size="log2FC", +): + + print("Compute significant programs") - print('Compute significant programs') - significant_programs = {} - + for samp in Sample: # Read perturbation results for each day df = pd.read_csv(f"{Perturbation_path}_{samp}.txt", sep="\t") - + # Filter for significant associations - significant = df[df['adj_pval'] < adj_pval_threshold] - + significant = df[df["adj_pval"] < adj_pval_threshold] + # Group by target_name and collect significant programs - for target in significant['target_name'].unique(): - target_data = significant[significant['target_name'] == target] - programs = target_data['program_name'].tolist() - + for target in significant["target_name"].unique(): + target_data = significant[significant["target_name"] == target] + programs = target_data["program_name"].tolist() + if target not in significant_programs: significant_programs[target] = {} - + significant_programs[target][samp] = programs - + return significant_programs -# get signifcant programs -def get_significant_programs_df(Perturbation_path, Sample = ["D0", "sample_D1","sample_D2","sample_D3"], adj_pval_threshold=0.05, effect_size='log2FC'): - print('Compile significant programs') +# get signifcant programs +def get_significant_programs_df( + Perturbation_path, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + adj_pval_threshold=0.05, + effect_size="log2FC", +): + + print("Compile significant programs") # Get significant programs - significant_programs = get_significant_programs(Perturbation_path, Sample, adj_pval_threshold, effect_size=effect_size) + significant_programs = get_significant_programs( + Perturbation_path, Sample, adj_pval_threshold, effect_size=effect_size + ) # Convert to DataFrame format for easier viewing sig_prog_data = [] for gene, days_data in significant_programs.items(): - row = {'target_name': gene} + row = {"target_name": gene} for samp in Sample: programs = days_data.get(samp, []) - row[f'significant programs {samp}'] = ', '.join(map(str, programs)) if programs else '' + row[f"significant programs {samp}"] = ( + ", ".join(map(str, programs)) if programs else "" + ) sig_prog_data.append(row) df_significant_programs = pd.DataFrame(sig_prog_data) - df_significant_programs = df_significant_programs.set_index('target_name') + df_significant_programs = df_significant_programs.set_index("target_name") for samp in Sample: - df_significant_programs[f'# programs {samp}'] = df_significant_programs[f'significant programs {samp}'].apply( - lambda x: str(len(x.split(','))) if x and x != '' else 0 - ) - + df_significant_programs[f"# programs {samp}"] = df_significant_programs[ + f"significant programs {samp}" + ].apply(lambda x: str(len(x.split(","))) if x and x != "" else 0) return df_significant_programs -# Get top correlation terms -def get_correlation_df(perturbation_path, Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], top_n=5, save_path = None, effect_size='log2FC'): + +# Get top correlation terms +def get_correlation_df( + perturbation_path, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + top_n=5, + save_path=None, + effect_size="log2FC", +): correlation_results_all_days = {} @@ -696,8 +852,10 @@ def get_correlation_df(perturbation_path, Sample=["D0", "sample_D1", "sample_D2" # Read and pivot directly. Real perturb files (U-test, CRT) have # target_name as the first column, so we must NOT use index_col=0 — # otherwise pivot_table can't find 'target_name'. - df = pd.read_csv(perturb_path, sep='\t') - pivot_df = df.pivot_table(index='target_name', columns='program_name', values=effect_size) + df = pd.read_csv(perturb_path, sep="\t") + pivot_df = df.pivot_table( + index="target_name", columns="program_name", values=effect_size + ) # Compute correlation matrix (genes x genes) — vectorized corr_matrix = pivot_df.T.corr() @@ -707,7 +865,11 @@ def get_correlation_df(perturbation_path, Sample=["D0", "sample_D1", "sample_D2" if save_path is not None: corr_matrix.to_csv(f"{save_path}/corr_gene_matrix_{day}.txt", sep="\t") - corr_matrix.to_csv(f"{save_path}/corr_gene_matrix_{day}.txt.gz", sep="\t", compression="gzip") + corr_matrix.to_csv( + f"{save_path}/corr_gene_matrix_{day}.txt.gz", + sep="\t", + compression="gzip", + ) genes = corr_matrix.index.values vals = corr_matrix.values # numpy array @@ -715,183 +877,313 @@ def get_correlation_df(perturbation_path, Sample=["D0", "sample_D1", "sample_D2" # Use argpartition for O(n) top/bottom selection instead of full sort rows = [] for i in range(len(genes)): - - # for each row, remove NaN and get valid genes + # for each row, remove NaN and get valid genes row = vals[i] valid = ~np.isnan(row) valid_genes = genes[valid] # Top positive - top_idx = np.argpartition(row[valid], -(top_n-1))[-top_n:] # get top 5 value's index (-1 for accessing k-th value) - top_idx = top_idx[np.argsort(row[valid][top_idx])[::-1]] # get top 5 values and sort them in right otder + top_idx = np.argpartition(row[valid], -(top_n - 1))[ + -top_n: + ] # get top 5 value's index (-1 for accessing k-th value) + top_idx = top_idx[ + np.argsort(row[valid][top_idx])[::-1] + ] # get top 5 values and sort them in right otder # Bottom negative - bot_idx = np.argpartition(row[valid], (top_n-1))[:top_n] # get bottom 5 value's index (-1 for accessing k-th value) - bot_idx = bot_idx[np.argsort(row[valid][bot_idx])] # get bottom 5 values and sort them in right otder - - rows.append({ - 'target_name': genes[i], - f'top 5 pos correls targets (program log2fc) {day}': '; '.join(str(g) for g in valid_genes[top_idx]), - f'top 5 pos correls (program log2fc) {day}': '; '.join(f"{v:.3f}" for v in row[valid][top_idx]), - f'top 5 neg correls targets (program log2fc) {day}': '; '.join(str(g) for g in valid_genes[bot_idx]), - f'top 5 neg correls (program log2fc) {day}': '; '.join(f"{v:.3f}" for v in row[valid][bot_idx]), - }) - - correlation_results_all_days[day] = pd.DataFrame(rows).set_index('target_name') - + bot_idx = np.argpartition(row[valid], (top_n - 1))[ + :top_n + ] # get bottom 5 value's index (-1 for accessing k-th value) + bot_idx = bot_idx[ + np.argsort(row[valid][bot_idx]) + ] # get bottom 5 values and sort them in right otder + + rows.append( + { + "target_name": genes[i], + f"top 5 pos correls targets (program log2fc) {day}": "; ".join( + str(g) for g in valid_genes[top_idx] + ), + f"top 5 pos correls (program log2fc) {day}": "; ".join( + f"{v:.3f}" for v in row[valid][top_idx] + ), + f"top 5 neg correls targets (program log2fc) {day}": "; ".join( + str(g) for g in valid_genes[bot_idx] + ), + f"top 5 neg correls (program log2fc) {day}": "; ".join( + f"{v:.3f}" for v in row[valid][bot_idx] + ), + } + ) + correlation_results_all_days[day] = pd.DataFrame(rows).set_index("target_name") return pd.concat(correlation_results_all_days.values(), axis=1) -#-------------- helper methods starget summary-------------- - -# final function to compile target Summary sheet -def Compile_Target_Summary_sheet(mdata, perturbation_path, Sample = ["D0", "sample_D1","sample_D2","sample_D3"], adj_pval_threshold= 0.05, -top_n=5, T= 0.5 , categorical_key = "sample", prog_key = 'cNMF', data_key = 'rna', guide_targets_key = "guide_targets", guide_assignment_key = "guide_assignment", save_path = None, effect_size='log2FC', -gene_names_key=None, control_target_name='non-targeting'): - - print('Generate target summary sheet') - - - df_mean_expr_per_day = get_guide_mean_expr_per_day(mdata, categorical_key=categorical_key, prog_key=prog_key, data_key=data_key, - guide_targets_key=guide_targets_key) - df_guide_days = get_guide_cells_per_days(mdata, categorical_key=categorical_key, guide_assignment_key=guide_assignment_key, prog_key=prog_key, - data_key=data_key, guide_targets_key=guide_targets_key) - df_significant_program = get_significant_programs_df(perturbation_path ,Sample=Sample, adj_pval_threshold=adj_pval_threshold, effect_size=effect_size) - df_specificity_program = get_specificity_program(perturbation_path, Sample=Sample, T = T, save_path = save_path, effect_size=effect_size) # save_path for saving specificity scores - df_correlation = get_correlation_df(perturbation_path ,Sample=Sample, top_n=top_n, save_path = save_path, effect_size=effect_size) - df_kd = compute_kd_efficiency(mdata, categorical_key=categorical_key, prog_key=prog_key, data_key=data_key, - gene_names_key=gene_names_key, guide_targets_key=guide_targets_key, - guide_assignment_key=guide_assignment_key, control_target_name=control_target_name, save_path=save_path) +# -------------- helper methods starget summary-------------- + + +# final function to compile target Summary sheet +def Compile_Target_Summary_sheet( + mdata, + perturbation_path, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + adj_pval_threshold=0.05, + top_n=5, + T=0.5, + categorical_key="sample", + prog_key="cNMF", + data_key="rna", + guide_targets_key="guide_targets", + guide_assignment_key="guide_assignment", + save_path=None, + effect_size="log2FC", + gene_names_key=None, + control_target_name="non-targeting", +): + + print("Generate target summary sheet") + + df_mean_expr_per_day = get_guide_mean_expr_per_day( + mdata, + categorical_key=categorical_key, + prog_key=prog_key, + data_key=data_key, + guide_targets_key=guide_targets_key, + ) + df_guide_days = get_guide_cells_per_days( + mdata, + categorical_key=categorical_key, + guide_assignment_key=guide_assignment_key, + prog_key=prog_key, + data_key=data_key, + guide_targets_key=guide_targets_key, + ) + df_significant_program = get_significant_programs_df( + perturbation_path, + Sample=Sample, + adj_pval_threshold=adj_pval_threshold, + effect_size=effect_size, + ) + df_specificity_program = get_specificity_program( + perturbation_path, + Sample=Sample, + T=T, + save_path=save_path, + effect_size=effect_size, + ) # save_path for saving specificity scores + df_correlation = get_correlation_df( + perturbation_path, + Sample=Sample, + top_n=top_n, + save_path=save_path, + effect_size=effect_size, + ) + df_kd = compute_kd_efficiency( + mdata, + categorical_key=categorical_key, + prog_key=prog_key, + data_key=data_key, + gene_names_key=gene_names_key, + guide_targets_key=guide_targets_key, + guide_assignment_key=guide_assignment_key, + control_target_name=control_target_name, + save_path=save_path, + ) - dfs = [df_mean_expr_per_day, df_guide_days, df_significant_program, df_specificity_program, df_correlation, df_kd] + dfs = [ + df_mean_expr_per_day, + df_guide_days, + df_significant_program, + df_specificity_program, + df_correlation, + df_kd, + ] final_merged_df = reduce( - lambda left, right: left.merge(right, left_index=True, right_index=True, how='outer'), - dfs - ).fillna('') + lambda left, right: left.merge( + right, left_index=True, right_index=True, how="outer" + ), + dfs, + ).fillna("") return final_merged_df +# load summary +# -------------- helper methods summary-------------- -# load summary +# get simply items ready in the summary sheet +def simple_Summary_cols( + df, + k, + df_GO, + df_Perturbation, + df_Program_loading, + df_Explained_Variance=None, + specicicity_path=None, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + non_tagerting_key=None, + effect_size="log2FC", + adjusted_pval_key="Adjusted P-value", +): + + programs = list(df.index) # use actual program indices from df, not range(k) + + # Fix #2: Cast programs to str for consistent comparisons with program_name columns + programs_str = [str(p) for p in programs] + + # create GO summary col + if df_GO is not None: + df_GO_enriched = df_GO.loc[df_GO[adjusted_pval_key] <= 0.05] + df["Total Enriched GO Terms"] = [ + df_GO_enriched[df_GO_enriched["program_name"].astype(str) == i].shape[0] + for i in programs_str + ] + + # remove non-targeting off the list of perturbed genes + if df_Perturbation is not None: + if non_tagerting_key is not None: + df_Perturbation = df_Perturbation[ + ~df_Perturbation["target_name"].isin(non_tagerting_key) + ] + + # Fix #2: Cast program_name to str for consistent comparisons + df_Perturbation = df_Perturbation.copy() + df_Perturbation["program_name"] = df_Perturbation["program_name"].astype(str) + + # compute # regulator per condition for + and - + conditions = (df_Perturbation["Sample"]).unique() + for condition in conditions: + df_Perturbation_ = df_Perturbation[df_Perturbation["Sample"] == condition] + + # create perturbation program summary col + df_Perturbation_enriched = df_Perturbation_.loc[ + df_Perturbation_["adj_pval"] <= 0.05 + ] + df_Perturbation_positive = df_Perturbation_enriched.loc[ + df_Perturbation_enriched[effect_size] > 0 + ] + df_Perturbation_negative = df_Perturbation_enriched.loc[ + df_Perturbation_enriched[effect_size] < 0 + ] + + df[f"Significant regulators with positive effect {condition}"] = [ + df_Perturbation_positive[ + df_Perturbation_positive["program_name"] == i + ].shape[0] + for i in programs_str + ] + df[f"Significant regulators with negative effect {condition}"] = [ + df_Perturbation_negative[ + df_Perturbation_negative["program_name"] == i + ].shape[0] + for i in programs_str + ] + + # create perturbation gene summary col + for condition in conditions: + df_Perturbation_enriched = df_Perturbation.loc[ + df_Perturbation["adj_pval"] <= 0.05 + ] + df_Perturbation_D = df_Perturbation_enriched.loc[ + df_Perturbation_enriched["Sample"] == condition + ] + + targets_list = [] + for i in programs_str: + matching = df_Perturbation_D.loc[df_Perturbation_D["program_name"] == i] + unique_indices = matching["target_name"].unique() + joined = ";".join([str(x) for x in unique_indices]) + targets_list.append(joined) + + df[f"sigfdr0.05_targets_sorted_abslog2fcd_{condition}"] = targets_list + + # Fix #5: Moved specificity block outside df_Perturbation guard + if specicicity_path is not None: + if df_Perturbation is not None: + conditions = df_Perturbation["Sample"].unique() + else: + conditions = Sample + + for condition in conditions: + specificity_score_df = pd.read_csv( + os.path.join(specicicity_path, f"specificity_score_{condition}.txt"), + sep="\t", + index_col=0, + ) + + col = {} + for prog in specificity_score_df.columns: + if df_Perturbation is not None: + df_Perturbation_cond = df_Perturbation[ + df_Perturbation["Sample"] == condition + ] + prog_perturb = df_Perturbation_cond[ + df_Perturbation_cond["program_name"] == (prog) + ] + sig_genes = prog_perturb[prog_perturb["adj_pval"] < 0.1][ + "target_name" + ].values + filtered_scores = specificity_score_df.loc[ + specificity_score_df.index.isin(sig_genes), prog + ] + else: + filtered_scores = specificity_score_df[prog].dropna() + + top5 = filtered_scores.nlargest(5).index.tolist() + # Fix #3: Cast key to match df.index dtype + try: + col_key = type(programs[0])(prog) + except (ValueError, TypeError): + col_key = prog + col[col_key] = ";".join(str(g) for g in top5) + # Fix #4: Typo FRD -> FDR + df[f"Top 5 specific regulators (FDR<0.1) {condition}"] = col + + # Fix #1: Use len(programs) instead of k + df["Total Enriched Enhancer Motifs"] = [""] * len(programs) + df["Total Enriched Promoter Motifs"] = [""] * len(programs) + + # create top gene summary col + if df_Program_loading is not None: + df["top10_loaded_genes"] = [ + ";".join(df_Program_loading.loc[str(i)][:10]) for i in programs + ] + df["top30_loaded_genes"] = [ + ";".join(df_Program_loading.loc[str(i)][:30]) for i in programs + ] + + # create explained variance summary col + if df_Explained_Variance is not None: + df["variance_explained"] = df_Explained_Variance["VarianceExplained"] -#-------------- helper methods summary-------------- -# get simply items ready in the summary sheet -def simple_Summary_cols(df, k, df_GO, df_Perturbation, df_Program_loading, df_Explained_Variance = None, specicicity_path = None, Sample = ["D0", - "sample_D1","sample_D2","sample_D3" ], - non_tagerting_key = None, effect_size='log2FC', adjusted_pval_key='Adjusted P-value'): - - programs = list(df.index) # use actual program indices from df, not range(k) - - # Fix #2: Cast programs to str for consistent comparisons with program_name columns - programs_str = [str(p) for p in programs] - - # create GO summary col - if df_GO is not None: - df_GO_enriched = df_GO.loc[df_GO[adjusted_pval_key]<=0.05] - df['Total Enriched GO Terms'] = [df_GO_enriched[df_GO_enriched['program_name'].astype(str)==i].shape[0] for i in programs_str] - - # remove non-targeting off the list of perturbed genes - if df_Perturbation is not None: - if non_tagerting_key is not None: - df_Perturbation = df_Perturbation[~df_Perturbation['target_name'].isin(non_tagerting_key)] - - # Fix #2: Cast program_name to str for consistent comparisons - df_Perturbation = df_Perturbation.copy() - df_Perturbation['program_name'] = df_Perturbation['program_name'].astype(str) - - # compute # regulator per condition for + and - - conditions = (df_Perturbation['Sample']).unique() - for condition in conditions: - df_Perturbation_ = df_Perturbation[df_Perturbation['Sample'] == condition] - - # create perturbation program summary col - df_Perturbation_enriched = df_Perturbation_.loc[df_Perturbation_['adj_pval']<=0.05] - df_Perturbation_positive = df_Perturbation_enriched.loc[df_Perturbation_enriched[effect_size] > 0 ] - df_Perturbation_negative = df_Perturbation_enriched.loc[df_Perturbation_enriched[effect_size] < 0 ] - - df[f'Significant regulators with positive effect {condition}'] = [df_Perturbation_positive[df_Perturbation_positive['program_name']==i].shape[0] for i in programs_str] - df[f'Significant regulators with negative effect {condition}'] = [df_Perturbation_negative[df_Perturbation_negative['program_name']==i].shape[0] for i in programs_str] - - # create perturbation gene summary col - for condition in conditions: - df_Perturbation_enriched = df_Perturbation.loc[df_Perturbation['adj_pval']<=0.05] - df_Perturbation_D = df_Perturbation_enriched.loc[df_Perturbation_enriched['Sample'] == condition] - - targets_list = [] - for i in programs_str: - matching = df_Perturbation_D.loc[df_Perturbation_D['program_name'] == i] - unique_indices = matching["target_name"].unique() - joined = ';'.join([str(x) for x in unique_indices]) - targets_list.append(joined) - - df[f'sigfdr0.05_targets_sorted_abslog2fcd_{condition}'] = targets_list - - # Fix #5: Moved specificity block outside df_Perturbation guard - if specicicity_path is not None: - if df_Perturbation is not None: - conditions = df_Perturbation['Sample'].unique() - else: - conditions = Sample - - for condition in conditions: - specificity_score_df = pd.read_csv(os.path.join(specicicity_path, f"specificity_score_{condition}.txt"), sep="\t", index_col=0) - - col = {} - for prog in specificity_score_df.columns: - if df_Perturbation is not None: - df_Perturbation_cond = df_Perturbation[df_Perturbation['Sample'] == condition] - prog_perturb = df_Perturbation_cond[df_Perturbation_cond['program_name'] == (prog)] - sig_genes = prog_perturb[prog_perturb['adj_pval'] < 0.1]['target_name'].values - filtered_scores = specificity_score_df.loc[specificity_score_df.index.isin(sig_genes), prog] - else: - filtered_scores = specificity_score_df[prog].dropna() - - top5 = filtered_scores.nlargest(5).index.tolist() - # Fix #3: Cast key to match df.index dtype - try: - col_key = type(programs[0])(prog) - except (ValueError, TypeError): - col_key = prog - col[col_key] = ";".join(str(g) for g in top5) - # Fix #4: Typo FRD -> FDR - df[f"Top 5 specific regulators (FDR<0.1) {condition}"] = col - - # Fix #1: Use len(programs) instead of k - df['Total Enriched Enhancer Motifs'] = [''] * len(programs) - df['Total Enriched Promoter Motifs'] = [''] * len(programs) - - # create top gene summary col - if df_Program_loading is not None: - df['top10_loaded_genes'] = [';'.join(df_Program_loading.loc[str(i)][:10]) for i in programs] - df['top30_loaded_genes'] = [';'.join(df_Program_loading.loc[str(i)][:30]) for i in programs] - - # create explained variance summary col - if df_Explained_Variance is not None: - df['variance_explained'] = df_Explained_Variance['VarianceExplained'] - # make the program info in summary sheet -def get_program_info_Summary_cols(mdata, categorical_key = "sample"): +def get_program_info_Summary_cols(mdata, categorical_key="sample"): # create cell info col summary - df_cell = pd.DataFrame(data=mdata['cNMF'].X, index=mdata['cNMF'].obs_names, columns=mdata['cNMF'].var_names) + df_cell = pd.DataFrame( + data=mdata["cNMF"].X, + index=mdata["cNMF"].obs_names, + columns=mdata["cNMF"].var_names, + ) results = [] # use actual program names from mdata var_names - programs = [int(v) for v in mdata['cNMF'].var_names] + programs = [int(v) for v in mdata["cNMF"].var_names] # Loop through all programs using label-based access for prog in programs: - df_cell_program = pd.DataFrame({ - "expression": df_cell[str(prog)], - "cell_type": mdata['rna'].obs[categorical_key].values - }) + df_cell_program = pd.DataFrame( + { + "expression": df_cell[str(prog)], + "cell_type": mdata["rna"].obs[categorical_key].values, + } + ) df_mean = df_cell_program.groupby("cell_type")["expression"].mean() df_frac = df_cell_program.groupby("cell_type")["expression"].apply( @@ -899,81 +1191,118 @@ def get_program_info_Summary_cols(mdata, categorical_key = "sample"): ) # Store as row with program as index - results.append({ - 'program_name': prog, - **{f'Mean program score {ct}': df_mean[ct] for ct in df_mean.index}, - **{f'Fra cells above mean program score {ct}': df_frac[ct] for ct in df_frac.index} - }) + results.append( + { + "program_name": prog, + **{f"Mean program score {ct}": df_mean[ct] for ct in df_mean.index}, + **{ + f"Fra cells above mean program score {ct}": df_frac[ct] + for ct in df_frac.index + }, + } + ) - df_cell_info = pd.DataFrame(results).set_index('program_name') + df_cell_info = pd.DataFrame(results).set_index("program_name") return df_cell_info + # make top terms for the summray sheet -def get_top_terms_Summary_cols(df_GO, df_Geneset, adjusted_pval_key='Adjusted P-value'): +def get_top_terms_Summary_cols(df_GO, df_Geneset, adjusted_pval_key="Adjusted P-value"): top_GO = None top_Geneset = None if df_GO is not None: df_GO_filtered = df_GO[df_GO[adjusted_pval_key] < 0.05] - top_GO = df_GO_filtered.groupby('program_name').apply( - lambda x: ';'.join(x.sort_values(adjusted_pval_key).index[:10]) + top_GO = df_GO_filtered.groupby("program_name").apply( + lambda x: ";".join(x.sort_values(adjusted_pval_key).index[:10]) ) if df_Geneset is not None: df_Geneset_filtered = df_Geneset[df_Geneset[adjusted_pval_key] < 0.05] - top_Geneset = df_Geneset_filtered.groupby('program_name').apply( - lambda x: ';'.join(x.sort_values(adjusted_pval_key).index[:10]) - ) + top_Geneset = df_Geneset_filtered.groupby("program_name").apply( + lambda x: ";".join(x.sort_values(adjusted_pval_key).index[:10]) + ) # When both inputs are None, pd.DataFrame({col: None, col: None}) raises # "If using all scalar values, you must pass an index". Return an empty # frame with the expected columns instead. if top_GO is None and top_Geneset is None: - return pd.DataFrame(columns=['top10_enriched_genesets', 'top10_enriched_go_terms']) + return pd.DataFrame( + columns=["top10_enriched_genesets", "top10_enriched_go_terms"] + ) - top_terms = pd.DataFrame({ - 'top10_enriched_genesets': top_Geneset, - 'top10_enriched_go_terms': top_GO - }) + top_terms = pd.DataFrame( + {"top10_enriched_genesets": top_Geneset, "top10_enriched_go_terms": top_GO} + ) return top_terms -#-------------- helper methods summary-------------- -# compile summry sheet -def Compile_Summary_sheet(mdata, df_GO, df_Geneset, df_Perturbation, df_Program_loading, df_Explained_Variance , specicicity_path = None, Sample = ["D0", "sample_D1","sample_D2","sample_D3"], -categorical_key = "sample",non_tagerting_key=None, effect_size='log2FC', adjusted_pval_key='Adjusted P-value'): +# -------------- helper methods summary-------------- - print('complie summary sheet') - # set program # from mdata (works even if df_GO is None) - k = mdata['cNMF'].varm["loadings"].shape[0] - programs = sorted(int(v) for v in mdata['cNMF'].var_names) +# compile summry sheet +def Compile_Summary_sheet( + mdata, + df_GO, + df_Geneset, + df_Perturbation, + df_Program_loading, + df_Explained_Variance, + specicicity_path=None, + Sample=["D0", "sample_D1", "sample_D2", "sample_D3"], + categorical_key="sample", + non_tagerting_key=None, + effect_size="log2FC", + adjusted_pval_key="Adjusted P-value", +): + + print("complie summary sheet") - df = pd.DataFrame({ - 'manual_annotation_label': [''] * k, - 'manual_timepoint': [''] * k, - 'Notes': [''] * k, - 'Automatic Timepoint': [''] * k }, index=pd.Index(programs, name='program_name')) + # set program # from mdata (works even if df_GO is None) + k = mdata["cNMF"].varm["loadings"].shape[0] + programs = sorted(int(v) for v in mdata["cNMF"].var_names) + + df = pd.DataFrame( + { + "manual_annotation_label": [""] * k, + "manual_timepoint": [""] * k, + "Notes": [""] * k, + "Automatic Timepoint": [""] * k, + }, + index=pd.Index(programs, name="program_name"), + ) - simple_Summary_cols(df, k, df_GO, df_Perturbation, df_Program_loading, df_Explained_Variance, specicicity_path = specicicity_path, Sample = Sample, non_tagerting_key=non_tagerting_key, effect_size=effect_size, adjusted_pval_key=adjusted_pval_key) - df_cell_info_cols = get_program_info_Summary_cols(mdata,categorical_key) - df_top_terms = get_top_terms_Summary_cols(df_GO, df_Geneset, adjusted_pval_key=adjusted_pval_key) + simple_Summary_cols( + df, + k, + df_GO, + df_Perturbation, + df_Program_loading, + df_Explained_Variance, + specicicity_path=specicicity_path, + Sample=Sample, + non_tagerting_key=non_tagerting_key, + effect_size=effect_size, + adjusted_pval_key=adjusted_pval_key, + ) + df_cell_info_cols = get_program_info_Summary_cols(mdata, categorical_key) + df_top_terms = get_top_terms_Summary_cols( + df_GO, df_Geneset, adjusted_pval_key=adjusted_pval_key + ) # refill automatic time point - col_names = [f'Mean program score {samp}' for samp in Sample] - col_mapping = {f'Mean program score {samp}': samp for samp in Sample} # Create mapping from column names to condition A, B, C + col_names = [f"Mean program score {samp}" for samp in Sample] + col_mapping = { + f"Mean program score {samp}": samp for samp in Sample + } # Create mapping from column names to condition A, B, C df_mean = df_cell_info_cols[col_names] - - df['Automatic Timepoint'] = df_mean.idxmax(axis=1).map(col_mapping) + + df["Automatic Timepoint"] = df_mean.idxmax(axis=1).map(col_mapping) merged_df = pd.merge( - df, - df_cell_info_cols, - left_index=True, - right_index=True, - how='outer' + df, df_cell_info_cols, left_index=True, right_index=True, how="outer" ) # df_top_terms comes from a groupby('program_name') on string-typed @@ -988,24 +1317,14 @@ def Compile_Summary_sheet(mdata, df_GO, df_Geneset, df_Perturbation, df_Program_ df_top_terms.index = df_top_terms.index.astype(str) final_merged_df = pd.merge( - merged_df, - df_top_terms, - left_index=True, - right_index=True, - how='outer' + merged_df, df_top_terms, left_index=True, right_index=True, how="outer" ) # move top 10 loaded gene left of notes - cols = final_merged_df.columns.tolist() - cols.remove('top10_loaded_genes') - i = cols.index('Notes') - cols.insert(i, 'top10_loaded_genes') + cols = final_merged_df.columns.tolist() + cols.remove("top10_loaded_genes") + i = cols.index("Notes") + cols.insert(i, "top10_loaded_genes") final_merged_df = final_merged_df[cols] return final_merged_df - - - - - - diff --git a/src/Stage3_Interpretation/B_Summarization/src/__init__.py b/src/Stage3_Interpretation/B_Summarization/src/__init__.py index 303f049..eacccb5 100644 --- a/src/Stage3_Interpretation/B_Summarization/src/__init__.py +++ b/src/Stage3_Interpretation/B_Summarization/src/__init__.py @@ -1,5 +1,17 @@ # Excel compilation helpers -from .Compile_excel_sheet import compile_Program_loading_score_sheet_long, compile_Program_loading_score_sheet_flat, Compile_GO_sheet, \ - Compile_Geneset_sheet, Compile_Trait_sheet, Compile_Perturbation_sheet, Compile_Association_sheet, Compile_Explained_variance, \ - Compile_Target_Summary_sheet, Compile_Summary_sheet, load_simple_sheets, add_specificity_scores_file, check_program_name_match, \ - compute_kd_efficiency +from .Compile_excel_sheet import ( + compile_Program_loading_score_sheet_long, + compile_Program_loading_score_sheet_flat, + Compile_GO_sheet, + Compile_Geneset_sheet, + Compile_Trait_sheet, + Compile_Perturbation_sheet, + Compile_Association_sheet, + Compile_Explained_variance, + Compile_Target_Summary_sheet, + Compile_Summary_sheet, + load_simple_sheets, + add_specificity_scores_file, + check_program_name_match, + compute_kd_efficiency, +) diff --git a/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py b/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py index 228a952..323fbe9 100644 --- a/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py +++ b/src/Stage3_Interpretation/C_Annotation/0.Prepare/Extract_program_information.py @@ -10,6 +10,7 @@ import pandas as pd import scipy.sparse as sp + # loaders def _check_file(path): path = Path(path) @@ -21,14 +22,20 @@ def _check_file(path): def load_gene_spectra(prog_adata, loadings_key="loadings", gene_names_key="var_names"): """programs x genes spectra from varm[loadings_key], gene names from uns[gene_names_key].""" if loadings_key not in prog_adata.varm: - raise KeyError(f"varm['{loadings_key}'] not found (have {list(prog_adata.varm)})") + raise KeyError( + f"varm['{loadings_key}'] not found (have {list(prog_adata.varm)})" + ) if gene_names_key not in prog_adata.uns: - raise KeyError(f"uns['{gene_names_key}'] not found (have {list(prog_adata.uns)})") + raise KeyError( + f"uns['{gene_names_key}'] not found (have {list(prog_adata.uns)})" + ) loadings = np.asarray(prog_adata.varm[loadings_key]) genes = np.asarray(prog_adata.uns[gene_names_key]).astype(str) if loadings.shape[1] != len(genes): - raise ValueError(f"varm['{loadings_key}'] has {loadings.shape[1]} genes but " - f"uns['{gene_names_key}'] has {len(genes)} names") + raise ValueError( + f"varm['{loadings_key}'] has {loadings.shape[1]} genes but " + f"uns['{gene_names_key}'] has {len(genes)} names" + ) return pd.DataFrame(loadings, index=prog_adata.var_names.astype(int), columns=genes) @@ -36,19 +43,27 @@ def load_cell_usage(prog_adata): """cells x programs usage matrix.""" X = prog_adata.X X = X.toarray() if sp.issparse(X) else np.asarray(X) - return pd.DataFrame(X, index=prog_adata.obs_names, columns=prog_adata.var_names.astype(int)) + return pd.DataFrame( + X, index=prog_adata.obs_names, columns=prog_adata.var_names.astype(int) + ) def load_condition(prog_adata, categorical_key="sample"): """Cell condition labels and their sorted unique levels (cells with a missing label are kept as NaN, so groupby leaves them out of the per-condition stats).""" if categorical_key not in prog_adata.obs: - raise KeyError(f"obs['{categorical_key}'] not found (have {list(prog_adata.obs.columns)})") + raise KeyError( + f"obs['{categorical_key}'] not found (have {list(prog_adata.obs.columns)})" + ) raw = prog_adata.obs[categorical_key] n_missing = int(raw.isna().sum()) if n_missing: - print(f" [warn] {n_missing:,} cells have no obs['{categorical_key}'] label; left out of condition stats") - labels = raw.astype(object).where(raw.notna()).map(lambda x: x if pd.isna(x) else str(x)) + print( + f" [warn] {n_missing:,} cells have no obs['{categorical_key}'] label; left out of condition stats" + ) + labels = ( + raw.astype(object).where(raw.notna()).map(lambda x: x if pd.isna(x) else str(x)) + ) return labels, sorted(str(c) for c in labels.dropna().unique()) @@ -74,13 +89,17 @@ def perturbation_paths(path_base, conditions): paths = {cond: base.parent / f"{base.name}_{cond}.txt" for cond in conditions} missing = [str(p) for p in paths.values() if not p.exists()] if missing: - raise FileNotFoundError(f"No perturbation file for {len(missing)} of {len(conditions)} " - f"conditions: {missing}") + raise FileNotFoundError( + f"No perturbation file for {len(missing)} of {len(conditions)} " + f"conditions: {missing}" + ) found = sorted(base.parent.glob(f"{base.name}_*.txt")) if len(found) != len(conditions): extra = sorted(set(map(str, found)) - set(map(str, paths.values()))) - raise ValueError(f"{len(found)} perturbation files match '{base.name}_*.txt' but mdata has " - f"{len(conditions)} conditions {conditions}; unmatched files: {extra}") + raise ValueError( + f"{len(found)} perturbation files match '{base.name}_*.txt' but mdata has " + f"{len(conditions)} conditions {conditions}; unmatched files: {extra}" + ) return paths @@ -109,15 +128,19 @@ def compute_program_condition_stats(usage, cond_labels): "per_condition": {cond: {"pct_cells_expressed", "mean_usage"}}}}. """ grouped = usage.groupby(cond_labels.values) - mean = grouped.mean() # compuate program mean usage in condition - pct = (usage > 0).groupby(cond_labels.values).mean() * 100 # compuate program used in % cell in condition + mean = grouped.mean() # compuate program mean usage in condition + pct = (usage > 0).groupby( + cond_labels.values + ).mean() * 100 # compuate program used in % cell in condition stats = {} for pid in usage.columns: stats[pid] = { "top_condition": str(mean[pid].idxmax()), "per_condition": { - str(cond): {"pct_cells_expressed": round(float(pct.at[cond, pid]), 1), - "mean_usage": float(f"{mean.at[cond, pid]:.3g}")} + str(cond): { + "pct_cells_expressed": round(float(pct.at[cond, pid]), 1), + "mean_usage": float(f"{mean.at[cond, pid]:.3g}"), + } for cond in mean.index }, } @@ -142,22 +165,33 @@ def load_top_GO(go_df, pid, n=10, fdr=0.05): """Top-N significant GO terms (by adjusted p, then Combined Score), with any parenthesized text such as the GO ID removed.""" sub = go_df[(go_df["program_name"] == pid) & (go_df["Adjusted P-value"] < fdr)] - sub = sub.sort_values(["Adjusted P-value", "Combined Score"], - ascending=[True, False], kind="mergesort") + sub = sub.sort_values( + ["Adjusted P-value", "Combined Score"], + ascending=[True, False], + kind="mergesort", + ) return [re.sub(r"\s*\([^)]*\)", "", t).strip() for t in sub["Term"].head(n)] def load_regulator(assoc_df, pid, n=6, fdr=0.05, lfc_col="log2FC"): """Top-N significant regulators (adj_pval < fdr), ranked by adj_pval then |lfc_col|.""" if lfc_col not in assoc_df.columns: - raise KeyError(f"log2FC column '{lfc_col}' not found (have {list(assoc_df.columns)})") - sub = assoc_df[(assoc_df["program_name"] == pid) & (assoc_df["adj_pval"] < fdr)].copy() + raise KeyError( + f"log2FC column '{lfc_col}' not found (have {list(assoc_df.columns)})" + ) + sub = assoc_df[ + (assoc_df["program_name"] == pid) & (assoc_df["adj_pval"] < fdr) + ].copy() sub["_abs_lfc"] = sub[lfc_col].abs() - sub = sub.sort_values(["adj_pval", "_abs_lfc"], ascending=[True, False], kind="mergesort") + sub = sub.sort_values( + ["adj_pval", "_abs_lfc"], ascending=[True, False], kind="mergesort" + ) return [ - {"gene": str(r["target_name"]), - "log2fc": round(float(r[lfc_col]), 2), - "adj_pval": float(f"{r['adj_pval']:.3g}")} + { + "gene": str(r["target_name"]), + "log2fc": round(float(r[lfc_col]), 2), + "adj_pval": float(f"{r['adj_pval']:.3g}"), + } for _, r in sub.head(n).iterrows() ] @@ -172,27 +206,43 @@ def load_regulator_gene_pairs(program_genes, distinctive_genes, regulators): reg_stats = {} for cond, regs in regulators.items(): for r in regs: - reg_stats.setdefault(r["gene"], {})[cond] = {"log2fc": r["log2fc"], "adj_pval": r["adj_pval"]} - genes = [(g, "program_gene") for g in program_genes] + [(g, "distinctive_gene") for g in distinctive_genes] + reg_stats.setdefault(r["gene"], {})[cond] = { + "log2fc": r["log2fc"], + "adj_pval": r["adj_pval"], + } + genes = [(g, "program_gene") for g in program_genes] + [ + (g, "distinctive_gene") for g in distinctive_genes + ] pairs = {} for reg, stats in reg_stats.items(): for gene, category in genes: if gene != reg: - pairs[f"{gene}-{reg}"] = {"gene": gene, "regulator": reg, "gene_category": category, - "regulator_stats": stats} - return {"n_pairs": len(pairs), "n_regulators": len(reg_stats), "n_genes": len(genes), "pairs": pairs} + pairs[f"{gene}-{reg}"] = { + "gene": gene, + "regulator": reg, + "gene_category": category, + "regulator_stats": stats, + } + return { + "n_pairs": len(pairs), + "n_regulators": len(reg_stats), + "n_genes": len(genes), + "pairs": pairs, + } # research brief -def build_research_brief(label, organism, cell_type, conditions, has_regulators, - top_condition=None): +def build_research_brief( + label, organism, cell_type, conditions, has_regulators, top_condition=None +): """Short instruction referencing the bundle's fields (adapted from GPI bundle.py).""" subject = cell_type or organism or "cell" role = f"{cell_type} biologist" if cell_type else "cell biologist" reg_clause = ( " and the genes in `perturbation_regulators` (research these the same way as the " "program genes)" - if has_regulators else "" + if has_regulators + else "" ) lines = [ f"# Program {label} — {organism} {subject} gene program", @@ -233,15 +283,18 @@ def build_research_brief(label, organism, cell_type, conditions, has_regulators, def build_bundle(pid, spectra, uniqueness, go_df, assoc_by_cond, cond_stats, args): label = f"P{pid}" program_genes = load_top_gene(spectra, pid, args.top_gene) - distinctive_genes = load_top_unique_gene(uniqueness, pid, args.top_unique_gene, - exclude=program_genes) + distinctive_genes = load_top_unique_gene( + uniqueness, pid, args.top_unique_gene, exclude=program_genes + ) go_terms = load_top_GO(go_df, pid, args.top_GO, args.fdr) if not go_terms: print(f" [warn] {label}: no GO term with adj p < {args.fdr}") regulators = {} for cond, assoc_df in assoc_by_cond.items(): - regs = load_regulator(assoc_df, pid, args.top_regulator, args.fdr, args.log2fc_key) + regs = load_regulator( + assoc_df, pid, args.top_regulator, args.fdr, args.log2fc_key + ) if regs: regulators[cond] = regs @@ -258,45 +311,109 @@ def build_bundle(pid, spectra, uniqueness, go_df, assoc_by_cond, cond_stats, arg } if regulators: bundle["perturbation_regulators"] = regulators - bundle["regulator_gene"] = load_regulator_gene_pairs(program_genes, distinctive_genes, regulators) + bundle["regulator_gene"] = load_regulator_gene_pairs( + program_genes, distinctive_genes, regulators + ) bundle["research_brief"] = build_research_brief( - label, args.organism, args.cell_type, conditions, bool(regulators), - top_condition=cond_stats[pid]["top_condition"]) + label, + args.organism, + args.cell_type, + conditions, + bool(regulators), + top_condition=cond_stats[pid]["top_condition"], + ) return bundle def build_parser(): - p = argparse.ArgumentParser(description="Extract annotation bundles from PerturbNMF outputs.") + p = argparse.ArgumentParser( + description="Extract annotation bundles from PerturbNMF outputs." + ) # IO - p.add_argument("--mdata_path", required=True, help="Path to the cNMF MuData (.h5mu) file.") + p.add_argument( + "--mdata_path", required=True, help="Path to the cNMF MuData (.h5mu) file." + ) p.add_argument("--GO_path", required=True, help="GO_term_enrichment.txt.") - p.add_argument("--perturbation_path_base", required=True, help="Path base of the perturbation association files; '_.txt' is read for each level of --categorical_key.") + p.add_argument( + "--perturbation_path_base", + required=True, + help="Path base of the perturbation association files; '_.txt' is read for each level of --categorical_key.", + ) p.add_argument("--out_dir", required=True, help="Output directory.") # keys - p.add_argument("--data_key", default="rna", help="Key of the gene expression modality in MuData (recorded in meta only).") - p.add_argument("--prog_key", default="cNMF", help="Key of the cNMF program modality in MuData.") - p.add_argument("--loadings_key", default="loadings", help="varm key of the program x gene spectra in the program modality.") - p.add_argument("--gene_name_key", default="var_names", help="uns key of the gene names matching the loadings columns.") - p.add_argument("--categorical_key", default="sample", help="obs key of the condition labels.") - p.add_argument("--log2fc_key", default="log2FC", help="log2FC column in the perturbation files (e.g. 'approx_log2FC' for CRT results).") + p.add_argument( + "--data_key", + default="rna", + help="Key of the gene expression modality in MuData (recorded in meta only).", + ) + p.add_argument( + "--prog_key", default="cNMF", help="Key of the cNMF program modality in MuData." + ) + p.add_argument( + "--loadings_key", + default="loadings", + help="varm key of the program x gene spectra in the program modality.", + ) + p.add_argument( + "--gene_name_key", + default="var_names", + help="uns key of the gene names matching the loadings columns.", + ) + p.add_argument( + "--categorical_key", default="sample", help="obs key of the condition labels." + ) + p.add_argument( + "--log2fc_key", + default="log2FC", + help="log2FC column in the perturbation files (e.g. 'approx_log2FC' for CRT results).", + ) # context info - p.add_argument("--programs", type=int, nargs="+", help="Program ids, space separated (e.g. 1 2 3). Default: all.") - p.add_argument("--cell_type", default="", help="Cell type, e.g. 'aortic endothelial cell'.") + p.add_argument( + "--programs", + type=int, + nargs="+", + help="Program ids, space separated (e.g. 1 2 3). Default: all.", + ) + p.add_argument( + "--cell_type", default="", help="Cell type, e.g. 'aortic endothelial cell'." + ) p.add_argument("--organism", default="human") # select top items of the program - p.add_argument("--top_gene", type=int, default=15, help="Top-N genes by loading score.") - p.add_argument("--top_unique_gene", type=int, default=8, help="Top-M genes by uniqueness, excluding the top-loading N.") - p.add_argument("--membership_top", type=int, default=300, help="Top genes per program counted as members for uniqueness IDF.") + p.add_argument( + "--top_gene", type=int, default=15, help="Top-N genes by loading score." + ) + p.add_argument( + "--top_unique_gene", + type=int, + default=8, + help="Top-M genes by uniqueness, excluding the top-loading N.", + ) + p.add_argument( + "--membership_top", + type=int, + default=300, + help="Top genes per program counted as members for uniqueness IDF.", + ) p.add_argument("--top_GO", type=int, default=10, help="Top GO terms per program.") - p.add_argument("--top_regulator", type=int, default=6, help="Top significant regulators per program per condition.") + p.add_argument( + "--top_regulator", + type=int, + default=6, + help="Top significant regulators per program per condition.", + ) # significance cutoff shared by every adjusted p-value filter (GO terms, regulators) - p.add_argument("--fdr", type=float, default=0.05, help="Adjusted p-value cutoff for GO terms and regulators.") + p.add_argument( + "--fdr", + type=float, + default=0.05, + help="Adjusted p-value cutoff for GO terms and regulators.", + ) return p @@ -307,7 +424,9 @@ def main(): prog_adata = mu.read_h5ad(str(_check_file(args.mdata_path)), mod=args.prog_key) cond_labels, conditions = load_condition(prog_adata, args.categorical_key) - print(f"[load] {len(conditions)} conditions in obs['{args.categorical_key}']: {conditions}") + print( + f"[load] {len(conditions)} conditions in obs['{args.categorical_key}']: {conditions}" + ) spectra = load_gene_spectra(prog_adata, args.loadings_key, args.gene_name_key) print(f"[load] spectra: {spectra.shape[0]} programs x {spectra.shape[1]} genes") @@ -333,7 +452,9 @@ def main(): programs = args.programs or available missing = sorted(set(programs) - set(available)) if missing: - raise ValueError(f"Programs not in spectra (available {available[0]}..{available[-1]}): {missing}") + raise ValueError( + f"Programs not in spectra (available {available[0]}..{available[-1]}): {missing}" + ) programs = list(dict.fromkeys(programs)) # dedupe, keep order # make output folder @@ -344,15 +465,23 @@ def main(): # one JSON per program labels, n_pairs = [], {} for pid in programs: - bundle = build_bundle(pid, spectra, uniqueness, go_df, assoc_by_cond, cond_stats, args) - (bundle_dir / f"{bundle['program_id']}.json").write_text(json.dumps(bundle, indent=2)) + bundle = build_bundle( + pid, spectra, uniqueness, go_df, assoc_by_cond, cond_stats, args + ) + (bundle_dir / f"{bundle['program_id']}.json").write_text( + json.dumps(bundle, indent=2) + ) regs = bundle.get("perturbation_regulators", {}) reg_summary = ", ".join(f"{c}={len(regs.get(c, []))}" for c in conditions) - n_pairs[bundle["program_id"]] = bundle.get("regulator_gene", {}).get("n_pairs", 0) - print(f" {bundle['program_id']}: genes={len(bundle['program_genes'])} " - f"distinctive={len(bundle['distinctive_genes'])} GO={len(bundle['GO'])} " - f"regulators[{reg_summary}] pairs={n_pairs[bundle['program_id']]} " - f"top={bundle['program_specificity']['top_condition']}") + n_pairs[bundle["program_id"]] = bundle.get("regulator_gene", {}).get( + "n_pairs", 0 + ) + print( + f" {bundle['program_id']}: genes={len(bundle['program_genes'])} " + f"distinctive={len(bundle['distinctive_genes'])} GO={len(bundle['GO'])} " + f"regulators[{reg_summary}] pairs={n_pairs[bundle['program_id']]} " + f"top={bundle['program_specificity']['top_condition']}" + ) labels.append(bundle["program_id"]) meta = { @@ -363,11 +492,28 @@ def main(): "perturbation_path_base": str(args.perturbation_path_base), "perturbation_path": {c: str(p) for c, p in perturb_paths.items()}, }, - "keys": {k: getattr(args, k) for k in ( - "data_key", "prog_key", "loadings_key", "gene_name_key", "categorical_key", "log2fc_key")}, - "params": {k: getattr(args, k) for k in ( - "top_gene", "top_unique_gene", "membership_top", "top_GO", - "top_regulator", "fdr")}, + "keys": { + k: getattr(args, k) + for k in ( + "data_key", + "prog_key", + "loadings_key", + "gene_name_key", + "categorical_key", + "log2fc_key", + ) + }, + "params": { + k: getattr(args, k) + for k in ( + "top_gene", + "top_unique_gene", + "membership_top", + "top_GO", + "top_regulator", + "fdr", + ) + }, "organism": args.organism, "cell_type": args.cell_type, "conditions": conditions, @@ -376,10 +522,12 @@ def main(): "n_regulator_gene_pairs_total": sum(n_pairs.values()), } (bundle_dir / "meta.json").write_text(json.dumps(meta, indent=2)) - print(f"[done] {len(labels)} bundle(s) -> {bundle_dir}; meta -> {bundle_dir / 'meta.json'}") + print( + f"[done] {len(labels)} bundle(s) -> {bundle_dir}; meta -> {bundle_dir / 'meta.json'}" + ) print("Extraction of programs finished") return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_MyGene.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_MyGene.py index db99a2d..12e67ed 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_MyGene.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_MyGene.py @@ -2,22 +2,37 @@ import sys from datetime import datetime -from src import (as_list, default_out_dir, load_program_JSON, report_query, - request_with_retry, split_cached, write_outputs) +from src import ( + as_list, + default_out_dir, + load_program_JSON, + report_query, + request_with_retry, + split_cached, + write_outputs, +) SOURCE = "MyGene" MYGENE_URL = "https://mygene.info/v3/query" # query -def query_mygene(genes, species="human", fields="symbol,name,alias,summary,go.BP.term,entrezgene,generif", - batch_size=1000): +def query_mygene( + genes, + species="human", + fields="symbol,name,alias,summary,go.BP.term,entrezgene,generif", + batch_size=1000, +): """Batched POST to MyGene.info, matching each gene against symbol and alias (so renamed genes such as KIAA1429 -> VIRMA still resolve). Returns the raw hits.""" hits = [] for i in range(0, len(genes), batch_size): - data = {"q": ",".join(genes[i:i + batch_size]), "scopes": "symbol,alias", - "species": species, "fields": fields} + data = { + "q": ",".join(genes[i : i + batch_size]), + "scopes": "symbol,alias", + "species": species, + "fields": fields, + } hits += request_with_retry("POST", MYGENE_URL, data=data).json() return hits @@ -39,8 +54,11 @@ def format_gene_info(hit, top_go_bp=10, top_generif=10): """Aliases and function (summary + GO biological process terms) of one hit, plus GeneRIF statements with the PMID of the paper supporting each.""" go_bp = [t["term"] for t in as_list(hit.get("go", {}).get("BP")) if "term" in t] - generif = [{"pmid": str(r["pubmed"]), "text": r.get("text")} - for r in as_list(hit.get("generif")) if "pubmed" in r] + generif = [ + {"pmid": str(r["pubmed"]), "text": r.get("text")} + for r in as_list(hit.get("generif")) + if "pubmed" in r + ] return { "found": True, "symbol": hit.get("symbol"), @@ -54,22 +72,61 @@ def format_gene_info(hit, top_go_bp=10, top_generif=10): def build_parser(): - p = argparse.ArgumentParser(description="Extend program bundles with MyGene.info aliases and gene function.") + p = argparse.ArgumentParser( + description="Extend program bundles with MyGene.info aliases and gene function." + ) # IO - p.add_argument("--bundle_dir", required=True, help="program_bundles folder written by Extract_program_information.py.") - p.add_argument("--out_dir", default=None, help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.") + p.add_argument( + "--bundle_dir", + required=True, + help="program_bundles folder written by Extract_program_information.py.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.", + ) p.add_argument("--species", default="human", help="MyGene species (name or taxid).") # query - p.add_argument("--fields", default="symbol,name,alias,summary,go.BP.term,entrezgene,generif", help="MyGene fields to request.") - p.add_argument("--top_go_bp", type=int, default=10, help="Max GO biological process terms kept per gene.") - p.add_argument("--top_generif", type=int, default=10, help="Max GeneRIF statements (PMID + text) kept per gene; needs 'generif' in --fields.") - p.add_argument("--batch_size", type=int, default=1000, help="Genes per MyGene request (API max 1000).") - p.add_argument("--overwrite", action="store_true", help="Re-query genes that already have this source's info in --out_dir.") + p.add_argument( + "--fields", + default="symbol,name,alias,summary,go.BP.term,entrezgene,generif", + help="MyGene fields to request.", + ) + p.add_argument( + "--top_go_bp", + type=int, + default=10, + help="Max GO biological process terms kept per gene.", + ) + p.add_argument( + "--top_generif", + type=int, + default=10, + help="Max GeneRIF statements (PMID + text) kept per gene; needs 'generif' in --fields.", + ) + p.add_argument( + "--batch_size", + type=int, + default=1000, + help="Genes per MyGene request (API max 1000).", + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-query genes that already have this source's info in --out_dir.", + ) return p @@ -84,15 +141,23 @@ def main(): info = {} if query: print(f"[{SOURCE}] querying {len(query)} genes -> {MYGENE_URL}") - hits = pick_hits(query_mygene(query, args.species, args.fields, args.batch_size)) - info = {g: format_gene_info(h, args.top_go_bp, args.top_generif) for g, h in hits.items()} + hits = pick_hits( + query_mygene(query, args.species, args.fields, args.batch_size) + ) + info = { + g: format_gene_info(h, args.top_go_bp, args.top_generif) + for g, h in hits.items() + } renamed, not_found = report_query(SOURCE, query, info) meta = { "created": datetime.now().isoformat(timespec="seconds"), "inputs": {"bundle_dir": str(args.bundle_dir)}, "source": MYGENE_URL, - "params": {k: getattr(args, k) for k in ("species", "fields", "top_go_bp", "top_generif", "batch_size")}, + "params": { + k: getattr(args, k) + for k in ("species", "fields", "top_go_bp", "top_generif", "batch_size") + }, "programs": list(bundles), "n_genes": len(genes), "n_skipped": len(cached), @@ -105,5 +170,5 @@ def main(): return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_NCBI.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_NCBI.py index 3ee55f3..fe34a40 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_NCBI.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_NCBI.py @@ -4,8 +4,14 @@ import time from datetime import datetime -from src import (default_out_dir, load_program_JSON, report_query, - request_with_retry, split_cached, write_outputs) +from src import ( + default_out_dir, + load_program_JSON, + report_query, + request_with_retry, + split_cached, + write_outputs, +) SOURCE = "NCBI" NCBI_BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils" @@ -25,8 +31,11 @@ def _ncbi_params(): def _eutils(endpoint, data): """POST to an E-utility, pausing to stay under NCBI's rate limit (3/s, 10/s with a key).""" - resp = request_with_retry("POST", f"{NCBI_BASE_URL}/{endpoint}.fcgi", - data={**_ncbi_params(), **data, "retmode": "json"}) + resp = request_with_retry( + "POST", + f"{NCBI_BASE_URL}/{endpoint}.fcgi", + data={**_ncbi_params(), **data, "retmode": "json"}, + ) time.sleep(0.1 if os.getenv("NCBI_API_KEY") else 0.34) return resp.json() @@ -36,9 +45,15 @@ def search_gene_ids(genes, taxid=9606, batch_size=200): """NCBI Gene ids whose symbol or alias ([sym]) matches any query gene.""" ids = [] for i in range(0, len(genes), batch_size): - terms = " OR ".join(f"{g}[sym]" for g in genes[i:i + batch_size]) - res = _eutils("esearch", {"db": "gene", "term": f"({terms}) AND {taxid}[taxid] AND alive[prop]", - "retmax": 10000}) + terms = " OR ".join(f"{g}[sym]" for g in genes[i : i + batch_size]) + res = _eutils( + "esearch", + { + "db": "gene", + "term": f"({terms}) AND {taxid}[taxid] AND alive[prop]", + "retmax": 10000, + }, + ) ids += res["esearchresult"]["idlist"] return list(dict.fromkeys(ids)) @@ -47,7 +62,9 @@ def fetch_summaries(ids, batch_size=200): """{gene_id: esummary record}.""" records = {} for i in range(0, len(ids), batch_size): - res = _eutils("esummary", {"db": "gene", "id": ",".join(ids[i:i + batch_size])})["result"] + res = _eutils( + "esummary", {"db": "gene", "id": ",".join(ids[i : i + batch_size])} + )["result"] records.update({uid: res[uid] for uid in res["uids"]}) return records @@ -57,8 +74,12 @@ def fetch_generif_pmids(ids, batch_size=20): batches: well-studied genes link thousands of PMIDs and NCBI cuts off large responses.""" pmids = {} for i in range(0, len(ids), batch_size): - data = {"dbfrom": "gene", "db": "pubmed", "linkname": "gene_pubmed_rif", - "id": ids[i:i + batch_size]} # repeated id= gives one linkset per gene + data = { + "dbfrom": "gene", + "db": "pubmed", + "linkname": "gene_pubmed_rif", + "id": ids[i : i + batch_size], + } # repeated id= gives one linkset per gene for ls in _eutils("elink", data).get("linksets", []): links = [l for db in ls.get("linksetdbs", []) for l in db.get("links", [])] pmids[str(ls["ids"][0])] = [str(l) for l in links] @@ -102,21 +123,55 @@ def format_gene_info(uid, record, generif_pmids, top_generif=10): def build_parser(): - p = argparse.ArgumentParser(description="Extend program bundles with NCBI Gene aliases, summary and GeneRIF PMIDs.") + p = argparse.ArgumentParser( + description="Extend program bundles with NCBI Gene aliases, summary and GeneRIF PMIDs." + ) # IO - p.add_argument("--bundle_dir", required=True, help="program_bundles folder written by Extract_program_information.py.") - p.add_argument("--out_dir", default=None, help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.") + p.add_argument( + "--bundle_dir", + required=True, + help="program_bundles folder written by Extract_program_information.py.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.") - p.add_argument("--taxid", type=int, default=9606, help="NCBI taxonomy id (9606 = human).") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.", + ) + p.add_argument( + "--taxid", type=int, default=9606, help="NCBI taxonomy id (9606 = human)." + ) # query - p.add_argument("--top_generif", type=int, default=10, help="Max GeneRIF PMIDs kept per gene.") - p.add_argument("--batch_size", type=int, default=200, help="Genes/ids per esearch/esummary request.") - p.add_argument("--elink_batch_size", type=int, default=20, help="Gene ids per GeneRIF elink request (large batches get cut off by NCBI).") - p.add_argument("--overwrite", action="store_true", help="Re-query genes that already have this source's info in --out_dir.") + p.add_argument( + "--top_generif", type=int, default=10, help="Max GeneRIF PMIDs kept per gene." + ) + p.add_argument( + "--batch_size", + type=int, + default=200, + help="Genes/ids per esearch/esummary request.", + ) + p.add_argument( + "--elink_batch_size", + type=int, + default=20, + help="Gene ids per GeneRIF elink request (large batches get cut off by NCBI).", + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-query genes that already have this source's info in --out_dir.", + ) return p @@ -134,16 +189,23 @@ def main(): ids = search_gene_ids(query, args.taxid, args.batch_size) records = fetch_summaries(ids, args.batch_size) picked = pick_records(query, records) - generif = fetch_generif_pmids(list(dict.fromkeys(picked.values())), args.elink_batch_size) - info = {g: format_gene_info(uid, records[uid], generif, args.top_generif) - for g, uid in picked.items()} + generif = fetch_generif_pmids( + list(dict.fromkeys(picked.values())), args.elink_batch_size + ) + info = { + g: format_gene_info(uid, records[uid], generif, args.top_generif) + for g, uid in picked.items() + } renamed, not_found = report_query(SOURCE, query, info) meta = { "created": datetime.now().isoformat(timespec="seconds"), "inputs": {"bundle_dir": str(args.bundle_dir)}, "source": NCBI_BASE_URL, - "params": {k: getattr(args, k) for k in ("taxid", "top_generif", "batch_size", "elink_batch_size")}, + "params": { + k: getattr(args, k) + for k in ("taxid", "top_generif", "batch_size", "elink_batch_size") + }, "programs": list(bundles), "n_genes": len(genes), "n_skipped": len(cached), @@ -156,5 +218,5 @@ def main(): return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_UniProt.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_UniProt.py index 9fac367..df537fb 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_UniProt.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/search_UniProt.py @@ -2,13 +2,21 @@ import sys from datetime import datetime -from src import (default_out_dir, load_program_JSON, report_query, - request_with_retry, split_cached, write_outputs) +from src import ( + default_out_dir, + load_program_JSON, + report_query, + request_with_retry, + split_cached, + write_outputs, +) SOURCE = "UniProt" UNIPROT_URL = "https://rest.uniprot.org/uniprotkb/search" -FIELDS = ("accession,gene_primary,gene_synonym,protein_name,cc_function," - "cc_subcellular_location,cc_tissue_specificity,cc_pathway") +FIELDS = ( + "accession,gene_primary,gene_synonym,protein_name,cc_function," + "cc_subcellular_location,cc_tissue_specificity,cc_pathway" +) # query @@ -17,14 +25,21 @@ def query_uniprot(genes, taxid=9606, batch_size=100): query gene, following the Link: rel="next" pagination.""" entries = [] for i in range(0, len(genes), batch_size): - terms = " OR ".join(f"gene_exact:{g}" for g in genes[i:i + batch_size]) + terms = " OR ".join(f"gene_exact:{g}" for g in genes[i : i + batch_size]) url = UNIPROT_URL - params = {"query": f"({terms}) AND organism_id:{taxid} AND reviewed:true", - "fields": FIELDS, "format": "json", "size": 500} + params = { + "query": f"({terms}) AND organism_id:{taxid} AND reviewed:true", + "fields": FIELDS, + "format": "json", + "size": 500, + } while url: resp = request_with_retry("GET", url, params=params) entries += resp.json().get("results", []) - url, params = resp.links.get("next", {}).get("url"), None # next url carries the query + url, params = ( + resp.links.get("next", {}).get("url"), + None, + ) # next url carries the query return entries @@ -47,7 +62,11 @@ def pick_entries(genes, entries): by_primary.setdefault(n, e) for n in synonyms: by_synonym.setdefault(n, e) - return {g: by_primary.get(g) or by_synonym[g] for g in genes if g in by_primary or g in by_synonym} + return { + g: by_primary.get(g) or by_synonym[g] + for g in genes + if g in by_primary or g in by_synonym + } def _comments(entry, ctype): @@ -55,15 +74,29 @@ def _comments(entry, ctype): def _texts(entry, ctype): - return " ".join(t["value"] for c in _comments(entry, ctype) for t in c.get("texts", [])) or None + return ( + " ".join( + t["value"] for c in _comments(entry, ctype) for t in c.get("texts", []) + ) + or None + ) def format_gene_info(entry): primary, synonyms = _gene_names(entry) - function_pmids = [ev["id"] for c in _comments(entry, "FUNCTION") for t in c.get("texts", []) - for ev in t.get("evidences", []) if ev.get("source") == "PubMed"] - locations = [loc["location"]["value"] for c in _comments(entry, "SUBCELLULAR LOCATION") - for loc in c.get("subcellularLocations", []) if "location" in loc] + function_pmids = [ + ev["id"] + for c in _comments(entry, "FUNCTION") + for t in c.get("texts", []) + for ev in t.get("evidences", []) + if ev.get("source") == "PubMed" + ] + locations = [ + loc["location"]["value"] + for c in _comments(entry, "SUBCELLULAR LOCATION") + for loc in c.get("subcellularLocations", []) + if "location" in loc + ] rec_name = entry.get("proteinDescription", {}).get("recommendedName", {}) return { "found": True, @@ -80,19 +113,43 @@ def format_gene_info(entry): def build_parser(): - p = argparse.ArgumentParser(description="Extend program bundles with UniProtKB protein function, location and tissue specificity.") + p = argparse.ArgumentParser( + description="Extend program bundles with UniProtKB protein function, location and tissue specificity." + ) # IO - p.add_argument("--bundle_dir", required=True, help="program_bundles folder written by Extract_program_information.py.") - p.add_argument("--out_dir", default=None, help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.") + p.add_argument( + "--bundle_dir", + required=True, + help="program_bundles folder written by Extract_program_information.py.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.") - p.add_argument("--taxid", type=int, default=9606, help="NCBI taxonomy id (9606 = human).") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.", + ) + p.add_argument( + "--taxid", type=int, default=9606, help="NCBI taxonomy id (9606 = human)." + ) # query - p.add_argument("--batch_size", type=int, default=100, help="Genes per UniProt query.") - p.add_argument("--overwrite", action="store_true", help="Re-query genes that already have this source's info in --out_dir.") + p.add_argument( + "--batch_size", type=int, default=100, help="Genes per UniProt query." + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-query genes that already have this source's info in --out_dir.", + ) return p @@ -115,8 +172,12 @@ def main(): "created": datetime.now().isoformat(timespec="seconds"), "inputs": {"bundle_dir": str(args.bundle_dir)}, "source": UNIPROT_URL, - "params": {"taxid": args.taxid, "batch_size": args.batch_size, "fields": FIELDS, - "reviewed_only": True}, + "params": { + "taxid": args.taxid, + "batch_size": args.batch_size, + "fields": FIELDS, + "reviewed_only": True, + }, "programs": list(bundles), "n_genes": len(genes), "n_skipped": len(cached), @@ -129,5 +190,5 @@ def main(): return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/src.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/src.py index 656dbd3..382eb2a 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/src.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene/src.py @@ -8,6 +8,7 @@ when it already exists (so earlier sources are kept) and '/P.json' otherwise. The original bundles are never modified. """ + import copy import json import time @@ -44,9 +45,13 @@ def load_program_JSON(bundle_dir, out_dir, programs): paths[pid] = extended if extended.exists() else bundle_dir / f"P{pid}.json" missing = [str(p) for p in paths.values() if not p.exists()] if missing: - raise FileNotFoundError(f"No bundle for {len(missing)} of {len(programs)} programs: {missing}") + raise FileNotFoundError( + f"No bundle for {len(missing)} of {len(programs)} programs: {missing}" + ) n_ext = sum(p.parent == out_dir for p in paths.values()) - print(f"[load] {len(paths)} bundle(s): {n_ext} from {out_dir}, {len(paths) - n_ext} from {bundle_dir}") + print( + f"[load] {len(paths)} bundle(s): {n_ext} from {out_dir}, {len(paths) - n_ext} from {bundle_dir}" + ) bundles = {} for pid, p in paths.items(): bundles[f"P{pid}"] = json.loads(p.read_text()) @@ -56,27 +61,37 @@ def load_program_JSON(bundle_dir, out_dir, programs): def validate_bundle(bundle, path): """Raise if the bundle lacks the gene lists or has malformed gene entries.""" + def _valid_entry(e): - return isinstance(e, str) or (isinstance(e, dict) and isinstance(e.get("gene"), str) - and isinstance(e.get("gene_info", {}), dict)) + return isinstance(e, str) or ( + isinstance(e, dict) + and isinstance(e.get("gene"), str) + and isinstance(e.get("gene_info", {}), dict) + ) for key in GENE_LISTS: if not isinstance(bundle.get(key), list): raise ValueError(f"{path}: '{key}' missing or not a list") bad = [e for e in bundle[key] if not _valid_entry(e)] if bad: - raise ValueError(f"{path}: '{key}' has {len(bad)} entries that are neither a gene " - f"name nor {{'gene': ...}}: {bad[:3]}") + raise ValueError( + f"{path}: '{key}' has {len(bad)} entries that are neither a gene " + f"name nor {{'gene': ...}}: {bad[:3]}" + ) regulators = bundle.get("perturbation_regulators", {}) if not isinstance(regulators, dict): - raise ValueError(f"{path}: 'perturbation_regulators' must be {{condition: [entries]}}") + raise ValueError( + f"{path}: 'perturbation_regulators' must be {{condition: [entries]}}" + ) for cond, regs in regulators.items(): if not isinstance(regs, list): raise ValueError(f"{path}: perturbation_regulators['{cond}'] is not a list") bad = [r for r in regs if not (isinstance(r, dict) and _valid_entry(r))] if bad: - raise ValueError(f"{path}: perturbation_regulators['{cond}'] has {len(bad)} entries " - f"without a 'gene' name: {bad[:3]}") + raise ValueError( + f"{path}: perturbation_regulators['{cond}'] has {len(bad)} entries " + f"without a 'gene' name: {bad[:3]}" + ) def gene_name(entry): @@ -109,8 +124,10 @@ def split_cached(bundles, source, overwrite=False): if isinstance(e, dict) and source in e.get("gene_info", {}): cached.setdefault(gene_name(e), e["gene_info"][source]) query = [g for g in genes if g not in cached] - print(f"[{source}] {len(genes)} unique genes: {len(cached)} already have {source} info " - f"(skipped{'' if cached else '; none'}), {len(query)} to query") + print( + f"[{source}] {len(genes)} unique genes: {len(cached)} already have {source} info " + f"(skipped{'' if cached else '; none'}), {len(query)} to query" + ) return genes, cached, query @@ -121,13 +138,17 @@ def request_with_retry(method, url, retries=3, **kwargs): for attempt in range(retries + 1): try: resp = requests.request(method, url, timeout=60, **kwargs) - except (requests.ConnectionError, requests.Timeout, requests.exceptions.ChunkedEncodingError): + except ( + requests.ConnectionError, + requests.Timeout, + requests.exceptions.ChunkedEncodingError, + ): if attempt == retries: raise - time.sleep(2 ** attempt) + time.sleep(2**attempt) continue if resp.status_code in TRANSIENT_STATUS and attempt < retries: - time.sleep(2 ** attempt) + time.sleep(2**attempt) continue resp.raise_for_status() return resp @@ -157,7 +178,9 @@ def report_query(source, genes, info): print(f"[{source}] nothing to query") return {}, [] not_found = [g for g in genes if g not in info] - renamed = {g: i["symbol"] for g, i in info.items() if i.get("symbol") and i["symbol"] != g} + renamed = { + g: i["symbol"] for g, i in info.items() if i.get("symbol") and i["symbol"] != g + } print(f"[{source}] found {len(info)}/{len(genes)} queried; renamed {len(renamed)}") if not_found: print(f" [warn] {len(not_found)} genes not found in {source}: {not_found}") @@ -169,9 +192,13 @@ def write_outputs(out_dir, bundles, source, info, meta): out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) for label, bundle in bundles.items(): - (out_dir / f"{label}.json").write_text(json.dumps(attach_source(bundle, source, info), indent=2)) + (out_dir / f"{label}.json").write_text( + json.dumps(attach_source(bundle, source, info), indent=2) + ) genes = collect_genes(bundle) - print(f" {label}: genes={len(genes)} found={sum(info.get(g, NOT_FOUND)['found'] for g in genes)}") + print( + f" {label}: genes={len(genes)} found={sum(info.get(g, NOT_FOUND)['found'] for g in genes)}" + ) meta_path = out_dir / f"meta_{source}.json" meta_path.write_text(json.dumps(meta, indent=2)) print(f"[done] {len(bundles)} bundle(s) -> {out_dir}; meta -> {meta_path}") diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene_interaction/search_OmniPath.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene_interaction/search_OmniPath.py index 019ab10..f18ef9b 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene_interaction/search_OmniPath.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.0.Search_database/search_gene_interaction/search_OmniPath.py @@ -5,26 +5,56 @@ from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "search_gene")) -from src import as_list, default_out_dir, gene_name, load_program_JSON, request_with_retry # noqa: E402 +from src import ( + as_list, + default_out_dir, + gene_name, + load_program_JSON, + request_with_retry, +) # noqa: E402 SOURCE = "OmniPath" OMNIPATH_URL = "https://omnipathdb.org/interactions" -ALL_DATASETS = ("omnipath", "pathwayextra", "kinaseextra", "ligrecextra", "collectri", "dorothea", - "tf_target", "mirnatarget", "lncrna_mrna", "tf_mirna", "small_molecule") +ALL_DATASETS = ( + "omnipath", + "pathwayextra", + "kinaseextra", + "ligrecextra", + "collectri", + "dorothea", + "tf_target", + "mirnatarget", + "lncrna_mrna", + "tf_mirna", + "small_molecule", +) FIELDS = "sources,references,curation_effort,type,datasets,extra_attrs,dorothea_level" GENE_INFO_SOURCES = ("MyGene", "NCBI", "UniProt") # stored in params so results written in an older layout are re-queried, not skipped PAIRS_FORMAT = "regulator_gene_v3" # gene_category in the extract step's regulator_gene block -> query category used in the pairs -CATEGORY_NAMES = {"program_gene": "program_gene-regulator", "distinctive_gene": "distinctive_gene-regulator"} +CATEGORY_NAMES = { + "program_gene": "program_gene-regulator", + "distinctive_gene": "distinctive_gene-regulator", +} # loaders def select_from_pair_list(bundle): """Query sets from the bundle's regulator_gene block (Extract_program_information.py): ({category: [gene entries]}, {regulator: {cond: stats}}, {regulator: entry}, {(gene, regulator)}).""" - entries = {gene_name(e): e for key in ("program_genes", "distinctive_genes") for e in bundle[key]} - entries.update({gene_name(r): r for regs in bundle.get("perturbation_regulators", {}).values() for r in regs}) + entries = { + gene_name(e): e + for key in ("program_genes", "distinctive_genes") + for e in bundle[key] + } + entries.update( + { + gene_name(r): r + for regs in bundle.get("perturbation_regulators", {}).values() + for r in regs + } + ) gene_sets, regulators, reg_entries, listed = {}, {}, {}, set() for p in bundle["regulator_gene"]["pairs"].values(): gene, reg = p["gene"], p["regulator"] @@ -38,7 +68,9 @@ def select_from_pair_list(bundle): return gene_sets, regulators, reg_entries, listed -def select_sets(bundle, top_gene=15, top_regulator=6, include_unique=False, top_unique_gene=8): +def select_sets( + bundle, top_gene=15, top_regulator=6, include_unique=False, top_unique_gene=8 +): """Query sets of one program: {category: [gene entries]}, {regulator: {cond: stats}}, {regulator: entry} and the listed (gene, regulator) pairs (None = every combination). @@ -49,12 +81,17 @@ def select_sets(bundle, top_gene=15, top_regulator=6, include_unique=False, top_ return select_from_pair_list(bundle) gene_sets = {"program_gene-regulator": bundle["program_genes"][:top_gene]} if include_unique: - gene_sets["distinctive_gene-regulator"] = bundle["distinctive_genes"][:top_unique_gene] + gene_sets["distinctive_gene-regulator"] = bundle["distinctive_genes"][ + :top_unique_gene + ] regulators, reg_entries = {}, {} for cond, regs in bundle.get("perturbation_regulators", {}).items(): for r in regs[:top_regulator]: g = gene_name(r) - regulators.setdefault(g, {})[cond] = {"log2fc": r.get("log2fc"), "adj_pval": r.get("adj_pval")} + regulators.setdefault(g, {})[cond] = { + "log2fc": r.get("log2fc"), + "adj_pval": r.get("adj_pval"), + } reg_entries.setdefault(g, r) return gene_sets, regulators, reg_entries, None @@ -68,7 +105,9 @@ def name_candidates(entry, use_aliases=True): info = entry.get("gene_info", {}).get(src, {}) if not info.get("found"): continue - names += [info.get("symbol"), info.get("accession")] + as_list(info.get("alias")) + names += [info.get("symbol"), info.get("accession")] + as_list( + info.get("alias") + ) return list(dict.fromkeys(str(n) for n in names if n)) @@ -76,7 +115,11 @@ def canonical_id(entry): """Stable id of a gene (UniProt accession, then NCBI/MyGene gene id, then its name), so an old and a new name of the same gene (MESDC1 / TLNRD1) are not paired with each other.""" info = entry.get("gene_info", {}) if isinstance(entry, dict) else {} - for src, key in (("UniProt", "accession"), ("NCBI", "gene_id"), ("MyGene", "entrezgene")): + for src, key in ( + ("UniProt", "accession"), + ("NCBI", "gene_id"), + ("MyGene", "entrezgene"), + ): if info.get(src, {}).get(key): return f"{src}:{info[src][key]}" return gene_name(entry) @@ -107,9 +150,16 @@ def query_omnipath(sources, targets, datasets, taxid=9606): without it OmniPath ORs the two lists). Unknown names are ignored by the server.""" if not sources or not targets: return [] - params = {"sources": ",".join(sources), "targets": ",".join(targets), - "source_target": "AND", "genesymbols": "yes", "organisms": taxid, - "datasets": ",".join(datasets), "fields": FIELDS, "format": "json"} + params = { + "sources": ",".join(sources), + "targets": ",".join(targets), + "source_target": "AND", + "genesymbols": "yes", + "organisms": taxid, + "datasets": ",".join(datasets), + "fields": FIELDS, + "format": "json", + } return request_with_retry("GET", OMNIPATH_URL, params=params).json() @@ -132,8 +182,11 @@ def _resolve(record, side, cand_map): def _sign(stim, inhib): - return {(True, True): "both", (True, False): "activation", - (False, True): "inhibition"}.get((bool(stim), bool(inhib)), "unknown") + return { + (True, True): "both", + (True, False): "activation", + (False, True): "inhibition", + }.get((bool(stim), bool(inhib)), "unknown") def _pmids(references): @@ -144,14 +197,20 @@ def _pmids(references): def format_interaction(record, direction, max_pmids=20): pmids = _pmids(record.get("references")) - mechanisms = [str(m) for k, v in (record.get("extra_attrs") or {}).items() - if k.endswith("_mechanism") for m in as_list(v)] + mechanisms = [ + str(m) + for k, v in (record.get("extra_attrs") or {}).items() + if k.endswith("_mechanism") + for m in as_list(v) + ] return { "source": record.get("source_genesymbol"), "target": record.get("target_genesymbol"), "direction": direction, "sign": _sign(record.get("is_stimulation"), record.get("is_inhibition")), - "consensus_sign": _sign(record.get("consensus_stimulation"), record.get("consensus_inhibition")), + "consensus_sign": _sign( + record.get("consensus_stimulation"), record.get("consensus_inhibition") + ), "type": record.get("type"), "datasets": [d for d in ALL_DATASETS if record.get(d) is True], "resources": as_list(record.get("sources")), @@ -163,7 +222,9 @@ def format_interaction(record, direction, max_pmids=20): } -def build_pairs(records, gene_category, regulators, cand_map, canon, max_pmids=20, listed=None): +def build_pairs( + records, gene_category, regulators, cand_map, canon, max_pmids=20, listed=None +): """{"GENE-REGULATOR": entry} for EVERY tested pair (the listed pairs, or every gene x regulator when listed is None): pairs with OmniPath records first (strongest first), then pairs with none (found: false), in selection order. A gene is never paired with itself, @@ -174,122 +235,250 @@ def build_pairs(records, gene_category, regulators, cand_map, canon, max_pmids=2 pairs = {} for r in records: - (src, src_as), (tgt, tgt_as) = _resolve(r, "source", cand_map), _resolve(r, "target", cand_map) + (src, src_as), (tgt, tgt_as) = ( + _resolve(r, "source", cand_map), + _resolve(r, "target", cand_map), + ) if src is None or tgt is None or canon[src] == canon[tgt]: continue if src in regulators and tgt in gene_category: - gene, reg, gene_as, reg_as, direction = tgt, src, tgt_as, src_as, "regulator->gene" + gene, reg, gene_as, reg_as, direction = ( + tgt, + src, + tgt_as, + src_as, + "regulator->gene", + ) elif tgt in regulators and src in gene_category: - gene, reg, gene_as, reg_as, direction = src, tgt, src_as, tgt_as, "gene->regulator" + gene, reg, gene_as, reg_as, direction = ( + src, + tgt, + src_as, + tgt_as, + "gene->regulator", + ) else: continue # gene-gene or regulator-regulator record if not r.get("is_directed"): direction = "undirected" - pair = pairs.setdefault((gene, reg), {"matched_as": {"gene": gene_as, "regulator": reg_as}, - "interactions": [], "pmids": []}) + pair = pairs.setdefault( + (gene, reg), + { + "matched_as": {"gene": gene_as, "regulator": reg_as}, + "interactions": [], + "pmids": [], + }, + ) pair["interactions"].append(format_interaction(r, direction, max_pmids)) pair["pmids"] += _pmids(r.get("references")) out = [] for gene, cat in gene_category.items(): for reg in regulators: - if canon[gene] == canon[reg] or (listed is not None and (gene, reg) not in listed): + if canon[gene] == canon[reg] or ( + listed is not None and (gene, reg) not in listed + ): continue - hit = pairs.get((gene, reg), {"matched_as": None, "interactions": [], "pmids": []}) + hit = pairs.get( + (gene, reg), {"matched_as": None, "interactions": [], "pmids": []} + ) inter = hit["interactions"] pmids = list(dict.fromkeys(hit["pmids"])) resources = list(dict.fromkeys(s for i in inter for s in i["resources"])) - out.append({ - "pair": f"{gene}-{reg}", - "gene": gene, "regulator": reg, - "query_category": cat, - "gene_category": cat.split("-")[0], - "found": bool(inter), - "regulator_stats": regulators[reg], - "tried_names": {"gene": tried.get(gene, []), "regulator": tried.get(reg, [])}, - "matched_as": hit["matched_as"], - "n_references": len(pmids), - "pmids": pmids[:max_pmids], - "n_resources": len(resources), - "resources": resources, - "curation_effort": sum(i["curation_effort"] or 0 for i in inter), - "interaction_categories": list(dict.fromkeys(i["type"] for i in inter)), - "datasets": list(dict.fromkeys(d for i in inter for d in i["datasets"])), - "directions": list(dict.fromkeys(i["direction"] for i in inter)), - "signs": list(dict.fromkeys(i["sign"] for i in inter)), - "interactions": inter, - }) + out.append( + { + "pair": f"{gene}-{reg}", + "gene": gene, + "regulator": reg, + "query_category": cat, + "gene_category": cat.split("-")[0], + "found": bool(inter), + "regulator_stats": regulators[reg], + "tried_names": { + "gene": tried.get(gene, []), + "regulator": tried.get(reg, []), + }, + "matched_as": hit["matched_as"], + "n_references": len(pmids), + "pmids": pmids[:max_pmids], + "n_resources": len(resources), + "resources": resources, + "curation_effort": sum(i["curation_effort"] or 0 for i in inter), + "interaction_categories": list( + dict.fromkeys(i["type"] for i in inter) + ), + "datasets": list( + dict.fromkeys(d for i in inter for d in i["datasets"]) + ), + "directions": list(dict.fromkeys(i["direction"] for i in inter)), + "signs": list(dict.fromkeys(i["sign"] for i in inter)), + "interactions": inter, + } + ) # stable sort: found pairs by evidence, not-found pairs keep selection order - out.sort(key=lambda p: (p["found"], p["n_references"], p["curation_effort"]), reverse=True) + out.sort( + key=lambda p: (p["found"], p["n_references"], p["curation_effort"]), + reverse=True, + ) return {p.pop("pair"): p for p in out} def search_program(bundle, params): """OmniPath block of one program, plus the alias candidates dropped as ambiguous.""" gene_sets, regulators, reg_entries, listed = select_sets( - bundle, params["top_gene"], params["top_regulator"], - params["include_unique"], params["top_unique_gene"]) + bundle, + params["top_gene"], + params["top_regulator"], + params["include_unique"], + params["top_unique_gene"], + ) gene_entries, gene_category = {}, {} for cat, entries in gene_sets.items(): for e in entries: gene_entries.setdefault(gene_name(e), e) gene_category.setdefault(gene_name(e), cat) - cand_map, dropped = build_candidate_map({**gene_entries, **reg_entries}, params["use_aliases"]) + cand_map, dropped = build_candidate_map( + {**gene_entries, **reg_entries}, params["use_aliases"] + ) gene_cands = [n for n, g in cand_map.items() if g in gene_category] reg_cands = [n for n, g in cand_map.items() if g in regulators] records = fetch_records(gene_cands, reg_cands, params["datasets"], params["taxid"]) canon = {g: canonical_id(e) for g, e in {**gene_entries, **reg_entries}.items()} - pairs = build_pairs(records, gene_category, regulators, cand_map, canon, params["max_pmids"], listed) + pairs = build_pairs( + records, gene_category, regulators, cand_map, canon, params["max_pmids"], listed + ) - tested = {cat: sum(p["query_category"] == cat for p in pairs.values()) for cat in gene_sets} - found = {cat: sum(p["query_category"] == cat and p["found"] for p in pairs.values()) for cat in gene_sets} - block = {"params": params, "pair_source": "regulator_gene" if listed is not None else "flags", - "n_regulators": len(regulators), - "n_pairs_tested": tested, "n_pairs_found": found, "pairs": pairs} + tested = { + cat: sum(p["query_category"] == cat for p in pairs.values()) + for cat in gene_sets + } + found = { + cat: sum(p["query_category"] == cat and p["found"] for p in pairs.values()) + for cat in gene_sets + } + block = { + "params": params, + "pair_source": "regulator_gene" if listed is not None else "flags", + "n_regulators": len(regulators), + "n_pairs_tested": tested, + "n_pairs_found": found, + "pairs": pairs, + } return block, dropped def build_parser(): - p = argparse.ArgumentParser(description="Query OmniPath for each program's regulator-gene pairs (the regulator_gene list from Extract_program_information.py; older bundles: top genes x top regulators).") + p = argparse.ArgumentParser( + description="Query OmniPath for each program's regulator-gene pairs (the regulator_gene list from Extract_program_information.py; older bundles: top genes x top regulators)." + ) # IO - p.add_argument("--bundle_dir", required=True, help="PerturbNMF_Info folder written by Extract_program_information.py.") - p.add_argument("--out_dir", default=None, help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.") + p.add_argument( + "--bundle_dir", + required=True, + help="PerturbNMF_Info folder written by Extract_program_information.py.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Shared output directory (existing P.json there are extended). Default: /../Gene_info_extended_PerturbNMF_Info.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.") - p.add_argument("--taxid", type=int, default=9606, help="NCBI taxonomy id (9606 = human).") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3); only these bundles are read.", + ) + p.add_argument( + "--taxid", type=int, default=9606, help="NCBI taxonomy id (9606 = human)." + ) # select top items of the program - p.add_argument("--top_gene", type=int, default=15, help="Top-N loaded genes (program_genes) paired with regulators. Only for bundles without a regulator_gene block.") - p.add_argument("--top_regulator", type=int, default=6, help="Top-N regulators per condition. Only for bundles without a regulator_gene block.") - p.add_argument("--include_unique", action="store_true", help="Also pair the top unique genes (distinctive_genes) with the regulators. Only for bundles without a regulator_gene block.") - p.add_argument("--top_unique_gene", type=int, default=8, help="Top-M unique genes, used with --include_unique.") + p.add_argument( + "--top_gene", + type=int, + default=15, + help="Top-N loaded genes (program_genes) paired with regulators. Only for bundles without a regulator_gene block.", + ) + p.add_argument( + "--top_regulator", + type=int, + default=6, + help="Top-N regulators per condition. Only for bundles without a regulator_gene block.", + ) + p.add_argument( + "--include_unique", + action="store_true", + help="Also pair the top unique genes (distinctive_genes) with the regulators. Only for bundles without a regulator_gene block.", + ) + p.add_argument( + "--top_unique_gene", + type=int, + default=8, + help="Top-M unique genes, used with --include_unique.", + ) # query - p.add_argument("--datasets", nargs="+", default=list(ALL_DATASETS), choices=ALL_DATASETS, help="OmniPath datasets to search. Default: all.") - p.add_argument("--no_aliases", action="store_true", help="Query bundle gene names only, not the MyGene/NCBI/UniProt aliases in gene_info.") - p.add_argument("--max_pmids", type=int, default=20, help="Max PMIDs listed per pair / interaction (counts are uncapped).") - p.add_argument("--overwrite", action="store_true", help="Re-query programs that already have OmniPath results with the same params.") + p.add_argument( + "--datasets", + nargs="+", + default=list(ALL_DATASETS), + choices=ALL_DATASETS, + help="OmniPath datasets to search. Default: all.", + ) + p.add_argument( + "--no_aliases", + action="store_true", + help="Query bundle gene names only, not the MyGene/NCBI/UniProt aliases in gene_info.", + ) + p.add_argument( + "--max_pmids", + type=int, + default=20, + help="Max PMIDs listed per pair / interaction (counts are uncapped).", + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-query programs that already have OmniPath results with the same params.", + ) return p def main(): args = build_parser().parse_args() out_dir = Path(args.out_dir or default_out_dir(args.bundle_dir)) - params = {"top_gene": args.top_gene, "top_regulator": args.top_regulator, - "include_unique": args.include_unique, "top_unique_gene": args.top_unique_gene, - "datasets": list(args.datasets), "use_aliases": not args.no_aliases, - "max_pmids": args.max_pmids, "taxid": args.taxid, - "pairs_format": PAIRS_FORMAT} + params = { + "top_gene": args.top_gene, + "top_regulator": args.top_regulator, + "include_unique": args.include_unique, + "top_unique_gene": args.top_unique_gene, + "datasets": list(args.datasets), + "use_aliases": not args.no_aliases, + "max_pmids": args.max_pmids, + "taxid": args.taxid, + "pairs_format": PAIRS_FORMAT, + } bundles = load_program_JSON(args.bundle_dir, out_dir, args.programs) - entries = [e for b in bundles.values() for e in b["program_genes"] + b["distinctive_genes"] - + [r for regs in b.get("perturbation_regulators", {}).values() for r in regs]] - if params["use_aliases"] and not any(isinstance(e, dict) and "gene_info" in e for e in entries): - print(" [warn] no gene_info in the bundles; only bundle gene names are queried " - "(run search_gene/search_MyGene.py, search_NCBI.py, search_UniProt.py first to use aliases)") + entries = [ + e + for b in bundles.values() + for e in b["program_genes"] + + b["distinctive_genes"] + + [r for regs in b.get("perturbation_regulators", {}).values() for r in regs] + ] + if params["use_aliases"] and not any( + isinstance(e, dict) and "gene_info" in e for e in entries + ): + print( + " [warn] no gene_info in the bundles; only bundle gene names are queried " + "(run search_gene/search_MyGene.py, search_NCBI.py, search_UniProt.py first to use aliases)" + ) out_dir.mkdir(parents=True, exist_ok=True) summary, skipped = {}, [] @@ -298,14 +487,23 @@ def main(): previous = bundle.get("gene_interactions", {}).get(SOURCE, {}) if previous.get("params") == params and not args.overwrite: skipped.append(label) - print(f" {label}: OmniPath results with the same params exist; skipped (use --overwrite)") + print( + f" {label}: OmniPath results with the same params exist; skipped (use --overwrite)" + ) continue block, dropped = search_program(bundle, params) bundle.setdefault("gene_interactions", {})[SOURCE] = block (out_dir / f"{label}.json").write_text(json.dumps(bundle, indent=2)) - summary[label] = {"n_regulators": block["n_regulators"], "n_pairs_tested": block["n_pairs_tested"], - "n_pairs_found": block["n_pairs_found"], "dropped_aliases": dropped} - found = ", ".join(f"{c}={block['n_pairs_found'][c]}/{block['n_pairs_tested'][c]}" for c in block["n_pairs_found"]) + summary[label] = { + "n_regulators": block["n_regulators"], + "n_pairs_tested": block["n_pairs_tested"], + "n_pairs_found": block["n_pairs_found"], + "dropped_aliases": dropped, + } + found = ", ".join( + f"{c}={block['n_pairs_found'][c]}/{block['n_pairs_tested'][c]}" + for c in block["n_pairs_found"] + ) print(f" {label}: regulators={block['n_regulators']} pairs found [{found}]") if dropped: print(f" [warn] ambiguous aliases dropped: {dropped}") @@ -321,9 +519,11 @@ def main(): } meta_path = out_dir / f"meta_{SOURCE}.json" meta_path.write_text(json.dumps(meta, indent=2)) - print(f"[done] {len(summary)} searched, {len(skipped)} skipped -> {out_dir}; meta -> {meta_path}") + print( + f"[done] {len(summary)} searched, {len(skipped)} skipped -> {out_dir}; meta -> {meta_path}" + ) return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_arxiv.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_arxiv.py index 042a5fd..cbc15d1 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_arxiv.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_arxiv.py @@ -2,16 +2,16 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # llm_summary_agent / pdf_download +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # llm_summary_agent / pdf_download from llm_summary_agent import Summarize_Agent from pdf_download import download_paper, report_downloads + # ArXiv API def search_Arxiv( - query: str = "", - max_results: int = 3, - mini_handler = None, - out_dir: str = None + query: str = "", max_results: int = 3, mini_handler=None, out_dir: str = None ) -> list[dict]: """Search ArXiv for papers related to the query @@ -27,7 +27,7 @@ def search_Arxiv( retriever = ArxivRetriever( top_k_results=max_results, load_max_docs=max_results, - load_all_available_meta=True + load_all_available_meta=True, ) # Search for papers related to the query @@ -37,43 +37,74 @@ def search_Arxiv( data_list = [] downloads = [] for paper in related_papers: + url = paper.metadata.get("Entry ID", "No URL") - url = paper.metadata.get('Entry ID', 'No URL') - content = "No abstract available." if paper.page_content: content = Summarize_Agent(paper.page_content, mini_handler) - - title = paper.metadata.get('Title', paper.metadata.get('title', 'No Title')) - authors = paper.metadata.get('Authors', paper.metadata.get('authors', 'Unknown Authors')) - + + title = paper.metadata.get("Title", paper.metadata.get("title", "No Title")) + authors = paper.metadata.get( + "Authors", paper.metadata.get("authors", "Unknown Authors") + ) + if isinstance(authors, list): - authors = ', '.join(authors) - - pdf_line = '' + authors = ", ".join(authors) + + pdf_line = "" if out_dir: # http://arxiv.org/abs/2101.00001v1 -> https://arxiv.org/pdf/2101.00001v1 - arxiv_id = url.rstrip('/').rsplit('/abs/', 1)[-1] if '/abs/' in url else None + arxiv_id = ( + url.rstrip("/").rsplit("/abs/", 1)[-1] if "/abs/" in url else None + ) pdf_urls = [f"https://arxiv.org/pdf/{arxiv_id}"] if arxiv_id else [] - pdf_path, msg = download_paper(out_dir, f"arXiv_{arxiv_id}" if arxiv_id else title, - doi=paper.metadata.get('DOI'), pdf_urls=pdf_urls) + pdf_path, msg = download_paper( + out_dir, + f"arXiv_{arxiv_id}" if arxiv_id else title, + doi=paper.metadata.get("DOI"), + pdf_urls=pdf_urls, + ) downloads.append((title, pdf_path, msg)) - pdf_line = '\n' + 'PDF: ' + (pdf_path or f'Not downloaded ({msg})') + pdf_line = "\n" + "PDF: " + (pdf_path or f"Not downloaded ({msg})") - data_list.append('Title: ' + title + '\n' + 'Url: ' + url + pdf_line + '\n' + 'Summary: ' + content) + data_list.append( + "Title: " + + title + + "\n" + + "Url: " + + url + + pdf_line + + "\n" + + "Summary: " + + content + ) if out_dir: report_downloads("ArXiv", query, downloads, out_dir) - return "\n".join(data_list) + '\n' + "Data source: ArXiv" + return "\n".join(data_list) + "\n" + "Data source: ArXiv" if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_Arxiv on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max papers per query (capped at 15).") - parser.add_argument("--out_dir", default=None, help="Download open-access PDFs here. Default: no downloads.") + + parser = argparse.ArgumentParser( + description="Run search_Arxiv on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", + type=int, + default=3, + help="Max papers per query (capped at 15).", + ) + parser.add_argument( + "--out_dir", + default=None, + help="Download open-access PDFs here. Default: no downloads.", + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_crossref.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_crossref.py index ce4787f..dff953f 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_crossref.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_crossref.py @@ -8,7 +8,9 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # llm_summary_agent / pdf_download +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # llm_summary_agent / pdf_download from llm_summary_agent import Summarize_Agent from pdf_download import download_paper, report_downloads @@ -29,7 +31,7 @@ def normalize_doi(value): doi = str(value).strip().lower() for prefix in ("https://doi.org/", "http://doi.org/", "http://dx.doi.org/", "doi:"): if doi.startswith(prefix): - doi = doi[len(prefix):] + doi = doi[len(prefix) :] break doi = doi.rstrip(".,;)") return doi if _DOI_RE.fullmatch(doi) else None @@ -62,7 +64,9 @@ def _get(url, params=None, timeout=30, retries=1): last = None for attempt in range(retries + 1): try: - resp = requests.get(url, params=params, timeout=timeout, headers={"User-Agent": USER_AGENT}) + resp = requests.get( + url, params=params, timeout=timeout, headers={"User-Agent": USER_AGENT} + ) if resp.status_code in TRANSIENT_STATUS and attempt < retries: time.sleep(0.5 * (attempt + 1)) continue @@ -81,34 +85,42 @@ def _crossref_record(item): return None titles = item.get("title") or [] title = titles[0] if isinstance(titles, list) and titles else str(titles or "") - date_parts = ((item.get("published") or {}).get("date-parts") - or (item.get("issued") or {}).get("date-parts") or []) + date_parts = ( + (item.get("published") or {}).get("date-parts") + or (item.get("issued") or {}).get("date-parts") + or [] + ) year = date_parts[0][0] if date_parts and date_parts[0] else None container = item.get("container-title") or [] relation = item.get("relation") or {} - authors = [a.get("family") or a.get("name") or "" for a in (item.get("author") or [])] + authors = [ + a.get("family") or a.get("name") or "" for a in (item.get("author") or []) + ] authors = [a for a in authors if a] return { "doi": normalize_doi(item.get("DOI")), "title": _clean(title) or "No Title", - "authors": ", ".join(authors[:3]) + (" et al." if len(authors) > 3 else "") if authors else None, + "authors": ", ".join(authors[:3]) + (" et al." if len(authors) > 3 else "") + if authors + else None, "year": year, "journal": container[0] if container else None, "study_type": item.get("type"), "abstract": _clean(item.get("abstract")), - "is_preprint": str(item.get("subtype") or "").casefold() == "preprint" or item.get("type") == "posted-content", + "is_preprint": str(item.get("subtype") or "").casefold() == "preprint" + or item.get("type") == "posted-content", "is_retracted": "is-retracted-by" in relation, # publisher-deposited PDF links, tried first by download_paper - "pdf_urls": [link.get("URL") for link in (item.get("link") or []) - if link.get("URL") and "pdf" in str(link.get("content-type", "")).lower()], + "pdf_urls": [ + link.get("URL") + for link in (item.get("link") or []) + if link.get("URL") and "pdf" in str(link.get("content-type", "")).lower() + ], } def search_Crossref( - query: str = "", - max_results: int = 3, - mini_handler = None, - out_dir: str = None + query: str = "", max_results: int = 3, mini_handler=None, out_dir: str = None ) -> str: """Search Crossref for works related to the query, or resolve a single DOI @@ -138,7 +150,9 @@ def search_Crossref( doi = normalize_doi(query) if doi: - payload = _get(f"{CROSSREF_BASE_URL}/works/{quote(doi, safe='')}", params=params) + payload = _get( + f"{CROSSREF_BASE_URL}/works/{quote(doi, safe='')}", params=params + ) items = [(payload or {}).get("message", {})] else: payload = _get( @@ -146,7 +160,9 @@ def search_Crossref( params={"query.bibliographic": query, "rows": max_results, **params}, ) items = (payload or {}).get("message", {}).get("items", []) - records = [r for r in (_crossref_record(item) for item in items) if r][:max_results] + records = [r for r in (_crossref_record(item) for item in items) if r][ + :max_results + ] except Exception as e: return f"Error during Crossref search: {e}" @@ -158,34 +174,79 @@ def search_Crossref( downloads = [] for rec in records: url = f"https://doi.org/{rec['doi']}" if rec["doi"] else "No URL" - tags = "".join(t for t, flag in ((" [RETRACTED]", rec["is_retracted"]), (" [PREPRINT]", rec["is_preprint"])) if flag) + tags = "".join( + t + for t, flag in ( + (" [RETRACTED]", rec["is_retracted"]), + (" [PREPRINT]", rec["is_preprint"]), + ) + if flag + ) info = f"Authors: {rec['authors']} | Year: {rec['year']} | Journal: {rec['journal']} | Type: {rec['study_type']}" content = "No abstract available." if rec["abstract"]: - content = Summarize_Agent(rec["abstract"], mini_handler) if mini_handler else rec["abstract"] + content = ( + Summarize_Agent(rec["abstract"], mini_handler) + if mini_handler + else rec["abstract"] + ) - pdf_line = '' + pdf_line = "" if out_dir: - pdf_path, msg = download_paper(out_dir, rec["doi"] or rec["title"], doi=rec["doi"], pdf_urls=rec["pdf_urls"]) + pdf_path, msg = download_paper( + out_dir, + rec["doi"] or rec["title"], + doi=rec["doi"], + pdf_urls=rec["pdf_urls"], + ) downloads.append((rec["title"], pdf_path, msg)) - pdf_line = '\n' + 'PDF: ' + (pdf_path or f'Not downloaded ({msg})') - - data_list.append('Title: ' + rec["title"] + tags + '\n' + 'Url: ' + url + '\n' + info + pdf_line + '\n' + 'Summary: ' + content) + pdf_line = "\n" + "PDF: " + (pdf_path or f"Not downloaded ({msg})") + + data_list.append( + "Title: " + + rec["title"] + + tags + + "\n" + + "Url: " + + url + + "\n" + + info + + pdf_line + + "\n" + + "Summary: " + + content + ) if out_dir: report_downloads("Crossref", query, downloads, out_dir) - return "\n".join(data_list) + '\n' + "Data source: Crossref" + return "\n".join(data_list) + "\n" + "Data source: Crossref" if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_Crossref on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max papers per query (capped at 15).") - parser.add_argument("--out_dir", default=None, help="Download open-access PDFs here. Default: no downloads.") + + parser = argparse.ArgumentParser( + description="Run search_Crossref on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", + type=int, + default=3, + help="Max papers per query (capped at 15).", + ) + parser.add_argument( + "--out_dir", + default=None, + help="Download open-access PDFs here. Default: no downloads.", + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") - print(search_Crossref(query=q, max_results=args.max_results, out_dir=args.out_dir)) + print( + search_Crossref(query=q, max_results=args.max_results, out_dir=args.out_dir) + ) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_generif.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_generif.py index 469b558..cf59bde 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_generif.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_generif.py @@ -13,6 +13,7 @@ Env: ANTHROPIC_API_KEY, NCBI_EMAIL / NCBI_API_KEY, UNPAYWALL_EMAIL; read from AGeneTic/.env when present. """ + import argparse import json import re @@ -22,18 +23,30 @@ from pathlib import Path HERE = Path(__file__).resolve().parent -LIT_DIR = HERE.parent # 1.1.Search_literature: src / llm_summary_agent / pdf_download +LIT_DIR = HERE.parent # 1.1.Search_literature: src / llm_summary_agent / pdf_download sys.path[:0] = [str(HERE), str(LIT_DIR)] from llm_summary_agent import Summarize_GeneRIF_Agent # noqa: E402 from pdf_download import download_paper # noqa: E402 from search_ncbi import NCBI_BASE_URL, _get, _ncbi_params, _parse_pubmed_xml # noqa: E402 -from src import (attach_gene_block, default_out_dir, describe_gene, gene_entries, # noqa: E402 - load_env, load_program_JSON, make_handler, write_bundles, write_meta) +from src import ( + attach_gene_block, + default_out_dir, + describe_gene, + gene_entries, # noqa: E402 + load_env, + load_program_JSON, + make_handler, + write_bundles, + write_meta, +) SOURCE = "GeneRIF" STOP_WORDS = {"cell", "cells", "line", "human", "mouse", "primary"} -_METHODS_RE = re.compile(r"\n\s*(materials?\s+and\s+methods|methods|experimental procedures)\s*\n", re.IGNORECASE) +_METHODS_RE = re.compile( + r"\n\s*(materials?\s+and\s+methods|methods|experimental procedures)\s*\n", + re.IGNORECASE, +) # query @@ -41,24 +54,35 @@ def fetch_records(pmids): """{pmid: PubMed record} via one efetch call.""" if not pmids: return {} - xml_text = _get(f"{NCBI_BASE_URL}/efetch.fcgi", - params={"db": "pubmed", "id": ",".join(pmids), "retmode": "xml", **_ncbi_params()}, - as_json=False) + xml_text = _get( + f"{NCBI_BASE_URL}/efetch.fcgi", + params={ + "db": "pubmed", + "id": ",".join(pmids), + "retmode": "xml", + **_ncbi_params(), + }, + as_json=False, + ) time.sleep(0.34) return {r["pmid"]: r for r in _parse_pubmed_xml(str(xml_text))} def context_terms(cell_type, extra): words = re.findall(r"[A-Za-z0-9-]+", cell_type or "") - terms = [w for w in words if len(w) > 3 and w.lower() not in STOP_WORDS] + list(extra or []) + terms = [w for w in words if len(w) > 3 and w.lower() not in STOP_WORDS] + list( + extra or [] + ) return [t.lower().rstrip("*") for t in dict.fromkeys(terms)] def rank_records(records, terms): """Cell-type-matching papers first, then primary research over reviews; stable otherwise.""" + def score(r): text = f"{r.get('title') or ''} {r.get('abstract') or ''}".lower() return (-any(t in text for t in terms), r.get("study_type") == "review") + return sorted(records, key=score) @@ -66,6 +90,7 @@ def pdf_text(path, max_chars): """Methods-first excerpt of a PDF (the part that states cell lines, species and conditions).""" try: import fitz + with fitz.open(path) as doc: text = "\n".join(page.get_text() for page in doc) except Exception as e: # a broken PDF must not stop the gene @@ -73,13 +98,18 @@ def pdf_text(path, max_chars): return None m = _METHODS_RE.search(text) start = m.start() if m else 0 - return text[start:start + max_chars].strip() or None + return text[start : start + max_chars].strip() or None # per gene def cache_key(pmids, cell_type, args): - return {"pmids": pmids, "cell_type": cell_type, "model": args.model, - "download_pdfs": args.download_pdfs, "pdf_chars": args.pdf_chars} + return { + "pmids": pmids, + "cell_type": cell_type, + "model": args.model, + "download_pdfs": args.download_pdfs, + "pdf_chars": args.pdf_chars, + } def summarize_gene(gene, entry, cell_type, terms, handler, out_dir, args): @@ -89,10 +119,14 @@ def summarize_gene(gene, entry, cell_type, terms, handler, out_dir, args): return {"found": False, "reason": "no GeneRIF PMIDs in gene_info.NCBI"}, False records = fetch_records([str(p) for p in stored]) - ranked = rank_records([records[str(p)] for p in stored if str(p) in records], terms)[:args.top_n] + ranked = rank_records( + [records[str(p)] for p in stored if str(p) in records], terms + )[: args.top_n] pmids = [r["pmid"] for r in ranked] - cache_file = out_dir / "gene_cache" / f"{re.sub(r'[^A-Za-z0-9._-]+', '_', gene)}.json" + cache_file = ( + out_dir / "gene_cache" / f"{re.sub(r'[^A-Za-z0-9._-]+', '_', gene)}.json" + ) key = cache_key(pmids, cell_type, args) if cache_file.exists() and not args.overwrite: cached = json.loads(cache_file.read_text()) @@ -101,22 +135,48 @@ def summarize_gene(gene, entry, cell_type, terms, handler, out_dir, args): papers = [] for r in ranked: - paper = {k: r.get(k) for k in ("pmid", "title", "year", "journal", "doi", "pmcid", "study_type", - "is_retracted", "abstract")} + paper = { + k: r.get(k) + for k in ( + "pmid", + "title", + "year", + "journal", + "doi", + "pmcid", + "study_type", + "is_retracted", + "abstract", + ) + } paper.update(pdf_path=None, pdf_status="not requested", fulltext=None) if args.download_pdfs: - path, msg = download_paper(out_dir / "gene_pdfs" / gene, r["pmid"], - pmid=r["pmid"], pmcid=r.get("pmcid"), doi=r.get("doi")) + path, msg = download_paper( + out_dir / "gene_pdfs" / gene, + r["pmid"], + pmid=r["pmid"], + pmcid=r.get("pmcid"), + doi=r.get("doi"), + ) paper.update(pdf_path=path, pdf_status=msg) if path: paper["fulltext"] = pdf_text(path, args.pdf_chars) papers.append(paper) - summary = Summarize_GeneRIF_Agent(gene, describe_gene(gene, entry), papers, cell_type, handler, - effort=args.effort, max_tokens=args.max_tokens) + summary = Summarize_GeneRIF_Agent( + gene, + describe_gene(gene, entry), + papers, + cell_type, + handler, + effort=args.effort, + max_tokens=args.max_tokens, + ) extracted = {p["pmid"]: p for p in (summary or {}).get("papers", [])} for paper in papers: - paper.update({k: v for k, v in extracted.get(paper["pmid"], {}).items() if k != "pmid"}) + paper.update( + {k: v for k, v in extracted.get(paper["pmid"], {}).items() if k != "pmid"} + ) paper["text_source"] = "abstract+pdf" if paper.pop("fulltext") else "abstract" paper.pop("abstract") block = { @@ -128,34 +188,86 @@ def summarize_gene(gene, entry, cell_type, terms, handler, out_dir, args): "llm_error": summary is None, "cache_file": str(cache_file), } - if summary is not None: # cache only successful extractions + if summary is not None: # cache only successful extractions cache_file.parent.mkdir(parents=True, exist_ok=True) cache_file.write_text(json.dumps({"key": key, "block": block}, indent=2)) return block, True def build_parser(): - p = argparse.ArgumentParser(description="Summarize the GeneRIF papers of every program gene (claim, species, cell line, condition).") + p = argparse.ArgumentParser( + description="Summarize the GeneRIF papers of every program gene (claim, species, cell line, condition)." + ) # IO - p.add_argument("--info_dir", required=True, help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.") - p.add_argument("--out_dir", default=None, help=f"Output folder. Default: /../{default_out_dir('x').name}.") + p.add_argument( + "--info_dir", + required=True, + help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.", + ) + p.add_argument( + "--out_dir", + default=None, + help=f"Output folder. Default: /../{default_out_dir('x').name}.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3).") - p.add_argument("--cell_type", default=None, help="Target cell type for context matching. Default: the bundle's cell_type.") - p.add_argument("--context_terms", nargs="*", default=[], help="Extra words that mark a paper as cell-type matched when ranking (e.g. endothelial HUVEC HAEC).") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3).", + ) + p.add_argument( + "--cell_type", + default=None, + help="Target cell type for context matching. Default: the bundle's cell_type.", + ) + p.add_argument( + "--context_terms", + nargs="*", + default=[], + help="Extra words that mark a paper as cell-type matched when ranking (e.g. endothelial HUVEC HAEC).", + ) # query - p.add_argument("--top_n", type=int, default=5, help="GeneRIF papers summarized per gene, after ranking. The pool is gene_info.NCBI.generif_pmids, capped upstream by search_NCBI.py --top_generif.") - p.add_argument("--download_pdfs", action="store_true", help="Download open-access PDFs into /gene_pdfs// and give their text to the summary agent.") - p.add_argument("--pdf_chars", type=int, default=15000, help="Max PDF characters (from the Methods section when found) given to the agent per paper.") + p.add_argument( + "--top_n", + type=int, + default=5, + help="GeneRIF papers summarized per gene, after ranking. The pool is gene_info.NCBI.generif_pmids, capped upstream by search_NCBI.py --top_generif.", + ) + p.add_argument( + "--download_pdfs", + action="store_true", + help="Download open-access PDFs into /gene_pdfs// and give their text to the summary agent.", + ) + p.add_argument( + "--pdf_chars", + type=int, + default=15000, + help="Max PDF characters (from the Methods section when found) given to the agent per paper.", + ) # agent - p.add_argument("--model", default="claude-sonnet-5", help="Claude model for the summary agent.") - p.add_argument("--effort", default="medium", choices=["low", "medium", "high", "xhigh", "max"], help="Effort level of the summary agent.") - p.add_argument("--max_tokens", type=int, default=16000, help="Max output tokens per gene call.") - p.add_argument("--overwrite", action="store_true", help="Ignore gene_cache/ and re-summarize every gene.") + p.add_argument( + "--model", default="claude-sonnet-5", help="Claude model for the summary agent." + ) + p.add_argument( + "--effort", + default="medium", + choices=["low", "medium", "high", "xhigh", "max"], + help="Effort level of the summary agent.", + ) + p.add_argument( + "--max_tokens", type=int, default=16000, help="Max output tokens per gene call." + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Ignore gene_cache/ and re-summarize every gene.", + ) return p @@ -168,21 +280,29 @@ def main(): blocks, n_llm = {}, 0 for label, bundle in bundles.items(): - cell_type = args.cell_type if args.cell_type is not None else bundle.get("cell_type", "") + cell_type = ( + args.cell_type + if args.cell_type is not None + else bundle.get("cell_type", "") + ) terms = context_terms(cell_type, args.context_terms) entries = gene_entries(bundle) print(f"[{label}] {len(entries)} genes; context terms: {terms}") for gene, entry in entries.items(): - if gene in blocks: # same gene in an earlier program of this run + if gene in blocks: # same gene in an earlier program of this run continue - block, ran = summarize_gene(gene, entry, cell_type, terms, handler, out_dir, args) + block, ran = summarize_gene( + gene, entry, cell_type, terms, handler, out_dir, args + ) blocks[gene] = block n_llm += ran if block["found"]: n_pdf = sum(bool(p["pdf_path"]) for p in block["papers"]) - print(f" {gene}: {len(block['papers'])} papers, {n_pdf} PDFs, " - f"in cell type={block['evidence_in_cell_type']}" - f"{' (cached)' if not ran else ''}{' [LLM error]' if block['llm_error'] else ''}") + print( + f" {gene}: {len(block['papers'])} papers, {n_pdf} PDFs, " + f"in cell type={block['evidence_in_cell_type']}" + f"{' (cached)' if not ran else ''}{' [LLM error]' if block['llm_error'] else ''}" + ) else: print(f" {gene}: {block['reason']}") bundles[label] = attach_gene_block(bundle, SOURCE, blocks) @@ -192,8 +312,19 @@ def main(): meta = { "created": datetime.now().isoformat(timespec="seconds"), "inputs": {"info_dir": str(args.info_dir)}, - "params": {k: getattr(args, k) for k in ("cell_type", "context_terms", "top_n", "download_pdfs", - "pdf_chars", "model", "effort", "max_tokens")}, + "params": { + k: getattr(args, k) + for k in ( + "cell_type", + "context_terms", + "top_n", + "download_pdfs", + "pdf_chars", + "model", + "effort", + "max_tokens", + ) + }, "programs": list(bundles), "n_genes": len(blocks), "n_with_generif": len(found), diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_nature.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_nature.py index ee9f40e..ce5155b 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_nature.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_nature.py @@ -7,7 +7,9 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # llm_summary_agent / pdf_download +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # llm_summary_agent / pdf_download from llm_summary_agent import Summarize_Agent from pdf_download import download_paper, report_downloads @@ -27,7 +29,9 @@ TRANSIENT_STATUS = {408, 425, 429, 500, 502, 503, 504} _TAG_RE = re.compile(r"<[^>]+>") -_DESCRIPTION_RE = re.compile(r' s41586-020-2012-7 + article_id = doi.split("/", 1)[1] # 10.1038/s41586-020-2012-7 -> s41586-020-2012-7 titles = item.get("title") or [] - date_parts = ((item.get("published") or {}).get("date-parts") - or (item.get("issued") or {}).get("date-parts") or []) + date_parts = ( + (item.get("published") or {}).get("date-parts") + or (item.get("issued") or {}).get("date-parts") + or [] + ) container = item.get("container-title") or [] - authors = [a.get("family") or a.get("name") or "" for a in (item.get("author") or [])] + authors = [ + a.get("family") or a.get("name") or "" for a in (item.get("author") or []) + ] authors = [a for a in authors if a] return { "doi": doi, @@ -98,7 +111,9 @@ def _nature_record(item): "url": f"{NATURE_BASE_URL}/articles/{article_id}", "pdf_url": f"{NATURE_BASE_URL}/articles/{article_id}.pdf", "title": _clean(titles[0] if titles else "") or "No Title", - "authors": ", ".join(authors[:3]) + (" et al." if len(authors) > 3 else "") if authors else None, + "authors": ", ".join(authors[:3]) + (" et al." if len(authors) > 3 else "") + if authors + else None, "year": date_parts[0][0] if date_parts and date_parts[0] else None, "journal": container[0] if container else None, "study_type": item.get("type"), @@ -108,10 +123,7 @@ def _nature_record(item): def search_Nature( - query: str = "", - max_results: int = 3, - mini_handler = None, - out_dir: str = None + query: str = "", max_results: int = 3, mini_handler=None, out_dir: str = None ) -> str: """Search Nature portfolio journals (Nature, Nature Genetics, Nature Communications, ...) for articles related to the query and optionally download their PDFs @@ -135,14 +147,19 @@ def search_Nature( try: query = _bounded_query(query) max_results = _bounded_limit(max_results, MAX_SEARCH_RESULTS) - params = {"query.bibliographic": query, "rows": max_results, - "filter": f"prefix:{NATURE_DOI_PREFIX},type:journal-article"} + params = { + "query.bibliographic": query, + "rows": max_results, + "filter": f"prefix:{NATURE_DOI_PREFIX},type:journal-article", + } email = os.getenv("CROSSREF_MAILTO") or os.getenv("PUBMED_EMAIL") if email: params["mailto"] = email payload = _get(f"{CROSSREF_BASE_URL}/works", params=params) items = (payload or {}).get("message", {}).get("items", []) - records = [r for r in (_nature_record(item) for item in items) if r][:max_results] + records = [r for r in (_nature_record(item) for item in items) if r][ + :max_results + ] except Exception as e: return f"Error during Nature search: {e}" @@ -159,34 +176,75 @@ def search_Nature( abstract = rec["abstract"] or _landing_page_abstract(rec["url"]) content = "No abstract available." if abstract: - content = Summarize_Agent(abstract, mini_handler) if mini_handler else abstract + content = ( + Summarize_Agent(abstract, mini_handler) if mini_handler else abstract + ) - pdf_line = '' + pdf_line = "" if out_dir: if i: - time.sleep(PDF_PAUSE_SECONDS) # polite spacing under the institutional license - pdf_path, msg = download_paper(out_dir, f"nature_{rec['article_id']}", - doi=rec["doi"], pdf_urls=[rec["pdf_url"]]) + time.sleep( + PDF_PAUSE_SECONDS + ) # polite spacing under the institutional license + pdf_path, msg = download_paper( + out_dir, + f"nature_{rec['article_id']}", + doi=rec["doi"], + pdf_urls=[rec["pdf_url"]], + ) if not pdf_path: msg += "; is this running on the Stanford network (Sherlock / campus / VPN)?" downloads.append((rec["title"], pdf_path, msg)) - pdf_line = '\n' + 'PDF: ' + (pdf_path or f'Not downloaded ({msg})') - - data_list.append('Title: ' + rec["title"] + tags + '\n' + 'Url: ' + rec["url"] + '\n' + info + pdf_line + '\n' + 'Summary: ' + content) + pdf_line = "\n" + "PDF: " + (pdf_path or f"Not downloaded ({msg})") + + data_list.append( + "Title: " + + rec["title"] + + tags + + "\n" + + "Url: " + + rec["url"] + + "\n" + + info + + pdf_line + + "\n" + + "Summary: " + + content + ) if out_dir: report_downloads("Nature", query, downloads, out_dir) - return "\n".join(data_list) + '\n' + "Data source: Nature (nature.com, Stanford subscription)" + return ( + "\n".join(data_list) + + "\n" + + "Data source: Nature (nature.com, Stanford subscription)" + ) if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_Nature on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max papers per query (capped at 15).") - parser.add_argument("--out_dir", default=None, help="Download open-access PDFs here. Default: no downloads.") + + parser = argparse.ArgumentParser( + description="Run search_Nature on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", + type=int, + default=3, + help="Max papers per query (capped at 15).", + ) + parser.add_argument( + "--out_dir", + default=None, + help="Download open-access PDFs here. Default: no downloads.", + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") - print(search_Nature(query=q, max_results=args.max_results, out_dir=args.out_dir)) + print( + search_Nature(query=q, max_results=args.max_results, out_dir=args.out_dir) + ) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_ncbi.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_ncbi.py index b6e2dfc..9df6a0d 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_ncbi.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_ncbi.py @@ -7,7 +7,9 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # llm_summary_agent / pdf_download +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # llm_summary_agent / pdf_download from llm_summary_agent import Summarize_Agent from pdf_download import download_paper, report_downloads @@ -29,7 +31,7 @@ def normalize_doi(value): doi = str(value).strip().lower() for prefix in ("https://doi.org/", "http://doi.org/", "http://dx.doi.org/", "doi:"): if doi.startswith(prefix): - doi = doi[len(prefix):] + doi = doi[len(prefix) :] break doi = doi.rstrip(".,;)") return doi if _DOI_RE.fullmatch(doi) else None @@ -73,7 +75,9 @@ def _get(url, params=None, as_json=True, timeout=30, retries=1): last = None for attempt in range(retries + 1): try: - resp = requests.get(url, params=params, timeout=timeout, headers={"User-Agent": USER_AGENT}) + resp = requests.get( + url, params=params, timeout=timeout, headers={"User-Agent": USER_AGENT} + ) if resp.status_code in TRANSIENT_STATUS and attempt < retries: time.sleep(0.5 * (attempt + 1)) continue @@ -94,8 +98,11 @@ def _node_text(node): def _pubmed_year(article): - for path in (".//ArticleDate/Year", ".//Journal/JournalIssue/PubDate/Year", - ".//Journal/JournalIssue/PubDate/MedlineDate"): + for path in ( + ".//ArticleDate/Year", + ".//Journal/JournalIssue/PubDate/Year", + ".//Journal/JournalIssue/PubDate/MedlineDate", + ): text = _node_text(article.find(path)) if text: m = _YEAR_RE.search(text) @@ -106,11 +113,19 @@ def _pubmed_year(article): def _study_type(pub_types): lowered = [t.casefold() for t in pub_types] - for key, label in (("review", "review"), ("clinical trial", "clinical trial"), - ("meta-analysis", "meta-analysis"), ("randomized", "randomized trial")): + for key, label in ( + ("review", "review"), + ("clinical trial", "clinical trial"), + ("meta-analysis", "meta-analysis"), + ("randomized", "randomized trial"), + ): if any(key in t for t in lowered): return label - non_generic = [t for t in pub_types if t.casefold() not in {"journal article", "research support"}] + non_generic = [ + t + for t in pub_types + if t.casefold() not in {"journal article", "research support"} + ] return non_generic[0] if non_generic else (pub_types[0] if pub_types else None) @@ -138,33 +153,47 @@ def _parse_pubmed_xml(xml_text): if text: label = node.attrib.get("Label") abstract_parts.append(f"{label}: {text}" if label else text) - pub_types = [t for t in (_node_text(n) for n in article.findall(".//PublicationTypeList/PublicationType")) if t] - comments = [node.attrib.get("RefType", "") - for node in medline.findall(".//CommentsCorrectionsList/CommentsCorrections")] + pub_types = [ + t + for t in ( + _node_text(n) + for n in article.findall(".//PublicationTypeList/PublicationType") + ) + if t + ] + comments = [ + node.attrib.get("RefType", "") + for node in medline.findall( + ".//CommentsCorrectionsList/CommentsCorrections" + ) + ] journal_node = article.find(".//Journal") - journal = _node_text(journal_node.find("Title")) if journal_node is not None else None - retracted = any(t.casefold() in {"retracted publication", "retraction of publication"} for t in pub_types) \ - or any("retraction" in c.casefold() for c in comments) - records.append({ - "pmid": pmid, - "doi": normalize_doi(ids.get("doi")), - "pmcid": ids.get("pmc") or None, - "title": _node_text(article.find("ArticleTitle")) or "No Title", - "year": _pubmed_year(article), - "journal": journal, - "study_type": _study_type(pub_types), - "abstract": " ".join(abstract_parts) or None, - "is_preprint": any("preprint" in t.casefold() for t in pub_types), - "is_retracted": retracted, - }) + journal = ( + _node_text(journal_node.find("Title")) if journal_node is not None else None + ) + retracted = any( + t.casefold() in {"retracted publication", "retraction of publication"} + for t in pub_types + ) or any("retraction" in c.casefold() for c in comments) + records.append( + { + "pmid": pmid, + "doi": normalize_doi(ids.get("doi")), + "pmcid": ids.get("pmc") or None, + "title": _node_text(article.find("ArticleTitle")) or "No Title", + "year": _pubmed_year(article), + "journal": journal, + "study_type": _study_type(pub_types), + "abstract": " ".join(abstract_parts) or None, + "is_preprint": any("preprint" in t.casefold() for t in pub_types), + "is_retracted": retracted, + } + ) return records def search_NCBI( - query: str = "", - max_results: int = 3, - mini_handler = None, - out_dir: str = None + query: str = "", max_results: int = 3, mini_handler=None, out_dir: str = None ) -> str: """Search PubMed through NCBI E-utilities (esearch -> efetch) for articles related to the query @@ -196,18 +225,31 @@ def search_NCBI( # Search for PMIDs related to the query payload = _get( f"{NCBI_BASE_URL}/esearch.fcgi", - params={"db": "pubmed", "term": query, "retmax": max_results, - "retmode": "json", "sort": "relevance", **_ncbi_params()}, + params={ + "db": "pubmed", + "term": query, + "retmax": max_results, + "retmode": "json", + "sort": "relevance", + **_ncbi_params(), + }, ) result = payload.get("esearchresult", {}) if isinstance(payload, dict) else {} - pmids = [p for p in (normalize_pmid(x) for x in result.get("idlist", [])) if p][:max_results] + pmids = [p for p in (normalize_pmid(x) for x in result.get("idlist", [])) if p][ + :max_results + ] if not pmids: return "No results found on NCBI PubMed." # Fetch canonical metadata for the PMIDs xml_text = _get( f"{NCBI_BASE_URL}/efetch.fcgi", - params={"db": "pubmed", "id": ",".join(pmids), "retmode": "xml", **_ncbi_params()}, + params={ + "db": "pubmed", + "id": ",".join(pmids), + "retmode": "xml", + **_ncbi_params(), + }, as_json=False, ) records = _parse_pubmed_xml(str(xml_text)) @@ -219,33 +261,77 @@ def search_NCBI( downloads = [] for rec in records: url = "https://pubmed.ncbi.nlm.nih.gov/" + rec["pmid"] - tags = "".join(t for t, flag in ((" [RETRACTED]", rec["is_retracted"]), (" [PREPRINT]", rec["is_preprint"])) if flag) + tags = "".join( + t + for t, flag in ( + (" [RETRACTED]", rec["is_retracted"]), + (" [PREPRINT]", rec["is_preprint"]), + ) + if flag + ) info = f"Year: {rec['year']} | Journal: {rec['journal']} | DOI: {rec['doi']} | Type: {rec['study_type']}" content = "No abstract available." if rec["abstract"]: - content = Summarize_Agent(rec["abstract"], mini_handler) if mini_handler else rec["abstract"] + content = ( + Summarize_Agent(rec["abstract"], mini_handler) + if mini_handler + else rec["abstract"] + ) - pdf_line = '' + pdf_line = "" if out_dir: - pdf_path, msg = download_paper(out_dir, rec["pmcid"] or rec["pmid"], pmid=rec["pmid"], pmcid=rec["pmcid"], doi=rec["doi"]) + pdf_path, msg = download_paper( + out_dir, + rec["pmcid"] or rec["pmid"], + pmid=rec["pmid"], + pmcid=rec["pmcid"], + doi=rec["doi"], + ) downloads.append((rec["title"], pdf_path, msg)) - pdf_line = '\n' + 'PDF: ' + (pdf_path or f'Not downloaded ({msg})') - - data_list.append('Title: ' + rec["title"] + tags + '\n' + 'Url: ' + url + '\n' + info + pdf_line + '\n' + 'Summary: ' + content) + pdf_line = "\n" + "PDF: " + (pdf_path or f"Not downloaded ({msg})") + + data_list.append( + "Title: " + + rec["title"] + + tags + + "\n" + + "Url: " + + url + + "\n" + + info + + pdf_line + + "\n" + + "Summary: " + + content + ) if out_dir: report_downloads("NCBI", query, downloads, out_dir) - return "\n".join(data_list) + '\n' + "Data source: NCBI PubMed (E-utilities)" + return "\n".join(data_list) + "\n" + "Data source: NCBI PubMed (E-utilities)" if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_NCBI on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max papers per query (capped at 15).") - parser.add_argument("--out_dir", default=None, help="Download open-access PDFs here. Default: no downloads.") + + parser = argparse.ArgumentParser( + description="Run search_NCBI on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", + type=int, + default=3, + help="Max papers per query (capped at 15).", + ) + parser.add_argument( + "--out_dir", + default=None, + help="Download open-access PDFs here. Default: no downloads.", + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_openalex.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_openalex.py index 3194a4d..3e55c33 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_openalex.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_openalex.py @@ -6,7 +6,9 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # llm_summary_agent / pdf_download +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # llm_summary_agent / pdf_download from llm_summary_agent import Summarize_Agent from pdf_download import download_paper, report_downloads @@ -27,7 +29,7 @@ def normalize_doi(value): doi = str(value).strip().lower() for prefix in ("https://doi.org/", "http://doi.org/", "http://dx.doi.org/", "doi:"): if doi.startswith(prefix): - doi = doi[len(prefix):] + doi = doi[len(prefix) :] break doi = doi.rstrip(".,;)") return doi if _DOI_RE.fullmatch(doi) else None @@ -58,7 +60,9 @@ def _bounded_limit(value, maximum): def _openalex_params(): key = os.getenv("OPENALEX_API_KEY") if not key: - raise RuntimeError("OPENALEX_API_KEY is not set; use search_NCBI / search_Crossref instead") + raise RuntimeError( + "OPENALEX_API_KEY is not set; use search_NCBI / search_Crossref instead" + ) params = {"api_key": key} email = os.getenv("OPENALEX_EMAIL") or os.getenv("OPENALEX_MAILTO") if email: @@ -71,7 +75,9 @@ def _get(url, params=None, timeout=30, retries=1): last = None for attempt in range(retries + 1): try: - resp = requests.get(url, params=params, timeout=timeout, headers={"User-Agent": USER_AGENT}) + resp = requests.get( + url, params=params, timeout=timeout, headers={"User-Agent": USER_AGENT} + ) if resp.status_code in TRANSIENT_STATUS and attempt < retries: time.sleep(0.5 * (attempt + 1)) continue @@ -96,7 +102,11 @@ def _rebuild_abstract(inverted_index): def _openalex_record(item): ids = item.get("ids") or {} pmid_raw = ids.get("pmid") - pmid = normalize_pmid(str(pmid_raw).rstrip("/").rsplit("/", 1)[-1]) if pmid_raw else None + pmid = ( + normalize_pmid(str(pmid_raw).rstrip("/").rsplit("/", 1)[-1]) + if pmid_raw + else None + ) source = (item.get("primary_location") or {}).get("source") or {} return { "openalex_id": item.get("id"), @@ -111,17 +121,19 @@ def _openalex_record(item): "is_retracted": bool(item.get("is_retracted")), "cited_by_count": int(item.get("cited_by_count") or 0), # open-access PDF locations, tried first by download_paper - "pdf_urls": list(dict.fromkeys( - loc.get("pdf_url") for loc in [item.get("best_oa_location") or {}] + (item.get("locations") or []) - if loc and loc.get("pdf_url"))), + "pdf_urls": list( + dict.fromkeys( + loc.get("pdf_url") + for loc in [item.get("best_oa_location") or {}] + + (item.get("locations") or []) + if loc and loc.get("pdf_url") + ) + ), } def search_OpenAlex( - query: str = "", - max_results: int = 3, - mini_handler = None, - out_dir: str = None + query: str = "", max_results: int = 3, mini_handler=None, out_dir: str = None ) -> str: """Search OpenAlex (cross-publisher, includes preprints) for works related to the query @@ -147,7 +159,9 @@ def search_OpenAlex( f"{OPENALEX_BASE_URL}/works", params={"search": query, "per-page": max_results, **_openalex_params()}, ) - records = [_openalex_record(item) for item in payload.get("results", [])][:max_results] + records = [_openalex_record(item) for item in payload.get("results", [])][ + :max_results + ] except Exception as e: return f"Error during OpenAlex search: {e}" @@ -158,36 +172,88 @@ def search_OpenAlex( data_list = [] downloads = [] for rec in records: - url = f"https://doi.org/{rec['doi']}" if rec["doi"] else (rec["openalex_id"] or "No URL") - tags = "".join(t for t, flag in ((" [RETRACTED]", rec["is_retracted"]), (" [PREPRINT]", rec["is_preprint"])) if flag) - info = (f"Year: {rec['year']} | Journal: {rec['journal']} | PMID: {rec['pmid']} | " - f"Type: {rec['study_type']} | Cited by: {rec['cited_by_count']}") + url = ( + f"https://doi.org/{rec['doi']}" + if rec["doi"] + else (rec["openalex_id"] or "No URL") + ) + tags = "".join( + t + for t, flag in ( + (" [RETRACTED]", rec["is_retracted"]), + (" [PREPRINT]", rec["is_preprint"]), + ) + if flag + ) + info = ( + f"Year: {rec['year']} | Journal: {rec['journal']} | PMID: {rec['pmid']} | " + f"Type: {rec['study_type']} | Cited by: {rec['cited_by_count']}" + ) content = "No abstract available." if rec["abstract"]: - content = Summarize_Agent(rec["abstract"], mini_handler) if mini_handler else rec["abstract"] + content = ( + Summarize_Agent(rec["abstract"], mini_handler) + if mini_handler + else rec["abstract"] + ) - pdf_line = '' + pdf_line = "" if out_dir: - pdf_path, msg = download_paper(out_dir, rec["doi"] or rec["pmid"] or rec["openalex_id"], pmid=rec["pmid"], doi=rec["doi"], pdf_urls=rec["pdf_urls"]) + pdf_path, msg = download_paper( + out_dir, + rec["doi"] or rec["pmid"] or rec["openalex_id"], + pmid=rec["pmid"], + doi=rec["doi"], + pdf_urls=rec["pdf_urls"], + ) downloads.append((rec["title"], pdf_path, msg)) - pdf_line = '\n' + 'PDF: ' + (pdf_path or f'Not downloaded ({msg})') - - data_list.append('Title: ' + rec["title"] + tags + '\n' + 'Url: ' + url + '\n' + info + pdf_line + '\n' + 'Summary: ' + content) + pdf_line = "\n" + "PDF: " + (pdf_path or f"Not downloaded ({msg})") + + data_list.append( + "Title: " + + rec["title"] + + tags + + "\n" + + "Url: " + + url + + "\n" + + info + + pdf_line + + "\n" + + "Summary: " + + content + ) if out_dir: report_downloads("OpenAlex", query, downloads, out_dir) - return "\n".join(data_list) + '\n' + "Data source: OpenAlex" + return "\n".join(data_list) + "\n" + "Data source: OpenAlex" if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_OpenAlex on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max papers per query (capped at 15).") - parser.add_argument("--out_dir", default=None, help="Download open-access PDFs here. Default: no downloads.") + + parser = argparse.ArgumentParser( + description="Run search_OpenAlex on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", + type=int, + default=3, + help="Max papers per query (capped at 15).", + ) + parser.add_argument( + "--out_dir", + default=None, + help="Download open-access PDFs here. Default: no downloads.", + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") - print(search_OpenAlex(query=q, max_results=args.max_results, out_dir=args.out_dir)) + print( + search_OpenAlex(query=q, max_results=args.max_results, out_dir=args.out_dir) + ) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_pubmed.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_pubmed.py index 3b066a5..1440ee7 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_pubmed.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.1.Search_literature/search_pubmed.py @@ -2,17 +2,17 @@ import sys from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # llm_summary_agent / pdf_download +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # llm_summary_agent / pdf_download from llm_summary_agent import Summarize_Agent from pdf_download import download_paper, report_downloads + # PubMed API # @tool def search_PubMed( - query: str = "", - max_results: int = 3, - mini_handler = None, - out_dir: str = None + query: str = "", max_results: int = 3, mini_handler=None, out_dir: str = None ) -> list[dict]: """Search PubMed for articles related to the query @@ -29,7 +29,7 @@ def search_PubMed( 2. Boolean Operators: - AND - - OR + - OR - NOT - Use parentheses () for complex queries @@ -59,37 +59,62 @@ def search_PubMed( downloads = [] for article in related_articles: # ipdb.set_trace() - url = "https://pubmed.ncbi.nlm.nih.gov/" + article.metadata.get('uid', '') + url = "https://pubmed.ncbi.nlm.nih.gov/" + article.metadata.get("uid", "") content = "No abstract available." if article.page_content: content = Summarize_Agent(article.page_content, mini_handler) - title = article.metadata.get('Title', 'No Title') + title = article.metadata.get("Title", "No Title") if type(title) == dict: - title = title['#text'] + title = title["#text"] - pdf_line = '' + pdf_line = "" if out_dir: - pmid = article.metadata.get('uid', '') + pmid = article.metadata.get("uid", "") pdf_path, msg = download_paper(out_dir, pmid or title, pmid=pmid) downloads.append((title, pdf_path, msg)) - pdf_line = '\n ' + 'PDF: ' + (pdf_path or f'Not downloaded ({msg})') - - data_list.append('Title: ' + title + '\n '+ "Url: "+ url + pdf_line + '\n ' + 'Summary: ' + content) + pdf_line = "\n " + "PDF: " + (pdf_path or f"Not downloaded ({msg})") + + data_list.append( + "Title: " + + title + + "\n " + + "Url: " + + url + + pdf_line + + "\n " + + "Summary: " + + content + ) if out_dir: report_downloads("PubMed", query, downloads, out_dir) - return "\n".join(data_list) + '\n' + "Data source: PubMed" - + return "\n".join(data_list) + "\n" + "Data source: PubMed" if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_PubMed on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max papers per query (capped at 15).") - parser.add_argument("--out_dir", default=None, help="Download open-access PDFs here. Default: no downloads.") + + parser = argparse.ArgumentParser( + description="Run search_PubMed on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", + type=int, + default=3, + help="Max papers per query (capped at 15).", + ) + parser.add_argument( + "--out_dir", + default=None, + help="Download open-access PDFs here. Default: no downloads.", + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") - print(search_PubMed(query=q, max_results=args.max_results, out_dir=args.out_dir)) + print( + search_PubMed(query=q, max_results=args.max_results, out_dir=args.out_dir) + ) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_bing.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_bing.py index 8e53c0d..bde4a83 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_bing.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_bing.py @@ -8,12 +8,20 @@ from selenium.webdriver.chrome.options import Options from tools.page_fetch import fetch_page_content_and_summarize -os.environ['DISPLAY'] = ':99' + +os.environ["DISPLAY"] = ":99" ua = UserAgent() -def BingSearchTool(args, query: str, mini_handler, read_content: bool = True, return_num: int = 5, screenshot: bool = False): +def BingSearchTool( + args, + query: str, + mini_handler, + read_content: bool = True, + return_num: int = 5, + screenshot: bool = False, +): """ Use Selenium to get search results from Bing. Optionally, fetch content from the links and summarize. @@ -31,13 +39,12 @@ def BingSearchTool(args, query: str, mini_handler, read_content: bool = True, re driver = None results = [] - url = f"https://www.bing.com/search?q={query}" user_agent = ua.random options = Options() - if not getattr(args, 'visualize', False): + if not getattr(args, "visualize", False): options.add_argument("--headless") options.add_argument(f"user-agent={user_agent}") options.add_argument("--no-sandbox") @@ -48,7 +55,7 @@ def BingSearchTool(args, query: str, mini_handler, read_content: bool = True, re options.add_argument("--disable-blink-features=AutomationControlled") options.add_argument("--incognito") options.add_argument("--disable-dev-shm-usage") - options.add_argument("--blink-settings=imagesEnabled=false") + options.add_argument("--blink-settings=imagesEnabled=false") service = Service(args.chrome_driver) driver = webdriver.Chrome(service=service, options=options) @@ -81,25 +88,32 @@ def BingSearchTool(args, query: str, mini_handler, read_content: bool = True, re link = link_elem["href"] - if not (link.startswith('http://') or link.startswith('https://')): + if not (link.startswith("http://") or link.startswith("https://")): continue if read_content: try: - page_content = fetch_page_content_and_summarize(args, link, mini_handler, False) - results.append(f"Title: {title}\nURL: {link}\n\n Content:{page_content}") + page_content = fetch_page_content_and_summarize( + args, link, mini_handler, False + ) + results.append( + f"Title: {title}\nURL: {link}\n\n Content:{page_content}" + ) except: continue else: snippet_elem = result.find("p") - snippet = snippet_elem.get_text().strip() if snippet_elem else "No snippet available." + snippet = ( + snippet_elem.get_text().strip() + if snippet_elem + else "No snippet available." + ) results.append(f"Title: {title}\nSnippet: {snippet}\nURL: {link}") success_count += 1 if success_count >= return_num: break - if results: return "\n\n-----------------------\n\n".join(results) else: @@ -108,9 +122,24 @@ def BingSearchTool(args, query: str, mini_handler, read_content: bool = True, re if __name__ == "__main__": import argparse + parser = argparse.ArgumentParser(description="Bing Search Tool") - parser.add_argument('--chrome_driver', type=str, default="/usr/local/bin/chromedriver", help='Path to ChromeDriver') - parser.add_argument('--visualize', action='store_true', help='Visualize the search results') + parser.add_argument( + "--chrome_driver", + type=str, + default="/usr/local/bin/chromedriver", + help="Path to ChromeDriver", + ) + parser.add_argument( + "--visualize", action="store_true", help="Visualize the search results" + ) args = parser.parse_args() - result = BingSearchTool(args, "Genetic disease", mini_handler=None, read_content=False, return_num=5, screenshot=False) + result = BingSearchTool( + args, + "Genetic disease", + mini_handler=None, + read_content=False, + return_num=5, + screenshot=False, + ) print(result) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_duckduckgo.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_duckduckgo.py index d2d5cb2..b1d925e 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_duckduckgo.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_duckduckgo.py @@ -4,7 +4,9 @@ # @tool -def DuckDuckGoSearchTool(args, query: str, read_content: bool = True, return_num: int = 5, mini_handler=None): +def DuckDuckGoSearchTool( + args, query: str, read_content: bool = True, return_num: int = 5, mini_handler=None +): """ Use the DDGS (DuckDuckGo Search) to get search results. Optionally, fetch content from the links and summarize. @@ -21,27 +23,31 @@ def DuckDuckGoSearchTool(args, query: str, read_content: bool = True, return_num try: # Perform a DuckDuckGo search using DDGS ddgs = DDGS() - results = ddgs.text(query, max_results = return_num) # Modify max_results if needed + results = ddgs.text( + query, max_results=return_num + ) # Modify max_results if needed # If no results found if not results: return "No results found on DuckDuckGo." - search_results = [] for result in results: - title = result.get('title') - snippet = result.get('snippet', 'No snippet available.') - link = result.get('link') - + title = result.get("title") + snippet = result.get("snippet", "No snippet available.") + link = result.get("link") # If read_content is True, fetch and parse the content of the linked page if read_content: - page_content = fetch_page_content_and_summarize(args, link, mini_handler, False) + page_content = fetch_page_content_and_summarize( + args, link, mini_handler, False + ) search_results.append(page_content) else: - search_results.append(f"Title: {title}\nSnippet: {snippet}\nURL: {link}") + search_results.append( + f"Title: {title}\nSnippet: {snippet}\nURL: {link}" + ) return "\n\n".join(search_results) @@ -50,5 +56,7 @@ def DuckDuckGoSearchTool(args, query: str, read_content: bool = True, return_num if __name__ == "__main__": - result = DuckDuckGoSearchTool(None, "Genetic disease", read_content=False, return_num=5) + result = DuckDuckGoSearchTool( + None, "Genetic disease", read_content=False, return_num=5 + ) print(result) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_google.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_google.py index 16bbf37..cf08a92 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_google.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_google.py @@ -12,7 +12,7 @@ def GoogleSearchTool( read_content: bool = True, return_num: int = 5, screenshot: bool = False, - start: int = 1 + start: int = 1, ): """ Use API to get search results from Google. Optionally, fetch content from the links and summarize. @@ -44,23 +44,24 @@ def GoogleSearchTool( if not (1 <= return_num <= 10): return "Error: return_num must be between 1 and 10." - if start < 1 or start > 91: # The official limit is usually around 100, but we conservatively use 91 to ensure start+num-1 <= 100 + if ( + start < 1 or start > 91 + ): # The official limit is usually around 100, but we conservatively use 91 to ensure start+num-1 <= 100 return "Error: start must be between 1 and 91." try: api_key = args.google_api search_engine_id = args.search_engine_id - url = 'https://www.googleapis.com/customsearch/v1' + url = "https://www.googleapis.com/customsearch/v1" params = { - 'key': api_key, - 'cx': search_engine_id, - 'q': query.strip(), - 'num': return_num, - 'start': start + "key": api_key, + "cx": search_engine_id, + "q": query.strip(), + "num": return_num, + "start": start, } - response = requests.get(url, params=params, timeout=20) if response.status_code != 200: @@ -71,27 +72,33 @@ def GoogleSearchTool( return f"Error: {response.status_code}\n{err}" results_json = response.json() - items = results_json.get('items', []) + items = results_json.get("items", []) if not items: return "No results found." search_results = [] for i, item in enumerate(items): - title = item.get('title') - link = item.get('link') - snippet = item.get('snippet', 'No snippet available.') + title = item.get("title") + link = item.get("link") + snippet = item.get("snippet", "No snippet available.") # print(f"Result {i}: {title}\nSnippet: {snippet}\nURL: {link}\n") # ipdb.set_trace() if read_content and link: try: - page_content = fetch_page_content_and_summarize(args, link, mini_handler, False) + page_content = fetch_page_content_and_summarize( + args, link, mini_handler, False + ) search_results.append(page_content) except Exception as e: - search_results.append(f"Title: {title}\nSnippet: {snippet}\nURL: {link}") - + search_results.append( + f"Title: {title}\nSnippet: {snippet}\nURL: {link}" + ) + # print(f"Fetched and summarized content from: {link}\n") else: - search_results.append(f"Title: {title}\nSnippet: {snippet}\nURL: {link}") + search_results.append( + f"Title: {title}\nSnippet: {snippet}\nURL: {link}" + ) return "\n\n".join(search_results) @@ -103,9 +110,27 @@ def GoogleSearchTool( if __name__ == "__main__": import argparse + parser = argparse.ArgumentParser(description="Google Search Tool") - parser.add_argument('--google_api', type=str, default=os.getenv("GOOGLE_API_KEY"), help='Google API key') - parser.add_argument('--search_engine_id', type=str, default=os.getenv("GOOGLE_CSE_ID"), help='Google CSE id (cx)') + parser.add_argument( + "--google_api", + type=str, + default=os.getenv("GOOGLE_API_KEY"), + help="Google API key", + ) + parser.add_argument( + "--search_engine_id", + type=str, + default=os.getenv("GOOGLE_CSE_ID"), + help="Google CSE id (cx)", + ) args = parser.parse_args() - result = GoogleSearchTool(args, "Genetic disease", mini_handler=None, read_content=False, return_num=5, screenshot=False) + result = GoogleSearchTool( + args, + "Genetic disease", + mini_handler=None, + read_content=False, + return_num=5, + screenshot=False, + ) print(result) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_uptodate.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_uptodate.py index 5bbb6dd..0db24f8 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_uptodate.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_uptodate.py @@ -29,43 +29,53 @@ def UptodateSearchTool(args, query: str): try: url = f"https://www.uptodate.com/login" - headers = {'User-Agent': ua.random} + headers = {"User-Agent": ua.random} options = Options() if not args.visualize: options.add_argument("--headless") # Run in headless mode (no GUI) # options.add_argument("--no-sandbox") # This is helpful in certain environments (e.g., Docker) # options.add_argument("--disable-dev-shm-usage") # Helps in environments with limited shared memory - options.add_argument(f'user-agent={headers["User-Agent"]}') + options.add_argument(f"user-agent={headers['User-Agent']}") - service = Service(args.chrome_driver) # Update with the path to your chromedriver + service = Service( + args.chrome_driver + ) # Update with the path to your chromedriver driver = webdriver.Chrome(service=service, options=options) driver.get(url) time.sleep(1) - + # ipdb.set_trace() # click //*[@id="onetrust-accept-btn-handler"] try: element = WebDriverWait(driver, 10).until( - EC.element_to_be_clickable((By.XPATH, '//*[@id="onetrust-accept-btn-handler"]')) + EC.element_to_be_clickable( + (By.XPATH, '//*[@id="onetrust-accept-btn-handler"]') + ) ) element.click() except Exception as e: print(f"Error: {str(e)}") - + # input the username into //*[@id="username"] - driver.find_element(By.XPATH, '//*[@id="userName"]').send_keys(args.uptodate_user) - + driver.find_element(By.XPATH, '//*[@id="userName"]').send_keys( + args.uptodate_user + ) + # click continue //*[@id="loginForm"]/fieldset/div[2]/span - driver.find_element(By.XPATH, '//*[@id="loginForm"]/fieldset/div[2]/span').click() - + driver.find_element( + By.XPATH, '//*[@id="loginForm"]/fieldset/div[2]/span' + ).click() + # input the password into //*[@id="password"] - driver.find_element(By.XPATH, '//*[@id="password"]').send_keys(args.uptodate_pwd) - + driver.find_element(By.XPATH, '//*[@id="password"]').send_keys( + args.uptodate_pwd + ) + # find the login button - login_button = driver.find_element(By.ID, 'btnLoginSubmit') - + login_button = driver.find_element(By.ID, "btnLoginSubmit") + # maximum the window driver.maximize_window() @@ -75,42 +85,48 @@ def UptodateSearchTool(args, query: str): # input the query into //*[@id="tbSearch"] time.sleep(1) - driver.find_element(By.ID, 'tbSearch').send_keys(query) - + driver.find_element(By.ID, "tbSearch").send_keys(query) + # click //*[@id="search-combobox"]/div[1]/div/form/span[3] time.sleep(1) - driver.find_element(By.XPATH, '//*[@id="search-combobox"]/div[1]/div/form/span[3]').click() - + driver.find_element( + By.XPATH, '//*[@id="search-combobox"]/div[1]/div/form/span[3]' + ).click() + # wait for the search results to be fully rendered # WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CLASS_NAME, "search-result"))) - + # click first result try: element = WebDriverWait(driver, 10).until( - EC.element_to_be_clickable((By.CLASS_NAME, 'search-result')) + EC.element_to_be_clickable((By.CLASS_NAME, "search-result")) ) element.click() except Exception as e: print(f"Error: {str(e)}") - + element = WebDriverWait(driver, 10).until( - EC.element_to_be_clickable((By.CSS_SELECTOR, "a.searchResultLink.search-result-primary-link[role='heading']")) + EC.element_to_be_clickable( + ( + By.CSS_SELECTOR, + "a.searchResultLink.search-result-primary-link[role='heading']", + ) + ) ) element.click() - + time.sleep(1) - + # get all paragraphs all_paragraph = [] # WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CLASS_NAME, "topic-title"))) - soup = BeautifulSoup(driver.page_source, 'html.parser') - paragraphs = soup.find_all('p') + soup = BeautifulSoup(driver.page_source, "html.parser") + paragraphs = soup.find_all("p") for paragraph in paragraphs: all_paragraph.append(paragraph.text) - - all_paragraph = '\n'.join(all_paragraph) - + all_paragraph = "\n".join(all_paragraph) + return all_paragraph except Exception as e: return str(e) @@ -118,5 +134,5 @@ def UptodateSearchTool(args, query: str): driver.quit() -if __name__ == '__main__': - print(UptodateSearchTool('COVID-19')) +if __name__ == "__main__": + print(UptodateSearchTool("COVID-19")) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_wiki.py b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_wiki.py index a320658..a53cb14 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_wiki.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/1.2.Search_web/search_wiki.py @@ -3,27 +3,29 @@ import requests -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # 1.Search: llm_summary_agent +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1]) +) # 1.Search: llm_summary_agent from llm_summary_agent import Summarize_Agent # noqa: E402 # Wikipedia API (MediaWiki action API; no langchain_community needed) WIKI_API = "https://en.wikipedia.org/w/api.php" USER_AGENT = "PerturbNMF-AGeneTic/0.1 (literature search)" -MAX_CHARS = 4000 # page text kept per article (the intro and first sections) +MAX_CHARS = 4000 # page text kept per article (the intro and first sections) def _wiki_get(params): - resp = requests.get(WIKI_API, params={**params, "format": "json", "formatversion": 2}, - headers={"User-Agent": USER_AGENT}, timeout=30) + resp = requests.get( + WIKI_API, + params={**params, "format": "json", "formatversion": 2}, + headers={"User-Agent": USER_AGENT}, + timeout=30, + ) resp.raise_for_status() return resp.json() -def search_Wiki( - query: str = "", - max_results: int = 3, - mini_handler = None -) -> str: +def search_Wiki(query: str = "", max_results: int = 3, mini_handler=None) -> str: """Search Wikipedia for articles related to the query Args: @@ -33,13 +35,25 @@ def search_Wiki( """ try: # Search for article titles related to the query - hits = _wiki_get({"action": "query", "list": "search", "srsearch": query, - "srlimit": max_results})["query"]["search"] + hits = _wiki_get( + { + "action": "query", + "list": "search", + "srsearch": query, + "srlimit": max_results, + } + )["query"]["search"] if not hits: - return '' + return "" # Fetch plain-text page content for the titles - pages = _wiki_get({"action": "query", "prop": "extracts", "explaintext": 1, - "titles": "|".join(h["title"] for h in hits)})["query"]["pages"] + pages = _wiki_get( + { + "action": "query", + "prop": "extracts", + "explaintext": 1, + "titles": "|".join(h["title"] for h in hits), + } + )["query"]["pages"] except Exception as e: return f"Error during Wikipedia search: {e}" text = {p["title"]: p.get("extract") or "" for p in pages} @@ -50,18 +64,35 @@ def search_Wiki( title = hit["title"] url = "https://en.wikipedia.org/wiki/" + title.replace(" ", "_") page = text.get(title, "")[:MAX_CHARS] - content = Summarize_Agent(page, mini_handler) if page else 'No content available.' - data_list.append('Title: ' + title + '\n' + \ - ' Url: ' + url + '\n' + - 'Summary: ' + content) + content = ( + Summarize_Agent(page, mini_handler) if page else "No content available." + ) + data_list.append( + "Title: " + + title + + "\n" + + " Url: " + + url + + "\n" + + "Summary: " + + content + ) + + return "\n".join(data_list) + "\n" + "Data source: Wikipedia" - return '\n'.join(data_list) + '\n' + "Data source: Wikipedia" if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Run search_Wiki on one or more queries.") - parser.add_argument("--query", nargs="+", required=True, help="One or more queries, each quoted.") - parser.add_argument("--max_results", type=int, default=3, help="Max articles per query.") + + parser = argparse.ArgumentParser( + description="Run search_Wiki on one or more queries." + ) + parser.add_argument( + "--query", nargs="+", required=True, help="One or more queries, each quoted." + ) + parser.add_argument( + "--max_results", type=int, default=3, help="Max articles per query." + ) args = parser.parse_args() for q in args.query: print(f"===== {q}") diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/llm_query_agent.py b/src/Stage3_Interpretation/C_Annotation/1.Search/llm_query_agent.py index b6169b6..b07933d 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/llm_query_agent.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/llm_query_agent.py @@ -16,6 +16,7 @@ Env: ANTHROPIC_API_KEY (or AGeneTic/.env). """ + import argparse import math import sys @@ -28,8 +29,18 @@ HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) -from src import (default_out_dir, describe_gene, gene_aliases, gene_entries, load_env, # noqa: E402 - load_program_JSON, make_handler, regulator_stats, write_bundles, write_meta) +from src import ( + default_out_dir, + describe_gene, + gene_aliases, + gene_entries, + load_env, # noqa: E402 + load_program_JSON, + make_handler, + regulator_stats, + write_bundles, + write_meta, +) SOURCE = "QueryAgent" @@ -50,30 +61,50 @@ class QueryRung(BaseModel): - scope: str = Field(description="What this rung covers, e.g. 'gene + regulator + cell type'.") - pubmed_query: str = Field(description="PubMed E-utilities syntax: field tags ([tiab], [mesh]), AND/OR/NOT, aliases OR-ed in parentheses.") - text_query: str = Field(description="Plain keyword query for Crossref / Nature / OpenAlex / web search (no field tags).") + scope: str = Field( + description="What this rung covers, e.g. 'gene + regulator + cell type'." + ) + pubmed_query: str = Field( + description="PubMed E-utilities syntax: field tags ([tiab], [mesh]), AND/OR/NOT, aliases OR-ed in parentheses." + ) + text_query: str = Field( + description="Plain keyword query for Crossref / Nature / OpenAlex / web search (no field tags)." + ) class PlannedQuery(BaseModel): id: str = Field(description="Short unique id, e.g. Q1.") question_type: QuestionType - hypothesis: str = Field(description="The specific, testable statement this query looks for evidence on.") - genes: list[str] = Field(description="Gene symbols involved, exactly as named in the digest.") - priority: Literal[1, 2, 3] = Field(description="1 = most informative for the program's function.") - ladder: list[QueryRung] = Field(description="2-4 rungs from most specific to broadest; searched in order.") + hypothesis: str = Field( + description="The specific, testable statement this query looks for evidence on." + ) + genes: list[str] = Field( + description="Gene symbols involved, exactly as named in the digest." + ) + priority: Literal[1, 2, 3] = Field( + description="1 = most informative for the program's function." + ) + ladder: list[QueryRung] = Field( + description="2-4 rungs from most specific to broadest; searched in order." + ) class CurationQuestion(BaseModel): id: str = Field(description="Short unique id, e.g. C1.") question_type: QuestionType - question: str = Field(description="A self-contained question answerable from retrieved papers (names genes and cell type explicitly).") + question: str = Field( + description="A self-contained question answerable from retrieved papers (names genes and cell type explicitly)." + ) genes: list[str] - query_ids: list[str] = Field(description="Ids of the queries whose papers answer this question.") + query_ids: list[str] = Field( + description="Ids of the queries whose papers answer this question." + ) class LiteraturePlan(BaseModel): - interpreted_goal: str = Field(description="One or two sentences: what the literature search should establish.") + interpreted_goal: str = Field( + description="One or two sentences: what the literature search should establish." + ) queries: list[PlannedQuery] curation_questions: list[CurationQuestion] @@ -99,58 +130,82 @@ class LiteraturePlan(BaseModel): # digest def _fmt_stats(stats): - return "; ".join(f"{c}: log2FC={s.get('log2fc')}, adj_p={s.get('adj_pval')}" for c, s in stats.items()) + return "; ".join( + f"{c}: log2FC={s.get('log2fc')}, adj_p={s.get('adj_pval')}" + for c, s in stats.items() + ) def bundle_digest(bundle, max_genes=40): entries = gene_entries(bundle) regs = regulator_stats(bundle) - lines = [f"Program {bundle.get('program_id')} | organism: {bundle.get('organism')} | " - f"cell type: {bundle.get('cell_type')} | conditions: {', '.join(bundle.get('conditions', []))}", - f"Enriched GO terms: {', '.join(bundle.get('GO', [])) or 'none'}", - f"Program specificity: {bundle.get('program_specificity', {}).get('per_condition')}", ""] + lines = [ + f"Program {bundle.get('program_id')} | organism: {bundle.get('organism')} | " + f"cell type: {bundle.get('cell_type')} | conditions: {', '.join(bundle.get('conditions', []))}", + f"Enriched GO terms: {', '.join(bundle.get('GO', [])) or 'none'}", + f"Program specificity: {bundle.get('program_specificity', {}).get('per_condition')}", + "", + ] def gene_lines(gene): entry = entries.get(gene, {}) text = describe_gene(gene, entry) grif = entry.get("gene_info", {}).get("GeneRIF") or {} if grif.get("found"): - text += (f"\n GeneRIF context (evidence_in_cell_type={grif.get('evidence_in_cell_type')}): " - f"{grif.get('context_summary')}") + text += ( + f"\n GeneRIF context (evidence_in_cell_type={grif.get('evidence_in_cell_type')}): " + f"{grif.get('context_summary')}" + ) return text - for key, title in (("program_genes", "Program genes (top loadings)"), - ("distinctive_genes", "Distinctive genes")): - genes = [e["gene"] for e in bundle.get(key, []) if isinstance(e, dict)][:max_genes] + for key, title in ( + ("program_genes", "Program genes (top loadings)"), + ("distinctive_genes", "Distinctive genes"), + ): + genes = [e["gene"] for e in bundle.get(key, []) if isinstance(e, dict)][ + :max_genes + ] lines += [f"{title}:"] + [gene_lines(g) for g in genes] + [""] lines.append("Perturbation regulators (knockdown effect on this program):") - lines += [f"{gene_lines(r)}\n effect: {_fmt_stats(s)}" for r, s in regs.items()] + [""] + lines += [ + f"{gene_lines(r)}\n effect: {_fmt_stats(s)}" for r, s in regs.items() + ] + [""] omnipath = bundle.get("gene_interactions", {}).get("OmniPath") if omnipath: pairs = omnipath.get("pairs", {}).values() found = [f"{p['regulator']}-{p['gene']}" for p in pairs if p.get("found")] missing = [f"{p['regulator']}-{p['gene']}" for p in pairs if not p.get("found")] - lines += [f"OmniPath regulator-gene pairs with a database interaction: {', '.join(found) or 'none'}", - f"OmniPath pairs with NO database interaction: {', '.join(missing) or 'none'}"] + lines += [ + f"OmniPath regulator-gene pairs with a database interaction: {', '.join(found) or 'none'}", + f"OmniPath pairs with NO database interaction: {', '.join(missing) or 'none'}", + ] return "\n".join(lines) def build_plan(bundle, goal, handler, args): system = SYSTEM_PROMPT.format( types="\n".join(f" {k}: {v}" for k, v in QUESTION_TYPES.items()), - disease=(f"Disease context for D9: {args.disease_context}." if args.disease_context - else "No disease context was given: do not plan D9 questions."), - max_queries=args.max_queries) + disease=( + f"Disease context for D9: {args.disease_context}." + if args.disease_context + else "No disease context was given: do not plan D9 questions." + ), + max_queries=args.max_queries, + ) prompt = f"Research goal:\n{goal}\n\nProgram digest:\n{bundle_digest(bundle)}" - out = handler.get_json(system, prompt, LiteraturePlan, effort=args.effort, max_tokens=args.max_tokens) + out = handler.get_json( + system, prompt, LiteraturePlan, effort=args.effort, max_tokens=args.max_tokens + ) return out.model_dump() if out is not None else None def from_research_brief(bundle, handler, args): brief = bundle.get("research_brief") if not brief: - raise ValueError(f"{bundle.get('program_id')}: no research_brief; pass --question instead") + raise ValueError( + f"{bundle.get('program_id')}: no research_brief; pass --question instead" + ) return build_plan(bundle, brief, handler, args) @@ -165,7 +220,9 @@ def regulator_strength(genes, regs): for g in genes: for s in regs.get(g, {}).values(): if s.get("log2fc") is not None and s.get("adj_pval"): - best = max(best, abs(s["log2fc"]) * -math.log10(max(s["adj_pval"], 1e-300))) + best = max( + best, abs(s["log2fc"]) * -math.log10(max(s["adj_pval"], 1e-300)) + ) return best @@ -201,9 +258,9 @@ def resolve(genes, owner): q["regulator_strength"] = round(regulator_strength(q["genes"], regs), 3) queries.append(q) queries.sort(key=lambda q: (q["priority"], -q["regulator_strength"])) - for q in queries[args.max_queries:]: + for q in queries[args.max_queries :]: dropped.append(f"{q['id']}: over --max_queries; query dropped") - queries = queries[:args.max_queries] + queries = queries[: args.max_queries] ids = {q["id"]: f"{plan_name}:{q['id']}" for q in queries} for q in queries: @@ -212,7 +269,9 @@ def resolve(genes, owner): for c in plan["curation_questions"]: c["genes"] = resolve(c["genes"], c["id"]) c["query_ids"] = [ids[i] for i in c["query_ids"] if i in ids] - if not c["query_ids"] and c["question_type"] != "C7_cell_type": # C7 can also read gene_pdfs/ + if ( + not c["query_ids"] and c["question_type"] != "C7_cell_type" + ): # C7 can also read gene_pdfs/ dropped.append(f"{c['id']}: none of its queries kept; question dropped") continue c["id"] = f"{plan_name}:{c['id']}" @@ -221,26 +280,72 @@ def resolve(genes, owner): def build_parser(): - p = argparse.ArgumentParser(description="Turn a program's research_brief (or --question) into literature queries and curation questions.") + p = argparse.ArgumentParser( + description="Turn a program's research_brief (or --question) into literature queries and curation questions." + ) # IO - p.add_argument("--info_dir", required=True, help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.") - p.add_argument("--out_dir", default=None, help="Literature_info_extended_PerturbNMF_Info folder. Default: /../Literature_info_extended_PerturbNMF_Info.") + p.add_argument( + "--info_dir", + required=True, + help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Literature_info_extended_PerturbNMF_Info folder. Default: /../Literature_info_extended_PerturbNMF_Info.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3).") - p.add_argument("--question", default=None, help="Free-text question to plan for instead of the research_brief (vague -> specific mode).") - p.add_argument("--plan_name", default=None, help="Key under literature_plan (other plans are kept). Default: 'brief', or 'question' with --question.") - p.add_argument("--disease_context", default=None, help="Enables D9 disease questions, e.g. 'coronary artery disease'.") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3).", + ) + p.add_argument( + "--question", + default=None, + help="Free-text question to plan for instead of the research_brief (vague -> specific mode).", + ) + p.add_argument( + "--plan_name", + default=None, + help="Key under literature_plan (other plans are kept). Default: 'brief', or 'question' with --question.", + ) + p.add_argument( + "--disease_context", + default=None, + help="Enables D9 disease questions, e.g. 'coronary artery disease'.", + ) # plan - p.add_argument("--max_queries", type=int, default=12, help="Max queries kept per program (ranked by priority, then regulator |log2FC| x -log10 adj_p).") + p.add_argument( + "--max_queries", + type=int, + default=12, + help="Max queries kept per program (ranked by priority, then regulator |log2FC| x -log10 adj_p).", + ) # agent - p.add_argument("--model", default="claude-sonnet-5", help="Claude model for the query agent.") - p.add_argument("--effort", default="high", choices=["low", "medium", "high", "xhigh", "max"], help="Effort level of the query agent.") - p.add_argument("--max_tokens", type=int, default=16000, help="Max output tokens per program.") - p.add_argument("--overwrite", action="store_true", help="Re-plan programs that already have literature_plan[].") + p.add_argument( + "--model", default="claude-sonnet-5", help="Claude model for the query agent." + ) + p.add_argument( + "--effort", + default="high", + choices=["low", "medium", "high", "xhigh", "max"], + help="Effort level of the query agent.", + ) + p.add_argument( + "--max_tokens", type=int, default=16000, help="Max output tokens per program." + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-plan programs that already have literature_plan[].", + ) return p @@ -255,36 +360,68 @@ def main(): summary = {} for label, bundle in bundles.items(): if plan_name in bundle.get("literature_plan", {}) and not args.overwrite: - print(f"[{label}] literature_plan['{plan_name}'] exists; skipped (use --overwrite)") + print( + f"[{label}] literature_plan['{plan_name}'] exists; skipped (use --overwrite)" + ) continue - plan = (from_question(args.question, bundle, handler, args) if args.question - else from_research_brief(bundle, handler, args)) + plan = ( + from_question(args.question, bundle, handler, args) + if args.question + else from_research_brief(bundle, handler, args) + ) if plan is None: - print(f"[{label}] [warn] query agent returned nothing; bundle left unchanged") + print( + f"[{label}] [warn] query agent returned nothing; bundle left unchanged" + ) summary[label] = {"error": True} continue plan, dropped = finalize_plan(plan, bundle, plan_name, args) - plan.update(source="question" if args.question else "research_brief", goal=args.question or bundle.get("research_brief"), - model=args.model, created=datetime.now().isoformat(timespec="seconds"), dropped=dropped) + plan.update( + source="question" if args.question else "research_brief", + goal=args.question or bundle.get("research_brief"), + model=args.model, + created=datetime.now().isoformat(timespec="seconds"), + dropped=dropped, + ) bundle.setdefault("literature_plan", {})[plan_name] = plan types = {} for q in plan["queries"]: types[q["question_type"]] = types.get(q["question_type"], 0) + 1 - summary[label] = {"n_queries": len(plan["queries"]), "n_questions": len(plan["curation_questions"]), - "question_types": types, "n_dropped": len(dropped)} - print(f"[{label}] {len(plan['queries'])} queries, {len(plan['curation_questions'])} curation questions; " - f"types {types}; {len(dropped)} dropped") + summary[label] = { + "n_queries": len(plan["queries"]), + "n_questions": len(plan["curation_questions"]), + "question_types": types, + "n_dropped": len(dropped), + } + print( + f"[{label}] {len(plan['queries'])} queries, {len(plan['curation_questions'])} curation questions; " + f"types {types}; {len(dropped)} dropped" + ) for d in dropped: print(f" [drop] {d}") write_bundles(out_dir, bundles) - write_meta(out_dir, SOURCE, { - "created": datetime.now().isoformat(timespec="seconds"), - "inputs": {"info_dir": str(args.info_dir)}, - "plan_name": plan_name, - "params": {k: getattr(args, k) for k in ("question", "disease_context", "max_queries", "model", "effort", "max_tokens")}, - "programs": summary, - }) + write_meta( + out_dir, + SOURCE, + { + "created": datetime.now().isoformat(timespec="seconds"), + "inputs": {"info_dir": str(args.info_dir)}, + "plan_name": plan_name, + "params": { + k: getattr(args, k) + for k in ( + "question", + "disease_context", + "max_queries", + "model", + "effort", + "max_tokens", + ) + }, + "programs": summary, + }, + ) return 0 diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/llm_summary_agent.py b/src/Stage3_Interpretation/C_Annotation/1.Search/llm_summary_agent.py index 6ffc294..ebc6c7c 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/llm_summary_agent.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/llm_summary_agent.py @@ -9,6 +9,7 @@ The structured agents take papers as {pmid, title, year, journal, study_type, abstract, fulltext} and a handler from api/LLM_interface.claude_api; they return a dict, or None when the call failed. """ + from typing import Literal, Optional from pydantic import BaseModel, Field @@ -29,25 +30,55 @@ def Summarize_Agent(text, handler): return text ask = getattr(handler, "get_completion", handler) - return ask("You are a molecular biologist. Summarize this article in one paragraph: the genes studied, what " - "they were shown to do, and the species, cell types and conditions used. Keep only key findings.", - text) or "" + return ( + ask( + "You are a molecular biologist. Summarize this article in one paragraph: the genes studied, what " + "they were shown to do, and the species, cell types and conditions used. Keep only key findings.", + text, + ) + or "" + ) # shared -EvidenceType = Literal["knockout", "knockdown", "CRISPR_screen", "overexpression", "mutation", "binding", - "structural", "expression_correlation", "genetic_association", "clinical", "review", "other"] +EvidenceType = Literal[ + "knockout", + "knockdown", + "CRISPR_screen", + "overexpression", + "mutation", + "binding", + "structural", + "expression_correlation", + "genetic_association", + "clinical", + "review", + "other", +] # mirrors 2.Evidence_curation/qa.py CATEGORIES (paper-qa / GeneProgramExplorer vocabulary) -Relationship = Literal["physical_interaction", "enzymatic_modification", "transcriptional_regulation", - "same_pathway", "same_cellular_phenotype", "same_disease", "none"] +Relationship = Literal[ + "physical_interaction", + "enzymatic_modification", + "transcriptional_regulation", + "same_pathway", + "same_cellular_phenotype", + "same_disease", + "none", +] class ExperimentalContext(BaseModel): pmid: str evidence_type: EvidenceType - species: Optional[str] = Field(description="Species studied, e.g. human, mouse; null if not stated.") - cell_line_or_tissue: Optional[str] = Field(description="Cell lines / primary cells / tissues used; null if not stated.") - condition_or_stimulus: Optional[str] = Field(description="Treatment, stimulus, disease model or condition tested; null if none.") + species: Optional[str] = Field( + description="Species studied, e.g. human, mouse; null if not stated." + ) + cell_line_or_tissue: Optional[str] = Field( + description="Cell lines / primary cells / tissues used; null if not stated." + ) + condition_or_stimulus: Optional[str] = Field( + description="Treatment, stimulus, disease model or condition tested; null if none." + ) context_match: Literal["same_cell_type", "related_cell_type", "unrelated"] @@ -74,14 +105,22 @@ def _ask(handler, system, prompt, schema, effort, max_tokens): # GeneRIF provenance (1.1.Search_literature/search_generif.py) class GeneRIFPaper(ExperimentalContext): claim: str = Field(description="One sentence: what this paper shows the gene does.") - direction: Optional[str] = Field(description="Effect of the gene, e.g. 'promotes X', 'represses Y', 'required for Z'.") - partner_genes: list[str] = Field(description="Genes the paper links to this gene (targets, binders, pathway members).") + direction: Optional[str] = Field( + description="Effect of the gene, e.g. 'promotes X', 'represses Y', 'required for Z'." + ) + partner_genes: list[str] = Field( + description="Genes the paper links to this gene (targets, binders, pathway members)." + ) class GeneRIFSummary(BaseModel): papers: list[GeneRIFPaper] - context_summary: str = Field(description="2-3 sentences: the gene's function in the target cell type vs elsewhere, from these papers only.") - evidence_in_cell_type: bool = Field(description="True if at least one paper studies the gene in the target cell type or a closely related one.") + context_summary: str = Field( + description="2-3 sentences: the gene's function in the target cell type vs elsewhere, from these papers only." + ) + evidence_in_cell_type: bool = Field( + description="True if at least one paper studies the gene in the target cell type or a closely related one." + ) GENERIF_SYSTEM = f"""You are a molecular biologist curating gene-function evidence for a CRISPR perturbation screen. @@ -89,7 +128,9 @@ class GeneRIFSummary(BaseModel): {PAPER_RULES}""" -def Summarize_GeneRIF_Agent(gene, gene_desc, papers, cell_type, handler, effort="medium", max_tokens=16000): +def Summarize_GeneRIF_Agent( + gene, gene_desc, papers, cell_type, handler, effort="medium", max_tokens=16000 +): """Extract claim / species / cell line / condition / evidence type for a gene's GeneRIF papers. Args: @@ -102,27 +143,50 @@ def Summarize_GeneRIF_Agent(gene, gene_desc, papers, cell_type, handler, effort= Returns: dict | None: GeneRIFSummary as a dict, or None when the call failed. """ - prompt = (f"Gene:\n{gene_desc}\n\nTarget cell type: {cell_type or 'unspecified'}\n\n" + _format_papers(papers)) + prompt = ( + f"Gene:\n{gene_desc}\n\nTarget cell type: {cell_type or 'unspecified'}\n\n" + + _format_papers(papers) + ) return _ask(handler, GENERIF_SYSTEM, prompt, GeneRIFSummary, effort, max_tokens) # gene-gene interactions (papers found for a regulator-gene pair) class InteractionPaper(ExperimentalContext): - relationship: Relationship = Field(description="Most direct relationship between the two genes this paper shows; bare co-mention is 'none'.") - source_gene: Optional[str] = Field(description="Upstream gene when the paper shows a direction, else null.") - target_gene: Optional[str] = Field(description="Downstream gene when the paper shows a direction, else null.") + relationship: Relationship = Field( + description="Most direct relationship between the two genes this paper shows; bare co-mention is 'none'." + ) + source_gene: Optional[str] = Field( + description="Upstream gene when the paper shows a direction, else null." + ) + target_gene: Optional[str] = Field( + description="Downstream gene when the paper shows a direction, else null." + ) directed: bool - sign: Literal["activates", "represses", "unclear"] = Field(description="Effect of source on target.") - mechanism: Optional[str] = Field(description="One sentence: how the two genes are linked (e.g. m6A methylation stabilizes the mRNA).") - stance: Literal["supports", "contradicts", "context_only"] = Field(description="Does the paper support a link between the two genes, argue against one, or only give background?") + sign: Literal["activates", "represses", "unclear"] = Field( + description="Effect of source on target." + ) + mechanism: Optional[str] = Field( + description="One sentence: how the two genes are linked (e.g. m6A methylation stabilizes the mRNA)." + ) + stance: Literal["supports", "contradicts", "context_only"] = Field( + description="Does the paper support a link between the two genes, argue against one, or only give background?" + ) class InteractionSummary(BaseModel): papers: list[InteractionPaper] - best_relationship: Relationship = Field(description="Most direct relationship supported across the papers ('none' if no paper supports one).") - summary: str = Field(description="2-3 sentences on how the two genes are related, from these papers only.") - evidence_in_cell_type: bool = Field(description="True if a supporting paper is in the target cell type or a closely related one.") - consistent_with_perturbation: Literal["yes", "no", "unclear"] = Field(description="Does the literature direction agree with the observed knockdown effect on the gene's program?") + best_relationship: Relationship = Field( + description="Most direct relationship supported across the papers ('none' if no paper supports one)." + ) + summary: str = Field( + description="2-3 sentences on how the two genes are related, from these papers only." + ) + evidence_in_cell_type: bool = Field( + description="True if a supporting paper is in the target cell type or a closely related one." + ) + consistent_with_perturbation: Literal["yes", "no", "unclear"] = Field( + description="Does the literature direction agree with the observed knockdown effect on the gene's program?" + ) INTERACTION_SYSTEM = f"""You are a molecular biologist curating gene-gene interaction evidence for a CRISPR perturbation screen. @@ -132,8 +196,18 @@ class InteractionSummary(BaseModel): {PAPER_RULES}""" -def Summarize_Interaction_Agent(gene, regulator, gene_desc, regulator_desc, papers, cell_type, regulator_stats, - handler, effort="medium", max_tokens=16000): +def Summarize_Interaction_Agent( + gene, + regulator, + gene_desc, + regulator_desc, + papers, + cell_type, + regulator_stats, + handler, + effort="medium", + max_tokens=16000, +): """Relationship / direction / mechanism / stance for the papers found for a regulator-gene pair. Args: @@ -148,26 +222,45 @@ def Summarize_Interaction_Agent(gene, regulator, gene_desc, regulator_desc, pape Returns: dict | None: InteractionSummary as a dict, or None when the call failed. """ - effect = "; ".join(f"{c}: log2FC={s.get('log2fc')}, adj_p={s.get('adj_pval')}" - for c, s in (regulator_stats or {}).items()) or "not given" - prompt = (f"Program gene:\n{gene_desc}\n\nPerturbed regulator:\n{regulator_desc}\n\n" - f"Knockdown effect of {regulator} on the program containing {gene}: {effect}\n" - f"Target cell type: {cell_type or 'unspecified'}\n\n" + _format_papers(papers)) - return _ask(handler, INTERACTION_SYSTEM, prompt, InteractionSummary, effort, max_tokens) + effect = ( + "; ".join( + f"{c}: log2FC={s.get('log2fc')}, adj_p={s.get('adj_pval')}" + for c, s in (regulator_stats or {}).items() + ) + or "not given" + ) + prompt = ( + f"Program gene:\n{gene_desc}\n\nPerturbed regulator:\n{regulator_desc}\n\n" + f"Knockdown effect of {regulator} on the program containing {gene}: {effect}\n" + f"Target cell type: {cell_type or 'unspecified'}\n\n" + _format_papers(papers) + ) + return _ask( + handler, INTERACTION_SYSTEM, prompt, InteractionSummary, effort, max_tokens + ) # literature_plan queries (papers found for one hypothesis) class HypothesisPaper(ExperimentalContext): - stance: Literal["supports", "contradicts", "mixed", "context_only"] = Field(description="The paper's stance on the hypothesis.") - finding: str = Field(description="One sentence: what the paper shows that bears on the hypothesis.") - genes_mentioned: list[str] = Field(description="Genes of the hypothesis that the paper studies.") + stance: Literal["supports", "contradicts", "mixed", "context_only"] = Field( + description="The paper's stance on the hypothesis." + ) + finding: str = Field( + description="One sentence: what the paper shows that bears on the hypothesis." + ) + genes_mentioned: list[str] = Field( + description="Genes of the hypothesis that the paper studies." + ) class HypothesisSummary(BaseModel): papers: list[HypothesisPaper] verdict: Literal["supported", "contradicted", "mixed", "insufficient"] - summary: str = Field(description="2-3 sentences weighing the evidence, from these papers only.") - open_questions: list[str] = Field(description="Gaps the papers leave open that a follow-up search could target.") + summary: str = Field( + description="2-3 sentences weighing the evidence, from these papers only." + ) + open_questions: list[str] = Field( + description="Gaps the papers leave open that a follow-up search could target." + ) QUERY_SYSTEM = f"""You are a molecular biologist weighing literature evidence for one hypothesis about a gene program @@ -177,7 +270,9 @@ class HypothesisSummary(BaseModel): {PAPER_RULES}""" -def Summarize_Query_Agent(query, papers, gene_descs, cell_type, handler, effort="medium", max_tokens=16000): +def Summarize_Query_Agent( + query, papers, gene_descs, cell_type, handler, effort="medium", max_tokens=16000 +): """Stance / finding / context per paper and an overall verdict for one literature_plan query. Args: @@ -191,6 +286,8 @@ def Summarize_Query_Agent(query, papers, gene_descs, cell_type, handler, effort= dict | None: HypothesisSummary as a dict, or None when the call failed. """ genes = "\n".join(gene_descs.get(g, f"- {g}") for g in query.get("genes", [])) - prompt = (f"Hypothesis ({query.get('question_type')}): {query['hypothesis']}\n\nGenes:\n{genes}\n\n" - f"Target cell type: {cell_type or 'unspecified'}\n\n" + _format_papers(papers)) + prompt = ( + f"Hypothesis ({query.get('question_type')}): {query['hypothesis']}\n\nGenes:\n{genes}\n\n" + f"Target cell type: {cell_type or 'unspecified'}\n\n" + _format_papers(papers) + ) return _ask(handler, QUERY_SYSTEM, prompt, HypothesisSummary, effort, max_tokens) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/pdf_download.py b/src/Stage3_Interpretation/C_Annotation/1.Search/pdf_download.py index 01b6f16..7652619 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/pdf_download.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/pdf_download.py @@ -19,6 +19,7 @@ Env: NCBI_EMAIL / PUBMED_EMAIL (NCBI courtesy), UNPAYWALL_EMAIL (required by Unpaywall; falls back to NCBI_EMAIL / PUBMED_EMAIL, and Unpaywall is skipped if none is set). """ + import os import re from pathlib import Path @@ -55,8 +56,13 @@ def _download(url, dest, timeout=DEFAULT_TIMEOUT, params=None): if dest.exists() and dest.stat().st_size > 0: return str(dest.resolve()), "cached" try: - resp = requests.get(url, params=params, timeout=timeout, allow_redirects=True, - headers={"User-Agent": _user_agent()}) + resp = requests.get( + url, + params=params, + timeout=timeout, + allow_redirects=True, + headers={"User-Agent": _user_agent()}, + ) except requests.RequestException as e: return None, f"network error ({type(e).__name__})" if resp.status_code != 200 or not resp.content: @@ -88,9 +94,15 @@ def pmc_pdf_url(pmcid, timeout=DEFAULT_TIMEOUT): """Return (pdf_url, reason) for a PMCID from the PMC open-access S3 bucket.""" pmcid = str(pmcid).upper() latest = None - for version in (1, 2, 3): # articles are versioned PMCxxxx.1, .2, ...; keep the newest + for version in ( + 1, + 2, + 3, + ): # articles are versioned PMCxxxx.1, .2, ...; keep the newest try: - resp = requests.get(f"{PMC_S3}/metadata/{pmcid}.{version}.json", timeout=timeout) + resp = requests.get( + f"{PMC_S3}/metadata/{pmcid}.{version}.json", timeout=timeout + ) except requests.RequestException as e: return None, f"PMC error ({type(e).__name__})" if resp.status_code != 200: @@ -100,7 +112,10 @@ def pmc_pdf_url(pmcid, timeout=DEFAULT_TIMEOUT): return None, "not in PMC open-access dataset" pdf_url = latest.get("pdf_url") if not pdf_url: - return None, f"PMC has no open-access PDF (license: {latest.get('license_code')})" + return ( + None, + f"PMC has no open-access PDF (license: {latest.get('license_code')})", + ) return pdf_url.replace("s3://pmc-oa-opendata", PMC_S3).split("?")[0], "pmc" @@ -110,16 +125,21 @@ def unpaywall_pdf_url(doi, timeout=DEFAULT_TIMEOUT): if not email: return None, "Unpaywall skipped (set UNPAYWALL_EMAIL)" try: - resp = requests.get(UNPAYWALL.format(doi=doi), params={"email": email}, timeout=timeout) + resp = requests.get( + UNPAYWALL.format(doi=doi), params={"email": email}, timeout=timeout + ) data = resp.json() if resp.status_code == 200 else {} except (requests.RequestException, ValueError) as e: return None, f"Unpaywall error ({type(e).__name__})" pdf_url = (data.get("best_oa_location") or {}).get("url_for_pdf") - return (pdf_url, "unpaywall") if pdf_url else (None, "no open-access PDF on Unpaywall") + return ( + (pdf_url, "unpaywall") if pdf_url else (None, "no open-access PDF on Unpaywall") + ) -def download_paper(out_dir, name, pmid=None, pmcid=None, doi=None, pdf_urls=(), - timeout=DEFAULT_TIMEOUT): +def download_paper( + out_dir, name, pmid=None, pmcid=None, doi=None, pdf_urls=(), timeout=DEFAULT_TIMEOUT +): """Try every open-access route for one paper. Returns (pdf_path, message): pdf_path is None when nothing worked, and message then diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/run_literature_search.py b/src/Stage3_Interpretation/C_Annotation/1.Search/run_literature_search.py index 8e82813..d9a5096 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/run_literature_search.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/run_literature_search.py @@ -21,6 +21,7 @@ Env: ANTHROPIC_API_KEY (or a Claude Code login), NCBI_EMAIL / NCBI_API_KEY, UNPAYWALL_EMAIL, OPENALEX_API_KEY (only for the openalex source); read from AGeneTic/.env when present. """ + import argparse import asyncio import importlib @@ -32,12 +33,29 @@ HERE = Path(__file__).resolve().parent AGENETIC_DIR = HERE.parent -sys.path[:0] = [str(HERE), str(HERE / "1.1.Search_literature"), str(AGENETIC_DIR / "2.Evidence_curation")] - -from claude_agent_sdk import (AssistantMessage, ClaudeAgentOptions, ResultMessage, # noqa: E402 - ToolUseBlock, create_sdk_mcp_server, query, tool) +sys.path[:0] = [ + str(HERE), + str(HERE / "1.1.Search_literature"), + str(AGENETIC_DIR / "2.Evidence_curation"), +] + +from claude_agent_sdk import ( + AssistantMessage, + ClaudeAgentOptions, + ResultMessage, # noqa: E402 + ToolUseBlock, + create_sdk_mcp_server, + query, + tool, +) from evidence_cache import pair_key # noqa: E402 -from src import default_out_dir, describe_gene, gene_entries, load_env, load_program_JSON # noqa: E402 +from src import ( + default_out_dir, + describe_gene, + gene_entries, + load_env, + load_program_JSON, +) # noqa: E402 SERVER = "literature" # source -> (module, function); pubmed / arxiv need langchain_community @@ -54,25 +72,36 @@ "type": "object", "properties": { "queries": {"type": "array", "items": {"type": "string"}}, - "rung_reached": {"type": ["integer", "null"]}, # plan mode: 1-based ladder rung that found evidence - "papers": {"type": "array", "items": { - "type": "object", - "properties": { - "title": {"type": "string"}, - "pmid": {"type": ["string", "null"]}, - "doi": {"type": ["string", "null"]}, - "url": {"type": ["string", "null"]}, - "pdf_path": {"type": ["string", "null"]}, - "relevance": {"type": "string"}, + "rung_reached": { + "type": ["integer", "null"] + }, # plan mode: 1-based ladder rung that found evidence + "papers": { + "type": "array", + "items": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "pmid": {"type": ["string", "null"]}, + "doi": {"type": ["string", "null"]}, + "url": {"type": ["string", "null"]}, + "pdf_path": {"type": ["string", "null"]}, + "relevance": {"type": "string"}, + }, + "required": ["title", "relevance"], }, - "required": ["title", "relevance"], - }}, - "web_hits": {"type": "array", "items": { - "type": "object", - "properties": {"title": {"type": "string"}, "url": {"type": "string"}, - "snippet": {"type": "string"}}, - "required": ["title", "url"], - }}, + }, + "web_hits": { + "type": "array", + "items": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "url": {"type": "string"}, + "snippet": {"type": "string"}, + }, + "required": ["title", "url"], + }, + }, "notes": {"type": "string"}, }, "required": ["queries", "papers", "web_hits", "notes"], @@ -120,14 +149,20 @@ def not_found_pairs(bundle): omnipath = bundle.get("gene_interactions", {}).get("OmniPath") if omnipath is None: - raise ValueError(f"{bundle.get('program_id')}: no gene_interactions.OmniPath block " - "(run 1.0.Search_database/search_gene_interaction/search_OmniPath.py first)") + raise ValueError( + f"{bundle.get('program_id')}: no gene_interactions.OmniPath block " + "(run 1.0.Search_database/search_gene_interaction/search_OmniPath.py first)" + ) return [p for p in omnipath["pairs"].values() if not p["found"]] def plan_queries(bundle): """Every query of every literature_plan in the bundle (llm_query_agent.py), in plan order.""" - return [q for plan in bundle.get("literature_plan", {}).values() for q in plan.get("queries", [])] + return [ + q + for plan in bundle.get("literature_plan", {}).values() + for q in plan.get("queries", []) + ] def task_dir_name(query_id): @@ -135,40 +170,48 @@ def task_dir_name(query_id): def build_plan_prompt(q, entries, bundle): - ladder = "\n".join(f" rung {i}: {r['scope']}\n PubMed: {r['pubmed_query']}\n text: {r['text_query']}" - for i, r in enumerate(q["ladder"], 1)) - return "\n".join([ - f"Program {bundle.get('program_id')} ({bundle.get('organism', 'human')}, " - f"cell type: {bundle.get('cell_type') or 'unspecified'}).", - f"Program GO terms: {', '.join(bundle.get('GO', [])[:8]) or 'none'}", - "", - f"Hypothesis ({q['question_type']}): {q['hypothesis']}", - "", - "Genes involved:", - *[describe_gene(g, entries.get(g)) for g in q["genes"]], - "", - "Query ladder:", - ladder, - ]) + ladder = "\n".join( + f" rung {i}: {r['scope']}\n PubMed: {r['pubmed_query']}\n text: {r['text_query']}" + for i, r in enumerate(q["ladder"], 1) + ) + return "\n".join( + [ + f"Program {bundle.get('program_id')} ({bundle.get('organism', 'human')}, " + f"cell type: {bundle.get('cell_type') or 'unspecified'}).", + f"Program GO terms: {', '.join(bundle.get('GO', [])[:8]) or 'none'}", + "", + f"Hypothesis ({q['question_type']}): {q['hypothesis']}", + "", + "Genes involved:", + *[describe_gene(g, entries.get(g)) for g in q["genes"]], + "", + "Query ladder:", + ladder, + ] + ) def build_prompt(pair, entries, bundle): gene, reg = pair["gene"], pair["regulator"] - stats = "; ".join(f"{c}: log2FC={s.get('log2fc')}, adj_p={s.get('adj_pval')}" - for c, s in pair.get("regulator_stats", {}).items()) - return "\n".join([ - f"Program {bundle.get('program_id')} ({bundle.get('organism', 'human')}, " - f"cell type: {bundle.get('cell_type') or 'unspecified'}).", - f"Program GO terms: {', '.join(bundle.get('GO', [])[:8]) or 'none'}", - "", - f"Program gene ({pair.get('query_category')}):", - describe_gene(gene, entries.get(gene)), - "", - f"Perturbed regulator ({stats}):", - describe_gene(reg, entries.get(reg)), - "", - f"Find literature relating {gene} and {reg}.", - ]) + stats = "; ".join( + f"{c}: log2FC={s.get('log2fc')}, adj_p={s.get('adj_pval')}" + for c, s in pair.get("regulator_stats", {}).items() + ) + return "\n".join( + [ + f"Program {bundle.get('program_id')} ({bundle.get('organism', 'human')}, " + f"cell type: {bundle.get('cell_type') or 'unspecified'}).", + f"Program GO terms: {', '.join(bundle.get('GO', [])[:8]) or 'none'}", + "", + f"Program gene ({pair.get('query_category')}):", + describe_gene(gene, entries.get(gene)), + "", + f"Perturbed regulator ({stats}):", + describe_gene(reg, entries.get(reg)), + "", + f"Find literature relating {gene} and {reg}.", + ] + ) # tools @@ -179,7 +222,9 @@ def load_search_functions(sources): try: funcs[name] = getattr(importlib.import_module(module), name) except ImportError as e: - raise ImportError(f"source '{s}' ({module}.py) cannot be imported: {e}") from e + raise ImportError( + f"source '{s}' ({module}.py) cannot be imported: {e}" + ) from e return funcs @@ -189,26 +234,42 @@ def list_pdfs(pair_dir): def make_tools(funcs, pair_dir, max_papers, calls): """One agent tool per search function; each downloads into pair_dir and logs to calls.""" - schema = {"type": "object", - "properties": {"query": {"type": "string"}, - "max_results": {"type": "integer", "minimum": 1, "maximum": 15}}, - "required": ["query"]} + schema = { + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "integer", "minimum": 1, "maximum": 15}, + }, + "required": ["query"], + } tools = [] for name, fn in funcs.items(): + def _make(name=name, fn=fn): @tool(name, (fn.__doc__ or name).strip(), schema) async def _run(args): n = min(int(args.get("max_results") or max_papers), max_papers) before = list_pdfs(pair_dir) try: - text = await asyncio.to_thread(fn, args["query"], n, None, str(pair_dir)) + text = await asyncio.to_thread( + fn, args["query"], n, None, str(pair_dir) + ) except Exception as e: # a failing source must not end the agent run text = f"Error during {name}: {e}" new = sorted(list_pdfs(pair_dir) - before) - calls.append({"tool": name, "query": args["query"], "max_results": n, - "new_pdfs": new, "error": text.startswith("Error")}) + calls.append( + { + "tool": name, + "query": args["query"], + "max_results": n, + "new_pdfs": new, + "error": text.startswith("Error"), + } + ) return {"content": [{"type": "text", "text": str(text)}]} + return _run + tools.append(_make()) return tools @@ -218,15 +279,19 @@ async def search_task(prompt, system_prompt, pair_dir, funcs, args, task_info): """Run one agent search into pair_dir; task_info (pair or query fields) is stored in the log.""" pair_dir.mkdir(parents=True, exist_ok=True) calls, web_queries = [], [] - server = create_sdk_mcp_server(SERVER, tools=make_tools(funcs, pair_dir, args.max_papers, calls)) - allowed = [f"mcp__{SERVER}__{n}" for n in funcs] + (["WebSearch"] if args.web_search else []) + server = create_sdk_mcp_server( + SERVER, tools=make_tools(funcs, pair_dir, args.max_papers, calls) + ) + allowed = [f"mcp__{SERVER}__{n}" for n in funcs] + ( + ["WebSearch"] if args.web_search else [] + ) options = ClaudeAgentOptions( system_prompt=system_prompt.format(max_calls=args.max_tool_calls), mcp_servers={SERVER: server}, - tools=["WebSearch"] if args.web_search else [], # no Bash / Read / Write + tools=["WebSearch"] if args.web_search else [], # no Bash / Read / Write allowed_tools=allowed, permission_mode="dontAsk", - setting_sources=[], # ignore user/project settings + setting_sources=[], # ignore user/project settings model=args.model, max_turns=args.max_turns, max_budget_usd=args.max_budget_usd, @@ -263,36 +328,107 @@ async def search_task(prompt, system_prompt, pair_dir, funcs, args, task_info): "is_error": getattr(final, "is_error", True), "stop_reason": getattr(final, "subtype", None), } - (pair_dir / "web_hits.json").write_text(json.dumps(result.get("web_hits", []), indent=2)) + (pair_dir / "web_hits.json").write_text( + json.dumps(result.get("web_hits", []), indent=2) + ) (pair_dir / "search_log.json").write_text(json.dumps(log, indent=2)) return log def build_parser(): - p = argparse.ArgumentParser(description="Claude Agent SDK literature search for literature_plan queries or OmniPath not-found gene-regulator pairs.") + p = argparse.ArgumentParser( + description="Claude Agent SDK literature search for literature_plan queries or OmniPath not-found gene-regulator pairs." + ) # IO - p.add_argument("--info_dir", required=True, help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.") - p.add_argument("--lit_dir", default=None, help="Literature_info_extended_PerturbNMF_Info folder (P.json there win over --info_dir). Default: /../Literature_info_extended_PerturbNMF_Info.") - p.add_argument("--out_dir", default=None, help="Search output folder. Default: /Literature_search.") + p.add_argument( + "--info_dir", + required=True, + help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.", + ) + p.add_argument( + "--lit_dir", + default=None, + help="Literature_info_extended_PerturbNMF_Info folder (P.json there win over --info_dir). Default: /../Literature_info_extended_PerturbNMF_Info.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Search output folder. Default: /Literature_search.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3).") - p.add_argument("--mode", default="auto", choices=["auto", "plan", "omnipath"], help="plan: literature_plan queries (llm_query_agent.py); omnipath: OmniPath not-found pairs; auto: plan when the bundle has one.") - p.add_argument("--max_queries", type=int, default=None, help="plan mode: max queries searched per program (in plan order). Default: all.") - p.add_argument("--max_pairs", type=int, default=None, help="omnipath mode: max not-found pairs searched per program (in OmniPath's stored order). Default: all.") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3).", + ) + p.add_argument( + "--mode", + default="auto", + choices=["auto", "plan", "omnipath"], + help="plan: literature_plan queries (llm_query_agent.py); omnipath: OmniPath not-found pairs; auto: plan when the bundle has one.", + ) + p.add_argument( + "--max_queries", + type=int, + default=None, + help="plan mode: max queries searched per program (in plan order). Default: all.", + ) + p.add_argument( + "--max_pairs", + type=int, + default=None, + help="omnipath mode: max not-found pairs searched per program (in OmniPath's stored order). Default: all.", + ) # search - p.add_argument("--sources", nargs="+", default=["ncbi", "crossref", "nature"], choices=list(SOURCES), help="Paper searchers given to the agent (openalex needs OPENALEX_API_KEY; pubmed/arxiv need langchain_community).") - p.add_argument("--max_papers", type=int, default=5, help="Max papers per tool call (searchers cap at 15).") - p.add_argument("--no_web_search", dest="web_search", action="store_false", help="Do not give the agent the WebSearch tool.") + p.add_argument( + "--sources", + nargs="+", + default=["ncbi", "crossref", "nature"], + choices=list(SOURCES), + help="Paper searchers given to the agent (openalex needs OPENALEX_API_KEY; pubmed/arxiv need langchain_community).", + ) + p.add_argument( + "--max_papers", + type=int, + default=5, + help="Max papers per tool call (searchers cap at 15).", + ) + p.add_argument( + "--no_web_search", + dest="web_search", + action="store_false", + help="Do not give the agent the WebSearch tool.", + ) # agent - p.add_argument("--model", default="claude-sonnet-5", help="Claude model for the search agent.") - p.add_argument("--max_turns", type=int, default=20, help="Max agent turns per task.") - p.add_argument("--max_tool_calls", type=int, default=8, help="Tool-call budget stated in the agent's instructions.") - p.add_argument("--max_budget_usd", type=float, default=1.0, help="Hard spend cap per task (USD).") - p.add_argument("--overwrite", action="store_true", help="Re-search tasks that already have a search_log.json.") + p.add_argument( + "--model", default="claude-sonnet-5", help="Claude model for the search agent." + ) + p.add_argument( + "--max_turns", type=int, default=20, help="Max agent turns per task." + ) + p.add_argument( + "--max_tool_calls", + type=int, + default=8, + help="Tool-call budget stated in the agent's instructions.", + ) + p.add_argument( + "--max_budget_usd", + type=float, + default=1.0, + help="Hard spend cap per task (USD).", + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-search tasks that already have a search_log.json.", + ) return p @@ -301,14 +437,37 @@ def program_tasks(bundle, mode, args): entries = gene_entries(bundle) info = {"program": bundle.get("program_id"), "cell_type": bundle.get("cell_type")} if mode == "plan": - return [(task_dir_name(q["id"]), build_plan_prompt(q, entries, bundle), PLAN_SYSTEM_PROMPT, - {**info, "mode": "plan", "query_id": q["id"], "question_type": q["question_type"], - "hypothesis": q["hypothesis"], "genes": q["genes"]}) - for q in plan_queries(bundle)[:args.max_queries]] - return [(pair_key(p["gene"], p["regulator"]), build_prompt(p, entries, bundle), SYSTEM_PROMPT, - {**info, "mode": "omnipath", "pair": pair_key(p["gene"], p["regulator"]), - "gene": p["gene"], "regulator": p["regulator"]}) - for p in not_found_pairs(bundle)[:args.max_pairs]] + return [ + ( + task_dir_name(q["id"]), + build_plan_prompt(q, entries, bundle), + PLAN_SYSTEM_PROMPT, + { + **info, + "mode": "plan", + "query_id": q["id"], + "question_type": q["question_type"], + "hypothesis": q["hypothesis"], + "genes": q["genes"], + }, + ) + for q in plan_queries(bundle)[: args.max_queries] + ] + return [ + ( + pair_key(p["gene"], p["regulator"]), + build_prompt(p, entries, bundle), + SYSTEM_PROMPT, + { + **info, + "mode": "omnipath", + "pair": pair_key(p["gene"], p["regulator"]), + "gene": p["gene"], + "regulator": p["regulator"], + }, + ) + for p in not_found_pairs(bundle)[: args.max_pairs] + ] async def run(args): @@ -320,40 +479,66 @@ async def run(args): summary, done = {}, set() for label, bundle in bundles.items(): - mode = args.mode if args.mode != "auto" else ("plan" if plan_queries(bundle) else "omnipath") + mode = ( + args.mode + if args.mode != "auto" + else ("plan" if plan_queries(bundle) else "omnipath") + ) tasks = program_tasks(bundle, mode, args) print(f"[{label}] {len(tasks)} {mode} task(s) to search") summary[label] = {"mode": mode, "tasks": []} for key, prompt, system_prompt, task_info in tasks: task_dir = out_dir / key summary[label]["tasks"].append(key) - if key in done or ((task_dir / "search_log.json").exists() and not args.overwrite): + if key in done or ( + (task_dir / "search_log.json").exists() and not args.overwrite + ): print(f" {key}: already searched; skipped (use --overwrite)") continue - log = await search_task(prompt, system_prompt, task_dir, funcs, args, task_info) + log = await search_task( + prompt, system_prompt, task_dir, funcs, args, task_info + ) done.add(key) - print(f" {key}: tool calls={len(log['tool_calls'])} web={len(log['web_search_queries'])} " - f"PDFs={log['n_pdfs']} rung={log['rung_reached']} cost=${log['total_cost_usd'] or 0:.3f} " - f"{'[agent error: ' + str(log['stop_reason']) + ']' if log['is_error'] else ''}") + print( + f" {key}: tool calls={len(log['tool_calls'])} web={len(log['web_search_queries'])} " + f"PDFs={log['n_pdfs']} rung={log['rung_reached']} cost=${log['total_cost_usd'] or 0:.3f} " + f"{'[agent error: ' + str(log['stop_reason']) + ']' if log['is_error'] else ''}" + ) meta = { "created": datetime.now().isoformat(timespec="seconds"), "inputs": {"info_dir": str(args.info_dir), "lit_dir": str(lit_dir)}, - "params": {k: getattr(args, k) for k in ( - "mode", "max_queries", "max_pairs", "sources", "max_papers", "web_search", "model", - "max_turns", "max_tool_calls", "max_budget_usd")}, + "params": { + k: getattr(args, k) + for k in ( + "mode", + "max_queries", + "max_pairs", + "sources", + "max_papers", + "web_search", + "model", + "max_turns", + "max_tool_calls", + "max_budget_usd", + ) + }, "programs": summary, "n_searched_this_run": len(done), } (out_dir / "meta.json").write_text(json.dumps(meta, indent=2)) - print(f"[done] {len(done)} task(s) searched -> {out_dir}; meta -> {out_dir / 'meta.json'}") + print( + f"[done] {len(done)} task(s) searched -> {out_dir}; meta -> {out_dir / 'meta.json'}" + ) def main(): args = build_parser().parse_args() load_env() if not os.environ.get("ANTHROPIC_API_KEY"): - print(" [warn] ANTHROPIC_API_KEY not set (env or AGeneTic/.env); relying on a Claude Code login") + print( + " [warn] ANTHROPIC_API_KEY not set (env or AGeneTic/.env); relying on a Claude Code login" + ) if "openalex" in args.sources and not os.environ.get("OPENALEX_API_KEY"): raise SystemExit("--sources openalex needs OPENALEX_API_KEY") asyncio.run(run(args)) diff --git a/src/Stage3_Interpretation/C_Annotation/1.Search/src.py b/src/Stage3_Interpretation/C_Annotation/1.Search/src.py index d3af47a..cf9b5c4 100644 --- a/src/Stage3_Interpretation/C_Annotation/1.Search/src.py +++ b/src/Stage3_Interpretation/C_Annotation/1.Search/src.py @@ -7,6 +7,7 @@ '/P.json' is read when it exists (an earlier literature step already extended it), '/P.json' otherwise. The 1.0 bundles are never modified. """ + import copy import json import os @@ -47,9 +48,13 @@ def load_program_JSON(info_dir, out_dir, programs): paths[f"P{pid}"] = extended if extended.exists() else info_dir / f"P{pid}.json" missing = [str(p) for p in paths.values() if not p.exists()] if missing: - raise FileNotFoundError(f"No bundle for {len(missing)} of {len(paths)} programs: {missing}") + raise FileNotFoundError( + f"No bundle for {len(missing)} of {len(paths)} programs: {missing}" + ) n_ext = sum(p.parent == out_dir for p in paths.values()) - print(f"[load] {len(paths)} bundle(s): {n_ext} from {out_dir}, {len(paths) - n_ext} from {info_dir}") + print( + f"[load] {len(paths)} bundle(s): {n_ext} from {out_dir}, {len(paths) - n_ext} from {info_dir}" + ) return {label: json.loads(p.read_text()) for label, p in paths.items()} @@ -70,9 +75,12 @@ def make_handler(model): """Anthropic API handler (api/LLM_interface.claude_api); needs ANTHROPIC_API_KEY.""" sys.path.insert(0, str(AGENETIC_DIR / "api")) from LLM_interface import claude_api + key = os.environ.get("ANTHROPIC_API_KEY") if not key: - raise SystemExit("ANTHROPIC_API_KEY not set (export it or add it to AGeneTic/.env)") + raise SystemExit( + "ANTHROPIC_API_KEY not set (export it or add it to AGeneTic/.env)" + ) return claude_api(api_key=key, model=model) @@ -100,7 +108,10 @@ def regulator_stats(bundle): stats = {} for cond, regs in bundle.get("perturbation_regulators", {}).items(): for r in regs: - stats.setdefault(r["gene"], {})[cond] = {"log2fc": r.get("log2fc"), "adj_pval": r.get("adj_pval")} + stats.setdefault(r["gene"], {})[cond] = { + "log2fc": r.get("log2fc"), + "adj_pval": r.get("adj_pval"), + } return stats @@ -122,12 +133,16 @@ def describe_gene(gene, entry): mg, up = info.get("MyGene", {}), info.get("UniProt", {}) if mg.get("name"): lines.append(f" name: {mg['name']}") - summary = mg.get("summary") or up.get("function") or info.get("NCBI", {}).get("summary") + summary = ( + mg.get("summary") or up.get("function") or info.get("NCBI", {}).get("summary") + ) if summary: lines.append(f" function: {summary[:600]}") if mg.get("go_bp"): lines.append(f" GO BP: {', '.join(mg['go_bp'][:6])}") - return "\n".join([f"- {gene} (names/aliases: {', '.join(gene_aliases(gene, entry))})"] + lines) + return "\n".join( + [f"- {gene} (names/aliases: {', '.join(gene_aliases(gene, entry))})"] + lines + ) def attach_gene_block(bundle, source, info): diff --git a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/curate_evidence.py b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/curate_evidence.py index df8b459..f6fdba1 100644 --- a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/curate_evidence.py +++ b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/curate_evidence.py @@ -18,6 +18,7 @@ Pair results go to gene_interactions.Literature, question results to literature_evidence (other blocks are kept), plus '/meta.json'. """ + import argparse import asyncio import json @@ -92,7 +93,9 @@ def paperqa_evidence(result, evidence_file, n_pdfs, search_log_file, web_hits_fi "category": result["category"], "directed": result["directed"], "short_excerpt": result["short_excerpt"], - "citation": {k: result.get(k) for k in ("title", "year", "doi", "url", "pmids")}, + "citation": { + k: result.get(k) for k in ("title", "year", "doi", "url", "pmids") + }, "n_pdfs": n_pdfs, "n_pdfs_parsed": result["n_pdfs"], "evidence_file": str(evidence_file) if evidence_file.exists() else None, @@ -110,49 +113,84 @@ async def curate_pair(pair, cell_type, lit_dir, evidence_dir, settings, overwrit n_pdfs = len(list(pair_dir.glob("*.pdf"))) log = load_search_log(pair_dir) if log is None or n_pdfs == 0: - return {"status": "no_literature" if log else "not_searched", "source": None, - "category": None, "n_pdfs": n_pdfs}, False + return { + "status": "no_literature" if log else "not_searched", + "source": None, + "category": None, + "n_pdfs": n_pdfs, + }, False evidence_file = cache_path(evidence_dir, gene, reg) result = None if overwrite else read_result(evidence_dir, gene, reg) if result is not None and result.get("n_pdfs", 0) == 0: - result = None # an earlier run could not index any PDF (e.g. API error); retry it + result = ( + None # an earlier run could not index any PDF (e.g. API error); retry it + ) ran = result is None if ran: metas = log.get("result", {}).get("papers", []) - result = await qa.classify_pair(gene, reg, pair_dir, settings, cell_type=cell_type, metas=metas) - if result["n_pdfs"]: # cache only when paper-qa indexed at least one PDF + result = await qa.classify_pair( + gene, reg, pair_dir, settings, cell_type=cell_type, metas=metas + ) + if result["n_pdfs"]: # cache only when paper-qa indexed at least one PDF write_result(evidence_dir, result) - return paperqa_evidence(result, evidence_file, n_pdfs, pair_dir / "search_log.json", - pair_dir / "web_hits.json"), ran + return paperqa_evidence( + result, + evidence_file, + n_pdfs, + pair_dir / "search_log.json", + pair_dir / "web_hits.json", + ), ran async def curate_program(bundle, args, settings): omnipath = bundle.get("gene_interactions", {}).get("OmniPath") if omnipath is None: - raise ValueError(f"{bundle.get('program_id')}: no gene_interactions.OmniPath block " - "(run 1.Search/1.0.Search_database/search_gene_interaction/search_OmniPath.py first)") - cell_type = args.cell_type if args.cell_type is not None else bundle.get("cell_type", "") + raise ValueError( + f"{bundle.get('program_id')}: no gene_interactions.OmniPath block " + "(run 1.Search/1.0.Search_database/search_gene_interaction/search_OmniPath.py first)" + ) + cell_type = ( + args.cell_type if args.cell_type is not None else bundle.get("cell_type", "") + ) pairs, n_ran = {}, 0 for name, pair in omnipath["pairs"].items(): - record, ran = await curate_pair(pair, cell_type, args.literature_dir, args.evidence_dir, - settings, args.overwrite) - pairs[name] = {"gene": pair["gene"], "regulator": pair["regulator"], - "query_category": pair["query_category"], **record} + record, ran = await curate_pair( + pair, + cell_type, + args.literature_dir, + args.evidence_dir, + settings, + args.overwrite, + ) + pairs[name] = { + "gene": pair["gene"], + "regulator": pair["regulator"], + "query_category": pair["query_category"], + **record, + } n_ran += ran - print(f" {name}: {record['status']}" - f"{' category=' + str(record.get('category')) if record.get('source') == 'paper-qa' else ''}" - f"{' (paper-qa run)' if ran else ''}") + print( + f" {name}: {record['status']}" + f"{' category=' + str(record.get('category')) if record.get('source') == 'paper-qa' else ''}" + f"{' (paper-qa run)' if ran else ''}" + ) counts = {} for p in pairs.values(): counts[p["status"]] = counts.get(p["status"], 0) + 1 - block = {"params": {"cell_type": cell_type, "qa": args.qa_params}, - "n_pairs": len(pairs), "n_by_status": counts, "pairs": pairs} + block = { + "params": {"cell_type": cell_type, "qa": args.qa_params}, + "n_pairs": len(pairs), + "n_by_status": counts, + "pairs": pairs, + } return block, n_ran def question_pdf_dirs(question, lit_dir, literature_dir): - dirs = [Path(literature_dir) / qid.replace(":", "__") for qid in question["query_ids"]] + dirs = [ + Path(literature_dir) / qid.replace(":", "__") for qid in question["query_ids"] + ] if question["question_type"] == "C7_cell_type": dirs += [Path(lit_dir) / "gene_pdfs" / g for g in question["genes"]] return [d for d in dirs if d.exists()] @@ -164,48 +202,119 @@ async def curate_questions(bundle, args, settings): for plan_name, plan in bundle.get("literature_plan", {}).items(): for q in plan.get("curation_questions", []): dirs = question_pdf_dirs(q, args.lit_dir, args.literature_dir) - cache = args.evidence_dir / "questions" / f"{q['id'].replace(':', '__')}.json" - result = None if args.overwrite or not cache.exists() else json.loads(cache.read_text()) + cache = ( + args.evidence_dir / "questions" / f"{q['id'].replace(':', '__')}.json" + ) + result = ( + None + if args.overwrite or not cache.exists() + else json.loads(cache.read_text()) + ) if result is not None and result.get("n_pdfs", 0) == 0: result = None ran = result is None if not dirs: - result, ran = {"verdict": None, "n_pdfs": 0, "status": "no_literature"}, False + result, ran = ( + {"verdict": None, "n_pdfs": 0, "status": "no_literature"}, + False, + ) elif ran: - metas = [p for d in dirs for p in ((load_search_log(d) or {}).get("result", {}).get("papers", []))] - result = await qa.answer_question(q["question"], dirs, settings, metas=metas) + metas = [ + p + for d in dirs + for p in ( + (load_search_log(d) or {}).get("result", {}).get("papers", []) + ) + ] + result = await qa.answer_question( + q["question"], dirs, settings, metas=metas + ) result["status"] = "answered" if result["n_pdfs"] else "no_parsed_pdf" if result["n_pdfs"]: cache.parent.mkdir(parents=True, exist_ok=True) cache.write_text(json.dumps(result, indent=2)) n_ran += ran - questions[q["id"]] = {"plan": plan_name, "question_type": q["question_type"], "question": q["question"], - "genes": q["genes"], "query_ids": q["query_ids"], - "pdf_dirs": [str(d) for d in dirs], **result} - print(f" {q['id']}: {result.get('status')} verdict={result.get('verdict')}" - f"{' (paper-qa run)' if ran else ''}") + questions[q["id"]] = { + "plan": plan_name, + "question_type": q["question_type"], + "question": q["question"], + "genes": q["genes"], + "query_ids": q["query_ids"], + "pdf_dirs": [str(d) for d in dirs], + **result, + } + print( + f" {q['id']}: {result.get('status')} verdict={result.get('verdict')}" + f"{' (paper-qa run)' if ran else ''}" + ) counts = {} for r in questions.values(): counts[str(r.get("verdict"))] = counts.get(str(r.get("verdict")), 0) + 1 - return {"params": {"qa": args.qa_params}, "n_questions": len(questions), - "n_by_verdict": counts, "questions": questions}, n_ran + return { + "params": {"qa": args.qa_params}, + "n_questions": len(questions), + "n_by_verdict": counts, + "questions": questions, + }, n_ran def build_parser(): - p = argparse.ArgumentParser(description="Curate gene-regulator evidence from OmniPath JSON and paper-qa over literature PDFs.") + p = argparse.ArgumentParser( + description="Curate gene-regulator evidence from OmniPath JSON and paper-qa over literature PDFs." + ) # IO - p.add_argument("--info_dir", required=True, help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.") - p.add_argument("--lit_dir", default=None, help=f"{LIT_FOLDER} folder (read first, and written). Default: /../{LIT_FOLDER}.") - p.add_argument("--literature_dir", default=None, help="Literature_search folder from 1.1.Search_literature/run_literature_search.py. Default: /Literature_search.") - p.add_argument("--evidence_dir", default=None, help="paper-qa cache folder. Default: /../Evidence_curation.") - p.add_argument("--config", default=str(HERE / "config.yaml"), help="paper-qa settings (qa: block).") + p.add_argument( + "--info_dir", + required=True, + help="Gene_info_extended_PerturbNMF_Info folder from 1.0.Search_database.", + ) + p.add_argument( + "--lit_dir", + default=None, + help=f"{LIT_FOLDER} folder (read first, and written). Default: /../{LIT_FOLDER}.", + ) + p.add_argument( + "--literature_dir", + default=None, + help="Literature_search folder from 1.1.Search_literature/run_literature_search.py. Default: /Literature_search.", + ) + p.add_argument( + "--evidence_dir", + default=None, + help="paper-qa cache folder. Default: /../Evidence_curation.", + ) + p.add_argument( + "--config", + default=str(HERE / "config.yaml"), + help="paper-qa settings (qa: block).", + ) # context info - p.add_argument("--programs", type=int, nargs="+", required=True, help="Program ids, space separated (e.g. 1 2 3).") - p.add_argument("--cell_type", default=None, help="Cell-type context for paper-qa. Default: the bundle's cell_type.") - p.add_argument("--targets", nargs="+", default=["pairs", "questions"], choices=["pairs", "questions"], help="pairs: OmniPath pairs -> gene_interactions.Literature; questions: literature_plan curation questions -> literature_evidence.") - p.add_argument("--overwrite", action="store_true", help="Re-run paper-qa for pairs / questions that already have a cached result.") + p.add_argument( + "--programs", + type=int, + nargs="+", + required=True, + help="Program ids, space separated (e.g. 1 2 3).", + ) + p.add_argument( + "--cell_type", + default=None, + help="Cell-type context for paper-qa. Default: the bundle's cell_type.", + ) + p.add_argument( + "--targets", + nargs="+", + default=["pairs", "questions"], + choices=["pairs", "questions"], + help="pairs: OmniPath pairs -> gene_interactions.Literature; questions: literature_plan curation questions -> literature_evidence.", + ) + p.add_argument( + "--overwrite", + action="store_true", + help="Re-run paper-qa for pairs / questions that already have a cached result.", + ) return p @@ -214,11 +323,17 @@ def main(): load_env() info_dir = Path(args.info_dir) args.lit_dir = Path(args.lit_dir or info_dir.resolve().parent / LIT_FOLDER) - args.literature_dir = Path(args.literature_dir or args.lit_dir / "Literature_search") - args.evidence_dir = Path(args.evidence_dir or info_dir.resolve().parent / "Evidence_curation") + args.literature_dir = Path( + args.literature_dir or args.lit_dir / "Literature_search" + ) + args.evidence_dir = Path( + args.evidence_dir or info_dir.resolve().parent / "Evidence_curation" + ) if not args.literature_dir.exists(): - raise FileNotFoundError(f"Literature folder not found: {args.literature_dir} " - "(run 1.Search/1.1.Search_literature/run_literature_search.py first)") + raise FileNotFoundError( + f"Literature folder not found: {args.literature_dir} " + "(run 1.Search/1.1.Search_literature/run_literature_search.py first)" + ) cfg = load_config(args.config) args.qa_params = cfg.get("qa", {}) settings = qa.build_settings(cfg) @@ -229,15 +344,21 @@ def main(): paths[f"P{p}"] = extended if extended.exists() else info_dir / f"P{p}.json" missing = [str(p) for p in paths.values() if not p.exists()] if missing: - raise FileNotFoundError(f"No extended bundle for {len(missing)} program(s): {missing}") + raise FileNotFoundError( + f"No extended bundle for {len(missing)} program(s): {missing}" + ) args.evidence_dir.mkdir(parents=True, exist_ok=True) summary, total_ran = {}, 0 for label, path in paths.items(): bundle = json.loads(path.read_text()) summary[label] = {} - if "pairs" in args.targets and "OmniPath" in bundle.get("gene_interactions", {}): - print(f"[{label}] curating {len(bundle['gene_interactions']['OmniPath']['pairs'])} pair(s)") + if "pairs" in args.targets and "OmniPath" in bundle.get( + "gene_interactions", {} + ): + print( + f"[{label}] curating {len(bundle['gene_interactions']['OmniPath']['pairs'])} pair(s)" + ) block, n_ran = asyncio.run(curate_program(bundle, args, settings)) bundle["gene_interactions"][SOURCE] = block summary[label]["pairs"] = block["n_by_status"] @@ -249,7 +370,9 @@ def main(): bundle["literature_evidence"] = block summary[label]["questions"] = block["n_by_verdict"] total_ran += n_ran - print(f" {label}: questions {block['n_by_verdict']} ({n_ran} paper-qa run(s))") + print( + f" {label}: questions {block['n_by_verdict']} ({n_ran} paper-qa run(s))" + ) if not summary[label]: print(f" [warn] {label}: no OmniPath pairs or literature_plan to curate") args.lit_dir.mkdir(parents=True, exist_ok=True) @@ -257,15 +380,21 @@ def main(): meta = { "created": datetime.now().isoformat(timespec="seconds"), - "inputs": {"info_dir": str(info_dir), "lit_dir": str(args.lit_dir), "literature_dir": str(args.literature_dir), - "config": str(args.config)}, + "inputs": { + "info_dir": str(info_dir), + "lit_dir": str(args.lit_dir), + "literature_dir": str(args.literature_dir), + "config": str(args.config), + }, "qa": args.qa_params, "per_program": summary, "n_paperqa_runs": total_ran, } (args.evidence_dir / "meta.json").write_text(json.dumps(meta, indent=2)) - print(f"[done] {len(paths)} program(s) -> {args.lit_dir} (gene_interactions.{SOURCE}, literature_evidence); " - f"cache/meta -> {args.evidence_dir}") + print( + f"[done] {len(paths)} program(s) -> {args.lit_dir} (gene_interactions.{SOURCE}, literature_evidence); " + f"cache/meta -> {args.evidence_dir}" + ) return 0 diff --git a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/evidence_cache.py b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/evidence_cache.py index f32dbdd..5fbc488 100644 --- a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/evidence_cache.py +++ b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/evidence_cache.py @@ -4,6 +4,7 @@ one JSON file per pair (filename = sorted, upper-cased symbols, e.g. ``ACVR2B__SMAD2.json``). The GPE augment step later reads these to attach literature evidence to existing edges. """ + from __future__ import annotations import json diff --git a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/gpe_export.py b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/gpe_export.py index e7b859b..0df462a 100644 --- a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/gpe_export.py +++ b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/gpe_export.py @@ -5,6 +5,7 @@ edge whose pair has a cached result, we append a ``literature`` entry to the edge's ``evidence[]``. We never create new edges here. """ + from __future__ import annotations import json @@ -26,6 +27,7 @@ def _load_gpe(src_dir: str = DEFAULT_GPE_SRC): sys.path.insert(0, src_dir) from gene_program_explorer.literature_edges import CATEGORY_MAP # noqa: E402 from gene_program_explorer.model import Citation, Evidence # noqa: E402 + return Evidence, Citation, CATEGORY_MAP @@ -37,7 +39,9 @@ def _already_attached(edge: dict) -> bool: ) -def attach_cached_evidence(contract: dict, evidence_dir, src_dir: str = DEFAULT_GPE_SRC) -> int: +def attach_cached_evidence( + contract: dict, evidence_dir, src_dir: str = DEFAULT_GPE_SRC +) -> int: """Append cached literature evidence to existing gene-gene edges. Returns #edges updated.""" Evidence, Citation, CATEGORY_MAP = _load_gpe(src_dir) @@ -51,7 +55,10 @@ def attach_cached_evidence(contract: dict, evidence_dir, src_dir: str = DEFAULT_ for edge in contract.get("edges", []): if (edge.get("endpointType") or "node-node") != "node-node": continue - a, b = symbol_by_id.get(edge.get("source")), symbol_by_id.get(edge.get("target")) + a, b = ( + symbol_by_id.get(edge.get("source")), + symbol_by_id.get(edge.get("target")), + ) if not a or not b: continue if _already_attached(edge): @@ -65,13 +72,17 @@ def attach_cached_evidence(contract: dict, evidence_dir, src_dir: str = DEFAULT_ pmids = result.get("pmids") or [] citation = Citation( pmid=pmids[0] if pmids else None, - doi=result.get("doi"), title=result.get("title"), - year=result.get("year"), url=result.get("url"), + doi=result.get("doi"), + title=result.get("title"), + year=result.get("year"), + url=result.get("url"), ) has_cite = any([citation.pmid, citation.doi, citation.title, citation.url]) evidence = Evidence( - kind="literature", sourceDb="paper-qa", - interactionType=label, excerpt=result.get("short_excerpt"), + kind="literature", + sourceDb="paper-qa", + interactionType=label, + excerpt=result.get("short_excerpt"), citation=citation if has_cite else None, ) edge.setdefault("evidence", []).append(evidence.to_dict()) @@ -92,4 +103,6 @@ def validate(contract: dict, schema_path: str = DEFAULT_SCHEMA) -> None: for edge in contract.get("edges", []): for endpoint in (edge["source"], edge["target"]): if endpoint not in valid: - raise ValueError(f"Edge {edge['id']} references unknown endpoint {endpoint}") + raise ValueError( + f"Edge {edge['id']} references unknown endpoint {endpoint}" + ) diff --git a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/qa.py b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/qa.py index 0d40434..da207e7 100644 --- a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/qa.py +++ b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/qa.py @@ -9,6 +9,7 @@ The controlled vocabulary mirrors GeneProgramExplorer's literature_edges.EVIDENCE_CATEGORIES so the GPE PaperQA provider can map a category straight onto a contract edge. """ + from __future__ import annotations import re @@ -30,7 +31,9 @@ ] _CATEGORY_RE = re.compile(r"CATEGORY\s*[:=]\s*([a-z_]+)", re.IGNORECASE) -_VERDICT_RE = re.compile(r"VERDICT\s*[:=]\s*(supported|contradicted|mixed|insufficient)", re.IGNORECASE) +_VERDICT_RE = re.compile( + r"VERDICT\s*[:=]\s*(supported|contradicted|mixed|insufficient)", re.IGNORECASE +) _DIRECTED_RE = re.compile(r"DIRECTED\s*[:=]\s*(yes|no|true|false)", re.IGNORECASE) # Edge excerpts must be short — the viewer renders evidence.excerpt unclamped. @@ -58,7 +61,7 @@ def short_excerpt_from(answer_text: str, fallback: str = "") -> str: tail = "" dir_match = _DIRECTED_RE.search(answer_text or "") if dir_match: - tail = answer_text[dir_match.end():] + tail = answer_text[dir_match.end() :] # Drop any leftover header tokens / leading punctuation before the prose. tail = re.sub(r"^[\s:>\-]*", "", tail) candidate = first_sentence(tail) or first_sentence(fallback) @@ -93,7 +96,9 @@ async def _build_docs(pdf_dir, settings: Settings) -> Docs: try: await docs.aadd(str(pdf), settings=settings) except Exception as e: # a single unparsable PDF shouldn't sink the whole run - print(f" [warn] paper-qa could not index {pdf.name}: {type(e).__name__}: {str(e)[:200]}") + print( + f" [warn] paper-qa could not index {pdf.name}: {type(e).__name__}: {str(e)[:200]}" + ) continue return docs @@ -130,7 +135,7 @@ def _top_context(session: PQASession): """Return the highest-scoring context, or None.""" if not session.contexts: return None - return max(session.contexts, key=lambda c: (c.score if c.score is not None else -1)) + return max(session.contexts, key=lambda c: c.score if c.score is not None else -1) def _citation_from_context(ctx, metas: Optional[list] = None) -> dict: @@ -152,24 +157,51 @@ def _citation_from_context(ctx, metas: Optional[list] = None) -> dict: break if not url and doi: url = f"https://doi.org/{doi}" - return {"title": title, "titles": titles, "year": year, "doi": doi, "url": url, "pmids": pmids} + return { + "title": title, + "titles": titles, + "year": year, + "doi": doi, + "url": url, + "pmids": pmids, + } -async def classify_pair(gene_a: str, gene_b: str, pdf_dir: Path, settings: Settings, - cell_type: str = "", metas: Optional[list] = None) -> dict: +async def classify_pair( + gene_a: str, + gene_b: str, + pdf_dir: Path, + settings: Settings, + cell_type: str = "", + metas: Optional[list] = None, +) -> dict: """Classify the A-B relationship from the PDFs into the controlled vocabulary.""" docs = await _build_docs(pdf_dir, settings) n_pdfs = len(docs.docs) result = { - "source_symbol": gene_a, "target_symbol": gene_b, "cell_type": cell_type, - "category": None, "directed": False, "excerpt": None, "short_excerpt": None, - "answer": "", "title": None, "titles": [], "year": None, "doi": None, "url": None, - "pmids": [], "source_db": "paper-qa", "n_pdfs": n_pdfs, + "source_symbol": gene_a, + "target_symbol": gene_b, + "cell_type": cell_type, + "category": None, + "directed": False, + "excerpt": None, + "short_excerpt": None, + "answer": "", + "title": None, + "titles": [], + "year": None, + "doi": None, + "url": None, + "pmids": [], + "source_db": "paper-qa", + "n_pdfs": n_pdfs, } if n_pdfs == 0: return result - session = await docs.aquery(classification_prompt(gene_a, gene_b, cell_type), settings=settings) + session = await docs.aquery( + classification_prompt(gene_a, gene_b, cell_type), settings=settings + ) answer_text = session.answer or session.formatted_answer or "" result["answer"] = answer_text @@ -182,14 +214,18 @@ async def classify_pair(gene_a: str, gene_b: str, pdf_dir: Path, settings: Setti result["category"] = category dir_match = _DIRECTED_RE.search(answer_text) - result["directed"] = bool(dir_match and dir_match.group(1).lower() in ("yes", "true")) + result["directed"] = bool( + dir_match and dir_match.group(1).lower() in ("yes", "true") + ) top = _top_context(session) if top is not None: result["excerpt"] = top.context result.update({k: v for k, v in _citation_from_context(top, metas).items()}) # Short, one-sentence text for the contract edge (viewer renders excerpt unclamped). - result["short_excerpt"] = short_excerpt_from(answer_text, fallback=result.get("excerpt") or "") + result["short_excerpt"] = short_excerpt_from( + answer_text, fallback=result.get("excerpt") or "" + ) return result @@ -202,12 +238,24 @@ def question_prompt(question: str) -> str: ) -async def answer_question(question: str, pdf_dirs, settings: Settings, metas: Optional[list] = None, - max_citations: int = 3) -> dict: +async def answer_question( + question: str, + pdf_dirs, + settings: Settings, + metas: Optional[list] = None, + max_citations: int = 3, +) -> dict: """Answer one curation question over the PDFs of several folders; verdict + top citations.""" docs = await _build_docs(pdf_dirs, settings) - result = {"question": question, "verdict": None, "answer": "", "citations": [], - "short_excerpt": None, "n_pdfs": len(docs.docs), "source_db": "paper-qa"} + result = { + "question": question, + "verdict": None, + "answer": "", + "citations": [], + "short_excerpt": None, + "n_pdfs": len(docs.docs), + "source_db": "paper-qa", + } if not docs.docs: return result session = await docs.aquery(question_prompt(question), settings=settings) @@ -215,8 +263,14 @@ async def answer_question(question: str, pdf_dirs, settings: Settings, metas: Op result["answer"] = answer_text m = _VERDICT_RE.search(answer_text) result["verdict"] = m.group(1).lower() if m else None - ranked = sorted(session.contexts, key=lambda c: (c.score if c.score is not None else -1), reverse=True) - result["citations"] = [{**_citation_from_context(c, metas), "excerpt": first_sentence(c.context)} - for c in ranked[:max_citations]] + ranked = sorted( + session.contexts, + key=lambda c: c.score if c.score is not None else -1, + reverse=True, + ) + result["citations"] = [ + {**_citation_from_context(c, metas), "excerpt": first_sentence(c.context)} + for c in ranked[:max_citations] + ] result["short_excerpt"] = first_sentence(_VERDICT_RE.sub("", answer_text).strip()) return result diff --git a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/run_qa.py b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/run_qa.py index 97e9054..e4ff182 100644 --- a/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/run_qa.py +++ b/src/Stage3_Interpretation/C_Annotation/2.Evidence_curation/run_qa.py @@ -12,6 +12,7 @@ For whole programs use curate_evidence.py instead. """ + import argparse import asyncio import json @@ -37,9 +38,12 @@ def run_question(args, settings) -> None: "answer": session.answer, "n_pdfs": len(list(args.pdf_dir.glob("*.pdf"))), "contexts": [ - {"score": c.score, "summary": c.context, - "citation": getattr(c.text.doc, "citation", None), - "title": getattr(c.text.doc, "title", None)} + { + "score": c.score, + "summary": c.context, + "citation": getattr(c.text.doc, "citation", None), + "title": getattr(c.text.doc, "title", None), + } for c in contexts ], } @@ -54,8 +58,14 @@ def run_question(args, settings) -> None: def _format_report(payload: dict) -> str: - lines = [f"# {payload['question']}", "", "## Answer", payload["answer"], "", - f"## Ranked evidence ({len(payload['contexts'])} contexts, {payload['n_pdfs']} PDFs)"] + lines = [ + f"# {payload['question']}", + "", + "## Answer", + payload["answer"], + "", + f"## Ranked evidence ({len(payload['contexts'])} contexts, {payload['n_pdfs']} PDFs)", + ] for i, c in enumerate(payload["contexts"], 1): cite = c.get("citation") or c.get("title") or "unknown source" lines.append(f"{i}. [score {c['score']}] {cite}\n {c['summary']}") @@ -66,8 +76,16 @@ def run_pair(args, settings) -> None: gene_a, gene_b = args.pair log = load_search_log(args.pdf_dir) or {} metas = log.get("result", {}).get("papers", []) - result = asyncio.run(qa.classify_pair(gene_a, gene_b, args.pdf_dir, settings, - cell_type=args.cell_type, metas=metas)) + result = asyncio.run( + qa.classify_pair( + gene_a, + gene_b, + args.pdf_dir, + settings, + cell_type=args.cell_type, + metas=metas, + ) + ) text = json.dumps(result, indent=2) if args.json_out: Path(args.json_out).write_text(text) @@ -77,12 +95,27 @@ def run_pair(args, settings) -> None: def main() -> None: parser = argparse.ArgumentParser(description="paper-qa over a local folder of PDFs") - parser.add_argument("--pdf_dir", required=True, type=Path, help="Folder of PDFs (e.g. Literature_search/__).") - parser.add_argument("--config", default=str(HERE / "config.yaml"), help="paper-qa settings (qa: block).") + parser.add_argument( + "--pdf_dir", + required=True, + type=Path, + help="Folder of PDFs (e.g. Literature_search/__).", + ) + parser.add_argument( + "--config", + default=str(HERE / "config.yaml"), + help="paper-qa settings (qa: block).", + ) parser.add_argument("--question", help="Free-form question to answer over the PDFs") - parser.add_argument("--pair", nargs=2, metavar=("GENE_A", "GENE_B"), - help="Two gene symbols to classify a relationship between") - parser.add_argument("--cell_type", default="", help="Optional cell-type context for pair mode") + parser.add_argument( + "--pair", + nargs=2, + metavar=("GENE_A", "GENE_B"), + help="Two gene symbols to classify a relationship between", + ) + parser.add_argument( + "--cell_type", default="", help="Optional cell-type context for pair mode" + ) parser.add_argument("--json_out", help="Where to write the JSON result") parser.add_argument("--report_out", help="(question mode) Markdown report path") args = parser.parse_args() diff --git a/src/Stage3_Interpretation/C_Annotation/3.Generate/html_report.py b/src/Stage3_Interpretation/C_Annotation/3.Generate/html_report.py index ec55159..d846b27 100644 --- a/src/Stage3_Interpretation/C_Annotation/3.Generate/html_report.py +++ b/src/Stage3_Interpretation/C_Annotation/3.Generate/html_report.py @@ -15,6 +15,7 @@ edge colors follow GeneProgramExplorer's viewer (viewer/src/lib/colorScale.js). Only the standard library is needed. """ + import argparse import json import re @@ -25,13 +26,19 @@ HERE = Path(__file__).resolve().parent TEMPLATE_PATH = HERE / "report_template.html" CYTOSCAPE_PATH = HERE / "cytoscape.min.js" -CYTOSCAPE_CDN = '' +CYTOSCAPE_CDN = ( + '' +) # GeneProgramExplorer color scales BLUE, WHITE, RED = (33, 102, 172), (247, 247, 247), (178, 24, 43) INDIGO = (67, 56, 202) -EDGE_COLORS = {"activation": "#16a34a", "inhibition": "#dc2626", "mixed": "#64748b", - "literature": "#d97706"} +EDGE_COLORS = { + "activation": "#16a34a", + "inhibition": "#dc2626", + "mixed": "#64748b", + "literature": "#d97706", +} # colors @@ -70,7 +77,9 @@ def load_bundles(info_dir, programs=None): paths = {pid: info_dir / f"P{pid}.json" for pid in dict.fromkeys(programs)} missing = [str(p) for p in paths.values() if not p.exists()] if missing: - raise FileNotFoundError(f"No program JSON for {len(missing)} program(s): {missing}") + raise FileNotFoundError( + f"No program JSON for {len(missing)} program(s): {missing}" + ) if not paths: raise FileNotFoundError(f"No P.json in {info_dir}") return {pid: json.loads(p.read_text()) for pid, p in paths.items()} @@ -100,12 +109,18 @@ def gene_card(entry): for block in (mg, nc, up): aliases += block.get("alias") or [] symbol = mg.get("symbol") or nc.get("symbol") or up.get("symbol") or entry["gene"] - pmids = list(dict.fromkeys((up.get("function_pmids") or []) + (nc.get("generif_pmids") or []))) + pmids = list( + dict.fromkeys( + (up.get("function_pmids") or []) + (nc.get("generif_pmids") or []) + ) + ) return { "gene": entry["gene"], "symbol": symbol, "name": mg.get("name") or nc.get("name") or up.get("protein_name"), - "aliases": [a for a in dict.fromkeys(aliases) if a not in (symbol, entry["gene"])], + "aliases": [ + a for a in dict.fromkeys(aliases) if a not in (symbol, entry["gene"]) + ], "summary": mg.get("summary") or nc.get("summary"), "go_bp": mg.get("go_bp") or [], "entrez": mg.get("entrezgene") or nc.get("gene_id"), @@ -115,7 +130,9 @@ def gene_card(entry): "function": up.get("function"), "location": up.get("subcellular_location") or [], "pmids": pmids[:12], - "sources": [s for s in ("MyGene", "NCBI", "UniProt") if info.get(s, {}).get("found")], + "sources": [ + s for s in ("MyGene", "NCBI", "UniProt") if info.get(s, {}).get("found") + ], } @@ -125,13 +142,24 @@ def select_regulators(bundle, n=5): for cond, regs in bundle.get("perturbation_regulators", {}).items(): for r in regs: g = r["gene"] - per_cond.setdefault(g, {})[cond] = {"log2fc": r.get("log2fc"), "adj_pval": r.get("adj_pval")} + per_cond.setdefault(g, {})[cond] = { + "log2fc": r.get("log2fc"), + "adj_pval": r.get("adj_pval"), + } key = (r.get("adj_pval", 1.0), -abs(r.get("log2fc") or 0)) if g not in best or key < best[g][0]: best[g] = (key, r, cond) ranked = sorted(best.items(), key=lambda kv: kv[1][0])[:n] - return [{"gene": g, "log2fc": r.get("log2fc"), "adj_pval": r.get("adj_pval"), "condition": cond, - "per_condition": per_cond[g]} for g, (_, r, cond) in ranked] + return [ + { + "gene": g, + "log2fc": r.get("log2fc"), + "adj_pval": r.get("adj_pval"), + "condition": cond, + "per_condition": per_cond[g], + } + for g, (_, r, cond) in ranked + ] # graph @@ -154,13 +182,17 @@ def omnipath_edge(pair): sign = _sign(pair.get("signs")) return { "id": f"e_op_{pair['regulator']}_{pair['gene']}", - "source": f"n_{pair['regulator']}", "target": f"n_{pair['gene']}", - "kind": "omnipath", "sign": sign, "color": EDGE_COLORS[sign], + "source": f"n_{pair['regulator']}", + "target": f"n_{pair['gene']}", + "kind": "omnipath", + "sign": sign, + "color": EDGE_COLORS[sign], "tarrow": "triangle" if "regulator->gene" in dirs else "none", "sarrow": "triangle" if "gene->regulator" in dirs else "none", "lstyle": "solid", "width": round(1.6 + min(pair.get("n_references", 0), 20) / 4, 2), - "regulator": pair["regulator"], "gene": pair["gene"], + "regulator": pair["regulator"], + "gene": pair["gene"], "categories": pair.get("interaction_categories", []), "directions": pair.get("directions", []), "resources": pair.get("resources", []), @@ -174,13 +206,21 @@ def omnipath_edge(pair): def literature_edge(pair): return { "id": f"e_lit_{pair['regulator']}_{pair['gene']}", - "source": f"n_{pair['regulator']}", "target": f"n_{pair['gene']}", - "kind": "literature", "sign": "literature", "color": EDGE_COLORS["literature"], - "tarrow": "triangle" if pair.get("directed") else "none", "sarrow": "none", - "lstyle": "dashed", "width": 2.2, - "regulator": pair["regulator"], "gene": pair["gene"], - "category": pair.get("category"), "short_excerpt": pair.get("short_excerpt"), - "citation": pair.get("citation") or {}, "n_pdfs": pair.get("n_pdfs"), + "source": f"n_{pair['regulator']}", + "target": f"n_{pair['gene']}", + "kind": "literature", + "sign": "literature", + "color": EDGE_COLORS["literature"], + "tarrow": "triangle" if pair.get("directed") else "none", + "sarrow": "none", + "lstyle": "dashed", + "width": 2.2, + "regulator": pair["regulator"], + "gene": pair["gene"], + "category": pair.get("category"), + "short_excerpt": pair.get("short_excerpt"), + "citation": pair.get("citation") or {}, + "n_pdfs": pair.get("n_pdfs"), } @@ -194,17 +234,36 @@ def build_graph(bundle, top_regulator=5, top_gene=15): nodes = [] for r in regulators: fill = diverging_color(r["log2fc"], m) - nodes.append({"id": f"n_{r['gene']}", "label": r["gene"], "kind": "regulator", - "w": node_width(r["gene"], regulator=True), - "fill": fill, "fg": text_on(fill), "log2fc": r["log2fc"], - "adj_pval": r["adj_pval"], "condition": r["condition"], - "per_condition": r["per_condition"], "in_program": r["gene"] in genes}) + nodes.append( + { + "id": f"n_{r['gene']}", + "label": r["gene"], + "kind": "regulator", + "w": node_width(r["gene"], regulator=True), + "fill": fill, + "fg": text_on(fill), + "log2fc": r["log2fc"], + "adj_pval": r["adj_pval"], + "condition": r["condition"], + "per_condition": r["per_condition"], + "in_program": r["gene"] in genes, + } + ) for i, g in enumerate(genes): if g in reg_names: - continue # a regulator that is also a program gene is drawn once, as a regulator + continue # a regulator that is also a program gene is drawn once, as a regulator fill = sequential_color(0.75 * (1 - i / max(len(genes) - 1, 1)) + 0.1) - nodes.append({"id": f"n_{g}", "label": g, "kind": "gene", "rank": i + 1, "w": node_width(g), - "fill": fill, "fg": text_on(fill)}) + nodes.append( + { + "id": f"n_{g}", + "label": g, + "kind": "gene", + "rank": i + 1, + "w": node_width(g), + "fill": fill, + "fg": text_on(fill), + } + ) node_ids = {n["id"] for n in nodes} gi = bundle.get("gene_interactions", {}) @@ -212,18 +271,30 @@ def build_graph(bundle, top_regulator=5, top_gene=15): edges = [] if omnipath: for pair in omnipath["pairs"].values(): - if pair["found"] and pair["regulator"] in reg_names and pair["gene"] in genes: + if ( + pair["found"] + and pair["regulator"] in reg_names + and pair["gene"] in genes + ): edge = omnipath_edge(pair) if edge["source"] in node_ids and edge["target"] in node_ids: edges.append(edge) literature = gi.get("Literature") if literature: for pair in literature["pairs"].values(): - if (pair.get("source") == "paper-qa" and pair.get("category") - and pair["regulator"] in reg_names and pair["gene"] in genes): + if ( + pair.get("source") == "paper-qa" + and pair.get("category") + and pair["regulator"] in reg_names + and pair["gene"] in genes + ): edges.append(literature_edge(pair)) - return {"nodes": nodes, "edges": edges, "has_omnipath": omnipath is not None, - "has_literature": literature is not None} + return { + "nodes": nodes, + "edges": edges, + "has_omnipath": omnipath is not None, + "has_literature": literature is not None, + } # program @@ -231,36 +302,74 @@ def omnipath_summary(bundle): op = bundle.get("gene_interactions", {}).get("OmniPath") if not op: return None - return {"n_regulators": op.get("n_regulators"), - "n_found": sum(op.get("n_pairs_found", {}).values()), - "n_tested": sum(op.get("n_pairs_tested", {}).values()), - "by_category": {c: [op["n_pairs_found"].get(c, 0), n] - for c, n in op.get("n_pairs_tested", {}).items()}, - "datasets": op.get("params", {}).get("datasets", [])} + return { + "n_regulators": op.get("n_regulators"), + "n_found": sum(op.get("n_pairs_found", {}).values()), + "n_tested": sum(op.get("n_pairs_tested", {}).values()), + "by_category": { + c: [op["n_pairs_found"].get(c, 0), n] + for c, n in op.get("n_pairs_tested", {}).items() + }, + "datasets": op.get("params", {}).get("datasets", []), + } def literature_rows(bundle): lit = bundle.get("gene_interactions", {}).get("Literature") if not lit: return None - rows = [{"pair": k, **{f: p.get(f) for f in ("gene", "regulator", "status", "source", "category", - "directed", "short_excerpt", "citation", "n_pdfs")}} - for k, p in lit["pairs"].items() if p.get("status") != "omnipath"] + rows = [ + { + "pair": k, + **{ + f: p.get(f) + for f in ( + "gene", + "regulator", + "status", + "source", + "category", + "directed", + "short_excerpt", + "citation", + "n_pdfs", + ) + }, + } + for k, p in lit["pairs"].items() + if p.get("status") != "omnipath" + ] return {"n_by_status": lit.get("n_by_status", {}), "rows": rows} def build_lead(bundle, regulators): spec = bundle.get("program_specificity", {}) top = spec.get("top_condition") - pct = spec.get("per_condition", {}).get(top, {}).get("pct_cells_expressed") if top else None - n_regs = len({r["gene"] for regs in bundle.get("perturbation_regulators", {}).values() for r in regs}) - parts = [f"{len(bundle.get('program_genes', []))} top-loading and " - f"{len(bundle.get('distinctive_genes', []))} distinctive genes"] - parts.append(f"{n_regs} significant regulator{'' if n_regs == 1 else 's'}" - + (f", led by {regulators[0]['gene']}" if regulators else "")) + pct = ( + spec.get("per_condition", {}).get(top, {}).get("pct_cells_expressed") + if top + else None + ) + n_regs = len( + { + r["gene"] + for regs in bundle.get("perturbation_regulators", {}).values() + for r in regs + } + ) + parts = [ + f"{len(bundle.get('program_genes', []))} top-loading and " + f"{len(bundle.get('distinctive_genes', []))} distinctive genes" + ] + parts.append( + f"{n_regs} significant regulator{'' if n_regs == 1 else 's'}" + + (f", led by {regulators[0]['gene']}" if regulators else "") + ) lead = "; ".join(parts) + "." if top: - lead += f" Usage is highest in {top}" + (f" ({pct}% of cells)." if pct is not None else ".") + lead += f" Usage is highest in {top}" + ( + f" ({pct}% of cells)." if pct is not None else "." + ) return lead @@ -280,9 +389,13 @@ def build_program(pid, bundle, top_regulator=5, top_gene=15): "specificity": bundle.get("program_specificity", {}), "go": [re.sub(r"\s*\(GO:\d+\)", "", t) for t in go], "program_genes": [gene_name(e) for e in bundle.get("program_genes", [])], - "distinctive_genes": [gene_name(e) for e in bundle.get("distinctive_genes", [])], - "regulators_by_condition": {c: [{k: r.get(k) for k in ("gene", "log2fc", "adj_pval")} for r in regs] - for c, regs in bundle.get("perturbation_regulators", {}).items()}, + "distinctive_genes": [ + gene_name(e) for e in bundle.get("distinctive_genes", []) + ], + "regulators_by_condition": { + c: [{k: r.get(k) for k in ("gene", "log2fc", "adj_pval")} for r in regs] + for c, regs in bundle.get("perturbation_regulators", {}).items() + }, "top_regulators": [r["gene"] for r in regulators], "genes": {g: gene_card(e) for g, e in entries.items()}, "graph": graph, @@ -301,64 +414,121 @@ def _html_esc(s): return str(s).replace("&", "&").replace("<", "<").replace(">", ">") -def render_html(programs, dataset_name="", inline_cytoscape=True, top_regulator=5, top_gene=15): +def render_html( + programs, dataset_name="", inline_cytoscape=True, top_regulator=5, top_gene=15 +): template = TEMPLATE_PATH.read_text() if inline_cytoscape: lib = CYTOSCAPE_PATH.read_text() if "{lib}" else: cyto = CYTOSCAPE_CDN data = _js_safe(json.dumps({p["id"]: p for p in programs}, ensure_ascii=False)) - plist = _js_safe(json.dumps([[p["id"], p["label"], p["title"]] for p in programs], ensure_ascii=False)) + plist = _js_safe( + json.dumps( + [[p["id"], p["label"], p["title"]] for p in programs], ensure_ascii=False + ) + ) # small placeholders first, then the data and library blobs (never re-scanned) - html = (template - .replace("__DATASET_SUB__", _html_esc(dataset_name or "AGeneTic")) - .replace("__DATASET_JSON__", json.dumps(dataset_name or "")) - .replace("__NUM_PROGRAMS__", str(len(programs))) - .replace("__GENERATED_ON__", datetime.now().strftime("%Y-%m-%d")) - .replace("__TOP_REGULATOR__", str(top_regulator)) - .replace("__TOP_GENE__", str(top_gene)) - .replace("__PROGRAM_LIST_JSON__", plist)) + html = ( + template.replace("__DATASET_SUB__", _html_esc(dataset_name or "AGeneTic")) + .replace("__DATASET_JSON__", json.dumps(dataset_name or "")) + .replace("__NUM_PROGRAMS__", str(len(programs))) + .replace("__GENERATED_ON__", datetime.now().strftime("%Y-%m-%d")) + .replace("__TOP_REGULATOR__", str(top_regulator)) + .replace("__TOP_GENE__", str(top_gene)) + .replace("__PROGRAM_LIST_JSON__", plist) + ) html = html.replace("__PROGRAMS_JSON__", data, 1) return html.replace("__CYTOSCAPE_TAG__", cyto, 1) def build_parser(): - p = argparse.ArgumentParser(description="Write an HTML report with per-program regulator-gene interaction graphs.") + p = argparse.ArgumentParser( + description="Write an HTML report with per-program regulator-gene interaction graphs." + ) # IO - p.add_argument("--info_dir", required=True, help="Gene_info_extended_PerturbNMF_Info folder (P.json).") - p.add_argument("--out", default=None, help="Output HTML. Default: /../Report/AGeneTic_report.html.") + p.add_argument( + "--info_dir", + required=True, + help="Gene_info_extended_PerturbNMF_Info folder (P.json).", + ) + p.add_argument( + "--out", + default=None, + help="Output HTML. Default: /../Report/AGeneTic_report.html.", + ) # context info - p.add_argument("--programs", type=int, nargs="+", default=None, help="Program ids, space separated (e.g. 1 2 3). Default: every P.json in --info_dir.") - p.add_argument("--dataset_name", default="", help="Dataset label shown in the top bar and hero.") + p.add_argument( + "--programs", + type=int, + nargs="+", + default=None, + help="Program ids, space separated (e.g. 1 2 3). Default: every P.json in --info_dir.", + ) + p.add_argument( + "--dataset_name", + default="", + help="Dataset label shown in the top bar and hero.", + ) # graph - p.add_argument("--top_regulator", type=int, default=5, help="Regulators per graph (best adj_pval over conditions).") - p.add_argument("--top_gene", type=int, default=15, help="Top-loading program genes per graph.") - p.add_argument("--cdn", action="store_true", help="Load Cytoscape.js from unpkg instead of inlining it (smaller file, needs internet).") + p.add_argument( + "--top_regulator", + type=int, + default=5, + help="Regulators per graph (best adj_pval over conditions).", + ) + p.add_argument( + "--top_gene", type=int, default=15, help="Top-loading program genes per graph." + ) + p.add_argument( + "--cdn", + action="store_true", + help="Load Cytoscape.js from unpkg instead of inlining it (smaller file, needs internet).", + ) return p def main(): args = build_parser().parse_args() bundles = load_bundles(args.info_dir, args.programs) - programs = [build_program(pid, b, args.top_regulator, args.top_gene) for pid, b in bundles.items()] + programs = [ + build_program(pid, b, args.top_regulator, args.top_gene) + for pid, b in bundles.items() + ] for p in programs: g = p["graph"] n_reg = sum(n["kind"] == "regulator" for n in g["nodes"]) - note = "" if g["has_omnipath"] else " [warn] no OmniPath block (run OmniPath.py); graph has no edges" - print(f" {p['label']}: regulators={n_reg} genes={len(g['nodes']) - n_reg} " - f"edges={len(g['edges'])}{note}") - - out = Path(args.out or Path(args.info_dir).resolve().parent / "Report" / "AGeneTic_report.html") + note = ( + "" + if g["has_omnipath"] + else " [warn] no OmniPath block (run OmniPath.py); graph has no edges" + ) + print( + f" {p['label']}: regulators={n_reg} genes={len(g['nodes']) - n_reg} " + f"edges={len(g['edges'])}{note}" + ) + + out = Path( + args.out + or Path(args.info_dir).resolve().parent / "Report" / "AGeneTic_report.html" + ) out.parent.mkdir(parents=True, exist_ok=True) - out.write_text(render_html(programs, args.dataset_name, not args.cdn, - args.top_regulator, args.top_gene)) - print(f"[done] {len(programs)} program(s) -> {out} ({out.stat().st_size / 1e6:.1f} MB)") + out.write_text( + render_html( + programs, args.dataset_name, not args.cdn, args.top_regulator, args.top_gene + ) + ) + print( + f"[done] {len(programs)} program(s) -> {out} ({out.stat().st_size / 1e6:.1f} MB)" + ) return 0 diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/Slurm_Version/run_literature_search.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/Slurm_Version/run_literature_search.py index ef4a287..545014c 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/Slurm_Version/run_literature_search.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/Slurm_Version/run_literature_search.py @@ -8,6 +8,7 @@ --interactions "regulates,induces,promotes,inhibits" \ --domain-keywords "proliferation,migration,differentiation" """ + from __future__ import annotations import argparse @@ -33,7 +34,9 @@ from Literature_search.src.paper_fetcher import PaperFetcher, PaperData, summarize_paper from Literature_search.src.verification import Verifier, VerificationResult from Literature_search.src.output_writer import ( - write_program_json, write_program_markdown, write_program_html, + write_program_json, + write_program_markdown, + write_program_html, ) logger = logging.getLogger(__name__) @@ -43,6 +46,7 @@ # Excel loading # ------------------------------------------------------------------ + def _match_column(df_columns: List[str], target: str) -> Optional[str]: """Case-insensitive column matching.""" target_lower = target.lower().replace(" ", "_") @@ -63,7 +67,9 @@ def _match_column(df_columns: List[str], target: str) -> Optional[str]: } -def load_excel(excel_path: Path, programs: Optional[List[int]] = None) -> List[Dict[str, Any]]: +def load_excel( + excel_path: Path, programs: Optional[List[int]] = None +) -> List[Dict[str, Any]]: """Load program rows from Excel.""" df = pd.read_excel(excel_path) logger.info("Loaded Excel: %d rows, %d columns", len(df), len(df.columns)) @@ -98,6 +104,7 @@ def load_excel(excel_path: Path, programs: Optional[List[int]] = None) -> List[D # Per-program processing # ------------------------------------------------------------------ + def _parse_gene_list(text: str) -> List[str]: """Parse comma/semicolon separated gene list, filtering Ensembl IDs.""" if not text: @@ -145,7 +152,9 @@ def process_program( logger.info("Loaded %d cached LLM queries", len(llm_queries)) else: llm_queries = generate_search_queries( - llm, row, max_queries=args.max_llm_queries, + llm, + row, + max_queries=args.max_llm_queries, interaction_verbs=interaction_verbs, ) query_cache.write_text(json.dumps(llm_queries, indent=2), encoding="utf-8") @@ -224,16 +233,22 @@ def process_program( synthesis_text=synthesis_text, ) - logger.info("Program %s complete: %d papers, %d verified", program_id, len(papers), len(verifications)) + logger.info( + "Program %s complete: %d papers, %d verified", + program_id, + len(papers), + len(verifications), + ) # ------------------------------------------------------------------ # Main # ------------------------------------------------------------------ + def run(args) -> None: """Run the literature search pipeline.""" - # initializing output dir + # initializing output dir cache_dir = args.output_dir / "cache" cache_dir.mkdir(parents=True, exist_ok=True) @@ -241,14 +256,14 @@ def run(args) -> None: interaction_verbs = [v.strip() for v in args.interactions.split(",") if v.strip()] domain_keywords = [k.strip() for k in args.domain_keywords.split(",") if k.strip()] - # Initialize llm model + # Initialize llm model llm = LLMBackend( provider=args.llm_provider, model=args.llm_model, max_tokens=args.llm_max_tokens, ) - # initializing useful agent + # initializing useful agent search_engine = SearchEngine() paper_fetcher = PaperFetcher(args.output_dir) verifier = Verifier(llm=llm if args.semantic_check else None) @@ -260,12 +275,18 @@ def run(args) -> None: rows = load_excel(args.excel, programs=programs) logger.info("Processing %d programs", len(rows)) - # annotate each program + # annotate each program for row in rows: try: process_program( - row, args, llm, search_engine, paper_fetcher, verifier, - interaction_verbs, domain_keywords, + row, + args, + llm, + search_engine, + paper_fetcher, + verifier, + interaction_verbs, + domain_keywords, ) except Exception as exc: pid = row.get("program_id", "?") @@ -283,36 +304,69 @@ def main(argv=None): description="Literature Search Agent for PerturbNMF program annotation", ) # Required - parser.add_argument("--excel", type=Path, required=True, help="Input Excel with program rows") - parser.add_argument("--output-dir", type=Path, required=True, help="Output directory") + parser.add_argument( + "--excel", type=Path, required=True, help="Input Excel with program rows" + ) + parser.add_argument( + "--output-dir", type=Path, required=True, help="Output directory" + ) # Program selection - parser.add_argument("--programs", type=str, default=None, - help="Comma-separated program IDs (default: all)") + parser.add_argument( + "--programs", + type=str, + default=None, + help="Comma-separated program IDs (default: all)", + ) # Search configuration - parser.add_argument("--interactions", type=str, - default="regulates,induces,promotes,inhibits,suppresses,activates,binds,phosphorylates,modulates,mediates,targets,controls,decreases,increases,blocks,triggers,catalyzes", - help="Comma-separated interaction verbs for query formulation") - parser.add_argument("--domain-keywords", type=str, - default="inflammation,proliferation,migration,adhesion,metabolism,hypoxia,differentiation,apoptosis,cell cycle,signaling,stress response,development", - help="Comma-separated domain keywords for evidence scoring") - parser.add_argument("--max-papers", type=int, default=30, - help="Max papers per program (default: 30)") - parser.add_argument("--max-pubtator-results", type=int, default=50, - help="Max results per PubTator query (default: 50)") - parser.add_argument("--max-llm-queries", type=int, default=8, - help="Max LLM-generated queries per program (default: 8)") + parser.add_argument( + "--interactions", + type=str, + default="regulates,induces,promotes,inhibits,suppresses,activates,binds,phosphorylates,modulates,mediates,targets,controls,decreases,increases,blocks,triggers,catalyzes", + help="Comma-separated interaction verbs for query formulation", + ) + parser.add_argument( + "--domain-keywords", + type=str, + default="inflammation,proliferation,migration,adhesion,metabolism,hypoxia,differentiation,apoptosis,cell cycle,signaling,stress response,development", + help="Comma-separated domain keywords for evidence scoring", + ) + parser.add_argument( + "--max-papers", + type=int, + default=30, + help="Max papers per program (default: 30)", + ) + parser.add_argument( + "--max-pubtator-results", + type=int, + default=50, + help="Max results per PubTator query (default: 50)", + ) + parser.add_argument( + "--max-llm-queries", + type=int, + default=8, + help="Max LLM-generated queries per program (default: 8)", + ) # LLM configuration - parser.add_argument("--llm-provider", default="stanford", - choices=["anthropic", "stanford", "openai", "deepseek", "gemini"]) + parser.add_argument( + "--llm-provider", + default="stanford", + choices=["anthropic", "stanford", "openai", "deepseek", "gemini"], + ) parser.add_argument("--llm-model", default=None) parser.add_argument("--llm-max-tokens", type=int, default=4096) # Verification - parser.add_argument("--semantic-check", action="store_true", default=False, - help="Enable LLM semantic verification (costs tokens, off by default)") + parser.add_argument( + "--semantic-check", + action="store_true", + default=False, + help="Enable LLM semantic verification (costs tokens, off by default)", + ) # Resume parser.add_argument("--resume", action="store_true", default=True) @@ -324,7 +378,9 @@ def main(argv=None): if not args.excel.exists(): parser.error(f"Excel file not found: {args.excel}") if not args.output_dir.parent.exists(): - parser.error(f"Parent directory for output does not exist: {args.output_dir.parent}") + parser.error( + f"Parent directory for output does not exist: {args.output_dir.parent}" + ) # Setup logging logging.basicConfig( @@ -332,23 +388,23 @@ def main(argv=None): format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", ) - # intalize folders to be used + # intalize folders to be used args.output_dir.mkdir(parents=True, exist_ok=True) # Save config slurm_info = { - 'job_id': os.environ.get('SLURM_JOB_ID'), - 'job_name': os.environ.get('SLURM_JOB_NAME'), - 'partition': os.environ.get('SLURM_JOB_PARTITION'), - 'node_list': os.environ.get('SLURM_JOB_NODELIST'), + "job_id": os.environ.get("SLURM_JOB_ID"), + "job_name": os.environ.get("SLURM_JOB_NAME"), + "partition": os.environ.get("SLURM_JOB_PARTITION"), + "node_list": os.environ.get("SLURM_JOB_NODELIST"), } - job_id = slurm_info['job_id'] or 'no_jobid' + job_id = slurm_info["job_id"] or "no_jobid" config_to_save = { - 'args': {k: str(v) for k, v in vars(args).items()}, - 'slurm_info': slurm_info, + "args": {k: str(v) for k, v in vars(args).items()}, + "slurm_info": slurm_info, } - config_path = args.output_dir / f'config_{job_id}.yml' - with open(config_path, 'w') as f: + config_path = args.output_dir / f"config_{job_id}.yml" + with open(config_path, "w") as f: yaml.dump(config_to_save, f, default_flow_style=False, width=1000) run(args) diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/__init__.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/__init__.py index c717c6d..f75324f 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/__init__.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/__init__.py @@ -2,12 +2,18 @@ Hybrid PubTator3 + LLM-generated query search with self-verification. """ + from .llm_backend import LLMBackend from .query_generator import generate_search_queries from .search_engine import SearchEngine from .paper_fetcher import PaperFetcher, PaperData, PaperCache, summarize_paper from .verification import Verifier, VerificationResult -from .output_writer import write_program_json, write_program_markdown, write_program_html +from .output_writer import ( + write_program_json, + write_program_markdown, + write_program_html, +) + __all__ = [ "LLMBackend", "generate_search_queries", diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/literature_mining.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/literature_mining.py index 9bdbd2e..69a5610 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/literature_mining.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/literature_mining.py @@ -3,6 +3,7 @@ Extracted from ProgramExplorer/src/literature_mining.py — only the functions needed by the Literature Search pipeline. """ + from __future__ import annotations import re @@ -16,15 +17,38 @@ EUTILS_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils" DEFAULT_DOMAIN_KEYWORDS = [ - "inflammation", "proliferation", "migration", "adhesion", "metabolism", - "hypoxia", "differentiation", "apoptosis", "cell cycle", "signaling", - "stress response", "development" + "inflammation", + "proliferation", + "migration", + "adhesion", + "metabolism", + "hypoxia", + "differentiation", + "apoptosis", + "cell cycle", + "signaling", + "stress response", + "development", ] DEFAULT_INTERACTION_VERBS = [ - "regulates", "induces", "promotes", "inhibits", "suppresses", "activates", - "binds", "phosphorylates", "modulates", "mediates", "targets", "controls", - "decreases", "increases", "blocks", "triggers", "catalyzes" + "regulates", + "induces", + "promotes", + "inhibits", + "suppresses", + "activates", + "binds", + "phosphorylates", + "modulates", + "mediates", + "targets", + "controls", + "decreases", + "increases", + "blocks", + "triggers", + "catalyzes", ] @@ -33,12 +57,16 @@ def split_text_into_sentences(text: str) -> List[str]: return [] protected = text subs = { - "et al.": "ET_AL_MARKER", "e.g.": "EG_MARKER", "i.e.": "IE_MARKER", - "Fig.": "FIG_MARKER", "Ref.": "REF_MARKER", "vs.": "VS_MARKER" + "et al.": "ET_AL_MARKER", + "e.g.": "EG_MARKER", + "i.e.": "IE_MARKER", + "Fig.": "FIG_MARKER", + "Ref.": "REF_MARKER", + "vs.": "VS_MARKER", } for k, v in subs.items(): protected = protected.replace(k, v) - parts = re.split(r'(?<=[.?!])\s+', protected) + parts = re.split(r"(?<=[.?!])\s+", protected) sentences = [] for p in parts: restored = p @@ -51,7 +79,10 @@ def split_text_into_sentences(text: str) -> List[str]: def extract_evidence_sentences( - abstract: str, title: str, target_genes: Set[str], context_genes: Set[str], + abstract: str, + title: str, + target_genes: Set[str], + context_genes: Set[str], domain_keywords: Optional[List[str]] = None, ) -> Dict[str, List[str]]: if domain_keywords is None: @@ -66,12 +97,19 @@ def extract_evidence_sentences( if re.search(rf"\b{re.escape(g)}\b", sent, re.IGNORECASE): found_targets.append(g) normalized_sent = sent.strip(" .") - existing_sents = [s[1].strip(" .") for s in [item for sublist in gene_to_sentences.values() for item in sublist]] + existing_sents = [ + s[1].strip(" .") + for s in [ + item for sublist in gene_to_sentences.values() for item in sublist + ] + ] if normalized_sent in existing_sents: continue score = 0 for m in context_genes: - if m not in found_targets and re.search(rf"\b{re.escape(m)}\b", sent, re.IGNORECASE): + if m not in found_targets and re.search( + rf"\b{re.escape(m)}\b", sent, re.IGNORECASE + ): score += 2 if any(k in sent_lower for k in domain_keywords): score += 1 @@ -115,9 +153,15 @@ def search_pubtator(query: str, max_results: int = 50) -> List[Dict[str, Any]]: return [] data = resp.json() results = data.get("results", [])[:max_results] - return [{'pmid': r.get('pmid'), 'title': r.get('title', ''), - 'score': r.get('score', 0), 'text_hl': r.get('text_hl', '')} - for r in results] + return [ + { + "pmid": r.get("pmid"), + "title": r.get("title", ""), + "score": r.get("score", 0), + "text_hl": r.get("text_hl", ""), + } + for r in results + ] except Exception as e: logger.error(f"PubTator search error: {e}") return [] @@ -142,43 +186,49 @@ def fetch_bioc_relations_with_text(pmids: List[int]) -> Dict[int, Dict[str, Any] docs = [data] result = {} for doc in docs: - pmid = doc.get('pmid') or doc.get('id') + pmid = doc.get("pmid") or doc.get("id") if not pmid: continue - title = '' - abstract = '' + title = "" + abstract = "" gene_mentions: Dict[str, set] = {} - for passage in doc.get('passages', []): - ptype = passage.get('infons', {}).get('type', '') - if ptype == 'title': - title = passage.get('text', '') - elif ptype == 'abstract': - abstract = passage.get('text', '') - for ann in passage.get('annotations', []): - if ann.get('infons', {}).get('type') == 'Gene': - gene_name = ann.get('infons', {}).get('name', '') - text_mention = ann.get('text', '') + for passage in doc.get("passages", []): + ptype = passage.get("infons", {}).get("type", "") + if ptype == "title": + title = passage.get("text", "") + elif ptype == "abstract": + abstract = passage.get("text", "") + for ann in passage.get("annotations", []): + if ann.get("infons", {}).get("type") == "Gene": + gene_name = ann.get("infons", {}).get("name", "") + text_mention = ann.get("text", "") if gene_name and text_mention: if gene_name.upper() not in gene_mentions: gene_mentions[gene_name.upper()] = set() gene_mentions[gene_name.upper()].add(text_mention) relations = [] - for rel in doc.get('relations', []): - infons = rel.get('infons', {}) - r1 = infons.get('role1', {}) - r2 = infons.get('role2', {}) - rel_type = infons.get('type', 'Unknown') - score = float(infons.get('score', 0)) - if r1.get('type') == 'Gene' and r2.get('type') == 'Gene': - if rel_type in ('Positive_Correlation', 'Negative_Correlation'): - relations.append({ - 'gene1': r1.get('name', ''), 'gene2': r2.get('name', ''), - 'type': rel_type, 'score': score - }) + for rel in doc.get("relations", []): + infons = rel.get("infons", {}) + r1 = infons.get("role1", {}) + r2 = infons.get("role2", {}) + rel_type = infons.get("type", "Unknown") + score = float(infons.get("score", 0)) + if r1.get("type") == "Gene" and r2.get("type") == "Gene": + if rel_type in ("Positive_Correlation", "Negative_Correlation"): + relations.append( + { + "gene1": r1.get("name", ""), + "gene2": r2.get("name", ""), + "type": rel_type, + "score": score, + } + ) gene_mentions_list = {k: list(v) for k, v in gene_mentions.items()} result[int(pmid)] = { - 'relations': relations, 'title': title, - 'abstract': abstract, 'gene_mentions': gene_mentions_list + "relations": relations, + "title": title, + "abstract": abstract, + "gene_mentions": gene_mentions_list, } return result except Exception as e: @@ -229,15 +279,17 @@ def search_pubmed(query: str, max_results: int = 20) -> List[Dict[str, Any]]: results = [] for pmid in pmids: info = data.get(pmid, {}) - results.append({ - "pmid": int(pmid), - "title": info.get("title", ""), - "source": info.get("source", ""), - "authors": ", ".join( - a.get("name", "") for a in info.get("authors", [])[:3] - ), - "pubdate": info.get("pubdate", ""), - }) + results.append( + { + "pmid": int(pmid), + "title": info.get("title", ""), + "source": info.get("source", ""), + "authors": ", ".join( + a.get("name", "") for a in info.get("authors", [])[:3] + ), + "pubdate": info.get("pubdate", ""), + } + ) return results except Exception as e: logger.error(f"PubMed ESummary error: {e}") diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/llm_backend.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/llm_backend.py index 4defd1f..c08d8e9 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/llm_backend.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/llm_backend.py @@ -3,6 +3,7 @@ Supports: anthropic, openai, deepseek, gemini. Provider classes adapted from DeepRare/api/interface.py. """ + from __future__ import annotations import json @@ -15,6 +16,7 @@ # Optional .env support try: from dotenv import load_dotenv + load_dotenv() except ImportError: pass @@ -24,9 +26,11 @@ # Provider classes (from DeepRare/api/interface.py) # ------------------------------------------------------------------ + class Openai_api: def __init__(self, api_key, model): from openai import OpenAI + self.client = OpenAI(api_key=api_key) self.model = model @@ -38,7 +42,7 @@ def get_completion(self, system_prompt, prompt, seed=42): messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, - ] + ], ) return str(completion.choices[0].message.content) except Exception as e: @@ -49,10 +53,11 @@ def get_completion(self, system_prompt, prompt, seed=42): class deepseek_api: def __init__(self, api_key, model): from openai import OpenAI + self.client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com") - if model == 'deepseek-v3-241226': + if model == "deepseek-v3-241226": self.model = "deepseek-chat" - elif model == 'deepseek-r1-250120': + elif model == "deepseek-r1-250120": self.model = "deepseek-reasoner" else: self.model = model @@ -65,7 +70,7 @@ def get_completion(self, system_prompt, prompt, seed=42): {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ], - stream=False + stream=False, ) return str(completion.choices[0].message.content) except Exception as e: @@ -76,6 +81,7 @@ def get_completion(self, system_prompt, prompt, seed=42): class gemini_api: def __init__(self, api_key, model): import google.generativeai as genai + genai.configure(api_key=api_key) self.model = genai.GenerativeModel(model) @@ -92,6 +98,7 @@ def get_completion(self, system_prompt, prompt, seed=42): class claude_api: def __init__(self, api_key, model): import anthropic + self.client = anthropic.Anthropic(api_key=api_key) self.model = model @@ -101,9 +108,7 @@ def get_completion(self, system_prompt, prompt, seed=42): model=self.model, max_tokens=4000, system=system_prompt, - messages=[ - {"role": "user", "content": prompt} - ] + messages=[{"role": "user", "content": prompt}], ) return str(message.content[0].text) except Exception as e: @@ -121,8 +126,10 @@ class stanford_api: POST https://aiapi-prod.stanford.edu/v1/images/generations Auth: Bearer token via STANFORD_API_KEY """ + def __init__(self, api_key, base_url, model): from openai import OpenAI + self.client = OpenAI(api_key=api_key, base_url=base_url) self._validate_model(model) self.model = model @@ -149,7 +156,7 @@ def get_completion(self, system_prompt, prompt, seed=42): messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, - ] + ], ) return str(completion.choices[0].message.content) except Exception as e: @@ -162,19 +169,19 @@ def get_completion(self, system_prompt, prompt, seed=42): # ------------------------------------------------------------------ PROVIDER_MAP = { - "anthropic": (claude_api, "ANTHROPIC_API_KEY"), - "stanford": (stanford_api, "STANFORD_API_KEY"), - "openai": (Openai_api, "OPENAI_API_KEY"), - "deepseek": (deepseek_api, "DEEPSEEK_API_KEY"), - "gemini": (gemini_api, "GOOGLE_API_KEY"), + "anthropic": (claude_api, "ANTHROPIC_API_KEY"), + "stanford": (stanford_api, "STANFORD_API_KEY"), + "openai": (Openai_api, "OPENAI_API_KEY"), + "deepseek": (deepseek_api, "DEEPSEEK_API_KEY"), + "gemini": (gemini_api, "GOOGLE_API_KEY"), } DEFAULT_MODELS = { "anthropic": "claude-sonnet-4-5-20250929", - "stanford": "claude-4-5-sonnet", # also: gpt-5, o1, o3-mini, gpt-5-mini, claude-3-haiku, claude-4-sonnet, claude-opus-4-6, Llama-4, gemini-2.5-pro - "openai": "gpt-4o", - "deepseek": "deepseek-v3-241226", - "gemini": "gemini-2.0-flash", + "stanford": "claude-4-5-sonnet", # also: gpt-5, o1, o3-mini, gpt-5-mini, claude-3-haiku, claude-4-sonnet, claude-opus-4-6, Llama-4, gemini-2.5-pro + "openai": "gpt-4o", + "deepseek": "deepseek-v3-241226", + "gemini": "gemini-2.0-flash", } @@ -182,6 +189,7 @@ def get_completion(self, system_prompt, prompt, seed=42): # LLMBackend wrapper # ------------------------------------------------------------------ + class LLMBackend: """Thin wrapper providing complete() and complete_json() over provider classes.""" @@ -224,7 +232,9 @@ def __init__( else: self._client = cls(api_key=api_key, model=self.model) - logger.info("LLMBackend initialised: provider=%s model=%s", self.provider, self.model) + logger.info( + "LLMBackend initialised: provider=%s model=%s", self.provider, self.model + ) def complete( self, @@ -261,4 +271,3 @@ def token_usage(self) -> dict: "total_input_tokens": self.total_input_tokens, "total_output_tokens": self.total_output_tokens, } - diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/output_writer.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/output_writer.py index 849332b..03c359e 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/output_writer.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/output_writer.py @@ -1,4 +1,5 @@ """JSON + Markdown + HTML output generation per program.""" + from __future__ import annotations import html as html_lib @@ -52,9 +53,7 @@ def write_program_json( for pmid, paper in papers.items(): paper_dict = paper.to_dict() - paper_dict["verification"] = [ - v.to_dict() for v in verifications.get(pmid, []) - ] + paper_dict["verification"] = [v.to_dict() for v in verifications.get(pmid, [])] # Compute overall badge vr_list = verifications.get(pmid, []) if vr_list: @@ -106,8 +105,13 @@ def write_program_markdown( # Context header lines.append("## Program Context\n") - for key in ("top_genes", "regulators", - "day_most_active", "cell_types", "go_enrichment"): + for key in ( + "top_genes", + "regulators", + "day_most_active", + "cell_types", + "go_enrichment", + ): val = program_context.get(key, "") if val: label = key.replace("_", " ").title() @@ -124,7 +128,9 @@ def write_program_markdown( badge = _overall_badge(vr_list) lines.append(f"### {rank}. {paper.title} {badge}\n") - lines.append(f"**PMID**: [{paper.pmid}](https://pubmed.ncbi.nlm.nih.gov/{paper.pmid}/)") + lines.append( + f"**PMID**: [{paper.pmid}](https://pubmed.ncbi.nlm.nih.gov/{paper.pmid}/)" + ) if paper.authors: lines.append(f" | **Authors**: {paper.authors}") if paper.journal: @@ -189,6 +195,7 @@ def _sort_papers( verifications: Dict[int, List[VerificationResult]], ) -> List[int]: """Sort PMIDs: verified first, then by evidence count.""" + def sort_key(pmid): vr_list = verifications.get(pmid, []) all_passed = all(v.passed for v in vr_list) if vr_list else False @@ -275,7 +282,7 @@ def _badge_html(badge: str) -> str: return ( f'' - f'{_esc(label)}' + f"{_esc(label)}" ) @@ -293,7 +300,8 @@ def write_program_html( sorted_pmids = _sort_papers(papers, verifications) n_verified = sum( - 1 for pm in sorted_pmids + 1 + for pm in sorted_pmids if _overall_badge(verifications.get(pm, [])) == "[VERIFIED]" ) @@ -301,10 +309,12 @@ def write_program_html( parts.append(_html_head(program_id)) # -- Header -- - parts.append(f'

Program {_esc(str(program_id))} — Literature Evidence

') + parts.append( + f"

Program {_esc(str(program_id))} — Literature Evidence

" + ) parts.append( f'

Generated {datetime.now().strftime("%Y-%m-%d %H:%M")} ' - f'• {len(papers)} papers • {n_verified} verified

' + f"• {len(papers)} papers • {n_verified} verified

" ) # -- Context card -- @@ -322,20 +332,26 @@ def write_program_html( if val: # Turn gene lists into pills if key in ("top_genes", "regulators"): - genes = [g.strip() for g in val.replace(";", ",").split(",") if g.strip()] - val_html = " ".join(f'{_esc(g)}' for g in genes) + genes = [ + g.strip() for g in val.replace(";", ",").split(",") if g.strip() + ] + val_html = " ".join( + f'{_esc(g)}' for g in genes + ) else: val_html = _esc(val) - parts.append(f'{_esc(label)}{val_html}') - parts.append('') + parts.append( + f'{_esc(label)}{val_html}' + ) + parts.append("") # -- Synthesis -- if synthesis_text and "unavailable" not in synthesis_text.lower(): parts.append('

Synthesis

') - parts.append(f'

{_esc(synthesis_text)}

') + parts.append(f"

{_esc(synthesis_text)}

") # -- Papers -- - parts.append('

Paper Evidence

') + parts.append("

Paper Evidence

") for rank, pmid in enumerate(sorted_pmids, 1): paper = papers[pmid] vr_list = verifications.get(pmid, []) @@ -343,61 +359,69 @@ def write_program_html( parts.append('') # .card.paper + parts.append("") # .card.paper - parts.append('') # close .container + parts.append("") # close .container content = "\n".join(parts) path.write_text(content, encoding="utf-8") diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/paper_fetcher.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/paper_fetcher.py index 227f50a..ec26166 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/paper_fetcher.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/paper_fetcher.py @@ -1,4 +1,5 @@ """Fetch, parse, and cache paper metadata from PubTator3 BioC-JSON.""" + from __future__ import annotations import json @@ -73,7 +74,9 @@ def put(self, paper: PaperData) -> None: p = self._path(paper.pmid) # Atomic write: write to temp then rename tmp = p.with_suffix(".tmp") - tmp.write_text(json.dumps(paper.to_dict(), ensure_ascii=False), encoding="utf-8") + tmp.write_text( + json.dumps(paper.to_dict(), ensure_ascii=False), encoding="utf-8" + ) tmp.rename(p) @@ -132,7 +135,10 @@ def fetch_papers( evidence = {} if abstract or title: evidence = extract_evidence_sentences( - abstract, title, target_set, context_set, + abstract, + title, + target_set, + context_set, domain_keywords=domain_keywords, ) diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/preprocess_input.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/preprocess_input.py index b2acd6e..da2f5a3 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/preprocess_input.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/preprocess_input.py @@ -9,6 +9,7 @@ --sheet-name "Summary" \ --Conditions "condA,condB" """ + from __future__ import annotations import argparse @@ -29,14 +30,18 @@ def _collect_regulators(row: pd.Series, conditions: List[str]) -> str: col = f"Top 5 specific regulators (FDR<0.1) {cond}" val = row.get(col, "") if pd.notna(val) and str(val).strip(): - genes = [g.strip() for g in str(val).replace(",", ";").split(";") if g.strip()] + genes = [ + g.strip() for g in str(val).replace(",", ";").split(";") if g.strip() + ] all_regs.extend(genes) # Also check sigfdr columns as fallback for cond in conditions: col = f"sigfdr0.05_targets_sorted_abslog2fcd_{cond}" val = row.get(col, "") if pd.notna(val) and str(val).strip() and not all_regs: - genes = [g.strip() for g in str(val).replace(",", ";").split(";") if g.strip()] + genes = [ + g.strip() for g in str(val).replace(",", ";").split(";") if g.strip() + ] all_regs.extend(genes) # Deduplicate preserving order seen = set() @@ -118,15 +123,17 @@ def convert_summary_to_lit_input( go_enrichment = _safe_str(row.get("top10_enriched_go_terms", "")) other_enrichment = _safe_str(row.get("top10_enriched_genesets", "")) - rows.append({ - "program_id": prog_id, - "top_genes": genes_raw, - "regulators": regulators, - "day_most_active": day, - "cell_types": cell_types, - "go_enrichment": go_enrichment, - "other_enrichment": other_enrichment, - }) + rows.append( + { + "program_id": prog_id, + "top_genes": genes_raw, + "regulators": regulators, + "day_most_active": day, + "cell_types": cell_types, + "go_enrichment": go_enrichment, + "other_enrichment": other_enrichment, + } + ) return pd.DataFrame(rows) @@ -136,37 +143,53 @@ def main(argv=None): description="Convert PerturbNMF summary table to literature search input" ) parser.add_argument( - "--summary-excel", type=Path, required=True, - help="Path to the summary table Excel file" + "--summary-excel", + type=Path, + required=True, + help="Path to the summary table Excel file", ) parser.add_argument( - "--output", type=Path, required=True, - help="Output Excel path for literature search input" + "--output", + type=Path, + required=True, + help="Output Excel path for literature search input", ) parser.add_argument( - "--sheet-name", type=str, default="Summary", - help="Sheet name in the summary Excel (default: Summary)" + "--sheet-name", + type=str, + default="Summary", + help="Sheet name in the summary Excel (default: Summary)", ) parser.add_argument( - "--top-genes-col", type=str, default="top30_loaded_genes", - help="Column name for top loaded genes (default: top30_loaded_genes)" + "--top-genes-col", + type=str, + default="top30_loaded_genes", + help="Column name for top loaded genes (default: top30_loaded_genes)", ) parser.add_argument( - "--num-genes", type=int, default=None, - help="Truncate gene list to this many genes (default: all)" + "--num-genes", + type=int, + default=None, + help="Truncate gene list to this many genes (default: all)", ) parser.add_argument( - "--Conditions", type=str, required=True, - help="Comma-separated condition/sample names (e.g. 'condA,condB')" + "--Conditions", + type=str, + required=True, + help="Comma-separated condition/sample names (e.g. 'condA,condB')", ) parser.add_argument( - "--programs", type=str, default=None, - help="Comma-separated program IDs to include (default: all)" + "--programs", + type=str, + default=None, + help="Comma-separated program IDs to include (default: all)", ) args = parser.parse_args(argv) - logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + logging.basicConfig( + level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s" + ) if not args.summary_excel.exists(): parser.error(f"Summary Excel not found: {args.summary_excel}") @@ -174,9 +197,15 @@ def main(argv=None): conditions = [c.strip() for c in args.Conditions.split(",")] # Load summary table - logger.info("Loading summary table from %s (sheet: %s)", args.summary_excel, args.sheet_name) - summary_df = pd.read_excel(args.summary_excel, sheet_name=args.sheet_name, index_col=0) - logger.info("Loaded %d programs, %d columns", len(summary_df), len(summary_df.columns)) + logger.info( + "Loading summary table from %s (sheet: %s)", args.summary_excel, args.sheet_name + ) + summary_df = pd.read_excel( + args.summary_excel, sheet_name=args.sheet_name, index_col=0 + ) + logger.info( + "Loaded %d programs, %d columns", len(summary_df), len(summary_df.columns) + ) # Filter programs if specified if args.programs: diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/query_generator.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/query_generator.py index 6c26909..77e0318 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/query_generator.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/query_generator.py @@ -1,4 +1,5 @@ """LLM-powered query generation from program context.""" + from __future__ import annotations import json @@ -63,7 +64,9 @@ def generate_search_queries( system_prompt = _load_prompt("query_generation.txt") context = _format_program_context(program_row) if interaction_verbs: - context += f"\nInteraction verbs to use in queries: {', '.join(interaction_verbs)}" + context += ( + f"\nInteraction verbs to use in queries: {', '.join(interaction_verbs)}" + ) user_prompt = ( f"Generate {max_queries} targeted PubMed search queries for the " f"following gene expression program:\n\n" @@ -84,11 +87,13 @@ def generate_search_queries( valid = [] for q in queries: if isinstance(q, dict) and "query" in q: - valid.append({ - "query": q["query"], - "strategy": q.get("strategy", "unknown"), - "target": q.get("target", ""), - }) + valid.append( + { + "query": q["query"], + "strategy": q.get("strategy", "unknown"), + "target": q.get("target", ""), + } + ) if not valid: logger.warning("No valid queries from LLM; using fallback") return _fallback_queries(program_row, interaction_verbs) @@ -113,20 +118,24 @@ def _fallback_queries( if genes: # Broad gene OR query (top 10) gene_or = " OR ".join(genes[:10]) - queries.append({ - "query": f"({gene_or})", - "strategy": "broad_context", - "target": "papers mentioning any top genes", - }) + queries.append( + { + "query": f"({gene_or})", + "strategy": "broad_context", + "target": "papers mentioning any top genes", + } + ) # Gene pair queries for top 3 pairs for i in range(0, min(6, len(genes)), 2): if i + 1 < len(genes): - queries.append({ - "query": f"({genes[i]} AND {genes[i+1]})", - "strategy": "gene_pair", - "target": f"co-occurrence of {genes[i]} and {genes[i+1]}", - }) + queries.append( + { + "query": f"({genes[i]} AND {genes[i + 1]})", + "strategy": "gene_pair", + "target": f"co-occurrence of {genes[i]} and {genes[i + 1]}", + } + ) # Regulator queries (single column, no direction) regs = row.get("regulators", "") @@ -136,18 +145,22 @@ def _fallback_queries( for reg in reg_list[:3]: if genes: # Basic regulator + target genes query - queries.append({ - "query": f"({reg} AND ({' OR '.join(genes[:5])}))", - "strategy": "regulator_target", - "target": f"regulator {reg} with target genes", - }) + queries.append( + { + "query": f"({reg} AND ({' OR '.join(genes[:5])}))", + "strategy": "regulator_target", + "target": f"regulator {reg} with target genes", + } + ) # Regulator + verb + context query (if verbs provided) if interaction_verbs and cell_type: verb = interaction_verbs[0] - queries.append({ - "query": f'("{reg} {verb}" AND {genes[0]} AND {cell_type})', - "strategy": "regulator_verb_context", - "target": f"{reg} {verb} {genes[0]} in {cell_type}", - }) + queries.append( + { + "query": f'("{reg} {verb}" AND {genes[0]} AND {cell_type})', + "strategy": "regulator_verb_context", + "target": f"{reg} {verb} {genes[0]} in {cell_type}", + } + ) return queries[:8] diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/search_engine.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/search_engine.py index 6b6c5a5..c44fc9d 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/search_engine.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/search_engine.py @@ -1,4 +1,5 @@ """PubTator3 search orchestration: baseline + LLM queries + dedup.""" + from __future__ import annotations import logging @@ -30,9 +31,7 @@ def search_baseline( logger.info("Baseline query (%d genes): %s...", len(genes), query[:120]) return search_pubtator(query, max_results=max_results) - def search_query( - self, query: str, max_results: int = 50 - ) -> List[Dict[str, Any]]: + def search_query(self, query: str, max_results: int = 50) -> List[Dict[str, Any]]: """Execute a single PubTator3 query.""" logger.debug("Executing query: %s", query[:120]) return search_pubtator(query, max_results=max_results) @@ -73,7 +72,9 @@ def search_all_for_program( pmid_info[pmid]["sources"].append("baseline") # 2) PubMed search (broader than PubTator3) - pubmed_results = search_pubmed(f"({' OR '.join(genes[:10])})", max_results=max_pubtator_results) + pubmed_results = search_pubmed( + f"({' OR '.join(genes[:10])})", max_results=max_pubtator_results + ) for r in pubmed_results: pmid = r.get("pmid") if pmid is None: diff --git a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/verification.py b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/verification.py index 2f8173b..8681ebb 100644 --- a/src/Stage3_Interpretation/C_Annotation/Literature_search/src/verification.py +++ b/src/Stage3_Interpretation/C_Annotation/Literature_search/src/verification.py @@ -3,6 +3,7 @@ Level 1 (citation check): Does the cited PMID contain the mentioned genes? Level 2 (semantic check): Does the evidence sentence support the claim? (LLM) """ + from __future__ import annotations import json @@ -75,8 +76,11 @@ def verify_citation( """ if not claimed_genes: return VerificationResult( - level=1, passed=True, confidence=1.0, - details="No genes to check", method="citation_check", + level=1, + passed=True, + confidence=1.0, + details="No genes to check", + method="citation_check", ) doc_genes_upper = set(paper.gene_mentions.keys()) @@ -92,13 +96,18 @@ def verify_citation( confidence = len(found) / len(claimed_genes) if claimed_genes else 1.0 if missing: - details = f"Found {found}; missing {missing} in PMID:{paper.pmid} annotations" + details = ( + f"Found {found}; missing {missing} in PMID:{paper.pmid} annotations" + ) else: details = f"All genes found in PMID:{paper.pmid} annotations" return VerificationResult( - level=1, passed=passed, confidence=confidence, - details=details, method="citation_check", + level=1, + passed=passed, + confidence=confidence, + details=details, + method="citation_check", ) # ------------------------------------------------------------------ @@ -114,7 +123,9 @@ def verify_semantic( """Use LLM to check if evidence sentence supports the claim.""" if self.llm is None: return VerificationResult( - level=2, passed=True, confidence=0.5, + level=2, + passed=True, + confidence=0.5, details="Semantic check skipped (no LLM configured)", method="semantic_check", ) @@ -131,7 +142,9 @@ def verify_semantic( except Exception as exc: logger.warning("Semantic check failed: %s", exc) return VerificationResult( - level=2, passed=True, confidence=0.5, + level=2, + passed=True, + confidence=0.5, details=f"Semantic check error: {exc}", method="semantic_check", ) @@ -141,8 +154,11 @@ def verify_semantic( reasoning = result.get("reasoning", "") return VerificationResult( - level=2, passed=supported, confidence=confidence, - details=reasoning, method="semantic_check", + level=2, + passed=supported, + confidence=confidence, + details=reasoning, + method="semantic_check", ) # ------------------------------------------------------------------ diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/Slurm_Version/run_annotation.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/Slurm_Version/run_annotation.py index 1c52b91..4c8e4d6 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/Slurm_Version/run_annotation.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/Slurm_Version/run_annotation.py @@ -4,6 +4,7 @@ Can be run directly or submitted via SLURM. Wraps the run_pipeline.py orchestrator with sys.path setup. """ + import sys import warnings from pathlib import Path diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/01_genes_to_string_enrichment.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/01_genes_to_string_enrichment.py index 3314f5f..3d1eb4e 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/01_genes_to_string_enrichment.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/01_genes_to_string_enrichment.py @@ -63,12 +63,19 @@ import sys as _sys from pathlib import Path as _Path + _sys.path.insert(0, str(_Path(__file__).resolve().parent)) -from column_mapper import ColumnMapper, standardize_gene_loading, standardize_celltype_enrichment +from column_mapper import ( + ColumnMapper, + standardize_gene_loading, + standardize_celltype_enrichment, +) -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) logger = logging.getLogger(__name__) """ @@ -151,13 +158,23 @@ def parse_topics(value: Optional[object]) -> Optional[Set[int]]: def apply_test_mode( args: argparse.Namespace, config: Dict[str, Any], cli_overrides: Set[str] ) -> argparse.Namespace: - test_cfg = config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + test_cfg = ( + config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + ) enabled = bool(test_cfg.get("enabled") or config.get("test_mode")) if not enabled: return args - topics = test_cfg.get("topics") or test_cfg.get("programs") or config.get("test_programs") - if hasattr(args, "topics") and "topics" not in cli_overrides and not getattr(args, "topics", None): + topics = ( + test_cfg.get("topics") + or test_cfg.get("programs") + or config.get("test_programs") + ) + if ( + hasattr(args, "topics") + and "topics" not in cli_overrides + and not getattr(args, "topics", None) + ): if topics is not None: args.topics = topics return args @@ -168,7 +185,9 @@ def apply_config_overrides( config: Dict[str, Any], cli_overrides: Set[str], ) -> argparse.Namespace: - steps_cfg = config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + steps_cfg = ( + config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + ) step_cfg = steps_cfg.get("string_enrichment", {}) if isinstance(step_cfg, dict) and args.command in step_cfg: step_cfg = step_cfg.get(args.command, {}) @@ -193,24 +212,24 @@ def apply_config_overrides( # ----------------------------- Cell-type summary ----------------------------- # Thresholds for categorizing cell-type enrichment (FDR < 0.05 required) CELLTYPE_THRESHOLDS = { - 'highly_cell_type_specific': {'log2fc_min': 3.0, 'log2fc_max': None}, - 'moderately_enriched': {'log2fc_min': 1.5, 'log2fc_max': 3.0}, - 'weakly_enriched': {'log2fc_min': 0.5, 'log2fc_max': 1.5}, - 'depleted': {'log2fc_min': None, 'log2fc_max': -0.5}, + "highly_cell_type_specific": {"log2fc_min": 3.0, "log2fc_max": None}, + "moderately_enriched": {"log2fc_min": 1.5, "log2fc_max": 3.0}, + "weakly_enriched": {"log2fc_min": 0.5, "log2fc_max": 1.5}, + "depleted": {"log2fc_min": None, "log2fc_max": -0.5}, } # Column order for cell-type summary output CELLTYPE_CATEGORIES = [ - 'highly_cell_type_specific', - 'moderately_enriched', - 'weakly_enriched', - 'depleted', + "highly_cell_type_specific", + "moderately_enriched", + "weakly_enriched", + "depleted", ] def extract_program_id(value: object) -> Optional[int]: """Extract numeric program ID from various naming formats. - + Handles formats like: - 'Program_1', 'program_1' -> 1 - 'Topic_1', 'topic_1' -> 1 @@ -218,39 +237,39 @@ def extract_program_id(value: object) -> Optional[int]: - 'X1', 'X_1' -> 1 (regulator file format) - '1', 1 -> 1 - 'Program1', 'Topic1' -> 1 - + Args: value: String or int containing program identifier - + Returns: Integer program ID or None if parsing fails """ import re - + if value is None: return None - + # If already an integer, return it if isinstance(value, (int, np.integer)): return int(value) - + # Convert to string and strip whitespace val_str = str(value).strip() - + # Try direct integer conversion first try: return int(val_str) except (ValueError, TypeError): pass - + # Try to extract digits from common patterns # Patterns: Program_X, Topic_X, P_X, X_X (regulator format), program_X, topic_X, p_X patterns = [ - r'^(?:program|topic|p|x)_(\d+)$', # program_1, topic_1, p_1, x_1 - r'^(?:program|topic|p|x)(\d+)$', # program1, topic1, p1, x1 - r'^(\d+)$', # just the number + r"^(?:program|topic|p|x)_(\d+)$", # program_1, topic_1, p_1, x_1 + r"^(?:program|topic|p|x)(\d+)$", # program1, topic1, p1, x1 + r"^(\d+)$", # just the number ] - + for pattern in patterns: match = re.match(pattern, val_str, re.IGNORECASE) if match: @@ -258,109 +277,133 @@ def extract_program_id(value: object) -> Optional[int]: return int(match.group(1)) except (ValueError, TypeError, IndexError): continue - + return None def validate_celltype_enrichment(df: pd.DataFrame, file_path: Path) -> bool: """Validate cell-type enrichment DataFrame format and content. - + Non-fatal validation that logs warnings instead of raising errors. This allows the pipeline to continue even with imperfect input data. - + Args: df: DataFrame to validate file_path: Path to the file (for warning messages) - + Returns: True if validation passed, False if issues were found """ warnings_list = [] has_critical_errors = False - + # 1. Check required columns - required_cols = {'cell_type', 'program', 'log2_fc_in_vs_out', 'fdr'} + required_cols = {"cell_type", "program", "log2_fc_in_vs_out", "fdr"} missing = required_cols - set(df.columns) if missing: warnings_list.append(f"Missing required columns: {sorted(missing)}") warnings_list.append(f" Found columns: {sorted(df.columns)}") warnings_list.append(f" Required: {sorted(required_cols)}") has_critical_errors = True - + # 2. Check DataFrame is not empty if df.empty: warnings_list.append("Cell-type enrichment file is empty") has_critical_errors = True - + # If critical errors, log and return early if has_critical_errors: logger.warning(f"Cell-type enrichment file has critical issues: {file_path}") for w in warnings_list: logger.warning(f" {w}") return False - + # 3. Validate 'program' column - flexible format checking df_temp = df.copy() - df_temp['program_id_parsed'] = df_temp['program'].apply(extract_program_id) - unparseable = df_temp[df_temp['program_id_parsed'].isna()] + df_temp["program_id_parsed"] = df_temp["program"].apply(extract_program_id) + unparseable = df_temp[df_temp["program_id_parsed"].isna()] if not unparseable.empty: - sample_invalid = unparseable['program'].head(5).tolist() - warnings_list.append(f"Could not parse {len(unparseable)} program identifiers: {sample_invalid}") - warnings_list.append(f" Supported formats: Program_X, program_X, Topic_X, topic_X, P_X, p_X, X_X (regulator), ProgramX, TopicX, X") - + sample_invalid = unparseable["program"].head(5).tolist() + warnings_list.append( + f"Could not parse {len(unparseable)} program identifiers: {sample_invalid}" + ) + warnings_list.append( + f" Supported formats: Program_X, program_X, Topic_X, topic_X, P_X, p_X, X_X (regulator), ProgramX, TopicX, X" + ) + # 4. Validate 'log2_fc_in_vs_out' is numeric try: - fc_numeric = pd.to_numeric(df['log2_fc_in_vs_out'], errors='coerce') + fc_numeric = pd.to_numeric(df["log2_fc_in_vs_out"], errors="coerce") non_numeric_fc = df[fc_numeric.isna()] if not non_numeric_fc.empty: - sample_bad = non_numeric_fc['log2_fc_in_vs_out'].head(5).tolist() - warnings_list.append(f"Non-numeric values in 'log2_fc_in_vs_out': {sample_bad}") - warnings_list.append(f" Found {len(non_numeric_fc)} non-numeric log2FC values (will be ignored)") + sample_bad = non_numeric_fc["log2_fc_in_vs_out"].head(5).tolist() + warnings_list.append( + f"Non-numeric values in 'log2_fc_in_vs_out': {sample_bad}" + ) + warnings_list.append( + f" Found {len(non_numeric_fc)} non-numeric log2FC values (will be ignored)" + ) except Exception as e: warnings_list.append(f"Failed to validate 'log2_fc_in_vs_out' column: {e}") - + # 5. Validate 'fdr' is numeric and in valid range [0, 1] try: - fdr_numeric = pd.to_numeric(df['fdr'], errors='coerce') + fdr_numeric = pd.to_numeric(df["fdr"], errors="coerce") non_numeric_fdr = df[fdr_numeric.isna()] if not non_numeric_fdr.empty: - sample_bad = non_numeric_fdr['fdr'].head(5).tolist() + sample_bad = non_numeric_fdr["fdr"].head(5).tolist() warnings_list.append(f"Non-numeric values in 'fdr': {sample_bad}") - warnings_list.append(f" Found {len(non_numeric_fdr)} non-numeric FDR values (will be ignored)") - + warnings_list.append( + f" Found {len(non_numeric_fdr)} non-numeric FDR values (will be ignored)" + ) + # Check FDR range (should be 0-1 for proper FDR values) valid_fdr = fdr_numeric.dropna() if len(valid_fdr) > 0: out_of_range_mask = (valid_fdr < 0) | (valid_fdr > 1) if out_of_range_mask.any(): out_of_range_vals = valid_fdr[out_of_range_mask].head(5).tolist() - warnings_list.append(f"FDR values out of range [0, 1]: {out_of_range_vals}") - warnings_list.append(f" Found {out_of_range_mask.sum()} out-of-range FDR values") + warnings_list.append( + f"FDR values out of range [0, 1]: {out_of_range_vals}" + ) + warnings_list.append( + f" Found {out_of_range_mask.sum()} out-of-range FDR values" + ) except Exception as e: warnings_list.append(f"Failed to validate 'fdr' column: {e}") - + # 6. Check for completely empty cell_type values - empty_celltypes = df[df['cell_type'].isna() | (df['cell_type'].astype(str).str.strip() == '')] + empty_celltypes = df[ + df["cell_type"].isna() | (df["cell_type"].astype(str).str.strip() == "") + ] if not empty_celltypes.empty: - warnings_list.append(f"Found {len(empty_celltypes)} rows with empty 'cell_type' values (will be ignored)") - + warnings_list.append( + f"Found {len(empty_celltypes)} rows with empty 'cell_type' values (will be ignored)" + ) + # 7. Log summary statistics (informational) - n_programs = df['program'].nunique() - n_celltypes = df['cell_type'].nunique() + n_programs = df["program"].nunique() + n_celltypes = df["cell_type"].nunique() logger.info(f"Cell-type enrichment file summary: {file_path}") logger.info(f" Total rows: {len(df)}") logger.info(f" Unique programs: {n_programs}") logger.info(f" Unique cell types: {n_celltypes}") - + # Check for reasonable data coverage (warnings only) if n_programs < 10: - warnings_list.append(f"Very few programs found ({n_programs}). Expected 50-100 for typical cNMF results.") + warnings_list.append( + f"Very few programs found ({n_programs}). Expected 50-100 for typical cNMF results." + ) if n_celltypes < 3: - warnings_list.append(f"Very few cell types found ({n_celltypes}). Expected multiple cell types.") - + warnings_list.append( + f"Very few cell types found ({n_celltypes}). Expected multiple cell types." + ) + # Log all warnings if warnings_list: - logger.warning(f"Cell-type enrichment validation found {len(warnings_list)} issue(s):") + logger.warning( + f"Cell-type enrichment validation found {len(warnings_list)} issue(s):" + ) for w in warnings_list: logger.warning(f" {w}") return False @@ -412,45 +455,49 @@ def generate_celltype_summary( validate_celltype_enrichment(df, enrichment_file) # Validate required columns (critical check) - required_cols = {'cell_type', 'program', 'log2_fc_in_vs_out', 'fdr'} + required_cols = {"cell_type", "program", "log2_fc_in_vs_out", "fdr"} missing = required_cols - set(df.columns) if missing: logger.error(f"Enrichment file missing required columns: {missing}") logger.error(f" Found columns: {sorted(df.columns)}") - logger.error(f" Cannot proceed without required columns. Please check your input file.") + logger.error( + f" Cannot proceed without required columns. Please check your input file." + ) raise ValueError(f"Enrichment file missing required columns: {missing}") # Extract program ID using flexible parser (handles Program_X, program_X, Topic_X, X, etc.) - df['program_id'] = df['program'].apply(extract_program_id) - + df["program_id"] = df["program"].apply(extract_program_id) + # Check for programs that couldn't be parsed - unparsed = df[df['program_id'].isna()] + unparsed = df[df["program_id"].isna()] if not unparsed.empty: - logger.warning(f"Could not parse program IDs for {len(unparsed)} rows (will be excluded):") + logger.warning( + f"Could not parse program IDs for {len(unparsed)} rows (will be excluded):" + ) logger.warning(f" Sample values: {unparsed['program'].head(5).tolist()}") - df = df[df['program_id'].notna()].copy() - + df = df[df["program_id"].notna()].copy() + # Ensure program_id is integer type - df['program_id'] = df['program_id'].astype(int) + df["program_id"] = df["program_id"].astype(int) # Filter to requested topics if specified if topics: - df = df[df['program_id'].isin(list(topics))].copy() + df = df[df["program_id"].isin(list(topics))].copy() logger.info(f"Filtered to {len(df)} rows for topics: {sorted(topics)}") # Filter to significant results - df_sig = df[df['fdr'] < fdr_threshold].copy() + df_sig = df[df["fdr"] < fdr_threshold].copy() logger.info(f"Found {len(df_sig)} significant rows (FDR < {fdr_threshold})") # Use cell_type values as-is (assume already has correct names) - df_sig['cell_type_display'] = df_sig['cell_type'] + df_sig["cell_type_display"] = df_sig["cell_type"] # Categorize each row by log2FC thresholds def categorize_row(row: pd.Series) -> Optional[str]: - log2fc = row['log2_fc_in_vs_out'] + log2fc = row["log2_fc_in_vs_out"] for cat, bounds in thresholds.items(): - min_val = bounds.get('log2fc_min') - max_val = bounds.get('log2fc_max') + min_val = bounds.get("log2fc_min") + max_val = bounds.get("log2fc_max") # Check if log2fc falls in this category if min_val is not None and max_val is not None: if min_val <= log2fc < max_val: @@ -463,28 +510,28 @@ def categorize_row(row: pd.Series) -> Optional[str]: return cat return None - df_sig['category'] = df_sig.apply(categorize_row, axis=1) - df_categorized = df_sig.dropna(subset=['category']) + df_sig["category"] = df_sig.apply(categorize_row, axis=1) + df_categorized = df_sig.dropna(subset=["category"]) logger.info(f"Categorized {len(df_categorized)} rows into enrichment categories") # Build summary: for each program, collect cell types per category - all_programs = sorted(df['program_id'].unique()) + all_programs = sorted(df["program_id"].unique()) records = [] for pid in all_programs: - program_data = df_categorized[df_categorized['program_id'] == pid] - row = {'program': f'Program_{pid}'} + program_data = df_categorized[df_categorized["program_id"] == pid] + row = {"program": f"Program_{pid}"} for cat in CELLTYPE_CATEGORIES: - cat_data = program_data[program_data['category'] == cat] - cell_types = sorted(cat_data['cell_type_display'].unique()) - row[cat] = '|'.join(cell_types) if cell_types else '' + cat_data = program_data[program_data["category"] == cat] + cell_types = sorted(cat_data["cell_type_display"].unique()) + row[cat] = "|".join(cell_types) if cell_types else "" records.append(row) # Create output DataFrame summary_df = pd.DataFrame(records) - summary_df = summary_df[['program'] + CELLTYPE_CATEGORIES] + summary_df = summary_df[["program"] + CELLTYPE_CATEGORIES] # Write output ensure_parent_dir(str(output_file)) @@ -493,6 +540,7 @@ def categorize_row(row: pd.Series) -> Optional[str]: return len(summary_df) + """ @description Default output path resolver for the STRING enrichment CLI. @@ -513,12 +561,18 @@ def apply_default_paths(args: argparse.Namespace) -> argparse.Namespace: n_top = getattr(args, "n_top", None) or 100 if hasattr(args, "json_out") and not getattr(args, "json_out", None): - args.json_out = str(DEFAULT_ENRICH_DIR / DEFAULT_GENES_JSON_TEMPLATE.format(n_top=n_top)) + args.json_out = str( + DEFAULT_ENRICH_DIR / DEFAULT_GENES_JSON_TEMPLATE.format(n_top=n_top) + ) if hasattr(args, "csv_out") and not getattr(args, "csv_out", None): - args.csv_out = str(DEFAULT_ENRICH_DIR / DEFAULT_GENES_OVERVIEW_TEMPLATE.format(n_top=n_top)) + args.csv_out = str( + DEFAULT_ENRICH_DIR / DEFAULT_GENES_OVERVIEW_TEMPLATE.format(n_top=n_top) + ) if hasattr(args, "out_csv_full") and not getattr(args, "out_csv_full", None): args.out_csv_full = str(DEFAULT_ENRICH_DIR / DEFAULT_STRING_FULL) - if hasattr(args, "out_csv_filtered") and not getattr(args, "out_csv_filtered", None): + if hasattr(args, "out_csv_filtered") and not getattr( + args, "out_csv_filtered", None + ): args.out_csv_filtered = str(DEFAULT_ENRICH_DIR / DEFAULT_STRING_FILTERED) return args @@ -544,7 +598,9 @@ def cache_path(cache_dir: Path, program_id: int) -> Path: return cache_dir / f"program_{program_id}_enrichment.json" -def load_cached_results(cache_dir: Path, program_id: int) -> Optional[List[Dict[str, Any]]]: +def load_cached_results( + cache_dir: Path, program_id: int +) -> Optional[List[Dict[str, Any]]]: cache_file = cache_path(cache_dir, program_id) if not cache_file.exists(): return None @@ -559,7 +615,9 @@ def load_cached_results(cache_dir: Path, program_id: int) -> Optional[List[Dict[ return data -def write_cached_results(cache_dir: Path, program_id: int, results: List[Dict[str, Any]]) -> None: +def write_cached_results( + cache_dir: Path, program_id: int, results: List[Dict[str, Any]] +) -> None: cache_dir.mkdir(parents=True, exist_ok=True) cache_file = cache_path(cache_dir, program_id) cache_file.write_text(json.dumps(results, indent=2), encoding="utf-8") @@ -567,6 +625,7 @@ def write_cached_results(cache_dir: Path, program_id: int, results: List[Dict[st # --------------------------- Extract top genes (CSV) -------------------------- + def ensure_parent_dir(path_str: str) -> None: path = Path(path_str) if path.parent and not path.parent.exists(): @@ -580,7 +639,7 @@ def resolve_program_id_column(df: pd.DataFrame) -> str: """ mapper = ColumnMapper(df) try: - actual_col = mapper.get_column('program_id', required=True) + actual_col = mapper.get_column("program_id", required=True) return actual_col except ValueError as e: raise ValueError( @@ -604,15 +663,15 @@ def extract_top_genes_by_program( Standardizes Gene/Name, Score/Loading columns automatically. """ mapper = ColumnMapper(df) - + # Get standardized column names try: - cols = mapper.get_columns(['gene', 'score'], required=True) - gene_col = cols['gene'] - score_col = cols['score'] + cols = mapper.get_columns(["gene", "score"], required=True) + gene_col = cols["gene"] + score_col = cols["score"] except ValueError as e: raise ValueError(f"Missing required columns for gene extraction: {e}") - + # Verify program ID column exists if id_col not in df.columns: raise ValueError(f"Program ID column '{id_col}' not found in DataFrame") @@ -665,26 +724,22 @@ def build_uniqueness_table(df: pd.DataFrame, id_col: str) -> pd.DataFrame: Build gene loading table with UniquenessScore using flexible column names. """ mapper = ColumnMapper(df) - + # Get standardized column names try: - cols = mapper.get_columns(['gene', 'score'], required=True) - gene_col = cols['gene'] - score_col = cols['score'] + cols = mapper.get_columns(["gene", "score"], required=True) + gene_col = cols["gene"] + score_col = cols["score"] except ValueError as e: raise ValueError(f"Missing required columns for uniqueness computation: {e}") - + if id_col not in df.columns: raise ValueError(f"Program ID column '{id_col}' not found") work = df.copy() - + # Standardize column names to Name, Score, program_id - rename_map = { - gene_col: 'Name', - score_col: 'Score', - id_col: 'program_id' - } + rename_map = {gene_col: "Name", score_col: "Score", id_col: "program_id"} work = work.rename(columns=rename_map) if "UniquenessScore" not in work.columns or work["UniquenessScore"].isna().all(): @@ -705,7 +760,9 @@ def build_uniqueness_table(df: pd.DataFrame, id_col: str) -> pd.DataFrame: columns = ["Name", "Score", "program_id", "UniquenessScore"] out_df = work[columns].copy() - out_df.dropna(subset=["Name", "Score", "program_id", "UniquenessScore"], inplace=True) + out_df.dropna( + subset=["Name", "Score", "program_id", "UniquenessScore"], inplace=True + ) return out_df @@ -717,12 +774,16 @@ def build_overview_long_table( for program_id_str, genes in top_map.items(): program_id = normalize_program_id(program_id_str) if program_id not in sub_indexed_cache: - sub_indexed_cache[program_id] = ( - df[df[id_col] == program_id].set_index("Name") + sub_indexed_cache[program_id] = df[df[id_col] == program_id].set_index( + "Name" ) sub_idx = sub_indexed_cache[program_id] for rank, gene in enumerate(genes, start=1): - score = float(sub_idx.loc[gene, "Score"]) if gene in sub_idx.index else float("nan") + score = ( + float(sub_idx.loc[gene, "Score"]) + if gene in sub_idx.index + else float("nan") + ) records.append( {"program_id": program_id, "rank": rank, "gene": gene, "score": score} ) @@ -741,7 +802,9 @@ def cmd_extract(args: argparse.Namespace) -> int: logger.error(f"Input not found: {input_path}") return 2 df = pd.read_csv(input_path) - logger.info(f"Loaded input: {input_path} with shape {df.shape} and columns {list(df.columns)}") + logger.info( + f"Loaded input: {input_path} with shape {df.shape} and columns {list(df.columns)}" + ) id_col = resolve_program_id_column(df) uniqueness_out = args.gene_loading_out or str(default_uniqueness_output(input_path)) @@ -779,9 +842,15 @@ def cmd_extract(args: argparse.Namespace) -> int: STRING_ENRICH_ENDPOINT = "https://string-db.org/api/json/enrichment" -def call_string_enrichment(genes: List[str], species: int, retries: int = 3, sleep_between: float = 0.6) -> List[Dict[str, Any]]: +def call_string_enrichment( + genes: List[str], species: int, retries: int = 3, sleep_between: float = 0.6 +) -> List[Dict[str, Any]]: identifiers_value = "\r".join(genes) - params = {"identifiers": identifiers_value, "species": species, "caller_identity": "topic_analysis_string_enrichment"} + params = { + "identifiers": identifiers_value, + "species": species, + "caller_identity": "topic_analysis_string_enrichment", + } attempt = 0 while attempt <= retries: @@ -795,9 +864,13 @@ def call_string_enrichment(genes: List[str], species: int, retries: int = 3, sle data = [] return data if isinstance(data, list) else [] else: - logger.warning(f"STRING returned status {response.status_code}: {response.text[:200]}") + logger.warning( + f"STRING returned status {response.status_code}: {response.text[:200]}" + ) except requests.RequestException as e: - logger.warning(f"HTTP error on STRING request (attempt {attempt+1}/{retries+1}): {e}") + logger.warning( + f"HTTP error on STRING request (attempt {attempt + 1}/{retries + 1}): {e}" + ) attempt += 1 time.sleep(min(2.0, sleep_between * (attempt + 1))) @@ -819,9 +892,13 @@ def build_full_csv(program_to_results: Dict[str, List[Dict[str, Any]]]) -> pd.Da "fdr": float(t.get("fdr", float("nan"))), "p_value": float(t.get("p_value", float("nan"))), "number_of_genes": int(t.get("number_of_genes", 0)), - "number_of_genes_in_background": int(t.get("number_of_genes_in_background", 0)), + "number_of_genes_in_background": int( + t.get("number_of_genes_in_background", 0) + ), "ncbiTaxonId": int(t.get("ncbiTaxonId", 0)), - "inputGenes": "|".join(t.get("inputGenes", [])) if t.get("inputGenes") else "", + "inputGenes": "|".join(t.get("inputGenes", [])) + if t.get("inputGenes") + else "", } ) df = pd.DataFrame(rows) @@ -849,7 +926,9 @@ def cmd_enrich(args: argparse.Namespace) -> int: logger.error("--figures-dir is required when --figures-only is set.") return 2 if not args.figures_only and (not args.out_csv_full or not args.out_csv_filtered): - logger.error("--out-csv-full and --out-csv-filtered are required unless --figures-only is set.") + logger.error( + "--out-csv-full and --out-csv-filtered are required unless --figures-only is set." + ) return 2 genes_path = Path(args.genes_json) @@ -857,7 +936,9 @@ def cmd_enrich(args: argparse.Namespace) -> int: logger.error(f"Genes JSON not found: {genes_path}") return 2 - program_to_genes: Dict[str, List[str]] = json.loads(genes_path.read_text(encoding="utf-8")) + program_to_genes: Dict[str, List[str]] = json.loads( + genes_path.read_text(encoding="utf-8") + ) allowed_topics = parse_topics(args.topics) if allowed_topics: program_to_genes = { @@ -865,7 +946,9 @@ def cmd_enrich(args: argparse.Namespace) -> int: for pid, genes in program_to_genes.items() if int(pid) in allowed_topics } - logger.info(f"Loaded gene lists for {len(program_to_genes)} programs from {genes_path}") + logger.info( + f"Loaded gene lists for {len(program_to_genes)} programs from {genes_path}" + ) existing_full_df: Optional[pd.DataFrame] = None existing_programs: Set[int] = set() @@ -891,11 +974,17 @@ def cmd_enrich(args: argparse.Namespace) -> int: program_to_results: Dict[str, List[Dict[str, Any]]] = {} total = len(program_to_genes) - for idx, program_id in enumerate(sorted(program_to_genes.keys(), key=lambda x: int(x)), start=1): + for idx, program_id in enumerate( + sorted(program_to_genes.keys(), key=lambda x: int(x)), start=1 + ): program_id_int = int(program_id) - genes = [g for g in program_to_genes[program_id] if isinstance(g, str) and g.strip()] + genes = [ + g for g in program_to_genes[program_id] if isinstance(g, str) and g.strip() + ] skip_enrichment = ( - args.resume and not args.force_refresh and program_id_int in existing_programs + args.resume + and not args.force_refresh + and program_id_int in existing_programs ) results: Optional[List[Dict[str, Any]]] = None @@ -922,7 +1011,9 @@ def cmd_enrich(args: argparse.Namespace) -> int: retries=args.retries, sleep_between=args.sleep, ) - logger.info("Program %s: retrieved %d enriched terms", program_id, len(results)) + logger.info( + "Program %s: retrieved %d enriched terms", program_id, len(results) + ) if cache_dir: write_cached_results(cache_dir, program_id_int, results) time.sleep(args.sleep) @@ -1071,18 +1162,28 @@ def download_string_enrichment_figure( # -------------------------------- Entry points ------------------------------- + def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description="Extract program gene lists and run STRING enrichment.") + parser = argparse.ArgumentParser( + description="Extract program gene lists and run STRING enrichment." + ) subparsers = parser.add_subparsers(dest="command", required=True) # extract - p_extract = subparsers.add_parser("extract", help="Extract top-N genes per program from loading CSV") + p_extract = subparsers.add_parser( + "extract", help="Extract top-N genes per program from loading CSV" + ) p_extract.add_argument("--config", help="Path to config file (YAML or JSON)") p_extract.add_argument( "--input", help="Path to loading CSV (columns: Name, Score, RowID or program_id)", ) - p_extract.add_argument("--n-top", type=int, default=100, help="Number of top genes per program to extract") + p_extract.add_argument( + "--n-top", + type=int, + default=100, + help="Number of top genes per program to extract", + ) p_extract.add_argument("--json-out", help="Output JSON {program_id: [genes...]}") p_extract.add_argument("--csv-out", help="Output overview CSV") p_extract.add_argument( @@ -1097,12 +1198,23 @@ def build_parser() -> argparse.ArgumentParser: p_extract.set_defaults(func=cmd_extract) # enrich - p_enrich = subparsers.add_parser("enrich", help="Run STRING enrichment for program gene lists from JSON") + p_enrich = subparsers.add_parser( + "enrich", help="Run STRING enrichment for program gene lists from JSON" + ) p_enrich.add_argument("--config", help="Path to config file (YAML or JSON)") - p_enrich.add_argument("--genes-json", help="Path to JSON mapping {program_id: [genes...]}") - p_enrich.add_argument("--species", type=int, default=10090, help="NCBI/STRING species id (default: 10090 mouse)") + p_enrich.add_argument( + "--genes-json", help="Path to JSON mapping {program_id: [genes...]}" + ) + p_enrich.add_argument( + "--species", + type=int, + default=10090, + help="NCBI/STRING species id (default: 10090 mouse)", + ) p_enrich.add_argument("--out-csv-full", help="Full unfiltered CSV output path") - p_enrich.add_argument("--out-csv-filtered", help="Filtered CSV (Process/KEGG only, background<500)") + p_enrich.add_argument( + "--out-csv-filtered", help="Filtered CSV (Process/KEGG only, background<500)" + ) p_enrich.add_argument("--figures-dir", help="Directory to save enrichment figures") p_enrich.add_argument( "--figures-only", @@ -1123,8 +1235,12 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Ignore cached results and re-query STRING", ) - p_enrich.add_argument("--sleep", type=float, default=0.6, help="Sleep seconds between API calls") - p_enrich.add_argument("--retries", type=int, default=3, help="Retries per program on HTTP failures") + p_enrich.add_argument( + "--sleep", type=float, default=0.6, help="Sleep seconds between API calls" + ) + p_enrich.add_argument( + "--retries", type=int, default=3, help="Retries per program on HTTP failures" + ) p_enrich.add_argument( "--topics", type=str, @@ -1171,27 +1287,32 @@ def run_all(args: argparse.Namespace) -> int: rc = cmd_extract(args) if rc != 0: return rc - + # Generate cell-type summary if enrichment file is provided - if getattr(args, 'celltype_enrichment', None): + if getattr(args, "celltype_enrichment", None): celltype_path = Path(args.celltype_enrichment) if not celltype_path.exists(): logger.error(f"Cell-type enrichment file not found: {celltype_path}") return 2 - + # Determine output path - if getattr(args, 'celltype_summary_out', None): + if getattr(args, "celltype_summary_out", None): summary_out = Path(args.celltype_summary_out) else: # Default: place next to enrichment file or in output dir - if getattr(args, 'out_csv_full', None): - summary_out = Path(args.out_csv_full).parent / "celltype_summary.csv" + if getattr(args, "out_csv_full", None): + summary_out = ( + Path(args.out_csv_full).parent / "celltype_summary.csv" + ) else: - summary_out = celltype_path.parent / "program_celltype_annotations_summary_generated.csv" - + summary_out = ( + celltype_path.parent + / "program_celltype_annotations_summary_generated.csv" + ) + # Parse topics topics = parse_topics(args.topics) - + try: n_programs = generate_celltype_summary( enrichment_file=celltype_path, @@ -1202,7 +1323,7 @@ def run_all(args: argparse.Namespace) -> int: except Exception as e: logger.error(f"Failed to generate cell-type summary: {e}") return 2 - + enrich_args = argparse.Namespace( genes_json=args.json_out, species=args.species, diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/02_fetch_ncbi_data.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/02_fetch_ncbi_data.py index 3be3a50..f1bd7c2 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/02_fetch_ncbi_data.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/02_fetch_ncbi_data.py @@ -54,7 +54,9 @@ ) from column_mapper import standardize_regulator_results -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) logger = logging.getLogger(__name__) """ @@ -120,18 +122,32 @@ def parse_topics_value(value: Optional[object]) -> Optional[List[int]]: def apply_test_mode( args: argparse.Namespace, config: Dict[str, Any], cli_overrides: Set[str] ) -> argparse.Namespace: - test_cfg = config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + test_cfg = ( + config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + ) enabled = bool(test_cfg.get("enabled") or config.get("test_mode")) if not enabled: return args - topics = test_cfg.get("topics") or test_cfg.get("programs") or config.get("test_programs") - if hasattr(args, "topics") and "topics" not in cli_overrides and not getattr(args, "topics", None): + topics = ( + test_cfg.get("topics") + or test_cfg.get("programs") + or config.get("test_programs") + ) + if ( + hasattr(args, "topics") + and "topics" not in cli_overrides + and not getattr(args, "topics", None) + ): if topics is not None: args.topics = topics num_programs = test_cfg.get("num_programs") or config.get("test_num_programs") - if hasattr(args, "num_programs") and "num_programs" not in cli_overrides and not getattr(args, "num_programs", None): + if ( + hasattr(args, "num_programs") + and "num_programs" not in cli_overrides + and not getattr(args, "num_programs", None) + ): if num_programs is not None: args.num_programs = num_programs return args @@ -140,7 +156,9 @@ def apply_test_mode( def apply_config_overrides( args: argparse.Namespace, config: Dict[str, Any], cli_overrides: Set[str] ) -> argparse.Namespace: - steps_cfg = config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + steps_cfg = ( + config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + ) step_cfg = steps_cfg.get("literature_fetch", {}) if not isinstance(step_cfg, dict): return args @@ -153,6 +171,7 @@ def apply_config_overrides( setattr(args, dest, value) return args + # Global uniqueness score helpers """ @description @@ -212,117 +231,153 @@ def add_global_uniqueness_scores(df: pd.DataFrame) -> pd.DataFrame: return updated -def ensure_global_uniqueness( - df: pd.DataFrame, logger: logging.Logger -) -> pd.DataFrame: +def ensure_global_uniqueness(df: pd.DataFrame, logger: logging.Logger) -> pd.DataFrame: if "UniquenessScore" in df.columns and not df["UniquenessScore"].isna().all(): return df logger.info("UniquenessScore missing; computing global uniqueness scores.") return add_global_uniqueness_scores(df) + # Domain Constants for Programmatic Scoring DOMAIN_KEYWORDS = [ - "inflammation", "proliferation", "migration", "adhesion", "metabolism", - "hypoxia", "differentiation", "apoptosis", "cell cycle", "signaling", - "stress response", "development" + "inflammation", + "proliferation", + "migration", + "adhesion", + "metabolism", + "hypoxia", + "differentiation", + "apoptosis", + "cell cycle", + "signaling", + "stress response", + "development", ] INTERACTION_VERBS = [ - "regulates", "induces", "promotes", "inhibits", "suppresses", "activates", - "binds", "phosphorylates", "modulates", "mediates", "targets", "controls", - "decreases", "increases", "blocks", "triggers", "catalyzes" + "regulates", + "induces", + "promotes", + "inhibits", + "suppresses", + "activates", + "binds", + "phosphorylates", + "modulates", + "mediates", + "targets", + "controls", + "decreases", + "increases", + "blocks", + "triggers", + "catalyzes", ] # Mechanistic verbs for sentence extraction (subset focused on causality) MECHANISTIC_VERBS = [ - "regulates", "induces", "promotes", "inhibits", "suppresses", "activates", - "mediates", "controls", "stimulates", "blocks", "enhances", "reduces", - "increases", "decreases", "triggers", "phosphorylates", "upregulates", - "downregulates", "attenuates", "potentiates", "modulates", "drives" + "regulates", + "induces", + "promotes", + "inhibits", + "suppresses", + "activates", + "mediates", + "controls", + "stimulates", + "blocks", + "enhances", + "reduces", + "increases", + "decreases", + "triggers", + "phosphorylates", + "upregulates", + "downregulates", + "attenuates", + "potentiates", + "modulates", + "drives", ] def extract_mechanistic_sentences( - text: str, - gene1: str, - gene2: str, - max_sentences: int = 2 + text: str, gene1: str, gene2: str, max_sentences: int = 2 ) -> List[str]: """Extract sentences mentioning both genes with mechanistic verbs. - + Args: text: Abstract or full text to search gene1: First gene symbol (e.g., regulator) gene2: Second gene symbol (e.g., target) max_sentences: Maximum sentences to return - + Returns: List of full sentences containing both genes and a mechanistic verb """ if not text: return [] - + sentences = split_text_into_sentences(text) mechanistic = [] - + gene1_lower = gene1.lower() gene2_lower = gene2.lower() - + for sent in sentences: sent_lower = sent.lower() # Check if both genes mentioned (or gene name variants) has_gene1 = gene1_lower in sent_lower has_gene2 = gene2_lower in sent_lower - + if has_gene1 and has_gene2: # Check for mechanistic verb for verb in MECHANISTIC_VERBS: if verb in sent_lower: mechanistic.append(sent.strip()) break - + return mechanistic[:max_sentences] def extract_any_mechanistic_sentences( - text: str, - gene: str, - max_sentences: int = 3 + text: str, gene: str, max_sentences: int = 3 ) -> List[str]: """Extract sentences mentioning a gene with mechanistic verbs. - + Less strict than extract_mechanistic_sentences - finds any sentence where the gene appears with a mechanistic verb, regardless of second gene. - + Args: text: Abstract or full text to search gene: Gene symbol to search for max_sentences: Maximum sentences to return - + Returns: List of full sentences containing the gene and a mechanistic verb """ if not text: return [] - + sentences = split_text_into_sentences(text) mechanistic = [] - + gene_lower = gene.lower() - + for sent in sentences: sent_lower = sent.lower() if gene_lower not in sent_lower: continue - + # Check for mechanistic verb for verb in MECHANISTIC_VERBS: if verb in sent_lower: mechanistic.append(sent.strip()) break - + return mechanistic[:max_sentences] + def split_text_into_sentences(text: str) -> List[str]: """ Robustly split abstract into sentences using Regex. @@ -330,7 +385,7 @@ def split_text_into_sentences(text: str) -> List[str]: """ if not text: return [] - + # Simplified robust splitter for biomedical text: # 1. Protect common abbreviations by replacing temporarily protected = text @@ -340,14 +395,14 @@ def split_text_into_sentences(text: str) -> List[str]: "i.e.": "IE_MARKER", "Fig.": "FIG_MARKER", "Ref.": "REF_MARKER", - "vs.": "VS_MARKER" + "vs.": "VS_MARKER", } for k, v in subs.items(): protected = protected.replace(k, v) - + # 2. Split by [.?!] followed by space or end of string - parts = re.split(r'(?<=[.?!])\s+', protected) - + parts = re.split(r"(?<=[.?!])\s+", protected) + # 3. Restore and clean sentences = [] for p in parts: @@ -355,16 +410,14 @@ def split_text_into_sentences(text: str) -> List[str]: for k, v in subs.items(): restored = restored.replace(v, k) s = restored.strip() - if len(s) > 10: # Filter noise + if len(s) > 10: # Filter noise sentences.append(s) - + return sentences + def extract_evidence_sentences( - abstract: str, - title: str, - target_genes: Set[str], - context_genes: Set[str] + abstract: str, title: str, target_genes: Set[str], context_genes: Set[str] ) -> Dict[str, List[str]]: """ Extracts high-quality sentences mentioning target genes. @@ -372,47 +425,53 @@ def extract_evidence_sentences( """ full_text = f"{title}. {abstract}" sentences = split_text_into_sentences(full_text) - + gene_to_sentences = {g: [] for g in target_genes} - + for sent in sentences: sent_lower = sent.lower() - + # Check for gene mentions found_targets = [] for g in target_genes: if re.search(rf"\b{re.escape(g)}\b", sent, re.IGNORECASE): found_targets.append(g) - - + # Deduplication check (robust) normalized_sent = sent.strip(" .") - existing_sents = [s[1].strip(" .") for s in [item for sublist in gene_to_sentences.values() for item in sublist]] + existing_sents = [ + s[1].strip(" .") + for s in [ + item for sublist in gene_to_sentences.values() for item in sublist + ] + ] if normalized_sent in existing_sents: # logger.info(f"DEBUG: Skipping duplicate: {normalized_sent[:20]}...") continue # Score Sentence score = 0 - + # +2 for Co-occurrence with OTHER context genes (members) for m in context_genes: - if m not in found_targets and re.search(rf"\b{re.escape(m)}\b", sent, re.IGNORECASE): + if m not in found_targets and re.search( + rf"\b{re.escape(m)}\b", sent, re.IGNORECASE + ): score += 2 - + # +1 for Domain Keywords if any(k in sent_lower for k in DOMAIN_KEYWORDS): score += 1 - + # +1 for Interaction Verbs if any(v in sent_lower for v in INTERACTION_VERBS): score += 1 - + # Assign sentence to all found targets for g in found_targets: if sent not in [s[1] for s in gene_to_sentences[g]]: gene_to_sentences[g].append((score, sent)) - + # Sort and filter top K per gene final_map = {} for g, items in gene_to_sentences.items(): @@ -421,22 +480,24 @@ def extract_evidence_sentences( items.sort(key=lambda x: x[0], reverse=True) # Keep top 2 unique sentences final_map[g] = [x[1] for x in items[:2]] - + return final_map + CACHE_DIR = Path("data/cache") CACHE_FILE = CACHE_DIR / "ncbi_bioc_cache.json" + def load_program_genes( - csv_path: Path, - top_n_loading: int = 20, + csv_path: Path, + top_n_loading: int = 20, top_n_unique: int = 10, - top_n_member: int = 100 + top_n_member: int = 100, ) -> Dict[int, Dict[str, List[str]]]: """ Load driver and member genes per program. Returns: {program_id: {"drivers": [...], "top_loading": [...], "top_unique": [...], "members": [...], "all_genes": [...]}} - + Drivers: Top N Loading + Top M Unique (non-overlapping) = exactly N + M genes Members (Verification): Top 100 Loading """ @@ -448,36 +509,42 @@ def load_program_genes( programs = {} for pid, group in df.groupby("program_id"): # 1. Top Loading (by Score) - sorted_loading = group.sort_values("Score", ascending=False)["Name"].astype(str).tolist() + sorted_loading = ( + group.sort_values("Score", ascending=False)["Name"].astype(str).tolist() + ) top_loading_genes = sorted_loading[:top_n_loading] top_loading_set = set(top_loading_genes) - + # 2. Top Unique (by UniquenessScore) - NON-OVERLAPPING with top loading top_unique_genes = [] if "UniquenessScore" in group.columns: - sorted_unique = group.sort_values("UniquenessScore", ascending=False)["Name"].astype(str).tolist() + sorted_unique = ( + group.sort_values("UniquenessScore", ascending=False)["Name"] + .astype(str) + .tolist() + ) # Filter out genes already in top_loading, then take top N for gene in sorted_unique: if gene not in top_loading_set: top_unique_genes.append(gene) if len(top_unique_genes) >= top_n_unique: break - + # Drivers: exactly top_n_loading + top_n_unique genes (no overlap) drivers = top_loading_genes + top_unique_genes - + # Members: Top 100 Loading members = sorted_loading[:top_n_member] - + # All genes: Full list for STRING validation (typically 300) all_genes = sorted_loading - + programs[int(pid)] = { "drivers": drivers, "top_loading": top_loading_genes, "top_unique": top_unique_genes, "members": members, - "all_genes": all_genes # For STRING regulator validation + "all_genes": all_genes, # For STRING regulator validation } return programs @@ -548,9 +615,7 @@ def resolve_gene_summaries( summary_type, len(all_drivers), ) - symbol_to_summary = harmonizome_client.get_gene_summaries( - sorted(all_drivers) - ) + symbol_to_summary = harmonizome_client.get_gene_summaries(sorted(all_drivers)) for pid in program_ids: drivers = programs[pid]["drivers"] summaries = { @@ -572,10 +637,12 @@ def load_cache() -> Dict[str, Any]: return {} return {} + def save_cache(cache: Dict[str, Any]): CACHE_DIR.mkdir(parents=True, exist_ok=True) CACHE_FILE.write_text(json.dumps(cache), encoding="utf-8") + def parse_bioc_abstract(doc: Dict[str, Any]) -> str: """Extract abstract text from BioC document.""" text_parts = [] @@ -583,10 +650,14 @@ def parse_bioc_abstract(doc: Dict[str, Any]) -> str: for p in passages: inf = p.get("infons", {}) # Abstract or Title - if inf.get("type") in ("title", "abstract") or inf.get("section_type") in ("TITLE", "ABSTRACT"): - text_parts.append(p.get("text", "")) + if inf.get("type") in ("title", "abstract") or inf.get("section_type") in ( + "TITLE", + "ABSTRACT", + ): + text_parts.append(p.get("text", "")) return " ".join(text_parts) + def parse_bioc_relations(doc: Dict[str, Any]) -> List[Dict[str, str]]: """Extract relation annotations if available (Pubtator 3 exports them).""" # Pubtator 3 JSON usually has "annotations" and "relations" keys in passages or doc level @@ -596,7 +667,7 @@ def parse_bioc_relations(doc: Dict[str, Any]) -> List[Dict[str, str]]: # Also check passages level for p in doc.get("passages", []): rels.extend(p.get("relations", [])) - + extracted = [] for r in rels: # Pubtator relations often: type="Chemical-Gene", etc. @@ -607,9 +678,10 @@ def parse_bioc_relations(doc: Dict[str, Any]) -> List[Dict[str, str]]: inf = r.get("infons", {}) rel_type = inf.get("type") if rel_type in ("Gene-Disease", "Chemical-Gene", "Gene-Chemical"): - extracted.append({"type": rel_type, "id": r.get("id")}) + extracted.append({"type": rel_type, "id": r.get("id")}) return extracted + def find_gene_mentions(doc: Dict[str, Any], member_genes: List[str]) -> List[str]: """ Check which member genes are mentioned in the doc. @@ -623,8 +695,8 @@ def find_gene_mentions(doc: Dict[str, Any], member_genes: List[str]) -> List[str """ found = set() # Collect all gene annotations from doc - doc_genes = [] - + doc_genes = [] + # 1. Collect annotated genes from Pubtator for p in doc.get("passages", []): for ann in p.get("annotations", []): @@ -632,10 +704,11 @@ def find_gene_mentions(doc: Dict[str, Any], member_genes: List[str]) -> List[str if inf.get("type") == "Gene": # Text mention txt = ann.get("text", "") - if txt: doc_genes.append(txt) + if txt: + doc_genes.append(txt) # Normalized ID - # ident = inf.get("identifier") - + # ident = inf.get("identifier") + # 2. Case-insensitive match against member_genes # Simple set intersection members_lower = {g.lower() for g in member_genes} @@ -646,7 +719,7 @@ def find_gene_mentions(doc: Dict[str, Any], member_genes: List[str]) -> List[str # Optimization: Pre-map lower -> original pass - # Fallback to Text Search? + # Fallback to Text Search? # If we requested BioC-XML/JSON from Pubtator, it SHOULD have annotations. # Text search is simpler for now given complex normalization. full_text = parse_bioc_abstract(doc).lower() @@ -655,12 +728,12 @@ def find_gene_mentions(doc: Dict[str, Any], member_genes: List[str]) -> List[str # e.g " " + g + " " -> simplistic # Let's rely on string containment for MVP, but acknowledge risk. # Better: if g in doc_genes list. - + # Let's assume Pubtator annotations are the gold standard. # We check if any annotation matches our gene symbol (case-insensitive) if any(g.lower() == dg.lower() for dg in doc_genes): found.add(g) - + return list(found) @@ -679,27 +752,27 @@ def load_regulator_data( significance_threshold: float = 0.05, ) -> Dict[int, pd.DataFrame]: """Load significant regulators from SCEPTRE results CSV. - + Returns dict mapping program_id -> DataFrame with columns: [grna_target, log_2_fold_change, p_value] """ if not csv_path or not csv_path.exists(): logger.warning(f"Regulator file not found: {csv_path}") return {} - + df = pd.read_csv(csv_path) df = standardize_regulator_results( df, significance_threshold=significance_threshold ) df = df[df["significant"] == True].copy() # type: ignore - + result = {} for pid, group in df.groupby("program_id"): keep_cols = ["grna_target", "log_2_fold_change", "p_value", "significant"] if "adj_p_value" in group.columns: keep_cols.append("adj_p_value") result[int(pid)] = group[keep_cols].copy() - + logger.info(f"Loaded regulators for {len(result)} programs") return result @@ -711,10 +784,10 @@ def get_top_regulators( top_n_positive: Optional[int] = None, top_n_negative: Optional[int] = None, use_all_significant: bool = False, - max_regulators: int = 20 + max_regulators: int = 20, ) -> Dict[str, List[Dict[str, Any]]]: """Get top positive and negative regulators for a program. - + Args: regulator_data: Dict mapping program_id to DataFrame program_id: Program to get regulators for @@ -723,7 +796,7 @@ def get_top_regulators( top_n_negative: Number of negative regulators to return use_all_significant: If True, return all significant regulators (up to max_regulators) max_regulators: Maximum regulators per category when using all significant - + Returns: { 'positive': [{'gene': 'Fzd4', 'log2fc': -0.62, 'pvalue': 1e-10}, ...], # negative log2FC = activator 'negative': [{'gene': 'Eng', 'log2fc': 0.49, 'pvalue': 1e-8}, ...] # positive log2FC = repressor @@ -731,34 +804,39 @@ def get_top_regulators( """ reg_df = regulator_data.get(program_id) if reg_df is None or len(reg_df) == 0: - return {'positive': [], 'negative': []} - + return {"positive": [], "negative": []} + top_n_positive = top_n if top_n_positive is None else top_n_positive top_n_negative = top_n if top_n_negative is None else top_n_negative - sig_df = reg_df[reg_df['significant'] == True].copy() - + sig_df = reg_df[reg_df["significant"] == True].copy() + sorted_df = sig_df.sort_values(by="log_2_fold_change") - + # Negative log2FC = positive regulators (activators) if use_all_significant: positive = sorted_df[sorted_df["log_2_fold_change"] < 0].head(max_regulators) - negative = sorted_df[sorted_df["log_2_fold_change"] > 0].tail(max_regulators).iloc[::-1] + negative = ( + sorted_df[sorted_df["log_2_fold_change"] > 0] + .tail(max_regulators) + .iloc[::-1] + ) else: positive = sorted_df[sorted_df["log_2_fold_change"] < 0].head(top_n_positive) - negative = sorted_df[sorted_df["log_2_fold_change"] > 0].tail(top_n_negative).iloc[::-1] - + negative = ( + sorted_df[sorted_df["log_2_fold_change"] > 0] + .tail(top_n_negative) + .iloc[::-1] + ) + def extract_regulator(row): - result = { - 'gene': row['grna_target'], - 'log2fc': row['log_2_fold_change'] - } - if 'p_value' in row: - result['pvalue'] = row['p_value'] + result = {"gene": row["grna_target"], "log2fc": row["log_2_fold_change"]} + if "p_value" in row: + result["pvalue"] = row["p_value"] return result - + return { - 'positive': [extract_regulator(row) for _, row in positive.iterrows()], - 'negative': [extract_regulator(row) for _, row in negative.iterrows()] + "positive": [extract_regulator(row) for _, row in positive.iterrows()], + "negative": [extract_regulator(row) for _, row in negative.iterrows()], } @@ -766,20 +844,20 @@ def validate_regulator_with_string( regulator: str, program_genes: List[str], species: int = 10090, # mouse - required_score: int = 400 + required_score: int = 400, ) -> Dict[str, Any]: """Validate regulator-program relationship via STRING-DB. - + Queries STRING for direct protein-protein interactions between the regulator and program genes. Returns only interactions with genes that are actually in the program. - + Args: regulator: Gene symbol of the regulator program_genes: List of program's gene symbols species: NCBI taxon ID (10090 for mouse) required_score: Minimum STRING combined score (0-1000) - + Returns: { 'regulator': 'Fzd4', 'n_program_targets': 5, @@ -790,51 +868,55 @@ def validate_regulator_with_string( } """ import time + time.sleep(1.0) # Rate limiting - + result = get_regulator_program_interactions( regulator=regulator, program_genes=program_genes, species=species, required_score=required_score, - top_n=10 + top_n=10, ) - + return { - 'regulator': regulator, - 'n_program_targets': result['n_interactions'], - 'string_interactions': result['interactions'] + "regulator": regulator, + "n_program_targets": result["n_interactions"], + "string_interactions": result["interactions"], } def search_pubtator(query: str, max_results: int = 50) -> List[Dict[str, Any]]: """Search PubTator3 for papers matching query. - + Returns list of {pmid, title, score, text_hl} dicts. - + Note: PubTator3 uses 'size' parameter for page size (max ~100 per request). """ url = f"{PUBTATOR_API_BASE}/search/" # Use 'size' parameter for page size (not 'page_size' or 'limit') params = {"text": query, "page": 1, "size": min(max_results, 100)} - + try: # Increased timeout to 120s for complex gene queries resp = requests.get(url, params=params, timeout=120) if resp.status_code != 200: logger.warning(f"PubTator search failed: {resp.status_code}") return [] - + data = resp.json() results = data.get("results", [])[:max_results] - - return [{ - 'pmid': r.get('pmid'), - 'title': r.get('title', ''), - 'score': r.get('score', 0), - 'text_hl': r.get('text_hl', '') - } for r in results] - + + return [ + { + "pmid": r.get("pmid"), + "title": r.get("title", ""), + "score": r.get("score", 0), + "text_hl": r.get("text_hl", ""), + } + for r in results + ] + except Exception as e: logger.error(f"PubTator search error: {e}") return [] @@ -842,7 +924,7 @@ def search_pubtator(query: str, max_results: int = 50) -> List[Dict[str, Any]]: def fetch_bioc_relations_with_text(pmids: List[int]) -> Dict[int, Dict[str, Any]]: """Fetch BioC-JSON and extract gene-gene relations plus abstract text. - + Returns: {pmid: { 'relations': [{'gene1': 'FZD4', 'gene2': 'CTNNB1', 'type': 'Positive_Correlation', 'score': 0.99}, ...], 'title': str, @@ -852,18 +934,18 @@ def fetch_bioc_relations_with_text(pmids: List[int]) -> Dict[int, Dict[str, Any] """ if not pmids: return {} - + url = f"{PUBTATOR_API_BASE}/publications/export/biocjson" payload = {"pmids": pmids} - + try: resp = requests.post(url, json=payload, timeout=60) if resp.status_code != 200: logger.warning(f"PubTator BioC fetch failed: {resp.status_code}") return {} - + data = resp.json() - + # Handle different response formats if isinstance(data, dict) and "PubTator3" in data: docs = data["PubTator3"] @@ -871,66 +953,70 @@ def fetch_bioc_relations_with_text(pmids: List[int]) -> Dict[int, Dict[str, Any] docs = data else: docs = [data] - + result = {} for doc in docs: - pmid = doc.get('pmid') or doc.get('id') + pmid = doc.get("pmid") or doc.get("id") if not pmid: continue - + # Extract title and abstract from passages, plus gene annotations - title = '' - abstract = '' - gene_mentions: Dict[str, Set[str]] = {} # normalized name -> set of text mentions - - for passage in doc.get('passages', []): - ptype = passage.get('infons', {}).get('type', '') - if ptype == 'title': - title = passage.get('text', '') - elif ptype == 'abstract': - abstract = passage.get('text', '') - + title = "" + abstract = "" + gene_mentions: Dict[ + str, Set[str] + ] = {} # normalized name -> set of text mentions + + for passage in doc.get("passages", []): + ptype = passage.get("infons", {}).get("type", "") + if ptype == "title": + title = passage.get("text", "") + elif ptype == "abstract": + abstract = passage.get("text", "") + # Extract gene annotations to build alias mapping - for ann in passage.get('annotations', []): - if ann.get('infons', {}).get('type') == 'Gene': - gene_name = ann.get('infons', {}).get('name', '') - text_mention = ann.get('text', '') + for ann in passage.get("annotations", []): + if ann.get("infons", {}).get("type") == "Gene": + gene_name = ann.get("infons", {}).get("name", "") + text_mention = ann.get("text", "") if gene_name and text_mention: if gene_name.upper() not in gene_mentions: gene_mentions[gene_name.upper()] = set() gene_mentions[gene_name.upper()].add(text_mention) - + # Extract relations relations = [] - for rel in doc.get('relations', []): - infons = rel.get('infons', {}) - r1 = infons.get('role1', {}) - r2 = infons.get('role2', {}) - rel_type = infons.get('type', 'Unknown') - score = float(infons.get('score', 0)) - + for rel in doc.get("relations", []): + infons = rel.get("infons", {}) + r1 = infons.get("role1", {}) + r2 = infons.get("role2", {}) + rel_type = infons.get("type", "Unknown") + score = float(infons.get("score", 0)) + # Only gene-gene relations with mechanistic types - if r1.get('type') == 'Gene' and r2.get('type') == 'Gene': - if rel_type in ('Positive_Correlation', 'Negative_Correlation'): - relations.append({ - 'gene1': r1.get('name', ''), - 'gene2': r2.get('name', ''), - 'type': rel_type, - 'score': score - }) - + if r1.get("type") == "Gene" and r2.get("type") == "Gene": + if rel_type in ("Positive_Correlation", "Negative_Correlation"): + relations.append( + { + "gene1": r1.get("name", ""), + "gene2": r2.get("name", ""), + "type": rel_type, + "score": score, + } + ) + # Convert sets to lists for JSON serialization gene_mentions_list = {k: list(v) for k, v in gene_mentions.items()} - + result[int(pmid)] = { - 'relations': relations, - 'title': title, - 'abstract': abstract, - 'gene_mentions': gene_mentions_list + "relations": relations, + "title": title, + "abstract": abstract, + "gene_mentions": gene_mentions_list, } - + return result - + except Exception as e: logger.error(f"PubTator BioC fetch error: {e}") return {} @@ -941,20 +1027,20 @@ def validate_regulator_program( program_genes: List[str], keyword: str = "", max_pmids: int = 50, - min_relation_score: float = 0.5 + min_relation_score: float = 0.5, ) -> Dict[str, Any]: """Validate regulator-program relationship via PubTator3. - + Extracts mechanistic gene-gene relations (Positive/Negative_Correlation only) and the full sentences describing those relationships from abstracts. - + Args: regulator: Gene symbol of the regulator program_genes: List of program's top genes keyword: Optional keyword terms (tissue/cell type) to restrict the search; empty = no filter max_pmids: Number of PMIDs to fetch for relation extraction (default 50) min_relation_score: Minimum score for relations - + Returns: { 'regulator': 'Fzd4', 'papers_found': 20, @@ -977,143 +1063,150 @@ def validate_regulator_program( # Optionally restricted to a tissue/cell-type keyword genes_or = " OR ".join(program_genes[:10]) # Use top 10 program genes query = f"({regulator}) AND ({genes_or})" + (f" AND ({keyword})" if keyword else "") - + logger.info(f" Validating {regulator}: {query[:80]}...") - + # Search time.sleep(PUBTATOR_RATE_LIMIT) search_results = search_pubtator(query, max_results=60) - + if not search_results: return { - 'regulator': regulator, - 'papers_found': 0, - 'papers_with_relations': 0, - 'mechanistic_relations': [], - 'top_papers': [] + "regulator": regulator, + "papers_found": 0, + "papers_with_relations": 0, + "mechanistic_relations": [], + "top_papers": [], } - + # Fetch BioC for top PMIDs (with abstract text) - pmids = [r['pmid'] for r in search_results[:max_pmids] if r['pmid']] - + pmids = [r["pmid"] for r in search_results[:max_pmids] if r["pmid"]] + time.sleep(PUBTATOR_RATE_LIMIT) pmid_data = fetch_bioc_relations_with_text(pmids) - + # Extract mechanistic relations involving the regulator # Only Positive_Correlation and Negative_Correlation (skip Association) # Only keep relations where the target gene is in the program's gene list mechanistic_relations = [] regulator_upper = regulator.upper() - program_genes_upper = {g.upper() for g in program_genes} # For filtering target genes + program_genes_upper = { + g.upper() for g in program_genes + } # For filtering target genes seen_gene_pmid = set() # Deduplicate by (gene, pmid) pair pmids_with_relations = set() - + for pmid, data in pmid_data.items(): - rels = data.get('relations', []) - abstract = data.get('abstract', '') - title = data.get('title', '') - + rels = data.get("relations", []) + abstract = data.get("abstract", "") + title = data.get("title", "") + for rel in rels: - if rel['score'] < min_relation_score: + if rel["score"] < min_relation_score: continue - - g1 = rel['gene1'].upper() - g2 = rel['gene2'].upper() - rel_type = rel['type'] - + + g1 = rel["gene1"].upper() + g2 = rel["gene2"].upper() + rel_type = rel["type"] + # Only mechanistic relation types - if rel_type not in ('Positive_Correlation', 'Negative_Correlation'): + if rel_type not in ("Positive_Correlation", "Negative_Correlation"): continue - + # Check if regulator is involved if regulator_upper not in (g1, g2): continue - + pmids_with_relations.add(pmid) - other_gene = rel['gene2'] if g1 == regulator_upper else rel['gene1'] + other_gene = rel["gene2"] if g1 == regulator_upper else rel["gene1"] dedup_key = (other_gene.upper(), pmid) - + if dedup_key in seen_gene_pmid: continue seen_gene_pmid.add(dedup_key) - + # Extract mechanistic sentences mentioning both genes sentences = extract_mechanistic_sentences( abstract, regulator, other_gene, max_sentences=2 ) - + # If no sentence with both genes, try finding any mechanistic sentence about regulator if not sentences: sentences = extract_any_mechanistic_sentences( abstract, regulator, max_sentences=1 ) - - mechanistic_relations.append({ - 'pmid': pmid, - 'target_gene': other_gene, - 'relation_type': rel_type, - 'score': rel['score'], - 'title': title, - 'sentences': sentences - }) - + + mechanistic_relations.append( + { + "pmid": pmid, + "target_gene": other_gene, + "relation_type": rel_type, + "score": rel["score"], + "title": title, + "sentences": sentences, + } + ) + # Fallback: If no formal relations found, still extract mechanistic sentences from abstracts # This catches cases where the regulator is mentioned but not in PubTator relations if not mechanistic_relations: for pmid, data in pmid_data.items(): - abstract = data.get('abstract', '') - title = data.get('title', '') - gene_mentions = data.get('gene_mentions', {}) - + abstract = data.get("abstract", "") + title = data.get("title", "") + gene_mentions = data.get("gene_mentions", {}) + if not abstract: continue - + # Get all text aliases for this regulator from PubTator annotations regulator_aliases = gene_mentions.get(regulator.upper(), [regulator]) if not regulator_aliases: regulator_aliases = [regulator] - + # Try to find any mechanistic sentence using the regulator or its aliases sentences = [] for alias in regulator_aliases: - sentences = extract_any_mechanistic_sentences(abstract, alias, max_sentences=2) + sentences = extract_any_mechanistic_sentences( + abstract, alias, max_sentences=2 + ) if sentences: break - + # If no alias worked, try the original name if not sentences: - sentences = extract_any_mechanistic_sentences(abstract, regulator, max_sentences=2) - + sentences = extract_any_mechanistic_sentences( + abstract, regulator, max_sentences=2 + ) + if sentences: - mechanistic_relations.append({ - 'pmid': pmid, - 'target_gene': 'unknown', - 'relation_type': 'Mentioned', - 'score': 0.5, # Lower confidence score - 'title': title, - 'sentences': sentences - }) - + mechanistic_relations.append( + { + "pmid": pmid, + "target_gene": "unknown", + "relation_type": "Mentioned", + "score": 0.5, # Lower confidence score + "title": title, + "sentences": sentences, + } + ) + # Sort by score descending, keep top 10 - mechanistic_relations.sort(key=lambda x: -x['score']) + mechanistic_relations.sort(key=lambda x: -x["score"]) mechanistic_relations = mechanistic_relations[:10] - + # Build top papers list (for context even without relations) top_papers = [] for r in search_results[:10]: - top_papers.append({ - 'pmid': r['pmid'], - 'title': r['title'] - }) - - papers_with_rels = len([p for p, d in pmid_data.items() if d.get('relations')]) - + top_papers.append({"pmid": r["pmid"], "title": r["title"]}) + + papers_with_rels = len([p for p, d in pmid_data.items() if d.get("relations")]) + return { - 'regulator': regulator, - 'papers_found': len(search_results), - 'papers_with_relations': papers_with_rels, - 'mechanistic_relations': mechanistic_relations, - 'top_papers': top_papers + "regulator": regulator, + "papers_found": len(search_results), + "papers_with_relations": papers_with_rels, + "mechanistic_relations": mechanistic_relations, + "top_papers": top_papers, } @@ -1130,113 +1223,116 @@ def validate_program_regulators( use_all_significant: bool = False, max_regulators: int = 20, use_batch: bool = True, - min_score: int = 400 + min_score: int = 400, ) -> Dict[str, Any]: """Validate all top regulators for a program. - + Uses STRING-DB for interaction validation (default) or PubTator3 for literature. With use_batch=True, makes ONE STRING query per program (much faster). - + Args: use_all_significant: If True, validate all significant regulators (not just top N) max_regulators: Maximum regulators per category when using all significant use_batch: If True, use single batch STRING query (faster) min_score: Minimum STRING combined score (0-1000, 400=medium confidence) - + Returns: { 'positive_regulators': [validation_result, ...], 'negative_regulators': [validation_result, ...] } """ top_regs = get_top_regulators( - regulator_data, - program_id, + regulator_data, + program_id, top_n=top_n_regulators, top_n_positive=top_n_positive_regulators, top_n_negative=top_n_negative_regulators, use_all_significant=use_all_significant, - max_regulators=max_regulators + max_regulators=max_regulators, ) - - result = { - 'positive_regulators': [], - 'negative_regulators': [] - } - - all_regulator_genes = [r['gene'] for r in top_regs['positive']] + [r['gene'] for r in top_regs['negative']] - + + result = {"positive_regulators": [], "negative_regulators": []} + + all_regulator_genes = [r["gene"] for r in top_regs["positive"]] + [ + r["gene"] for r in top_regs["negative"] + ] + # Use batch STRING query (ONE call for all regulators) if use_string and use_batch and all_regulator_genes: batch_results = batch_validate_regulators( regulator_genes=all_regulator_genes, program_genes=program_genes, - required_score=min_score + required_score=min_score, ) - + # Build results for positive regulators (activators) - for reg_info in top_regs['positive']: - gene = reg_info['gene'] + for reg_info in top_regs["positive"]: + gene = reg_info["gene"] interactions = batch_results.get(gene, []) - result['positive_regulators'].append({ - 'regulator': gene, - 'log2fc': reg_info['log2fc'], - 'n_program_targets': len(interactions), - 'string_interactions': interactions # List of {target, score} - }) - + result["positive_regulators"].append( + { + "regulator": gene, + "log2fc": reg_info["log2fc"], + "n_program_targets": len(interactions), + "string_interactions": interactions, # List of {target, score} + } + ) + # Build results for negative regulators (repressors) - for reg_info in top_regs['negative']: - gene = reg_info['gene'] + for reg_info in top_regs["negative"]: + gene = reg_info["gene"] interactions = batch_results.get(gene, []) - result['negative_regulators'].append({ - 'regulator': gene, - 'log2fc': reg_info['log2fc'], - 'n_program_targets': len(interactions), - 'string_interactions': interactions - }) - + result["negative_regulators"].append( + { + "regulator": gene, + "log2fc": reg_info["log2fc"], + "n_program_targets": len(interactions), + "string_interactions": interactions, + } + ) + return result - + # Fallback: individual queries (slower) - for reg_info in top_regs['positive']: + for reg_info in top_regs["positive"]: if use_string: logger.info(f" Validating activator {reg_info['gene']} with STRING...") validation = validate_regulator_with_string( - regulator=reg_info['gene'], - program_genes=program_genes + regulator=reg_info["gene"], program_genes=program_genes ) else: validation = validate_regulator_program( - regulator=reg_info['gene'], + regulator=reg_info["gene"], program_genes=program_genes, keyword=keyword, - max_pmids=max_pmids_per_regulator + max_pmids=max_pmids_per_regulator, ) - validation['log2fc'] = reg_info['log2fc'] - result['positive_regulators'].append(validation) - - for reg_info in top_regs['negative']: + validation["log2fc"] = reg_info["log2fc"] + result["positive_regulators"].append(validation) + + for reg_info in top_regs["negative"]: if use_string: logger.info(f" Validating repressor {reg_info['gene']} with STRING...") validation = validate_regulator_with_string( - regulator=reg_info['gene'], - program_genes=program_genes + regulator=reg_info["gene"], program_genes=program_genes ) else: validation = validate_regulator_program( - regulator=reg_info['gene'], + regulator=reg_info["gene"], program_genes=program_genes, keyword=keyword, - max_pmids=max_pmids_per_regulator + max_pmids=max_pmids_per_regulator, ) - validation['log2fc'] = reg_info['log2fc'] - result['negative_regulators'].append(validation) - + validation["log2fc"] = reg_info["log2fc"] + result["negative_regulators"].append(validation) + return result def main(): - parser = argparse.ArgumentParser(description="Fetch NCBI literature evidence for gene programs") + parser = argparse.ArgumentParser( + description="Fetch NCBI literature evidence for gene programs" + ) parser.add_argument("--config", help="Path to config file (YAML or JSON)") parser.add_argument( "--input", @@ -1247,8 +1343,8 @@ def main(): parser.add_argument( "--keyword", default="", - help='Optional keyword query restricting the PubMed search to your tissue/cell type, ' - 'e.g. \'("T cell" OR lymphocyte)\' (default: no keyword filter)', + help="Optional keyword query restricting the PubMed search to your tissue/cell type, " + "e.g. '(\"T cell\" OR lymphocyte)' (default: no keyword filter)", ) parser.add_argument("--api-key", help="NCBI API Key") parser.add_argument( @@ -1261,22 +1357,74 @@ def main(): "--full-summaries", action="store_true", help="Fetch full literature-based summaries from Harmonizome HTML pages (~3x longer). " - "Default is short API descriptions.", + "Default is short API descriptions.", + ) + parser.add_argument( + "--num-programs", type=int, help="Limit number of programs (for testing)" + ) + parser.add_argument( + "--topics", + type=str, + help="Comma-separated list of topic IDs to process (e.g. '6,7,8')", + ) + parser.add_argument( + "--top-loading", + type=int, + default=20, + help="Number of top loading genes (default 20)", + ) + parser.add_argument( + "--top-unique", + type=int, + default=10, + help="Number of top unique genes, non-overlapping with loading (default 10)", + ) + parser.add_argument( + "--top-member", type=int, default=100, help="Number of member genes for context" + ) + parser.add_argument( + "--regulator-file", type=str, help="CSV with regulator perturbation results" + ) + parser.add_argument( + "--top-regulators", + type=int, + default=3, + help="Backward-compatible default for positive/negative regulator counts", + ) + parser.add_argument( + "--top-positive-regulators", + type=int, + help="Number of top positive regulators to validate", + ) + parser.add_argument( + "--top-negative-regulators", + type=int, + help="Number of top negative regulators to validate", + ) + parser.add_argument( + "--regulator-significance-threshold", + type=float, + default=0.05, + help="Adjusted p-value threshold used when the regulator file has no 'significant' column (default 0.05)", + ) + parser.add_argument( + "--regulator-pmids", + type=int, + default=50, + help="Number of PMIDs to fetch per regulator (default 50)", ) - parser.add_argument("--num-programs", type=int, help="Limit number of programs (for testing)") - parser.add_argument("--topics", type=str, help="Comma-separated list of topic IDs to process (e.g. '6,7,8')") - parser.add_argument("--top-loading", type=int, default=20, help="Number of top loading genes (default 20)") - parser.add_argument("--top-unique", type=int, default=10, help="Number of top unique genes, non-overlapping with loading (default 10)") - parser.add_argument("--top-member", type=int, default=100, help="Number of member genes for context") - parser.add_argument("--regulator-file", type=str, help="CSV with regulator perturbation results") - parser.add_argument("--top-regulators", type=int, default=3, help="Backward-compatible default for positive/negative regulator counts") - parser.add_argument("--top-positive-regulators", type=int, help="Number of top positive regulators to validate") - parser.add_argument("--top-negative-regulators", type=int, help="Number of top negative regulators to validate") - parser.add_argument("--regulator-significance-threshold", type=float, default=0.05, help="Adjusted p-value threshold used when the regulator file has no 'significant' column (default 0.05)") - parser.add_argument("--regulator-pmids", type=int, default=50, help="Number of PMIDs to fetch per regulator (default 50)") - parser.add_argument("--all-significant", action="store_true", help="Validate ALL significant regulators (not just top N). Fast with STRING.") - parser.add_argument("--max-regulators", type=int, default=20, help="Maximum regulators per category when using --all-significant (default 20)") - + parser.add_argument( + "--all-significant", + action="store_true", + help="Validate ALL significant regulators (not just top N). Fast with STRING.", + ) + parser.add_argument( + "--max-regulators", + type=int, + default=20, + help="Maximum regulators per category when using --all-significant (default 20)", + ) + args = parser.parse_args() config = load_config(args.config) cli_overrides = get_cli_overrides(sys.argv) @@ -1284,21 +1432,23 @@ def main(): args = apply_test_mode(args, config, cli_overrides) if not args.input or not args.csv_out or not args.json_out: - raise SystemExit("--input, --csv-out, and --json-out are required (via CLI or config).") - + raise SystemExit( + "--input, --csv-out, and --json-out are required (via CLI or config)." + ) + # Init Client client = NcbiClient(api_key=args.api_key) - + # Load Programs logger.info(f"Loading programs from {args.input}...") programs = load_program_genes( - Path(args.input), - top_n_loading=args.top_loading, + Path(args.input), + top_n_loading=args.top_loading, top_n_unique=args.top_unique, - top_n_member=args.top_member + top_n_member=args.top_member, ) program_ids = sorted(programs.keys()) - + # Filter by specific topics if requested selected_topics = parse_topics_value(args.topics) if selected_topics: @@ -1306,36 +1456,33 @@ def main(): logger.info(f"Limiting to specific topics: {program_ids}") # Fallback to num-programs limit elif args.num_programs: - program_ids = program_ids[:args.num_programs] + program_ids = program_ids[: args.num_programs] logger.info(f"Limiting to first {args.num_programs} programs: {program_ids}") # ========================================================================= # Step 1: Search (Gather PMIDs) # ========================================================================= - + all_pmids = set() - program_pmid_map = {} # pid -> [pmids] - program_meta = {} # pid -> {query: str, total_hits: int} - + program_pmid_map = {} # pid -> [pmids] + program_meta = {} # pid -> {query: str, total_hits: int} + for pid in program_ids: drivers = programs[pid]["drivers"] if not drivers: continue - + # Construct Query # (Gene1 OR Gene2 ...) AND Keyword genes_or = " OR ".join(drivers) query = f"({genes_or})" + (f" AND {args.keyword}" if args.keyword else "") - + logger.info(f"[Program {pid}] Searching: {query}") pmids = client.search_literature(query, size=25) logger.info(f"[Program {pid}] Found {len(pmids)} articles") - + program_pmid_map[pid] = pmids - program_meta[pid] = { - "query": query, - "total_hits": len(pmids) - } + program_meta[pid] = {"query": query, "total_hits": len(pmids)} all_pmids.update(pmids) # ========================================================================= @@ -1355,22 +1502,22 @@ def main(): # ========================================================================= # Step 2: Fetch BioC Attributes (Fetch & Cache) # ========================================================================= - + logger.info(f"Fetching annotations for {len(all_pmids)} total PMIDs...") cache = load_cache() - + # Determine missing PMIDs missing_pmids = [p for p in all_pmids if str(p) not in cache] - + if missing_pmids: # Batch fetch BATCH_SIZE = 100 for i in range(0, len(missing_pmids), BATCH_SIZE): - batch = missing_pmids[i:i+BATCH_SIZE] - logger.info(f"Fetching batch {i}-{i+len(batch)}/{len(missing_pmids)}") - + batch = missing_pmids[i : i + BATCH_SIZE] + logger.info(f"Fetching batch {i}-{i + len(batch)}/{len(missing_pmids)}") + docs = client.fetch_bioc_annotations(batch) - + # Update cache with docs for doc in docs: # Find ID @@ -1378,117 +1525,125 @@ def main(): doc_id = doc.get("id") if doc_id: cache[str(doc_id)] = doc - + # Save incrementally save_cache(cache) - + logger.info(f"Cache updated. Total docs in cache: {len(cache)}") - final_context = {} # pid -> {literautre: ...} summary_rows = [] - + for pid in program_ids: pmids = program_pmid_map.get(pid, []) - members = set(programs[pid]["members"]) # All top 50 (Context) - drivers = set(programs[pid]["drivers"]) # Top 10 (Targets) - + members = set(programs[pid]["members"]) # All top 50 (Context) + drivers = set(programs[pid]["drivers"]) # Top 10 (Targets) + scored_papers = [] - + # Aggregation of Evidence Sentences # gene -> list of "sentence (PMID:123)" aggregated_snippets = {g: [] for g in drivers} - + for pmid in pmids: doc = cache.get(str(pmid)) if not doc: continue - + # 1. Parsing abstract = parse_bioc_abstract(doc) title = "" for p in doc.get("passages", []): - if p.get("infons", {}).get("type") == "title": - title = p.get("text", "") - break - if not title: title = abstract[:50] + "..." - + if p.get("infons", {}).get("type") == "title": + title = p.get("text", "") + break + if not title: + title = abstract[:50] + "..." + # 2. Identify Genes mentioned in this paper (for scoring the paper itself) mentioned_genes = find_gene_mentions(doc, list(members)) - + # 3. High-Quality Papers (Gold Standard) - Score >= 2 score = len(mentioned_genes) if score >= 2: - scored_papers.append({ - "pmid": pmid, - "title": title, - "score": score, - "mentions": mentioned_genes, - "abstract": abstract - }) - + scored_papers.append( + { + "pmid": pmid, + "title": title, + "score": score, + "mentions": mentioned_genes, + "abstract": abstract, + } + ) + # 4. Sentence-Level Evidence Extraction - # We extract sentences for ALL papers, even low scoring ones, + # We extract sentences for ALL papers, even low scoring ones, # if they contain a Driver Gene + useful context. - + evidence_map = extract_evidence_sentences(abstract, title, drivers, members) - + for g, sentences in evidence_map.items(): for s in sentences: snippet = f'"{s}" (PMID:{pmid})' aggregated_snippets[g].append(snippet) - + # Sort Papers by Score scored_papers.sort(key=lambda x: x["score"], reverse=True) top_papers = scored_papers[:5] # Keep top 5 - + # Build Context JSON Data paper_contexts = [] for p in top_papers: - paper_contexts.append({ - "pmid": p["pmid"], - "title": p["title"], - "score": p["score"], - "genes": p["mentions"] - }) - + paper_contexts.append( + { + "pmid": p["pmid"], + "title": p["title"], + "score": p["score"], + "genes": p["mentions"], + } + ) + # Add to Summary CSV - summary_rows.append({ - "program_id": pid, - "pmid": p["pmid"], - "score": p["score"], - "title": p["title"], - "genes_found": "|".join(p["mentions"]) - }) - + summary_rows.append( + { + "program_id": pid, + "pmid": p["pmid"], + "score": p["score"], + "title": p["title"], + "genes_found": "|".join(p["mentions"]), + } + ) + final_context[pid] = { "meta": program_meta.get(pid, {}), "top_papers": paper_contexts, "gene_summaries": program_gene_summaries.get(pid, {}), "gene_summaries_source": gene_summary_source, - "evidence_snippets": aggregated_snippets + "evidence_snippets": aggregated_snippets, } # ========================================================================= # Step 4: Regulator Validation (if regulator file provided) # ========================================================================= - + if args.regulator_file: logger.info("=" * 60) logger.info("Step 4: Validating regulator-program relationships via PubTator3") logger.info("=" * 60) - + regulator_data = load_regulator_data( Path(args.regulator_file), significance_threshold=args.regulator_significance_threshold, ) - + for pid in program_ids: # Use all 300 genes for STRING validation (not just drivers) program_genes = programs[pid].get("all_genes", programs[pid]["drivers"]) - + if args.all_significant: - logger.info(f"[Program {pid}] Validating ALL significant regulators against {len(program_genes)} program genes...") + logger.info( + f"[Program {pid}] Validating ALL significant regulators against {len(program_genes)} program genes..." + ) else: top_pos = args.top_positive_regulators or args.top_regulators top_neg = args.top_negative_regulators or args.top_regulators @@ -1496,7 +1651,7 @@ def main(): f"[Program {pid}] Validating top {top_pos} positive + " f"{top_neg} negative regulators against {len(program_genes)} program genes..." ) - + validation_result = validate_program_regulators( program_id=pid, regulator_data=regulator_data, @@ -1507,32 +1662,38 @@ def main(): top_n_negative_regulators=args.top_negative_regulators, max_pmids_per_regulator=args.regulator_pmids, use_all_significant=args.all_significant, - max_regulators=args.max_regulators + max_regulators=args.max_regulators, ) - + # Add to final context final_context[pid]["regulator_validation"] = validation_result - + # Log summary - pos_count = len(validation_result['positive_regulators']) - neg_count = len(validation_result['negative_regulators']) + pos_count = len(validation_result["positive_regulators"]) + neg_count = len(validation_result["negative_regulators"]) # Handle both STRING (n_program_targets) and PubTator (papers_found) formats - total_targets = sum(r.get('n_program_targets', r.get('papers_found', 0)) or 0 - for r in validation_result['positive_regulators']) - total_targets += sum(r.get('n_program_targets', r.get('papers_found', 0)) or 0 - for r in validation_result['negative_regulators']) - logger.info(f"[Program {pid}] Validated {pos_count} positive + {neg_count} negative regulators, {total_targets} total STRING interactions") + total_targets = sum( + r.get("n_program_targets", r.get("papers_found", 0)) or 0 + for r in validation_result["positive_regulators"] + ) + total_targets += sum( + r.get("n_program_targets", r.get("papers_found", 0)) or 0 + for r in validation_result["negative_regulators"] + ) + logger.info( + f"[Program {pid}] Validated {pos_count} positive + {neg_count} negative regulators, {total_targets} total STRING interactions" + ) # ========================================================================= # Output # ========================================================================= - + # JSON Path(args.json_out).parent.mkdir(parents=True, exist_ok=True) with open(args.json_out, "w") as f: json.dump(final_context, f, indent=2) logger.info(f"Written JSON context to {args.json_out}") - + # CSV if summary_rows: df_sum = pd.DataFrame(summary_rows) @@ -1544,5 +1705,6 @@ def main(): logger.warning("No high-quality papers found. Summary CSV is empty.") Path(args.csv_out).touch() + if __name__ == "__main__": main() diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/03_submit_and_monitor_batch.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/03_submit_and_monitor_batch.py index 60c467b..f746553 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/03_submit_and_monitor_batch.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/03_submit_and_monitor_batch.py @@ -46,6 +46,7 @@ import pandas as pd import sys as _sys from pathlib import Path as _Path + _sys.path.insert(0, str(_Path(__file__).resolve().parent)) from column_mapper import standardize_regulator_results @@ -53,6 +54,7 @@ # Load environment variables from .env file try: from dotenv import load_dotenv + load_dotenv() except ImportError: pass # dotenv not required if env vars are already set @@ -61,6 +63,7 @@ try: from google import genai # type: ignore from google.genai.types import CreateBatchJobConfig, JobState, HttpOptions # type: ignore + VERTEX_AVAILABLE = True except ImportError: genai = None # type: ignore @@ -72,19 +75,24 @@ # Anthropic imports (for direct Anthropic Batch API - default) try: import anthropic + ANTHROPIC_AVAILABLE = True except ImportError: anthropic = None # type: ignore ANTHROPIC_AVAILABLE = False -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) logger = logging.getLogger(__name__) MODEL = "claude-sonnet-4-5-20250929" # Anthropic model name DEFAULT_ANNOTATION_ROLE = "molecular and cell biologist" DEFAULT_ANNOTATION_CONTEXT = "a gene program extracted from single-cell Perturb-seq" -DEFAULT_SEARCH_KEYWORD = "(none specified)" # set --search-keyword to your cell type / tissue +DEFAULT_SEARCH_KEYWORD = ( + "(none specified)" # set --search-keyword to your cell type / tissue +) """ @description @@ -149,18 +157,32 @@ def parse_topics_value(value: Optional[object]) -> Optional[List[int]]: def apply_test_mode( args: argparse.Namespace, config: Dict[str, Any], cli_overrides: Set[str] ) -> argparse.Namespace: - test_cfg = config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + test_cfg = ( + config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + ) enabled = bool(test_cfg.get("enabled") or config.get("test_mode")) if not enabled: return args - topics = test_cfg.get("topics") or test_cfg.get("programs") or config.get("test_programs") - if hasattr(args, "topics") and "topics" not in cli_overrides and not getattr(args, "topics", None): + topics = ( + test_cfg.get("topics") + or test_cfg.get("programs") + or config.get("test_programs") + ) + if ( + hasattr(args, "topics") + and "topics" not in cli_overrides + and not getattr(args, "topics", None) + ): if topics is not None: args.topics = topics num_topics = test_cfg.get("num_topics") or test_cfg.get("num_programs") - if hasattr(args, "num_topics") and "num_topics" not in cli_overrides and not getattr(args, "num_topics", None): + if ( + hasattr(args, "num_topics") + and "num_topics" not in cli_overrides + and not getattr(args, "num_topics", None) + ): if num_topics is not None: args.num_topics = num_topics return args @@ -169,7 +191,9 @@ def apply_test_mode( def apply_config_overrides( args: argparse.Namespace, config: Dict[str, Any], cli_overrides: Set[str] ) -> argparse.Namespace: - steps_cfg = config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + steps_cfg = ( + config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + ) section_map = { "prepare": "batch_prepare", "submit-vertex": "vertex_submit", @@ -188,6 +212,7 @@ def apply_config_overrides( setattr(args, dest, value) return args + # Global uniqueness score helpers """ @description @@ -247,18 +272,21 @@ def add_global_uniqueness_scores(df: pd.DataFrame) -> pd.DataFrame: return updated -def ensure_global_uniqueness( - df: pd.DataFrame, logger: logging.Logger -) -> pd.DataFrame: +def ensure_global_uniqueness(df: pd.DataFrame, logger: logging.Logger) -> pd.DataFrame: if "UniquenessScore" in df.columns and not df["UniquenessScore"].isna().all(): return df logger.info("UniquenessScore missing; computing global uniqueness scores.") return add_global_uniqueness_scores(df) + # Vertex AI Configuration # Set these for your own Google Cloud project; there are no built-in defaults. -VERTEX_PROJECT_ID = os.environ.get("VERTEX_PROJECT_ID") or os.environ.get("GOOGLE_CLOUD_PROJECT") -VERTEX_LOCATION = os.environ.get("VERTEX_LOCATION") or os.environ.get("GOOGLE_CLOUD_LOCATION", "us-east5") +VERTEX_PROJECT_ID = os.environ.get("VERTEX_PROJECT_ID") or os.environ.get( + "GOOGLE_CLOUD_PROJECT" +) +VERTEX_LOCATION = os.environ.get("VERTEX_LOCATION") or os.environ.get( + "GOOGLE_CLOUD_LOCATION", "us-east5" +) VERTEX_BUCKET = os.environ.get("VERTEX_BUCKET") # e.g. gs:///batch VERTEX_MODEL_MAP = { @@ -309,7 +337,10 @@ def load_celltype_annotations( summary_path = celltype_file else: # Try common filenames in the directory - for filename in ["celltype_summary.csv", "program_celltype_annotations_summary.csv"]: + for filename in [ + "celltype_summary.csv", + "program_celltype_annotations_summary.csv", + ]: candidate = celltype_dir / filename if candidate.exists(): summary_path = candidate @@ -398,7 +429,9 @@ def prepare_enrichment_mapping( return {} df["category"] = df["category"].astype(str).str.strip() - df_sorted = df.sort_values(["program_id", "category", "fdr"], ascending=[True, True, True]) + df_sorted = df.sort_values( + ["program_id", "category", "fdr"], ascending=[True, True, True] + ) enrichment_by_program: Dict[int, Dict[str, List[dict]]] = {} for (pid, cat), sub in df_sorted.groupby(["program_id", "category"], sort=False): # type: ignore @@ -414,7 +447,7 @@ def load_ncbi_context(json_path: Optional[Path]) -> Dict[int, Dict[str, Any]]: if json_path: logger.warning("NCBI file not found: %s", json_path) return {} - + try: data = json.loads(json_path.read_text(encoding="utf-8")) # Parse keys as ints (JSON keys are always strings) @@ -424,13 +457,17 @@ def load_ncbi_context(json_path: Optional[Path]) -> Dict[int, Dict[str, Any]]: return {} -def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, allowed_genes: Optional[Set[str]] = None) -> str: +def format_ncbi_context( + ncbi_data: Dict[int, Dict[str, Any]], + program_id: int, + allowed_genes: Optional[Set[str]] = None, +) -> str: ctx = ncbi_data.get(program_id) if not ctx: return "Literature evidence: None available." - + lines = [] - + # 1. Official Gene Summaries (Filtered by allowed_genes) summaries = ctx.get("gene_summaries", {}) if summaries: @@ -444,36 +481,36 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a for gene in sorted_genes: if allowed_genes and gene not in allowed_genes: continue - + desc = summaries[gene] if source == "ncbi": # Remove citations like [provided by ...] desc = re.sub( - r'\s*\[provided by .*?\]\.?', - '', + r"\s*\[provided by .*?\]\.?", + "", desc, flags=re.IGNORECASE, ) lines.append(f"- {gene}: {desc}") count += 1 - + if count == 0: - lines.append("None available for selected genes.") + lines.append("None available for selected genes.") # 2. Aggregated Evidence (Snippets) # gene -> list of strings "Statement (PMID:123)" ev_map = ctx.get("evidence_snippets", {}) if ev_map: lines.append("\nAggregated Evidence (Contextual sentences from literature):") - + # Sort genes to be deterministic sorted_ev_genes = sorted(ev_map.keys()) has_snippets = False - + for sym in sorted_ev_genes: if allowed_genes and sym not in allowed_genes: continue - + s_list = ev_map[sym] if s_list: # Robust Deduplication: @@ -481,27 +518,27 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a # 2. Check PMID (Max 1 snippet per PMID per Gene) seen_normalized = set() seen_pmids = set() - + gene_snippets = [] - + for s in s_list: # s format: "Sentence text.." (PMID:12345) # Extract PMID using regex - pmid_match = re.search(r'\(PMID:(\d+)\)', s) + pmid_match = re.search(r"\(PMID:(\d+)\)", s) pmid = pmid_match.group(1) if pmid_match else None - + # Dedup by PMID (One sentence per paper per gene) if pmid: if pmid in seen_pmids: continue seen_pmids.add(pmid) - + # Dedup by Content (if PMID parsing fails or for safety) norm = s.strip(" .") if norm in seen_normalized: continue seen_normalized.add(norm) - + # Clean up display string # Remove double periods, ensure single period clean_s = s.strip() @@ -512,18 +549,19 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a # Just ensure it looks nice. if not clean_s.endswith("."): clean_s += "." - + gene_snippets.append(f"- {sym}: {clean_s}") - - if len(gene_snippets) >= 5: break # Max 5 per gene - + + if len(gene_snippets) >= 5: + break # Max 5 per gene + if gene_snippets: has_snippets = True lines.extend(gene_snippets) if not has_snippets: - lines.append("None found.") - + lines.append("None found.") + return "\n".join(lines) @@ -532,7 +570,7 @@ def load_regulator_data( significance_threshold: float = 0.05, ) -> Dict[int, pd.DataFrame]: """Load significant regulators from SCEPTRE results CSV. - + Returns a dict mapping program_id -> DataFrame with columns: [grna_target, log_2_fold_change, p_value] """ @@ -540,14 +578,14 @@ def load_regulator_data( if csv_path: logger.warning("Regulator file not found: %s", csv_path) return {} - + try: df = pd.read_csv(csv_path) df = standardize_regulator_results( df, significance_threshold=significance_threshold ) df = df[df["significant"] == True].copy() - + # Group by program result = {} for pid, group in df.groupby("program_id"): @@ -555,7 +593,7 @@ def load_regulator_data( if "adj_p_value" in group.columns: keep_cols.append("adj_p_value") result[pid] = group[keep_cols].copy() - + logger.info("Loaded regulators for %d programs", len(result)) return result except Exception as e: @@ -569,34 +607,40 @@ def format_regulator_context( top_n: int = 5, ) -> str: """Format basic regulator context for a program (Perturb-seq stats only). - + This is a fallback when no literature validation is available. Shows top positive and negative regulators with log2FC. """ reg_df = regulator_data.get(program_id) if reg_df is None or len(reg_df) == 0: return "Regulator perturbations: None significant." - + # Sort by log2FC sorted_df = reg_df.sort_values("log_2_fold_change") - + # Negative log2FC = positive regulators (knockdown reduces program, so gene activates it) positive_regs = sorted_df[sorted_df["log_2_fold_change"] < 0].head(top_n) # Positive log2FC = negative regulators (knockdown increases program, so gene represses it) negative_regs = sorted_df[sorted_df["log_2_fold_change"] > 0].tail(top_n).iloc[::-1] - - lines = ["Regulator perturbations (from Perturb-seq; log2FC indicates effect of knockdown on program activity):"] - + + lines = [ + "Regulator perturbations (from Perturb-seq; log2FC indicates effect of knockdown on program activity):" + ] + if len(positive_regs) > 0: lines.append("\nPositive regulators (activators; knockdown reduces program):") for _, row in positive_regs.iterrows(): - lines.append(f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}") - + lines.append( + f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}" + ) + if len(negative_regs) > 0: lines.append("\nNegative regulators (repressors; knockdown increases program):") for _, row in negative_regs.iterrows(): - lines.append(f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}") - + lines.append( + f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}" + ) + return "\n".join(lines) @@ -609,20 +653,20 @@ def format_regulator_analysis_context( min_score: int = 400, ) -> str: """Format comprehensive regulator analysis with compact STRING interactions. - + Shows top regulators (by effect size) with STRING interactions in compact format: - Gene (log2FC=X): Target1(score), Target2(score), ... - + Args: min_score: Minimum STRING score to display (default 400 = medium confidence) """ reg_df = regulator_data.get(program_id) ctx = ncbi_data.get(program_id, {}) validation = ctx.get("regulator_validation") - + if reg_df is None or len(reg_df) == 0: return "### Regulator evidence\nNo significant regulators identified from Perturb-seq." - + lines = ["### Regulator evidence"] lines.append( f"(Top {top_positive_regulators} activators + " @@ -630,13 +674,13 @@ def format_regulator_analysis_context( "STRING-DB interactions shown)" ) lines.append("") - + def format_compact_interactions(val: Dict[str, Any]) -> str: """Format STRING interactions compactly: Target1(score), Target2(score), ...""" string_ints = val.get("string_interactions", []) if not string_ints: return "" - + # Filter by min_score and format compactly parts = [] for si in string_ints: @@ -644,17 +688,17 @@ def format_compact_interactions(val: Dict[str, Any]) -> str: score = si.get("score", 0) if score >= min_score: parts.append(f"{target}({score})") - + if not parts: return "" return " → " + ", ".join(parts[:8]) # Limit to 8 targets per regulator - + # Get validated regulators from ncbi_data (includes all significant) activators = validation.get("positive_regulators", []) if validation else [] repressors = validation.get("negative_regulators", []) if validation else [] n_activators = min(len(activators), top_positive_regulators) n_repressors = min(len(repressors), top_negative_regulators) - + # Format activators if activators: lines.append("#### Activators (knockdown reduces program activity)") @@ -662,27 +706,27 @@ def format_compact_interactions(val: Dict[str, Any]) -> str: gene = reg.get("regulator", "") log2fc = reg.get("log2fc", 0) interactions_str = format_compact_interactions(reg) - + if interactions_str: lines.append(f"- **{gene}** (log2FC={log2fc:.3f}){interactions_str}") else: lines.append(f"- {gene} (log2FC={log2fc:.3f})") lines.append("") - - # Format repressors + + # Format repressors if repressors: lines.append("#### Repressors (knockdown increases program activity)") for reg in repressors[:n_repressors]: gene = reg.get("regulator", "") log2fc = reg.get("log2fc", 0) interactions_str = format_compact_interactions(reg) - + if interactions_str: lines.append(f"- **{gene}** (log2FC={log2fc:+.3f}){interactions_str}") else: lines.append(f"- {gene} (log2FC={log2fc:+.3f})") lines.append("") - + return "\n".join(lines) @@ -722,9 +766,13 @@ def select_program_genes( if program_df.empty: return [], [] top_loading_genes = ( - program_df.sort_values("Score", ascending=False)["Name"].head(top_loading).tolist() # type: ignore + program_df.sort_values("Score", ascending=False)["Name"] + .head(top_loading) + .tolist() # type: ignore ) - unique_ranked = program_df.sort_values("UniquenessScore", ascending=False)["Name"].tolist() # type: ignore + unique_ranked = program_df.sort_values("UniquenessScore", ascending=False)[ + "Name" + ].tolist() # type: ignore unique_genes = [gene for gene in unique_ranked if gene not in top_loading_genes] return top_loading_genes, unique_genes[:top_unique] @@ -763,9 +811,9 @@ def generate_prompt( top_enrichment=top_enrichment, genes_per_term=genes_per_term, ) - + celltype_context = format_celltype_context(celltype_map, program_id) - + # Build gene context gene_context = ( f"Top-loading genes (top {len(top_loading_genes)}):\n" @@ -776,10 +824,12 @@ def generate_prompt( f"\n\nUnique genes (top {len(unique_genes)} non-overlapping):\n" f"{', '.join(unique_genes)}" ) - + allowed_genes = set(top_loading_genes) | set(unique_genes) - ncbi_context = format_ncbi_context(ncbi_data, program_id, allowed_genes=allowed_genes) - + ncbi_context = format_ncbi_context( + ncbi_data, program_id, allowed_genes=allowed_genes + ) + regulator_analysis = format_regulator_analysis_context( regulator_data or {}, ncbi_data, @@ -886,7 +936,7 @@ def cmd_prepare(args: argparse.Namespace) -> int: return 2 program_ids = sorted(gene_df["program_id"].dropna().astype(int).unique().tolist()) - + # Filter by specific topics if requested selected_topics = parse_topics_value(args.topics) if selected_topics: @@ -958,21 +1008,22 @@ def cmd_prepare(args: argparse.Namespace) -> int: # Vertex AI Functions # ============================================================================= + def convert_to_vertex_jsonl(input_path: Path, output_path: Path) -> int: """Convert Anthropic batch JSON to JSONL format for Vertex AI. - + Anthropic direct API format: {"custom_id": "...", "params": {"model": "...", "messages": [...], "max_tokens": ...}} - + Vertex AI Claude batch format: {"custom_id": "...", "request": {"messages": [...], "anthropic_version": "vertex-2023-10-16", "max_tokens": ...}} """ - with open(input_path, 'r', encoding='utf-8') as f: + with open(input_path, "r", encoding="utf-8") as f: data = json.load(f) - + requests = data.get("requests", []) - - with open(output_path, 'w', encoding='utf-8') as f: + + with open(output_path, "w", encoding="utf-8") as f: for req in requests: # Convert from Anthropic format to Vertex AI format params = req.get("params", {}) @@ -982,16 +1033,18 @@ def convert_to_vertex_jsonl(input_path: Path, output_path: Path) -> int: "messages": params.get("messages", []), "anthropic_version": "vertex-2023-10-16", "max_tokens": params.get("max_tokens", 4096), - } + }, } # Optionally include system prompt if present if "system" in params: vertex_req["request"]["system"] = params["system"] - + line = json.dumps(vertex_req) f.write(line + "\n") - - logger.info(f"Converted {len(requests)} requests to Vertex AI JSONL format: {output_path}") + + logger.info( + f"Converted {len(requests)} requests to Vertex AI JSONL format: {output_path}" + ) return len(requests) @@ -999,13 +1052,13 @@ def upload_to_gcs(local_path: Path, gcs_uri: str) -> bool: """Upload a local file to GCS using gcloud CLI.""" cmd = ["gcloud", "storage", "cp", str(local_path), gcs_uri] logger.info(f"Uploading {local_path} to {gcs_uri}...") - + result = subprocess.run(cmd, capture_output=True, text=True) - + if result.returncode != 0: logger.error(f"GCS upload failed: {result.stderr}") return False - + logger.info(f"Upload successful: {gcs_uri}") return True @@ -1013,35 +1066,39 @@ def upload_to_gcs(local_path: Path, gcs_uri: str) -> bool: def cmd_submit_vertex(args: argparse.Namespace) -> int: """Submit a batch job to Vertex AI.""" if not VERTEX_AVAILABLE: - logger.error("google-genai package not installed. Run: pip install google-genai") + logger.error( + "google-genai package not installed. Run: pip install google-genai" + ) return 1 if not args.batch_file: logger.error("batch_file is required (via CLI or config).") return 1 - + batch_path = Path(args.batch_file) if not batch_path.exists(): logger.error(f"Batch file not found: {batch_path}") return 1 - + # Generate timestamp for unique naming timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S") base_name = batch_path.stem - + # Convert to JSONL format jsonl_path = batch_path.with_suffix(".jsonl") convert_to_vertex_jsonl(batch_path, jsonl_path) - + # Upload to GCS bucket = args.bucket or VERTEX_BUCKET if not bucket: - logger.error("No GCS bucket: pass --bucket gs:///batch or set VERTEX_BUCKET.") + logger.error( + "No GCS bucket: pass --bucket gs:///batch or set VERTEX_BUCKET." + ) return 1 gcs_input_uri = f"{bucket}/inputs/{base_name}_{timestamp}.jsonl" if not upload_to_gcs(jsonl_path, gcs_input_uri): return 2 - + # Resolve model name model = args.model or "claude-sonnet-4-5" if model in VERTEX_MODEL_MAP: @@ -1050,27 +1107,27 @@ def cmd_submit_vertex(args: argparse.Namespace) -> int: model_path = model else: model_path = f"publishers/anthropic/models/{model}" - + gcs_output_prefix = f"{bucket}/outputs/{base_name}_{timestamp}/" - + logger.info(f"Submitting batch job to Vertex AI...") logger.info(f" Model: {model_path}") logger.info(f" Input: {gcs_input_uri}") logger.info(f" Output: {gcs_output_prefix}") - + client = genai.Client(http_options=HttpOptions(api_version="v1")) # type: ignore - + job = client.batches.create( model=model_path, src=gcs_input_uri, config=CreateBatchJobConfig(dest=gcs_output_prefix), # type: ignore ) - + print(f"\nJob created!") print(f" Job name: {job.name}") print(f" Job state: {job.state}") print(f" Output: {gcs_output_prefix}") - + if args.wait: completed_states = { # type: ignore JobState.JOB_STATE_SUCCEEDED, # type: ignore @@ -1078,49 +1135,53 @@ def cmd_submit_vertex(args: argparse.Namespace) -> int: JobState.JOB_STATE_CANCELLED, # type: ignore JobState.JOB_STATE_PAUSED, # type: ignore } - + print("\nWaiting for job completion (checking every 30 seconds)...") - + while job.state not in completed_states: time.sleep(30) job = client.batches.get(name=job.name) print(f" Job state: {job.state}") - + print(f"\nFinal state: {job.state}") - + if job.state == JobState.JOB_STATE_SUCCEEDED: # type: ignore print(f"SUCCESS! Results available at: {gcs_output_prefix}") else: print(f"Job did not succeed.") return 3 else: - print(f"\nCheck status with: python {__file__} check-vertex --job-name {job.name}") - + print( + f"\nCheck status with: python {__file__} check-vertex --job-name {job.name}" + ) + return 0 def cmd_check_vertex(args: argparse.Namespace) -> int: """Check the status of a Vertex AI batch job.""" if not VERTEX_AVAILABLE: - logger.error("google-genai package not installed. Run: pip install google-genai") + logger.error( + "google-genai package not installed. Run: pip install google-genai" + ) return 1 if not args.job_name: logger.error("--job-name is required (via CLI or config).") return 1 - + client = genai.Client(http_options=HttpOptions(api_version="v1")) # type: ignore - + job = client.batches.get(name=args.job_name) - + print(f"Job name: {job.name}") print(f"Job state: {job.state}") - - if hasattr(job, 'create_time'): + + if hasattr(job, "create_time"): print(f"Created: {job.create_time}") - if hasattr(job, 'update_time'): + if hasattr(job, "update_time"): print(f"Updated: {job.update_time}") - + # Check if completed completed_states = { # type: ignore JobState.JOB_STATE_SUCCEEDED, # type: ignore @@ -1128,7 +1189,7 @@ def cmd_check_vertex(args: argparse.Namespace) -> int: JobState.JOB_STATE_CANCELLED, # type: ignore JobState.JOB_STATE_PAUSED, # type: ignore } - + if job.state == JobState.JOB_STATE_SUCCEEDED: # type: ignore print(f"\nJob completed successfully!") print(f"Download results with:") @@ -1137,7 +1198,7 @@ def cmd_check_vertex(args: argparse.Namespace) -> int: print(f"\nJob ended with state: {job.state}") else: print(f"\nJob is still processing...") - + return 0 @@ -1145,9 +1206,10 @@ def cmd_check_vertex(args: argparse.Namespace) -> int: # Anthropic Direct Batch API Commands (Default) # ============================================================================= + def cmd_submit_anthropic(args: argparse.Namespace) -> int: """Submit a batch job to Anthropic Batch API (default). - + This is the preferred method - uses Anthropic's native Batch API directly. Faster and simpler than Vertex AI for most use cases. """ @@ -1158,93 +1220,105 @@ def cmd_submit_anthropic(args: argparse.Namespace) -> int: if not args.batch_file: logger.error("batch_file is required (via CLI or config).") return 1 - + batch_path = Path(args.batch_file) if not batch_path.exists(): logger.error(f"Batch file not found: {batch_path}") return 1 - + # Load batch request with open(batch_path, "r", encoding="utf-8") as f: data = json.load(f) - + requests_list = data.get("requests", []) if not requests_list: logger.error("No requests found in batch file.") return 1 - + logger.info(f"Loaded {len(requests_list)} requests from {batch_path}") - + # Convert to Anthropic Batch API format # Anthropic expects: {"custom_id": "...", "params": {"model": ..., "max_tokens": ..., "messages": [...]}} anthropic_requests = [] model = args.model or MODEL max_tokens = args.max_tokens or 8192 - + for req in requests_list: custom_id = req.get("custom_id", f"request_{len(anthropic_requests)}") params = req.get("params", {}) messages = params.get("messages", []) - + if not messages: logger.warning(f"Skipping request {custom_id}: no messages") continue - - anthropic_requests.append({ - "custom_id": custom_id, - "params": { - "model": model, - "max_tokens": max_tokens, - "messages": messages, + + anthropic_requests.append( + { + "custom_id": custom_id, + "params": { + "model": model, + "max_tokens": max_tokens, + "messages": messages, + }, } - }) - - logger.info(f"Submitting {len(anthropic_requests)} requests to Anthropic Batch API...") + ) + + logger.info( + f"Submitting {len(anthropic_requests)} requests to Anthropic Batch API..." + ) logger.info(f" Model: {model}") logger.info(f" Max tokens: {max_tokens}") - + # Create Anthropic client (supports ANTHROPIC_AUTH_TOKEN for Bearer auth # or ANTHROPIC_API_KEY for x-api-key auth) auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") api_key = os.environ.get("ANTHROPIC_API_KEY") if not auth_token and not api_key: - logger.error("ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment. Set in .env file.") + logger.error( + "ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment. Set in .env file." + ) return 1 - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic(api_key=api_key) - + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic(api_key=api_key) + ) + # Submit batch batch = client.messages.batches.create(requests=anthropic_requests) - + print(f"\nBatch created!") print(f" Batch ID: {batch.id}") print(f" Status: {batch.processing_status}") print(f" Requests: {batch.request_counts}") - + # Save batch ID for later retrieval batch_id_file = batch_path.with_suffix(".batch_id") batch_id_file.write_text(batch.id, encoding="utf-8") logger.info(f"Saved batch ID to {batch_id_file}") - + if args.wait: print("\nWaiting for batch completion (checking every 30 seconds)...") - + while batch.processing_status == "in_progress": time.sleep(30) batch = client.messages.batches.retrieve(batch.id) - print(f" Status: {batch.processing_status} | Completed: {batch.request_counts.succeeded}/{batch.request_counts.processing + batch.request_counts.succeeded}") - + print( + f" Status: {batch.processing_status} | Completed: {batch.request_counts.succeeded}/{batch.request_counts.processing + batch.request_counts.succeeded}" + ) + print(f"\nFinal status: {batch.processing_status}") - + if batch.processing_status == "ended": # Download results output_file = batch_path.with_name(f"{batch_path.stem}_results.jsonl") logger.info(f"Downloading results to {output_file}...") - + with open(output_file, "w", encoding="utf-8") as f: for result in client.messages.batches.results(batch.id): f.write(json.dumps(result.model_dump()) + "\n") - + print(f"SUCCESS! Results saved to: {output_file}") print(f" Succeeded: {batch.request_counts.succeeded}") print(f" Errored: {batch.request_counts.errored}") @@ -1252,9 +1326,13 @@ def cmd_submit_anthropic(args: argparse.Namespace) -> int: print(f"Batch did not complete successfully.") return 3 else: - print(f"\nCheck status with: python {__file__} check-anthropic --batch-id {batch.id}") - print(f"Or retrieve results later with: python {__file__} results-anthropic --batch-id {batch.id}") - + print( + f"\nCheck status with: python {__file__} check-anthropic --batch-id {batch.id}" + ) + print( + f"Or retrieve results later with: python {__file__} results-anthropic --batch-id {batch.id}" + ) + return 0 @@ -1294,8 +1372,7 @@ def cmd_submit_single(args: argparse.Namespace) -> int: # Determine output path output_file = Path( - args.output - or str(batch_path.with_name(batch_path.stem + "_results.jsonl")) + args.output or str(batch_path.with_name(batch_path.stem + "_results.jsonl")) ) output_file.parent.mkdir(parents=True, exist_ok=True) @@ -1345,14 +1422,18 @@ def cmd_submit_single(args: argparse.Namespace) -> int: for idx, req in enumerate(requests_list, start=1): custom_id = req.get("custom_id", f"request_{idx}") if custom_id in completed_ids: - logger.info("[%d/%d] %s: skipped (already completed)", idx, total, custom_id) + logger.info( + "[%d/%d] %s: skipped (already completed)", idx, total, custom_id + ) succeeded += 1 continue params = req.get("params", {}) messages = params.get("messages", []) if not messages: - logger.warning("[%d/%d] %s: skipped (no messages)", idx, total, custom_id) + logger.warning( + "[%d/%d] %s: skipped (no messages)", idx, total, custom_id + ) continue try: @@ -1386,7 +1467,11 @@ def cmd_submit_single(args: argparse.Namespace) -> int: out_tok = response.usage.output_tokens if response.usage else "?" logger.info( "[%d/%d] %s: succeeded (%sin/%sout)", - idx, total, custom_id, in_tok, out_tok, + idx, + total, + custom_id, + in_tok, + out_tok, ) except Exception as exc: @@ -1420,18 +1505,26 @@ def cmd_check_anthropic(args: argparse.Namespace) -> int: batch_id_file = Path(args.batch_file).with_suffix(".batch_id") if batch_id_file.exists(): batch_id = batch_id_file.read_text(encoding="utf-8").strip() - + if not batch_id: - logger.error("--batch-id is required (or provide --batch-file with saved .batch_id)") + logger.error( + "--batch-id is required (or provide --batch-file with saved .batch_id)" + ) return 1 - + auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") api_key = os.environ.get("ANTHROPIC_API_KEY") if not auth_token and not api_key: - logger.error("ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment.") + logger.error( + "ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment." + ) return 1 - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic(api_key=api_key) + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic(api_key=api_key) + ) batch = client.messages.batches.retrieve(batch_id) print(f"Batch ID: {batch.id}") @@ -1442,11 +1535,11 @@ def cmd_check_anthropic(args: argparse.Namespace) -> int: print(f" Succeeded: {batch.request_counts.succeeded}") print(f" Errored: {batch.request_counts.errored}") print(f" Canceled: {batch.request_counts.canceled}") - + if batch.processing_status == "ended": print(f"\nBatch completed! Retrieve results with:") print(f" python {__file__} results-anthropic --batch-id {batch.id}") - + return 0 @@ -1462,36 +1555,42 @@ def cmd_results_anthropic(args: argparse.Namespace) -> int: batch_id_file = Path(args.batch_file).with_suffix(".batch_id") if batch_id_file.exists(): batch_id = batch_id_file.read_text(encoding="utf-8").strip() - + if not batch_id: logger.error("--batch-id is required") return 1 - + auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") api_key = os.environ.get("ANTHROPIC_API_KEY") if not auth_token and not api_key: - logger.error("ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment.") + logger.error( + "ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment." + ) return 1 - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic(api_key=api_key) + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic(api_key=api_key) + ) batch = client.messages.batches.retrieve(batch_id) if batch.processing_status != "ended": logger.error(f"Batch is still processing: {batch.processing_status}") return 1 - + output_file = Path(args.output or f"batch_{batch_id}_results.jsonl") - + logger.info(f"Downloading results to {output_file}...") - + with open(output_file, "w", encoding="utf-8") as f: for result in client.messages.batches.results(batch_id): f.write(json.dumps(result.model_dump()) + "\n") - + print(f"Results saved to: {output_file}") print(f" Succeeded: {batch.request_counts.succeeded}") print(f" Errored: {batch.request_counts.errored}") - + return 0 @@ -1612,43 +1711,40 @@ def add_shared_options(p: argparse.ArgumentParser) -> None: help="Adjusted p-value threshold used when the regulator file has no 'significant' column", ) - p_prepare = subparsers.add_parser("prepare", help="Build the batch request JSON only") + p_prepare = subparsers.add_parser( + "prepare", help="Build the batch request JSON only" + ) add_shared_options(p_prepare) p_prepare.set_defaults(func=cmd_prepare) # ========================================================================= # Anthropic Direct API Commands (Default - Recommended) # ========================================================================= - + # Anthropic submit command (DEFAULT) p_submit = subparsers.add_parser( - "submit", - help="Submit batch to Anthropic Batch API (default, recommended)" + "submit", help="Submit batch to Anthropic Batch API (default, recommended)" ) p_submit.add_argument( - "batch_file", - nargs="?", - help="Path to the prepared batch JSON file" + "batch_file", nargs="?", help="Path to the prepared batch JSON file" ) p_submit.add_argument( "--config", help="Path to config file (YAML or JSON)", ) p_submit.add_argument( - "--model", - default=MODEL, - help=f"Model to use (default: {MODEL})" + "--model", default=MODEL, help=f"Model to use (default: {MODEL})" ) p_submit.add_argument( "--max-tokens", type=int, default=8192, - help="Max tokens for response (default: 8192)" + help="Max tokens for response (default: 8192)", ) p_submit.add_argument( "--wait", action="store_true", - help="Wait for job completion and download results" + help="Wait for job completion and download results", ) p_submit.set_defaults(func=cmd_submit_anthropic) @@ -1684,17 +1780,10 @@ def add_shared_options(p: argparse.ArgumentParser) -> None: p_submit_single.set_defaults(func=cmd_submit_single) # Anthropic check command - p_check = subparsers.add_parser( - "check", - help="Check status of Anthropic batch job" - ) - p_check.add_argument( - "--batch-id", - help="Anthropic batch ID" - ) + p_check = subparsers.add_parser("check", help="Check status of Anthropic batch job") + p_check.add_argument("--batch-id", help="Anthropic batch ID") p_check.add_argument( - "--batch-file", - help="Path to batch JSON (will read .batch_id file)" + "--batch-file", help="Path to batch JSON (will read .batch_id file)" ) p_check.add_argument( "--config", @@ -1704,21 +1793,13 @@ def add_shared_options(p: argparse.ArgumentParser) -> None: # Anthropic results command p_results = subparsers.add_parser( - "results", - help="Download results from completed Anthropic batch" - ) - p_results.add_argument( - "--batch-id", - help="Anthropic batch ID" + "results", help="Download results from completed Anthropic batch" ) + p_results.add_argument("--batch-id", help="Anthropic batch ID") p_results.add_argument( - "--batch-file", - help="Path to batch JSON (will read .batch_id file)" - ) - p_results.add_argument( - "--output", - help="Output JSONL file path" + "--batch-file", help="Path to batch JSON (will read .batch_id file)" ) + p_results.add_argument("--output", help="Output JSONL file path") p_results.add_argument( "--config", help="Path to config file (YAML or JSON)", @@ -1731,13 +1812,10 @@ def add_shared_options(p: argparse.ArgumentParser) -> None: # Vertex AI submit command p_submit_vertex = subparsers.add_parser( - "submit-vertex", - help="Submit a prepared batch JSON to Vertex AI" + "submit-vertex", help="Submit a prepared batch JSON to Vertex AI" ) p_submit_vertex.add_argument( - "batch_file", - nargs="?", - help="Path to the prepared batch JSON file" + "batch_file", nargs="?", help="Path to the prepared batch JSON file" ) p_submit_vertex.add_argument( "--config", @@ -1746,28 +1824,28 @@ def add_shared_options(p: argparse.ArgumentParser) -> None: p_submit_vertex.add_argument( "--model", default="claude-sonnet-4-5", - help="Model to use (default: claude-sonnet-4-5). Options: " + ", ".join(VERTEX_MODEL_MAP.keys()) + help="Model to use (default: claude-sonnet-4-5). Options: " + + ", ".join(VERTEX_MODEL_MAP.keys()), ) p_submit_vertex.add_argument( "--bucket", default=VERTEX_BUCKET, - help=f"GCS bucket prefix (default: {VERTEX_BUCKET})" + help=f"GCS bucket prefix (default: {VERTEX_BUCKET})", ) p_submit_vertex.add_argument( "--wait", action="store_true", - help="Wait for job completion instead of returning immediately" + help="Wait for job completion instead of returning immediately", ) p_submit_vertex.set_defaults(func=cmd_submit_vertex) # Vertex AI check command p_check_vertex = subparsers.add_parser( - "check-vertex", - help="Check status of a Vertex AI batch job" + "check-vertex", help="Check status of a Vertex AI batch job" ) p_check_vertex.add_argument( "--job-name", - help="Full job name from Vertex AI (e.g., projects/.../batchPredictionJobs/...)" + help="Full job name from Vertex AI (e.g., projects/.../batchPredictionJobs/...)", ) p_check_vertex.add_argument( "--config", diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/04_parse_and_summarize.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/04_parse_and_summarize.py index 5ed39a0..7bc9953 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/04_parse_and_summarize.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/04_parse_and_summarize.py @@ -104,7 +104,9 @@ def get_cli_overrides(argv: List[str]) -> set[str]: def apply_config_overrides( args: argparse.Namespace, config: Dict[str, object], cli_overrides: set[str] ) -> argparse.Namespace: - steps_cfg = config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + steps_cfg = ( + config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + ) step_cfg = steps_cfg.get("parse_results", {}) if not isinstance(step_cfg, dict): return args @@ -120,24 +122,24 @@ def apply_config_overrides( def download_from_gcs(gcs_prefix: str, local_dir: str) -> List[str]: """Download all files from a GCS prefix to a local directory. - + Returns list of downloaded file paths. """ os.makedirs(local_dir, exist_ok=True) - + # Ensure prefix ends with * for glob-like behavior if not gcs_prefix.endswith("*"): gcs_prefix = gcs_prefix.rstrip("/") + "/*" - + print(f"Downloading from {gcs_prefix} to {local_dir}...") - + cmd = ["gcloud", "storage", "cp", gcs_prefix, local_dir] result = subprocess.run(cmd, capture_output=True, text=True) - + if result.returncode != 0: print(f" ✗ GCS download failed: {result.stderr}") return [] - + # Find downloaded files downloaded = glob.glob(os.path.join(local_dir, "*")) print(f" ✓ Downloaded {len(downloaded)} files") @@ -150,7 +152,7 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: Supports both Anthropic direct API and Vertex AI response formats: - Anthropic: {"custom_id": "...", "result": {"type": "succeeded", "message": {"content": [{"text": "..."}]}}} - Vertex AI: {"custom_id": "...", "response": {"content": [{"text": "..."}]}} - + Returns a list of topic IDs successfully written. """ if not os.path.exists(output_dir): @@ -164,7 +166,11 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: try: data = json.loads(line) custom_id = data.get("custom_id") - if custom_id and isinstance(custom_id, str) and custom_id.startswith("topic_"): + if ( + custom_id + and isinstance(custom_id, str) + and custom_id.startswith("topic_") + ): # expected format: topic__annotation m = re.match(r"topic_(\d+)", custom_id) if not m: @@ -172,15 +178,17 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: topic_number = int(m.group(1)) text_content = "" - + # Try Vertex AI format first: {"response": {"content": [{"text": "..."}]}} response = data.get("response", {}) if response: content = response.get("content", []) if isinstance(content, list) and content: first = content[0] - text_content = first.get("text", "") if isinstance(first, dict) else "" - + text_content = ( + first.get("text", "") if isinstance(first, dict) else "" + ) + # Fallback to Anthropic format: {"result": {"type": "succeeded", "message": {"content": [...]}}} if not text_content: result = data.get("result", {}) @@ -189,14 +197,22 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: content = message.get("content", []) if isinstance(content, list) and content: first = content[0] - text_content = first.get("text", "") if isinstance(first, dict) else "" + text_content = ( + first.get("text", "") + if isinstance(first, dict) + else "" + ) elif result.get("type") == "errored": error = result.get("error", {}) - print(f" ✗ Topic {topic_number} failed. Reason: {error.get('message', 'Unknown error')}") + print( + f" ✗ Topic {topic_number} failed. Reason: {error.get('message', 'Unknown error')}" + ) continue if text_content: - output_filename = os.path.join(output_dir, f"topic_{topic_number}_annotation.md") + output_filename = os.path.join( + output_dir, f"topic_{topic_number}_annotation.md" + ) with open(output_filename, "w", encoding="utf-8") as out_f: out_f.write(text_content) print(f" ✓ Saved Topic {topic_number}") @@ -211,7 +227,9 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: return saved_ids -def load_top_genes_by_topic(gene_loading_file: str, top_n: int = 10) -> Dict[int, List[str]]: +def load_top_genes_by_topic( + gene_loading_file: str, top_n: int = 10 +) -> Dict[int, List[str]]: """Load top genes by RowID from gene loading CSV.""" top_genes_by_topic: Dict[int, List[str]] = {} if not gene_loading_file or not os.path.exists(gene_loading_file): @@ -225,12 +243,15 @@ def load_top_genes_by_topic(gene_loading_file: str, top_n: int = 10) -> Dict[int elif "RowID" in gene_df.columns: group_col = "RowID" else: - print(" ✗ CSV missing 'program_id' or 'RowID' column") - return top_genes_by_topic + print(" ✗ CSV missing 'program_id' or 'RowID' column") + return top_genes_by_topic for topic_id, group in gene_df.groupby(group_col): genes = ( - group.sort_values("Score", ascending=False)["Name"].astype(str).head(top_n).tolist() + group.sort_values("Score", ascending=False)["Name"] + .astype(str) + .head(top_n) + .tolist() ) top_genes_by_topic[int(topic_id)] = genes print(f" ✓ Loaded gene data for {len(top_genes_by_topic)} topics") @@ -253,7 +274,9 @@ def generate_unique_topic_names( key=lambda x: int(re.search(r"topic_(\d+)_", x).group(1)), ) - top_genes_by_topic = load_top_genes_by_topic(gene_loading_file) if gene_loading_file else {} + top_genes_by_topic = ( + load_top_genes_by_topic(gene_loading_file) if gene_loading_file else {} + ) topic_rows: List[Dict[str, object]] = [] used_names = set() @@ -272,7 +295,9 @@ def generate_unique_topic_names( content, re.DOTALL, ) - keywords_match = re.search(r"\*\*(?:Program label|Three Key Words):\*\*(.*?)\n", content, re.DOTALL) + keywords_match = re.search( + r"\*\*(?:Program label|Three Key Words):\*\*(.*?)\n", content, re.DOTALL + ) summary = summary_match.group(1).strip() if summary_match else "" keywords_str = keywords_match.group(1).strip() if keywords_match else "" @@ -297,7 +322,9 @@ def generate_unique_topic_names( final_name = f"{base_name} ({signature_gene})" while final_name in used_names: - if topic_number in top_genes_by_topic and counter - 2 < len(top_genes_by_topic[topic_number]): + if topic_number in top_genes_by_topic and counter - 2 < len( + top_genes_by_topic[topic_number] + ): signature_gene = top_genes_by_topic[topic_number][counter - 2] final_name = f"{base_name} ({signature_gene})" else: @@ -307,7 +334,9 @@ def generate_unique_topic_names( re.IGNORECASE, ) if distinguishing_terms and counter - 2 < len(distinguishing_terms): - final_name = f"{base_name} ({distinguishing_terms[counter - 2]})" + final_name = ( + f"{base_name} ({distinguishing_terms[counter - 2]})" + ) else: final_name = f"{base_name} ({counter})" counter += 1 @@ -351,18 +380,27 @@ def build_parser() -> argparse.ArgumentParser: description="Parse Anthropic JSONL results to markdown and generate a summary CSV." ) parser.add_argument("--config", help="Path to config file (YAML or JSON)") - parser.add_argument("--results-jsonl", help="Path to the .jsonl results file from Anthropic Batch API") - parser.add_argument("--markdown-dir", help="Directory to write/read per-topic markdown files") + parser.add_argument( + "--results-jsonl", + help="Path to the .jsonl results file from Anthropic Batch API", + ) + parser.add_argument( + "--markdown-dir", help="Directory to write/read per-topic markdown files" + ) parser.add_argument("--summary-csv", help="Path to write the summary CSV") parser.add_argument( "--gene-loading-file", help="Optional gene loading CSV (columns: Name, Score, RowID) to enrich summary naming", ) - parser.add_argument("--no-parse", action="store_true", help="Skip parsing JSONL → markdown step") - parser.add_argument("--no-summary", action="store_true", help="Skip generating summary CSV step") + parser.add_argument( + "--no-parse", action="store_true", help="Skip parsing JSONL → markdown step" + ) + parser.add_argument( + "--no-summary", action="store_true", help="Skip generating summary CSV step" + ) parser.add_argument( "--gcs-prefix", - help="GCS prefix to download Vertex AI results from (e.g., gs://bucket/outputs/job_123/)" + help="GCS prefix to download Vertex AI results from (e.g., gs://bucket/outputs/job_123/)", ) return parser @@ -382,7 +420,7 @@ def main() -> None: # Download Vertex AI results from GCS temp_dir = tempfile.mkdtemp(prefix="vertex_results_") downloaded = download_from_gcs(args.gcs_prefix, temp_dir) - + # Find JSONL files in downloaded results jsonl_files = [f for f in downloaded if f.endswith(".jsonl")] if jsonl_files: @@ -396,7 +434,9 @@ def main() -> None: if not args.markdown_dir: raise SystemExit("--markdown-dir is required unless --no-parse is set") if not results_jsonl: - raise SystemExit("--results-jsonl or --gcs-prefix is required unless --no-parse is set") + raise SystemExit( + "--results-jsonl or --gcs-prefix is required unless --no-parse is set" + ) parse_final_results(results_jsonl, args.markdown_dir) # Step B: markdown → summary CSV @@ -405,7 +445,9 @@ def main() -> None: raise SystemExit("--markdown-dir is required unless --no-summary is set") if not args.summary_csv: raise SystemExit("--summary-csv is required unless --no-summary is set") - generate_unique_topic_names(args.markdown_dir, args.summary_csv, args.gene_loading_file) + generate_unique_topic_names( + args.markdown_dir, args.summary_csv, args.gene_loading_file + ) if __name__ == "__main__": diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/05_generate_html_report.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/05_generate_html_report.py index 9026cf7..422154e 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/05_generate_html_report.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/05_generate_html_report.py @@ -23,45 +23,61 @@ from markdown import markdown import sys as _sys from pathlib import Path as _Path + _sys.path.insert(0, str(_Path(__file__).resolve().parent)) from column_mapper import standardize_regulator_results PRIORITY_GENES = [ - "Hif1a", "Epas1", "Arnt", "Fzd4", "Fzd6", "Idh2", "Mdh2", "Ogdh", - "Hsp90ab1", "Hspa5", "Creb3l2", "Fkbp8", "Kdr", "Egln1", "Egln2", "Foxo1", "Foxo3" + "Hif1a", + "Epas1", + "Arnt", + "Fzd4", + "Fzd6", + "Idh2", + "Mdh2", + "Ogdh", + "Hsp90ab1", + "Hspa5", + "Creb3l2", + "Fkbp8", + "Kdr", + "Egln1", + "Egln2", + "Foxo1", + "Foxo3", ] def extract_program_stats(annotation_md: str) -> dict: """Extract stats from annotation markdown.""" - stats = {'top_loading': '', 'unique': '', 'celltype': ''} - + stats = {"top_loading": "", "unique": "", "celltype": ""} + # Extract top-loading genes - match = re.search(r'\*\*Top-loading genes?:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Top-loading genes?:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['top_loading'] = match.group(1).strip() - + stats["top_loading"] = match.group(1).strip() + # Extract unique genes - match = re.search(r'\*\*Unique genes?:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Unique genes?:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['unique'] = match.group(1).strip() - + stats["unique"] = match.group(1).strip() + # Extract cell-type enrichment - match = re.search(r'\*\*Cell-type enrichment:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Cell-type enrichment:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['celltype'] = match.group(1).strip() - + stats["celltype"] = match.group(1).strip() + # Extract brief summary - match = re.search(r'\*\*Brief Summary:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Brief Summary:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['summary'] = match.group(1).strip() - + stats["summary"] = match.group(1).strip() + # Extract program label - match = re.search(r'\*\*Program label:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Program label:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['label'] = match.group(1).strip() - + stats["label"] = match.group(1).strip() + return stats @@ -75,7 +91,7 @@ def generate_report( regulator_significance_threshold: float = 0.05, ): """Generate Design A: Scientific Minimal HTML report.""" - + # Load data summary_df = pd.read_csv(summary_csv) volcano_df = None @@ -84,94 +100,111 @@ def generate_report( pd.read_csv(volcano_csv), significance_threshold=regulator_significance_threshold, ) - + # Process volcano data by program volcano_by_program = {} if volcano_df is not None: pvalue_col = ( - 'adj_p_value' - if 'adj_p_value' in volcano_df.columns and not volcano_df['adj_p_value'].isna().all() - else 'p_value' + "adj_p_value" + if "adj_p_value" in volcano_df.columns + and not volcano_df["adj_p_value"].isna().all() + else "p_value" ) - volcano_df['neg_log10_pvalue'] = -np.log10(volcano_df[pvalue_col].replace(0, 1e-300)) - volcano_df.loc[np.isinf(volcano_df['neg_log10_pvalue']), 'neg_log10_pvalue'] = 300 - - for tid, group in volcano_df.groupby('program_id'): + volcano_df["neg_log10_pvalue"] = -np.log10( + volcano_df[pvalue_col].replace(0, 1e-300) + ) + volcano_df.loc[np.isinf(volcano_df["neg_log10_pvalue"]), "neg_log10_pvalue"] = ( + 300 + ) + + for tid, group in volcano_df.groupby("program_id"): volcano_by_program[int(tid)] = [ { - 'g': row['grna_target'], - 'fc': round(row['log_2_fold_change'], 3), - 'p': round(row['neg_log10_pvalue'], 2), - 's': bool(row['significant']) + "g": row["grna_target"], + "fc": round(row["log_2_fold_change"], 3), + "p": round(row["neg_log10_pvalue"], 2), + "s": bool(row["significant"]), } for _, row in group.iterrows() ] - + # Build per-program data programs_data = [] for _, row in summary_df.iterrows(): - topic_id = int(row['Topic']) - topic_name = row['Name'] - + topic_id = int(row["Topic"]) + topic_name = row["Name"] + # Read annotation ann_path = Path(annotations_dir) / f"topic_{topic_id}_annotation.md" - annotation_md = ann_path.read_text(encoding='utf-8') if ann_path.exists() else "" - + annotation_md = ( + ann_path.read_text(encoding="utf-8") if ann_path.exists() else "" + ) + # Extract stats stats = extract_program_stats(annotation_md) - + # Enrichment paths enr_rel = os.path.relpath(enrichment_dir, os.path.dirname(output_html)) - - programs_data.append({ - 'id': topic_id, - 'name': topic_name, - 'label': stats.get('label', topic_name), - 'summary': stats.get('summary', ''), - 'top_loading': stats.get('top_loading', ''), - 'unique': stats.get('unique', ''), - 'celltype': stats.get('celltype', ''), - 'annotation_html': markdown(annotation_md, extensions=['tables', 'fenced_code']), - 'annotation_text': annotation_md, # For full-text search - 'kegg_fig': f"{enr_rel}/program_{topic_id}_kegg_enrichment.png", - 'process_fig': f"{enr_rel}/program_{topic_id}_process_enrichment.png", - 'volcano': volcano_by_program.get(topic_id, []), - }) - - generated_on = datetime.now().strftime('%Y-%m-%d %H:%M') - + + programs_data.append( + { + "id": topic_id, + "name": topic_name, + "label": stats.get("label", topic_name), + "summary": stats.get("summary", ""), + "top_loading": stats.get("top_loading", ""), + "unique": stats.get("unique", ""), + "celltype": stats.get("celltype", ""), + "annotation_html": markdown( + annotation_md, extensions=["tables", "fenced_code"] + ), + "annotation_text": annotation_md, # For full-text search + "kegg_fig": f"{enr_rel}/program_{topic_id}_kegg_enrichment.png", + "process_fig": f"{enr_rel}/program_{topic_id}_process_enrichment.png", + "volcano": volcano_by_program.get(topic_id, []), + } + ) + + generated_on = datetime.now().strftime("%Y-%m-%d %H:%M") + # Generate HTML html = generate_design_a_html(programs_data, len(programs_data), generated_on) - + Path(output_html).parent.mkdir(parents=True, exist_ok=True) - with open(output_html, 'w', encoding='utf-8') as f: + with open(output_html, "w", encoding="utf-8") as f: f.write(html) - + print(f"✓ Design A (Scientific Minimal): {output_html}") print(f" - {len(programs_data)} programs") def generate_design_a_html(programs_data, num_programs, generated_on): """Generate Design A HTML.""" - + # Build TOC - toc_items = "\n".join([ - f'' - f'{p["id"]}. {p["name"]}' - for p in programs_data - ]) - + toc_items = "\n".join( + [ + f'' + f"{p['id']}. {p['name']}" + for p in programs_data + ] + ) + # Program selector options - options = "\n".join([ - f'' - for p in programs_data - ]) - + options = "\n".join( + [ + f'' + for p in programs_data + ] + ) + # Program data JS - programs_js = f"window.PROGRAMS = {json.dumps({p['id']: p for p in programs_data})};\n" + programs_js = ( + f"window.PROGRAMS = {json.dumps({p['id']: p for p in programs_data})};\n" + ) programs_js += f"window.PRIORITY_GENES = {json.dumps(PRIORITY_GENES)};\n" - - html = f''' + + html = f""" @@ -585,8 +618,8 @@ def generate_design_a_html(programs_data, num_programs, generated_on): }})(); -''' - +""" + return html @@ -643,7 +676,9 @@ def get_cli_overrides(argv: list[str]) -> set[str]: def apply_config_overrides( args: argparse.Namespace, config: dict, cli_overrides: set[str] ) -> argparse.Namespace: - steps_cfg = config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + steps_cfg = ( + config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + ) step_cfg = steps_cfg.get("html_report", {}) if not isinstance(step_cfg, dict): return args @@ -657,7 +692,7 @@ def apply_config_overrides( return args -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", help="Path to config file (YAML or JSON)") parser.add_argument("--summary-csv") @@ -690,7 +725,9 @@ def apply_config_overrides( raise SystemExit(f"Missing required arguments: {', '.join(missing)}") generate_report( - args.summary_csv, args.annotations_dir, args.enrichment_dir, + args.summary_csv, + args.annotations_dir, + args.enrichment_dir, args.volcano_csv, args.gene_loading_csv, args.output_html, diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/batch_utils.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/batch_utils.py index 6c7e2ee..fff6874 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/batch_utils.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/batch_utils.py @@ -50,8 +50,13 @@ def check_anthropic_batch( try: import os + auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic() + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic() + ) batch = client.messages.batches.retrieve(batch_id) except Exception as exc: logger.error("Failed to retrieve batch %s: %s", batch_id, exc) diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/column_mapper.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/column_mapper.py index 58a7bdd..7573743 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/column_mapper.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/column_mapper.py @@ -25,7 +25,7 @@ class ColumnMapper: """ Maps flexible column names to standardized names. - + Supports case-insensitive matching and common variants for: - Gene names: Name, Gene, gene_name, GeneName, etc. - Scores: Score, score, Loading, loading, etc. @@ -34,155 +34,248 @@ class ColumnMapper: - Log2 fold change: log2_fc, log2FC, log2_fold_change, etc. - FDR/p-values: fdr, FDR, p_value, pval, etc. """ - + # Define column aliases (all lowercase for matching) ALIASES: Dict[str, List[str]] = { - 'gene': [ - 'gene', 'name', 'gene_name', 'genename', 'gene_symbol', - 'genesymbol', 'symbol', 'gene_id', 'geneid' + "gene": [ + "gene", + "name", + "gene_name", + "genename", + "gene_symbol", + "genesymbol", + "symbol", + "gene_id", + "geneid", + ], + "score": [ + "score", + "loading", + "weight", + "value", + "loading_score", + "loadingscore", + "gene_score", + "genescore", ], - 'score': [ - 'score', 'loading', 'weight', 'value', 'loading_score', - 'loadingscore', 'gene_score', 'genescore' + "program_id": [ + "program_id", + "programid", + "rowid", + "row_id", + "topic", + "topic_id", + "topicid", + "program", + "program_number", + "topic_number", + "k", + "component", + "factor", + "response_id", + "response", + "program_name", + "program_label", + "topic_name", ], - 'program_id': [ - 'program_id', 'programid', 'rowid', 'row_id', - 'topic', 'topic_id', 'topicid', 'program', 'program_number', - 'topic_number', 'k', 'component', 'factor', 'response_id', - 'response', 'program_name', 'program_label', 'topic_name' + "cell_type": [ + "cell_type", + "celltype", + "cluster", + "cell_cluster", + "cellcluster", + "annotation", + "cell_annotation", + "cellannotation", + "type", + "celltype_id", + "cluster_id", ], - 'cell_type': [ - 'cell_type', 'celltype', 'cluster', 'cell_cluster', - 'cellcluster', 'annotation', 'cell_annotation', - 'cellannotation', 'type', 'celltype_id', 'cluster_id' + "log2_fc": [ + "log2_fc", + "log2fc", + "log2_fold_change", + "log2foldchange", + "log_2_fold_change", + "log2_fc_in_vs_out", + "log2fc_in_vs_out", + "log2_fold", + "log2fold", + "lfc", + "l2fc", + "fold_change", + "foldchange", + "fc", + "log2ratio", + "log2_ratio", ], - 'log2_fc': [ - 'log2_fc', 'log2fc', 'log2_fold_change', 'log2foldchange', - 'log_2_fold_change', 'log2_fc_in_vs_out', 'log2fc_in_vs_out', 'log2_fold', - 'log2fold', 'lfc', 'l2fc', 'fold_change', 'foldchange', - 'fc', 'log2ratio', 'log2_ratio' + "fdr": [ + "fdr", + "fdr_corrected", + "q_value", + "qvalue", + "qval", + "q", + "padj", + "p_adj", + "p_adjusted", + "adjusted_pvalue", + "p_value", + "pvalue", + "pval", + "p", ], - 'fdr': [ - 'fdr', 'fdr_corrected', 'q_value', 'qvalue', 'qval', 'q', - 'padj', 'p_adj', 'p_adjusted', 'adjusted_pvalue', - 'p_value', 'pvalue', 'pval', 'p' + "regulator_gene": [ + "grna_target", + "target_gene", + "target_gene_name", + "target_gene_names", + "perturbation_target", + "gene_target", + "regulator", + "regulator_gene", + "gene", ], - 'regulator_gene': [ - 'grna_target', 'target_gene', 'target_gene_name', - 'target_gene_names', 'perturbation_target', 'gene_target', - 'regulator', 'regulator_gene', 'gene' + "raw_p_value": [ + "p_value", + "pvalue", + "p_val", + "pval", + "raw_p_value", + "raw_pvalue", + "raw_pval", + "nominal_p_value", + "nominal_pvalue", + "nominal_pval", + "p", ], - 'raw_p_value': [ - 'p_value', 'pvalue', 'p_val', 'pval', 'raw_p_value', - 'raw_pvalue', 'raw_pval', 'nominal_p_value', - 'nominal_pvalue', 'nominal_pval', 'p' + "adj_p_value": [ + "adj_pval", + "adj_pvalue", + "adjusted_pval", + "adjusted_pvalue", + "adjusted_p_value", + "p_adj", + "padj", + "p_adjusted", + "fdr", + "fdr_corrected", + "q_value", + "qvalue", + "qval", + "q", ], - 'adj_p_value': [ - 'adj_pval', 'adj_pvalue', 'adjusted_pval', 'adjusted_pvalue', - 'adjusted_p_value', 'p_adj', 'padj', 'p_adjusted', - 'fdr', 'fdr_corrected', 'q_value', 'qvalue', 'qval', 'q' + "significant": [ + "significant", + "is_significant", + "significant_hit", + "is_sig", + "sig", + "passed_significance", ], - 'significant': [ - 'significant', 'is_significant', 'significant_hit', - 'is_sig', 'sig', 'passed_significance' - ] } - + def __init__(self, df: pd.DataFrame): """ Initialize mapper with a DataFrame. - + Args: df: DataFrame to map columns for """ self.df = df self.columns_lower = {col.lower(): col for col in df.columns} self._mapped: Dict[str, str] = {} - + def get_column(self, standard_name: str, required: bool = True) -> Optional[str]: """ Find the actual column name that matches the standard name. - + Args: standard_name: Standard column name (e.g., 'gene', 'score', 'program_id') required: If True, raise ValueError when column not found - + Returns: Actual column name in the DataFrame, or None if not found and not required - + Raises: ValueError: If required=True and column not found """ # Return cached mapping if available if standard_name in self._mapped: return self._mapped[standard_name] - + # Get aliases for this standard name if standard_name not in self.ALIASES: if required: raise ValueError(f"Unknown standard column name: {standard_name}") return None - + aliases = self.ALIASES[standard_name] - + # Try to find a matching column (case-insensitive) for alias in aliases: if alias in self.columns_lower: actual_col = self.columns_lower[alias] self._mapped[standard_name] = actual_col return actual_col - + if required: raise ValueError( f"Could not find column for '{standard_name}'. " f"Expected one of: {aliases[:5]}... " f"Found columns: {sorted(self.df.columns)}" ) - + return None - - def get_columns(self, standard_names: List[str], required: bool = True) -> Dict[str, Optional[str]]: + + def get_columns( + self, standard_names: List[str], required: bool = True + ) -> Dict[str, Optional[str]]: """ Find multiple columns at once. - + Args: standard_names: List of standard column names required: If True, raise ValueError for any missing required columns - + Returns: Dict mapping standard names to actual column names """ result = {} missing = [] - + for name in standard_names: try: result[name] = self.get_column(name, required=required) except ValueError: missing.append(name) result[name] = None - + if missing and required: - aliases_str = "\n".join([ - f" - {name}: {', '.join(self.ALIASES.get(name, [])[:5])}" - for name in missing - ]) + aliases_str = "\n".join( + [ + f" - {name}: {', '.join(self.ALIASES.get(name, [])[:5])}" + for name in missing + ] + ) raise ValueError( f"Missing required columns: {missing}\n" f"Expected aliases:\n{aliases_str}\n" f"Found columns: {sorted(self.df.columns)}" ) - + return result - - def rename_columns(self, standard_names: List[str], inplace: bool = False) -> pd.DataFrame: + + def rename_columns( + self, standard_names: List[str], inplace: bool = False + ) -> pd.DataFrame: """ Rename DataFrame columns to standard names. - + Args: standard_names: List of standard column names to rename inplace: If True, modify the DataFrame in place - + Returns: DataFrame with renamed columns """ @@ -191,7 +284,7 @@ def rename_columns(self, standard_names: List[str], inplace: bool = False) -> pd actual_col = self.get_column(std_name, required=False) if actual_col: column_mapping[actual_col] = std_name - + if inplace: self.df.rename(columns=column_mapping, inplace=True) return self.df @@ -202,25 +295,25 @@ def rename_columns(self, standard_names: List[str], inplace: bool = False) -> pd def standardize_gene_loading(df: pd.DataFrame) -> pd.DataFrame: """ Standardize a gene loading DataFrame to have columns: Name, Score, program_id. - + Args: df: Gene loading DataFrame - + Returns: DataFrame with standardized column names """ mapper = ColumnMapper(df) - + # Get the required columns - cols = mapper.get_columns(['gene', 'score', 'program_id'], required=True) - + cols = mapper.get_columns(["gene", "score", "program_id"], required=True) + # Rename to standard names rename_map = { - cols['gene']: 'Name', - cols['score']: 'Score', - cols['program_id']: 'program_id' + cols["gene"]: "Name", + cols["score"]: "Score", + cols["program_id"]: "program_id", } - + return df.rename(columns=rename_map) @@ -228,26 +321,28 @@ def standardize_celltype_enrichment(df: pd.DataFrame) -> pd.DataFrame: """ Standardize a cell-type enrichment DataFrame. Expected output columns: cell_type, program, log2_fc_in_vs_out, fdr - + Args: df: Cell-type enrichment DataFrame - + Returns: DataFrame with standardized column names """ mapper = ColumnMapper(df) - + # Get the required columns - cols = mapper.get_columns(['cell_type', 'program_id', 'log2_fc', 'fdr'], required=True) - + cols = mapper.get_columns( + ["cell_type", "program_id", "log2_fc", "fdr"], required=True + ) + # Rename to standard names rename_map = { - cols['cell_type']: 'cell_type', - cols['program_id']: 'program', - cols['log2_fc']: 'log2_fc_in_vs_out', - cols['fdr']: 'fdr' + cols["cell_type"]: "cell_type", + cols["program_id"]: "program", + cols["log2_fc"]: "log2_fc_in_vs_out", + cols["fdr"]: "fdr", } - + return df.rename(columns=rename_map) @@ -266,9 +361,9 @@ def extract_program_id(value: object) -> Optional[int]: pass patterns = [ - r'^(?:program|topic|p|x)_(\d+)$', - r'^(?:program|topic|p|x)(\d+)$', - r'^(\d+)$', + r"^(?:program|topic|p|x)_(\d+)$", + r"^(?:program|topic|p|x)(\d+)$", + r"^(\d+)$", ] for pattern in patterns: match = re.match(pattern, val_str, re.IGNORECASE) @@ -283,13 +378,13 @@ def _coerce_boolean_series(series: pd.Series) -> pd.Series: return series.fillna(False) text = series.astype(str).str.strip().str.lower() - truthy = {'true', 't', '1', 'yes', 'y'} - falsy = {'false', 'f', '0', 'no', 'n', 'nan', 'none', ''} + truthy = {"true", "t", "1", "yes", "y"} + falsy = {"false", "f", "0", "no", "n", "nan", "none", ""} mapped = pd.Series(pd.NA, index=series.index, dtype="boolean") mapped[text.isin(truthy)] = True mapped[text.isin(falsy)] = False - numeric = pd.to_numeric(series, errors='coerce') + numeric = pd.to_numeric(series, errors="coerce") numeric_bool = numeric.notna() & numeric.ne(0) mapped = mapped.where(mapped.notna(), numeric_bool) return mapped.fillna(False).astype(bool) @@ -314,12 +409,12 @@ def standardize_regulator_results( """ mapper = ColumnMapper(df) cols = mapper.get_columns( - ['program_id', 'regulator_gene', 'log2_fc'], + ["program_id", "regulator_gene", "log2_fc"], required=True, ) - raw_p_col = mapper.get_column('raw_p_value', required=False) - adj_p_col = mapper.get_column('adj_p_value', required=False) - significant_col = mapper.get_column('significant', required=False) + raw_p_col = mapper.get_column("raw_p_value", required=False) + adj_p_col = mapper.get_column("adj_p_value", required=False) + significant_col = mapper.get_column("significant", required=False) if raw_p_col is None and adj_p_col is None: raise ValueError( @@ -329,10 +424,10 @@ def standardize_regulator_results( standardized = df.copy() - program_ids = standardized[cols['program_id']].apply(extract_program_id) + program_ids = standardized[cols["program_id"]].apply(extract_program_id) if program_ids.isna().any(): bad_values = ( - standardized.loc[program_ids.isna(), cols['program_id']] + standardized.loc[program_ids.isna(), cols["program_id"]] .astype(str) .drop_duplicates() .head(5) @@ -343,35 +438,43 @@ def standardize_regulator_results( f"'{cols['program_id']}'. Example values: {bad_values}" ) - standardized['program_id'] = program_ids.astype(int) - standardized['grna_target'] = standardized[cols['regulator_gene']].astype(str) - standardized['log_2_fold_change'] = pd.to_numeric( - standardized[cols['log2_fc']], errors='coerce' + standardized["program_id"] = program_ids.astype(int) + standardized["grna_target"] = standardized[cols["regulator_gene"]].astype(str) + standardized["log_2_fold_change"] = pd.to_numeric( + standardized[cols["log2_fc"]], errors="coerce" ) - if standardized['log_2_fold_change'].isna().any(): + if standardized["log_2_fold_change"].isna().any(): raise ValueError( f"Could not parse numeric log2 fold changes from '{cols['log2_fc']}'." ) if raw_p_col is not None: - standardized['p_value'] = pd.to_numeric(standardized[raw_p_col], errors='coerce') + standardized["p_value"] = pd.to_numeric( + standardized[raw_p_col], errors="coerce" + ) else: - standardized['p_value'] = pd.to_numeric(standardized[adj_p_col], errors='coerce') # type: ignore[index] + standardized["p_value"] = pd.to_numeric( + standardized[adj_p_col], errors="coerce" + ) # type: ignore[index] - if standardized['p_value'].isna().all(): + if standardized["p_value"].isna().all(): raise ValueError("Regulator p-values could not be parsed as numeric values.") if adj_p_col is not None: - standardized['adj_p_value'] = pd.to_numeric( - standardized[adj_p_col], errors='coerce' + standardized["adj_p_value"] = pd.to_numeric( + standardized[adj_p_col], errors="coerce" ) if significant_col is not None: - standardized['significant'] = _coerce_boolean_series(standardized[significant_col]) + standardized["significant"] = _coerce_boolean_series( + standardized[significant_col] + ) elif adj_p_col is not None: - standardized['significant'] = standardized['adj_p_value'].le(significance_threshold) + standardized["significant"] = standardized["adj_p_value"].le( + significance_threshold + ) else: - standardized['significant'] = standardized['p_value'].le(significance_threshold) + standardized["significant"] = standardized["p_value"].le(significance_threshold) return standardized diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/compile_regulators.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/compile_regulators.py index 119cf75..e1d9afd 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/compile_regulators.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/compile_regulators.py @@ -4,7 +4,8 @@ Given a directory of per-timepoint regulator files, merges them into one combined file with a 'day' column and standardised column names. """ -#%% + +# %% import argparse import sys from pathlib import Path @@ -74,16 +75,21 @@ def main() -> int: ) parser.add_argument("data_dir", type=Path, help="Directory with per-day files") parser.add_argument( - "--days", nargs="+", default=["D0", "D1", "D2", "D3"], + "--days", + nargs="+", + default=["D0", "D1", "D2", "D3"], help="Time point labels (default: D0 D1 D2 D3)", ) parser.add_argument( - "--pattern", default="50_CRT_{day}.txt", + "--pattern", + default="50_CRT_{day}.txt", help="Filename pattern with {day} placeholder", ) parser.add_argument("-o", "--output", type=Path, default=None) parser.add_argument( - "--rename", nargs="*", default=["target_name=target_gene_name", "approx_log2FC=log2fc"], + "--rename", + nargs="*", + default=["target_name=target_gene_name", "approx_log2FC=log2fc"], help="Column renames as key=value pairs", ) args = parser.parse_args() diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/convert_spectra_to_csv.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/convert_spectra_to_csv.py index 4bc7562..3101997 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/convert_spectra_to_csv.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/convert_spectra_to_csv.py @@ -9,7 +9,8 @@ By default, only the top-N genes per program (by score) are kept. Use --n-top 0 to keep all genes. """ -#%% + +# %% import argparse import sys from pathlib import Path @@ -26,7 +27,8 @@ def main() -> int: help="Path to gene_spectra_score .txt (TSV) file", ) parser.add_argument( - "-o", "--output", + "-o", + "--output", help="Output CSV path (default: _top.csv in same dir)", ) parser.add_argument( @@ -48,9 +50,7 @@ def main() -> int: # Melt to long format: RowID (program), Name (gene), Score (value) df.index.name = "RowID" - long = df.reset_index().melt( - id_vars="RowID", var_name="Name", value_name="Score" - ) + long = df.reset_index().melt(id_vars="RowID", var_name="Name", value_name="Score") # Shift RowID to 0-based indexing long["RowID"] = long["RowID"] - 1 @@ -65,9 +65,9 @@ def main() -> int: print(f"Kept top {args.n_top} genes per program") # Sort: RowID ascending, Score descending within each program - long = long.sort_values( - ["RowID", "Score"], ascending=[True, False] - ).reset_index(drop=True) + long = long.sort_values(["RowID", "Score"], ascending=[True, False]).reset_index( + drop=True + ) # Reorder columns to match expected format: Name, Score, RowID long = long[["Name", "Score", "RowID"]] diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/gene_extraction.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/gene_extraction.py index fc6d9dd..c1350a0 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/gene_extraction.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/gene_extraction.py @@ -7,6 +7,7 @@ - Cell-type enrichment validation and summary generation - Shared config loading / CLI override utilities """ + from __future__ import annotations import json @@ -19,25 +20,29 @@ import numpy as np import pandas as pd -from .column_mapper import ColumnMapper, standardize_gene_loading, standardize_celltype_enrichment +from .column_mapper import ( + ColumnMapper, + standardize_gene_loading, + standardize_celltype_enrichment, +) logger = logging.getLogger(__name__) # ----------------------------- Cell-type summary ----------------------------- # Thresholds for categorizing cell-type enrichment (FDR < 0.05 required) CELLTYPE_THRESHOLDS = { - 'highly_cell_type_specific': {'log2fc_min': 3.0, 'log2fc_max': None}, - 'moderately_enriched': {'log2fc_min': 1.5, 'log2fc_max': 3.0}, - 'weakly_enriched': {'log2fc_min': 0.5, 'log2fc_max': 1.5}, - 'depleted': {'log2fc_min': None, 'log2fc_max': -0.5}, + "highly_cell_type_specific": {"log2fc_min": 3.0, "log2fc_max": None}, + "moderately_enriched": {"log2fc_min": 1.5, "log2fc_max": 3.0}, + "weakly_enriched": {"log2fc_min": 0.5, "log2fc_max": 1.5}, + "depleted": {"log2fc_min": None, "log2fc_max": -0.5}, } # Column order for cell-type summary output CELLTYPE_CATEGORIES = [ - 'highly_cell_type_specific', - 'moderately_enriched', - 'weakly_enriched', - 'depleted', + "highly_cell_type_specific", + "moderately_enriched", + "weakly_enriched", + "depleted", ] @@ -79,9 +84,9 @@ def extract_program_id(value: object) -> Optional[int]: # Try to extract digits from common patterns # Patterns: Program_X, Topic_X, P_X, X_X (regulator format), program_X, topic_X, p_X patterns = [ - r'^(?:program|topic|p|x)_(\d+)$', # program_1, topic_1, p_1, x_1 - r'^(?:program|topic|p|x)(\d+)$', # program1, topic1, p1, x1 - r'^(\d+)$', # just the number + r"^(?:program|topic|p|x)_(\d+)$", # program_1, topic_1, p_1, x_1 + r"^(?:program|topic|p|x)(\d+)$", # program1, topic1, p1, x1 + r"^(\d+)$", # just the number ] for pattern in patterns: @@ -112,7 +117,7 @@ def validate_celltype_enrichment(df: pd.DataFrame, file_path: Path) -> bool: has_critical_errors = False # 1. Check required columns - required_cols = {'cell_type', 'program', 'log2_fc_in_vs_out', 'fdr'} + required_cols = {"cell_type", "program", "log2_fc_in_vs_out", "fdr"} missing = required_cols - set(df.columns) if missing: warnings_list.append(f"Missing required columns: {sorted(missing)}") @@ -134,32 +139,42 @@ def validate_celltype_enrichment(df: pd.DataFrame, file_path: Path) -> bool: # 3. Validate 'program' column - flexible format checking df_temp = df.copy() - df_temp['program_id_parsed'] = df_temp['program'].apply(extract_program_id) - unparseable = df_temp[df_temp['program_id_parsed'].isna()] + df_temp["program_id_parsed"] = df_temp["program"].apply(extract_program_id) + unparseable = df_temp[df_temp["program_id_parsed"].isna()] if not unparseable.empty: - sample_invalid = unparseable['program'].head(5).tolist() - warnings_list.append(f"Could not parse {len(unparseable)} program identifiers: {sample_invalid}") - warnings_list.append(f" Supported formats: Program_X, program_X, Topic_X, topic_X, P_X, p_X, X_X (regulator), ProgramX, TopicX, X") + sample_invalid = unparseable["program"].head(5).tolist() + warnings_list.append( + f"Could not parse {len(unparseable)} program identifiers: {sample_invalid}" + ) + warnings_list.append( + f" Supported formats: Program_X, program_X, Topic_X, topic_X, P_X, p_X, X_X (regulator), ProgramX, TopicX, X" + ) # 4. Validate 'log2_fc_in_vs_out' is numeric try: - fc_numeric = pd.to_numeric(df['log2_fc_in_vs_out'], errors='coerce') + fc_numeric = pd.to_numeric(df["log2_fc_in_vs_out"], errors="coerce") non_numeric_fc = df[fc_numeric.isna()] if not non_numeric_fc.empty: - sample_bad = non_numeric_fc['log2_fc_in_vs_out'].head(5).tolist() - warnings_list.append(f"Non-numeric values in 'log2_fc_in_vs_out': {sample_bad}") - warnings_list.append(f" Found {len(non_numeric_fc)} non-numeric log2FC values (will be ignored)") + sample_bad = non_numeric_fc["log2_fc_in_vs_out"].head(5).tolist() + warnings_list.append( + f"Non-numeric values in 'log2_fc_in_vs_out': {sample_bad}" + ) + warnings_list.append( + f" Found {len(non_numeric_fc)} non-numeric log2FC values (will be ignored)" + ) except Exception as e: warnings_list.append(f"Failed to validate 'log2_fc_in_vs_out' column: {e}") # 5. Validate 'fdr' is numeric and in valid range [0, 1] try: - fdr_numeric = pd.to_numeric(df['fdr'], errors='coerce') + fdr_numeric = pd.to_numeric(df["fdr"], errors="coerce") non_numeric_fdr = df[fdr_numeric.isna()] if not non_numeric_fdr.empty: - sample_bad = non_numeric_fdr['fdr'].head(5).tolist() + sample_bad = non_numeric_fdr["fdr"].head(5).tolist() warnings_list.append(f"Non-numeric values in 'fdr': {sample_bad}") - warnings_list.append(f" Found {len(non_numeric_fdr)} non-numeric FDR values (will be ignored)") + warnings_list.append( + f" Found {len(non_numeric_fdr)} non-numeric FDR values (will be ignored)" + ) # Check FDR range (should be 0-1 for proper FDR values) valid_fdr = fdr_numeric.dropna() @@ -167,19 +182,27 @@ def validate_celltype_enrichment(df: pd.DataFrame, file_path: Path) -> bool: out_of_range_mask = (valid_fdr < 0) | (valid_fdr > 1) if out_of_range_mask.any(): out_of_range_vals = valid_fdr[out_of_range_mask].head(5).tolist() - warnings_list.append(f"FDR values out of range [0, 1]: {out_of_range_vals}") - warnings_list.append(f" Found {out_of_range_mask.sum()} out-of-range FDR values") + warnings_list.append( + f"FDR values out of range [0, 1]: {out_of_range_vals}" + ) + warnings_list.append( + f" Found {out_of_range_mask.sum()} out-of-range FDR values" + ) except Exception as e: warnings_list.append(f"Failed to validate 'fdr' column: {e}") # 6. Check for completely empty cell_type values - empty_celltypes = df[df['cell_type'].isna() | (df['cell_type'].astype(str).str.strip() == '')] + empty_celltypes = df[ + df["cell_type"].isna() | (df["cell_type"].astype(str).str.strip() == "") + ] if not empty_celltypes.empty: - warnings_list.append(f"Found {len(empty_celltypes)} rows with empty 'cell_type' values (will be ignored)") + warnings_list.append( + f"Found {len(empty_celltypes)} rows with empty 'cell_type' values (will be ignored)" + ) # 7. Log summary statistics (informational) - n_programs = df['program'].nunique() - n_celltypes = df['cell_type'].nunique() + n_programs = df["program"].nunique() + n_celltypes = df["cell_type"].nunique() logger.info(f"Cell-type enrichment file summary: {file_path}") logger.info(f" Total rows: {len(df)}") logger.info(f" Unique programs: {n_programs}") @@ -187,13 +210,19 @@ def validate_celltype_enrichment(df: pd.DataFrame, file_path: Path) -> bool: # Check for reasonable data coverage (warnings only) if n_programs < 10: - warnings_list.append(f"Very few programs found ({n_programs}). Expected 50-100 for typical cNMF results.") + warnings_list.append( + f"Very few programs found ({n_programs}). Expected 50-100 for typical cNMF results." + ) if n_celltypes < 3: - warnings_list.append(f"Very few cell types found ({n_celltypes}). Expected multiple cell types.") + warnings_list.append( + f"Very few cell types found ({n_celltypes}). Expected multiple cell types." + ) # Log all warnings if warnings_list: - logger.warning(f"Cell-type enrichment validation found {len(warnings_list)} issue(s):") + logger.warning( + f"Cell-type enrichment validation found {len(warnings_list)} issue(s):" + ) for w in warnings_list: logger.warning(f" {w}") return False @@ -245,45 +274,49 @@ def generate_celltype_summary( validate_celltype_enrichment(df, enrichment_file) # Validate required columns (critical check) - required_cols = {'cell_type', 'program', 'log2_fc_in_vs_out', 'fdr'} + required_cols = {"cell_type", "program", "log2_fc_in_vs_out", "fdr"} missing = required_cols - set(df.columns) if missing: logger.error(f"Enrichment file missing required columns: {missing}") logger.error(f" Found columns: {sorted(df.columns)}") - logger.error(f" Cannot proceed without required columns. Please check your input file.") + logger.error( + f" Cannot proceed without required columns. Please check your input file." + ) raise ValueError(f"Enrichment file missing required columns: {missing}") # Extract program ID using flexible parser (handles Program_X, program_X, Topic_X, X, etc.) - df['program_id'] = df['program'].apply(extract_program_id) + df["program_id"] = df["program"].apply(extract_program_id) # Check for programs that couldn't be parsed - unparsed = df[df['program_id'].isna()] + unparsed = df[df["program_id"].isna()] if not unparsed.empty: - logger.warning(f"Could not parse program IDs for {len(unparsed)} rows (will be excluded):") + logger.warning( + f"Could not parse program IDs for {len(unparsed)} rows (will be excluded):" + ) logger.warning(f" Sample values: {unparsed['program'].head(5).tolist()}") - df = df[df['program_id'].notna()].copy() + df = df[df["program_id"].notna()].copy() # Ensure program_id is integer type - df['program_id'] = df['program_id'].astype(int) + df["program_id"] = df["program_id"].astype(int) # Filter to requested topics if specified if topics: - df = df[df['program_id'].isin(list(topics))].copy() + df = df[df["program_id"].isin(list(topics))].copy() logger.info(f"Filtered to {len(df)} rows for topics: {sorted(topics)}") # Filter to significant results - df_sig = df[df['fdr'] < fdr_threshold].copy() + df_sig = df[df["fdr"] < fdr_threshold].copy() logger.info(f"Found {len(df_sig)} significant rows (FDR < {fdr_threshold})") # Use cell_type values as-is (assume already has correct names) - df_sig['cell_type_display'] = df_sig['cell_type'] + df_sig["cell_type_display"] = df_sig["cell_type"] # Categorize each row by log2FC thresholds def categorize_row(row: pd.Series) -> Optional[str]: - log2fc = row['log2_fc_in_vs_out'] + log2fc = row["log2_fc_in_vs_out"] for cat, bounds in thresholds.items(): - min_val = bounds.get('log2fc_min') - max_val = bounds.get('log2fc_max') + min_val = bounds.get("log2fc_min") + max_val = bounds.get("log2fc_max") # Check if log2fc falls in this category if min_val is not None and max_val is not None: if min_val <= log2fc < max_val: @@ -296,28 +329,28 @@ def categorize_row(row: pd.Series) -> Optional[str]: return cat return None - df_sig['category'] = df_sig.apply(categorize_row, axis=1) - df_categorized = df_sig.dropna(subset=['category']) + df_sig["category"] = df_sig.apply(categorize_row, axis=1) + df_categorized = df_sig.dropna(subset=["category"]) logger.info(f"Categorized {len(df_categorized)} rows into enrichment categories") # Build summary: for each program, collect cell types per category - all_programs = sorted(df['program_id'].unique()) + all_programs = sorted(df["program_id"].unique()) records = [] for pid in all_programs: - program_data = df_categorized[df_categorized['program_id'] == pid] - row = {'program': f'Program_{pid}'} + program_data = df_categorized[df_categorized["program_id"] == pid] + row = {"program": f"Program_{pid}"} for cat in CELLTYPE_CATEGORIES: - cat_data = program_data[program_data['category'] == cat] - cell_types = sorted(cat_data['cell_type_display'].unique()) - row[cat] = '|'.join(cell_types) if cell_types else '' + cat_data = program_data[program_data["category"] == cat] + cell_types = sorted(cat_data["cell_type_display"].unique()) + row[cat] = "|".join(cell_types) if cell_types else "" records.append(row) # Create output DataFrame summary_df = pd.DataFrame(records) - summary_df = summary_df[['program'] + CELLTYPE_CATEGORIES] + summary_df = summary_df[["program"] + CELLTYPE_CATEGORIES] # Write output ensure_parent_dir(str(output_file)) @@ -329,6 +362,7 @@ def categorize_row(row: pd.Series) -> Optional[str]: # --------------------------- File / path helpers ------------------------------ + def ensure_parent_dir(path_str: str) -> None: path = Path(path_str) if path.parent and not path.parent.exists(): @@ -342,7 +376,7 @@ def resolve_program_id_column(df: pd.DataFrame) -> str: """ mapper = ColumnMapper(df) try: - actual_col = mapper.get_column('program_id', required=True) + actual_col = mapper.get_column("program_id", required=True) return actual_col except ValueError as e: raise ValueError( @@ -360,6 +394,7 @@ def normalize_program_id(value: object) -> object: # --------------------------- Extract top genes (CSV) -------------------------- + def extract_top_genes_by_program( df: pd.DataFrame, n_top: int, id_col: str ) -> Dict[str, List[str]]: @@ -371,9 +406,9 @@ def extract_top_genes_by_program( # Get standardized column names try: - cols = mapper.get_columns(['gene', 'score'], required=True) - gene_col = cols['gene'] - score_col = cols['score'] + cols = mapper.get_columns(["gene", "score"], required=True) + gene_col = cols["gene"] + score_col = cols["score"] except ValueError as e: raise ValueError(f"Missing required columns for gene extraction: {e}") @@ -399,6 +434,7 @@ def extract_top_genes_by_program( # ----------------------------- Uniqueness table ------------------------------- + def default_uniqueness_output(input_path: Path) -> Path: suffix = input_path.suffix or ".csv" stem = input_path.stem @@ -415,9 +451,9 @@ def build_uniqueness_table(df: pd.DataFrame, id_col: str) -> pd.DataFrame: # Get standardized column names try: - cols = mapper.get_columns(['gene', 'score'], required=True) - gene_col = cols['gene'] - score_col = cols['score'] + cols = mapper.get_columns(["gene", "score"], required=True) + gene_col = cols["gene"] + score_col = cols["score"] except ValueError as e: raise ValueError(f"Missing required columns for uniqueness computation: {e}") @@ -427,11 +463,7 @@ def build_uniqueness_table(df: pd.DataFrame, id_col: str) -> pd.DataFrame: work = df.copy() # Standardize column names to Name, Score, program_id - rename_map = { - gene_col: 'Name', - score_col: 'Score', - id_col: 'program_id' - } + rename_map = {gene_col: "Name", score_col: "Score", id_col: "program_id"} work = work.rename(columns=rename_map) if "UniquenessScore" not in work.columns or work["UniquenessScore"].isna().all(): @@ -452,7 +484,9 @@ def build_uniqueness_table(df: pd.DataFrame, id_col: str) -> pd.DataFrame: columns = ["Name", "Score", "program_id", "UniquenessScore"] out_df = work[columns].copy() - out_df.dropna(subset=["Name", "Score", "program_id", "UniquenessScore"], inplace=True) + out_df.dropna( + subset=["Name", "Score", "program_id", "UniquenessScore"], inplace=True + ) return out_df @@ -464,12 +498,16 @@ def build_overview_long_table( for program_id_str, genes in top_map.items(): program_id = normalize_program_id(program_id_str) if program_id not in sub_indexed_cache: - sub_indexed_cache[program_id] = ( - df[df[id_col] == program_id].set_index("Name") + sub_indexed_cache[program_id] = df[df[id_col] == program_id].set_index( + "Name" ) sub_idx = sub_indexed_cache[program_id] for rank, gene in enumerate(genes, start=1): - score = float(sub_idx.loc[gene, "Score"]) if gene in sub_idx.index else float("nan") + score = ( + float(sub_idx.loc[gene, "Score"]) + if gene in sub_idx.index + else float("nan") + ) records.append( {"program_id": program_id, "rank": rank, "gene": gene, "score": score} ) @@ -535,13 +573,23 @@ def parse_topics(value: Optional[object]) -> Optional[Set[int]]: def apply_test_mode( args: object, config: Dict[str, Any], cli_overrides: Set[str] ) -> object: - test_cfg = config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + test_cfg = ( + config.get("test", {}) if isinstance(config.get("test", {}), dict) else {} + ) enabled = bool(test_cfg.get("enabled") or config.get("test_mode")) if not enabled: return args - topics = test_cfg.get("topics") or test_cfg.get("programs") or config.get("test_programs") - if hasattr(args, "topics") and "topics" not in cli_overrides and not getattr(args, "topics", None): + topics = ( + test_cfg.get("topics") + or test_cfg.get("programs") + or config.get("test_programs") + ) + if ( + hasattr(args, "topics") + and "topics" not in cli_overrides + and not getattr(args, "topics", None) + ): if topics is not None: args.topics = topics # type: ignore[attr-defined] return args @@ -552,9 +600,15 @@ def apply_config_overrides( config: Dict[str, Any], cli_overrides: Set[str], ) -> object: - steps_cfg = config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + steps_cfg = ( + config.get("steps", {}) if isinstance(config.get("steps", {}), dict) else {} + ) step_cfg = steps_cfg.get("string_enrichment", {}) - if isinstance(step_cfg, dict) and hasattr(args, "command") and getattr(args, "command") in step_cfg: + if ( + isinstance(step_cfg, dict) + and hasattr(args, "command") + and getattr(args, "command") in step_cfg + ): step_cfg = step_cfg.get(getattr(args, "command"), {}) if not isinstance(step_cfg, dict): return args @@ -572,12 +626,18 @@ def apply_default_paths(args: object) -> object: n_top = getattr(args, "n_top", None) or 100 if hasattr(args, "json_out") and not getattr(args, "json_out", None): - args.json_out = str(DEFAULT_ENRICH_DIR / DEFAULT_GENES_JSON_TEMPLATE.format(n_top=n_top)) # type: ignore[attr-defined] + args.json_out = str( + DEFAULT_ENRICH_DIR / DEFAULT_GENES_JSON_TEMPLATE.format(n_top=n_top) + ) # type: ignore[attr-defined] if hasattr(args, "csv_out") and not getattr(args, "csv_out", None): - args.csv_out = str(DEFAULT_ENRICH_DIR / DEFAULT_GENES_OVERVIEW_TEMPLATE.format(n_top=n_top)) # type: ignore[attr-defined] + args.csv_out = str( + DEFAULT_ENRICH_DIR / DEFAULT_GENES_OVERVIEW_TEMPLATE.format(n_top=n_top) + ) # type: ignore[attr-defined] if hasattr(args, "out_csv_full") and not getattr(args, "out_csv_full", None): args.out_csv_full = str(DEFAULT_ENRICH_DIR / DEFAULT_STRING_FULL) # type: ignore[attr-defined] - if hasattr(args, "out_csv_filtered") and not getattr(args, "out_csv_filtered", None): + if hasattr(args, "out_csv_filtered") and not getattr( + args, "out_csv_filtered", None + ): args.out_csv_filtered = str(DEFAULT_ENRICH_DIR / DEFAULT_STRING_FILTERED) # type: ignore[attr-defined] return args diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/harmonizome_api.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/harmonizome_api.py index 9a4cca7..44c8e6c 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/harmonizome_api.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/harmonizome_api.py @@ -43,11 +43,11 @@ def parse_summary_json(json_str: str) -> Tuple[str, List[Dict]]: """Parse Harmonizome summary JSON structure into plain text and references. - - The JSON has a tree structure with types: root, p (paragraph), t (text), + + The JSON has a tree structure with types: root, p (paragraph), t (text), fg (footnote group), fg_f (footnote ref), rg (reference group), r (reference), b (bold), i (italic), a (link). - + Returns: Tuple of (plain_text_summary, list_of_references) """ @@ -56,24 +56,26 @@ def parse_summary_json(json_str: str) -> Tuple[str, List[Dict]]: except json.JSONDecodeError as e: logger.warning("Failed to parse summary JSON: %s", e) return "", [] - + paragraphs = [] references = [] - + def extract_text(node: dict) -> str: """Recursively extract text from a node.""" if not isinstance(node, dict): return "" - + node_type = node.get("type", "") - + if node_type == "t": # Text node return node.get("text", "").strip() elif node_type in ("p", "root", "b", "i", "a"): # Container nodes - extract children children = node.get("children", []) - return " ".join(extract_text(c) for c in children if isinstance(c, dict)).strip() + return " ".join( + extract_text(c) for c in children if isinstance(c, dict) + ).strip() elif node_type == "fg": # Footnote group - skip for plain text return "" @@ -83,14 +85,16 @@ def extract_text(node: dict) -> str: else: # Unknown type - try to extract children children = node.get("children", []) - return " ".join(extract_text(c) for c in children if isinstance(c, dict)).strip() - + return " ".join( + extract_text(c) for c in children if isinstance(c, dict) + ).strip() + def extract_references(node: dict) -> List[Dict]: """Extract reference information from rg (reference group) nodes.""" refs = [] if not isinstance(node, dict): return refs - + if node.get("type") == "rg": for child in node.get("children", []): if child.get("type") == "r": @@ -123,9 +127,9 @@ def extract_references(node: dict) -> List[Dict]: # Recurse into children for child in node.get("children", []): refs.extend(extract_references(child)) - + return refs - + # Process the root node if data.get("type") == "root": for child in data.get("children", []): @@ -133,18 +137,18 @@ def extract_references(node: dict) -> List[Dict]: para_text = extract_text(child) if para_text: # Clean up whitespace - para_text = re.sub(r'\s+', ' ', para_text).strip() + para_text = re.sub(r"\s+", " ", para_text).strip() paragraphs.append(para_text) elif child.get("type") == "rg": references.extend(extract_references(child)) - + summary_text = "\n\n".join(paragraphs) return summary_text, references class HarmonizomeClient: """Client for Harmonizome gene summaries. - + By default, fetches short NCBI-like descriptions from the API. Optionally fetches full literature-based summaries from HTML pages when use_full_summaries=True. @@ -160,7 +164,7 @@ def __init__( use_full_summaries: bool = False, ) -> None: """Initialize the client. - + Args: base_url: Harmonizome base URL session: Optional requests session for connection pooling @@ -178,9 +182,11 @@ def __init__( self.include_references = include_references self.use_full_summaries = use_full_summaries - def get_gene_summary_from_html(self, symbol: str) -> Optional[Tuple[str, List[Dict]]]: + def get_gene_summary_from_html( + self, symbol: str + ) -> Optional[Tuple[str, List[Dict]]]: """Fetch full gene summary from HTML page. - + Returns: Tuple of (summary_text, references) or None if failed """ @@ -190,40 +196,40 @@ def get_gene_summary_from_html(self, symbol: str) -> Optional[Tuple[str, List[Di except requests.RequestException as exc: logger.warning("Harmonizome HTML request failed for %s: %s", symbol, exc) return None - + if resp.status_code != 200: - logger.warning("Harmonizome HTML returned %s for %s", resp.status_code, symbol) + logger.warning( + "Harmonizome HTML returned %s for %s", resp.status_code, symbol + ) return None - + # Extract summary JSON from HTML # Look for:
...
match = re.search( - r'
]*>\s*(.*?)\s*
', - resp.text, - re.DOTALL + r'
]*>\s*(.*?)\s*
', resp.text, re.DOTALL ) - + if not match: logger.debug("No summary-content div found for %s", symbol) return None - + # Decode HTML entities (" -> ") json_str = html.unescape(match.group(1)) - + summary_text, references = parse_summary_json(json_str) - + if not summary_text: return None - + return summary_text, references @staticmethod def _clean_source_annotations(text: str) -> str: """Remove source annotations like [provided by RefSeq, Jun 2016] to save tokens.""" # Pattern matches: [provided by X, Month Year] or [provided by X] - cleaned = re.sub(r'\s*\[provided by [^\]]+\]', '', text) + cleaned = re.sub(r"\s*\[provided by [^\]]+\]", "", text) # Clean up any double spaces - cleaned = re.sub(r' +', ' ', cleaned) + cleaned = re.sub(r" +", " ", cleaned) return cleaned.strip() def get_gene_summary_from_api(self, symbol: str) -> Optional[str]: @@ -236,7 +242,9 @@ def get_gene_summary_from_api(self, symbol: str) -> Optional[str]: return None if resp.status_code != 200: - logger.warning("Harmonizome API returned %s for %s", resp.status_code, symbol) + logger.warning( + "Harmonizome API returned %s for %s", resp.status_code, symbol + ) return None try: @@ -252,10 +260,10 @@ def get_gene_summary_from_api(self, symbol: str) -> Optional[str]: def get_gene_summary(self, symbol: str) -> Optional[str]: """Fetch gene summary. - + If use_full_summaries=True: tries HTML first for full literature-based summary, falls back to API if unavailable. - + If use_full_summaries=False (default): uses short API description only. """ if self.use_full_summaries: @@ -267,35 +275,37 @@ def get_gene_summary(self, symbol: str) -> Optional[str]: # Append reference PMIDs pmids = [r.get("pmid") for r in references if r.get("pmid")] if pmids: - summary_text += f"\n\nReferences: {', '.join(f'PMID:{p}' for p in pmids)}" + summary_text += ( + f"\n\nReferences: {', '.join(f'PMID:{p}' for p in pmids)}" + ) return summary_text - + # Fall back to API logger.debug("Falling back to API for %s", symbol) - + return self.get_gene_summary_from_api(symbol) def get_gene_summaries(self, symbols: Iterable[str]) -> Dict[str, str]: """Fetch summaries for multiple genes. - + Args: symbols: Iterable of gene symbols - + Returns: Dict mapping symbol to summary text """ results: Dict[str, str] = {} symbols_list = list(symbols) - + for i, symbol in enumerate(symbols_list): summary = self.get_gene_summary(symbol) if summary: results[symbol] = summary else: logger.warning("No summary found for %s", symbol) - + # Rate limiting if self.sleep_seconds and i < len(symbols_list) - 1: time.sleep(self.sleep_seconds) - + return results diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/html_report.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/html_report.py index 744fe76..031c7f8 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/html_report.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/html_report.py @@ -16,6 +16,7 @@ generate_report — build the full HTML report from data files generate_design_a_html — render the Design A HTML template """ + from __future__ import annotations import json @@ -30,39 +31,54 @@ from .column_mapper import standardize_regulator_results PRIORITY_GENES = [ - "Hif1a", "Epas1", "Arnt", "Fzd4", "Fzd6", "Idh2", "Mdh2", "Ogdh", - "Hsp90ab1", "Hspa5", "Creb3l2", "Fkbp8", "Kdr", "Egln1", "Egln2", "Foxo1", "Foxo3" + "Hif1a", + "Epas1", + "Arnt", + "Fzd4", + "Fzd6", + "Idh2", + "Mdh2", + "Ogdh", + "Hsp90ab1", + "Hspa5", + "Creb3l2", + "Fkbp8", + "Kdr", + "Egln1", + "Egln2", + "Foxo1", + "Foxo3", ] def extract_program_stats(annotation_md: str) -> dict: """Extract stats from annotation markdown.""" - stats = {'top_loading': '', 'unique': '', 'celltype': ''} + stats = {"top_loading": "", "unique": "", "celltype": ""} # Extract top-loading genes - match = re.search(r'\*\*Top-loading genes?:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Top-loading genes?:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['top_loading'] = match.group(1).strip() + stats["top_loading"] = match.group(1).strip() # Extract unique genes - match = re.search(r'\*\*Unique genes?:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Unique genes?:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['unique'] = match.group(1).strip() + stats["unique"] = match.group(1).strip() # Extract cell-type enrichment - match = re.search(r'\*\*Cell-type enrichment:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Cell-type enrichment:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['celltype'] = match.group(1).strip() + stats["celltype"] = match.group(1).strip() # Extract brief summary - match = re.search(r'\*\*Brief Summary:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Brief Summary:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['summary'] = match.group(1).strip() + stats["summary"] = match.group(1).strip() # Extract program label - match = re.search(r'\*\*Program label:\*\*\s*([^\n]+)', annotation_md, re.I) + match = re.search(r"\*\*Program label:\*\*\s*([^\n]+)", annotation_md, re.I) if match: - stats['label'] = match.group(1).strip() + stats["label"] = match.group(1).strip() return stats @@ -91,20 +107,25 @@ def generate_report( volcano_by_program = {} if volcano_df is not None: pvalue_col = ( - 'adj_p_value' - if 'adj_p_value' in volcano_df.columns and not volcano_df['adj_p_value'].isna().all() - else 'p_value' + "adj_p_value" + if "adj_p_value" in volcano_df.columns + and not volcano_df["adj_p_value"].isna().all() + else "p_value" + ) + volcano_df["neg_log10_pvalue"] = -np.log10( + volcano_df[pvalue_col].replace(0, 1e-300) + ) + volcano_df.loc[np.isinf(volcano_df["neg_log10_pvalue"]), "neg_log10_pvalue"] = ( + 300 ) - volcano_df['neg_log10_pvalue'] = -np.log10(volcano_df[pvalue_col].replace(0, 1e-300)) - volcano_df.loc[np.isinf(volcano_df['neg_log10_pvalue']), 'neg_log10_pvalue'] = 300 - for tid, group in volcano_df.groupby('program_id'): + for tid, group in volcano_df.groupby("program_id"): volcano_by_program[int(tid)] = [ { - 'g': row['grna_target'], - 'fc': round(row['log_2_fold_change'], 3), - 'p': round(row['neg_log10_pvalue'], 2), - 's': bool(row['significant']) + "g": row["grna_target"], + "fc": round(row["log_2_fold_change"], 3), + "p": round(row["neg_log10_pvalue"], 2), + "s": bool(row["significant"]), } for _, row in group.iterrows() ] @@ -112,12 +133,14 @@ def generate_report( # Build per-program data programs_data = [] for _, row in summary_df.iterrows(): - topic_id = int(row['Topic']) - topic_name = row['Name'] + topic_id = int(row["Topic"]) + topic_name = row["Name"] # Read annotation ann_path = Path(annotations_dir) / f"topic_{topic_id}_annotation.md" - annotation_md = ann_path.read_text(encoding='utf-8') if ann_path.exists() else "" + annotation_md = ( + ann_path.read_text(encoding="utf-8") if ann_path.exists() else "" + ) # Extract stats stats = extract_program_stats(annotation_md) @@ -125,28 +148,32 @@ def generate_report( # Enrichment paths enr_rel = os.path.relpath(enrichment_dir, os.path.dirname(output_html)) - programs_data.append({ - 'id': topic_id, - 'name': topic_name, - 'label': stats.get('label', topic_name), - 'summary': stats.get('summary', ''), - 'top_loading': stats.get('top_loading', ''), - 'unique': stats.get('unique', ''), - 'celltype': stats.get('celltype', ''), - 'annotation_html': markdown(annotation_md, extensions=['tables', 'fenced_code']), - 'annotation_text': annotation_md, # For full-text search - 'kegg_fig': f"{enr_rel}/program_{topic_id}_kegg_enrichment.png", - 'process_fig': f"{enr_rel}/program_{topic_id}_process_enrichment.png", - 'volcano': volcano_by_program.get(topic_id, []), - }) - - generated_on = datetime.now().strftime('%Y-%m-%d %H:%M') + programs_data.append( + { + "id": topic_id, + "name": topic_name, + "label": stats.get("label", topic_name), + "summary": stats.get("summary", ""), + "top_loading": stats.get("top_loading", ""), + "unique": stats.get("unique", ""), + "celltype": stats.get("celltype", ""), + "annotation_html": markdown( + annotation_md, extensions=["tables", "fenced_code"] + ), + "annotation_text": annotation_md, # For full-text search + "kegg_fig": f"{enr_rel}/program_{topic_id}_kegg_enrichment.png", + "process_fig": f"{enr_rel}/program_{topic_id}_process_enrichment.png", + "volcano": volcano_by_program.get(topic_id, []), + } + ) + + generated_on = datetime.now().strftime("%Y-%m-%d %H:%M") # Generate HTML html = generate_design_a_html(programs_data, len(programs_data), generated_on) Path(output_html).parent.mkdir(parents=True, exist_ok=True) - with open(output_html, 'w', encoding='utf-8') as f: + with open(output_html, "w", encoding="utf-8") as f: f.write(html) print(f"Design A (Scientific Minimal): {output_html}") @@ -157,23 +184,29 @@ def generate_design_a_html(programs_data, num_programs, generated_on): """Generate Design A HTML.""" # Build TOC - toc_items = "\n".join([ - f'' - f'{p["id"]}. {p["name"]}' - for p in programs_data - ]) + toc_items = "\n".join( + [ + f'' + f"{p['id']}. {p['name']}" + for p in programs_data + ] + ) # Program selector options - options = "\n".join([ - f'' - for p in programs_data - ]) + options = "\n".join( + [ + f'' + for p in programs_data + ] + ) # Program data JS - programs_js = f"window.PROGRAMS = {json.dumps({p['id']: p for p in programs_data})};\n" + programs_js = ( + f"window.PROGRAMS = {json.dumps({p['id']: p for p in programs_data})};\n" + ) programs_js += f"window.PRIORITY_GENES = {json.dumps(PRIORITY_GENES)};\n" - html = f''' + html = f""" @@ -587,6 +620,6 @@ def generate_design_a_html(programs_data, num_programs, generated_on): }})(); -''' +""" return html diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/literature_mining.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/literature_mining.py index a99fe6d..3a5def7 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/literature_mining.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/literature_mining.py @@ -4,6 +4,7 @@ Library module — no CLI/argparse code. """ + from __future__ import annotations import re @@ -31,6 +32,7 @@ # Private helpers (duplicated from gene_extraction to avoid circular imports) # --------------------------------------------------------------------------- + def _ensure_program_id_column(df: pd.DataFrame) -> pd.DataFrame: if "program_id" in df.columns: return df @@ -45,7 +47,9 @@ def _add_global_uniqueness_scores(df: pd.DataFrame) -> pd.DataFrame: required_cols = {"Name", "Score", "program_id"} missing = required_cols - set(df.columns) if missing: - raise ValueError(f"CSV missing required columns for uniqueness: {sorted(missing)}") + raise ValueError( + f"CSV missing required columns for uniqueness: {sorted(missing)}" + ) updated = df.copy() updated["Score"] = pd.to_numeric(updated["Score"], errors="coerce") updated["program_id"] = pd.to_numeric(updated["program_id"], errors="coerce") @@ -74,22 +78,63 @@ def _ensure_global_uniqueness(df: pd.DataFrame) -> pd.DataFrame: # --------------------------------------------------------------------------- DOMAIN_KEYWORDS = [ - "inflammation", "proliferation", "migration", "adhesion", "metabolism", - "hypoxia", "differentiation", "apoptosis", "cell cycle", "signaling", - "stress response", "development" + "inflammation", + "proliferation", + "migration", + "adhesion", + "metabolism", + "hypoxia", + "differentiation", + "apoptosis", + "cell cycle", + "signaling", + "stress response", + "development", ] INTERACTION_VERBS = [ - "regulates", "induces", "promotes", "inhibits", "suppresses", "activates", - "binds", "phosphorylates", "modulates", "mediates", "targets", "controls", - "decreases", "increases", "blocks", "triggers", "catalyzes" + "regulates", + "induces", + "promotes", + "inhibits", + "suppresses", + "activates", + "binds", + "phosphorylates", + "modulates", + "mediates", + "targets", + "controls", + "decreases", + "increases", + "blocks", + "triggers", + "catalyzes", ] MECHANISTIC_VERBS = [ - "regulates", "induces", "promotes", "inhibits", "suppresses", "activates", - "mediates", "controls", "stimulates", "blocks", "enhances", "reduces", - "increases", "decreases", "triggers", "phosphorylates", "upregulates", - "downregulates", "attenuates", "potentiates", "modulates", "drives" + "regulates", + "induces", + "promotes", + "inhibits", + "suppresses", + "activates", + "mediates", + "controls", + "stimulates", + "blocks", + "enhances", + "reduces", + "increases", + "decreases", + "triggers", + "phosphorylates", + "upregulates", + "downregulates", + "attenuates", + "potentiates", + "modulates", + "drives", ] CACHE_DIR = Path("data/cache") @@ -103,17 +148,22 @@ def _ensure_global_uniqueness(df: pd.DataFrame) -> pd.DataFrame: # Text processing # --------------------------------------------------------------------------- + def split_text_into_sentences(text: str) -> List[str]: if not text: return [] protected = text subs = { - "et al.": "ET_AL_MARKER", "e.g.": "EG_MARKER", "i.e.": "IE_MARKER", - "Fig.": "FIG_MARKER", "Ref.": "REF_MARKER", "vs.": "VS_MARKER" + "et al.": "ET_AL_MARKER", + "e.g.": "EG_MARKER", + "i.e.": "IE_MARKER", + "Fig.": "FIG_MARKER", + "Ref.": "REF_MARKER", + "vs.": "VS_MARKER", } for k, v in subs.items(): protected = protected.replace(k, v) - parts = re.split(r'(?<=[.?!])\s+', protected) + parts = re.split(r"(?<=[.?!])\s+", protected) sentences = [] for p in parts: restored = p @@ -178,12 +228,19 @@ def extract_evidence_sentences( if re.search(rf"\b{re.escape(g)}\b", sent, re.IGNORECASE): found_targets.append(g) normalized_sent = sent.strip(" .") - existing_sents = [s[1].strip(" .") for s in [item for sublist in gene_to_sentences.values() for item in sublist]] + existing_sents = [ + s[1].strip(" .") + for s in [ + item for sublist in gene_to_sentences.values() for item in sublist + ] + ] if normalized_sent in existing_sents: continue score = 0 for m in context_genes: - if m not in found_targets and re.search(rf"\b{re.escape(m)}\b", sent, re.IGNORECASE): + if m not in found_targets and re.search( + rf"\b{re.escape(m)}\b", sent, re.IGNORECASE + ): score += 2 if any(k in sent_lower for k in DOMAIN_KEYWORDS): score += 1 @@ -205,12 +262,16 @@ def extract_evidence_sentences( # BioC parsing # --------------------------------------------------------------------------- + def parse_bioc_abstract(doc: Dict[str, Any]) -> str: text_parts = [] passages = doc.get("passages", []) for p in passages: inf = p.get("infons", {}) - if inf.get("type") in ("title", "abstract") or inf.get("section_type") in ("TITLE", "ABSTRACT"): + if inf.get("type") in ("title", "abstract") or inf.get("section_type") in ( + "TITLE", + "ABSTRACT", + ): text_parts.append(p.get("text", "")) return " ".join(text_parts) @@ -248,6 +309,7 @@ def find_gene_mentions(doc: Dict[str, Any], member_genes: List[str]) -> List[str # Cache # --------------------------------------------------------------------------- + def load_cache() -> Dict[str, Any]: if CACHE_FILE.exists(): try: @@ -266,23 +328,30 @@ def save_cache(cache: Dict[str, Any]): # Gene loading # --------------------------------------------------------------------------- + def load_program_genes( csv_path: Path, top_n_loading: int = 20, top_n_unique: int = 10, - top_n_member: int = 100 + top_n_member: int = 100, ) -> Dict[int, Dict[str, List[str]]]: df = pd.read_csv(csv_path) df = _ensure_program_id_column(df) df = _ensure_global_uniqueness(df) programs = {} for pid, group in df.groupby("program_id"): - sorted_loading = group.sort_values("Score", ascending=False)["Name"].astype(str).tolist() + sorted_loading = ( + group.sort_values("Score", ascending=False)["Name"].astype(str).tolist() + ) top_loading_genes = sorted_loading[:top_n_loading] top_loading_set = set(top_loading_genes) top_unique_genes = [] if "UniquenessScore" in group.columns: - sorted_unique = group.sort_values("UniquenessScore", ascending=False)["Name"].astype(str).tolist() + sorted_unique = ( + group.sort_values("UniquenessScore", ascending=False)["Name"] + .astype(str) + .tolist() + ) for gene in sorted_unique: if gene not in top_loading_set: top_unique_genes.append(gene) @@ -296,7 +365,7 @@ def load_program_genes( "top_loading": top_loading_genes, "top_unique": top_unique_genes, "members": members, - "all_genes": all_genes + "all_genes": all_genes, } return programs @@ -305,6 +374,7 @@ def load_program_genes( # Gene summaries # --------------------------------------------------------------------------- + def resolve_gene_summaries( source: str, programs: Dict[int, Dict[str, List[str]]], @@ -320,7 +390,9 @@ def resolve_gene_summaries( if not all_drivers: return {pid: {} for pid in program_ids} if source == "ncbi": - logger.info("Resolving Entrez IDs for %d unique driver genes...", len(all_drivers)) + logger.info( + "Resolving Entrez IDs for %d unique driver genes...", len(all_drivers) + ) symbol_to_id = ncbi_client.normalize_genes(list(all_drivers)) valid_ids = sorted({gid for gid in symbol_to_id.values() if gid}) logger.info("Fetching summaries for %d gene IDs...", len(valid_ids)) @@ -341,12 +413,19 @@ def resolve_gene_summaries( use_full_summaries=use_full_summaries ) summary_type = "full HTML" if use_full_summaries else "short API" - logger.info("Fetching Harmonizome %s summaries for %d unique driver genes...", - summary_type, len(all_drivers)) + logger.info( + "Fetching Harmonizome %s summaries for %d unique driver genes...", + summary_type, + len(all_drivers), + ) symbol_to_summary = harmonizome_client.get_gene_summaries(sorted(all_drivers)) for pid in program_ids: drivers = programs[pid]["drivers"] - summaries = {sym: symbol_to_summary[sym] for sym in drivers if sym in symbol_to_summary} + summaries = { + sym: symbol_to_summary[sym] + for sym in drivers + if sym in symbol_to_summary + } program_gene_summaries[pid] = summaries return program_gene_summaries raise ValueError(f"Unsupported gene summary source: {source}") @@ -356,6 +435,7 @@ def resolve_gene_summaries( # Regulator loading # --------------------------------------------------------------------------- + def load_regulator_data( csv_path: Path, significance_threshold: float = 0.05, @@ -364,7 +444,9 @@ def load_regulator_data( logger.warning(f"Regulator file not found: {csv_path}") return {} df = pd.read_csv(csv_path) - df = standardize_regulator_results(df, significance_threshold=significance_threshold) + df = standardize_regulator_results( + df, significance_threshold=significance_threshold + ) df = df[df["significant"] == True].copy() result = {} for pid, group in df.groupby("program_id"): @@ -383,31 +465,39 @@ def get_top_regulators( top_n_positive: Optional[int] = None, top_n_negative: Optional[int] = None, use_all_significant: bool = False, - max_regulators: int = 20 + max_regulators: int = 20, ) -> Dict[str, List[Dict[str, Any]]]: reg_df = regulator_data.get(program_id) if reg_df is None or len(reg_df) == 0: - return {'positive': [], 'negative': []} + return {"positive": [], "negative": []} top_n_positive = top_n if top_n_positive is None else top_n_positive top_n_negative = top_n if top_n_negative is None else top_n_negative - sig_df = reg_df[reg_df['significant'] == True].copy() + sig_df = reg_df[reg_df["significant"] == True].copy() sorted_df = sig_df.sort_values(by="log_2_fold_change") if use_all_significant: positive = sorted_df[sorted_df["log_2_fold_change"] < 0].head(max_regulators) - negative = sorted_df[sorted_df["log_2_fold_change"] > 0].tail(max_regulators).iloc[::-1] + negative = ( + sorted_df[sorted_df["log_2_fold_change"] > 0] + .tail(max_regulators) + .iloc[::-1] + ) else: positive = sorted_df[sorted_df["log_2_fold_change"] < 0].head(top_n_positive) - negative = sorted_df[sorted_df["log_2_fold_change"] > 0].tail(top_n_negative).iloc[::-1] + negative = ( + sorted_df[sorted_df["log_2_fold_change"] > 0] + .tail(top_n_negative) + .iloc[::-1] + ) def extract_regulator(row): - result = {'gene': row['grna_target'], 'log2fc': row['log_2_fold_change']} - if 'p_value' in row: - result['pvalue'] = row['p_value'] + result = {"gene": row["grna_target"], "log2fc": row["log_2_fold_change"]} + if "p_value" in row: + result["pvalue"] = row["p_value"] return result return { - 'positive': [extract_regulator(row) for _, row in positive.iterrows()], - 'negative': [extract_regulator(row) for _, row in negative.iterrows()] + "positive": [extract_regulator(row) for _, row in positive.iterrows()], + "negative": [extract_regulator(row) for _, row in negative.iterrows()], } @@ -415,6 +505,7 @@ def extract_regulator(row): # PubTator3 search and BioC fetch # --------------------------------------------------------------------------- + def search_pubtator(query: str, max_results: int = 50) -> List[Dict[str, Any]]: url = f"{PUBTATOR_API_BASE}/search/" params = {"text": query, "page": 1, "size": min(max_results, 100)} @@ -425,9 +516,15 @@ def search_pubtator(query: str, max_results: int = 50) -> List[Dict[str, Any]]: return [] data = resp.json() results = data.get("results", [])[:max_results] - return [{'pmid': r.get('pmid'), 'title': r.get('title', ''), - 'score': r.get('score', 0), 'text_hl': r.get('text_hl', '')} - for r in results] + return [ + { + "pmid": r.get("pmid"), + "title": r.get("title", ""), + "score": r.get("score", 0), + "text_hl": r.get("text_hl", ""), + } + for r in results + ] except Exception as e: logger.error(f"PubTator search error: {e}") return [] @@ -452,43 +549,49 @@ def fetch_bioc_relations_with_text(pmids: List[int]) -> Dict[int, Dict[str, Any] docs = [data] result = {} for doc in docs: - pmid = doc.get('pmid') or doc.get('id') + pmid = doc.get("pmid") or doc.get("id") if not pmid: continue - title = '' - abstract = '' + title = "" + abstract = "" gene_mentions: Dict[str, Set[str]] = {} - for passage in doc.get('passages', []): - ptype = passage.get('infons', {}).get('type', '') - if ptype == 'title': - title = passage.get('text', '') - elif ptype == 'abstract': - abstract = passage.get('text', '') - for ann in passage.get('annotations', []): - if ann.get('infons', {}).get('type') == 'Gene': - gene_name = ann.get('infons', {}).get('name', '') - text_mention = ann.get('text', '') + for passage in doc.get("passages", []): + ptype = passage.get("infons", {}).get("type", "") + if ptype == "title": + title = passage.get("text", "") + elif ptype == "abstract": + abstract = passage.get("text", "") + for ann in passage.get("annotations", []): + if ann.get("infons", {}).get("type") == "Gene": + gene_name = ann.get("infons", {}).get("name", "") + text_mention = ann.get("text", "") if gene_name and text_mention: if gene_name.upper() not in gene_mentions: gene_mentions[gene_name.upper()] = set() gene_mentions[gene_name.upper()].add(text_mention) relations = [] - for rel in doc.get('relations', []): - infons = rel.get('infons', {}) - r1 = infons.get('role1', {}) - r2 = infons.get('role2', {}) - rel_type = infons.get('type', 'Unknown') - score = float(infons.get('score', 0)) - if r1.get('type') == 'Gene' and r2.get('type') == 'Gene': - if rel_type in ('Positive_Correlation', 'Negative_Correlation'): - relations.append({ - 'gene1': r1.get('name', ''), 'gene2': r2.get('name', ''), - 'type': rel_type, 'score': score - }) + for rel in doc.get("relations", []): + infons = rel.get("infons", {}) + r1 = infons.get("role1", {}) + r2 = infons.get("role2", {}) + rel_type = infons.get("type", "Unknown") + score = float(infons.get("score", 0)) + if r1.get("type") == "Gene" and r2.get("type") == "Gene": + if rel_type in ("Positive_Correlation", "Negative_Correlation"): + relations.append( + { + "gene1": r1.get("name", ""), + "gene2": r2.get("name", ""), + "type": rel_type, + "score": score, + } + ) gene_mentions_list = {k: list(v) for k, v in gene_mentions.items()} result[int(pmid)] = { - 'relations': relations, 'title': title, - 'abstract': abstract, 'gene_mentions': gene_mentions_list + "relations": relations, + "title": title, + "abstract": abstract, + "gene_mentions": gene_mentions_list, } return result except Exception as e: @@ -500,26 +603,34 @@ def fetch_bioc_relations_with_text(pmids: List[int]) -> Dict[int, Dict[str, Any] # Regulator validation # --------------------------------------------------------------------------- + def validate_regulator_with_string( - regulator: str, program_genes: List[str], - species: int = 10090, required_score: int = 400 + regulator: str, + program_genes: List[str], + species: int = 10090, + required_score: int = 400, ) -> Dict[str, Any]: time.sleep(1.0) result = get_regulator_program_interactions( - regulator=regulator, program_genes=program_genes, - species=species, required_score=required_score, top_n=10 + regulator=regulator, + program_genes=program_genes, + species=species, + required_score=required_score, + top_n=10, ) return { - 'regulator': regulator, - 'n_program_targets': result['n_interactions'], - 'string_interactions': result['interactions'] + "regulator": regulator, + "n_program_targets": result["n_interactions"], + "string_interactions": result["interactions"], } def validate_regulator_program( - regulator: str, program_genes: List[str], + regulator: str, + program_genes: List[str], keyword: str = "", - max_pmids: int = 50, min_relation_score: float = 0.5 + max_pmids: int = 50, + min_relation_score: float = 0.5, ) -> Dict[str, Any]: genes_or = " OR ".join(program_genes[:10]) query = f"({regulator}) AND ({genes_or})" + (f" AND ({keyword})" if keyword else "") @@ -527,9 +638,14 @@ def validate_regulator_program( time.sleep(PUBTATOR_RATE_LIMIT) search_results = search_pubtator(query, max_results=60) if not search_results: - return {'regulator': regulator, 'papers_found': 0, - 'papers_with_relations': 0, 'mechanistic_relations': [], 'top_papers': []} - pmids = [r['pmid'] for r in search_results[:max_pmids] if r['pmid']] + return { + "regulator": regulator, + "papers_found": 0, + "papers_with_relations": 0, + "mechanistic_relations": [], + "top_papers": [], + } + pmids = [r["pmid"] for r in search_results[:max_pmids] if r["pmid"]] time.sleep(PUBTATOR_RATE_LIMIT) pmid_data = fetch_bioc_relations_with_text(pmids) mechanistic_relations = [] @@ -538,37 +654,47 @@ def validate_regulator_program( seen_gene_pmid = set() pmids_with_relations = set() for pmid, data in pmid_data.items(): - rels = data.get('relations', []) - abstract = data.get('abstract', '') - title = data.get('title', '') + rels = data.get("relations", []) + abstract = data.get("abstract", "") + title = data.get("title", "") for rel in rels: - if rel['score'] < min_relation_score: + if rel["score"] < min_relation_score: continue - g1 = rel['gene1'].upper() - g2 = rel['gene2'].upper() - rel_type = rel['type'] - if rel_type not in ('Positive_Correlation', 'Negative_Correlation'): + g1 = rel["gene1"].upper() + g2 = rel["gene2"].upper() + rel_type = rel["type"] + if rel_type not in ("Positive_Correlation", "Negative_Correlation"): continue if regulator_upper not in (g1, g2): continue pmids_with_relations.add(pmid) - other_gene = rel['gene2'] if g1 == regulator_upper else rel['gene1'] + other_gene = rel["gene2"] if g1 == regulator_upper else rel["gene1"] dedup_key = (other_gene.upper(), pmid) if dedup_key in seen_gene_pmid: continue seen_gene_pmid.add(dedup_key) - sentences = extract_mechanistic_sentences(abstract, regulator, other_gene, max_sentences=2) + sentences = extract_mechanistic_sentences( + abstract, regulator, other_gene, max_sentences=2 + ) if not sentences: - sentences = extract_any_mechanistic_sentences(abstract, regulator, max_sentences=1) - mechanistic_relations.append({ - 'pmid': pmid, 'target_gene': other_gene, 'relation_type': rel_type, - 'score': rel['score'], 'title': title, 'sentences': sentences - }) + sentences = extract_any_mechanistic_sentences( + abstract, regulator, max_sentences=1 + ) + mechanistic_relations.append( + { + "pmid": pmid, + "target_gene": other_gene, + "relation_type": rel_type, + "score": rel["score"], + "title": title, + "sentences": sentences, + } + ) if not mechanistic_relations: for pmid, data in pmid_data.items(): - abstract = data.get('abstract', '') - title = data.get('title', '') - gene_mentions = data.get('gene_mentions', {}) + abstract = data.get("abstract", "") + title = data.get("title", "") + gene_mentions = data.get("gene_mentions", {}) if not abstract: continue regulator_aliases = gene_mentions.get(regulator.upper(), [regulator]) @@ -576,91 +702,128 @@ def validate_regulator_program( regulator_aliases = [regulator] sentences = [] for alias in regulator_aliases: - sentences = extract_any_mechanistic_sentences(abstract, alias, max_sentences=2) + sentences = extract_any_mechanistic_sentences( + abstract, alias, max_sentences=2 + ) if sentences: break if not sentences: - sentences = extract_any_mechanistic_sentences(abstract, regulator, max_sentences=2) + sentences = extract_any_mechanistic_sentences( + abstract, regulator, max_sentences=2 + ) if sentences: - mechanistic_relations.append({ - 'pmid': pmid, 'target_gene': 'unknown', 'relation_type': 'Mentioned', - 'score': 0.5, 'title': title, 'sentences': sentences - }) - mechanistic_relations.sort(key=lambda x: -x['score']) + mechanistic_relations.append( + { + "pmid": pmid, + "target_gene": "unknown", + "relation_type": "Mentioned", + "score": 0.5, + "title": title, + "sentences": sentences, + } + ) + mechanistic_relations.sort(key=lambda x: -x["score"]) mechanistic_relations = mechanistic_relations[:10] - top_papers = [{'pmid': r['pmid'], 'title': r['title']} for r in search_results[:10]] - papers_with_rels = len([p for p, d in pmid_data.items() if d.get('relations')]) + top_papers = [{"pmid": r["pmid"], "title": r["title"]} for r in search_results[:10]] + papers_with_rels = len([p for p, d in pmid_data.items() if d.get("relations")]) return { - 'regulator': regulator, 'papers_found': len(search_results), - 'papers_with_relations': papers_with_rels, - 'mechanistic_relations': mechanistic_relations, 'top_papers': top_papers + "regulator": regulator, + "papers_found": len(search_results), + "papers_with_relations": papers_with_rels, + "mechanistic_relations": mechanistic_relations, + "top_papers": top_papers, } def validate_program_regulators( - program_id: int, regulator_data: Dict[int, pd.DataFrame], - program_genes: List[str], keyword: str = "", - top_n_regulators: int = 3, top_n_positive_regulators: Optional[int] = None, + program_id: int, + regulator_data: Dict[int, pd.DataFrame], + program_genes: List[str], + keyword: str = "", + top_n_regulators: int = 3, + top_n_positive_regulators: Optional[int] = None, top_n_negative_regulators: Optional[int] = None, - max_pmids_per_regulator: int = 50, use_string: bool = True, - use_all_significant: bool = False, max_regulators: int = 20, - use_batch: bool = True, min_score: int = 400 + max_pmids_per_regulator: int = 50, + use_string: bool = True, + use_all_significant: bool = False, + max_regulators: int = 20, + use_batch: bool = True, + min_score: int = 400, ) -> Dict[str, Any]: top_regs = get_top_regulators( - regulator_data, program_id, top_n=top_n_regulators, + regulator_data, + program_id, + top_n=top_n_regulators, top_n_positive=top_n_positive_regulators, top_n_negative=top_n_negative_regulators, - use_all_significant=use_all_significant, max_regulators=max_regulators + use_all_significant=use_all_significant, + max_regulators=max_regulators, ) - result: Dict[str, Any] = {'positive_regulators': [], 'negative_regulators': []} - all_regulator_genes = [r['gene'] for r in top_regs['positive']] + [r['gene'] for r in top_regs['negative']] + result: Dict[str, Any] = {"positive_regulators": [], "negative_regulators": []} + all_regulator_genes = [r["gene"] for r in top_regs["positive"]] + [ + r["gene"] for r in top_regs["negative"] + ] if use_string and use_batch and all_regulator_genes: batch_results = batch_validate_regulators( regulator_genes=all_regulator_genes, program_genes=program_genes, - required_score=min_score + required_score=min_score, ) - for reg_info in top_regs['positive']: - gene = reg_info['gene'] + for reg_info in top_regs["positive"]: + gene = reg_info["gene"] interactions = batch_results.get(gene, []) - result['positive_regulators'].append({ - 'regulator': gene, 'log2fc': reg_info['log2fc'], - 'n_program_targets': len(interactions), - 'string_interactions': interactions - }) - for reg_info in top_regs['negative']: - gene = reg_info['gene'] + result["positive_regulators"].append( + { + "regulator": gene, + "log2fc": reg_info["log2fc"], + "n_program_targets": len(interactions), + "string_interactions": interactions, + } + ) + for reg_info in top_regs["negative"]: + gene = reg_info["gene"] interactions = batch_results.get(gene, []) - result['negative_regulators'].append({ - 'regulator': gene, 'log2fc': reg_info['log2fc'], - 'n_program_targets': len(interactions), - 'string_interactions': interactions - }) + result["negative_regulators"].append( + { + "regulator": gene, + "log2fc": reg_info["log2fc"], + "n_program_targets": len(interactions), + "string_interactions": interactions, + } + ) return result - for reg_info in top_regs['positive']: + for reg_info in top_regs["positive"]: if use_string: logger.info(f" Validating activator {reg_info['gene']} with STRING...") validation = validate_regulator_with_string( - regulator=reg_info['gene'], program_genes=program_genes) + regulator=reg_info["gene"], program_genes=program_genes + ) else: validation = validate_regulator_program( - regulator=reg_info['gene'], program_genes=program_genes, - keyword=keyword, max_pmids=max_pmids_per_regulator) - validation['log2fc'] = reg_info['log2fc'] - result['positive_regulators'].append(validation) - - for reg_info in top_regs['negative']: + regulator=reg_info["gene"], + program_genes=program_genes, + keyword=keyword, + max_pmids=max_pmids_per_regulator, + ) + validation["log2fc"] = reg_info["log2fc"] + result["positive_regulators"].append(validation) + + for reg_info in top_regs["negative"]: if use_string: logger.info(f" Validating repressor {reg_info['gene']} with STRING...") validation = validate_regulator_with_string( - regulator=reg_info['gene'], program_genes=program_genes) + regulator=reg_info["gene"], program_genes=program_genes + ) else: validation = validate_regulator_program( - regulator=reg_info['gene'], program_genes=program_genes, - keyword=keyword, max_pmids=max_pmids_per_regulator) - validation['log2fc'] = reg_info['log2fc'] - result['negative_regulators'].append(validation) + regulator=reg_info["gene"], + program_genes=program_genes, + keyword=keyword, + max_pmids=max_pmids_per_regulator, + ) + validation["log2fc"] = reg_info["log2fc"] + result["negative_regulators"].append(validation) return result diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/llm_submission.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/llm_submission.py index d8913cf..bdbe6d1 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/llm_submission.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/llm_submission.py @@ -24,6 +24,7 @@ # --------------------------------------------------------------------------- try: from dotenv import load_dotenv + load_dotenv() except ImportError: pass # dotenv not required if env vars are already set @@ -34,6 +35,7 @@ try: from google import genai # type: ignore from google.genai.types import CreateBatchJobConfig, JobState, HttpOptions # type: ignore + VERTEX_AVAILABLE = True except ImportError: genai = None # type: ignore @@ -47,6 +49,7 @@ # --------------------------------------------------------------------------- try: import anthropic + ANTHROPIC_AVAILABLE = True except ImportError: anthropic = None # type: ignore @@ -57,8 +60,12 @@ # --------------------------------------------------------------------------- # Set these for your own Google Cloud project; there are no built-in defaults. -VERTEX_PROJECT_ID = os.environ.get("VERTEX_PROJECT_ID") or os.environ.get("GOOGLE_CLOUD_PROJECT") -VERTEX_LOCATION = os.environ.get("VERTEX_LOCATION") or os.environ.get("GOOGLE_CLOUD_LOCATION", "us-east5") +VERTEX_PROJECT_ID = os.environ.get("VERTEX_PROJECT_ID") or os.environ.get( + "GOOGLE_CLOUD_PROJECT" +) +VERTEX_LOCATION = os.environ.get("VERTEX_LOCATION") or os.environ.get( + "GOOGLE_CLOUD_LOCATION", "us-east5" +) VERTEX_BUCKET = os.environ.get("VERTEX_BUCKET") # e.g. gs:///batch VERTEX_MODEL_MAP = { @@ -77,6 +84,7 @@ # Vertex AI Functions # ============================================================================= + def convert_to_vertex_jsonl(input_path: Path, output_path: Path) -> int: """Convert Anthropic batch JSON to JSONL format for Vertex AI. @@ -86,12 +94,12 @@ def convert_to_vertex_jsonl(input_path: Path, output_path: Path) -> int: Vertex AI Claude batch format: {"custom_id": "...", "request": {"messages": [...], "anthropic_version": "vertex-2023-10-16", "max_tokens": ...}} """ - with open(input_path, 'r', encoding='utf-8') as f: + with open(input_path, "r", encoding="utf-8") as f: data = json.load(f) requests = data.get("requests", []) - with open(output_path, 'w', encoding='utf-8') as f: + with open(output_path, "w", encoding="utf-8") as f: for req in requests: # Convert from Anthropic format to Vertex AI format params = req.get("params", {}) @@ -101,7 +109,7 @@ def convert_to_vertex_jsonl(input_path: Path, output_path: Path) -> int: "messages": params.get("messages", []), "anthropic_version": "vertex-2023-10-16", "max_tokens": params.get("max_tokens", 4096), - } + }, } # Optionally include system prompt if present if "system" in params: @@ -110,7 +118,9 @@ def convert_to_vertex_jsonl(input_path: Path, output_path: Path) -> int: line = json.dumps(vertex_req) f.write(line + "\n") - logger.info(f"Converted {len(requests)} requests to Vertex AI JSONL format: {output_path}") + logger.info( + f"Converted {len(requests)} requests to Vertex AI JSONL format: {output_path}" + ) return len(requests) @@ -132,7 +142,9 @@ def upload_to_gcs(local_path: Path, gcs_uri: str) -> bool: def cmd_submit_vertex(args: argparse.Namespace) -> int: """Submit a batch job to Vertex AI.""" if not VERTEX_AVAILABLE: - logger.error("google-genai package not installed. Run: pip install google-genai") + logger.error( + "google-genai package not installed. Run: pip install google-genai" + ) return 1 if not args.batch_file: @@ -155,7 +167,9 @@ def cmd_submit_vertex(args: argparse.Namespace) -> int: # Upload to GCS bucket = args.bucket or VERTEX_BUCKET if not bucket: - logger.error("No GCS bucket: pass --bucket gs:///batch or set VERTEX_BUCKET.") + logger.error( + "No GCS bucket: pass --bucket gs:///batch or set VERTEX_BUCKET." + ) return 1 gcs_input_uri = f"{bucket}/inputs/{base_name}_{timestamp}.jsonl" if not upload_to_gcs(jsonl_path, gcs_input_uri): @@ -213,7 +227,9 @@ def cmd_submit_vertex(args: argparse.Namespace) -> int: print(f"Job did not succeed.") return 3 else: - print(f"\nCheck status with: python {__file__} check-vertex --job-name {job.name}") + print( + f"\nCheck status with: python {__file__} check-vertex --job-name {job.name}" + ) return 0 @@ -221,7 +237,9 @@ def cmd_submit_vertex(args: argparse.Namespace) -> int: def cmd_check_vertex(args: argparse.Namespace) -> int: """Check the status of a Vertex AI batch job.""" if not VERTEX_AVAILABLE: - logger.error("google-genai package not installed. Run: pip install google-genai") + logger.error( + "google-genai package not installed. Run: pip install google-genai" + ) return 1 if not args.job_name: @@ -235,9 +253,9 @@ def cmd_check_vertex(args: argparse.Namespace) -> int: print(f"Job name: {job.name}") print(f"Job state: {job.state}") - if hasattr(job, 'create_time'): + if hasattr(job, "create_time"): print(f"Created: {job.create_time}") - if hasattr(job, 'update_time'): + if hasattr(job, "update_time"): print(f"Updated: {job.update_time}") # Check if completed @@ -264,6 +282,7 @@ def cmd_check_vertex(args: argparse.Namespace) -> int: # Anthropic Direct Batch API Commands (Default) # ============================================================================= + def cmd_submit_anthropic(args: argparse.Namespace) -> int: """Submit a batch job to Anthropic Batch API (default). @@ -309,16 +328,20 @@ def cmd_submit_anthropic(args: argparse.Namespace) -> int: logger.warning(f"Skipping request {custom_id}: no messages") continue - anthropic_requests.append({ - "custom_id": custom_id, - "params": { - "model": model, - "max_tokens": max_tokens, - "messages": messages, + anthropic_requests.append( + { + "custom_id": custom_id, + "params": { + "model": model, + "max_tokens": max_tokens, + "messages": messages, + }, } - }) + ) - logger.info(f"Submitting {len(anthropic_requests)} requests to Anthropic Batch API...") + logger.info( + f"Submitting {len(anthropic_requests)} requests to Anthropic Batch API..." + ) logger.info(f" Model: {model}") logger.info(f" Max tokens: {max_tokens}") @@ -327,10 +350,16 @@ def cmd_submit_anthropic(args: argparse.Namespace) -> int: auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") api_key = os.environ.get("ANTHROPIC_API_KEY") if not auth_token and not api_key: - logger.error("ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment. Set in .env file.") + logger.error( + "ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment. Set in .env file." + ) return 1 - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic(api_key=api_key) + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic(api_key=api_key) + ) # Submit batch batch = client.messages.batches.create(requests=anthropic_requests) @@ -351,7 +380,9 @@ def cmd_submit_anthropic(args: argparse.Namespace) -> int: while batch.processing_status == "in_progress": time.sleep(30) batch = client.messages.batches.retrieve(batch.id) - print(f" Status: {batch.processing_status} | Completed: {batch.request_counts.succeeded}/{batch.request_counts.processing + batch.request_counts.succeeded}") + print( + f" Status: {batch.processing_status} | Completed: {batch.request_counts.succeeded}/{batch.request_counts.processing + batch.request_counts.succeeded}" + ) print(f"\nFinal status: {batch.processing_status}") @@ -371,8 +402,12 @@ def cmd_submit_anthropic(args: argparse.Namespace) -> int: print(f"Batch did not complete successfully.") return 3 else: - print(f"\nCheck status with: python {__file__} check-anthropic --batch-id {batch.id}") - print(f"Or retrieve results later with: python {__file__} results-anthropic --batch-id {batch.id}") + print( + f"\nCheck status with: python {__file__} check-anthropic --batch-id {batch.id}" + ) + print( + f"Or retrieve results later with: python {__file__} results-anthropic --batch-id {batch.id}" + ) return 0 @@ -413,8 +448,7 @@ def cmd_submit_single(args: argparse.Namespace) -> int: # Determine output path output_file = Path( - args.output - or str(batch_path.with_name(batch_path.stem + "_results.jsonl")) + args.output or str(batch_path.with_name(batch_path.stem + "_results.jsonl")) ) output_file.parent.mkdir(parents=True, exist_ok=True) @@ -464,14 +498,18 @@ def cmd_submit_single(args: argparse.Namespace) -> int: for idx, req in enumerate(requests_list, start=1): custom_id = req.get("custom_id", f"request_{idx}") if custom_id in completed_ids: - logger.info("[%d/%d] %s: skipped (already completed)", idx, total, custom_id) + logger.info( + "[%d/%d] %s: skipped (already completed)", idx, total, custom_id + ) succeeded += 1 continue params = req.get("params", {}) messages = params.get("messages", []) if not messages: - logger.warning("[%d/%d] %s: skipped (no messages)", idx, total, custom_id) + logger.warning( + "[%d/%d] %s: skipped (no messages)", idx, total, custom_id + ) continue try: @@ -505,7 +543,11 @@ def cmd_submit_single(args: argparse.Namespace) -> int: out_tok = response.usage.output_tokens if response.usage else "?" logger.info( "[%d/%d] %s: succeeded (%sin/%sout)", - idx, total, custom_id, in_tok, out_tok, + idx, + total, + custom_id, + in_tok, + out_tok, ) except Exception as exc: @@ -541,16 +583,24 @@ def cmd_check_anthropic(args: argparse.Namespace) -> int: batch_id = batch_id_file.read_text(encoding="utf-8").strip() if not batch_id: - logger.error("--batch-id is required (or provide --batch-file with saved .batch_id)") + logger.error( + "--batch-id is required (or provide --batch-file with saved .batch_id)" + ) return 1 auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") api_key = os.environ.get("ANTHROPIC_API_KEY") if not auth_token and not api_key: - logger.error("ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment.") + logger.error( + "ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment." + ) return 1 - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic(api_key=api_key) + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic(api_key=api_key) + ) batch = client.messages.batches.retrieve(batch_id) print(f"Batch ID: {batch.id}") @@ -589,10 +639,16 @@ def cmd_results_anthropic(args: argparse.Namespace) -> int: auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN") api_key = os.environ.get("ANTHROPIC_API_KEY") if not auth_token and not api_key: - logger.error("ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment.") + logger.error( + "ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY not found in environment." + ) return 1 - client = anthropic.Anthropic(auth_token=auth_token) if auth_token else anthropic.Anthropic(api_key=api_key) + client = ( + anthropic.Anthropic(auth_token=auth_token) + if auth_token + else anthropic.Anthropic(api_key=api_key) + ) batch = client.messages.batches.retrieve(batch_id) if batch.processing_status != "ended": diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/ncbi_api.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/ncbi_api.py index 7bd9a89..9f17e80 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/ncbi_api.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/ncbi_api.py @@ -1,5 +1,5 @@ """ -@description +@description Unified NCBI API wrapper for Pubtator 3 and E-Utilities. Handles rate limiting, pagination, and data retrieval for gene-literaure integration. @@ -40,28 +40,31 @@ logger = logging.getLogger(__name__) + class NcbiClient: def __init__(self, api_key: Optional[str] = None): """ Initialize NBCI Client. - + Args: api_key: NCBI API Key. If None, checks 'NCBI_API_KEY' env var. """ self.api_key = api_key or os.environ.get("NCBI_API_KEY") self.sleep_time = RATE_LIMIT_WITH_KEY if self.api_key else RATE_LIMIT_NO_KEY self.session = requests.Session() - + if self.api_key: logger.info("NcbiClient initialized with API Key (High throughput)") else: logger.warning("NcbiClient initialized WITHOUT API Key (Low throughput)") - def _get(self, url: str, params: Dict[str, Any], retries: int = 3) -> Optional[requests.Response]: + def _get( + self, url: str, params: Dict[str, Any], retries: int = 3 + ) -> Optional[requests.Response]: """Internal GET with retry and rate limiting.""" if self.api_key: params["api_key"] = self.api_key - + for attempt in range(retries): try: time.sleep(self.sleep_time) @@ -75,13 +78,17 @@ def _get(self, url: str, params: Dict[str, Any], retries: int = 3) -> Optional[r logger.warning(f"Rate limit 429 hit. Sleeping {wait}s...") time.sleep(wait) else: - logger.warning(f"Request failed {url}: {resp.status_code} {resp.text[:100]}") + logger.warning( + f"Request failed {url}: {resp.status_code} {resp.text[:100]}" + ) except requests.RequestException as e: - logger.warning(f"Request error (attempt {attempt+1}): {e}") + logger.warning(f"Request error (attempt {attempt + 1}): {e}") time.sleep(1) return None - def _post(self, url: str, json_data: Dict[str, Any], retries: int = 3) -> Optional[requests.Response]: + def _post( + self, url: str, json_data: Dict[str, Any], retries: int = 3 + ) -> Optional[requests.Response]: """Internal POST with retry and rate limiting (API key passed in URL for POSTs sometimes needed, or params).""" params = {} if self.api_key: @@ -99,9 +106,11 @@ def _post(self, url: str, json_data: Dict[str, Any], retries: int = 3) -> Option logger.warning(f"Rate limit 429 hit. Sleeping {wait}s...") time.sleep(wait) else: - logger.warning(f"POST failed {url}: {resp.status_code} {resp.text[:100]}") + logger.warning( + f"POST failed {url}: {resp.status_code} {resp.text[:100]}" + ) except requests.RequestException as e: - logger.warning(f"POST error (attempt {attempt+1}): {e}") + logger.warning(f"POST error (attempt {attempt + 1}): {e}") time.sleep(1) return None @@ -113,7 +122,7 @@ def search_literature(self, query: str, page: int = 1, size: int = 20) -> List[i """ Search Pubtator 3 for PMIDs. Wraps: GET https://www.ncbi.nlm.nih.gov/research/pubtator3-api/search/ - + Args: query: Search query string page: Page number (1-indexed) @@ -124,7 +133,7 @@ def search_literature(self, query: str, page: int = 1, size: int = 20) -> List[i resp = self._get(url, params) if not resp: return [] - + try: data = resp.json() # Structure: {"results": [{"pmid": "123"}, ...], ...} @@ -139,7 +148,7 @@ def search_literature(self, query: str, page: int = 1, size: int = 20) -> List[i if pmid: pmids.append(int(pmid)) elif isinstance(item, (int, str)): - pmids.append(int(item)) + pmids.append(int(item)) return pmids except ValueError: logger.error("Failed to parse JSON from Pubtator search") @@ -154,11 +163,11 @@ def fetch_bioc_annotations(self, pmids: List[int]) -> List[Dict[str, Any]]: # API expects pmids as list in body? Or comma string? # Pubtator3 API docs usually say POST with body `{"pmids": [1, 2]}` payload = {"pmids": pmids, "full": True} - + resp = self._post(url, payload) if not resp: return [] - + # Responses is BioC-JSON (one JSON object usually, or lines?) # Typically it's a JSON object with "documents": [...] or just a list try: @@ -167,9 +176,9 @@ def fetch_bioc_annotations(self, pmids: List[int]) -> List[Dict[str, Any]]: data = resp.json() # If it's a wrapper, extract docs if isinstance(data, dict): - if "PubTator3" in data: - return data["PubTator3"] - return data.get("documents", []) if "documents" in data else [data] + if "PubTator3" in data: + return data["PubTator3"] + return data.get("documents", []) if "documents" in data else [data] elif isinstance(data, list): return data return [] @@ -181,55 +190,57 @@ def fetch_bioc_annotations(self, pmids: List[int]) -> List[Dict[str, Any]]: # E-Utilities # ------------------------------------------------------------------------- - def normalize_genes(self, symbols: List[str], organism: str = "mouse") -> Dict[str, str]: + def normalize_genes( + self, symbols: List[str], organism: str = "mouse" + ) -> Dict[str, str]: """ Map Gene Symbols to Entrez IDs using ESearch. Returns: {Symbol: EntrezID} """ if not symbols: return {} - + # Build query: (Sym1[Gene] OR Sym2[Gene]) AND organism[Orgn] # Chunking needed if list is huge, but usually caller handles chunks. # But ESearch GET url length limit ~2k chars. - + mapping = {} # We can try to do this effectively by querying one by one or small batches # Or simpler: "Symbol[Gene] AND Organism[Orgn]" - + # A robust way is 'esearch' then 'esummary', but that is slow for many genes. # Alternative: Use 'elinky' or simply trust the symbol if unique. - - # For this implementation, let's assume we search them one by one or small groups + + # For this implementation, let's assume we search them one by one or small groups # is too slow. # Better approach: return "symbol" as ID if normalization is hard sans-DB. # BUT, for Pubtator, Symbol is often NOT enough for ID-based queries if we want precision. # However, Pubtator /search/ accepts text. So Symbol is fine for Part 1 (Search). - - # Wait, the PLAN says: "Step 1: Search ... Query: (GENE1 OR ...)". + + # Wait, the PLAN says: "Step 1: Search ... Query: (GENE1 OR ...)". # So we actually just need the symbols for the query text. - + # For "Step 2: Fetch & Cache", we rely on PMIDs. - + # The E-Utilities part is for "Gene Summaries" which NEEDS IDs. - + # Let's implement a bulk lookup: # esearch term="Sym1[Sym] OR Sym2[Sym] ..." -> this loses the mapping (who is who). - + # Correct approach for mapping: - # Use E-Utils 'esearch' with 'term=Symbol[Gene] AND Organism[Orgn]' one-by-one? + # Use E-Utils 'esearch' with 'term=Symbol[Gene] AND Organism[Orgn]' one-by-one? # No, too slow. # Hack: The 'biotools' approach often downloads a gene_info.gz. # Given we are an agent, let's try a batched approach but parse carefully. # Actually, for < 1000 genes, single queries are okay-ish if parallel? No. - + # Let's use a simpler heuristic or just return the symbol for now if ID not needed for Context. # Implementation Plan says: "Fetch gene functional summaries (E-Utilities)". # This requires Entrez IDs. - + # Let's implement single lookup loop for now (safest). # Optimization: Only look up the *Driver* genes (Top 5-10/program). - + for sym in symbols: term = f"{sym}[Sym] AND {organism}[Orgn]" url = f"{EUTILS_BASE}/esearch.fcgi" @@ -255,34 +266,33 @@ def get_gene_summaries(self, gene_ids: List[str]) -> Dict[str, str]: """ if not gene_ids: return {} - + # ESummary allows batching comma-separated # Batch size 100 is safe results = {} - + # Chunking chunk_size = 50 for i in range(0, len(gene_ids), chunk_size): - chunk = gene_ids[i:i+chunk_size] + chunk = gene_ids[i : i + chunk_size] url = f"{EUTILS_BASE}/esummary.fcgi" params = {"db": "gene", "id": ",".join(chunk), "retmode": "json"} - + resp = self._get(url, params) if not resp: continue - + try: data = resp.json() # struct: result: { "uids": [...], "123": {summary: ...}, "456": ... } base = data.get("result", {}) uids = base.get("uids", []) for uid in uids: - item = base.get(uid, {}) - summary = item.get("summary", "") - if summary: - results[str(uid)] = summary + item = base.get(uid, {}) + summary = item.get("summary", "") + if summary: + results[str(uid)] = summary except Exception as e: logger.error(f"Error parsing ESummary JSON: {e}") - - return results + return results diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/pipeline_state.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/pipeline_state.py index 32772c6..92886c2 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/pipeline_state.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/pipeline_state.py @@ -96,7 +96,9 @@ def save_state(path: Path, state: PipelineState) -> None: path.write_text(json.dumps(state.to_dict(), indent=2), encoding="utf-8") -def mark_step(state: PipelineState, step: str, status: str, info: Optional[Dict[str, Any]] = None) -> None: +def mark_step( + state: PipelineState, step: str, status: str, info: Optional[Dict[str, Any]] = None +) -> None: s = state.steps.get(step) or StepState() s.status = status if status == "completed": diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/prompt_builder.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/prompt_builder.py index 7c83676..fc756f4 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/prompt_builder.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/prompt_builder.py @@ -27,7 +27,9 @@ MODEL = "claude-sonnet-4-5-20250929" # Anthropic model name DEFAULT_ANNOTATION_ROLE = "molecular and cell biologist" DEFAULT_ANNOTATION_CONTEXT = "a gene program extracted from single-cell Perturb-seq" -DEFAULT_SEARCH_KEYWORD = "(none specified)" # set --search-keyword to your cell type / tissue +DEFAULT_SEARCH_KEYWORD = ( + "(none specified)" # set --search-keyword to your cell type / tissue +) PROMPT_TEMPLATE = """ ## Edited prompt (evidence-first, low-speculation) @@ -140,9 +142,7 @@ def _add_global_uniqueness_scores(df: pd.DataFrame) -> pd.DataFrame: return updated -def _ensure_global_uniqueness( - df: pd.DataFrame, log: logging.Logger -) -> pd.DataFrame: +def _ensure_global_uniqueness(df: pd.DataFrame, log: logging.Logger) -> pd.DataFrame: if "UniquenessScore" in df.columns and not df["UniquenessScore"].isna().all(): return df log.info("UniquenessScore missing; computing global uniqueness scores.") @@ -193,7 +193,10 @@ def load_celltype_annotations( summary_path = celltype_file else: # Try common filenames in the directory - for filename in ["celltype_summary.csv", "program_celltype_annotations_summary.csv"]: + for filename in [ + "celltype_summary.csv", + "program_celltype_annotations_summary.csv", + ]: candidate = celltype_dir / filename if candidate.exists(): summary_path = candidate @@ -282,7 +285,9 @@ def prepare_enrichment_mapping( return {} df["category"] = df["category"].astype(str).str.strip() - df_sorted = df.sort_values(["program_id", "category", "fdr"], ascending=[True, True, True]) + df_sorted = df.sort_values( + ["program_id", "category", "fdr"], ascending=[True, True, True] + ) enrichment_by_program: Dict[int, Dict[str, List[dict]]] = {} for (pid, cat), sub in df_sorted.groupby(["program_id", "category"], sort=False): # type: ignore @@ -308,7 +313,11 @@ def load_ncbi_context(json_path: Optional[Path]) -> Dict[int, Dict[str, Any]]: return {} -def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, allowed_genes: Optional[Set[str]] = None) -> str: +def format_ncbi_context( + ncbi_data: Dict[int, Dict[str, Any]], + program_id: int, + allowed_genes: Optional[Set[str]] = None, +) -> str: ctx = ncbi_data.get(program_id) if not ctx: return "Literature evidence: None available." @@ -333,8 +342,8 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a if source == "ncbi": # Remove citations like [provided by ...] desc = re.sub( - r'\s*\[provided by .*?\]\.?', - '', + r"\s*\[provided by .*?\]\.?", + "", desc, flags=re.IGNORECASE, ) @@ -342,7 +351,7 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a count += 1 if count == 0: - lines.append("None available for selected genes.") + lines.append("None available for selected genes.") # 2. Aggregated Evidence (Snippets) # gene -> list of strings "Statement (PMID:123)" @@ -371,7 +380,7 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a for s in s_list: # s format: "Sentence text.." (PMID:12345) # Extract PMID using regex - pmid_match = re.search(r'\(PMID:(\d+)\)', s) + pmid_match = re.search(r"\(PMID:(\d+)\)", s) pmid = pmid_match.group(1) if pmid_match else None # Dedup by PMID (One sentence per paper per gene) @@ -399,14 +408,15 @@ def format_ncbi_context(ncbi_data: Dict[int, Dict[str, Any]], program_id: int, a gene_snippets.append(f"- {sym}: {clean_s}") - if len(gene_snippets) >= 5: break # Max 5 per gene + if len(gene_snippets) >= 5: + break # Max 5 per gene if gene_snippets: has_snippets = True lines.extend(gene_snippets) if not has_snippets: - lines.append("None found.") + lines.append("None found.") return "\n".join(lines) @@ -469,17 +479,23 @@ def format_regulator_context( # Positive log2FC = negative regulators (knockdown increases program, so gene represses it) negative_regs = sorted_df[sorted_df["log_2_fold_change"] > 0].tail(top_n).iloc[::-1] - lines = ["Regulator perturbations (from Perturb-seq; log2FC indicates effect of knockdown on program activity):"] + lines = [ + "Regulator perturbations (from Perturb-seq; log2FC indicates effect of knockdown on program activity):" + ] if len(positive_regs) > 0: lines.append("\nPositive regulators (activators; knockdown reduces program):") for _, row in positive_regs.iterrows(): - lines.append(f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}") + lines.append( + f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}" + ) if len(negative_regs) > 0: lines.append("\nNegative regulators (repressors; knockdown increases program):") for _, row in negative_regs.iterrows(): - lines.append(f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}") + lines.append( + f"- {row['grna_target']}: log2FC = {row['log_2_fold_change']:.3f}" + ) return "\n".join(lines) @@ -606,9 +622,13 @@ def select_program_genes( if program_df.empty: return [], [] top_loading_genes = ( - program_df.sort_values("Score", ascending=False)["Name"].head(top_loading).tolist() # type: ignore + program_df.sort_values("Score", ascending=False)["Name"] + .head(top_loading) + .tolist() # type: ignore ) - unique_ranked = program_df.sort_values("UniquenessScore", ascending=False)["Name"].tolist() # type: ignore + unique_ranked = program_df.sort_values("UniquenessScore", ascending=False)[ + "Name" + ].tolist() # type: ignore unique_genes = [gene for gene in unique_ranked if gene not in top_loading_genes] return top_loading_genes, unique_genes[:top_unique] @@ -662,7 +682,9 @@ def generate_prompt( ) allowed_genes = set(top_loading_genes) | set(unique_genes) - ncbi_context = format_ncbi_context(ncbi_data, program_id, allowed_genes=allowed_genes) + ncbi_context = format_ncbi_context( + ncbi_data, program_id, allowed_genes=allowed_genes + ) regulator_analysis = format_regulator_analysis_context( regulator_data or {}, diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/result_parser.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/result_parser.py index ec42001..f2d3995 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/result_parser.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/result_parser.py @@ -10,6 +10,7 @@ load_top_genes_by_topic — load top genes per program from a gene-loading CSV generate_unique_topic_names — build a summary CSV with unique topic names """ + from __future__ import annotations import glob @@ -70,7 +71,11 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: try: data = json.loads(line) custom_id = data.get("custom_id") - if custom_id and isinstance(custom_id, str) and custom_id.startswith("topic_"): + if ( + custom_id + and isinstance(custom_id, str) + and custom_id.startswith("topic_") + ): # expected format: topic__annotation m = re.match(r"topic_(\d+)", custom_id) if not m: @@ -85,7 +90,9 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: content = response.get("content", []) if isinstance(content, list) and content: first = content[0] - text_content = first.get("text", "") if isinstance(first, dict) else "" + text_content = ( + first.get("text", "") if isinstance(first, dict) else "" + ) # Fallback to Anthropic format: {"result": {"type": "succeeded", "message": {"content": [...]}}} if not text_content: @@ -95,14 +102,22 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: content = message.get("content", []) if isinstance(content, list) and content: first = content[0] - text_content = first.get("text", "") if isinstance(first, dict) else "" + text_content = ( + first.get("text", "") + if isinstance(first, dict) + else "" + ) elif result.get("type") == "errored": error = result.get("error", {}) - print(f" Topic {topic_number} failed. Reason: {error.get('message', 'Unknown error')}") + print( + f" Topic {topic_number} failed. Reason: {error.get('message', 'Unknown error')}" + ) continue if text_content: - output_filename = os.path.join(output_dir, f"topic_{topic_number}_annotation.md") + output_filename = os.path.join( + output_dir, f"topic_{topic_number}_annotation.md" + ) with open(output_filename, "w", encoding="utf-8") as out_f: out_f.write(text_content) print(f" Saved Topic {topic_number}") @@ -117,7 +132,9 @@ def parse_final_results(result_file: str, output_dir: str) -> List[int]: return saved_ids -def load_top_genes_by_topic(gene_loading_file: str, top_n: int = 10) -> Dict[int, List[str]]: +def load_top_genes_by_topic( + gene_loading_file: str, top_n: int = 10 +) -> Dict[int, List[str]]: """Load top genes by RowID from gene loading CSV.""" top_genes_by_topic: Dict[int, List[str]] = {} if not gene_loading_file or not os.path.exists(gene_loading_file): @@ -136,7 +153,10 @@ def load_top_genes_by_topic(gene_loading_file: str, top_n: int = 10) -> Dict[int for topic_id, group in gene_df.groupby(group_col): genes = ( - group.sort_values("Score", ascending=False)["Name"].astype(str).head(top_n).tolist() + group.sort_values("Score", ascending=False)["Name"] + .astype(str) + .head(top_n) + .tolist() ) top_genes_by_topic[int(topic_id)] = genes print(f" Loaded gene data for {len(top_genes_by_topic)} topics") @@ -159,7 +179,9 @@ def generate_unique_topic_names( key=lambda x: int(re.search(r"topic_(\d+)_", x).group(1)), ) - top_genes_by_topic = load_top_genes_by_topic(gene_loading_file) if gene_loading_file else {} + top_genes_by_topic = ( + load_top_genes_by_topic(gene_loading_file) if gene_loading_file else {} + ) topic_rows: List[Dict[str, object]] = [] used_names: set[str] = set() @@ -178,7 +200,9 @@ def generate_unique_topic_names( content, re.DOTALL, ) - keywords_match = re.search(r"\*\*(?:Program label|Three Key Words):\*\*(.*?)\n", content, re.DOTALL) + keywords_match = re.search( + r"\*\*(?:Program label|Three Key Words):\*\*(.*?)\n", content, re.DOTALL + ) summary = summary_match.group(1).strip() if summary_match else "" keywords_str = keywords_match.group(1).strip() if keywords_match else "" @@ -203,7 +227,9 @@ def generate_unique_topic_names( final_name = f"{base_name} ({signature_gene})" while final_name in used_names: - if topic_number in top_genes_by_topic and counter - 2 < len(top_genes_by_topic[topic_number]): + if topic_number in top_genes_by_topic and counter - 2 < len( + top_genes_by_topic[topic_number] + ): signature_gene = top_genes_by_topic[topic_number][counter - 2] final_name = f"{base_name} ({signature_gene})" else: @@ -213,7 +239,9 @@ def generate_unique_topic_names( re.IGNORECASE, ) if distinguishing_terms and counter - 2 < len(distinguishing_terms): - final_name = f"{base_name} ({distinguishing_terms[counter - 2]})" + final_name = ( + f"{base_name} ({distinguishing_terms[counter - 2]})" + ) else: final_name = f"{base_name} ({counter})" counter += 1 diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/run_pipeline.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/run_pipeline.py index 4255cf4..17699f7 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/run_pipeline.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/run_pipeline.py @@ -20,6 +20,7 @@ (01_genes_to_string_enrichment.py ... 05_generate_html_report.py) located in the same directory as this file. """ + from __future__ import annotations import argparse @@ -36,9 +37,11 @@ try: from dotenv import load_dotenv except ImportError: # python-dotenv is optional + def load_dotenv() -> None: # type: ignore[misc] pass + # Resolve the directory containing this script (and sibling step scripts). _SCRIPT_DIR = Path(__file__).resolve().parent @@ -63,19 +66,18 @@ def load_dotenv() -> None: # type: ignore[misc] load_dotenv() logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(levelname)s - %(message)s" + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # Pipeline steps in order STEPS = [ - "string_enrichment", # Step 1: Extract genes + STRING enrichment - "literature_fetch", # Step 2: Fetch NCBI/Harmonizome data - "batch_prepare", # Step 3a: Prepare LLM batch request - "batch_submit", # Step 3b: Submit to Vertex AI - "parse_results", # Step 4: Parse LLM responses - "html_report", # Step 5: Generate HTML report + "string_enrichment", # Step 1: Extract genes + STRING enrichment + "literature_fetch", # Step 2: Fetch NCBI/Harmonizome data + "batch_prepare", # Step 3a: Prepare LLM batch request + "batch_submit", # Step 3b: Submit to Vertex AI + "parse_results", # Step 4: Parse LLM responses + "html_report", # Step 5: Generate HTML report ] @@ -91,16 +93,21 @@ class PipelineConfig: All intermediate paths are auto-generated. """ + # Required inputs (user must provide these) gene_loading: Path - celltype_enrichment: Optional[Path] # Optional raw enrichment file (summary auto-generated) + celltype_enrichment: Optional[ + Path + ] # Optional raw enrichment file (summary auto-generated) regulator_file: Optional[Path] output_dir: Path # Optional settings with defaults topics: Optional[List[int]] = None # None = all topics species: int = 10090 # Mouse by default - keyword: str = "" # PubMed keyword for your tissue/cell type; empty = no keyword filter + keyword: str = ( + "" # PubMed keyword for your tissue/cell type; empty = no keyword filter + ) annotation_role: str = "molecular and cell biologist" annotation_context: str = "a gene program extracted from single-cell Perturb-seq" n_top_genes: int = 300 @@ -119,7 +126,9 @@ class PipelineConfig: llm_wait: bool = False # Default: do not wait; resume later # Vertex AI settings (only used if llm_backend="vertex") - vertex_bucket: str = os.environ.get("VERTEX_BUCKET", "") # e.g. gs:///batch + vertex_bucket: str = os.environ.get( + "VERTEX_BUCKET", "" + ) # e.g. gs:///batch # Gene summary options full_summaries: bool = False # Use full HTML summaries (3x longer) @@ -146,24 +155,22 @@ def __post_init__(self): "genes_json": out / "genes_top.json", "genes_overview": out / "genes_overview.csv", "enrichment_full": out / "string_enrichment" / "enrichment_full.csv", - "enrichment_filtered": out / "string_enrichment" / "enrichment_filtered.csv", + "enrichment_filtered": out + / "string_enrichment" + / "enrichment_filtered.csv", "enrichment_figures": out / "string_enrichment" / "figures", "enrichment_cache": out / "string_enrichment" / "cache", "gene_loading_with_uniqueness": out / "gene_loading_with_uniqueness.csv", "celltype_summary": out / "celltype_summary.csv", - # Step 2 outputs "ncbi_json": out / "literature_context.json", "ncbi_csv": out / "literature_context.csv", - # Step 3 outputs "batch_request_json": out / "llm_batches" / "batch_request.json", "batch_request_jsonl": out / "llm_batches" / "batch_request.jsonl", - # Step 4 outputs "annotations_dir": out / "annotations", "summary_csv": out / "annotations" / "summary.csv", - # Step 5 outputs "report_html": out / "annotations" / "report.html", } @@ -184,15 +191,30 @@ def from_dict(cls, data: Dict[str, Any]) -> "PipelineConfig": # Filter to only known fields known_fields = { - "gene_loading", "celltype_enrichment", "output_dir", - "regulator_file", "topics", "species", "keyword", - "annotation_role", "annotation_context", - "n_top_genes", "top_loading", "top_unique", - "top_enrichment", "genes_per_term", - "top_positive_regulators", "top_negative_regulators", + "gene_loading", + "celltype_enrichment", + "output_dir", + "regulator_file", + "topics", + "species", + "keyword", + "annotation_role", + "annotation_context", + "n_top_genes", + "top_loading", + "top_unique", + "top_enrichment", + "genes_per_term", + "top_positive_regulators", + "top_negative_regulators", "regulator_significance_threshold", - "llm_backend", "llm_model", "llm_max_tokens", "llm_wait", - "vertex_bucket", "full_summaries", "resume" + "llm_backend", + "llm_model", + "llm_max_tokens", + "llm_wait", + "vertex_bucket", + "full_summaries", + "resume", } if "keyword" not in data and "context" in data: data["keyword"] = data["context"] @@ -221,26 +243,44 @@ def run_step_1_string_enrichment(config: PipelineConfig) -> bool: script = _SCRIPT_DIR / "01_genes_to_string_enrichment.py" cmd = [ - sys.executable, str(script), "all", - "--input", str(config.gene_loading), - "--n-top", str(config.n_top_genes), - "--json-out", str(config.get_path("genes_json")), - "--csv-out", str(config.get_path("genes_overview")), - "--species", str(config.species), - "--out-csv-full", str(config.get_path("enrichment_full")), - "--out-csv-filtered", str(config.get_path("enrichment_filtered")), - "--figures-dir", str(config.get_path("enrichment_figures")), - "--cache-dir", str(config.get_path("enrichment_cache")), - "--gene-loading-out", str(config.get_path("gene_loading_with_uniqueness")), + sys.executable, + str(script), + "all", + "--input", + str(config.gene_loading), + "--n-top", + str(config.n_top_genes), + "--json-out", + str(config.get_path("genes_json")), + "--csv-out", + str(config.get_path("genes_overview")), + "--species", + str(config.species), + "--out-csv-full", + str(config.get_path("enrichment_full")), + "--out-csv-filtered", + str(config.get_path("enrichment_filtered")), + "--figures-dir", + str(config.get_path("enrichment_figures")), + "--cache-dir", + str(config.get_path("enrichment_cache")), + "--gene-loading-out", + str(config.get_path("gene_loading_with_uniqueness")), ] if config.celltype_enrichment: - cmd.extend([ - "--celltype-enrichment", str(config.celltype_enrichment), - "--celltype-summary-out", str(config.get_path("celltype_summary")), - ]) + cmd.extend( + [ + "--celltype-enrichment", + str(config.celltype_enrichment), + "--celltype-summary-out", + str(config.get_path("celltype_summary")), + ] + ) else: - logger.info("Cell-type enrichment file not provided; skipping cell-type summary generation.") + logger.info( + "Cell-type enrichment file not provided; skipping cell-type summary generation." + ) if config.topics: cmd.extend(["--topics", ",".join(map(str, config.topics))]) @@ -272,15 +312,23 @@ def run_step_2_literature_fetch(config: PipelineConfig) -> bool: gene_input = config.gene_loading cmd = [ - sys.executable, str(script), - "--input", str(gene_input), - "--keyword", config.keyword, - "--json-out", str(config.get_path("ncbi_json")), - "--csv-out", str(config.get_path("ncbi_csv")), - "--gene-summary-source", "harmonizome", + sys.executable, + str(script), + "--input", + str(gene_input), + "--keyword", + config.keyword, + "--json-out", + str(config.get_path("ncbi_json")), + "--csv-out", + str(config.get_path("ncbi_csv")), + "--gene-summary-source", + "harmonizome", # Fetch summaries for exactly top_loading + top_unique genes - "--top-loading", str(config.top_loading), - "--top-unique", str(config.top_unique), + "--top-loading", + str(config.top_loading), + "--top-unique", + str(config.top_unique), ] # Optional: Use full HTML summaries (3x longer) @@ -289,11 +337,16 @@ def run_step_2_literature_fetch(config: PipelineConfig) -> bool: if config.regulator_file and config.regulator_file.exists(): cmd.extend(["--regulator-file", str(config.regulator_file)]) - cmd.extend([ - "--top-positive-regulators", str(config.top_positive_regulators), - "--top-negative-regulators", str(config.top_negative_regulators), - "--regulator-significance-threshold", str(config.regulator_significance_threshold), - ]) + cmd.extend( + [ + "--top-positive-regulators", + str(config.top_positive_regulators), + "--top-negative-regulators", + str(config.top_negative_regulators), + "--regulator-significance-threshold", + str(config.regulator_significance_threshold), + ] + ) if config.topics: cmd.extend(["--topics", ",".join(map(str, config.topics))]) @@ -331,31 +384,51 @@ def run_step_3a_batch_prepare(config: PipelineConfig) -> bool: search_keyword = config.keyword cmd = [ - sys.executable, str(script), "prepare", - "--gene-file", str(gene_input), - "--celltype-dir", str(celltype_summary.parent), # Fallback directory - "--enrichment-file", str(config.get_path("enrichment_filtered")), - "--ncbi-file", str(config.get_path("ncbi_json")), - "--output-file", str(config.get_path("batch_request_json")), - "--top-loading", str(config.top_loading), - "--top-unique", str(config.top_unique), - "--top-enrichment", str(config.top_enrichment), - "--genes-per-term", str(config.genes_per_term), - "--annotation-role", str(config.annotation_role), - "--annotation-context", str(config.annotation_context), - "--search-keyword", str(search_keyword), + sys.executable, + str(script), + "prepare", + "--gene-file", + str(gene_input), + "--celltype-dir", + str(celltype_summary.parent), # Fallback directory + "--enrichment-file", + str(config.get_path("enrichment_filtered")), + "--ncbi-file", + str(config.get_path("ncbi_json")), + "--output-file", + str(config.get_path("batch_request_json")), + "--top-loading", + str(config.top_loading), + "--top-unique", + str(config.top_unique), + "--top-enrichment", + str(config.top_enrichment), + "--genes-per-term", + str(config.genes_per_term), + "--annotation-role", + str(config.annotation_role), + "--annotation-context", + str(config.annotation_context), + "--search-keyword", + str(search_keyword), ] if celltype_summary.exists(): cmd.extend(["--celltype-file", str(celltype_summary)]) if config.regulator_file: - cmd.extend([ - "--regulator-file", str(config.regulator_file), - "--top-positive-regulators", str(config.top_positive_regulators), - "--top-negative-regulators", str(config.top_negative_regulators), - "--regulator-significance-threshold", str(config.regulator_significance_threshold), - ]) + cmd.extend( + [ + "--regulator-file", + str(config.regulator_file), + "--top-positive-regulators", + str(config.top_positive_regulators), + "--top-negative-regulators", + str(config.top_negative_regulators), + "--regulator-significance-threshold", + str(config.regulator_significance_threshold), + ] + ) if config.topics: cmd.extend(["--topics", ",".join(map(str, config.topics))]) @@ -370,7 +443,9 @@ def run_step_3a_batch_prepare(config: PipelineConfig) -> bool: return True -def run_step_3b_batch_submit(config: PipelineConfig) -> Optional[tuple[str, Optional[str], bool]]: +def run_step_3b_batch_submit( + config: PipelineConfig, +) -> Optional[tuple[str, Optional[str], bool]]: """Step 3b: Submit batch to LLM backend. Returns (batch_id, results_path, waiting) @@ -445,16 +520,20 @@ def short_model_name(full: str) -> str: if backend == "vertex": # Read batch ID (Vertex stores job name) batch_id_file = config.get_path("batch_request_json").with_suffix(".batch_id") - batch_id = batch_id_file.read_text(encoding="utf-8").strip() if batch_id_file.exists() else None + batch_id = ( + batch_id_file.read_text(encoding="utf-8").strip() + if batch_id_file.exists() + else None + ) waiting = not config.llm_wait # Vertex AI: look for GCS prefix output_base = None for line in result.stdout.splitlines(): if "Results available at:" in line or "Output:" in line: - match = re.search(r'(gs://[^\s]+)', line) + match = re.search(r"(gs://[^\s]+)", line) if match: - output_base = match.group(1).rstrip('/') + output_base = match.group(1).rstrip("/") break output_path = None @@ -465,7 +544,7 @@ def short_model_name(full: str) -> str: if list_result.returncode == 0: for gcs_line in list_result.stdout.splitlines(): if "prediction-model-" in gcs_line: - output_path = gcs_line.rstrip('/') + output_path = gcs_line.rstrip("/") break if not output_path: output_path = output_base @@ -499,7 +578,9 @@ def short_model_name(full: str) -> str: return batch_id or "", results_path, waiting -def run_step_4_parse_results(config: PipelineConfig, results_path: Optional[str] = None) -> bool: +def run_step_4_parse_results( + config: PipelineConfig, results_path: Optional[str] = None +) -> bool: """Step 4: Parse LLM responses and generate summary. Args: @@ -520,10 +601,14 @@ def run_step_4_parse_results(config: PipelineConfig, results_path: Optional[str] gene_input = config.gene_loading cmd = [ - sys.executable, str(script), - "--markdown-dir", str(config.get_path("annotations_dir")), - "--summary-csv", str(config.get_path("summary_csv")), - "--gene-loading-file", str(gene_input), + sys.executable, + str(script), + "--markdown-dir", + str(config.get_path("annotations_dir")), + "--summary-csv", + str(config.get_path("summary_csv")), + "--gene-loading-file", + str(gene_input), ] if results_path: @@ -559,19 +644,29 @@ def run_step_5_html_report(config: PipelineConfig) -> bool: gene_input = config.gene_loading cmd = [ - sys.executable, str(script), - "--summary-csv", str(config.get_path("summary_csv")), - "--annotations-dir", str(config.get_path("annotations_dir")), - "--enrichment-dir", str(config.get_path("enrichment_figures")), - "--gene-loading-csv", str(gene_input), - "--output-html", str(config.get_path("report_html")), + sys.executable, + str(script), + "--summary-csv", + str(config.get_path("summary_csv")), + "--annotations-dir", + str(config.get_path("annotations_dir")), + "--enrichment-dir", + str(config.get_path("enrichment_figures")), + "--gene-loading-csv", + str(gene_input), + "--output-html", + str(config.get_path("report_html")), ] if config.regulator_file and config.regulator_file.exists(): - cmd.extend([ - "--volcano-csv", str(config.regulator_file), - "--regulator-significance-threshold", str(config.regulator_significance_threshold), - ]) + cmd.extend( + [ + "--volcano-csv", + str(config.regulator_file), + "--regulator-significance-threshold", + str(config.regulator_significance_threshold), + ] + ) logger.info(f"Running: {' '.join(cmd)}") result = subprocess.run(cmd, capture_output=False) @@ -628,7 +723,9 @@ def run_pipeline( ) return False if state.config_hash != config_hash and force_restart: - logger.warning("Config changed; restarting pipeline and overwriting prior outputs after warning.") + logger.warning( + "Config changed; restarting pipeline and overwriting prior outputs after warning." + ) state = init_state(config_hash, config.topics) save_state(state_path, state) else: @@ -759,7 +856,9 @@ def run_pipeline( # If no results path but batch_id exists, check/fetch batch status if not results_path and state.batch.get("batch_id"): batch_id = state.batch["batch_id"] - logger.info("No results path found. Checking batch status for %s...", batch_id) + logger.info( + "No results path found. Checking batch status for %s...", batch_id + ) from batch_utils import check_anthropic_batch, check_vertex_batch @@ -773,10 +872,14 @@ def run_pipeline( results_dest = config.get_path("batch_request_json").with_name( config.get_path("batch_request_json").stem + "_results.jsonl" ) - status, fetched_path = check_anthropic_batch(batch_id, str(results_dest)) + status, fetched_path = check_anthropic_batch( + batch_id, str(results_dest) + ) if status == "in_progress": - logger.info("Batch %s still processing. Rerun later to continue.", batch_id) + logger.info( + "Batch %s still processing. Rerun later to continue.", batch_id + ) save_state(state_path, state) return True # Exit cleanly, not an error @@ -794,7 +897,9 @@ def run_pipeline( logger.info("Batch results downloaded to: %s", fetched_path) if not results_path: - logger.error("No batch results available. Provide --gcs-prefix or ensure batch completed.") + logger.error( + "No batch results available. Provide --gcs-prefix or ensure batch completed." + ) mark_step(state, step, "failed") save_state(state_path, state) return False @@ -864,69 +969,52 @@ def main(): batch_submit Step 3b: Submit to LLM backend (Anthropic/Vertex) parse_results Step 4: Parse LLM responses html_report Step 5: Generate HTML report -""" +""", ) # Config file - parser.add_argument( - "--config", "-c", - help="Path to YAML config file" - ) + parser.add_argument("--config", "-c", help="Path to YAML config file") # Required inputs (can be overridden by CLI) - parser.add_argument( - "--gene-loading", - help="Path to gene loading matrix CSV" - ) + parser.add_argument("--gene-loading", help="Path to gene loading matrix CSV") parser.add_argument( "--celltype-enrichment", - help="Path to raw cell-type enrichment CSV (summary auto-generated)" - ) - parser.add_argument( - "--output-dir", - help="Output directory for all results" + help="Path to raw cell-type enrichment CSV (summary auto-generated)", ) + parser.add_argument("--output-dir", help="Output directory for all results") # Optional inputs parser.add_argument( - "--regulator-file", - help="Path to SCEPTRE regulator results CSV" + "--regulator-file", help="Path to SCEPTRE regulator results CSV" ) # Step control parser.add_argument( - "--start-from", - choices=STEPS, - help="Step to start from (skips earlier steps)" + "--start-from", choices=STEPS, help="Step to start from (skips earlier steps)" ) parser.add_argument( - "--stop-after", - choices=STEPS, - help="Step to stop after (skips later steps)" + "--stop-after", choices=STEPS, help="Step to stop after (skips later steps)" ) parser.add_argument( "--gcs-prefix", - help="GCS prefix for batch results (for resuming at parse_results)" + help="GCS prefix for batch results (for resuming at parse_results)", ) # Optional settings parser.add_argument( - "--topics", - help="Comma-separated list of topic IDs to process (default: all)" + "--topics", help="Comma-separated list of topic IDs to process (default: all)" ) parser.add_argument( "--species", type=int, default=10090, - help="NCBI taxonomy ID (default: 10090 for mouse)" + help="NCBI taxonomy ID (default: 10090 for mouse)", ) parser.add_argument( - "--keyword", - help="PubMed search keyword for your tissue/cell type" + "--keyword", help="PubMed search keyword for your tissue/cell type" ) parser.add_argument( - "--annotation-role", - help="Specialist role used in the LLM prompt header" + "--annotation-role", help="Specialist role used in the LLM prompt header" ) parser.add_argument( "--annotation-context", @@ -953,21 +1041,22 @@ def main(): parser.add_argument( "--no-resume", action="store_true", - help="Disable resume/caching (re-query all APIs)" + help="Disable resume/caching (re-query all APIs)", ) parser.add_argument( "--wait", action="store_true", - help="Wait for batch completion (default: submit and exit; resume later)") + help="Wait for batch completion (default: submit and exit; resume later)", + ) parser.add_argument( "--force-restart", action="store_true", - help="Ignore existing state (warn) and restart pipeline, potentially overwriting outputs" + help="Ignore existing state (warn) and restart pipeline, potentially overwriting outputs", ) parser.add_argument( "--restart-from", choices=STEPS, - help="Re-run from the specified step (warns about overwriting downstream outputs)" + help="Re-run from the specified step (warns about overwriting downstream outputs)", ) args = parser.parse_args() @@ -983,12 +1072,15 @@ def main(): ) config = PipelineConfig( gene_loading=Path(args.gene_loading), - celltype_enrichment=Path(args.celltype_enrichment) if args.celltype_enrichment else None, + celltype_enrichment=Path(args.celltype_enrichment) + if args.celltype_enrichment + else None, regulator_file=Path(args.regulator_file) if args.regulator_file else None, output_dir=Path(args.output_dir), top_positive_regulators=args.top_positive_regulators or 3, top_negative_regulators=args.top_negative_regulators or 3, - regulator_significance_threshold=args.regulator_significance_threshold or 0.05, + regulator_significance_threshold=args.regulator_significance_threshold + or 0.05, ) # Apply CLI overrides diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_api.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_api.py index 3568f4e..f9ce57d 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_api.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_api.py @@ -37,37 +37,35 @@ def get_string_ids( - identifiers: List[str], - species: int = MOUSE_TAXON, - limit: int = 1 + identifiers: List[str], species: int = MOUSE_TAXON, limit: int = 1 ) -> Dict[str, str]: """Map gene names to STRING identifiers. - + Args: identifiers: List of gene names/symbols species: NCBI taxon ID (10090 for mouse) limit: Max matches per identifier - + Returns: Dict mapping input name -> STRING ID (e.g., "Fzd4" -> "10090.ENSMUSP00000029156") """ if not identifiers: return {} - + request_url = f"{STRING_API_URL}/json/get_string_ids" - + params = { "identifiers": "\r".join(identifiers), "species": species, "limit": limit, - "caller_identity": CALLER_IDENTITY + "caller_identity": CALLER_IDENTITY, } - + try: response = requests.post(request_url, data=params, timeout=30) response.raise_for_status() results = response.json() - + # Map input name to STRING ID mapping = {} for r in results: @@ -76,9 +74,9 @@ def get_string_ids( string_id = r.get("stringId", "") if pref_name and string_id: mapping[pref_name] = string_id - + return mapping - + except Exception as e: logger.warning(f"STRING ID mapping failed: {e}") return {} @@ -88,37 +86,37 @@ def get_network_interactions( identifiers: List[str], species: int = MOUSE_TAXON, required_score: int = 400, - network_type: str = "functional" + network_type: str = "functional", ) -> List[Dict[str, Any]]: """Get interactions between a set of proteins. - + Args: identifiers: List of gene names species: NCBI taxon ID required_score: Minimum combined score (0-1000, default 400 = medium confidence) network_type: "functional" or "physical" - + Returns: List of interaction dicts with scores for each evidence channel """ if len(identifiers) < 2: return [] - + request_url = f"{STRING_API_URL}/json/network" - + params = { "identifiers": "\r".join(identifiers), "species": species, "required_score": required_score, "network_type": network_type, - "caller_identity": CALLER_IDENTITY + "caller_identity": CALLER_IDENTITY, } - + try: response = requests.post(request_url, data=params, timeout=60) response.raise_for_status() return response.json() - + except Exception as e: logger.warning(f"STRING network query failed: {e}") return [] @@ -128,34 +126,34 @@ def get_interaction_partners( identifiers: List[str], species: int = MOUSE_TAXON, limit: int = 50, - required_score: int = 400 + required_score: int = 400, ) -> List[Dict[str, Any]]: """Get all STRING interaction partners of proteins. - + Args: identifiers: List of gene names - species: NCBI taxon ID + species: NCBI taxon ID limit: Max partners per protein required_score: Minimum combined score - + Returns: List of interaction dicts """ request_url = f"{STRING_API_URL}/json/interaction_partners" - + params = { "identifiers": "\r".join(identifiers), "species": species, "limit": limit, "required_score": required_score, - "caller_identity": CALLER_IDENTITY + "caller_identity": CALLER_IDENTITY, } - + try: response = requests.post(request_url, data=params, timeout=60) response.raise_for_status() return response.json() - + except Exception as e: logger.warning(f"STRING interaction_partners query failed: {e}") return [] @@ -167,14 +165,14 @@ def get_regulator_program_interactions( species: int = MOUSE_TAXON, required_score: int = 400, top_n: int = 10, - partner_limit: int = 500 + partner_limit: int = 500, ) -> Dict[str, Any]: """Query STRING for interactions between a regulator and program genes. - + This is the main function for validating regulator-program relationships. - It queries STRING for the regulator and finds which program genes it + It queries STRING for the regulator and finds which program genes it directly interacts with according to STRING evidence. - + Args: regulator: Gene symbol of the regulator (e.g., "Fzd4") program_genes: List of program's gene symbols (can be up to 300) @@ -182,7 +180,7 @@ def get_regulator_program_interactions( required_score: Minimum STRING combined score (0-1000) top_n: Max interactions to return in output partner_limit: Max partners to fetch from STRING (default 500) - + Returns: Dict with: - regulator: str @@ -194,25 +192,25 @@ def get_regulator_program_interactions( "regulator": regulator, "interactions": [], "n_interactions": 0, - "program_genes_with_evidence": [] + "program_genes_with_evidence": [], } - + # Query STRING for regulator's interaction partners # Use higher limit to search against larger program gene lists partners = get_interaction_partners( identifiers=[regulator], species=species, limit=partner_limit, - required_score=required_score + required_score=required_score, ) - + if not partners: logger.debug(f"No STRING partners found for {regulator}") return result - + # Create lowercase set of program genes for matching program_set = {g.lower() for g in program_genes} - + # Filter to only interactions with program genes program_interactions = [] for p in partners: @@ -221,22 +219,26 @@ def get_regulator_program_interactions( if partner_name.lower() in program_set: interaction = { "target_gene": partner_name, - "score": int(p.get("score", 0) * 1000) if p.get("score", 0) <= 1 else int(p.get("score", 0)), + "score": int(p.get("score", 0) * 1000) + if p.get("score", 0) <= 1 + else int(p.get("score", 0)), "experimental_score": int(float(p.get("escore", 0)) * 1000), "database_score": int(float(p.get("dscore", 0)) * 1000), "textmining_score": int(float(p.get("tscore", 0)) * 1000), "coexpression_score": int(float(p.get("ascore", 0)) * 1000), } program_interactions.append(interaction) - + # Sort by score and take top N program_interactions.sort(key=lambda x: -x["score"]) program_interactions = program_interactions[:top_n] - + result["interactions"] = program_interactions result["n_interactions"] = len(program_interactions) - result["program_genes_with_evidence"] = [i["target_gene"] for i in program_interactions] - + result["program_genes_with_evidence"] = [ + i["target_gene"] for i in program_interactions + ] + return result @@ -244,18 +246,18 @@ def batch_validate_regulators( regulator_genes: List[str], program_genes: List[str], species: int = MOUSE_TAXON, - required_score: int = 400 + required_score: int = 400, ) -> Dict[str, List[Dict[str, Any]]]: """Batch query STRING for all regulator-program interactions in ONE call. - + Much faster than individual queries - one API call per program. - + Args: regulator_genes: List of regulator gene symbols program_genes: List of program gene symbols species: NCBI taxon ID required_score: Minimum combined score (0-1000, 400 = medium confidence) - + Returns: Dict mapping regulator -> list of interactions with program genes { @@ -265,39 +267,43 @@ def batch_validate_regulators( } """ result = {reg: [] for reg in regulator_genes} - + if not regulator_genes or not program_genes: return result - + # Combine all genes for single query all_genes = list(set(regulator_genes + program_genes)) - - logger.info(f" Batch STRING query: {len(regulator_genes)} regulators x {len(program_genes)} program genes...") - + + logger.info( + f" Batch STRING query: {len(regulator_genes)} regulators x {len(program_genes)} program genes..." + ) + # Query STRING network for all interactions interactions = get_network_interactions( - identifiers=all_genes, - species=species, - required_score=required_score + identifiers=all_genes, species=species, required_score=required_score ) - + if not interactions: logger.warning(" No STRING interactions returned") return result - + # Create sets for fast lookup regulator_set = {g.lower() for g in regulator_genes} program_set = {g.lower() for g in program_genes} - + # Parse interactions: keep only regulator <-> program gene pairs for i in interactions: gene_a = i.get("preferredName_A", "") gene_b = i.get("preferredName_B", "") - score = int(i.get("score", 0) * 1000) if i.get("score", 0) <= 1 else int(i.get("score", 0)) - + score = ( + int(i.get("score", 0) * 1000) + if i.get("score", 0) <= 1 + else int(i.get("score", 0)) + ) + # Check if this is a regulator-program interaction a_lower, b_lower = gene_a.lower(), gene_b.lower() - + if a_lower in regulator_set and b_lower in program_set: # gene_a is regulator, gene_b is program gene for reg in regulator_genes: @@ -310,15 +316,17 @@ def batch_validate_regulators( if reg.lower() == b_lower: result[reg].append({"target": gene_a, "score": score}) break - + # Sort each regulator's interactions by score for reg in result: result[reg].sort(key=lambda x: -x["score"]) - + n_total = sum(len(v) for v in result.values()) n_with_int = sum(1 for v in result.values() if v) - logger.info(f" Found {n_total} interactions for {n_with_int}/{len(regulator_genes)} regulators") - + logger.info( + f" Found {n_total} interactions for {n_with_int}/{len(regulator_genes)} regulators" + ) + return result @@ -327,115 +335,120 @@ def validate_regulators_with_string( program_genes: List[str], species: int = MOUSE_TAXON, top_regulators: int = 3, - required_score: int = 400 + required_score: int = 400, ) -> Dict[str, Any]: """Validate multiple regulators against program genes using STRING. - + Args: regulators: List of regulator dicts with 'gene', 'log2FC', etc. program_genes: List of program gene symbols species: NCBI taxon ID top_regulators: Number of top regulators per category to validate required_score: Minimum STRING score - + Returns: Dict with 'activators' and 'repressors', each containing validated regulator info with STRING interactions. """ - result = { - "activators": [], - "repressors": [] - } - + result = {"activators": [], "repressors": []} + # Split into activators (negative log2FC) and repressors (positive log2FC) activators = [r for r in regulators if r.get("log2FC", 0) < 0] repressors = [r for r in regulators if r.get("log2FC", 0) > 0] - + # Sort by absolute effect size activators.sort(key=lambda x: x.get("log2FC", 0)) # Most negative first repressors.sort(key=lambda x: -x.get("log2FC", 0)) # Most positive first - + # Validate top activators for i, reg in enumerate(activators[:top_regulators]): gene = reg.get("gene", "") logger.info(f" Validating activator {gene} with STRING...") - + time.sleep(RATE_LIMIT_SECONDS) interactions = get_regulator_program_interactions( regulator=gene, program_genes=program_genes, species=species, - required_score=required_score + required_score=required_score, ) - - result["activators"].append({ - "gene": gene, - "log2FC": reg.get("log2FC", 0), - "pvalue": reg.get("pvalue"), - "string_interactions": interactions["interactions"], - "n_program_targets": interactions["n_interactions"] - }) - + + result["activators"].append( + { + "gene": gene, + "log2FC": reg.get("log2FC", 0), + "pvalue": reg.get("pvalue"), + "string_interactions": interactions["interactions"], + "n_program_targets": interactions["n_interactions"], + } + ) + # Add remaining activators without validation - for reg in activators[top_regulators:top_regulators + 5]: - result["activators"].append({ - "gene": reg.get("gene", ""), - "log2FC": reg.get("log2FC", 0), - "pvalue": reg.get("pvalue"), - "string_interactions": None, # Not validated - "n_program_targets": None - }) - + for reg in activators[top_regulators : top_regulators + 5]: + result["activators"].append( + { + "gene": reg.get("gene", ""), + "log2FC": reg.get("log2FC", 0), + "pvalue": reg.get("pvalue"), + "string_interactions": None, # Not validated + "n_program_targets": None, + } + ) + # Validate top repressors for i, reg in enumerate(repressors[:top_regulators]): gene = reg.get("gene", "") logger.info(f" Validating repressor {gene} with STRING...") - + time.sleep(RATE_LIMIT_SECONDS) interactions = get_regulator_program_interactions( regulator=gene, program_genes=program_genes, species=species, - required_score=required_score + required_score=required_score, + ) + + result["repressors"].append( + { + "gene": gene, + "log2FC": reg.get("log2FC", 0), + "pvalue": reg.get("pvalue"), + "string_interactions": interactions["interactions"], + "n_program_targets": interactions["n_interactions"], + } ) - - result["repressors"].append({ - "gene": gene, - "log2FC": reg.get("log2FC", 0), - "pvalue": reg.get("pvalue"), - "string_interactions": interactions["interactions"], - "n_program_targets": interactions["n_interactions"] - }) - - # Add remaining repressors without validation - for reg in repressors[top_regulators:top_regulators + 5]: - result["repressors"].append({ - "gene": reg.get("gene", ""), - "log2FC": reg.get("log2FC", 0), - "pvalue": reg.get("pvalue"), - "string_interactions": None, - "n_program_targets": None - }) - + + # Add remaining repressors without validation + for reg in repressors[top_regulators : top_regulators + 5]: + result["repressors"].append( + { + "gene": reg.get("gene", ""), + "log2FC": reg.get("log2FC", 0), + "pvalue": reg.get("pvalue"), + "string_interactions": None, + "n_program_targets": None, + } + ) + return result # Test function if __name__ == "__main__": logging.basicConfig(level=logging.INFO) - + # Test with Fzd4 and some known Wnt pathway genes test_genes = ["Ctnnb1", "Lef1", "Wnt7a", "Wnt7b", "Axin2", "Vegfa", "Kdr"] - + print("Testing STRING API...") print(f"Query: Fzd4 interactions with {test_genes}") - + result = get_regulator_program_interactions( - regulator="Fzd4", - program_genes=test_genes, - species=10090 + regulator="Fzd4", program_genes=test_genes, species=10090 ) - + print(f"\nFound {result['n_interactions']} interactions:") for i in result["interactions"]: - print(f" Fzd4 -- {i['target_gene']} (score={i['score']}, exp={i['experimental_score']}, db={i['database_score']})") + print( + f" Fzd4 -- {i['target_gene']} (score={i['score']}, exp={i['experimental_score']}, db={i['database_score']})" + ) diff --git a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_enrichment.py b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_enrichment.py index 47b8db5..3024f57 100644 --- a/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_enrichment.py +++ b/src/Stage3_Interpretation/C_Annotation/ProgramExplorer/src/string_enrichment.py @@ -7,6 +7,7 @@ - Full and filtered (Process/KEGG) CSV construction - Enrichment figure download from STRING API """ + from __future__ import annotations import json @@ -31,7 +32,9 @@ def cache_path(cache_dir: Path, program_id: int) -> Path: return cache_dir / f"program_{program_id}_enrichment.json" -def load_cached_results(cache_dir: Path, program_id: int) -> Optional[List[Dict[str, Any]]]: +def load_cached_results( + cache_dir: Path, program_id: int +) -> Optional[List[Dict[str, Any]]]: cache_file = cache_path(cache_dir, program_id) if not cache_file.exists(): return None @@ -46,7 +49,9 @@ def load_cached_results(cache_dir: Path, program_id: int) -> Optional[List[Dict[ return data -def write_cached_results(cache_dir: Path, program_id: int, results: List[Dict[str, Any]]) -> None: +def write_cached_results( + cache_dir: Path, program_id: int, results: List[Dict[str, Any]] +) -> None: cache_dir.mkdir(parents=True, exist_ok=True) cache_file = cache_path(cache_dir, program_id) cache_file.write_text(json.dumps(results, indent=2), encoding="utf-8") @@ -77,9 +82,13 @@ def call_string_enrichment( data = [] return data if isinstance(data, list) else [] else: - logger.warning(f"STRING returned status {response.status_code}: {response.text[:200]}") + logger.warning( + f"STRING returned status {response.status_code}: {response.text[:200]}" + ) except requests.RequestException as e: - logger.warning(f"HTTP error on STRING request (attempt {attempt+1}/{retries+1}): {e}") + logger.warning( + f"HTTP error on STRING request (attempt {attempt + 1}/{retries + 1}): {e}" + ) attempt += 1 time.sleep(min(2.0, sleep_between * (attempt + 1))) @@ -104,9 +113,13 @@ def build_full_csv(program_to_results: Dict[str, List[Dict[str, Any]]]) -> pd.Da "fdr": float(t.get("fdr", float("nan"))), "p_value": float(t.get("p_value", float("nan"))), "number_of_genes": int(t.get("number_of_genes", 0)), - "number_of_genes_in_background": int(t.get("number_of_genes_in_background", 0)), + "number_of_genes_in_background": int( + t.get("number_of_genes_in_background", 0) + ), "ncbiTaxonId": int(t.get("ncbiTaxonId", 0)), - "inputGenes": "|".join(t.get("inputGenes", [])) if t.get("inputGenes") else "", + "inputGenes": "|".join(t.get("inputGenes", [])) + if t.get("inputGenes") + else "", } ) df = pd.DataFrame(rows) diff --git a/src/Stage3_Interpretation/C_Annotation/api/LLM_interface.py b/src/Stage3_Interpretation/C_Annotation/api/LLM_interface.py index 5710cc4..c5dd4f3 100644 --- a/src/Stage3_Interpretation/C_Annotation/api/LLM_interface.py +++ b/src/Stage3_Interpretation/C_Annotation/api/LLM_interface.py @@ -1,4 +1,4 @@ -#%% +# %% import anthropic from pydantic import ValidationError from pydantic import BaseModel @@ -10,38 +10,45 @@ class ProgramAnnotation(BaseModel): confidence: float rationale: str + class claude_api: def __init__(self, api_key, model): self.client = anthropic.Anthropic(api_key=api_key, max_retries=3, timeout=600) self.model = model def get_completion(self, system_prompt, prompt, effort="medium", max_tokens=4000): - try: + try: msg = self.client.messages.create( - model=self.model, - max_tokens=max_tokens, + model=self.model, + max_tokens=max_tokens, system=system_prompt, messages=[{"role": "user", "content": prompt}], - output_config = {"effort": effort}, # "low" | "medium" | "high" | "xhigh" | "max" + output_config={ + "effort": effort + }, # "low" | "medium" | "high" | "xhigh" | "max" ) if msg.stop_reason == "max_tokens": print("Warning: answer was cut off") - + return "".join(b.text for b in msg.content if b.type == "text") - + except anthropic.APIError as e: print(e) return None - def get_json(self, system_prompt, prompt, schema_model, effort="medium", max_tokens=4000): - try: + def get_json( + self, system_prompt, prompt, schema_model, effort="medium", max_tokens=4000 + ): + try: msg = self.client.messages.parse( - model=self.model, - max_tokens=max_tokens, + model=self.model, + max_tokens=max_tokens, system=system_prompt, messages=[{"role": "user", "content": prompt}], output_format=schema_model, - output_config = {"effort": effort}, # "low" | "medium" | "high" | "xhigh" | "max" + output_config={ + "effort": effort + }, # "low" | "medium" | "high" | "xhigh" | "max" ) except ValidationError as e: @@ -60,7 +67,7 @@ def get_json(self, system_prompt, prompt, schema_model, effort="medium", max_tok print("Claude refused this request") return None + return msg.parsed_output - return msg.parsed_output # %% diff --git a/tests/Script/Stage1_Inference/conftest.py b/tests/Script/Stage1_Inference/conftest.py index d32aa9a..f2bd689 100644 --- a/tests/Script/Stage1_Inference/conftest.py +++ b/tests/Script/Stage1_Inference/conftest.py @@ -108,13 +108,20 @@ def check_data_format( print(f"Found adata.obsm['{guide_assignment_key}']\n") try: import scipy.sparse as sp + if sp.issparse(guide_assignment): - print(f"WARNING: '{guide_assignment_key}' is sparse. Converting to dense array...") + print( + f"WARNING: '{guide_assignment_key}' is sparse. Converting to dense array..." + ) dense_array = guide_assignment.toarray() adata.obsm[guide_assignment_key] = dense_array - print(f"'{guide_assignment_key}' converted to dense array (shape: {dense_array.shape})\n") + print( + f"'{guide_assignment_key}' converted to dense array (shape: {dense_array.shape})\n" + ) else: - print(f"'{guide_assignment_key}' is already dense (shape: {guide_assignment.shape})\n") + print( + f"'{guide_assignment_key}' is already dense (shape: {guide_assignment.shape})\n" + ) except Exception as e: print(f"WARNING: Error checking '{guide_assignment_key}' sparsity: {e}\n") is_valid = False diff --git a/tests/Script/Stage1_Inference/test_edge_cases.py b/tests/Script/Stage1_Inference/test_edge_cases.py index 176bcc9..8096875 100644 --- a/tests/Script/Stage1_Inference/test_edge_cases.py +++ b/tests/Script/Stage1_Inference/test_edge_cases.py @@ -113,12 +113,13 @@ def _assert_basic_outputs(output_dir): # Tests # --------------------------------------------------------------------------- -class TestEdgeCases: +class TestEdgeCases: def test_minibatch_shuffle(self, mini_h5ad_path): """Minibatch mode with shuffle=True completes.""" output_dir = _run_minimal( - mini_h5ad_path, "minibatch_shuffle", + mini_h5ad_path, + "minibatch_shuffle", mode="minibatch", minibatch_shuffle=True, minibatch_size=500, @@ -132,7 +133,8 @@ def test_minibatch_shuffle(self, mini_h5ad_path): def test_minibatch_no_shuffle(self, mini_h5ad_path): """Minibatch mode with shuffle=False completes.""" output_dir = _run_minimal( - mini_h5ad_path, "minibatch_no_shuffle", + mini_h5ad_path, + "minibatch_no_shuffle", mode="minibatch", minibatch_shuffle=False, minibatch_size=500, @@ -151,7 +153,8 @@ def test_custom_seeds(self, mini_h5ad_path): np.save(seeds_path, seeds) output_dir = _run_minimal( - mini_h5ad_path, "custom_seeds", + mini_h5ad_path, + "custom_seeds", nmf_seeds=seeds, ) _assert_basic_outputs(output_dir) @@ -159,7 +162,8 @@ def test_custom_seeds(self, mini_h5ad_path): def test_cpu_only(self, mini_h5ad_path): """use_gpu=False falls back to CPU without error.""" output_dir = _run_minimal( - mini_h5ad_path, "cpu_only", + mini_h5ad_path, + "cpu_only", use_gpu=False, ) _assert_basic_outputs(output_dir) @@ -173,7 +177,9 @@ def test_remove_noncoding(self, mini_h5ad_path): n_before = adata.n_vars adata_filtered = adata[:, mask].copy() n_after = adata_filtered.n_vars - print(f" Filtered genes: {n_before} → {n_after} (removed {n_before - n_after} ENSG)") + print( + f" Filtered genes: {n_before} → {n_after} (removed {n_before - n_after} ENSG)" + ) # Save filtered data filtered_dir = os.path.join(OUTPUT_BASE, "remove_noncoding") @@ -182,14 +188,16 @@ def test_remove_noncoding(self, mini_h5ad_path): adata_filtered.write(filtered_path) output_dir = _run_minimal( - filtered_path, "remove_noncoding", + filtered_path, + "remove_noncoding", ) _assert_basic_outputs(output_dir) def test_minibatch_size_larger_than_data(self, mini_h5ad_path): """minibatch_size >> n_cells doesn't crash.""" output_dir = _run_minimal( - mini_h5ad_path, "large_minibatch", + mini_h5ad_path, + "large_minibatch", mode="minibatch", minibatch_size=999999, minibatch_max_epoch=3, @@ -202,7 +210,8 @@ def test_minibatch_size_larger_than_data(self, mini_h5ad_path): def test_numhvgenes_larger_than_total(self, mini_h5ad_path): """num_highvar_genes > total genes clamps gracefully.""" output_dir = _run_minimal( - mini_h5ad_path, "large_numhvgenes", + mini_h5ad_path, + "large_numhvgenes", num_highvar_genes=999999, ) _assert_basic_outputs(output_dir) @@ -210,7 +219,8 @@ def test_numhvgenes_larger_than_total(self, mini_h5ad_path): def test_densify(self, mini_h5ad_path): """densify=True works with dense matrix input.""" output_dir = _run_minimal( - mini_h5ad_path, "densify", + mini_h5ad_path, + "densify", densify=True, ) _assert_basic_outputs(output_dir) @@ -218,7 +228,8 @@ def test_densify(self, mini_h5ad_path): def test_dataloader_mode_edge(self, mini_h5ad_path): """Dataloader mode with small minibatch_size completes.""" output_dir = _run_minimal( - mini_h5ad_path, "dataloader_small_batch", + mini_h5ad_path, + "dataloader_small_batch", mode="dataloader", minibatch_size=100, minibatch_max_epoch=3, @@ -270,16 +281,21 @@ def test_multiple_thresholds(self, mini_h5ad_path): for k in EDGE_K: for thresh in thresholds: ts = str(thresh).replace(".", "_") - usage = os.path.join(inference_dir, f"Inference.usages.k_{k}.dt_{ts}.consensus.txt") + usage = os.path.join( + inference_dir, f"Inference.usages.k_{k}.dt_{ts}.consensus.txt" + ) if not os.path.exists(usage): missing.append(f"usages k={k} thresh={thresh}") - assert not missing, f"Missing outputs for thresholds:\n" + "\n".join(f" {f}" for f in missing) + assert not missing, f"Missing outputs for thresholds:\n" + "\n".join( + f" {f}" for f in missing + ) print(f" All {len(thresholds)} thresholds produced output") def test_regularization(self, mini_h5ad_path): """Non-zero regularization completes without error.""" output_dir = _run_minimal( - mini_h5ad_path, "regularization", + mini_h5ad_path, + "regularization", alpha_usage=0.1, alpha_spectra=0.1, l1_ratio_usage=0.5, @@ -290,7 +306,8 @@ def test_regularization(self, mini_h5ad_path): def test_double_precision(self, mini_h5ad_path): """fp_precision='double' (float64) completes.""" output_dir = _run_minimal( - mini_h5ad_path, "double_precision", + mini_h5ad_path, + "double_precision", fp_precision="double", ) _assert_basic_outputs(output_dir) @@ -310,7 +327,8 @@ def test_reproducibility(self, mini_h5ad_path): data_b = np.load(path_b, allow_pickle=True) np.testing.assert_array_almost_equal( - data_a['data'], data_b['data'], + data_a["data"], + data_b["data"], decimal=5, err_msg=f"Spectra mismatch for k={k} iter={i}", ) diff --git a/tests/Script/Stage1_Inference/test_inference_parallel_sklearn.py b/tests/Script/Stage1_Inference/test_inference_parallel_sklearn.py index c32eeba..0904800 100644 --- a/tests/Script/Stage1_Inference/test_inference_parallel_sklearn.py +++ b/tests/Script/Stage1_Inference/test_inference_parallel_sklearn.py @@ -23,13 +23,21 @@ import pandas as pd import cnmf -from .conftest import TEST_K, TEST_NUMITER, TEST_NUMHVGENES, TEST_SEL_THRESH, TEST_SEED, PERSISTENT_OUTPUT +from .conftest import ( + TEST_K, + TEST_NUMITER, + TEST_NUMHVGENES, + TEST_SEL_THRESH, + TEST_SEED, + PERSISTENT_OUTPUT, +) from Stage1_Inference.src.run_cNMF import rename_all_NMF # ---------- shared state ---------- + class _State: per_k_dirs = {} @@ -41,6 +49,7 @@ class _State: # ---------- tests ---------- + class TestSkCNMFParallel: """Test parallel K factorization and merge.""" @@ -76,15 +85,21 @@ def test_01_prepare_and_factorize_per_k(self, mini_h5ad_path): assert os.path.isdir(cnmf_tmp), f"cnmf_tmp not found for K={k}" for i in range(TEST_NUMITER): - spectra = os.path.join(cnmf_tmp, f"{run_name}.spectra.k_{k}.iter_{i}.df.npz") - assert os.path.exists(spectra), f"Missing spectra K={k} iter={i}: {spectra}" + spectra = os.path.join( + cnmf_tmp, f"{run_name}.spectra.k_{k}.iter_{i}.df.npz" + ) + assert os.path.exists(spectra), ( + f"Missing spectra K={k} iter={i}: {spectra}" + ) state.per_k_dirs[k] = k_dir print(f" K={k}: factorized ({TEST_NUMITER} iters)") def test_02_combine_parallel_results(self): """Combine per-K spectra files into Inference_all/cnmf_tmp/.""" - assert len(state.per_k_dirs) == len(TEST_K), "All K values must be factorized first" + assert len(state.per_k_dirs) == len(TEST_K), ( + "All K values must be factorized first" + ) combined_name = "Inference_all" combined_cnmf_tmp = os.path.join(OUTPUT_DIR, combined_name, "cnmf_tmp") @@ -102,9 +117,15 @@ def test_02_combine_parallel_results(self): missing = [] for k in TEST_K: for i in range(TEST_NUMITER): - f = os.path.join(combined_cnmf_tmp, f"{combined_name}.spectra.k_{k}.iter_{i}.df.npz") + f = os.path.join( + combined_cnmf_tmp, f"{combined_name}.spectra.k_{k}.iter_{i}.df.npz" + ) if not os.path.exists(f): missing.append(f"Inference_all.spectra.k_{k}.iter_{i}.df.npz") - assert not missing, f"Missing {len(missing)} combined spectra:\n" + "\n".join(f" {f}" for f in missing) - print(f" Combined {len(TEST_K)} K values × {TEST_NUMITER} iters into {combined_cnmf_tmp}") + assert not missing, f"Missing {len(missing)} combined spectra:\n" + "\n".join( + f" {f}" for f in missing + ) + print( + f" Combined {len(TEST_K)} K values × {TEST_NUMITER} iters into {combined_cnmf_tmp}" + ) diff --git a/tests/Script/Stage1_Inference/test_inference_parallel_torch.py b/tests/Script/Stage1_Inference/test_inference_parallel_torch.py index 64be84e..1919dd4 100644 --- a/tests/Script/Stage1_Inference/test_inference_parallel_torch.py +++ b/tests/Script/Stage1_Inference/test_inference_parallel_torch.py @@ -24,7 +24,14 @@ import pytest import numpy as np -from .conftest import TEST_K, TEST_NUMITER, TEST_NUMHVGENES, TEST_SEL_THRESH, TEST_SEED, PERSISTENT_OUTPUT +from .conftest import ( + TEST_K, + TEST_NUMITER, + TEST_NUMHVGENES, + TEST_SEL_THRESH, + TEST_SEED, + PERSISTENT_OUTPUT, +) from torch_cnmf import cNMF from Stage1_Inference.src.run_cNMF import rename_all_NMF @@ -32,6 +39,7 @@ # ---------- shared state ---------- + class _State: per_k_dirs = {} @@ -43,6 +51,7 @@ class _State: # ---------- tests ---------- + class TestTorchCNMFParallel: """Test parallel K factorization and merge with torch-cNMF.""" @@ -83,15 +92,21 @@ def test_01_prepare_and_factorize_per_k(self, mini_h5ad_path): assert os.path.isdir(cnmf_tmp), f"cnmf_tmp not found for K={k}" for i in range(TEST_NUMITER): - spectra = os.path.join(cnmf_tmp, f"{run_name}.spectra.k_{k}.iter_{i}.df.npz") - assert os.path.exists(spectra), f"Missing spectra K={k} iter={i}: {spectra}" + spectra = os.path.join( + cnmf_tmp, f"{run_name}.spectra.k_{k}.iter_{i}.df.npz" + ) + assert os.path.exists(spectra), ( + f"Missing spectra K={k} iter={i}: {spectra}" + ) state.per_k_dirs[k] = k_dir print(f" K={k}: factorized ({TEST_NUMITER} iters)") def test_02_combine_parallel_results(self): """Combine per-K spectra files into Inference_all/cnmf_tmp/.""" - assert len(state.per_k_dirs) == len(TEST_K), "All K values must be factorized first" + assert len(state.per_k_dirs) == len(TEST_K), ( + "All K values must be factorized first" + ) combined_name = "Inference_all" combined_cnmf_tmp = os.path.join(OUTPUT_DIR, combined_name, "cnmf_tmp") @@ -109,9 +124,15 @@ def test_02_combine_parallel_results(self): missing = [] for k in TEST_K: for i in range(TEST_NUMITER): - f = os.path.join(combined_cnmf_tmp, f"{combined_name}.spectra.k_{k}.iter_{i}.df.npz") + f = os.path.join( + combined_cnmf_tmp, f"{combined_name}.spectra.k_{k}.iter_{i}.df.npz" + ) if not os.path.exists(f): missing.append(f"Inference_all.spectra.k_{k}.iter_{i}.df.npz") - assert not missing, f"Missing {len(missing)} combined spectra:\n" + "\n".join(f" {f}" for f in missing) - print(f" Combined {len(TEST_K)} K values × {TEST_NUMITER} iters into {combined_cnmf_tmp}") + assert not missing, f"Missing {len(missing)} combined spectra:\n" + "\n".join( + f" {f}" for f in missing + ) + print( + f" Combined {len(TEST_K)} K values × {TEST_NUMITER} iters into {combined_cnmf_tmp}" + ) diff --git a/tests/Script/Stage1_Inference/test_inference_sklearn.py b/tests/Script/Stage1_Inference/test_inference_sklearn.py index 5b2e72a..14f9b29 100644 --- a/tests/Script/Stage1_Inference/test_inference_sklearn.py +++ b/tests/Script/Stage1_Inference/test_inference_sklearn.py @@ -15,7 +15,14 @@ import mudata as mu import cnmf -from .conftest import TEST_K, TEST_NUMITER, TEST_NUMHVGENES, TEST_SEL_THRESH, TEST_SEED, check_data_format +from .conftest import ( + TEST_K, + TEST_NUMITER, + TEST_NUMHVGENES, + TEST_SEL_THRESH, + TEST_SEED, + check_data_format, +) from Stage1_Inference.src.run_cNMF import run_cnmf_consensus, compile_results from Stage1_Inference.src.plot_diagnostics import generate_all_plots @@ -23,8 +30,10 @@ # ---------- shared state across ordered tests ---------- + class _State: """Mutable container shared across tests in this module.""" + cnmf_obj = None run_dir = None @@ -34,6 +43,7 @@ class _State: # ---------- tests (run in order) ---------- + class TestSkCNMFInference: """Ordered end-to-end test for sk-cNMF inference.""" @@ -183,9 +193,13 @@ def test_08_diagnostic_plots(self): thresh_str = str(thresh).replace(".", "_") elbow = os.path.join(plots_dir, f"elbow_curves_k{k}_dt{thresh_str}.pdf") assert os.path.exists(elbow), f"Missing elbow plot: {elbow}" - usage = os.path.join(plots_dir, f"usage_heatmap_k{k}_dt{thresh_str}.pdf") + usage = os.path.join( + plots_dir, f"usage_heatmap_k{k}_dt{thresh_str}.pdf" + ) assert os.path.exists(usage), f"Missing usage heatmap: {usage}" - violin = os.path.join(plots_dir, f"loading_violins_k{k}_dt{thresh_str}.pdf") + violin = os.path.join( + plots_dir, f"loading_violins_k{k}_dt{thresh_str}.pdf" + ) assert os.path.exists(violin), f"Missing violin plot: {violin}" print(f"\nDiagnostic plots saved to: {plots_dir}") @@ -227,12 +241,18 @@ def test_09_assert_all_outputs(self): # K-selection plot if not os.path.exists(os.path.join(inference_dir, "Inference.k_selection.png")): missing.append("Inference.k_selection.png") - if not os.path.exists(os.path.join(inference_dir, "diagnosis_plots", "k_selection.png")): + if not os.path.exists( + os.path.join(inference_dir, "diagnosis_plots", "k_selection.png") + ): missing.append("diagnosis_plots/k_selection.png") # Logs if not os.path.exists(os.path.join(inference_dir, "logs", "test_run.log")): missing.append("logs/test_run.log") - assert not missing, f"Missing {len(missing)} output files:\n" + "\n".join(f" {f}" for f in missing) - print(f"\nAll output files verified ({len(TEST_K) * len(TEST_SEL_THRESH)} K×thresh combos)") + assert not missing, f"Missing {len(missing)} output files:\n" + "\n".join( + f" {f}" for f in missing + ) + print( + f"\nAll output files verified ({len(TEST_K) * len(TEST_SEL_THRESH)} K×thresh combos)" + ) diff --git a/tests/Script/Stage1_Inference/test_inference_torch.py b/tests/Script/Stage1_Inference/test_inference_torch.py index c4f4535..939cd24 100644 --- a/tests/Script/Stage1_Inference/test_inference_torch.py +++ b/tests/Script/Stage1_Inference/test_inference_torch.py @@ -21,8 +21,13 @@ import mudata as mu from .conftest import ( - TEST_K, TEST_NUMITER, TEST_NUMHVGENES, TEST_SEL_THRESH, TEST_SEED, - PERSISTENT_OUTPUT, check_data_format, + TEST_K, + TEST_NUMITER, + TEST_NUMHVGENES, + TEST_SEL_THRESH, + TEST_SEED, + PERSISTENT_OUTPUT, + check_data_format, ) from torch_cnmf import cNMF @@ -34,6 +39,7 @@ # Helper: run the full inference pipeline for a given mode and return output dir # --------------------------------------------------------------------------- + def _make_torch_output_dir(mode): """Create output dir for a torch-cNMF mode test. @@ -160,19 +166,32 @@ def _validate_outputs(output_dir, adata_dir): assert "rna" in mdata.mod, f"k={k}: Missing 'rna' modality" assert "cNMF" in mdata.mod, f"k={k}: Missing 'cNMF' modality" assert mdata["cNMF"].n_obs > 0, f"k={k}: No cells in cNMF modality" - assert mdata["cNMF"].n_vars <= k, f"k={k}: Expected <= {k} programs, got {mdata['cNMF'].n_vars}" - assert "loadings" in mdata["cNMF"].varm, f"k={k}: Missing varm['loadings']" - assert "guide_names" in mdata["cNMF"].uns, f"k={k}: Missing uns['guide_names']" - assert "guide_targets" in mdata["cNMF"].uns, f"k={k}: Missing uns['guide_targets']" - - assert not missing, f"Missing {len(missing)} output files:\n" + "\n".join(f" {f}" for f in missing) - print(f"\nAll output files verified ({len(TEST_K) * len(TEST_SEL_THRESH)} K×thresh combos)") + assert mdata["cNMF"].n_vars <= k, ( + f"k={k}: Expected <= {k} programs, got {mdata['cNMF'].n_vars}" + ) + assert "loadings" in mdata["cNMF"].varm, ( + f"k={k}: Missing varm['loadings']" + ) + assert "guide_names" in mdata["cNMF"].uns, ( + f"k={k}: Missing uns['guide_names']" + ) + assert "guide_targets" in mdata["cNMF"].uns, ( + f"k={k}: Missing uns['guide_targets']" + ) + + assert not missing, f"Missing {len(missing)} output files:\n" + "\n".join( + f" {f}" for f in missing + ) + print( + f"\nAll output files verified ({len(TEST_K) * len(TEST_SEL_THRESH)} K×thresh combos)" + ) # --------------------------------------------------------------------------- # Test: data validation (shared across modes) # --------------------------------------------------------------------------- + class TestDataValid: def test_data_valid(self, mini_h5ad_path): adata = ad.read_h5ad(mini_h5ad_path) @@ -187,6 +206,7 @@ def test_data_valid(self, mini_h5ad_path): # Test: batch mode (full HALS, default) # --------------------------------------------------------------------------- + class TestBatchMode: """End-to-end test for torch-cNMF in batch mode.""" @@ -207,6 +227,7 @@ def test_batch_inference(self, mini_h5ad_path): # Test: minibatch mode # --------------------------------------------------------------------------- + class TestMinibatchMode: """End-to-end test for torch-cNMF in minibatch mode.""" @@ -230,6 +251,7 @@ def test_minibatch_inference(self, mini_h5ad_path): # Test: dataloader mode # --------------------------------------------------------------------------- + class TestDataloaderMode: """End-to-end test for torch-cNMF in dataloader mode.""" diff --git a/tests/Script/Stage2_Evaluation/conftest.py b/tests/Script/Stage2_Evaluation/conftest.py index 1168242..620a9e9 100644 --- a/tests/Script/Stage2_Evaluation/conftest.py +++ b/tests/Script/Stage2_Evaluation/conftest.py @@ -10,7 +10,9 @@ import pandas as pd import mudata -PIPELINE_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +PIPELINE_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) sys.path.insert(0, os.path.join(PIPELINE_ROOT, "src")) _TEST_OUTPUT = os.path.join(PIPELINE_ROOT, "tests", "output") @@ -22,12 +24,14 @@ os.path.join(_TEST_OUTPUT, "sk-cNMF", "Inference"), ] + def _find_inference_path(): for path in _INFERENCE_CANDIDATES: if os.path.isdir(path): return path return _INFERENCE_CANDIDATES[0] # default even if missing (will skip at fixture) + DEFAULT_INFERENCE_PATH = _find_inference_path() EVAL_OUTPUT_DIR = os.path.join(os.path.dirname(DEFAULT_INFERENCE_PATH), "Evaluation") @@ -111,20 +115,22 @@ def eval_output_dir_per_k(mdata_copy_per_k): def cnmf_obj(inference_path): """Create a cNMF object pointing to the test inference output.""" from torch_cnmf import cNMF + run_dir = os.path.dirname(inference_path) # parent of Inference/ - return cNMF(output_dir=run_dir, name='Inference') + return cNMF(output_dir=run_dir, name="Inference") @pytest.fixture(scope="session") def X_norm(inference_path): """Load the normalized counts matrix from inference output.""" import scanpy as sc + norm_path = os.path.join(inference_path, "cnmf_tmp", "Inference.norm_counts.h5ad") if not os.path.isfile(norm_path): pytest.skip(f"Normalized counts not found: {norm_path}") adata = sc.read_h5ad(norm_path) X = adata.X - if hasattr(X, 'toarray'): + if hasattr(X, "toarray"): X = X.toarray() return X @@ -149,13 +155,17 @@ def reference_targets_from_annotation(guide_annotation_path, test_mdata): df = pd.read_csv(guide_annotation_path, sep="\t", index_col=0) non_targeting_guides = df[df["targeting"] == False].index.values # Intersect annotation guide names with mdata guide_names, then get target group names - guide_metadata = pd.DataFrame({ - "guide": test_mdata["cNMF"].uns["guide_names"], - "target": test_mdata["cNMF"].uns["guide_targets"], - }) + guide_metadata = pd.DataFrame( + { + "guide": test_mdata["cNMF"].uns["guide_names"], + "target": test_mdata["cNMF"].uns["guide_targets"], + } + ) matched = guide_metadata[guide_metadata["guide"].isin(non_targeting_guides)] - print(f"[conftest] annotation non-targeting guides: {len(non_targeting_guides)}, " - f"matched in mdata: {len(matched)}") + print( + f"[conftest] annotation non-targeting guides: {len(non_targeting_guides)}, " + f"matched in mdata: {len(matched)}" + ) ref_targets = matched["target"].unique().tolist() print(f"[conftest] resolved reference target groups: {ref_targets}") if len(ref_targets) == 0: @@ -171,6 +181,7 @@ def reference_targets_from_annotation(guide_annotation_path, test_mdata): # Calibration fixtures (use same real data as metrics tests) # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def calibration_output_dir(): """Output directory for calibration tests — saved inside Evaluation/.""" diff --git a/tests/Script/Stage2_Evaluation/test_crt.py b/tests/Script/Stage2_Evaluation/test_crt.py index 267fa68..6e0f64c 100644 --- a/tests/Script/Stage2_Evaluation/test_crt.py +++ b/tests/Script/Stage2_Evaluation/test_crt.py @@ -22,14 +22,24 @@ # CRT import (skip entire module if sceptre unavailable) # --------------------------------------------------------------------------- -CRT_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", "src", - "Stage2_Evaluation", "B_Calibration", "Slurm_version", "CRT") +CRT_DIR = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "src", + "Stage2_Evaluation", + "B_Calibration", + "Slurm_version", + "CRT", +) CRT_DIR = os.path.abspath(CRT_DIR) try: sys.path.insert(0, CRT_DIR) import CRT as crt_mod # noqa: F401 — needs sceptre (programDE env) from CRT import reformat_data_for_CRT, run_CRT, save_result, main as crt_main + _HAS_CRT = True except ImportError: _HAS_CRT = False @@ -39,6 +49,7 @@ # Small helpers # --------------------------------------------------------------------------- + def _make_crt_args( *, categorical_key="batch", @@ -74,6 +85,7 @@ def _make_crt_args( def _subset_to_one_batch(mdata, batch_col="batch", batch=None): """Return a MuData restricted to a single batch — keeps run_CRT iterations to 1.""" import mudata + rna = mdata["rna"].copy() cnmf = mdata["cNMF"].copy() if batch is None: @@ -90,10 +102,12 @@ def _build_fake_crt_out(n_targets=4, n_programs=3, seed=0): rng = np.random.default_rng(seed) targets = [f"gene{i}" for i in range(n_targets)] programs = [f"K{i}" for i in range(n_programs)] - pvals = pd.DataFrame(rng.uniform(1e-6, 1, (n_targets, n_programs)), - index=targets, columns=programs) - betas = pd.DataFrame(rng.normal(0, 0.5, (n_targets, n_programs)), - index=targets, columns=programs) + pvals = pd.DataFrame( + rng.uniform(1e-6, 1, (n_targets, n_programs)), index=targets, columns=programs + ) + betas = pd.DataFrame( + rng.normal(0, 0.5, (n_targets, n_programs)), index=targets, columns=programs + ) return {"pvals_skew_df": pvals, "betas_df": betas} @@ -101,7 +115,10 @@ def _build_fake_crt_out(n_targets=4, n_programs=3, seed=0): # Reformat: structural tests # =========================================================================== -@pytest.mark.skipif(not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)") + +@pytest.mark.skipif( + not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)" +) class TestCRTReformat: """Test the reformat_data_for_CRT function from the CRT calibration script.""" @@ -141,14 +158,18 @@ def test_reformat_cnmf_usage_dense_float(self, test_mdata): # Reformat: covariate variants # =========================================================================== -@pytest.mark.skipif(not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)") + +@pytest.mark.skipif( + not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)" +) class TestCRTReformatCovariates: """Vary the `covariates` / `log_covariates` arguments and check covar matrix.""" def test_reformat_no_covariates_empty_covar(self, test_mdata): """With no covariates, covar should be a DataFrame with 0 columns.""" - adata = reformat_data_for_CRT(test_mdata.copy(), - covariates=None, log_covariates=None) + adata = reformat_data_for_CRT( + test_mdata.copy(), covariates=None, log_covariates=None + ) covar = adata.obsm["covar"] assert isinstance(covar, pd.DataFrame) assert covar.shape == (adata.n_obs, 0) @@ -175,8 +196,9 @@ def test_reformat_multiple_raw_covariates(self, test_mdata): def test_reformat_log_covariates_prefix(self, test_mdata): """log_covariates should be prefixed `log_` and equal log1p(obs[key]).""" - adata = reformat_data_for_CRT(test_mdata.copy(), - log_covariates=["guide_umi_counts"]) + adata = reformat_data_for_CRT( + test_mdata.copy(), log_covariates=["guide_umi_counts"] + ) covar = adata.obsm["covar"] assert list(covar.columns) == ["log_guide_umi_counts"] expected = np.log1p(adata.obs["guide_umi_counts"].values) @@ -200,7 +222,10 @@ def test_reformat_combined_covariates_and_log(self, test_mdata): # Reformat: parametrized across K values # =========================================================================== -@pytest.mark.skipif(not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)") + +@pytest.mark.skipif( + not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)" +) class TestCRTReformatPerK: """Verify reformat works across the inference K values (5, 10, 15).""" @@ -216,14 +241,19 @@ def test_reformat_per_k_structure(self, mdata_copy_per_k): # save_result: FDR method variants # =========================================================================== -@pytest.mark.skipif(not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)") + +@pytest.mark.skipif( + not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)" +) class TestCRTSaveResult: """save_result should write a TSV with adj_pval correctly computed for each FDR method.""" def test_save_result_BH_writes_tsv(self, tmp_path): out = _build_fake_crt_out(n_targets=5, n_programs=3) args = _make_crt_args(FDR_method="BH") - df = save_result(out, k=5, output_folder=str(tmp_path), condition="D0", args=args) + df = save_result( + out, k=5, output_folder=str(tmp_path), condition="D0", args=args + ) # File written to expected path expected = tmp_path / "5_CRT_D0.txt" @@ -246,7 +276,9 @@ def test_save_result_StoreyQ_writes_tsv(self, tmp_path): out = _build_fake_crt_out(n_targets=5, n_programs=3, seed=1) args = _make_crt_args(FDR_method="StoreyQ") - df = save_result(out, k=10, output_folder=str(tmp_path), condition="D1", args=args) + df = save_result( + out, k=10, output_folder=str(tmp_path), condition="D1", args=args + ) expected = tmp_path / "10_CRT_D1.txt" assert expected.exists() @@ -260,9 +292,15 @@ def test_save_result_column_order(self, tmp_path): """save_result column order should be stable: target, program, log2FC, p-value, adj_pval.""" out = _build_fake_crt_out(n_targets=3, n_programs=2) args = _make_crt_args(FDR_method="BH") - df = save_result(out, k=5, output_folder=str(tmp_path), condition="D2", args=args) + df = save_result( + out, k=5, output_folder=str(tmp_path), condition="D2", args=args + ) assert list(df.columns) == [ - "target_name", "program_name", "log2FC", "p-value", "adj_pval", + "target_name", + "program_name", + "log2FC", + "p-value", + "adj_pval", ] @@ -274,16 +312,18 @@ def test_save_result_column_order(self, tmp_path): slow = pytest.mark.slow -@pytest.mark.skipif(not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)") +@pytest.mark.skipif( + not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)" +) class TestCRTRunCRT: """Exercise the full run_CRT pipeline on a single batch with tiny B.""" def _prepare_adata(self, mdata, covariates=None, log_covariates=None): """Subset to one batch and run reformat + flooring (mirrors main()).""" mdata_one = _subset_to_one_batch(mdata) - adata = reformat_data_for_CRT(mdata_one, - covariates=covariates, - log_covariates=log_covariates) + adata = reformat_data_for_CRT( + mdata_one, covariates=covariates, log_covariates=log_covariates + ) U = adata.obsm["cnmf_usage"].copy() U = np.maximum(U, 1e-8) U /= U.sum(axis=1, keepdims=True) @@ -294,6 +334,7 @@ def _prepare_adata(self, mdata, covariates=None, log_covariates=None): def test_run_CRT_small_guide_group(self, test_mdata, tmp_path): """Tiny number_guide (2) — smoke test that pipeline runs and writes outputs.""" import matplotlib + matplotlib.use("Agg") adata = self._prepare_adata(test_mdata) args = _make_crt_args(number_guide=2, number_permutations=8, FDR_method="BH") @@ -318,6 +359,7 @@ def test_run_CRT_large_guide_group(self, test_mdata, tmp_path): Test data has 103 non-targeting guides, so a group size of 10 is well within range. """ import matplotlib + matplotlib.use("Agg") adata = self._prepare_adata(test_mdata) args = _make_crt_args(number_guide=10, number_permutations=8, FDR_method="BH") @@ -329,6 +371,7 @@ def test_run_CRT_large_guide_group(self, test_mdata, tmp_path): def test_run_CRT_with_covariates(self, test_mdata, tmp_path): """Run with raw + log covariates included — covar matrix flows into prepare_crt_inputs.""" import matplotlib + matplotlib.use("Agg") adata = self._prepare_adata( test_mdata, @@ -346,6 +389,7 @@ def test_run_CRT_with_covariates(self, test_mdata, tmp_path): def test_run_CRT_more_permutations(self, test_mdata, tmp_path): """Higher number_permutations — verifies the B parameter is wired through correctly.""" import matplotlib + matplotlib.use("Agg") adata = self._prepare_adata(test_mdata) args = _make_crt_args(number_guide=2, number_permutations=32, FDR_method="BH") @@ -361,9 +405,12 @@ def test_run_CRT_StoreyQ_fdr(self, test_mdata, tmp_path): pytest.skip("multipy not installed — StoreyQ branch unavailable") import matplotlib + matplotlib.use("Agg") adata = self._prepare_adata(test_mdata) - args = _make_crt_args(number_guide=2, number_permutations=8, FDR_method="StoreyQ") + args = _make_crt_args( + number_guide=2, number_permutations=8, FDR_method="StoreyQ" + ) run_CRT(adata, k=5, output_folder=str(tmp_path), args=args) txts = list(tmp_path.glob("5_CRT_*.txt")) @@ -376,7 +423,10 @@ def test_run_CRT_StoreyQ_fdr(self, test_mdata, tmp_path): # main(): CLI integration via sys.argv # =========================================================================== -@pytest.mark.skipif(not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)") + +@pytest.mark.skipif( + not _HAS_CRT, reason="CRT requires programDE env (sceptre not available)" +) class TestCRTMain: """Exercise `main()` end-to-end with a constructed argv. Heavy — marked slow.""" @@ -389,6 +439,7 @@ def test_main_smoke_BH(self, test_mdata, tmp_path, monkeypatch): via sys.argv. All outputs routed to tmp_path/results via --save_dir. """ import matplotlib + matplotlib.use("Agg") # Subset to one batch so the run_CRT inner loop iterates exactly once. @@ -406,15 +457,24 @@ def test_main_smoke_BH(self, test_mdata, tmp_path, monkeypatch): argv = [ "CRT.py", - "--out_dir", str(out_dir), - "--run_name", run_name, - "--components", "5", - "--sel_thresh", "2.0", - "--categorical_key", "batch", - "--number_guide", "2", - "--number_permutations", "8", - "--FDR_method", "BH", - "--save_dir", str(save_dir), + "--out_dir", + str(out_dir), + "--run_name", + run_name, + "--components", + "5", + "--sel_thresh", + "2.0", + "--categorical_key", + "batch", + "--number_guide", + "2", + "--number_permutations", + "8", + "--FDR_method", + "BH", + "--save_dir", + str(save_dir), ] monkeypatch.setattr(sys, "argv", argv) @@ -426,4 +486,6 @@ def test_main_smoke_BH(self, test_mdata, tmp_path, monkeypatch): assert list(save_dir.glob("CRT_*.png")) assert list(save_dir.glob("5_CRT_*.txt")) # And a config_*.yml at save_config_dir (which == save_dir here) - assert list(save_dir.glob("config_*.yml")), "main() should have written a config_*.yml" + assert list(save_dir.glob("config_*.yml")), ( + "main() should have written a config_*.yml" + ) diff --git a/tests/Script/Stage2_Evaluation/test_metrics.py b/tests/Script/Stage2_Evaluation/test_metrics.py index 5b1ed8c..2d76b8e 100644 --- a/tests/Script/Stage2_Evaluation/test_metrics.py +++ b/tests/Script/Stage2_Evaluation/test_metrics.py @@ -33,114 +33,217 @@ # Results saved to Evaluation/{K}_2_0/ # =========================================================================== -class TestCategoricalPerK: +class TestCategoricalPerK: def test_compute_categorical_dunn(self, mdata_copy_per_k, eval_output_dir_per_k): - from Stage2_Evaluation.A_Metrics.src.association_categorical import compute_categorical_association + from Stage2_Evaluation.A_Metrics.src.association_categorical import ( + compute_categorical_association, + ) + k = mdata_copy_per_k.uns["test_k"] results_df, posthoc_df = compute_categorical_association( - mdata_copy_per_k, prog_key="cNMF", categorical_key="batch", - test="dunn", n_jobs=1, inplace=False, pseudobulk_key=None + mdata_copy_per_k, + prog_key="cNMF", + categorical_key="batch", + test="dunn", + n_jobs=1, + inplace=False, + pseudobulk_key=None, ) assert "batch_kruskall_wallis_stat" in results_df.columns assert len(results_df) == mdata_copy_per_k["cNMF"].n_vars - results_df.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_categorical_association_results.txt"), sep="\t") - posthoc_df.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_categorical_association_posthoc.txt"), sep="\t", index=False) + results_df.to_csv( + os.path.join( + eval_output_dir_per_k, f"{k}_categorical_association_results.txt" + ), + sep="\t", + ) + posthoc_df.to_csv( + os.path.join( + eval_output_dir_per_k, f"{k}_categorical_association_posthoc.txt" + ), + sep="\t", + index=False, + ) class TestPerturbationPerK: - def test_compute_perturbation(self, mdata_copy_per_k, eval_output_dir_per_k): - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association, + ) + k = mdata_copy_per_k.uns["test_k"] for batch in mdata_copy_per_k["rna"].obs["batch"].unique(): - mdata_batch = mdata_copy_per_k[mdata_copy_per_k["rna"].obs["batch"] == batch] + mdata_batch = mdata_copy_per_k[ + mdata_copy_per_k["rna"].obs["batch"] == batch + ] result = compute_perturbation_association( - mdata_batch, prog_key="cNMF", - reference_targets=["non-targeting"], pseudobulk=False, - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + mdata_batch, + prog_key="cNMF", + reference_targets=["non-targeting"], + pseudobulk=False, + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) assert isinstance(result, pd.DataFrame) assert all(result["pval"].between(0, 1)) - result.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_perturbation_association_results_{batch}.txt"), sep="\t", index=False) + result.to_csv( + os.path.join( + eval_output_dir_per_k, + f"{k}_perturbation_association_results_{batch}.txt", + ), + sep="\t", + index=False, + ) class TestGenesetPerK: - def test_get_gene_loadings(self, mdata_copy_per_k, eval_output_dir_per_k): - from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import get_program_gene_loadings + from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import ( + get_program_gene_loadings, + ) + k = mdata_copy_per_k.uns["test_k"] loadings = get_program_gene_loadings( mdata_copy_per_k, prog_key="cNMF", data_key="rna", gene_names_key="symbol" ) assert loadings.shape[1] == mdata_copy_per_k["cNMF"].n_vars - loadings.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_gene_loadings.txt"), sep="\t") + loadings.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_gene_loadings.txt"), sep="\t" + ) def test_reactome_enrichment(self, mdata_copy_per_k, eval_output_dir_per_k): - from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import compute_geneset_enrichment + from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import ( + compute_geneset_enrichment, + ) + k = mdata_copy_per_k.uns["test_k"] result = compute_geneset_enrichment( - mdata_copy_per_k, prog_key="cNMF", data_key="rna", prog_name=None, - organism="human", library="Reactome_2022", method="fisher", - database="enrichr", n_top=300, n_jobs=1, - inplace=False, user_geneset=None, use_loadings_gene=False, - gene_names_key="symbol", use_cache=True + mdata_copy_per_k, + prog_key="cNMF", + data_key="rna", + prog_name=None, + organism="human", + library="Reactome_2022", + method="fisher", + database="enrichr", + n_top=300, + n_jobs=1, + inplace=False, + user_geneset=None, + use_loadings_gene=False, + gene_names_key="symbol", + use_cache=True, ) assert isinstance(result, pd.DataFrame) - result.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_geneset_enrichment.txt"), sep="\t", index=False) - + result.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_geneset_enrichment.txt"), + sep="\t", + index=False, + ) class TestGOEnrichmentPerK: - def test_go_enrichment(self, mdata_copy_per_k, eval_output_dir_per_k): - from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import compute_geneset_enrichment + from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import ( + compute_geneset_enrichment, + ) + k = mdata_copy_per_k.uns["test_k"] result = compute_geneset_enrichment( - mdata_copy_per_k, prog_key="cNMF", data_key="rna", prog_name=None, - organism="human", library="GO_Biological_Process_2023", method="fisher", - database="enrichr", n_top=300, n_jobs=1, - inplace=False, user_geneset=None, use_loadings_gene=False, - gene_names_key="symbol", use_cache=True + mdata_copy_per_k, + prog_key="cNMF", + data_key="rna", + prog_name=None, + organism="human", + library="GO_Biological_Process_2023", + method="fisher", + database="enrichr", + n_top=300, + n_jobs=1, + inplace=False, + user_geneset=None, + use_loadings_gene=False, + gene_names_key="symbol", + use_cache=True, ) assert isinstance(result, pd.DataFrame) - result.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_GO_term_enrichment.txt"), sep="\t", index=False) + result.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_GO_term_enrichment.txt"), + sep="\t", + index=False, + ) class TestTraitEnrichmentPerK: - def test_trait_enrichment(self, mdata_copy_per_k, eval_output_dir_per_k): - from Stage2_Evaluation.A_Metrics.src.enrichment_trait import compute_trait_enrichment + from Stage2_Evaluation.A_Metrics.src.enrichment_trait import ( + compute_trait_enrichment, + ) + k = mdata_copy_per_k.uns["test_k"] gwas_path = os.path.join( - os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), - "src", "Stage2_Evaluation", "Resources", "OpenTargets_L2G_Filtered.csv.gz" + os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + ), + "src", + "Stage2_Evaluation", + "Resources", + "OpenTargets_L2G_Filtered.csv.gz", ) if not os.path.isfile(gwas_path): pytest.skip(f"GWAS data not found: {gwas_path}") result = compute_trait_enrichment( - mdata_copy_per_k, gwas_data=gwas_path, - prog_key="cNMF", prog_name=None, data_key="rna", - library="OT_GWAS", n_jobs=1, inplace=False, - key_column="trait_efos", gene_column="gene_name", - method="fisher", n_top=300, use_loadings_gene=True, - gene_names_key="symbol" + mdata_copy_per_k, + gwas_data=gwas_path, + prog_key="cNMF", + prog_name=None, + data_key="rna", + library="OT_GWAS", + n_jobs=1, + inplace=False, + key_column="trait_efos", + gene_column="gene_name", + method="fisher", + n_top=300, + use_loadings_gene=True, + gene_names_key="symbol", ) assert isinstance(result, pd.DataFrame) - result.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_trait_enrichment.txt"), sep="\t", index=False) + result.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_trait_enrichment.txt"), + sep="\t", + index=False, + ) + class TestExplainedVariancePerK: + def test_compute_explained_variance( + self, mdata_copy_per_k, eval_output_dir_per_k, cnmf_obj, X_norm + ): + from Stage2_Evaluation.A_Metrics.src.explained_variance import ( + compute_explained_variance, + ) - def test_compute_explained_variance(self, mdata_copy_per_k, eval_output_dir_per_k, cnmf_obj, X_norm): - from Stage2_Evaluation.A_Metrics.src.explained_variance import compute_explained_variance k = mdata_copy_per_k.uns["test_k"] program_names = mdata_copy_per_k["cNMF"].var_names - compute_explained_variance(cnmf_obj, X_norm, k, - output_folder=eval_output_dir_per_k, - thre='2.0', - program_name=program_names) + compute_explained_variance( + cnmf_obj, + X_norm, + k, + output_folder=eval_output_dir_per_k, + thre="2.0", + program_name=program_names, + ) ev_path = os.path.join(eval_output_dir_per_k, f"{k}_Explained_Variance.txt") - summary_path = os.path.join(eval_output_dir_per_k, f"{k}_Explained_Variance_Summary.txt") + summary_path = os.path.join( + eval_output_dir_per_k, f"{k}_Explained_Variance_Summary.txt" + ) assert os.path.isfile(ev_path) assert os.path.isfile(summary_path) ev_df = pd.read_csv(ev_path, sep="\t") @@ -149,29 +252,54 @@ def test_compute_explained_variance(self, mdata_copy_per_k, eval_output_dir_per_ class TestIntegrationPerK: - - def test_categorical_then_perturbation(self, mdata_copy_per_k, eval_output_dir_per_k): + def test_categorical_then_perturbation( + self, mdata_copy_per_k, eval_output_dir_per_k + ): """Run categorical + perturbation per K on real MuData.""" - from Stage2_Evaluation.A_Metrics.src.association_categorical import compute_categorical_association - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association + from Stage2_Evaluation.A_Metrics.src.association_categorical import ( + compute_categorical_association, + ) + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association, + ) + k = mdata_copy_per_k.uns["test_k"] cat_results, cat_posthoc = compute_categorical_association( - mdata_copy_per_k, prog_key="cNMF", categorical_key="batch", - test="dunn", n_jobs=1, inplace=False + mdata_copy_per_k, + prog_key="cNMF", + categorical_key="batch", + test="dunn", + n_jobs=1, + inplace=False, ) assert len(cat_results) > 0 perturb_results = compute_perturbation_association( - mdata_copy_per_k, prog_key="cNMF", + mdata_copy_per_k, + prog_key="cNMF", reference_targets=["non-targeting"], - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) assert len(perturb_results) > 0 - cat_results.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_integration_categorical.txt"), sep="\t") - cat_posthoc.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_integration_posthoc.txt"), sep="\t", index=False) - perturb_results.to_csv(os.path.join(eval_output_dir_per_k, f"{k}_integration_perturbation.txt"), sep="\t", index=False) + cat_results.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_integration_categorical.txt"), + sep="\t", + ) + cat_posthoc.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_integration_posthoc.txt"), + sep="\t", + index=False, + ) + perturb_results.to_csv( + os.path.join(eval_output_dir_per_k, f"{k}_integration_perturbation.txt"), + sep="\t", + index=False, + ) # =========================================================================== @@ -180,6 +308,7 @@ def test_categorical_then_perturbation(self, mdata_copy_per_k, eval_output_dir_p # --- association_perturbation.py: guide annotation table vs key --- + class TestPerturbationReferenceTargets: """Test that using a guide annotation table produces the same reference_targets as passing the key string directly.""" @@ -191,15 +320,27 @@ def test_annotation_table_extracts_nontargeting(self, guide_annotation_path): reference_targets = df_non.index.values.tolist() assert len(reference_targets) > 0 # All extracted targets should be non-targeting type - assert all(df.loc[t, "type"] in ("non-targeting", "safe-targeting") for t in reference_targets) + assert all( + df.loc[t, "type"] in ("non-targeting", "safe-targeting") + for t in reference_targets + ) - def test_perturbation_with_annotation_table(self, mdata_copy, reference_targets_from_annotation): + def test_perturbation_with_annotation_table( + self, mdata_copy, reference_targets_from_annotation + ): """Run perturbation association using reference_targets extracted from annotation table.""" - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association, + ) + result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", + mdata_copy, + prog_key="cNMF", reference_targets=reference_targets_from_annotation, - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) assert isinstance(result, pd.DataFrame) assert "adj_pval" in result.columns @@ -207,43 +348,70 @@ def test_perturbation_with_annotation_table(self, mdata_copy, reference_targets_ def test_perturbation_with_key_string(self, mdata_copy): """Run perturbation association using the direct key string.""" - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association, + ) + result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", + mdata_copy, + prog_key="cNMF", reference_targets=["non-targeting"], - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) assert isinstance(result, pd.DataFrame) assert "adj_pval" in result.columns assert all(result["pval"].between(0, 1)) - def test_both_methods_same_results(self, mdata_copy, reference_targets_from_annotation): + def test_both_methods_same_results( + self, mdata_copy, reference_targets_from_annotation + ): """Both methods should produce identical perturbation results.""" - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association, + ) + result_key = compute_perturbation_association( - mdata_copy, prog_key="cNMF", + mdata_copy, + prog_key="cNMF", reference_targets=["non-targeting"], - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) # Need a fresh copy since mdata_copy is function-scoped but same object within one test result_table = compute_perturbation_association( - mdata_copy, prog_key="cNMF", + mdata_copy, + prog_key="cNMF", reference_targets=reference_targets_from_annotation, - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) # Same shape and target names assert set(result_key["target_name"]) == set(result_table["target_name"]) # Merge and compare p-values - merged = result_key.merge(result_table, on=["target_name", "program_name"], suffixes=("_key", "_table")) + merged = result_key.merge( + result_table, + on=["target_name", "program_name"], + suffixes=("_key", "_table"), + ) assert np.allclose(merged["pval_key"], merged["pval_table"], atol=1e-10) # --- association_perturbation.py: unit tests --- -class TestPerturbationAssociation: +class TestPerturbationAssociation: def test_get_guide_metadata(self, test_mdata): - from Stage2_Evaluation.A_Metrics.src.association_perturbation import get_guide_metadata + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + get_guide_metadata, + ) + gm = get_guide_metadata(test_mdata, prog_key="cNMF") assert "Target" in gm.columns assert len(gm) == 1859 @@ -251,31 +419,43 @@ def test_get_guide_metadata(self, test_mdata): def test_compute_perturbation_single_unit(self): """Test the inner per-program Mann-Whitney function with synthetic data.""" - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association_ + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association_, + ) + rng = np.random.default_rng(42) test_data = ad.AnnData( X=sparse.csr_matrix(rng.random((30, 3))), - var=pd.DataFrame(index=["p0", "p1", "p2"]) + var=pd.DataFrame(index=["p0", "p1", "p2"]), ) ref_data = ad.AnnData( X=sparse.csr_matrix(rng.random((50, 3))), - var=pd.DataFrame(index=["p0", "p1", "p2"]) + var=pd.DataFrame(index=["p0", "p1", "p2"]), ) results = [] - compute_perturbation_association_(test_data, ref_data, "p0", "target_X", results) + compute_perturbation_association_( + test_data, ref_data, "p0", "target_X", results + ) assert len(results) == 1 row = results[0] assert row[0] == "target_X" assert row[1] == "p0" assert np.isfinite(row[5]) # stat - assert 0 <= row[6] <= 1 # pval + assert 0 <= row[6] <= 1 # pval def test_pval_sanitization(self, mdata_copy): - from Stage2_Evaluation.A_Metrics.src.association_perturbation import compute_perturbation_association + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + compute_perturbation_association, + ) + result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", + mdata_copy, + prog_key="cNMF", reference_targets=["non-targeting"], - collapse_targets=True, n_jobs=1, inplace=False, FDR_method="BH" + collapse_targets=True, + n_jobs=1, + inplace=False, + FDR_method="BH", ) assert all(np.isfinite(result["pval"])) assert all(result["pval"] >= 0) @@ -284,10 +464,11 @@ def test_pval_sanitization(self, mdata_copy): # --- Utilities.py --- -class TestUtilities: +class TestUtilities: def test_rename_gene_dictionary(self): from Stage2_Evaluation.A_Metrics.src.Utilities import rename_gene_dictionary + rng = np.random.default_rng(42) n_cells, n_genes = 10, 5 ens_ids = [f"ENS{i:05d}" for i in range(n_genes)] @@ -295,7 +476,7 @@ def test_rename_gene_dictionary(self): rna_ad = ad.AnnData( X=sparse.csr_matrix(rng.random((n_cells, n_genes))), var=pd.DataFrame(index=ens_ids), - obs=pd.DataFrame(index=[f"cell_{i}" for i in range(n_cells)]) + obs=pd.DataFrame(index=[f"cell_{i}" for i in range(n_cells)]), ) mdata = mu.MuData({"rna": rna_ad}) with tempfile.NamedTemporaryFile(mode="w", suffix=".tsv", delete=False) as f: @@ -312,13 +493,14 @@ def test_rename_gene_dictionary(self): def test_rename_gene_dictionary_partial(self): """Unmapped IDs stay as-is.""" from Stage2_Evaluation.A_Metrics.src.Utilities import rename_gene_dictionary + rng = np.random.default_rng(42) n_cells, n_genes = 10, 5 ens_ids = [f"ENS{i:05d}" for i in range(n_genes)] rna_ad = ad.AnnData( X=sparse.csr_matrix(rng.random((n_cells, n_genes))), var=pd.DataFrame(index=ens_ids), - obs=pd.DataFrame(index=[f"cell_{i}" for i in range(n_cells)]) + obs=pd.DataFrame(index=[f"cell_{i}" for i in range(n_cells)]), ) mdata = mu.MuData({"rna": rna_ad}) with tempfile.NamedTemporaryFile(mode="w", suffix=".tsv", delete=False) as f: diff --git a/tests/Script/Stage2_Evaluation/test_utest.py b/tests/Script/Stage2_Evaluation/test_utest.py index 5cc16ed..7219f1c 100644 --- a/tests/Script/Stage2_Evaluation/test_utest.py +++ b/tests/Script/Stage2_Evaluation/test_utest.py @@ -20,7 +20,9 @@ def _get_nontargeting_indices(guide_annotation_df, mdata, prog_key="cNMF"): """Get mdata column indices for non-targeting guides by intersecting annotation table guide names with mdata guide_names.""" - nt_guide_names = guide_annotation_df[guide_annotation_df["targeting"] == False].index.values + nt_guide_names = guide_annotation_df[ + guide_annotation_df["targeting"] == False + ].index.values mdata_guide_names = mdata[prog_key].uns["guide_names"] nt_in_mdata = np.isin(mdata_guide_names, nt_guide_names) return np.where(nt_in_mdata)[0] @@ -30,13 +32,21 @@ def _get_nontargeting_indices(guide_annotation_df, mdata, prog_key="cNMF"): # Helpers to call the original U-test pipeline script # =========================================================================== + def _load_utest_module(): """Import the U-test calibration module (has hyphens in path).""" import importlib.util + script_path = os.path.join( - os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), - "src", "Stage2_Evaluation", "B_Calibration", "Slurm_version", - "U-test_perturbation_calibration", "U-test_perturbation_calibration.py", + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ), + "src", + "Stage2_Evaluation", + "B_Calibration", + "Slurm_version", + "U-test_perturbation_calibration", + "U-test_perturbation_calibration.py", ) spec = importlib.util.spec_from_file_location("utest_calibration", script_path) mod = importlib.util.module_from_spec(spec) @@ -44,10 +54,17 @@ def _load_utest_module(): return mod -def _make_utest_args(inference_path, components, categorical_key="batch", - sel_thresh=None, number_run=1, number_guide=3): +def _make_utest_args( + inference_path, + components, + categorical_key="batch", + sel_thresh=None, + number_run=1, + number_guide=3, +): """Create a mock args namespace for the U-test module.""" from types import SimpleNamespace + out_dir = os.path.dirname(os.path.dirname(inference_path)) run_name = os.path.basename(os.path.dirname(inference_path)) return SimpleNamespace( @@ -74,10 +91,13 @@ def _make_utest_args(inference_path, components, categorical_key="batch", # Tests for perturbation association (core of U-test calibration) # =========================================================================== + class TestPerturbationAssociation: """Run perturbation tests per K using the original U-test pipeline functions.""" - def test_real_perturbation_per_k(self, mdata_copy_per_k, eval_output_dir_per_k, inference_path): + def test_real_perturbation_per_k( + self, mdata_copy_per_k, eval_output_dir_per_k, inference_path + ): """Call compute_real_perturbation_tests from the original script, per K per batch. Verifies: @@ -86,6 +106,7 @@ def test_real_perturbation_per_k(self, mdata_copy_per_k, eval_output_dir_per_k, - plot_calibration_comparison saves a PNG into the per-(K, sel_thresh) folder """ import matplotlib + matplotlib.use("Agg") utest_mod = _load_utest_module() @@ -111,14 +132,22 @@ def test_real_perturbation_per_k(self, mdata_copy_per_k, eval_output_dir_per_k, eval_output_dir_per_k, f"{k}_perturbation_association_results_{samp}.txt", ) - assert os.path.exists(expected_txt), f"Missing per-(K, sample) file: {expected_txt}" + assert os.path.exists(expected_txt), ( + f"Missing per-(K, sample) file: {expected_txt}" + ) # Violin plot via original function (saves per-(K, sel_thresh) folder) utest_mod.plot_calibration_comparison(result_df) - expected_png = os.path.join(eval_output_dir_per_k, "U_test_perturbation_association_calibration.png") - assert os.path.exists(expected_png), f"Missing per-folder calibration plot: {expected_png}" + expected_png = os.path.join( + eval_output_dir_per_k, "U_test_perturbation_association_calibration.png" + ) + assert os.path.exists(expected_png), ( + f"Missing per-folder calibration plot: {expected_png}" + ) - def test_fake_perturbation_per_k(self, mdata_copy_per_k, eval_output_dir_per_k, inference_path): + def test_fake_perturbation_per_k( + self, mdata_copy_per_k, eval_output_dir_per_k, inference_path + ): """Call compute_fake_perturbation_tests from the original script, per K per batch. Verifies per-(K, sample) fake result files are saved into Evaluation/{K}_{thresh}/. @@ -129,7 +158,10 @@ def test_fake_perturbation_per_k(self, mdata_copy_per_k, eval_output_dir_per_k, k = mdata_copy_per_k.uns["test_k"] utest_mod.args = _make_utest_args( - inference_path, components=[k], number_run=1, number_guide=3, + inference_path, + components=[k], + number_run=1, + number_guide=3, ) utest_mod.mdata_guide = None @@ -149,9 +181,13 @@ def test_fake_perturbation_per_k(self, mdata_copy_per_k, eval_output_dir_per_k, eval_output_dir_per_k, f"{k}_fake_perturbation_association_results_{samp}.txt", ) - assert os.path.exists(expected_txt), f"Missing per-(K, sample) fake file: {expected_txt}" + assert os.path.exists(expected_txt), ( + f"Missing per-(K, sample) fake file: {expected_txt}" + ) - def test_plots_per_k_real_eval_folder(self, mdata_copy_per_k, eval_output_dir_per_k, inference_path): + def test_plots_per_k_real_eval_folder( + self, mdata_copy_per_k, eval_output_dir_per_k, inference_path + ): """Run real + fake tests, then call both plot_qq_comparison and plot_calibration_comparison. Verifies that BOTH ``U_test_perturbation_association_qqplot.png`` and @@ -161,6 +197,7 @@ def test_plots_per_k_real_eval_folder(self, mdata_copy_per_k, eval_output_dir_pe target_name=='targeting'; the violin plot compares real vs fake side-by-side). """ import matplotlib + matplotlib.use("Agg") np.random.seed(0) @@ -168,7 +205,10 @@ def test_plots_per_k_real_eval_folder(self, mdata_copy_per_k, eval_output_dir_pe k = mdata_copy_per_k.uns["test_k"] utest_mod.args = _make_utest_args( - inference_path, components=[k], number_run=1, number_guide=3, + inference_path, + components=[k], + number_run=1, + number_guide=3, ) utest_mod.mdata_guide = None @@ -195,30 +235,44 @@ def test_plots_per_k_real_eval_folder(self, mdata_copy_per_k, eval_output_dir_pe utest_mod.plot_qq_comparison(combined) utest_mod.plot_calibration_comparison(combined) - qq_png = os.path.join(eval_output_dir_per_k, "U_test_perturbation_association_qqplot.png") - cal_png = os.path.join(eval_output_dir_per_k, "U_test_perturbation_association_calibration.png") + qq_png = os.path.join( + eval_output_dir_per_k, "U_test_perturbation_association_qqplot.png" + ) + cal_png = os.path.join( + eval_output_dir_per_k, "U_test_perturbation_association_calibration.png" + ) assert os.path.exists(qq_png), f"Missing per-folder QQ plot: {qq_png}" assert os.path.getsize(qq_png) > 1000, f"QQ plot looks empty: {qq_png}" - assert os.path.exists(cal_png), f"Missing per-folder calibration plot: {cal_png}" - assert os.path.getsize(cal_png) > 1000, f"Calibration plot looks empty: {cal_png}" + assert os.path.exists(cal_png), ( + f"Missing per-folder calibration plot: {cal_png}" + ) + assert os.path.getsize(cal_png) > 1000, ( + f"Calibration plot looks empty: {cal_png}" + ) # =========================================================================== # Tests for guide metadata handling # =========================================================================== -class TestGuideMetadata: +class TestGuideMetadata: def test_get_guide_metadata(self, mdata_copy): - from Stage2_Evaluation.A_Metrics.src.association_perturbation import get_guide_metadata + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + get_guide_metadata, + ) + guide_meta = get_guide_metadata(mdata_copy, prog_key="cNMF") assert isinstance(guide_meta, pd.DataFrame) assert "Target" in guide_meta.columns assert len(guide_meta) == len(mdata_copy["cNMF"].uns["guide_names"]) def test_guide_metadata_has_nontargeting(self, mdata_copy): - from Stage2_Evaluation.A_Metrics.src.association_perturbation import get_guide_metadata + from Stage2_Evaluation.A_Metrics.src.association_perturbation import ( + get_guide_metadata, + ) + guide_meta = get_guide_metadata(mdata_copy, prog_key="cNMF") nt_count = (guide_meta["Target"] == "non-targeting").sum() assert nt_count > 0 @@ -236,15 +290,20 @@ def test_guide_assignment_shape(self, test_mdata): # Tests for output file structure # =========================================================================== -class TestCalibrationOutput: +class TestCalibrationOutput: def test_save_perturbation_results(self, mdata_copy, calibration_output_dir): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) out_path = os.path.join(calibration_output_dir, "test_perturbation_results.txt") result.to_csv(out_path, sep="\t", index=False) @@ -253,16 +312,23 @@ def test_save_perturbation_results(self, mdata_copy, calibration_output_dir): assert len(loaded) == len(result) assert "target_name" in loaded.columns - def _run_multiple_iterations(self, mdata_copy, guide_annotation_df, use_file_access): + def _run_multiple_iterations( + self, mdata_copy, guide_annotation_df, use_file_access + ): """Helper: run 3 calibration iterations with key or file-based access.""" from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + rng = np.random.default_rng(0) all_results = [] if use_file_access: - mdata_copy["cNMF"].uns["guide_targets"] = mdata_copy["cNMF"].uns["guide_names"].copy() + mdata_copy["cNMF"].uns["guide_targets"] = ( + mdata_copy["cNMF"].uns["guide_names"].copy() + ) nt_indices = _get_nontargeting_indices(guide_annotation_df, mdata_copy) - nt_guide_names = guide_annotation_df[guide_annotation_df["targeting"] == False].index.values + nt_guide_names = guide_annotation_df[ + guide_annotation_df["targeting"] == False + ].index.values for iteration in range(3): _mdata = mdata_copy.copy() @@ -284,15 +350,21 @@ def _run_multiple_iterations(self, mdata_copy, guide_annotation_df, use_file_acc _mdata["rna"].uns["guide_targets"] = new_targets if use_file_access: - ref_targets = [g for g in nt_guide_names if g in prog.uns["guide_names"]] + ref_targets = [ + g for g in nt_guide_names if g in prog.uns["guide_names"] + ] else: ref_targets = ["non-targeting"] result = compute_perturbation_association( - _mdata, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + _mdata, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=ref_targets, - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) result["run"] = iteration all_results.append(result) @@ -303,30 +375,41 @@ def _run_multiple_iterations(self, mdata_copy, guide_annotation_df, use_file_acc def test_multiple_iterations_concatenate_key(self, mdata_copy, guide_annotation_df): """Multiple iterations using key-based reference_targets=['non-targeting'].""" - self._run_multiple_iterations(mdata_copy, guide_annotation_df, use_file_access=False) + self._run_multiple_iterations( + mdata_copy, guide_annotation_df, use_file_access=False + ) - def test_multiple_iterations_concatenate_file(self, mdata_copy, guide_annotation_df): + def test_multiple_iterations_concatenate_file( + self, mdata_copy, guide_annotation_df + ): """Multiple iterations using file-based reference_targets (individual guide names).""" - self._run_multiple_iterations(mdata_copy, guide_annotation_df, use_file_access=True) + self._run_multiple_iterations( + mdata_copy, guide_annotation_df, use_file_access=True + ) # =========================================================================== # Tests for edge cases # =========================================================================== -class TestEdgeCases: +class TestEdgeCases: def test_sparse_guide_assignment(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + prog = mdata_copy["cNMF"] assignment = prog.obsm["guide_assignment"] if not sparse.issparse(assignment): prog.obsm["guide_assignment"] = sparse.csr_matrix(assignment) result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) assert isinstance(result, pd.DataFrame) assert len(result) > 0 @@ -334,11 +417,16 @@ def test_sparse_guide_assignment(self, mdata_copy): def test_all_pvals_finite_after_fdr(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) assert all(np.isfinite(result["adj_pval"])) assert all(result["adj_pval"] >= 0) @@ -346,11 +434,16 @@ def test_all_pvals_finite_after_fdr(self, mdata_copy): def test_inplace_mode_updates_mdata(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + ret = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=True, + FDR_method="BH", + n_jobs=1, + inplace=True, ) assert ret is None assert "perturbation_association_target_stat" in mdata_copy["cNMF"].varm @@ -362,10 +455,11 @@ def test_inplace_mode_updates_mdata(self, mdata_copy): # Tests for condition-stratified calibration # =========================================================================== -class TestConditionStratifiedCalibration: +class TestConditionStratifiedCalibration: def test_per_batch_results_differ(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + batches = mdata_copy["rna"].obs["batch"].unique() assert len(batches) >= 2 @@ -374,27 +468,39 @@ def test_per_batch_results_differ(self, mdata_copy): mask = mdata_copy["rna"].obs["batch"] == batch mdata_sub = mdata_copy[mask] result = compute_perturbation_association( - mdata_sub, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_sub, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) results_by_batch[batch] = result b1, b2 = list(results_by_batch.keys()) - assert not results_by_batch[b1]["pval"].values.tolist() == results_by_batch[b2]["pval"].values.tolist() + assert ( + not results_by_batch[b1]["pval"].values.tolist() + == results_by_batch[b2]["pval"].values.tolist() + ) def test_per_batch_combine_with_metadata(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + all_results = [] for batch in mdata_copy["rna"].obs["batch"].unique(): mask = mdata_copy["rna"].obs["batch"] == batch mdata_sub = mdata_copy[mask] result = compute_perturbation_association( - mdata_sub, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_sub, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) result["sample"] = batch result["K"] = 5 @@ -410,15 +516,20 @@ def test_per_batch_combine_with_metadata(self, mdata_copy): # Tests for guide-level (non-collapsed) perturbation association # =========================================================================== -class TestGuideLevelAssociation: +class TestGuideLevelAssociation: def test_guide_level_returns_guide_name_col(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=False, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=False, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) assert isinstance(result, pd.DataFrame) assert "guide_name" in result.columns @@ -426,17 +537,26 @@ def test_guide_level_returns_guide_name_col(self, mdata_copy): def test_guide_level_more_rows_than_target_level(self, mdata_copy): from Stage2_Evaluation.A_Metrics.src import compute_perturbation_association + target_result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=True, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=True, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) guide_result = compute_perturbation_association( - mdata_copy, prog_key="cNMF", - collapse_targets=False, pseudobulk=False, + mdata_copy, + prog_key="cNMF", + collapse_targets=False, + pseudobulk=False, reference_targets=["non-targeting"], - FDR_method="BH", n_jobs=1, inplace=False, + FDR_method="BH", + n_jobs=1, + inplace=False, ) assert len(guide_result) > len(target_result) @@ -452,17 +572,19 @@ def _assert_valid_png(path, min_bytes=1000): """Assert that `path` exists, is non-empty, and starts with the PNG magic header.""" assert path.exists(), f"Missing plot: {path}" size = path.stat().st_size - assert size > min_bytes, f"Plot too small ({size} bytes), likely empty figure: {path}" + assert size > min_bytes, ( + f"Plot too small ({size} bytes), likely empty figure: {path}" + ) with open(path, "rb") as fh: header = fh.read(len(PNG_MAGIC)) assert header == PNG_MAGIC, f"File at {path} is not a valid PNG (header={header!r})" class TestCalibrationVisualization: - def test_qq_plot_saves_per_folder(self, tmp_path): """plot_qq_comparison saves one valid PNG per (K, sel_thresh) into Evaluation/{K}_{thresh}/.""" import matplotlib + matplotlib.use("Agg") from types import SimpleNamespace @@ -484,9 +606,25 @@ def test_qq_plot_saves_per_folder(self, tmp_path): for k in components: for st in sel_thresh_list: for p in rng.uniform(0, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": True, "target_name": "geneX", "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": True, + "target_name": "geneX", + "pval": p, + } + ) for p in rng.uniform(0, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": False, "target_name": "targeting", "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": False, + "target_name": "targeting", + "pval": p, + } + ) test_stats_dfs = pd.DataFrame(rows) utest_mod.plot_qq_comparison(test_stats_dfs) @@ -496,7 +634,13 @@ def test_qq_plot_saves_per_folder(self, tmp_path): for k in components: for st in sel_thresh_list: thresh_str = str(st).replace(".", "_") - expected = tmp_path / run_name / "Evaluation" / f"{k}_{thresh_str}" / "U_test_perturbation_association_qqplot.png" + expected = ( + tmp_path + / run_name + / "Evaluation" + / f"{k}_{thresh_str}" + / "U_test_perturbation_association_qqplot.png" + ) _assert_valid_png(expected) produced_paths.append(expected) @@ -512,6 +656,7 @@ def test_qq_plot_saves_per_folder(self, tmp_path): def test_calibration_plot_saves_per_folder(self, tmp_path): """plot_calibration_comparison saves one valid PNG per (K, sel_thresh) into Evaluation/{K}_{thresh}/.""" import matplotlib + matplotlib.use("Agg") from types import SimpleNamespace @@ -534,9 +679,25 @@ def test_calibration_plot_saves_per_folder(self, tmp_path): for st in sel_thresh_list: for samp in ["d0", "d1"]: for p in rng.uniform(1e-6, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": True, "sample": samp, "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": True, + "sample": samp, + "pval": p, + } + ) for p in rng.uniform(1e-6, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": False, "sample": samp, "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": False, + "sample": samp, + "pval": p, + } + ) test_stats_dfs = pd.DataFrame(rows) utest_mod.plot_calibration_comparison(test_stats_dfs) @@ -546,7 +707,13 @@ def test_calibration_plot_saves_per_folder(self, tmp_path): for k in components: for st in sel_thresh_list: thresh_str = str(st).replace(".", "_") - expected = tmp_path / run_name / "Evaluation" / f"{k}_{thresh_str}" / "U_test_perturbation_association_calibration.png" + expected = ( + tmp_path + / run_name + / "Evaluation" + / f"{k}_{thresh_str}" + / "U_test_perturbation_association_calibration.png" + ) _assert_valid_png(expected) produced_paths.append(expected) @@ -560,6 +727,7 @@ def test_calibration_plot_saves_per_folder(self, tmp_path): def test_qq_plot_skips_missing_k(self, tmp_path): """plot_qq_comparison should skip (no file) for K values with no data, but write plots for the rest.""" import matplotlib + matplotlib.use("Agg") from types import SimpleNamespace @@ -582,9 +750,25 @@ def test_qq_plot_skips_missing_k(self, tmp_path): for k in [5, 15]: for st in sel_thresh_list: for p in rng.uniform(0, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": True, "target_name": "geneX", "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": True, + "target_name": "geneX", + "pval": p, + } + ) for p in rng.uniform(0, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": False, "target_name": "targeting", "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": False, + "target_name": "targeting", + "pval": p, + } + ) test_stats_dfs = pd.DataFrame(rows) utest_mod.plot_qq_comparison(test_stats_dfs) @@ -592,17 +776,32 @@ def test_qq_plot_skips_missing_k(self, tmp_path): for k in [5, 15]: for st in sel_thresh_list: thresh_str = str(st).replace(".", "_") - expected = tmp_path / run_name / "Evaluation" / f"{k}_{thresh_str}" / "U_test_perturbation_association_qqplot.png" + expected = ( + tmp_path + / run_name + / "Evaluation" + / f"{k}_{thresh_str}" + / "U_test_perturbation_association_qqplot.png" + ) _assert_valid_png(expected) # K=10 had no data — no QQ PNG should have been written. thresh_str = str(sel_thresh_list[0]).replace(".", "_") - missing = tmp_path / run_name / "Evaluation" / f"10_{thresh_str}" / "U_test_perturbation_association_qqplot.png" - assert not missing.exists(), f"QQ plot should not exist for empty K=10: {missing}" + missing = ( + tmp_path + / run_name + / "Evaluation" + / f"10_{thresh_str}" + / "U_test_perturbation_association_qqplot.png" + ) + assert not missing.exists(), ( + f"QQ plot should not exist for empty K=10: {missing}" + ) def test_calibration_plot_skips_missing_k(self, tmp_path): """plot_calibration_comparison should skip K values with no data.""" import matplotlib + matplotlib.use("Agg") from types import SimpleNamespace @@ -625,9 +824,25 @@ def test_calibration_plot_skips_missing_k(self, tmp_path): for st in sel_thresh_list: for samp in ["d0", "d1"]: for p in rng.uniform(1e-6, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": True, "sample": samp, "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": True, + "sample": samp, + "pval": p, + } + ) for p in rng.uniform(1e-6, 1, 50): - rows.append({"K": k, "sel_thresh": st, "real": False, "sample": samp, "pval": p}) + rows.append( + { + "K": k, + "sel_thresh": st, + "real": False, + "sample": samp, + "pval": p, + } + ) test_stats_dfs = pd.DataFrame(rows) utest_mod.plot_calibration_comparison(test_stats_dfs) @@ -635,23 +850,46 @@ def test_calibration_plot_skips_missing_k(self, tmp_path): for k in [5, 15]: for st in sel_thresh_list: thresh_str = str(st).replace(".", "_") - expected = tmp_path / run_name / "Evaluation" / f"{k}_{thresh_str}" / "U_test_perturbation_association_calibration.png" + expected = ( + tmp_path + / run_name + / "Evaluation" + / f"{k}_{thresh_str}" + / "U_test_perturbation_association_calibration.png" + ) _assert_valid_png(expected) thresh_str = str(sel_thresh_list[0]).replace(".", "_") - missing = tmp_path / run_name / "Evaluation" / f"10_{thresh_str}" / "U_test_perturbation_association_calibration.png" - assert not missing.exists(), f"Calibration plot should not exist for empty K=10: {missing}" + missing = ( + tmp_path + / run_name + / "Evaluation" + / f"10_{thresh_str}" + / "U_test_perturbation_association_calibration.png" + ) + assert not missing.exists(), ( + f"Calibration plot should not exist for empty K=10: {missing}" + ) # =========================================================================== # Tests for load_*_perturbation_tests (discover per-(K, sel_thresh, sample) files) # =========================================================================== + class TestLoadFunctions: """Tests for load_real_perturbation_tests and load_fake_perturbation_tests.""" - def _write_dummy_results(self, tmp_path, run_name, components, sel_thresh_list, - samples, fake=False, n_rows=10): + def _write_dummy_results( + self, + tmp_path, + run_name, + components, + sel_thresh_list, + samples, + fake=False, + n_rows=10, + ): """Create dummy per-(K, sample) result files in Evaluation/{K}_{thresh}/.""" prefix = "fake_" if fake else "" rng = np.random.default_rng(0) @@ -661,11 +899,13 @@ def _write_dummy_results(self, tmp_path, run_name, components, sel_thresh_list, folder = tmp_path / run_name / "Evaluation" / f"{k}_{thresh_str}" folder.mkdir(parents=True, exist_ok=True) for samp in samples: - df = pd.DataFrame({ - "target_name": [f"g{i}" for i in range(n_rows)], - "pval": rng.uniform(0, 1, n_rows), - "adj_pval": rng.uniform(0, 1, n_rows), - }) + df = pd.DataFrame( + { + "target_name": [f"g{i}" for i in range(n_rows)], + "pval": rng.uniform(0, 1, n_rows), + "adj_pval": rng.uniform(0, 1, n_rows), + } + ) if fake: df["K"] = k df["sel_thresh"] = st @@ -677,6 +917,7 @@ def _write_dummy_results(self, tmp_path, run_name, components, sel_thresh_list, def test_load_real_perturbation_tests(self, tmp_path): """load_real_perturbation_tests discovers samples and stacks per-(K, thresh, sample).""" from types import SimpleNamespace + utest_mod = _load_utest_module() run_name = "test_run" @@ -686,7 +927,13 @@ def test_load_real_perturbation_tests(self, tmp_path): n_rows = 10 self._write_dummy_results( - tmp_path, run_name, components, sel_thresh_list, samples, fake=False, n_rows=n_rows, + tmp_path, + run_name, + components, + sel_thresh_list, + samples, + fake=False, + n_rows=n_rows, ) utest_mod.args = SimpleNamespace( @@ -709,6 +956,7 @@ def test_load_real_perturbation_tests(self, tmp_path): def test_load_fake_perturbation_tests(self, tmp_path): """load_fake_perturbation_tests discovers samples and stacks per-(K, thresh, sample).""" from types import SimpleNamespace + utest_mod = _load_utest_module() run_name = "test_run" @@ -718,7 +966,13 @@ def test_load_fake_perturbation_tests(self, tmp_path): n_rows = 10 self._write_dummy_results( - tmp_path, run_name, components, sel_thresh_list, samples, fake=True, n_rows=n_rows, + tmp_path, + run_name, + components, + sel_thresh_list, + samples, + fake=True, + n_rows=n_rows, ) utest_mod.args = SimpleNamespace( diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/Gene/conftest.py b/tests/Script/Stage3_Interpretation/A_Plotting/Gene/conftest.py index 0294027..56bb074 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/Gene/conftest.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/Gene/conftest.py @@ -11,14 +11,26 @@ import numpy as np import pandas as pd import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from pathlib import Path PIPELINE_ROOT = Path(__file__).resolve().parents[5] sys.path.insert(0, str(PIPELINE_ROOT / "src")) -DEFAULT_INFERENCE_PATH = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" -PLOT_OUTPUT = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Interpretation" / "Plotting" / "Gene" +DEFAULT_INFERENCE_PATH = ( + PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" +) +PLOT_OUTPUT = ( + PIPELINE_ROOT + / "tests" + / "output" + / "torch-cNMF" + / "batch" + / "Interpretation" + / "Plotting" + / "Gene" +) @pytest.fixture(scope="session") @@ -32,6 +44,7 @@ def gene_output_dir(): def test_mdata(): """Load real h5mu (k=5) from inference output.""" import mudata + h5mu_path = DEFAULT_INFERENCE_PATH / "adata" / "cNMF_5_2_0.h5mu" if not h5mu_path.exists(): pytest.skip(f"h5mu not found: {h5mu_path}") @@ -48,19 +61,21 @@ def synthetic_gene_perturbation_tsv(gene_output_dir): """ rng = np.random.default_rng(42) rows = [] - for target in [f'gene_{i}' for i in range(10)]: + for target in [f"gene_{i}" for i in range(10)]: for prog in range(5): - rows.append({ - 'target_name': target, - 'program_name': str(prog), - 'log2FC': rng.normal(0, 0.5), - 'adj_pval': rng.uniform(0.001, 0.2), - 'pval': rng.uniform(0.0005, 0.15), - 'stat': rng.normal(0, 2), - }) + rows.append( + { + "target_name": target, + "program_name": str(prog), + "log2FC": rng.normal(0, 0.5), + "adj_pval": rng.uniform(0.001, 0.2), + "pval": rng.uniform(0.0005, 0.15), + "stat": rng.normal(0, 2), + } + ) df = pd.DataFrame(rows) - path = os.path.join(gene_output_dir, 'synthetic_gene_perturbation.txt') - df.to_csv(path, sep='\t', index=True) + path = os.path.join(gene_output_dir, "synthetic_gene_perturbation.txt") + df.to_csv(path, sep="\t", index=True) return path @@ -68,7 +83,7 @@ def synthetic_gene_perturbation_tsv(gene_output_dir): def synthetic_gene_correlation(): """Synthetic gene-gene correlation matrix for waterfall tests.""" rng = np.random.default_rng(42) - genes = [f'gene_{i}' for i in range(10)] + genes = [f"gene_{i}" for i in range(10)] mat = rng.uniform(-0.5, 1.0, (10, 10)) mat = (mat + mat.T) / 2 np.fill_diagonal(mat, np.nan) diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/Gene/test_perturbed_gene_qc_plots.py b/tests/Script/Stage3_Interpretation/A_Plotting/Gene/test_perturbed_gene_qc_plots.py index 4bbb6f8..b82b8a0 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/Gene/test_perturbed_gene_qc_plots.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/Gene/test_perturbed_gene_qc_plots.py @@ -16,7 +16,8 @@ import pandas as pd import pytest import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") import matplotlib.pyplot as plt from Stage3_Interpretation.A_Plotting.src.Perturbed_gene_QC_plots import ( @@ -33,10 +34,9 @@ class TestPlotTopProgramPerGene: - def _pick_unique_symbol(self, test_mdata): """Return a gene symbol that appears exactly once in rna.var['symbol'].""" - symbols = test_mdata['rna'].var['symbol'].astype(str) + symbols = test_mdata["rna"].var["symbol"].astype(str) counts = symbols.value_counts() unique = counts[counts == 1].index if len(unique) == 0: @@ -53,24 +53,26 @@ def test_returns_axes_and_saves(self, test_mdata, gene_output_dir): save_path=gene_output_dir, save_name="test_top_program_per_gene", ) - plt.close('all') + plt.close("all") assert ax is not None - assert os.path.isfile(os.path.join(gene_output_dir, "test_top_program_per_gene.svg")) + assert os.path.isfile( + os.path.join(gene_output_dir, "test_top_program_per_gene.svg") + ) def test_missing_gene(self, test_mdata, gene_output_dir): """plot_top_program_per_gene returns None for a missing gene (standalone).""" result = plot_top_program_per_gene( test_mdata, - Target_Gene='NONEXISTENT_GENE_XYZ', + Target_Gene="NONEXISTENT_GENE_XYZ", save_path=gene_output_dir, save_name="test_top_program_missing", ) - plt.close('all') + plt.close("all") assert result is None def test_duplicate_symbol_raises(self, test_mdata, gene_output_dir): """Symbols that map from multiple Ensembl IDs should raise ValueError.""" - symbols = test_mdata['rna'].var['symbol'].astype(str) + symbols = test_mdata["rna"].var["symbol"].astype(str) dup_counts = symbols.value_counts() dups = dup_counts[dup_counts > 1].index if len(dups) == 0: @@ -87,7 +89,7 @@ def test_duplicate_symbol_raises(self, test_mdata, gene_output_dir): def test_top_n_exceeds_programs_raises(self, test_mdata, gene_output_dir): """top_n_programs larger than n_programs should raise ValueError.""" gene = self._pick_unique_symbol(test_mdata) - n_programs = test_mdata['cNMF'].varm['loadings'].shape[0] + n_programs = test_mdata["cNMF"].varm["loadings"].shape[0] with pytest.raises(ValueError, match="exceeds the number of programs"): plot_top_program_per_gene( test_mdata, @@ -99,18 +101,19 @@ def test_top_n_exceeds_programs_raises(self, test_mdata, gene_output_dir): class TestPlotLog2FC: - - def test_returns_axes_and_df(self, synthetic_gene_perturbation_tsv, gene_output_dir): + def test_returns_axes_and_df( + self, synthetic_gene_perturbation_tsv, gene_output_dir + ): """plot_log2FC returns (Axes, DataFrame) and saves SVG.""" ax, df = plot_log2FC( synthetic_gene_perturbation_tsv, - Target='gene_0', + Target="gene_0", num_item=3, significance_threshold=0.5, save_path=gene_output_dir, save_name="test_gene_log2fc", ) - plt.close('all') + plt.close("all") assert ax is not None assert isinstance(df, pd.DataFrame) assert os.path.isfile(os.path.join(gene_output_dir, "test_gene_log2fc.svg")) @@ -120,27 +123,30 @@ def test_missing_target(self, synthetic_gene_perturbation_tsv, gene_output_dir): fig, ax = plt.subplots() result_ax, df = plot_log2FC( synthetic_gene_perturbation_tsv, - Target='NONEXISTENT', + Target="NONEXISTENT", ax=ax, ) - fig.savefig(os.path.join(gene_output_dir, "test_gene_log2fc_missing.png"), dpi=100) - plt.close('all') + fig.savefig( + os.path.join(gene_output_dir, "test_gene_log2fc_missing.png"), dpi=100 + ) + plt.close("all") assert isinstance(df, pd.DataFrame) assert len(df) == 0 class TestPlotVolcano: - - def test_returns_axes_df_texts(self, synthetic_gene_perturbation_tsv, gene_output_dir): + def test_returns_axes_df_texts( + self, synthetic_gene_perturbation_tsv, gene_output_dir + ): """plot_volcano returns (Axes, DataFrame, list) and saves SVG.""" ax, df, texts = plot_volcano( synthetic_gene_perturbation_tsv, - Target='gene_0', + Target="gene_0", significance_threshold=0.5, save_path=gene_output_dir, save_name="test_gene_volcano", ) - plt.close('all') + plt.close("all") assert ax is not None assert isinstance(df, pd.DataFrame) assert isinstance(texts, list) @@ -148,12 +154,11 @@ def test_returns_axes_df_texts(self, synthetic_gene_perturbation_tsv, gene_outpu class TestComputeGeneCorrelationMatrix: - def test_returns_square_gene_corr(self, test_mdata, gene_output_dir): """compute_gene_correlation_matrix returns a G x G row-on-demand corr and writes the .npz.""" save_path = os.path.join(gene_output_dir, "test_gene_loading_corr.npz") result = compute_gene_correlation_matrix(test_mdata, save_path=save_path) - n_genes = test_mdata['rna'].n_vars + n_genes = test_mdata["rna"].n_vars assert result.shape == (n_genes, n_genes) g = result.columns[0] row = result.loc[g] @@ -164,10 +169,9 @@ def test_returns_square_gene_corr(self, test_mdata, gene_output_dir): class TestAnalyzeCorrelations: - def test_returns_axes(self, gene_output_dir): """analyze_correlations returns Axes for a known gene.""" - genes = [f'gene_{i}' for i in range(10)] + genes = [f"gene_{i}" for i in range(10)] rng = np.random.default_rng(42) mat = rng.uniform(-0.5, 1.0, (10, 10)) mat = (mat + mat.T) / 2 @@ -176,32 +180,34 @@ def test_returns_axes(self, gene_output_dir): ax = analyze_correlations( corr, - Target='gene_0', + Target="gene_0", top_corr_genes=3, save_path=gene_output_dir, save_name="test_gene_correlations", ) - plt.close('all') + plt.close("all") assert ax is not None - assert os.path.isfile(os.path.join(gene_output_dir, "test_gene_correlations.svg")) + assert os.path.isfile( + os.path.join(gene_output_dir, "test_gene_correlations.svg") + ) def test_missing_gene(self, gene_output_dir): """analyze_correlations returns None for missing gene (standalone).""" - corr = pd.DataFrame( - np.eye(3), index=['a', 'b', 'c'], columns=['a', 'b', 'c'] - ) + corr = pd.DataFrame(np.eye(3), index=["a", "b", "c"], columns=["a", "b", "c"]) result = analyze_correlations( - corr, Target='MISSING', + corr, + Target="MISSING", save_path=gene_output_dir, save_name="test_gene_corr_missing", ) - plt.close('all') + plt.close("all") assert result is None class TestComputeGeneWaterfallCor: - - def test_returns_correlation_matrix(self, synthetic_gene_perturbation_tsv, gene_output_dir): + def test_returns_correlation_matrix( + self, synthetic_gene_perturbation_tsv, gene_output_dir + ): """compute_gene_waterfall_cor returns a T x T corr (self = 1 per row) and writes the .npz.""" save_path = os.path.join(gene_output_dir, "test_regulator_corr.npz") result = compute_gene_waterfall_cor( @@ -221,34 +227,32 @@ def test_returns_correlation_matrix(self, synthetic_gene_perturbation_tsv, gene_ class TestCreateGeneCorrelationWaterfall: - def test_returns_axes_and_texts(self, synthetic_gene_correlation, gene_output_dir): """create_gene_correlation_waterfall returns (Axes, list of texts).""" ax, texts = create_gene_correlation_waterfall( synthetic_gene_correlation, - Target_Gene='gene_0', + Target_Gene="gene_0", top_corr_genes=2, save_path=gene_output_dir, save_name="test_gene_waterfall", ) - plt.close('all') + plt.close("all") assert ax is not None assert isinstance(texts, list) assert os.path.isfile(os.path.join(gene_output_dir, "test_gene_waterfall.svg")) class TestPlotPerturbationVsControl: - def test_with_real_mdata(self, test_mdata, gene_output_dir): """plot_perturbation_vs_control works if guide_assignment exists.""" - if 'guide_assignment' not in test_mdata['cNMF'].obsm: + if "guide_assignment" not in test_mdata["cNMF"].obsm: pytest.skip("guide_assignment not in mdata") - if 'guide_targets' not in test_mdata['cNMF'].uns: + if "guide_targets" not in test_mdata["cNMF"].uns: pytest.skip("guide_targets not in mdata") # Pick a target gene that actually exists in guide_targets - targets = np.array(test_mdata['cNMF'].uns['guide_targets']) - valid = [t for t in np.unique(targets) if t != 'non-targeting'] + targets = np.array(test_mdata["cNMF"].uns["guide_targets"]) + valid = [t for t in np.unique(targets) if t != "non-targeting"] if not valid: pytest.skip("No targeting guides found") @@ -259,22 +263,25 @@ def test_with_real_mdata(self, test_mdata, gene_output_dir): target_gene=target_gene, ax=ax, ) - fig.savefig(os.path.join(gene_output_dir, "test_perturb_vs_control.png"), dpi=100) - plt.close('all') - assert os.path.isfile(os.path.join(gene_output_dir, "test_perturb_vs_control.png")) + fig.savefig( + os.path.join(gene_output_dir, "test_perturb_vs_control.png"), dpi=100 + ) + plt.close("all") + assert os.path.isfile( + os.path.join(gene_output_dir, "test_perturb_vs_control.png") + ) class TestPlotPerturbationVsControlByCondition: - def test_with_real_mdata(self, test_mdata, gene_output_dir): """plot_perturbation_vs_control_by_condition works if guide data exists.""" - if 'guide_assignment' not in test_mdata['cNMF'].obsm: + if "guide_assignment" not in test_mdata["cNMF"].obsm: pytest.skip("guide_assignment not in mdata") - if 'guide_targets' not in test_mdata['cNMF'].uns: + if "guide_targets" not in test_mdata["cNMF"].uns: pytest.skip("guide_targets not in mdata") - targets = np.array(test_mdata['cNMF'].uns['guide_targets']) - valid = [t for t in np.unique(targets) if t != 'non-targeting'] + targets = np.array(test_mdata["cNMF"].uns["guide_targets"]) + valid = [t for t in np.unique(targets) if t != "non-targeting"] if not valid: pytest.skip("No targeting guides found") @@ -283,12 +290,17 @@ def test_with_real_mdata(self, test_mdata, gene_output_dir): result = plot_perturbation_vs_control_by_condition( test_mdata, target_gene=target_gene, - condition_key='batch', + condition_key="batch", ax=ax, ) - fig.savefig(os.path.join(gene_output_dir, "test_perturb_vs_control_by_condition.png"), dpi=100) - plt.close('all') - assert os.path.isfile(os.path.join(gene_output_dir, "test_perturb_vs_control_by_condition.png")) + fig.savefig( + os.path.join(gene_output_dir, "test_perturb_vs_control_by_condition.png"), + dpi=100, + ) + plt.close("all") + assert os.path.isfile( + os.path.join(gene_output_dir, "test_perturb_vs_control_by_condition.png") + ) if __name__ == "__main__": diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/conftest.py b/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/conftest.py index 8c36462..337be50 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/conftest.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/conftest.py @@ -11,7 +11,8 @@ import numpy as np import pandas as pd import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from pathlib import Path PIPELINE_ROOT = Path(__file__).resolve().parents[5] @@ -52,10 +53,14 @@ def eval_dir(): # Real data fixtures — loaded once per session from batch output # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def real_stability_stats(inference_dir): """Load stability/error stats from real torch-cNMF batch output.""" - from Stage3_Interpretation.A_Plotting.src.k_selection_plots import load_stablity_error_data + from Stage3_Interpretation.A_Plotting.src.k_selection_plots import ( + load_stablity_error_data, + ) + output_directory = str(BATCH_OUTPUT) return load_stablity_error_data( output_directory=output_directory, @@ -68,14 +73,21 @@ def real_stability_stats(inference_dir): def real_enrichment_df(eval_dir): """Load enrichment data (GO, genesets, traits) from real evaluation output.""" # Verify all required enrichment files exist before loading - sel_str = str(SEL_THRESH).replace('.', '_') + sel_str = str(SEL_THRESH).replace(".", "_") for k in COMPONENTS: k_folder = EVAL_DIR / f"{k}_{sel_str}" - for pattern in ['{k}_GO_term_enrichment.txt', '{k}_geneset_enrichment.txt', '{k}_trait_enrichment.txt']: + for pattern in [ + "{k}_GO_term_enrichment.txt", + "{k}_geneset_enrichment.txt", + "{k}_trait_enrichment.txt", + ]: fpath = k_folder / pattern.format(k=k) if not fpath.exists(): pytest.skip(f"Missing enrichment file: {fpath}") - from Stage3_Interpretation.A_Plotting.src.k_selection_plots import load_enrichment_data + from Stage3_Interpretation.A_Plotting.src.k_selection_plots import ( + load_enrichment_data, + ) + return load_enrichment_data( folder=str(EVAL_DIR), components=COMPONENTS, @@ -87,14 +99,17 @@ def real_enrichment_df(eval_dir): @pytest.fixture(scope="session") def real_perturbation_df(eval_dir): """Load perturbation association results from real evaluation output.""" - sel_str = str(SEL_THRESH).replace('.', '_') + sel_str = str(SEL_THRESH).replace(".", "_") for k in COMPONENTS: k_folder = EVAL_DIR / f"{k}_{sel_str}" for samp in SAMPLES: fpath = k_folder / f"{k}_perturbation_association_results_{samp}.txt" if not fpath.exists(): pytest.skip(f"Missing perturbation file: {fpath}") - from Stage3_Interpretation.A_Plotting.src.k_selection_plots import load_perturbation_data + from Stage3_Interpretation.A_Plotting.src.k_selection_plots import ( + load_perturbation_data, + ) + return load_perturbation_data( folder=str(EVAL_DIR), components=COMPONENTS, @@ -106,12 +121,15 @@ def real_perturbation_df(eval_dir): @pytest.fixture(scope="session") def real_explained_var(eval_dir): """Load explained variance data from real evaluation output.""" - sel_str = str(SEL_THRESH).replace('.', '_') + sel_str = str(SEL_THRESH).replace(".", "_") for k in COMPONENTS: fpath = EVAL_DIR / f"{k}_{sel_str}" / f"{k}_Explained_Variance_Summary.txt" if not fpath.exists(): pytest.skip(f"Missing variance file: {fpath}") - from Stage3_Interpretation.A_Plotting.src.k_selection_plots import load_explained_variance_data + from Stage3_Interpretation.A_Plotting.src.k_selection_plots import ( + load_explained_variance_data, + ) + return load_explained_variance_data( folder=str(EVAL_DIR), components=COMPONENTS, @@ -126,8 +144,8 @@ def real_spectra_matrices(inference_dir): k10 = INFERENCE_DIR / "Inference.gene_spectra_score.k_10.dt_2_0.txt" if not k5.exists() or not k10.exists(): pytest.skip("Spectra score files not found") - m1 = pd.read_csv(str(k5), sep='\t', index_col=0) - m2 = pd.read_csv(str(k10), sep='\t', index_col=0) + m1 = pd.read_csv(str(k5), sep="\t", index_col=0) + m2 = pd.read_csv(str(k10), sep="\t", index_col=0) return m1, m2 @@ -135,27 +153,33 @@ def real_spectra_matrices(inference_dir): # Synthetic fixtures — for pure logic tests that don't need real data # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def synthetic_stability_stats(): """Synthetic DataFrame mimicking load_stablity_error_data output.""" ks = [5, 10, 15, 20, 30] - return pd.DataFrame({ - 'k': ks, - 'local_density_threshold': [2.0] * 5, - 'silhouette': [0.8, 0.75, 0.7, 0.65, 0.6], - 'prediction_error': [1e6, 8e5, 6e5, 5e5, 4e5], - }) + return pd.DataFrame( + { + "k": ks, + "local_density_threshold": [2.0] * 5, + "silhouette": [0.8, 0.75, 0.7, 0.65, 0.6], + "prediction_error": [1e6, 8e5, 6e5, 5e5, 4e5], + } + ) @pytest.fixture(scope="session") def synthetic_count_df(): """Synthetic enrichment count DataFrame.""" ks = [5, 10, 15, 20, 30] - return pd.DataFrame({ - 'go_terms': [3, 5, 8, 12, 10], - 'genesets': [2, 4, 6, 9, 7], - 'traits': [1, 2, 3, 5, 4], - }, index=ks) + return pd.DataFrame( + { + "go_terms": [3, 5, 8, 12, 10], + "genesets": [2, 4, 6, 9, 7], + "traits": [1, 2, 3, 5, 4], + }, + index=ks, + ) @pytest.fixture(scope="session") @@ -164,13 +188,17 @@ def synthetic_test_stats_df(): rng = np.random.default_rng(42) rows = [] for k in [5, 10, 15, 20, 30]: - for samp in ['D0', 'sample_D1', 'sample_D2']: - for target in [f'gene_{i}' for i in range(5)]: - rows.append({ - 'K': k, 'sample': samp, 'target_name': target, - 'adj_pval': rng.uniform(0.0001, 0.1), - 'log2FC': rng.normal(0, 0.5), - }) + for samp in ["D0", "sample_D1", "sample_D2"]: + for target in [f"gene_{i}" for i in range(5)]: + rows.append( + { + "K": k, + "sample": samp, + "target_name": target, + "adj_pval": rng.uniform(0.0001, 0.1), + "log2FC": rng.normal(0, 0.5), + } + ) return pd.DataFrame(rows) @@ -184,8 +212,8 @@ def synthetic_explained_var(): def synthetic_spectra_matrices(): """Two synthetic programs x genes DataFrames for logic tests.""" rng = np.random.default_rng(42) - genes = [f'gene_{i}' for i in range(100)] - programs = [f'prog_{i}' for i in range(5)] + genes = [f"gene_{i}" for i in range(100)] + programs = [f"prog_{i}" for i in range(5)] m1 = pd.DataFrame(rng.random((5, 100)), index=programs, columns=genes) m2 = pd.DataFrame(rng.random((5, 100)), index=programs, columns=genes) return m1, m2 diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/test_k_selection_plots.py b/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/test_k_selection_plots.py index 6179137..6129496 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/test_k_selection_plots.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/KSelection/test_k_selection_plots.py @@ -16,7 +16,8 @@ import pandas as pd import pytest import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") import matplotlib.pyplot as plt from Stage3_Interpretation.A_Plotting.src.k_selection_plots import ( @@ -40,38 +41,42 @@ # Helper / pure-logic tests (no real data needed) # --------------------------------------------------------------------------- -class TestHelpers: +class TestHelpers: def test_style_ax(self): """_style_ax applies styling without crashing.""" fig, ax = plt.subplots() - _style_ax(ax, xlabel='X', ylabel='Y', title='Test') - assert not ax.spines['top'].get_visible() - assert not ax.spines['right'].get_visible() - plt.close('all') + _style_ax(ax, xlabel="X", ylabel="Y", title="Test") + assert not ax.spines["top"].get_visible() + assert not ax.spines["right"].get_visible() + plt.close("all") def test_filter_components(self): """_filter_components keeps only requested K values.""" - stats = pd.DataFrame({ - 'k': [5, 10, 15, 20], - 'silhouette': [0.8, 0.7, 0.6, 0.5], - 'prediction_error': [1e6, 8e5, 6e5, 4e5], - }) + stats = pd.DataFrame( + { + "k": [5, 10, 15, 20], + "silhouette": [0.8, 0.7, 0.6, 0.5], + "prediction_error": [1e6, 8e5, 6e5, 4e5], + } + ) result = _filter_components(stats, [5, 15, 99]) - assert list(result['k']) == [5, 15] + assert list(result["k"]) == [5, 15] assert len(result) == 2 def test_load_stats_file_tsv(self, kselection_output_dir): """_load_stats_file reads a TSV file correctly.""" - df = pd.DataFrame({ - 'k': [5, 10], - 'silhouette': [0.8, 0.7], - 'prediction_error': [1e6, 8e5], - }) - path = os.path.join(kselection_output_dir, 'test_stats.tsv') - df.to_csv(path, sep='\t', index=False) + df = pd.DataFrame( + { + "k": [5, 10], + "silhouette": [0.8, 0.7], + "prediction_error": [1e6, 8e5], + } + ) + path = os.path.join(kselection_output_dir, "test_stats.tsv") + df.to_csv(path, sep="\t", index=False) result = _load_stats_file(path) - assert list(result.columns) == ['k', 'silhouette', 'prediction_error'] + assert list(result.columns) == ["k", "silhouette", "prediction_error"] assert len(result) == 2 @@ -79,34 +84,33 @@ def test_load_stats_file_tsv(self, kselection_output_dir): # Data loader tests (real data) # --------------------------------------------------------------------------- -class TestLoadStabilityError: +class TestLoadStabilityError: def test_returns_valid_dataframe(self, real_stability_stats): """load_stablity_error_data returns DataFrame with expected columns.""" df = real_stability_stats assert isinstance(df, pd.DataFrame) - assert 'k' in df.columns - assert 'silhouette' in df.columns - assert 'prediction_error' in df.columns + assert "k" in df.columns + assert "silhouette" in df.columns + assert "prediction_error" in df.columns def test_contains_expected_k_values(self, real_stability_stats): """Loaded stats should contain K=5,10,15.""" - ks = set(real_stability_stats['k'].tolist()) + ks = set(real_stability_stats["k"].tolist()) assert {5, 10, 15}.issubset(ks) def test_silhouette_in_valid_range(self, real_stability_stats): """Silhouette scores should be in [-1, 1].""" - vals = real_stability_stats['silhouette'] + vals = real_stability_stats["silhouette"] assert vals.min() >= -1.0 assert vals.max() <= 1.0 def test_prediction_error_positive(self, real_stability_stats): """Prediction errors should be non-negative.""" - assert (real_stability_stats['prediction_error'] >= 0).all() + assert (real_stability_stats["prediction_error"] >= 0).all() class TestLoadEnrichment: - def test_returns_valid_dataframe(self, real_enrichment_df): """load_enrichment_data returns DataFrame indexed by K.""" df = real_enrichment_df @@ -120,23 +124,21 @@ def test_counts_are_non_negative(self, real_enrichment_df): class TestLoadPerturbation: - def test_returns_valid_dataframe(self, real_perturbation_df): """load_perturbation_data returns DataFrame with sample and K columns.""" df = real_perturbation_df assert isinstance(df, pd.DataFrame) - assert 'sample' in df.columns - assert 'K' in df.columns + assert "sample" in df.columns + assert "K" in df.columns assert len(df) > 0 def test_contains_expected_samples(self, real_perturbation_df): """Loaded data should contain at least some of the expected samples.""" - samples = set(real_perturbation_df['sample'].unique()) + samples = set(real_perturbation_df["sample"].unique()) assert len(samples) > 0 class TestLoadExplainedVariance: - def test_returns_valid_dict(self, real_explained_var): """load_explained_variance_data returns a dict mapping K to float.""" ev = real_explained_var @@ -156,8 +158,8 @@ def test_variance_values_reasonable(self, real_explained_var): # Plot tests (real data) # --------------------------------------------------------------------------- -class TestPlotStabilityError: +class TestPlotStabilityError: def test_saves_files(self, real_stability_stats, kselection_output_dir): """plot_stablity_error saves stability and error PNG/SVG files.""" plot_stablity_error( @@ -166,9 +168,13 @@ def test_saves_files(self, real_stability_stats, kselection_output_dir): file_name="test_stability_error", selected_k=10, ) - plt.close('all') - assert os.path.isfile(os.path.join(kselection_output_dir, "test_stability_error_stability.png")) - assert os.path.isfile(os.path.join(kselection_output_dir, "test_stability_error_error.png")) + plt.close("all") + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_stability_error_stability.png") + ) + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_stability_error_error.png") + ) def test_no_selected_k(self, real_stability_stats, kselection_output_dir): """plot_stablity_error works without selected_k.""" @@ -177,12 +183,13 @@ def test_no_selected_k(self, real_stability_stats, kselection_output_dir): folder_name=kselection_output_dir, file_name="test_stability_no_vline", ) - plt.close('all') - assert os.path.isfile(os.path.join(kselection_output_dir, "test_stability_no_vline_stability.png")) + plt.close("all") + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_stability_no_vline_stability.png") + ) class TestPlotEnrichment: - def test_saves_files(self, real_enrichment_df, kselection_output_dir): """plot_enrichment saves one file per enrichment category.""" plot_enrichment( @@ -191,14 +198,17 @@ def test_saves_files(self, real_enrichment_df, kselection_output_dir): file_name="test_enrichment", selected_k=10, ) - plt.close('all') + plt.close("all") for col in real_enrichment_df.columns: - assert os.path.isfile(os.path.join(kselection_output_dir, f"test_enrichment_{col}.png")) + assert os.path.isfile( + os.path.join(kselection_output_dir, f"test_enrichment_{col}.png") + ) class TestPlotPerturbation: - - def test_saves_files_and_returns_df(self, real_perturbation_df, kselection_output_dir): + def test_saves_files_and_returns_df( + self, real_perturbation_df, kselection_output_dir + ): """plot_perturbation saves per-sample and aggregated plots, returns DataFrame.""" result = plot_perturbation( real_perturbation_df, @@ -207,14 +217,17 @@ def test_saves_files_and_returns_df(self, real_perturbation_df, kselection_outpu file_name="test_perturbation", selected_k=10, ) - plt.close('all') + plt.close("all") assert isinstance(result, pd.DataFrame) - assert os.path.isfile(os.path.join(kselection_output_dir, "test_perturbation_per_sample.png")) - assert os.path.isfile(os.path.join(kselection_output_dir, "test_perturbation_all_samples.png")) + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_perturbation_per_sample.png") + ) + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_perturbation_all_samples.png") + ) class TestPlotExplainedVariance: - def test_saves_file(self, real_explained_var, kselection_output_dir): """plot_explained_variance saves PNG file.""" plot_explained_variance( @@ -223,15 +236,21 @@ def test_saves_file(self, real_explained_var, kselection_output_dir): file_name="test_explained_var", selected_k=10, ) - plt.close('all') - assert os.path.isfile(os.path.join(kselection_output_dir, "test_explained_var.png")) + plt.close("all") + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_explained_var.png") + ) class TestPlotPanel: - - def test_k_selection_panel(self, real_stability_stats, real_enrichment_df, - real_perturbation_df, real_explained_var, - kselection_output_dir): + def test_k_selection_panel( + self, + real_stability_stats, + real_enrichment_df, + real_perturbation_df, + real_explained_var, + kselection_output_dir, + ): """plot_k_selection_panel saves the combined panel.""" plot_k_selection_panel( real_stability_stats, @@ -243,12 +262,17 @@ def test_k_selection_panel(self, real_stability_stats, real_enrichment_df, file_name="test_panel", selected_k=10, ) - plt.close('all') + plt.close("all") assert os.path.isfile(os.path.join(kselection_output_dir, "test_panel.png")) - def test_k_selection_panel_no_traits(self, real_stability_stats, real_enrichment_df, - real_perturbation_df, real_explained_var, - kselection_output_dir): + def test_k_selection_panel_no_traits( + self, + real_stability_stats, + real_enrichment_df, + real_perturbation_df, + real_explained_var, + kselection_output_dir, + ): """plot_k_selection_panel_no_traits saves the combined panel without traits.""" plot_k_selection_panel_no_traits( real_stability_stats, @@ -260,8 +284,10 @@ def test_k_selection_panel_no_traits(self, real_stability_stats, real_enrichment file_name="test_panel_no_traits", selected_k=10, ) - plt.close('all') - assert os.path.isfile(os.path.join(kselection_output_dir, "test_panel_no_traits.png")) + plt.close("all") + assert os.path.isfile( + os.path.join(kselection_output_dir, "test_panel_no_traits.png") + ) if __name__ == "__main__": diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/Program/conftest.py b/tests/Script/Stage3_Interpretation/A_Plotting/Program/conftest.py index e710092..42fc3bf 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/Program/conftest.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/Program/conftest.py @@ -11,14 +11,26 @@ import numpy as np import pandas as pd import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from pathlib import Path PIPELINE_ROOT = Path(__file__).resolve().parents[5] sys.path.insert(0, str(PIPELINE_ROOT / "src")) -DEFAULT_INFERENCE_PATH = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" -PLOT_OUTPUT = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Interpretation" / "Plotting" / "Program" +DEFAULT_INFERENCE_PATH = ( + PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" +) +PLOT_OUTPUT = ( + PIPELINE_ROOT + / "tests" + / "output" + / "torch-cNMF" + / "batch" + / "Interpretation" + / "Plotting" + / "Program" +) @pytest.fixture(scope="session") @@ -32,6 +44,7 @@ def program_output_dir(): def test_mdata(): """Load real h5mu (k=5) from inference output.""" import mudata + h5mu_path = DEFAULT_INFERENCE_PATH / "adata" / "cNMF_5_2_0.h5mu" if not h5mu_path.exists(): pytest.skip(f"h5mu not found: {h5mu_path}") @@ -49,18 +62,20 @@ def synthetic_perturbation_tsv(program_output_dir): rng = np.random.default_rng(42) rows = [] for prog in range(5): - for target in [f'gene_{i}' for i in range(10)]: - rows.append({ - 'target_name': str(prog), - 'program_name': target, - 'log2FC': rng.normal(0, 0.5), - 'adj_pval': rng.uniform(0.001, 0.2), - 'pval': rng.uniform(0.0005, 0.15), - 'stat': rng.normal(0, 2), - }) + for target in [f"gene_{i}" for i in range(10)]: + rows.append( + { + "target_name": str(prog), + "program_name": target, + "log2FC": rng.normal(0, 0.5), + "adj_pval": rng.uniform(0.001, 0.2), + "pval": rng.uniform(0.0005, 0.15), + "stat": rng.normal(0, 2), + } + ) df = pd.DataFrame(rows) - path = os.path.join(program_output_dir, 'synthetic_perturbation.txt') - df.to_csv(path, sep='\t', index=False) + path = os.path.join(program_output_dir, "synthetic_perturbation.txt") + df.to_csv(path, sep="\t", index=False) return path @@ -71,21 +86,23 @@ def synthetic_perturbation_base(program_output_dir): Returns the base path (without _SAMPLE.txt suffix). """ rng = np.random.default_rng(42) - base = os.path.join(program_output_dir, 'synthetic_perturb') - for samp in ['D0', 'sample_D1']: + base = os.path.join(program_output_dir, "synthetic_perturb") + for samp in ["D0", "sample_D1"]: rows = [] for prog in range(5): - for target in [f'gene_{i}' for i in range(8)]: - rows.append({ - 'target_name': str(prog), - 'program_name': target, - 'log2FC': rng.normal(0, 0.5), - 'adj_pval': rng.uniform(0.001, 0.2), - 'pval': rng.uniform(0.0005, 0.15), - 'stat': rng.normal(0, 2), - }) + for target in [f"gene_{i}" for i in range(8)]: + rows.append( + { + "target_name": str(prog), + "program_name": target, + "log2FC": rng.normal(0, 0.5), + "adj_pval": rng.uniform(0.001, 0.2), + "pval": rng.uniform(0.0005, 0.15), + "stat": rng.normal(0, 2), + } + ) df = pd.DataFrame(rows) - df.to_csv(f'{base}_{samp}.txt', sep='\t', index=False) + df.to_csv(f"{base}_{samp}.txt", sep="\t", index=False) return base @@ -107,14 +124,16 @@ def synthetic_go_tsv(program_output_dir): rows = [] for prog in range(5): for i in range(8): - rows.append({ - 'program_name': prog, - 'Term': f'GO:00{prog}{i:02d} regulation of process {i}', - 'Adjusted P-value': np.random.uniform(0.001, 0.1), - }) + rows.append( + { + "program_name": prog, + "Term": f"GO:00{prog}{i:02d} regulation of process {i}", + "Adjusted P-value": np.random.uniform(0.001, 0.1), + } + ) df = pd.DataFrame(rows) - df.index = df['program_name'] - df = df.drop(columns=['program_name']) - path = os.path.join(program_output_dir, 'synthetic_GO.txt') - df.to_csv(path, sep='\t') + df.index = df["program_name"] + df = df.drop(columns=["program_name"]) + path = os.path.join(program_output_dir, "synthetic_GO.txt") + df.to_csv(path, sep="\t") return path diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_expression_weighted_plots.py b/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_expression_weighted_plots.py index 54fa9b7..6b1eebc 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_expression_weighted_plots.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_expression_weighted_plots.py @@ -15,7 +15,8 @@ import pandas as pd import pytest import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") import matplotlib.pyplot as plt from Stage3_Interpretation.A_Plotting.src.Program_expression_weighted_plots import ( @@ -26,12 +27,13 @@ class TestComputeProgramExpressionByCondition: - def test_returns_series(self, test_mdata): """compute_program_expression_by_condition returns a Series indexed by condition.""" - prog_name = test_mdata['cNMF'].var_names[0] + prog_name = test_mdata["cNMF"].var_names[0] result = compute_program_expression_by_condition( - test_mdata, Target_Program=prog_name, groupby='batch', + test_mdata, + Target_Program=prog_name, + groupby="batch", ) assert isinstance(result, pd.Series) assert len(result) > 0 @@ -39,44 +41,50 @@ def test_returns_series(self, test_mdata): class TestPlotProgramHeatmapWeighted: - - def test_saves_file(self, test_mdata, synthetic_perturbation_base, program_output_dir): + def test_saves_file( + self, test_mdata, synthetic_perturbation_base, program_output_dir + ): """plot_program_heatmap_weighted produces a plot (or returns None if no sig regulators).""" - prog_name = test_mdata['cNMF'].var_names[0] + prog_name = test_mdata["cNMF"].var_names[0] result = plot_program_heatmap_weighted( perturb_path_base=synthetic_perturbation_base, mdata=test_mdata, Target_Program=prog_name, - sample=['D0', 'sample_D1'], - groupby='batch', + sample=["D0", "sample_D1"], + groupby="batch", p_value=0.5, # relaxed threshold to ensure something passes save_path=program_output_dir, save_name="test_heatmap_weighted", ) - plt.close('all') + plt.close("all") # Returns Axes if regulators found, None if empty if result is not None: - assert os.path.isfile(os.path.join(program_output_dir, "test_heatmap_weighted.pdf")) + assert os.path.isfile( + os.path.join(program_output_dir, "test_heatmap_weighted.pdf") + ) class TestPlotProgramHeatmapExpressionScaled: - - def test_saves_file(self, test_mdata, synthetic_perturbation_base, program_output_dir): + def test_saves_file( + self, test_mdata, synthetic_perturbation_base, program_output_dir + ): """plot_program_heatmap_expression_scaled produces a scaled heatmap.""" - prog_name = test_mdata['cNMF'].var_names[0] + prog_name = test_mdata["cNMF"].var_names[0] result = plot_program_heatmap_expression_scaled( perturb_path_base=synthetic_perturbation_base, mdata=test_mdata, Target_Program=prog_name, - sample=['D0', 'sample_D1'], - groupby='batch', + sample=["D0", "sample_D1"], + groupby="batch", p_value=0.5, save_path=program_output_dir, save_name="test_heatmap_scaled", ) - plt.close('all') + plt.close("all") if result is not None: - assert os.path.isfile(os.path.join(program_output_dir, "test_heatmap_scaled.pdf")) + assert os.path.isfile( + os.path.join(program_output_dir, "test_heatmap_scaled.pdf") + ) if __name__ == "__main__": diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_qc_plots.py b/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_qc_plots.py index 142a7c7..933010b 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_qc_plots.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_program_qc_plots.py @@ -16,7 +16,8 @@ import pandas as pd import pytest import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") import matplotlib.pyplot as plt from Stage3_Interpretation.A_Plotting.src.Program_QC_plots import ( @@ -33,20 +34,19 @@ class TestComputeProgramCorrelation: - def test_returns_symmetric_dataframe(self, test_mdata, program_output_dir): """compute_program_correlation_matrix returns a symmetric DataFrame equal to pandas .corr().""" save_path = os.path.join(program_output_dir, "test_program_corr.npz") result = compute_program_correlation_matrix(test_mdata, save_path=save_path) assert isinstance(result, pd.DataFrame) assert result.shape[0] == result.shape[1] - n_programs = test_mdata['cNMF'].n_vars + n_programs = test_mdata["cNMF"].n_vars assert result.shape[0] == n_programs # Check symmetry np.testing.assert_allclose(result.values, result.values.T, atol=1e-10) # Same values as the legacy pandas path - X = test_mdata['cNMF'].X - X = X.toarray() if hasattr(X, 'toarray') else X + X = test_mdata["cNMF"].X + X = X.toarray() if hasattr(X, "toarray") else X expected = pd.DataFrame(X).corr().fillna(0).values np.testing.assert_allclose(result.values, expected, atol=1e-8) with np.load(save_path) as f: @@ -55,7 +55,6 @@ def test_returns_symmetric_dataframe(self, test_mdata, program_output_dir): class TestAnalyzeProgramCorrelations: - def test_returns_axes(self, synthetic_program_correlation, program_output_dir): """analyze_program_correlations returns an Axes object.""" ax = analyze_program_correlations( @@ -65,7 +64,7 @@ def test_returns_axes(self, synthetic_program_correlation, program_output_dir): save_path=program_output_dir, save_name="test_program_correlations", ) - plt.close('all') + plt.close("all") assert ax is not None def test_missing_program(self, synthetic_program_correlation, program_output_dir): @@ -76,16 +75,15 @@ def test_missing_program(self, synthetic_program_correlation, program_output_dir save_path=program_output_dir, save_name="test_program_corr_missing", ) - plt.close('all') + plt.close("all") # Returns None in standalone mode when program not found assert result is None class TestPlotTopGenePerProgram: - def test_returns_axes_and_saves(self, test_mdata, program_output_dir): """plot_top_gene_per_program returns Axes and saves SVG.""" - prog_name = test_mdata['cNMF'].var_names[0] + prog_name = test_mdata["cNMF"].var_names[0] ax = plot_top_gene_per_program( test_mdata, Target_Program=prog_name, @@ -93,9 +91,11 @@ def test_returns_axes_and_saves(self, test_mdata, program_output_dir): save_path=program_output_dir, save_name="test_top_gene_per_program", ) - plt.close('all') + plt.close("all") assert ax is not None - assert os.path.isfile(os.path.join(program_output_dir, "test_top_gene_per_program.svg")) + assert os.path.isfile( + os.path.join(program_output_dir, "test_top_gene_per_program.svg") + ) def test_missing_program_raises(self, test_mdata, program_output_dir): """plot_top_gene_per_program raises ValueError for a program id that doesn't exist.""" @@ -110,8 +110,8 @@ def test_missing_program_raises(self, test_mdata, program_output_dir): def test_num_gene_exceeds_genes_raises(self, test_mdata, program_output_dir): """num_gene larger than n_genes should raise ValueError.""" - prog_name = test_mdata['cNMF'].var_names[0] - n_genes = test_mdata['cNMF'].varm['loadings'].shape[1] + prog_name = test_mdata["cNMF"].var_names[0] + n_genes = test_mdata["cNMF"].varm["loadings"].shape[1] with pytest.raises(ValueError, match="exceeds the number of genes"): plot_top_gene_per_program( test_mdata, @@ -123,80 +123,90 @@ def test_num_gene_exceeds_genes_raises(self, test_mdata, program_output_dir): class TestPlotViolin: - def test_returns_axes_and_saves(self, test_mdata, program_output_dir): """plot_violin returns Axes and saves SVG.""" - prog_name = test_mdata['cNMF'].var_names[0] + prog_name = test_mdata["cNMF"].var_names[0] ax = plot_violin( test_mdata, Target_Program=prog_name, - groupby='batch', + groupby="batch", save_path=program_output_dir, save_name="test_violin", ) - plt.close('all') + plt.close("all") assert ax is not None assert os.path.isfile(os.path.join(program_output_dir, "test_violin.svg")) class TestPlotProgramLog2FC: - def test_returns_axes_and_df(self, synthetic_perturbation_tsv, program_output_dir): """plot_program_log2FC returns (Axes, DataFrame).""" fig, ax = plt.subplots() ax, df = plot_program_log2FC( synthetic_perturbation_tsv, - Target='0', - tagert_col_name='target_name', - plot_col_name='program_name', + Target="0", + tagert_col_name="target_name", + plot_col_name="program_name", num_item=3, p_value=0.5, ax=ax, ) - fig.savefig(os.path.join(program_output_dir, "test_program_log2fc.png"), dpi=100) - plt.close('all') + fig.savefig( + os.path.join(program_output_dir, "test_program_log2fc.png"), dpi=100 + ) + plt.close("all") assert ax is not None assert isinstance(df, pd.DataFrame) - assert os.path.isfile(os.path.join(program_output_dir, "test_program_log2fc.png")) + assert os.path.isfile( + os.path.join(program_output_dir, "test_program_log2fc.png") + ) def test_missing_target(self, synthetic_perturbation_tsv, program_output_dir): """plot_program_log2FC handles missing target gracefully.""" fig, ax = plt.subplots() result_ax, df = plot_program_log2FC( synthetic_perturbation_tsv, - Target='NONEXISTENT', + Target="NONEXISTENT", ax=ax, ) - fig.savefig(os.path.join(program_output_dir, "test_program_log2fc_missing.png"), dpi=100) - plt.close('all') + fig.savefig( + os.path.join(program_output_dir, "test_program_log2fc_missing.png"), dpi=100 + ) + plt.close("all") assert isinstance(df, pd.DataFrame) assert len(df) == 0 class TestPlotProgramVolcano: - - def test_returns_axes_df_texts(self, synthetic_perturbation_tsv, program_output_dir): + def test_returns_axes_df_texts( + self, synthetic_perturbation_tsv, program_output_dir + ): """plot_program_volcano returns (Axes, DataFrame, list).""" fig, ax = plt.subplots() result_ax, df, texts = plot_program_volcano( synthetic_perturbation_tsv, - Target='0', - tagert_col_name='target_name', - plot_col_name='program_name', + Target="0", + tagert_col_name="target_name", + plot_col_name="program_name", p_value=0.5, ax=ax, ) - fig.savefig(os.path.join(program_output_dir, "test_program_volcano.png"), dpi=100) - plt.close('all') + fig.savefig( + os.path.join(program_output_dir, "test_program_volcano.png"), dpi=100 + ) + plt.close("all") assert result_ax is not None assert isinstance(df, pd.DataFrame) assert isinstance(texts, list) - assert os.path.isfile(os.path.join(program_output_dir, "test_program_volcano.png")) + assert os.path.isfile( + os.path.join(program_output_dir, "test_program_volcano.png") + ) class TestComputeProgramWaterfallCor: - - def test_returns_correlation_matrix(self, synthetic_perturbation_tsv, program_output_dir): + def test_returns_correlation_matrix( + self, synthetic_perturbation_tsv, program_output_dir + ): """compute_program_waterfall_cor returns a DataFrame with diagonal 1 and writes the .npz.""" save_path = os.path.join(program_output_dir, "test_program_waterfall_corr.npz") result = compute_program_waterfall_cor( @@ -206,7 +216,9 @@ def test_returns_correlation_matrix(self, synthetic_perturbation_tsv, program_ou assert isinstance(result, pd.DataFrame) assert result.shape[0] == result.shape[1] d = np.diag(result.values) - assert np.allclose(d[np.isfinite(d)], 1.0) # self kept; the waterfall plots drop it + assert np.allclose( + d[np.isfinite(d)], 1.0 + ) # self kept; the waterfall plots drop it with np.load(save_path) as f: full = f["corr"] assert list(f["program"]) == list(result.index) @@ -217,8 +229,9 @@ def test_returns_correlation_matrix(self, synthetic_perturbation_tsv, program_ou class TestCreateProgramCorrelationWaterfall: - - def test_returns_axes_and_texts(self, synthetic_program_correlation, program_output_dir): + def test_returns_axes_and_texts( + self, synthetic_program_correlation, program_output_dir + ): """create_program_correlation_waterfall returns (Axes, list of texts) and skips self.""" # Diagonal 1, as compute_program_waterfall_cor produces; the plot drops self corr = synthetic_program_correlation.copy() @@ -229,16 +242,21 @@ def test_returns_axes_and_texts(self, synthetic_program_correlation, program_out top_num=2, ax=ax, ) - fig.savefig(os.path.join(program_output_dir, "test_program_waterfall.png"), dpi=100) - plt.close('all') + fig.savefig( + os.path.join(program_output_dir, "test_program_waterfall.png"), dpi=100 + ) + plt.close("all") assert result_ax is not None assert isinstance(texts, list) - assert "0" not in [t.get_text() for t in texts] # the program itself is not labeled - assert os.path.isfile(os.path.join(program_output_dir, "test_program_waterfall.png")) + assert "0" not in [ + t.get_text() for t in texts + ] # the program itself is not labeled + assert os.path.isfile( + os.path.join(program_output_dir, "test_program_waterfall.png") + ) class TestTopGOPerProgram: - def test_returns_axes_and_labels(self, synthetic_go_tsv, program_output_dir): """top_GO_per_program returns (Axes, list of wrapped labels).""" ax, labels = top_GO_per_program( @@ -248,11 +266,13 @@ def test_returns_axes_and_labels(self, synthetic_go_tsv, program_output_dir): save_path=program_output_dir, save_name="test_go_per_program", ) - plt.close('all') + plt.close("all") assert ax is not None assert isinstance(labels, list) assert len(labels) == 3 - assert os.path.isfile(os.path.join(program_output_dir, "test_go_per_program.svg")) + assert os.path.isfile( + os.path.join(program_output_dir, "test_go_per_program.svg") + ) if __name__ == "__main__": diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/conftest.py b/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/conftest.py index 81db104..f68e851 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/conftest.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/conftest.py @@ -18,15 +18,29 @@ import numpy as np import pandas as pd import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from pathlib import Path PIPELINE_ROOT = Path(__file__).resolve().parents[5] sys.path.insert(0, str(PIPELINE_ROOT / "src")) -DEFAULT_INFERENCE_PATH = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" -EVAL_5_2_0 = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Evaluation" / "5_2_0" -HTML_OUTPUT = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Interpretation" / "Plotting" / "gene_html" +DEFAULT_INFERENCE_PATH = ( + PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" +) +EVAL_5_2_0 = ( + PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Evaluation" / "5_2_0" +) +HTML_OUTPUT = ( + PIPELINE_ROOT + / "tests" + / "output" + / "torch-cNMF" + / "batch" + / "Interpretation" + / "Plotting" + / "gene_html" +) SAMPLES = ["D0", "D1", "D2", "D3"] @@ -48,6 +62,7 @@ def html_share_path(): def test_mdata(): """Real k=5 MuData from batch inference output.""" import mudata + h5mu_path = DEFAULT_INFERENCE_PATH / "adata" / "cNMF_5_2_0.h5mu" if not h5mu_path.exists(): pytest.skip(f"h5mu not found: {h5mu_path}") @@ -108,7 +123,10 @@ def gene_loading_corr_matrix(test_mdata, target_gene): @pytest.fixture(scope="session") def perturb_corr_by_sample(perturb_path_base, available_samples): """Real per-sample gene x gene waterfall correlation matrices.""" - from Stage3_Interpretation.A_Plotting.src.Perturbed_gene_QC_plots import compute_gene_waterfall_cor + from Stage3_Interpretation.A_Plotting.src.Perturbed_gene_QC_plots import ( + compute_gene_waterfall_cor, + ) + out = {} for samp in available_samples: path = f"{perturb_path_base}_{samp}.txt" diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene.py b/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene.py index 78e571b..c22a558 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene.py @@ -14,6 +14,7 @@ import pytest import matplotlib + matplotlib.use("Agg") from Stage3_Interpretation.A_Plotting.src.html_Perturbed_gene_QC_plots import ( @@ -32,11 +33,16 @@ # JSON-builder helpers # --------------------------------------------------------------------------- -class TestBuilders: +class TestBuilders: def test_build_top_programs(self, test_mdata, target_gene): - d = _build_top_programs(test_mdata, target_gene, top_n=3, - ensembl_to_symbol_file=None, gene_name_key="symbol") + d = _build_top_programs( + test_mdata, + target_gene, + top_n=3, + ensembl_to_symbol_file=None, + gene_name_key="symbol", + ) assert set(d.keys()) == {"programs", "loadings"} assert len(d["programs"]) <= 3 @@ -47,31 +53,46 @@ def test_build_gene_correlations(self, gene_loading_corr_matrix, target_gene): def test_build_log2fc_gene(self, perturb_path_base, target_gene, available_samples): d = _build_log2fc_gene( - f"{perturb_path_base}_{available_samples[0]}.txt", target_gene, - target_col="target_name", program_col="program_name", - log2fc_col="log2FC", num_item=3, p_value=0.5, + f"{perturb_path_base}_{available_samples[0]}.txt", + target_gene, + target_col="target_name", + program_col="program_name", + log2fc_col="log2FC", + num_item=3, + p_value=0.5, ) assert "regulators" in d and "log2fc" in d and "adj_pval" in d - def test_build_volcano_gene(self, perturb_path_base, target_gene, available_samples): + def test_build_volcano_gene( + self, perturb_path_base, target_gene, available_samples + ): d = _build_volcano_gene( - f"{perturb_path_base}_{available_samples[0]}.txt", target_gene, - target_col="target_name", program_col="program_name", + f"{perturb_path_base}_{available_samples[0]}.txt", + target_gene, + target_col="target_name", + program_col="program_name", log2fc_col="log2FC", - down_thred_log=-0.05, up_thred_log=0.05, p_value=0.5, + down_thred_log=-0.05, + up_thred_log=0.05, + p_value=0.5, ) assert "category" in d assert all(c in {"up", "down", "ns"} for c in d["category"]) - def test_build_gene_waterfall(self, perturb_corr_by_sample, target_gene, available_samples): - d = _build_gene_waterfall(perturb_corr_by_sample[available_samples[0]], - target_gene, top_num=2) + def test_build_gene_waterfall( + self, perturb_corr_by_sample, target_gene, available_samples + ): + d = _build_gene_waterfall( + perturb_corr_by_sample[available_samples[0]], target_gene, top_num=2 + ) assert set(d.keys()) == {"programs", "r", "labeled"} def test_build_kd_vs_control(self, test_mdata, target_gene): d = _build_kd_vs_control( - test_mdata, target_gene, - condition_key="batch", control_target_name="non-targeting", + test_mdata, + target_gene, + condition_key="batch", + control_target_name="non-targeting", gene_name_key="symbol", ) assert "groups" in d @@ -81,12 +102,19 @@ def test_build_kd_vs_control(self, test_mdata, target_gene): # End-to-end: export_gene_html writes the full subtree # --------------------------------------------------------------------------- -class TestExportGeneHTML: +class TestExportGeneHTML: @pytest.fixture(scope="class") - def exported_gene(self, test_mdata, target_gene, perturb_path_base, - gene_loading_corr_matrix, perturb_corr_by_sample, - html_share_path, available_samples): + def exported_gene( + self, + test_mdata, + target_gene, + perturb_path_base, + gene_loading_corr_matrix, + perturb_corr_by_sample, + html_share_path, + available_samples, + ): export_gene_html( mdata=test_mdata, perturb_path_base=perturb_path_base, @@ -107,7 +135,11 @@ def exported_gene(self, test_mdata, target_gene, perturb_path_base, position_index=1, position_total=1, ) - return Path(html_share_path) / f"gene_{target_gene}", target_gene, available_samples + return ( + Path(html_share_path) / f"gene_{target_gene}", + target_gene, + available_samples, + ) def test_html_page_written(self, exported_gene): gene_dir, gene, _ = exported_gene @@ -155,12 +187,19 @@ def test_metadata_json(self, exported_gene): # write_gene_share_index # --------------------------------------------------------------------------- -class TestWriteGeneShareIndex: +class TestWriteGeneShareIndex: @pytest.fixture - def ensure_one_gene_exported(self, test_mdata, target_gene, perturb_path_base, - gene_loading_corr_matrix, perturb_corr_by_sample, - html_share_path, available_samples): + def ensure_one_gene_exported( + self, + test_mdata, + target_gene, + perturb_path_base, + gene_loading_corr_matrix, + perturb_corr_by_sample, + html_share_path, + available_samples, + ): gene_dir = Path(html_share_path) / f"gene_{target_gene}" if not (gene_dir / f"gene_{target_gene}.html").exists(): export_gene_html( @@ -183,8 +222,12 @@ def ensure_one_gene_exported(self, test_mdata, target_gene, perturb_path_base, ) return target_gene - def test_share_index_written(self, ensure_one_gene_exported, html_share_path, target_gene): - write_gene_share_index(html_share_path, [target_gene], {"test": "config", "k": 5}) + def test_share_index_written( + self, ensure_one_gene_exported, html_share_path, target_gene + ): + write_gene_share_index( + html_share_path, [target_gene], {"test": "config", "k": 5} + ) share = Path(html_share_path) assert (share / "shared" / "style.css").is_file() diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene_no_perturb.py b/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene_no_perturb.py index 8c52aed..fd8d8eb 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene_no_perturb.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/gene_html/test_html_gene_no_perturb.py @@ -17,13 +17,23 @@ import pytest import matplotlib + matplotlib.use("Agg") -from Stage3_Interpretation.A_Plotting.src.html_Perturbed_gene_QC_plots import export_gene_html +from Stage3_Interpretation.A_Plotting.src.html_Perturbed_gene_QC_plots import ( + export_gene_html, +) -def _export(tmp_path, test_mdata, target_gene, gene_loading_corr_matrix, - available_samples, perturb_path_base=None, perturb_corr_by_sample=None): +def _export( + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, + perturb_path_base=None, + perturb_corr_by_sample=None, +): export_gene_html( mdata=test_mdata, perturb_path_base=perturb_path_base, @@ -46,13 +56,22 @@ def _export(tmp_path, test_mdata, target_gene, gene_loading_corr_matrix, class TestNoPerturbPathBase: - def test_page_written_without_perturb_files( - self, tmp_path, test_mdata, target_gene, gene_loading_corr_matrix, available_samples + self, + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, ): """perturb_path_base=None still produces a page; no per-condition JSON.""" - gene_dir = _export(tmp_path, test_mdata, target_gene, - gene_loading_corr_matrix, available_samples) + gene_dir = _export( + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, + ) assert (gene_dir / f"gene_{target_gene}.html").is_file() # header panels still there assert (gene_dir / "data" / "top_programs.json").is_file() @@ -63,10 +82,20 @@ def test_page_written_without_perturb_files( assert not list((gene_dir / "data").glob("waterfall_*.json")) def test_metadata_reports_zero_significant( - self, tmp_path, test_mdata, target_gene, gene_loading_corr_matrix, available_samples + self, + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, ): - gene_dir = _export(tmp_path, test_mdata, target_gene, - gene_loading_corr_matrix, available_samples) + gene_dir = _export( + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, + ) with open(gene_dir / "metadata.json") as f: meta = json.load(f) assert meta["n_significant_program_perturbations_total"] == 0 @@ -74,7 +103,6 @@ def test_metadata_reports_zero_significant( class TestRequiredArgs: - def test_missing_loading_corr_raises( self, tmp_path, test_mdata, target_gene, available_samples ): @@ -82,10 +110,21 @@ def test_missing_loading_corr_raises( _export(tmp_path, test_mdata, target_gene, None, available_samples) def test_perturb_base_without_corr_raises( - self, tmp_path, test_mdata, target_gene, gene_loading_corr_matrix, - available_samples, perturb_path_base + self, + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, + perturb_path_base, ): with pytest.raises(ValueError, match="perturb_corr_by_sample is required"): - _export(tmp_path, test_mdata, target_gene, gene_loading_corr_matrix, - available_samples, perturb_path_base=perturb_path_base, - perturb_corr_by_sample=None) + _export( + tmp_path, + test_mdata, + target_gene, + gene_loading_corr_matrix, + available_samples, + perturb_path_base=perturb_path_base, + perturb_corr_by_sample=None, + ) diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/program_html/conftest.py b/tests/Script/Stage3_Interpretation/A_Plotting/program_html/conftest.py index ada4a8b..0c27080 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/program_html/conftest.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/program_html/conftest.py @@ -19,16 +19,30 @@ import numpy as np import pandas as pd import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from pathlib import Path PIPELINE_ROOT = Path(__file__).resolve().parents[5] sys.path.insert(0, str(PIPELINE_ROOT / "src")) # Default to the batch run (matches the rest of the test suite). -DEFAULT_INFERENCE_PATH = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" -EVAL_5_2_0 = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Evaluation" / "5_2_0" -HTML_OUTPUT = PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Interpretation" / "Plotting" / "program_html" +DEFAULT_INFERENCE_PATH = ( + PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Inference" +) +EVAL_5_2_0 = ( + PIPELINE_ROOT / "tests" / "output" / "torch-cNMF" / "batch" / "Evaluation" / "5_2_0" +) +HTML_OUTPUT = ( + PIPELINE_ROOT + / "tests" + / "output" + / "torch-cNMF" + / "batch" + / "Interpretation" + / "Plotting" + / "program_html" +) # Sample names match the per-sample perturb TSVs produced by the U-test pipeline. @@ -51,6 +65,7 @@ def html_share_path(): def test_mdata(): """Real k=5 MuData from batch inference output.""" import mudata + h5mu_path = DEFAULT_INFERENCE_PATH / "adata" / "cNMF_5_2_0.h5mu" if not h5mu_path.exists(): pytest.skip(f"h5mu not found: {h5mu_path}") @@ -91,14 +106,20 @@ def perturbed_gene_found(test_mdata): @pytest.fixture(scope="session") def program_correlation(test_mdata): """Real program x program correlation matrix from cNMF loadings.""" - from Stage3_Interpretation.A_Plotting.src.Program_QC_plots import compute_program_correlation_matrix + from Stage3_Interpretation.A_Plotting.src.Program_QC_plots import ( + compute_program_correlation_matrix, + ) + return compute_program_correlation_matrix(test_mdata) @pytest.fixture(scope="session") def waterfall_correlation(perturb_path_base): """Real per-sample waterfall correlation matrices computed from perturb TSVs.""" - from Stage3_Interpretation.A_Plotting.src.Program_QC_plots import compute_program_waterfall_cor + from Stage3_Interpretation.A_Plotting.src.Program_QC_plots import ( + compute_program_waterfall_cor, + ) + out = {} for samp in SAMPLES: path = f"{perturb_path_base}_{samp}.txt" diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program.py b/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program.py index 639abf7..05d5214 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program.py @@ -17,6 +17,7 @@ import pytest import matplotlib + matplotlib.use("Agg") from Stage3_Interpretation.A_Plotting.src.html_Program_QC_plots import ( @@ -33,20 +34,30 @@ # Tests for individual JSON-builder helpers (cheap, no plot rendering) # --------------------------------------------------------------------------- -class TestBuilders: +class TestBuilders: def test_build_top_genes(self, test_mdata): target = str(test_mdata["cNMF"].var_names[0]) - d = _build_top_genes(test_mdata, target, num_gene=5, - file_to_dictionary=None, gene_name_key="symbol") + d = _build_top_genes( + test_mdata, + target, + num_gene=5, + file_to_dictionary=None, + gene_name_key="symbol", + ) assert set(d.keys()) == {"genes", "loadings"} assert len(d["genes"]) == 5 assert len(d["loadings"]) == 5 def test_build_go_terms(self, go_path, test_mdata): target = str(test_mdata["cNMF"].var_names[0]) - d = _build_go_terms(go_path, target, num_term=3, - p_value_name="Adjusted P-value", term_col="Term") + d = _build_go_terms( + go_path, + target, + num_term=3, + p_value_name="Adjusted P-value", + term_col="Term", + ) assert "terms" in d and "adj_pval" in d and "neglog10p" in d def test_build_correlations(self, program_correlation, test_mdata): @@ -67,12 +78,20 @@ def test_build_violin(self, test_mdata): # End-to-end: export_program_html writes the full subtree # --------------------------------------------------------------------------- -class TestExportProgramHTML: +class TestExportProgramHTML: @pytest.fixture(scope="class") - def exported_program(self, test_mdata, perturb_path_base, go_path, - program_correlation, waterfall_correlation, - perturbed_gene_found, html_share_path, available_samples): + def exported_program( + self, + test_mdata, + perturb_path_base, + go_path, + program_correlation, + waterfall_correlation, + perturbed_gene_found, + html_share_path, + available_samples, + ): target = str(test_mdata["cNMF"].var_names[0]) export_program_html( mdata=test_mdata, @@ -115,8 +134,13 @@ def test_per_panel_json_written(self, exported_program): prog_dir, _, samples = exported_program data_dir = prog_dir / "data" # Header panels - for fname in ["top_genes.json", "go_terms.json", "correlations.json", - "violin.json", "heatmap.json"]: + for fname in [ + "top_genes.json", + "go_terms.json", + "correlations.json", + "violin.json", + "heatmap.json", + ]: p = data_dir / fname assert p.is_file(), f"Missing per-panel JSON: {p}" with open(p) as f: @@ -145,13 +169,20 @@ def test_metadata_json(self, exported_program): # write_share_index # --------------------------------------------------------------------------- -class TestWriteShareIndex: +class TestWriteShareIndex: @pytest.fixture - def ensure_one_program_exported(self, test_mdata, perturb_path_base, go_path, - program_correlation, waterfall_correlation, - perturbed_gene_found, html_share_path, - available_samples): + def ensure_one_program_exported( + self, + test_mdata, + perturb_path_base, + go_path, + program_correlation, + waterfall_correlation, + perturbed_gene_found, + html_share_path, + available_samples, + ): target = str(test_mdata["cNMF"].var_names[0]) prog_dir = Path(html_share_path) / f"program_{target}" if not (prog_dir / f"program_{target}.html").exists(): @@ -175,7 +206,9 @@ def ensure_one_program_exported(self, test_mdata, perturb_path_base, go_path, ) return target - def test_share_index_written(self, ensure_one_program_exported, html_share_path, test_mdata): + def test_share_index_written( + self, ensure_one_program_exported, html_share_path, test_mdata + ): program_ids = [str(p) for p in test_mdata["cNMF"].var_names] write_share_index(html_share_path, program_ids, {"test": "config", "k": 5}) diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program_no_perturb.py b/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program_no_perturb.py index c6e801a..7eb8201 100644 --- a/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program_no_perturb.py +++ b/tests/Script/Stage3_Interpretation/A_Plotting/program_html/test_html_program_no_perturb.py @@ -17,9 +17,12 @@ import pytest import matplotlib + matplotlib.use("Agg") -from Stage3_Interpretation.A_Plotting.src.html_Program_QC_plots import export_program_html +from Stage3_Interpretation.A_Plotting.src.html_Program_QC_plots import ( + export_program_html, +) @pytest.fixture(scope="module") @@ -27,9 +30,17 @@ def target_program(test_mdata): return str(test_mdata["cNMF"].var_names[0]) -def _export(tmp_path, test_mdata, go_path, target, program_correlation, - perturbed_gene_found, available_samples, - perturb_path_base=None, waterfall_correlation=None): +def _export( + tmp_path, + test_mdata, + go_path, + target, + program_correlation, + perturbed_gene_found, + available_samples, + perturb_path_base=None, + waterfall_correlation=None, +): export_program_html( mdata=test_mdata, perturb_path_base=perturb_path_base, @@ -52,14 +63,26 @@ def _export(tmp_path, test_mdata, go_path, target, program_correlation, class TestNoPerturbPathBase: - def test_page_written_without_perturb_files( - self, tmp_path, test_mdata, go_path, target_program, program_correlation, - perturbed_gene_found, available_samples + self, + tmp_path, + test_mdata, + go_path, + target_program, + program_correlation, + perturbed_gene_found, + available_samples, ): """perturb_path_base=None still produces a page; no per-sample JSON, no heatmap.""" - prog_dir = _export(tmp_path, test_mdata, go_path, target_program, - program_correlation, perturbed_gene_found, available_samples) + prog_dir = _export( + tmp_path, + test_mdata, + go_path, + target_program, + program_correlation, + perturbed_gene_found, + available_samples, + ) assert (prog_dir / f"program_{target_program}.html").is_file() # header panels still there assert (prog_dir / "data" / "top_genes.json").is_file() @@ -72,31 +95,70 @@ def test_page_written_without_perturb_files( assert not (prog_dir / "data" / "heatmap.json").exists() def test_metadata_reports_zero_significant( - self, tmp_path, test_mdata, go_path, target_program, program_correlation, - perturbed_gene_found, available_samples + self, + tmp_path, + test_mdata, + go_path, + target_program, + program_correlation, + perturbed_gene_found, + available_samples, ): - prog_dir = _export(tmp_path, test_mdata, go_path, target_program, - program_correlation, perturbed_gene_found, available_samples) + prog_dir = _export( + tmp_path, + test_mdata, + go_path, + target_program, + program_correlation, + perturbed_gene_found, + available_samples, + ) with open(prog_dir / "metadata.json") as f: meta = json.load(f) assert meta["n_significant_regulators_total"] == 0 class TestRequiredArgs: - def test_missing_program_correlation_raises( - self, tmp_path, test_mdata, go_path, target_program, - perturbed_gene_found, available_samples + self, + tmp_path, + test_mdata, + go_path, + target_program, + perturbed_gene_found, + available_samples, ): with pytest.raises(ValueError, match="program_correlation is required"): - _export(tmp_path, test_mdata, go_path, target_program, None, - perturbed_gene_found, available_samples) + _export( + tmp_path, + test_mdata, + go_path, + target_program, + None, + perturbed_gene_found, + available_samples, + ) def test_perturb_base_without_waterfall_raises( - self, tmp_path, test_mdata, go_path, target_program, program_correlation, - perturbed_gene_found, available_samples, perturb_path_base + self, + tmp_path, + test_mdata, + go_path, + target_program, + program_correlation, + perturbed_gene_found, + available_samples, + perturb_path_base, ): with pytest.raises(ValueError, match="waterfall_correlation is required"): - _export(tmp_path, test_mdata, go_path, target_program, program_correlation, - perturbed_gene_found, available_samples, - perturb_path_base=perturb_path_base, waterfall_correlation=None) + _export( + tmp_path, + test_mdata, + go_path, + target_program, + program_correlation, + perturbed_gene_found, + available_samples, + perturb_path_base=perturb_path_base, + waterfall_correlation=None, + ) diff --git a/tests/Script/Stage3_Interpretation/B_Summarization/conftest.py b/tests/Script/Stage3_Interpretation/B_Summarization/conftest.py index 8fe8c2d..8f5661d 100644 --- a/tests/Script/Stage3_Interpretation/B_Summarization/conftest.py +++ b/tests/Script/Stage3_Interpretation/B_Summarization/conftest.py @@ -3,6 +3,7 @@ Provides synthetic MuData, enrichment DataFrames, and perturbation files so tests run without real data. """ + import os import numpy as np import pandas as pd @@ -11,7 +12,9 @@ from scipy import sparse TESTS_ROOT = Path(__file__).resolve().parent.parent.parent.parent -OUTPUT_DIR = TESTS_ROOT / "output" / "torch-cNMF" / "batch" / "Interpretation" / "Summarization" +OUTPUT_DIR = ( + TESTS_ROOT / "output" / "torch-cNMF" / "batch" / "Interpretation" / "Summarization" +) @pytest.fixture(scope="session") @@ -24,6 +27,7 @@ def output_dir(): # Synthetic MuData # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def synthetic_mdata(): """Build a minimal MuData with 50 cells, 20 genes, 3 programs, 3 guides.""" @@ -34,7 +38,9 @@ def synthetic_mdata(): rng = np.random.default_rng(42) # RNA modality - X_rna = sparse.random(n_cells, n_genes, density=0.3, format="csr", random_state=42).astype(np.float32) + X_rna = sparse.random( + n_cells, n_genes, density=0.3, format="csr", random_state=42 + ).astype(np.float32) gene_names = [f"Gene{i}" for i in range(n_genes)] cell_ids = [f"Cell{i}" for i in range(n_cells)] samples = rng.choice(["S1", "S2", "S3"], size=n_cells) @@ -75,19 +81,22 @@ def mdata_copy(synthetic_mdata): # Synthetic enrichment DataFrames # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def synthetic_go_df(): """Minimal GO enrichment DataFrame matching Compile_GO_sheet output.""" rows = [] for prog in range(3): for t in range(5): - rows.append({ - "program_name": str(prog), - "Term": f"GO:00{prog}{t}", - "Adjusted P-value": 0.001 * (t + 1), - "Genes": ";".join([f"Gene{j}" for j in range(7)]), - "Combined Score": 100 - t * 10, - }) + rows.append( + { + "program_name": str(prog), + "Term": f"GO:00{prog}{t}", + "Adjusted P-value": 0.001 * (t + 1), + "Genes": ";".join([f"Gene{j}" for j in range(7)]), + "Combined Score": 100 - t * 10, + } + ) df = pd.DataFrame(rows).set_index("Term") return df @@ -97,12 +106,14 @@ def synthetic_geneset_df(): rows = [] for prog in range(3): for t in range(3): - rows.append({ - "program_name": str(prog), - "Term": f"GENESET_{prog}_{t}", - "Adjusted P-value": 0.01 * (t + 1), - "Genes": ";".join([f"Gene{j}" for j in range(4)]), - }) + rows.append( + { + "program_name": str(prog), + "Term": f"GENESET_{prog}_{t}", + "Adjusted P-value": 0.01 * (t + 1), + "Genes": ";".join([f"Gene{j}" for j in range(4)]), + } + ) df = pd.DataFrame(rows).set_index("Term") return df @@ -112,16 +123,19 @@ def synthetic_explained_variance_df(): # Match the real `{K}_Explained_Variance.txt` format produced by Stage 2, # which has columns: VarianceExplained, ProgramID, program_name. # Compile_Summary_sheet's simple_Summary_cols looks up "VarianceExplained". - return pd.DataFrame({ - "program_name": [0, 1, 2], - "VarianceExplained": [0.15, 0.10, 0.08], - }).set_index("program_name") + return pd.DataFrame( + { + "program_name": [0, 1, 2], + "VarianceExplained": [0.15, 0.10, 0.08], + } + ).set_index("program_name") # --------------------------------------------------------------------------- # Synthetic perturbation files on disk # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def perturbation_dir(output_dir): """Write synthetic perturbation TSV files and return base path (without _Sample.txt).""" @@ -137,12 +151,14 @@ def perturbation_dir(output_dir): rows = [] for tgt in targets: for prog in programs: - rows.append({ - "target_name": tgt, - "program_name": prog, - "log2FC": rng.normal(0, 1), - "adj_pval": rng.uniform(0, 0.2), - }) + rows.append( + { + "target_name": tgt, + "program_name": prog, + "log2FC": rng.normal(0, 1), + "adj_pval": rng.uniform(0, 0.2), + } + ) df = pd.DataFrame(rows) df.to_csv(f"{base}_{samp}.txt", sep="\t", index=True) @@ -158,6 +174,7 @@ def sample_list(): # Synthetic program loading flat DataFrame # --------------------------------------------------------------------------- + @pytest.fixture(scope="session") def synthetic_program_loading_flat(synthetic_mdata): """Mimics output of compile_Program_loading_score_sheet_flat.""" diff --git a/tests/Script/Stage3_Interpretation/B_Summarization/test_summary_table.py b/tests/Script/Stage3_Interpretation/B_Summarization/test_summary_table.py index e5438ea..800ef2a 100644 --- a/tests/Script/Stage3_Interpretation/B_Summarization/test_summary_table.py +++ b/tests/Script/Stage3_Interpretation/B_Summarization/test_summary_table.py @@ -1,4 +1,5 @@ """Unit tests for Interpretation.Summary_table.src.Compile_excel_sheet.""" + import os import sys import numpy as np @@ -41,6 +42,7 @@ # Test: compile_Program_loading_score_sheet_flat # ========================================================================== + class TestProgramLoadingFlat: def test_returns_dataframe(self, synthetic_mdata): df = compile_Program_loading_score_sheet_flat(synthetic_mdata, num_gene=5) @@ -49,7 +51,9 @@ def test_returns_dataframe(self, synthetic_mdata): def test_shape(self, synthetic_mdata): n_programs = synthetic_mdata["cNMF"].n_vars num_gene = 5 - df = compile_Program_loading_score_sheet_flat(synthetic_mdata, num_gene=num_gene) + df = compile_Program_loading_score_sheet_flat( + synthetic_mdata, num_gene=num_gene + ) assert df.shape == (n_programs, num_gene) def test_index_name(self, synthetic_mdata): @@ -58,7 +62,9 @@ def test_index_name(self, synthetic_mdata): def test_columns_are_ranks(self, synthetic_mdata): num_gene = 10 - df = compile_Program_loading_score_sheet_flat(synthetic_mdata, num_gene=num_gene) + df = compile_Program_loading_score_sheet_flat( + synthetic_mdata, num_gene=num_gene + ) assert list(df.columns) == list(range(1, num_gene + 1)) def test_genes_are_from_var_names(self, synthetic_mdata): @@ -73,16 +79,19 @@ def test_genes_are_from_var_names(self, synthetic_mdata): # Test: Compile_GO_sheet / Compile_Geneset_sheet / Compile_Trait_sheet # ========================================================================== + class TestEnrichmentSheets: @pytest.fixture def go_file(self, output_dir): path = output_dir / "test_GO.txt" - df = pd.DataFrame({ - "Term": ["GO:001", "GO:002"], - "program_name": ["0", "1"], - "Adjusted P-value": [0.001, 0.05], - "Genes": ["A;B;C;D;E;F;G", "X;Y;Z"], - }) + df = pd.DataFrame( + { + "Term": ["GO:001", "GO:002"], + "program_name": ["0", "1"], + "Adjusted P-value": [0.001, 0.05], + "Genes": ["A;B;C;D;E;F;G", "X;Y;Z"], + } + ) df.to_csv(path, sep="\t") return str(path) @@ -109,6 +118,7 @@ def test_compile_trait_sheet(self, go_file): # Test: Compile_Perturbation_sheet # ========================================================================== + class TestPerturbationSheet: def test_concatenates_samples(self, perturbation_dir, sample_list): df = Compile_Perturbation_sheet(perturbation_dir, Sample=sample_list) @@ -131,15 +141,18 @@ def test_row_count(self, perturbation_dir, sample_list): # Test: Compile_Association_sheet # ========================================================================== + class TestAssociationSheet: @pytest.fixture def association_file(self, output_dir): path = output_dir / "test_association.txt" - df = pd.DataFrame({ - "program_name": [0, 1, 2], - "chi2": [10.5, 20.3, 5.1], - "pval": [0.001, 0.0001, 0.05], - }) + df = pd.DataFrame( + { + "program_name": [0, 1, 2], + "chi2": [10.5, 20.3, 5.1], + "pval": [0.001, 0.0001, 0.05], + } + ) df.to_csv(path, sep="\t") return str(path) @@ -153,14 +166,17 @@ def test_loads_and_reindexes(self, association_file): # Test: Compile_Explained_variance # ========================================================================== + class TestExplainedVariance: @pytest.fixture def variance_file(self, output_dir): path = output_dir / "test_variance.txt" - df = pd.DataFrame({ - "program_name": [0, 1, 2], - "variance_explained": [0.15, 0.10, 0.08], - }) + df = pd.DataFrame( + { + "program_name": [0, 1, 2], + "variance_explained": [0.15, 0.10, 0.08], + } + ) df.to_csv(path, sep="\t", index=False) return str(path) @@ -176,6 +192,7 @@ def test_loads_correctly(self, variance_file): # Test: Specificity score helpers # ========================================================================== + class TestSpecificityHelpers: @pytest.fixture def perturb_df(self): @@ -184,12 +201,14 @@ def perturb_df(self): rows = [] for tgt in ["GeneA", "GeneB", "GeneC"]: for prog in ["0", "1", "2"]: - rows.append({ - "target_name": tgt, - "program_name": prog, - "log2FC": rng.normal(0, 1), - "adj_pval": rng.uniform(0, 0.1), - }) + rows.append( + { + "target_name": tgt, + "program_name": prog, + "log2FC": rng.normal(0, 1), + "adj_pval": rng.uniform(0, 0.1), + } + ) return pd.DataFrame(rows) def test_compute_regulator_zscore_shape(self, perturb_df): @@ -240,6 +259,7 @@ def test_pmi_is_finite(self, perturb_df): # Test: get_specificity_program (end-to-end with files) # ========================================================================== + class TestGetSpecificityProgram: def test_returns_dataframe(self, perturbation_dir, sample_list): df = get_specificity_program(perturbation_dir, Sample=sample_list, T=0.5) @@ -256,6 +276,7 @@ def test_columns_per_sample(self, perturbation_dir, sample_list): # Test: get_significant_programs / get_significant_programs_df # ========================================================================== + class TestSignificantPrograms: def test_returns_dict(self, perturbation_dir, sample_list): result = get_significant_programs(perturbation_dir, Sample=sample_list) @@ -272,6 +293,7 @@ def test_df_has_count_columns(self, perturbation_dir, sample_list): # Test: MuData helper functions # ========================================================================== + class TestMuDataHelpers: def test_get_guide_cells_per_days(self, synthetic_mdata): df = get_guide_cells_per_days(synthetic_mdata, categorical_key="sample") @@ -302,6 +324,7 @@ def test_get_program_info_summary_cols(self, synthetic_mdata): # Test: get_top_terms_Summary_cols # ========================================================================== + class TestTopTermsSummaryCols: def test_returns_dataframe(self, synthetic_go_df, synthetic_geneset_df): df = get_top_terms_Summary_cols(synthetic_go_df, synthetic_geneset_df) @@ -318,6 +341,7 @@ def test_handles_none_inputs(self): # Test: check_program_name_match # ========================================================================== + class TestCheckProgramNameMatch: def test_no_warning_on_match(self, synthetic_mdata, capsys): df = pd.DataFrame({"program_name": ["0", "1", "2"], "value": [1, 2, 3]}) @@ -340,13 +364,21 @@ def test_skips_none(self, synthetic_mdata): # Test: Compile_Summary_sheet (integration) # ========================================================================== + class TestCompileSummarySheet: def test_returns_dataframe( - self, synthetic_mdata, synthetic_go_df, synthetic_geneset_df, - synthetic_program_loading_flat, synthetic_explained_variance_df, - perturbation_dir, sample_list, + self, + synthetic_mdata, + synthetic_go_df, + synthetic_geneset_df, + synthetic_program_loading_flat, + synthetic_explained_variance_df, + perturbation_dir, + sample_list, ): - df_perturbation = Compile_Perturbation_sheet(perturbation_dir, Sample=sample_list) + df_perturbation = Compile_Perturbation_sheet( + perturbation_dir, Sample=sample_list + ) df = Compile_Summary_sheet( mdata=synthetic_mdata, df_GO=synthetic_go_df, @@ -362,11 +394,18 @@ def test_returns_dataframe( assert len(df) == synthetic_mdata["cNMF"].n_vars def test_has_manual_columns( - self, synthetic_mdata, synthetic_go_df, synthetic_geneset_df, - synthetic_program_loading_flat, synthetic_explained_variance_df, - perturbation_dir, sample_list, + self, + synthetic_mdata, + synthetic_go_df, + synthetic_geneset_df, + synthetic_program_loading_flat, + synthetic_explained_variance_df, + perturbation_dir, + sample_list, ): - df_perturbation = Compile_Perturbation_sheet(perturbation_dir, Sample=sample_list) + df_perturbation = Compile_Perturbation_sheet( + perturbation_dir, Sample=sample_list + ) df = Compile_Summary_sheet( mdata=synthetic_mdata, df_GO=synthetic_go_df, @@ -377,14 +416,25 @@ def test_has_manual_columns( Sample=sample_list, categorical_key="sample", ) - for col in ["manual_annotation_label", "manual_timepoint", "Notes", "Automatic Timepoint"]: + for col in [ + "manual_annotation_label", + "manual_timepoint", + "Notes", + "Automatic Timepoint", + ]: assert col in df.columns def test_works_with_none_enrichment( - self, synthetic_mdata, synthetic_program_loading_flat, - synthetic_explained_variance_df, perturbation_dir, sample_list, + self, + synthetic_mdata, + synthetic_program_loading_flat, + synthetic_explained_variance_df, + perturbation_dir, + sample_list, ): - df_perturbation = Compile_Perturbation_sheet(perturbation_dir, Sample=sample_list) + df_perturbation = Compile_Perturbation_sheet( + perturbation_dir, Sample=sample_list + ) df = Compile_Summary_sheet( mdata=synthetic_mdata, df_GO=None, diff --git a/tests/Script/create_mini_dataset.py b/tests/Script/create_mini_dataset.py index c629a98..6287946 100644 --- a/tests/Script/create_mini_dataset.py +++ b/tests/Script/create_mini_dataset.py @@ -26,8 +26,9 @@ def main(): parser = argparse.ArgumentParser(description=__doc__.splitlines()[1]) - parser.add_argument("--full_data", required=True, - help="Path to the full counts .h5ad to subsample") + parser.add_argument( + "--full_data", required=True, help="Path to the full counts .h5ad to subsample" + ) args = parser.parse_args() full_data = args.full_data diff --git a/tools/check_no_lab_specific_content.py b/tools/check_no_lab_specific_content.py index 521db7e..0d1cd1b 100755 --- a/tools/check_no_lab_specific_content.py +++ b/tools/check_no_lab_specific_content.py @@ -47,11 +47,16 @@ r"@stanford\.edu", r"tony_method", ] -FORBIDDEN_RE = re.compile("|".join(f"(?:{p})" for p in FORBIDDEN_PATTERNS), re.IGNORECASE) +FORBIDDEN_RE = re.compile( + "|".join(f"(?:{p})" for p in FORBIDDEN_PATTERNS), re.IGNORECASE +) REPO_ROOT = Path( subprocess.run( - ["git", "rev-parse", "--show-toplevel"], capture_output=True, text=True, check=True + ["git", "rev-parse", "--show-toplevel"], + capture_output=True, + text=True, + check=True, ).stdout.strip() ) ALLOWLIST_PATH = REPO_ROOT / "tools" / "lab_specific_allowlist.txt" @@ -74,7 +79,9 @@ def load_allowlist(path: Path) -> tuple[list[str], list[str]]: elif kind == "text" and value: literals.append(value) else: - sys.exit(f"{path}: cannot parse allowlist line: {raw!r} (expected 'path: ' or 'text: ')") + sys.exit( + f"{path}: cannot parse allowlist line: {raw!r} (expected 'path: ' or 'text: ')" + ) return path_globs, literals @@ -118,11 +125,15 @@ def scan_tracked(path_globs: list[str], literals: list[str]) -> list[str]: continue for lineno, line in enumerate(lines, 1): for hit in find_hits(line, literals): - problems.append(f"{file_path}:{lineno}: matches '{hit}': {line.strip()[:160]}") + problems.append( + f"{file_path}:{lineno}: matches '{hit}': {line.strip()[:160]}" + ) return problems -def scan_diff(diff_args: list[str], path_globs: list[str], literals: list[str]) -> list[str]: +def scan_diff( + diff_args: list[str], path_globs: list[str], literals: list[str] +) -> list[str]: """Scan only added lines (and added/renamed file paths) in a diff.""" problems = [] diff = git("diff", "--no-color", "--unified=0", "--no-ext-diff", *diff_args) @@ -131,7 +142,13 @@ def scan_diff(diff_args: list[str], path_globs: list[str], literals: list[str]) for line in diff.splitlines(): if line.startswith("+++ "): target = line[4:] - current_file = None if target == "/dev/null" else target[2:] if target.startswith("b/") else target + current_file = ( + None + if target == "/dev/null" + else target[2:] + if target.startswith("b/") + else target + ) if current_file and not is_path_allowed(current_file, path_globs): for hit in find_hits(current_file, literals): problems.append(f"{current_file}: [path] matches '{hit}'") @@ -143,7 +160,9 @@ def scan_diff(diff_args: list[str], path_globs: list[str], literals: list[str]) if line.startswith("+") and current_file: if not is_path_allowed(current_file, path_globs): for hit in find_hits(line[1:], literals): - problems.append(f"{current_file}:{lineno}: matches '{hit}': {line[1:].strip()[:160]}") + problems.append( + f"{current_file}:{lineno}: matches '{hit}': {line[1:].strip()[:160]}" + ) lineno += 1 return problems @@ -151,8 +170,14 @@ def scan_diff(diff_args: list[str], path_globs: list[str], literals: list[str]) def main() -> int: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) mode = parser.add_mutually_exclusive_group() - mode.add_argument("--staged", action="store_true", help="Scan lines added in the index (for pre-commit)") - mode.add_argument("--diff", metavar="REF", help="Scan lines added since REF (e.g. origin/main)") + mode.add_argument( + "--staged", + action="store_true", + help="Scan lines added in the index (for pre-commit)", + ) + mode.add_argument( + "--diff", metavar="REF", help="Scan lines added since REF (e.g. origin/main)" + ) args = parser.parse_args() path_globs, literals = load_allowlist(ALLOWLIST_PATH) From 540cefca9268e94d4db36c49ce7f752558c35eda Mon Sep 17 00:00:00 2001 From: OlgaPushkarev Date: Thu, 8 Oct 2026 17:09:11 -0700 Subject: [PATCH 3/5] Ignore the ruff format commit in git blame Co-Authored-By: Claude Opus 5.5 (1M context) --- .git-blame-ignore-revs | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .git-blame-ignore-revs diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000..6aae21d --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,5 @@ +# Format-only commits, ignored by git blame (GitHub reads this file automatically). +# Locally: git config blame.ignoreRevsFile .git-blame-ignore-revs + +# Format all Python files with ruff +e77572acd6e6c5ff1e02e85572b62d03419eb5ac From 213c3240b943e70506e9a82578074e38285a6072 Mon Sep 17 00:00:00 2001 From: OlgaPushkarev Date: Fri, 9 Oct 2026 11:51:24 -0700 Subject: [PATCH 4/5] Support non-perturbed scRNA-seq and mouse data Non-perturbed data (no guide information): - K-selection: new --no_trait / --no_perturbation flags skip the trait and perturbation inputs and plots; the shared regulator row is factored into _plot_regulator_row and hidden when there is no perturbation data. - Program analysis: tolerate a missing uns['guide_targets']. - Excel summary: skip the Targets Summary sheet (and the specificity scores it writes) when there are no guides; fall back to var_names when the gene symbol column is absent. - Evaluation: fail early with a clear message when --Perform_perturbation is requested without guide information. Mouse data: - Mouse genes are matched to human gene-set libraries and gene summaries through curated orthologs (union of MGI and Ensembl, new orthologs.py) instead of upper-casing symbols. The table is downloaded once and cached in src/Stage2_Evaluation/Resources/; --ortholog_table pins a fixed copy. Genes without a human ortholog are excluded from enrichment; one-to-many orthologs keep all human genes. MSigDB mouse gene sets keep mouse symbols. - Excel summary: --organism mouse annotates genes with their human orthologs' MyGene summaries and a 'Human ortholog' column. - Evaluation: trait (GWAS) enrichment is human-only and is skipped with a message for other organisms. - Plotting utilities: honor the species argument of the MyGene helpers. Human runs are unchanged: existing Stage 2/3 tests give identical per-test outcomes before and after, and human evaluation outputs match apart from gene order within the Genes column and last-digit rounding. Tests: tests/Script/test_non_perturbed.py and tests/Script/Stage2_Evaluation/test_orthologs.py (synthetic data, offline). Docs: README section on non-perturbed and mouse data, parameter tables. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../references/02-evaluation.md | 3 +- .../references/04-visualization.md | 2 + .../references/05-annotation-summary.md | 1 + .../references/parameter-catalog.md | 7 +- README.md | 22 ++ src/Stage2_Evaluation/A_Metrics/README.md | 3 +- .../Slurm_Version/cNMF_evaluation_pipeline.py | 26 ++ .../A_Metrics/src/enrichment_geneset.py | 40 +- .../A_Metrics/src/orthologs.py | 189 +++++++++ .../A_Plotting/README.md | 2 + .../Slurm_Version/cNMF_k_selection.py | 77 +++- .../Slurm_Version/cNMF_program_analysis.py | 22 +- .../A_Plotting/src/k_selection_plots.py | 280 ++++++------- .../A_Plotting/src/utilities.py | 6 +- .../Slurm_Version/cNMF_excel_summary.py | 67 +++- .../src/Compile_excel_sheet.py | 80 +++- .../Stage2_Evaluation/test_orthologs.py | 157 ++++++++ tests/Script/test_non_perturbed.py | 368 ++++++++++++++++++ 18 files changed, 1116 insertions(+), 236 deletions(-) create mode 100644 src/Stage2_Evaluation/A_Metrics/src/orthologs.py create mode 100644 tests/Script/Stage2_Evaluation/test_orthologs.py create mode 100644 tests/Script/test_non_perturbed.py diff --git a/.claude/skills/perturbNMF-runner/references/02-evaluation.md b/.claude/skills/perturbNMF-runner/references/02-evaluation.md index 89f8ce8..c0eef7c 100644 --- a/.claude/skills/perturbNMF-runner/references/02-evaluation.md +++ b/.claude/skills/perturbNMF-runner/references/02-evaluation.md @@ -37,7 +37,8 @@ Reconstruction error and stability are computed automatically. | `--guide_annotation_path` | None | TSV with a `targeting` column to identify non-targeting controls (alternative to `--guide_annotation_key`) | | `--gene_names_key` | `symbol` | Column in data_guide["rna"].var with gene names | | `--FDR_method` | `StoreyQ` | FDR correction: `StoreyQ` or `BH` | -| `--organism` | `human` | Species for enrichment | +| `--organism` | `human` | Species for enrichment (`human` or `mouse`; mouse genes are mapped to human orthologs) | +| `--ortholog_table` | None | Mouse-human ortholog TSV; default MGI + Ensembl, cached in `Resources/` | | `--n_top` | `300` | Top genes for enrichment tests | | `--guide_annotation_key` | `["non-targeting"]` | Non-targeting guide identifiers (accepts multiple values) | | `--use_cache` | flag | Load enrichr gene set libraries from cached JSON in `Resources/` instead of downloading; falls back to download + cache on miss | diff --git a/.claude/skills/perturbNMF-runner/references/04-visualization.md b/.claude/skills/perturbNMF-runner/references/04-visualization.md index ddc50c0..b3dd193 100644 --- a/.claude/skills/perturbNMF-runner/references/04-visualization.md +++ b/.claude/skills/perturbNMF-runner/references/04-visualization.md @@ -29,6 +29,8 @@ | `--selected_k` | None | K value to highlight with a red dashed line | | `--stability_file` | None | Pre-computed stability/error TSV/NPZ (bypasses `cnmf.consensus()`; for torch-cNMF runs) | | `--run_program_dotplot` | off | Enable per-(K, sel_thresh) program dotplots (default off; requires the inference h5mu artifact per K/threshold) | +| `--no_trait` | off | Skip trait (GWAS) enrichment inputs and plots (non-human data) | +| `--no_perturbation` | off | Skip perturbation inputs and regulator plots (non-perturbed data, no guides) | Enrichment file patterns use `{k}` placeholder (and `{sample}` for perturbation). See `references/parameter-catalog.md` Section 4 for all optional params. diff --git a/.claude/skills/perturbNMF-runner/references/05-annotation-summary.md b/.claude/skills/perturbNMF-runner/references/05-annotation-summary.md index c8c0c18..d0634ca 100644 --- a/.claude/skills/perturbNMF-runner/references/05-annotation-summary.md +++ b/.claude/skills/perturbNMF-runner/references/05-annotation-summary.md @@ -120,6 +120,7 @@ other stage (no notebook required). Submit one job per K to cover multiple K val | `--perturbation_file_name` | `perturbation_association_results` | Perturbation file stem (e.g. `CRT`) | | `--effect_size` | `log2FC` | Effect-size column (e.g. `approx_log2FC`) | | `--gene_names_key` | `symbol` | var column with gene symbols | +| `--organism` | `human` | Species of the data; `mouse` annotates genes with their human orthologs' summaries | | `--non_targeting_key` | `non-targeting` | Negative-control target label(s) | See `references/parameter-catalog.md` (Section 12) for the full flag list, including diff --git a/.claude/skills/perturbNMF-runner/references/parameter-catalog.md b/.claude/skills/perturbNMF-runner/references/parameter-catalog.md index 8d5e58d..4d8b1b0 100644 --- a/.claude/skills/perturbNMF-runner/references/parameter-catalog.md +++ b/.claude/skills/perturbNMF-runner/references/parameter-catalog.md @@ -219,7 +219,8 @@ Same as sk-cNMF except: | `--gwas_data_path` | str | None | Path to GWAS data (required only for trait enrichment) | | `--X_normalized_path` | str | None | Normalized counts h5ad (needed for explained variance). Path pattern: `//Inference/cnmf_tmp/Inference.norm_counts.h5ad` (prefix is always `Inference`, not the run_name) | | `--guide_annotation_path` | str | None | Guide annotation TSV | -| `--organism` | str | `human` | Species for enrichment | +| `--organism` | str | `human` | Species for enrichment (`human` or `mouse`; mouse genes are mapped to human orthologs) | +| `--ortholog_table` | str | None | Mouse-human ortholog TSV; default MGI + Ensembl, cached in `Resources/` | | `--FDR_method` | str | `StoreyQ` | FDR method: `StoreyQ` or `BH` | | `--n_top` | int | `300` | Top genes for enrichment tests | | `--gene_names_key` | str | `symbol` | Column in data_guide["rna"].var with gene names | @@ -271,6 +272,8 @@ Same as sk-cNMF except: | `--variance_col` | str | `Total` | No | Variance column name | | `--stability_file` | str | None | No | Pre-computed stability/error file (TSV or NPZ) | | `--run_program_dotplot` | flag | False | No | Enable per-(K, sel_thresh) program dotplots; requires `{output_dir}/{run_name}/Inference/adata/cNMF_{K}_{thresh}.h5mu` to exist for each pair | +| `--no_trait` | flag | False | No | Skip trait (GWAS) enrichment files and plots (non-human data) | +| `--no_perturbation` | flag | False | No | Skip perturbation files and regulator plots (non-perturbed data, no guides) | --- @@ -616,6 +619,8 @@ cover multiple K values. | `--save_path` | str | `{out_dir}/{run_name}/Interpretation/Summary_table/{K}_{thresh}` | Output dir for the `.xlsx`, TSVs, and specificity/correlation/KD sidecars | | `--mdata_path` | str | `{out_dir}/{run_name}/Inference/adata/cNMF_{K}_{thresh}.h5mu` | Input MuData | | `--num_gene` | int | `300` | Top genes per program / per enriched term to keep | +| `--organism` | str | `human` | Species of the data; with `mouse`, each gene gets the MyGene summaries of its human orthologs (listed in a `Human ortholog` column) | +| `--ortholog_table` | str | None | Mouse-human ortholog TSV; default MGI + Ensembl, cached in `Resources/` | | `--perturbation_file_name` | str | `perturbation_association_results` | Perturbation result file stem (between `{K}_` and `_{Condition}.txt`); e.g. `CRT` | | `--Conditions` | str (nargs=\*) | all labels found in the data | Condition labels (values of obs[categorical_key]); `--Sample` is a deprecated alias | | `--effect_size` | str | `log2FC` | Effect-size column in perturbation files (e.g. `approx_log2FC`) | diff --git a/README.md b/README.md index 8491c7a..67c49a6 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,28 @@ See [`src/Stage2_Evaluation/B_Calibration/README.md`](src/Stage2_Evaluation/B_Ca See [`src/Stage3_Interpretation/README.md`](src/Stage3_Interpretation/README.md) for detailed parameters and output format. +## Non-perturbed and mouse data + +The pipeline also runs on scRNA-seq data without perturbations (no guide information in the +input) and on mouse data. Per stage: + +| Stage | Non-perturbed data (no guides) | Mouse data | +|---|---|---| +| 1. Inference | Leave out the `--guide_*` options | `--species mouse` | +| 2a. Metrics | Leave out `--Perform_perturbation` and `--guide_annotation_path` | `--organism mouse`: genes are mapped to human orthologs for gene set and GO enrichment. Trait (GWAS) enrichment is human-only and is skipped with a message if requested | +| 2b. Calibration | Not applicable (needs non-targeting guides); skip | No change | +| 3a. K-selection plots | `--no_perturbation` | `--no_trait` | +| 3a. Program plots | Leave out `--perturb_path_base`; perturbed-gene annotations are left empty | No change | +| 3a. Perturbed-gene plots | Not applicable; skip | No change | +| 3b. Excel summary | The Targets Summary sheet is skipped automatically | `--organism mouse`: genes are annotated with their human orthologs' summaries | +| 3c. Annotation | No change | `--taxid 10090` (NCBI, UniProt, OmniPath search), `--species mouse` (MyGene search), `--organism mouse` (program extraction) | + +Mouse genes are matched to human gene-set libraries and gene summaries through curated +orthologs (union of MGI and Ensembl), not by upper-casing symbols. The table is downloaded +on first use and cached in `src/Stage2_Evaluation/Resources/orthologs_mouse_human.tsv`; +pass `--ortholog_table` to use a fixed copy. Mouse genes without a human ortholog are left +out of enrichment, and a gene with several human orthologs contributes to all of them. + ## Claude Code Skills This repo ships four Claude Code skills under `.claude/skills/` for guided pipeline execution, validation, and maintenance. Open the linked `SKILL.md` for trigger phrases, detailed usage, and helper script references. diff --git a/src/Stage2_Evaluation/A_Metrics/README.md b/src/Stage2_Evaluation/A_Metrics/README.md index 301f502..61927cf 100644 --- a/src/Stage2_Evaluation/A_Metrics/README.md +++ b/src/Stage2_Evaluation/A_Metrics/README.md @@ -40,7 +40,8 @@ The evaluation pipeline tests cNMF programs against various criteria to assess t | X_normalized_path | str | None | Path to normalized cell x gene matrix (.h5ad), required for explained variance | | guide_annotation_path | str | None | Path to tab-separated guide annotation file with "targeting" column | | gwas_data_path | str | None | Path to GWAS data file for trait enrichment (required when --Perform_trait is set) | -| organism | str | "human" | Organism/species for enrichment analysis | +| organism | str | "human" | Organism/species for enrichment analysis (`human` or `mouse`; mouse genes are mapped to human orthologs to match the human gene-set libraries) | +| ortholog_table | str | None | For mouse: mouse-human ortholog TSV (columns `mouse_symbol`, `mouse_ensembl`, `human_symbol`, `source`). Default: MGI + Ensembl table, downloaded once and cached in `src/Stage2_Evaluation/Resources/` | | FDR_method | str | "StoreyQ" | FDR correction method for perturbation association | | n_top | int | 300 | Number of top loaded genes to use for enrichment tests | | guide_annotation_key | list of str | ["non-targeting"] | Name(s) of non-targeting guide target labels | diff --git a/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py b/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py index d41296c..6c8443a 100644 --- a/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py +++ b/src/Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py @@ -207,6 +207,15 @@ def main(): type=str, default="human", ) + parser.add_argument( + "--ortholog_table", + type=str, + default=None, + help="For --organism mouse: TSV of mouse-human orthologs used to match genes " + "to the human gene-set libraries (columns: mouse_symbol, mouse_ensembl, " + "human_symbol, source). Default: MGI + Ensembl table, downloaded once and " + "cached in src/Stage2_Evaluation/Resources/", + ) parser.add_argument( "--FDR_method", help="Method for FDR correction in perturbation association (default: StoreyQ)", @@ -235,6 +244,15 @@ def main(): args = parser.parse_args() + # The bundled GWAS resource (OpenTargets L2G) is human-only. + if args.Perform_trait and args.organism.lower() != "human": + print( + f"--organism {args.organism}: skipping trait (GWAS) enrichment, which is " + "only available for human data. Pass --no_trait to K-selection plotting.", + flush=True, + ) + args.Perform_trait = False + if args.Perform_trait and args.gwas_data_path is None: parser.error("--gwas_data_path is required when --Perform_trait is set") @@ -387,6 +405,12 @@ def main(): # Run perturbation assocation if args.Perform_perturbation: + if args.guide_targets_key not in mdata[args.prog_key].uns: + raise ValueError( + f"--Perform_perturbation needs guide information, but uns['" + f"{args.guide_targets_key}'] is missing from '{args.prog_key}' " + f"(K={k}). For non-perturbed data, drop --Perform_perturbation." + ) # Validate reference_targets against mdata guide_targets mdata_targets = set(mdata[args.prog_key].uns[args.guide_targets_key]) matched_ref = mdata_targets.intersection(reference_targets) @@ -441,6 +465,7 @@ def main(): data_key=args.data_key, prog_name=None, organism=args.organism, + ortholog_table=args.ortholog_table, library="Reactome_2022", method="fisher", database="enrichr", @@ -461,6 +486,7 @@ def main(): data_key=args.data_key, prog_name=None, organism=args.organism, + ortholog_table=args.ortholog_table, library="GO_Biological_Process_2023", method="fisher", database="enrichr", diff --git a/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py b/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py index 209dc5c..e63a66a 100644 --- a/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py +++ b/src/Stage2_Evaluation/A_Metrics/src/enrichment_geneset.py @@ -12,6 +12,8 @@ from gseapy import Biomart, Msigdb from tqdm.auto import tqdm +from .orthologs import collapse_to_human, load_ortholog_table + logging.basicConfig(level=logging.INFO) _RESOURCES_DIR = os.path.join( @@ -78,6 +80,18 @@ def get_idconversion(var_names, organism="human", chunk_size=100): return gene_names +def _mouse_gene_ids(mdata, prog_key, data_key, gene_names_key): + """Mouse symbols or Ensembl IDs for the loadings columns, unchanged in case.""" + if gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: + genes = mdata[data_key].var[gene_names_key].astype(str).tolist() + elif "var_names" in mdata[prog_key].uns.keys(): + genes = [str(v) for v in mdata[prog_key].uns["var_names"]] + else: + genes = [str(v) for v in mdata[data_key].var_names] + # "Symbol:ENSMUSG..." style names: keep the symbol + return [g.split(":", 1)[0] if ":ens" in g.lower() else g for g in genes] + + def get_program_gene_loadings( mdata, prog_key="prog", @@ -85,10 +99,22 @@ def get_program_gene_loadings( data_key="rna", organism="human", gene_names_key=None, + ortholog_table=None, + map_to_human=True, ): - """Get gene loadings for each program in the mudata object.""" + """Get gene loadings for each program in the mudata object. - if gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: + For mouse data with map_to_human=True (gene sets in human symbols), genes are + mapped to human orthologs (see orthologs.py); ortholog_table is an optional + path to a pre-built ortholog table. With map_to_human=False, mouse symbols are + kept unchanged (gene sets in mouse symbols, e.g. MSigDB mouse). + """ + if organism not in ("human", "mouse"): + raise ValueError(f"Unsupported organism '{organism}'; use 'human' or 'mouse'.") + + if organism == "mouse": + gene_names = _mouse_gene_ids(mdata, prog_key, data_key, gene_names_key) + elif gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: gene_names = ( mdata[data_key].var[gene_names_key].astype(str).str.upper().tolist() ) @@ -113,6 +139,9 @@ def get_program_gene_loadings( loadings["gene_names"] = gene_names loadings.set_index("gene_names", inplace=True) + if organism == "mouse" and map_to_human: + loadings = collapse_to_human(loadings, load_ortholog_table(ortholog_table)) + # with open('var_names.txt', 'w') as fil: # for nam in mdata[data_key].var_names: # fil.write(nam+'\n') @@ -403,6 +432,7 @@ def compute_geneset_enrichment( use_loadings_gene=False, # False = use all oncology gene, True = use interact expresseed gene gene_names_key: Optional[str] = None, use_cache: bool = False, + ortholog_table: Optional[str] = None, **kwargs, ) -> Optional[pd.DataFrame]: """ @@ -442,6 +472,9 @@ def compute_geneset_enrichment( number of threads to run processes on. inplace : bool (default: True) whether to insert the results back into the mudata object. + ortholog_table : str (default: None) + For organism='mouse': path to a mouse-human ortholog table. By default the + MGI + Ensembl table is downloaded once and cached in Resources/. Returns ------- @@ -484,6 +517,9 @@ def compute_geneset_enrichment( data_key=data_key, organism=organism, gene_names_key=gene_names_key, + ortholog_table=ortholog_table, + # Enrichr and user gene sets use human symbols; MSigDB mouse sets do not. + map_to_human=not (database == "msigdb" and user_geneset is None), ) # run enrichment diff --git a/src/Stage2_Evaluation/A_Metrics/src/orthologs.py b/src/Stage2_Evaluation/A_Metrics/src/orthologs.py new file mode 100644 index 0000000..65e83ba --- /dev/null +++ b/src/Stage2_Evaluation/A_Metrics/src/orthologs.py @@ -0,0 +1,189 @@ +"""Mouse to human ortholog mapping. + +Human gene-set libraries (Enrichr, MSigDB human) and human gene summaries are matched +to mouse genes through curated orthologs rather than by upper-casing mouse symbols. +Upper-casing only works when the two species share a symbol, misses orthologs with +different names (e.g. Zfp451 -> ZNF451) and can pick an unrelated human gene. + +The ortholog table is the union of two sources, since each has gaps the other fills: +- MGI HOM_MouseHumanSequence.rpt (mouse-human homology classes) +- Ensembl BioMart mouse -> human homologs + +The combined table is downloaded on first use and cached as a TSV in +src/Stage2_Evaluation/Resources/, so later runs are offline and reproducible. A +pre-built table can also be passed explicitly (columns: mouse_symbol, mouse_ensembl, +human_symbol, source). +""" + +from __future__ import annotations + +import io +import logging +import os +import urllib.parse +import urllib.request +from collections.abc import Iterable + +import pandas as pd + +logger = logging.getLogger(__name__) + +RESOURCES_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), + "Resources", +) +DEFAULT_TABLE_PATH = os.path.join(RESOURCES_DIR, "orthologs_mouse_human.tsv") +TABLE_COLUMNS = ["mouse_symbol", "mouse_ensembl", "human_symbol", "source"] + +MGI_URL = "https://www.informatics.jax.org/downloads/reports/HOM_MouseHumanSequence.rpt" +# www.ensembl.org redirects BioMart requests; the mirrors answer directly. +ENSEMBL_MARTS = [ + "https://useast.ensembl.org/biomart/martservice", + "https://asia.ensembl.org/biomart/martservice", + "https://www.ensembl.org/biomart/martservice", +] +ENSEMBL_QUERY = ( + '' + '' + '' + '' + '' + '' + "" +) +DOWNLOAD_TIMEOUT_S = 600 +MOUSE_TAXON = "10090" +HUMAN_TAXON = "9606" + + +def _download(url: str) -> str: + with urllib.request.urlopen(url, timeout=DOWNLOAD_TIMEOUT_S) as response: + return response.read().decode("utf-8") + + +def _fetch_mgi() -> pd.DataFrame: + hom = pd.read_csv(io.StringIO(_download(MGI_URL)), sep="\t", dtype=str) + key, taxon = "DB Class Key", "NCBI Taxon ID" + mouse = hom.loc[hom[taxon] == MOUSE_TAXON, [key, "Symbol"]] + human = hom.loc[hom[taxon] == HUMAN_TAXON, [key, "Symbol"]] + pairs = mouse.merge(human, on=key, suffixes=("_mouse", "_human")) + if pairs.empty: + raise RuntimeError(f"No mouse-human pairs parsed from {MGI_URL}") + return pd.DataFrame( + { + "mouse_symbol": pairs["Symbol_mouse"], + "mouse_ensembl": pd.NA, + "human_symbol": pairs["Symbol_human"], + "source": "MGI", + } + ) + + +def _fetch_ensembl() -> pd.DataFrame: + query = urllib.parse.urlencode({"query": ENSEMBL_QUERY}) + errors = [] + for mart in ENSEMBL_MARTS: + try: + text = _download(f"{mart}?{query}") + df = pd.read_csv(io.StringIO(text), sep="\t", dtype=str) + if df.shape[1] != 3 or df.empty: + raise RuntimeError(f"unexpected response: {text[:200]!r}") + break + except (OSError, ValueError, RuntimeError) as err: # try the next mirror + errors.append(f"{mart}: {err}") + else: + raise RuntimeError("Ensembl BioMart query failed:\n" + "\n".join(errors)) + df.columns = ["mouse_ensembl", "mouse_symbol", "human_symbol"] + # Rows without a human ortholog are kept: they map Ensembl IDs to mouse symbols. + df["source"] = "Ensembl" + return df[TABLE_COLUMNS] + + +def build_ortholog_table() -> pd.DataFrame: + """Download MGI and Ensembl orthologs and return their union.""" + table = pd.concat([_fetch_mgi(), _fetch_ensembl()], ignore_index=True) + table = table.dropna(subset=["mouse_symbol"]).drop_duplicates() + return table[TABLE_COLUMNS].reset_index(drop=True) + + +def load_ortholog_table(path: str | None = None) -> pd.DataFrame: + """Load the ortholog table from `path`, the cache, or download and cache it.""" + if path is not None: + if not os.path.isfile(path): + raise FileNotFoundError(f"Ortholog table not found: {path}") + table = pd.read_csv(path, sep="\t", dtype=str) + elif os.path.isfile(DEFAULT_TABLE_PATH): + logger.info(f"Loading cached ortholog table from {DEFAULT_TABLE_PATH}") + table = pd.read_csv(DEFAULT_TABLE_PATH, sep="\t", dtype=str) + else: + logger.info("Downloading mouse-human ortholog table (MGI + Ensembl)") + table = build_ortholog_table() + os.makedirs(RESOURCES_DIR, exist_ok=True) + table.to_csv(DEFAULT_TABLE_PATH, sep="\t", index=False) + logger.info(f"Saved ortholog table to {DEFAULT_TABLE_PATH}") + + missing = set(TABLE_COLUMNS) - set(table.columns) + if missing: + raise ValueError(f"Ortholog table is missing columns: {sorted(missing)}") + return table + + +def map_mouse_to_human( + genes: Iterable[str], table: pd.DataFrame | None = None +) -> dict[str, list[str]]: + """Map mouse gene symbols or Ensembl IDs (ENSMUSG...) to human ortholog symbols. + + Returns a dict from each input gene to a sorted list of human symbols; genes with + no human ortholog map to an empty list. A gene with several human orthologs keeps + all of them. + """ + if table is None: + table = load_ortholog_table() + pairs = table.dropna(subset=["human_symbol"]) + symbol_to_human = pairs.groupby("mouse_symbol")["human_symbol"].apply(set) + id_to_symbol = ( + table.dropna(subset=["mouse_ensembl"]) + .groupby("mouse_ensembl")["mouse_symbol"] + .apply(set) + ) + + mapping = {} + for gene in genes: + gene = str(gene) + if gene.upper().startswith("ENSMUSG"): + symbols = id_to_symbol.get(gene.split(".", 1)[0], set()) # drop version + else: + symbols = {gene} + human = set() + for symbol in symbols: + human |= symbol_to_human.get(symbol, set()) + mapping[gene] = sorted(human) + return mapping + + +def collapse_to_human( + df: pd.DataFrame, table: pd.DataFrame | None = None +) -> pd.DataFrame: + """Re-index a (mouse genes x columns) DataFrame by human ortholog symbols. + + Each human ortholog takes the values of its mouse gene; genes without a human + ortholog are dropped, and a human gene with several mouse orthologs keeps the + per-column maximum. + """ + df = df.groupby(level=0).max() # one row per mouse gene + mapping = map_mouse_to_human(df.index, table) + rows = [(mouse, human) for mouse, hs in mapping.items() for human in hs] + n_unmapped = sum(1 for hs in mapping.values() if not hs) + logger.info( + f"Ortholog mapping: {len(mapping) - n_unmapped}/{len(mapping)} mouse genes " + f"have a human ortholog; the rest are excluded." + ) + if not rows: + raise ValueError( + "No mouse genes could be mapped to human orthologs. Check that gene " + "names are mouse symbols or Ensembl IDs (e.g. set gene_names_key)." + ) + mouse_idx, human_idx = zip(*rows) + expanded = df.loc[list(mouse_idx)].copy() + expanded.index = pd.Index(human_idx, name=df.index.name) + return expanded.groupby(level=0).max() diff --git a/src/Stage3_Interpretation/A_Plotting/README.md b/src/Stage3_Interpretation/A_Plotting/README.md index 7ce17ce..4c60668 100644 --- a/src/Stage3_Interpretation/A_Plotting/README.md +++ b/src/Stage3_Interpretation/A_Plotting/README.md @@ -64,6 +64,8 @@ Generates a multi-panel figure to guide K selection. Reads pre-computed Stage 2 | `--variance_col` | str | `Total` | Column name for variance values. Use `Total` for summary files, or a per-program column name (will be summed) | | `--stability_file` | str | `None` | Path to a pre-computed stability/error file (TSV or NPZ). Bypasses `cnmf.consensus()` — useful for torch-cNMF runs where the cnmf package isn't installed | | `--run_program_dotplot` | flag | off | If set, also generate per-(K, sel_thresh) program dotplots. Requires the inference `cNMF_{K}_{thresh}.h5mu` to exist for each pair; off by default because the loop is slow | +| `--no_trait` | flag | off | Do not require or read trait (GWAS) enrichment files and drop the trait plots. Use when evaluation ran without trait enrichment (e.g. mouse data) | +| `--no_perturbation` | flag | off | Do not require or read perturbation association files and drop the regulator plots. Use for non-perturbed data (no guides); `--Conditions` is then not needed | ### Outputs diff --git a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py index dc2392a..00b1cf1 100644 --- a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py +++ b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py @@ -20,6 +20,7 @@ plot_explained_variance, programs_dotplots, plot_k_selection_panel, + plot_k_selection_panel_no_traits, read_condition_labels_from_h5mu, ) @@ -142,6 +143,20 @@ def main(): help="If set, also generate per-(K, sel_thresh) program dotplots. Requires the inference " "h5mu (cNMF_{K}_{thresh}.h5mu) to exist for each (K, sel_thresh). Default off.", ) + parser.add_argument( + "--no_trait", + action="store_true", + help="Skip trait (GWAS) enrichment: do not require or read the trait files and " + "drop the trait plots. Use when evaluation ran without --Perform_trait " + "(e.g. mouse data). Default off.", + ) + parser.add_argument( + "--no_perturbation", + action="store_true", + help="Skip perturbation results: do not require or read the perturbation " + "association files and drop the regulator plots. Use for datasets without " + "guides (non-perturbed scRNA-seq). Default off.", + ) args = parser.parse_args() @@ -154,7 +169,7 @@ def main(): yaml.dump(args_dict, f, default_flow_style=False, width=1000) # ---- Condition labels: default to the values of --groupby in the inference h5mu ---- - if args.Conditions is None: + if args.Conditions is None and not args.no_perturbation: first_h5mu = ( f"{args.output_directory}/{args.run_name}/Inference/adata/" f"cNMF_{args.K[0]}_{str(args.sel_threshs[0]).replace('.', '_')}.h5mu" @@ -175,15 +190,21 @@ def main(): ) _ev_pat = args.variance_file or "{k}_Explained_Variance_Summary.txt" + required_pats = [_go_pat, _gs_pat, _ev_pat] + if not args.no_trait: + required_pats.append(_tr_pat) + missing = [] for sel_thresh in args.sel_threshs: sel_str = str(sel_thresh).replace(".", "_") for k in args.K: k_folder = os.path.join(args.eval_folder_name, f"{k}_{sel_str}") - for pat in [_go_pat, _gs_pat, _tr_pat, _ev_pat]: + for pat in required_pats: fpath = os.path.join(k_folder, pat.format(k=k)) if not os.path.isfile(fpath): missing.append(fpath) + if args.no_perturbation: + continue for cond in args.Conditions: fpath = os.path.join(k_folder, _pt_pat.format(k=k, sample=cond)) if not os.path.isfile(fpath): @@ -218,31 +239,40 @@ def main(): trait_file=args.trait_file, term_col=args.term_col, adjpval_col=args.adjpval_col, + include_traits=not args.no_trait, ) plot_enrichment( count_df, folder_name=args.save_folder_name, file_name=f"Enrichment_{sel_thresh}", + include_traits=not args.no_trait, ) # Perturbation - test_stats_df = load_perturbation_data( - folder=args.eval_folder_name, - components=args.K, - sel_thresh=sel_thresh, - conditions=args.Conditions, - pval=args.pval, - perturbation_file=args.perturbation_file, - perturb_adjpval_col=args.perturb_adjpval_col, - perturb_target_col=args.perturb_target_col, - perturb_log2fc_col=args.perturb_log2fc_col, - ) - plot_perturbation( - test_stats_df, - folder_name=args.save_folder_name, - pval=args.pval, - file_name=f"Perturbation_{sel_thresh}", - ) + if args.no_perturbation: + print( + "--no_perturbation set: skipping perturbation results and regulator plots.", + flush=True, + ) + test_stats_df = None + else: + test_stats_df = load_perturbation_data( + folder=args.eval_folder_name, + components=args.K, + sel_thresh=sel_thresh, + conditions=args.Conditions, + pval=args.pval, + perturbation_file=args.perturbation_file, + perturb_adjpval_col=args.perturb_adjpval_col, + perturb_target_col=args.perturb_target_col, + perturb_log2fc_col=args.perturb_log2fc_col, + ) + plot_perturbation( + test_stats_df, + folder_name=args.save_folder_name, + pval=args.pval, + file_name=f"Perturbation_{sel_thresh}", + ) # Explained Variance stats_EV = load_explained_variance_data( @@ -260,8 +290,13 @@ def main(): # Motif (working in progress) - # K-selection panel - plot_k_selection_panel( + # K-selection panel (perturbation row is hidden when test_stats_df is None) + panel_fn = ( + plot_k_selection_panel_no_traits + if args.no_trait + else plot_k_selection_panel + ) + panel_fn( stats_SE, count_df, test_stats_df, diff --git a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py index 5324231..4866e16 100644 --- a/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py +++ b/src/Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_program_analysis.py @@ -235,12 +235,24 @@ def main(): print(f"there are {program_len} Program found") # found detected perturbed gene (use gene symbols from var column when var_names are Ensembl IDs) - perturbed_gene = np.unique(mdata[args.prog_key].uns["guide_targets"]) - if args.gene_name_key in mdata[args.data_key].var.columns: - gene_symbols = mdata[args.data_key].var[args.gene_name_key].astype(str).tolist() + # uns['guide_targets'] only exists for perturbation data (written by inference + # when guides are present); without it there are no perturbed genes to mark. + if "guide_targets" in mdata[args.prog_key].uns: + perturbed_gene = np.unique(mdata[args.prog_key].uns["guide_targets"]) + if args.gene_name_key in mdata[args.data_key].var.columns: + gene_symbols = ( + mdata[args.data_key].var[args.gene_name_key].astype(str).tolist() + ) + else: + gene_symbols = mdata[args.data_key].var_names.tolist() + perturbed_gene_found = sorted(set(gene_symbols) & set(perturbed_gene.tolist())) else: - gene_symbols = mdata[args.data_key].var_names.tolist() - perturbed_gene_found = sorted(set(gene_symbols) & set(perturbed_gene.tolist())) + perturbed_gene_found = [] + print( + f"No uns['guide_targets'] in '{args.prog_key}': treating the data as " + "non-perturbed; perturbed-gene annotations will be empty.", + flush=True, + ) # compute correlations (both K x K; full matrices written to corr_dir every run) corr_dir = args.corr_matrix_path or args.save_path diff --git a/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py b/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py index 1b31c61..a235d22 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py +++ b/src/Stage3_Interpretation/A_Plotting/src/k_selection_plots.py @@ -413,7 +413,13 @@ def load_enrichment_data( trait_file=None, term_col="Term", adjpval_col="Adjusted P-value", + include_traits=True, ): + """Count unique significant terms per K for GO, gene set and trait enrichment. + + Set include_traits=False when trait enrichment was not run (non-human data); + the trait files are then not read and the 'traits' column is filled with 0. + """ # loading function def load(k, term, folder, term_col, adjpval_col): @@ -465,15 +471,16 @@ def load(k, term, folder, term_col, adjpval_col): adjpval_col, ) ) - term_df.append( - load( - k, - "traits", - "{}/{}".format(k_folder, trait_file.format(k=k)), - term_col, - adjpval_col, + if include_traits: + term_df.append( + load( + k, + "traits", + "{}/{}".format(k_folder, trait_file.format(k=k)), + term_col, + adjpval_col, + ) ) - ) term_df = pd.concat(term_df, ignore_index=True) @@ -513,7 +520,9 @@ def load(k, term, folder, term_col, adjpval_col): # plot loaded df — 3 independent square figures -def plot_enrichment(count_df, folder_name=None, file_name=None, selected_k=None): +def plot_enrichment( + count_df, folder_name=None, file_name=None, selected_k=None, include_traits=True +): enrichment_config = [ ("go_terms", "Unique GO terms", "GO term enrichment", _COLORS["go_terms"], "o"), ( @@ -526,6 +535,9 @@ def plot_enrichment(count_df, folder_name=None, file_name=None, selected_k=None) ("traits", "Unique traits", "Trait enrichment", _COLORS["traits"], "D"), ] + if not include_traits: + enrichment_config = [c for c in enrichment_config if c[0] != "traits"] + for col, ylabel, title, color, marker in enrichment_config: fig, ax = plt.subplots(figsize=(4, 3.5)) @@ -800,6 +812,88 @@ def plot_explained_variance(stats, folder_name=None, file_name=None, selected_k= plt.show() +def _plot_regulator_row(ax_all, ax_per_condition, test_stats_df, pval, _add_vline): + """Plot unique significant regulators per K, pooled and per condition.""" + # --- Row 3, Col 0: Regulators (all conditions) --- + ax = ax_all + plotting_df = ( + test_stats_df.loc[test_stats_df.adj_pval <= pval, ["K", "target_name"]] + .drop_duplicates() + .groupby(["K"]) + .count() + .reset_index() + ) + ax.plot( + plotting_df["K"], + plotting_df["target_name"], + color=_COLORS["primary"], + linewidth=1.5, + marker="o", + markersize=4, + markerfacecolor=_COLORS["perturbation"], + markeredgecolor=_COLORS["primary"], + markeredgewidth=0.8, + zorder=3, + ) + ax.fill_between( + plotting_df["K"], + plotting_df["target_name"], + alpha=0.08, + color=_COLORS["perturbation"], + ) + _add_vline(ax) + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title="Unique regulators for all conditions", + ) + ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) + + # --- Row 3, Col 1: Regulators (per sample) --- + ax = ax_per_condition + plotting_df_sample = ( + test_stats_df.loc[ + test_stats_df.adj_pval <= pval, ["K", "condition", "target_name"] + ] + .drop_duplicates() + .groupby(["K", "condition"]) + .count() + .reset_index() + ) + sns.lineplot( + x="K", + y="target_name", + hue="condition", + data=plotting_df_sample, + palette=_PALETTE[: plotting_df_sample["condition"].nunique()], + linewidth=1.5, + marker="o", + markersize=4, + ax=ax, + legend="brief", + ) + _add_vline(ax) + _style_ax( + ax, + xlabel="Number of components (k)", + ylabel="No. unique regulators", + title="Unique regulators per condition", + ) + ax.legend( + title="Condition", + fontsize=7, + title_fontsize=8, + frameon=True, + fancybox=False, + edgecolor="#cccccc", + bbox_to_anchor=(1.02, 1), + loc="upper left", + borderaxespad=0, + ) + ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) + + # Combined panel figure: 3 rows x 3 columns # Row 1: Stability, Error, Explained variance # Row 2: GO terms, Gene sets, Traits @@ -945,84 +1039,12 @@ def _add_vline(ax): _style_ax(ax, xlabel="Number of components (k)", ylabel=ylabel, title=title) ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) - # --- Row 3, Col 0: Regulators (all conditions) --- - ax = axes[2, 0] - plotting_df = ( - test_stats_df.loc[test_stats_df.adj_pval <= pval, ["K", "target_name"]] - .drop_duplicates() - .groupby(["K"]) - .count() - .reset_index() - ) - ax.plot( - plotting_df["K"], - plotting_df["target_name"], - color=_COLORS["primary"], - linewidth=1.5, - marker="o", - markersize=4, - markerfacecolor=_COLORS["perturbation"], - markeredgecolor=_COLORS["primary"], - markeredgewidth=0.8, - zorder=3, - ) - ax.fill_between( - plotting_df["K"], - plotting_df["target_name"], - alpha=0.08, - color=_COLORS["perturbation"], - ) - _add_vline(ax) - _style_ax( - ax, - xlabel="Number of components (k)", - ylabel="No. unique regulators", - title="Unique regulators for all conditions", - ) - ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) - - # --- Row 3, Col 1: Regulators (per sample) --- - ax = axes[2, 1] - plotting_df_sample = ( - test_stats_df.loc[ - test_stats_df.adj_pval <= pval, ["K", "condition", "target_name"] - ] - .drop_duplicates() - .groupby(["K", "condition"]) - .count() - .reset_index() - ) - sns.lineplot( - x="K", - y="target_name", - hue="condition", - data=plotting_df_sample, - palette=_PALETTE[: plotting_df_sample["condition"].nunique()], - linewidth=1.5, - marker="o", - markersize=4, - ax=ax, - legend="brief", - ) - _add_vline(ax) - _style_ax( - ax, - xlabel="Number of components (k)", - ylabel="No. unique regulators", - title="Unique regulators per condition", - ) - ax.legend( - title="Condition", - fontsize=7, - title_fontsize=8, - frameon=True, - fancybox=False, - edgecolor="#cccccc", - bbox_to_anchor=(1.02, 1), - loc="upper left", - borderaxespad=0, - ) - ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) + # --- Row 3: Regulators (hidden when there is no perturbation data) --- + if test_stats_df is None: + axes[2, 0].set_visible(False) + axes[2, 1].set_visible(False) + else: + _plot_regulator_row(axes[2, 0], axes[2, 1], test_stats_df, pval, _add_vline) # --- Row 3, Col 2: hide empty panel --- axes[2, 2].set_visible(False) @@ -1198,84 +1220,12 @@ def _add_vline(ax): # --- Row 2, Col 2: hide empty panel --- axes[1, 2].set_visible(False) - # --- Row 3, Col 0: Regulators (all conditions) --- - ax = axes[2, 0] - plotting_df = ( - test_stats_df.loc[test_stats_df.adj_pval <= pval, ["K", "target_name"]] - .drop_duplicates() - .groupby(["K"]) - .count() - .reset_index() - ) - ax.plot( - plotting_df["K"], - plotting_df["target_name"], - color=_COLORS["primary"], - linewidth=1.5, - marker="o", - markersize=4, - markerfacecolor=_COLORS["perturbation"], - markeredgecolor=_COLORS["primary"], - markeredgewidth=0.8, - zorder=3, - ) - ax.fill_between( - plotting_df["K"], - plotting_df["target_name"], - alpha=0.08, - color=_COLORS["perturbation"], - ) - _add_vline(ax) - _style_ax( - ax, - xlabel="Number of components (k)", - ylabel="No. unique regulators", - title="Unique regulators for all conditions", - ) - ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) - - # --- Row 3, Col 1: Regulators (per sample) --- - ax = axes[2, 1] - plotting_df_sample = ( - test_stats_df.loc[ - test_stats_df.adj_pval <= pval, ["K", "condition", "target_name"] - ] - .drop_duplicates() - .groupby(["K", "condition"]) - .count() - .reset_index() - ) - sns.lineplot( - x="K", - y="target_name", - hue="condition", - data=plotting_df_sample, - palette=_PALETTE[: plotting_df_sample["condition"].nunique()], - linewidth=1.5, - marker="o", - markersize=4, - ax=ax, - legend="brief", - ) - _add_vline(ax) - _style_ax( - ax, - xlabel="Number of components (k)", - ylabel="No. unique regulators", - title="Unique regulators per condition", - ) - ax.legend( - title="Condition", - fontsize=7, - title_fontsize=8, - frameon=True, - fancybox=False, - edgecolor="#cccccc", - bbox_to_anchor=(1.02, 1), - loc="upper left", - borderaxespad=0, - ) - ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=5)) + # --- Row 3: Regulators (hidden when there is no perturbation data) --- + if test_stats_df is None: + axes[2, 0].set_visible(False) + axes[2, 1].set_visible(False) + else: + _plot_regulator_row(axes[2, 0], axes[2, 1], test_stats_df, pval, _add_vline) # --- Row 3, Col 2: hide empty panel --- axes[2, 2].set_visible(False) diff --git a/src/Stage3_Interpretation/A_Plotting/src/utilities.py b/src/Stage3_Interpretation/A_Plotting/src/utilities.py index 9fa8bf3..061b253 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/utilities.py +++ b/src/Stage3_Interpretation/A_Plotting/src/utilities.py @@ -30,7 +30,7 @@ def convert_with_mygene(dataframe, species="human", index=True): dataframe.index, scopes="ensembl.gene", fields="symbol,name", - species="human", + species=species, ) else: # Query multiple genes at once @@ -38,7 +38,7 @@ def convert_with_mygene(dataframe, species="human", index=True): dataframe.columns, scopes="ensembl.gene", fields="symbol,name", - species="human", + species=species, ) # Create mapping dictionary @@ -64,7 +64,7 @@ def convert_adata_with_mygene(adata, species="human"): mg = mygene.MyGeneInfo() gene_list = adata.var_names.tolist() annotations = mg.querymany( - gene_list, scopes="ensembl.gene", fields="symbol", species="human" + gene_list, scopes="ensembl.gene", fields="symbol", species=species ) # Process the results to create mapping diff --git a/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py b/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py index ef48ea4..94d0457 100644 --- a/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py +++ b/src/Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py @@ -155,6 +155,22 @@ def main(): default="guide_targets", help="uns key for guide target names. Default: guide_targets", ) + parser.add_argument( + "--organism", + type=str, + default="human", + help="Species of the data, for MyGene gene summaries in the Program Loadings " + "sheet. With mouse, each gene gets the summaries of its human orthologs " + "(MGI + Ensembl), listed in a 'Human ortholog' column. Default: human", + ) + parser.add_argument( + "--ortholog_table", + type=str, + default=None, + help="For --organism mouse: TSV of mouse-human orthologs (columns: " + "mouse_symbol, mouse_ensembl, human_symbol, source). Default: the table " + "cached in src/Stage2_Evaluation/Resources/, downloaded if missing", + ) parser.add_argument( "--gene_names_key", type=str, @@ -289,6 +305,8 @@ def main(): num_gene=args.num_gene, data_key=args.data_key, gene_names_key=args.gene_names_key, + species=args.organism, + ortholog_table=args.ortholog_table, ) df_Program_loading_flat = compile_Program_loading_score_sheet_flat( mdata, @@ -370,19 +388,29 @@ def main(): ) # ── Target Summary (writes specificity / correlation / KD-efficiency sidecars to save_path) ── - df_Target_Summary = Compile_Target_Summary_sheet( - mdata, - Perturbation_path_base, - Sample=args.Conditions, - categorical_key=args.categorical_key, - prog_key=args.prog_key, - data_key=args.data_key, - guide_targets_key=args.guide_targets_key, - save_path=args.save_path, - effect_size=args.effect_size, - control_target_name=args.control_target_name, - gene_names_key=args.gene_names_key, - ) + # Needs guide information; skipped for non-perturbed data. + has_guides = args.guide_targets_key in mdata[args.prog_key].uns + if has_guides: + df_Target_Summary = Compile_Target_Summary_sheet( + mdata, + Perturbation_path_base, + Sample=args.Conditions, + categorical_key=args.categorical_key, + prog_key=args.prog_key, + data_key=args.data_key, + guide_targets_key=args.guide_targets_key, + save_path=args.save_path, + effect_size=args.effect_size, + control_target_name=args.control_target_name, + gene_names_key=args.gene_names_key, + ) + else: + print( + f"No uns['{args.guide_targets_key}'] in '{args.prog_key}': " + "skipping the Targets Summary sheet (non-perturbed data).", + flush=True, + ) + df_Target_Summary = None # ── Program Summary ── df_Summary = Compile_Summary_sheet( @@ -393,7 +421,8 @@ def main(): df_Program_loading_flat, df_Explained_Variance, Sample=args.Conditions, - specicicity_path=args.save_path, + # specificity scores are written by the Targets Summary step + specicicity_path=args.save_path if has_guides else None, categorical_key=args.categorical_key, non_tagerting_key=args.non_targeting_key, effect_size=args.effect_size, @@ -405,9 +434,10 @@ def main(): df_Program_loading_long.to_csv( f"{args.save_path}/Program_Loadings_{args.K}_{thresh_str}.tsv", sep="\t" ) - df_Target_Summary.to_csv( - f"{args.save_path}/Targets_Summary_{args.K}_{thresh_str}.tsv", sep="\t" - ) + if df_Target_Summary is not None: + df_Target_Summary.to_csv( + f"{args.save_path}/Targets_Summary_{args.K}_{thresh_str}.tsv", sep="\t" + ) print(f"Saved separate TSV files to {args.save_path}") # ── Write Excel ── @@ -418,7 +448,8 @@ def main(): df_Program_loading_long.to_excel( writer, sheet_name="Program Loadings", index=True ) - df_Target_Summary.to_excel(writer, sheet_name="Targets Summary", index=True) + if df_Target_Summary is not None: + df_Target_Summary.to_excel(writer, sheet_name="Targets Summary", index=True) if df_Association is not None: df_Association.to_excel(writer, sheet_name="Sample Association", index=True) diff --git a/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py b/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py index 924b338..a4da4ff 100644 --- a/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py +++ b/src/Stage3_Interpretation/B_Summarization/src/Compile_excel_sheet.py @@ -112,8 +112,19 @@ def add_specificity_scores_file( # -------------- helper methods for loading sheets-------------- def compile_Program_loading_score_sheet_long( - mdata, num_gene=300, data_key="rna", gene_names_key=None + mdata, + num_gene=300, + data_key="rna", + gene_names_key=None, + species="human", + ortholog_table=None, ): + """Top genes per program with MyGene summaries (long form). + + For species='mouse', each gene is annotated with the summaries of its human + orthologs (MGI + Ensembl table, see Stage2 orthologs.py), listed in a 'Human + ortholog' column; ortholog_table is an optional path to a pre-built table. + """ print("Load program loadings data in long form") @@ -124,7 +135,9 @@ def compile_Program_loading_score_sheet_long( ) # Rename columns to curated gene names - if gene_names_key is not None: + # Fall back to var_names when the symbol column is absent (e.g. var_names are + # already symbols), as compute_kd_efficiency does. + if gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: rename_map = dict( zip(mdata[data_key].var_names, mdata[data_key].var[gene_names_key]) ) @@ -146,6 +159,15 @@ def compile_Program_loading_score_sheet_long( mg = MyGeneInfo() long_data = [] + use_orthologs = species == "mouse" + if use_orthologs: + from Stage2_Evaluation.A_Metrics.src.orthologs import ( + load_ortholog_table, + map_mouse_to_human, + ) + + table = load_ortholog_table(ortholog_table) + for program_idx, row in result_df.iterrows(): # Query all genes for this program at once genes_list = row.dropna().astype(str).str.strip().tolist() @@ -154,13 +176,24 @@ def compile_Program_loading_score_sheet_long( if not genes_list: continue + if use_orthologs: + orthologs = map_mouse_to_human(genes_list, table) + query_genes = sorted({h for hs in orthologs.values() for h in hs}) + query_species = "human" + else: + query_genes, query_species = genes_list, species + # Get results for all genes in this program - results = mg.querymany( - genes_list, - scopes="symbol", - fields="summary", - species="human", - verbose=False, + results = ( + mg.querymany( + query_genes, + scopes="symbol", + fields="summary", + species=query_species, + verbose=False, + ) + if query_genes + else [] ) # Create a lookup dictionary for faster access @@ -175,16 +208,23 @@ def compile_Program_loading_score_sheet_long( # Skip NaN values if pd.notna(gene) and gene != "": gene_clean = str(gene).strip() - annotation = summary_dict.get(gene_clean, "N/A") - - long_data.append( - { - "Program": program_idx, - "Rank": rank, - "Gene": gene_clean, - "Annotation": annotation, - } - ) + record = { + "Program": program_idx, + "Rank": rank, + "Gene": gene_clean, + } + if use_orthologs: + human = orthologs.get(gene_clean, []) + summaries = [ + f"{h}: {summary_dict[h]}" + for h in human + if summary_dict.get(h, "N/A") != "N/A" + ] + record["Human ortholog"] = ";".join(human) + record["Annotation"] = " | ".join(summaries) or "N/A" + else: + record["Annotation"] = summary_dict.get(gene_clean, "N/A") + long_data.append(record) # Convert to DataFrame annotation_df = pd.DataFrame(long_data) @@ -205,7 +245,9 @@ def compile_Program_loading_score_sheet_flat( ) # Rename columns to curated gene names - if gene_names_key is not None: + # Fall back to var_names when the symbol column is absent (e.g. var_names are + # already symbols), as compute_kd_efficiency does. + if gene_names_key is not None and gene_names_key in mdata[data_key].var.columns: rename_map = dict( zip(mdata[data_key].var_names, mdata[data_key].var[gene_names_key]) ) diff --git a/tests/Script/Stage2_Evaluation/test_orthologs.py b/tests/Script/Stage2_Evaluation/test_orthologs.py new file mode 100644 index 0000000..9ec18c4 --- /dev/null +++ b/tests/Script/Stage2_Evaluation/test_orthologs.py @@ -0,0 +1,157 @@ +"""Unit tests for mouse to human ortholog mapping (Stage2 orthologs.py). + +Uses a small synthetic ortholog table, so no network access is needed. + +Usage (from the repo root): + python -m pytest tests/Script/Stage2_Evaluation/test_orthologs.py -v +""" + +import anndata as ad +import mudata as md +import numpy as np +import pandas as pd +import pytest + +from Stage2_Evaluation.A_Metrics.src.enrichment_geneset import ( + get_program_gene_loadings, +) +from Stage2_Evaluation.A_Metrics.src.orthologs import ( + TABLE_COLUMNS, + collapse_to_human, + load_ortholog_table, + map_mouse_to_human, +) + +# mouse_symbol, mouse_ensembl, human_symbol, source +ORTHOLOGS = [ + ("Cdh5", "ENSMUSG00000031871", "CDH5", "Ensembl"), # same name, one-to-one + ("Cdh5", np.nan, "CDH5", "MGI"), # duplicate from the second source + ("Zfp451", np.nan, "ZNF451", "MGI"), # different name + ("H2-Ab1", np.nan, "HLA-DQB1", "Ensembl"), # one-to-many + ("H2-Ab1", np.nan, "HLA-DQB2", "Ensembl"), + ("Cd59a", np.nan, "CD59", "Ensembl"), # many-to-one + ("Cd59b", np.nan, "CD59", "Ensembl"), + ("Gm12345", "ENSMUSG00000099999", np.nan, "Ensembl"), # no human ortholog +] +# Upper-cases to an existing human symbol but is not its ortholog in the table. +CASE_DECOY = "Ccdc115" + + +@pytest.fixture +def table(): + return pd.DataFrame(ORTHOLOGS, columns=TABLE_COLUMNS) + + +@pytest.fixture +def table_path(table, tmp_path): + path = tmp_path / "orthologs.tsv" + table.to_csv(path, sep="\t", index=False) + return str(path) + + +def test_map_mouse_to_human(table): + mapping = map_mouse_to_human( + ["Cdh5", "Zfp451", "H2-Ab1", "Cd59a", "Gm12345", CASE_DECOY], table + ) + assert mapping == { + "Cdh5": ["CDH5"], + "Zfp451": ["ZNF451"], + "H2-Ab1": ["HLA-DQB1", "HLA-DQB2"], + "Cd59a": ["CD59"], + "Gm12345": [], + CASE_DECOY: [], # no fallback to upper-casing + } + + +def test_map_ensembl_ids(table): + mapping = map_mouse_to_human( + ["ENSMUSG00000031871", "ENSMUSG00000031871.7", "ENSMUSG00000099999"], table + ) + assert mapping["ENSMUSG00000031871"] == ["CDH5"] + assert mapping["ENSMUSG00000031871.7"] == ["CDH5"] # version suffix ignored + assert mapping["ENSMUSG00000099999"] == [] + + +def test_collapse_to_human(table): + loadings = pd.DataFrame( + {"p0": [5.0, 1.0, 2.0, 3.0, 4.0, 9.0], "p1": [0.5, 0.1, 0.2, 0.3, 0.6, 0.9]}, + index=["Cdh5", "Zfp451", "H2-Ab1", "Cd59a", "Cd59b", "Gm12345"], + ) + human = collapse_to_human(loadings, table) + + assert set(human.index) == {"CDH5", "ZNF451", "HLA-DQB1", "HLA-DQB2", "CD59"} + # one-to-many: both human orthologs take the mouse loading + assert human.loc["HLA-DQB1", "p0"] == human.loc["HLA-DQB2", "p0"] == 2.0 + # many-to-one: per-program maximum over the mouse orthologs + assert human.loc["CD59", "p0"] == 4.0 and human.loc["CD59", "p1"] == 0.6 + assert human.loc["ZNF451", "p1"] == 0.1 + + +def test_collapse_without_any_ortholog_fails(table): + with pytest.raises(ValueError, match="No mouse genes"): + collapse_to_human(pd.DataFrame({"p0": [1.0]}, index=["Gm12345"]), table) + + +def test_load_ortholog_table_validates_columns(tmp_path): + bad = tmp_path / "bad.tsv" + pd.DataFrame({"mouse": ["Cdh5"], "human": ["CDH5"]}).to_csv( + bad, sep="\t", index=False + ) + with pytest.raises(ValueError, match="missing columns"): + load_ortholog_table(str(bad)) + + +def _mdata(genes): + rng = np.random.default_rng(0) + rna = ad.AnnData( + X=rng.random((6, len(genes))).astype(np.float32), + var=pd.DataFrame(index=genes), + ) + prog = ad.AnnData( + X=rng.random((6, 2)).astype(np.float32), var=pd.DataFrame(index=["0", "1"]) + ) + prog.varm["loadings"] = rng.random((2, len(genes))).astype(np.float32) + prog.uns["var_names"] = np.array(genes) + return md.MuData({"rna": rna, "cNMF": prog}) + + +def test_program_loadings_mouse_uses_orthologs(table_path): + genes = ["Cdh5", "Zfp451", "H2-Ab1", "Gm12345", CASE_DECOY] + loadings = get_program_gene_loadings( + _mdata(genes), + prog_key="cNMF", + data_key="rna", + organism="mouse", + ortholog_table=table_path, + ) + assert set(loadings.index) == {"CDH5", "ZNF451", "HLA-DQB1", "HLA-DQB2"} + assert list(loadings.columns) == ["0", "1"] + + +def test_program_loadings_mouse_native_symbols(table_path): + """map_to_human=False (MSigDB mouse gene sets) keeps mouse symbols unchanged.""" + genes = ["Cdh5", "Zfp451"] + loadings = get_program_gene_loadings( + _mdata(genes), + prog_key="cNMF", + data_key="rna", + organism="mouse", + ortholog_table=table_path, + map_to_human=False, + ) + assert list(loadings.index) == genes + + +def test_program_loadings_human_unchanged(): + genes = ["CDH5", "znf451"] + loadings = get_program_gene_loadings( + _mdata(genes), prog_key="cNMF", data_key="rna", organism="human" + ) + assert list(loadings.index) == ["CDH5", "ZNF451"] + + +def test_unsupported_organism(): + with pytest.raises(ValueError, match="Unsupported organism"): + get_program_gene_loadings( + _mdata(["Cdh5"]), prog_key="cNMF", data_key="rna", organism="rat" + ) diff --git a/tests/Script/test_non_perturbed.py b/tests/Script/test_non_perturbed.py new file mode 100644 index 0000000..818ee88 --- /dev/null +++ b/tests/Script/test_non_perturbed.py @@ -0,0 +1,368 @@ +"""Tests for running the pipeline on non-perturbed (no guide) and non-human data. + +Covers the code paths that previously assumed CRISPR guide information or human data: + +- Stage 2 evaluation skips trait (GWAS) enrichment for non-human organisms and fails + early with a clear message when perturbation association is requested without guides. +- Stage 3 K-selection runs without trait or perturbation files (--no_trait, + --no_perturbation). +- Stage 3 Excel summary runs on an h5mu without uns['guide_targets'] and, for + mouse data, annotates genes with their human orthologs' summaries. + +All inputs are small synthetic files; no network access or inference output is needed. + +Usage (from the repo root, in the pipeline environment): + python -m pytest tests/Script/test_non_perturbed.py -v +""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path +from typing import ClassVar + +import matplotlib + +matplotlib.use("Agg") + +import anndata as ad +import mudata as md +import numpy as np +import pandas as pd +import pytest + +PIPELINE_ROOT = Path(__file__).resolve().parents[2] +SRC = PIPELINE_ROOT / "src" +sys.path.insert(0, str(SRC)) + +EVAL_SCRIPT = ( + SRC / "Stage2_Evaluation/A_Metrics/Slurm_Version/cNMF_evaluation_pipeline.py" +) +K_SELECTION_SCRIPT = ( + SRC / "Stage3_Interpretation/A_Plotting/Slurm_Version/cNMF_k_selection.py" +) +EXCEL_SCRIPT = ( + SRC / "Stage3_Interpretation/B_Summarization/Slurm_Version/cNMF_excel_summary.py" +) + +K_VALUES = [3, 5] +SEL_THRESH = 2.0 +SEL_STR = "2_0" +RUN_NAME = "non_perturbed_run" +N_CELLS = 40 +N_GENES = 25 +SEED = 0 + + +def _load_module(path, name): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _make_mdata(k, rng): + """MuData with 'rna' and 'cNMF' modalities and no guide information.""" + cells = [f"cell{i}" for i in range(N_CELLS)] + genes = [f"Gene{i}" for i in range(N_GENES)] # mouse-style capitalization + rna = ad.AnnData( + X=rng.poisson(2.0, size=(N_CELLS, N_GENES)).astype(np.float32), + obs=pd.DataFrame( + {"sample": pd.Categorical(rng.choice(["ctrl", "treated"], N_CELLS))}, + index=cells, + ), + # Symbols as var_names and no 'symbol' column, as in Seurat-converted data + var=pd.DataFrame(index=genes), + ) + prog = ad.AnnData( + X=rng.random((N_CELLS, k)).astype(np.float32), + obs=pd.DataFrame(index=cells), + var=pd.DataFrame(index=[str(p) for p in range(k)]), + ) + prog.varm["loadings"] = rng.random((k, N_GENES)).astype(np.float32) + prog.uns["var_names"] = np.array(genes) + return md.MuData({"rna": rna, "cNMF": prog}) + + +@pytest.fixture(scope="module") +def run_dir(tmp_path_factory): + """Synthetic run directory: inference h5mu files plus Stage 2 outputs without + trait or perturbation results, as produced for non-perturbed mouse data.""" + rng = np.random.default_rng(SEED) + out_dir = tmp_path_factory.mktemp("out") + run = out_dir / RUN_NAME + adata_dir = run / "Inference" / "adata" + adata_dir.mkdir(parents=True) + + stability = [] + for k in K_VALUES: + _make_mdata(k, rng).write(adata_dir / f"cNMF_{k}_{SEL_STR}.h5mu") + + eval_dir = run / "Evaluation" / f"{k}_{SEL_STR}" + eval_dir.mkdir(parents=True) + programs = [str(p) for p in range(k)] + for fname in [f"{k}_GO_term_enrichment.txt", f"{k}_geneset_enrichment.txt"]: + pd.DataFrame( + { + "Term": [f"TERM_{p}" for p in programs], + "Adjusted P-value": rng.uniform(0, 0.1, k), + "Genes": ["Gene0;Gene1"] * k, + "program_name": programs, + } + ).to_csv(eval_dir / fname, sep="\t", index=False) + pd.DataFrame({"Total": [rng.uniform(0.2, 0.6)]}, index=["all"]).to_csv( + eval_dir / f"{k}_Explained_Variance_Summary.txt", sep="\t" + ) + pd.DataFrame( + { + "VarianceExplained": rng.uniform(0, 0.1, k), + "ProgramID": programs, + "program_name": programs, + } + ).to_csv(eval_dir / f"{k}_Explained_Variance.txt", sep="\t", index=False) + stability.append( + { + "k": k, + "silhouette": rng.uniform(0.5, 1.0), + "prediction_error": rng.uniform(1, 2), + } + ) + + stability_file = run / "Inference" / "k_selection_stats.tsv" + pd.DataFrame(stability).to_csv(stability_file, sep="\t", index=False) + return {"out_dir": out_dir, "run": run, "stability_file": stability_file} + + +# Stage 2: evaluation + + +@pytest.fixture(scope="module") +def evaluation_main(): + return _load_module(EVAL_SCRIPT, "cNMF_evaluation_pipeline").main + + +def _eval_argv(run_dir, *extra): + return [ + str(EVAL_SCRIPT), + "--out_dir", + str(run_dir["out_dir"]), + "--run_name", + RUN_NAME, + "--K", + str(K_VALUES[0]), + "--sel_threshs", + str(SEL_THRESH), + *extra, + ] + + +def test_trait_enrichment_skipped_for_mouse( + run_dir, evaluation_main, monkeypatch, capsys +): + """--Perform_trait with a non-human organism is skipped with a message, + without requiring --gwas_data_path.""" + monkeypatch.setattr( + sys, "argv", _eval_argv(run_dir, "--Perform_trait", "--organism", "mouse") + ) + assert evaluation_main() == 0 + + assert "skipping trait (GWAS) enrichment" in capsys.readouterr().out + eval_dir = run_dir["run"] / "Evaluation" / f"{K_VALUES[0]}_{SEL_STR}" + assert not (eval_dir / f"{K_VALUES[0]}_trait_enrichment.txt").exists() + + +def test_trait_enrichment_still_requires_gwas_for_human( + run_dir, evaluation_main, monkeypatch +): + """Human behavior is unchanged: --Perform_trait needs --gwas_data_path.""" + monkeypatch.setattr( + sys, "argv", _eval_argv(run_dir, "--Perform_trait", "--organism", "human") + ) + with pytest.raises(SystemExit): + evaluation_main() + + +def test_perturbation_without_guides_fails_clearly( + run_dir, evaluation_main, monkeypatch +): + monkeypatch.setattr(sys, "argv", _eval_argv(run_dir, "--Perform_perturbation")) + with pytest.raises(ValueError, match="drop --Perform_perturbation"): + evaluation_main() + + +# Stage 3: K-selection plotting + + +def test_k_selection_without_trait_and_perturbation(run_dir, tmp_path): + save_dir = tmp_path / "K_selection" + cmd = [ + sys.executable, + str(K_SELECTION_SCRIPT), + "--output_directory", + str(run_dir["out_dir"]), + "--run_name", + RUN_NAME, + "--save_folder_name", + str(save_dir), + "--eval_folder_name", + str(run_dir["run"] / "Evaluation"), + "--K", + *[str(k) for k in K_VALUES], + "--sel_threshs", + str(SEL_THRESH), + "--stability_file", + str(run_dir["stability_file"]), + "--no_trait", + "--no_perturbation", + ] + env = {**os.environ, "MPLBACKEND": "Agg"} + result = subprocess.run(cmd, capture_output=True, text=True, env=env, check=False) + assert result.returncode == 0, result.stdout[-3000:] + result.stderr[-3000:] + + produced = {p.name for p in save_dir.iterdir()} + assert f"K-selection_panel_{SEL_THRESH}.png" in produced + assert f"Enrichment_{SEL_THRESH}_go_terms.png" in produced + assert f"Enrichment_{SEL_THRESH}_traits.png" not in produced + assert not any(name.startswith("Perturbation_") for name in produced) + + +def test_k_selection_requires_files_by_default(run_dir, tmp_path): + """Without the new flags, missing trait/perturbation files are still an error.""" + cmd = [ + sys.executable, + str(K_SELECTION_SCRIPT), + "--output_directory", + str(run_dir["out_dir"]), + "--run_name", + RUN_NAME, + "--save_folder_name", + str(tmp_path / "K_selection"), + "--eval_folder_name", + str(run_dir["run"] / "Evaluation"), + "--K", + *[str(k) for k in K_VALUES], + "--sel_threshs", + str(SEL_THRESH), + "--Conditions", + "ctrl", + "--stability_file", + str(run_dir["stability_file"]), + ] + result = subprocess.run(cmd, capture_output=True, text=True, check=False) + assert result.returncode != 0 + assert "required evaluation file(s) not found" in result.stdout + + +def test_k_selection_panel_hides_regulator_row(tmp_path): + from Stage3_Interpretation.A_Plotting.src import k_selection_plots as ksp + + stats = pd.DataFrame( + {"k": K_VALUES, "silhouette": [0.8, 0.7], "prediction_error": [1.5, 1.2]} + ) + counts = pd.DataFrame( + {"go_terms": [3, 4], "genesets": [2, 3], "traits": [0, 0]}, index=K_VALUES + ) + variance = {3: 0.3, 5: 0.4} + figures = [] + original_subplots = ksp.plt.subplots + + def capture_subplots(*args, **kwargs): + fig, axes = original_subplots(*args, **kwargs) + figures.append(axes) + return fig, axes + + ksp.plt.subplots = capture_subplots + try: + for panel_fn in [ + ksp.plot_k_selection_panel, + ksp.plot_k_selection_panel_no_traits, + ]: + panel_fn( + stats, counts, None, variance, folder_name=str(tmp_path), file_name="p" + ) + finally: + ksp.plt.subplots = original_subplots + + for axes in figures: + assert not axes[2, 0].get_visible() + assert not axes[2, 1].get_visible() + + +# Stage 3: Excel summary + + +class _FakeMyGene: + """Offline stand-in for mygene.MyGeneInfo that records the species queried.""" + + species_seen: ClassVar[list] = [] + + def querymany(self, genes, species=None, **kwargs): + _FakeMyGene.species_seen.append(species) + return [{"query": g, "summary": f"summary of {g}"} for g in genes] + + +def test_excel_summary_without_guides(run_dir, monkeypatch, capsys, tmp_path): + """Non-perturbed mouse data: no Targets Summary sheet, and gene summaries come + from human orthologs (synthetic table: even-numbered genes have an ortholog).""" + excel = _load_module(EXCEL_SCRIPT, "cNMF_excel_summary") + from Stage3_Interpretation.B_Summarization.src import Compile_excel_sheet + + monkeypatch.setattr(Compile_excel_sheet, "MyGeneInfo", _FakeMyGene) + _FakeMyGene.species_seen.clear() + ortholog_table = tmp_path / "orthologs.tsv" + pd.DataFrame( + { + "mouse_symbol": [f"Gene{i}" for i in range(0, N_GENES, 2)], + "mouse_ensembl": pd.NA, + "human_symbol": [f"HSGENE{i}" for i in range(0, N_GENES, 2)], + "source": "test", + } + ).to_csv(ortholog_table, sep="\t", index=False) + + k = K_VALUES[1] + monkeypatch.setattr( + sys, + "argv", + [ + str(EXCEL_SCRIPT), + "--out_dir", + str(run_dir["out_dir"]), + "--run_name", + RUN_NAME, + "--K", + str(k), + "--sel_thresh", + str(SEL_THRESH), + "--organism", + "mouse", + "--categorical_key", + "sample", + "--num_gene", + "5", + "--ortholog_table", + str(ortholog_table), + ], + ) + assert excel.main() == 0 + + assert "skipping the Targets Summary sheet" in capsys.readouterr().out + assert set(_FakeMyGene.species_seen) == {"human"} + save_path = run_dir["run"] / "Interpretation" / "Summary_table" / f"{k}_{SEL_STR}" + sheets = pd.ExcelFile(save_path / f"cNMF_{k}_{SEL_STR}.xlsx").sheet_names + assert "Summary" in sheets and "Program Loadings" in sheets + assert "Targets Summary" not in sheets + assert not (save_path / f"Targets_Summary_{k}_{SEL_STR}.tsv").exists() + + loadings = pd.read_csv( + save_path / f"Program_Loadings_{k}_{SEL_STR}.tsv", + sep="\t", + keep_default_na=False, # keep the literal "N/A" + ) + for _, row in loadings.iterrows(): + n = int(row["Gene"].removeprefix("Gene")) + if n % 2 == 0: + assert row["Human ortholog"] == f"HSGENE{n}" + assert row["Annotation"] == f"HSGENE{n}: summary of HSGENE{n}" + else: + assert row["Human ortholog"] == "" and row["Annotation"] == "N/A" From 9370349f588c92266bf01dd54bdbb4b624ca3406 Mon Sep 17 00:00:00 2001 From: OlgaPushkarev Date: Fri, 9 Oct 2026 13:25:19 -0700 Subject: [PATCH 5/5] Reuse Seurat-style UMAP/PCA embeddings in ensure_umap Data converted from Seurat stores its embeddings as obsm['UMAP'] / obsm['PCA'] (often as DataFrames), which scanpy does not see, so program and perturbed-gene plots recomputed a different UMAP. ensure_umap now copies such embeddings to 'X_umap' (first two columns) / 'X_pca' before deciding whether to compute one. Existing X_umap/X_pca entries and the original keys are left untouched; files without these keys behave as before. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../A_Plotting/src/utilities.py | 32 ++++++- .../A_Plotting/Program/test_ensure_umap.py | 87 +++++++++++++++++++ 2 files changed, 118 insertions(+), 1 deletion(-) create mode 100644 tests/Script/Stage3_Interpretation/A_Plotting/Program/test_ensure_umap.py diff --git a/src/Stage3_Interpretation/A_Plotting/src/utilities.py b/src/Stage3_Interpretation/A_Plotting/src/utilities.py index 061b253..4258e14 100644 --- a/src/Stage3_Interpretation/A_Plotting/src/utilities.py +++ b/src/Stage3_Interpretation/A_Plotting/src/utilities.py @@ -225,8 +225,38 @@ def _natural_sort_key(filepath): ] +# obsm keys used by other tools (e.g. Seurat conversions) for the scanpy embeddings +ALT_EMBEDDING_KEYS = {"X_umap": ("UMAP", "umap"), "X_pca": ("PCA", "pca")} + + +def _adopt_alt_embedding_keys(adata): + """Copy an embedding stored under a non-scanpy obsm key to 'X_umap'/'X_pca'. + + Reuses the embedding the data came with instead of recomputing a different + one. Existing 'X_umap'/'X_pca' entries and the original keys are left as is. + """ + for std_key, alt_keys in ALT_EMBEDDING_KEYS.items(): + if std_key in adata.obsm: + continue + for alt_key in alt_keys: + if alt_key not in adata.obsm: + continue + emb = adata.obsm[alt_key] + # scanpy plotting needs an array; some converters store a DataFrame + emb = emb.to_numpy() if hasattr(emb, "to_numpy") else np.asarray(emb) + adata.obsm[std_key] = emb[:, :2] if std_key == "X_umap" else emb + print(f"ensure_umap: using obsm['{alt_key}'] as '{std_key}'.", flush=True) + break + + def ensure_umap(mdata, data_key, prog_key, n_top_genes=2000, n_comps=50): - """Compute PCA/UMAP from top-variance genes and write to both modalities; no-op if X_umap already exists in both.""" + """Compute PCA/UMAP from top-variance genes and write to both modalities; no-op if X_umap already exists in both. + + Embeddings stored under 'UMAP'/'PCA' (e.g. Seurat conversions) are reused first. + """ + _adopt_alt_embedding_keys(mdata[data_key]) + _adopt_alt_embedding_keys(mdata[prog_key]) + has_prog = "X_umap" in mdata[prog_key].obsm has_data = "X_umap" in mdata[data_key].obsm if has_prog and has_data: diff --git a/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_ensure_umap.py b/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_ensure_umap.py new file mode 100644 index 0000000..a740d29 --- /dev/null +++ b/tests/Script/Stage3_Interpretation/A_Plotting/Program/test_ensure_umap.py @@ -0,0 +1,87 @@ +"""Tests for ensure_umap reusing embeddings stored under non-scanpy obsm keys. + +Seurat conversions store embeddings as obsm['UMAP'] / obsm['PCA'] (often as +DataFrames); ensure_umap should reuse them instead of computing a new UMAP. + +Usage (from the repo root): + python -m pytest tests/Script/Stage3_Interpretation/A_Plotting/Program/test_ensure_umap.py -v +""" + +import anndata as ad +import mudata as md +import numpy as np +import pandas as pd + +from Stage3_Interpretation.A_Plotting.src.utilities import ensure_umap + +N_CELLS = 60 +N_GENES = 40 +N_PROGRAMS = 3 +SEED = 0 + + +def _mdata(): + """MuData with 'rna' and 'cNMF' modalities and no embeddings.""" + rng = np.random.default_rng(SEED) + cells = [f"cell{i}" for i in range(N_CELLS)] + rna = ad.AnnData( + X=rng.poisson(3.0, size=(N_CELLS, N_GENES)).astype(np.float32), + obs=pd.DataFrame(index=cells), + var=pd.DataFrame(index=[f"Gene{i}" for i in range(N_GENES)]), + ) + prog = ad.AnnData( + X=rng.random((N_CELLS, N_PROGRAMS)).astype(np.float32), + obs=pd.DataFrame(index=cells), + var=pd.DataFrame(index=[str(p) for p in range(N_PROGRAMS)]), + ) + return md.MuData({"rna": rna, "cNMF": prog}) + + +def _seurat_embedding(mdata, n_cols, prefix): + rng = np.random.default_rng(SEED + n_cols) + return pd.DataFrame( + rng.normal(size=(N_CELLS, n_cols)), + index=mdata["rna"].obs_names, + columns=[f"{prefix}_{i + 1}" for i in range(n_cols)], + ) + + +def test_adopts_seurat_keys_and_mirrors_to_programs(capsys): + mdata = _mdata() + umap = _seurat_embedding(mdata, 3, "UMAP") # some converters keep 3 dims + pca = _seurat_embedding(mdata, 10, "PC") + mdata["rna"].obsm["UMAP"] = umap + mdata["rna"].obsm["PCA"] = pca + + ensure_umap(mdata, "rna", "cNMF") + + for key in ["rna", "cNMF"]: + x_umap = mdata[key].obsm["X_umap"] + assert isinstance(x_umap, np.ndarray) + np.testing.assert_array_equal(x_umap, umap.to_numpy()[:, :2]) + np.testing.assert_array_equal(mdata[key].obsm["X_pca"], pca.to_numpy()) + assert "UMAP" in mdata["rna"].obsm and "PCA" in mdata["rna"].obsm + assert "using obsm['UMAP'] as 'X_umap'" in capsys.readouterr().out + + +def test_existing_x_umap_is_not_overwritten(): + mdata = _mdata() + existing = np.zeros((N_CELLS, 2)) + for key in ["rna", "cNMF"]: + mdata[key].obsm["X_umap"] = existing + mdata["rna"].obsm["UMAP"] = _seurat_embedding(mdata, 2, "UMAP") + + ensure_umap(mdata, "rna", "cNMF") + + for key in ["rna", "cNMF"]: + np.testing.assert_array_equal(mdata[key].obsm["X_umap"], existing) + + +def test_computes_umap_when_no_embedding_exists(): + mdata = _mdata() + + ensure_umap(mdata, "rna", "cNMF", n_top_genes=N_GENES, n_comps=5) + + for key in ["rna", "cNMF"]: + assert mdata[key].obsm["X_umap"].shape == (N_CELLS, 2) + assert mdata[key].obsm["X_pca"].shape == (N_CELLS, 5)
') parts.append( - f'

{rank}. {_badge_html(badge)} ' + f"

{rank}. {_badge_html(badge)} " f'' - f'{_esc(paper.title or "Untitled")}

' + f"{_esc(paper.title or 'Untitled')}" ) parts.append( f'

PMID:{pmid}' - + (f' • {_esc(paper.authors)}' if paper.authors else '') - + (f' • {_esc(paper.journal)} ({paper.year})' if paper.journal else '') - + '

' + + (f" • {_esc(paper.authors)}" if paper.authors else "") + + (f" • {_esc(paper.journal)} ({paper.year})" if paper.journal else "") + + "

" ) # LLM summary if paper.summary: - parts.append(f'

Summary: {_esc(paper.summary)}

') + parts.append(f"

Summary: {_esc(paper.summary)}

") # Evidence sentences if paper.evidence_sentences: - parts.append('
Evidence sentences
    ') + parts.append( + '
    Evidence sentences
      ' + ) for gene, sents in paper.evidence_sentences.items(): for sent in sents: parts.append( f'
    • {_esc(gene)} ' - f'“{_esc(sent)}”
      • ' + f"“{_esc(sent)}”" ) - parts.append('
    ') + parts.append("
") # Gene-gene relations if paper.relations: - parts.append('
Gene-gene relations
    ') + parts.append( + '
    Gene-gene relations
      ' + ) for rel in paper.relations[:5]: rtype = rel["type"].replace("_", " ") score = rel.get("score", "") - score_str = f' ({score:.2f})' if score else "" + score_str = ( + f' ({score:.2f})' if score else "" + ) parts.append( f'
    • {_esc(rel["gene1"])} ' - f'→ {_esc(rtype)} → ' + f"→ {_esc(rtype)} → " f'{_esc(rel["gene2"])}{score_str}
      • ' ) - parts.append('
    ') + parts.append("
") # Verification if vr_list: - parts.append('
Verification details
    ') + parts.append( + '
    Verification details
      ' + ) for vr in vr_list: icon = "✓" if vr.passed else "✗" cls = "pass" if vr.passed else "fail" parts.append( f'
    • {icon} L{vr.level} ({_esc(vr.method)}): ' - f'{_esc(vr.details)}
      • ' + f"{_esc(vr.details)}" ) - parts.append('
    ') + parts.append("
") - parts.append('