Add files using upload-large-folder tool

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/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/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/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/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/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/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

flax_model/alphagenome/finetuning/finetune_test.py

@@ -0,0 +1,179 @@
+# 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.
+
+import os
+
+from absl.testing import absltest
+from flax_model.alphagenome.finetuning import finetune
+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
+import jax.numpy as jnp
+import numpy as np
+import optax
+import orbax.checkpoint as ocp
+import pandas as pd
+
+
+def _create_mock_df(modality_name: str, num_tracks: int) -> pd.DataFrame | None:
+  if num_tracks == 0:
+    return None
+  return pd.DataFrame({
+      'file_path': [f'path{i}' for i in range(num_tracks)],
+      'name': [f'{modality_name}_track{i}' for i in range(num_tracks)],
+      'strand': ['.'] * num_tracks,
+  })
+
+
+def get_mock_metadata(
+    num_atac_tracks: int = 0, num_dnase_tracks: int = 0, rna_seq_tracks: int = 0
+) -> metadata_lib.AlphaGenomeOutputMetadata:
+  return metadata_lib.AlphaGenomeOutputMetadata(
+      atac=_create_mock_df('atac', num_atac_tracks),
+      dnase=_create_mock_df('dnase', num_dnase_tracks),
+      rna_seq=_create_mock_df('rna_seq', rna_seq_tracks),
+  )
+
+
+class FinetuneTest(absltest.TestCase):
+
+  def test_finetune_train_step(self):
+    seq_length, batch_size, key = 131072, 1, jax.random.key(0)
+
+    # Setup base model.
+    base_metadata = {
+        dna_model.Organism.HOMO_SAPIENS: get_mock_metadata(
+            num_atac_tracks=(atac_tracks_base := 2),
+            rna_seq_tracks=(rna_seq_tracks_base := 1),
+        )
+    }
+
+    @hk.transform_with_state
+    def base_forward(batch: schemas.DataBatch):
+      return model_lib.AlphaGenome(base_metadata).loss(batch)
+
+    # Batch contains ATAC and RNA-seq data.
+    base_batch = schemas.DataBatch(
+        dna_sequence=np.zeros((batch_size, seq_length, 4), dtype=np.float32),
+        organism_index=np.zeros((batch_size,), dtype=np.int32),
+        atac=jnp.zeros(
+            (batch_size, seq_length, atac_tracks_base), dtype=np.float32
+        ),
+        atac_mask=np.ones((batch_size, 1, atac_tracks_base), dtype=bool),
+        rna_seq=np.zeros(
+            (batch_size, seq_length, rna_seq_tracks_base), dtype=np.float32
+        ),
+        rna_seq_mask=np.ones((batch_size, 1, rna_seq_tracks_base), dtype=bool),
+    )
+    base_params, base_state = jax.eval_shape(base_forward.init, key, base_batch)
+
+    # Setup fine-tuning model with different metadata.
+    ft_metadata = {
+        dna_model.Organism.HOMO_SAPIENS: get_mock_metadata(
+            num_atac_tracks=(atac_tracks_ft := 3),  # More ATAC tracks.
+            rna_seq_tracks=0,  # No RNA-seq tracks.
+            num_dnase_tracks=(dnase_tracks_ft := 7),  # New DNase tracks.
+        )
+    }
+    batch_ft = schemas.DataBatch(
+        dna_sequence=np.zeros((batch_size, seq_length, 4), dtype=np.float32),
+        organism_index=np.zeros((batch_size,), dtype=np.int32),
+        atac=jnp.zeros(
+            (batch_size, seq_length, atac_tracks_ft), dtype=np.float32
+        ),
+        atac_mask=np.ones((batch_size, 1, atac_tracks_ft), dtype=bool),
+        dnase=jnp.zeros(
+            (batch_size, seq_length, dnase_tracks_ft), dtype=np.float32
+        ),
+        dnase_mask=np.ones((batch_size, 1, dnase_tracks_ft), dtype=bool),
+    )
+
+    forward_fn = finetune.get_forward_fn(ft_metadata)
+    params_ft, _ = jax.eval_shape(forward_fn.init, key, batch_ft)
+
+    def merge(frozen, trainable):
+      trainable = hk.data_structures.filter(
+          lambda module_name, *_: 'head' in module_name, trainable
+      )
+      merged = hk.data_structures.merge(frozen, trainable)
+      return merged
+
+    params = merge(base_params, params_ft)
+
+    optimizer = optax.adam(learning_rate=1e-3)
+    opt_state = jax.eval_shape(optimizer.init, params)
+    train_step = finetune.get_train_step(forward_fn.apply, optimizer)
+    _, _, _, scalars = jax.eval_shape(
+        train_step, params, base_state, opt_state, batch_ft
+    )
+    self.assertIn('loss', scalars)
+    self.assertIn('atac_loss', scalars)
+    self.assertIn('dnase_loss', scalars)
+    self.assertNotIn('rna_seq_loss', scalars)  # RNA-seq head not in FT model.
+
+  def test_create_finetuned_dna_sequence_model(self):
+    seq_length, batch_size, key = 131072, 1, jax.random.key(0)
+    ft_metadata = {
+        dna_model.Organism.HOMO_SAPIENS: get_mock_metadata(
+            num_atac_tracks=(atac_tracks_ft := 3),  # More ATAC tracks.
+            rna_seq_tracks=0,  # No RNA-seq tracks.
+            num_dnase_tracks=(dnase_tracks_ft := 7),  # New DNase tracks.
+        )
+    }
+    batch_ft = schemas.DataBatch(
+        dna_sequence=np.zeros((batch_size, seq_length, 4), dtype=np.float32),
+        organism_index=np.zeros((batch_size,), dtype=np.int32),
+        atac=jnp.zeros(
+            (batch_size, seq_length, atac_tracks_ft), dtype=np.float32
+        ),
+        atac_mask=np.ones((batch_size, 1, atac_tracks_ft), dtype=bool),
+        dnase=jnp.zeros(
+            (batch_size, seq_length, dnase_tracks_ft), dtype=np.float32
+        ),
+        dnase_mask=np.ones((batch_size, 1, dnase_tracks_ft), dtype=bool),
+    )
+
+    forward_fn = finetune.get_forward_fn(ft_metadata)
+    params_ft, state_ft = jax.eval_shape(forward_fn.init, key, batch_ft)
+
+    checkpointer = ocp.StandardCheckpointer()
+    ckpt_dir = self.create_tempdir().full_path
+    ckpt_path = os.path.join(ckpt_dir, 'checkpoint')
+    checkpointer.save(
+        ckpt_path,
+        jax.tree_util.tree_map(
+            lambda x: jnp.empty(x.shape, x.dtype), (params_ft, state_ft)
+        ),
+    )
+    checkpointer.wait_until_finished()
+
+    ft_organism_settings = {
+        k: dna_model.OrganismSettings(metadata=v)
+        for k, v in ft_metadata.items()
+    }
+    finetuned_dna_sequence_model = dna_model.create(
+        ckpt_path,
+        organism_settings=ft_organism_settings,
+        device=jax.devices('cpu')[0],
+    )
+    self.assertIsInstance(
+        finetuned_dna_sequence_model, dna_model.AlphaGenomeModel
+    )
+
+
+if __name__ == '__main__':
+  absltest.main()

