Skills › Scientific Agent Skills › alphagenome › _common.py
_common.py
524 lines · skills/alphagenome/scripts/_common.py at 5f2bc03
1"""Shared helpers for the alphagenome skill scripts.
3Standard library only, so variant parsing, interval arithmetic, API-key
4discovery, Phred conversion, and table export all work (and are testable)
5without the ``alphagenome`` package installed. The scripts import the SDK
6lazily through :func:`require_alphagenome`.
7"""
9from __future__ import annotations
11import csv
12import io
13import json
14import math
15import os
16import re
17import sys
18from collections.abc import Iterable, Mapping, Sequence
19from dataclasses import dataclass, field
20from pathlib import Path
22# --------------------------------------------------------------------------
23# constants
24# --------------------------------------------------------------------------
26#: Environment variables checked, in order, for the AlphaGenome API key.
27#: The first spelling is what DeepMind's own agent skills use; the second is
28#: what ``alphagenome.colab_utils.get_api_key`` reads by default.
29API_KEY_ENV_VARS: tuple[str, ...] = ("ALPHAGENOME_API_KEY", "ALPHA_GENOME_API_KEY")
31GET_KEY_URL = "https://deepmind.google.com/science/alphagenome"
32ATLAS_BASE_URL = "https://deepmind.google.com/science/alphagenome/atlas"
33ATLAS_TRACK_PREDICTIONS_URL = f"{ATLAS_BASE_URL}/track-predictions"
34DOCS_URL = "https://www.alphagenomedocs.com/"
36#: GENCODE v46 annotation used by the Atlas and the model, as a Feather file
37#: (about 318 MB; needs roughly 4 GB of RAM to load with pandas).
38GTF_FEATHER_URL = (
39 "https://storage.googleapis.com/alphagenome/reference/gencode/hg38/"
40 "gencode.v46.annotation.gtf.gz.feather"
41)
43#: Atlas scorer names for the composite AlphaGenome Variant Impact score and
44#: its 18-way SHAP feature attribution.
45AVI_SCORER = "AVI_SCORE"
46AVI_FEATURES_SCORER = "AVI_SCORE_FEATURE_IMPORTANCE"
47AVI_SCORERS: tuple[str, str] = (AVI_SCORER, AVI_FEATURES_SCORER)
49#: The 18 AVI input features, keyed by the ``var['name']`` values returned by
50#: the ``AVI_SCORE_FEATURE_IMPORTANCE`` scorer. Value: (display name, category,
51#: Atlas track scorers that feed the feature; empty for annotation features).
52AVI_FEATURES: dict[str, tuple[str, str, tuple[str, ...]]] = {
53 "MERGED_SPLICING": (
54 "Splicing",
55 "Splicing",
56 ("SPLICE_SITES", "SPLICE_SITE_USAGE", "SPLICE_JUNCTIONS"),
57 ),
58 "MAX_ABS_RNA_SEQ": ("RNA-seq", "Transcription", ("RNA_SEQ",)),
59 "MAX_ABS_ATAC": ("ATAC-seq", "Chromatin accessibility", ("ATAC",)),
60 "MAX_ABS_DNASE": ("DNase-seq", "Chromatin accessibility", ("DNASE",)),
61 "MAX_ABS_CHIP_TF": ("ChIP-TF", "Transcription factor binding", ("CHIP_TF",)),
62 "MAX_ABS_CHIP_HISTONE": ("ChIP-Histone", "Histone modification", ("CHIP_HISTONE",)),
63 "MAX_ABS_CAGE": ("CAGE", "Transcription initiation", ("CAGE",)),
64 "MAX_ABS_PROCAP": ("PRO-cap", "Transcription initiation", ("PROCAP",)),
65 "MAX_ABS_POLYADENYLATION": ("Polyadenylation", "Transcription", ("POLYADENYLATION",)),
66 "MAX_ABS_CONTACT_MAPS": ("3D genome contacts", "3D genome organization", ("CONTACT_MAPS",)),
67 "ALPHAMISSENSE": ("AlphaMissense", "Protein impact", ()),
68 "CACTUS_241_WAY": ("Zoonomia Cactus 241-way", "Evolutionary conservation", ()),
69 "PHASTCONS_470_WAY": ("PhastCons 470-way", "Evolutionary conservation", ()),
70 "PROTEIN_TERMINATION": ("Protein termination", "Coding consequence", ()),
71 "START_LOST": ("Start lost", "Coding consequence", ()),
72 "STOP_LOST": ("Stop lost", "Coding consequence", ()),
73 "IS_INSERTION": ("Insertion", "Indel type", ()),
74 "IS_DELETION": ("Deletion", "Indel type", ()),
75}
77#: Atlas track-level scorers (one AnnData each; tracks on the ``var`` axis).
78#: ``scorer_metadata()`` is the authority on what the server currently serves.
79ATLAS_TRACK_SCORERS: tuple[str, ...] = (
80 "ATAC",
81 "DNASE",
82 "CHIP_TF",
83 "CHIP_HISTONE",
84 "CAGE",
85 "PROCAP",
86 "RNA_SEQ",
87 "POLYADENYLATION",
88 "SPLICE_SITES",
89 "SPLICE_SITE_USAGE",
90 "SPLICE_JUNCTIONS",
91 "CONTACT_MAPS",
92)
94#: Per-position substitution count the Atlas stores (three alternates per base).
95SNVS_PER_BASE = 3
97INSTALL_HINT = (
98 "alphagenome is not installed. Run: uv pip install alphagenome "
99 "(or invoke with: uv run --with alphagenome <script>)"
100)
102# --------------------------------------------------------------------------
103# variants and intervals
104# --------------------------------------------------------------------------
106_VALID_BASES = frozenset("ACGTN")
109@dataclass(frozen=True)
110class VariantSpec:
111 """A variant in VCF terms: 1-based position, explicit REF and ALT."""
113 chromosome: str
114 position: int
115 ref: str
116 alt: str
117 name: str = field(default="", compare=False)
119 def __str__(self) -> str:
120 return f"{self.chromosome}:{self.position}:{self.ref}>{self.alt}"
122 @property
123 def is_snv(self) -> bool:
124 return len(self.ref) == 1 and len(self.alt) == 1
126 def to_row(self) -> dict[str, object]:
127 return {
128 "variant": str(self),
129 "chromosome": self.chromosome,
130 "position": self.position,
131 "ref": self.ref,
132 "alt": self.alt,
133 "name": self.name,
134 }
137def normalize_chromosome(contig: str) -> str:
138 """Return the ``chr``-prefixed spelling the Atlas and model expect.
140 ``MT`` and ``M`` become ``chrM``; ``chr`` is added when missing. Nothing
141 else is rewritten: an unplaced contig stays as given.
142 """
143 text = str(contig).strip()
144 if not text:
145 raise ValueError("empty chromosome name")
146 if text.lower().startswith("chr"):
147 text = text[3:]
148 if text.upper() in {"MT", "M"}:
149 return "chrM"
150 return f"chr{text}"
153# Accepted spellings. Group order is always chrom, pos, ref, alt.
154_VARIANT_PATTERNS: tuple[tuple[str, re.Pattern[str]], ...] = (
155 # chr22:1024:A>C (AlphaGenome default) and chr22:1024:A:C (Open Targets BigQuery)
156 ("default", re.compile(r"^(\w+):(\d+):([ACGTNacgtn]+)[>:]([ACGTNacgtn]+)$")),
157 # 22-1024-A-C (gnomAD)
158 ("gnomad", re.compile(r"^(\w+)-(\d+)-([ACGTNacgtn]+)-([ACGTNacgtn]+)$")),
159 # chr22_1024_A_C_b38 (GTEx, build suffix optional) and 22_1024_A_C (Open Targets)
160 ("gtex", re.compile(r"^(\w+?)_(\d+)_([ACGTNacgtn]+)_([ACGTNacgtn]+)(?:_b38)?$")),
161)
164def parse_variant_string(text: str, name: str = "") -> VariantSpec:
165 """Parse ``chr:pos:ref>alt`` and the gnomAD, GTEx, and Open Targets spellings.
167 Positions are 1-based, as in VCF and on the Atlas website. Bases are
168 upper-cased. Raises ``ValueError`` on anything else, including rsIDs,
169 which neither the Atlas API nor the portal accept.
170 """
171 candidate = text.strip()
172 for _, pattern in _VARIANT_PATTERNS:
173 match = pattern.match(candidate)
174 if match:
175 chrom, pos, ref, alt = match.groups()
176 return VariantSpec(
177 chromosome=normalize_chromosome(chrom),
178 position=int(pos),
179 ref=ref.upper(),
180 alt=alt.upper(),
181 name=name,
182 )
183 raise ValueError(
184 f"cannot parse variant {text!r}; use chr:pos:ref>alt with a 1-based "
185 "position (also accepted: 22-1024-A-C, chr22_1024_A_C_b38, 22:1024:A:C). "
186 "rsIDs are not supported - resolve them to coordinates first."
187 )
190def parse_interval_string(text: str) -> tuple[str, int, int]:
191 """Parse a 1-based closed ``chr:start-end`` into 0-based half-open bounds.
193 Returns ``(chromosome, start0, end)`` ready for ``genome.Interval``.
194 Thousands separators are tolerated.
195 """
196 candidate = text.strip().replace(",", "")
197 match = re.match(r"^(\w+):(\d+)-(\d+)$", candidate)
198 if not match:
199 raise ValueError(f"cannot parse interval {text!r}; expected chr:start-end (1-based, closed)")
200 chrom, start, end = match.group(1), int(match.group(2)), int(match.group(3))
201 if start < 1:
202 raise ValueError(f"interval start must be >= 1 (got {start}); coordinates are 1-based")
203 if end < start:
204 raise ValueError(f"interval end {end} is before start {start}")
205 return normalize_chromosome(chrom), start - 1, end
208def interval_width(start0: int, end: int) -> int:
209 return end - start0
212def _split_delimited(path: Path) -> tuple[list[str], list[list[str]]]:
213 text = path.read_text(encoding="utf-8")
214 sample = text[:4096]
215 if path.suffix.lower() == ".csv":
216 delimiter = ","
217 elif "\t" in sample:
218 delimiter = "\t"
219 else:
220 try:
221 delimiter = csv.Sniffer().sniff(sample, delimiters=",\t;").delimiter
222 except csv.Error:
223 delimiter = ","
224 reader = csv.reader(io.StringIO(text), delimiter=delimiter)
225 rows = [row for row in reader if row and any(cell.strip() for cell in row)]
226 if not rows:
227 return [], []
228 header = [cell.strip() for cell in rows[0]]
229 return header, rows[1:]
232_SYMBOLIC_ALT = re.compile(r"^(<.*>|\*|\.)$")
235def read_vcf(path: Path) -> tuple[list[VariantSpec], list[str]]:
236 """Read the CHROM/POS/ID/REF/ALT columns of a VCF, splitting multi-allelic ALTs.
238 Returns ``(variants, warnings)``. Symbolic alleles (``<DEL>``, ``*``, ``.``)
239 and breakends are skipped with a warning; the Atlas has no entry for them.
240 """
241 variants: list[VariantSpec] = []
242 warnings: list[str] = []
243 with path.open(encoding="utf-8") as handle:
244 for line_number, line in enumerate(handle, start=1):
245 if not line.strip() or line.startswith("#"):
246 continue
247 fields = line.rstrip("\n").split("\t")
248 if len(fields) < 5:
249 warnings.append(f"line {line_number}: fewer than 5 columns, skipped")
250 continue
251 chrom, pos, identifier, ref, alts = fields[:5]
252 for alt in alts.split(","):
253 if _SYMBOLIC_ALT.match(alt) or "[" in alt or "]" in alt:
254 warnings.append(f"line {line_number}: symbolic allele {alt!r} skipped")
255 continue
256 if not set(ref.upper()) <= _VALID_BASES or not set(alt.upper()) <= _VALID_BASES:
257 warnings.append(f"line {line_number}: non-ACGTN allele {ref}>{alt} skipped")
258 continue
259 variants.append(
260 VariantSpec(
261 chromosome=normalize_chromosome(chrom),
262 position=int(pos),
263 ref=ref.upper(),
264 alt=alt.upper(),
265 name="" if identifier in {".", ""} else identifier,
266 )
267 )
268 return variants, warnings
271def read_variant_table(path: str | Path) -> tuple[list[VariantSpec], list[str]]:
272 """Read variants from a VCF or a delimited table.
274 Delimited files need either a ``variant`` (or ``variant_id`` / ``id``)
275 column holding ``chr:pos:ref>alt`` strings, or ``CHROM``/``POS``/``REF``/
276 ``ALT`` columns (case-insensitive; ``chromosome``/``position`` also work),
277 matching the layout of the official batch-scoring notebook.
278 """
279 file_path = Path(path)
280 if not file_path.exists():
281 raise FileNotFoundError(file_path)
282 if file_path.suffix.lower() == ".vcf" or file_path.name.lower().endswith(".vcf.gz"):
283 if file_path.name.lower().endswith(".gz"):
284 raise ValueError("gzipped VCF is not supported by this helper; decompress it first")
285 return read_vcf(file_path)
287 header, rows = _split_delimited(file_path)
288 if not header:
289 return [], [f"{file_path.name}: empty file"]
290 lookup = {name.lower(): index for index, name in enumerate(header)}
292 def column(*names: str) -> int | None:
293 for name in names:
294 if name in lookup:
295 return lookup[name]
296 return None
298 chrom_col = column("chrom", "chromosome", "chr", "#chrom")
299 pos_col = column("pos", "position")
300 ref_col = column("ref", "reference_bases", "reference")
301 alt_col = column("alt", "alternate_bases", "alternate")
302 id_col = column("variant_id", "id", "name", "rsid")
303 variant_col = column("variant", "variant_str", "variant_string")
305 variants: list[VariantSpec] = []
306 warnings: list[str] = []
307 if None not in (chrom_col, pos_col, ref_col, alt_col):
308 for index, row in enumerate(rows, start=2):
309 try:
310 variants.append(
311 VariantSpec(
312 chromosome=normalize_chromosome(row[chrom_col]),
313 position=int(row[pos_col]),
314 ref=row[ref_col].strip().upper(),
315 alt=row[alt_col].strip().upper(),
316 name=row[id_col].strip() if id_col is not None and id_col < len(row) else "",
317 )
318 )
319 except (ValueError, IndexError) as error:
320 warnings.append(f"line {index}: {error}")
321 elif variant_col is not None or id_col is not None:
322 source = variant_col if variant_col is not None else id_col
323 for index, row in enumerate(rows, start=2):
324 try:
325 text = row[source]
326 name = row[id_col].strip() if id_col is not None and id_col != source and id_col < len(row) else ""
327 variants.append(parse_variant_string(text, name=name))
328 except (ValueError, IndexError) as error:
329 warnings.append(f"line {index}: {error}")
330 else:
331 raise ValueError(
332 f"{file_path.name}: need CHROM/POS/REF/ALT columns or a 'variant' column; "
333 f"found {header}"
334 )
335 return variants, warnings
338def collect_variants(
339 variant_args: Sequence[str] | None,
340 input_path: str | None,
341) -> tuple[list[VariantSpec], list[str]]:
342 """Merge ``--variant`` strings and an ``--input`` file into one list."""
343 variants: list[VariantSpec] = []
344 warnings: list[str] = []
345 for text in variant_args or ():
346 variants.append(parse_variant_string(text))
347 if input_path:
348 from_file, file_warnings = read_variant_table(input_path)
349 variants.extend(from_file)
350 warnings.extend(file_warnings)
351 return variants, warnings
354# --------------------------------------------------------------------------
355# credentials and SDK
356# --------------------------------------------------------------------------
359def load_api_key(env_var: str | None = None) -> str:
360 """Return the API key from the environment, or exit with instructions.
362 Never accepts the key on the command line: it would land in shell history
363 and process listings.
364 """
365 names = (env_var,) if env_var else API_KEY_ENV_VARS
366 for name in names:
367 value = os.environ.get(name, "").strip()
368 if value:
369 return value
370 raise SystemExit(
371 f"No AlphaGenome API key found in {', '.join(names)}. Request a key at "
372 f"{GET_KEY_URL} (free for non-commercial use) and export it, for example:\n"
373 f" export {names[0]}=...\n"
374 "Never paste the key into a command line or commit it."
375 )
378def require_alphagenome():
379 """Import and return the ``alphagenome`` package, or exit with an install hint."""
380 try:
381 import alphagenome # noqa: F401 (import check only)
382 except ImportError:
383 print(INSTALL_HINT, file=sys.stderr)
384 sys.exit(2)
385 return alphagenome
388# --------------------------------------------------------------------------
389# score arithmetic
390# --------------------------------------------------------------------------
393def unit_interval(text: str) -> float:
394 """Validate a finite quantile threshold for argparse."""
395 import argparse
397 try:
398 value = float(text)
399 except ValueError as error:
400 raise argparse.ArgumentTypeError("expected a number in [0, 1]") from error
401 if not math.isfinite(value) or not 0 <= value <= 1:
402 raise argparse.ArgumentTypeError("expected a finite number in [0, 1]")
403 return value
406def cdf_to_tail_and_phred(cdf_quantile: float, floor: float = 1e-7) -> tuple[float, float]:
407 """Convert the Atlas ``quantiles`` layer (a CDF value) to (tail quantile, Phred).
409 The Atlas stores the cumulative quantile of the AVI raw score against all
410 genome-wide SNVs. ``tail = 1 - cdf`` is the fraction of SNVs scoring at
411 least this high; ``phred = -10 * log10(tail)``, so Phred 20 is the top 1%.
412 The tail is floored so a saturated quantile does not become infinity.
413 """
414 if not math.isfinite(cdf_quantile):
415 return float("nan"), float("nan")
416 tail = max(floor, 1.0 - float(cdf_quantile))
417 return tail, -10.0 * math.log10(tail)
420def phred_to_top_percent(phred: float) -> float:
421 """Phred 20 -> 1.0 (top 1%); Phred 30 -> 0.1 (top 0.1%)."""
422 if not math.isfinite(phred):
423 return float("nan")
424 return (10.0 ** (-phred / 10.0)) * 100.0
427def feature_display_name(key: str) -> str:
428 entry = AVI_FEATURES.get(key)
429 return entry[0] if entry else key
432def portal_variant_url(variant: str) -> str:
433 """Deep link to one variant on the Atlas website (AVI track shown).
435 ``atlas_link.py`` builds the full range of portal URLs; this is the
436 minimal form the query scripts attach to every scored variant.
437 """
438 from urllib.parse import quote # noqa: PLC0415
440 return f"{ATLAS_BASE_URL}?q={quote(variant, safe=':')}&m=variant&lItems=avi"
443# --------------------------------------------------------------------------
444# output
445# --------------------------------------------------------------------------
447_FORMATS = ("tsv", "csv", "json", "parquet")
450def _stringify(value: object) -> object:
451 if value is None:
452 return ""
453 if isinstance(value, float):
454 if math.isnan(value):
455 return ""
456 return f"{value:.6g}"
457 return value
460def write_rows(
461 rows: Sequence[Mapping[str, object]],
462 output: str | None,
463 fmt: str = "tsv",
464 columns: Sequence[str] | None = None,
465) -> None:
466 """Write rows to ``output`` (or stdout) as TSV, CSV, JSON, or Parquet.
468 Column order is the first-seen key order unless ``columns`` is given.
469 Parquet needs pandas and pyarrow, both of which ``alphagenome`` installs.
470 """
471 if fmt not in _FORMATS:
472 raise ValueError(f"unknown format {fmt!r}; choose from {_FORMATS}")
473 if columns is None:
474 seen: dict[str, None] = {}
475 for row in rows:
476 for key in row:
477 seen.setdefault(key, None)
478 columns = list(seen)
480 if fmt == "parquet":
481 if not output:
482 raise ValueError("parquet output needs -o/--output")
483 import pandas as pd # noqa: PLC0415 (optional dependency)
485 pd.DataFrame(list(rows), columns=list(columns)).to_parquet(output, index=False)
486 return
488 if fmt == "json":
489 text = json.dumps([{key: row.get(key) for key in columns} for row in rows], indent=2, default=str)
490 text += "\n"
491 else:
492 buffer = io.StringIO()
493 writer = csv.writer(buffer, delimiter="\t" if fmt == "tsv" else ",", lineterminator="\n")
494 writer.writerow(columns)
495 for row in rows:
496 writer.writerow([_stringify(row.get(key)) for key in columns])
497 text = buffer.getvalue()
499 if output:
500 Path(output).write_text(text, encoding="utf-8")
501 else:
502 sys.stdout.write(text)
505def warn(message: str) -> None:
506 print(f"warning: {message}", file=sys.stderr)
509def format_from_output(output: str | None, explicit: str | None) -> str:
510 """Pick an output format: explicit flag, else file extension, else TSV."""
511 if explicit:
512 return explicit
513 if output:
514 suffix = Path(output).suffix.lower().lstrip(".")
515 if suffix in _FORMATS:
516 return suffix
517 if suffix == "txt":
518 return "tsv"
519 return "tsv"
522def chunked(items: Sequence, size: int) -> Iterable[Sequence]:
523 for start in range(0, len(items), size):
524 yield items[start : start + size]