Skip to content

just_dna_enricher.atlas_client

just_dna_enricher.atlas_client

The AlphaGenome Atlas client — precomputed variant scores over gRPC, on two packages (RM192).

The [atlas] extra, not core. uv add alphagenome costs 550 MB and 47 packages against a tier whose entire runtime list is httpx/tenacity/huggingface-hub, and six of the twenty dependencies that wheel declares are never imported on any scoring path. The .proto sources are Apache-2.0, so grpcio + protobuf reach every Atlas RPC — 19 MB (pyproject.toml carries the measurement; 22 MB was the grpcio release current at the design round), with score payloads decoding through struct.unpack from the standard library. Measured in ALPHAGENOME_ATLAS.md § 6.2, and pinned by test_imports_stay_within_the_declared_floor rather than left as a claim in prose.

The bindings are generated, not committed, and neither are the sources they come from: just-dna-enricher atlas generate fetches the .proto files from google-deepmind/alphagenome at the commit and per-file sha256 atlas_protos.py pins, then builds them into a git-ignored generated/ package. So this module's import is guarded, and a checkout that has not run the generator gets a message naming the command instead of a traceback from protobuf. The repository carries the pin and the distribution carries the files — hatch_build.py runs the same fetch at wheel-build time, which is what made an installed wheel work at all. That is RM196, and docs/vendor/alphagenome_protos/README.md is where the vendored copy used to be.

Three house rules shape the code rather than the wire format:

  • A client leaking its transport library's exception has no contract (@client-exception-contract). Every grpc.RpcError is translated at the boundary into an AtlasError subclass, and the distinction that matters is which kind of no: a transport failure is retryable, a refusal is not, and "this variant is not in the precomputed set" is neither — it is an answer.
  • The house algebra is three-valued and None is never False. An indel is not scored zero, it is not scored at all, and score_variant says so by raising rather than returning a number.
  • A verdict function with several arms owes a reason function with the same arms (@answered-is-not-absent), which is why every refusal carries the server's own words.

One retry layer, and it is ours (RM280). The vendored grpc_service_config.json asks grpc to retry, and grpc 1.83.1 caps service-config retries at five attempts, which is what upstream already asks for — so JUST_DNA_HTTP_RETRY_ATTEMPTS could not raise it, and a tenacity layer stacked on top would have multiplied attempts (5 x N) rather than set a floor. The channel therefore gets grpc's retry switched off (channel_options), and AtlasClient._call retries with upstream's own budget and backoff under net.attempt_floor, like every other client in this tier (@retry-attempt-floor). What it retries is derived from _translate: exactly the statuses that become AtlasUnavailable, so INTERNAL, which the contract always called retryable and upstream's config never retried, is now asked again.

A gate, and its interval is zero because the service measured none (RM307). The Atlas publishes no rate budget, so it was measured on 2026-09-29 with the stub called directly (no retry layer to hide a throttle): 479 sequential calls in 90 s (5.3/s), then 1,493 calls from four threads in 70 s (21/s, about 1,270 a minute) came back OK, with no RESOURCE_EXHAUSTED anywhere. What did move was concurrency: sixteen threads on one channel got less through (16.7/s) and 32 of 1,600 calls sat until the 30 s deadline. So ATLAS_REQUEST_INTERVAL is 0.0, stated rather than defaulted, and the gate is still there for the two things a gate does besides waiting: it counts spent, one per attempt, so a host metering egress sees Atlas calls like every other upstream's (S95), and a host running several clients can inject one shared gate= with whatever interval it wants.

ListDenseVariantScores is here — score_interval — and the reason it took a second attempt is recorded in _interval: the blocker was Interval.strand having no zero member, not the x-goog-fieldmask header or 32 bp chunking that this docstring blamed for a day.

AtlasError

Bases: RuntimeError

Base for every failure this client reports. No grpc.RpcError escapes past it.

AtlasUnavailable

Bases: AtlasError