flax_model/alphagenome/model/metadata/metadata_test.py

@@ -0,0 +1,215 @@
+# 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 collections.abc import Mapping, Sequence
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax_model.alphagenome._sdk.data import ontology
+from flax_model.alphagenome._sdk.models import dna_model
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome.model.metadata import metadata as metadata_lib
+import chex
+import jax
+import numpy as np
+import pandas as pd
+
+
+class MetadataTest(parameterized.TestCase):
+
+  @parameterized.parameters([
+      dict(organism=dna_model.Organism.HOMO_SAPIENS),
+      dict(organism=dna_model.Organism.MUS_MUSCULUS),
+  ])
+  def test_load(self, organism: dna_model.Organism):
+    metadata = metadata_lib.load(organism)
+    self.assertIsInstance(metadata, dna_output.OutputMetadata)
+    reindexing = metadata.strand_reindexing
+    self.assertIsInstance(reindexing, dict)
+    self.assertCountEqual(dna_output.OutputType, reindexing.keys())
+
+  def test_organism_shapes_match(self):
+    metadata = tuple(metadata_lib.load(o) for o in dna_model.Organism)
+    for output_type in dna_output.OutputType:
+      with self.subTest(f'{output_type=}'):
+        track_metadata = tuple(m.get(output_type).values for m in metadata)
+        # Only assert the first track dimension matches, as fields may differ.
+        chex.assert_equal_shape_prefix(track_metadata, 1)
+
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='Normal',
+          output_type=dna_output.OutputType.ATAC,
+          names=[f'track_{i}' for i in range(6)],
+          strands=['+', '+', '+', '-', '-', '-'],
+          expected=[3, 4, 5, 0, 1, 2],
+      ),
+      dict(
+          testcase_name='Unstranded',
+          output_type=dna_output.OutputType.RNA_SEQ,
+          names=[f'track_{i}' for i in range(6)],
+          strands=['.', '.', '.', '.', '.', '.'],
+          expected=[0, 1, 2, 3, 4, 5],
+      ),
+      dict(
+          testcase_name='Interleaved',
+          output_type=dna_output.OutputType.RNA_SEQ,
+          names=[f'track_{i}' for i in range(6)],
+          strands=['+', '-', '+', '-', '+', '-'],
+          expected=[1, 0, 3, 2, 5, 4],
+      ),
+      dict(
+          testcase_name='Padding',
+          output_type=dna_output.OutputType.DNASE,
+          names=[f'track_{i}' for i in range(6)] + ['Padding', 'padding'],
+          strands=['+', '-', '+', '-', '+', '-', '.', '.'],
+          expected=[1, 0, 3, 2, 5, 4, 6, 7],
+      ),
+      dict(
+          testcase_name='SpliceJunctions',
+          output_type=dna_output.OutputType.SPLICE_JUNCTIONS,
+          names=[f'track_{i}' for i in range(3)],
+          strands=['.'] * 3,
+          expected=[3, 4, 5, 0, 1, 2],
+      ),
+      dict(
+          testcase_name='SpliceJunctionPadding',
+          output_type=dna_output.OutputType.SPLICE_JUNCTIONS,
+          names=[f'track_{i}' for i in range(3)] + ['PADDING'],
+          strands=['.'] * 4,
+          expected=[4, 5, 6, 7, 0, 1, 2, 3],
+      ),
+  ])
+  def test_strand_reindexing(
+      self,
+      output_type: dna_output.OutputType,
+      names: Sequence[str],
+      strands: Sequence[str],
+      expected: Sequence[int],
+  ):
+    metadata = metadata_lib.AlphaGenomeOutputMetadata(
+        **{
+            output_type.name.lower(): pd.DataFrame(
+                {'name': names, 'strand': strands}
+            ),
+        },
+    )
+    result = metadata.strand_reindexing.get(output_type)
+    result = result.tolist() if result is not None else None
+    self.assertEqual(result, expected)
+
+  @parameterized.named_parameters(
+      (
+          'Normal',
+          dna_output.OutputType.ATAC,
+          [f'track_{i}' for i in range(6)],
+          ['.'] * 6,
+          [False] * 6,
+      ),
+      (
+          'Padding',
+          dna_output.OutputType.DNASE,
+          [f'track_{i}' for i in range(3)] + ['Padding', 'Padding'],
+          ['.'] * 3 + ['.', '.'],
+          [False] * 3 + [True, True],
+      ),
+  )
+  def test_padding(
+      self,
+      output_type: dna_output.OutputType,
+      names: Sequence[str],
+      strands: Sequence[str],
+      expected: Sequence[int],
+  ):
+    metadata = metadata_lib.AlphaGenomeOutputMetadata(
+        **{
+            output_type.name.lower(): pd.DataFrame(
+                {'name': names, 'strand': strands}
+            )
+        },
+    )
+    np.testing.assert_array_equal(metadata.padding.get(output_type), expected)
+
+  @parameterized.named_parameters(
+      (
+          'SingleOutput',
+          [dna_output.OutputType.ATAC],
+          None,
+          {dna_output.OutputType.ATAC: [True, True, True]},
+      ),
+      (
+          'SingleOntology',
+          [dna_output.OutputType.ATAC],
+          [ontology.from_curie('CL:0000084')],
+          {dna_output.OutputType.ATAC: [True, True, False]},
+      ),
+      (
+          'MultipleOutputsMissingOntology',
+          [dna_output.OutputType.ATAC, dna_output.OutputType.DNASE],
+          [ontology.from_curie('CL:0000001')],
+          {
+              dna_output.OutputType.ATAC: [False, False, False],
+              dna_output.OutputType.DNASE: [True, True, True, False, False],
+          },
+      ),
+      (
+          'SpliceSites',
+          [dna_output.OutputType.SPLICE_SITES],
+          [ontology.from_curie('CL:0000001')],
+          {
+              dna_output.OutputType.SPLICE_SITES: [True, True, True],
+          },
+      ),
+      (
+          'AllOutputs',
+          list(dna_output.OutputType),
+          None,
+          {
+              dna_output.OutputType.ATAC: [True, True, True],
+              dna_output.OutputType.DNASE: [True, True, True, False, False],
+              dna_output.OutputType.SPLICE_SITES: [True, True, True],
+          },
+      ),
+  )
+  def test_create_track_masks(
+      self,
+      requested_outputs: Sequence[dna_output.OutputType],
+      requested_ontologies: Sequence[str] | None,
+      expected: Mapping[dna_output.OutputType, ...],
+  ):
+    example_metadata = metadata_lib.AlphaGenomeOutputMetadata(
+        atac=pd.DataFrame({
+            'name': ['track_1', 'track_2', 'track_3'],
+            'strand': ['+', '-', '.'],
+            'ontology_curie': ['CL:0000084', 'CL:0000084', 'UBERONE:0000001'],
+        }),
+        dnase=pd.DataFrame({
+            'name': [f'track_{i}' for i in range(3)] + ['Padding', 'Padding'],
+            'strand': ['+'] * 3 + ['.', '.'],
+            'ontology_curie': ['CL:0000001'] * 3 + ['', ''],
+        }),
+        splice_sites=pd.DataFrame({
+            'name': [f'track_{i}' for i in range(3)],
+            'strand': ['.'] * 3,
+        }),
+    )
+    track_masks = metadata_lib.create_track_masks(
+        example_metadata,
+        requested_outputs=requested_outputs,
+        requested_ontologies=requested_ontologies,
+    )
+    jax.tree.map(np.testing.assert_array_equal, track_masks, expected)
+
+
+if __name__ == '__main__':
+  absltest.main()

