diff --git a/annotate_clean.py b/annotate_clean.py index 5a37a94..b2e1ed7 100644 --- a/annotate_clean.py +++ b/annotate_clean.py @@ -54,14 +54,17 @@ def load_token(path="~/.config/oncokb/token"): return p.read_text().strip() return os.environ.get("ONCOKB_TOKEN", "").strip() -def vep_annotate(df, reference, gtf_path=None): +def vep_annotate(df, reference, gtf_path=None, offline=False): """Try VEP if installed and cache exists, else Ensembl REST, else fallback.""" import shutil vep_bin = shutil.which("vep") if vep_bin and gtf_path and Path(gtf_path).is_file(): pass rows = [] - for _, r in df.iterrows(): + total = len(df) + for idx, (_, r) in enumerate(df.iterrows()): + if idx % 500 == 0 and total > 500: + print(f" annotate {idx}/{total} ...", flush=True) chrom = str(r["chrom"]) pos1 = int(r["position"]) + 1 ref = str(r["ref"]); alt = str(r["alt"]) @@ -77,11 +80,21 @@ def vep_annotate(df, reference, gtf_path=None): "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel", }) continue + if offline: + rows.append({ + "_orig_idx": r.name, + "Ген": "intergenic", + "HGVS_c": f"c.{pos1}{ref}>{alt}", + "HGVS_p": "p.(?)", + "HGVS": f"c.{pos1}{ref}>{alt} p.(?) ({hgvs_g})", + "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel", + }) + continue # Real hg38 - try Ensembl REST, expand all transcripts try: hgvs_ens = f"{chrom.replace('chr','')}:g.{pos1}{ref}>{alt}" url = f"https://rest.ensembl.org/vep/homo_sapiens/hgvs/{hgvs_ens}?content-type=application/json" - resp = requests.get(url, headers={"Content-Type": "application/json", "Accept": "application/json"}, timeout=15) + resp = requests.get(url, headers={"Content-Type": "application/json", "Accept": "application/json"}, timeout=(3, 5)) if resp.status_code == 200: j = resp.json() tcs = j[0].get("transcript_consequences", []) if j and isinstance(j, list) and j[0] else [] @@ -118,22 +131,24 @@ def vep_annotate(df, reference, gtf_path=None): }) return pd.DataFrame(rows) -def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors"): +def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors", offline=False): """Batch query OncoKB byHgvsVariant. Returns dict hgvs_g -> {amp, oncogenic}.""" - if not token: + if not token or offline: + return {} + if len(hgvs_g_list) > 500: + print(f" [oncokb] {len(hgvs_g_list)} variants - skipping OncoKB (use --offline to silence, or annotate subset)", flush=True) return {} headers = {"Authorization": f"Bearer {token}"} out = {} - for hgvs in hgvs_g_list: + for i, hgvs in enumerate(hgvs_g_list): + if i % 100 == 0 and len(hgvs_g_list) > 100: + print(f" [oncokb] {i}/{len(hgvs_g_list)} ...", flush=True) try: - # OncoKB expects e.g. 7:g.140453136A>T params = {"hgvsg": hgvs, "tumorType": tumor_type} - resp = requests.get(ONCOKB_URL, headers=headers, params=params, timeout=10) + resp = requests.get(ONCOKB_URL, headers=headers, params=params, timeout=(3, 5)) if resp.status_code == 200: j = resp.json() - # j contains levelOfEvidence, oncogenic, etc. amp = j.get("highestSensitiveLevel") or j.get("levelAssociated") or "" - # Map OncoKB levels to AMP Tier level_map = {"LEVEL_1": "Tier I", "LEVEL_2": "Tier II", "LEVEL_3A": "Tier II", "LEVEL_3B": "Tier III", "LEVEL_4": "Tier IV", "LEVEL_R1": "Tier I", "LEVEL_R2": "Tier II"} amp_tier = level_map.get(amp, amp) oncogenic = j.get("oncogenic", "") @@ -142,7 +157,7 @@ def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors"): out[hgvs] = {"AMP": "", "ONCO": ""} except Exception: out[hgvs] = {"AMP": "", "ONCO": ""} - time.sleep(0.2) + time.sleep(0.05) return out def fetch_civic(): @@ -251,6 +266,7 @@ def main(): ap.add_argument("--out", default=None, help="output xlsx (default .annotated.xlsx)") ap.add_argument("--tumor-type", default="All Solid Tumors", help="OncoKB tumor type (pan-cancer default)") ap.add_argument("--token", default="~/.config/oncokb/token", help="OncoKB token file or env") + ap.add_argument("--offline", action="store_true", help="skip Ensembl/OncoKB network calls (fast, offline)") args = ap.parse_args() clean_path = Path(args.clean) @@ -258,13 +274,15 @@ def main(): if df_clean.empty: print("clean_variants.csv is empty (or only header) - nothing to annotate") out = args.out or str(clean_path).replace(".csv", ".annotated.xlsx") - # Still create empty Excel with headers df_empty = pd.DataFrame(columns=["Ген","HGVS_c","HGVS_p","HGVS","Тип варианта и эффект","_orig_idx"]) build_excel(df_clean, df_empty, out) return + if len(df_clean) > 500 and not args.offline: + print(f"[warn] {len(df_clean)} variants - network annotation will be slow. Use --offline for fast placeholder.", flush=True) + # HGVS / gene / effect - df_annot = vep_annotate(df_clean, args.reference) + df_annot = vep_annotate(df_clean, args.reference, offline=args.offline) # OncoKB token = load_token(args.token) @@ -273,7 +291,7 @@ def main(): hgvs_g_for_oncokb = [] for _, r in df_clean.iterrows(): hgvs_g_for_oncokb.append(f"{str(r['chrom']).replace('chr','')}:g.{int(r['position'])+1}{r['ref']}>{r['alt']}") - oncokb_map = fetch_oncokb(hgvs_g_for_oncokb, token, args.tumor_type) if token else {} + oncokb_map = fetch_oncokb(hgvs_g_for_oncokb, token, args.tumor_type, offline=args.offline) if token else {} # Attach ACMG/AMP/ONCO to annot rows acmg_list = annotate_acmg(df_clean)