Skills › Scientific Agent Skills › alphagenome › score_variants.py
score_variants.py
235 lines · skills/alphagenome/scripts/score_variants.py at 5f2bc03
1#!/usr/bin/env python3
2# /// script
3# requires-python = ">=3.10"
4# dependencies = ["alphagenome>=0.9.0"]
5# ///
6"""Score variants on demand with the AlphaGenome model API (human or mouse).
8Use this when a variant is not in the Atlas (indels, mouse, non-reference
9REF alleles, a haplotype background) or when you need a scorer configuration
10the Atlas did not precompute. For hg38 SNVs, ``atlas_query.py`` is faster and
11has a larger quota.
13Each variant is scored inside a window centred on it (default 1 Mb, the
14model's full context) with the recommended scorers, and the result is the
15official tidy long table: one row per variant x scorer x track (x gene).
17Examples:
18 python score_variants.py --variant "chr22:36201698:A>C" -o scores.tsv
19 python score_variants.py --input variants.vcf --scorers RNA_SEQ SPLICE_SITE_USAGE \\
20 --ontology UBERON:0001157 --min-abs-quantile 0.99 -o colon.tsv
21 python score_variants.py --organism mouse --variant "chr7:45000000:A>G" --sequence-length 500KB
22 python score_variants.py --list-scorers
23 python score_variants.py --list-tracks --output-type RNA_SEQ --query liver -o tracks.tsv
24"""
26from __future__ import annotations
28import argparse
29import sys
30from collections.abc import Sequence
31from typing import Any
33import _common as common
35#: Supported input windows. Longer context is better; 1 Mb is the default.
36SEQUENCE_LENGTHS: dict[str, int] = {
37 "16KB": 2**14,
38 "100KB": 2**17,
39 "500KB": 2**19,
40 "1MB": 2**20,
41}
43ORGANISMS = ("human", "mouse")
44MAX_SCORERS_PER_REQUEST = 20
47def parse_sequence_length(text: str) -> int:
48 key = text.strip().upper().replace(" ", "")
49 if key in SEQUENCE_LENGTHS:
50 return SEQUENCE_LENGTHS[key]
51 if key.isdigit() and int(key) in SEQUENCE_LENGTHS.values():
52 return int(key)
53 raise ValueError(f"sequence length must be one of {', '.join(SEQUENCE_LENGTHS)} (got {text!r})")
56def select_scorers(names: Sequence[str] | None, organism: str, include_active: bool = False) -> list[Any]:
57 """Pick recommended scorer configurations by name, validated for the organism.
59 With no names, every recommended scorer is used except the ``*_ACTIVE``
60 variants (absolute activity of the allele, not the REF-ALT difference),
61 which are only useful when explicitly asked for.
62 """
63 from alphagenome.models import dna_client, variant_scorers # noqa: PLC0415
65 recommended = variant_scorers.RECOMMENDED_VARIANT_SCORERS
66 organism_enum = dna_client.Organism.HOMO_SAPIENS if organism == "human" else dna_client.Organism.MUS_MUSCULUS
67 if names:
68 unknown = [name for name in names if name not in recommended]
69 if unknown:
70 raise ValueError(f"unknown scorer(s) {unknown}; choose from {sorted(recommended)}")
71 chosen = list(dict.fromkeys(names))
72 else:
73 chosen = [name for name in recommended if include_active or not name.endswith("_ACTIVE")]
74 selected = []
75 for name in chosen:
76 scorer = recommended[name]
77 supported = variant_scorers.SUPPORTED_ORGANISMS[scorer.base_variant_scorer]
78 if organism_enum.to_proto() not in supported:
79 common.warn(f"scorer {name} is not available for {organism}; skipped")
80 continue
81 selected.append(scorer)
82 if len(selected) > MAX_SCORERS_PER_REQUEST:
83 raise ValueError(f"at most {MAX_SCORERS_PER_REQUEST} scorers per request; {len(selected)} selected")
84 return selected
87def make_model(args: argparse.Namespace):
88 api_key = common.load_api_key(args.api_key_env)
89 common.require_alphagenome()
90 import grpc # noqa: PLC0415
91 from alphagenome.models import dna_client # noqa: PLC0415
93 try:
94 return dna_client.create(api_key, timeout=args.timeout)
95 except grpc.FutureTimeoutError as error: # pragma: no cover - network
96 raise SystemExit(
97 f"could not reach the AlphaGenome service within {args.timeout}s "
98 "(gdmscience.googleapis.com:443); check network access and proxies"
99 ) from error
102def organism_enum(name: str):
103 from alphagenome.models import dna_client # noqa: PLC0415
105 return dna_client.Organism.HOMO_SAPIENS if name == "human" else dna_client.Organism.MUS_MUSCULUS
108def list_tracks(model, args: argparse.Namespace) -> int:
109 metadata = model.output_metadata(organism_enum(args.organism)).concatenate()
110 frame = metadata
111 if args.output_type:
112 wanted = {name.upper() for name in args.output_type}
113 column = "output_type" if "output_type" in frame.columns else None
114 if column:
115 frame = frame[frame[column].astype(str).str.upper().str.replace("OUTPUTTYPE.", "", regex=False).isin(wanted)]
116 if args.query:
117 needle = args.query.lower()
118 mask = frame.astype(str).apply(lambda column: column.str.lower().str.contains(needle, regex=False)).any(axis=1)
119 frame = frame[mask]
120 rows = frame.to_dict("records")
121 if not args.output and len(rows) > 200 and not args.force_stdout:
122 raise SystemExit(f"{len(rows)} tracks match; narrow with --output-type/--query or write them with -o FILE")
123 common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
124 return 0
127def list_scorers() -> int:
128 from alphagenome.models import variant_scorers # noqa: PLC0415
130 rows = []
131 for name, scorer in variant_scorers.RECOMMENDED_VARIANT_SCORERS.items():
132 row: dict[str, Any] = {"name": name, "scorer": type(scorer).__name__}
133 for attribute in ("requested_output", "width", "aggregation_type"):
134 value = getattr(scorer, attribute, None)
135 if value is not None:
136 row[attribute] = getattr(value, "name", value)
137 rows.append(row)
138 common.write_rows(rows, None, "tsv")
139 return 0
142def score(model, args: argparse.Namespace) -> int:
143 from alphagenome.models import variant_scorers # noqa: PLC0415
145 specs, warnings = common.collect_variants(args.variant, args.input)
146 for message in warnings:
147 common.warn(message)
148 if not specs:
149 raise SystemExit("no variants given; use --variant and/or --input")
151 length = parse_sequence_length(args.sequence_length)
152 scorers = select_scorers(args.scorers, args.organism, include_active=args.include_active)
153 if not scorers:
154 raise SystemExit("no scorers left after organism filtering")
156 from alphagenome.data import genome # noqa: PLC0415
158 variants = [
159 genome.Variant(spec.chromosome, spec.position, spec.ref, spec.alt, name=spec.name or str(spec)) for spec in specs
160 ]
161 intervals = [variant.reference_interval.resize(length) for variant in variants]
162 results = model.score_variants(
163 intervals,
164 variants,
165 scorers,
166 organism=organism_enum(args.organism),
167 progress_bar=False,
168 max_workers=args.workers,
169 )
170 frame = variant_scorers.tidy_scores(results, match_gene_strand=True)
171 if frame is None or frame.empty:
172 raise SystemExit("the model returned no scores")
174 if args.ontology and "ontology_curie" in frame.columns:
175 frame = frame[frame["ontology_curie"].isin(set(args.ontology))]
176 if args.biosample and "biosample_name" in frame.columns:
177 frame = frame[frame["biosample_name"].astype(str).str.contains(args.biosample, case=False, regex=False)]
178 if args.gene and "gene_name" in frame.columns:
179 frame = frame[frame["gene_name"].isin(set(args.gene)) | frame["gene_name"].isna()]
180 if args.min_abs_quantile is not None:
181 if "quantile_score" not in frame.columns or frame["quantile_score"].isna().all():
182 raise SystemExit("no calibrated quantiles returned; omit --min-abs-quantile and inspect raw scores")
183 frame = frame[frame["quantile_score"].abs() >= args.min_abs_quantile]
185 frame = frame.sort_values("raw_score", key=lambda column: column.abs(), ascending=False)
186 rows = frame.to_dict("records")
187 if not args.output and len(rows) > 200 and not args.force_stdout:
188 raise SystemExit(f"{len(rows)} rows; write them with -o FILE (or pass --force-stdout)")
189 common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
190 return 0
193def build_parser() -> argparse.ArgumentParser:
194 parser = argparse.ArgumentParser(
195 prog="score_variants.py",
196 description="Score variants on demand with the AlphaGenome model (recommended scorers, tidy output).",
197 epilog=__doc__.split("Examples:", 1)[-1],
198 formatter_class=argparse.RawDescriptionHelpFormatter,
199 )
200 parser.add_argument("--variant", nargs="+", metavar="CHR:POS:REF>ALT", help="one or more variants (1-based)")
201 parser.add_argument("--input", metavar="FILE", help="VCF, or TSV/CSV with CHROM/POS/REF/ALT or a 'variant' column")
202 parser.add_argument("--scorers", nargs="+", metavar="NAME", help="recommended scorer names (default: all non-ACTIVE); see --list-scorers")
203 parser.add_argument("--include-active", action="store_true", help="also run the *_ACTIVE scorers when --scorers is omitted")
204 parser.add_argument("--sequence-length", default="1MB", help="16KB, 100KB, 500KB, or 1MB (default 1MB)")
205 parser.add_argument("--organism", choices=ORGANISMS, default="human", help="human (hg38) or mouse (mm10)")
206 parser.add_argument("--ontology", nargs="+", metavar="CURIE", help="keep tracks with these ontology CURIEs (post-hoc filter)")
207 parser.add_argument("--biosample", metavar="TEXT", help="keep tracks whose biosample name contains this text")
208 parser.add_argument("--gene", nargs="+", metavar="SYMBOL", help="keep gene-centric rows for these symbols")
209 parser.add_argument("--min-abs-quantile", type=common.unit_interval, help="keep rows with |quantile_score| >= this in [0,1]; requires calibrated scores (human recommended scorers)")
210 parser.add_argument("--workers", type=int, default=5, help="parallel requests (default 5)")
211 parser.add_argument("--list-scorers", action="store_true", help="print the recommended scorer configurations and exit")
212 parser.add_argument("--list-tracks", action="store_true", help="print the model's track metadata (ontology CURIEs, biosamples) and exit")
213 parser.add_argument("--output-type", nargs="+", metavar="TYPE", help="with --list-tracks: restrict to output types, e.g. RNA_SEQ DNASE")
214 parser.add_argument("--query", metavar="TEXT", help="with --list-tracks: case-insensitive substring filter")
215 parser.add_argument("-o", "--output", help="write here instead of stdout (extension picks the format)")
216 parser.add_argument("--format", choices=("tsv", "csv", "json", "parquet"), help="override the output format")
217 parser.add_argument("--api-key-env", metavar="NAME", help="environment variable holding the key (default: ALPHAGENOME_API_KEY, then ALPHA_GENOME_API_KEY)")
218 parser.add_argument("--timeout", type=float, default=30.0, help="seconds to wait for the gRPC channel (default 30)")
219 parser.add_argument("--force-stdout", action="store_true", help="print large tables to stdout anyway")
220 return parser
223def main(argv: Sequence[str] | None = None) -> int:
224 args = build_parser().parse_args(argv)
225 if args.list_scorers:
226 common.require_alphagenome()
227 return list_scorers()
228 model = make_model(args)
229 if args.list_tracks:
230 return list_tracks(model, args)
231 return score(model, args)
234if __name__ == "__main__":
235 sys.exit(main())