flax_model/alphagenome/model/variant_scoring/__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.
+
+"""AlphaGenome Variant Scoring."""

flax_model/alphagenome/model/variant_scoring/splice_junction.py

@@ -0,0 +1,306 @@
+# 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.
+
+"""Implementation of splice junction variant scoring."""
+
+from flax_model.alphagenome._sdk import typing
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome._sdk.data import junction_data
+from flax_model.alphagenome._sdk.models import dna_output
+from flax_model.alphagenome._sdk.models import variant_scorers
+from flax_model.alphagenome.model.variant_scoring import gene_mask_extractor
+from flax_model.alphagenome.model.variant_scoring import variant_scoring
+import anndata
+import einshape
+import jax
+import jax.numpy as jnp
+from jaxtyping import Float, Int32, Shaped  # pylint: disable=g-multiple-import, g-importing-member
+import numpy as np
+import pandas as pd
+import pyranges
+
+MAX_SPLICE_SITES = 256
+PAD_VALUE = -1
+
+
+def _create_empty(mask_metadata: pd.DataFrame, track_metadata: pd.DataFrame):
+  """Create empty AnnData object for splice junction scoring."""
+  junction_columns = [
+      'junction_Start',
+      'junction_End',
+  ]
+  return variant_scoring.create_anndata(
+      np.zeros((0, len(track_metadata['name'])), dtype=np.float32),
+      obs=pd.DataFrame(columns=list(mask_metadata.columns) + junction_columns),
+      var=pd.DataFrame({
+          'strand': '.',
+          'name': track_metadata['name'],
+          'gtex_tissue': track_metadata.get('gtex_tissue'),
+          'ontology_curie': track_metadata.get('ontology_curie'),
+          'biosample_type': track_metadata.get('biosample_type'),
+          'biosample_name': track_metadata.get('biosample_name'),
+          'biosample_life_stage': track_metadata.get('biosample_life_stage'),
+          'data_source': track_metadata.get('data_source'),
+          'Assay title': track_metadata.get('Assay title'),
+      }),
+  )
+
+
+def _create(
+    junction_scores: pd.DataFrame,
+    mask_metadata: pd.DataFrame,
+    track_metadata: pd.DataFrame,
+) -> anndata.AnnData:
+  """Converts a dataframe of junction scores to an AnnData object."""
+  if mask_metadata.empty or junction_scores.empty:
+    raise ValueError('Both junction_scores and mask_metadata must be non-empty')
+
+  junction_scores = junction_scores[
+      junction_scores['gene_id'].isin(mask_metadata['gene_id'])
+  ]
+
+  gene_max_scores = []
+  track_names = track_metadata['name']
+  for gene_id in junction_scores['gene_id'].unique():
+    gene_junction_scores = junction_scores[
+        junction_scores.gene_id == gene_id
+    ].reset_index(drop=True)
+    gene_junction_scores = gene_junction_scores.iloc[
+        gene_junction_scores[track_names].values.argmax(0)
+    ]
+    gene_max_scores.append(gene_junction_scores)
+  junction_scores = pd.concat(gene_max_scores)
+  score_values = junction_scores[track_names].values
+
+  # Merge junction information with mask metadata.
+  junctions_all_genes = junction_scores[['gene_id', 'Start', 'End']]
+  junctions_all_genes.columns = ['gene_id', 'junction_Start', 'junction_End']
+  mask_metadata = junctions_all_genes.merge(
+      mask_metadata, on='gene_id', sort=False
+  )
+  # Create the final track metadata.
+  track_metadata = pd.DataFrame({
+      'strand': '.',  # We already matched prediction by strand.
+      'name': track_names,
+      'gtex_tissue': track_metadata['gtex_tissue'],
+      'ontology_curie': track_metadata.get('ontology_curie'),
+      'biosample_type': track_metadata.get('biosample_type'),
+      'biosample_name': track_metadata.get('biosample_name'),
+      'biosample_life_stage': track_metadata.get('biosample_life_stage'),
+      'data_source': track_metadata.get('data_source'),
+      'Assay title': track_metadata.get('Assay title'),
+  })
+  ann_data = variant_scoring.create_anndata(
+      score_values,
+      obs=mask_metadata,
+      var=track_metadata,
+  )
+  # Remove duplicated junctions. Per gene, we report junctions that has maximum
+  # score in any tissue. For the reported junctions, we return their predictions
+  # in all tissues.
+  return ann_data[~ann_data.obs.duplicated()].copy()
+
+
+@typing.jaxtyped
+def unstack_junction_predictions(
+    splice_junction_prediction: Float[np.ndarray, 'D D _'],
+    splice_site_positions: Int32[np.ndarray, '4 D'],
+    interval: genome.Interval | None = None,
+) -> tuple[
+    Float[np.ndarray, 'num_junctions num_tracks'],
+    Shaped[np.ndarray, 'num_junctions'],
+    Int32[np.ndarray, 'num_junctions'],
+    Int32[np.ndarray, 'num_junctions'],
+]:
+  """Unstack splice junction predictions to long format."""
+  # Unpack splice junction predictions.
+  splice_junction_prediction = einshape.numpy_einshape(
+      'da(st)->dast', splice_junction_prediction, s=2
+  )
+  # Convert splice site positions.
+  remove_padding_fn = lambda x: x[x != PAD_VALUE]
+  pos_donors = remove_padding_fn(splice_site_positions[0])
+  pos_acceptors = remove_padding_fn(splice_site_positions[1])
+  neg_donors = remove_padding_fn(splice_site_positions[2])
+  neg_acceptors = remove_padding_fn(splice_site_positions[3])
+  junction_pred_pos = splice_junction_prediction[
+      : len(pos_donors), : len(pos_acceptors), 0
+  ]
+  junction_pred_pos = einshape.numpy_einshape('dat->(da)t', junction_pred_pos)
+  num_pos_donors = len(pos_donors)
+  pos_donors = np.repeat(pos_donors, len(pos_acceptors))
+  pos_acceptors = np.tile(pos_acceptors, num_pos_donors)
+  junction_pred_neg = splice_junction_prediction[
+      : len(neg_donors), : len(neg_acceptors), 1
+  ]
+  junction_pred_neg = einshape.numpy_einshape('dat->(da)t', junction_pred_neg)
+  num_neg_donors = len(neg_donors)
+  neg_donors = np.repeat(neg_donors, len(neg_acceptors))
+  neg_acceptors = np.tile(neg_acceptors, num_neg_donors)
+  # Combine into  a single output.
+  junction_predictions = np.concatenate(
+      [junction_pred_pos, junction_pred_neg], axis=0
+  )
+  # Junction start and end positions.
+  starts = np.concatenate([pos_donors, neg_acceptors]) + 1
+  starts += interval.start if interval is not None else 0
+  ends = np.concatenate([pos_acceptors, neg_donors])
+  ends += interval.start if interval is not None else 0
+  strands = np.array(['+'] * len(pos_donors) + ['-'] * len(neg_donors))
+  filter_mask = (starts < ends) & (starts > 0)
+
+  return (
+      junction_predictions[filter_mask],
+      strands[filter_mask],
+      starts[filter_mask],
+      ends[filter_mask],
+  )
+
+
+def junction_predictions_to_dataframe(
+    splice_junction_prediction: Float[np.ndarray, 'D D _'],
+    splice_site_positions: Int32[np.ndarray, 'T_mul_4 D'],
+    metadata: junction_data.JunctionMetadata,
+    interval: genome.Interval,
+) -> pd.DataFrame:
+  """Convert splice junction predictions to a dataframe."""
+  junction_predictions, strands, starts, ends = unstack_junction_predictions(
+      splice_junction_prediction, splice_site_positions, interval
+  )
+  junctions = pd.DataFrame({
+      'Chromosome': interval.chromosome,
+      'Start': starts,
+      'End': ends,
+      'Strand': strands,
+  })
+  predictions = pd.DataFrame(junction_predictions, columns=metadata['name'])
+  return pd.concat([junctions, predictions], axis=1)
+
+
+class SpliceJunctionVariantScorer(variant_scoring.VariantScorer):
+  """Implements the SpliceJunction variant scoring strategy.
+
+  Scores variants by the maximum of absolute delta pair counts of junctions
+  within the input interval. Junctions are annotated by overlapping with the
+  gtf gene intervals.
+  """
+
+  def __init__(self, gtf: pd.DataFrame):
+    self._gene_mask_extractor = gene_mask_extractor.GeneMaskExtractor(
+        gtf,
+        gene_mask_extractor.GeneMaskType.BODY,
+        gene_query_type=gene_mask_extractor.GeneQueryType.VARIANT_OVERLAPPING,
+        filter_protein_coding=True,
+    )
+
+  def get_masks_and_metadata(
+      self,
+      interval: genome.Interval,
+      variant: genome.Variant,
+      *,
+      settings: variant_scorers.SpliceJunctionScorer,
+      track_metadata: dna_output.OutputMetadata,
+  ) -> tuple[None, pd.DataFrame]:
+    """See base class."""
+    del settings, track_metadata
+    _, metadata = self._gene_mask_extractor.extract(interval, variant)
+    metadata['interval'] = interval
+    return None, metadata
+
+  def score_variant(
+      self,
+      ref: variant_scoring.ScoreVariantInput,
+      alt: variant_scoring.ScoreVariantInput,
+      *,
+      masks: None,
+      settings: variant_scorers.SpliceJunctionScorer,
+      variant: genome.Variant | None = None,
+      interval: genome.Interval | None = None,
+  ) -> variant_scoring.ScoreVariantOutput:
+    """See base class."""
+    del variant, interval, masks
+    ref_junctions = ref[settings.requested_output]['predictions']
+    alt_junctions = alt[settings.requested_output]['predictions']
+
+    splice_site_positions = ref[settings.requested_output][
+        'splice_site_positions'
+    ]
+
+    # JAX dynamic slicing does not work with transfer_guard.
+    with jax.transfer_guard('allow'):
+      # Ignore splice sites beyond the max_splice_sites specified. This works
+      # because padding splice sites are always at the end of the array.
+      ref_junctions = ref_junctions[:MAX_SPLICE_SITES, :MAX_SPLICE_SITES]
+      alt_junctions = alt_junctions[:MAX_SPLICE_SITES, :MAX_SPLICE_SITES]
+      splice_site_positions = splice_site_positions[:, :MAX_SPLICE_SITES]
+
+    @jax.jit
+    def _apply_log_offset(x):
+      return jnp.log(x + 1e-7)
+
+    ref_junctions = _apply_log_offset(ref_junctions)
+    alt_junctions = _apply_log_offset(alt_junctions)
+
+    return {
+        'delta_counts': (alt_junctions - ref_junctions).astype(jnp.float16),
+        'splice_site_positions': splice_site_positions,
+    }
+
+  def finalize_variant(
+      self,
+      scores: variant_scoring.ScoreVariantResult,
+      *,
+      track_metadata: dna_output.OutputMetadata,
+      mask_metadata: pd.DataFrame,
+      settings: variant_scorers.SpliceJunctionScorer,
+  ) -> anndata.AnnData:
+    """See base class."""
+    track_metadata = track_metadata.get(settings.requested_output)
+
+    if mask_metadata.empty:
+      return _create_empty(mask_metadata, track_metadata)
+
+    delta_counts = scores['delta_counts']
+
+    interval = mask_metadata['interval'].values[0]
+    mask_metadata = mask_metadata.drop(columns=['interval'])
+
+    delta_counts = junction_predictions_to_dataframe(
+        np.abs(delta_counts, dtype=np.float32),
+        scores['splice_site_positions'],
+        metadata=track_metadata,
+        interval=interval,
+    )
+    if delta_counts.empty:
+      return _create_empty(mask_metadata, track_metadata)
+
+    junction_scores = (
+        pyranges.PyRanges(delta_counts)
+        .join(
+            pyranges.PyRanges(
+                mask_metadata.rename(columns={'strand': 'Strand'})
+            ),
+            strandedness='same',
+        )
+        .df
+    )
+
+    if not junction_scores.empty:
+      junction_scores = junction_scores[
+          (junction_scores['Start'] > junction_scores['Start_b'])
+          & (junction_scores['End'] < junction_scores['End_b'])
+      ]
+      return _create(junction_scores, mask_metadata, track_metadata)
+    else:
+      return _create_empty(mask_metadata, track_metadata)