The service could not be reached or did not answer. Retryable; says nothing about the variant.

AtlasRefused

Bases: AtlasError

The service understood the request and declined it. Not retryable.

AtlasRefMismatch

Bases: AtlasRefused

REF does not match the assembly at that position.

Worth its own type because it is a finding about the caller's data, not about the service — the Atlas validates against GRCh38 and names the real base, which is the check a local file lookup cannot make (@va-omits-ref: a VA does not encode ref).

AtlasNotScored

Bases: AtlasError

The variant is outside the precomputed set — indels, today.

Deliberately not an AtlasRefused: the request was legal and the answer is "unknown", which is the third state. Callers that treat this as zero are the bug this type exists to prevent (@unreachable-not-absent).

GeneRef dataclass

GeneRef(gene_id: str, name: str | None)

One gene a scorer attributed a score to — the source's own attribution, not ours.

gene_id is kept verbatim, version included. Upstream's proto comment says the field is the "ENSEMBL gene identifier without version number, e.g. ENSG00000100342" and the live service returns ENSG00000040608.14, so the comment is wrong and the bytes are not (@verbatim-except-order). A caller joining on bare accessions truncates deliberately; one who believed the comment would have written a join that silently matches nothing.

VariantScore dataclass

VariantScore(
    scorer: str,
    raw: tuple[float, ...] | None,
    quantile: tuple[float, ...] | None,
    shape: tuple[int, ...],
    genes: tuple[GeneRef, ...] = (),
)

One scorer's answer for one variant.

raw is the magnitude on the scorer's own scale; quantile is its empirical rank against a background of common variants (MAF > 0.01 in any gnomAD v3 population). Both are None when the server returned the block but not that field — absent is not zero here either.

phred property

phred: float | None

The Phred-scaled quantile, or None when there is no quantile to scale.

Reproduces the PHRED column of the published AVI artifact from its calibrated_scores. Scalar scorers only — a multi-valued block has no single Phred value to report.

IntervalScore dataclass

IntervalScore(
    chrom: str,
    position: int,
    ref: str,
    alt: str,
    scores: tuple[VariantScore, ...],
)

One variant inside an interval, with the scorer block that came back for it.

AtlasClient

AtlasClient(
    stub, *, api_key: str, gate: PacingGate | None = None
)

The three Atlas RPCs, with the transport's exceptions kept inside.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def __init__(self, stub, *, api_key: str, gate: PacingGate | None = None) -> None:
    self._stub = stub
    self._metadata = (("x-goog-api-key", api_key),)
    self._gate = gate if gate is not None else PacingGate(ATLAS_REQUEST_INTERVAL)

gate property

gate: PacingGate

The gate every attempt waits on; its spent is this client's upstream attempts.

score_variant

score_variant(
    chrom: str,
    position: int,
    ref: str,
    alt: str,
    *,
    scorers: tuple[str, ...] = (),
) -> tuple[VariantScore, ...]

Precomputed scores for one SNV. Raises rather than inventing a number.

position is the 1-based VCF position, passed through unchanged (@start-1based).

Source code in enricher/src/just_dna_enricher/atlas_client.py
def score_variant(
    self, chrom: str, position: int, ref: str, alt: str, *, scorers: tuple[str, ...] = ()
) -> tuple[VariantScore, ...]:
    """Precomputed scores for one SNV. Raises rather than inventing a number.

    `position` is the 1-based VCF position, passed through unchanged (`@start-1based`).
    """
    chrom = wire_contig(chrom)
    label = f"{chrom}:{position} {ref}>{alt}"
    request = atlas_service_pb2.GetDenseVariantScoresRequest(
        variant=dna_model_pb2.Variant(
            chromosome=chrom,
            position=position,
            reference_bases=ref,
            alternate_bases=alt,
        ),
        organism=dna_model_pb2.ORGANISM_HOMO_SAPIENS,
        filter=scorer_filter(*scorers),
    )
    try:
        response = self._call(self._stub.GetDenseVariantScores, request, self._metadata)
    except grpc.RpcError as exc:
        raise _translate(exc, variant=label) from exc
    return tuple(
        VariantScore(
            scorer=block.variant_scorer.name,
            raw=unpack_float32(block.scores),
            quantile=unpack_float32(block.calibrated_scores),
            shape=tuple(block.shape),
            genes=_gene_refs(block),
        )
        for block in response.scores
    )

