Skills › Scientific Agent Skills › alphagenome › atlas_query.py
atlas_query.py
536 lines · skills/alphagenome/scripts/atlas_query.py at 5f2bc03
1#!/usr/bin/env python3
2# /// script
3# requires-python = ">=3.10"
4# dependencies = ["alphagenome>=0.9.0"]
5# ///
6"""Query the AlphaGenome Atlas: precomputed variant effects for every hg38 SNV.
8Four subcommands, all against the ``alphagenome.atlas`` gRPC client:
10 avi AlphaGenome Variant Impact (AVI) score, Phred, and the 18 SHAP
11 feature attributions for variants or for every SNV in a window.
12 scores Raw and quantile scores from the track-level scorers (RNA_SEQ,
13 DNASE, CHIP_TF, SPLICE_SITE_USAGE, ...) as one tidy row per
14 variant x track (x gene, for gene-centric scorers).
15 scorers The scorers the server currently serves, with track counts.
16 tracks The track catalogue behind a scorer (biosample, ontology CURIE,
17 assay, TF, histone mark) for choosing ontology filters.
19Coordinates: variants are 1-based ``chr:pos:ref>alt``; ``--interval`` is a
201-based closed ``chr:start-end``. The Atlas is GRCh38 (hg38) only and the REF
21allele must match the reference: swapped REF/ALT can cause a lookup miss.
22Returned variant identities are checked against the request. Reads the API
23key from ALPHAGENOME_API_KEY (or ALPHA_GENOME_API_KEY).
25Examples:
26 python atlas_query.py avi --variant "chr22:36201698:A>C" "chr9:128225994:G>A"
27 python atlas_query.py avi --input candidates.vcf --min-phred 20 -o avi.tsv
28 python atlas_query.py avi --interval chr11:5225727-5226575 --top-k 25
29 python atlas_query.py scores --variant "chr22:36201698:A>C" \\
30 --scorers RNA_SEQ SPLICE_SITE_USAGE --ontology UBERON:0001157 -o scores.tsv
31 python atlas_query.py scorers
32 python atlas_query.py tracks --scorer CHIP_TF --query GATA1
33"""
35from __future__ import annotations
37import argparse
38import concurrent.futures
39import math
40import sys
41from collections.abc import Mapping, Sequence
42from typing import Any
44import _common as common
46DEFAULT_MAX_WINDOW_BP = 1_000
47DEFAULT_WORKERS = 8
48TRACK_METADATA_COLUMNS = (
49 "name",
50 "strand",
51 "ontology_curie",
52 "biosample_name",
53 "biosample_type",
54 "Assay title",
55 "assay_title",
56 "gtex_tissue",
57 "transcription_factor",
58 "histone_mark",
59 "data_source",
60)
63# --------------------------------------------------------------------------
64# pure conversion helpers (unit-tested with synthetic AnnData)
65# --------------------------------------------------------------------------
68def _obs_variant_strings(adata) -> list[str]:
69 if adata.obs is None or "variant" not in adata.obs.columns:
70 return [str(index) for index in adata.obs_names]
71 return [str(value) for value in adata.obs["variant"].tolist()]
74def _feature_names(adata) -> list[str]:
75 if adata.var is not None and "name" in adata.var.columns:
76 return [str(value) for value in adata.var["name"].tolist()]
77 return [str(value) for value in adata.var_names]
80def _layer(adata, name: str):
81 try:
82 return adata.layers[name] if name in adata.layers else None
83 except (KeyError, TypeError):
84 return None
87def summarize_avi(avi_adata, fi_adata=None) -> list[dict[str, Any]]:
88 """One row per variant from the ``AVI_SCORE`` and feature-importance AnnData.
90 Columns: variant, avi_raw, avi_cdf_quantile, avi_tail_quantile, avi_phred,
91 avi_top_percent, top_feature_key, top_feature, top_feature_value, and one
92 ``fi_<KEY>`` column per attribution. Missing attributions leave NaN.
93 """
94 import numpy as np # noqa: PLC0415
96 rows: list[dict[str, Any]] = []
97 if avi_adata is None or avi_adata.X is None or avi_adata.X.shape[0] == 0:
98 return rows
100 variants = _obs_variant_strings(avi_adata)
101 raw = np.asarray(avi_adata.X, dtype=float).reshape(len(variants), -1)
102 quantiles = _layer(avi_adata, "quantiles")
103 if quantiles is not None:
104 quantiles = np.asarray(quantiles, dtype=float).reshape(len(variants), -1)
106 fi_by_variant: dict[str, np.ndarray] = {}
107 fi_names: list[str] = []
108 if fi_adata is not None and fi_adata.X is not None and fi_adata.X.shape[0] > 0:
109 fi_names = _feature_names(fi_adata)
110 fi_matrix = np.asarray(fi_adata.X, dtype=float).reshape(fi_adata.X.shape[0], -1)
111 for variant, values in zip(_obs_variant_strings(fi_adata), fi_matrix, strict=True):
112 fi_by_variant[variant] = values
114 for index, variant in enumerate(variants):
115 row: dict[str, Any] = {"variant": variant}
116 row["avi_raw"] = float(raw[index, 0]) if raw.shape[1] else float("nan")
117 if quantiles is not None and quantiles.shape[1]:
118 cdf = float(quantiles[index, 0])
119 tail, phred = common.cdf_to_tail_and_phred(cdf)
120 else:
121 cdf, tail, phred = float("nan"), float("nan"), float("nan")
122 row["avi_cdf_quantile"] = cdf
123 row["avi_tail_quantile"] = tail
124 row["avi_phred"] = phred
125 row["avi_top_percent"] = common.phred_to_top_percent(phred)
127 values = fi_by_variant.get(variant)
128 if values is not None and len(fi_names) == len(values):
129 top = int(np.argmax(np.abs(values)))
130 row["top_feature_key"] = fi_names[top]
131 row["top_feature"] = common.feature_display_name(fi_names[top])
132 row["top_feature_value"] = float(values[top])
133 for name, value in zip(fi_names, values, strict=True):
134 row[f"fi_{name}"] = float(value)
135 else:
136 row["top_feature_key"] = ""
137 row["top_feature"] = ""
138 row["top_feature_value"] = float("nan")
139 row["atlas_url"] = common.portal_variant_url(variant)
140 rows.append(row)
141 return rows
144def tidy_atlas_scores(scores: Mapping[str, Any]) -> list[dict[str, Any]]:
145 """Flatten ``{scorer: AnnData}`` into one row per (obs row, track).
147 ``obs`` rows are variants, or variant x gene for the gene-centric scorers
148 (RNA_SEQ, POLYADENYLATION, SPLICE_*); ``var`` rows are tracks. Raw scores
149 come from ``X`` and, when the server sent them, quantile scores from the
150 ``quantiles`` layer.
151 """
152 import numpy as np # noqa: PLC0415
154 rows: list[dict[str, Any]] = []
155 for scorer, adata in scores.items():
156 if adata is None or adata.X is None or 0 in adata.X.shape:
157 continue
158 matrix = np.asarray(adata.X, dtype=float)
159 quantiles = _layer(adata, "quantiles")
160 if quantiles is not None:
161 quantiles = np.asarray(quantiles, dtype=float)
162 obs = adata.obs if adata.obs is not None else None
163 var = adata.var if adata.var is not None else None
164 var_columns = [column for column in TRACK_METADATA_COLUMNS if var is not None and column in var.columns]
165 obs_columns = [
166 column
167 for column in ("gene_id", "gene_name", "strand", "junction_Start", "junction_End")
168 if obs is not None and column in obs.columns
169 ]
170 variants = _obs_variant_strings(adata)
171 for i in range(matrix.shape[0]):
172 base: dict[str, Any] = {"variant": variants[i], "scorer": scorer}
173 for column in obs_columns:
174 key = {"strand": "gene_strand", "junction_Start": "junction_start", "junction_End": "junction_end"}.get(
175 column, column
176 )
177 base[key] = obs.iloc[i][column]
178 for j in range(matrix.shape[1]):
179 row = dict(base)
180 row["track_index"] = j
181 for column in var_columns:
182 key = "assay_title" if column == "Assay title" else column
183 row[f"track_{key}" if key in {"name", "strand"} else key] = var.iloc[j][column]
184 row["raw_score"] = float(matrix[i, j])
185 if quantiles is not None:
186 row["quantile_score"] = float(quantiles[i, j])
187 rows.append(row)
188 return rows
191def max_abs_track(adata) -> dict[str, Any] | None:
192 """The single largest |score| entry of a track-scorer AnnData, with its metadata."""
193 import numpy as np # noqa: PLC0415
195 if adata is None or adata.X is None or 0 in adata.X.shape:
196 return None
197 matrix = np.asarray(adata.X, dtype=float)
198 flat = int(np.nanargmax(np.abs(matrix)))
199 i, j = divmod(flat, matrix.shape[1])
200 result: dict[str, Any] = {"raw_score": float(matrix[i, j]), "track_index": j}
201 if adata.var is not None:
202 for column in ("name", "biosample_name", "ontology_curie", "transcription_factor", "histone_mark"):
203 if column in adata.var.columns:
204 result[column] = adata.var.iloc[j][column]
205 if adata.obs is not None:
206 for column in ("gene_name", "gene_id"):
207 if column in adata.obs.columns:
208 result[column] = adata.obs.iloc[i][column]
209 return result
212def scorer_rows(metadata: Mapping[str, Any]) -> list[dict[str, Any]]:
213 rows = []
214 for name, meta in metadata.items():
215 frame = getattr(meta, "track_metadata", None)
216 row: dict[str, Any] = {
217 "scorer": name,
218 "is_signed": bool(getattr(meta, "is_signed", False)),
219 "n_tracks": int(len(frame)) if frame is not None else 0,
220 }
221 if frame is not None and "biosample_name" in frame.columns:
222 row["n_biosamples"] = int(frame["biosample_name"].nunique())
223 if frame is not None and "ontology_curie" in frame.columns:
224 row["n_ontology_terms"] = int(frame["ontology_curie"].nunique())
225 rows.append(row)
226 return sorted(rows, key=lambda item: item["scorer"])
229def track_rows(metadata: Mapping[str, Any], scorer: str | None, query: str | None) -> list[dict[str, Any]]:
230 rows: list[dict[str, Any]] = []
231 needle = query.lower() if query else None
232 for name, meta in metadata.items():
233 if scorer and name != scorer:
234 continue
235 frame = getattr(meta, "track_metadata", None)
236 if frame is None or frame.empty:
237 continue
238 for index, (_, record) in enumerate(frame.iterrows()):
239 row: dict[str, Any] = {"scorer": name, "track_index": index}
240 for column in TRACK_METADATA_COLUMNS:
241 if column in frame.columns:
242 key = "assay_title" if column == "Assay title" else column
243 row[key] = record[column]
244 if needle and not any(needle in str(value).lower() for value in row.values()):
245 continue
246 rows.append(row)
247 return rows
250# --------------------------------------------------------------------------
251# client plumbing
252# --------------------------------------------------------------------------
255def make_client(args: argparse.Namespace):
256 api_key = common.load_api_key(args.api_key_env)
257 common.require_alphagenome()
258 import grpc # noqa: PLC0415
259 from alphagenome.atlas import atlas # noqa: PLC0415
261 try:
262 return atlas.create(api_key, timeout=args.timeout)
263 except grpc.FutureTimeoutError as error: # pragma: no cover - network
264 raise SystemExit(
265 f"could not reach the Atlas service within {args.timeout}s "
266 "(gdmscience.googleapis.com:443); check network access and proxies"
267 ) from error
270def to_genome_variant(spec: common.VariantSpec):
271 from alphagenome.data import genome # noqa: PLC0415
273 return genome.Variant(
274 chromosome=spec.chromosome,
275 position=spec.position,
276 reference_bases=spec.ref,
277 alternate_bases=spec.alt,
278 name=spec.name,
279 )
282def to_genome_interval(text: str, max_window: int):
283 from alphagenome.data import genome # noqa: PLC0415
285 chromosome, start0, end = common.parse_interval_string(text)
286 width = common.interval_width(start0, end)
287 if width > max_window:
288 raise SystemExit(
289 f"interval is {width:,} bp, above --max-window {max_window:,} bp. Each base "
290 f"expands to {common.SNVS_PER_BASE} variants; raise --max-window deliberately."
291 )
292 return genome.Interval(chromosome=chromosome, start=start0, end=end)
295def _query_each(
296 client,
297 specs: Sequence[common.VariantSpec],
298 requested_scorers: Sequence[str],
299 workers: int,
300 **filters: Any,
301) -> list[tuple[common.VariantSpec, Mapping[str, Any] | None, str]]:
302 """Query variants one by one, keeping per-variant failures as messages."""
304 def one(spec: common.VariantSpec):
305 try:
306 result = client.query_variant(to_genome_variant(spec), requested_scorers=list(requested_scorers), **filters)
307 for scorer, adata in result.items():
308 if adata.obs is not None and "variant" in adata.obs:
309 returned = set(_obs_variant_strings(adata))
310 if returned - {str(spec)}:
311 raise ValueError(
312 f"{scorer} returned a different variant: {sorted(returned)}; "
313 f"requested {spec}. Check REF and genome build."
314 )
315 return spec, result, ""
316 except Exception as error: # noqa: BLE001 - surfaced per variant
317 return spec, None, f"{type(error).__name__}: {error}"
319 if workers <= 1 or len(specs) <= 1:
320 return [one(spec) for spec in specs]
321 with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
322 return list(pool.map(one, specs))
325# --------------------------------------------------------------------------
326# subcommands
327# --------------------------------------------------------------------------
330def cmd_avi(args: argparse.Namespace) -> int:
331 client = make_client(args)
332 rows: list[dict[str, Any]] = []
334 if args.interval:
335 interval = to_genome_interval(args.interval, args.max_window)
336 scores = client.query_interval(interval, requested_scorers=list(common.AVI_SCORERS), progress_bar=False)
337 rows = summarize_avi(scores.get(common.AVI_SCORER), scores.get(common.AVI_FEATURES_SCORER))
338 for row in rows:
339 row["error"] = ""
340 else:
341 specs, warnings = common.collect_variants(args.variant, args.input)
342 for message in warnings:
343 common.warn(message)
344 if not specs:
345 raise SystemExit("no variants given; use --variant, --input, or --interval")
346 indels = [spec for spec in specs if not spec.is_snv]
347 if indels:
348 common.warn(
349 f"{len(indels)} indel(s) requested; the public Atlas API currently serves "
350 "genome-wide SNVs, so indels may come back as lookup misses"
351 )
352 for spec, scores, error in _query_each(client, specs, common.AVI_SCORERS, args.workers):
353 if scores is None:
354 row = spec.to_row()
355 row.update({"avi_raw": float("nan"), "avi_phred": float("nan"), "error": error})
356 rows.append(row)
357 continue
358 summary = summarize_avi(scores.get(common.AVI_SCORER), scores.get(common.AVI_FEATURES_SCORER))
359 if not summary:
360 row = spec.to_row()
361 row.update({"avi_raw": float("nan"), "avi_phred": float("nan"), "error": "no AVI_SCORE returned"})
362 rows.append(row)
363 continue
364 row = spec.to_row()
365 row.update(summary[0])
366 row["variant"] = str(spec)
367 row["error"] = ""
368 rows.append(row)
370 if args.with_tracks:
371 for row in rows:
372 key = row.get("top_feature_key") or ""
373 scorers = common.AVI_FEATURES.get(key, ("", "", ()))[2]
374 if not scorers or row.get("error"):
375 continue
376 try:
377 spec = common.parse_variant_string(row["variant"])
378 detail = client.query_variant(to_genome_variant(spec), requested_scorers=list(scorers))
379 except Exception as error: # noqa: BLE001
380 row["top_track_error"] = f"{type(error).__name__}: {error}"
381 continue
382 best = None
383 for scorer in scorers:
384 hit = max_abs_track(detail.get(scorer))
385 if hit and (best is None or abs(hit["raw_score"]) > abs(best[1]["raw_score"])):
386 best = (scorer, hit)
387 if best:
388 scorer, hit = best
389 row["top_track_scorer"] = scorer
390 row["top_track_name"] = hit.get("name", "")
391 row["top_track_biosample"] = hit.get("biosample_name", "")
392 row["top_track_ontology"] = hit.get("ontology_curie", "")
393 row["top_track_gene"] = hit.get("gene_name", "")
394 row["top_track_raw_score"] = hit["raw_score"]
396 if args.min_phred is not None:
397 rows = [row for row in rows if row.get("error") or (isinstance(row.get("avi_phred"), float) and row["avi_phred"] >= args.min_phred)]
399 rows.sort(key=lambda row: -(row.get("avi_phred") if isinstance(row.get("avi_phred"), float) and math.isfinite(row["avi_phred"]) else -1))
400 if args.top_k:
401 rows = rows[: args.top_k]
403 if not args.output and len(rows) > 50 and not args.force_stdout:
404 raise SystemExit(f"{len(rows)} rows; write them with -o FILE (or pass --force-stdout)")
405 common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
406 return 0
409def cmd_scores(args: argparse.Namespace) -> int:
410 client = make_client(args)
411 requested = list(args.scorers)
412 filters = {
413 "ontology_terms": args.ontology or None,
414 "gene_names": args.gene or None,
415 "gene_ids": args.gene_id or None,
416 }
417 rows: list[dict[str, Any]] = []
418 if args.interval:
419 interval = to_genome_interval(args.interval, args.max_window)
420 scores = client.query_interval(interval, requested_scorers=requested, progress_bar=False, **filters)
421 rows = tidy_atlas_scores(scores)
422 else:
423 specs, warnings = common.collect_variants(args.variant, args.input)
424 for message in warnings:
425 common.warn(message)
426 if not specs:
427 raise SystemExit("no variants given; use --variant, --input, or --interval")
428 for spec, scores, error in _query_each(client, specs, requested, args.workers, **filters):
429 if scores is None:
430 common.warn(f"{spec}: {error}")
431 rows.append({"variant": str(spec), "scorer": "", "error": error})
432 continue
433 rows.extend(tidy_atlas_scores(scores))
435 if args.min_abs_quantile is not None:
436 calibrated = [row for row in rows if "quantile_score" in row]
437 if not calibrated and any(not row.get("error") for row in rows):
438 raise SystemExit("no calibrated quantiles returned; omit --min-abs-quantile and inspect raw scores")
439 rows = [
440 row
441 for row in rows
442 if row.get("error") or (
443 "quantile_score" in row and abs(row["quantile_score"]) >= args.min_abs_quantile
444 )
445 ]
446 if not args.output and len(rows) > 200 and not args.force_stdout:
447 raise SystemExit(f"{len(rows)} rows; write them with -o FILE (or pass --force-stdout)")
448 common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
449 return 0
452def cmd_scorers(args: argparse.Namespace) -> int:
453 client = make_client(args)
454 rows = scorer_rows(client.scorer_metadata())
455 common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
456 return 0
459def cmd_tracks(args: argparse.Namespace) -> int:
460 client = make_client(args)
461 rows = track_rows(client.scorer_metadata(), args.scorer, args.query)
462 if not args.output and len(rows) > 200 and not args.force_stdout:
463 raise SystemExit(
464 f"{len(rows)} tracks match; narrow with --scorer/--query or write them with -o FILE"
465 )
466 common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
467 return 0
470# --------------------------------------------------------------------------
471# CLI
472# --------------------------------------------------------------------------
475def _add_common(parser: argparse.ArgumentParser) -> None:
476 parser.add_argument("-o", "--output", help="write here instead of stdout (extension picks the format)")
477 parser.add_argument("--format", choices=("tsv", "csv", "json", "parquet"), help="override the output format")
478 parser.add_argument("--api-key-env", metavar="NAME", help="environment variable holding the key (default: ALPHAGENOME_API_KEY, then ALPHA_GENOME_API_KEY)")
479 parser.add_argument("--timeout", type=float, default=30.0, help="seconds to wait for the gRPC channel (default 30)")
480 parser.add_argument("--force-stdout", action="store_true", help="print large tables to stdout anyway")
483def _add_variant_inputs(parser: argparse.ArgumentParser) -> None:
484 parser.add_argument("--variant", nargs="+", metavar="CHR:POS:REF>ALT", help="one or more variants (1-based; gnomAD/GTEx spellings accepted)")
485 parser.add_argument("--input", metavar="FILE", help="VCF, or TSV/CSV with CHROM/POS/REF/ALT or a 'variant' column")
486 parser.add_argument("--interval", metavar="CHR:START-END", help="1-based closed window: every SNV in it (3 per base)")
487 parser.add_argument("--max-window", type=int, default=DEFAULT_MAX_WINDOW_BP, help=f"refuse --interval wider than this (default {DEFAULT_MAX_WINDOW_BP} bp)")
488 parser.add_argument("--workers", type=int, default=DEFAULT_WORKERS, help=f"parallel variant queries (default {DEFAULT_WORKERS})")
491def build_parser() -> argparse.ArgumentParser:
492 parser = argparse.ArgumentParser(
493 prog="atlas_query.py",
494 description="Query the AlphaGenome Atlas (precomputed variant effects, hg38).",
495 epilog=__doc__.split("Examples:", 1)[-1] if "Examples:" in __doc__ else None,
496 formatter_class=argparse.RawDescriptionHelpFormatter,
497 )
498 sub = parser.add_subparsers(dest="command", required=True)
500 avi = sub.add_parser("avi", help="AVI score, Phred, and feature attributions")
501 _add_variant_inputs(avi)
502 avi.add_argument("--min-phred", type=float, help="keep variants with AVI Phred >= this (20 = top 1%%)")
503 avi.add_argument("--top-k", type=int, help="keep the K highest-Phred variants")
504 avi.add_argument("--with-tracks", action="store_true", help="also report the strongest track behind the top feature (one extra query per variant)")
505 _add_common(avi)
506 avi.set_defaults(func=cmd_avi)
508 scores = sub.add_parser("scores", help="raw + quantile scores from the track-level scorers")
509 _add_variant_inputs(scores)
510 scores.add_argument("--scorers", nargs="+", default=["RNA_SEQ"], metavar="NAME", help="Atlas scorer names (see `scorers`); default RNA_SEQ")
511 scores.add_argument("--ontology", nargs="+", metavar="CURIE", help="keep tracks for these ontology terms, e.g. UBERON:0001157 CL:0000084")
512 scores.add_argument("--gene", nargs="+", metavar="SYMBOL", help="gene-centric scorers: keep these gene symbols")
513 scores.add_argument("--gene-id", nargs="+", metavar="ENSG", help="gene-centric scorers: keep these Ensembl gene IDs")
514 scores.add_argument("--min-abs-quantile", type=common.unit_interval, help="keep |quantile_score| >= threshold in [0,1]; signed scorers use [-1,1], unsigned use [0,1]")
515 _add_common(scores)
516 scores.set_defaults(func=cmd_scores)
518 scorers = sub.add_parser("scorers", help="list the scorers the Atlas currently serves")
519 _add_common(scorers)
520 scorers.set_defaults(func=cmd_scorers)
522 tracks = sub.add_parser("tracks", help="track catalogue (biosample, ontology CURIE, assay) per scorer")
523 tracks.add_argument("--scorer", metavar="NAME", help="restrict to one scorer")
524 tracks.add_argument("--query", metavar="TEXT", help="case-insensitive substring over every column")
525 _add_common(tracks)
526 tracks.set_defaults(func=cmd_tracks)
527 return parser
530def main(argv: Sequence[str] | None = None) -> int:
531 args = build_parser().parse_args(argv)
532 return args.func(args)
535if __name__ == "__main__":
536 sys.exit(main())