flax_model/alphagenome/model/variant_scoring/variant_scoring_test.py

@@ -0,0 +1,76 @@
+# 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 absl.testing import absltest
+from absl.testing import parameterized
+from flax_model.alphagenome._sdk.data import genome
+from flax_model.alphagenome.model.variant_scoring import variant_scoring
+import chex
+import numpy as np
+
+
+class AlignAlternateTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='snp',
+          variant=genome.Variant('chr1', 2, 'C', 'T'),
+          expected_aligned_alt=np.arange(10, dtype=np.float32).reshape(-1, 1),
+      ),
+      dict(
+          testcase_name='insertion',
+          variant=genome.Variant('chr1', 2, 'C', 'TGA'),
+          expected_aligned_alt=np.array(
+              [0, 3, 4, 5, 6, 7, 8, 9, 0, 0], dtype=np.float32
+          ).reshape(-1, 1),
+      ),
+      dict(
+          testcase_name='multi_insertion',
+          variant=genome.Variant('chr1', 2, 'CC', 'CCTGA'),
+          expected_aligned_alt=np.array(
+              [0, 1, 5, 6, 7, 8, 9, 0, 0, 0], dtype=np.float32
+          ).reshape(-1, 1),
+      ),
+      dict(
+          testcase_name='multi_insertion_long_insertion',
+          variant=genome.Variant('chr1', 5, 'A', 'AAAAAAA'),
+          expected_aligned_alt=np.array(
+              [0, 1, 2, 3, 9, 0, 0, 0, 0, 0], dtype=np.float32
+          ).reshape(-1, 1),
+      ),
+      dict(
+          testcase_name='deletion',
+          variant=genome.Variant('chr1', 2, 'CGA', 'C'),
+          expected_aligned_alt=np.array(
+              [0, 1, 0, 0, 2, 3, 4, 5, 6, 7], dtype=np.float32
+          ).reshape(-1, 1),
+      ),
+      dict(
+          testcase_name='multi_deletion',
+          variant=genome.Variant('chr1', 2, 'CCCGA', 'CCC'),
+          expected_aligned_alt=np.array(
+              [0, 1, 2, 3, 0, 0, 4, 5, 6, 7], dtype=np.float32
+          ).reshape(-1, 1),
+      ),
+  )
+  def test_align_alt(self, variant, expected_aligned_alt):
+    interval = genome.Interval('chr1', 0, 10)
+    alt = np.arange(10, dtype=np.float32).reshape(-1, 1)
+    aligned_alt = variant_scoring.align_alternate(alt, variant, interval)
+    chex.assert_shape(aligned_alt, (10, 1))
+    np.testing.assert_array_equal(aligned_alt, expected_aligned_alt)
+
+
+if __name__ == '__main__':
+  absltest.main()

