Alphagenome

scripts/atlas_query.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["alphagenome>=0.9.0"]
# ///
"""Query the AlphaGenome Atlas: precomputed variant effects for every hg38 SNV.

Four subcommands, all against the ``alphagenome.atlas`` gRPC client:

  avi       AlphaGenome Variant Impact (AVI) score, Phred, and the 18 SHAP
            feature attributions for variants or for every SNV in a window.
  scores    Raw and quantile scores from the track-level scorers (RNA_SEQ,
            DNASE, CHIP_TF, SPLICE_SITE_USAGE, ...) as one tidy row per
            variant x track (x gene, for gene-centric scorers).
  scorers   The scorers the server currently serves, with track counts.
  tracks    The track catalogue behind a scorer (biosample, ontology CURIE,
            assay, TF, histone mark) for choosing ontology filters.

Coordinates: variants are 1-based ``chr:pos:ref>alt``; ``--interval`` is a
1-based closed ``chr:start-end``. The Atlas is GRCh38 (hg38) only and the REF
allele must match the reference: swapped REF/ALT can cause a lookup miss.
Returned variant identities are checked against the request. Reads the API
key from ALPHAGENOME_API_KEY (or ALPHA_GENOME_API_KEY).

Examples:
  python atlas_query.py avi --variant "chr22:36201698:A>C" "chr9:128225994:G>A"
  python atlas_query.py avi --input candidates.vcf --min-phred 20 -o avi.tsv
  python atlas_query.py avi --interval chr11:5225727-5226575 --top-k 25
  python atlas_query.py scores --variant "chr22:36201698:A>C" \\
      --scorers RNA_SEQ SPLICE_SITE_USAGE --ontology UBERON:0001157 -o scores.tsv
  python atlas_query.py scorers
  python atlas_query.py tracks --scorer CHIP_TF --query GATA1
"""

from __future__ import annotations

import argparse
import concurrent.futures
import math
import sys
from collections.abc import Mapping, Sequence
from typing import Any

import _common as common

DEFAULT_MAX_WINDOW_BP = 1_000
DEFAULT_WORKERS = 8
TRACK_METADATA_COLUMNS = (
    "name",
    "strand",
    "ontology_curie",
    "biosample_name",
    "biosample_type",
    "Assay title",
    "assay_title",
    "gtex_tissue",
    "transcription_factor",
    "histone_mark",
    "data_source",
)


# --------------------------------------------------------------------------
# pure conversion helpers (unit-tested with synthetic AnnData)
# --------------------------------------------------------------------------


def _obs_variant_strings(adata) -> list[str]:
    if adata.obs is None or "variant" not in adata.obs.columns:
        return [str(index) for index in adata.obs_names]
    return [str(value) for value in adata.obs["variant"].tolist()]


def _feature_names(adata) -> list[str]:
    if adata.var is not None and "name" in adata.var.columns:
        return [str(value) for value in adata.var["name"].tolist()]
    return [str(value) for value in adata.var_names]


def _layer(adata, name: str):
    try:
        return adata.layers[name] if name in adata.layers else None
    except (KeyError, TypeError):
        return None