score_interval

score_interval(
    chrom: str,
    start: int,
    end: int,
    *,
    scorers: tuple[str, ...] = (),
    gene_names: tuple[str, ...] = (),
    field_mask: bool = True,
) -> tuple[IntervalScore, ...]

Every scored variant in [start, end), following next_page_token to the last page.

start/end are 1-based VCF positions, converted to the proto's 0-based interval here so this module speaks one coordinate convention throughout (@start-1based).

A filter is effectively required, and that is a size limit rather than a rule. Measured: an unfiltered 32 bp interval answers with a 43 MB message and dies on the 4 MB client default — RESOURCE_EXHAUSTED: Received message larger than max. So an unfiltered request is not "slower", it fails, and the caller is told which knob to turn rather than left to read a transport error.

Pagination, not chunking, is what bounds a page: a 1,024 bp interval comes back as 512 scores and a token. The SDK's 32 bp sub-intervals are its parallelism strategy, not a protocol requirement — a 128 bp interval answers in one call.

The token cannot be trusted on its own, and that is upstream's bug rather than a precaution: the server returns one on an exactly-full final page, and the next request comes back INVALID_ARGUMENT. So the walk also stops once the requested interval is covered. The loop below has the measurement; a caller reading only this docstring would otherwise re-derive the crash.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def score_interval(
    self,
    chrom: str,
    start: int,
    end: int,
    *,
    scorers: tuple[str, ...] = (),
    gene_names: tuple[str, ...] = (),
    field_mask: bool = True,
) -> tuple[IntervalScore, ...]:
    """Every scored variant in `[start, end)`, following `next_page_token` to the last page.

    `start`/`end` are **1-based VCF positions**, converted to the proto's 0-based interval here
    so this module speaks one coordinate convention throughout (`@start-1based`).

    **A filter is effectively required, and that is a size limit rather than a rule.** Measured:
    an unfiltered 32 bp interval answers with a 43 MB message and dies on the 4 MB client
    default — `RESOURCE_EXHAUSTED: Received message larger than max`. So an unfiltered request is
    not "slower", it fails, and the caller is told which knob to turn rather than left to read
    a transport error.

    Pagination, not chunking, is what bounds a page: a 1,024 bp interval comes back as 512
    scores and a token. The SDK's 32 bp sub-intervals are its *parallelism* strategy, not a
    protocol requirement — a 128 bp interval answers in one call.

    **The token cannot be trusted on its own**, and that is upstream's bug rather than a
    precaution: the server returns one on an *exactly-full final page*, and the next request
    comes back `INVALID_ARGUMENT`. So the walk also stops once the requested interval is
    covered. The loop below has the measurement; a caller reading only this docstring would
    otherwise re-derive the crash.
    """
    if not scorers and not gene_names:
        raise AtlasRefused(
            f"{chrom}:{start}-{end}: an interval query needs a filter. Unfiltered, the Atlas "
            "answers with every scorer for every variant — measured at 43 MB for 32 bp, against "
            "a 4 MB default receive limit. Pass `scorers` and/or `gene_names`."
        )
    chrom = wire_contig(chrom)
    label = f"{chrom}:{start}-{end}"
    metadata = self._metadata
    if field_mask:
        metadata = (*metadata, ("x-goog-fieldmask", ",".join(LIST_FIELD_MASK)))

    request = atlas_service_pb2.ListDenseVariantScoresRequest(
        interval=_interval(chrom, start - 1, end),
        organism=dna_model_pb2.ORGANISM_HOMO_SAPIENS,
        filter=interval_filter(scorers=scorers, gene_names=gene_names),
    )
    out: list[IntervalScore] = []
    while True:
        try:
            response = self._call(self._stub.ListDenseVariantScores, request, metadata)
        except grpc.RpcError as exc:
            raise _translate(exc, variant=label) from exc
        for entry in response.variant_scores:
            out.append(
                IntervalScore(
                    chrom=entry.variant.chromosome,
                    position=entry.variant.position,
                    ref=entry.variant.reference_bases,
                    alt=entry.variant.alternate_bases,
                    scores=tuple(
                        VariantScore(
                            scorer=block.variant_scorer.name,
                            raw=unpack_float32(block.scores),
                            quantile=unpack_float32(block.calibrated_scores),
                            shape=tuple(block.shape),
                            genes=_gene_refs(block),
                        )
                        for block in entry.scores
                    ),
                )
            )
        if not response.next_page_token:
            return tuple(out)
        # **The server hands back a token on an exactly-full final page, and following it 400s.**
        # Measured on 2026-09-10: a 1,000 bp interval is 3,000 variants over six pages and the
        # short last page correctly omits the token, while a 1,024 bp interval is 3,072 — exactly
        # six full pages of 512 — and page six carries a token whose seventh request comes back
        # `INVALID_ARGUMENT`. That is AIP-158 violated in the one place a faithful client cannot
        # survive it, and the SDK's own loop has the same shape; it never trips because the SDK
        # only ever sends 32 bp sub-intervals, which cannot fill a page.
        #
        # So the token is followed but not *trusted* as the sole terminator: reaching the last
        # base the caller asked for ends the walk too. Deriving it from the request rather than
        # from a page-size constant keeps it right if the server's page size ever moves.
        if out and max(score.position for score in out) >= end:
            return tuple(out)
        request.page_token = response.next_page_token