scripts/inference.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+if [[ -n "${ONESCIENCE_DATASETS_DIR:-}" ]]; then
+  DATA_ROOT_DIR="${ONESCIENCE_DATASETS_DIR}/AlphaGenome"
+else
+  DATA_ROOT_DIR="${PROJECT_DIR}/data"
+fi
+
+if [[ -n "${ONESCIENCE_MODELS_DIR:-}" ]]; then
+  MODEL_ROOT_DIR="${ONESCIENCE_MODELS_DIR}/AlphaGenome"
+else
+  MODEL_ROOT_DIR="${PROJECT_DIR}/weight"
+fi
+
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+
+python "${SCRIPT_DIR}/run_inference.py" \
+    --fasta_path "${DATA_ROOT_DIR}/reference/HOMO_SAPIENS/GRCh38.p13.genome.fa" \
+    --model_dir "${MODEL_ROOT_DIR}/alphagenome-all-folds" \
+    --chromosome chr19 \
+    --start 10587331 \
+    --end 11635907 \
+    --output_dir "${PROJECT_DIR}/outputs"

weight/README.md

@@ -0,0 +1,13 @@
+# AlphaGenome Weights
+
+Place local AlphaGenome Orbax checkpoints here when not using the shared OneScience model directory.
+
+Expected local layout:
+
+```text
+weight/
+└── alphagenome-all-folds/
+    ├── _CHECKPOINT_METADATA
+    ├── _METADATA
+    └── ...
+```