def summarize_avi(avi_adata, fi_adata=None) -> list[dict[str, Any]]:
    """One row per variant from the ``AVI_SCORE`` and feature-importance AnnData.

    Columns: variant, avi_raw, avi_cdf_quantile, avi_tail_quantile, avi_phred,
    avi_top_percent, top_feature_key, top_feature, top_feature_value, and one
    ``fi_<KEY>`` column per attribution. Missing attributions leave NaN.
    """
    import numpy as np  # noqa: PLC0415

    rows: list[dict[str, Any]] = []
    if avi_adata is None or avi_adata.X is None or avi_adata.X.shape[0] == 0:
        return rows

    variants = _obs_variant_strings(avi_adata)
    raw = np.asarray(avi_adata.X, dtype=float).reshape(len(variants), -1)
    quantiles = _layer(avi_adata, "quantiles")
    if quantiles is not None:
        quantiles = np.asarray(quantiles, dtype=float).reshape(len(variants), -1)

    fi_by_variant: dict[str, np.ndarray] = {}
    fi_names: list[str] = []
    if fi_adata is not None and fi_adata.X is not None and fi_adata.X.shape[0] > 0:
        fi_names = _feature_names(fi_adata)
        fi_matrix = np.asarray(fi_adata.X, dtype=float).reshape(fi_adata.X.shape[0], -1)
        for variant, values in zip(_obs_variant_strings(fi_adata), fi_matrix, strict=True):
            fi_by_variant[variant] = values

    for index, variant in enumerate(variants):
        row: dict[str, Any] = {"variant": variant}
        row["avi_raw"] = float(raw[index, 0]) if raw.shape[1] else float("nan")
        if quantiles is not None and quantiles.shape[1]:
            cdf = float(quantiles[index, 0])
            tail, phred = common.cdf_to_tail_and_phred(cdf)
        else:
            cdf, tail, phred = float("nan"), float("nan"), float("nan")
        row["avi_cdf_quantile"] = cdf
        row["avi_tail_quantile"] = tail
        row["avi_phred"] = phred
        row["avi_top_percent"] = common.phred_to_top_percent(phred)

        values = fi_by_variant.get(variant)
        if values is not None and len(fi_names) == len(values):
            top = int(np.argmax(np.abs(values)))
            row["top_feature_key"] = fi_names[top]
            row["top_feature"] = common.feature_display_name(fi_names[top])
            row["top_feature_value"] = float(values[top])
            for name, value in zip(fi_names, values, strict=True):
                row[f"fi_{name}"] = float(value)
        else:
            row["top_feature_key"] = ""
            row["top_feature"] = ""
            row["top_feature_value"] = float("nan")
        row["atlas_url"] = common.portal_variant_url(variant)
        rows.append(row)
    return rows


def tidy_atlas_scores(scores: Mapping[str, Any]) -> list[dict[str, Any]]:
    """Flatten ``{scorer: AnnData}`` into one row per (obs row, track).

    ``obs`` rows are variants, or variant x gene for the gene-centric scorers
    (RNA_SEQ, POLYADENYLATION, SPLICE_*); ``var`` rows are tracks. Raw scores
    come from ``X`` and, when the server sent them, quantile scores from the
    ``quantiles`` layer.
    """
    import numpy as np  # noqa: PLC0415

    rows: list[dict[str, Any]] = []
    for scorer, adata in scores.items():
        if adata is None or adata.X is None or 0 in adata.X.shape:
            continue
        matrix = np.asarray(adata.X, dtype=float)
        quantiles = _layer(adata, "quantiles")
        if quantiles is not None:
            quantiles = np.asarray(quantiles, dtype=float)
        obs = adata.obs if adata.obs is not None else None
        var = adata.var if adata.var is not None else None
        var_columns = [column for column in TRACK_METADATA_COLUMNS if var is not None and column in var.columns]
        obs_columns = [
            column
            for column in ("gene_id", "gene_name", "strand", "junction_Start", "junction_End")
            if obs is not None and column in obs.columns
        ]
        variants = _obs_variant_strings(adata)
        for i in range(matrix.shape[0]):
            base: dict[str, Any] = {"variant": variants[i], "scorer": scorer}
            for column in obs_columns:
                key = {"strand": "gene_strand", "junction_Start": "junction_start", "junction_End": "junction_end"}.get(
                    column, column
                )
                base[key] = obs.iloc[i][column]
            for j in range(matrix.shape[1]):
                row = dict(base)
                row["track_index"] = j
                for column in var_columns:
                    key = "assay_title" if column == "Assay title" else column
                    row[f"track_{key}" if key in {"name", "strand"} else key] = var.iloc[j][column]
                row["raw_score"] = float(matrix[i, j])
                if quantiles is not None:
                    row["quantile_score"] = float(quantiles[i, j])
                rows.append(row)
    return rows


def max_abs_track(adata) -> dict[str, Any] | None:
    """The single largest |score| entry of a track-scorer AnnData, with its metadata."""
    import numpy as np  # noqa: PLC0415

    if adata is None or adata.X is None or 0 in adata.X.shape:
        return None
    matrix = np.asarray(adata.X, dtype=float)
    flat = int(np.nanargmax(np.abs(matrix)))
    i, j = divmod(flat, matrix.shape[1])
    result: dict[str, Any] = {"raw_score": float(matrix[i, j]), "track_index": j}
    if adata.var is not None:
        for column in ("name", "biosample_name", "ontology_curie", "transcription_factor", "histone_mark"):
            if column in adata.var.columns:
                result[column] = adata.var.iloc[j][column]
    if adata.obs is not None:
        for column in ("gene_name", "gene_id"):
            if column in adata.obs.columns:
                result[column] = adata.obs.iloc[i][column]
    return result