scorer_names

scorer_names() -> tuple[str, ...]

Every scorer the Atlas serves. 22 of them at the time of writing.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def scorer_names(self) -> tuple[str, ...]:
    """Every scorer the Atlas serves. 22 of them at the time of writing."""
    request = atlas_service_pb2.ListVariantScoresMetadataRequest(
        organism=dna_model_pb2.ORGANISM_HOMO_SAPIENS
    )
    try:
        response = self._call(self._stub.ListVariantScoresMetadata, request, self._metadata)
    except grpc.RpcError as exc:
        raise _translate(exc, variant="<metadata>") from exc
    return tuple(entry.variant_scorer.name for entry in response.variant_scorer_metadata)

channel_options

channel_options() -> tuple[tuple[str, object], ...]

The options connect opens its channel with: upstream's config minus its retry policy.

Two retry layers multiply, they do not floor. grpc retries under the service config and caps that at five attempts, so a tenacity layer over it would make 5 x N attempts per call while the knob could still not raise grpc's five. So grpc's retry is off twice over: grpc.enable_retries is 0, and the config the channel receives carries no retryPolicy at all, so neither a flipped grpc default nor a dropped option turns it back on unnoticed. Everything else in the method config — the per-attempt timeout in particular, which is a deadline and not a retry — is kept verbatim.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def channel_options() -> tuple[tuple[str, object], ...]:
    """The options `connect` opens its channel with: upstream's config minus its retry policy.

    **Two retry layers multiply, they do not floor.** grpc retries under the service config and caps
    that at five attempts, so a `tenacity` layer over it would make 5 x N attempts per call while the
    knob could still not raise grpc's five. So grpc's retry is off twice over: `grpc.enable_retries`
    is `0`, and the config the channel receives carries no `retryPolicy` at all, so neither a flipped
    grpc default nor a dropped option turns it back on unnoticed. Everything else in the method config
    — the per-attempt `timeout` in particular, which is a deadline and not a retry — is kept verbatim.
    """
    config = {
        **_SERVICE_CONFIG,
        "methodConfig": [
            {key: value for key, value in method.items() if key != "retryPolicy"}
            for method in _SERVICE_CONFIG["methodConfig"]
        ],
    }
    return (("grpc.service_config", json.dumps(config)), ("grpc.enable_retries", 0))