weight/alphagenome-all-folds/README.md

@@ -0,0 +1,334 @@
+---
+license: other
+license_name: alphagenome
+license_link: https://deepmind.google.com/science/alphagenome/model-terms
+extra_gated_heading: Access AlphaGenome on Hugging Face
+extra_gated_prompt: >-
+  AlphaGenome is provided for non-commercial use only and is subject to the
+  [Model Terms of Use](https://deepmind.google.com/science/alphagenome/model-terms).
+  To accept terms, please login, complete the required fields and click Accept.
+  Requests are processed immediately.
+extra_gated_button_content: Accept and continue
+extra_gated_fields:
+  Organization, university, or other affiliation(s): text
+language:
+- en
+tags:
+- biology
+---
+
+## Description
+
+AlphaGenome is a unified DNA sequence model designed to advance regulatory
+variant-effect prediction and shed light on genome function. It analyzes DNA
+sequences of up to 1 million base pairs to deliver predictions at single
+base-pair resolution across diverse modalities, including gene expression,
+splicing patterns, chromatin features, and contact maps.
+
+The core architecture features a U-Net-style design that combines an encoder for
+downsampling, transformers with inter-device communication to capture long-range
+interactions, and a decoder for upsampling. These components feed into
+task-specific output heads that generate predictions at their respective
+assay-specific resolutions.
+
+By achieving state-of-the-art performance across diverse genomic benchmarks,
+AlphaGenome provides a robust framework for understanding the molecular function
+of DNA sequences and interpreting non-coding variation.
+
+## Inputs and outputs
+
+*   **Input:** Up to 1 Mb (2\*\*20) one-hot encoded DNA sequence, and an
+    organism type index (representing human or mouse).
+*   **Output:** 11 diverse modalities, including:
+    *   RNA expression (RNA-Seq, CAGE-seq and PRO-cap).
+    *   Chromatin accessibility (DNase-seq and ATAC-seq).
+    *   Histone modifications.
+    *   Transcription factor binding.
+    *   Chromatin contact maps.
+    *   Splice sites and their usage.
+    *   Splice junction coordinates and strength.
+
+## Citation
+
+```
+@article{alphagenome,
+  title={Advancing regulatory variant effect prediction with {AlphaGenome}},
+  author={Avsec, {\v Z}iga and Latysheva, Natasha and Cheng, Jun and Novati, Guido and Taylor, Kyle R. and Ward, Tom and Bycroft, Clare and Nicolaisen, Lauren and Arvaniti, Eirini and Pan, Joshua and Thomas, Raina and Dutordoir, Vincent and Perino, Matteo and De, Soham and Karollus, Alexander and Gayoso, Adam and Sargeant, Toby and Mottram, Anne and Wong, Lai Hong and Drot{\'a}r, Pavol and Kosiorek, Adam and Senior, Andrew and Tanburn, Richard and Applebaum, Taylor and Basu, Souradeep and Hassabis, Demis and Kohli, Pushmeet},
+  journal={Nature},
+  volume={649},
+  number={8099},
+  year={2026},
+  doi={10.1038/s41586-025-10014-0},
+  publisher={Nature Publishing Group UK London}
+}
+```
+
+## Installation
+
+To install the accompanying code necessary to run the model, please run the
+following:
+
+```shell
+$ pip install git+https://github.com/google-deepmind/alphagenome_research.git
+```
+
+## Usage
+
+In addition to the model, we provide a DNA model class that wraps the core model
+and provides a more intuitive set of functions for creating predictions, scoring
+variants, performing in silico mutagenesis (ISM) and more.
+
+Here's an example of making a variant prediction:
+
+```python
+from alphagenome.data import genome
+from alphagenome.visualization import plot_components
+from alphagenome_research.model import dna_model
+import matplotlib.pyplot as plt
+
+model = dna_model.create_from_huggingface('all_folds')
+
+interval = genome.Interval(chromosome='chr22', start=35677410, end=36725986)
+variant = genome.Variant(
+    chromosome='chr22',
+    position=36201698,
+    reference_bases='A',
+    alternate_bases='C',
+)
+
+outputs = model.predict_variant(
+    interval=interval,
+    variant=variant,
+    ontology_terms=['UBERON:0001157'],
+    requested_outputs=[dna_model.OutputType.RNA_SEQ],
+)
+
+plot_components.plot(
+    [
+        plot_components.OverlaidTracks(
+            tdata={
+                'REF': outputs.reference.rna_seq,
+                'ALT': outputs.alternate.rna_seq,
+            },
+            colors={'REF': 'dimgrey', 'ALT': 'red'},
+        ),
+    ],
+    interval=outputs.reference.rna_seq.interval.resize(2**15),
+    # Annotate the location of the variant as a vertical line.
+    annotations=[plot_components.VariantAnnotation([variant], alpha=0.8)],
+)
+plt.show()
+```
+
+## Model Data
+
+AlphaGenome was trained to predict read coverage for a wide range of different
+functional genomics assays – including RNA-seq, DNase-seq, CAGE, ChIP-seq
+– directly from human and mouse reference genome sequences. The training
+data was sourced from large public consortia including
+[ENCODE](http://encodeproject.org/), [GTEx](https://www.gtexportal.org/),
+[4D Nucleome](https://4dnucleome.org/) and
+[FANTOM5](https://fantom.gsc.riken.jp/5/), encompassing experimental
+measurements of key regulatory modalities across hundreds of cell types and
+tissues. For more information, please refer to the
+[AlphaGenome manuscript](https://www.biorxiv.org/content/10.1101/2025.06.25.661532v2).
+
+Along with the release of the weights and model inference code, we are releasing
+the complete training, validation and test datasets used for AlphaGenome. The
+data is stored in compressed TFRecord format and can be loaded as follows:
+
+```python
+from alphagenome.data import fold_intervals
+from alphagenome_research.io import dataset
+from alphagenome_research.model import dna_model
+
+ds_iter = dataset.create_dataset(
+    organism=dna_model.Organism.HOMO_SAPIENS,
+    fold_split=dna_model.ModelVersion.ALL_FOLDS,
+    subset=fold_intervals.Subset.TRAIN,
+).as_numpy_iterator()
+element = next(ds_iter)
+```
+
+Each element in the dataset is a tuple, where each item corresponds to a
+specific data bundle. The example below illustrates this structure, showing the
+ATAC and DNase bundles alongside the input DNA sequence, masks, and interval
+metadata:
+
+```python
+({"atac": "bfloat16[1052672,256]",
+  "atac_mask": "bool[1,256]",
+  "dna_sequence": "float32[1052672,4]",
+  "interval/chromosome": b"chr7",
+  "interval/end": "int64[]",
+  "interval/start": "int64[]"},
+ {"dna_sequence": "float32[1052672,4]",
+  "dnase": "bfloat16[1052672,384]",
+  "dnase_mask": "bool[1,384]",
+  "interval/chromosome": b"chr7",
+  "interval/end": "int64[]",
+  "interval/start": "int64[]"},
+...
+)
+```
+
+**Notes:**
+
+*   **Interval extension**: To support data augmentation via sequence shifting
+    during training, the 1 Mb intervals were extended by 4,096 base pairs (2,048
+    bp on each side).
+*   **GTEx data exclusion**: The released dataset does not contain GTEx tissue
+    data due to licensing restrictions, though the corresponding column headers
+    remain in the bundle’s metadata. However, users can easily extend the
+    dataset with GTEx or other external sources using the provided interval
+    metadata (chromosome, start, end).
+
+## Implementation Information
+
+### Hardware
+
+AlphaGenome training involved a two-stage process: pre-training and
+distillation. Pre-training was carried out on 256
+[Tensor Processing Units (TPUv3)](https://docs.cloud.google.com/tpu/docs/v3)
+using sequence parallelism across groups of 4 interconnected chips. We leverage
+TPU pods (large clusters of TPUs) to provide a scalable solution for training,
+to enable large batch sizes which can lead to better model quality. For
+distillation, AlphaGenome was trained on 64 NVIDIA H100 GPUs without sequence
+parallelism. The evaluation of all models was carried out on NVIDIA H100 GPUs
+without sequence parallelism.
+
+### Software
+
+Training was done using [JAX](https://github.com/google/jax) and
+[JAXline](https://github.com/google-deepmind/jaxline).
+
+JAX allows researchers to take advantage of the latest generation of hardware,
+including TPUs, for faster and more efficient training of large models.
+
+JAXline is a distributed JAX training and evaluation framework. It is designed
+to be forked, covering only the most general aspects of experiment boilerplate.
+
+## Evaluation
+
+AlphaGenome demonstrates state-of-the-art performance across a diverse set of
+genomic prediction tasks. The model was evaluated using two primary approaches:
+
+1.  **Genome Track Prediction**: Assessing the ability to predict functional
+    genomic signals (read coverage) on previously unseen DNA sequences (held-out
+    test intervals).
+2.  **Variant Effect Prediction (VEP):** Assessing the ability to predict the
+    molecular consequences of genetic variants (e.g., single nucleotide
+    variants) by comparing predictions for reference and alternative alleles
+    against ground-truth datasets (e.g., experimental QTL effect sizes, readouts
+    from reporter assays).
+
+Key Highlights:
+
+*   **Broad SOTA Performance:** AlphaGenome matched or outperformed the best
+    available external models on **22 out of 24** genome track prediction
+    evaluations.
+*   **Variant Interpretation:** For variant effect prediction, AlphaGenome
+    matched or exceeded top-performing external models on **25 out of 26**
+    evaluations.
+*   **Multimodal Capability:** Unlike specialized models, AlphaGenome jointly
+    predicts all assessed modalities – including splicing, expression,
+    accessibility, and 3D contact maps – within a single framework.
+*   **Single-Pass Efficiency:** The distilled student model achieves this
+    performance and broad coverage with a single inference pass, eliminating the
+    need for complex model ensembling.
+
+The tables below detail the performance metrics for specific modalities and
+tasks.
+
+## Benchmark Results
+
+The following table focuses on the accuracy of the pre-trained, non-distilled
+model in predicting genomic tracks on unseen sequences:
+
+Modality              | Evaluation                       | Metric      | Resolution | Value | Baseline Model | Relative Improvement (%)
+:-------------------- | :------------------------------- | :---------- | :--------- | :---- | :------------- | :-----------------------
+**Splicing**          | Splice site classification       | auPRC       | 1 bp       | 0.79  | DeltaSplice    | 1.0
+                 | Splice site usage                | Pearson r   | 1 bp       | 0.86  | DeltaSplice    | 6.7
+**RNA expression**    | RNA-seq coverage                 | Pearson r   | 1 bp       | 0.59  | Borzoi         | 28.2
+                 | RNA-seq coverage                 | Pearson r   | 32 bp      | 0.78  | Borzoi         | 4.6
+                 | RNA-seq gene expr. LFC           | Pearson r   | Gene       | 0.57  | Borzoi         | 14.7
+                 | CAGE coverage                    | Pearson r   | 32 bp      | 0.74  | Borzoi         | 4.4
+                 | CAGE coverage                    | Pearson r   | 128 bp     | 0.71  | Enformer       | \-0.3
+                 | Alternative PA                   | Spearman r  | Gene       | 0.87  | Borzoi         | 13.1
+**DNA accessibility** | DNase-seq coverage               | Profile JSD | 1 bp       | 0.51  | ChromBPNet     | 6.4
+                 | DNase-seq coverage               | Pearson r   | 32 bp      | 0.86  | Borzoi         | 4.7
+                 | DNase-seq coverage               | Pearson r   | 128 bp     | 0.87  | Enformer       | 2.4
+                 | ATAC-seq coverage                | Profile JSD | 1 bp       | 0.46  | ChromBPNet     | 1.6
+                 | ATAC-seq coverage                | Pearson r   | 32 bp      | 0.57  | Borzoi         | 3.4
+                 | ATAC-seq coverage                | Pearson r   | 128 bp     | 0.72  | Enformer       | 2.7
+**Histone mods**      | Histone ChIP-seq coverage        | Pearson r   | 32 bp      | 0.69  | Borzoi         | 3.1
+                 | Histone ChIP-seq coverage        | Pearson r   | 128 bp     | 0.71  | Enformer       | 2.4
+**TF binding**        | TF ChIP-seq coverage             | Pearson r   | 32 bp      | 0.55  | Borzoi         | 5.0
+                 | TF ChIP-seq coverage             | Pearson r   | 128 bp     | 0.58  | Enformer       | 1.4
+**DNA contact maps**  | Orca Contact maps                | Pearson r   | 4000 bp    | 0.79  | Orca           | 6.3
+                 | Orca Contact maps cell type diff | Pearson r   | 4000 bp    | 0.42  | Orca           | 42.3
+
+And the next table details the performance of the distilled model in predicting
+the functional effects of genetic variants:
+
+Modality              | Evaluation                    | Type      | Metric     | Value | Baseline Model | Relative Improvement (%)
+:-------------------- | :---------------------------- | :-------- | :--------- | :---- | :------------- | :-----------------------
+**Splicing**          | ClinVar splice site region    | Causality | auPRC      | 0.57  | Pangolin       | 3.7
+                 | ClinVar noncoding             | \-        | auPRC      | 0.66  | Pangolin       | 2.9
+                 | ClinVar missense              | \-        | auPRC      | 0.18  | DeltaSplice    | 13.7
+                 | Splicing outlier (zero-shot)  | \-        | auPRC      | 0.22  | Pangolin       | 59.1
+                 | Splicing outlier (supervised) | \-        | auPRC      | 0.28  | AbSplice       | 13.0
+                 | sQTL                          | Causality | auPRC      | 0.76  | Pangolin       | 13.9
+                 | MFASS                         | \-        | auPRC      | 0.51  | Pangolin       | \-5.7
+**RNA expression**    | eQTL                          | Direction | Spearman r | 0.49  | Borzoi         | 25.5
+                 | eQTL (zero-shot)              | Causality | auROC      | 0.71  | Borzoi         | 5.4
+                 | eQTL (supervised)             | Causality | auROC      | 0.80  | Borzoi         | 15.6
+                 | ENCODE E2G (zero-shot)        | \-        | auPRC      | 0.75  | Borzoi         | 13.0
+                 | paQTL                         | \-        | auPRC      | 0.63  | Borzoi         | 7.3
+**DNA accessibility** | CAGI5 MPRA                    | Causality | Pearson r  | 0.65  | Borzoi         | 6.3
+                 | ds/caQTL                      | Direction | Pearson r  | 0.70  | ChromBPNet     | 7.7
+                 | ds/caQTL                      | Causality | auPRC      | 0.52  | Borzoi         | 18.0
+**TF binding**        | bQTL                          | Direction | Pearson r  | 0.55  | Borzoi         | 2.8
+                 | bQTL                          | Causality | auPRC      | 0.50  | Borzoi         | 6.0
+
+## Usage and Limitations
+
+### License
+
+Unless required by applicable law or agreed to in writing, all software and
+materials distributed here (under the
+[Apache 2.0 License](https://github.com/google-deepmind/alphagenome_research/blob/main/README.md)
+with respect to model code and
+[non-commercial terms](https://deepmind.google.com/science/alphagenome/model-terms)
+for the model parameters) are distributed on an "AS IS" BASIS, WITHOUT
+WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+licenses for the specific language governing permissions and limitations under
+those licenses.
+
+### Intended usage
+
+*   Non-Commercial use only: The model parameters are restricted to
+    non-commercial use by non-commercial organizations (e.g., universities,
+    non-profits, research institutes, and journalism). It must not be used for
+    any commercial activities or on behalf of commercial entities.
+*   Model derivatives: If you fine-tune or modify this model, the resulting
+    model is classified as a "Derivative". All derivatives are subject to these
+    exact same terms, meaning strictly no commercial use is permitted for
+    fine-tuned versions.
+*   Distillation: Training a new model using the outputs or predictions of
+    AlphaGenome is also restricted; resulting models must be governed by the
+    AlphaGenome Model Parameters Terms of Use
+
+### Limitations
+
+*   Like other sequence-based models, accurately capturing the influence of very
+    distant regulatory elements, like those over 100,000 DNA letters away, is
+    still an ongoing challenge. Another priority for future work is further
+    increasing the model’s ability to capture cell- and tissue-specific
+    patterns. We haven't designed or validated AlphaGenome for personal genome
+    prediction, a known challenge for AI models. Instead, we focused more on
+    characterising the performance on individual genetic variants. And while
+    AlphaGenome can predict molecular outcomes, it doesn't give the full picture
+    of how genetic variations lead to complex traits or diseases. These often
+    involve broader biological processes, like developmental and environmental
+    factors, that are beyond the direct scope of our model.

weight/alphagenome-all-folds/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1759332751356271135, "commit_timestamp_nsecs": 1759332758258030815, "custom_metadata": {}}

weight/alphagenome-all-folds/gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+ocdbt.process_0/d/48900811ab32b5ad3873d40e96d63846 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/6b08a83c9d6a2ade925dbcdd91299ce2 filter=lfs diff=lfs merge=lfs -text
+ocdbt.process_0/d/bda173d4f00d3afcef0a043614df05a6 filter=lfs diff=lfs merge=lfs -text