def scorer_rows(metadata: Mapping[str, Any]) -> list[dict[str, Any]]:
    rows = []
    for name, meta in metadata.items():
        frame = getattr(meta, "track_metadata", None)
        row: dict[str, Any] = {
            "scorer": name,
            "is_signed": bool(getattr(meta, "is_signed", False)),
            "n_tracks": int(len(frame)) if frame is not None else 0,
        }
        if frame is not None and "biosample_name" in frame.columns:
            row["n_biosamples"] = int(frame["biosample_name"].nunique())
        if frame is not None and "ontology_curie" in frame.columns:
            row["n_ontology_terms"] = int(frame["ontology_curie"].nunique())
        rows.append(row)
    return sorted(rows, key=lambda item: item["scorer"])


def track_rows(metadata: Mapping[str, Any], scorer: str | None, query: str | None) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    needle = query.lower() if query else None
    for name, meta in metadata.items():
        if scorer and name != scorer:
            continue
        frame = getattr(meta, "track_metadata", None)
        if frame is None or frame.empty:
            continue
        for index, (_, record) in enumerate(frame.iterrows()):
            row: dict[str, Any] = {"scorer": name, "track_index": index}
            for column in TRACK_METADATA_COLUMNS:
                if column in frame.columns:
                    key = "assay_title" if column == "Assay title" else column
                    row[key] = record[column]
            if needle and not any(needle in str(value).lower() for value in row.values()):
                continue
            rows.append(row)
    return rows


# --------------------------------------------------------------------------
# client plumbing
# --------------------------------------------------------------------------


def make_client(args: argparse.Namespace):
    api_key = common.load_api_key(args.api_key_env)
    common.require_alphagenome()
    import grpc  # noqa: PLC0415
    from alphagenome.atlas import atlas  # noqa: PLC0415

    try:
        return atlas.create(api_key, timeout=args.timeout)
    except grpc.FutureTimeoutError as error:  # pragma: no cover - network
        raise SystemExit(
            f"could not reach the Atlas service within {args.timeout}s "
            "(gdmscience.googleapis.com:443); check network access and proxies"
        ) from error


def to_genome_variant(spec: common.VariantSpec):
    from alphagenome.data import genome  # noqa: PLC0415

    return genome.Variant(
        chromosome=spec.chromosome,
        position=spec.position,
        reference_bases=spec.ref,
        alternate_bases=spec.alt,
        name=spec.name,
    )


def to_genome_interval(text: str, max_window: int):
    from alphagenome.data import genome  # noqa: PLC0415

    chromosome, start0, end = common.parse_interval_string(text)
    width = common.interval_width(start0, end)
    if width > max_window:
        raise SystemExit(
            f"interval is {width:,} bp, above --max-window {max_window:,} bp. Each base "
            f"expands to {common.SNVS_PER_BASE} variants; raise --max-window deliberately."
        )
    return genome.Interval(chromosome=chromosome, start=start0, end=end)


def _query_each(
    client,
    specs: Sequence[common.VariantSpec],
    requested_scorers: Sequence[str],
    workers: int,
    **filters: Any,
) -> list[tuple[common.VariantSpec, Mapping[str, Any] | None, str]]:
    """Query variants one by one, keeping per-variant failures as messages."""

    def one(spec: common.VariantSpec):
        try:
            result = client.query_variant(to_genome_variant(spec), requested_scorers=list(requested_scorers), **filters)
            for scorer, adata in result.items():
                if adata.obs is not None and "variant" in adata.obs:
                    returned = set(_obs_variant_strings(adata))
                    if returned - {str(spec)}:
                        raise ValueError(
                            f"{scorer} returned a different variant: {sorted(returned)}; "
                            f"requested {spec}. Check REF and genome build."
                        )
            return spec, result, ""
        except Exception as error:  # noqa: BLE001 - surfaced per variant
            return spec, None, f"{type(error).__name__}: {error}"

    if workers <= 1 or len(specs) <= 1:
        return [one(spec) for spec in specs]
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
        return list(pool.map(one, specs))


# --------------------------------------------------------------------------
# subcommands
# --------------------------------------------------------------------------