wire_contig

wire_contig(chrom: str) -> str

22 → chr22, and chr22 unchanged — the spelling the Atlas wire expects.

This module already owns one coordinate convention and this is the second. _interval's docstring says the 0-based/1-based conversion happens here so callers pass VCF positions throughout; contig spelling is exactly the same class of fact, and leaving it to callers is what let a pass send 22 and get back a bare NOT_FOUND: Chromosome 22 not found — a live failure that no offline test could have produced, because a stub answers whatever it is asked.

One normalizer, not two (@one-normalizer-two-spellings). alphagenome_check had a private _chr doing this for the snapshot join, which is the shape where a private name keeps the second caller from finding the first: the same conversion was needed at the RPC boundary and nothing pointed there. VariantRow normalizes through vrs.normalize_chrom and stores 22; AlphaGenome is UCSC-style chr22 on the wire and in its tabix index alike. One direction only, at the boundary — the prefixed spelling is the source's, so the conversion belongs to the code crossing into it rather than to either model.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def wire_contig(chrom: str) -> str:
    """`22` → `chr22`, and `chr22` unchanged — the spelling the Atlas wire expects.

    **This module already owns one coordinate convention and this is the second.** `_interval`'s
    docstring says the 0-based/1-based conversion happens here so callers pass VCF positions
    throughout; contig spelling is exactly the same class of fact, and leaving it to callers is what
    let a pass send `22` and get back a bare `NOT_FOUND: Chromosome 22 not found` — a live failure
    that no offline test could have produced, because a stub answers whatever it is asked.

    **One normalizer, not two** (`@one-normalizer-two-spellings`). `alphagenome_check` had a private
    `_chr` doing this for the snapshot join, which is the shape where a private name keeps the second
    caller from finding the first: the same conversion was needed at the RPC boundary and nothing
    pointed there. `VariantRow` normalizes through `vrs.normalize_chrom` and stores `22`; AlphaGenome
    is UCSC-style `chr22` on the wire and in its tabix index alike. One direction only, at the
    boundary — the prefixed spelling is the source's, so the conversion belongs to the code crossing
    into it rather than to either model.
    """
    value = str(chrom).strip()
    return value if value.lower().startswith("chr") else f"chr{value}"

phred_from_quantile

phred_from_quantile(quantile: float) -> float

-10 log10(1 - q), the transform the published AVI file's PHRED column is.

A saturated quantile — exactly 1.0, which is what the API returns for the most extreme variants — raises rather than returning infinity or a clamped stand-in. The honest answer there is "this surface cannot tell you", and a number would be a worse answer than a refusal: the file has the real value and the API does not (@unreachable-not-absent, at the resolution of one float).

Source code in enricher/src/just_dna_enricher/atlas_client.py
def phred_from_quantile(quantile: float) -> float:
    """`-10 log10(1 - q)`, the transform the published AVI file's `PHRED` column is.

    A saturated quantile — exactly 1.0, which is what the API returns for the most extreme variants
    — raises rather than returning infinity or a clamped stand-in. The honest answer there is "this
    surface cannot tell you", and a number would be a worse answer than a refusal: the file has the
    real value and the API does not (`@unreachable-not-absent`, at the resolution of one float).
    """
    if not 0.0 <= quantile <= 1.0:
        raise ValueError(f"quantile must be in [0, 1], got {quantile!r}")
    if quantile >= 1.0:
        raise AtlasNotScored(
            "quantile saturated at 1.0: the float32 wire format caps a derived Phred at "
            f"{MAX_REPRESENTABLE_PHRED:.3f} and the published artifact goes above it. "
            "Read PHRED from the downloaded file for this variant."
        )
    return -10.0 * math.log10(1.0 - quantile)

unpack_float32

