Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+weight/alphagenome-all-folds/ocdbt.process_0/d/48900811ab32b5ad3873d40e96d63846 filter=lfs diff=lfs merge=lfs -text
+weight/alphagenome-all-folds/ocdbt.process_0/d/6b08a83c9d6a2ade925dbcdd91299ce2 filter=lfs diff=lfs merge=lfs -text
+weight/alphagenome-all-folds/ocdbt.process_0/d/bda173d4f00d3afcef0a043614df05a6 filter=lfs diff=lfs merge=lfs -text

.mv

@@ -0,0 +1 @@
+Revision:master,CreatedAt:1784688508

LICENSE

@@ -0,0 +1,203 @@
+Copyright 2025 Onescience Authors. All rights reserved.
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2025 Onescience Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,248 @@
+---
+license: apache-2.0
+tasks:
+  - genomic-sequence-modeling
+frameworks:
+  - jax
+language:
+  - en
+  - zh
+tags:
+  - OneScience
+  - Life Sciences
+  - Genomics
+  - DNA Sequence Model
+  - Variant Effect Prediction
+  - AlphaGenome
+datasets:
+  - OneScience/alphagenome_dataset
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">AlphaGenome</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+AlphaGenome is a DNA sequence model proposed by Google DeepMind. It can take DNA intervals up to 1 Mbp as input and predict multiple classes of genomic functional signals for track prediction and regulatory variant effect scoring.
+
+Paper: Advancing regulatory variant effect prediction with AlphaGenome  
+https://www.nature.com/articles/s41586-025-10014-0
+
+# Model Description
+
+AlphaGenome is implemented based on JAX / Flax and supports genomic interval inference, variant effect scoring, track evaluation, and fine-tuning examples. This model package is accompanied by the ModelScope dataset `OneScience/alphagenome_dataset`, which can be used for quick local verification.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :--- |
+| Genomic interval prediction | Input a reference genome FASTA, chromosome, and interval coordinates, and output predicted tracks for ATAC, DNase, CAGE, RNA-seq, ChIP, and other signals |
+| Variant effect scoring | Input a VCF or built-in example variants, compare prediction differences between reference and variant sequences, and generate a variant scoring table |
+| Track prediction evaluation | Use validation data from the AlphaGenome dataset to calculate regression evaluation metrics for different assay bundles |
+| Fine-tuning experiments | Use a custom reference genome, interval CSV, and BigWig signal files to verify the fine-tuning workflow |
+| ModelScope / OneCode runtime | After downloading the model project and accompanying dataset, quickly verify script connectivity in a biology-domain runtime environment |
+
+
+
+# Usage Instructions
+
+## 1. OneCode Usage
+
+You can experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to experience intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU runtime is recommended.
+- CPU can be used for import checks and small-configuration connectivity verification, but full training and inference are slow.
+- DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version matching the current cluster.
+
+
+
+
+
+**Environment Check**
+
+- NVIDIA GPU:
+
+```bash
+nvidia-smi
+```
+
+- Hygon DCU:
+
+```bash
+hy-smi
+```
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/alphagenome --local_dir ./alphagenome
+cd alphagenome
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# First activate DTK and Conda
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Supports uv installation
+pip install onescience[bio-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+After installation, return to the model package directory:
+
+```bash
+cd ./alphagenome
+```
+
+### Training and Inference Data Introduction
+
+The OneScience community has uploaded the data required for AlphaGenome inference, evaluation, and fine-tuning to ModelScope: [OneScience/alphagenome_dataset](https://modelscope.cn/datasets/OneScience/alphagenome_dataset). After downloading, place the data in the `data/` directory of the model package.
+
+```bash
+modelscope download --dataset OneScience/alphagenome_dataset --local_dir ./data
+```
+
+### Training Weights
+
+The repository includes `weight/alphagenome-all-folds`, and all scripts also support specifying model weights through `--model_dir`.
+
+### Prepare Weights
+
+If using local weights, place the AlphaGenome Orbax checkpoint in the following directory:
+
+```text
+weight/
+  alphagenome-all-folds/
+    _CHECKPOINT_METADATA
+    _METADATA
+    ...
+```
+
+If running in a shared runtime environment, you can also reuse a unified directory through environment variables:
+
+```bash
+export ONESCIENCE_MODELS_DIR=/path/to/onescience/models
+export ONESCIENCE_DATASETS_DIR=/path/to/onescience/datasets
+```
+
+The scripts will preferentially read:
+
+- `${ONESCIENCE_MODELS_DIR}/AlphaGenome/alphagenome-all-folds`
+- `${ONESCIENCE_DATASETS_DIR}/AlphaGenome`
+
+If the above environment variables are not set, the defaults under the current model package are read:
+
+- `weight/alphagenome-all-folds`
+- `data/`
+
+### Interval Inference
+
+```bash
+bash scripts/inference.sh
+```
+
+Equivalent Python command example:
+
+```bash
+python scripts/run_inference.py \
+  --fasta_path ./data/reference/HOMO_SAPIENS/GRCh38.p13.genome.fa \
+  --model_dir ./weight/alphagenome-all-folds \
+  --chromosome chr19 \
+  --start 10587331 \
+  --end 11635907 \
+  --output_dir ./outputs
+```
+
+Inference results will be saved to `outputs/`.
+
+### Variant Effect Scoring
+
+```bash
+bash scripts/run_variant.sh
+```
+
+When specifying VCF input, you can use:
+
+```bash
+python scripts/run_variant_scoring.py \
+  --vcf_path ./data/example.vcf \
+  --fasta_path ./data/reference/HOMO_SAPIENS/GRCh38.p13.genome.fa \
+  --model_dir ./weight/alphagenome-all-folds \
+  --output_dir ./outputs_variant
+```
+
+Scoring results will be saved as CSV files.
+
+### Track Prediction Evaluation
+
+```bash
+bash scripts/run_track.sh
+```
+
+You can also explicitly specify the data and output paths:
+
+```bash
+python scripts/run_track_prediction_eval.py \
+  --model_dir ./weight/alphagenome-all-folds \
+  --model_version ALL_FOLDS \
+  --data_dir ./data/v1/train \
+  --output_path ./outputs_track/eval_results.csv
+```
+
+### Fine-tuning Example
+
+```bash
+python scripts/run_finetuning.py \
+  --fasta_path ./data/reference/HOMO_SAPIENS/GRCh38.p13.genome.fa \
+  --regions_csv ./data/finetune_regions.csv \
+  --bigwig_paths ./data/sample_atac.bw \
+  --output_dir ./finetuned_model \
+  --num_steps 1000 \
+  --batch_size 2
+```
+
+# Data Format
+
+The ModelScope dataset `OneScience/alphagenome_dataset` is recommended to be downloaded to `data/` under the model package. The default structure is as follows:
+
+```text
+data/
+  reference/
+    HOMO_SAPIENS/
+      GRCh38.p13.genome.fa
+      GRCh38.p13.genome.fa.fai
+  v1/
+    train/
+      ...
+```
+
+Where:
+
+- `reference/HOMO_SAPIENS/GRCh38.p13.genome.fa` is the human reference genome FASTA.
+- `.fai` is the FASTA index file.
+- `v1/train/` is the data directory used for track prediction evaluation.
+- Custom fine-tuning also requires preparing an interval CSV file with column names `chromosome,start,end`, as well as one or more BigWig signal files.
+
+# Official OneScience Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- This repository is based on the AlphaGenome open-source model and provides DCU adaptation. The related source code uses the Apache License 2.0.
+- For scientific research, cite the original paper: [Advancing regulatory variant effect prediction with AlphaGenome](https://www.nature.com/articles/s41586-025-10014-0).

flax_model/__init__.py

@@ -0,0 +1 @@
+"""Local flax_models copied from OneScience for this standalone package."""

flax_model/alphagenome/__init__.py

@@ -0,0 +1,39 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AlphaGenome - 统一DNA序列基因组模型，已集成至OneScience。
+
+AlphaGenome 是 Google DeepMind 发布的统一DNA序列模型，能够分析最长100万碱基对的
+DNA序列，以单碱基对分辨率预测基因表达、剪接模式、染色质特征和接触图谱。
+
+模块结构:
+    model/         - 核心模型组件 (AlphaGenome, SequenceEncoder, 各类Head)
+    io/            - 数据输入输出 (BundleName, Dataset, Fasta)
+    finetuning/    - 微调工具 (训练步骤、数据管道)
+    evals/         - 评估指标 (track预测、回归指标)
+    _sdk/          - AlphaGenome SDK (vendored alphagenome v0.6.1)
+                     包含数据类型、gRPC客户端、变异/区间评分、可视化和ISM解释工具
+
+使用方式:
+    from flax_model.alphagenome.model import model as alphagenome_model
+    from flax_model.alphagenome.io import dataset
+    from flax_model.alphagenome.finetuning import finetune
+
+    # SDK 数据类型 (原 from alphagenome.data import genome):
+    from flax_model.alphagenome._sdk.data import genome
+    from flax_model.alphagenome._sdk.models import dna_model
+    from flax_model.alphagenome._sdk.models import dna_output
+"""
+
+__version__ = '0.1.0'

flax_model/alphagenome/_sdk/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AlphaGenome Python SDK (vendored into OneScience).
+
+Original package: alphagenome v0.6.1
+Provides data types, gRPC client, variant/interval scoring,
+visualization, and ISM interpretation utilities for genomic models.
+
+This is a vendored copy integrated into OneScience to avoid requiring
+a separate alphagenome package installation.
+"""
+
+__version__ = '0.6.1'

flax_model/alphagenome/_sdk/colab_utils.py

@@ -0,0 +1,62 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for Google Colab."""
+
+import os
+
+
+def get_api_key(secret: str = 'ALPHA_GENOME_API_KEY'):
+  """Returns API key from environment variable or Colab secrets.
+
+  Tries to retrieve the API key from the environment first. If not found,
+  attempts to retrieve it from Colab secrets (if running in Colab).
+
+  Args:
+    secret: The name of the environment variable or Colab secret key to
+      retrieve.
+
+  Raises:
+    ValueError: If the API key cannot be found in the environment or Colab
+      secrets.
+  """
+
+  if api_key := os.environ.get(secret):
+    return api_key
+
+  try:
+    # pylint: disable=g-import-not-at-top, import-outside-toplevel
+    from google.colab import userdata  # pytype: disable=import-error
+    # pylint: enable=g-import-not-at-top, import-outside-toplevel
+
+    try:
+      api_key = userdata.get(secret)
+      return api_key
+    except (
+        userdata.NotebookAccessError,
+        userdata.SecretNotFoundError,
+        userdata.TimeoutException,
+    ) as e:
+      raise ValueError(
+          f'Cannot find or access API key in Colab secrets with {secret=}. Make'
+          ' sure you have added the API key to Colab secrets and enabled'
+          ' access. See'
+          ' https://www.alphagenomedocs.com/installation.html#add-api-key-to-secrets'
+          ' for more details.'
+      ) from e
+  except ImportError:
+    # Not running in Colab.
+    pass
+
+  raise ValueError(f'Cannot find API key with {secret=}.')

flax_model/alphagenome/_sdk/data/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Data classes for interacting and visualizing genomic models."""

flax_model/alphagenome/_sdk/data/fold_intervals.py

@@ -0,0 +1,115 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Genomics intervals used for training model folds."""
+
+import enum
+
+from flax_model.alphagenome._sdk.models import dna_client
+import immutabledict
+import pandas as pd
+
+
+_DEFAULT_EXAMPLE_REGIONS = immutabledict.immutabledict({
+    dna_client.Organism.HOMO_SAPIENS: (
+        'https://github.com/calico/borzoi/raw/'
+        '5c9358222b5026abb733ed5fb84f3f6c77239b37/data/sequences_human.bed.gz'
+    ),
+    dna_client.Organism.MUS_MUSCULUS: (
+        'https://github.com/calico/borzoi/raw/'
+        '5c9358222b5026abb733ed5fb84f3f6c77239b37/data/sequences_mouse.bed.gz'
+    ),
+})
+
+
+class Subset(enum.Enum):
+  """Subset of the data."""
+
+  TRAIN = 0
+  VALID = 1
+  TEST = 2
+
+
+# Fold ONE is aligned with all trained Borzoi checkpoints: 3 and 4 are held out.
+_VALID_FOLD = immutabledict.immutabledict({
+    0: 'fold0',
+    1: 'fold3',
+    2: 'fold2',
+    3: 'fold6',
+    -1: 'fold0',
+})
+
+_TEST_FOLD = immutabledict.immutabledict({
+    0: 'fold1',
+    1: 'fold4',
+    2: 'fold5',
+    3: 'fold7',
+    -1: 'fold1',
+})
+
+_MODEL_VERSION_TO_FOLD = immutabledict.immutabledict({
+    dna_client.ModelVersion.FOLD_0: 0,
+    dna_client.ModelVersion.FOLD_1: 1,
+    dna_client.ModelVersion.FOLD_2: 2,
+    dna_client.ModelVersion.FOLD_3: 3,
+    dna_client.ModelVersion.ALL_FOLDS: -1,
+})
+
+
+def get_all_folds() -> list[str]:
+  """Returns the names of all data folds."""
+  return [f'fold{i}' for i in range(8)]
+
+
+def get_fold_names(
+    model_version: dna_client.ModelVersion, subset: Subset
+) -> list[str]:
+  """Returns the names of the folds for a given model version and subset."""
+  match subset:
+    case Subset.VALID:
+      return [_VALID_FOLD[_MODEL_VERSION_TO_FOLD[model_version]]]
+    case Subset.TEST:
+      return [_TEST_FOLD[_MODEL_VERSION_TO_FOLD[model_version]]]
+    case Subset.TRAIN:
+      all_folds = get_all_folds()
+      if _MODEL_VERSION_TO_FOLD[model_version] == -1:
+        return all_folds
+      remove_folds = get_fold_names(
+          model_version, Subset.VALID
+      ) + get_fold_names(model_version, Subset.TEST)
+      for fold in remove_folds:
+        all_folds.remove(fold)
+      return all_folds
+    case _:
+      raise ValueError(f'Unknown {subset=}')
+
+
+def get_fold_intervals(
+    model_version: dna_client.ModelVersion,
+    organism: dna_client.Organism,
+    subset: Subset,
+    example_regions_path: str | None = None,
+) -> pd.DataFrame:
+  """Returns the intervals for a given model version and subset."""
+  if example_regions_path is None:
+    example_regions_path = _DEFAULT_EXAMPLE_REGIONS[organism]
+
+  example_regions = pd.read_csv(
+      example_regions_path,
+      sep='\t',
+      names=['chromosome', 'start', 'end', 'fold'],
+  )
+  return example_regions[
+      example_regions.fold.isin(get_fold_names(model_version, subset))
+  ]

flax_model/alphagenome/_sdk/data/gene_annotation.py

@@ -0,0 +1,445 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for working with gene annotations (e.g., GTFs)."""
+
+from collections.abc import Sequence
+import enum
+
+from flax_model.alphagenome._sdk.data import genome
+import numpy as np
+import pandas as pd
+
+
+@enum.unique
+class TranscriptType(enum.Enum):
+  """Valid Transcript types available in the GENCODE GTF."""
+
+  IG_C_GENE = 'IG_C_gene'
+  IG_C_PSEUDOGENE = 'IG_C_pseudogene'
+  IG_D_GENE = 'IG_D_gene'
+  IG_J_GENE = 'IG_J_gene'
+  IG_J_PSEUDOGENE = 'IG_J_pseudogene'
+  IG_V_GENE = 'IG_V_gene'
+  IG_V_PSEUDOGENE = 'IG_V_pseudogene'
+  IG_PSEUDOGENE = 'IG_pseudogene'
+  MT_RRNA = 'Mt_rRNA'
+  MT_TRNA = 'Mt_tRNA'
+  TEC = 'TEC'
+  TR_C_GENE = 'TR_C_gene'
+  TR_D_GENE = 'TR_D_gene'
+  TR_J_GENE = 'TR_J_gene'
+  TR_J_PSEUDOGENE = 'TR_J_pseudogene'
+  TR_V_GENE = 'TR_V_gene'
+  TR_V_PSEUDOGENE = 'TR_V_pseudogene'
+  ARTIFACT = 'artifact'
+  LNCRNA = 'lncRNA'
+  MIRNA = 'miRNA'
+  MISC_RNA = 'misc_RNA'
+  NON_STOP_DECAY = 'non_stop_decay'
+  NONSENSE_MEDIATED_DECAY = 'nonsense_mediated_decay'
+  PROCESSED_PSEUDOGENE = 'processed_pseudogene'
+  PROCESSED_TRANSCRIPT = 'processed_transcript'
+  PROTEIN_CODING = 'protein_coding'
+  PROTEIN_CODING_CDS_NOT_DEFINED = 'protein_coding_CDS_not_defined'
+  PROTEIN_CODING_LOF = 'protein_coding_LoF'
+  RRNA = 'rRNA'
+  RRNA_PSEUDOGENE = 'rRNA_pseudogene'
+  RETAINED_INTRON = 'retained_intron'
+  RIBOZYME = 'ribozyme'
+  SRNA = 'sRNA'
+  SCRNA = 'scRNA'
+  SCARNA = 'scaRNA'
+  SNRNA = 'snRNA'
+  SNORNA = 'snoRNA'
+  TRANSCRIBED_PROCESSED_PSEUDOGENE = 'transcribed_processed_pseudogene'
+  TRANSCRIBED_UNITARY_PSEUDOGENE = 'transcribed_unitary_pseudogene'
+  TRANSCRIBED_UNPROCESSED_PSEUDOGENE = 'transcribed_unprocessed_pseudogene'
+  TRANSLATED_PROCESSED_PSEUDOGENE = 'translated_processed_pseudogene'
+  UNITARY_PSEUDOGENE = 'unitary_pseudogene'
+  UNPROCESSED_PSEUDOGENE = 'unprocessed_pseudogene'
+  VAULT_RNA = 'vault_RNA'
+
+
+def extract_tss(gtf: pd.DataFrame, feature: str = 'transcript') -> pd.DataFrame:
+  """Extract transcription start sites (TSS) from a DataFrame.
+
+  Args:
+    gtf: pd.DataFrame containing gene annotation.
+    feature: Feature in the GTF file to use (either transcript or gene).
+
+  Returns:
+    pd.DataFrame containing transcription start sites as zero-width point
+    intervals (Start == End, 0-based).
+  """
+  tss = gtf[(gtf.Feature == feature)].copy()
+
+  # Remove the extra base to make it width=0.
+  # .....[)TRANSCRIPT (strand = +)
+  # TPIRCSNART[)..... (strand = -)
+  new_start = np.where(tss.Strand == '-', tss.End, tss.Start)
+  tss.Start = new_start
+  tss.End = new_start
+
+  return tss
+
+
+def filter_transcript_type(
+    gtf: pd.DataFrame,
+    transcript_types: tuple[TranscriptType, ...] | None = None,
+) -> pd.DataFrame:
+  """Filter GTF entries by transcript types.
+
+  This function takes a GTF DataFrame and a list of transcript types and returns
+  a new DataFrame containing only the transcripts with the specified types.
+
+  The GTF DataFrame must contain a column named 'transcript_type' or
+  'transcript_biotype'. The function will raise a ValueError if neither of these
+  columns is present.
+
+  Args:
+    gtf: pd.DataFrame or pyranges.PyRanges.
+    transcript_types: List of valid transcript types to use for filtering.
+
+  Returns:
+    pd.DataFrame of GENCODE GTF entries subset to rows with the requested
+    transcript types.
+  """
+  if transcript_types is not None:
+    transcript_types_str = [x.value for x in transcript_types]
+    if 'transcript_type' in gtf.columns:
+      gtf = gtf[gtf.transcript_type.isin(transcript_types_str)]
+    elif 'transcript_biotype' in gtf.columns:
+      gtf = gtf[gtf.transcript_biotype.isin(transcript_types_str)]
+    else:
+      raise ValueError('transcript_type or transcript_biotype not in gtf.')
+  return gtf
+
+
+def filter_protein_coding(
+    gtf: pd.DataFrame, include_gene_entries: bool = False
+) -> pd.DataFrame:
+  """Filter GTF entries to only protein-coding genes.
+
+  Args:
+    gtf: pd.DataFrame of GENCODE GTF entries. This data frame must contain a
+      column named 'transcript_type' or 'transcript_biotype'.
+    include_gene_entries: Whether to include gene entries in addition to
+      transcript entries.
+
+  Returns:
+    pd.DataFrame of GENCODE GTF entries subset to rows with protein-coding
+    genes.
+  """
+  if include_gene_entries:
+    if 'gene_type' in gtf.columns:
+      gtf = gtf[gtf.gene_type == TranscriptType.PROTEIN_CODING.value]
+    else:
+      raise ValueError('gene_type not in gtf.')
+  else:
+    gtf = filter_transcript_type(gtf, (TranscriptType.PROTEIN_CODING,))
+  return gtf
+
+
+def filter_to_longest_transcript(
+    gtf: pd.DataFrame,
+) -> pd.DataFrame:
+  """Filter GTF entries to only the longest transcript per gene.
+
+  Args:
+    gtf: pd.DataFrame of GENCODE GTF entries. Must contain columns 'Feature',
+      'End', 'Start', 'gene_id', and 'transcript_id'.
+
+  Returns:
+    pd.DataFrame of GENCODE GTF entries subset to rows with the longest
+    transcript per gene.
+  """
+  lengths = gtf[gtf['Feature'] == 'transcript'].reset_index(drop=True)
+  lengths['transcript_length'] = lengths['End'] - lengths['Start'] + 1
+
+  # Identify longest transcripts per gene_id.
+  longest_transcripts = lengths.loc[
+      lengths.groupby('gene_id')['transcript_length'].idxmax()
+  ]
+
+  return gtf[gtf['transcript_id'].isin(longest_transcripts['transcript_id'])]
+
+
+def filter_to_mane_select_transcript(gtf: pd.DataFrame) -> pd.DataFrame:
+  """Filter GTF entries to only the MANE select transcript.
+
+  Note that the MANE_Select tag only exists for the human GTF file.
+
+  Args:
+    gtf: pd.DataFrame of GENCODE GTF entries. Must contain columns 'tag'.
+
+  Returns:
+    pd.DataFrame of GENCODE GTF entries subset to rows representing MANE
+    select transcripts, which are transcripts that are well-supported,
+    conserved, and expressed.
+  """
+  filtered_gtf = gtf[gtf['tag'].fillna('').str.contains('MANE_Select')]
+  if filtered_gtf.empty:
+    raise ValueError(
+        'No MANE_Select transcripts found in the GTF, possibly due to non-human'
+        ' GTF.'
+    )
+  return filtered_gtf
+
+
+def filter_transcript_support_level(
+    gtf: pd.DataFrame,
+    transcript_support_levels: str | Sequence[str],
+) -> pd.DataFrame:
+  """Filter GTF to only transcripts with specific GENCODE support levels.
+
+  As documented in the [Ensembl
+  glossary](https://www.ensembl.org/Help/Glossary),
+  the transcript support level (TSL) indicates the degree of evidence that was
+  used to construct the transcript.
+
+  As taken from the glossary, the levels are:
+
+  | Transcript support level | Description |  |
+  |---|---|---|
+  | 1 | A transcript where all splice junctions are supported by at least one
+  non-suspect mRNA. |
+  | 2 | A transcript where the best supporting mRNA is flagged as suspect or the
+  support is from multiple ESTs |
+  | 3 | A transcript where the only support is from a single EST |
+  | 4 | A transcript where the best supporting EST is flagged as suspect |
+  | 5 | A transcript where no single transcript supports the model structure. |
+  | NA | A transcript that was not analysed for TSL. |
+
+  Args:
+    gtf: pd.DataFrame of GENCODE GTF entries. Must contain column
+      'transcript_support_level'.
+    transcript_support_levels: List of valid transcript support levels to use
+      for filtering. This must be a subset of ['1', '2', '3', '4', '5']. Can
+      also be single string.
+
+  Returns:
+    pd.DataFrame exactly as provided, but subset to rows with the specified
+    support level(s).
+
+  Transcripts are scored by GENCODE according to how well mRNA and EST
+  alignments
+  match over its full length. Valid levels are:
+  '1': All splice junctions of the transcript are supported by at least one
+  non-suspect mRNA.
+  '2': The best supporting mRNA is flagged as suspect or the support is from
+  multiple ESTs.
+  '3': The only support is from a single EST.
+  '4': The best supporting EST is flagged as suspect.
+  '5': No single transcript supports the model structure.
+  'NA': The transcript was not analyzed (not supported by this filter function).
+
+  See GENCODE GTF format documentation for further details:
+  https://www.gencodegenes.org/pages/data_format.html
+  """
+  if isinstance(transcript_support_levels, str):
+    transcript_support_levels = list(transcript_support_levels)
+
+  supported_tsls = {'1', '2', '3', '4', '5'}
+  if not set(transcript_support_levels).issubset(supported_tsls):
+    raise ValueError(
+        f'transcript_support_level must be one of {supported_tsls}, but was'
+        f' {transcript_support_levels}'
+    )
+  return gtf[gtf.transcript_support_level.isin(transcript_support_levels)]
+
+
+def upgrade_annotation_ids(
+    old_ids: pd.Series, new_ids: pd.Series, patchless: bool = False
+) -> pd.Series:
+  """Upgrade or add transcript id patch version to Ensembl IDs.
+
+  This function works by
+
+  1. Dropping the patch version from `old_ids` and `new_ids`
+  2. Merging the two on the patch-less ids.
+  3. Returning the result of the merge as a pd.Series, with the 'old_ids' as
+  the index and the 'new_ids' as the values.
+
+  Ensembl patch versions have two formats: ENST####.<patch>_PAR_Y or
+  ENST####.<patch>. Both are handled.
+
+  The function will raise a ValueError if either `old_ids` or `new_ids` result
+  in duplicates after dropping the patch version.
+
+  Examples:
+  * If the old ids are ENST00010.1 and the new ids are ENST00010.3,
+  then the mapping will be ENST00010.1 -> ENST00010.3.
+  * If the old ids are ENST00010 and the new ids are ENST00010.3, then the
+  mapping will be ENST00010 -> ENST00010.3.
+
+  Args:
+    old_ids: A pd.Series of Ensembl transcript or gene ids with older or missing
+      version/patch numbers. The index of the series is ignored.
+    new_ids: A pd.Series of transcript or gene ids with newer version/patch
+      numbers. The index of the series is ignored.
+    patchless: whether old_ids are missing patch.
+
+  Returns:
+    A pd.Series, with the 'old_ids' as the index and the 'new_ids' as the
+    values.
+  """
+  new_ids = new_ids.drop_duplicates()
+
+  def drop_version(x):
+    """Drop the patch version from an Ensembl ID."""
+    if not x.str.contains('.', regex=False).all():
+      raise ValueError('All ids need to contain the patch version.')
+
+    id_split = x.str.partition('.')
+
+    # Retain anything after _ such as PAR_Y for ENST####.<patch>_PAR_Y.
+    return id_split[0] + id_split[2].str.partition('_')[2]
+
+  old_ids_nopatch = old_ids if patchless else drop_version(old_ids)
+  new_ids_nopatch = drop_version(new_ids)
+  assert (
+      not old_ids_nopatch.duplicated().any()
+  ), 'old_ids not unique without version'
+  assert (
+      not new_ids_nopatch.duplicated().any()
+  ), 'new_ids not unique without version'
+  df = pd.merge(
+      pd.DataFrame({'old': old_ids.values, 'no_version': old_ids_nopatch}),
+      pd.DataFrame({
+          'new': new_ids.values,
+          'no_version': new_ids_nopatch,
+      }),
+      on='no_version',
+      how='left',
+  )
+  return pd.Series(
+      df.set_index('old').loc[old_ids].new.values, index=old_ids.index
+  )
+
+
+def get_gene_intervals(
+    gtf: pd.DataFrame,
+    gene_symbols: Sequence[str] | None = None,
+    gene_ids: Sequence[str] | None = None,
+) -> list[genome.Interval]:
+  """Returns a list of stranded `genome.Interval`s for the given identifiers.
+
+  Args:
+    gtf: pd.DataFrame of GENCODE GTF entries. Must contain columns 'Feature',
+      'gene_name', 'gene_id', 'Chromosome', 'Start', 'End', and 'Strand'.
+    gene_symbols: A sequence of gene names or gene symbols (e.g., ['EGFR',
+      'TNF', 'TP53']). Matching is case-insensitive.
+    gene_ids: A sequence of Ensembl gene IDs, which can be patched (e.g.
+      ['ENSG00000141510.17']) or unpatched (e.g., ['ENSG00000141510']). Matching
+      is done on unpatched IDs.
+
+  Returns:
+    A list of `genome.Interval`s for the given identifiers. The
+    returned list of intervals is in the same order as the input gene
+    identifiers.
+
+  Raises:
+    ValueError: If neither or both gene_symbols and gene_ids are set, or if no
+      interval or multiple intervals are found for any of the given gene
+      identifiers.
+  """
+  if (gene_symbols is None) == (gene_ids is None):
+    raise ValueError('Exactly one of gene_symbols or gene_ids must be set.')
+
+  gtf_genes = gtf[gtf['Feature'] == 'gene'].copy()
+
+  if gene_symbols is not None:
+    id_col = 'gene_name'
+    input_ids = gene_symbols
+    process_fn = lambda s: s.str.upper()
+  else:
+    id_col = 'gene_id'
+    input_ids = gene_ids
+    process_fn = lambda s: s.str.split('.', n=1).str[0]
+
+  processed_input_ids = process_fn(pd.Series(input_ids, dtype=str))
+  gtf_genes['processed_id'] = process_fn(gtf_genes[id_col])
+
+  # Filter the GTF to only the genes that are in the input IDs.
+  gtf_subset = gtf_genes[
+      gtf_genes['processed_id'].isin(processed_input_ids.unique())
+  ]
+
+  dup_mask = gtf_subset['processed_id'].duplicated(keep=False)
+  if dup_mask.any():
+    offending_ids = gtf_subset.loc[dup_mask, id_col].unique()
+    raise ValueError(
+        'Multiple intervals found for gene(s):'
+        f' {", ".join(sorted(offending_ids))}.'
+    )
+
+  # Create a lookup map from processed_id to GTF data.
+  # Use reindex to order genes by input and insert NaNs for missing genes.
+  gtf_map = gtf_subset.set_index('processed_id')
+  result_df = gtf_map.reindex(processed_input_ids)
+
+  missing_mask = result_df['Chromosome'].isnull()
+  if missing_mask.any():
+    missing_ids = pd.Series(input_ids)[missing_mask.values].unique()
+    raise ValueError(
+        f'No interval found for gene(s): {", ".join(sorted(missing_ids))}.'
+    )
+
+  # Add original identifiers to the result for the 'name' field of the interval.
+  result_df[id_col] = list(input_ids)
+
+  return [
+      genome.Interval(
+          chromosome=row.Chromosome,
+          start=row.Start,
+          end=row.End,
+          strand=row.Strand,
+          name=getattr(row, id_col),
+      )
+      for row in result_df.itertuples()
+  ]
+
+
+def get_gene_interval(
+    gtf: pd.DataFrame,
+    gene_symbol: str | None = None,
+    gene_id: str | None = None,
+) -> genome.Interval:
+  """Returns a stranded `genome.Interval` given a gene identifier.
+
+  Either gene_symbol or gene_id must be set, but not both.
+
+  Args:
+    gtf: pd.DataFrame of GENCODE GTF entries. Must contain columns 'Feature',
+      'gene_name', 'gene_id', 'Chromosome', 'Start', 'End', and 'Strand'.
+    gene_symbol: A gene name or gene symbol (e.g., 'EGFR', 'TNF', 'TP53')
+    gene_id: An Ensembl gene ID, which can be patched (e.g.
+      'ENSG00000141510.17') or unpatched (e.g., 'ENSG00000141510').
+
+  Returns:
+    A `genome.Interval` for the given gene identifier.
+
+  Raises:
+    ValueError: If neither or both gene_symbol and gene_id are set, or if no
+      interval or multiple intervals are found for the given gene identifier.
+  """
+  if sum(x is not None for x in [gene_symbol, gene_id]) != 1:
+    raise ValueError('Exactly one of gene_symbol or gene_id must be set.')
+
+  return get_gene_intervals(
+      gtf,
+      [gene_symbol] if gene_symbol else None,
+      [gene_id] if gene_id else None,
+  )[0]

flax_model/alphagenome/_sdk/data/genome.py

@@ -0,0 +1,1181 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for working with genome-related objects such as intervals."""
+
+import collections
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+import copy
+import dataclasses
+import enum
+import re
+import sys
+from typing import Any, Protocol
+
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+import numpy as np
+from typing_extensions import Self
+
+STRAND_POSITIVE = '+'  # Also called forward strand, 5'->3' direction.
+STRAND_NEGATIVE = '-'  # Also called negative strand, 3'->5' direction.
+STRAND_UNSTRANDED = '.'
+STRAND_OPTIONS = (STRAND_POSITIVE, STRAND_NEGATIVE, STRAND_UNSTRANDED)
+_INTERVAL_START_END_REGEX = re.compile(r'(-?\d+)-(-?\d+)')
+VALID_VARIANT_BASES = frozenset('ACGTN')
+
+
+class Strand(enum.IntEnum):
+  """Represents the strand of a DNA sequence.
+
+  This enum defines the possible strands for a DNA sequence:
+
+    * `POSITIVE`:  The forward strand (5' to 3').
+    * `NEGATIVE`: The reverse strand (3' to 5').
+    * `UNSTRANDED`:  The strand is not specified.
+  """
+
+  POSITIVE = enum.auto()
+  NEGATIVE = enum.auto()
+  UNSTRANDED = enum.auto()
+
+  def __str__(self):
+    match self:
+      case Strand.POSITIVE:
+        return STRAND_POSITIVE
+      case Strand.NEGATIVE:
+        return STRAND_NEGATIVE
+      case Strand.UNSTRANDED:
+        return STRAND_UNSTRANDED
+
+  @classmethod
+  def from_str(cls, strand: str) -> Self:
+    match strand:
+      case '+':
+        return cls.POSITIVE
+      case '-':
+        return cls.NEGATIVE
+      case '.':
+        return cls.UNSTRANDED
+      case _:
+        raise ValueError(f'Strand needs to be in {STRAND_OPTIONS}')
+
+  def to_proto(self) -> dna_model_pb2.Strand:
+    match self:
+      case Strand.POSITIVE:
+        return dna_model_pb2.Strand.STRAND_POSITIVE
+      case Strand.NEGATIVE:
+        return dna_model_pb2.Strand.STRAND_NEGATIVE
+      case Strand.UNSTRANDED:
+        return dna_model_pb2.Strand.STRAND_UNSTRANDED
+
+  @classmethod
+  def from_proto(cls, strand: dna_model_pb2.Strand) -> Self:
+    match strand:
+      case dna_model_pb2.Strand.STRAND_POSITIVE:
+        return cls.POSITIVE
+      case dna_model_pb2.Strand.STRAND_NEGATIVE:
+        return cls.NEGATIVE
+      case dna_model_pb2.Strand.STRAND_UNSTRANDED:
+        return cls.UNSTRANDED
+      case _:
+        raise ValueError(f'Strand needs to be in {STRAND_OPTIONS}')
+
+
+PYRANGES_INTERVAL_COLUMNS = ('Chromosome', 'Start', 'End', 'Strand', 'Name')
+
+
+@dataclasses.dataclass(order=True)
+class Interval:
+  """Represents a genomic interval.
+
+  A genomic interval is a region on a chromosome defined by a start and end
+  position. This class provides methods for manipulating and comparing
+  intervals, and for calculating coverage and overlap.
+
+  Attributes:
+    chromosome: The chromosome name (e.g., 'chr1', '1').
+    start: The 0-based start position.
+    end: The 0-based end position (must be greater than or equal to start).
+    strand: The strand of the interval ('+', '-', or '.'). Defaults to '.'
+      (unstranded).
+    name: An optional name for the interval.
+    info: An optional dictionary to store additional information.
+    negative_strand: True if the interval is on the negative strand, False
+      otherwise.
+    width: The width of the interval (end - start).
+  """
+
+  chromosome: str
+  start: int
+  end: int
+  strand: str = STRAND_UNSTRANDED
+  name: str = dataclasses.field(default='', compare=False, hash=False)
+  info: dict[str, Any] = dataclasses.field(
+      default_factory=dict, repr=False, compare=False, hash=False
+  )
+
+  def __post_init__(self):
+    if self.end < self.start:
+      raise ValueError('end < start. Interval: ' + str(self))
+    if self.strand not in STRAND_OPTIONS:
+      raise ValueError(
+          f'Strand needs to be in {STRAND_OPTIONS}, found {self.strand}.'
+      )
+
+  @property
+  def negative_strand(self) -> bool:
+    """Returns True if interval is on the negative strand, False otherwise."""
+    return self.strand == STRAND_NEGATIVE
+
+  @property
+  def width(self) -> int:
+    """Returns the width of the interval."""
+    return self.end - self.start
+
+  def copy(self) -> Self:
+    """Returns a deep copy of the interval."""
+    return copy.deepcopy(self)
+
+  def __str__(self) -> str:
+    """Returns a string representation of the interval."""
+    return f'{self.chromosome}:{self.start}-{self.end}:{self.strand}'
+
+  @classmethod
+  def from_str(cls, string: str) -> Self:
+    """Creates an Interval from a string (e.g., 'chr1:100-200:+')."""
+    chromosome, interval, *strand = string.split(':', maxsplit=2)
+    if strand:
+      strand = strand[0]
+    else:
+      strand = STRAND_UNSTRANDED
+    # Get start and end from the interval string.
+    match = _INTERVAL_START_END_REGEX.fullmatch(interval)
+    if match:
+      start, end = int(match.group(1)), int(match.group(2))
+    else:
+      raise ValueError(f'Invalid interval: {string}')
+    return cls(
+        chromosome=chromosome, start=int(start), end=int(end), strand=strand
+    )
+
+  def to_proto(self) -> dna_model_pb2.Interval:
+    """Converts the interval to a protobuf message."""
+    return dna_model_pb2.Interval(
+        chromosome=self.chromosome,
+        start=self.start,
+        end=self.end,
+        strand=Strand.from_str(self.strand).to_proto(),
+    )
+
+  @classmethod
+  def from_proto(cls, proto: dna_model_pb2.Interval) -> Self:
+    """Creates an Interval from a protobuf message."""
+    return cls(
+        chromosome=proto.chromosome,
+        start=proto.start,
+        end=proto.end,
+        strand=str(Strand.from_proto(proto.strand)),
+    )
+
+  def to_interval_dict(self) -> dict[str, str | int]:
+    """Converts the interval to a dictionary."""
+    return dict(
+        chrom=self.chromosome,
+        start=self.start,
+        end=self.end,
+        strand=self.strand,
+    )
+
+  @classmethod
+  def from_interval_dict(cls, interval: Mapping[str, str | int]) -> Self:
+    """Creates an Interval from a dictionary."""
+    return cls(
+        chromosome=str(interval['chrom']),
+        start=int(interval['start']),
+        end=int(interval['end']),
+        strand=str(interval.get('strand', STRAND_UNSTRANDED)),
+    )
+
+  @classmethod
+  def from_pyranges_dict(
+      cls, row: Mapping[str, Any], ignore_info: bool = False
+  ) -> 'Interval':
+    """Creates an Interval from a pyranges-like dictionary.
+
+    This method constructs an `Interval` object from a dictionary that follows
+    the pyranges format, such as a row from a :class:`pandas.DataFrame`
+    converted to a dict.
+
+    The dictionary should have the following keys:
+
+      * 'Chromosome': The chromosome name.
+      * 'Start': The start position.
+      * 'End': The end position.
+      * 'Strand': The strand (optional, defaults to unstranded).
+      * 'Name': The interval name (optional).
+
+    Any other keys in the dictionary will be added to the `info` attribute of
+    the `Interval` object, unless `ignore_info` is set to True.
+
+    Args:
+      row: A dictionary containing interval data.
+      ignore_info: If True, any keys in the dictionary that are not part of the
+        standard pyranges columns ('Chromosome', 'Start', 'End', 'Strand',
+        'Name') will not be added to the `info` attribute.
+
+    Returns:
+      An `Interval` object created from the input dictionary.
+    """
+    if ignore_info:
+      info = {}
+    else:
+      info = {
+          k: v for k, v in row.items() if k not in PYRANGES_INTERVAL_COLUMNS
+      }
+
+    return cls(
+        chromosome=str(row['Chromosome']),
+        start=int(row['Start']),
+        end=int(row['End']),
+        strand=str(row.get('Strand', STRAND_UNSTRANDED)),
+        name=str(row.get('Name', '')),
+        info=info,
+    )
+
+  def to_pyranges_dict(self) -> dict[str, int | str]:
+    """Converts the interval to a pyranges-like dictionary."""
+    return {
+        'Chromosome': self.chromosome,
+        'Start': self.start,
+        'End': self.end,
+        'Name': self.name,
+        'Strand': self.strand,
+        **self.info,
+    }
+
+  def swap_strand(self) -> Self:
+    """Swaps the strand of the interval."""
+    obj = self.copy()
+    if obj.strand == STRAND_POSITIVE:
+      obj.strand = STRAND_NEGATIVE
+    elif obj.strand == STRAND_NEGATIVE:
+      obj.strand = STRAND_POSITIVE
+    elif obj.strand == STRAND_UNSTRANDED:
+      raise ValueError('Cannot swap unstranded intervals.')
+    return obj
+
+  def as_unstranded(self) -> Self:
+    """Returns an unstranded copy of the interval."""
+    obj = self.copy()
+    obj.strand = STRAND_UNSTRANDED
+    return obj
+
+  def within_reference(self, reference_length: int = sys.maxsize) -> bool:
+    """Checks if the interval is within the valid reference range."""
+    return self.start >= 0 and self.end <= reference_length
+
+  def truncate(self, reference_length: int = sys.maxsize) -> Self:
+    """Truncates the interval to fit within the valid reference range."""
+    obj = self.copy()
+    if reference_length <= 0:
+      raise ValueError('Reference length should be larger than 0.')
+    if self.within_reference(reference_length):
+      return obj
+    else:
+      obj.start = max(self.start, 0)
+      obj.end = min(self.end, reference_length)
+      return obj
+
+  def center(self, use_strand: bool = True) -> int:
+    """Computes the center of the interval.
+
+    For intervals with an odd width, the center is rounded up for
+    positive/unstranded intervals and rounded down for negative strand
+    intervals.
+
+    If `use_strand` is True and the interval is on the negative strand, the
+    center is calculated differently to maintain consistency when stacking
+    sequences from different intervals oriented in the forward strand direction.
+    This ensures that the relative distance between the interval's upstream
+    boundary and its center is preserved.
+
+    Args:
+      use_strand: If True, the strand of the interval is considered when
+        calculating the center.
+
+    Returns:
+      The integer representing the center position of the interval.
+
+    Examples:
+      >>> Interval('1', 1, 3, '+').center()
+      2
+      >>> Interval('1', 1, 3, '-').center()  # Strand doesn't matter.
+      2
+      >>> Interval('1', 1, 4, '+').center()
+      3
+      >>> Interval('1', 1, 4, '-').center()
+      2
+      >>> Interval('1', 1, 4, '-').center()
+      2
+      >>> Interval('1', 1, 4, '+').center(use_strand=False)
+      3
+      >>> Interval('1', 1, 2, '-').center()
+      1
+    """
+    center = (self.start + self.end) // 2
+    if use_strand and self.negative_strand:
+      return center
+    else:
+      return center + self.width % 2
+
+  def shift(self, offset: int, use_strand: bool = True) -> Self:
+    """Shifts the interval by the given offset.
+
+    Args:
+      offset: The amount to shift the interval.
+      use_strand: If True, the shift direction is reversed for negative strand
+        intervals.
+
+    Returns:
+      A new shifted interval.
+    """
+    obj = self.copy()
+    if use_strand and self.negative_strand:
+      offset = -offset
+    obj.start = self.start + offset
+    obj.end = self.end + offset
+    return obj
+
+  def boundary_shift(
+      self, start_offset: int = 0, end_offset: int = 0, use_strand: bool = True
+  ) -> Self:
+    """Extends or shrinks the interval by adjusting the positions with padding.
+
+    Args:
+      start_offset: The amount to shift the start position.
+      end_offset: The amount to shift the end position.
+      use_strand: If True, the offsets are applied in reverse for negative
+        strand intervals.
+
+    Returns:
+      A new interval with adjusted boundaries.
+    """
+    return self.pad(-start_offset, end_offset, use_strand=use_strand)
+
+  def pad(
+      self, start_pad: int, end_pad: int, *, use_strand: bool = True
+  ) -> Self:
+    """Pads the interval by adding the specified padding to the start and end.
+
+    Args:
+      start_pad: The amount of padding to add to the start.
+      end_pad: The amount of padding to add to the end.
+      use_strand: If True, padding is applied in reverse for negative strand
+        intervals.
+
+    Returns:
+      A new padded interval.
+    """
+    obj = self.copy()
+    obj.pad_inplace(start_pad, end_pad, use_strand=use_strand)
+    return obj
+
+  def pad_inplace(
+      self, start_pad: int, end_pad: int, *, use_strand: bool = True
+  ):
+    """Pads the interval in place by adding padding to the start and end.
+
+    Args:
+      start_pad: The amount of padding to add to the start.
+      end_pad: The amount of padding to add to the end.
+      use_strand: If True, padding is applied in reverse for negative strand
+        intervals.
+    """
+    if use_strand and self.strand == '-':
+      start_pad, end_pad = end_pad, start_pad
+    self.start -= start_pad
+    self.end += end_pad
+    if self.width < 0:
+      raise ValueError('Resulting interval has negative length')
+
+  def resize(self, width: int, use_strand: bool = True) -> Self:
+    """Resizes the interval to a new width, centered around the original center.
+
+    Args:
+      width: The new width of the interval.
+      use_strand: If True, resizing considers the strand orientation.
+
+    Returns:
+      A new resized interval.
+    """
+    obj = self.copy()
+    obj.resize_inplace(width, use_strand)
+    return obj
+
+  def resize_inplace(self, width: int, use_strand: bool = True) -> None:
+    """Resizes the interval in place, centered around the original center.
+
+    Args:
+      width: The new width of the interval.
+      use_strand: If True, resizing considers the strand orientation.
+    """
+    if width < 0:
+      raise ValueError(f'Width needs to be > 0. Found: {width}.')
+
+    if width is None or self.width == width:
+      return
+
+    center = self.center()
+    if use_strand and self.negative_strand:
+      self.start = center - width // 2
+      self.end = center + (width + 1) // 2
+    else:
+      self.start = center - (width + 1) // 2
+      self.end = center + width // 2
+    assert self.width == width
+
+  def overlaps(self, interval: Self) -> bool:
+    """Checks if this interval overlaps with another interval."""
+    return (
+        self.chromosome == interval.chromosome
+        and self.start < interval.end
+        and interval.start < self.end
+    )
+
+  def contains(self, interval: Self) -> bool:
+    """Checks if this interval completely contains another interval."""
+    return (
+        self.chromosome == interval.chromosome
+        and self.start <= interval.start
+        and self.end >= interval.end
+    )
+
+  def intersect(self, interval: Self) -> Self | None:
+    """Returns the intersection of this interval with another interval."""
+    output = self.copy()
+    if not self.overlaps(interval):
+      return None
+    output.start = max(self.start, interval.start)
+    output.end = min(self.end, interval.end)
+    return output
+
+  def coverage(
+      self, intervals: Sequence[Self], *, bin_size: int = 1
+  ) -> np.ndarray:
+    """Computes coverage track from sequence of intervals overlapping interval.
+
+    This method calculates the coverage of this interval by a set of other
+    intervals. The coverage is defined as the number of intervals that overlap
+    each position within this interval.
+
+    The `bin_size` parameter allows you to bin the coverage into equal-sized
+    windows. This can be useful for summarizing coverage over larger regions.
+    If `bin_size` is 1, the coverage is calculated at single-base resolution.
+
+    Args:
+      intervals: A sequence of `Interval` objects that may overlap this
+        interval.
+      bin_size: The size of the bins used to calculate coverage. Must be a
+        positive integer that divides the width of the interval.
+
+    Returns:
+      A 1D numpy array representing the coverage track. The length of the array
+      is `self.width // bin_size`. Each element in the array represents the
+      summed coverage within the corresponding bin.
+
+    Raises:
+      ValueError: If `bin_size` is not a positive integer or if the interval
+        width is not divisible by `bin_size`.
+    """
+    if bin_size <= 0:
+      raise ValueError('bin_size needs to be larger or equal to 1.')
+    if self.width % bin_size != 0:
+      raise ValueError(
+          f'interval width {self.width} needs to be divisible '
+          f'by bin_size {bin_size}.'
+      )
+    output = np.zeros((self.width,), dtype=np.int32)
+    for interval in intervals:
+      if not self.overlaps(interval):
+        continue
+      relative_start = max(interval.start - self.start, 0)
+      relative_end = min(interval.end - self.start, self.width)
+      output[relative_start:relative_end] += 1
+    if bin_size > 1:
+      return output.reshape((self.width // bin_size, bin_size)).sum(axis=-1)
+    else:
+      return output
+
+  def overlap_ranges(
+      self,
+      intervals: Sequence[Self],
+  ) -> np.ndarray:
+    """Returns overlapping ranges from intervals overlapping this interval.
+
+    Args:
+      intervals: Sequence of candidate intervals to test for overlap.
+
+    Returns:
+      2D numpy array indicating the start and end of the overlapping ranges.
+    """
+    output = []
+    for interval in intervals:
+      if not self.overlaps(interval):
+        continue
+      relative_start = max(interval.start - self.start, 0)
+      relative_end = min(interval.end - self.start, self.width)
+      output.append([relative_start, relative_end])
+
+    return (
+        np.asarray(output, dtype=np.int32)
+        if output
+        else np.empty((0, 2), dtype=np.int32)
+    )
+
+  def binary_mask(
+      self, intervals: Sequence[Self], bin_size: int = 1
+  ) -> np.ndarray:
+    """Boolean mask True if any interval overlaps the bin: coverage > 0."""
+    return self.coverage(intervals, bin_size=bin_size) > 0
+
+  def coverage_stranded(
+      self, intervals: Sequence[Self], *, bin_size: int = 1
+  ) -> np.ndarray:
+    """Computes a coverage track from intervals overlapping this interval.
+
+    This method considers the strand information of both self and intervals.
+
+    Args:
+      intervals: Sequence of intervals possibly overlapping self.
+      bin_size: Resolution at which to bin the output coverage track. Coverage
+        within each bin (if larger than 1) will be summarized using sum().
+
+    Returns:
+      Numpy array of shape (self.width // bin_size, 2) where output[:, 0]
+      represents coverage for intervals on the same strand as self and
+      output[:, 1] represents coverage of intervals on the opposite strand.
+    """
+    # Split intervals based on strand.
+    forward_intervals = []
+    reverse_intervals = []
+    for interval in intervals:
+      if interval.negative_strand:
+        reverse_intervals.append(interval)
+      else:
+        forward_intervals.append(interval)
+
+    coverage = np.stack(
+        [
+            self.coverage(forward_intervals, bin_size=bin_size),
+            self.coverage(reverse_intervals, bin_size=bin_size),
+        ],
+        axis=-1,
+    )
+    if self.negative_strand:
+      return coverage[::-1, ::-1]
+    else:
+      return coverage
+
+  def binary_mask_stranded(
+      self, intervals: Sequence[Self], bin_size: int = 1
+  ) -> np.ndarray:
+    """Boolean mask True if any interval overlaps the bin: coverage > 0."""
+    return self.coverage_stranded(intervals, bin_size=bin_size) > 0
+
+
+_DEFAULT_REGEX = re.compile(r'(chr(?:X|Y|M|\d+)):(\d+):([ACGTN]*)>([ACGTN]*)')
+_GTEX_REGEX = re.compile(
+    r'(chr(?:X|Y|M|\d+))_(\d+)_([ACGTN]*)_([ACGTN]*)_?[a-zA-Z0-9]*'
+)
+_OPEN_TARGETS_REGEX = re.compile(r'((?:X|Y|M|\d+))_(\d+)_([ACGTN]*)_([ACGTN]*)')
+_OPEN_TARGETS_BIGQUERY_REGEX = re.compile(
+    r'((?:X|Y|M|\d+)):(\d+):([ACGTN]*):([ACGTN]*)'
+)
+_GNOMAD_REGEX = re.compile(r'((?:X|Y|M|\d+))-(\d+)-([ACGTN]*)-([ACGTN]*)')
+
+
+class VariantFormat(enum.Enum):
+  """A format for parsing a string into a Variant object.
+
+  This is used to convert from a string to a formal Variant object.
+  Note that it does not perform any validation (e.g. it does not verify that the
+  reference allele corresponds to the reference genome or that the
+  position is valid).
+
+  Example formats:
+    DEFAULT: chr22:1024:A>C
+    GTEX: chr22_1024_A_C_b38 (build suffix is optional)
+    OPEN_TARGETS: 22_1024_A_C (chr prefix is omitted)
+    OPEN_TARGETS_BIGQUERY: 22:1024:A:C (chr prefix is omitted)
+    GNOMAD: 22-1024-A-C (chr prefix is omitted)
+  """
+
+  DEFAULT = 'default'
+  GTEX = 'gtex'
+  OPEN_TARGETS = 'open_targets'
+  OPEN_TARGETS_BIGQUERY = 'open_targets_bigquery'
+  GNOMAD = 'gnomad'
+
+  def to_regex(self) -> re.Pattern[str]:
+    """Returns a regular expression for the variant format."""
+    match self:
+      case VariantFormat.DEFAULT:
+        return _DEFAULT_REGEX
+      case VariantFormat.GTEX:
+        return _GTEX_REGEX
+      case VariantFormat.OPEN_TARGETS:
+        return _OPEN_TARGETS_REGEX
+      case VariantFormat.OPEN_TARGETS_BIGQUERY:
+        return _OPEN_TARGETS_BIGQUERY_REGEX
+      case VariantFormat.GNOMAD:
+        return _GNOMAD_REGEX
+
+
+@dataclasses.dataclass
+class Variant:
+  """Represents a genomic variant/mutation.
+
+  Differs from the Variant definition in a VCF file, which allows
+  for multiple alternative bases and contains sample information. This
+  `Variant` class does not include sample information or variant call
+  quality information.
+
+  Attributes:
+    chromosome: The chromosome name (e.g., 'chr1', '1').
+    position: The 1-based position of the variant on the chromosome.
+    reference_bases: The reference base(s) at the variant position. Most
+      frequently (not always!), these correspond to the sequence in the
+      reference genome at positions: [position, ..., position +
+      len(reference_bases) - 1]
+    alternate_bases: The alternate base(s) that replace the reference. For
+      example, if sequence='ACT', position=2, reference_bases='C',
+      alternate_bases='TG', then the actual (alternate) sequence would be ATGT.
+    name: An optional name for the variant (e.g., a dbSNP ID like rs206437).
+    info: An optional dictionary for additional variant information.
+  """
+
+  chromosome: str
+  position: int
+  reference_bases: str
+  alternate_bases: str
+  name: str = dataclasses.field(default='', compare=False, hash=False)
+  info: dict[str, Any] = dataclasses.field(
+      default_factory=dict, repr=False, compare=False, hash=False
+  )
+
+  def __post_init__(self):
+    """Validates the variant's position."""
+    if self.position < 1:
+      raise ValueError(f'Position has to be >=1. Found: {self.position}.')
+    if not set(self.reference_bases).issubset(VALID_VARIANT_BASES):
+      raise ValueError(
+          f'Invalid reference bases: "{self.reference_bases}". Must only'
+          ' contain "ACGTN".'
+      )
+    if not set(self.alternate_bases).issubset(VALID_VARIANT_BASES):
+      raise ValueError(
+          f'Invalid alternate bases: "{self.alternate_bases}". Must only'
+          ' contain "ACGTN".'
+      )
+
+  def __str__(self):
+    """Returns a string representation of the variant."""
+    ref_alt = f'{self.reference_bases}>{self.alternate_bases}'
+    return f'{self.chromosome}:{self.position}:{ref_alt}'
+
+  def as_truncated_str(self, max_length: int = 50):
+    """Truncates the variant str's ref and alt bases to the given max length."""
+
+    def _truncate(s: str):
+      if len(s) <= max_length:
+        return s
+      else:
+        return s[: max_length // 2] + '...' + s[-max_length // 2 :]
+
+    return (
+        f'{self.chromosome}:{self.position}:'
+        f'{_truncate(self.reference_bases)}>{_truncate(self.alternate_bases)}'
+    )
+
+  @property
+  def start(self) -> int:
+    """Returns the 0-based start position of the variant."""
+    return self.position - 1
+
+  @property
+  def end(self) -> int:
+    """Returns the 0-based end position of the variant."""
+    return self.start + len(self.reference_bases)
+
+  @property
+  def reference_interval(self) -> Interval:
+    """Returns an `Interval` for the variant's reference sequence."""
+    return Interval(self.chromosome, self.start, self.end)
+
+  def reference_overlaps(self, interval: Interval) -> bool:
+    """Checks if the variant's reference overlaps with the interval."""
+    return interval.overlaps(Interval(self.chromosome, self.start, self.end))
+
+  def alternate_overlaps(self, interval: Interval) -> bool:
+    """Checks if the variant's alternate overlaps with the interval."""
+    return interval.overlaps(
+        Interval(
+            self.chromosome, self.start, self.start + len(self.alternate_bases)
+        )
+    )
+
+  @property
+  def is_snv(self) -> bool:
+    """Return if the variant is a Single Nucleotide Variant (SNV)."""
+    return len(self.reference_bases) == 1 and len(self.alternate_bases) == 1
+
+  @property
+  def is_deletion(self) -> bool:
+    """Return if the variant is a deletion."""
+    return len(self.reference_bases) > len(self.alternate_bases)
+
+  @property
+  def is_insertion(self) -> bool:
+    """Return if the variant is an insertion."""
+    return len(self.reference_bases) < len(self.alternate_bases)
+
+  @property
+  def is_frameshift(self) -> bool:
+    """Return if the variant is a frameshift."""
+    indel_size = abs(len(self.reference_bases) - len(self.alternate_bases))
+    return indel_size > 0 and indel_size % 3 != 0
+
+  @property
+  def is_indel(self) -> bool:
+    """Return if the variant is an insertion or deletion."""
+    return self.is_insertion or self.is_deletion
+
+  @property
+  def is_structural(self) -> bool:
+    """Return if the variant is a structural variant."""
+    indel_size = abs(len(self.reference_bases) - len(self.alternate_bases))
+    return indel_size >= 50
+
+  def copy(self) -> Self:
+    """Returns a deep copy of the variant."""
+    return copy.deepcopy(self)
+
+  @classmethod
+  def from_str(
+      cls, string: str, variant_format: VariantFormat = VariantFormat.DEFAULT
+  ) -> Self:
+    """Creates a `Variant` from a string representation.
+
+    Args:
+      string: The string representation.
+      variant_format: The format of the variant string. By default, this uses
+        "chromosome:position:ref>alt" (for example, "chr1:1024:A>C"). See
+        VariantFormat for alternate formatting options.
+
+    Returns:
+      A `Variant` object.
+    """
+    result = re.fullmatch(variant_format.to_regex(), string)
+    if result is None:
+      raise ValueError(f'Invalid format for variant string: {string}')
+
+    chromosome, position, reference, alternate = result.groups()
+    # Add chr prefix if not already present.
+    if not chromosome.startswith('chr'):
+      chromosome = f'chr{chromosome}'
+    return cls(
+        chromosome=chromosome,
+        position=int(position),
+        reference_bases=reference,
+        alternate_bases=alternate,
+    )
+
+  def to_dict(self) -> dict[str, Any]:
+    """Converts the variant to a dictionary."""
+    return dataclasses.asdict(self)
+
+  @classmethod
+  def from_dict(cls, dictionary: Mapping[str, Any] | Self) -> Self:
+    """Creates a `Variant` from a dictionary."""
+    return cls(**dictionary)  # pytype: disable=bad-return-type
+
+  def to_proto(self) -> dna_model_pb2.Variant:
+    """Converts the variant to a protobuf message."""
+    return dna_model_pb2.Variant(
+        chromosome=self.chromosome,
+        position=self.position,
+        reference_bases=self.reference_bases,
+        alternate_bases=self.alternate_bases,
+    )
+
+  @classmethod
+  def from_proto(cls, variant: dna_model_pb2.Variant) -> Self:
+    """Creates a `Variant` from a protobuf message."""
+    return cls(
+        chromosome=variant.chromosome,
+        position=variant.position,
+        reference_bases=variant.reference_bases,
+        alternate_bases=variant.alternate_bases,
+    )
+
+  def split(self, anchor: int) -> tuple[Self | None, Self | None]:
+    """Splits the variant into two at the anchor point.
+
+    If the anchor point falls within the variant's reference sequence, the
+    variant is split into two new variants: one upstream of the anchor and one
+    downstream. If the anchor is outside the variant's reference sequence,
+    the original variant is returned on the appropriate side, and None on the
+    other.
+
+    Example:
+      position=       3
+      ref:       ...[ A C ]...
+      alt:       .....T G T C
+      anchor=3         |
+      returns:    (chr1:3:A>T, chr1:4:C>GTC)
+
+    Args:
+      anchor: The 0-based anchor point to split the variant.
+
+    Returns:
+      A tuple of the upstream and downstream variants. If the variant is only
+      on one side of the anchorpoint, then None is returned.
+    """
+    if anchor <= self.start:
+      return None, self.copy()
+    elif anchor >= self.end:
+      return self.copy(), None
+    else:
+      mid = anchor - self.start
+      upstream, downstream = self.copy(), self.copy()
+      upstream.reference_bases = self.reference_bases[:mid]
+      upstream.alternate_bases = self.alternate_bases[:mid]
+
+      downstream.position = anchor + 1
+      downstream.reference_bases = self.reference_bases[mid:]
+      downstream.alternate_bases = self.alternate_bases[mid:]
+      return upstream, downstream
+
+
+@dataclasses.dataclass
+class Junction(Interval):
+  """Represents a splice junction.
+
+  A splice junction is a point in a pre-mRNA transcript where an intron is
+  removed and exons are joined during RNA splicing. This class inherits from
+  `Interval` and adds properties and methods specific to splice junctions.
+
+  Attributes:
+    chromosome: The chromosome name.
+    start: The 0-based start position of the junction.
+    end: The 0-based end position of the junction.
+    strand: The strand of the junction ('+' or '-').
+    name: An optional name for the junction.
+    info: An optional dictionary to store additional information.
+    k: An optional integer representing the number of reads supporting the
+      splice junction.
+
+  Raises:
+    ValueError: If the strand is unstranded.
+  """
+
+  k: int | None = None
+
+  def __post_init__(self):
+    """Validates that the junction is stranded."""
+    super().__post_init__()
+    if self.strand == STRAND_UNSTRANDED:
+      raise ValueError('Junctions must be stranded.')
+
+  @property
+  def acceptor(self) -> int:
+    """Returns the acceptor site position."""
+    return self.start if self.strand == STRAND_NEGATIVE else self.end
+
+  @property
+  def donor(self) -> int:
+    """Returns the donor site position."""
+    return self.end if self.strand == STRAND_NEGATIVE else self.start
+
+  def dinucleotide_region(self) -> tuple[Interval, Interval]:
+    """Returns the dinucleotide regions around acceptor and donor sites."""
+    return (
+        Interval(
+            self.chromosome, self.start, self.start + 2, strand=self.strand
+        ),
+        Interval(self.chromosome, self.end - 2, self.end, strand=self.strand),
+    )
+
+  def acceptor_region(self, overhang: tuple[int, int] = (250, 250)) -> Interval:
+    """Returns the region around the acceptor site with overhang."""
+    return Interval(
+        self.chromosome, self.acceptor, self.acceptor, strand=self.strand
+    ).pad(start_pad=overhang[0], end_pad=overhang[1])
+
+  def donor_region(self, overhang: tuple[int, int] = (250, 250)) -> Interval:
+    """Returns the region around the donor site with overhang."""
+    return Interval(
+        self.chromosome, self.donor, self.donor, strand=self.strand
+    ).pad(start_pad=overhang[0], end_pad=overhang[1])
+
+
+class _FastaExtractorType(Protocol):
+  """Protocol definition for extracting intervals from a Fasta file."""
+
+  def extract(self, interval: Interval) -> str:
+    """Extract and return the DNA sequence for a given interval."""
+
+
+def _prefix_length(*sequences) -> int:
+  """Returns the length of the common prefix for a sequence of strings."""
+  i = 0
+  for chars in zip(*sequences, strict=False):
+    if all(c == chars[0] for c in chars):
+      i += 1
+    else:
+      break
+  return i
+
+
+def normalize_variant(
+    variant: Variant, extractor: _FastaExtractorType
+) -> Variant:
+  """Normalize a Variant by left-aligning the reference and alternate bases.
+
+  Normalization applied following algorithm described in
+  https://doi.org/10.1093/bioinformatics/btv112.
+
+  Args:
+    variant: The Variant to normalize.
+    extractor: The FastaExtractor to use for extracting the reference sequence.
+
+  Returns:
+    The normalized Variant.
+  """
+  if variant.is_snv:
+    return variant
+
+  chromosome = variant.chromosome
+  genome_sequence = extractor.extract(
+      Interval(
+          chromosome,
+          variant.start,
+          variant.start
+          + max(len(variant.reference_bases), len(variant.alternate_bases)),
+      )
+  )
+
+  position = variant.position
+  alleles = [genome_sequence, variant.reference_bases, variant.alternate_bases]
+
+  # Remove any common suffix from the alleles.
+  finished = False
+  while not finished:
+    suffix_length = _prefix_length(*[reversed(a) for a in alleles])
+    if suffix_length > 0:
+      alleles = [a[:-suffix_length] for a in alleles]
+
+    # If any alleles are empty, extend all alleles by 1 nucleotide to the left.
+    if not all(alleles):
+      position -= 1
+      base = extractor.extract(Interval(chromosome, position - 1, position))[0]
+      alleles = [base + a for a in alleles]
+    else:
+      # Finished iff we haven't shifted alleles and don't have a common suffix.
+      finished = suffix_length == 0
+
+  # Left-align the variant to the genome, ensuring at least 1 bp of context.
+  max_prefix_length = len(min(alleles)) - 1
+  prefix_length = _prefix_length(*alleles)
+  i = min(max_prefix_length, prefix_length)
+
+  _, reference_bases, alternate_bases = alleles
+  return Variant(
+      chromosome=chromosome,
+      position=position + i,
+      reference_bases=reference_bases[i:],
+      alternate_bases=alternate_bases[i:],
+  )
+
+
+def _split_intervals(
+    intervals: Iterable[Interval], marker: int, bounds: list[tuple[int, int]]
+):
+  """Splits intervals into start and end points with markers."""
+  for i in intervals:
+    bounds.append((i.start, +marker))
+    bounds.append((i.end, -marker))
+
+
+def _group_by_chromosome(
+    intervals: Iterable[Interval],
+) -> dict[str, list[Interval]]:
+  """Groups intervals by chromosome."""
+  interval_map = collections.defaultdict(list)
+  for i in intervals:
+    interval_map[i.chromosome].append(i)
+  return dict(interval_map)
+
+
+def intersect_intervals(
+    lhs: Iterable[Interval],
+    rhs: Iterable[Interval],
+    *,
+    result_strand: str = '.',
+) -> Iterator[Interval]:
+  """Generates the intersection of two interval sets.
+
+  Point ranges (width == 0) are considered to intersect any range which contains
+  them.
+
+  In these examples, the intersection is a point range (2,2)
+  ```
+      1  2  3  4
+    ..|..|..|..|..
+          <>        (start=2, end=2)
+    ------->        (..., end=2)
+  ```
+
+  ```
+      1  2  3  4
+    ..|..|..|..|..
+          <>        (start=2, end=2)
+          <-------  (start=2, end=...)
+  ```
+
+  For consistency, this means that the following results in a point
+  intersection:
+
+  ```
+      1  2  3  4
+    ..|..|..|..|..
+    ------->        (start=..., end=2)
+          <-------  (start=2, end=...)
+  ```
+
+  Args:
+    lhs: A set of intervals.
+    rhs: A set of intervals.
+    result_strand: The strand for the resulting intervals.
+
+  Yields:
+    The intersection of intervals. Overlapping intervals within either `lhs`
+      or `rhs` are implicitly unioned.
+  """
+
+  def _intersect(lhs, rhs, chrom):
+    """Calculates the intersection for a specific chromosome.
+
+    Deconstructs two sets of intervals into (start, +k) and (end, -k) positions,
+    where k is different for each set, then sorts all interval endpoints. By
+    walking the sorted endpoints and accumulating the k's, we can work out when
+    we enter or exit an interval from either set.
+
+    This allows emitting intervals corresponding to the intersection, and with
+    appropriate k's, detecting when we enter an interval in either set before
+    exiting the previous.
+
+    Args:
+      lhs: The first list of `Interval` objects.
+      rhs: The second list of `Interval` objects.
+      chrom: The chromosome name.
+
+    Yields:
+      `Interval` objects representing the intersections.
+    """
+    bounds = []
+    _split_intervals(lhs, 0x00001, bounds)
+    _split_intervals(rhs, 0x10000, bounds)
+    accum = 0
+    start = None
+    for pos, delta in sorted(bounds, key=lambda x: (x[0], -x[1])):
+      old_accum = accum
+      accum += delta
+      if accum & 0xFFFF and accum & 0xFFFF0000:
+        if start is None:
+          start = pos
+      elif old_accum & 0xFFFF and old_accum & 0xFFFF0000:
+        yield Interval(chrom, start, pos, result_strand)
+        start = None
+
+  lhs = _group_by_chromosome(lhs)
+  rhs = _group_by_chromosome(rhs)
+  for chromosome in set(lhs) & set(rhs):
+    yield from _intersect(lhs[chromosome], rhs[chromosome], chromosome)
+
+
+def union_intervals(
+    lhs: Iterable[Interval],
+    rhs: Iterable[Interval],
+    *,
+    result_strand: str = '.',
+) -> Iterator[Interval]:
+  """Generates the union of two interval sets.
+
+  Args:
+    lhs: A non-overlapping set of intervals.
+    rhs: A non-overlapping set of intervals.
+    result_strand: The strand for the resulting intervals.
+
+  Yields:
+    The union of intervals. Any position covered by an interval in `lhs` or
+    `rhs` is covered in the result.
+  """
+
+  def _union(lhs, rhs, chrom):
+    """Calculates the union for a specific chromosome, analog of _intersect."""
+    bounds = []
+    _split_intervals(lhs, 0x00001, bounds)
+    _split_intervals(rhs, 0x10000, bounds)
+    accum = 0
+    start = end = None
+    for pos, delta in sorted(bounds, key=lambda x: (x[0], -x[1])):
+      if accum == 0 and end is not None and end < pos:
+        # Delay generating an interval until after the observed end point.
+        # This merges abutting ranges.
+        yield Interval(chrom, start, end, result_strand)
+        start = end = None
+      accum += delta
+      if accum:
+        if start is None:
+          start = pos
+      else:
+        assert start is not None
+        end = pos
+    if end is not None:
+      yield Interval(chrom, start, end, result_strand)
+
+  lhs = _group_by_chromosome(lhs)
+  rhs = _group_by_chromosome(rhs)
+  for chromosome in set(lhs) | set(rhs):
+    yield from _union(
+        lhs.get(chromosome, []), rhs.get(chromosome, []), chromosome
+    )
+
+
+def merge_overlapping_intervals(
+    intervals: Sequence[Interval],
+) -> list[Interval]:
+  """Merges overlapping intervals and returns a sorted list.
+
+  Args:
+    intervals: A sequence of intervals with the same strand.
+
+  Returns:
+    A new sorted list of merged intervals.
+  """
+  if not intervals:
+    return []
+  assert all(i.strand == intervals[0].strand for i in intervals)
+  return list(union_intervals(intervals, [], result_strand=intervals[0].strand))

flax_model/alphagenome/_sdk/data/junction_data.py

@@ -0,0 +1,272 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Splice junction data container.
+
+`JunctionData` stores splice junction data for a given transcript or interval.
+"""
+
+from collections.abc import Sequence
+import dataclasses
+from typing import Any
+
+from flax_model.alphagenome._sdk import typing
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import ontology
+from jaxtyping import Float, Shaped  # pylint: disable=g-multiple-import, g-importing-member
+import numpy as np
+import pandas as pd
+
+JunctionMetadata = pd.DataFrame
+
+
+@typing.jaxtyped
+@dataclasses.dataclass(frozen=True)
+class JunctionData:
+  """Container for storing splice junction data.
+
+  Attributes:
+    junctions: A numpy array representing the splice junctions.
+    values: A numpy array of floats representing the values associated with each
+      junction for each track.
+    metadata: A pandas DataFrame containing metadata for each track.
+    interval: An optional `Interval` object representing the genomic region
+      containing the junctions.
+    uns: An optional dictionary to store additional unstructured data.
+
+  Raises:
+    ValueError: If the number of tracks in `values` does not match the number
+      of rows in `metadata`, or if `metadata` contains duplicate names.
+  """
+
+  junctions: Shaped[np.ndarray, 'num_junctions']
+  values: Float[np.ndarray, 'num_junctions num_tracks']
+  metadata: JunctionMetadata
+  interval: genome.Interval | None = None
+  uns: dict[str, Any] | None = None
+
+  def __post_init__(self):
+    """Validates the consistency of the data."""
+    if self.values.shape[1] != len(self.metadata):
+      raise ValueError(
+          f'Number of tracks {self.values.shape[1]} and '
+          f'metadata {len(self.metadata)} do not match.'
+      )
+
+    if self.metadata['name'].duplicated().any():
+      raise ValueError('Metadata contain duplicated names.')
+
+  def __len__(self):
+    """Returns the number of junctions."""
+    return len(self.junctions)
+
+  @property
+  def num_tracks(self) -> int:
+    """Returns the number of tracks."""
+    return len(self.metadata)
+
+  @property
+  def names(self) -> np.ndarray:
+    """Returns an array of track names (not necessarily unique)."""
+    return self.metadata['name'].values
+
+  @property
+  def strands(self) -> np.ndarray:
+    """Returns an array of junction strands."""
+    return np.array([j.strand for j in self.junctions])
+
+  @property
+  def possible_strands(self) -> np.ndarray:
+    """All possible strands."""
+    return np.unique(self.strands)
+
+  @property
+  def ontology_terms(self) -> Sequence[ontology.OntologyTerm | None] | None:
+    """Returns a list of ontology terms (if available)."""
+    if 'ontology_curie' in self.metadata.columns:
+      return [
+          ontology.from_curie(curie) if curie is not None else None
+          for curie in self.metadata['ontology_curie'].values
+      ]
+    else:
+      return None
+
+  # Track order / subset wrangling:
+  def filter_tracks(self, mask: np.ndarray | list[bool]) -> 'JunctionData':
+    """Filters tracks by a boolean mask.
+
+    Args:
+      mask: A boolean mask to select tracks.
+
+    Returns:
+      A new `JunctionData` object with the filtered tracks.
+    """
+    return JunctionData(
+        junctions=self.junctions,
+        values=self.values[:, mask],
+        metadata=self.metadata.loc[mask],
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def filter_to_strand(self, strand: str) -> 'JunctionData':
+    """Filters junctions to a specific DNA strand.
+
+    Args:
+      strand: The strand to filter by ('+' or '-').
+
+    Returns:
+      A new `JunctionData` object with junctions on the specified strand.
+    """
+    mask = self.strands == strand
+    return JunctionData(
+        junctions=self.junctions[mask],
+        values=self.values[mask, :],
+        metadata=self.metadata,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def normalize_values(self, total_k: float = 10.0) -> 'JunctionData':
+    """Normalizes the values by the k value."""
+    values = self.values * total_k / self.values.sum()
+    return JunctionData(
+        junctions=self.junctions,
+        values=values,
+        metadata=self.metadata,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def filter_to_positive_strand(self) -> 'JunctionData':
+    """Filters junctions to the positive DNA strand."""
+    return self.filter_to_strand(genome.STRAND_POSITIVE)
+
+  def filter_to_negative_strand(self) -> 'JunctionData':
+    """Filters junctions to the negative DNA strand."""
+    return self.filter_to_strand(genome.STRAND_NEGATIVE)
+
+  def filter_by_tissue(self, tissue: str) -> 'JunctionData':
+    """Filters tracks by GTEx tissue type.
+
+    Args:
+      tissue: The GTEx tissue type to filter by.
+
+    Returns:
+      A new `JunctionData` object with tracks from the specified tissue.
+
+    Raises:
+      ValueError: If the metadata does not contain a 'gtex_tissue' column.
+    """
+    if 'gtex_tissue' not in self.metadata.columns:
+      raise ValueError(
+          'Metadata does not contain gtex_tissue column. '
+          f'Got {set(self.metadata.columns)}.'
+      )
+    return self.filter_tracks(self.metadata['gtex_tissue'] == tissue)
+
+  def filter_by_name(self, name: str) -> 'JunctionData':
+    """Filters tracks by name."""
+    return self.filter_tracks(self.metadata['name'] == name)
+
+  def filter_by_ontology(self, ontology_curie: str) -> 'JunctionData':
+    """Filters tracks by ontology term.
+
+    Args:
+      ontology_curie: The ontology term CURIE to filter by.
+
+    Returns:
+      A new `JunctionData` object with tracks associated with the specified
+        ontology term.
+
+    Raises:
+      ValueError: If the metadata does not contain an 'ontology_curie' column.
+    """
+    if 'ontology_curie' not in self.metadata.columns:
+      raise ValueError(
+          'Metadata does not contain ontology_curie column. '
+          f'Got {set(self.metadata.columns)}.'
+      )
+    return self.filter_tracks(self.metadata['ontology_curie'] == ontology_curie)
+
+  def intersect_with_interval(
+      self, interval: genome.Interval
+  ) -> 'JunctionData':
+    """Returns the intersection of the junctions and the interval."""
+    mask = np.array([j.overlaps(interval) for j in self.junctions])
+    return JunctionData(
+        junctions=self.junctions[mask],
+        values=self.values[mask, :],
+        metadata=self.metadata,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+
+def get_junctions_to_plot(
+    *,
+    predictions: JunctionData,
+    name: str,
+    strand: str,
+    k_threshold: float | None = 0.0,
+) -> list[genome.Junction]:
+  """Gets a list of junctions to plot.
+
+  Filters the junctions in the `predictions` by name and strand, and
+  applies a threshold on the `k` value (read count).
+
+  Args:
+    predictions: A `JunctionData` object containing junction predictions.
+    name: The name to filter by.
+    strand: The strand to filter by ('+' or '-').
+    k_threshold: The minimum `k` value for a junction to be included. If None,
+      use 5% of the maximum value.
+
+  Returns:
+    A list of `Junction` objects to plot.
+
+  Raises:
+    ValueError: If more than one track is found for the specified name.
+  """
+  filtered = predictions.filter_by_name(name)
+  if filtered.num_tracks > 1:
+    raise ValueError(
+        f'Expected only one ontology term, got {filtered.num_tracks}.'
+    )
+  filtered_junctions = []
+  if filtered.values.size == 0:
+    return filtered_junctions
+
+  if k_threshold is None:
+    k_threshold = filtered.values.max() * 0.05
+
+  # Round k for better visualization.
+  for interval, k in zip(filtered.junctions, filtered.values, strict=True):
+    # Filter by threshold and strand.
+    if interval.strand != strand:
+      continue
+    k = k.item()
+    if k >= k_threshold:
+      k = round(k, 2)
+
+      filtered_junctions.append(
+          genome.Junction(
+              interval.chromosome,
+              interval.start,
+              interval.end,
+              interval.strand,
+              k=k,
+          )
+      )
+  return filtered_junctions

flax_model/alphagenome/_sdk/data/ontology.py

@@ -0,0 +1,122 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Handling of biological ontologies."""
+
+from collections.abc import Sequence
+import dataclasses
+import enum
+
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+import immutabledict
+
+
+class OntologyType(enum.IntEnum):
+  """Supported ontology types.
+
+  CLO: Cell Line Ontology.
+  UBERON: Uber-anatomy ontology.
+  CL: Cell Ontology.
+  EFO: Experimental Factor Ontology.
+  NTR: New Term Requested.
+  """
+
+  CLO = 1
+  UBERON = 2
+  CL = 3
+  EFO = 4
+  NTR = 5
+
+
+_ONTOLOGY_TYPE_TO_PROTO_ENUM = immutabledict.immutabledict({
+    OntologyType.CLO: dna_model_pb2.OntologyType.ONTOLOGY_TYPE_CLO,
+    OntologyType.UBERON: dna_model_pb2.OntologyType.ONTOLOGY_TYPE_UBERON,
+    OntologyType.CL: dna_model_pb2.OntologyType.ONTOLOGY_TYPE_CL,
+    OntologyType.EFO: dna_model_pb2.OntologyType.ONTOLOGY_TYPE_EFO,
+    OntologyType.NTR: dna_model_pb2.OntologyType.ONTOLOGY_TYPE_NTR,
+})
+
+
+@dataclasses.dataclass(frozen=True)
+class OntologyTerm:
+  """A single biological ontology term.
+
+  Attributes:
+    type: The ontology type.
+    id: The ID of the term within the ontology.
+  """
+
+  type: OntologyType
+  id: int
+
+  @property
+  def ontology_curie(self) -> str:
+    """Returns the CURIE (Compact Uniform Resource Identifier) for the term."""
+    return f'{self.type.name}:{self.id:07d}'
+
+  def to_proto(self) -> dna_model_pb2.OntologyTerm:
+    """Converts the ontology term to a protobuf message."""
+    return dna_model_pb2.OntologyTerm(
+        ontology_type=_ONTOLOGY_TYPE_TO_PROTO_ENUM[self.type], id=self.id
+    )
+
+
+def from_curie(ontology_curie: str) -> OntologyTerm:
+  """Creates an `OntologyTerm` from a CURIE.
+
+  Args:
+    ontology_curie: The CURIE (Compact Uniform Resource Identifier) for the
+      ontology term.
+
+  Returns:
+    An `OntologyTerm` object.
+  """
+  try:
+    ontology_type, local_id = ontology_curie.split(':')
+  except ValueError as e:
+    raise ValueError(
+        f'Invalid {ontology_curie=}. CURIEs must be of the form <type>:<id>.'
+    ) from e
+  try:
+    ontology_type = OntologyType[ontology_type]
+  except KeyError as e:
+    raise ValueError(f'Cannot parse {ontology_type=}') from e
+  return OntologyTerm(ontology_type, int(local_id))
+
+
+def from_curies(ontology_curies: Sequence[str]) -> Sequence[OntologyTerm]:
+  """Creates a list of `OntologyTerm` objects from a list of CURIEs.
+
+  Args:
+    ontology_curies: A list of CURIEs (Compact Uniform Resource Identifiers).
+
+  Returns:
+    A list of `OntologyTerm` objects.
+  """
+  return [from_curie(curie) for curie in ontology_curies]
+
+
+def from_proto(proto: dna_model_pb2.OntologyTerm) -> OntologyTerm:
+  """Creates an `OntologyTerm` from a protobuf message.
+
+  Args:
+    proto: An `OntologyTerm` protobuf message.
+
+  Returns:
+    An `OntologyTerm` object.
+  """
+  return OntologyTerm(
+      OntologyType(proto.ontology_type),
+      proto.id,
+  )

flax_model/alphagenome/_sdk/data/track_data.py

@@ -0,0 +1,918 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Track data container analogous to AnnData."""
+
+from collections.abc import Sequence
+import copy
+import dataclasses
+import enum
+from typing import Any, Union
+
+from flax_model.alphagenome._sdk import typing
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import ontology
+from jaxtyping import Bool, Float32, Int32  # pylint: disable=g-multiple-import, g-importing-member
+import numpy as np
+import pandas as pd
+
+# Required columns: name, strand.
+# Optional standardized columns: cell_type, assay, padding.
+TrackMetadata = pd.DataFrame
+PositionalIndex = slice | genome.Interval | int
+TrackIndex = np.ndarray | Sequence[int] | Sequence[str] | slice | int | str
+Index = PositionalIndex | tuple[PositionalIndex, TrackIndex] | TrackIndex
+
+
+@enum.unique
+class AggregationType(enum.Enum):
+  """Aggregation types for downsampling/upsampling track resolutions.
+
+  SUM: Sum pooling, where values within a bin are summed. This is recommended
+    for continuous tracks where the total value within a bin is meaningful (e.g.
+    read counts, coverage).
+  MAX: Max pooling, where the maximum value within a bin is selected. This is
+    recommended for binary tracks (e.g., gene masks, regions of interest).
+  """
+
+  SUM = 'sum'
+  MAX = 'max'
+
+
+@typing.jaxtyped
+@dataclasses.dataclass(frozen=True)
+class TrackData:
+  """Container for storing track values and metadata.
+
+  `TrackData` stores multiple genomic tracks at the same resolution, stacked
+  into an ND matrix of shape (positional_bins, num_tracks). It also contains
+  metadata information as a pandas DataFrame with `num_tracks` rows.
+
+  Metadata DataFrame has two main required columns:
+
+    * name: The name of the track.
+    * strand: The strand of the track ('+', '-', or '.').
+
+  Other columns are optional.
+
+  Valid shapes of `TrackData.values` are:
+
+    * [num_tracks]
+    * [positional_bins, num_tracks]
+    * [positional_bins, positional_bins, num_tracks]
+    * ...
+
+  `TrackData` can store both model predictions and raw data. It can
+  optionally hold information about the `genome.Interval` from which the data
+  were derived and `.uns` for storing additional unstructured data.
+
+  In addition to being a container, `TrackData` provides functionality for
+  common aggregation and slicing operations.
+
+  Attributes:
+    values: A numpy array of floats or integers representing the track values.
+      Positional axes have the same length. Example valid shapes are:
+      [num_tracks], [positional_bins, num_tracks], and [positional_bins,
+      positional_bins, num_tracks].
+    metadata: A pandas DataFrame containing metadata for each track. The
+      DataFrame must have at least two columns: 'name' and 'strand'.
+    resolution: The resolution of the track data in base pairs.
+    interval: An optional `Interval` object representing the genomic region.
+    uns: An optional dictionary to store additional unstructured data.
+
+  Raises:
+    ValueError: If the number of tracks in `values` does not match the  number
+      of rows in `metadata`, or if `metadata` contains duplicate (name, strand)
+      pairs, or if the positional axes have different lengths, or if the
+      interval width does not match the expected width.
+  """
+
+  # Use Union due to https://github.com/patrick-kidger/jaxtyping/issues/73.
+  values: Union[
+      Float32[np.ndarray, '*positional_bins num_tracks'],
+      Int32[np.ndarray, '*positional_bins num_tracks'],
+      Bool[np.ndarray, '*positional_bins num_tracks'],
+  ]
+  metadata: TrackMetadata
+  resolution: int = 1
+  interval: genome.Interval | None = None
+  uns: dict[str, Any] | None = (
+      # Unstructured data dict, analagous to anndata.AnnData.uns.
+      None
+  )
+
+  def __post_init__(self):
+    """Validates the consistency of the data."""
+    if self.values.shape[-1] != len(self.metadata):
+      raise ValueError(
+          f'value number of tracks {self.values.shape[-1]} and '
+          f'metadata {len(self.metadata)} do not match.'
+      )
+
+    if self.positional_axes:
+      if len(set(np.array(self.values.shape)[self.positional_axes])) != 1:
+        raise ValueError('All positional axes must have the same length.')
+
+      if self.interval and self.interval.width != self.width:
+        raise ValueError(
+            f'Interval width must match expected width. {self.interval.width=},'
+            f' {self.width=}'
+        )
+
+    if not {'name', 'strand'}.issubset(self.metadata.columns):
+      raise ValueError('Metadata must contain columns "name" and "strand".')
+
+    if self.metadata[['name', 'strand']].duplicated().any():
+      raise ValueError(
+          'Metadata contain duplicated values for (name, strand) tuples.'
+      )
+
+  @property
+  def positional_axes(self) -> list[int]:
+    """Returns a list of the positional axes."""
+    return list(range(self.values.ndim - 1))
+
+  @property
+  def num_tracks(self) -> int:
+    """Returns the number of tracks."""
+    return self.values.shape[-1]
+
+  @property
+  def width(self) -> int:
+    """Returns the interval width covered by the tracks."""
+    if self.positional_axes:
+      return self.values.shape[0] * self.resolution
+    else:
+      return 0
+
+  @property
+  def names(self) -> np.ndarray:
+    """Returns an array of track names (not necessarily unique)."""
+    return self.metadata['name'].values
+
+  @property
+  def strands(self) -> np.ndarray:
+    """Returns an array of track strands."""
+    return self.metadata['strand'].values
+
+  @property
+  def ontology_terms(self) -> Sequence[ontology.OntologyTerm | None] | None:
+    """Returns a list of ontology terms (if available)."""
+    if 'ontology_curie' in self.metadata.columns:
+      return [
+          ontology.from_curie(curie) if curie is not None else None
+          for curie in self.metadata['ontology_curie'].values
+      ]
+    else:
+      return None
+
+  def copy(self) -> 'TrackData':
+    """Returns a deep copy of the `TrackData` object."""
+    if self.interval:
+      interval = self.interval.copy()
+    else:
+      interval = None
+    return TrackData(
+        self.values.copy(),
+        resolution=self.resolution,
+        metadata=self.metadata.copy(),
+        interval=interval,
+        uns=copy.deepcopy(self.uns),
+    )
+
+  def bin_index(self, relative_position: int) -> int:
+    """Returns the bin index for a relative position.
+
+    Args:
+      relative_position: The relative position within the interval.
+
+    Returns:
+      The corresponding bin index.
+    """
+    return relative_position // self.resolution
+
+  def slice_by_positions(self, start: int, end: int) -> 'TrackData':
+    """Slices the track data along the positional axes.
+
+    The slicing follows Python slicing conventions (0 indexed, and includes
+    elements up to end-1).
+
+    Args:
+      start: The 1-bp resolution start position for slicing.
+      end: The 1-bp resolution end position for slicing.
+
+    Returns:
+      A new `TrackData` object with the sliced values.
+
+    Raises:
+      ValueError: If (end - start) is greater than the width, or if (end -
+      start) is not divisible by the resolution.
+    """
+    if (end - start) > self.width:
+      raise ValueError(
+          'When slicing track data, (end - start) must be less than or '
+          'equal to width.'
+      )
+
+    if (end - start) % self.resolution != 0:
+      raise ValueError(
+          f'end - start needs to be to be divisible by {self.resolution=}'
+      )
+
+    sl = slice(self.bin_index(start), self.bin_index(end))
+    slice_list = [slice(None)] * self.values.ndim
+    for i in self.positional_axes:
+      slice_list[i] = sl
+
+    interval = self.interval
+    if interval:
+      interval = genome.Interval(
+          interval.chromosome,
+          interval.start + start,
+          interval.start + end,
+          strand=interval.strand,
+          name=interval.name,
+          info=interval.info,
+      )
+
+    return TrackData(
+        self.values[tuple(slice_list)],
+        resolution=self.resolution,
+        metadata=self.metadata,
+        interval=interval,
+        uns=self.uns,
+    )
+
+  def slice_by_interval(
+      self, interval: genome.Interval, match_resolution: bool = False
+  ) -> 'TrackData':
+    """Slices the track data using a `genome.Interval`.
+
+    Args:
+      interval: The interval to slice to.
+      match_resolution: If True, the interval will first be extended to make
+        sure the width is divisible by resolution.
+
+    Returns:
+      A new `TrackData` object sliced to the interval.
+
+    Raises:
+      ValueError: If `.interval` is not specified or if the specified interval
+        is not fully contained within the current interval.
+    """
+    if self.interval is None:
+      raise ValueError(
+          '.interval is needs to be specified for slice_by_interval.'
+      )
+    if not self.interval.contains(interval):
+      raise ValueError(
+          f'Interval {self.interval=} does not fully contain {interval=}.'
+      )
+    start = interval.start - self.interval.start
+    end = interval.end - self.interval.start
+
+    if match_resolution and self.resolution != 1:
+      start = int(np.floor(start / self.resolution) * self.resolution)
+      end = int(np.ceil(end / self.resolution) * self.resolution)
+    return self.slice_by_positions(start, end)
+
+  def pad(self, start_pad: int, end_pad: int) -> 'TrackData':
+    """Pads the track data along positional axes.
+
+    Args:
+      start_pad: The amount of padding to add at the beginning.
+      end_pad: The amount of padding to add at the end.
+
+    Returns:
+      A new `TrackData` object with padded values.
+
+    Raises:
+      ValueError: If `start_pad` or `end_pad` is not divisible by the
+      resolution.
+    """
+    if start_pad == 0 and end_pad == 0:
+      return self
+    if start_pad % self.resolution != 0:
+      raise ValueError(f'start_pad needs to be divisible by {self.resolution}')
+    if end_pad % self.resolution != 0:
+      raise ValueError(f'end_pad needs to be divisible by {self.resolution}')
+
+    pad = [(0, 0)] * self.values.ndim
+    for axis in self.positional_axes:
+      pad[axis] = (start_pad // self.resolution, end_pad // self.resolution)
+
+    return TrackData(
+        np.pad(self.values, tuple(pad)),
+        resolution=self.resolution,
+        metadata=self.metadata,
+        interval=None,  # Padding invalidates the interval.
+        uns=self.uns,
+    )
+
+  def resize(self, width: int) -> 'TrackData':
+    """Resizes the track data by cropping or padding with a fixed center.
+
+    Args:
+      width: The desired width in base pairs.
+
+    Returns:
+      A new `TrackData` object with resized values.
+
+    Raises:
+      ValueError: If `width` is not divisible by the resolution.
+    """
+    if width == self.width:
+      return self
+    elif width > self.width:
+      if width % self.resolution != 0:
+        raise ValueError(f'width needs to be divisible by {self.resolution}')
+      pad_amount = (width - self.width) // self.resolution
+      pad_start = (pad_amount // 2 + pad_amount % 2) * self.resolution
+      pad_end = (pad_amount // 2) * self.resolution
+      return self.pad(pad_start, pad_end)
+    else:
+      crop_amount = (self.width - width) // self.resolution
+      start = (crop_amount // 2 + crop_amount % 2) * self.resolution
+      return self.slice_by_positions(start, start + width)
+
+  def upsample(
+      self,
+      resolution: int,
+      aggregation_type: AggregationType = AggregationType.SUM,
+  ) -> 'TrackData':
+    """Upsamples the track data to a higher resolution.
+
+    Args:
+      resolution: The desired resolution in base pairs.
+      aggregation_type: The aggregation method to use for pooling the values.
+
+    Returns:
+      A new `TrackData` object with upsampled values.
+
+    Raises:
+      ValueError: If `resolution` is not lower than the current resolution
+        or not divisible by the current resolution.
+    """
+    if resolution == self.resolution:
+      return self
+    if resolution > self.resolution:
+      raise ValueError(f'Resolution must be lower than {self.resolution}')
+    repeat = self.resolution // resolution
+    if self.resolution % resolution != 0:
+      raise ValueError(f'Resolution not divisible by {resolution}')
+
+    values = self.values
+    for axis in self.positional_axes:
+      values = np.repeat(values, repeat, axis=axis)
+      match aggregation_type:
+        case AggregationType.SUM:
+          values = values / repeat
+        case AggregationType.MAX:
+          pass
+    return TrackData(
+        values,
+        resolution=resolution,
+        metadata=self.metadata,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def downsample(
+      self,
+      resolution: int,
+      aggregation_type: AggregationType = AggregationType.SUM,
+  ) -> 'TrackData':
+    """Downsamples the track data to a lower resolution.
+
+    Args:
+      resolution: The desired resolution in base pairs.
+      aggregation_type: The aggregation method to use for pooling the values.
+
+    Returns:
+      A new `TrackData` object with downsampled values.
+
+    Raises:
+      ValueError: If `resolution` is not greater than the current resolution
+        or not divisible by the current resolution.
+    """
+    if resolution == self.resolution:
+      return self
+    if resolution < self.resolution:
+      raise ValueError(f'Resolution must be greater than {self.resolution}')
+    if resolution % self.resolution != 0:
+      raise ValueError(f'Resolution not divisible by {resolution}')
+    pool_width = resolution // self.resolution
+
+    values = self.values
+    for axis in self.positional_axes:
+      # Bring axis of interest to the front, reshape and aggregate, and reswap
+      values = np.swapaxes(values, 0, axis)
+      shape = list(values.shape)
+      reshaped_values = values.reshape(
+          [shape[0] // pool_width, pool_width] + shape[1:]
+      )
+      match aggregation_type:
+        case AggregationType.SUM:
+          values = reshaped_values.sum(axis=1)
+        case AggregationType.MAX:
+          values = reshaped_values.max(axis=1)
+      values = np.swapaxes(values, 0, axis)
+
+    return TrackData(
+        values,
+        resolution=resolution,
+        metadata=self.metadata,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def change_resolution(
+      self,
+      resolution: int,
+      aggregation_type: AggregationType = AggregationType.SUM,
+  ) -> 'TrackData':
+    """Changes the resolution of the track data.
+
+    Args:
+      resolution: The desired resolution in base pairs.
+      aggregation_type: The aggregation method to use for pooling the values.
+
+    Returns:
+      A new `TrackData` object with the new resolution.
+    """
+    if resolution >= self.resolution:
+      return self.downsample(resolution, aggregation_type)
+    else:
+      return self.upsample(resolution, aggregation_type)
+
+  def filter_tracks(self, mask: np.ndarray | list[bool]) -> 'TrackData':
+    """Filters tracks by a boolean mask.
+
+    Args:
+      mask: A boolean mask to select tracks.
+
+    Returns:
+      A new `TrackData` object with the filtered tracks.
+    """
+    return TrackData(
+        self.values[..., mask],
+        resolution=self.resolution,
+        metadata=self.metadata.iloc[mask],
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def filter_to_positive_strand(self) -> 'TrackData':
+    """Filters tracks to the positive DNA strand."""
+    return self.filter_tracks(self.strands == genome.STRAND_POSITIVE)
+
+  def filter_to_negative_strand(self) -> 'TrackData':
+    """Filters tracks to the negative DNA strand."""
+    return self.filter_tracks(self.strands == genome.STRAND_NEGATIVE)
+
+  def filter_to_nonnegative_strand(self) -> 'TrackData':
+    """Filters tracks to the non-negative DNA strands (positive and unstranded)."""
+    return self.filter_tracks(self.strands != genome.STRAND_NEGATIVE)
+
+  def filter_to_nonpositive_strand(self) -> 'TrackData':
+    """Filters tracks to the non-positive DNA strands (negative and unstranded)."""
+    return self.filter_tracks(self.strands != genome.STRAND_POSITIVE)
+
+  def filter_to_stranded(self) -> 'TrackData':
+    """Filters tracks to stranded tracks (excluding unstranded)."""
+    return self.filter_tracks(self.strands != genome.STRAND_UNSTRANDED)
+
+  def filter_to_unstranded(self) -> 'TrackData':
+    """Filters tracks to unstranded tracks."""
+    return self.filter_tracks(self.strands == genome.STRAND_UNSTRANDED)
+
+  def select_tracks_by_index(
+      self, idx: np.ndarray | Sequence[int]
+  ) -> 'TrackData':
+    """Selects tracks by numerical index.
+
+    Args:
+      idx: A list or array of numerical indices to select tracks.
+
+    Returns:
+      A new `TrackData` object with the selected tracks.
+    """
+    return TrackData(
+        self.values[..., idx],
+        resolution=self.resolution,
+        metadata=self.metadata.iloc[idx],
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def select_tracks_by_name(
+      self, names: np.ndarray | Sequence[str]
+  ) -> 'TrackData':
+    """Selects tracks by name.
+
+    Args:
+      names: A list or array of track names to select.
+
+    Returns:
+      A new `TrackData` object with the selected tracks.
+    """
+    track_idx = pd.Series(np.arange(self.num_tracks), index=self.names)
+    return self.select_tracks_by_index(track_idx.loc[names].values)
+
+  def __getitem__(self, index: Index) -> 'TrackData':
+    """Retrieves a subset of TrackData using positional and/or track indices.
+
+    This method allows slicing `TrackData` similar to numpy arrays or pandas
+    DataFrames. The index can be a single value or a tuple.
+
+    Args:
+      index: A single index or a tuple of indices. If a single index, it's
+        treated as a positional index if the `TrackData` has positional axes,
+        otherwise as a track index. If a tuple, the first element specifies the
+        positional slice, and the second element specifies the track index.
+        Positional indices can be `int`, `slice`, or `genome.Interval`, and this
+        slice is applied to all positional axes of the `values` array. Track
+        indices can be `int`, `str` (track name), `slice`, `Sequence[int]`, or
+        `Sequence[str]` (track names).
+
+    Returns:
+      A new `TrackData` object containing the selected subset.
+
+    Raises:
+      IndexError: If a slice step is not 1 for positional indexing.
+      IndexError: If an unsupported index type is provided.
+    """
+    if isinstance(index, tuple):
+      position_index, track_index = index
+    elif self.positional_axes:
+      position_index, track_index = index, None
+    else:
+      position_index, track_index = None, index
+      if isinstance(track_index, genome.Interval):
+        raise IndexError(
+            'Track indexing by interval is supported only when there are'
+            ' positional axes.'
+        )
+
+    tdata = self
+    match position_index:
+      case None:
+        pass
+      case int():
+        tdata = tdata.slice_by_positions(position_index, position_index + 1)
+      case slice():
+        if position_index.step is not None and position_index.step != 1:
+          raise IndexError('Slice step must be 1 for positional indexing.')
+        if position_index != slice(None):
+          tdata = tdata.slice_by_positions(
+              position_index.start, position_index.stop
+          )
+      case genome.Interval():
+        tdata = tdata.slice_by_interval(position_index)
+      case _:
+        raise IndexError(
+            f'Unsupported positional index type: {type(position_index)}'
+        )
+
+    match track_index:
+      case None:
+        pass
+      case str():
+        tdata = tdata.select_tracks_by_name([track_index])
+      case int():
+        tdata = tdata.select_tracks_by_index([track_index])
+      case slice():
+        if track_index != slice(None):
+          indices = np.arange(tdata.num_tracks)[track_index]
+          tdata = tdata.select_tracks_by_index(indices)
+      case np.ndarray() if np.issubdtype(track_index.dtype, np.character):
+        tdata = tdata.select_tracks_by_name(track_index)
+      case np.ndarray():
+        tdata = tdata.select_tracks_by_index(track_index)
+      case Sequence():
+        track_index_arr = np.asarray(track_index)
+        if np.issubdtype(track_index_arr.dtype, np.character):
+          tdata = tdata.select_tracks_by_name(track_index_arr)
+        else:
+          tdata = tdata.select_tracks_by_index(track_index_arr)
+      case _:
+        raise IndexError(f'Unsupported track index type: {type(track_index)}')
+    return tdata
+
+  def groupby(self, column: str) -> dict[str, 'TrackData']:
+    """Splits tracks into groups based on a metadata column.
+
+    This method splits the tracks in the `TrackData` object into separate
+    `TrackData` objects based on the unique values in the specified metadata
+    column. It returns a dictionary where the keys are the  unique values in
+    the column, and the values are new `TrackData` objects containing the
+    tracks corresponding to each key.
+
+    Args:
+      column: The name of the metadata column to split by.
+
+    Returns:
+      A dictionary mapping unique values in the column to `TrackData` objects
+      containing the corresponding tracks.
+    """
+    output = {}
+    for key in self.metadata[column].unique():
+      mask = (self.metadata[column] == key).values
+      output[key] = self.filter_tracks(mask)
+    return output
+
+  def _reverse_complement_idx(self) -> np.ndarray:
+    """Gets indices for reverse complementing the tracks.
+
+    Returns:
+      An array of indices that reorders the tracks to achieve reverse
+      complementation.
+
+    Raises:
+      ValueError: If not all stranded tracks have both '+' and '-' strands,
+         or if the number of '+' and '-' stranded tracks is not equal.
+    """
+    df_strands = pd.DataFrame({
+        'name': self.names,
+        'strand': self.strands,
+        'old_idx': np.arange(self.num_tracks),
+    })
+    df_strands = df_strands[df_strands.strand != genome.STRAND_UNSTRANDED]
+    df_strands.sort_values(['strand', 'name'], inplace=True)
+    if np.all(df_strands.groupby('name').size() != 2):
+      raise ValueError('Not all stranded tracks have both + and - strand.')
+    if (df_strands.strand == genome.STRAND_POSITIVE).sum() != (
+        df_strands.strand == genome.STRAND_NEGATIVE
+    ).sum():
+      raise ValueError(
+          'We need to have the exact same number of + and - stranded tracks'
+      )
+    new_idx = df_strands.old_idx.values.reshape((2, -1))[::-1].ravel()
+    # Swap strands by idx.
+    idx = np.arange(self.num_tracks)
+    idx[df_strands.old_idx.values] = new_idx
+    return idx
+
+  def reverse_complement(self) -> 'TrackData':
+    """Reverse complements the track data and interval if present.
+
+    Returns:
+      A new `TrackData` object with reverse complemented tracks.
+    """
+    if self.interval:
+      # Note that the interval needs to be stranded in order to perform
+      # this operation.
+      interval = self.interval.swap_strand()
+    else:
+      interval = None
+
+    idx = self._reverse_complement_idx()
+    slices = [slice(None)] * self.values.ndim
+    slices[-1] = idx
+    for axis in self.positional_axes:
+      slices[axis] = slice(None, None, -1)
+
+    return TrackData(
+        self.values[tuple(slices)],
+        resolution=self.resolution,
+        metadata=self.metadata.iloc[idx],
+        interval=interval,
+        uns=self.uns,
+    )
+
+  def _check_track_data_compatibility(self, other: 'TrackData') -> None:
+    """Checks if two `TrackData` objects are compatible for sum/diff.
+
+    Args:
+      other: The other `TrackData` object to compare.
+
+    Raises:
+      TypeError: If `other` is not a `TrackData` object.
+      ValueError: If the intervals, resolutions, shapes, or metadata
+        shapes don't match between the two objects.
+    """
+    if not isinstance(other, TrackData):
+      raise TypeError(
+          f'Unsupported type "{type(other)}". Must be a TrackData object'
+      )
+    if self.interval != other.interval:
+      raise ValueError('Intervals must match for the two TrackData objects.')
+    if self.resolution != other.resolution:
+      raise ValueError('Resolutions must match for the two TrackData objects.')
+    if self.values.shape != other.values.shape:
+      raise ValueError('Shapes must match for the two TrackData objects.')
+    if self.metadata.shape != other.metadata.shape:
+      raise ValueError(
+          'Metadata shapes must match for the two TrackData objects.'
+      )
+
+  def __add__(self, other: 'TrackData') -> 'TrackData':
+    """Adds the values of two `TrackData` objects.
+
+    Args:
+      other: The `TrackData` object to add.
+
+    Returns:
+      A new `TrackData` object with the summed values.
+
+    Raises:
+      ValueError: If the objects are not compatible (see
+        `_check_track_data_compatibility`).
+      TypeError: If `other` is not a `TrackData` object.
+    """
+    self._check_track_data_compatibility(other)
+    new_values = self.values + other.values
+    return TrackData(
+        values=new_values,
+        metadata=self.metadata,
+        resolution=self.resolution,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+  def __sub__(self, other: 'TrackData') -> 'TrackData':
+    """Subtracts the values of two `TrackData` objects.
+
+    Args:
+      other: The `TrackData` object to subtract.
+
+    Returns:
+      A new `TrackData` object with the difference of the values.
+
+    Raises:
+      ValueError: If the objects are not compatible (see
+        `_check_track_data_compatibility`).
+      TypeError: If `other` is not a `TrackData` object.
+    """
+    self._check_track_data_compatibility(other)
+    new_values = self.values - other.values
+    return TrackData(
+        values=new_values,
+        metadata=self.metadata,
+        resolution=self.resolution,
+        interval=self.interval,
+        uns=self.uns,
+    )
+
+
+def concat(
+    track_datas: Sequence[TrackData],
+    extra_metadata_name_and_keys: (
+        tuple[str, Sequence[str | int | float]] | None
+    ) = None,
+) -> TrackData:
+  """Concatenates multiple `TrackData` objects along the track dimension.
+
+  This function combines multiple `TrackData` objects into a single object
+  by concatenating their values and metadata. The resulting `TrackData`
+  object will have the same resolution and interval as the input objects.
+
+  Args:
+    track_datas: A sequence of `TrackData` objects to concatenate. All objects
+      must have the same resolution, interval, and width.
+    extra_metadata_name_and_keys: An optional tuple specifying a new metadata
+      column to add. The first element is the column name, and the second is a
+      sequence of values to populate the column.
+
+  Returns:
+    A new `TrackData` object containing the concatenated data.
+
+  Raises:
+    ValueError: If the input `TrackData` objects have different resolutions,
+      intervals, or widths, or if the length of
+      `extra_metadata_name_and_keys[1]` does not match the length of
+      `track_datas`.
+  """
+  if len(set(x.resolution for x in track_datas)) != 1:
+    raise ValueError('Track data contain multiple resolutions')
+  if len(set(str(x.interval) for x in track_datas)) != 1:
+    raise ValueError('Track data contain multiple intervals')
+  if len(set(x.width for x in track_datas)) != 1:
+    raise ValueError('Track data are of different width')
+  if extra_metadata_name_and_keys:
+    if len(extra_metadata_name_and_keys[1]) != len(track_datas):
+      raise ValueError(
+          'Second element of new_metadata_name_and_keys must be the same'
+          ' length as track_datas'
+      )
+
+  concatenated_metadata = (
+      pd.concat(
+          [x.metadata for x in track_datas],
+          keys=extra_metadata_name_and_keys[1]
+          if extra_metadata_name_and_keys
+          else None,
+          names=[extra_metadata_name_and_keys[0]]
+          if extra_metadata_name_and_keys
+          else None,
+      )
+      .reset_index(level=0, drop=extra_metadata_name_and_keys is None)
+      .reset_index(drop=True)
+  )
+  return TrackData(
+      np.concatenate(
+          [x.values for x in track_datas], axis=track_datas[0].values.ndim - 1
+      ),
+      resolution=track_datas[0].resolution,
+      metadata=concatenated_metadata,
+      interval=track_datas[0].interval,
+      uns=None,
+  )
+
+
+def interleave(
+    track_datas: Sequence[TrackData], name_prefixes: Sequence[str]
+) -> TrackData:
+  """Interleaves multiple `TrackData` objects by alternating rows.
+
+  This function combines multiple `TrackData` objects into a single object
+  by interleaving their rows and metadata. This interleaves operation alternates
+  between the trackdatas, like shuffling cards, i.e., 'abcd' interleaved with
+  'efgh' would be "aebfcgdh". The resulting `TrackData` object will have the
+  same resolution and interval as the input objects, but the number of tracks
+  will be the sum of the tracks in the input objects.
+
+  Args:
+    track_datas: A sequence of `TrackData` objects to interleave. All objects
+      must have the same shape, resolution, and interval. The order in this list
+      will determine the interleaving order.
+    name_prefixes: A sequence of name prefixes to add to the track names in the
+      metadata to ensure uniqueness of (name, strand) pairs.
+
+  Returns:
+    A new `TrackData` object containing the interleaved data.
+
+  Raises:
+    ValueError: If the input `TrackData` objects have different shapes,
+      resolutions, or intervals.
+  """
+  # Checks on the track data.
+  shapes = set(data.values.shape for data in track_datas)
+  if len(shapes) != 1:
+    raise ValueError(
+        'Cannot interleave track data which have different shapes. '
+        f'Detected shapes: {shapes}'
+    )
+
+  if any(data.resolution != track_datas[0].resolution for data in track_datas):
+    raise ValueError(
+        'Cannot interleave track data which have different resolutions. '
+        f'Detected shapes: {shapes}'
+    )
+
+  if any(data.interval != track_datas[0].interval for data in track_datas):
+    raise ValueError(
+        'Cannot interleave track data which have different intervals. '
+    )
+
+  # Interleave arrays.
+  shape = list(track_datas[0].values.shape)
+  shape[-1] = shape[-1] * len(track_datas)
+  interleaved_data = np.empty(tuple(shape), dtype=track_datas[0].values.dtype)
+
+  for i, data in enumerate(track_datas):
+    interleaved_data[..., i :: len(track_datas)] = data.values
+
+  # Interleave metadata.
+  metadatas = []
+  for prefix, data in zip(name_prefixes, track_datas):
+    metadata = data.metadata.copy()
+    metadata['name'] = prefix + metadata['name']
+    metadatas.append(metadata)
+
+  # Assign a new index idx that, when sorted, will produce an interleave.
+  interleaved_metadata = (
+      pd.concat([
+          metadata.assign(
+              idx=np.arange(
+                  stop=(len(metadata) * len(track_datas)),
+                  step=len(track_datas),
+              )
+              + i
+          )
+          for i, metadata in enumerate(metadatas)
+      ])
+      .sort_values('idx')
+      .drop('idx', axis=1)
+      .reset_index(drop=True)
+  )
+
+  return TrackData(
+      values=interleaved_data,
+      metadata=interleaved_metadata,
+      resolution=track_datas[0].resolution,
+      interval=track_datas[0].interval,
+      uns={'num_interleaved_trackdatas': len(track_datas)},
+  )

flax_model/alphagenome/_sdk/data/transcript.py

@@ -0,0 +1,755 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for working with transcripts."""
+
+import collections
+import copy
+import dataclasses
+import functools
+import sys
+from typing import Any
+
+from flax_model.alphagenome._sdk.data import genome
+import pandas as pd
+
+
+MITOCHONDRIAL_CHROMS = ['M', 'chrM', 'MT']
+
+
+@dataclasses.dataclass(frozen=True)
+class Transcript:
+  """Represents transcript object containing attributes from a GTF file.
+
+  A transcript is a region of DNA that encodes a single RNA molecule. The
+  Transcript dataclass contains attributes that describe the structure and
+  content of a transcript, namely:
+
+  Attributes:
+    exons: A list of `genome.Interval`s representing exons within transcript.
+      Each `Transcript` must contain exons.
+    cds: An optional list of `genome.Interval`s representing coding sequences
+      (CDS) within a transcript. CDS include start codon and exclude stop codon.
+    start_codon: An optional list of `genome.Interval`s representing a single
+      start codon. Start codons can be split by introns, therefore might have
+      more than one genomic interval. Some coding transcripts are missing start
+      codons, e.g., ENST00000455638.6.
+    stop_codon: An optional list of `genome.Interval`s representing a single
+      stop codon. Stop codon can be split by introns, therefore might have more
+      than one genomic interval. Some transcripts coding transcripts are missing
+      stop codons, e.g., ENST00000574051.5.
+    transcript_id: An optional string representing a transcript id.
+    gene_id: An optional string representing a gene id.
+    protein_id: An optional string representing a protein id which is encoded by
+      the transcript.
+    uniprot_id: An optional UniprotKB-AC id string.
+    info: a dictionary of additional information on a transcript.
+    chromosome: chromosome name on which the transcript is present. Must be the
+      same for all genomic intervals within a transcript.
+    is_mitochondrial: whether the transcript is on the mitochondria chromosome.
+    strand_int: strand on which transcript is present as an int. -1 for negative
+      strand +1 for positive strand
+    strand: strand (positive or negative) on which transcript is present. Must
+      be the same for all genomic intervals within a transcript.
+    is_negative_strand: a boolean value indicating whether transcript is on
+      negative strand.
+    is_positive_strand: a boolean value indicating whether transcript is on
+      positive strand.
+    transcript_interval: a genomic interval of a transcript.
+    selenocysteines: a list of intervals where selenocysteines are present
+      within a transcript.
+    selenocysteine_pos_in_protein: a list of 0-based positions of
+      selenocysteines in protein encoded by the transcript.
+    is_coding: a value indicating whether a `Transcript` contains coding
+      sequences (CDS) or not.
+    cds_including_stop_codon: a list of CDS and stop_codon intervals with
+      overlapping intervals merged.
+    utr5: A list of genomic intervals representing 5' untranslated region. 5'
+      UTR doesn't include start codon. There may be no 5' UTR present in the
+      transcript or UTRs can be split by introns.
+    utr3: A list of genomic intervals representing 3' untranslated region. 3'
+      UTR doesn't include stop codon. There may be no 3' UTR present in the
+      transcript or UTRs can be split by introns.
+    splice_regions: a list of splice regions within a transcript.
+    splice_donor_sites: a list of splice donor sites. Commonly, the RNA sequence
+      that is removed begins with the dinucleotide GU at its 5′ end.
+    splice_acceptor_sites: a list of splice acceptor sites. Commonly, the RNA
+      sequence that is removed ends with AG at its 3′ end.
+    splice_donors: a list of splice donors. The first nucleotide of the intron
+      (0-based).
+    splice_acceptors: a list of splice acceptors. The last nucleotide of the
+      intron (0-based).
+  """
+
+  exons: list[genome.Interval]
+  cds: list[genome.Interval] | None = None
+  start_codon: list[genome.Interval] | None = None
+  stop_codon: list[genome.Interval] | None = None
+  transcript_id: str | None = dataclasses.field(compare=False, default=None)
+  gene_id: str | None = dataclasses.field(compare=False, default=None)
+  protein_id: str | None = dataclasses.field(compare=False, default=None)
+  uniprot_id: str | None = dataclasses.field(compare=False, default=None)
+  info: dict[str, Any] = dataclasses.field(
+      default_factory=dict, repr=False, compare=False, hash=False
+  )
+
+  def offset_in_cds(self, genome_position: int) -> int | None:
+    """Return the offset within the set of CDS exons of `genome_position`.
+
+    Args:
+      genome_position: A coordinate presumed to be on the same chromosome as
+        this transcript.
+
+    Returns:
+      The offset of `genome_position` from the start of the CDS, accounting for
+      strand, or None, if `genome_position` does not overlap the CDS.
+    """
+    offset = 0
+    for cds_exon in self.cds_including_stop_codon[:: self.strand_int]:
+      if cds_exon.start <= genome_position < cds_exon.end:
+        if self.is_positive_strand:
+          return offset + genome_position - cds_exon.start
+        else:
+          return offset + cds_exon.end - genome_position - 1
+      offset += cds_exon.width
+    return None
+
+  @property
+  def chromosome(self) -> str:
+    """Gets the chromosome name on which the transcript is present.
+
+    Returns:
+      The chromosome name.
+    """
+    return self.exons[0].chromosome
+
+  @property
+  def is_mitochondrial(self) -> bool:
+    """Gets whether the transcript is on the mitochondria chromosome.
+
+    Returns:
+      True if the transcript is on the mitochondria chromosome, False otherwise.
+    """
+    return self.chromosome in MITOCHONDRIAL_CHROMS
+
+  # TODO: b/376466056 - Unify strand representations.
+  @property
+  def strand_int(self) -> int:
+    """Gets the strand as an integer.
+
+    Returns:
+      -1 for negative strand, +1 for positive strand, 0 for unknown.
+    """
+    return {genome.STRAND_NEGATIVE: -1, genome.STRAND_POSITIVE: +1}.get(
+        self.strand, 0
+    )
+
+  @property
+  def strand(self) -> str:
+    """Gets the strand on which the transcript is present.
+
+    Returns:
+      The strand (positive or negative).
+    """
+    return self.exons[0].strand
+
+  @property
+  def is_positive_strand(self) -> bool:
+    return self.strand == genome.STRAND_POSITIVE
+
+  @property
+  def is_negative_strand(self) -> bool:
+    return self.strand == genome.STRAND_NEGATIVE
+
+  @functools.cached_property
+  def transcript_interval(self) -> genome.Interval:
+    """Gets a genomic interval of a transcript.
+
+    Returns:
+      A genomic interval of a transcript where transcript start is equal to the
+      first exon start and transcript end is equal to the last exon end.
+    """
+    return genome.Interval(
+        self.chromosome,
+        self.exons[0].start,
+        self.exons[-1].end,
+        strand=self.strand,
+    )
+
+  @functools.cached_property
+  def selenocysteines(self) -> list[genome.Interval]:
+    if 'selenocysteines' not in self.info:
+      return []
+    return self.info['selenocysteines']
+
+  @functools.cached_property
+  def selenocysteine_pos_in_protein(self) -> list[int]:
+    # 0-based
+    selenocystein_pos = []
+    for selenocysteine in self.selenocysteines:
+      if selenocysteine.info['cds_offset'] is None:
+        raise ValueError(
+            'Transcript cannot be translated due to '
+            'bad input data of selenocysteines.'
+        )
+      selenocystein_pos.append(selenocysteine.info['cds_offset'] // 3)
+    return selenocystein_pos
+
+  @functools.cached_property
+  def introns(self) -> list[genome.Junction]:
+    """Get a list of intron intervals.
+
+    Returns:
+      A list of genomic intervals representing introns, where a single intron
+      junction is an interval spanning between two adjacent exons.
+    """
+    intron_intervals = []
+    for i in range(1, len(self.exons)):
+      intron_intervals.append(
+          genome.Junction(
+              self.chromosome,
+              self.exons[i - 1].end,
+              self.exons[i].start,
+              strand=self.strand,
+          )
+      )
+    return intron_intervals
+
+  @functools.cached_property
+  def is_coding(self) -> bool:
+    return bool(self.cds)
+
+  @functools.cached_property
+  def cds_including_stop_codon(self) -> list[genome.Interval]:
+    """Obtains coding sequences including stop codon.
+
+    By default gtf files exclude stop codons from CDS while gff include stop
+    codons within coding sequences.
+    """
+    if not self.is_coding:
+      return []
+    return genome.merge_overlapping_intervals(
+        self.cds + (self.stop_codon or [])
+    )
+
+  @functools.cached_property
+  def utr5(self) -> list[genome.Interval]:
+    return self._get_utr(self.strand != genome.STRAND_NEGATIVE)
+
+  @functools.cached_property
+  def utr3(self) -> list[genome.Interval]:
+    return self._get_utr(self.strand == genome.STRAND_NEGATIVE)
+
+  def _get_utr(self, before: bool) -> list[genome.Interval]:
+    """Gets the UTRs located before/after first/last coding sequence."""
+    utrs = []
+    if not self.cds:
+      return utrs
+
+    merged_cds_stop = genome.merge_overlapping_intervals(
+        self.cds + (self.stop_codon or [])
+    )
+
+    if before:
+      start, end = 0, merged_cds_stop[0].start
+    else:
+      start, end = merged_cds_stop[-1].end, sys.maxsize
+
+    valid_interval = genome.Interval(
+        self.chromosome, start, end, strand=self.strand
+    )
+
+    for exon in filter(lambda x: x.overlaps(valid_interval), self.exons):
+      intersect = valid_interval.intersect(exon)
+      if intersect:
+        utrs.append(intersect)
+    return utrs
+
+  # TODO: b/376465275 - deal with cases where exon shorter than 3 bp length
+  # TODO: b/376465275 - deal with cases where intron is shorther than 4 bp
+  @functools.cached_property
+  def splice_regions(self) -> list[genome.Interval]:
+    """Obtains and returns splice regions of a transcript.
+
+    splice region (SO:0001630) is "within 1-3 bases of the exon or 3-8 bases of
+    the intron.
+    """
+    if not self.introns:
+      return []
+    splice_regions = []
+
+    for intron, prev_exon, next_exon in zip(
+        self.introns, self.exons[:-1], self.exons[1:]
+    ):
+      if prev_exon.width > 2:
+        splice_regions.append(
+            genome.Interval(
+                self.chromosome,
+                intron.start - 3,
+                intron.start,
+                strand=self.strand,
+            )
+        )
+
+      if next_exon.width > 2:
+        splice_regions.append(
+            genome.Interval(
+                self.chromosome, intron.end, intron.end + 3, strand=self.strand
+            )
+        )
+
+      if intron.width > 4:
+        splice_regions.append(
+            genome.Interval(
+                self.chromosome,
+                intron.start + 2,
+                min(intron.start + 8, intron.end - 2),
+                strand=self.strand,
+            )
+        )
+        splice_regions.append(
+            genome.Interval(
+                self.chromosome,
+                max(intron.end - 8, intron.start + 2),
+                intron.end - 2,
+                strand=self.strand,
+            )
+        )
+    return genome.merge_overlapping_intervals(splice_regions)
+
+  @functools.cached_property
+  def splice_donor_sites(self) -> list[genome.Interval]:
+    return self._get_splice_sites(False, intron_overhang=2, exon_overhang=0)
+
+  @functools.cached_property
+  def splice_acceptor_sites(self) -> list[genome.Interval]:
+    return self._get_splice_sites(True, intron_overhang=2, exon_overhang=0)
+
+  @functools.cached_property
+  def splice_donors(self) -> list[genome.Interval]:
+    # To be consistent with the splice sites defined by intron start and end,
+    # the overhang for donor and acceptor are different.
+    return self._get_splice_sites(  # pytype: disable=bad-return-type  # enable-cached-property
+        False, intron_overhang=1, exon_overhang=0
+    )
+
+  @functools.cached_property
+  def splice_acceptors(self) -> list[genome.Interval]:
+    return self._get_splice_sites(  # pytype: disable=bad-return-type  # enable-cached-property
+        True, intron_overhang=0, exon_overhang=1
+    )
+
+  # TODO: b/376465275 - deal with cases where intron shorter than 4 bp length.
+  def _get_splice_sites(
+      self, acceptor: bool, intron_overhang: int, exon_overhang: int
+  ) -> list[genome.Interval]:
+    """Obtains splice acceptor/donor intervals.
+
+    https://www.nature.com/scitable/topicpage/rna-splicing-introns-exons-and-spliceosome-12375/#:~:text=Introns%20are%20removed%20from%20primary,AG%20at%20its%203%E2%80%B2%20end
+    Introns are removed from primary transcripts by cleavage at conserved
+    sequences called splice sites. These sites are found at the 5′ and 3′
+    ends of introns. Most commonly, the RNA sequence that is removed begins with
+    the dinucleotide GU at its 5′ end, and ends with AG at its 3′ end.
+
+    if - strand, the end two bases of intron are splice donor bases,
+    if +, then the start two bases.
+
+    Args:
+      acceptor: value indicating whether splice acceptor or donor should be
+        obtained.
+      intron_overhang: bases into the intron.
+      exon_overhang: bases into the exon.
+
+    Returns:
+      List of intervals of splice acceptor/donor sites.
+    """
+    if not self.introns:
+      return []
+    splice_sites = []
+    for intron in self.introns:  # pylint:disable=not-an-iterable
+      if intron.width < 4:
+        continue
+      if self.is_negative_strand != acceptor:
+        splice = genome.Interval(
+            self.chromosome,
+            intron.end - intron_overhang,
+            intron.end + exon_overhang,
+            strand=self.strand,
+        )
+      else:
+        splice = genome.Interval(
+            self.chromosome,
+            intron.start - exon_overhang,
+            intron.start + intron_overhang,
+            strand=self.strand,
+        )
+      splice_sites.append(splice)
+    return splice_sites
+
+  def __post_init__(self):
+    if not self.exons:
+      raise ValueError('Transcript must contain at least one exon.')
+
+    for exon in self.exons:
+      if exon.strand != self.strand or exon.chromosome != self.chromosome:
+        raise ValueError(
+            'Transcript intervals are inconsistent. All intervals of a '
+            'transcript should have same strand and chromosome.'
+        )
+
+    if self.cds:
+      # first exons can be part of UTR. Searching for the first coding exon.
+      index = 0
+      for exon in self.exons:
+        # if overlaps, will check whether exon contains it in the latter loop.
+        if exon.overlaps(self.cds[0]):
+          break
+        index += 1
+      if index == len(self.exons) or len(self.cds) + index > len(self.exons):
+        raise ValueError(
+            'The number of coding exons must be the same as CDS '
+            'and CDS cannot be outside of the exon intervals.'
+        )
+
+      # checks each subsequent exon is coding
+      for seq, exon in zip(self.cds, self.exons[index : index + len(self.cds)]):
+        if not exon.contains(seq):
+          raise ValueError(
+              'The number of coding exons must be the same as CDS '
+              'and CDS cannot be outside of the exon intervals.'
+          )
+        if seq.strand != self.strand:
+          raise ValueError(
+              'Transcript intervals are inconsistent. All intervals of a '
+              'transcript should have same strand and chromosome.'
+          )
+
+    for sc in self.selenocysteines:  # pylint:disable=not-an-iterable
+      sc_pos = sc.end - 1 if sc.negative_strand else sc.start
+      sc.info['cds_offset'] = self.offset_in_cds(sc_pos)
+
+  def __len__(self):
+    return self.transcript_interval.width
+
+  @classmethod
+  def from_gtf_df(
+      cls,
+      transcript_df: pd.DataFrame,
+      ignore_info: bool = True,
+      fix_truncation: bool = False,
+  ) -> 'Transcript':
+    """Initialises Trancript object from a given transcript dataframe.
+
+    Args:
+      transcript_df: Dataframe representing a transcript. The dataframe must
+        contain a single transcript.
+      ignore_info: If True, other columns in transcript_df won't be added to the
+        info field, except transcript_type and selenocysteines.
+      fix_truncation: Whether or not apply truncation fixation to CDS.
+
+    Returns:
+      Initialised Transcript object.
+    Raises:
+      ValueError: if the dataframe provided is invalid (no or more than one
+      transcript, transcript has inconsistent strand or chromosome,
+      transcript doesn't contain exons, CDS are not within exons, etc.)
+    """
+    if transcript_df.empty:
+      raise ValueError('transcript_df is empty')
+    if 'Feature' not in transcript_df:
+      raise ValueError('transcript_df must contain Feature column.')
+
+    if (
+        'transcript_id' in transcript_df
+        and len(transcript_df.transcript_id.unique()) > 1
+    ):
+      raise ValueError('transcript_df should only contain a single transcript.')
+
+    # Convert rows to genome.Interval list.
+    transcript_df = transcript_df.sort_values(by='Start')
+    intervals_per_feature = collections.defaultdict(list)
+    exon_row = None
+    for _, row in transcript_df.iterrows():
+      interval = genome.Interval.from_pyranges_dict(
+          row, ignore_info=True
+      )  # pytype: disable=wrong-arg-types  # pandas-drop-duplicates-overloads
+      if row.Feature in ['CDS', 'stop_codon']:
+        interval.info['frame'] = int(row.Frame)
+      if exon_row is None and row.Feature == 'exon':
+        exon_row = row
+      intervals_per_feature[row.Feature].append(interval)
+
+    if exon_row is None:
+      raise ValueError('transcript_df must contain at least one exon')
+
+    # Seed info.
+    if ignore_info:
+      info = {}
+    else:
+      skip = list(genome.PYRANGES_INTERVAL_COLUMNS) + [
+          'Feature',
+          'transcript_type',
+          'Selenocysteines',
+          'gene_type',
+      ]
+      info = {k: v for k, v in exon_row.items() if k not in skip}
+
+    if 'transcript_type' in exon_row:
+      info['transcript_type'] = exon_row['transcript_type']
+
+    if 'Selenocysteine' in intervals_per_feature:
+      info['selenocysteines'] = intervals_per_feature['Selenocysteine']
+
+    if 'gene_type' in exon_row:
+      info['gene_type'] = exon_row['gene_type']
+
+    transcript_obj = cls(
+        exons=intervals_per_feature['exon'],
+        cds=intervals_per_feature.get('CDS', None),
+        start_codon=intervals_per_feature.get('start_codon', None),
+        stop_codon=intervals_per_feature.get('stop_codon', None),
+        transcript_id=exon_row.get('transcript_id', None),
+        gene_id=exon_row.get('gene_id', None),
+        protein_id=exon_row.get('protein_id', None),
+        uniprot_id=exon_row.get('uniprot_id', None),
+        info=info,
+    )
+    if fix_truncation:
+      return Transcript.fix_truncation(transcript_obj)
+    return transcript_obj
+
+  @classmethod
+  def fix_truncation(cls, transcript: 'Transcript') -> 'Transcript':
+    """Fixes CDS start and stop positions to be within coding frame.
+
+    Args:
+      transcript: a transcript to fix.
+
+    Returns:
+      New transcript with set start/stop codons and fixed CDS if the total
+      length of CDS is > 6. Returns a copy of original transcript otherwise.
+    """
+    cds = sorted(
+        (transcript.cds or []) + (transcript.stop_codon or []),
+        key=lambda x: x.start,
+    )
+    cdna_len = sum(seq.width for seq in cds)
+    if cdna_len < 7:
+      return copy.deepcopy(transcript)
+
+    positive_strand = transcript.is_positive_strand
+    try:
+      frame = cds[0 if positive_strand else -1].info['frame']
+    except KeyError as key_error:
+      raise KeyError(
+          'CDS intervals are missing frame information,'
+          ' truncations cannot be deduced.'
+      ) from key_error
+    frame_last = (cdna_len - frame) % 3
+
+    cds, start_codon = cls._fix_coding_frame(
+        five_prime=True,
+        beginning=positive_strand,
+        frame=frame,
+        cds_transcript=cds,
+    )
+
+    cds, stop_codon = cls._fix_coding_frame(
+        five_prime=False,
+        beginning=not positive_strand,
+        frame=frame_last,
+        cds_transcript=cds,
+    )
+
+    return cls(
+        exons=[exon.copy() for exon in transcript.exons],
+        cds=cds,
+        start_codon=start_codon,
+        stop_codon=stop_codon,
+        transcript_id=transcript.transcript_id,
+        gene_id=transcript.gene_id,
+        protein_id=transcript.protein_id,
+        uniprot_id=transcript.uniprot_id,
+        info={**transcript.info, 'truncation_fixed': True},
+    )
+
+  @classmethod
+  def _fix_coding_frame(
+      cls,
+      five_prime: bool,
+      beginning: bool,
+      frame: int,
+      cds_transcript: list[genome.Interval],
+  ) -> tuple[list[genome.Interval], list[genome.Interval]]:
+    """Fixes coding frame for a transcript and returns a new start/stop codon.
+
+    Args:
+      five_prime: a value indicating whether a 5' end to be fixed.
+      beginning: indicates whether the beginning or the end of the transript to
+        be fixed.
+      frame: a current coding frame (0, 1, 2). For fixing 5' end, it is a
+        position at which the first full codon starts within a cds. For fixing
+        3' end, this indicates where the last full codon ends within a cds.
+      cds_transcript: a list of coding sequence intervals.
+
+    Returns:
+      A pair of lists where the first item is a list of CDS with fixed coding
+      frame and the second item is a start/stop codon.
+    """
+
+    def shorten_intervals(
+        cds_transcript: list[genome.Interval],
+        beginning: bool,
+        frame: int,
+    ) -> list[genome.Interval]:
+      cds = [interval.copy() for interval in cds_transcript]
+      index = 0 if beginning else -1
+      while frame > 0:
+        if cds[index].width > frame:
+          if beginning:
+            cds[index].start += frame
+          else:
+            cds[index].end -= frame
+          frame = 0
+        else:
+          frame -= cds[index].width
+          del cds[index]
+      return cds
+
+    # Fix CDS coding frame.
+    fixed_cds = shorten_intervals(cds_transcript, beginning, frame)
+
+    # Set codon.
+    codon_bases = 3
+    fixed_codon = []
+    for seq in fixed_cds if beginning else fixed_cds[::-1]:
+      if codon_bases == 0:
+        break
+      start, end = seq.start, seq.end
+      if seq.width > codon_bases and beginning:
+        end = seq.start + codon_bases
+      elif seq.width > codon_bases:
+        start = seq.end - codon_bases
+      interval = genome.Interval(seq.chromosome, start, end, strand=seq.strand)
+      fixed_codon.append(interval)
+      codon_bases -= interval.width
+
+    if five_prime:
+      fixed_cds[0 if beginning else -1].info['frame'] = 0
+    else:
+      # Remove newly set stop interval from cds.
+      fixed_cds = shorten_intervals(fixed_cds, beginning, 3)
+    return fixed_cds, fixed_codon
+
+
+class _RangeExtractor:
+  """Range extractor from gtf df."""
+
+  def __init__(self, df: pd.DataFrame):
+    self._df_start_end = {
+        chromosome: (dfc, dfc['Start'].values, dfc['End'].values)
+        for chromosome, dfc in df.groupby('Chromosome')
+    }
+    self._df_empty = df.iloc[:0]
+
+  def extract(self, interval: genome.Interval) -> pd.DataFrame:
+    """Finds all rows that contain the input Interval.
+
+    Args:
+      interval: query Interval
+
+    Returns:
+    a dataframe containing genome intervals that contain the query Interval
+    """
+    if interval.chromosome not in self._df_start_end:
+      return self._df_empty
+    else:
+      dfc, start, end = self._df_start_end[interval.chromosome]
+      start_contained = (interval.start <= start) & (start <= interval.end)
+      end_contained = (interval.start <= end) & (end <= interval.end)
+      interval_contained = (interval.start >= start) & (end >= interval.end)
+      return dfc[start_contained | end_contained | interval_contained]
+
+
+class TranscriptExtractor:
+  """Transcript extractor from gtf."""
+
+  def __init__(self, gtf_df: pd.DataFrame) -> None:
+    """Init.
+
+    Args:
+      gtf_df: pd.DataFrame of GENCODE GTF entries containing transcript
+        annotation. Must contain columns 'Chromosome', 'Start', 'End', 'Strand',
+        'Feature', and 'transcript_id'.
+    """
+    self._transcript_extractor = _RangeExtractor(
+        gtf_df[gtf_df.Feature == 'transcript'][
+            ['Chromosome', 'Start', 'End', 'Strand', 'transcript_id']
+        ]
+    )
+    self._transcript_indexed_gtf = gtf_df.set_index('transcript_id')
+    self._transcript_from_id_cache = None
+
+  def cache_transcripts(self) -> None:
+    """Speed up extract() by converting GTF to dictionary of Transcripts.
+
+    This may take ca 11 minutes on the full human genome GTF of 84k protein
+    coding transcripts and 15 s on chr22 (1.5k transcripts).
+
+    Running cache_transcripts() will speed up .extract() by ca 5-10x:
+    (11 ms vs 65 ms tested on chr22, or 15 ms vs 160 ms on whole genome).
+    """
+    self._transcript_from_id_cache = self._transcripts_from_gtf(
+        self._transcript_indexed_gtf.reset_index()
+    )
+
+  def _transcripts_from_gtf(
+      self,
+      gtf_df: pd.DataFrame,
+  ) -> dict[str, Transcript]:
+    return (
+        {  # pytype: disable=bad-return-type  # pandas-drop-duplicates-overloads
+            transcript_id: (
+                Transcript.fix_truncation(
+                    Transcript.from_gtf_df(gtf_subset, ignore_info=False)
+                )
+            )
+            for transcript_id, gtf_subset in gtf_df.groupby('transcript_id')
+        }
+    )
+
+  def extract(self, interval: genome.Interval) -> list[Transcript]:
+    """Extract transcripts overlapping an interval.
+
+    Args:
+      interval: Interval used to overlap with transcripts.
+
+    Returns:
+      List of transcript overlapping `interval`.
+    """
+    gtf_df_within_interval = self._transcript_extractor.extract(interval)
+    if gtf_df_within_interval.empty:
+      return []
+
+    transcript_ids = gtf_df_within_interval.transcript_id.dropna().unique()
+    if self._transcript_from_id_cache is not None:
+      return [
+          self._transcript_from_id_cache[transcript_id]
+          for transcript_id in transcript_ids
+      ]
+    else:
+      transcript_gtfs = self._transcript_indexed_gtf.loc[
+          transcript_ids
+      ].reset_index()
+      return list(self._transcripts_from_gtf(transcript_gtfs).values())

flax_model/alphagenome/_sdk/interpretation/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Library of tools for interpreting sequences and model predictions."""

flax_model/alphagenome/_sdk/interpretation/ism.py

@@ -0,0 +1,150 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""In-silico mutagenesis (ISM) functions for sequence interpretation."""
+
+from collections.abc import Sequence
+
+from flax_model.alphagenome._sdk.data import genome
+import numpy as np
+
+
+def ism_variants(
+    interval: genome.Interval,
+    sequence: str,
+    vocabulary: str = 'ACGT',
+    skip_n: bool = False,
+) -> list[genome.Variant]:
+  """Create a list of all possible single nucleotide variants for an interval.
+
+  Args:
+    interval: Interval for which to generate the variants.
+    sequence: Sequence extracted from the reference genome at the `interval`.
+      Needs to have the same length as `interval.width`.
+    vocabulary: Vocabulary of possible alternative bases contained in
+      `sequence`.
+    skip_n: If True, skip the N bases.
+
+  Returns:
+    List of all possible single nucleotide variants for a genomic interval.
+  """
+  if len(sequence) != interval.width:
+    raise ValueError('Sequence must have the same length as interval.')
+  variants = []
+  for position in range(interval.width):
+    reference_base = sequence[position]
+    if skip_n and reference_base == 'N':
+      continue
+    for alternative_base in vocabulary:
+      if reference_base == alternative_base:
+        continue
+      variants.append(
+          genome.Variant(
+              chromosome=interval.chromosome,
+              reference_bases=reference_base,
+              alternate_bases=alternative_base,
+              position=interval.start + position + 1,
+          )
+      )
+  return variants
+
+
+def ism_matrix(
+    variant_scores: Sequence[float],
+    variants: Sequence[genome.Variant],
+    interval: genome.Interval | None = None,
+    multiply_by_sequence: bool = True,
+    vocabulary: str = 'ACGT',
+    require_fully_filled: bool = True,
+) -> np.ndarray:
+  """Construct the ISM (position, base) matrix from individual ISM scores.
+
+  This function returns the relative effect of the variants compared to the
+  per-position average: score[position, base] - mean(score[position, :]]).
+
+  Args:
+    variant_scores: Variant effect scores corresponding to the variants. These
+      could be obtained from the score_variants() output summarised to a single
+      scalar. Summarisation could be for example be obtained by selecting a
+      specific variant scorer output and extract a specific value from the
+      (variant, track) matrix.
+    variants: Sequence of variants used to transform into the ISM matrix.
+    interval: Interval for which to get the contribution scores. All variants
+      need to be contained within that interval. If None, it will be
+      automatically inferred from variants.
+    multiply_by_sequence: If True, only return non-zero values at
+      one-hot-encoded reference genome sequence bases.
+    vocabulary: Vocabulary of possible alternative bases contained in
+      `sequence`. The order determines the column order of the returned matrix.
+    require_fully_filled: If True, raise an error if not all positions are
+      covered by variants.
+
+  Returns:
+    Matrix of shape (interval.width, 4) containing variant scores.
+  """
+  if len(variants) != len(variant_scores):
+    raise ValueError(
+        'Variants and variant_scores need to have the same length.'
+    )
+  if interval is None:
+    interval = genome.Interval(
+        chromosome=variants[0].chromosome,
+        start=min(variant.start for variant in variants),
+        end=max(variant.end for variant in variants),
+    )
+  scores = np.zeros((interval.width, 4), dtype=np.float32)
+  filled = np.zeros((interval.width, 4), dtype=bool)
+  base_index = {base: i for i, base in enumerate(vocabulary)}
+
+  for variant, score in zip(variants, variant_scores, strict=True):
+    if len(variant.alternate_bases) != 1 or len(variant.reference_bases) != 1:
+      # Only looking for single nucleotide variants.
+      continue
+    if not interval.contains(variant.reference_interval):
+      continue
+    position = variant.start - interval.start
+    scores[position, base_index[variant.alternate_bases]] = score
+    filled[position, base_index[variant.alternate_bases]] = True
+
+  # Check that all positions were covered by variants.
+  if (filled.sum(axis=-1) == 0).any() and require_fully_filled:
+    missing_positions = list(
+        np.where(filled.sum(axis=-1) == 0)[0] + interval.start + 1
+    )
+    raise ValueError(
+        'No variants were found for the (1-based) positions on chromosome '
+        f'{interval.chromosome}: {missing_positions}'
+    )
+  elif (filled.sum(axis=-1) == 4).any():
+    full_positions = list(
+        np.where(filled.sum(axis=-1) == 4)[0] + interval.start + 1
+    )
+    raise ValueError(
+        'Some variants were found with 4 different alternative bases for the'
+        ' position instead of 3. List of positions on chromosome '
+        f'{interval.chromosome}: {full_positions}'
+    )
+  elif (filled.sum(axis=-1) != 3).any() and require_fully_filled:
+    missing_positions = list(
+        np.where(filled.sum(axis=-1) != 3)[0] + interval.start + 1
+    )
+    raise ValueError(
+        'Some variants were found with only 1 or 2 alternative bases for the'
+        ' position instead of 3. List of positions on chromosome '
+        f'{interval.chromosome}: {missing_positions}'
+    )
+
+  scores -= np.sum(scores, axis=-1, keepdims=True) / (len(vocabulary) - 1)
+  if multiply_by_sequence:
+    scores = scores * (~filled)
+  return scores

flax_model/alphagenome/_sdk/models/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Models library for interacting with AlphaGenome models."""

flax_model/alphagenome/_sdk/models/dna_client.py

@@ -0,0 +1,907 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Client implementation for interacting with a DNA model server."""
+
+from collections.abc import Container, Iterable, Iterator, Mapping, Sequence
+import concurrent.futures
+import functools
+import random
+import time
+from typing import TypeVar
+
+from flax_model.alphagenome._sdk import tensor_utils
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import junction_data
+from flax_model.alphagenome._sdk.data import ontology
+from flax_model.alphagenome._sdk.data import track_data
+from flax_model.alphagenome._sdk.models import dna_model
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome._sdk.models import interval_scorers as interval_scorers_lib
+from flax_model.alphagenome._sdk.models import junction_data_utils
+from flax_model.alphagenome._sdk.models import track_data_utils
+from flax_model.alphagenome._sdk.models import variant_scorers as variant_scorers_lib
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+from flax_model.alphagenome._sdk.protos import dna_model_service_pb2
+from flax_model.alphagenome._sdk.protos import dna_model_service_pb2_grpc
+from flax_model.alphagenome._sdk.protos import tensor_pb2
+import anndata
+import grpc
+import numpy as np
+import pandas as pd
+import tqdm.auto
+
+
+# Supported DNA sequence lengths.
+SEQUENCE_LENGTH_16KB = 2**14  # 16_384
+SEQUENCE_LENGTH_100KB = 2**17  # 131_072
+SEQUENCE_LENGTH_500KB = 2**19  # 524_288
+SEQUENCE_LENGTH_1MB = 2**20  # 1_048_576
+
+SUPPORTED_SEQUENCE_LENGTHS = {
+    'SEQUENCE_LENGTH_16KB': SEQUENCE_LENGTH_16KB,
+    'SEQUENCE_LENGTH_100KB': SEQUENCE_LENGTH_100KB,
+    'SEQUENCE_LENGTH_500KB': SEQUENCE_LENGTH_500KB,
+    'SEQUENCE_LENGTH_1MB': SEQUENCE_LENGTH_1MB,
+}
+
+# Maximum width of an in silico mutagenesis (ISM) interval request.
+# This controls the behavior of `score_ism_variants` in the client.
+# When a user inputs an ISM interval that is wider than this value,
+# the client will automatically split the interval into chunks of this width.
+MAX_ISM_INTERVAL_WIDTH = 10
+
+# Maximum number of variant scorers per request.
+MAX_VARIANT_SCORERS_PER_REQUEST = 20
+
+_VALID_SEQUENCE_CHARACTERS = frozenset('ACGTN')
+
+RetryableFunction = TypeVar('RetryableFunction')
+
+
+def retry_rpc(
+    function: RetryableFunction,
+    *,
+    max_attempts: int = 5,
+    initial_backoff: float = 1.25,
+    backoff_multiplier: float = 1.5,
+    retry_status_codes: Container[grpc.StatusCode] = frozenset(
+        [grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.UNAVAILABLE]
+    ),
+    jitter: float = 0.2,
+) -> RetryableFunction:
+  """Decorator that retries when an RPC fails.
+
+  gRPC currently doesn't support retries for streaming RPCs. This decorator is
+  an implementation of the retry logic from https://grpc.io/docs/guides/retry.
+
+  Arguments:
+    function: Callable to retry.
+    max_attempts: Maximum number of attempts to make.
+    initial_backoff: Initial backoff time in seconds.
+    backoff_multiplier: Backoff multiplier to apply on each retry.
+    retry_status_codes: Set of status codes to retry on.
+    jitter: Jitter to apply to the backoff time. For example, a jitter of 0.2
+      and an initial backoff of 1.5 will result in a backoff time between 1.2
+      and 1.8 seconds.
+
+  Returns:
+    A decorator that retries the function on retryable RPC errors.
+  """
+  if backoff_multiplier < 1.0:
+    raise ValueError('Backoff multiplier must be >= 1.0.')
+
+  @functools.wraps(function)
+  def wrapper(*args, **kwargs):
+    attempt = 0
+    current_backoff = initial_backoff + (
+        random.uniform(-jitter, jitter) * initial_backoff
+    )
+    while True:
+      try:
+        attempt += 1
+        return function(*args, **kwargs)
+      except grpc.RpcError as e:
+        error_code = e.code()  # pytype: disable=attribute-error
+        if error_code not in retry_status_codes or attempt >= max_attempts:
+          raise e
+        time.sleep(current_backoff)
+        current_backoff *= backoff_multiplier
+
+  return wrapper
+
+
+ModelVersion = dna_model.ModelVersion
+Output = dna_output.Output
+OutputMetadata = dna_output.OutputMetadata
+OutputType = dna_output.OutputType
+VariantOutput = dna_output.VariantOutput
+Organism = dna_model.Organism
+
+PredictResponse = TypeVar(
+    'PredictResponse',
+    dna_model_service_pb2.PredictSequenceResponse,
+    dna_model_service_pb2.PredictIntervalResponse,
+)
+
+
+def _read_tensor_chunks(
+    responses: Iterator[
+        PredictResponse
+        | dna_model_service_pb2.PredictVariantResponse
+        | dna_model_service_pb2.ScoreVariantResponse
+        | dna_model_service_pb2.ScoreIntervalResponse
+    ],
+    chunk_count: int,
+) -> Iterable[tensor_pb2.TensorChunk]:
+  """Helper to read a sequence of tensor chunks from a response iterator."""
+  for _ in range(chunk_count):
+    try:
+      response = next(responses)
+    except StopIteration as e:
+      raise ValueError('Expected tensor_chunk, got end of stream') from e
+    if response.WhichOneof('payload') != 'tensor_chunk':
+      raise ValueError(
+          f'Expected tensor_chunk, got "{response.WhichOneof("payload")}"'
+          ' payload'
+      )
+    yield response.tensor_chunk
+
+
+def _make_output_data(
+    output: dna_model_pb2.Output,
+    responses: Iterator[
+        PredictResponse | dna_model_service_pb2.PredictVariantResponse
+    ],
+    interval: genome.Interval | None = None,
+) -> track_data.TrackData | np.ndarray | junction_data.JunctionData:
+  """Helper to create an output data object from an output proto."""
+  match output.WhichOneof('payload'):
+    case 'track_data':
+      return track_data_utils.from_protos(
+          output.track_data,
+          _read_tensor_chunks(responses, output.track_data.values.chunk_count),
+          interval=interval,
+      )
+    case 'data':
+      chunks = _read_tensor_chunks(responses, output.data.chunk_count)
+      values = tensor_utils.unpack_proto(output.data, chunks)
+      return tensor_utils.upcast_floating(values)
+    case 'junction_data':
+      return junction_data_utils.from_protos(
+          output.junction_data,
+          _read_tensor_chunks(
+              responses, output.junction_data.values.chunk_count
+          ),
+          interval=interval,
+      )
+    case _:
+      raise ValueError(
+          f'Unsupported output type: {output.WhichOneof("payload")}'
+      )
+
+
+def construct_output_metadata(
+    responses: Iterator[dna_model_service_pb2.MetadataResponse],
+) -> OutputMetadata:
+  """Constructs an OutputMetadata from a stream of responses."""
+  metadata = {}
+  for response in responses:
+    for metadata_proto in response.output_metadata:
+      match metadata_proto.WhichOneof('payload'):
+        case 'tracks':
+          metadata[metadata_proto.output_type] = (
+              track_data_utils.metadata_from_proto(metadata_proto.tracks)
+          )
+        case 'junctions':
+          metadata[metadata_proto.output_type] = (
+              junction_data_utils.metadata_from_proto(metadata_proto.junctions)
+          )
+        case _:
+          raise ValueError(
+              'Unsupported metadata type:'
+              f' {metadata_proto.WhichOneof("payload")}'
+          )
+
+  return OutputMetadata(
+      atac=metadata.get(dna_model_pb2.OUTPUT_TYPE_ATAC),
+      cage=metadata.get(dna_model_pb2.OUTPUT_TYPE_CAGE),
+      dnase=metadata.get(dna_model_pb2.OUTPUT_TYPE_DNASE),
+      rna_seq=metadata.get(dna_model_pb2.OUTPUT_TYPE_RNA_SEQ),
+      chip_histone=metadata.get(dna_model_pb2.OUTPUT_TYPE_CHIP_HISTONE),
+      chip_tf=metadata.get(dna_model_pb2.OUTPUT_TYPE_CHIP_TF),
+      splice_sites=metadata.get(dna_model_pb2.OUTPUT_TYPE_SPLICE_SITES),
+      splice_site_usage=metadata.get(
+          dna_model_pb2.OUTPUT_TYPE_SPLICE_SITE_USAGE
+      ),
+      splice_junctions=metadata.get(dna_model_pb2.OUTPUT_TYPE_SPLICE_JUNCTIONS),
+      contact_maps=metadata.get(dna_model_pb2.OUTPUT_TYPE_CONTACT_MAPS),
+      procap=metadata.get(dna_model_pb2.OUTPUT_TYPE_PROCAP),
+  )
+
+
+def _construct_output(
+    output_dict: Mapping[
+        dna_model_pb2.OutputType,
+        track_data.TrackData | np.ndarray | junction_data.JunctionData,
+    ],
+):
+  """Helper to construct an Output dataclass from a mapping of output types."""
+  output = Output(
+      atac=output_dict.get(dna_model_pb2.OUTPUT_TYPE_ATAC),
+      cage=output_dict.get(dna_model_pb2.OUTPUT_TYPE_CAGE),
+      dnase=output_dict.get(dna_model_pb2.OUTPUT_TYPE_DNASE),
+      rna_seq=output_dict.get(dna_model_pb2.OUTPUT_TYPE_RNA_SEQ),
+      chip_histone=output_dict.get(dna_model_pb2.OUTPUT_TYPE_CHIP_HISTONE),
+      chip_tf=output_dict.get(dna_model_pb2.OUTPUT_TYPE_CHIP_TF),
+      splice_sites=output_dict.get(dna_model_pb2.OUTPUT_TYPE_SPLICE_SITES),
+      splice_site_usage=output_dict.get(
+          dna_model_pb2.OUTPUT_TYPE_SPLICE_SITE_USAGE
+      ),
+      splice_junctions=output_dict.get(
+          dna_model_pb2.OUTPUT_TYPE_SPLICE_JUNCTIONS
+      ),
+      contact_maps=output_dict.get(dna_model_pb2.OUTPUT_TYPE_CONTACT_MAPS),
+      procap=output_dict.get(dna_model_pb2.OUTPUT_TYPE_PROCAP),
+  )
+  return output
+
+
+def _make_output(
+    responses: Iterator[PredictResponse],
+    *,
+    interval: genome.Interval | None = None,
+) -> Output:
+  """Helper to load an Output dataclass from a stream of response protos."""
+  outputs = {}
+
+  for response in responses:
+    match response.WhichOneof('payload'):
+      case 'output':
+        outputs[response.output.output_type] = _make_output_data(
+            response.output, responses, interval=interval
+        )
+      case 'tensor_chunk':
+        raise ValueError('Received tensor chunk before output proto.')
+      case _:
+        raise ValueError(
+            f'Unsupported response type: {response.WhichOneof("payload")}'
+        )
+
+  return _construct_output(outputs)
+
+
+def _make_variant_output(
+    responses: Iterator[dna_model_service_pb2.PredictVariantResponse],
+) -> VariantOutput:
+  """Helper to load an Output dataclass from a stream of response protos."""
+  ref_outputs = {}
+  alt_outputs = {}
+
+  for response in responses:
+    match response.WhichOneof('payload'):
+      case 'reference_output':
+        ref_outputs[response.reference_output.output_type] = _make_output_data(
+            response.reference_output, responses
+        )
+      case 'alternate_output':
+        alt_outputs[response.alternate_output.output_type] = _make_output_data(
+            response.alternate_output, responses
+        )
+      case 'tensor_chunk':
+        raise ValueError('Received tensor chunk before output proto.')
+      case _:
+        raise ValueError(
+            f'Unsupported response type: {response.WhichOneof("payload")}'
+        )
+
+  return VariantOutput(
+      reference=_construct_output(ref_outputs),
+      alternate=_construct_output(alt_outputs),
+  )
+
+
+def _construct_anndata_from_proto(
+    scorer_metadata: Sequence[dna_model_pb2.GeneScorerMetadata] | None,
+    track_metadata: Sequence[dna_model_pb2.TrackMetadata],
+    values: np.ndarray,
+    interval: genome.Interval,
+    variant: dna_model_pb2.Variant | None = None,
+) -> anndata.AnnData:
+  """Helper to construct `AnnData` from a scoring proto."""
+  obs = None
+  if scorer_metadata:
+    metadata = []
+    for gene_proto in scorer_metadata:
+      scorer_metadata = {
+          'gene_id': gene_proto.gene_id,
+      }
+      if gene_proto.HasField('strand'):
+        scorer_metadata['strand'] = str(
+            genome.Strand.from_proto(gene_proto.strand)
+        )
+
+      if gene_proto.HasField('name'):
+        scorer_metadata['gene_name'] = gene_proto.name
+      if gene_proto.HasField('type'):
+        scorer_metadata['gene_type'] = gene_proto.type
+      if gene_proto.HasField('junction_start'):
+        scorer_metadata['junction_Start'] = gene_proto.junction_start
+      if gene_proto.HasField('junction_end'):
+        scorer_metadata['junction_End'] = gene_proto.junction_end
+
+      metadata.append(scorer_metadata)
+
+    obs = pd.DataFrame(metadata)
+    obs.index = obs.index.map(str)
+
+  var = track_data_utils.metadata_from_proto(
+      dna_model_pb2.TracksMetadata(metadata=track_metadata)
+  )
+  var.index = var.index.map(str)
+
+  uns = {'interval': interval}
+  if variant is not None:
+    uns['variant'] = genome.Variant.from_proto(variant)
+  layers = None
+  if values.shape[0] > 1:
+    layers = {'quantiles': values[1]}
+
+  return anndata.AnnData(
+      X=values[0],
+      obs=obs,
+      var=var,
+      uns=uns,
+      layers=layers,
+  )
+
+
+def _construct_score_variant(
+    output: dna_model_pb2.ScoreVariantOutput,
+    responses: Iterator[
+        dna_model_service_pb2.ScoreVariantResponse
+        | dna_model_service_pb2.ScoreIsmVariantResponse
+    ],
+    interval: genome.Interval,
+) -> anndata.AnnData:
+  """Returns `AnnData`  of variant scores from a ScoreVariantOutput proto."""
+  # dna_model_pb2.ScoreVariantOutput currently has ONLY VariantData with values
+  # and metadata.
+  chunks = _read_tensor_chunks(
+      responses, output.variant_data.values.chunk_count
+  )
+  values = tensor_utils.unpack_proto(output.variant_data.values, chunks)
+  values = tensor_utils.upcast_floating(values)
+  anndata_scores = _construct_anndata_from_proto(
+      output.variant_data.metadata.gene_metadata,
+      output.variant_data.metadata.track_metadata,
+      values,
+      interval=interval,
+      variant=output.variant_data.metadata.variant,
+  )
+  return anndata_scores
+
+
+def _make_score_variant_output(
+    responses: Iterator[
+        dna_model_service_pb2.ScoreVariantResponse
+        | dna_model_service_pb2.ScoreIsmVariantResponse
+    ],
+    interval: genome.Interval,
+) -> list[anndata.AnnData]:
+  """Returns a list of `AnnData` variant scores from a stream of responses."""
+
+  anndata_scores = []
+  for response in responses:
+    match response.WhichOneof('payload'):
+      case 'output':
+        anndata_scores.append(
+            _construct_score_variant(
+                response.output, responses, interval=interval
+            )
+        )
+      case 'tensor_chunk':
+        raise ValueError('Received tensor chunk before output proto.')
+      case _:
+        raise ValueError(
+            f'Unsupported response type: {response.WhichOneof("payload")}'
+        )
+  return anndata_scores
+
+
+def _construct_score_interval(
+    output: dna_model_pb2.ScoreIntervalOutput,
+    responses: Iterator[dna_model_service_pb2.ScoreIntervalResponse],
+    interval: genome.Interval,
+) -> anndata.AnnData:
+  """Returns `AnnData` of interval scores from a ScoreIntervalOutput proto."""
+  chunks = _read_tensor_chunks(
+      responses, output.interval_data.values.chunk_count
+  )
+  values = tensor_utils.unpack_proto(output.interval_data.values, chunks)
+  values = tensor_utils.upcast_floating(values)
+  anndata_scores = _construct_anndata_from_proto(
+      output.interval_data.metadata.gene_metadata,
+      output.interval_data.metadata.track_metadata,
+      values,
+      interval=interval,
+  )
+  return anndata_scores
+
+
+def _make_interval_output(
+    responses: Iterator[dna_model_service_pb2.ScoreIntervalResponse],
+    interval: genome.Interval,
+) -> list[anndata.AnnData]:
+  """Returns a list of `AnnData` interval scores from a stream of responses."""
+
+  anndata_scores = []
+  for response in responses:
+    match response.WhichOneof('payload'):
+      case 'output':
+        anndata_scores.append(
+            _construct_score_interval(
+                response.output,
+                responses,
+                interval=interval,
+            )
+        )
+      case 'tensor_chunk':
+        raise ValueError('Received tensor chunk before output proto.')
+      case _:
+        raise ValueError(
+            f'Unsupported response type: {response.WhichOneof("payload")}'
+        )
+  return anndata_scores
+
+
+def _convert_ontologies_to_protos(
+    ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+) -> Sequence[dna_model_pb2.OntologyTerm] | None:
+  """Convert ontology terms or curies to unique list of OntologyTerm protos."""
+  if ontology_terms is None:
+    return None
+  else:
+    protos = []
+    for term in dict.fromkeys(ontology_terms):
+      if isinstance(term, ontology.OntologyTerm):
+        protos.append(term.to_proto())
+      else:
+        protos.append(ontology.from_curie(term).to_proto())
+    return protos
+
+
+def validate_sequence_length(length: int):
+  """Validate that the input sequence length is supported by the model."""
+  if length not in SUPPORTED_SEQUENCE_LENGTHS.values():
+    raise ValueError(
+        f'Sequence length {length} not supported by the model.'
+        f' Supported lengths: {[*SUPPORTED_SEQUENCE_LENGTHS.values()]}'
+    )
+
+
+class DnaClient(dna_model.DnaModel):
+  """Client for interacting with a DNA model server.
+
+  Attributes:
+    channel: gRPC channel to the DNA model server.
+    metadata: Metadata to send with each request.
+    model_version: Optional model version to use for the DNA model server. If
+      none provided, the default model version will be used.
+  """
+
+  def __init__(
+      self,
+      *,
+      channel: grpc.Channel,
+      model_version: ModelVersion | None = None,
+      metadata: Sequence[tuple[str, str]] = (),
+  ):
+    self._channel = channel
+    self._metadata = metadata
+    self._model_version = (
+        model_version.name if model_version is not None else None
+    )
+
+  @retry_rpc
+  def predict_sequence(
+      self,
+      sequence: str,
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+      interval: genome.Interval | None = None,
+  ) -> Output:
+    """Generate predictions for a given DNA sequence.
+
+    Args:
+      sequence: DNA sequence to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+      interval: Optional interval from which the sequence was derived. This is
+        used as the interval in the output TrackData.
+
+    Returns:
+      Output for the provided DNA sequence.
+    """
+    if not (unique_values := set(sequence)).issubset(
+        _VALID_SEQUENCE_CHARACTERS
+    ):
+      invalid_characters = ','.join(unique_values - _VALID_SEQUENCE_CHARACTERS)
+      raise ValueError(
+          'Invalid sequence, must only contain the characters "ACGTN". Found'
+          f' invalid characters: "{invalid_characters}"'
+      )
+    validate_sequence_length(len(sequence))
+    requested_outputs = [o.to_proto() for o in dict.fromkeys(requested_outputs)]
+    request = dna_model_service_pb2.PredictSequenceRequest(
+        sequence=sequence,
+        organism=organism.to_proto(),
+        ontology_terms=_convert_ontologies_to_protos(ontology_terms),
+        requested_outputs=requested_outputs,
+        model_version=self._model_version,
+    )
+    responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+        self._channel
+    ).PredictSequence(iter([request]), metadata=self._metadata)
+    return _make_output(responses, interval=interval)
+
+  @retry_rpc
+  def predict_interval(
+      self,
+      interval: genome.Interval,
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+  ) -> Output:
+    """Generate predictions for a given DNA interval.
+
+    Args:
+      interval: DNA interval to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+
+    Returns:
+      Output for the provided DNA interval.
+    """
+    validate_sequence_length(interval.width)
+    requested_outputs = [o.to_proto() for o in dict.fromkeys(requested_outputs)]
+    request = dna_model_service_pb2.PredictIntervalRequest(
+        interval=interval.to_proto(),
+        organism=organism.to_proto(),
+        ontology_terms=_convert_ontologies_to_protos(ontology_terms),
+        requested_outputs=requested_outputs,
+        model_version=self._model_version,
+    )
+    responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+        self._channel
+    ).PredictInterval(iter([request]), metadata=self._metadata)
+    return _make_output(responses)
+
+  @retry_rpc
+  def predict_variant(
+      self,
+      interval: genome.Interval,
+      variant: genome.Variant,
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+  ) -> VariantOutput:
+    """Generate predictions for a given DNA variant.
+
+    Args:
+      interval: DNA interval to make prediction for.
+      variant: DNA variant to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+
+    Returns:
+      Variant output for the provided DNA interval and variant.
+    """
+    validate_sequence_length(interval.width)
+    requested_outputs = [o.to_proto() for o in dict.fromkeys(requested_outputs)]
+    request = dna_model_service_pb2.PredictVariantRequest(
+        interval=interval.to_proto(),
+        variant=variant.to_proto(),
+        organism=organism.to_proto(),
+        ontology_terms=_convert_ontologies_to_protos(ontology_terms),
+        requested_outputs=requested_outputs,
+        model_version=self._model_version,
+    )
+    responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+        self._channel
+    ).PredictVariant(iter([request]), metadata=self._metadata)
+    return _make_variant_output(responses)
+
+  @retry_rpc
+  def score_interval(
+      self,
+      interval: genome.Interval,
+      interval_scorers: Sequence[interval_scorers_lib.IntervalScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+  ) -> list[anndata.AnnData]:
+    """Generate interval scores for a single given interval.
+
+    Args:
+      interval: Interval to make prediction for.
+      interval_scorers: Sequence of interval scorers to use for scoring. If no
+        interval scorers are provided, the recommended interval scorers for the
+        organism will be used.
+      organism: Organism to use for the prediction.
+
+    Returns:
+      List of `AnnData` interval scores.
+    """
+    if not interval_scorers:
+      interval_scorers = list(
+          interval_scorers_lib.RECOMMENDED_INTERVAL_SCORERS.values()
+      )
+
+    validate_sequence_length(interval.width)
+    if len(interval_scorers) != len(set(interval_scorers)):
+      raise ValueError(
+          f'Duplicate interval scorers requested: {interval_scorers}.'
+      )
+    if len(interval_scorers) > MAX_VARIANT_SCORERS_PER_REQUEST:
+      raise ValueError(
+          'Too many interval scorers requested, '
+          f'maximum allowed: {MAX_VARIANT_SCORERS_PER_REQUEST}.'
+      )
+    if any(
+        scorer.width is not None and scorer.width > interval.width
+        for scorer in interval_scorers
+    ):
+      raise ValueError('Interval scorers must have widths <= interval width.')
+    request = dna_model_service_pb2.ScoreIntervalRequest(
+        interval=interval.to_proto(),
+        interval_scorers=[scorer.to_proto() for scorer in interval_scorers],
+        organism=organism.to_proto(),
+        model_version=self._model_version,
+    )
+    responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+        self._channel
+    ).ScoreInterval(iter([request]), metadata=self._metadata)
+    interval_outputs = _make_interval_output(responses, interval)
+    for ix in range(len(interval_scorers)):
+      interval_outputs[ix].uns['interval_scorer'] = interval_scorers[ix]
+    return interval_outputs
+
+  @retry_rpc
+  def score_variant(
+      self,
+      interval: genome.Interval,
+      variant: genome.Variant,
+      variant_scorers: Sequence[variant_scorers_lib.VariantScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+  ) -> list[anndata.AnnData]:
+    """Generate variant scores for a single given DNA variant.
+
+    Args:
+      interval: DNA interval to make prediction for.
+      variant: DNA variant to make prediction for.
+      variant_scorers: Sequence of variant scorers to use for scoring. If no
+        variant scorers are provided, the recommended variant scorers for the
+        organism will be used.
+      organism: Organism to use for the prediction.
+
+    Returns:
+      List of `AnnData` variant scores.
+    """
+    if not variant_scorers:
+      variant_scorers = variant_scorers_lib.get_recommended_scorers(
+          organism.to_proto()
+      )
+
+    validate_sequence_length(interval.width)
+    for variant_scorer in variant_scorers:
+      supported_organisms = variant_scorers_lib.SUPPORTED_ORGANISMS[
+          variant_scorer.base_variant_scorer
+      ]
+      if organism.to_proto() not in supported_organisms:
+        supported_organisms = [
+            dna_model_pb2.Organism.Name(o).removeprefix('ORGANISM_')
+            for o in supported_organisms
+        ]
+        raise ValueError(
+            f'Unsupported organism: {organism.name} for scorer'
+            f' {variant_scorer}. Supported organisms:'
+            f' {supported_organisms}'
+        )
+    if len(variant_scorers) != len(set(variant_scorers)):
+      raise ValueError(
+          f'Duplicate variant scorers requested: {variant_scorers}.'
+      )
+    if len(variant_scorers) > MAX_VARIANT_SCORERS_PER_REQUEST:
+      raise ValueError(
+          'Too many variant scorers requested, '
+          f'maximum allowed: {MAX_VARIANT_SCORERS_PER_REQUEST}.'
+      )
+    request = dna_model_service_pb2.ScoreVariantRequest(
+        interval=interval.to_proto(),
+        variant=variant.to_proto(),
+        organism=organism.to_proto(),
+        variant_scorers=[vs.to_proto() for vs in variant_scorers],
+        model_version=self._model_version,
+    )
+    responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+        self._channel
+    ).ScoreVariant(iter([request]), metadata=self._metadata)
+    variant_outputs = _make_score_variant_output(responses, interval)
+    for ix in range(len(variant_scorers)):
+      variant_outputs[ix].uns['variant_scorer'] = variant_scorers[ix]
+    return variant_outputs
+
+  def score_ism_variants(
+      self,
+      interval: genome.Interval,
+      ism_interval: genome.Interval,
+      variant_scorers: Sequence[variant_scorers_lib.VariantScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      interval_variant: genome.Variant | None = None,
+      progress_bar: bool = True,
+      max_workers: int = dna_model.DEFAULT_MAX_WORKERS,
+  ) -> list[list[anndata.AnnData]]:
+    """Generate in-silico mutagenesis (ISM) variant scores for a given interval.
+
+    Args:
+      interval: DNA interval to make the prediction for.
+      ism_interval: Interval to perform ISM.
+      variant_scorers: Sequence of variant scorers to use for scoring each
+        variant. If no variant scorers are provided, the recommended variant
+        scorers for the organism will be used.
+      organism: Organism to use for the prediction.
+      interval_variant: Optional variant to apply to the sequence. If provided,
+        the alternate allele is used for in-silico mutagenesis, otherwise the
+        unaltered reference sequence is used.
+      progress_bar: If True, show a progress bar.
+      max_workers: Number of parallel workers to use.
+
+    Returns:
+      List of variant scores for each variant in the ISM interval.
+    """
+    if not variant_scorers:
+      variant_scorers = variant_scorers_lib.get_recommended_scorers(
+          organism.to_proto()
+      )
+
+    validate_sequence_length(interval.width)
+    if ism_interval.negative_strand:
+      raise ValueError('ISM interval must be on the positive strand.')
+
+    ism_intervals = []
+    for position in range(0, ism_interval.width, MAX_ISM_INTERVAL_WIDTH):
+      ism_intervals.append(
+          genome.Interval(
+              ism_interval.chromosome,
+              ism_interval.start + position,
+              min(
+                  ism_interval.start + position + MAX_ISM_INTERVAL_WIDTH,
+                  ism_interval.end,
+              ),
+          )
+      )
+
+    @retry_rpc
+    def _score_ism_variant(
+        ism_interval: genome.Interval,
+    ) -> list[anndata.AnnData]:
+      request = dna_model_service_pb2.ScoreIsmVariantRequest(
+          interval=interval.to_proto(),
+          ism_interval=ism_interval.to_proto(),
+          organism=organism.to_proto(),
+          variant_scorers=[vs.to_proto() for vs in variant_scorers],
+          model_version=self._model_version,
+          interval_variant=interval_variant.to_proto()
+          if interval_variant is not None
+          else None,
+      )
+      responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+          self._channel
+      ).ScoreIsmVariant(iter([request]), metadata=self._metadata)
+
+      return _make_score_variant_output(responses, interval)
+
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    ) as executor:
+      futures = [
+          executor.submit(_score_ism_variant, ism_interval)
+          for ism_interval in ism_intervals
+      ]
+      ism_scores = []
+      for future in tqdm.auto.tqdm(
+          concurrent.futures.as_completed(futures),
+          total=len(futures),
+          disable=not progress_bar,
+      ):
+        ism_scores.extend(future.result())
+
+    # Group scores by variant.
+    ism_scores_by_variant: dict[str, list[anndata.AnnData]] = {}
+    for variant_score in ism_scores:
+      ism_scores_by_variant.setdefault(
+          str(variant_score.uns['variant']), []
+      ).append(variant_score)
+
+    return list(ism_scores_by_variant.values())
+
+  def output_metadata(
+      self, organism: Organism = Organism.HOMO_SAPIENS
+  ) -> OutputMetadata:
+    """Get the metadata for a given organism.
+
+    Args:
+      organism: Organism to get metadata for.
+
+    Returns:
+      OutputMetadata for the provided organism.
+    """
+    request = dna_model_service_pb2.MetadataRequest(
+        organism=organism.to_proto()
+    )
+    responses = dna_model_service_pb2_grpc.DnaModelServiceStub(
+        self._channel
+    ).GetMetadata(request, metadata=self._metadata)
+    return construct_output_metadata(responses)
+
+
+def create(
+    api_key: str,
+    *,
+    model_version: ModelVersion | None = None,
+    timeout: float | None = None,
+    address: str | None = None,
+) -> DnaClient:
+  """Creates a model client for a given API key.
+
+  Args:
+    api_key: API key to use for authentication.
+    model_version: Optional model version to use for the DNA model server. If
+      none is provided, the default model will be used.
+    timeout: Optional timeout for waiting for the channel to be ready.
+    address: Optional server address to connect to.
+
+  Returns:
+   `DnaClient` instance.
+  """
+  options = [
+      ('grpc.max_send_message_length', -1),
+      ('grpc.max_receive_message_length', -1),
+  ]
+  address = address or 'dns:///gdmscience.googleapis.com:443'
+  channel = grpc.secure_channel(
+      address, grpc.ssl_channel_credentials(), options=options
+  )
+  grpc.channel_ready_future(channel).result(timeout)
+
+  return DnaClient(
+      channel=channel,
+      model_version=model_version,
+      metadata=[('x-goog-api-key', api_key)],
+  )

flax_model/alphagenome/_sdk/models/dna_model.py

@@ -0,0 +1,507 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Abstract base class for AlphaGenome DNA models."""
+
+import abc
+from collections.abc import Iterable, Sequence
+import concurrent.futures
+import enum
+
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import ontology
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome._sdk.models import interval_scorers as interval_scorers_lib
+from flax_model.alphagenome._sdk.models import variant_scorers as variant_scorers_lib
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+import anndata
+import tqdm.auto
+
+# Default maximum number of workers to use for parallel requests.
+DEFAULT_MAX_WORKERS = 5
+
+
+@enum.unique
+class ModelVersion(enum.Enum):
+  """Enumeration of all available model versions.
+
+  A fold is a part of the genome that is held out from training, and is used for
+  model validation.
+
+  Folds are numbered from 0 to 3, and represent disjoint subsets of the genome.
+
+  A model version that is designated FOLD_#, where # is an integer from 0 to 3,
+  indicates that the model was trained with that particular fold held out.
+
+  The ALL_FOLDS version refers to the distilled model that was trained on
+  an ensemble of teacher models trained on all folds.
+  """
+
+  ALL_FOLDS = enum.auto()  # Default model version if None explicitly set.
+  FOLD_0 = enum.auto()
+  FOLD_1 = enum.auto()
+  FOLD_2 = enum.auto()
+  FOLD_3 = enum.auto()
+
+
+class Organism(enum.Enum):
+  """Enumeration of all the available organisms."""
+
+  HOMO_SAPIENS = dna_model_pb2.ORGANISM_HOMO_SAPIENS
+  MUS_MUSCULUS = dna_model_pb2.ORGANISM_MUS_MUSCULUS
+
+  def to_proto(self) -> dna_model_pb2.Organism:
+    return self.value
+
+  def __lt__(self, other: 'Organism') -> bool:
+    if not isinstance(other, Organism):
+      return NotImplemented
+    return self.value < other.value
+
+
+class DnaModel(metaclass=abc.ABCMeta):
+  """Abstract base class for AlphaGenome DNA models."""
+
+  @abc.abstractmethod
+  def predict_sequence(
+      self,
+      sequence: str,
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[dna_output.OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+      interval: genome.Interval | None = None,
+  ) -> dna_output.Output:
+    """Generate predictions for a given DNA sequence.
+
+    Args:
+      sequence: DNA sequence to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+      interval: Optional interval from which the sequence was derived. This is
+        used as the interval in the output TrackData.
+
+    Returns:
+      Output for the provided DNA sequence.
+    """
+
+  def predict_sequences(
+      self,
+      sequences: Sequence[str],
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[dna_output.OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+      progress_bar: bool = True,
+      max_workers: int = DEFAULT_MAX_WORKERS,
+      intervals: Sequence[genome.Interval] | None = None,
+  ) -> list[dna_output.Output]:
+    """Generate predictions for a given DNA sequence.
+
+    Args:
+      sequences: DNA sequences to make predictions for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+      progress_bar: If True, show a progress bar.
+      max_workers: Number of parallel workers to use.
+      intervals: Optional intervals from which the sequences were derived. This
+        is used as the interval in the output TrackData. Must be the same length
+        as `sequences` if provided.
+
+    Returns:
+      Outputs for the provided DNA sequences.
+    """
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    ) as executor:
+      futures = [
+          executor.submit(
+              self.predict_sequence,
+              sequence=sequence,
+              organism=organism,
+              ontology_terms=ontology_terms,
+              requested_outputs=requested_outputs,
+              interval=interval,
+          )
+          for sequence, interval in zip(
+              sequences, intervals or [None] * len(sequences), strict=True
+          )
+      ]
+      futures_as_completed = tqdm.auto.tqdm(
+          concurrent.futures.as_completed(futures),
+          total=len(futures),
+          disable=not progress_bar,
+      )
+      for future in futures_as_completed:
+        if (exception := future.exception()) is not None:
+          executor.shutdown(wait=False, cancel_futures=True)
+          raise exception
+
+      return [future.result() for future in futures]
+
+  @abc.abstractmethod
+  def predict_interval(
+      self,
+      interval: genome.Interval,
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[dna_output.OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+  ) -> dna_output.Output:
+    """Generate predictions for a given DNA interval.
+
+    Args:
+      interval: DNA interval to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+
+    Returns:
+      Output for the provided DNA interval.
+    """
+
+  def predict_intervals(
+      self,
+      intervals: Sequence[genome.Interval],
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[dna_output.OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+      progress_bar: bool = True,
+      max_workers: int = DEFAULT_MAX_WORKERS,
+  ) -> list[dna_output.Output]:
+    """Generate predictions for a sequence of DNA intervals.
+
+    Args:
+      intervals: DNA intervals to make predictions for.
+      organism: Organism to use for the predictions.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+      progress_bar: If True, show a progress bar.
+      max_workers: Number of parallel workers to use.
+
+    Returns:
+      Outputs for the provided DNA intervals.
+    """
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    ) as executor:
+      futures = [
+          executor.submit(
+              self.predict_interval,
+              interval=interval,
+              organism=organism,
+              ontology_terms=ontology_terms,
+              requested_outputs=requested_outputs,
+          )
+          for interval in intervals
+      ]
+      futures_as_completed = tqdm.auto.tqdm(
+          concurrent.futures.as_completed(futures),
+          total=len(futures),
+          disable=not progress_bar,
+      )
+      for future in futures_as_completed:
+        if (exception := future.exception()) is not None:
+          executor.shutdown(wait=False, cancel_futures=True)
+          raise exception
+
+      return [future.result() for future in futures]
+
+  @abc.abstractmethod
+  def predict_variant(
+      self,
+      interval: genome.Interval,
+      variant: genome.Variant,
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[dna_output.OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+  ) -> dna_output.VariantOutput:
+    """Generate predictions for a given DNA variant.
+
+    Args:
+      interval: DNA interval to make prediction for.
+      variant: DNA variant to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+
+    Returns:
+      Variant output for the provided DNA interval and variant.
+    """
+
+  def predict_variants(
+      self,
+      intervals: genome.Interval | Sequence[genome.Interval],
+      variants: Sequence[genome.Variant],
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      requested_outputs: Iterable[dna_output.OutputType],
+      ontology_terms: Iterable[ontology.OntologyTerm | str] | None,
+      progress_bar: bool = True,
+      max_workers: int = DEFAULT_MAX_WORKERS,
+  ) -> list[dna_output.VariantOutput]:
+    """Generate predictions for a given DNA variant.
+
+    Args:
+      intervals: DNA interval(s) to make predictions for.
+      variants: DNA variants to make prediction for.
+      organism: Organism to use for the prediction.
+      requested_outputs: Iterable of OutputTypes indicating which subsets of
+        predictions to return.
+      ontology_terms: Iterable of ontology terms or curies to generate
+        predictions for. If None returns all ontologies.
+      progress_bar: If True, show a progress bar.
+      max_workers: Number of parallel workers to use.
+
+    Returns:
+      Variant outputs for each DNA interval and variant pair.
+    """
+    if not isinstance(intervals, Sequence):
+      intervals = [intervals] * len(variants)
+    elif len(intervals) != len(variants):
+      raise ValueError(
+          'Intervals and variants must have the same length.'
+          f'Got {len(intervals)} intervals and {len(variants)} variants.'
+      )
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    ) as executor:
+      futures = [
+          executor.submit(
+              self.predict_variant,
+              interval=interval,
+              variant=variant,
+              organism=organism,
+              ontology_terms=ontology_terms,
+              requested_outputs=requested_outputs,
+          )
+          for interval, variant in zip(intervals, variants, strict=True)
+      ]
+      futures_as_completed = tqdm.auto.tqdm(
+          concurrent.futures.as_completed(futures),
+          total=len(futures),
+          disable=not progress_bar,
+      )
+      for future in futures_as_completed:
+        if (exception := future.exception()) is not None:
+          executor.shutdown(wait=False, cancel_futures=True)
+          raise exception
+
+      return [future.result() for future in futures]
+
+  @abc.abstractmethod
+  def score_interval(
+      self,
+      interval: genome.Interval,
+      interval_scorers: Sequence[interval_scorers_lib.IntervalScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+  ) -> list[anndata.AnnData]:
+    """Generate interval scores for a single given interval.
+
+    Args:
+      interval: Interval to make prediction for.
+      interval_scorers: Sequence of interval scorers to use for scoring. If no
+        interval scorers are provided, the recommended interval scorers for the
+        organism will be used.
+      organism: Organism to use for the prediction.
+
+    Returns:
+      List of `AnnData` interval scores.
+    """
+
+  def score_intervals(
+      self,
+      intervals: Sequence[genome.Interval],
+      interval_scorers: Sequence[interval_scorers_lib.IntervalScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      progress_bar: bool = True,
+      max_workers: int = DEFAULT_MAX_WORKERS,
+  ) -> list[list[anndata.AnnData]]:
+    """Generate interval scores for a sequence of intervals.
+
+    Args:
+      intervals: Sequence of DNA intervals to make prediction for.
+      interval_scorers: Sequence of interval scorers to use for scoring. If no
+        interval scorers are provided, the recommended interval scorers for the
+        organism will be used.
+      organism: Organism to use for the prediction.
+      progress_bar: If True, show a progress bar.
+      max_workers: Number of parallel workers to use.
+
+    Returns:
+      List of `AnnData` lists corresponding to score_interval
+      outputs for each interval.
+    """
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    ) as executor:
+      futures = [
+          executor.submit(
+              self.score_interval,
+              interval=interval,
+              interval_scorers=interval_scorers,
+              organism=organism,
+          )
+          for interval in intervals
+      ]
+      futures_as_completed = tqdm.auto.tqdm(
+          concurrent.futures.as_completed(futures),
+          total=len(futures),
+          disable=not progress_bar,
+      )
+      for future in futures_as_completed:
+        if (exception := future.exception()) is not None:
+          executor.shutdown(wait=False, cancel_futures=True)
+          raise exception
+
+      results = [future.result() for future in futures]
+    return results
+
+  @abc.abstractmethod
+  def score_variant(
+      self,
+      interval: genome.Interval,
+      variant: genome.Variant,
+      variant_scorers: Sequence[variant_scorers_lib.VariantScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+  ) -> list[anndata.AnnData]:
+    """Generate variant scores for a single given DNA variant.
+
+    Args:
+      interval: DNA interval to make prediction for.
+      variant: DNA variant to make prediction for.
+      variant_scorers: Sequence of variant scorers to use for scoring. If no
+        variant scorers are provided, the recommended variant scorers for the
+        organism will be used.
+      organism: Organism to use for the prediction.
+
+    Returns:
+      List of `AnnData` variant scores.
+    """
+
+  def score_variants(
+      self,
+      intervals: genome.Interval | Sequence[genome.Interval],
+      variants: Sequence[genome.Variant],
+      variant_scorers: Sequence[variant_scorers_lib.VariantScorerTypes] = (),
+      *,
+      organism: Organism = Organism.HOMO_SAPIENS,
+      progress_bar: bool = True,
+      max_workers: int = DEFAULT_MAX_WORKERS,
+  ) -> list[list[anndata.AnnData]]:
+    """Generate variant scores for a sequence of variants.
+
+    Args:
+      intervals: DNA interval(s) to make prediction for. If a single interval is
+        provided, then the same interval will be used for all variants.
+      variants: Sequence of DNA variants to make prediction for.
+      variant_scorers: Sequence of variant scorers to use for scoring. If no
+        variant scorers are provided, the recommended variant scorers for the
+        organism will be used.
+      organism: Organism to use for the prediction.
+      progress_bar: If True, show a progress bar.
+      max_workers: Number of parallel workers to use.
+
+    Returns:
+      List of `AnnData` lists corresponding to score_variant outputs for each
+      variant.
+    """
+    if not isinstance(intervals, Sequence):
+      intervals = [intervals] * len(variants)
+    if len(intervals) != len(variants):
+      raise ValueError(
+          'Intervals and variants must have the same length.'
+          f'Got {len(intervals)} intervals and {len(variants)} variants.'
+      )
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    ) as executor:
+      futures = [
+          executor.submit(
+              self.score_variant,
+              interval=interval,
+              variant=variant,
+              variant_scorers=variant_scorers,
+              organism=organism,
+          )
+          for interval, variant in zip(intervals, variants, strict=True)
+      ]
+      futures_as_completed = tqdm.auto.tqdm(
+          concurrent.futures.as_completed(futures),
+          total=len(futures),
+          disable=not progress_bar,
+      )
+      for future in futures_as_completed:
+        if (exception := future.exception()) is not None:
+          executor.shutdown(wait=False, cancel_futures=True)
+          raise exception
+
+      return [future.result() for future in futures]
+
+  @abc.abstractmethod
+  def score_ism_variants(
+      self,
+      interval: genome.Interval,
+      ism_interval: genome.Interval,
+      variant_scorers: Sequence[variant_scorers_lib.VariantScorerTypes] = (),
+      *,
+      interval_variant: genome.Variant | None = None,
+      organism: Organism = Organism.HOMO_SAPIENS,
+  ) -> list[list[anndata.AnnData]]:
+    """Generate in-silico mutagenesis (ISM) variant scores for a given interval.
+
+    Args:
+      interval: DNA interval to make the prediction for.
+      ism_interval: Interval to perform ISM.
+      variant_scorers: Sequence of variant scorers to use for scoring each
+        variant. If no variant scorers are provided, the recommended variant
+        scorers for the organism will be used.
+      interval_variant: Optional variant to apply to the sequence. If provided,
+        the alternate allele is used for in-silico mutagenesis, otherwise the
+        unaltered reference sequence is used.
+      organism: Organism to use for the prediction.
+
+    Returns:
+      List of variant scores for each variant in the ISM interval.
+    """
+
+  @abc.abstractmethod
+  def output_metadata(
+      self, organism: Organism = Organism.HOMO_SAPIENS
+  ) -> dna_output.OutputMetadata:
+    """Get the metadata for a given organism.
+
+    Args:
+      organism: Organism to get metadata for.
+
+    Returns:
+      OutputMetadata for the provided organism.
+    """

flax_model/alphagenome/_sdk/models/dna_output.py

@@ -0,0 +1,435 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for AlphaGenome model outputs."""
+
+from collections.abc import Callable, Iterable, Mapping
+import dataclasses
+import enum
+from typing import Literal
+
+from flax_model.alphagenome._sdk import typing
+from flax_model.alphagenome._sdk.data import junction_data
+from flax_model.alphagenome._sdk.data import ontology
+from flax_model.alphagenome._sdk.data import track_data
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+import pandas as pd
+
+
+class OutputType(enum.Enum):
+  """Enumeration of all the available types of outputs.
+
+  Attributes:
+    ATAC: ATAC-seq tracks capturing chromatin accessibility.
+    CAGE: CAGE (Cap Analysis of Gene Expression) tracks capturing gene
+      expression.
+    DNASE: DNase I hypersensitive site tracks capturing chromatin accessibility.
+    RNA_SEQ: RNA sequencing tracks capturing gene expression.
+    CHIP_HISTONE: ChIP-seq tracks capturing histone modifications.
+    CHIP_TF: ChIP-seq tracks capturing transcription factor binding.
+    SPLICE_SITES: Splice site tracks capturing donor and acceptor splice sites.
+    SPLICE_SITE_USAGE: Splice site usage tracks capturing the fraction of the
+      time that each splice site is used.
+    SPLICE_JUNCTIONS: Splice junction tracks capturing split read RNA-seq counts
+      for each junction.
+    CONTACT_MAPS: Contact map tracks capturing 3D DNA-DNA contact probabilities.
+    PROCAP: Precision Run-On sequencing and capping, used to measure gene
+      expression.
+  """
+
+  ATAC = dna_model_pb2.OUTPUT_TYPE_ATAC
+  CAGE = dna_model_pb2.OUTPUT_TYPE_CAGE
+  DNASE = dna_model_pb2.OUTPUT_TYPE_DNASE
+  RNA_SEQ = dna_model_pb2.OUTPUT_TYPE_RNA_SEQ
+  CHIP_HISTONE = dna_model_pb2.OUTPUT_TYPE_CHIP_HISTONE
+  CHIP_TF = dna_model_pb2.OUTPUT_TYPE_CHIP_TF
+  SPLICE_SITES = dna_model_pb2.OUTPUT_TYPE_SPLICE_SITES
+  SPLICE_SITE_USAGE = dna_model_pb2.OUTPUT_TYPE_SPLICE_SITE_USAGE
+  SPLICE_JUNCTIONS = dna_model_pb2.OUTPUT_TYPE_SPLICE_JUNCTIONS
+  CONTACT_MAPS = dna_model_pb2.OUTPUT_TYPE_CONTACT_MAPS
+  PROCAP = dna_model_pb2.OUTPUT_TYPE_PROCAP
+
+  def __lt__(self, other: 'OutputType'):
+    """Compares if an other `OutputType` enum value is less than this one."""
+    return self.value < other.value
+
+  def to_proto(self) -> dna_model_pb2.OutputType:
+    """Converts the `OutputType` enum to a protobuf enum."""
+    return self.value
+
+  def __repr__(self) -> str:
+    """Returns name of the `OutputType` enum as the string representation."""
+    return self.name
+
+
+@typing.jaxtyped
+@dataclasses.dataclass(frozen=True)
+class Output:
+  """Model outputs for a single prediction.
+
+  Attributes:
+    atac: TrackData of type OutputType.ATAC.
+    cage: TrackData of type OutputType.CAGE.
+    dnase: TrackData of type OutputType.DNASE.
+    rna_seq: TrackData of type OutputType.RNA_SEQ.
+    chip_histone: TrackData of type OutputType.CHIP_HISTONE.
+    chip_tf: TrackData of type OutputType.CHIP_TF.
+    splice_sites: TrackData of type OutputType.SPLICE_SITES.
+    splice_site_usage: TrackData of type OutputType.SPLICE_SITE_USAGE.
+    splice_junctions: TrackData of type OutputType.SPLICE_JUNCTIONS.
+    contact_maps: TrackData of type OutputType.CONTACT_MAPS.
+    procap: TrackData of type OutputType.PROCAP.
+  """
+
+  atac: track_data.TrackData | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.ATAC}
+  )
+  cage: track_data.TrackData | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.CAGE}
+  )
+  dnase: track_data.TrackData | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.DNASE}
+  )
+  rna_seq: track_data.TrackData | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.RNA_SEQ}
+  )
+  chip_histone: track_data.TrackData | None = dataclasses.field(
+      default=None,
+      metadata={'output_type': OutputType.CHIP_HISTONE},
+  )
+  chip_tf: track_data.TrackData | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.CHIP_TF}
+  )
+
+  splice_sites: track_data.TrackData | None = dataclasses.field(
+      default=None,
+      metadata={'output_type': OutputType.SPLICE_SITES},
+  )
+  splice_site_usage: track_data.TrackData | None = dataclasses.field(
+      default=None,
+      metadata={'output_type': OutputType.SPLICE_SITE_USAGE},
+  )
+  splice_junctions: junction_data.JunctionData | None = dataclasses.field(
+      default=None,
+      metadata={'output_type': OutputType.SPLICE_JUNCTIONS},
+  )
+
+  contact_maps: track_data.TrackData | None = dataclasses.field(
+      default=None,
+      metadata={'output_type': OutputType.CONTACT_MAPS},
+  )
+  procap: track_data.TrackData | None = dataclasses.field(
+      default=None,
+      metadata={'output_type': OutputType.PROCAP},
+  )
+
+  def get(
+      self, output_type: OutputType
+  ) -> track_data.TrackData | junction_data.JunctionData | None:
+    """Gets the track data for the specified output type.
+
+    Args:
+      output_type: The type of output to retrieve.
+
+    Returns:
+      The track data for the specified output type, or None if no such data
+      exists.
+    """
+    for field in dataclasses.fields(self):
+      if field.metadata['output_type'] == output_type:
+        return getattr(self, field.name)
+
+  def map_track_data(
+      self,
+      fn: Callable[
+          [track_data.TrackData, OutputType],
+          track_data.TrackData | None,
+      ],
+  ) -> 'Output':
+    """Applies a transformation function to each `TrackData`.
+
+    Args:
+      fn: The function to apply to each `TrackData`. It should take a
+        `TrackData` object and an `OutputType` enum as input and return a
+        `TrackData` object or None.
+
+    Returns:
+      A new `Output` object with the transformed track data.
+    """
+    output_dict = {}
+    for field in dataclasses.fields(self):
+      output_type = field.metadata['output_type']
+      value = self.get(output_type)
+      if isinstance(value, track_data.TrackData):
+        output_dict[field.name] = fn(value, output_type)
+      else:
+        output_dict[field.name] = value
+    return Output(**output_dict)
+
+  def filter_to_strand(self, strand: Literal['+', '-', '.']) -> 'Output':
+    """Filters tracks by DNA strand.
+
+    Args:
+      strand: The strand to filter by ('+', '-', or '.').
+
+    Returns:
+      A new `Output` object with only the tracks on the specified strand.
+    """
+
+    def _filter_to_strand(
+        tdata: track_data.TrackData, output_type: OutputType
+    ) -> track_data.TrackData | None:
+      del output_type  # Unused.
+      return tdata.filter_tracks(tdata.strands == strand)
+
+    output = self.map_track_data(_filter_to_strand)
+    if output.splice_junctions is not None:
+      output = dataclasses.replace(
+          output,
+          splice_junctions=output.splice_junctions.filter_to_strand(strand),
+      )
+    return output
+
+  def filter_ontology_terms(
+      self, ontology_terms: Iterable[ontology.OntologyTerm]
+  ) -> 'Output':
+    """Filters tracks to specific ontology terms.
+
+    Args:
+      ontology_terms: An iterable of `OntologyTerm` objects to filter to.
+
+    Returns:
+      A new `Output` object with only the tracks associated with the specified
+      ontology terms.
+    """
+
+    def _filter_ontology(
+        tdata: track_data.TrackData, output_type: OutputType
+    ) -> track_data.TrackData | None:
+      del output_type  # Unused.
+      if track_ontologies := tdata.ontology_terms:
+        return tdata.filter_tracks(
+            [o in ontology_terms for o in track_ontologies]
+        )
+      else:
+        return tdata
+
+    return self.map_track_data(_filter_ontology)
+
+  def filter_output_type(self, output_types: Iterable[OutputType]) -> 'Output':
+    """Filters tracks to specific output type.
+
+    Args:
+      output_types: An iterable of `OutputType` enums to filter by.
+
+    Returns:
+      A new `Output` object with only the tracks of the specified output types.
+    """
+    output_dict = {}
+    for field in dataclasses.fields(self):
+      output_type = field.metadata['output_type']
+      if output_type in output_types:
+        output_dict[field.name] = self.get(output_type)
+      else:
+        output_dict[field.name] = None
+
+    return Output(**output_dict)
+
+  def resize(self, width: int) -> 'Output':
+    """Resizes all track data to a specified width.
+
+    Args:
+      width: The desired width in base pairs.
+
+    Returns:
+      A new `Output` object with resized track data.
+    """
+    return self.map_track_data(lambda tdata, _: tdata.resize(width))
+
+  def __add__(self, other: 'Output') -> 'Output':
+    """Adds the values of two `Output` objects element-wise.
+
+    Args:
+      other: The `Output` object to add.
+
+    Returns:
+      A new `Output` object with the summed values.
+    """
+
+    def add_track_data(
+        track_data1,
+        output_type: OutputType,
+    ):
+      """Adds two `TrackData` objects for a specific output type."""
+      track_data2 = other.get(output_type)
+      if track_data1 is None or track_data2 is None:
+        return None
+      return track_data1 + track_data2
+
+    return self.map_track_data(add_track_data)
+
+  def __sub__(self, other: 'Output') -> 'Output':
+    """Subtracts the values of two `Output` objects element-wise.
+
+    Args:
+      other: The `Output` object to subtract.
+
+    Returns:
+      A new `Output` object with the difference of the values.
+    """
+
+    def sub_track_data(track_data1, output_type: OutputType):
+      """Subtracts two `TrackData` objects for a specific output type."""
+      track_data2 = other.get(output_type)
+      if track_data1 is None or track_data2 is None:
+        return None
+      return track_data1 - track_data2
+
+    return self.map_track_data(sub_track_data)
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class OutputMetadata:
+  """Metadata detailing the content of model output.
+
+  Attributes:
+    atac: Metadata for ATAC-seq tracks.
+    cage: Metadata for CAGE tracks.
+    dnase: Metadata for DNase I hypersensitive site tracks.
+    rna_seq: Metadata for RNA sequencing tracks.
+    chip_histone: Metadata for ChIP-seq tracks capturing histone modifications.
+    chip_tf: Metadata for ChIP-seq tracks capturing transcription factor
+      binding.
+    splice_sites: Metadata for splice site tracks.
+    splice_site_usage: Metadata for splice site usage tracks.
+    splice_junctions: Metadata for splice junction tracks.
+    contact_maps: Metadata for contact map tracks.
+    procap: Metadata for procap tracks.
+  """
+
+  atac: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.ATAC}
+  )
+  cage: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.CAGE}
+  )
+  dnase: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.DNASE}
+  )
+  rna_seq: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.RNA_SEQ}
+  )
+  chip_histone: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.CHIP_HISTONE}
+  )
+  chip_tf: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.CHIP_TF}
+  )
+  splice_sites: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.SPLICE_SITES}
+  )
+  splice_site_usage: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.SPLICE_SITE_USAGE}
+  )
+  splice_junctions: junction_data.JunctionMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.SPLICE_JUNCTIONS}
+  )
+  contact_maps: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.CONTACT_MAPS}
+  )
+  procap: track_data.TrackMetadata | None = dataclasses.field(
+      default=None, metadata={'output_type': OutputType.PROCAP}
+  )
+
+  def get(self, output: OutputType) -> track_data.TrackMetadata | None:
+    """Gets the track metadata for a given output type.
+
+    Args:
+      output: The `OutputType` enum value.
+
+    Returns:
+      The corresponding track metadata, or None if it doesn't exist.
+    """
+    match output:
+      case OutputType.ATAC:
+        return self.atac
+      case OutputType.CAGE:
+        return self.cage
+      case OutputType.DNASE:
+        return self.dnase
+      case OutputType.RNA_SEQ:
+        return self.rna_seq
+      case OutputType.CHIP_HISTONE:
+        return self.chip_histone
+      case OutputType.CHIP_TF:
+        return self.chip_tf
+      case OutputType.SPLICE_SITES:
+        return self.splice_sites
+      case OutputType.SPLICE_JUNCTIONS:
+        return self.splice_junctions
+      case OutputType.SPLICE_SITE_USAGE:
+        return self.splice_site_usage
+      case OutputType.CONTACT_MAPS:
+        return self.contact_maps
+      case OutputType.PROCAP:
+        return self.procap
+
+  @classmethod
+  def from_outputs(
+      cls,
+      outputs: Mapping[
+          OutputType, track_data.TrackData | junction_data.JunctionData
+      ],
+  ) -> 'OutputMetadata':
+    """Creates an `OutputMetadata` from a mapping of output types to data.
+
+    Args:
+      outputs: A mapping from `OutputType` to `TrackData` or `JunctionData`.
+
+    Returns:
+      An `OutputMetadata` object.
+    """
+    kwargs = {}
+    for field in dataclasses.fields(cls):
+      output_type = field.metadata['output_type']
+      if data := outputs.get(output_type):
+        kwargs[field.name] = data.metadata
+    return cls(**kwargs)
+
+  def concatenate(self) -> track_data.TrackMetadata:
+    """Concatenates all metadata into a single DataFrame.
+
+    Returns:
+      A pandas DataFrame containing all the track metadata, with an additional
+      column 'output_type' specifying the type of each track.
+    """
+    df_list = []
+    for output_type in OutputType:
+      if (df := self.get(output_type)) is not None:
+        assert isinstance(df, pd.DataFrame)
+        df_list.append(df.assign(output_type=output_type))
+    return pd.concat(df_list)
+
+
+@typing.jaxtyped
+@dataclasses.dataclass(frozen=True)
+class VariantOutput:
+  """Model outputs for a variant prediction.
+
+  Attributes:
+    reference: The model output for the reference sequence.
+    alternate: The model output for the alternate sequence.
+  """
+
+  reference: Output
+  alternate: Output

flax_model/alphagenome/_sdk/models/interval_scorers.py

@@ -0,0 +1,145 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module containing interval scorer dataclasses for interval scoring."""
+
+import dataclasses
+import enum
+from typing import TypeVar
+
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+import immutabledict
+
+
+class IntervalAggregationType(enum.Enum):
+  """Enum indicating the type of interval scorer aggregation.
+
+  Attributes:
+    MEAN: Mean across positions.
+    SUM: Sum across positions.
+
+  Methods:
+    to_proto: Converts the aggregation type to its corresponding proto enum.
+  """
+
+  MEAN = dna_model_pb2.IntervalAggregationType.INTERVAL_AGGREGATION_TYPE_MEAN
+  SUM = dna_model_pb2.IntervalAggregationType.INTERVAL_AGGREGATION_TYPE_SUM
+
+  def to_proto(self) -> dna_model_pb2.IntervalAggregationType:
+    return self.value
+
+  def __repr__(self) -> str:
+    return self.name
+
+
+class BaseIntervalScorer(enum.Enum):
+  """Enum indicating the type of interval scoring.
+
+  Attributes:
+    GENE_MASK: Gene-mask scoring (e.g., aggregation of predictions within genes)
+  """
+
+  GENE_MASK = enum.auto()
+
+
+SUPPORTED_OUTPUT_TYPES = immutabledict.immutabledict({
+    BaseIntervalScorer.GENE_MASK: [
+        dna_output.OutputType.ATAC,
+        dna_output.OutputType.CAGE,
+        dna_output.OutputType.CHIP_HISTONE,
+        dna_output.OutputType.CHIP_TF,
+        dna_output.OutputType.DNASE,
+        dna_output.OutputType.RNA_SEQ,
+        dna_output.OutputType.SPLICE_SITES,
+        dna_output.OutputType.SPLICE_SITE_USAGE,
+    ],
+})
+
+SUPPORTED_WIDTHS = immutabledict.immutabledict({
+    BaseIntervalScorer.GENE_MASK: [None, 501, 2001, 10_001, 100_001, 200_001],
+})
+
+
+@dataclasses.dataclass(frozen=True)
+class GeneMaskScorer:
+  """Interval scorer for gene-mask scoring.
+
+  Attributes:
+    requested_output: The requested output type (e.g. ATAC, DNASE, etc.)
+    width: The width of the target interval to include in the aggregation.
+    aggregation_type: The type of aggregation to perform.
+    base_interval_scorer: The base interval scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+
+  Raises:
+    ValueError: If the requested output is not supported.
+  """
+
+  requested_output: dna_output.OutputType
+  width: int | None
+  aggregation_type: IntervalAggregationType
+
+  @property
+  def base_interval_scorer(self) -> BaseIntervalScorer:
+    return BaseIntervalScorer.GENE_MASK
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  def __post_init__(self):
+    if (
+        self.requested_output
+        not in SUPPORTED_OUTPUT_TYPES[self.base_interval_scorer]
+    ):
+      raise ValueError(
+          f'Unsupported requested output: {self.requested_output}. Supported'
+          f' output types: {SUPPORTED_OUTPUT_TYPES[self.base_interval_scorer]}'
+      )
+    if self.width not in SUPPORTED_WIDTHS[self.base_interval_scorer]:
+      raise ValueError(
+          f'Unsupported width: {self.width}. Supported widths:'
+          f' {SUPPORTED_WIDTHS[self.base_interval_scorer]}'
+      )
+
+  def to_proto(self) -> dna_model_pb2.IntervalScorer:
+    return dna_model_pb2.IntervalScorer(
+        gene_mask=dna_model_pb2.GeneMaskIntervalScorer(
+            requested_output=self.requested_output.to_proto(),
+            width=self.width,
+            aggregation_type=self.aggregation_type.to_proto(),
+        )
+    )
+
+
+# TypeVar for all interval scorer types.
+IntervalScorerTypes = TypeVar(
+    'IntervalScorerTypes',
+    bound=GeneMaskScorer,
+)
+
+# A dict of interval scorers, with our recommended settings for a wide range
+# of different use cases and settings.
+RECOMMENDED_INTERVAL_SCORERS = {
+    'RNA_SEQ': GeneMaskScorer(
+        requested_output=dna_output.OutputType.RNA_SEQ,
+        width=200_001,
+        aggregation_type=IntervalAggregationType.MEAN,
+    ),
+}

flax_model/alphagenome/_sdk/models/junction_data_utils.py

@@ -0,0 +1,216 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utils for converting JunctionData to/from protos."""
+
+from collections.abc import Iterable, Sequence
+
+from flax_model.alphagenome._sdk import tensor_utils
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import junction_data
+from flax_model.alphagenome._sdk.data import ontology
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+from flax_model.alphagenome._sdk.protos import tensor_pb2
+import numpy as np
+import pandas as pd
+
+
+def to_protos(
+    data: junction_data.JunctionData,
+    *,
+    bytes_per_chunk: int = 0,
+    compression_type: tensor_pb2.CompressionType = (
+        tensor_pb2.CompressionType.COMPRESSION_TYPE_NONE
+    ),
+) -> tuple[dna_model_pb2.JunctionData, Sequence[tensor_pb2.TensorChunk]]:
+  """Converts the `JunctionData` to protobuf messages.
+
+  Args:
+    data: The `JunctionData` object to convert to protos.
+    bytes_per_chunk: The maximum number of bytes per tensor chunk.
+    compression_type: The compression type to use for the tensor chunks.
+
+  Returns:
+    A tuple containing the `JunctionData` protobuf message and a sequence of
+      `TensorChunk` protobuf messages.
+  """
+  tensor, chunks = tensor_utils.pack_tensor(
+      data.values,
+      bytes_per_chunk=bytes_per_chunk,
+      compression_type=compression_type,
+  )
+
+  return (
+      dna_model_pb2.JunctionData(
+          junctions=[j.to_proto() for j in data.junctions],
+          values=tensor,
+          metadata=metadata_to_proto(data.metadata).metadata,
+          interval=data.interval.to_proto() if data.interval else None,
+      ),
+      chunks,
+  )
+
+
+def from_protos(
+    proto: dna_model_pb2.JunctionData,
+    chunks: Iterable[tensor_pb2.TensorChunk] = (),
+    *,
+    interval: genome.Interval | None = None,
+) -> junction_data.JunctionData:
+  """Converts a `JunctionData` protobuf to a `JunctionData` object.
+
+  Args:
+    proto: A `JunctionData` protobuf message.
+    chunks: A sequence of `TensorChunk` protobuf messages.
+    interval: Optional `Interval` object representing the genomic region
+      containing the junctions. Only used if the proto does not have an
+      interval.
+
+  Returns:
+    A `JunctionData` object.
+  """
+  values = tensor_utils.unpack_proto(proto.values, chunks)
+  values = tensor_utils.upcast_floating(values)
+
+  metadata = metadata_from_proto(
+      dna_model_pb2.JunctionsMetadata(metadata=proto.metadata)
+  )
+
+  if proto.HasField('interval'):
+    interval = genome.Interval.from_proto(proto.interval)
+
+  return junction_data.JunctionData(
+      junctions=np.array(
+          [genome.Interval.from_proto(j) for j in proto.junctions]
+      ),
+      values=values,
+      metadata=metadata,
+      interval=interval,
+  )
+
+
+def metadata_to_proto(
+    metadata: junction_data.JunctionMetadata,
+) -> dna_model_pb2.JunctionsMetadata:
+  """Converts junction metadata to a JunctionsMetadata.
+
+  Args:
+    metadata: A pandas DataFrame containing junction metadata.
+
+  Returns:
+    A `JunctionsMetadata` protobuf message.
+  """
+  names = metadata['name']
+  default_values = [None] * len(names)
+
+  columns = zip(
+      metadata['name'],
+      metadata.get('ontology_curie', default_values),
+      metadata.get('biosample_type', default_values),
+      metadata.get('biosample_name', default_values),
+      metadata.get('biosample_life_stage', default_values),
+      metadata.get('gtex_tissue', default_values),
+      metadata.get('data_source', default_values),
+      metadata.get('Assay title', default_values),
+      strict=True,
+  )
+
+  metadata_protos = []
+
+  for (
+      name,
+      ontology_curie,
+      biosample_type,
+      biosample_name,
+      biosample_life_stage,
+      gtex_tissue,
+      data_source,
+      assay,
+  ) in columns:
+    if biosample_type is not None:
+      biosample_proto = dna_model_pb2.Biosample(
+          name=biosample_name,
+          type=dna_model_pb2.BiosampleType.Value(
+              f'BIOSAMPLE_TYPE_{biosample_type.upper()}'
+          ),
+          stage=biosample_life_stage,
+      )
+    else:
+      biosample_proto = None
+
+    metadata_protos.append(
+        dna_model_pb2.JunctionMetadata(
+            name=name,
+            ontology_term=ontology.from_curie(ontology_curie).to_proto()
+            if ontology_curie
+            else None,
+            biosample=biosample_proto,
+            gtex_tissue=gtex_tissue,
+            data_source=data_source,
+            assay=assay,
+        )
+    )
+
+  return dna_model_pb2.JunctionsMetadata(metadata=metadata_protos)
+
+
+def metadata_from_proto(
+    proto: dna_model_pb2.JunctionsMetadata,
+) -> junction_data.JunctionMetadata:
+  """Create JunctionMetadata from a dna_model_pb2.JunctionsMetadata.
+
+  Args:
+    proto: A `JunctionsMetadata` protobuf message.
+
+  Returns:
+    A pandas DataFrame containing junction metadata.
+  """
+  metadata = []
+  for junction_proto in proto.metadata:
+    junction_metadata = {
+        'name': junction_proto.name,
+    }
+
+    if junction_proto.HasField('ontology_term'):
+      junction_metadata['ontology_curie'] = ontology.from_proto(
+          junction_proto.ontology_term
+      ).ontology_curie
+
+    if junction_proto.HasField('biosample'):
+      junction_metadata['biosample_name'] = junction_proto.biosample.name
+      junction_metadata['biosample_type'] = (
+          dna_model_pb2.BiosampleType.Name(junction_proto.biosample.type)
+          .removeprefix('BIOSAMPLE_TYPE_')
+          .lower()
+      )
+      if junction_proto.biosample.HasField('stage'):
+        junction_metadata['biosample_life_stage'] = (
+            junction_proto.biosample.stage
+        )
+
+    if junction_proto.HasField('gtex_tissue'):
+      junction_metadata['gtex_tissue'] = junction_proto.gtex_tissue
+
+    if junction_proto.HasField('data_source'):
+      junction_metadata['data_source'] = junction_proto.data_source
+
+    if junction_proto.HasField('assay'):
+      junction_metadata['Assay title'] = junction_proto.assay
+
+    metadata.append(junction_metadata)
+
+  if metadata:
+    return pd.DataFrame(metadata)
+  else:
+    return pd.DataFrame(columns=['name'])

flax_model/alphagenome/_sdk/models/track_data_utils.py

@@ -0,0 +1,260 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Track data utils for converting to/from protos."""
+
+from collections.abc import Iterable, Sequence
+
+from flax_model.alphagenome._sdk import tensor_utils
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import ontology
+from flax_model.alphagenome._sdk.data import track_data
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+from flax_model.alphagenome._sdk.protos import tensor_pb2
+import pandas as pd
+
+
+def to_protos(
+    data: track_data.TrackData,
+    *,
+    bytes_per_chunk: int = 0,
+    compression_type: tensor_pb2.CompressionType = (
+        tensor_pb2.CompressionType.COMPRESSION_TYPE_NONE
+    ),
+) -> tuple[dna_model_pb2.TrackData, Sequence[tensor_pb2.TensorChunk]]:
+  """Serializes `TrackData` to protobuf messages.
+
+  Args:
+    data: The `TrackData` object to serialize.
+    bytes_per_chunk: The maximum number of bytes per tensor chunk.
+    compression_type: The compression type to use for the tensor chunks.
+
+  Returns:
+    A tuple containing the `TrackData` protobuf message and a sequence of
+    `TensorChunk` protobuf messages.
+  """
+  tensor, chunks = tensor_utils.pack_tensor(
+      data.values,
+      bytes_per_chunk=bytes_per_chunk,
+      compression_type=compression_type,
+  )
+  return (
+      dna_model_pb2.TrackData(
+          values=tensor,
+          metadata=metadata_to_proto(data.metadata).metadata,
+          resolution=data.resolution,
+          interval=data.interval.to_proto() if data.interval else None,
+      ),
+      chunks,
+  )
+
+
+def from_protos(
+    proto: dna_model_pb2.TrackData,
+    chunks: Iterable[tensor_pb2.TensorChunk] = (),
+    *,
+    interval: genome.Interval | None = None,
+) -> track_data.TrackData:
+  """Creates a `TrackData` object from protobuf messages.
+
+  Args:
+    proto: A `TrackData` protobuf message.
+    chunks: A sequence of `TensorChunk` protobuf messages.
+    interval: Optional `Interval` object representing the genomic region
+      containing the tracks. Only used if the proto does not have an interval.
+
+  Returns:
+    A `TrackData` object.
+  """
+  metadata = metadata_from_proto(
+      dna_model_pb2.TracksMetadata(metadata=proto.metadata)
+  )
+
+  values = tensor_utils.unpack_proto(proto.values, chunks)
+  values = tensor_utils.upcast_floating(values)
+  resolution = proto.resolution if proto.HasField('resolution') else 1
+  if proto.HasField('interval'):
+    interval = genome.Interval.from_proto(proto.interval)
+
+  return track_data.TrackData(
+      values, metadata, resolution=resolution, interval=interval
+  )
+
+
+def metadata_to_proto(
+    metadata: track_data.TrackMetadata,
+) -> dna_model_pb2.TracksMetadata:
+  """Converts track metadata to a `TracksMetadata` protobuf message.
+
+  Args:
+    metadata: A pandas DataFrame containing track metadata, with the following
+      required columns: name, strand.
+
+  Returns:
+    A `TracksMetadata` protobuf message.
+  """
+  names = metadata['name']
+  default_values = [None] * len(names)
+
+  columns = zip(
+      metadata['name'],
+      metadata['strand'],
+      metadata.get('ontology_curie', default_values),
+      metadata.get('biosample_type', default_values),
+      metadata.get('biosample_name', default_values),
+      metadata.get('biosample_life_stage', default_values),
+      metadata.get('transcription_factor', default_values),
+      metadata.get('histone_mark', default_values),
+      metadata.get('gtex_tissue', default_values),
+      metadata.get('Assay title', default_values),
+      metadata.get('data_source', default_values),
+      metadata.get('genetically_modified', default_values),
+      metadata.get('endedness', default_values),
+      metadata.get('nonzero_mean', default_values),
+      strict=True,
+  )
+
+  metadata_protos = []
+  for (
+      name,
+      strand,
+      ontology_curie,
+      biosample_type,
+      biosample_name,
+      biosample_life_stage,
+      transcription_factor,
+      histone_mark,
+      gtex_tissue,
+      assay,
+      data_source,
+      genetically_modified,
+      endedness,
+      nonzero_mean,
+  ) in columns:
+    if biosample_type:
+      biosample = dna_model_pb2.Biosample(
+          name=biosample_name,
+          type=dna_model_pb2.BiosampleType.Value(
+              f'BIOSAMPLE_TYPE_{biosample_type.upper()}'
+          ),
+          stage=biosample_life_stage,
+      )
+    else:
+      biosample = None
+    if endedness is not None:
+      match endedness:
+        case 'paired':
+          endedness = dna_model_pb2.Endedness.ENDEDNESS_PAIRED
+        case 'single':
+          endedness = dna_model_pb2.Endedness.ENDEDNESS_SINGLE
+        case _:
+          raise ValueError(f'Unknown endedness: {endedness}')
+
+    metadata_protos.append(
+        dna_model_pb2.TrackMetadata(
+            name=name,
+            strand=genome.Strand.from_str(strand).to_proto()
+            if strand
+            else None,
+            ontology_term=ontology.from_curie(ontology_curie).to_proto()
+            if ontology_curie
+            else None,
+            biosample=biosample,
+            transcription_factor_code=transcription_factor
+            if isinstance(transcription_factor, str)
+            else None,
+            histone_mark_code=histone_mark
+            if isinstance(histone_mark, str)
+            else None,
+            gtex_tissue=gtex_tissue,
+            assay=assay,
+            data_source=data_source,
+            genetically_modified=genetically_modified,
+            endedness=endedness,
+            nonzero_mean=nonzero_mean,
+        )
+    )
+
+  return dna_model_pb2.TracksMetadata(metadata=metadata_protos)
+
+
+def metadata_from_proto(
+    proto: dna_model_pb2.TracksMetadata,
+) -> track_data.TrackMetadata:
+  """Creates track metadata from a `TracksMetadata` protobuf message.
+
+  Args:
+    proto: A `TracksMetadata` protobuf message.
+
+  Returns:
+    A pandas DataFrame containing track metadata.
+  """
+  metadata = []
+  for track_proto in proto.metadata:
+    track_metadata = {
+        'name': track_proto.name,
+        'strand': str(genome.Strand.from_proto(track_proto.strand)),
+    }
+
+    if track_proto.HasField('assay'):
+      track_metadata['Assay title'] = track_proto.assay
+
+    if track_proto.HasField('ontology_term'):
+      track_metadata['ontology_curie'] = ontology.from_proto(
+          track_proto.ontology_term
+      ).ontology_curie
+
+    if track_proto.HasField('biosample'):
+      track_metadata['biosample_name'] = track_proto.biosample.name
+      track_metadata['biosample_type'] = (
+          dna_model_pb2.BiosampleType.Name(track_proto.biosample.type)
+          .removeprefix('BIOSAMPLE_TYPE_')
+          .lower()
+      )
+      if track_proto.biosample.HasField('stage'):
+        track_metadata['biosample_life_stage'] = track_proto.biosample.stage
+
+    if track_proto.HasField('transcription_factor_code'):
+      track_metadata['transcription_factor'] = (
+          track_proto.transcription_factor_code
+      )
+
+    if track_proto.HasField('histone_mark_code'):
+      track_metadata['histone_mark'] = track_proto.histone_mark_code
+
+    if track_proto.HasField('gtex_tissue'):
+      track_metadata['gtex_tissue'] = track_proto.gtex_tissue
+
+    if track_proto.HasField('data_source'):
+      track_metadata['data_source'] = track_proto.data_source
+
+    if track_proto.HasField('endedness'):
+      track_metadata['endedness'] = (
+          dna_model_pb2.Endedness.Name(track_proto.endedness)
+          .removeprefix('ENDEDNESS_')
+          .lower()
+      )
+
+    if track_proto.HasField('genetically_modified'):
+      track_metadata['genetically_modified'] = track_proto.genetically_modified
+
+    if track_proto.HasField('nonzero_mean'):
+      track_metadata['nonzero_mean'] = track_proto.nonzero_mean
+
+    metadata.append(track_metadata)
+  if metadata:
+    return pd.DataFrame(metadata)
+  else:
+    return pd.DataFrame(columns=['name', 'strand'])

flax_model/alphagenome/_sdk/models/variant_scorers.py

@@ -0,0 +1,878 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module containing variant scorer dataclasses for variant scoring."""
+
+from collections.abc import Sequence
+import dataclasses
+import enum
+import itertools
+import math
+from typing import TypeVar
+
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome._sdk.protos import dna_model_pb2
+import anndata
+import immutabledict
+import pandas as pd
+
+
+class AggregationType(enum.Enum):
+  """Enum indicating the type of variant scorer aggregation.
+
+  Attributes:
+    DIFF_MEAN: Difference of means, i.e., mean(`ALT`) - mean(`REF`).
+    DIFF_SUM: Difference of sums, i.e., sum(`ALT`) - sum(`REF`).
+    DIFF_SUM_LOG2: Log scales predictions, then takes the sum and then the
+      difference, i.e., sum(log2(`ALT`)) - sum(log2(`REF`)).
+    DIFF_LOG2_SUM: Takes the sum of predictions, applies a log transform, then
+      takes the difference between predictions, i.e., log2(sum(`ALT`)) -
+      log2(sum(`REF`)).
+    L2_DIFF: Takes the difference of `ALT` and `REF` predictions, then computes
+      the L2 norm, i.e., l2_norm(`ALT` - `REF`).
+    L2_DIFF_LOG1P: Log scales the predictions + 1, takes the difference of `ALT`
+      and `REF` predictions, then computes the L2 norm, i.e.,
+      l2_norm(log1p(`ALT`) - log1p(`REF`)).
+    ACTIVE_MEAN: Maximum of means, i.e., max(mean(`ALT`), mean(`REF`)).
+    ACTIVE_SUM: Maximum of sums, i.e., max(sum(`ALT), sum(`REF`)).
+
+  Methods:
+    to_proto: Converts the aggregation type to its corresponding proto enum.
+  """
+
+  DIFF_MEAN = dna_model_pb2.AggregationType.AGGREGATION_TYPE_DIFF_MEAN
+  DIFF_SUM = dna_model_pb2.AggregationType.AGGREGATION_TYPE_DIFF_SUM
+  DIFF_SUM_LOG2 = dna_model_pb2.AggregationType.AGGREGATION_TYPE_DIFF_SUM_LOG2
+  DIFF_LOG2_SUM = dna_model_pb2.AggregationType.AGGREGATION_TYPE_DIFF_LOG2_SUM
+  L2_DIFF = dna_model_pb2.AggregationType.AGGREGATION_TYPE_L2_DIFF
+  L2_DIFF_LOG1P = dna_model_pb2.AggregationType.AGGREGATION_TYPE_L2_DIFF_LOG1P
+  ACTIVE_MEAN = dna_model_pb2.AggregationType.AGGREGATION_TYPE_ACTIVE_MEAN
+  ACTIVE_SUM = dna_model_pb2.AggregationType.AGGREGATION_TYPE_ACTIVE_SUM
+
+  def to_proto(self) -> dna_model_pb2.AggregationType:
+    return self.value
+
+  def __repr__(self) -> str:
+    return self.name
+
+
+class BaseVariantScorer(enum.Enum):
+  """Enum indicating the type of variant scoring.
+
+  Attributes:
+    CENTER_MASK: Center mask scorer.
+    CONTACT_MAP: Contact map scorer, a center mask scorer that handles 2D
+      tensors.
+    GENE_MASK_LFC: Gene-mask scoring based on log fold change.
+    GENE_MASK_ACTIVE: Gene-mask scoring based on active allele.
+    GENE_MASK_SPLICING: Gene-mask scoring for splicing.
+    PA_QTL: Polyadenylation QTL scoring.
+    SPLICE_JUNCTION: Splice junction scoring.
+  """
+
+  CENTER_MASK = enum.auto()
+  CONTACT_MAP = enum.auto()
+  GENE_MASK_LFC = enum.auto()
+  GENE_MASK_ACTIVE = enum.auto()
+  GENE_MASK_SPLICING = enum.auto()
+  PA_QTL = enum.auto()
+  SPLICE_JUNCTION = enum.auto()
+
+
+SUPPORTED_ORGANISMS = immutabledict.immutabledict({
+    BaseVariantScorer.CENTER_MASK: dna_model_pb2.Organism.values(),
+    BaseVariantScorer.CONTACT_MAP: dna_model_pb2.Organism.values(),
+    BaseVariantScorer.GENE_MASK_LFC: dna_model_pb2.Organism.values(),
+    BaseVariantScorer.GENE_MASK_ACTIVE: dna_model_pb2.Organism.values(),
+    BaseVariantScorer.GENE_MASK_SPLICING: dna_model_pb2.Organism.values(),
+    BaseVariantScorer.PA_QTL: [dna_model_pb2.Organism.ORGANISM_HOMO_SAPIENS],
+    BaseVariantScorer.SPLICE_JUNCTION: dna_model_pb2.Organism.values(),
+})
+
+SUPPORTED_OUTPUT_TYPES = immutabledict.immutabledict({
+    BaseVariantScorer.CENTER_MASK: [
+        dna_output.OutputType.ATAC,
+        dna_output.OutputType.CAGE,
+        dna_output.OutputType.DNASE,
+        dna_output.OutputType.PROCAP,
+        dna_output.OutputType.RNA_SEQ,
+        dna_output.OutputType.CHIP_HISTONE,
+        dna_output.OutputType.CHIP_TF,
+        dna_output.OutputType.SPLICE_SITES,
+        dna_output.OutputType.SPLICE_SITE_USAGE,
+    ],
+    BaseVariantScorer.CONTACT_MAP: [dna_output.OutputType.CONTACT_MAPS],
+    BaseVariantScorer.GENE_MASK_LFC: [
+        dna_output.OutputType.ATAC,
+        dna_output.OutputType.CAGE,
+        dna_output.OutputType.DNASE,
+        dna_output.OutputType.PROCAP,
+        dna_output.OutputType.RNA_SEQ,
+        dna_output.OutputType.SPLICE_SITES,
+        dna_output.OutputType.SPLICE_SITE_USAGE,
+    ],
+    BaseVariantScorer.GENE_MASK_ACTIVE: [
+        dna_output.OutputType.ATAC,
+        dna_output.OutputType.CAGE,
+        dna_output.OutputType.DNASE,
+        dna_output.OutputType.PROCAP,
+        dna_output.OutputType.RNA_SEQ,
+        dna_output.OutputType.SPLICE_SITES,
+        dna_output.OutputType.SPLICE_SITE_USAGE,
+    ],
+    BaseVariantScorer.GENE_MASK_SPLICING: [
+        dna_output.OutputType.SPLICE_SITES,
+        dna_output.OutputType.SPLICE_SITE_USAGE,
+    ],
+})
+
+SUPPORTED_WIDTHS = immutabledict.immutabledict({
+    BaseVariantScorer.CENTER_MASK: [None, 501, 2001, 10_001, 100_001, 200_001],
+    BaseVariantScorer.GENE_MASK_SPLICING: [None, 101, 1_001, 10_001],
+})
+
+SUPPORTED_AGGREGATIONS = immutabledict.immutabledict({
+    BaseVariantScorer.CENTER_MASK: [
+        AggregationType.DIFF_MEAN,
+        AggregationType.DIFF_SUM,
+        AggregationType.DIFF_SUM_LOG2,
+        AggregationType.DIFF_LOG2_SUM,
+        AggregationType.L2_DIFF,
+        AggregationType.L2_DIFF_LOG1P,
+        AggregationType.ACTIVE_MEAN,
+        AggregationType.ACTIVE_SUM,
+    ],
+})
+
+
+@dataclasses.dataclass(frozen=True)
+class CenterMaskScorer:
+  """Variant scorer for center mask scorers.
+
+  Variant scorer that aggregates ALT and REF values using a spatial mask
+  centered around the variant before computing the difference. This scorer
+  aggregates over the spatial axis and returns one score per output track.
+
+  Attributes:
+    requested_output: The requested output type (e.g., ATAC, DNASE, etc.)
+    width: The width of the mask around the variant. If None, the score is
+      computed over the entire sequence.
+    aggregation_type: The aggregation type.
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+
+  Raises:
+    ValueError: If the requested output, width, or aggregation type is not
+      supported.
+  """
+
+  requested_output: dna_output.OutputType
+  width: int | None
+  aggregation_type: AggregationType
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.CENTER_MASK
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    # The 'active' scorers track the maximum of the ALT and REF, and are
+    # non-directional. All other scorers track ALT - REF, with L2_DIFF
+    # nondirectional and the remaining directional.
+    return self.aggregation_type in [
+        AggregationType.DIFF_MEAN,
+        AggregationType.DIFF_SUM,
+        AggregationType.DIFF_SUM_LOG2,
+        AggregationType.DIFF_LOG2_SUM,
+    ]
+
+  def __post_init__(self):
+    if (
+        self.requested_output
+        not in SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]
+    ):
+      raise ValueError(
+          f'Unsupported requested output: {self.requested_output}. Supported'
+          f' output types: {SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]}'
+      )
+    if self.width not in SUPPORTED_WIDTHS[self.base_variant_scorer]:
+      raise ValueError(
+          f'Unsupported width: {self.width}. Supported widths:'
+          f' {SUPPORTED_WIDTHS[self.base_variant_scorer]}'
+      )
+    if (
+        self.aggregation_type
+        not in SUPPORTED_AGGREGATIONS[self.base_variant_scorer]
+    ):
+      raise ValueError(
+          f'Unsupported aggregation type: {self.aggregation_type}. Supported'
+          ' aggregation types:'
+          f' {SUPPORTED_AGGREGATIONS[self.base_variant_scorer]}'
+      )
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        center_mask=dna_model_pb2.CenterMaskScorer(
+            requested_output=self.requested_output.to_proto(),
+            width=self.width,
+            aggregation_type=self.aggregation_type.to_proto(),
+        )
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class ContactMapScorer:
+  """Variant scorer for contact map scorers.
+
+  Variant scorer that quantifies local contact disruption between ALT and REF
+  alleles. Uses a 1MB window centered around the variant to calculate the
+  mean absolute difference of contact frequencies, for all interactions
+  involving the variant-containing bin.
+
+  Attributes:
+    requested_output: The requested output type (e.g., CONTACT_MAPS).
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+  """
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.CONTACT_MAP
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    return False
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        contact_map=dna_model_pb2.ContactMapScorer()
+    )
+
+  @property
+  def requested_output(self) -> dna_output.OutputType:
+    return dna_output.OutputType.CONTACT_MAPS
+
+
+@dataclasses.dataclass(frozen=True)
+class GeneMaskLFCScorer:
+  """Variant scorer for gene-mask log fold change scoring.
+
+  Variant scorer that quantifies impact on overall gene transcript abundance.
+  Calculates the log fold change of gene expression level between ALT and REF
+  alleles using a gene exon mask.
+
+  Attributes:
+    requested_output: The requested output type (e.g. ATAC, DNASE, etc.)
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+
+  Raises:
+    ValueError: If the requested output is not supported.
+  """
+
+  requested_output: dna_output.OutputType
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.GENE_MASK_LFC
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    return True
+
+  def __post_init__(self):
+    if (
+        self.requested_output
+        not in SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]
+    ):
+      raise ValueError(
+          f'Unsupported requested output: {self.requested_output}. Supported'
+          f' output types: {SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]}'
+      )
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        gene_mask=dna_model_pb2.GeneMaskLFCScorer(
+            requested_output=self.requested_output.to_proto(),
+        )
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class GeneMaskActiveScorer:
+  """Variant scorer for gene-mask active allele scoring.
+
+  Variant scorer that captures absolute activity level associated with one of
+  the alleles, rather than quantifying the difference between the ALT and REF.
+  Active allele gene scores are calculated by taking the maximum of the
+  aggregated ALT and REF signals across exons for a gene of interest.
+
+  Attributes:
+    requested_output: The requested output type (e.g. ATAC, DNASE, etc.)
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+
+  Raises:
+    ValueError: If the requested output is not supported.
+  """
+
+  requested_output: dna_output.OutputType
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.GENE_MASK_ACTIVE
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    return False
+
+  def __post_init__(self):
+    if (
+        self.requested_output
+        not in SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]
+    ):
+      raise ValueError(
+          f'Unsupported requested output: {self.requested_output}. Supported'
+          f' output types: {SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]}'
+      )
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        gene_mask_active=dna_model_pb2.GeneMaskActiveScorer(
+            requested_output=self.requested_output.to_proto(),
+        )
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class GeneMaskSplicingScorer:
+  """Variant scorer for gene-mask scoring for splicing.
+
+  Variant scorer that quantifies changes in class assignment probabilties
+  (dna_output.OutputType.SPLICE_SITES) or changes in the usage of splice sites
+  (dna_output.OutputType.SPLICE_SITE_USAGE) between ALT and REF alleles using
+  a gene exon mask.
+
+  Attributes:
+    requested_output: The requested output type (e.g. ATAC, DNASE, etc.)
+    width: The width of the mask around the variant.
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+
+  Raises:
+    ValueError: If the requested output or width is not supported.
+  """
+
+  requested_output: dna_output.OutputType
+  width: int | None
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.GENE_MASK_SPLICING
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    return False
+
+  def __post_init__(self):
+    if (
+        self.requested_output
+        not in SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]
+    ):
+      raise ValueError(
+          f'Unsupported requested output: {self.requested_output}. Supported'
+          f' output types: {SUPPORTED_OUTPUT_TYPES[self.base_variant_scorer]}'
+      )
+    if self.width not in SUPPORTED_WIDTHS[self.base_variant_scorer]:
+      raise ValueError(
+          f'Unsupported width: {self.width}. Supported widths:'
+          f' {SUPPORTED_WIDTHS[self.base_variant_scorer]}'
+      )
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        gene_mask_splicing=dna_model_pb2.GeneMaskSplicingScorer(
+            requested_output=self.requested_output.to_proto(),
+            width=self.width,
+        )
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class PolyadenylationScorer:
+  """Variant scorer for polyadenylation QTLs (paQTLs).
+
+  Variant scorer for polyadenylation quantitative trait loci (paQTLs) to capture
+  a variant's impact on RNA isoform production.
+
+  The scoring approach quantifies the maximum log fold change in expression
+  between the set of proximal polyadenylation sites (PAS)s vs. the set of
+  distal PASs, where the sets of proximal and distal PASs are the PASs upstream
+  or downstream of any given 3' cleavage site respectively.
+
+  Attributes:
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+  """
+
+  @property
+  def requested_output(self) -> dna_output.OutputType:
+    """Returns the scorer's underlying output type."""
+    return dna_output.OutputType.RNA_SEQ
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.PA_QTL
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    return False
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        pa_qtl=dna_model_pb2.PolyadenylationScorer()
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class SpliceJunctionScorer:
+  """Variant scorer for splice junction scoring.
+
+  Attributes:
+    base_variant_scorer: The base variant scorer.
+    name: The name of the scorer (a composite of the above attributes that
+      uniquely identifies a scorer combination).
+    is_signed: Whether this variant scorer is directional (such that scores may
+      be negative or positive) or non-directional (scores are always positive).
+
+  Methods:
+    to_proto: Converts the scorer to its corresponding proto message.
+  """
+
+  @property
+  def requested_output(self) -> dna_output.OutputType:
+    """Returns the scorer's underlying output type."""
+    return dna_output.OutputType.SPLICE_JUNCTIONS
+
+  @property
+  def base_variant_scorer(self) -> BaseVariantScorer:
+    return BaseVariantScorer.SPLICE_JUNCTION
+
+  @property
+  def name(self) -> str:
+    return str(self)
+
+  @property
+  def is_signed(self) -> bool:
+    return False
+
+  def to_proto(self) -> dna_model_pb2.VariantScorer:
+    return dna_model_pb2.VariantScorer(
+        splice_junction=dna_model_pb2.SpliceJunctionScorer()
+    )
+
+
+# TypeVar for all variant scorer types.
+VariantScorerTypes = TypeVar(
+    'VariantScorerTypes',
+    CenterMaskScorer,
+    ContactMapScorer,
+    GeneMaskLFCScorer,
+    GeneMaskActiveScorer,
+    GeneMaskSplicingScorer,
+    PolyadenylationScorer,
+    SpliceJunctionScorer,
+)
+
+# A dict of variant scorers, with our recommended settings for a wide range
+# of different use cases and settings.
+RECOMMENDED_VARIANT_SCORERS = immutabledict.immutabledict({
+    'ATAC': CenterMaskScorer(
+        requested_output=dna_output.OutputType.ATAC,
+        width=501,
+        aggregation_type=AggregationType.DIFF_LOG2_SUM,
+    ),
+    'CONTACT_MAPS': ContactMapScorer(),
+    'DNASE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.DNASE,
+        width=501,
+        aggregation_type=AggregationType.DIFF_LOG2_SUM,
+    ),
+    'CHIP_TF': CenterMaskScorer(
+        requested_output=dna_output.OutputType.CHIP_TF,
+        width=501,
+        aggregation_type=AggregationType.DIFF_LOG2_SUM,
+    ),
+    'CHIP_HISTONE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.CHIP_HISTONE,
+        width=2001,
+        aggregation_type=AggregationType.DIFF_LOG2_SUM,
+    ),
+    'CAGE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.CAGE,
+        width=501,
+        aggregation_type=AggregationType.DIFF_LOG2_SUM,
+    ),
+    'PROCAP': CenterMaskScorer(
+        requested_output=dna_output.OutputType.PROCAP,
+        width=501,
+        aggregation_type=AggregationType.DIFF_LOG2_SUM,
+    ),
+    # Expression log fold change.
+    'RNA_SEQ': GeneMaskLFCScorer(
+        requested_output=dna_output.OutputType.RNA_SEQ
+    ),
+    'RNA_SEQ_ACTIVE': GeneMaskActiveScorer(
+        requested_output=dna_output.OutputType.RNA_SEQ
+    ),
+    'SPLICE_SITES': GeneMaskSplicingScorer(
+        requested_output=dna_output.OutputType.SPLICE_SITES,
+        width=None,
+    ),
+    'SPLICE_SITE_USAGE': GeneMaskSplicingScorer(
+        requested_output=dna_output.OutputType.SPLICE_SITE_USAGE,
+        width=None,
+    ),
+    'SPLICE_JUNCTIONS': SpliceJunctionScorer(),
+    'POLYADENYLATION': PolyadenylationScorer(),
+    'ATAC_ACTIVE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.ATAC,
+        width=501,
+        aggregation_type=AggregationType.ACTIVE_SUM,
+    ),
+    'DNASE_ACTIVE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.DNASE,
+        width=501,
+        aggregation_type=AggregationType.ACTIVE_SUM,
+    ),
+    'CHIP_TF_ACTIVE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.CHIP_TF,
+        width=501,
+        aggregation_type=AggregationType.ACTIVE_SUM,
+    ),
+    'CHIP_HISTONE_ACTIVE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.CHIP_HISTONE,
+        width=2001,
+        aggregation_type=AggregationType.ACTIVE_SUM,
+    ),
+    'CAGE_ACTIVE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.CAGE,
+        width=501,
+        aggregation_type=AggregationType.ACTIVE_SUM,
+    ),
+    'PROCAP_ACTIVE': CenterMaskScorer(
+        requested_output=dna_output.OutputType.PROCAP,
+        width=501,
+        aggregation_type=AggregationType.ACTIVE_SUM,
+    ),
+})
+
+
+def get_recommended_scorers(
+    organism: dna_model_pb2.Organism,
+) -> list[VariantScorerTypes]:
+  """Returns the recommended variant scorers for a given organism."""
+  return [
+      scorer
+      for scorer in RECOMMENDED_VARIANT_SCORERS.values()
+      if organism in SUPPORTED_ORGANISMS[scorer.base_variant_scorer]
+  ]
+
+
+def tidy_anndata(
+    adata: anndata.AnnData,
+    match_gene_strand: bool = True,
+    include_extended_metadata: bool = True,
+) -> pd.DataFrame:
+  """Formats an :class:`~anndata.AnnData` score as a tidy DataFrame.
+
+  This function converts the score output from an AnnData object into a
+  long-format pandas DataFrame, where each row represents:
+
+    - For non-gene-centric variant scorers: Score for a variant-track pair.
+    - For non-gene-centric interval scorers: Score for an interval-track pair.
+    - For gene-centric variant/interval scoring: Score for a
+      variant/interval-gene-track combination.
+
+  Args:
+    adata: An AnnData object containing scores.
+    match_gene_strand: If True (and using gene-centric scoring), rows with
+      mismatched gene and track strands are removed.
+    include_extended_metadata: If True, includes additional columns derived from
+      metadata specific to the output type, such as biosample name and type,
+      gtex tissue, transcription factor, and histone mark, if available. If
+      False, only includes minimal metadata columns required to unique identify
+      a track withing a given output type: track_name and track_strand.
+
+  Returns:
+    A pandas DataFrame with one score per row. The DataFrame includes
+    columns for variant ID (if applicable), scored interval, gene information
+    (if applicable), output type, variant/interval scorer, track name,
+    ontology term, assay type, track strand, and raw score. Additional metadata
+    such as biosample name and type, gtex tissue are also returned (where
+    available). See :func:`full_path_to.tidy_scores` for more details on the
+    returned columns.
+
+  Raises:
+    ValueError: If the input is not an AnnData object.
+  """
+  if not isinstance(adata, anndata.AnnData):
+    raise ValueError('Invalid input type. Must be an AnnData object.')
+
+  # Columns to include from the gene metadata. If the column did not
+  # exist in the original metadata (or if the scores do not have a concept
+  # of a gene), a column of all Nones is added.
+  gene_columns = [
+      'gene_id',
+      'gene_name',
+      'gene_type',
+      'gene_strand',
+      'junction_Start',
+      'junction_End',
+  ]
+  if math.prod(adata.X.shape) == 0:
+    # Scores are empty, so we return an empty dataframe.
+    return pd.DataFrame()
+  elif 'gene_id' in adata.obs and 'strand' in adata.obs:
+    # Scores are for a gene-based scorer.
+    obs = adata.obs.rename({'strand': 'gene_strand'}, axis=1)
+
+    # Remove patch number from gene_id.
+    obs['gene_id'] = obs['gene_id'].str.split('.', expand=True).get(0)
+    for col in gene_columns:
+      if col not in obs:
+        obs[col] = None
+  elif adata.X.shape[0] == 1 and adata.obs.empty:
+    # Scores are for a non-gene-based scorer.
+    obs = pd.DataFrame([[None] * len(gene_columns)], columns=gene_columns)
+  else:
+    raise ValueError('Scores contain gene metadata but no gene_id/strand.')
+
+  # Columns to include from the track metadata. Only added if the column
+  # appears in the original metadata.
+  var = adata.var.rename(
+      {'strand': 'track_strand', 'name': 'track_name'}, axis=1
+  )
+  track_columns = ['track_name', 'track_strand']
+  if include_extended_metadata:
+    extended_metadata_columns = [
+        'Assay title',
+        'ontology_curie',
+        'biosample_name',
+        'biosample_type',
+        'biosample_life_stage',
+        'gtex_tissue',
+        'data_source',
+        'endedness',
+        'genetically_modified',
+        'transcription_factor',
+        'histone_mark',
+    ]
+    track_columns.extend([c for c in extended_metadata_columns if c in var])
+
+  # Construct score dataframe.
+  df = pd.merge(var, obs, how='cross')
+  df['scored_interval'] = adata.uns['interval']
+
+  # Add id and scorer information depending on the type of score.
+  if 'variant_scorer' in adata.uns:
+    df['variant_id'] = adata.uns['variant']
+    df['output_type'] = adata.uns['variant_scorer'].requested_output.name
+    df['variant_scorer'] = str(adata.uns['variant_scorer'])
+    id_columns = ['variant_id', 'scored_interval']
+    scorer_columns = ['output_type', 'variant_scorer']
+  elif 'interval_scorer' in adata.uns:
+    df['output_type'] = adata.uns['interval_scorer'].requested_output.name
+    df['interval_scorer'] = str(adata.uns['interval_scorer'])
+    id_columns = ['scored_interval']
+    scorer_columns = ['output_type', 'interval_scorer']
+  else:
+    raise ValueError(
+        'Invalid scores, expected either variant_scorer or interval_scorer.'
+    )
+
+  # Formatting final touches.
+  df = df.loc[:, id_columns + gene_columns + scorer_columns + track_columns]
+  df['raw_score'] = adata.X.T.reshape((-1,))
+  if 'quantiles' in adata.layers:
+    df['quantile_score'] = adata.layers['quantiles'].T.reshape((-1,))
+
+  # Remove entries where DNA strands are mismatched. Note that we still retain
+  # all tracks that have strand '.' which indicates the data is unstranded.
+  if match_gene_strand:
+    mismatched_strands = (df['gene_strand'] == '+') & (
+        df['track_strand'] == '-'
+    ) | (df['gene_strand'] == '-') & (df['track_strand'] == '+')
+    df = df[~mismatched_strands].reset_index(drop=True)
+  return df.reset_index(drop=True)
+
+
+def tidy_scores(
+    scores: Sequence[anndata.AnnData] | Sequence[Sequence[anndata.AnnData]],
+    match_gene_strand: bool = True,
+    include_extended_metadata: bool = True,
+) -> pd.DataFrame | None:
+  """Formats scores into a tidy (long) pandas DataFrame.
+
+  This function reformats variant scores into a more readable DataFrame with one
+  score per row. This function supports both scores generated from variant
+  scorers (e.g., the output of `score_variant` or `score_variants`) or interval
+  scorers (e.g., the output of `score_interval` or `score_intervals`).
+
+  It handles both variant/interval-centric scoring (producing one score row per
+  variant/interval-track pair) and gene-centric scoring (one score row per
+  variant/interval-gene-track combination).
+
+  The function accepts these score input types:
+
+  - A sequence of AnnData objects (e.g., output of `score_variant` with one or
+    more scorers).
+  - A nested sequence of AnnData objects (e.g., output of `score_variants`
+    with multiple variants).
+
+  Scores from multiple scorers or multiple variants/intervals are concatenated
+  together into a single pandas DataFrame, containing the union of all
+  applicable columns across scorers.
+
+  Args:
+    scores: Scoring output as either a sequence of AnnData objects, or a nested
+      sequence of AnnData objects (for example, the outputs of `score_variant`
+      and `score_variants`, respectively).
+    match_gene_strand: If True (and using gene-centric scoring), rows with
+      mismatched gene and track strands are removed.
+    include_extended_metadata: Argument passed to `tidy_anndata` to include
+      additional metadata columns where available.
+
+  Returns:
+    pd.DataFrame with columns:
+
+      - variant_id (when applicable): Variant of interest (e.g.
+        chr22:36201698:A>C).
+      - scored_interval: Genomic interval scored (e.g. chr22:36100000-36300000).
+      - gene_id: ENSEMBL gene identifier without version number.
+        (e.g. ENSG00000100342), or None if not applicable.
+      - gene_name: HGNC gene symbol (e.g. APOL1), or None if not applicable.
+      - gene_type: Gene biotype (e.g. protein_coding, lncRNA), or None if not
+        applicable.
+      - gene_strand: Strand of the gene ('+', '-', '.'), or None if not
+        applicable.
+      - output_type: Type of the output from the model (e.g. RNA_SEQ, DNASE).
+      - interval_scorer (when applicable): Name of the interval scorer used.
+      - variant_scorer (when applicable): Name of the variant scorer used.
+      - track_name: Name of the output track (e.g. UBERON:0036149 total
+        RNA-seq).
+      - track_strand: Strand of the track ('+', '-', or '.').
+      - ontology_curie: Ontology term for the cell type or tissue of the track
+        (e.g. UBERON:0036149), or NaN if not applicable.
+      - gtex_tissue: Name of the gtex tissue (e.g. Liver), or NaN if not
+        applicable.
+      - Assay title: Subtype of the assay (e.g., total RNA-seq), or NaN if not
+        applicable.
+      - biosample_name: Name of the biosample (e.g. liver), or NaN if not
+        applicable.
+      - biosample_type: Type of biosample (e.g. 'tissue' or 'primary cell'), or
+        NaN if not applicable.
+      - transcription_factor: Name of the transcription factor (e.g. 'CTCF'), or
+        NaN if not applicable.
+      - histone_mark: Name of the histological mark (e.g. 'H3K4ME3'), or NaN if
+        not applicable.
+      - raw_score: Raw variant score.
+      - quantile_score (when applicable): Quantile score.
+
+  Raises:
+    ValueError: If the input is not a valid type (sequence of AnnData or nested
+      sequence of AnnDatas).
+  """
+  if isinstance(scores, Sequence):
+    tidied_anndata = [
+        tidy_anndata(adata, match_gene_strand, include_extended_metadata)
+        for adata in itertools.chain.from_iterable(scores)
+    ]
+    if not tidied_anndata:
+      return None
+
+    concatenated_df = pd.concat(tidied_anndata, axis=0, ignore_index=True)
+    # Put the raw_score (and quantile_score when applicable) columns last.
+    last_columns = ['raw_score']
+    if 'quantile_score' in concatenated_df.columns:
+      last_columns.append('quantile_score')
+    return concatenated_df[
+        [col for col in concatenated_df.columns if col not in last_columns]
+        + last_columns
+    ]
+  else:
+    raise ValueError(
+        'Invalid input type. Must be sequence of AnnDatas or nested sequence of'
+        ' AnnDatas.'
+    )

flax_model/alphagenome/_sdk/protos/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Protocol buffers for interacting with genomic models."""

flax_model/alphagenome/_sdk/protos/dna_model.proto

@@ -0,0 +1,559 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package google.gdm.gdmscience.alphagenome.v1main;
+
+import "alphagenome/protos/tensor.proto";
+
+option go_package = "google.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenome";
+option java_multiple_files = true;
+option java_outer_classname = "DnaModelProto";
+option java_package = "com.google.gdm.gdmscience.alphagenome.v1main";
+
+// Message that represents a genomic interval.
+message Interval {
+  // The chromosome name, e.g., "chr1".
+  string chromosome = 1;
+
+  // 0-based start position.
+  int64 start = 2;
+
+  // 0-based end position. Must be greater than or equal to start.
+  int64 end = 3;
+
+  // The strand of the interval.
+  Strand strand = 4;
+}
+
+// Message that represents a genomic variant/mutation.
+message Variant {
+  // The chromosome name, e.g., "chr1".
+  string chromosome = 1;
+
+  // The 1-based position of the variant on the chromosome.
+  int64 position = 2;
+
+  // The reference base(s) at the variant position. Normally this corresponds to
+  // the sequence in the reference genome. Must only contain characters "ACGTN".
+  string reference_bases = 3;
+
+  // The alternate base(s) that replace the reference. For example:
+  // If sequence="ACT", position=2, reference_bases="C", alternate_bases="TG",
+  // then the alternate sequence would be "ATGT'.
+  string alternate_bases = 4;
+}
+
+// A single biological ontology term.
+message OntologyTerm {
+  // The type of ontology this term belongs to.
+  OntologyType ontology_type = 1;
+
+  // The ID of the term within the ontology.
+  int64 id = 2;
+}
+
+// Message that represents a biological sample that is used in an experiment.
+message Biosample {
+  // Biosample type matching ENCODE conventions.
+  BiosampleType type = 1;
+
+  // The name of the biosample, e.g. "T-cell".
+  string name = 2;
+
+  // Optional stage of the biosample, e.g. "adult".
+  optional string stage = 3;
+}
+
+// Message containing metadata for a gene score.
+message GeneScorerMetadata {
+  // ENSEMBL gene identifier without version number, e.g. "ENSG00000100342".
+  string gene_id = 1;
+
+  // Optional HGNC gene symbol, e.g. "APOL1".
+  optional string name = 3;
+
+  // Optional strand of the gene.
+  optional Strand strand = 2;
+
+  // Optional gene biotype, e.g. "protein_coding" or "lncRNA".
+  optional string type = 4;
+
+  // Optional start position for junction scores.
+  optional int64 junction_start = 5;
+
+  // Optional end position for junction scores.
+  optional int64 junction_end = 6;
+}
+
+// Message containing metadata for a track.
+message TrackMetadata {
+  // Name of the track.
+  string name = 1;
+
+  // Strand for the track.
+  Strand strand = 2;
+
+  // Ontology term for the track.
+  OntologyTerm ontology_term = 3;
+
+  // Biosample for the track.
+  Biosample biosample = 4;
+
+  // Optional assay for the track. e.g. "polyA plus RNA-seq".
+  optional string assay = 5;
+
+  // Optional histone mark for the track.
+  optional string histone_mark_code = 9;
+
+  // Optional transcription factor for the track.
+  optional string transcription_factor_code = 10;
+
+  // Optional GTEx tissue for the track.
+  optional string gtex_tissue = 8;
+
+  // Optional data source for the track. e.g. "encode".
+  optional string data_source = 11;
+
+  // Optional sequence reading endedness.
+  optional Endedness endedness = 12;
+
+  // Optional boolean indicating whether the track is genetically modified.
+  optional bool genetically_modified = 13;
+
+  // Optional mean of the non-zero values in the track.
+  optional float nonzero_mean = 14;
+}
+
+// Container for TrackMetadata.
+message TracksMetadata {
+  // Ordered list of metadata for each track.
+  repeated TrackMetadata metadata = 1;
+}
+
+// Message containing metadata for a splice junctions.
+message JunctionMetadata {
+  // Name of the junction track, e.g. "Brain_Cerebellum".
+  string name = 1;
+
+  // Ontology term for the junction track.
+  OntologyTerm ontology_term = 2;
+
+  // Biosample for the junction track.
+  Biosample biosample = 3;
+
+  // Optional GTEx tissue for the junction track.
+  optional string gtex_tissue = 4;
+
+  // Optional data source for the track. e.g. "encode".
+  optional string data_source = 5;
+
+  // Optional assay for the track. e.g. "polyA plus RNA-seq".
+  optional string assay = 6;
+}
+
+// Container for list of JunctionMetadata.
+message JunctionsMetadata {
+  // Ordered list of metadata for each junction.
+  repeated JunctionMetadata metadata = 1;
+}
+
+// Message for storing track values and metadata.
+message TrackData {
+  // Values for the track.
+  Tensor values = 1;
+
+  // Metadata for the track. The number of metadata elements must match the last
+  // dimension of the values tensor.
+  repeated TrackMetadata metadata = 2;
+
+  // Resolution of the track data in base pairs.
+  optional int64 resolution = 3;
+
+  // Optional Interval representing the genomic region.
+  Interval interval = 4;
+}
+
+// Message for storing splice junction values and metadata.
+message JunctionData {
+  // Values for the splice junction for each track.
+  Tensor values = 1;
+
+  // Metadata for the splice junction. The number of elements must match the
+  // last dimension of the values tensor.
+  repeated JunctionMetadata metadata = 2;
+
+  // Predicted splice junctions. The number of elements must match the first
+  // dimensions of the values tensor.
+  repeated Interval junctions = 3;
+
+  // Optional Interval representing the genomic region.
+  Interval interval = 4;
+}
+
+// Message containing metadata for an interval prediction.
+message IntervalMetadata {
+  // Genomic interval.
+  Interval interval = 1;
+
+  // Track metadata for the interval prediction.
+  repeated TrackMetadata track_metadata = 2;
+
+  // Gene metadata for the interval prediction.
+  repeated GeneScorerMetadata gene_metadata = 3;
+}
+
+// Message for storing interval values and metadata.
+message IntervalData {
+  // Values for the interval.
+  Tensor values = 1;
+
+  // Metadata for the interval.
+  IntervalMetadata metadata = 2;
+}
+
+// Message containing metadata for a variant prediction.
+message VariantMetadata {
+  // Genomic variant.
+  Variant variant = 1;
+
+  // Track metadata for the variant prediction.
+  repeated TrackMetadata track_metadata = 2;
+
+  // Gene metadata for the variant prediction.
+  repeated GeneScorerMetadata gene_metadata = 3;
+}
+
+// Message for storing variant values and metadata.
+message VariantData {
+  // Values for the variant.
+  Tensor values = 1;
+
+  // Metadata for the variant.
+  VariantMetadata metadata = 2;
+}
+
+// Message for storing model outputs for a single output type.
+message Output {
+  // The type of output.
+  OutputType output_type = 1;
+
+  // The payload for the output.
+  oneof payload {
+    // Track data for the output. Used in all output types except
+    // OUTPUT_TYPE_SPLICE_JUNCTIONS.
+    TrackData track_data = 2;
+
+    // Raw predictions with no metadata.
+    Tensor data = 3;
+
+    // Junction data for the output. This is only set for
+    // OUTPUT_TYPE_SPLICE_JUNCTIONS.
+    JunctionData junction_data = 4;
+  }
+}
+
+// Message for storing score interval results.
+message ScoreIntervalOutput {
+  // Score interval data.
+  IntervalData interval_data = 1;
+}
+
+// Message for storing score variant results.
+message ScoreVariantOutput {
+  // Score variant data.
+  VariantData variant_data = 1;
+}
+
+// Interval scorer message for gene-mask scoring.
+message GeneMaskIntervalScorer {
+  // Requested output type.
+  OutputType requested_output = 1;
+
+  // Requested width.
+  optional int64 width = 2;
+
+  // Requested aggregation type.
+  IntervalAggregationType aggregation_type = 3;
+}
+
+// Container message for all interval scorers.
+message IntervalScorer {
+  // The interval scorer payload.
+  oneof scorer {
+    // Gene mask scorer.
+    GeneMaskIntervalScorer gene_mask = 1;
+  }
+}
+
+// Variant scorer message for center mask scoring.
+message CenterMaskScorer {
+  // The width of the mask around the variant.
+  optional int64 width = 1;
+
+  // The aggregation type.
+  AggregationType aggregation_type = 2;
+
+  // Requested output type.
+  OutputType requested_output = 3;
+}
+
+// Variant scorer message for gene-mask scoring based on log fold change.
+message GeneMaskLFCScorer {
+  // Requested output type.
+  OutputType requested_output = 1;
+}
+
+// Variant scorer message for gene-mask scoring based on active allele counts.
+message GeneMaskActiveScorer {
+  // Requested output type.
+  OutputType requested_output = 1;
+}
+
+// Variant scorer message for gene-mask scoring for splicing.
+message GeneMaskSplicingScorer {
+  // The width of the mask around the variant.
+  optional int64 width = 1;
+
+  // Requested output type.
+  OutputType requested_output = 2;
+}
+
+// Variant scorer message for polyadenylation QTLs (paQTLs).
+message PolyadenylationScorer {}
+
+// Variant scorer message for splice junction scoring.
+message SpliceJunctionScorer {}
+
+// Variant scorer message for contact map scoring.
+message ContactMapScorer {}
+
+// Container message for all variant scorers.
+message VariantScorer {
+  // The variant scorer payload.
+  oneof scorer {
+    // Center mask scorer.
+    CenterMaskScorer center_mask = 1;
+
+    // Gene mask scorer based on log fold change.
+    GeneMaskLFCScorer gene_mask = 2;
+
+    // Gene mask scorer for splicing.
+    GeneMaskSplicingScorer gene_mask_splicing = 3;
+
+    // Polyadenylation scorer.
+    PolyadenylationScorer pa_qtl = 4;
+
+    // Splice junction scorer.
+    SpliceJunctionScorer splice_junction = 5;
+
+    // Contact map scorer.
+    ContactMapScorer contact_map = 6;
+
+    // Gene mask scorer based on active allele counts.
+    GeneMaskActiveScorer gene_mask_active = 7;
+  }
+}
+
+// Message containing metadata for an output type.
+message OutputMetadata {
+  // The output type.
+  OutputType output_type = 1;
+
+  // The payload for the output metadata.
+  oneof payload {
+    // Track metadata for the output. Used in all output types except
+    // OUTPUT_TYPE_SPLICE_JUNCTIONS.
+    TracksMetadata tracks = 2;
+
+    // Junction metadata for the output. This is only set for
+    // OUTPUT_TYPE_SPLICE_JUNCTIONS.
+    JunctionsMetadata junctions = 3;
+  }
+}
+
+// Defines the possible strands for a DNA sequence.
+enum Strand {
+  // Unspecified strand.
+  STRAND_UNSPECIFIED = 0;
+
+  // The forward strand (5' to 3')
+  STRAND_POSITIVE = 1;
+
+  // The reverse strand (3' to 5')
+  STRAND_NEGATIVE = 2;
+
+  // The strand is not specified.
+  STRAND_UNSTRANDED = 3;
+}
+
+// Supported ontology types.
+enum OntologyType {
+  // Unspecified ontology type.
+  ONTOLOGY_TYPE_UNSPECIFIED = 0;
+
+  // Cell Line Ontology.
+  ONTOLOGY_TYPE_CLO = 1;
+
+  // Uber-anatomy ontology.
+  ONTOLOGY_TYPE_UBERON = 2;
+
+  // Cell Ontology.
+  ONTOLOGY_TYPE_CL = 3;
+
+  // Experimental Factor Ontology.
+  ONTOLOGY_TYPE_EFO = 4;
+
+  // New Term Requested.
+  ONTOLOGY_TYPE_NTR = 5;
+}
+
+// Biosample type matching ENCODE conventions. See
+// https://www.encodeproject.org/profiles/biosample_type.
+enum BiosampleType {
+  // Unspecified biosample type.
+  BIOSAMPLE_TYPE_UNSPECIFIED = 0;
+
+  // Primary cell sample.
+  BIOSAMPLE_TYPE_PRIMARY_CELL = 1;
+
+  // In-vitro differentiated cell sample.
+  BIOSAMPLE_TYPE_IN_VITRO_DIFFERENTIATED_CELLS = 2;
+
+  // Cell line sample.
+  BIOSAMPLE_TYPE_CELL_LINE = 3;
+
+  // Tissue sample.
+  BIOSAMPLE_TYPE_TISSUE = 4;
+
+  // Technical sample.
+  BIOSAMPLE_TYPE_TECHNICAL_SAMPLE = 5;
+
+  // Organoid sample.
+  BIOSAMPLE_TYPE_ORGANOID = 6;
+}
+
+// Enumeration of all the available types of outputs.
+enum OutputType {
+  // Unspecified output type.
+  OUTPUT_TYPE_UNSPECIFIED = 0;
+
+  // ATAC-seq tracks capturing chromatin accessibility.
+  OUTPUT_TYPE_ATAC = 1;
+
+  // CAGE (Cap Analysis of Gene Expression) tracks capturing gene expression.
+  OUTPUT_TYPE_CAGE = 2;
+
+  // DNase I hypersensitive site tracks capturing chromatin accessibility.
+  OUTPUT_TYPE_DNASE = 3;
+
+  // RNA sequencing tracks capturing gene expression.
+  OUTPUT_TYPE_RNA_SEQ = 4;
+
+  // ChIP-seq tracks capturing histone modifications.
+  OUTPUT_TYPE_CHIP_HISTONE = 5;
+
+  // ChIP-seq tracks capturing transcription factor binding.
+  OUTPUT_TYPE_CHIP_TF = 6;
+
+  // Splice site tracks capturing donor and acceptor splice sites.
+  OUTPUT_TYPE_SPLICE_SITES = 7;
+
+  // Splice site usage tracks capturing the fraction of the time that each
+  // splice site is used.
+  OUTPUT_TYPE_SPLICE_SITE_USAGE = 8;
+
+  // Splice junction tracks capturing split read RNA-seq counts for each
+  // junction.
+  OUTPUT_TYPE_SPLICE_JUNCTIONS = 9;
+
+  // Contact map tracks capturing 3D chromatin contact probabilities.
+  OUTPUT_TYPE_CONTACT_MAPS = 11;
+
+  // Precision Run-On sequencing and capping, used to measure gene expression.
+  OUTPUT_TYPE_PROCAP = 12;
+}
+
+// Enumeration of all the available organisms. Values relate to the NCBI
+// taxonomy ID.
+enum Organism {
+  // Unspecified organism.
+  ORGANISM_UNSPECIFIED = 0;
+
+  // Human (Homo sapiens).
+  ORGANISM_HOMO_SAPIENS = 9606;
+
+  // Mouse (Mus musculus).
+  ORGANISM_MUS_MUSCULUS = 10090;
+}
+
+// Enum indicating the type of interval scorer aggregation.
+enum IntervalAggregationType {
+  // Unspecified aggregation type.
+  INTERVAL_AGGREGATION_TYPE_UNSPECIFIED = 0;
+
+  // Mean across positions.
+  INTERVAL_AGGREGATION_TYPE_MEAN = 1;
+
+  // Sum across positions.
+  INTERVAL_AGGREGATION_TYPE_SUM = 2;
+}
+
+// Enum indicating the type of variant scorer aggregation.
+enum AggregationType {
+  // Unspecified aggregation type.
+  AGGREGATION_TYPE_UNSPECIFIED = 0;
+
+  // Difference of means, i.e., mean(ALT) - mean(REF).
+  AGGREGATION_TYPE_DIFF_MEAN = 1;
+
+  // Difference of sums, i.e., sum(ALT) - sum(REF).
+  AGGREGATION_TYPE_DIFF_SUM = 2;
+
+  // Log scales predictions, then takes the sum and then the difference, i.e.
+  // sum(log2(ALT)) - sum(log2(REF)).
+  AGGREGATION_TYPE_DIFF_SUM_LOG2 = 3;
+
+  // Takes the difference of ALT and REF predictions, then computes the L2 norm,
+  // i.e., l2_norm(ALT - REF).
+  AGGREGATION_TYPE_L2_DIFF = 4;
+
+  // Log scales the predictions + 1, takes the difference of `ALT` and `REF`
+  // predictions, then computes the L2 norm, i.e., l2_norm(log1p(`ALT`) -
+  // log1p(`REF`)).
+  AGGREGATION_TYPE_L2_DIFF_LOG1P = 8;
+
+  // Takes the sum of predictions, applies a log transform, then takes the
+  // difference between predictions, i.e., log2(sum(ALT)) - log2(sum(REF)).
+  AGGREGATION_TYPE_DIFF_LOG2_SUM = 5;
+
+  // Maximum of means, i.e., max(mean(ALT), mean(REF)).
+  AGGREGATION_TYPE_ACTIVE_MEAN = 6;
+
+  // Maximum of sums, i.e., max(sum(ALT), sum(REF)).
+  AGGREGATION_TYPE_ACTIVE_SUM = 7;
+}
+
+// Enum indicating the endedness of a track experiment.
+enum Endedness {
+  // Unspecified endedness.
+  ENDEDNESS_UNSPECIFIED = 0;
+
+  // Single end.
+  ENDEDNESS_SINGLE = 1;
+
+  // Paired end.
+  ENDEDNESS_PAIRED = 2;
+}

flax_model/alphagenome/_sdk/protos/dna_model_pb2.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: alphagenome/protos/dna_model.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from flax_model.alphagenome._sdk.protos import tensor_pb2 as alphagenome_dot_protos_dot_tensor__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\"alphagenome/protos/dna_model.proto\x12(google.gdm.gdmscience.alphagenome.v1main\x1a\x1f\x61lphagenome/protos/tensor.proto\"|\n\x08Interval\x12\x12\n\nchromosome\x18\x01 \x01(\t\x12\r\n\x05start\x18\x02 \x01(\x03\x12\x0b\n\x03\x65nd\x18\x03 \x01(\x03\x12@\n\x06strand\x18\x04 \x01(\x0e\x32\x30.google.gdm.gdmscience.alphagenome.v1main.Strand\"a\n\x07Variant\x12\x12\n\nchromosome\x18\x01 \x01(\t\x12\x10\n\x08position\x18\x02 \x01(\x03\x12\x17\n\x0freference_bases\x18\x03 \x01(\t\x12\x17\n\x0f\x61lternate_bases\x18\x04 \x01(\t\"i\n\x0cOntologyTerm\x12M\n\rontology_type\x18\x01 \x01(\x0e\x32\x36.google.gdm.gdmscience.alphagenome.v1main.OntologyType\x12\n\n\x02id\x18\x02 \x01(\x03\"~\n\tBiosample\x12\x45\n\x04type\x18\x01 \x01(\x0e\x32\x37.google.gdm.gdmscience.alphagenome.v1main.BiosampleType\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x12\n\x05stage\x18\x03 \x01(\tH\x00\x88\x01\x01\x42\x08\n\x06_stage\"\x8b\x02\n\x12GeneScorerMetadata\x12\x0f\n\x07gene_id\x18\x01 \x01(\t\x12\x11\n\x04name\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x45\n\x06strand\x18\x02 \x01(\x0e\x32\x30.google.gdm.gdmscience.alphagenome.v1main.StrandH\x01\x88\x01\x01\x12\x11\n\x04type\x18\x04 \x01(\tH\x02\x88\x01\x01\x12\x1b\n\x0ejunction_start\x18\x05 \x01(\x03H\x03\x88\x01\x01\x12\x19\n\x0cjunction_end\x18\x06 \x01(\x03H\x04\x88\x01\x01\x42\x07\n\x05_nameB\t\n\x07_strandB\x07\n\x05_typeB\x11\n\x0f_junction_startB\x0f\n\r_junction_end\"\xa7\x05\n\rTrackMetadata\x12\x0c\n\x04name\x18\x01 \x01(\t\x12@\n\x06strand\x18\x02 \x01(\x0e\x32\x30.google.gdm.gdmscience.alphagenome.v1main.Strand\x12M\n\rontology_term\x18\x03 \x01(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.OntologyTerm\x12\x46\n\tbiosample\x18\x04 \x01(\x0b\x32\x33.google.gdm.gdmscience.alphagenome.v1main.Biosample\x12\x12\n\x05\x61ssay\x18\x05 \x01(\tH\x00\x88\x01\x01\x12\x1e\n\x11histone_mark_code\x18\t \x01(\tH\x01\x88\x01\x01\x12&\n\x19transcription_factor_code\x18\n \x01(\tH\x02\x88\x01\x01\x12\x18\n\x0bgtex_tissue\x18\x08 \x01(\tH\x03\x88\x01\x01\x12\x18\n\x0b\x64\x61ta_source\x18\x0b \x01(\tH\x04\x88\x01\x01\x12K\n\tendedness\x18\x0c \x01(\x0e\x32\x33.google.gdm.gdmscience.alphagenome.v1main.EndednessH\x05\x88\x01\x01\x12!\n\x14genetically_modified\x18\r \x01(\x08H\x06\x88\x01\x01\x12\x19\n\x0cnonzero_mean\x18\x0e \x01(\x02H\x07\x88\x01\x01\x42\x08\n\x06_assayB\x14\n\x12_histone_mark_codeB\x1c\n\x1a_transcription_factor_codeB\x0e\n\x0c_gtex_tissueB\x0e\n\x0c_data_sourceB\x0c\n\n_endednessB\x17\n\x15_genetically_modifiedB\x0f\n\r_nonzero_mean\"[\n\x0eTracksMetadata\x12I\n\x08metadata\x18\x01 \x03(\x0b\x32\x37.google.gdm.gdmscience.alphagenome.v1main.TrackMetadata\"\xa9\x02\n\x10JunctionMetadata\x12\x0c\n\x04name\x18\x01 \x01(\t\x12M\n\rontology_term\x18\x02 \x01(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.OntologyTerm\x12\x46\n\tbiosample\x18\x03 \x01(\x0b\x32\x33.google.gdm.gdmscience.alphagenome.v1main.Biosample\x12\x18\n\x0bgtex_tissue\x18\x04 \x01(\tH\x00\x88\x01\x01\x12\x18\n\x0b\x64\x61ta_source\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05\x61ssay\x18\x06 \x01(\tH\x02\x88\x01\x01\x42\x0e\n\x0c_gtex_tissueB\x0e\n\x0c_data_sourceB\x08\n\x06_assay\"a\n\x11JunctionsMetadata\x12L\n\x08metadata\x18\x01 \x03(\x0b\x32:.google.gdm.gdmscience.alphagenome.v1main.JunctionMetadata\"\x86\x02\n\tTrackData\x12@\n\x06values\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.Tensor\x12I\n\x08metadata\x18\x02 \x03(\x0b\x32\x37.google.gdm.gdmscience.alphagenome.v1main.TrackMetadata\x12\x17\n\nresolution\x18\x03 \x01(\x03H\x00\x88\x01\x01\x12\x44\n\x08interval\x18\x04 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.IntervalB\r\n\x0b_resolution\"\xab\x02\n\x0cJunctionData\x12@\n\x06values\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.Tensor\x12L\n\x08metadata\x18\x02 \x03(\x0b\x32:.google.gdm.gdmscience.alphagenome.v1main.JunctionMetadata\x12\x45\n\tjunctions\x18\x03 \x03(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12\x44\n\x08interval\x18\x04 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\"\xfe\x01\n\x10IntervalMetadata\x12\x44\n\x08interval\x18\x01 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12O\n\x0etrack_metadata\x18\x02 \x03(\x0b\x32\x37.google.gdm.gdmscience.alphagenome.v1main.TrackMetadata\x12S\n\rgene_metadata\x18\x03 \x03(\x0b\x32<.google.gdm.gdmscience.alphagenome.v1main.GeneScorerMetadata\"\x9e\x01\n\x0cIntervalData\x12@\n\x06values\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.Tensor\x12L\n\x08metadata\x18\x02 \x01(\x0b\x32:.google.gdm.gdmscience.alphagenome.v1main.IntervalMetadata\"\xfb\x01\n\x0fVariantMetadata\x12\x42\n\x07variant\x18\x01 \x01(\x0b\x32\x31.google.gdm.gdmscience.alphagenome.v1main.Variant\x12O\n\x0etrack_metadata\x18\x02 \x03(\x0b\x32\x37.google.gdm.gdmscience.alphagenome.v1main.TrackMetadata\x12S\n\rgene_metadata\x18\x03 \x03(\x0b\x32<.google.gdm.gdmscience.alphagenome.v1main.GeneScorerMetadata\"\x9c\x01\n\x0bVariantData\x12@\n\x06values\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.Tensor\x12K\n\x08metadata\x18\x02 \x01(\x0b\x32\x39.google.gdm.gdmscience.alphagenome.v1main.VariantMetadata\"\xbc\x02\n\x06Output\x12I\n\x0boutput_type\x18\x01 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\x12I\n\ntrack_data\x18\x02 \x01(\x0b\x32\x33.google.gdm.gdmscience.alphagenome.v1main.TrackDataH\x00\x12@\n\x04\x64\x61ta\x18\x03 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.TensorH\x00\x12O\n\rjunction_data\x18\x04 \x01(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.JunctionDataH\x00\x42\t\n\x07payload\"d\n\x13ScoreIntervalOutput\x12M\n\rinterval_data\x18\x01 \x01(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.IntervalData\"a\n\x12ScoreVariantOutput\x12K\n\x0cvariant_data\x18\x01 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.VariantData\"\xe3\x01\n\x16GeneMaskIntervalScorer\x12N\n\x10requested_output\x18\x01 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\x12\x12\n\x05width\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12[\n\x10\x61ggregation_type\x18\x03 \x01(\x0e\x32\x41.google.gdm.gdmscience.alphagenome.v1main.IntervalAggregationTypeB\x08\n\x06_width\"q\n\x0eIntervalScorer\x12U\n\tgene_mask\x18\x01 \x01(\x0b\x32@.google.gdm.gdmscience.alphagenome.v1main.GeneMaskIntervalScorerH\x00\x42\x08\n\x06scorer\"\xd5\x01\n\x10\x43\x65nterMaskScorer\x12\x12\n\x05width\x18\x01 \x01(\x03H\x00\x88\x01\x01\x12S\n\x10\x61ggregation_type\x18\x02 \x01(\x0e\x32\x39.google.gdm.gdmscience.alphagenome.v1main.AggregationType\x12N\n\x10requested_output\x18\x03 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputTypeB\x08\n\x06_width\"c\n\x11GeneMaskLFCScorer\x12N\n\x10requested_output\x18\x01 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\"f\n\x14GeneMaskActiveScorer\x12N\n\x10requested_output\x18\x01 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\"\x86\x01\n\x16GeneMaskSplicingScorer\x12\x12\n\x05width\x18\x01 \x01(\x03H\x00\x88\x01\x01\x12N\n\x10requested_output\x18\x02 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputTypeB\x08\n\x06_width\"\x17\n\x15PolyadenylationScorer\"\x16\n\x14SpliceJunctionScorer\"\x12\n\x10\x43ontactMapScorer\"\xfb\x04\n\rVariantScorer\x12Q\n\x0b\x63\x65nter_mask\x18\x01 \x01(\x0b\x32:.google.gdm.gdmscience.alphagenome.v1main.CenterMaskScorerH\x00\x12P\n\tgene_mask\x18\x02 \x01(\x0b\x32;.google.gdm.gdmscience.alphagenome.v1main.GeneMaskLFCScorerH\x00\x12^\n\x12gene_mask_splicing\x18\x03 \x01(\x0b\x32@.google.gdm.gdmscience.alphagenome.v1main.GeneMaskSplicingScorerH\x00\x12Q\n\x06pa_qtl\x18\x04 \x01(\x0b\x32?.google.gdm.gdmscience.alphagenome.v1main.PolyadenylationScorerH\x00\x12Y\n\x0fsplice_junction\x18\x05 \x01(\x0b\x32>.google.gdm.gdmscience.alphagenome.v1main.SpliceJunctionScorerH\x00\x12Q\n\x0b\x63ontact_map\x18\x06 \x01(\x0b\x32:.google.gdm.gdmscience.alphagenome.v1main.ContactMapScorerH\x00\x12Z\n\x10gene_mask_active\x18\x07 \x01(\x0b\x32>.google.gdm.gdmscience.alphagenome.v1main.GeneMaskActiveScorerH\x00\x42\x08\n\x06scorer\"\x84\x02\n\x0eOutputMetadata\x12I\n\x0boutput_type\x18\x01 \x01(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\x12J\n\x06tracks\x18\x02 \x01(\x0b\x32\x38.google.gdm.gdmscience.alphagenome.v1main.TracksMetadataH\x00\x12P\n\tjunctions\x18\x03 \x01(\x0b\x32;.google.gdm.gdmscience.alphagenome.v1main.JunctionsMetadataH\x00\x42\t\n\x07payload*a\n\x06Strand\x12\x16\n\x12STRAND_UNSPECIFIED\x10\x00\x12\x13\n\x0fSTRAND_POSITIVE\x10\x01\x12\x13\n\x0fSTRAND_NEGATIVE\x10\x02\x12\x15\n\x11STRAND_UNSTRANDED\x10\x03*\xa2\x01\n\x0cOntologyType\x12\x1d\n\x19ONTOLOGY_TYPE_UNSPECIFIED\x10\x00\x12\x15\n\x11ONTOLOGY_TYPE_CLO\x10\x01\x12\x18\n\x14ONTOLOGY_TYPE_UBERON\x10\x02\x12\x14\n\x10ONTOLOGY_TYPE_CL\x10\x03\x12\x15\n\x11ONTOLOGY_TYPE_EFO\x10\x04\x12\x15\n\x11ONTOLOGY_TYPE_NTR\x10\x05*\xfd\x01\n\rBiosampleType\x12\x1e\n\x1a\x42IOSAMPLE_TYPE_UNSPECIFIED\x10\x00\x12\x1f\n\x1b\x42IOSAMPLE_TYPE_PRIMARY_CELL\x10\x01\x12\x30\n,BIOSAMPLE_TYPE_IN_VITRO_DIFFERENTIATED_CELLS\x10\x02\x12\x1c\n\x18\x42IOSAMPLE_TYPE_CELL_LINE\x10\x03\x12\x19\n\x15\x42IOSAMPLE_TYPE_TISSUE\x10\x04\x12#\n\x1f\x42IOSAMPLE_TYPE_TECHNICAL_SAMPLE\x10\x05\x12\x1b\n\x17\x42IOSAMPLE_TYPE_ORGANOID\x10\x06*\xd5\x02\n\nOutputType\x12\x1b\n\x17OUTPUT_TYPE_UNSPECIFIED\x10\x00\x12\x14\n\x10OUTPUT_TYPE_ATAC\x10\x01\x12\x14\n\x10OUTPUT_TYPE_CAGE\x10\x02\x12\x15\n\x11OUTPUT_TYPE_DNASE\x10\x03\x12\x17\n\x13OUTPUT_TYPE_RNA_SEQ\x10\x04\x12\x1c\n\x18OUTPUT_TYPE_CHIP_HISTONE\x10\x05\x12\x17\n\x13OUTPUT_TYPE_CHIP_TF\x10\x06\x12\x1c\n\x18OUTPUT_TYPE_SPLICE_SITES\x10\x07\x12!\n\x1dOUTPUT_TYPE_SPLICE_SITE_USAGE\x10\x08\x12 \n\x1cOUTPUT_TYPE_SPLICE_JUNCTIONS\x10\t\x12\x1c\n\x18OUTPUT_TYPE_CONTACT_MAPS\x10\x0b\x12\x16\n\x12OUTPUT_TYPE_PROCAP\x10\x0c*\\\n\x08Organism\x12\x18\n\x14ORGANISM_UNSPECIFIED\x10\x00\x12\x1a\n\x15ORGANISM_HOMO_SAPIENS\x10\x86K\x12\x1a\n\x15ORGANISM_MUS_MUSCULUS\x10\xeaN*\x8b\x01\n\x17IntervalAggregationType\x12)\n%INTERVAL_AGGREGATION_TYPE_UNSPECIFIED\x10\x00\x12\"\n\x1eINTERVAL_AGGREGATION_TYPE_MEAN\x10\x01\x12!\n\x1dINTERVAL_AGGREGATION_TYPE_SUM\x10\x02*\xbf\x02\n\x0f\x41ggregationType\x12 \n\x1c\x41GGREGATION_TYPE_UNSPECIFIED\x10\x00\x12\x1e\n\x1a\x41GGREGATION_TYPE_DIFF_MEAN\x10\x01\x12\x1d\n\x19\x41GGREGATION_TYPE_DIFF_SUM\x10\x02\x12\"\n\x1e\x41GGREGATION_TYPE_DIFF_SUM_LOG2\x10\x03\x12\x1c\n\x18\x41GGREGATION_TYPE_L2_DIFF\x10\x04\x12\"\n\x1e\x41GGREGATION_TYPE_L2_DIFF_LOG1P\x10\x08\x12\"\n\x1e\x41GGREGATION_TYPE_DIFF_LOG2_SUM\x10\x05\x12 \n\x1c\x41GGREGATION_TYPE_ACTIVE_MEAN\x10\x06\x12\x1f\n\x1b\x41GGREGATION_TYPE_ACTIVE_SUM\x10\x07*R\n\tEndedness\x12\x19\n\x15\x45NDEDNESS_UNSPECIFIED\x10\x00\x12\x14\n\x10\x45NDEDNESS_SINGLE\x10\x01\x12\x14\n\x10\x45NDEDNESS_PAIRED\x10\x02\x42\x94\x01\n,com.google.gdm.gdmscience.alphagenome.v1mainB\rDnaModelProtoP\x01ZSgoogle.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenomeb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flax_model.alphagenome._sdk.protos.dna_model_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  _globals['DESCRIPTOR']._options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n,com.google.gdm.gdmscience.alphagenome.v1mainB\rDnaModelProtoP\001ZSgoogle.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenome'
+  _globals['_STRAND']._serialized_start=5808
+  _globals['_STRAND']._serialized_end=5905
+  _globals['_ONTOLOGYTYPE']._serialized_start=5908
+  _globals['_ONTOLOGYTYPE']._serialized_end=6070
+  _globals['_BIOSAMPLETYPE']._serialized_start=6073
+  _globals['_BIOSAMPLETYPE']._serialized_end=6326
+  _globals['_OUTPUTTYPE']._serialized_start=6329
+  _globals['_OUTPUTTYPE']._serialized_end=6670
+  _globals['_ORGANISM']._serialized_start=6672
+  _globals['_ORGANISM']._serialized_end=6764
+  _globals['_INTERVALAGGREGATIONTYPE']._serialized_start=6767
+  _globals['_INTERVALAGGREGATIONTYPE']._serialized_end=6906
+  _globals['_AGGREGATIONTYPE']._serialized_start=6909
+  _globals['_AGGREGATIONTYPE']._serialized_end=7228
+  _globals['_ENDEDNESS']._serialized_start=7230
+  _globals['_ENDEDNESS']._serialized_end=7312
+  _globals['_INTERVAL']._serialized_start=113
+  _globals['_INTERVAL']._serialized_end=237
+  _globals['_VARIANT']._serialized_start=239
+  _globals['_VARIANT']._serialized_end=336
+  _globals['_ONTOLOGYTERM']._serialized_start=338
+  _globals['_ONTOLOGYTERM']._serialized_end=443
+  _globals['_BIOSAMPLE']._serialized_start=445
+  _globals['_BIOSAMPLE']._serialized_end=571
+  _globals['_GENESCORERMETADATA']._serialized_start=574
+  _globals['_GENESCORERMETADATA']._serialized_end=841
+  _globals['_TRACKMETADATA']._serialized_start=844
+  _globals['_TRACKMETADATA']._serialized_end=1523
+  _globals['_TRACKSMETADATA']._serialized_start=1525
+  _globals['_TRACKSMETADATA']._serialized_end=1616
+  _globals['_JUNCTIONMETADATA']._serialized_start=1619
+  _globals['_JUNCTIONMETADATA']._serialized_end=1916
+  _globals['_JUNCTIONSMETADATA']._serialized_start=1918
+  _globals['_JUNCTIONSMETADATA']._serialized_end=2015
+  _globals['_TRACKDATA']._serialized_start=2018
+  _globals['_TRACKDATA']._serialized_end=2280
+  _globals['_JUNCTIONDATA']._serialized_start=2283
+  _globals['_JUNCTIONDATA']._serialized_end=2582
+  _globals['_INTERVALMETADATA']._serialized_start=2585
+  _globals['_INTERVALMETADATA']._serialized_end=2839
+  _globals['_INTERVALDATA']._serialized_start=2842
+  _globals['_INTERVALDATA']._serialized_end=3000
+  _globals['_VARIANTMETADATA']._serialized_start=3003
+  _globals['_VARIANTMETADATA']._serialized_end=3254
+  _globals['_VARIANTDATA']._serialized_start=3257
+  _globals['_VARIANTDATA']._serialized_end=3413
+  _globals['_OUTPUT']._serialized_start=3416
+  _globals['_OUTPUT']._serialized_end=3732
+  _globals['_SCOREINTERVALOUTPUT']._serialized_start=3734
+  _globals['_SCOREINTERVALOUTPUT']._serialized_end=3834
+  _globals['_SCOREVARIANTOUTPUT']._serialized_start=3836
+  _globals['_SCOREVARIANTOUTPUT']._serialized_end=3933
+  _globals['_GENEMASKINTERVALSCORER']._serialized_start=3936
+  _globals['_GENEMASKINTERVALSCORER']._serialized_end=4163
+  _globals['_INTERVALSCORER']._serialized_start=4165
+  _globals['_INTERVALSCORER']._serialized_end=4278
+  _globals['_CENTERMASKSCORER']._serialized_start=4281
+  _globals['_CENTERMASKSCORER']._serialized_end=4494
+  _globals['_GENEMASKLFCSCORER']._serialized_start=4496
+  _globals['_GENEMASKLFCSCORER']._serialized_end=4595
+  _globals['_GENEMASKACTIVESCORER']._serialized_start=4597
+  _globals['_GENEMASKACTIVESCORER']._serialized_end=4699
+  _globals['_GENEMASKSPLICINGSCORER']._serialized_start=4702
+  _globals['_GENEMASKSPLICINGSCORER']._serialized_end=4836
+  _globals['_POLYADENYLATIONSCORER']._serialized_start=4838
+  _globals['_POLYADENYLATIONSCORER']._serialized_end=4861
+  _globals['_SPLICEJUNCTIONSCORER']._serialized_start=4863
+  _globals['_SPLICEJUNCTIONSCORER']._serialized_end=4885
+  _globals['_CONTACTMAPSCORER']._serialized_start=4887
+  _globals['_CONTACTMAPSCORER']._serialized_end=4905
+  _globals['_VARIANTSCORER']._serialized_start=4908
+  _globals['_VARIANTSCORER']._serialized_end=5543
+  _globals['_OUTPUTMETADATA']._serialized_start=5546
+  _globals['_OUTPUTMETADATA']._serialized_end=5806
+# @@protoc_insertion_point(module_scope)

flax_model/alphagenome/_sdk/protos/dna_model_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

flax_model/alphagenome/_sdk/protos/dna_model_service.proto

@@ -0,0 +1,268 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package google.gdm.gdmscience.alphagenome.v1main;
+
+import "alphagenome/protos/dna_model.proto";
+import "alphagenome/protos/tensor.proto";
+
+option go_package = "google.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenome";
+option java_multiple_files = true;
+option java_outer_classname = "DnaModelServiceProto";
+option java_package = "com.google.gdm.gdmscience.alphagenome.v1main";
+
+// Service for making predictions with DNA models.
+service DnaModelService {
+  // Makes a prediction for DNA sequence.
+  rpc PredictSequence(stream PredictSequenceRequest)
+      returns (stream PredictSequenceResponse) {}
+
+  // Make prediction for a single genomic interval.
+  rpc PredictInterval(stream PredictIntervalRequest)
+      returns (stream PredictIntervalResponse) {}
+
+  // Make variant effect predictions for a genomic interval.
+  rpc PredictVariant(stream PredictVariantRequest)
+      returns (stream PredictVariantResponse) {}
+
+  // Score a genomic interval.
+  rpc ScoreInterval(stream ScoreIntervalRequest)
+      returns (stream ScoreIntervalResponse) {}
+
+  // Score a variant for a genomic interval.
+  rpc ScoreVariant(stream ScoreVariantRequest)
+      returns (stream ScoreVariantResponse) {}
+
+  // Score ISM variant effect predictions for a genomic interval.
+  rpc ScoreIsmVariant(stream ScoreIsmVariantRequest)
+      returns (stream ScoreIsmVariantResponse) {}
+
+  // Get metadata for the model.
+  rpc GetMetadata(MetadataRequest) returns (stream MetadataResponse) {}
+}
+
+// Request message for predicting a sequence.
+message PredictSequenceRequest {
+  // DNA sequence to make prediction for. Must only contain characters "ACGTN".
+  string sequence = 1;
+
+  // Organism to use for the prediction.
+  Organism organism = 2;
+
+  // Ontology terms to generate predictions for. If empty returns all
+  // ontologies.
+  repeated OntologyTerm ontology_terms = 3;
+
+  // Output types to generate predictions for.
+  repeated OutputType requested_outputs = 4;
+
+  // Model version to use.
+  string model_version = 5;
+}
+
+// Response message for predicting a sequence.
+message PredictSequenceResponse {
+  // The payload for the response.
+  oneof payload {
+    // Output for a single output type.
+    Output output = 1;
+
+    // Tensor chunk related to the last output message received. It is an error
+    // to receive a tensor chunk without a previous output message.
+    TensorChunk tensor_chunk = 2;
+  }
+}
+
+// Request message for predicting an interval.
+message PredictIntervalRequest {
+  // DNA interval to make prediction for.
+  Interval interval = 1;
+
+  // Organism to use for the prediction.
+  Organism organism = 2;
+
+  // Output types to generate predictions for.
+  repeated OutputType requested_outputs = 3;
+
+  // Ontology terms to generate predictions for. If empty returns all
+  // ontologies.
+  repeated OntologyTerm ontology_terms = 4;
+
+  // Model version to use.
+  string model_version = 5;
+}
+
+// Response message for predicting an interval.
+message PredictIntervalResponse {
+  // The payload for the response.
+  oneof payload {
+    // Output for a single output type.
+    Output output = 1;
+
+    // Tensor chunk related to the last output message received. It is an error
+    // to receive a tensor chunk without a previous output message.
+    TensorChunk tensor_chunk = 2;
+  }
+}
+
+// Request message for predicting a variant.
+message PredictVariantRequest {
+  // DNA interval to make prediction for.
+  Interval interval = 1;
+
+  // DNA variant to make prediction for.
+  Variant variant = 2;
+
+  // Organism to use for the prediction.
+  Organism organism = 3;
+
+  // Output types to generate predictions for.
+  repeated OutputType requested_outputs = 4;
+
+  // Ontology terms to generate predictions for. If empty returns all
+  // ontologies.
+  repeated OntologyTerm ontology_terms = 5;
+
+  // Model version to use.
+  string model_version = 6;
+}
+
+// Response message for predicting a variant.
+message PredictVariantResponse {
+  // The payload for the response.
+  oneof payload {
+    // Reference output for a single output type.
+    Output reference_output = 1;
+
+    // Alternate output for a single output type.
+    Output alternate_output = 2;
+
+    // Tensor chunk related to the last reference or alternate output message
+    // received. It is an error to receive a tensor chunk without a previous
+    // reference or alternate output message.
+    TensorChunk tensor_chunk = 3;
+  }
+}
+
+// Request message for scoring an interval.
+message ScoreIntervalRequest {
+  // DNA interval to make prediction for.
+  Interval interval = 1;
+
+  // Organism to use for the prediction.
+  Organism organism = 2;
+
+  // Interval scorers to use for the prediction.
+  repeated IntervalScorer interval_scorers = 3;
+
+  // Model version to use.
+  string model_version = 4;
+}
+
+// Response message for scoring an interval.
+message ScoreIntervalResponse {
+  // The payload for the response.
+  oneof payload {
+    // Score interval output.
+    ScoreIntervalOutput output = 1;
+
+    // Tensor chunk related to the last score interval output message received.
+    // It is an error to receice a tensor chunk without a previous output
+    // message.
+    TensorChunk tensor_chunk = 2;
+  }
+}
+
+// Request message for scoring a variant.
+message ScoreVariantRequest {
+  // DNA interval to make prediction for.
+  Interval interval = 1;
+
+  // DNA variant to make prediction for.
+  Variant variant = 2;
+
+  // Organism to use for the prediction.
+  Organism organism = 3;
+
+  // Variant scorers to use for the prediction.
+  repeated VariantScorer variant_scorers = 4;
+
+  // Model version to use.
+  string model_version = 5;
+}
+
+// Response message for scoring a variant.
+message ScoreVariantResponse {
+  // The payload for the response.
+  oneof payload {
+    // Score variant output.
+    ScoreVariantOutput output = 1;
+
+    // Tensor chunk related to the last score variant output message received.
+    // It is an error to receive a tensor chunk without a previous output
+    // message.
+    TensorChunk tensor_chunk = 2;
+  }
+}
+
+// Request message for scoring an in-silico mutagenesis (ISM) interval.
+message ScoreIsmVariantRequest {
+  // DNA interval to make the prediction for.
+  Interval interval = 1;
+
+  // ISM interval to make the prediction for.
+  Interval ism_interval = 2;
+
+  // Organism to use for the prediction.
+  Organism organism = 3;
+
+  // Variant scorers to use for the prediction.
+  repeated VariantScorer variant_scorers = 4;
+
+  // Optional variant applied to the reference interval. If provided, the
+  // alternate allele is used for in-silico mutagenesis, otherwise the
+  // unaltered reference sequence is used.
+  optional Variant interval_variant = 6;
+
+  // Model version to use.
+  string model_version = 5;
+}
+
+// Response message for scoring an in-silico mutagenesis (ISM) interval.
+message ScoreIsmVariantResponse {
+  // The payload for the response.
+  oneof payload {
+    // Score variant output.
+    ScoreVariantOutput output = 1;
+
+    // Tensor chunk related to the last score variant output message received.
+    // It is an error to receive a tensor chunk without a previous output
+    // message.
+    TensorChunk tensor_chunk = 2;
+  }
+}
+
+// Request message for getting metadata for an organism.
+message MetadataRequest {
+  // The organism we should return metadata for.
+  Organism organism = 1;
+}
+
+// Response message for getting metadata for an organism.
+message MetadataResponse {
+  // The metadata for each output type.
+  repeated OutputMetadata output_metadata = 1;
+}

flax_model/alphagenome/_sdk/protos/dna_model_service_pb2.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: alphagenome/protos/dna_model_service.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from flax_model.alphagenome._sdk.protos import dna_model_pb2 as alphagenome_dot_protos_dot_dna__model__pb2
+from flax_model.alphagenome._sdk.protos import tensor_pb2 as alphagenome_dot_protos_dot_tensor__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n*alphagenome/protos/dna_model_service.proto\x12(google.gdm.gdmscience.alphagenome.v1main\x1a\"alphagenome/protos/dna_model.proto\x1a\x1f\x61lphagenome/protos/tensor.proto\"\xa8\x02\n\x16PredictSequenceRequest\x12\x10\n\x08sequence\x18\x01 \x01(\t\x12\x44\n\x08organism\x18\x02 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\x12N\n\x0eontology_terms\x18\x03 \x03(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.OntologyTerm\x12O\n\x11requested_outputs\x18\x04 \x03(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\x12\x15\n\rmodel_version\x18\x05 \x01(\t\"\xb7\x01\n\x17PredictSequenceResponse\x12\x42\n\x06output\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.OutputH\x00\x12M\n\x0ctensor_chunk\x18\x02 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x42\t\n\x07payload\"\xdc\x02\n\x16PredictIntervalRequest\x12\x44\n\x08interval\x18\x01 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12\x44\n\x08organism\x18\x02 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\x12O\n\x11requested_outputs\x18\x03 \x03(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\x12N\n\x0eontology_terms\x18\x04 \x03(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.OntologyTerm\x12\x15\n\rmodel_version\x18\x05 \x01(\t\"\xb7\x01\n\x17PredictIntervalResponse\x12\x42\n\x06output\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.OutputH\x00\x12M\n\x0ctensor_chunk\x18\x02 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x42\t\n\x07payload\"\x9f\x03\n\x15PredictVariantRequest\x12\x44\n\x08interval\x18\x01 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12\x42\n\x07variant\x18\x02 \x01(\x0b\x32\x31.google.gdm.gdmscience.alphagenome.v1main.Variant\x12\x44\n\x08organism\x18\x03 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\x12O\n\x11requested_outputs\x18\x04 \x03(\x0e\x32\x34.google.gdm.gdmscience.alphagenome.v1main.OutputType\x12N\n\x0eontology_terms\x18\x05 \x03(\x0b\x32\x36.google.gdm.gdmscience.alphagenome.v1main.OntologyTerm\x12\x15\n\rmodel_version\x18\x06 \x01(\t\"\x8e\x02\n\x16PredictVariantResponse\x12L\n\x10reference_output\x18\x01 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.OutputH\x00\x12L\n\x10\x61lternate_output\x18\x02 \x01(\x0b\x32\x30.google.gdm.gdmscience.alphagenome.v1main.OutputH\x00\x12M\n\x0ctensor_chunk\x18\x03 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x42\t\n\x07payload\"\x8d\x02\n\x14ScoreIntervalRequest\x12\x44\n\x08interval\x18\x01 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12\x44\n\x08organism\x18\x02 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\x12R\n\x10interval_scorers\x18\x03 \x03(\x0b\x32\x38.google.gdm.gdmscience.alphagenome.v1main.IntervalScorer\x12\x15\n\rmodel_version\x18\x04 \x01(\t\"\xc2\x01\n\x15ScoreIntervalResponse\x12O\n\x06output\x18\x01 \x01(\x0b\x32=.google.gdm.gdmscience.alphagenome.v1main.ScoreIntervalOutputH\x00\x12M\n\x0ctensor_chunk\x18\x02 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x42\t\n\x07payload\"\xce\x02\n\x13ScoreVariantRequest\x12\x44\n\x08interval\x18\x01 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12\x42\n\x07variant\x18\x02 \x01(\x0b\x32\x31.google.gdm.gdmscience.alphagenome.v1main.Variant\x12\x44\n\x08organism\x18\x03 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\x12P\n\x0fvariant_scorers\x18\x04 \x03(\x0b\x32\x37.google.gdm.gdmscience.alphagenome.v1main.VariantScorer\x12\x15\n\rmodel_version\x18\x05 \x01(\t\"\xc0\x01\n\x14ScoreVariantResponse\x12N\n\x06output\x18\x01 \x01(\x0b\x32<.google.gdm.gdmscience.alphagenome.v1main.ScoreVariantOutputH\x00\x12M\n\x0ctensor_chunk\x18\x02 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x42\t\n\x07payload\"\xbe\x03\n\x16ScoreIsmVariantRequest\x12\x44\n\x08interval\x18\x01 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12H\n\x0cism_interval\x18\x02 \x01(\x0b\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Interval\x12\x44\n\x08organism\x18\x03 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\x12P\n\x0fvariant_scorers\x18\x04 \x03(\x0b\x32\x37.google.gdm.gdmscience.alphagenome.v1main.VariantScorer\x12P\n\x10interval_variant\x18\x06 \x01(\x0b\x32\x31.google.gdm.gdmscience.alphagenome.v1main.VariantH\x00\x88\x01\x01\x12\x15\n\rmodel_version\x18\x05 \x01(\tB\x13\n\x11_interval_variant\"\xc3\x01\n\x17ScoreIsmVariantResponse\x12N\n\x06output\x18\x01 \x01(\x0b\x32<.google.gdm.gdmscience.alphagenome.v1main.ScoreVariantOutputH\x00\x12M\n\x0ctensor_chunk\x18\x02 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x42\t\n\x07payload\"W\n\x0fMetadataRequest\x12\x44\n\x08organism\x18\x01 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.Organism\"e\n\x10MetadataResponse\x12Q\n\x0foutput_metadata\x18\x01 \x03(\x0b\x32\x38.google.gdm.gdmscience.alphagenome.v1main.OutputMetadata2\xc4\x08\n\x0f\x44naModelService\x12\x9c\x01\n\x0fPredictSequence\x12@.google.gdm.gdmscience.alphagenome.v1main.PredictSequenceRequest\x1a\x41.google.gdm.gdmscience.alphagenome.v1main.PredictSequenceResponse\"\x00(\x01\x30\x01\x12\x9c\x01\n\x0fPredictInterval\x12@.google.gdm.gdmscience.alphagenome.v1main.PredictIntervalRequest\x1a\x41.google.gdm.gdmscience.alphagenome.v1main.PredictIntervalResponse\"\x00(\x01\x30\x01\x12\x99\x01\n\x0ePredictVariant\x12?.google.gdm.gdmscience.alphagenome.v1main.PredictVariantRequest\x1a@.google.gdm.gdmscience.alphagenome.v1main.PredictVariantResponse\"\x00(\x01\x30\x01\x12\x96\x01\n\rScoreInterval\x12>.google.gdm.gdmscience.alphagenome.v1main.ScoreIntervalRequest\x1a?.google.gdm.gdmscience.alphagenome.v1main.ScoreIntervalResponse\"\x00(\x01\x30\x01\x12\x93\x01\n\x0cScoreVariant\x12=.google.gdm.gdmscience.alphagenome.v1main.ScoreVariantRequest\x1a>.google.gdm.gdmscience.alphagenome.v1main.ScoreVariantResponse\"\x00(\x01\x30\x01\x12\x9c\x01\n\x0fScoreIsmVariant\x12@.google.gdm.gdmscience.alphagenome.v1main.ScoreIsmVariantRequest\x1a\x41.google.gdm.gdmscience.alphagenome.v1main.ScoreIsmVariantResponse\"\x00(\x01\x30\x01\x12\x88\x01\n\x0bGetMetadata\x12\x39.google.gdm.gdmscience.alphagenome.v1main.MetadataRequest\x1a:.google.gdm.gdmscience.alphagenome.v1main.MetadataResponse\"\x00\x30\x01\x42\x9b\x01\n,com.google.gdm.gdmscience.alphagenome.v1mainB\x14\x44naModelServiceProtoP\x01ZSgoogle.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenomeb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flax_model.alphagenome._sdk.protos.dna_model_service_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  _globals['DESCRIPTOR']._options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n,com.google.gdm.gdmscience.alphagenome.v1mainB\024DnaModelServiceProtoP\001ZSgoogle.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenome'
+  _globals['_PREDICTSEQUENCEREQUEST']._serialized_start=158
+  _globals['_PREDICTSEQUENCEREQUEST']._serialized_end=454
+  _globals['_PREDICTSEQUENCERESPONSE']._serialized_start=457
+  _globals['_PREDICTSEQUENCERESPONSE']._serialized_end=640
+  _globals['_PREDICTINTERVALREQUEST']._serialized_start=643
+  _globals['_PREDICTINTERVALREQUEST']._serialized_end=991
+  _globals['_PREDICTINTERVALRESPONSE']._serialized_start=994
+  _globals['_PREDICTINTERVALRESPONSE']._serialized_end=1177
+  _globals['_PREDICTVARIANTREQUEST']._serialized_start=1180
+  _globals['_PREDICTVARIANTREQUEST']._serialized_end=1595
+  _globals['_PREDICTVARIANTRESPONSE']._serialized_start=1598
+  _globals['_PREDICTVARIANTRESPONSE']._serialized_end=1868
+  _globals['_SCOREINTERVALREQUEST']._serialized_start=1871
+  _globals['_SCOREINTERVALREQUEST']._serialized_end=2140
+  _globals['_SCOREINTERVALRESPONSE']._serialized_start=2143
+  _globals['_SCOREINTERVALRESPONSE']._serialized_end=2337
+  _globals['_SCOREVARIANTREQUEST']._serialized_start=2340
+  _globals['_SCOREVARIANTREQUEST']._serialized_end=2674
+  _globals['_SCOREVARIANTRESPONSE']._serialized_start=2677
+  _globals['_SCOREVARIANTRESPONSE']._serialized_end=2869
+  _globals['_SCOREISMVARIANTREQUEST']._serialized_start=2872
+  _globals['_SCOREISMVARIANTREQUEST']._serialized_end=3318
+  _globals['_SCOREISMVARIANTRESPONSE']._serialized_start=3321
+  _globals['_SCOREISMVARIANTRESPONSE']._serialized_end=3516
+  _globals['_METADATAREQUEST']._serialized_start=3518
+  _globals['_METADATAREQUEST']._serialized_end=3605
+  _globals['_METADATARESPONSE']._serialized_start=3607
+  _globals['_METADATARESPONSE']._serialized_end=3708
+  _globals['_DNAMODELSERVICE']._serialized_start=3711
+  _globals['_DNAMODELSERVICE']._serialized_end=4803
+# @@protoc_insertion_point(module_scope)

flax_model/alphagenome/_sdk/protos/dna_model_service_pb2_grpc.py

@@ -0,0 +1,274 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+from flax_model.alphagenome._sdk.protos import dna_model_service_pb2 as alphagenome_dot_protos_dot_dna__model__service__pb2
+
+
+class DnaModelServiceStub(object):
+    """Service for making predictions with DNA models.
+    """
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.PredictSequence = channel.stream_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/PredictSequence',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictSequenceRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictSequenceResponse.FromString,
+                )
+        self.PredictInterval = channel.stream_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/PredictInterval',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictIntervalRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictIntervalResponse.FromString,
+                )
+        self.PredictVariant = channel.stream_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/PredictVariant',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictVariantRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictVariantResponse.FromString,
+                )
+        self.ScoreInterval = channel.stream_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/ScoreInterval',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIntervalRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIntervalResponse.FromString,
+                )
+        self.ScoreVariant = channel.stream_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/ScoreVariant',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreVariantRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreVariantResponse.FromString,
+                )
+        self.ScoreIsmVariant = channel.stream_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/ScoreIsmVariant',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIsmVariantRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIsmVariantResponse.FromString,
+                )
+        self.GetMetadata = channel.unary_stream(
+                '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/GetMetadata',
+                request_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.MetadataRequest.SerializeToString,
+                response_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.MetadataResponse.FromString,
+                )
+
+
+class DnaModelServiceServicer(object):
+    """Service for making predictions with DNA models.
+    """
+
+    def PredictSequence(self, request_iterator, context):
+        """Makes a prediction for DNA sequence.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def PredictInterval(self, request_iterator, context):
+        """Make prediction for a single genomic interval.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def PredictVariant(self, request_iterator, context):
+        """Make variant effect predictions for a genomic interval.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def ScoreInterval(self, request_iterator, context):
+        """Score a genomic interval.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def ScoreVariant(self, request_iterator, context):
+        """Score a variant for a genomic interval.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def ScoreIsmVariant(self, request_iterator, context):
+        """Score ISM variant effect predictions for a genomic interval.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetMetadata(self, request, context):
+        """Get metadata for the model.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_DnaModelServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'PredictSequence': grpc.stream_stream_rpc_method_handler(
+                    servicer.PredictSequence,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictSequenceRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictSequenceResponse.SerializeToString,
+            ),
+            'PredictInterval': grpc.stream_stream_rpc_method_handler(
+                    servicer.PredictInterval,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictIntervalRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictIntervalResponse.SerializeToString,
+            ),
+            'PredictVariant': grpc.stream_stream_rpc_method_handler(
+                    servicer.PredictVariant,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictVariantRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.PredictVariantResponse.SerializeToString,
+            ),
+            'ScoreInterval': grpc.stream_stream_rpc_method_handler(
+                    servicer.ScoreInterval,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIntervalRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIntervalResponse.SerializeToString,
+            ),
+            'ScoreVariant': grpc.stream_stream_rpc_method_handler(
+                    servicer.ScoreVariant,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreVariantRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreVariantResponse.SerializeToString,
+            ),
+            'ScoreIsmVariant': grpc.stream_stream_rpc_method_handler(
+                    servicer.ScoreIsmVariant,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIsmVariantRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIsmVariantResponse.SerializeToString,
+            ),
+            'GetMetadata': grpc.unary_stream_rpc_method_handler(
+                    servicer.GetMetadata,
+                    request_deserializer=alphagenome_dot_protos_dot_dna__model__service__pb2.MetadataRequest.FromString,
+                    response_serializer=alphagenome_dot_protos_dot_dna__model__service__pb2.MetadataResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'google.gdm.gdmscience.alphagenome.v1main.DnaModelService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class DnaModelService(object):
+    """Service for making predictions with DNA models.
+    """
+
+    @staticmethod
+    def PredictSequence(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/PredictSequence',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.PredictSequenceRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.PredictSequenceResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def PredictInterval(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/PredictInterval',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.PredictIntervalRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.PredictIntervalResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def PredictVariant(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/PredictVariant',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.PredictVariantRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.PredictVariantResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def ScoreInterval(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/ScoreInterval',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIntervalRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIntervalResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def ScoreVariant(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/ScoreVariant',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreVariantRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreVariantResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def ScoreIsmVariant(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(request_iterator, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/ScoreIsmVariant',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIsmVariantRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.ScoreIsmVariantResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def GetMetadata(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_stream(request, target, '/google.gdm.gdmscience.alphagenome.v1main.DnaModelService/GetMetadata',
+            alphagenome_dot_protos_dot_dna__model__service__pb2.MetadataRequest.SerializeToString,
+            alphagenome_dot_protos_dot_dna__model__service__pb2.MetadataResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

flax_model/alphagenome/_sdk/protos/tensor.proto

@@ -0,0 +1,104 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package google.gdm.gdmscience.alphagenome.v1main;
+
+option go_package = "google.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenome";
+option java_multiple_files = true;
+option java_outer_classname = "TensorProto";
+option java_package = "com.google.gdm.gdmscience.alphagenome.v1main";
+
+// Protocol buffer representing an arbitrarily large multi-dimensional array of
+// data, laid out in row-major format. To support tensors larger than the
+// maximum protocol buffer message size, the tensor payload may be split into
+// multiple chunks.
+message Tensor {
+  // The shape of the tensor. If empty, the data will be treated as a scalar.
+  repeated int32 shape = 1;
+
+  // Data type for the elements in this tensor.
+  DataType data_type = 2;
+
+  // The payload for the tensor.
+  oneof payload {
+    // The raw data for the tensor. If present, the tensor is not split into
+    // chunks.
+    TensorChunk array = 3;
+
+    // The number of chunks the tensor is split into.
+    int64 chunk_count = 4;
+  }
+}
+
+// A single chunk of a tensor.
+message TensorChunk {
+  // Flattened tensor data. Data is laid out in row-major order.
+  bytes data = 1;
+
+  // How the data is compressed. Compression is applied to each chunk
+  // independently.
+  CompressionType compression_type = 2;
+}
+
+// The data type of the tensor.
+enum DataType {
+  // Unspecified data type.
+  DATA_TYPE_UNSPECIFIED = 0;
+
+  // 16-bit "Brain Floating Point". See
+  // https://en.wikipedia.org/wiki/Bfloat16_floating-point_format for more
+  // details.
+  DATA_TYPE_BFLOAT16 = 1;
+
+  // 16-bit floating point.
+  DATA_TYPE_FLOAT16 = 11;
+
+  // 32-bit floating point.
+  DATA_TYPE_FLOAT32 = 2;
+
+  // 64-bit floating point.
+  DATA_TYPE_FLOAT64 = 3;
+
+  // 8-bit signed integer.
+  DATA_TYPE_INT8 = 4;
+
+  // 32-bit signed integer.
+  DATA_TYPE_INT32 = 5;
+
+  // 64-bit signed integer.
+  DATA_TYPE_INT64 = 6;
+
+  // 8-bit unsigned integer.
+  DATA_TYPE_UINT8 = 7;
+
+  // 32-bit unsigned integer.
+  DATA_TYPE_UINT32 = 8;
+
+  // 64-bit unsigned integer.
+  DATA_TYPE_UINT64 = 9;
+
+  // 8-bit boolean.
+  DATA_TYPE_BOOL = 10;
+}
+
+// Compression type for the tensor data.
+enum CompressionType {
+  // No compression.
+  COMPRESSION_TYPE_NONE = 0;
+
+  // ZSTD compression.
+  COMPRESSION_TYPE_ZSTD = 1;
+}

flax_model/alphagenome/_sdk/protos/tensor_pb2.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: alphagenome/protos/tensor.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1f\x61lphagenome/protos/tensor.proto\x12(google.gdm.gdmscience.alphagenome.v1main\"\xc8\x01\n\x06Tensor\x12\r\n\x05shape\x18\x01 \x03(\x05\x12\x45\n\tdata_type\x18\x02 \x01(\x0e\x32\x32.google.gdm.gdmscience.alphagenome.v1main.DataType\x12\x46\n\x05\x61rray\x18\x03 \x01(\x0b\x32\x35.google.gdm.gdmscience.alphagenome.v1main.TensorChunkH\x00\x12\x15\n\x0b\x63hunk_count\x18\x04 \x01(\x03H\x00\x42\t\n\x07payload\"p\n\x0bTensorChunk\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\x12S\n\x10\x63ompression_type\x18\x02 \x01(\x0e\x32\x39.google.gdm.gdmscience.alphagenome.v1main.CompressionType*\x95\x02\n\x08\x44\x61taType\x12\x19\n\x15\x44\x41TA_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12\x44\x41TA_TYPE_BFLOAT16\x10\x01\x12\x15\n\x11\x44\x41TA_TYPE_FLOAT16\x10\x0b\x12\x15\n\x11\x44\x41TA_TYPE_FLOAT32\x10\x02\x12\x15\n\x11\x44\x41TA_TYPE_FLOAT64\x10\x03\x12\x12\n\x0e\x44\x41TA_TYPE_INT8\x10\x04\x12\x13\n\x0f\x44\x41TA_TYPE_INT32\x10\x05\x12\x13\n\x0f\x44\x41TA_TYPE_INT64\x10\x06\x12\x13\n\x0f\x44\x41TA_TYPE_UINT8\x10\x07\x12\x14\n\x10\x44\x41TA_TYPE_UINT32\x10\x08\x12\x14\n\x10\x44\x41TA_TYPE_UINT64\x10\t\x12\x12\n\x0e\x44\x41TA_TYPE_BOOL\x10\n*G\n\x0f\x43ompressionType\x12\x19\n\x15\x43OMPRESSION_TYPE_NONE\x10\x00\x12\x19\n\x15\x43OMPRESSION_TYPE_ZSTD\x10\x01\x42\x92\x01\n,com.google.gdm.gdmscience.alphagenome.v1mainB\x0bTensorProtoP\x01ZSgoogle.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenomeb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flax_model.alphagenome._sdk.protos.tensor_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  _globals['DESCRIPTOR']._options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n,com.google.gdm.gdmscience.alphagenome.v1mainB\013TensorProtoP\001ZSgoogle.golang.org/genproto/googleapis/gdm/gdmscience/alphagenome/v1main;alphagenome'
+  _globals['_DATATYPE']._serialized_start=395
+  _globals['_DATATYPE']._serialized_end=672
+  _globals['_COMPRESSIONTYPE']._serialized_start=674
+  _globals['_COMPRESSIONTYPE']._serialized_end=745
+  _globals['_TENSOR']._serialized_start=78
+  _globals['_TENSOR']._serialized_end=278
+  _globals['_TENSORCHUNK']._serialized_start=280
+  _globals['_TENSORCHUNK']._serialized_end=392
+# @@protoc_insertion_point(module_scope)

flax_model/alphagenome/_sdk/protos/tensor_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

flax_model/alphagenome/_sdk/tensor_utils.py

@@ -0,0 +1,171 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for converting NumPy arrays to Tensor protocol buffers."""
+
+from collections.abc import Iterable, Sequence
+
+from flax_model.alphagenome._sdk.protos import tensor_pb2
+import immutabledict
+import ml_dtypes
+import numpy as np
+import zstandard
+
+
+_TENSOR_DTYPE_TO_NUMPY_DTYPE = immutabledict.immutabledict({
+    tensor_pb2.DataType.DATA_TYPE_BFLOAT16: np.dtype(ml_dtypes.bfloat16),
+    tensor_pb2.DataType.DATA_TYPE_FLOAT16: np.dtype(np.float16),
+    tensor_pb2.DataType.DATA_TYPE_FLOAT32: np.dtype(np.float32),
+    tensor_pb2.DataType.DATA_TYPE_FLOAT64: np.dtype(np.float64),
+    tensor_pb2.DataType.DATA_TYPE_INT8: np.dtype(np.int8),
+    tensor_pb2.DataType.DATA_TYPE_INT32: np.dtype(np.int32),
+    tensor_pb2.DataType.DATA_TYPE_INT64: np.dtype(np.int64),
+    tensor_pb2.DataType.DATA_TYPE_UINT8: np.dtype(np.uint8),
+    tensor_pb2.DataType.DATA_TYPE_UINT32: np.dtype(np.uint32),
+    tensor_pb2.DataType.DATA_TYPE_UINT64: np.dtype(np.uint64),
+    tensor_pb2.DataType.DATA_TYPE_BOOL: np.dtype(bool),
+})
+
+_NUMPY_DTYPE_TO_TENSOR_DTYPE = immutabledict.immutabledict(
+    {value: key for key, value in _TENSOR_DTYPE_TO_NUMPY_DTYPE.items()}
+)
+
+
+def _compress_bytes(
+    array: np.ndarray, compression_type: tensor_pb2.CompressionType
+):
+  """Compresses a c-contiguous array to the specified compression type."""
+  assert array.flags.c_contiguous
+  array = array.view(np.uint8)
+  match compression_type:
+    case tensor_pb2.CompressionType.COMPRESSION_TYPE_ZSTD:
+      return zstandard.compress(array.data)
+    case tensor_pb2.CompressionType.COMPRESSION_TYPE_NONE:
+      return bytes(array.data)
+
+
+def _decompress_bytes(
+    data: bytes, compression_type: tensor_pb2.CompressionType
+):
+  """Decompress bytes using the specified compression type."""
+  match compression_type:
+    case tensor_pb2.CompressionType.COMPRESSION_TYPE_ZSTD:
+      return zstandard.decompress(data)
+    case tensor_pb2.CompressionType.COMPRESSION_TYPE_NONE:
+      return data
+
+
+def pack_tensor(
+    value: ...,
+    *,
+    bytes_per_chunk: int = 0,
+    compression_type: tensor_pb2.CompressionType = (
+        tensor_pb2.CompressionType.COMPRESSION_TYPE_NONE
+    ),
+) -> tuple[tensor_pb2.Tensor, Sequence[tensor_pb2.TensorChunk]]:
+  """Encodes the value as a Tensor and optional sequence of chunks.
+
+  Args:
+    value: An array-like object to pack. For example, scalar (float, int, bool,
+      etc.), NumPy array, or nested lists of scalars.
+    bytes_per_chunk: The number of bytes to include in each chunk. If 0, the
+      entire value will be packed into the Tensor proto, otherwise the value
+      will be split into chunks of this size.
+    compression_type: The type of compression to apply to the data. This is
+      applied to each chunk separately.
+
+  Returns:
+    Tuple of Tensor protocol buffer and, if items_per_chunk is greater than 0, a
+    sequence of TensorChunk protos.
+  """
+  packed = tensor_pb2.Tensor()
+  value = np.ascontiguousarray(value)
+
+  packed.shape[:] = value.shape
+  packed.data_type = _NUMPY_DTYPE_TO_TENSOR_DTYPE[value.dtype]
+
+  chunks = []
+  if bytes_per_chunk > 0:
+    items_per_chunk = bytes_per_chunk // value.itemsize
+    if bytes_per_chunk < value.itemsize:
+      raise ValueError(f'{bytes_per_chunk=} must be >= {value.itemsize=}.')
+    for chunk in np.split(
+        value.ravel(), range(items_per_chunk, value.size, items_per_chunk)
+    ):
+      chunks.append(
+          tensor_pb2.TensorChunk(
+              data=_compress_bytes(chunk, compression_type),
+              compression_type=compression_type,
+          )
+      )
+    packed.chunk_count = len(chunks)
+  else:
+    packed.array.data = _compress_bytes(value, compression_type)
+    packed.array.compression_type = compression_type
+
+  return packed, chunks
+
+
+def unpack_proto(
+    proto: tensor_pb2.Tensor,
+    chunks: Iterable[tensor_pb2.TensorChunk] = (),
+) -> np.ndarray:
+  """Converts a Tensor proto and any chunks into a NumPy array.
+
+  Args:
+    proto: Tensor proto to unpack.
+    chunks: Optional sequence of TensorChunk protos to unpack.
+
+  Returns:
+    NumPy array of the unpacked data.
+  """
+  dtype = _TENSOR_DTYPE_TO_NUMPY_DTYPE[proto.data_type]
+  match proto.WhichOneof('payload'):
+    case 'array':
+      data = _decompress_bytes(proto.array.data, proto.array.compression_type)
+      array = np.frombuffer(data, dtype=dtype).reshape(proto.shape)
+    case 'chunk_count':
+      array = np.empty(np.prod(proto.shape) * dtype.itemsize, dtype=np.uint8)
+      bytes_received = 0
+      for chunk in chunks:
+        chunk_data = np.frombuffer(
+            _decompress_bytes(chunk.data, chunk.compression_type),
+            dtype=np.uint8,
+        )
+        array[bytes_received : bytes_received + chunk_data.nbytes] = chunk_data
+        bytes_received += chunk_data.nbytes
+
+      if bytes_received != array.nbytes:
+        raise ValueError(
+            f'Expected {array.nbytes} bytes but only received {bytes_received} '
+            'bytes.'
+        )
+      array = array.view(dtype).reshape(proto.shape)
+    case _:
+      raise ValueError(
+          f'Unsupported payload type: {proto.WhichOneof("payload")}'
+      )
+
+  return array
+
+
+def upcast_floating(x: np.ndarray) -> np.ndarray:
+  """Helper to upcast low-precision floating point arrays to float32."""
+  dtype = np.result_type(x)
+  if (
+      np.issubdtype(dtype, np.floating) or dtype == ml_dtypes.bfloat16
+  ) and dtype.itemsize < 4:
+    return x.astype(np.float32)
+  else:
+    return x

flax_model/alphagenome/_sdk/typing.py

@@ -0,0 +1,39 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for type annotations."""
+
+import importlib.metadata
+from typing import TypeVar
+
+import jaxtyping
+import typeguard
+
+_T = TypeVar('_T')
+
+
+def jaxtyped(fn: _T) -> _T:
+  """Wrapper around jaxtyping.jaxtyped that uses typeguard iff typeguard < 3."""
+  try:
+    major, *_ = importlib.metadata.version('typeguard').split('.')
+  except importlib.metadata.PackageNotFoundError:
+    major = -1
+
+  # Only use jaxtyping if typeguard is < 3. See
+  # https://docs.kidger.site/jaxtyping/api/runtime-type-checking/#runtime-type-checking
+  # for more details.
+  if int(major) < 3:
+    return jaxtyping.jaxtyped(fn, typechecker=typeguard.typechecked)
+  else:
+    return fn

flax_model/alphagenome/_sdk/visualization/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Library of tools for visualizing with genomic model predictions."""

flax_model/alphagenome/_sdk/visualization/plot.py

@@ -0,0 +1,569 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Plotting functions."""
+
+from collections.abc import Callable, Mapping, Sequence
+import functools
+import math
+from typing import Any, Literal
+
+from absl import logging
+from flax_model.alphagenome._sdk.data import genome
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+
+
+@functools.lru_cache(maxsize=8)
+def _get_text_path(
+    letter: str, letter_colors_scheme: str
+) -> tuple[mpl.text.TextPath, str]:
+  """Returns a memoized TextPath and color for the given letter and scheme."""
+  letters_width_colors = dict(
+      # Consistent with `classic` scheme from
+      # https://github.com/gecrooks/weblogo/blob/master/weblogo/logo.py#L136
+      default={
+          'A': (0.35, 'darkgreen'),
+          'C': (0.366, 'blue'),
+          'G': (0.384, 'orange'),
+          'T': (0.305, 'red'),
+      },
+      # Consistent with factorbook scheme from https://www.factorbook.org/
+      factorbook={
+          'A': (0.305, 'red'),
+          'C': (0.366, 'blue'),
+          'G': (0.384, 'orange'),
+          'T': (0.35, 'darkgreen'),
+      },
+  )
+  if (
+      letters_width_color := letters_width_colors.get(letter_colors_scheme)
+  ) is None:
+    raise ValueError(
+        f'letter_colors_scheme must be one of {letters_width_colors.keys()}.'
+        f'Got "{letter_colors_scheme}".'
+    )
+  else:
+    if (letter_width_color := letters_width_color.get(letter)) is None:
+      raise ValueError(
+          f'letter must be one of {letters_width_color.keys()}. Got "{letter}".'
+      )
+    else:
+      letter_width, letter_color = letter_width_color
+      font = mpl.font_manager.FontProperties(weight='bold')
+      return (
+          mpl.text.TextPath((-letter_width, 0), letter, size=1, prop=font),
+          letter_color,
+      )
+
+
+def seqlogo(
+    letter_heights: np.ndarray,
+    alphabet: str = 'ACGT',
+    one_based: bool = True,
+    start: int = 0,
+    ax: mpl.axes.Axes | None = None,
+    letter_colors: Literal['default', 'factorbook'] = 'default',
+):
+  """Sequence logo plot.
+
+  Args:
+    letter_heights: 2D array with shape (sequence length, len(alphabet)) shape.
+      Positive values indicate that a logo element extends above the horizontal
+      axis, and negative values indicate that the element extends below.
+    alphabet: Alphabet corresponding to the channel axis of `letter_heights`.
+    one_based: If True, plot letters at 1-based coordinates: first letter will
+      be centered at 1. If False, the first letter will be between 0 and 1.
+    start: Interval start.
+    ax: matplotlib axis to add figures to.
+    letter_colors: Name of the color scheme to use for coloring each letter.
+      Must be one of 'default', a scheme consistent with weblogo
+      (https://github.com/gecrooks/weblogo/blob/master/weblogo/logo.py#L136) or
+      'factorbook', a scheme consistent with Factorbook
+      (https://www.factorbook.org/).
+  """
+  ax = ax if ax is not None else plt.gca()
+
+  if letter_heights.ndim != 2:
+    raise ValueError(
+        'Expecting a 2D matrix of shape (sequence_length, len(alphabed))'
+    )
+  if letter_heights.shape[1] != len(alphabet):
+    raise ValueError('Last axis needs to match len(alphabet)')
+
+  globscale = 1.35
+  paths = []
+  facecolors = []
+
+  for x_offset, heights in enumerate(letter_heights):
+    last_positive_y = 0.0
+    last_negative_y = 0.0
+    x = start + x_offset + 0.5 + 0.5 * one_based
+
+    for height, letter in sorted(zip(heights, alphabet)):
+      if height == 0:
+        continue
+      # Start with lowest height and keep track of the last y coordinate.
+      if height > 0:
+        y_position = last_positive_y
+        last_positive_y += height
+      else:
+        y_position = last_negative_y
+        last_negative_y += height
+
+      base_path, letter_color = _get_text_path(letter, letter_colors)
+
+      transform = (
+          mpl.transforms.Affine2D()
+          .scale(globscale, height * globscale)
+          .translate(x, y_position)
+      )
+      paths.append(base_path.transformed(transform))
+      facecolors.append(letter_color)
+
+  if paths:
+    collection = mpl.collections.PathCollection(
+        paths,
+        facecolors=facecolors,
+        edgecolors='none',
+        linewidths=0,
+        transform=ax.transData,
+    )
+    ax.add_collection(collection)
+
+  # Make sure all letters are displayed.
+  ax.autoscale_view()
+
+
+def plot_contact_map(
+    contact_map: pd.DataFrame,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    square: bool = True,
+    cbar_shrink: float = 0.4,
+    ax=None,
+    **kwargs,
+):
+  """Visualize a contact map.
+
+  A contact map is a 2D array of values, where each value represents the
+  probability that two DNA bases are in contact.
+
+  The array is symmetric, i.e., the (i, j) values and the (j, i) values of
+  `contact_map` are equal.
+
+  Args:
+    contact_map: Contact map to visualize.
+    vmin: Minimum value for the colorbar.
+    vmax: Maximum value for the colorbar.
+    square: If `True`, plots the contact map as a square. If `False`, plots the
+      contact map as a rectangle with a shorter height than width.
+    cbar_shrink: Fraction by which to multiply the size of the colorbar.
+    ax: Matplotlib axis to add the heatmap to.
+    **kwargs: Additional keyword arguments passed to `sns.heatmap`.
+  """
+
+  def tuple2string(interval_tuple):
+    chromosome, start, end = interval_tuple
+    return f'{chromosome}:{start:,}-{end:,}'
+
+  if isinstance(contact_map, pd.DataFrame):
+    interval_string = (
+        tuple2string(contact_map.index[0])
+        + ' - '
+        + tuple2string(contact_map.index[-1])
+    )
+    contact_map = contact_map.values
+  else:
+    interval_string = ''
+
+  if vmin is None:
+    vmin = np.nanmin(contact_map[contact_map > 0])
+  if vmax is None:
+    vmax = np.nanmax(contact_map)
+  fall_cmap = mpl.colors.LinearSegmentedColormap.from_list(
+      'fall',
+      colors=[
+          'white',
+          (245 / 256, 166 / 256, 35 / 256),
+          (208 / 256, 2 / 256, 27 / 256),
+          'black',
+      ],
+  )
+  log_norm = mpl.colors.LogNorm(vmin=vmin, vmax=vmax)
+  cbar_min = math.floor(math.log10(log_norm.vmin))
+  cbar_max = 1 + math.ceil(math.log10(log_norm.vmax))
+  cbar_ticks = [math.pow(10, i) for i in range(cbar_min, cbar_max)]
+
+  if ax is None:
+    ax = plt.gca()
+  sns.heatmap(
+      contact_map + 1e-10,
+      cmap=fall_cmap,
+      norm=log_norm,
+      cbar_kws={'ticks': cbar_ticks, 'shrink': cbar_shrink},
+      vmin=log_norm.vmin,
+      vmax=log_norm.vmax,
+      xticklabels=False,
+      yticklabels=False,
+      square=square,
+      ax=ax,
+      **kwargs,
+  )
+
+  ax.set_xlabel(interval_string)
+
+
+def plot_track(
+    arr: np.ndarray,
+    ax: plt.Axes,
+    x: np.ndarray | None = None,
+    legend: bool = False,
+    ylim: float | None = None,
+    color: str | Sequence[str] | None = None,
+    filled: bool = False,
+) -> None:
+  """Plot a single track on the axis after inferring track type from array.
+
+  This function infers the type of track plot depending on channel
+  dimensionality of `arr`:
+
+    - If `arr` is a single-dimensional array, plot a single track line.
+    - If `arr` is a 2D array, and interpreting the second dimension as the
+      channel dimension:
+
+      - Plot a sequence logo if the number of channels is 4 (which suggests one
+        channel per DNA base).
+      - Plot two overlapping line plots if the number of channels is 2 (which
+        suggests one channel per DNA strand).
+
+  Args:
+    arr: array of values to plot.
+    ax: matplotlib axis to use for plotting the values on.
+    x: optional array of values to use for setting the x axis values. If absent,
+      then sequential values starting from 1 to the length of the array arr will
+      be used.
+    legend: whether to draw a legend or not on the double stranded DNA plot.
+    ylim: y axis limit.
+    color: color(s) used for line plots.
+    filled: whether to fill the space between the x axis and the line (or leave
+      it as whitespace).
+
+  Returns:
+    None. Adds a plot to the provided axis.
+  """
+  sequence_length = len(arr)
+
+  if x is None:
+    x = np.arange(1, sequence_length + 1)  # One-based indexing.
+
+  if arr.ndim == 1 or arr.shape[1] == 1:
+    if color is not None:
+      if isinstance(color, (list, tuple)):
+        logging.warning(
+            'A sequence of colors was passed to plot_track but the second '
+            'array dimension is 1 suggesting single stranded DNA. Using the '
+            'first entry %s as the line color.',
+            color[0],
+        )
+        color = color[0]
+
+    # In the case of a boolean array, we fill between values and the x axis and
+    # simplify the y axis components.
+    if arr.dtype == bool:
+      ax.fill_between(x, arr, step='mid', color=color)
+      ax.yaxis.set_ticks_position('none')
+      ax.set_yticklabels([])
+      ax.spines['left'].set_visible(False)
+      ax.set_ylim([0, 1.5])
+
+    else:
+      if filled:
+        ax.fill_between(x, np.ravel(arr), color=color)
+      ax.plot(x, np.ravel(arr), color=color)
+
+  elif arr.shape[1] == 4:
+    seqlogo(arr, ax=ax)
+
+  elif arr.shape[1] == 2:
+    if color is not None:
+      if not isinstance(color, Sequence) or len(color) != 2:
+        raise ValueError(
+            'Must pass sequence of 2 colors if plotting 2 dimensional array.'
+        )
+      color_1, color_2 = color
+    else:
+      color_1, color_2 = None, None
+    ax.plot(x, arr[:, 0], label='pos', color=color_1)
+    ax.plot(x, arr[:, 1], label='neg', color=color_2)
+    if legend:
+      ax.legend()
+  else:
+    raise ValueError(
+        f'Do not know how to plot array with shape[1] != {arr.shape[1]}. '
+        'Valid values are: 1, 2, or 4.'
+    )
+  if ylim is not None:
+    ax.set_ylim(ylim)
+
+
+def plot_tracks(
+    tracks: Mapping[str, Any],
+    x: np.ndarray | None = None,
+    title: str | None = None,
+    legend: bool = False,
+    fig_width: float = 20,
+    fig_track_height: float | Mapping[str, float] = 1.5,
+    ylim: str | Mapping[str, tuple[int, int]] | tuple[int, int] | None = 'auto',
+    yticks_min_max_only: bool = False,
+    ylab: bool = True,
+    color: str | Mapping[str, str] | None = None,
+    horizontal_ylab: bool = True,
+    filled_tracks: Sequence[str] | None = None,
+    despine: bool = True,
+    despine_keep_bottom: bool = False,
+    plot_track_fn: Callable[..., Any] = plot_track,
+) -> mpl.figure.Figure:
+  """Plot multiple tracks as subplots within one matplotlib figure.
+
+  Custom plotting functions may be passed as `plot_track_fn`. Otherwise,
+  `plot_track` will be used by default.
+
+  Args:
+    tracks: dictionary of tracks to plot. Example input: tracks = {'a': [1,2,3],
+      'b': [4,5,4]}. The arrays in the dict must have the same length.
+    x: optional array of x axis values. If absent, these will be inferred to be
+      from 1 to the length of the arrays in tracks.
+    title: Optional title to be added to the first subplot in the figure.
+    legend: Whether to plot a legend or not.
+    fig_width: Figure width.
+    fig_track_height: Either a scalar specifying the total height of the figure,
+      or a dictionary specifying the height of the subplot for each track.
+    ylim: y axis limit. Must be either "same" (shared y limit across subplots)
+      or "auto" (free scales, indicating separate limits per subplot).
+    yticks_min_max_only: whether the only y axis ticks should be the min and max
+      values.
+    ylab: y axis label.
+    color:  Either a string or a dict indicating a color per track.
+    horizontal_ylab: Whether to display the y axis label horizontally (instead
+      of the default vertical orientation).
+    filled_tracks: List of track names that should have filled line plots drawn
+      (instead of the default whitespace between the values and x axis).
+    despine: Whether to remove top, right, and bottom spines.
+    despine_keep_bottom: Whether to remove top and right spines, but keep the
+      bottom spine.
+    plot_track_fn: Optional custom function for plotting tracks. If absent,
+      defaults to `plot_track`.
+
+  Returns:
+    Matplotlib figure.
+  """
+  # Set figure height (either total height, or per subplot if passed a dict).
+  if isinstance(fig_track_height, dict):
+    fig_height = sum(fig_track_height.values())
+    gridspec_kw = {
+        'height_ratios': [
+            height / fig_height for height in fig_track_height.values()
+        ]
+    }
+  else:
+    fig_height = fig_track_height
+    gridspec_kw = dict()
+
+  # Generate figure axes.
+  fig, axes = plt.subplots(
+      nrows=len(tracks),
+      ncols=1,
+      figsize=(fig_width, fig_height),
+      gridspec_kw=gridspec_kw,
+      sharex=True,
+  )
+  if len(tracks) == 1:
+    axes = [axes]
+
+  # Set y axis limits.
+  if ylim == 'same':
+    ylim = (
+        0,
+        max(v.max() for v in tracks.values() if isinstance(v, np.ndarray)),
+    )
+  elif ylim == 'auto':
+    ylim = None
+  elif isinstance(ylim, str):
+    raise ValueError('Only "same" and "auto" are valid strings for ylim.')
+
+  # Iterate over tracks and plot one track per axis in the figure.
+  for i, (ax, (track, arr)) in enumerate(zip(axes, tracks.items())):
+    # Only set title in the first subplot.
+    if i == 0 and title is not None:
+      ax.set_title(title)
+
+    # Only set x axis ticks in the final subplot.
+    if i != len(tracks) - 1:
+      ax.xaxis.set_ticks_position('none')
+
+    track_ylim = ylim[track] if isinstance(ylim, dict) else ylim
+    track_color = color[track] if isinstance(color, dict) else color
+
+    # If the array is in fact a callable, then it is intended to be used as a
+    # custom plotting function (most commonly, for plotting transcripts).
+    if hasattr(arr, '__call__'):
+      # TODO: b/377225766 - Allow custom plot functions as dictionary elements.
+      arr(ax, ylim=ylim, color=color)
+    else:
+      filled_tracks = filled_tracks or []
+      plot_track_fn(
+          arr,
+          ax=ax,
+          x=x,
+          legend=legend,
+          ylim=track_ylim,
+          color=track_color,
+          filled=track in filled_tracks,
+      )
+
+    # Set y axis label.
+    if ylab:
+      if horizontal_ylab:
+        ax.set_ylabel(
+            track,
+            rotation=0,
+            multialignment='center',
+            va='center',
+            ha='right',
+            labelpad=5,
+        )
+      else:
+        ax.set_ylabel(track)
+
+    # Only draw the y axis ticks on the min and max value locations.
+    if yticks_min_max_only:
+      if track_ylim is None:
+        track_ylim = ax.get_yticks()[[0, -1]]
+      # TODO: b/377226499 - Preventing negative values might be sub-optimal.
+      minimum = max(track_ylim[0], 0)
+      maximum = track_ylim[1]
+      ax.spines['left'].set_bounds(minimum, maximum)
+      ax.set_yticks([minimum, maximum])
+
+    if despine:
+      ax.spines['top'].set_visible(False)
+      ax.spines['right'].set_visible(False)
+      if i != len(tracks) - 1 and despine_keep_bottom:
+        ax.spines['bottom'].set_visible(True)
+      else:
+        ax.spines['bottom'].set_visible(False)
+
+  # Enable default tick locator for the final subplot.
+  axes[-1].xaxis.set_major_locator(mpl.ticker.AutoLocator())
+
+  # No height for whitespace between subplots.
+  fig.subplots_adjust(hspace=0)
+
+  return fig
+
+
+def sashimi_plot(
+    junctions: Sequence[genome.Junction],
+    ax: plt.Axes,
+    interval: genome.Interval | None = None,
+    filter_threshold: float = 0.01,
+    annotate_counts: bool = True,
+    rng: np.random.Generator | None = None,
+):
+  """Plot splice junctions as a Sashimi plot.
+
+  Sashimi plots (first described [here](https://arxiv.org/abs/1306.3466)
+  visualize splice junctions from an RNA-seq experiment.
+
+  According to the authors,
+  * genomic reads are converted into read densities
+  * junction reads are plotted as arcs whose width is determined by the number
+  of reads aligned to the junction spanning the exons connected by the arc.
+
+  Args:
+    junctions: Splice junctions to plot.
+    ax: Matplotlib axis to add the plot to.
+    interval: Interval to plot, used to filter text annotations.
+    filter_threshold: Junctions with number of reads below this threshold will
+      be filtered out.
+    annotate_counts: Whether to annotate the junctions with read counts.
+    rng: Optional random number generator to use for jittering junction paths.
+      If unset will use NumPy's default random number generator.
+  """
+  rng = rng or np.random.default_rng()
+  total = np.sum([junction.k for junction in junctions])
+  # Random jitter position to avoid overlap.
+  jitters = rng.uniform(low=0.05, high=0.15, size=len(junctions))
+  for junction, jt in zip(junctions, jitters):
+    if junction.k < filter_threshold:
+      continue
+    k = junction.k / total
+    verts = [
+        (junction.start, 0.0),
+        (junction.start, jt),
+        (junction.end, jt),
+        (junction.end, 0.0),
+    ]
+
+    path = mpl.path.Path(
+        verts,
+        [
+            mpl.path.Path.MOVETO,
+            mpl.path.Path.CURVE4,
+            mpl.path.Path.CURVE4,
+            mpl.path.Path.CURVE4,
+        ],
+    )
+    patch = mpl.patches.PathPatch(path, facecolor='none', lw=min(k * 30, 5))
+    ax.add_patch(patch)
+    if annotate_counts:
+      text_pos = junction.center()
+      if interval is not None:
+        if text_pos < interval.start or text_pos > interval.end:
+          continue
+      ax.text(text_pos, 0.8 * jt, junction.k, horizontalalignment='center')
+  ax.set_ylim(0, 0.15)
+
+
+def pad_track(track: np.ndarray, new_len: int, value: int = 0) -> np.ndarray:
+  """Pad a track with `value` to the desired length.
+
+  If new_len - len(track) is an even number, the same amount of padding is added
+  to both sides.
+
+  If new_len - len(track) is an odd number, the extra padded element will be
+  added at the end of the array.
+
+  Args:
+    track: Track to pad.
+    new_len: Desired length of the padded track.
+    value: Value to use for padding.
+
+  Returns:
+    Padded track.
+  """
+  assert track.ndim == 2
+  assert new_len >= len(track)
+
+  out = np.empty((new_len, track.shape[1]))
+  out[:] = value
+  delta = new_len - len(track)
+  i = delta // 2
+  j = i + len(track)
+  out[i:j] = track
+  return out

flax_model/alphagenome/_sdk/visualization/plot_components.py

@@ -0,0 +1,1555 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Module containing the main plotting code for AlphaGenome model outputs.
+
+Three main elements are:
+
+  * `plot` function: The primary function for visualizing model outputs.
+  * Component classes: Implement the visualization components (e.g., tracks,
+    contact maps).
+  * Annotation classes: Implement the visualization annotations (e.g.,
+    intervals, variants).
+"""
+
+import abc
+from collections.abc import Mapping, Sequence
+
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import junction_data
+from flax_model.alphagenome._sdk.data import track_data
+from flax_model.alphagenome._sdk.data import transcript as transcript_utils
+from flax_model.alphagenome._sdk.visualization import plot as plot_lib
+from flax_model.alphagenome._sdk.visualization import plot_transcripts
+from jaxtyping import Float32  # pylint: disable=g-importing-member
+import matplotlib
+from matplotlib import colors as plt_colors
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+# String, RGB or RGBA color.
+_ColorType = (
+    str | tuple[float, float, float] | tuple[float, float, float, float]
+)
+
+
+def plot(
+    components: Sequence['AbstractComponent'],
+    interval: genome.Interval,
+    fig_width: int = 20,
+    fig_height_scale: float = 1.0,
+    title: str | None = None,
+    despine: bool = True,
+    despine_keep_bottom: bool = False,
+    annotations: Sequence['AbstractAnnotation'] | None = None,
+    annotation_offset_range: tuple[float, float] = (0.1, 0.6),
+    hspace: float = 0.3,
+    xlabel: str | None = None,
+) -> matplotlib.figure.Figure:
+  """Plots AlphaGenome model outputs as individual panels of 'components'.
+
+  This function generates a visualization of AlphaGenome model outputs
+  using a combination of components (e.g., tracks, contact maps) and
+  annotations (e.g., intervals, variants).
+
+  Args:
+    components: A sequence of components to visualize.
+    interval: The genomic interval to focus on (similar to setting xlim in
+      matplotlib).
+    fig_width: The total figure width.
+    fig_height_scale: Height of the individual track unit. Total plot height is
+      determined as a sum of the individual component heights.
+    title: An optional title for the overall plot.
+    despine: Whether to remove top, right, and bottom spines from the axes.
+    despine_keep_bottom: Whether to remove top and right spines, but keep the
+      bottom spine. Does not apply to transcript components, which are always
+      despined.
+    annotations: Sequence of annotations to visualise across all components.
+    annotation_offset_range: Relative positions in y-axis to place labels for
+      annotations. Each set of labels in the 'annotations' list are spaced
+      evenly within this range.
+    hspace: Vertical whitespace between subplots to avoid tick label overlap
+      (relative fraction).
+    xlabel: If a non-empty string is provided, it is used as the x-axis label.
+      If an empty string is provided, the x-axis label is removed. If None, the
+      default x-axis label is used.
+
+  Returns:
+    A matplotlib figure.
+  """
+
+  # If we are adding text labels for any of the annotation components,
+  # we need to add an extra empty component at the top.
+  add_label_axis = False
+  if annotations is not None:
+    add_label_axis = any(annot.has_labels for annot in annotations)
+    components = (
+        [EmptyComponent()] + list(components)
+        if add_label_axis
+        else list(components)
+    )
+
+  num_axes = sum(component.num_axes for component in components)
+  fig_height = sum(
+      component.total_height * fig_height_scale for component in components
+  )
+
+  offset = 0
+  gridspec_kw = {'height_ratios': []}
+  for component in components:
+    for i in range(component.num_axes):
+      gridspec_kw['height_ratios'].append(
+          component.get_ax_height(i) * fig_height_scale / fig_height
+      )
+      offset += 1
+
+  fig, axes = plt.subplots(
+      nrows=num_axes,
+      ncols=1,
+      figsize=(fig_width, fig_height),
+      gridspec_kw=gridspec_kw,
+      sharex=True,
+  )
+  # Handle the case of a single axis (e.g. single REF / ALT plot).
+  if not isinstance(axes, np.ndarray):
+    axes = [axes]
+
+  offset = 0
+  for component in components:
+    for i in range(component.num_axes):
+      ax = axes[offset]
+      component.plot_ax(ax, axis_index=i, interval=interval)
+
+      if offset == 0 and title is not None:
+        ax.set_title(title)
+
+      if despine:
+        ax.spines['top'].set_visible(False)
+        ax.spines['right'].set_visible(False)
+        if (
+            i != num_axes - 1
+            and despine_keep_bottom
+            and (not isinstance(component, TranscriptAnnotation))
+        ):
+          ax.spines['bottom'].set_visible(True)
+        else:
+          ax.spines['bottom'].set_visible(False)
+
+      # Remove tick marks which become visible due to a subplots_adjust() call.
+      if offset != num_axes - 1:
+        ax.tick_params(axis='x', which='both', bottom=False, labelbottom=False)
+
+      ax.set_xlim(interval.start, interval.end)
+
+      offset += 1
+
+  # Add annotations across all plotted components.
+  if annotations is not None:
+    # Add small vertical offset for each set of labels to avoid overlap.
+    num_labelled_annotations = sum(
+        annotation.has_labels for annotation in annotations
+    )
+    height_offsets = np.linspace(
+        annotation_offset_range[0],
+        annotation_offset_range[1],
+        num_labelled_annotations,
+    )
+    # Identify which axes are transcripts and which we want to annotate.
+    transcript_axes_idx = [
+        i
+        for i, comp in enumerate(components)
+        if isinstance(comp, TranscriptAnnotation)
+    ]
+    axes_to_annotate_idx = (
+        range(1, len(axes)) if add_label_axis else range(len(axes))
+    )
+    label_index = 0
+    for annotation in annotations:
+      for annotate_index in axes_to_annotate_idx:
+        if not annotation.is_variant and annotate_index in transcript_axes_idx:
+          # Do not add interval annotations to axes that involves transcripts.
+          continue
+        annotation.plot_ax(axes[annotate_index], interval, hspace)
+      if annotation.has_labels:
+        # Add labels to the empty axis at the top. All labels are added to
+        # the top axis, in the order they are supplied in annotations list.
+        annotation.plot_labels(axes[0], interval, height_offsets[label_index])
+        label_index += 1
+
+  # Enable default tick locator for the final subplot.
+  axes[-1].xaxis.set_major_locator(matplotlib.ticker.AutoLocator())
+
+  if xlabel is not None:
+    axes[-1].set_xlabel(xlabel)
+  else:
+    axes[-1].set_xlabel(f'Chromosome position; interval={interval}')
+
+  # Slight whitespace between subplots to avoid tick label overlap.
+  fig.subplots_adjust(hspace=hspace)
+
+  return fig
+
+
+class AbstractComponent(abc.ABC):
+  """Abstract base class for plot components."""
+
+  @abc.abstractmethod
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the plot height for the individual axis.
+
+    Args:
+      axis_index: The index of the axis.
+
+    Returns:
+      The height of the axis.
+    """
+
+  @property
+  def total_height(self) -> float:
+    """Returns the total figure height."""
+    return sum(self.get_ax_height(i) for i in range(self.num_axes))
+
+  @property
+  @abc.abstractmethod
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+
+  @abc.abstractmethod
+  def plot_ax(
+      self,
+      ax: matplotlib.axes.Axes,
+      axis_index: int,
+      interval: genome.Interval,
+  ):
+    """Plots the component on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+
+
+class Tracks(AbstractComponent):
+  """Component for visualizing tracks."""
+
+  def __init__(
+      self,
+      tdata: track_data.TrackData,
+      cmap: str = 'viridis',
+      truncate_cmap: bool = True,
+      track_height: float = 1.0,
+      filled: bool = False,
+      ylabel_template: str = '{name}:{strand}',
+      ylabel_horizontal: bool = True,
+      shared_y_scale: bool = False,
+      global_ylims: tuple[float, float] | None = None,
+      max_num_tracks: int = 50,
+      track_colors: Sequence[_ColorType] | str | None = None,
+      **kwargs,
+  ):
+    """Initializes the `Tracks` component.
+
+    Args:
+      tdata: The `TrackData` object to visualize.
+      cmap: The colormap to use for the tracks.
+      truncate_cmap: Whether to slightly truncate the colormap to avoid extreme
+        values.
+      track_height: The height of each track.
+      filled: Whether to fill the area under the tracks.
+      ylabel_template: A template for the y-axis labels.
+      ylabel_horizontal: Whether to make the y-axis labels horizontal.
+      shared_y_scale: Whether to use the same y-axis scale for all tracks.
+      global_ylims: Optional global y-axis limits (min and max).
+      max_num_tracks: The maximum number of tracks to plot.
+      track_colors: An optional sequence of colors to use for the tracks. If a
+        string is passed, it is used as a single color for all tracks.
+      **kwargs: Additional keyword arguments to pass to the plotting function.
+
+    Raises:
+      ValueError: If the number of tracks exceeds `max_num_tracks` or if the
+        track data has more than one positional axis or if the interval is not
+        set within the track data.
+    """
+    if tdata.num_tracks > max_num_tracks:
+      raise ValueError(
+          f'Too many tracks to plot: {tdata.num_tracks} > {max_num_tracks}.'
+      )
+    self._tdata = tdata
+    self._num_tracks = tdata.values.shape[-1]
+    self._track_height = track_height
+    self._filled = filled
+    self._ylabel_horizontal = ylabel_horizontal
+    self._ylabel_template = ylabel_template
+    self._shared_y_scale = shared_y_scale
+    self._global_ylims = global_ylims
+    self._kwargs = kwargs
+
+    if len(self._tdata.positional_axes) != 1:
+      raise ValueError(
+          'Only track_data with 1 positional axis is supported in Tracks.'
+      )
+    if self._tdata.interval is None:
+      raise ValueError('.interval needs to be set in track_data.')
+
+    # Set up color per set of interleaved tracks.
+    if getattr(self._tdata, 'uns') is not None:
+      self._num_tdata = self._tdata.uns['num_interleaved_trackdatas']
+    else:
+      self._num_tdata = 1
+
+    cmap = plt.get_cmap(cmap)
+    num_track_sets = self._num_tracks // self._num_tdata
+    if truncate_cmap:
+      # We do *1.2 to make the color change more gradual. This means that the
+      # the upper range of the cmap is never displayed, which often tends to
+      # achieve nicer results aesthetically.
+      self._colors = cmap(np.linspace(0, 1, round(num_track_sets * 1.2)))
+    else:
+      self._colors = cmap(np.linspace(0, 1, num_track_sets))
+
+    if track_colors is not None:
+      if isinstance(track_colors, str):
+        self._colors = [track_colors] * num_track_sets
+      elif len(track_colors) == 1:
+        self._colors = track_colors * num_track_sets
+      elif len(track_colors) != num_track_sets:
+        raise ValueError(
+            f'track_colors argument (length: {len(track_colors)}) must be'
+            ' either a single color, or the same number of track sets provided'
+            f' in the tdata ({num_track_sets}).'
+        )
+      else:
+        self._colors = track_colors
+
+  def _get_ylimits(self, tdata: track_data.TrackData):
+    """Computes y-axis limits for track sets."""
+    return [
+        (
+            tdata.values[:, i : i + self._num_tdata].min(),
+            tdata.values[:, i : i + self._num_tdata].max(),
+        )
+        for i in range(0, self._num_tracks, self._num_tdata)
+    ]
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._track_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    return self._tdata.num_tracks
+
+  def plot_ax(
+      self,
+      ax: matplotlib.axes.Axes,
+      axis_index: int,
+      interval: genome.Interval,
+  ):
+    """Plots the tracks on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+    tdata = self._tdata
+    assert tdata.interval is not None
+
+    # If an interval is passed, zoom in on that specific sub-interval.
+    if interval is not None:
+      tdata = tdata.slice_by_interval(interval, match_resolution=True)
+      del interval
+    x = (
+        np.arange(tdata.values.shape[0]) * tdata.resolution
+        + tdata.interval.start
+        + tdata.resolution / 2
+    )
+    arr = tdata.values[:, axis_index]
+
+    # Same shared y-axis limits across all tracks.
+    if self._shared_y_scale:
+      if self._global_ylims is not None:
+        ax.set_ylim(self._global_ylims)
+      else:
+        ax.set_ylim(tdata.values.min(), tdata.values.max())
+    else:
+      ax.set_ylim(arr.min(), arr.max())
+
+    # Set the colour of this track.
+    track_color = self._colors[axis_index // self._num_tdata]
+
+    # Draw the line-plot, or filled line-plot if filled=True.
+    if self._filled:
+      ax.fill_between(x, np.ravel(arr), color=track_color, **self._kwargs)
+    else:
+      ax.plot(x, arr, c=track_color, **self._kwargs)
+
+    if self._ylabel_template:
+      _set_ylabel(ax, self._get_ylabel(axis_index), self._ylabel_horizontal)
+
+  def _get_ylabel(self, axis_index: int) -> str:
+    """Returns the y-axis label for the given axis index."""
+    row = self._tdata.metadata.iloc[axis_index]
+    return self._ylabel_template.format(**row.to_dict())
+
+
+class OverlaidTracks(AbstractComponent):
+  """Component for visualizing overlaid track pairs, such as REF/ALT tracks."""
+
+  def __init__(
+      self,
+      tdata: Mapping[str, track_data.TrackData],
+      colors: Mapping[str, str] | None = None,
+      cmap: str | None = 'viridis',
+      track_height: float = 1.0,
+      ylabel_template: str = '{name}:{strand}',
+      ylabel_horizontal: bool = True,
+      shared_y_scale: bool = False,
+      global_ylims: tuple[float, float] | None = None,
+      yticks: Sequence[float] | None = None,
+      yticklabels: Sequence[str] | None = None,
+      alpha: float = 0.8,
+      order_tdata_by_mean: bool = True,
+      max_num_tracks: int = 50,
+      legend_loc: str = 'upper right',
+      **kwargs,
+  ):
+    """Initializes the `OverlaidTracks` component.
+
+    Args:
+      tdata: A dictionary mapping track names to `TrackData` objects.
+      colors: An optional dictionary mapping track names to colors.
+      cmap: The colormap to use if `colors` is not provided.
+      track_height: The height of each track.
+      ylabel_template: A template for the y-axis labels.
+      ylabel_horizontal: Whether to make the y-axis labels horizontal.
+      shared_y_scale: Whether to use the same y-axis scale for all tracks. This
+        is inferred from the min/max data values across all tracks.
+      global_ylims: Optional global y-axis limits (min and max).
+      yticks: Optional set y-axis tick values manually. If not provided, the
+        tick values will be automatically determined. If provided, the length of
+        yticks must match the length of yticklabels.
+      yticklabels: Optional set y-axis tick labels manually. If not provided,
+        the tick values will be automatically determined. If provided, the
+        length of yticklabels must match the length of yticks.
+      alpha: The transparency of the tracks.
+      order_tdata_by_mean: Whether to order the tracks by their mean value (in
+        descending order).
+      max_num_tracks: The maximum number of tracks to plot.
+      legend_loc: The location of the legend (such as 'upper left' or 'best').
+        See the matplotlib Axes legend documentation for more details and
+        options. If None, no legend is shown.
+      **kwargs: Additional keyword arguments to pass to the plotting function.
+
+    Raises:
+      ValueError: If the shapes or metadata of the track data do not match,
+                  or if the number of tracks exceeds `max_num_tracks`, or if
+                  the track data has more than one positional axis, or if the
+                  interval is not set, or if colors are passed but do not match
+                  the track data names.
+    """
+    self._tdata = tdata
+    self._colors = colors
+    self._cmap = cmap
+    self._track_height = track_height
+    self._ylabel_template = ylabel_template
+    self._ylabel_horizontal = ylabel_horizontal
+    self._shared_y_scale = shared_y_scale
+    self._global_ylims = global_ylims
+    self._yticks = yticks
+    self._yticklabels = yticklabels
+    self._alpha = alpha
+    self._order_tdata_by_mean = order_tdata_by_mean
+    self._kwargs = kwargs
+    self._first_tdata = list(tdata.values())[0]
+    self._legend_loc = legend_loc
+
+    if (
+        self._yticks is not None
+        and self._yticklabels is not None
+        and len(self._yticks) != len(self._yticklabels)
+    ):
+      raise ValueError(
+          'If passing yticks and yticklabels, the length of yticks must match'
+          ' the length of yticklabels.'
+      )
+
+    if len(set(data.values.shape for data in tdata.values())) != 1:
+      raise ValueError('Shapes of track data values must be the same.')
+
+    if not all(
+        self._first_tdata.metadata.equals(data.metadata)
+        for data in tdata.values()
+    ):
+      raise ValueError('Metadata of track data must be the same.')
+
+    if self._first_tdata.num_tracks > max_num_tracks:
+      raise ValueError(
+          f'Too many tracks to plot: {self._first_tdata.num_tracks} >'
+          f' {max_num_tracks}.'
+      )
+
+    # If a colors dict is passed, then there should be a matching color for each
+    # track data name.
+    if self._colors is not None:
+      if self._tdata.keys() != self._colors.keys():
+        raise ValueError(
+            f'If passing colors, each tdata name {list(self._tdata.keys())} '
+            'must have an associated color.'
+        )
+    # Otherwise, we define a color from a cmap.
+    else:
+      cmap = plt.get_cmap(self._cmap)
+      colors = cmap(np.linspace(0, 1, len(tdata)))
+      self._colors = dict(zip(self._tdata.keys(), colors))
+
+    if len(self._first_tdata.positional_axes) != 1:
+      raise ValueError(
+          'Only track_data with 1 positional axis is supported in'
+          ' OverlaidTracks.'
+      )
+    if self._first_tdata.interval is None:
+      raise ValueError('.interval needs to be set in track_data.')
+
+    if self._shared_y_scale:
+      # We compute min and max over all the track data arrays in the tdata dict.
+      all_values = np.stack([arr.values for arr in self._tdata.values()])
+      self._vmin = np.min(all_values)
+      self._vmax = np.max(all_values)
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._track_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    return self._first_tdata.num_tracks
+
+  def plot_ax(
+      self,
+      ax: matplotlib.axes.Axes,
+      axis_index: int,
+      interval: genome.Interval,
+  ):
+    """Plots the overlaid tracks on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+
+    def _maybe_slice_tdata(td):
+      """Slices the track data to the interval if passed."""
+      # If an interval is passed, zoom in on that specific sub-interval.
+      if interval is not None:
+        return td.slice_by_interval(interval, match_resolution=True)
+      else:
+        return td
+
+    def _make_ordered_dict_by_mean(tdata):
+      """Reorders track data dict by mean (descending) for better plotting."""
+      mean_tuples = [
+          (name, np.mean(td.values, dtype=np.float64))
+          for name, td in tdata.items()
+      ]
+      sorted_mean_tuples = sorted(
+          mean_tuples, key=lambda item: item[1], reverse=True
+      )
+
+      # Extract names in the sorted order and return ordered tdata.
+      sorted_names = [name for name, _ in sorted_mean_tuples]
+      ordered_tdata = {name: tdata[name] for name in sorted_names}
+      return ordered_tdata
+
+    # Maybe slice the track data to the interval, and reorder by mean.
+    tdata_sliced = {
+        name: _maybe_slice_tdata(td) for name, td in self._tdata.items()
+    }
+    self._tdata_ordered = (
+        _make_ordered_dict_by_mean(tdata_sliced)
+        if self._order_tdata_by_mean
+        else tdata_sliced
+    )
+
+    for name, tdata in self._tdata_ordered.items():
+      assert tdata.interval is not None
+      x = (
+          np.arange(tdata.values.shape[0]) * tdata.resolution
+          + tdata.interval.start
+          + tdata.resolution / 2
+      )
+      arr = tdata.values[:, axis_index]
+
+      if self._global_ylims is not None:
+        ax.set_ylim(self._global_ylims)
+      elif self._shared_y_scale:
+        ax.set_ylim(self._vmin, self._vmax)
+
+      # Plot the two tracks on the same axis. We plot the larger values first so
+      # that the overlap is more visible.
+      ax.plot(
+          x,
+          arr,
+          alpha=self._alpha,
+          **self._kwargs,
+          c=self._colors[name],
+      )
+      if axis_index == 0 and self._legend_loc is not None:
+        ax.legend(self._tdata_ordered.keys(), loc=self._legend_loc)
+
+      if self._yticks is not None:
+        ax.set_yticks(self._yticks)
+      if self._yticklabels is not None:
+        ax.set_yticklabels(self._yticklabels)
+
+    if self._ylabel_template:
+      _set_ylabel(ax, self._get_ylabel(axis_index), self._ylabel_horizontal)
+
+  def _get_ylabel(self, axis_index: int) -> str:
+    """Returns the y-axis label for the given track."""
+    # Metadata equality has already been checked so here we grab the first one.
+    metadata = self._tdata[list(self._tdata.keys())[0]].metadata
+    row = metadata.iloc[axis_index]
+    return self._ylabel_template.format(**row.to_dict())
+
+
+class ContactMaps(AbstractComponent):
+  """Component for visualizing contact maps.
+
+  The `vmin` and `vmax` parameters control the color scaling of the heatmap.
+  Values outside this range will be clipped to `vmin` or `vmax`.
+  """
+
+  def __init__(
+      self,
+      tdata: track_data.TrackData,
+      track_height: float = 10.0,
+      vmin: float | None = -1.0,
+      vmax: float | None = 2.0,
+      norm: matplotlib.colors.TwoSlopeNorm | None = None,
+      ylabel_horizontal: bool = True,
+      ylabel_template: str = '{name}',
+      cmap: matplotlib.colors.LinearSegmentedColormap | None = None,
+      max_num_tracks: int = 10,
+      **kwargs,
+  ):
+    """Initializes the `ContactMaps` component.
+
+    Args:
+      tdata: The `TrackData` object containing the contact maps.
+      track_height: The height of each contact map.
+      vmin: The minimum value for the color scale.
+      vmax: The maximum value for the color scale.
+      norm: An optional normalization for the color scale.
+      ylabel_horizontal: Whether to make the y-axis labels horizontal.
+      ylabel_template: A template for the y-axis labels.
+      cmap: The colormap to use for the contact maps.
+      max_num_tracks: The maximum number of tracks to plot.
+      **kwargs: Additional keyword arguments to pass to the plotting function.
+
+    Raises:
+      ValueError: If the number of tracks exceeds `max_num_tracks`, or if the
+        track data does not have 2 positional axes, or if the contact maps are
+        not square, or if the interval is not set in the track data.
+    """
+    if tdata.num_tracks > max_num_tracks:
+      raise ValueError(
+          f'Too many tracks to plot: {tdata.num_tracks} > {max_num_tracks}.'
+      )
+    self._tdata = tdata
+    self._resolution = tdata.resolution
+    self._track_height = track_height
+    self._vmin = vmin
+    self._vmax = vmax
+    self._norm = norm
+    self._ylabel_horizontal = ylabel_horizontal
+    self._ylabel_template = ylabel_template
+    # TODO: b/377292012 - Add orca_color_map.
+    self._cmap = (
+        cmap if cmap is not None else matplotlib.pyplot.get_cmap('autumn_r')
+    )
+    self._kwargs = kwargs
+    if len(self._tdata.positional_axes) != 2:
+      raise ValueError(
+          'Only track_data with 2 positional axes is supported in ContactMaps.'
+      )
+    if self._tdata.values.shape[0] != self._tdata.values.shape[1]:
+      raise ValueError('Contact maps must be square.')
+
+    if self._tdata.interval is None:
+      raise ValueError('.interval needs to be set in track_data.')
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._track_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    return self._tdata.num_tracks
+
+  def _get_bin_positions(
+      self, interval: genome.Interval, resolution: int
+  ) -> np.ndarray:
+    """Gets the positions of contact map bins in chromosome coordinates."""
+    bin_indices = np.arange(interval.width // resolution)
+    return interval.start + (resolution * bin_indices)
+
+  def _plot_pcolormesh(
+      self,
+      ax: matplotlib.axes.Axes,
+      x: np.ndarray,
+      y: np.ndarray,
+      arr: np.ndarray,
+      vmin: float | None = None,
+      vmax: float | None = None,
+      cmap: matplotlib.colors.Colormap | None = None,
+      norm: matplotlib.colors.Normalize | None = None,
+      **kwargs,
+  ) -> matplotlib.collections.QuadMesh:
+    """Plots the contact map heatmap using `pcolormesh`.
+
+    Note that upsampling the contact maps to single base pair resolution and
+    using something like .imshow() is infeasible since the upsampling blows up
+    memory.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      x: The x-axis coordinates.
+      y: The y-axis coordinates.
+      arr: The contact map data.
+      vmin: The minimum value for the color scale.
+      vmax: The maximum value for the color scale.
+      cmap: The colormap to use.
+      norm: An optional normalization for the color scale.
+      **kwargs: Additional keyword arguments to pass to `pcolormesh`.
+
+    Returns:
+      The `matplotlib.collections.QuadMesh` object representing the plot.
+    """
+    if not norm:  # Cannot pass both norm and vmin/vmax simultaneously.
+      return ax.pcolormesh(
+          x,
+          y,
+          arr,
+          vmin=vmin,
+          vmax=vmax,
+          cmap=cmap,
+          **kwargs,
+      )
+    else:
+      return ax.pcolormesh(
+          x,
+          y,
+          arr,
+          norm=norm,
+          **kwargs,
+      )
+
+  def plot_ax(
+      self,
+      ax: matplotlib.axes.Axes,
+      axis_index: int,
+      interval: genome.Interval,
+  ):
+    """Plots the contact map on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+    tdata = self._tdata
+    assert tdata.interval is not None
+
+    # If an interval is passed, zoom in on that specific sub-interval.
+    if interval is not None:
+      tdata = tdata.slice_by_interval(interval, match_resolution=True)
+      del interval
+
+    arr = tdata.values[:, :, axis_index]
+    x = self._get_bin_positions(tdata.interval, self._resolution)
+
+    # We shift the bin edges by half a step since pcolormesh will misalign
+    # the x-axis by half a step, since the plot values are centered within each
+    # bin, rather than at the bin edges.
+    half_bin_width = self._resolution // 2
+    x = x + half_bin_width
+
+    # The -1 reverse ordering ensures that contact maps are plotted with
+    # the diagonal going down from left to right.
+    y = np.arange(arr.shape[0])[::-1]
+
+    if not self._vmin:
+      self._vmin = np.min(arr)
+
+    if not self._vmax:
+      self._vmax = np.max(arr)
+
+    self._plot_pcolormesh(
+        ax=ax,
+        x=x,
+        y=y,
+        arr=arr,
+        vmin=self._vmin,
+        vmax=self._vmax,
+        cmap=self._cmap,
+        norm=self._norm,
+        **self._kwargs,
+    )
+
+    if self._ylabel_template:
+      _set_ylabel(ax, self._get_ylabel(axis_index), self._ylabel_horizontal)
+
+  def _get_ylabel(self, axis_index: int) -> str:
+    """Returns the y-axis label for the given contact map."""
+    row = self._tdata.metadata.iloc[axis_index]
+    return self._ylabel_template.format(**row.to_dict())
+
+
+class ContactMapsDiff(ContactMaps):
+  """Component for visualizing contact map differences.
+
+  This component visualizes the difference between two contact maps. It uses
+  a diverging red-blue color map with the center white color pinned to a
+  value of zero, with negative values being blue and positive values being red.
+
+  The `vmin` and `vmax` parameters control the color scaling of the heatmap.
+  Values outside this range will be clipped to `vmin` or `vmax`.
+  """
+
+  def __init__(
+      self,
+      tdata: track_data.TrackData,
+      track_height: float = 10.0,
+      vmin: float | None = -1.0,
+      vmax: float | None = 1.0,
+      ylabel_horizontal: bool = True,
+      ylabel_template: str = '{name}',
+      cmap: matplotlib.colors.LinearSegmentedColormap | str | None = 'RdBu_r',
+      max_num_tracks: int = 10,
+      **kwargs,
+  ):
+    """Initializes the `ContactMapsDiff` component.
+
+    Args:
+      tdata: The `TrackData` object containing the contact map differences.
+      track_height: The height of each contact map.
+      vmin: The minimum value for the color scale.
+      vmax: The maximum value for the color scale.
+      ylabel_horizontal: Whether to make the y-axis labels horizontal.
+      ylabel_template: A template for the y-axis labels.
+      cmap: The colormap to use for the contact maps.
+      max_num_tracks: The maximum number of tracks to plot.
+      **kwargs: Additional keyword arguments to pass to the plotting function.
+    """
+    self._norm = plt_colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax)
+
+    super().__init__(
+        tdata=tdata,
+        track_height=track_height,
+        vmin=vmin,
+        vmax=vmax,
+        ylabel_horizontal=ylabel_horizontal,
+        ylabel_template=ylabel_template,
+        cmap=cmap,
+        max_num_tracks=max_num_tracks,
+        **kwargs,
+    )
+
+
+def _set_ylabel(ax: matplotlib.axes.Axes, ylabel: str, horizontal: bool):
+  """Sets the y-axis label.
+
+  Args:
+    ax: The matplotlib axis to set the label on.
+    ylabel: The label text.
+    horizontal: Whether to make the label horizontal.
+  """
+  if ylabel:
+    if horizontal:
+      ax.set_ylabel(
+          ylabel,
+          rotation=0,
+          multialignment='center',
+          va='center',
+          ha='right',
+          labelpad=5,
+      )
+    else:
+      ax.set_ylabel(ylabel)
+
+
+class TranscriptAnnotation(AbstractComponent):
+  """Visualizes transcript annotations."""
+
+  def __init__(
+      self,
+      transcripts: Sequence[transcript_utils.Transcript],
+      adaptive_fig_height: bool = True,
+      fig_height: float = 1.0,
+      transcript_style: plot_transcripts.TranscriptStyle = (
+          plot_transcripts.TranscriptStylePreset.MINIMAL.value
+      ),
+      plot_labels_once: bool = True,
+      label_name: str = 'gene_name',
+      **kwargs,
+  ):
+    """Initializes the `TranscriptAnnotation` component.
+
+    Args:
+      transcripts: A sequence of `Transcript` objects to visualize.
+      adaptive_fig_height: Whether to adjust the figure height based on the
+        number of transcripts.
+      fig_height: The base figure height.
+      transcript_style: The style to use for plotting transcripts. The options
+        are defined in `plot_transcripts.TranscriptStylePreset`.
+      plot_labels_once: Whether to plot labels only once per transcript.
+      label_name: The attribute of the transcript to use for labels.
+      **kwargs: Additional keyword arguments to pass to the plotting function.
+    """
+    self._transcripts = transcripts
+    self._adaptive_fig_height = adaptive_fig_height
+    self._fig_height = fig_height
+    self._kwargs = kwargs
+    self._kwargs['label_name'] = label_name
+    self._kwargs['transcript_style'] = transcript_style
+    self._kwargs['plot_labels_once'] = plot_labels_once
+
+    self._num_transcripts = len(self._transcripts)
+    # TODO(b/377291518): adaptive fig height should in theory scale with the
+    # number of transcripts in the sub-interval, not the total number of
+    # transcripts passed, but this is a bit tricky to implement in the code and
+    # this approach seems to work well enough for now.
+    if self._adaptive_fig_height:
+      self._fig_height = max(0.05 * self._num_transcripts * self._fig_height, 1)
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._fig_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    return 1
+
+  def plot_ax(
+      self, ax: matplotlib.axes.Axes, axis_index: int, interval: genome.Interval
+  ):
+    """Plots the transcript annotations on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+    # Update transcripts to only those overlapping interval.
+    transcripts = [
+        t for t in self._transcripts if t.transcript_interval.overlaps(interval)
+    ]
+    ax.set_yticklabels([])
+    ax.set_yticks([])
+    ax.spines['left'].set_visible(False)
+    plot_transcripts.plot_transcripts(ax, transcripts, interval, **self._kwargs)
+
+
+class SeqLogo(AbstractComponent):
+  """Visualizes a sequence logo."""
+
+  def __init__(
+      self,
+      scores: Float32[np.ndarray, 'S A'],
+      scores_interval: genome.Interval,
+      fig_height: float = 1.0,
+      alphabet: str = 'ACGT',
+      max_width: int = 1000,
+      ylabel: str = '',
+      ylabel_horizontal: bool = True,
+      ylim: tuple[float, float] | None = None,
+      **kwargs,
+  ):
+    """Initializes the `SeqLogo` component.
+
+    Args:
+      scores: A numpy array of shape (sequence_length, alphabet_size) containing
+        the sequence logo scores.
+      scores_interval: The genomic interval corresponding to the scores.
+      fig_height: The height of the figure.
+      alphabet: The alphabet used in the sequence logo.
+      max_width: The maximum width of the sequence logo to plot.
+      ylabel: An optional label for the y-axis.
+      ylabel_horizontal: Whether to make the y-axis label horizontal.
+      ylim: An optional range to set for the y-axis.
+      **kwargs: Additional keyword arguments to pass to the plotting function.
+    """
+    self._scores = scores
+    self._scores_interval = scores_interval
+    self._fig_height = fig_height
+    self._alphabet = alphabet
+    self._max_width = max_width
+    self._ylabel = ylabel
+    self._ylabel_horizontal = ylabel_horizontal
+    self._ylim = ylim
+    self._kwargs = kwargs
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._fig_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    return 1
+
+  def plot_ax(
+      self, ax: matplotlib.axes.Axes, axis_index: int, interval: genome.Interval
+  ):
+    """Plots the sequence logo on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+    intersection = self._scores_interval.intersect(interval)
+    if intersection is None or intersection.width > self._max_width:
+      return
+    relative_start = intersection.start - self._scores_interval.start
+    scores = self._scores[
+        relative_start : (relative_start + intersection.width)
+    ]
+
+    plot_lib.seqlogo(
+        scores,
+        ax=ax,
+        alphabet=self._alphabet,
+        start=intersection.start,
+        one_based=False,
+        **self._kwargs,
+    )
+
+    _set_ylabel(ax, self._ylabel, self._ylabel_horizontal)
+    if self._ylim is not None:
+      ax.set_ylim(self._ylim)
+
+
+class Sashimi(AbstractComponent):
+  """Visualizes splice junctions as a Sashimi plot."""
+
+  def __init__(
+      self,
+      junction_track: junction_data.JunctionData,
+      fig_height: float = 1.0,
+      filter_threshold: float | None = None,
+      ylabel_template: str = '{name}',
+      ylabel_horizontal: bool = True,
+      annotate_counts: bool = True,
+      normalize_values: bool = True,
+      interval_contained: bool = True,
+      rng: np.random.Generator | None = None,
+  ):
+    """Initializes the `Sashimi` component.
+
+    Args:
+      junction_track: A `JunctionData` object to visualize.
+      fig_height: The height of the figure.
+      filter_threshold: The minimum value for a junction to be included in the
+        plot. This is typically based on the normalized read count. If None,
+        filter out junction values below 5% of the maximum value.
+      ylabel_template: A template for the y-axis labels.
+      ylabel_horizontal: Whether to make the y-axis label horizontal.
+      annotate_counts: Whether to annotate the junctions with read counts.
+      normalize_values: Whether to normalize the values to a constant sum.
+      interval_contained: Whether to only plot junctions contained in the
+        interval.
+      rng: Optional random number generator to use for jittering junction paths.
+        If unset will use NumPy's default random number generator.
+    """
+    if normalize_values:
+      self._junction_track = junction_track.normalize_values()
+    else:
+      self._junction_track = junction_track
+    self._fig_height = fig_height
+    self._filter_threshold = filter_threshold
+    self._ylabel_template = ylabel_template
+    self._ylabel_horizontal = ylabel_horizontal
+    self._annotate_counts = annotate_counts
+    self._interval_contained = interval_contained
+    self._rng = rng or np.random.default_rng()
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._fig_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    # Metadata for JunctionData do not have strand information.
+    # Strand information is in JunctionData.intervals.
+    return self._junction_track.num_tracks * len(
+        self._junction_track.possible_strands
+    )
+
+  def _get_strand_and_metadata_index(self, axis_index: int) -> tuple[str, int]:
+    """Returns the strand and metadata index for the given axis index."""
+    if len(self._junction_track.possible_strands) == 1:
+      strand = self._junction_track.possible_strands[0]
+      metadata_index = axis_index
+    else:
+      strand = '+' if axis_index % 2 == 0 else '-'
+      metadata_index = axis_index // 2
+    return strand, metadata_index
+
+  def plot_ax(
+      self, ax: matplotlib.axes.Axes, axis_index: int, interval: genome.Interval
+  ):
+    """Plots the Sashimi plot on the given axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+    strand, metadata_index = self._get_strand_and_metadata_index(axis_index)
+    track_name = self._junction_track.metadata.iloc[metadata_index]['name']
+    junction_track = self._junction_track.intersect_with_interval(interval)
+    junctions = junction_data.get_junctions_to_plot(
+        predictions=junction_track,
+        strand=strand,
+        name=track_name,
+        k_threshold=self._filter_threshold,
+    )
+    if self._interval_contained:
+      junctions = [j for j in junctions if interval.contains(j)]
+    else:
+      junctions = [j for j in junctions if j.overlaps(interval)]
+
+    plot_lib.sashimi_plot(
+        junctions,
+        ax=ax,
+        interval=interval,
+        filter_threshold=0,
+        annotate_counts=self._annotate_counts,
+        rng=self._rng,
+    )
+    ax.set_yticklabels([])
+    ax.set_yticks([])
+    ax.spines['left'].set_visible(False)
+    if self._ylabel_template:
+      _set_ylabel(ax, self._get_ylabel(axis_index), self._ylabel_horizontal)
+
+  def _get_ylabel(self, axis_index: int) -> str:
+    """Returns the y-axis label for the given axis index."""
+    strand, metadata_index = self._get_strand_and_metadata_index(axis_index)
+    row = self._junction_track.metadata.iloc[metadata_index]
+    row = row.to_dict()
+    row['strand'] = strand
+    return self._ylabel_template.format(**row)
+
+
+class EmptyComponent(AbstractComponent):
+  """An empty plotting component."""
+
+  def __init__(self, fig_height: float = 1.0):
+    """Initializes the `EmptyComponent`.
+
+    Args:
+      fig_height: The height of the figure.
+    """
+    self._fig_height = fig_height
+
+  def get_ax_height(self, axis_index: int) -> float:
+    """Returns the height of the axis."""
+    return self._fig_height
+
+  @property
+  def num_axes(self) -> int:
+    """Returns the number of matplotlib axes required by the component."""
+    return 1
+
+  def plot_ax(
+      self, ax: matplotlib.axes.Axes, axis_index: int, interval: genome.Interval
+  ):
+    """Plot an empty axis, removing all labels and spines from the axis.
+
+    Args:
+      ax: The matplotlib axis to plot on.
+      axis_index: The index of the axis.
+      interval: The genomic interval to plot.
+    """
+    ax.set_yticklabels([])
+    ax.set_yticks([])
+    ax.spines['left'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+
+
+class AbstractAnnotation(abc.ABC):
+  """Abstract base class for plot annotations.
+
+  Annotations are visual elements that can be added to plots to highlight
+  specific features or regions. This class defines the common interface
+  for all annotations.
+
+  Attributes:
+    annotations: A sequence of `Variant` or `Interval` objects representing the
+      annotations.
+    colors: An optional string or sequence of strings specifying the colors of
+      the annotations.
+    labels: An optional sequence of strings to use as labels for the
+      annotations.
+    use_default_labels: Whether to use default labels for the annotations if
+      `labels` is not provided.
+  """
+
+  def __init__(
+      self,
+      annotations: Sequence[genome.Variant] | Sequence[genome.Interval],
+      colors: str | Sequence[str] | None,
+      labels: Sequence[str] | None,
+      use_default_labels: bool,
+  ):
+    """Initializes the `AbstractAnnotation` class.
+
+    Args:
+      annotations: A sequence of `Variant` or `Interval` objects.
+      colors: An optional string or sequence of strings specifying colors.
+      labels: An optional sequence of strings to use as labels.
+      use_default_labels: Whether to use default labels if `labels` is not
+        provided.
+
+    Raises:
+      ValueError: If the length of `colors` or `labels` does not match the
+        length of `annotations`.
+    """
+    self._annotations = annotations
+    self._colors = colors
+    self._labels = labels
+    self._use_default_labels = use_default_labels
+
+    # Pre-processing / validation of inputs.
+    num_annotations = len(self._annotations)
+    if (self._colors is not None) and (not isinstance(self._colors, str)):
+      if len(self._colors) != num_annotations:
+        raise ValueError(
+            'Colors must have the same length as intervals/variants or just'
+            ' a single color string.'
+        )
+    if self._labels is not None:
+      if len(self._labels) != num_annotations:
+        raise ValueError(
+            'Labels must have the same length as intervals/variants.'
+        )
+
+  @abc.abstractmethod
+  def plot_ax(
+      self, ax: matplotlib.axes.Axes, interval: genome.Interval, hspace: float
+  ):
+    """Adds the annotation to an individual axis.
+
+    Args:
+      ax: The matplotlib axis to add the annotation to.
+      interval: The genomic interval to plot.
+      hspace: The vertical space between subplots.
+    """
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def plot_labels(
+      self,
+      ax: matplotlib.axes.Axes,
+      interval: genome.Interval,
+      label_height_factor: float,
+  ):
+    """Adds labels for the annotation to an axis.
+
+    Args:
+      ax: The matplotlib axis to add the labels to.
+      interval: The genomic interval to plot.
+      label_height_factor: A scaling factor for the label height.
+    """
+    raise NotImplementedError
+
+  @property
+  def is_variant(self) -> bool:
+    """Returns True if the annotation is a variant annotation."""
+    return isinstance(self._annotations[0], genome.Variant)
+
+  @property
+  def has_labels(self) -> bool:
+    """Returns True if the annotation has labels."""
+    return (self._labels is not None) | (
+        self.is_variant and self._use_default_labels
+    )
+
+  def add_label(
+      self,
+      ax: matplotlib.axes.Axes,
+      label_x_position: float,
+      label: str,
+      angle: float,
+      label_height_factor: float,
+      label_position: str = 'left',
+  ):
+    """Adds a single angled label to an axis.
+
+    Args:
+      ax: The matplotlib axis to add the label to.
+      label_x_position: The x position of the label.
+      label: The label text.
+      angle: The angle of the label.
+      label_height_factor: A scaling factor for the label height.
+      label_position: The (horizontal) placement of the label, relative to the x
+        position. Can be any position string accepted by the horizontalalignment
+        argument of matplotlib.axes.Axes.text.
+    """
+    ylims = ax.get_ylim()
+    label_height = ylims[0] + label_height_factor * np.diff(ylims)[0]
+    ax.text(
+        label_x_position,
+        label_height,
+        label,
+        color='black',
+        fontsize=10,
+        rotation=angle,
+        ha=label_position,
+        va='bottom',
+    )
+    ax.axvline(
+        label_x_position, ymax=label_height * 0.95, color='black', alpha=0.1
+    )
+
+
+class IntervalAnnotation(AbstractAnnotation):
+  """Visualizes intervals as rectangles across all plot components.
+
+  A rectangle is drawn for each interval and overlaid on top of the final plot,
+  spanning all plot components.
+  """
+
+  def __init__(
+      self,
+      intervals: Sequence[genome.Interval],
+      colors: str | Sequence[str] = 'darkgray',
+      alpha: float = 0.2,
+      labels: Sequence[str] | None = None,
+      use_default_labels: bool = True,
+      label_angle: float = 15,
+  ):
+    """Initializes the `IntervalAnnotation` class.
+
+    Args:
+      intervals: A sequence of `Interval` objects to annotate.
+      colors: An optional string or sequence of strings specifying the colors of
+        the interval annotation.
+      alpha: The transparency of the interval annotation.
+      labels: An optional sequence of strings to use as labels for the
+        intervals.
+      use_default_labels: Whether to use default labels for the intervals if
+        `labels` is not provided.
+      label_angle: The angle of the interval labels.
+    """
+    super().__init__(intervals, colors, labels, use_default_labels)
+    self._label_angle = label_angle
+    self._alpha = alpha
+    self._intervals = intervals
+
+  def plot_ax(
+      self,
+      ax: matplotlib.axes.Axes,
+      interval: genome.Interval,
+      hspace: float = 0.0,
+  ):
+    """Adds the interval annotation to an individual axis.
+
+    Args:
+      ax: The matplotlib axis to add the annotation to.
+      interval: The genomic interval to plot.
+      hspace: The vertical space between subplots.
+    """
+    for i, interval_i in enumerate(self._intervals):
+      if isinstance(self._colors, str):
+        color = self._colors
+      else:
+        color = self._colors[i]
+      # Only plot the piece of the annotation that intersects with the plotting
+      # interval.
+      intersection = interval_i.intersect(interval)
+      if intersection is None:
+        continue
+      ax.axvspan(
+          interval_i.start,
+          interval_i.end,
+          # Set y-maximum to be the top of the axis, including hspace.
+          ymax=(1 + hspace) * 1.03,
+          alpha=self._alpha,
+          facecolor=color,
+          # Set edgecolor to None for intervals to avoid horizongal lines
+          # between axes, but keep it for variants. As variants are typically
+          # a very thin rectangle, removing the edges results in the rectangle
+          # not being visible.
+          edgecolor=None,
+          clip_on=False,
+      )
+
+  def plot_labels(
+      self,
+      ax: matplotlib.axes.Axes,
+      interval: genome.Interval,
+      label_height_factor: float,
+  ):
+    """Adds interval labels to an axis.
+
+    Args:
+      ax: The matplotlib axis to add the labels to.
+      interval: The genomic interval to plot.
+      label_height_factor: A scaling factor for the label height.
+    """
+    # Only add labels if they are provided.
+    if self.has_labels:
+      for i, interval_i in enumerate(self._intervals):
+        label = self._labels[i]
+        # Place label in the middle of the rectangle.
+        # Only plot the piece of the annotation that intersects with the
+        # plotting interval.
+        intersection = interval_i.intersect(interval)
+        if intersection is None:
+          continue
+
+        self.add_label(
+            ax,
+            label_x_position=np.mean((interval_i.start, interval_i.end)),
+            label=label,
+            angle=self._label_angle,
+            label_height_factor=label_height_factor,
+        )
+
+
+class VariantAnnotation(AbstractAnnotation):
+  """Visualizes variants as thin line-like rectangles across plot components."""
+
+  def __init__(
+      self,
+      variants: Sequence[genome.Variant],
+      colors: str | Sequence[str] = 'orange',
+      alpha: float = 0.8,
+      labels: Sequence[str] | None = None,
+      use_default_labels: bool = True,
+      label_angle: float = 15,
+      label_position: str = 'left',
+  ):
+    """Initializes the `VariantAnnotation` class.
+
+    Args:
+      variants: A sequence of `Variant` objects to annotate.
+      colors: An optional string or sequence of strings specifying the colors of
+        the variant annotation.
+      alpha: The transparency of the variant annotation.
+      labels: An optional sequence of strings to use as labels for the variants.
+      use_default_labels: Whether to use default labels for the variants if
+        `labels` is not provided.
+      label_angle: The angle of the variant labels.
+      label_position: The (horizontal) placement of the variant label, relative
+        to the variant position. Can be any position string accepted by the
+        horizontalalignment argument of matplotlib.axes.Axes.text.
+    """
+    super().__init__(variants, colors, labels, use_default_labels)
+    self._label_angle = label_angle
+    self._alpha = alpha
+    self._variants = variants
+    self._label_position = label_position
+
+  def plot_ax(
+      self,
+      ax: matplotlib.axes.Axes,
+      interval: genome.Interval,
+      hspace: float = 0.0,
+  ):
+    """Adds a variant annotation to an individual axis.
+
+    Args:
+      ax: The matplotlib axis to add the annotation to.
+      interval: The genomic interval to plot.
+      hspace: The vertical space between subplots.
+    """
+    for i, variant in enumerate(self._variants):
+      if isinstance(self._colors, str):
+        color = self._colors
+      else:
+        color = self._colors[i]
+      interval_i = variant.reference_interval
+      # Only plot the piece of the annotation that intersects with the plotting
+      # interval.
+      intersection = interval_i.intersect(interval)
+      if intersection is None:
+        continue
+      ax.axvspan(
+          interval_i.start,
+          interval_i.end,
+          # Set y-maximum to be the top of the axis, including hspace.
+          ymax=(1 + hspace) * 1.03,
+          alpha=self._alpha,
+          facecolor=color,
+          # Set edgecolor to None for intervals to avoid horizongal lines
+          # between axes, but keep it for variants. As variants are typically
+          # a very thin rectanle, removing the edges results in the rectangle
+          # not being visible.
+          edgecolor=color,
+          clip_on=False,
+      )
+
+  def plot_labels(
+      self,
+      ax: matplotlib.axes.Axes,
+      interval: genome.Interval,
+      label_height_factor: float,
+  ):
+    """Adds variant labels to an axis.
+
+    Args:
+      ax: The matplotlib axis to add the labels to.
+      interval: The genomic interval to plot.
+      label_height_factor: A scaling factor for the label height.
+    """
+    # Only add labels if they are provided.
+    if self.has_labels:
+      for i, variant in enumerate(self._variants):
+        interval_i = variant.reference_interval
+        # Using truncated string method for genome.Variant class to get default
+        # labels for variants.
+        label = (
+            variant.as_truncated_str(max_length=20)
+            if self._use_default_labels
+            else self._labels[i]
+        )
+        # Place label in the middle of the rectangle.
+        # Only plot the piece of the annotation that intersects with the
+        # plotting interval.
+        intersection = interval_i.intersect(interval)
+        if intersection is None:
+          continue
+
+        self.add_label(
+            ax,
+            label_x_position=np.mean((interval_i.start, interval_i.end)),
+            label=label,
+            angle=self._label_angle,
+            label_height_factor=label_height_factor,
+            label_position=self._label_position,
+        )

flax_model/alphagenome/_sdk/visualization/plot_transcripts.py

@@ -0,0 +1,480 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Visualize transcripts/gene annotation in matplotlib."""
+
+from collections.abc import Sequence
+import dataclasses
+import enum
+from typing import Any
+
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import transcript as transcript_utils
+import intervaltree
+import matplotlib as mpl
+import matplotlib.figure
+import matplotlib.path
+import matplotlib.pyplot as plt
+
+
+@dataclasses.dataclass
+class TranscriptStyle:
+  """Style specification for a transcript plot.
+
+  CDS = protein coding sequence.
+  UTR = untranslated region.
+
+  Attributes:
+    cds_height: CDS height.
+    utr_height: UTR height.
+    cds_color: CDS color.
+    utr5_color: 5' UTR color.
+    utr3_color: 3' UTR color.
+    first_noncoding_exon_color: Color of the first non-coding exon. This helps
+      to indicate transcript directionality.
+    label_color: Label color.
+    xlim_pad: Controls amount of whitespace on the region flanks.
+  """
+
+  cds_height: float
+  utr_height: float
+  cds_color: str
+  utr5_color: str
+  utr3_color: str
+  first_noncoding_exon_color: str
+  label_color: str
+  xlim_pad: float
+
+
+class TranscriptStylePreset(enum.Enum):
+  """Style enum for transcript plots.
+
+  Attributes:
+    STANDARD: Standard transcript style.
+    MINIMAL: Minimal transcript style.
+  """
+
+  STANDARD = TranscriptStyle(
+      cds_height=0.7,
+      utr_height=0.35,
+      cds_color='#7f7f7f',  # Grey from tab10 palette.
+      utr5_color='#ff7f0e',  # Orange.
+      utr3_color='#1f77b4',  # Blue.
+      first_noncoding_exon_color='#2ca02c',  # Green.
+      label_color='#7f7f7f',
+      xlim_pad=0.01,
+  )
+
+  MINIMAL = TranscriptStyle(
+      cds_height=0.4,
+      utr_height=0.22,
+      cds_color='black',
+      utr5_color='black',
+      utr3_color='black',
+      # TODO(b/377291432): find a nice way of specifying strand orientation.
+      first_noncoding_exon_color='black',
+      label_color='black',
+      xlim_pad=0.05,
+  )
+
+
+def plot_transcripts(
+    ax: plt.Axes,
+    transcripts: Sequence[transcript_utils.Transcript],
+    interval: genome.Interval,
+    zero_origin: bool = False,
+    label_name: str | None = None,
+    transcript_style: TranscriptStyle = TranscriptStylePreset.STANDARD.value,
+    plot_labels_once: bool = False,
+    **kwargs,
+) -> mpl.figure.Figure:
+  """Plot transcripts.
+
+  Loops over each transcript in `transcripts` and calls `draw_transcript`.
+
+  Args:
+    ax: Matplotlib axis onto which to plot transcript annotations.
+    transcripts: Sequence of transcripts returned by a
+      transcript.TranscriptExtractor.
+    interval: Genomic interval at which to visualize the transcripts.
+    zero_origin: If True, the beginning of the interval will start with 0.
+    label_name: Which label in transcript.info to draw next to the transcript.
+    transcript_style: specification of transcript styling details.
+    plot_labels_once: If True, labels will only be plotted once.
+    **kwargs: kwargs passed to draw_transcript.
+
+  Returns:
+    Matplotlib figure object.
+  """
+  if not transcripts:
+    return
+
+  # Slightly pad x limits for nicer spacing, and shift x axis limits if needed.
+  shift = -interval.start if zero_origin else 0
+  xlim_pad = interval.width * transcript_style.xlim_pad
+  ax.set_xlim(
+      [interval.start + shift - xlim_pad, interval.end + shift + xlim_pad]
+  )
+
+  # Get typical label width and transcript heights.
+  text_width = _get_text_width(transcripts[0].info[label_name], ax=ax)
+  heights = _get_placement_heights(
+      transcripts, extend_fraction=1.0, front_padding=text_width
+  )
+
+  labels_already_drawn = []
+  for transcript in transcripts:
+    # Add transcript labels.
+    if label_name is not None:
+      label = transcript.info[label_name]
+    else:
+      label = None
+
+    draw_transcript(
+        ax=ax,
+        transcript=transcript,
+        interval=interval,
+        y=heights[transcript.transcript_id],
+        cds_height=transcript_style.cds_height,
+        utr_height=transcript_style.utr_height,
+        cds_color=transcript_style.cds_color,
+        utr5_color=transcript_style.utr5_color,
+        utr3_color=transcript_style.utr3_color,
+        first_noncoding_exon_color=transcript_style.first_noncoding_exon_color,
+        label_color=transcript_style.label_color,
+        shift=shift,
+        label=None
+        if (label in labels_already_drawn and plot_labels_once)
+        else label,
+        num_transcripts=len(transcripts),
+        **kwargs,
+    )
+
+    labels_already_drawn.append(label)
+
+  ax.set_ylim([min(heights.values()) - 1, max(heights.values()) + 1])
+
+
+def draw_transcript(
+    ax: plt.Axes,
+    transcript: transcript_utils.Transcript,
+    interval: genome.Interval,
+    y: float,
+    cds_height: float = 0.7,
+    utr_height: float = 0.35,
+    cds_color: str = '#7f7f7f',  # Grey from tab10 palette.
+    utr5_color: str = '#ff7f0e',  # Orange.
+    utr3_color: str = '#1f77b4',  # Blue.
+    first_noncoding_exon_color: str = '#2ca02c',  # Green.
+    shift: int = 0,
+    label: str | None = None,
+    label_color: str = '#7f7f7f',
+    num_transcripts: int = 1,
+    **kwargs,
+) -> None:
+  """Draw an individual transcript as rectangular components on an axis.
+
+  CDS = protein coding sequence.
+  UTR = untranslated region.
+
+  This function is used by `plot_transcripts`.
+
+  Args:
+    ax: Matplotlib axis onto which to draw the transcript.
+    transcript: Transcript to draw.
+    interval: Genomic interval at which to visualize the transcript.
+    y: Vertical position at which to draw the transcript.
+    cds_height: CDS height.
+    utr_height: UTR height.
+    cds_color: CDS color in hex string format.
+    utr5_color: 5' UTR color in hex string format.
+    utr3_color: 3' UTR color  in hex string format.
+    first_noncoding_exon_color: Color of the first non-coding exon. This helps
+      to indicate transcript directionality. Hex string format.
+    shift: X-axis shift.
+    label: Optional label to draw next to the transcript.
+    label_color: Label color.
+    num_transcripts: Total number of transcripts being drawn, used for dynamic
+      arrow sizing.
+    **kwargs: Additional keyword arguments passed to matplotlib plotting
+      functions.
+  """
+  ax.set_yticklabels([])
+  ax.set_yticks([])
+
+  def draw_exons_and_introns(exons, color, exon_height):
+    if not exons:
+      return
+    # 1. Draw all exons.
+    for exon in exons:
+      # TODO: b/377291432 - Skip drawing an exon if it will be drawn below
+      # separately to avoid overlap if alpha<1 and overlaps in vector format.
+      draw_interval(
+          ax=ax,
+          interval=exon,
+          y=y,
+          shift=shift,
+          height=exon_height,
+          color=color,
+          **kwargs,
+      )
+
+    # 2. Draw all introns.
+    for intron in transcript_utils.Transcript(exons).introns:
+      ax.plot([intron.start, intron.end], [y, y], color=color, linewidth=0.5)
+
+  # First draw all exons and introns with UTR height.
+  draw_exons_and_introns(
+      transcript.exons, color=cds_color, exon_height=utr_height
+  )
+  draw_interval(
+      ax=ax,
+      interval=transcript.exons[0],
+      y=y,
+      shift=shift,
+      label=label,
+      height=utr_height,
+      color=cds_color,
+      label_color=label_color,
+      **kwargs,
+  )
+
+  # Draw the first non-coding exon with a special color.
+  first_exon_index = 0 if transcript.is_negative_strand else -1
+  draw_interval(
+      ax=ax,
+      interval=transcript.exons[first_exon_index],
+      y=y,
+      height=utr_height,
+      shift=shift,
+      color=first_noncoding_exon_color,
+      **kwargs,
+  )
+
+  # Add UTRs for coding transcripts.
+  if transcript.cds is not None:
+    draw_exons_and_introns(
+        transcript.utr5, color=utr5_color, exon_height=utr_height
+    )
+    draw_exons_and_introns(
+        transcript.cds, color=cds_color, exon_height=cds_height
+    )
+    draw_exons_and_introns(
+        transcript.utr3, color=utr3_color, exon_height=utr_height
+    )
+
+  # Draw strand arrows across the full transcript span.
+  draw_strand_arrows(
+      ax=ax,
+      transcript=transcript,
+      interval=interval,
+      y=y,
+      color=cds_color,
+      cds_height=cds_height,
+      num_transcripts=num_transcripts,
+  )
+
+
+def draw_strand_arrows(
+    ax: plt.Axes,
+    transcript: transcript_utils.Transcript,
+    interval: genome.Interval,
+    y: float,
+    color: str,
+    *,
+    cds_height: float = 0.22,
+    num_transcripts: int = 1,
+    max_arrows_per_intron: int = 5,
+) -> None:
+  """Draw strand direction arrows on intron lines.
+
+  Arrow count per intron is computed dynamically based on the intron's width
+  relative to the visible interval. Marker size is derived from the UTR height
+  so arrows are always visually smaller than UTR exons.
+
+  Args:
+    ax: Matplotlib axis.
+    transcript: The transcript being drawn.
+    interval: The visible genomic interval.
+    y: Vertical position of the transcript.
+    color: Arrow color.
+    cds_height: CDS height in data coordinates, used to scale arrows.
+    num_transcripts: Total number of transcripts being drawn.
+    max_arrows_per_intron: Maximum number of arrows per intron.
+  """
+  introns = transcript_utils.Transcript(transcript.exons).introns
+  if not introns:
+    return
+
+  fig = ax.get_figure()
+  if fig is not None:
+    _, fig_height_inches = (
+        fig.get_size_inches()  # pytype: disable=attribute-error
+    )
+    ax_height_inches = ax.get_position().height * fig_height_inches
+    y_range = num_transcripts + 2
+    if y_range > 0:
+      pts_per_data = (ax_height_inches * 72) / y_range
+      markersize = min(4.0, cds_height * pts_per_data * 2)
+    else:
+      markersize = 4.0
+  else:
+    markersize = 4.0
+
+  # Custom chevron path: two line segments forming > or < shape.
+  if transcript.is_negative_strand:
+    chevron = matplotlib.path.Path(
+        [(0.5, 0.5), (-0.5, 0.0), (0.5, -0.5)],
+        [
+            matplotlib.path.Path.MOVETO,
+            matplotlib.path.Path.LINETO,
+            matplotlib.path.Path.LINETO,
+        ],
+    )
+  else:
+    chevron = matplotlib.path.Path(
+        [(-0.5, 0.5), (0.5, 0.0), (-0.5, -0.5)],
+        [
+            matplotlib.path.Path.MOVETO,
+            matplotlib.path.Path.LINETO,
+            matplotlib.path.Path.LINETO,
+        ],
+    )
+
+  arrow_positions = []
+  for intron in introns:
+    intron_to_interval_fraction = intron.width / interval.width
+    # Skip arrows for introns that are too small.
+    if intron_to_interval_fraction < 0.01:
+      continue
+    # Use sqrt scaling so large introns don't get overwhelmed with arrows.
+    num_arrows = min(
+        max(1, round(intron_to_interval_fraction**0.5 * max_arrows_per_intron)),
+        max_arrows_per_intron,
+    )
+    space = intron.width / (num_arrows + 1)
+    for i in range(1, num_arrows + 1):
+      arrow_pos = intron.start + i * space
+      # Skip arrows too close to interval edges.
+      if arrow_pos < interval.start + 10 or arrow_pos > interval.end - 10:
+        continue
+      arrow_positions.append(arrow_pos)
+
+  if arrow_positions:
+    ax.plot(
+        arrow_positions,
+        [y] * len(arrow_positions),
+        marker=chevron,
+        markersize=markersize,
+        color=color,
+        fillstyle='none',
+        markeredgewidth=0.8,
+        linestyle='none',
+        clip_on=True,
+    )
+
+
+def draw_interval(
+    ax: plt.Axes,
+    interval: genome.Interval,
+    y: float,
+    label: str | None = None,
+    height: float = 0.5,
+    shift: int = 0,
+    label_color: str = '#7f7f7f',
+    **kwargs,
+):
+  """Draw rectangle patch on the axis given a genomic interval.
+
+  Args:
+    ax: Matplotlib axis onto which to draw the interval.
+    interval: Genomic interval to draw.
+    y: Vertical position at which to draw the interval.
+    label: Optional label to draw next to the interval.
+    height: Height of the interval.
+    shift: X-axis shift.
+    label_color: Label color in hex string format.
+    **kwargs: Additional keyword arguments passed to matplotlib plotting
+      functions.
+  """
+  xy = (interval.start + shift, y - height / 2)
+  ax.add_patch(
+      mpl.patches.Rectangle(
+          xy=xy,
+          width=interval.width,
+          height=height,
+          clip_on=True,
+          linewidth=0,
+          **kwargs,
+      )
+  )
+
+  # Add center-aligned text label.
+  if label is not None:
+    ax.text(
+        x=max(xy[0], ax.get_xlim()[0]),
+        y=y,
+        s=label,
+        color=label_color,
+        horizontalalignment='right',
+        verticalalignment='center',
+    )
+
+
+def _get_placement_heights(
+    transcripts: Sequence[transcript_utils.Transcript],
+    extend_fraction: float = 1.0,
+    front_padding: float = 0.0,
+) -> dict[Any, int]:
+  """Get heights at which to place the transcripts."""
+  # TODO: b/376672690 - Implement simpler packing algorithm.
+  levels = [intervaltree.IntervalTree()]
+  # Sort transcripts by length and start placing longest transcripts first.
+  sorted_transcripts = sorted(
+      transcripts, key=lambda x: x.transcript_interval.width, reverse=True
+  )
+  transcript_levels = {}
+  for transcript in sorted_transcripts:
+    placed = False
+    level_idx = 0
+    while not placed:
+      if level_idx >= len(levels):
+        levels.append(intervaltree.IntervalTree())
+      if levels[level_idx].overlaps(
+          transcript.transcript_interval.start - front_padding,
+          transcript.transcript_interval.end,
+      ):
+        # Overlaps an existing interval -> increase the level.
+        level_idx += 1
+      else:
+        # Doesn't overlap. Remember the interval.
+        levels[level_idx].addi(
+            transcript.transcript_interval.start - front_padding,
+            int(transcript.transcript_interval.end * extend_fraction),
+        )
+        transcript_levels[transcript.transcript_id] = level_idx
+        placed = True
+  return {
+      transcript_id: len(levels) - 1 - level_idx
+      for transcript_id, level_idx in transcript_levels.items()
+  }
+
+
+def _get_text_width(label: str, ax: plt.Axes, **kwargs) -> float:
+  """Get text width in data coordinates."""
+  text = ax.text(0, 0, label, **kwargs)
+  plt.gcf().canvas.draw()
+  bb = text.get_window_extent().transformed(ax.transData.inverted())
+  text.remove()  # Remove text.
+  return bb.x1 - bb.x0

flax_model/alphagenome/evals/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""AlphaGenome Evaluation Utilities."""

flax_model/alphagenome/evals/regression_metrics.py

@@ -0,0 +1,204 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for experiments."""
+
+from typing import Sequence
+from flax_model.alphagenome._sdk import typing
+import chex
+import jax
+import jax.numpy as jnp
+from jaxtyping import ArrayLike, Float, PyTree  # pylint: disable=g-importing-member, g-multiple-import
+import numpy as np
+
+
+@chex.dataclass
+class _PearsonRState:
+  """State to compute PearsonR correlation coefficient."""
+
+  xy_sum: jax.Array
+  x_sum: jax.Array
+  xx_sum: jax.Array
+  y_sum: jax.Array
+  yy_sum: jax.Array
+  count: jax.Array
+
+  def __add__(self, other: '_PearsonRState') -> '_PearsonRState':
+    return jax.tree.map(lambda x, y: x + y, self, other)
+
+
+def _pearsonr_initialize() -> '_PearsonRState':
+  """Initialize PearsonrState with zeros."""
+  return _PearsonRState(
+      xy_sum=np.zeros(()),
+      x_sum=np.zeros(()),
+      xx_sum=np.zeros(()),
+      y_sum=np.zeros(()),
+      yy_sum=np.zeros(()),
+      count=np.zeros(()),
+  )
+
+
+def _pearsonr_update(
+    x: jax.Array,
+    y: jax.Array,
+    axis: Sequence[int] | int | None = None,
+    mask: jax.Array | None = None,
+) -> _PearsonRState:
+  """Construct PearsonrState by correlating two arrays."""
+  if mask is not None:
+    mask = jnp.astype(mask, bool)
+  return _PearsonRState(
+      xy_sum=jnp.sum(x * y, axis=axis, where=mask, dtype=jnp.float32),
+      x_sum=jnp.sum(x, axis=axis, where=mask, dtype=jnp.float32),
+      xx_sum=jnp.sum(jnp.square(x), axis=axis, where=mask, dtype=jnp.float32),
+      y_sum=jnp.sum(y, axis=axis, where=mask, dtype=jnp.float32),
+      yy_sum=jnp.sum(jnp.square(y), axis=axis, where=mask, dtype=jnp.float32),
+      count=jnp.sum(jnp.ones_like(x), axis=axis, where=mask, dtype=jnp.float32),
+  )
+
+
+def _pearsonr_result(state: _PearsonRState) -> jax.Array:
+  """Get PearsonR correlation coeficient."""
+  x_mean = state.x_sum / state.count
+  y_mean = state.y_sum / state.count
+
+  covariance = state.xy_sum - state.count * x_mean * y_mean
+
+  x_var = state.xx_sum - state.count * x_mean * x_mean
+  y_var = state.yy_sum - state.count * y_mean * y_mean
+  variance = x_var**0.5 * y_var**0.5
+  eps = jnp.finfo(variance.dtype).eps  # Avoid division by zero.
+  return covariance / (variance + eps)
+
+
+@chex.dataclass
+class RegressionState:
+  """State for accumulating regression statistics."""
+
+  pearsonr: _PearsonRState
+  pearsonr_log1p: _PearsonRState
+  sq_error: jax.Array
+  abs_error: jax.Array
+  count: jax.Array
+
+  def __add__(self, other: 'RegressionState') -> 'RegressionState':
+    return jax.tree.map(
+        lambda a, b: a + b,
+        self,
+        other,
+    )
+
+
+def initialize_regression_metrics() -> RegressionState:
+  """Initialize metric state."""
+  return RegressionState(
+      pearsonr=_pearsonr_initialize(),
+      pearsonr_log1p=_pearsonr_initialize(),
+      sq_error=np.zeros(()),
+      abs_error=np.zeros(()),
+      count=np.zeros(()),
+  )
+
+
+def update_regression_metrics(
+    y_true: jax.Array,
+    y_pred: jax.Array,
+    mask: jax.Array | None = None,
+) -> RegressionState:
+  y_true = jnp.astype(y_true, jnp.float32)
+  y_pred = jnp.astype(y_pred, jnp.float32)
+  return RegressionState(
+      pearsonr=_pearsonr_update(y_true, y_pred, mask=mask, axis=(-2, -3)),
+      pearsonr_log1p=_pearsonr_update(
+          jnp.log1p(y_true), jnp.log1p(y_pred), mask=mask, axis=(-2, -3)
+      ),
+      sq_error=jnp.sum(
+          jnp.square(y_true - y_pred),
+          axis=(-2, -3),
+          where=mask,
+          dtype=jnp.float32,
+      ),
+      abs_error=jnp.sum(
+          jnp.abs(y_true - y_pred),
+          axis=(-2, -3),
+          where=mask,
+          dtype=jnp.float32,
+      ),
+      count=jnp.sum(
+          jnp.ones_like(y_true), axis=(-2, -3), where=mask, dtype=jnp.float32
+      ),
+  )
+
+
+def finalize_regression_metrics(
+    state: PyTree[RegressionState],
+) -> PyTree[jax.Array]:
+  """Compute final metrics from accumulated state."""
+
+  def _finalize(state: RegressionState) -> PyTree[jax.Array]:
+    return {
+        'pearsonr': (
+            _pearsonr_result(state.pearsonr).mean(
+                where=state.pearsonr.count > 0
+            )
+        ),
+        'pearsonr_log1p': (
+            _pearsonr_result(state.pearsonr_log1p).mean(
+                where=state.pearsonr_log1p.count > 0
+            )
+        ),
+        'mse': jnp.mean(state.sq_error / state.count, where=state.count > 0),
+        'mae': jnp.mean(state.abs_error / state.count, where=state.count > 0),
+    }
+
+  return jax.tree.map(
+      _finalize, state, is_leaf=lambda x: isinstance(x, RegressionState)
+  )
+
+
+def reduce_regression_metrics(
+    previous_metrics: PyTree[RegressionState],
+    current_metrics: PyTree[RegressionState],
+) -> RegressionState:
+  """Reduce metrics from a single device to a single scalar."""
+  return jax.tree.map(
+      lambda x, y: x + y,
+      previous_metrics,
+      current_metrics,
+      is_leaf=lambda x: isinstance(x, RegressionState),
+  )
+
+
+@typing.jaxtyped
+def crop_sequence_length(
+    x: Float[ArrayLike, '... S D'], *, target_length: int
+) -> Float[ArrayLike, '... {target_length} D']:
+  """Crops an array to match the target length along the sequence dimension."""
+  sequence_axis = -2
+  if x.shape[sequence_axis] < target_length:
+    raise ValueError(
+        f'Input length {x.shape[sequence_axis]} is shorter than the requested'
+        f' cropped length of {target_length}.'
+    )
+  elif x.shape[sequence_axis] == target_length:
+    return x
+  else:
+    ltrim = (x.shape[sequence_axis] - target_length) // 2
+    rtrim = x.shape[sequence_axis] - target_length - ltrim
+    slices = [
+        slice(None),
+    ] * len(x.shape)
+    slices[sequence_axis] = slice(ltrim, -rtrim)
+    return x[tuple(slices)]

flax_model/alphagenome/evals/track_prediction.py

@@ -0,0 +1,202 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Evaluates AlphaGenome track prediction performance."""
+
+from collections.abc import Iterator
+import pprint
+from typing import Callable, Sequence
+from absl import app
+from absl import logging
+from flax_model.alphagenome._sdk.data import fold_intervals
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome.evals import regression_metrics
+from flax_model.alphagenome.io import bundles as bundles_lib
+from flax_model.alphagenome.io import dataset
+from flax_model.alphagenome.model import dna_model
+from flax_model.alphagenome.model import model as model_lib
+from flax_model.alphagenome.model import schemas
+from flax_model.alphagenome.model.metadata import metadata as metadata_lib
+import haiku as hk
+import jax
+from jax import sharding
+from jax.experimental import mesh_utils
+from jaxtyping import PyTree  # pylint: disable=g-importing-member
+import jmp
+import kagglehub
+import orbax.checkpoint as ocp
+import pandas as pd
+import tensorflow as tf
+
+
+PS = sharding.PartitionSpec
+PredictFn = Callable[
+    [
+        hk.Params,
+        hk.State,
+        jax.Array,
+        jax.Array,
+    ],
+    PyTree[jax.Array],
+]
+
+_SUBSET = fold_intervals.Subset.VALID
+_LOG_FREQUENCY = 5
+_EVAL_BUNDLES = [
+    bundles_lib.BundleName.ATAC,
+    bundles_lib.BundleName.CAGE,
+    bundles_lib.BundleName.CHIP_HISTONE,
+    bundles_lib.BundleName.CHIP_TF,
+    bundles_lib.BundleName.DNASE,
+    bundles_lib.BundleName.PROCAP,
+    bundles_lib.BundleName.RNA_SEQ,
+]
+
+
+def load_model(
+    model_version: dna_model.ModelVersion = dna_model.ModelVersion.FOLD_0,
+) -> tuple[hk.Params, hk.State, PredictFn]:
+  """Loads the model consiting of params, state and predict function."""
+  checkpoint_path = kagglehub.model_download(
+      f'google/alphagenome/jax/{model_version.name.lower()}'
+  )
+  params, state = ocp.StandardCheckpointer().restore(checkpoint_path)
+  metadata = {
+      organism: metadata_lib.load(organism) for organism in dna_model.Organism
+  }
+
+  @hk.transform_with_state
+  def forward(dna_sequence, organism_index):
+    policy = jmp.get_policy('params=float32,compute=bfloat16,output=bfloat16')
+    with hk.mixed_precision.push_policy(model_lib.AlphaGenome, policy):
+      return model_lib.AlphaGenome(metadata)(dna_sequence, organism_index)
+
+  @jax.jit(
+      in_shardings=(PS(), PS(), PS('data'), PS('data')),
+      out_shardings=PS('data'),
+  )
+  def predict(params, state, dna_sequence, organism_index) -> PyTree[jax.Array]:
+    (predictions, _), _ = forward.apply(
+        params, state, None, dna_sequence, organism_index
+    )
+    predictions = dna_model.extract_predictions(predictions)
+    return predictions
+
+  return params, state, predict
+
+
+def create_eval_step(
+    predict_fn: PredictFn, bundles: Sequence[bundles_lib.BundleName]
+):
+  """Returns the eval step function."""
+
+  @jax.jit(
+      in_shardings=(PS(), PS(), PS('data')),
+      out_shardings=PS(),
+  )
+  def eval_step(params, state, batch: schemas.DataBatch):
+    predictions = predict_fn(
+        params, state, batch.dna_sequence, batch.organism_index
+    )
+    metrics_step = {}
+    for bundle in bundles:
+      targets_true, mask = batch.get_genome_tracks(bundle)
+      targets_pred = predictions[dna_output.OutputType[bundle.name]]
+      targets_pred = regression_metrics.crop_sequence_length(
+          targets_pred, target_length=targets_true.shape[-2]
+      )
+      metrics_step[bundle.name] = regression_metrics.update_regression_metrics(
+          targets_true, targets_pred, mask
+      )
+    return metrics_step
+
+  return eval_step
+
+
+def evaluate(
+    params: hk.Params,
+    state: hk.State,
+    predict_fn: PredictFn,
+    bundles: Sequence[bundles_lib.BundleName],
+    dataset_iterator: Iterator[tuple[schemas.DataBatch, dataset.BatchMetadata]],
+):
+  """Evaluates the model."""
+  # Setup Mesh.
+  devices = mesh_utils.create_device_mesh((jax.local_device_count(),))
+  mesh = jax.sharding.Mesh(devices, axis_names=('data',))
+  sharding_rep = sharding.NamedSharding(mesh, PS())
+  sharding_data = sharding.NamedSharding(mesh, PS('data'))
+
+  # Replicate params and state.
+  params = jax.device_put(params, sharding_rep)
+  state = jax.device_put(state, sharding_rep)
+
+  eval_step = create_eval_step(predict_fn, bundles)
+  metrics = {
+      b.name: regression_metrics.initialize_regression_metrics()
+      for b in bundles
+  }
+  num_elements = 0
+
+  for i, (batch, _) in enumerate(dataset_iterator):
+    num_elements += batch.dna_sequence.shape[0]
+    if i % _LOG_FREQUENCY == 1:
+      m = pprint.pformat(
+          regression_metrics.finalize_regression_metrics(metrics)
+      )
+      logging.info('step %d: %s', i, m)
+
+    with jax.set_mesh(mesh):
+      batch = jax.device_put(batch, sharding_data)
+      step_metrics = eval_step(params, state, batch)
+
+    # Accumulate metrics.
+    step_metrics = jax.device_get(step_metrics)
+    metrics = regression_metrics.reduce_regression_metrics(
+        metrics, step_metrics
+    )
+  logging.info('num_elements: %d', num_elements)
+  return regression_metrics.finalize_regression_metrics(metrics)
+
+
+def run(
+    organism: dna_model.Organism = dna_model.Organism.HOMO_SAPIENS,
+    model_version: dna_model.ModelVersion = dna_model.ModelVersion.FOLD_0,
+) -> pd.DataFrame:
+  """Runs the track prediction experiment."""
+  logging.info('Starting track prediction experiment.')
+  params, state, predict_fn = load_model(model_version)
+  dataset_iterator = dataset.get_numpy_dataset_iterator(
+      batch_size=jax.local_device_count(),
+      organism=organism,
+      model_version=model_version,
+      bundles=_EVAL_BUNDLES,
+      subset=_SUBSET,
+  )
+  results = evaluate(params, state, predict_fn, _EVAL_BUNDLES, dataset_iterator)
+  flattened_results = {}
+  for bundle, result in results.items():
+    for metric, value in result.items():
+      logging.info('bundle: %s, metric: %s, value: %s', bundle, metric, value)
+      flattened_results[f'{bundle}_{metric}'] = value
+  df = pd.DataFrame(
+      {'metric': flattened_results.keys(), 'value': flattened_results.values()}
+  )
+  return df
+
+
+if __name__ == '__main__':
+  # Hide local GPUs from TF. TF is only used for data loading.
+  tf.config.set_visible_devices([], 'GPU')
+  app.run(lambda _: run())

flax_model/alphagenome/finetuning/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for finetuning."""

flax_model/alphagenome/finetuning/dataset.py

@@ -0,0 +1,236 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Data pipeline for reading sequence and tracks for fine-tuning."""
+
+from collections.abc import Iterator, Sequence
+import concurrent.futures
+from typing import Any, Mapping
+
+from absl import logging
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.models import dna_model
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome.io import fasta
+from flax_model.alphagenome.model import one_hot_encoder
+from flax_model.alphagenome.model.metadata import metadata as metadata_lib
+import numpy as np
+import pandas as pd
+import pyBigWig
+import tensorflow as tf
+
+
+class BigWigExtractor:
+  """BigWig file extractor using pyBigWig."""
+
+  def __init__(self, file_path: str):
+    self._bw = pyBigWig.open(file_path)
+    self._chromosomes = set(self._bw.chroms().keys())
+
+  def __del__(self):
+    self.close()
+
+  @property
+  def chromosomes(self) -> set[str]:
+    return self._chromosomes
+
+  def extract(self, interval: genome.Interval) -> tf.Tensor:
+    """Extracts values from a BigWig file for a given interval."""
+    if interval.chromosome not in self._chromosomes:
+      raise ValueError(
+          f'Chromosome {interval.chromosome} not found in BigWig. '
+          'Check self._bw.chroms() for available chromosomes. '
+      )
+    start = max(0, interval.start)
+    end = min(self._bw.chroms()[interval.chromosome], interval.end)
+
+    if end <= start:
+      return tf.cast(np.zeros(interval.width), tf.bfloat16)
+
+    values = self._bw.values(interval.chromosome, start, end, numpy=True)
+    if values.shape[0] != interval.width:
+      pad_left = start - interval.start
+      pad_right = interval.end - end
+      values = np.pad(
+          values, pad_width=(pad_left, pad_right), constant_values=0
+      )
+    return tf.cast(np.nan_to_num(values, nan=0.0), tf.bfloat16)
+
+  def close(self):
+    if self._bw is not None:
+      self._bw.close()
+
+
+class MultiTrackExtractor:
+  """Multi-track BigWig file extractor.
+
+  Returns tracks per output type. Tracks are ordered as per the metadata.
+  """
+
+  def __init__(
+      self,
+      output_metadata: metadata_lib.AlphaGenomeOutputMetadata,
+      sequence_length: int,
+      max_workers=32,
+  ):
+    """Initializes the MultiTrackExtractor.
+
+    Args:
+      output_metadata: The output metadata containing the track information.
+      sequence_length: The length of the sequence to extract.
+      max_workers: The maximum number of workers to use for parallel extraction.
+    """
+    self._output_metadata = output_metadata
+    self._sequence_length = sequence_length
+    self._executor = concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_workers
+    )
+
+    # group bigwig extractors by output type.
+    self._bw_extractors: dict[
+        dna_output.OutputType, Sequence[BigWigExtractor]
+    ] = {}
+    for output_type in dna_output.OutputType:
+      if (metadata := output_metadata.get(output_type)) is not None:
+        self._bw_extractors[output_type] = [
+            BigWigExtractor(file_path) for file_path in metadata['file_path']
+        ]
+
+    self._track_masks = {
+        f'{output_type.name.lower()}_mask': np.logical_not(mask).reshape(1, -1)
+        for output_type, mask in output_metadata.padding.items()
+    }
+
+  def get_output_signature(self):
+    """Returns the output signature of the dataset."""
+    signature = {}
+    for output_type, extractors in self._bw_extractors.items():
+      num_tracks = len(extractors)
+      output_name = output_type.name.lower()
+      signature[output_name] = tf.TensorSpec(
+          shape=(self._sequence_length, num_tracks),
+          dtype=tf.bfloat16,
+      )
+      signature[f'{output_name}_mask'] = tf.TensorSpec(
+          shape=(1, num_tracks),
+          dtype=tf.bool,
+      )
+    return signature
+
+  def extract(self, interval: genome.Interval) -> Mapping[str, tf.Tensor]:
+    """Extracts all tracks for all groups in parallel."""
+
+    # Submit
+    future_map = {}
+    for output_type, extractors in self._bw_extractors.items():
+      future_map[output_type] = [
+          self._executor.submit(bw.extract, interval) for bw in extractors
+      ]
+
+    # Collect
+    data = {}
+    for output_type, futures in future_map.items():
+      try:
+        results = [f.result() for f in futures]
+        data[output_type.name.lower()] = np.stack(results, axis=-1)
+      except Exception as e:
+        raise RuntimeError(
+            f"Failed to extract tracks for output group '{output_type.name}'"
+        ) from e
+    return data | self._track_masks
+
+  def close(self):
+    for bw_extractor in self._bw_extractors.values():
+      for bw in bw_extractor:
+        bw.close()
+
+
+class DataPipeline:
+  """Data pipeline for reading sequence and tracks."""
+
+  def __init__(
+      self,
+      *,
+      fasta_path: str,
+      intervals: pd.DataFrame,
+      output_metadata: metadata_lib.AlphaGenomeOutputMetadata,
+      organism: dna_model.Organism,
+      sequence_length: int,
+  ):
+    if organism != dna_model.Organism.HOMO_SAPIENS:
+      raise NotImplementedError('Only HOMO_SAPIENS is currently supported.')
+
+    self._fasta_extractor = fasta.FastaExtractor(fasta_path)
+    self._intervals = intervals
+    self._organsim = organism
+    self._organism_index = 0  # Only one organism is supported.
+    self._sequence_length = sequence_length
+    self._one_hot_encoder = one_hot_encoder.DNAOneHotEncoder()
+    self._multi_track_extractor = MultiTrackExtractor(
+        output_metadata=output_metadata,
+        sequence_length=sequence_length,
+    )
+
+  def get_element(self, idx: int) -> Mapping[str, Any]:
+    """Returns a single element of the dataset."""
+    interval = self._intervals.iloc[idx]
+    interval = genome.Interval(
+        chromosome=interval['chromosome'],
+        start=int(interval['start']),
+        end=int(interval['end']),
+    )
+    interval = interval.resize(self._sequence_length)
+    seq_str = self._fasta_extractor.extract(interval)
+    seq_one_hot = self._one_hot_encoder.encode(seq_str)
+    track_bundles = self._multi_track_extractor.extract(interval)
+    return {
+        'dna_sequence': seq_one_hot,
+        'organism_index': self._organism_index,
+        'bundles': track_bundles,
+    }
+
+  def get_output_signature(self):
+    """Returns the output signature of the dataset."""
+    return {
+        'dna_sequence': tf.TensorSpec(
+            shape=(self._sequence_length, 4), dtype=tf.float32
+        ),
+        'organism_index': tf.TensorSpec(shape=[], dtype=tf.int32),
+        'bundles': self._multi_track_extractor.get_output_signature(),
+    }
+
+  def get_generator(
+      self, num_epochs: int = -1, shuffle: bool = True, seed: int = 0
+  ) -> Iterator[Mapping[str, Any]]:
+    """Returns a generator for the dataset."""
+    rng = np.random.default_rng(seed=seed)
+    num_epochs = num_epochs if num_epochs > 0 else float('inf')
+    epoch_idx = 0
+    while epoch_idx < num_epochs:
+      epoch_idx += 1
+      if shuffle:
+        self._intervals = self._intervals.sample(
+            frac=1, random_state=rng
+        ).reset_index(drop=True)
+      for idx in range(len(self._intervals)):
+        try:
+          yield self.get_element(idx)
+        except Exception as e:  # pylint: disable=broad-except
+          logging.warning(
+              'Failed to get interval %s. With error: %s',
+              self._intervals.iloc[idx].to_dict(),
+              e,
+          )
+
+  def close(self):
+    self._multi_track_extractor.close()

flax_model/alphagenome/finetuning/dataset_test.py

@@ -0,0 +1,198 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.models import dna_client
+from flax_model.alphagenome.finetuning import dataset
+from flax_model.alphagenome.io import fasta
+from flax_model.alphagenome.model.metadata import metadata as metadata_lib
+import chex
+import numpy as np
+import pandas as pd
+import pyBigWig
+import tensorflow as tf
+
+MOCK_CHROM_SIZES = {'chr1': 100}
+
+
+def get_mock_metadata(num_atac_tracks: int = 2, num_dnase_tracks: int = 1):
+  return metadata_lib.AlphaGenomeOutputMetadata(
+      atac=pd.DataFrame({
+          'file_path': [f'path{i}' for i in range(num_atac_tracks)],
+          'name': [f'atac_track{i}' for i in range(num_atac_tracks)],
+          'strand': ['.'] * num_atac_tracks,
+      }),
+      dnase=pd.DataFrame({
+          'file_path': [f'path{i}' for i in range(num_dnase_tracks)],
+          'name': [f'dnase_track{i}' for i in range(num_dnase_tracks)],
+          'strand': ['.'] * num_dnase_tracks,
+      }),
+  )
+
+
+def get_mock_bw():
+  mock_bw = mock.Mock()
+  mock_bw.chroms.return_value = MOCK_CHROM_SIZES
+
+  def mock_values(chrom, start, end, numpy=True):
+    del chrom, numpy
+    return np.full(end - start, 1.0)
+
+  mock_bw.values = mock_values
+  return mock_bw
+
+
+class MockFastaExtractor:
+
+  def __init__(self, path):
+    pass
+
+  def extract(self, interval):
+    return 'A' * interval.width
+
+
+class DataTest(parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='valid_interval',
+          interval=genome.Interval('chr1', 10, 20),
+          expected_values=np.ones(10),
+      ),
+      dict(
+          testcase_name='interval_start_negative',
+          interval=genome.Interval('chr1', -10, 10),
+          expected_values=np.concatenate([np.zeros(10), np.ones(10)]),
+      ),
+      dict(
+          testcase_name='interval_end_beyond_chrom_length',
+          interval=genome.Interval('chr1', 90, 110),
+          expected_values=np.concatenate([np.ones(10), np.zeros(10)]),
+      ),
+      dict(
+          testcase_name='interval_fully_outside_negative',
+          interval=genome.Interval('chr1', -20, -10),
+          expected_values=np.zeros(10),
+      ),
+      dict(
+          testcase_name='interval_fully_outside_positive',
+          interval=genome.Interval('chr1', 110, 120),
+          expected_values=np.zeros(10),
+      ),
+      dict(
+          testcase_name='zero_length_interval',
+          interval=genome.Interval('chr1', 10, 10),
+          expected_values=np.array([]),
+      ),
+  ])
+  @mock.patch.object(pyBigWig, 'open')
+  def test_bigwig_extractor(self, mock_bigwig_open, interval, expected_values):
+    mock_bigwig_open.return_value = get_mock_bw()
+    extractor = dataset.BigWigExtractor('mock_path')
+    values = extractor.extract(interval)
+    self.assertEqual(values.shape, expected_values.shape)
+    np.testing.assert_array_equal(values, expected_values)
+    extractor.close()
+    mock_bigwig_open.assert_called_once_with('mock_path')
+
+  @mock.patch.object(pyBigWig, 'open')
+  def test_bigwig_extractor_wrong_chromosome(self, mock_bigwig_open):
+    mock_bigwig_open.return_value = get_mock_bw()
+    extractor = dataset.BigWigExtractor('mock_path')
+    with self.assertRaisesRegex(ValueError, 'Chromosome chr2 not found'):
+      extractor.extract(genome.Interval('chr2', 0, 10))
+    extractor.close()
+    mock_bigwig_open.assert_called_once_with('mock_path')
+
+  @mock.patch.object(pyBigWig, 'open')
+  def test_multi_track_extractor(self, mock_bigwig_open):
+    mock_bigwig_open.return_value = get_mock_bw()
+    extractor = dataset.MultiTrackExtractor(
+        output_metadata=get_mock_metadata(
+            num_atac_tracks=3, num_dnase_tracks=4
+        ),
+        sequence_length=10,
+    )
+    interval = genome.Interval('chr1', 10, 20)
+    result = extractor.extract(interval)
+
+    chex.assert_shape(result['atac'], (10, 3))
+    self.assertEqual(result['atac'].dtype, tf.bfloat16)
+    chex.assert_shape(result['atac_mask'], (1, 3))
+    self.assertEqual(result['atac_mask'].dtype, tf.bool)
+    chex.assert_shape(result['dnase'], (10, 4))
+    chex.assert_shape(result['dnase_mask'], (1, 4))
+    extractor.close()
+    self.assertEqual(mock_bigwig_open.call_count, 7)
+
+  @parameterized.named_parameters([
+      dict(testcase_name='one_epoch', num_epochs=1),
+      dict(testcase_name='three_epochs', num_epochs=3),
+  ])
+  @mock.patch.object(fasta, 'FastaExtractor', MockFastaExtractor)
+  @mock.patch.object(pyBigWig, 'open')
+  def test_data_pipeline(self, mock_bigwig_open, num_epochs):
+    mock_bigwig_open.return_value = get_mock_bw()
+    intervals = pd.DataFrame({
+        'chromosome': ['chr1', 'chr1', 'chr1'],
+        'start': [10, 50, 90],
+        'end': [30, 70, 110],
+    })
+    sequence_length = 30
+    pipeline = dataset.DataPipeline(
+        fasta_path='mock_fasta',
+        intervals=intervals,
+        output_metadata=get_mock_metadata(
+            num_atac_tracks=2, num_dnase_tracks=1
+        ),
+        organism=dna_client.Organism.HOMO_SAPIENS,
+        sequence_length=sequence_length,
+    )
+
+    generator = pipeline.get_generator(num_epochs=num_epochs, shuffle=False)
+    elements = list(generator)
+    self.assertLen(elements, len(intervals) * num_epochs)
+
+    element = elements[0]
+    chex.assert_shape(element['dna_sequence'], (sequence_length, 4))
+    np.testing.assert_array_equal(
+        element['dna_sequence'],
+        np.tile(np.array([[1, 0, 0, 0]]), (sequence_length, 1)),  # 'A'
+    )
+    self.assertEqual(element['organism_index'], 0)
+    bundles = element['bundles']
+    chex.assert_shape(bundles['atac'], (sequence_length, 2))
+    chex.assert_shape(bundles['dnase'], (sequence_length, 1))
+
+    # The third interval (non shuffled) is half out of bounds.
+    element = elements[2]
+    bundles = element['bundles']
+    expected_track = np.concatenate([np.ones(15), np.zeros(15)])
+    np.testing.assert_array_equal(
+        bundles['atac'], np.stack([expected_track] * 2, axis=-1)
+    )
+    np.testing.assert_array_equal(
+        bundles['dnase'], expected_track.reshape(-1, 1)
+    )
+
+    pipeline.close()
+    self.assertEqual(mock_bigwig_open.call_count, 3)
+
+
+if __name__ == '__main__':
+  absltest.main()

flax_model/alphagenome/finetuning/finetune.py

@@ -0,0 +1,153 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""AlphaGenome finetuning script."""
+
+from collections.abc import Iterator, Mapping
+from typing import Any, Callable
+from flax_model.alphagenome._sdk.data import fold_intervals
+from flax_model.alphagenome._sdk.models import dna_model
+from flax_model.alphagenome.finetuning import dataset as dataset_lib
+from flax_model.alphagenome.model import model
+from flax_model.alphagenome.model import schemas
+from flax_model.alphagenome.model.metadata import metadata as metadata_lib
+import haiku as hk
+import jax
+import jmp
+import optax
+import tensorflow as tf
+
+_FASTA_PATH = (
+    'https://storage.googleapis.com/alphagenome/reference/gencode/'
+    'hg38/GRCh38.p13.genome.fa'
+)
+
+
+def get_dataset_iterator(
+    *,
+    batch_size: int,
+    sequence_length: int,
+    output_metadata: metadata_lib.AlphaGenomeOutputMetadata,
+    model_version: dna_model.ModelVersion,
+    subset: fold_intervals.Subset,
+    organism: dna_model.Organism = dna_model.Organism.HOMO_SAPIENS,
+    fasta_path: str = _FASTA_PATH,
+    example_regions_path: str | None = None,
+) -> Iterator[schemas.DataBatch]:
+  """Converts pipeline output dict to a DataBatch schema.
+
+  Args:
+    batch_size: The batch size of the dataset.
+    sequence_length: The sequence length of the dataset.
+    output_metadata: Metadata for the output tracks for each organism.
+    model_version: The model version to use.
+    subset: The subset of the dataset.
+    organism: The organism to use.
+    fasta_path: The path to the reference genome FASTA file.
+    example_regions_path: The path to the example regions BED file.
+
+  Returns:
+    A schemas.DataBatch object.
+  """
+  intervals = fold_intervals.get_fold_intervals(
+      model_version,
+      organism,
+      subset,
+      example_regions_path=example_regions_path,
+  )
+  pipeline = dataset_lib.DataPipeline(
+      fasta_path=fasta_path,
+      intervals=intervals,
+      output_metadata=output_metadata,
+      organism=organism,
+      sequence_length=sequence_length,
+  )
+  dataset = tf.data.Dataset.from_generator(
+      pipeline.get_generator,
+      output_signature=pipeline.get_output_signature(),
+  )
+  dataset = (
+      dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE).as_numpy_iterator()
+  )
+
+  def iterator():
+    for batch in dataset:
+      yield schemas.DataBatch(
+          dna_sequence=batch['dna_sequence'],
+          organism_index=batch['organism_index'],
+          **batch['bundles'],
+      )
+
+  return iterator()
+
+
+def get_forward_fn(
+    output_metadata: Mapping[
+        dna_model.Organism, metadata_lib.AlphaGenomeOutputMetadata
+    ],
+    jmp_policy: str = 'params=float32,compute=bfloat16,output=bfloat16',
+) -> hk.TransformedWithState:
+  """Creates a Haiku transformed function for the AlphaGenome model.
+
+  Args:
+    output_metadata: Metadata for the output tracks for each organism.
+    jmp_policy: The JMP policy to use for mixed precision.
+
+  Returns:
+    A `hk.TransformedWithState` object representing the forward pass.
+  """
+  jmp_policy = jmp.get_policy(jmp_policy)
+
+  @hk.transform_with_state
+  def forward(batch: schemas.DataBatch):
+    with hk.mixed_precision.push_policy(model.AlphaGenome, jmp_policy):
+      return model.AlphaGenome(
+          output_metadata, freeze_trunk_embeddings=True
+      ).loss(batch)
+
+  return forward
+
+
+def get_train_step(
+    predict_fn: Callable[..., Any],
+    optimizer: optax.GradientTransformation,
+):
+  """Creates a jitted training step function using AlphaGenome.loss.
+
+  Args:
+    predict_fn: The Haiku transformed forward function.
+    optimizer: An Optax optimizer to apply gradients.
+
+  Returns:
+    A jitted function `train_step` that takes `params`, `state`, `opt_state`,
+    and a `batch` as input and returns the updated `params`, `next_state`,
+    `new_opt_state`, and a dictionary of `metrics`.
+  """
+
+  @jax.jit
+  def train_step(params, state, opt_state, batch):
+    def loss_fn(params, state, batch):
+      (loss, scalars, predictions), new_state = predict_fn(
+          params, state, None, batch
+      )
+      del predictions  # Unused.
+      return loss, (new_state, scalars)
+
+    loss_grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    (loss, (next_state, scalars)), grads = loss_grad_fn(params, state, batch)
+    scalars['loss'] = loss
+    updates, new_opt_state = optimizer.update(grads, opt_state, params)
+    new_params = optax.apply_updates(params, updates)
+    return new_params, next_state, new_opt_state, scalars
+
+  return train_step