def cmd_avi(args: argparse.Namespace) -> int:
    client = make_client(args)
    rows: list[dict[str, Any]] = []

    if args.interval:
        interval = to_genome_interval(args.interval, args.max_window)
        scores = client.query_interval(interval, requested_scorers=list(common.AVI_SCORERS), progress_bar=False)
        rows = summarize_avi(scores.get(common.AVI_SCORER), scores.get(common.AVI_FEATURES_SCORER))
        for row in rows:
            row["error"] = ""
    else:
        specs, warnings = common.collect_variants(args.variant, args.input)
        for message in warnings:
            common.warn(message)
        if not specs:
            raise SystemExit("no variants given; use --variant, --input, or --interval")
        indels = [spec for spec in specs if not spec.is_snv]
        if indels:
            common.warn(
                f"{len(indels)} indel(s) requested; the public Atlas API currently serves "
                "genome-wide SNVs, so indels may come back as lookup misses"
            )
        for spec, scores, error in _query_each(client, specs, common.AVI_SCORERS, args.workers):
            if scores is None:
                row = spec.to_row()
                row.update({"avi_raw": float("nan"), "avi_phred": float("nan"), "error": error})
                rows.append(row)
                continue
            summary = summarize_avi(scores.get(common.AVI_SCORER), scores.get(common.AVI_FEATURES_SCORER))
            if not summary:
                row = spec.to_row()
                row.update({"avi_raw": float("nan"), "avi_phred": float("nan"), "error": "no AVI_SCORE returned"})
                rows.append(row)
                continue
            row = spec.to_row()
            row.update(summary[0])
            row["variant"] = str(spec)
            row["error"] = ""
            rows.append(row)

    if args.with_tracks:
        for row in rows:
            key = row.get("top_feature_key") or ""
            scorers = common.AVI_FEATURES.get(key, ("", "", ()))[2]
            if not scorers or row.get("error"):
                continue
            try:
                spec = common.parse_variant_string(row["variant"])
                detail = client.query_variant(to_genome_variant(spec), requested_scorers=list(scorers))
            except Exception as error:  # noqa: BLE001
                row["top_track_error"] = f"{type(error).__name__}: {error}"
                continue
            best = None
            for scorer in scorers:
                hit = max_abs_track(detail.get(scorer))
                if hit and (best is None or abs(hit["raw_score"]) > abs(best[1]["raw_score"])):
                    best = (scorer, hit)
            if best:
                scorer, hit = best
                row["top_track_scorer"] = scorer
                row["top_track_name"] = hit.get("name", "")
                row["top_track_biosample"] = hit.get("biosample_name", "")
                row["top_track_ontology"] = hit.get("ontology_curie", "")
                row["top_track_gene"] = hit.get("gene_name", "")
                row["top_track_raw_score"] = hit["raw_score"]

    if args.min_phred is not None:
        rows = [row for row in rows if row.get("error") or (isinstance(row.get("avi_phred"), float) and row["avi_phred"] >= args.min_phred)]

    rows.sort(key=lambda row: -(row.get("avi_phred") if isinstance(row.get("avi_phred"), float) and math.isfinite(row["avi_phred"]) else -1))
    if args.top_k:
        rows = rows[: args.top_k]

    if not args.output and len(rows) > 50 and not args.force_stdout:
        raise SystemExit(f"{len(rows)} rows; write them with -o FILE (or pass --force-stdout)")
    common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
    return 0


def cmd_scores(args: argparse.Namespace) -> int:
    client = make_client(args)
    requested = list(args.scorers)
    filters = {
        "ontology_terms": args.ontology or None,
        "gene_names": args.gene or None,
        "gene_ids": args.gene_id or None,
    }
    rows: list[dict[str, Any]] = []
    if args.interval:
        interval = to_genome_interval(args.interval, args.max_window)
        scores = client.query_interval(interval, requested_scorers=requested, progress_bar=False, **filters)
        rows = tidy_atlas_scores(scores)
    else:
        specs, warnings = common.collect_variants(args.variant, args.input)
        for message in warnings:
            common.warn(message)
        if not specs:
            raise SystemExit("no variants given; use --variant, --input, or --interval")
        for spec, scores, error in _query_each(client, specs, requested, args.workers, **filters):
            if scores is None:
                common.warn(f"{spec}: {error}")
                rows.append({"variant": str(spec), "scorer": "", "error": error})
                continue
            rows.extend(tidy_atlas_scores(scores))

    if args.min_abs_quantile is not None:
        calibrated = [row for row in rows if "quantile_score" in row]
        if not calibrated and any(not row.get("error") for row in rows):
            raise SystemExit("no calibrated quantiles returned; omit --min-abs-quantile and inspect raw scores")
        rows = [
            row
            for row in rows
            if row.get("error") or (
                "quantile_score" in row and abs(row["quantile_score"]) >= args.min_abs_quantile
            )
        ]
    if not args.output and len(rows) > 200 and not args.force_stdout:
        raise SystemExit(f"{len(rows)} rows; write them with -o FILE (or pass --force-stdout)")
    common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
    return 0