unpack_float32(payload: bytes) -> tuple[float, ...] | None

Decode a score field. None for an absent field, never an empty tuple silently.

The wire format is little-endian float32, unpacked with the standard library because the whole point of this blueprint is that no array package is required to read a score.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def unpack_float32(payload: bytes) -> tuple[float, ...] | None:
    """Decode a score field. `None` for an absent field, never an empty tuple silently.

    The wire format is little-endian `float32`, unpacked with the standard library because the
    whole point of this blueprint is that no array package is required to read a score.
    """
    if not payload:
        return None
    if len(payload) % 4:
        raise AtlasError(f"score payload is not a whole number of float32: {len(payload)} bytes")
    return struct.unpack(f"<{len(payload) // 4}f", payload)

scorer_filter

scorer_filter(*scorers: str) -> str

The AIP-160 filter string selecting one or more scorers.

A string, not a message — the one piece of the request that does not come from the protos, and the reason a hand-built client is possible at all.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def scorer_filter(*scorers: str) -> str:
    """The AIP-160 filter string selecting one or more scorers.

    A string, not a message — the one piece of the request that does not come from the protos, and
    the reason a hand-built client is possible at all.
    """
    if not scorers:
        return ""
    return " OR ".join(f'scores.variant_scorer.name = "{s}"' for s in scorers)

interval_filter

interval_filter(
    *,
    scorers: tuple[str, ...] = (),
    gene_names: tuple[str, ...] = (),
) -> str

The AIP-160 filter an interval query needs, over scorers and/or attributed genes.

Two clauses ANDed, each an OR over its own members — the shape the SDK builds, reproduced here because the filter is a string and is therefore the one part of the request a hand-built client has to know rather than derive from the protos.

The gene clause is what makes a gene-scoped slice possible at all: the Atlas attributes a variant to genes across the model's whole input window, so filtering by gene name is a server-side selection rather than a coordinate range the caller guesses at.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def interval_filter(*, scorers: tuple[str, ...] = (), gene_names: tuple[str, ...] = ()) -> str:
    """The AIP-160 filter an interval query needs, over scorers and/or attributed genes.

    Two clauses ANDed, each an OR over its own members — the shape the SDK builds, reproduced here
    because the filter is a *string* and is therefore the one part of the request a hand-built client
    has to know rather than derive from the protos.

    The gene clause is what makes a gene-scoped slice possible at all: the Atlas attributes a variant
    to genes across the model's whole input window, so filtering by gene name is a server-side
    selection rather than a coordinate range the caller guesses at.
    """
    clauses = []
    if scorers:
        clauses.append(" OR ".join(f'scores.variant_scorer.name = "{s}"' for s in scorers))
    if gene_names:
        clauses.append(" OR ".join(f'scores.metadata.gene_scorers.metadata.name = "{g}"' for g in gene_names))
    return " AND ".join(f"({c})" for c in clauses)

connect

connect(
    api_key: str,
    *,
    address: str = DEFAULT_ADDRESS,
    timeout: float = 30.0,
    gate: PacingGate | None = None,
) -> AtlasClient

Open a channel and hand back a client, translating a failed handshake like any other error.

Source code in enricher/src/just_dna_enricher/atlas_client.py
def connect(
    api_key: str,
    *,
    address: str = DEFAULT_ADDRESS,
    timeout: float = 30.0,
    gate: PacingGate | None = None,
) -> AtlasClient:
    """Open a channel and hand back a client, translating a failed handshake like any other error."""
    channel = grpc.secure_channel(
        address,
        grpc.ssl_channel_credentials(),
        options=channel_options(),
    )
    try:
        grpc.channel_ready_future(channel).result(timeout)
    except grpc.FutureTimeoutError as exc:
        raise AtlasUnavailable(f"channel to {address} not ready within {timeout}s") from exc
    return AtlasClient(atlas_service_pb2_grpc.AtlasServiceStub(channel=channel), api_key=api_key, gate=gate)