def cmd_scorers(args: argparse.Namespace) -> int:
    client = make_client(args)
    rows = scorer_rows(client.scorer_metadata())
    common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
    return 0


def cmd_tracks(args: argparse.Namespace) -> int:
    client = make_client(args)
    rows = track_rows(client.scorer_metadata(), args.scorer, args.query)
    if not args.output and len(rows) > 200 and not args.force_stdout:
        raise SystemExit(
            f"{len(rows)} tracks match; narrow with --scorer/--query or write them with -o FILE"
        )
    common.write_rows(rows, args.output, common.format_from_output(args.output, args.format))
    return 0


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------


def _add_common(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("-o", "--output", help="write here instead of stdout (extension picks the format)")
    parser.add_argument("--format", choices=("tsv", "csv", "json", "parquet"), help="override the output format")
    parser.add_argument("--api-key-env", metavar="NAME", help="environment variable holding the key (default: ALPHAGENOME_API_KEY, then ALPHA_GENOME_API_KEY)")
    parser.add_argument("--timeout", type=float, default=30.0, help="seconds to wait for the gRPC channel (default 30)")
    parser.add_argument("--force-stdout", action="store_true", help="print large tables to stdout anyway")


def _add_variant_inputs(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--variant", nargs="+", metavar="CHR:POS:REF>ALT", help="one or more variants (1-based; gnomAD/GTEx spellings accepted)")
    parser.add_argument("--input", metavar="FILE", help="VCF, or TSV/CSV with CHROM/POS/REF/ALT or a 'variant' column")
    parser.add_argument("--interval", metavar="CHR:START-END", help="1-based closed window: every SNV in it (3 per base)")
    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)")
    parser.add_argument("--workers", type=int, default=DEFAULT_WORKERS, help=f"parallel variant queries (default {DEFAULT_WORKERS})")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="atlas_query.py",
        description="Query the AlphaGenome Atlas (precomputed variant effects, hg38).",
        epilog=__doc__.split("Examples:", 1)[-1] if "Examples:" in __doc__ else None,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="command", required=True)

    avi = sub.add_parser("avi", help="AVI score, Phred, and feature attributions")
    _add_variant_inputs(avi)
    avi.add_argument("--min-phred", type=float, help="keep variants with AVI Phred >= this (20 = top 1%%)")
    avi.add_argument("--top-k", type=int, help="keep the K highest-Phred variants")
    avi.add_argument("--with-tracks", action="store_true", help="also report the strongest track behind the top feature (one extra query per variant)")
    _add_common(avi)
    avi.set_defaults(func=cmd_avi)

    scores = sub.add_parser("scores", help="raw + quantile scores from the track-level scorers")
    _add_variant_inputs(scores)
    scores.add_argument("--scorers", nargs="+", default=["RNA_SEQ"], metavar="NAME", help="Atlas scorer names (see `scorers`); default RNA_SEQ")
    scores.add_argument("--ontology", nargs="+", metavar="CURIE", help="keep tracks for these ontology terms, e.g. UBERON:0001157 CL:0000084")
    scores.add_argument("--gene", nargs="+", metavar="SYMBOL", help="gene-centric scorers: keep these gene symbols")
    scores.add_argument("--gene-id", nargs="+", metavar="ENSG", help="gene-centric scorers: keep these Ensembl gene IDs")
    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]")
    _add_common(scores)
    scores.set_defaults(func=cmd_scores)

    scorers = sub.add_parser("scorers", help="list the scorers the Atlas currently serves")
    _add_common(scorers)
    scorers.set_defaults(func=cmd_scorers)

    tracks = sub.add_parser("tracks", help="track catalogue (biosample, ontology CURIE, assay) per scorer")
    tracks.add_argument("--scorer", metavar="NAME", help="restrict to one scorer")
    tracks.add_argument("--query", metavar="TEXT", help="case-insensitive substring over every column")
    _add_common(tracks)
    tracks.set_defaults(func=cmd_tracks)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())