notebooks_tutorials/04_model_interpretation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BXtGC0FwDK7i"
   },
   "source": [
    "# 🧬 NTv3 Multi-Scale Interpretation Tutorial\n",
    "\n",
    "This notebook uses [`InstaDeepAI/NTv3_tracks`](https://huggingface.co/InstaDeepAI/NTv3_tracks), the NTv3 model post-trained on functional tracks, for multi-scale interpretation of regulatory logic, enabling analysis from individual nucleotides to enhancer elements and full gene loci.\n",
    "\n",
    "> 📝 **Note for Google Colab users:** This notebook is compatible with Colab! For faster inference, make sure to enable GPU: Runtime → Change runtime type → GPU (T4 or better recommended).\n",
    "\n",
    "## 📋 Overview\n",
    "\n",
    "Given a genomic window from the **human genome (hg38)**, we demonstrate three complementary interpretability modalities:\n",
    "\n",
    "1. **Predicted functional tracks** - Regulatory activity across thousands of assays\n",
    "2. **Transformer attention maps** - How the model integrates information across long genomic distances\n",
    "3. **Gradient-based saliency maps** - Sequence positions most influential for predictions\n",
    "\n",
    "> 💡 **Note:** Functional tracks are predicted only for the center 37.5% of the input sequence.\n",
    "\n",
    "`NTv3_tracks` predicts functional tracks, not gene annotations. Gene annotation is served by the separate NTv3-annotate checkpoints.\n",
    "\n",
    "## 📚 Notebook Structure\n",
    "\n",
    "1. **Setup**: Install dependencies and configure the analysis\n",
    "2. **Data Loading**: Fetch the genomic sequence from UCSC\n",
    "3. **Model Loading**: Load the NTv3 tracks model, its tokenizer and track metadata\n",
    "4. **Inference**: Run the model to generate predictions and extract attention\n",
    "5. **Functional Tracks**: Visualize predicted regulatory tracks\n",
    "6. **Attention Maps**: Analyze long-range information\n",
    "7. **Saliency Maps**: Compute gradient-based sequence attributions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MS_mIORoDK7k"
   },
   "source": [
    "## 0) 📦 Imports + Setup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vKTwnGhmVYFu"
   },
   "outputs": [],
   "source": [
    "# The NTv3 models are gated: request access once on any model page, then log in.\n",
    "from huggingface_hub import get_token, notebook_login\n",
    "\n",
    "if get_token() is None:\n",
    "    notebook_login()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JRc8918_DK7l"
   },
   "outputs": [],
   "source": [
    "# Colab ships torch and numpy. Upgrading them breaks the preinstalled scipy.\n",
    "!pip -q install --upgrade 'transformers>=5.17,<6'\n",
    "!pip -q install fast-logomaker\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qdXtceUiDK7l"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import warnings\n",
    "import logging\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "import requests\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "\n",
    "# Device setup\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "if device == \"cuda\" and torch.cuda.get_device_capability(0)[0] >= 8:\n",
    "    dtype = torch.bfloat16\n",
    "else:\n",
    "    dtype = torch.float16\n",
    "print(f\"Device: {device}, dtype: {dtype}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gOCdTnRADK7m"
   },
   "source": [
    "## 1) 📦 Configuration\n",
    "\n",
    "Set your NTv3 model, species and genomic window here. Coordinates are in the assembly of the chosen species (hg38 for human, the default), and the tracks you can pick come from that assembly's track list.\n",
    "\n",
    "> 💡 **Note:** For a T4 GPU, we recommend using an input sequence length of 32kb or shorter.\n",
    "\n",
    "> The knobs are `species`, `assembly`, `chrom`, `center_position` and `sequence_length` in the cell below. `assembly` must be the one `SPECIES_TO_ASSEMBLY` gives for `species`, or empty to use it; a mismatch raises an error. `start` and `end` are derived from `center_position` and `sequence_length`. The tracks to plot are chosen later, in `tracks_to_plot` (section 5) and `saliency_track_id` (section 7).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iBmudDzIDK7m"
   },
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# User Configuration\n",
    "# -----------------------------\n",
    "model_name = \"InstaDeepAI/NTv3_tracks\"\n",
    "\n",
    "# Species and its assembly; leave assembly empty to use the species' assembly\n",
    "species = \"human\"  #@param [\"human\", \"mouse\", \"drosophila\", \"arabidopsis\", \"maize\", \"rice\", \"soybean\", \"caenorhabditis_elegans\", \"zebrafish\", \"chicken\", \"horse\", \"sheep\"]\n",
    "assembly = \"hg38\"  #@param [\"\", \"hg38\", \"mm10\", \"dm6\", \"TAIR10\", \"Zm-B73-REFERENCE-NAM-5.0\", \"IRGSP-1.0\", \"Glycine_max_v2.1\", \"WBcel235\", \"GRCz11\", \"GRCg6a\", \"EquCab3.0\", \"Oar_v3.1\"]\n",
    "\n",
    "# Genomic window (must be multiple of 128 due to model downsampling)\n",
    "chrom = \"chr11\"  #@param {type:\"string\"}\n",
    "# Center of region of interest (HBE1 locus)\n",
    "center_position = 5269945  #@param {type:\"integer\"}\n",
    "# use 131kb window on A100, 32_768 on T4 GPU\n",
    "sequence_length = 32768  #@param [16384, 32768, 65536, 131072] {type:\"raw\"}\n",
    "assert sequence_length % 128 == 0, \"sequence_length must be a multiple of 128\"\n",
    "\n",
    "# Calculate start/end from center\n",
    "start = center_position - sequence_length // 2\n",
    "end = center_position + sequence_length // 2\n",
    "\n",
    "# Analysis parameters\n",
    "token_resolution = 128  # bp per token\n",
    "\n",
    "# Track lists are keyed by assembly, so each species maps to one assembly.\n",
    "# The names match metadata[\"bigwigs_per_file_assembly\"], printed once the model loads.\n",
    "SPECIES_TO_ASSEMBLY = {\n",
    "    \"human\": \"hg38\",\n",
    "    \"mouse\": \"mm10\",\n",
    "    \"drosophila\": \"dm6\",\n",
    "    \"arabidopsis\": \"TAIR10\",\n",
    "    \"maize\": \"Zm-B73-REFERENCE-NAM-5.0\",\n",
    "    \"rice\": \"IRGSP-1.0\",\n",
    "    \"soybean\": \"Glycine_max_v2.1\",\n",
    "    \"caenorhabditis_elegans\": \"WBcel235\",\n",
    "    \"zebrafish\": \"GRCz11\",\n",
    "    \"chicken\": \"GRCg6a\",\n",
    "    \"horse\": \"EquCab3.0\",\n",
    "    \"sheep\": \"Oar_v3.1\",\n",
    "}\n",
    "if species not in SPECIES_TO_ASSEMBLY:\n",
    "    raise ValueError(f\"Unknown species {species!r}. Choose one of {list(SPECIES_TO_ASSEMBLY)}\")\n",
    "if not assembly:\n",
    "    assembly = SPECIES_TO_ASSEMBLY[species]\n",
    "elif assembly != SPECIES_TO_ASSEMBLY[species]:\n",
    "    raise ValueError(\n",
    "        f\"assembly={assembly!r} does not match species={species!r}, \"\n",
    "        f\"whose assembly is {SPECIES_TO_ASSEMBLY[species]!r}\"\n",
    "    )\n",
    "\n",
    "print(f\"Model: {model_name}\")\n",
    "print(f\"Region: {chrom}:{start:,}-{end:,} ({sequence_length:,} bp)\")\n",
    "print(f\"Center: {center_position:,}\")\n",
    "print(f\"Species: {species} -> Assembly: {assembly}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yytbImqlDK7m"
   },
   "source": [
    "## 2) 📥 Fetch Chromosome Sequence\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cqXIi7mYDK7m"
   },
   "outputs": [],
   "source": [
    "# Get the sequence from the UCSC API\n",
    "url = (\n",
    "    f\"https://api.genome.ucsc.edu/getData/sequence?\"\n",
    "    f\"genome={assembly};chrom={chrom};start={start};end={end}\"\n",
    ")\n",
    "seq = requests.get(url, timeout=60).json()[\"dna\"]\n",
    "print(f\"Original sequence length: {len(seq):,} bp\")\n",
    "\n",
    "# Crop to multiple of 128 (ensures clean tokenization)\n",
    "seq = seq[:int(len(seq) // 128) * 128]\n",
    "n_tokens = len(seq) / 128\n",
    "print(\n",
    "    f\"Cropped sequence length: {len(seq):,} bp \"\n",
    "    f\"({n_tokens:.0f} transformer tokens)\"\n",
    ")\n",
    "\n",
    "# UCSC returns soft-masked sequence (repeats in lowercase). NTv3_tracks was trained\n",
    "# on soft-masked genomes, so we keep the case as it is.\n",
    "soft_masked_fraction = sum(b.islower() for b in seq) / len(seq)\n",
    "print(f\"Soft-masked (lowercase) fraction: {soft_masked_fraction:.1%}\")\n",
    "print(f\"First 100 bp: {seq[:100]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jtJMea-1DK7m"
   },
   "source": [
    "## 3) 🤖 Load NTv3 Model + Tokenizer\n",
    "\n",
    "The model uses custom code, so we pass `trust_remote_code=True`. The species vocabulary and the track order ship with the model config, under `track_metadata`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L8RFTmRaDK7m"
   },
   "outputs": [],
   "source": [
    "# Load model\n",
    "print(\"Loading model...\")\n",
    "model = AutoModel.from_pretrained(\n",
    "    model_name,\n",
    "    trust_remote_code=True,\n",
    ").to(device)\n",
    "model.eval()\n",
    "\n",
    "# Load tokenizer\n",
    "print(\"Loading tokenizer...\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_name, trust_remote_code=True\n",
    ")\n",
    "\n",
    "# Species ids and track order ship with the model config\n",
    "metadata = model.config.track_metadata\n",
    "species_vocab = metadata[\"species_vocabulary\"]\n",
    "\n",
    "print(f\"\\n✅ Model loaded successfully\")\n",
    "supported_species = [s for s in species_vocab if not s.startswith(\"<\")]\n",
    "print(f\"   Supported species: {supported_species}\")\n",
    "assemblies = list(metadata[\"bigwigs_per_file_assembly\"])\n",
    "print(f\"   Available assemblies: {assemblies}\")\n",
    "if species not in supported_species:\n",
    "    raise ValueError(f\"Species {species!r} is not supported by {model_name}\")\n",
    "if assembly not in assemblies:\n",
    "    raise ValueError(f\"Assembly {assembly!r} has no tracks in {model_name}\")\n",
    "n_bigwigs = len(metadata[\"bigwigs_per_file_assembly\"][assembly])\n",
    "print(f\"   Number of {assembly} BigWig tracks: {n_bigwigs}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_XZ3xXIMDK7n"
   },
   "outputs": [],
   "source": [
    "# Inspect available outputs: channel i of the output is track i of this list\n",
    "bigwig_names = metadata[\"bigwigs_per_file_assembly\"][assembly]\n",
    "print(\n",
    "    f\"{len(bigwig_names)} functional tracks for {assembly}. \"\n",
    "    f\"First 10: {bigwig_names[:10]}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rFdiD_xZDK7n"
   },
   "source": [
    "## 4) ⚡ Run Inference\n",
    "\n",
    "Following the manuscript methods: The functional-track module produces per-base predictions for all species-matched regulatory assays. Supervision is applied only to the central 37.5% of the input sequence.\n",
    "\n",
    "We pass:\n",
    "- the tokenized DNA window, positionally: the forward takes the token tensor, not the tokenizer dict\n",
    "- `conditions_tokens`: the species id, as a list with one `(batch,)` tensor\n",
    "\n",
    "`bigwig_tracks_logits` comes back on the forward output. The trunk uses scaled-dot-product attention and does not return weights, so the maps below are computed from each attention block.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bAidYEvWDK7n"
   },
   "outputs": [],
   "source": [
    "# Tokenize DNA sequence (its length is already a multiple of 128)\n",
    "batch = tokenizer([seq], add_special_tokens=False, return_tensors=\"pt\")\n",
    "input_ids = batch[\"input_ids\"].to(device)\n",
    "# Characters outside the alphabet (e.g. IUPAC codes) become <unk> without changing the length\n",
    "assert tokenizer.unk_token_id not in input_ids, \"the sequence contains characters outside the alphabet\"\n",
    "\n",
    "# Encode species\n",
    "species_ids = torch.tensor([species_vocab[species]], device=device)\n",
    "\n",
    "print(f\"Input IDs shape: {input_ids.shape}\")\n",
    "print(f\"Species IDs shape: {species_ids.shape}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yWDD1Z0zDK7n"
   },
   "outputs": [],
   "source": [
    "# Run inference. Attention weights are not a forward argument on this model.\n",
    "import importlib\n",
    "\n",
    "print(\"Running forward pass...\")\n",
    "blocks = model.model.mha_blocks\n",
    "modeling = importlib.import_module(type(blocks[0]).__module__)\n",
    "captured = []\n",
    "originals = [block.forward for block in blocks]\n",
    "\n",
    "def _capturing_forward(block, original):\n",
    "    def forward(x, attention_mask=None, conditions=None, conditions_masks=None):\n",
    "        batch, seq_len, _ = x.shape\n",
    "        if block.conditioned:\n",
    "            hidden = block.norm_in(x, conditions, conditions_masks)\n",
    "        else:\n",
    "            hidden = block.norm_in(x)\n",
    "        query = block.norm_q(\n",
    "            block.q_layer(hidden).view(batch, seq_len, block.attention_heads, -1)\n",
    "        )\n",
    "        key = block.norm_k(\n",
    "            block.k_layer(hidden).view(batch, seq_len, block.num_kv_heads, -1)\n",
    "        )\n",
    "        inv_freq = modeling.build_rope_inv_freq(block.key_size, device=x.device)\n",
    "        query = modeling.apply_rope(query, inv_freq)\n",
    "        key = modeling.apply_rope(key, inv_freq)\n",
    "        group = block.attention_heads // block.num_kv_heads\n",
    "        query = query.transpose(1, 2)\n",
    "        key = key.transpose(1, 2).repeat_interleave(group, dim=1)\n",
    "        scale = 1.0 / math.sqrt(block.key_size)\n",
    "        logits = torch.matmul(query.float(), key.float().transpose(-2, -1)) * scale\n",
    "        if attention_mask is not None:\n",
    "            logits = logits.masked_fill(~attention_mask, float(\"-inf\"))\n",
    "        weights = torch.softmax(logits, dim=-1)\n",
    "        captured.append(weights[0].cpu().numpy())\n",
    "        return original(\n",
    "            x,\n",
    "            attention_mask=attention_mask,\n",
    "            conditions=conditions,\n",
    "            conditions_masks=conditions_masks,\n",
    "        )\n",
    "\n",
    "    return forward\n",
    "\n",
    "for block, original in zip(blocks, originals):\n",
    "    block.forward = _capturing_forward(block, original)\n",
    "try:\n",
    "    with torch.no_grad():\n",
    "        outputs = model(input_ids, conditions_tokens=[species_ids])\n",
    "finally:\n",
    "    for block, original in zip(blocks, originals):\n",
    "        block.forward = original\n",
    "\n",
    "# [batch, 0.375 * seq_len, num_tracks]\n",
    "bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
    "# (layers, heads, blocks, blocks), one 128 bp block per position\n",
    "attention_maps = np.stack(captured, axis=0)\n",
    "\n",
    "print(\"\\n✅ Forward pass completed\")\n",
    "print(f\"   BigWig logits: {tuple(bigwig_logits.shape)}\")\n",
    "print(f\"   Attention maps: {attention_maps.shape} (layers, heads, seq, seq)\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WpKxn1juDK7n"
   },
   "source": [
    "## 5) 📊 Visualize Functional Tracks\n",
    "\n",
    "Following the manuscript methods: Functional-track predictions are obtained directly from the model output tensor. Predictions are center-cropped to the supervised region (49,152 bp for 131 kb inputs).\n",
    "\n",
    "Metadata for all available tracks can be downloaded [here](https://huggingface.co/spaces/InstaDeepAI/ntv3_tracks/blob/main/data/functional_tracks_metadata.csv)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u3pzwP7VDK7n"
   },
   "outputs": [],
   "source": [
    "# Select functional tracks to plot\n",
    "tracks_to_plot = {\n",
    "    \"K562 RNA-seq\": \"ENCSR056HPM\",\n",
    "    \"K562 CAGE\": \"CNhs11250_P\",\n",
    "    \"K562 DNAse\": \"ENCSR921NMD\",\n",
    "    \"K562 GATA1 ChIP-seq\": \"ENCSR000EFT\"\n",
    "}\n",
    "\n",
    "# Every track must exist for this assembly\n",
    "missing_tracks = [v for v in tracks_to_plot.values() if v not in bigwig_names]\n",
    "if missing_tracks:\n",
    "    raise ValueError(f\"Tracks not available for {assembly}: {missing_tracks}\")\n",
    "\n",
    "print(f\"Plotting {len(tracks_to_plot)} functional tracks\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gtF2KGoQDK7n"
   },
   "outputs": [],
   "source": [
    "def plot_tracks(tracks, start_pos, end_pos, chrom_name, assembly_name,\n",
    "                height=1.0):\n",
    "    \"\"\"Plot multiple tracks in a stacked view.\n",
    "\n",
    "    Args:\n",
    "        tracks: Dictionary mapping track names to 1D arrays of values.\n",
    "        start_pos: Genomic start position for x-axis.\n",
    "        end_pos: Genomic end position for x-axis.\n",
    "        chrom_name: Chromosome name for labeling.\n",
    "        assembly_name: Assembly name for labeling.\n",
    "        height: Height multiplier for each track subplot.\n",
    "\n",
    "    Returns:\n",
    "        matplotlib.figure.Figure: The created figure object.\n",
    "    \"\"\"\n",
    "    n_tracks = len(tracks)\n",
    "    fig, axes = plt.subplots(\n",
    "        n_tracks, 1, figsize=(16, height * n_tracks), sharex=True\n",
    "    )\n",
    "    if n_tracks == 1:\n",
    "        axes = [axes]\n",
    "\n",
    "    for ax, (title, y) in zip(axes, tracks.items()):\n",
    "        x_positions = np.linspace(start_pos, end_pos, num=len(y))\n",
    "        ax.fill_between(x_positions, y, alpha=0.7)\n",
    "        ax.set_title(title, fontsize=8)\n",
    "        ax.yaxis.set_label_coords(-0.01, 0.5)\n",
    "        sns.despine(ax=ax, top=True, right=True, bottom=True)\n",
    "        ax.tick_params(axis='y', labelsize=7)\n",
    "\n",
    "    xlabel = f\"{chrom_name}:{start_pos:,}-{end_pos:,} ({assembly_name})\"\n",
    "    axes[-1].set_xlabel(xlabel, fontsize=9)\n",
    "    plt.tight_layout()\n",
    "    return fig\n",
    "\n",
    "# Extract and plot functional tracks\n",
    "bigwig = bigwig_logits.detach().float().cpu().numpy()[0]\n",
    "bigwig_tracks = {}\n",
    "for track_name, track_id in tracks_to_plot.items():\n",
    "    if bigwig_names and track_id in bigwig_names:\n",
    "        track_idx = bigwig_names.index(track_id)\n",
    "        bigwig_tracks[track_name] = bigwig[:, track_idx]\n",
    "\n",
    "if not bigwig_tracks:\n",
    "    print(\"⚠️ No tracks to plot (track names may not be available)\")\n",
    "else:\n",
    "    # Calculate prediction region (center 37.5%)\n",
    "    window_len = end - start\n",
    "    prediction_start = start + int(window_len * 0.3125)\n",
    "    prediction_end = prediction_start + int(window_len * 0.375)\n",
    "\n",
    "    fig = plot_tracks(\n",
    "        bigwig_tracks, prediction_start, prediction_end, chrom, assembly\n",
    "    )\n",
    "    plt.suptitle(\"Predicted Functional Tracks\", fontsize=11, y=1.02)\n",
    "    plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **How to read this plot.** The x axis is the genomic coordinate on `chrom` in `assembly`. It covers the predicted span, the center 37.5% of the input window, not the full input. Each panel is one track at 1 bp resolution, on its own y scale. A high value means the model predicts a strong signal for that assay at that position.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iqyiyLPzDK7o"
   },
   "source": [
    "## 6) 🔗 Analyze Attention Maps\n",
    "\n",
    "Following the manuscript methods: Transformer self-attention reveals how NTv3 integrates information across distal regions. We extract multi-head self-attention weights and average them across heads. Matrices are symmetrized as (A + A^T)/2.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "StvwdIcUDK7o"
   },
   "outputs": [],
   "source": [
    "# Extract attention from last layer\n",
    "layer_index = -1  # Last layer\n",
    "# [heads, seq_len, seq_len]\n",
    "attention_last_layer = attention_maps[layer_index]\n",
    "\n",
    "# Average over heads and symmetrize\n",
    "# [seq_len, seq_len]\n",
    "attention_mean = attention_last_layer.mean(axis=0)\n",
    "attention_mean = 0.5 * (attention_mean + attention_mean.T)  # Symmetrize\n",
    "\n",
    "print(f\"Attention map shape: {attention_mean.shape}\")\n",
    "attn_min = attention_mean.min()\n",
    "attn_max = attention_mean.max()\n",
    "print(f\"Attention value range: [{attn_min:.6f}, {attn_max:.6f}]\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qRxVY-6VDK7o"
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# Pyramid-style Attention Map Plotting (45-degree rotated triangular)\n",
    "# ============================================================\n",
    "from matplotlib.colors import TwoSlopeNorm\n",
    "\n",
    "# Global settings\n",
    "LINE_WIDTH = 0.5\n",
    "\n",
    "def _get_45deg_mesh(mat):\n",
    "    \"\"\"Create 45-degree rotated mesh for triangular attention visualization.\n",
    "\n",
    "    Args:\n",
    "        mat: Input matrix to create mesh for.\n",
    "\n",
    "    Returns:\n",
    "        tuple: (X, Y) mesh coordinates rotated 45 degrees.\n",
    "    \"\"\"\n",
    "    theta = -np.pi / 4\n",
    "    R = np.array([\n",
    "        [np.cos(theta), -np.sin(theta)],\n",
    "        [np.sin(theta), np.cos(theta)]\n",
    "    ])\n",
    "    K = len(mat) + 1\n",
    "    g = np.arange(0, K) - 0.5\n",
    "    X = np.tile(g[:, None], (1, K))\n",
    "    Y = np.tile(g[None, :], (K, 1))\n",
    "    xy = np.array([X.ravel(), Y.ravel()])\n",
    "    xy_rot = R @ xy\n",
    "    return xy_rot[0, :].reshape(K, K), xy_rot[1, :].reshape(K, K).T\n",
    "\n",
    "def plot_attention_panel(mat, filename=None, *, cmap=\"Blues\", vmin=0.0001,\n",
    "                         vmax=0.005, colorbar=True, dpi=600, figsize=(10, 10),\n",
    "                         variant_id=\"\", show_titles=True, show_xlabel=True,\n",
    "                         show_xticks=True, show_yticks=False,\n",
    "                         positions='TOKEN', token_resolution=128, show=True):\n",
    "    \"\"\"Plot triangular attention heatmap with 45-degree rotation (pyramid style).\n",
    "\n",
    "    Args:\n",
    "        mat: Attention matrix to plot.\n",
    "        filename: Output filename. If None and show=True, displays plot.\n",
    "        cmap: Colormap name. Defaults to \"Blues\".\n",
    "        vmin: Minimum value for color scale. Defaults to 0.0001.\n",
    "        vmax: Maximum value for color scale. Defaults to 0.005.\n",
    "        colorbar: Whether to show colorbar. Defaults to True.\n",
    "        dpi: Figure DPI. Defaults to 600.\n",
    "        figsize: Figure size tuple. Defaults to (10, 10).\n",
    "        variant_id: Title for the plot. Defaults to \"\".\n",
    "        show_titles: Whether to show title. Defaults to True.\n",
    "        show_xlabel: Whether to show x-axis label. Defaults to True.\n",
    "        show_xticks: Whether to show x-axis ticks. Defaults to True.\n",
    "        show_yticks: Whether to show y-axis ticks. Defaults to False.\n",
    "        positions: 'TOKEN' or 'BP' for position units. Defaults to 'TOKEN'.\n",
    "        token_resolution: Base pairs per token. Defaults to 128.\n",
    "        show: Whether to display plot. Defaults to True.\n",
    "\n",
    "    Returns:\n",
    "        matplotlib.figure.Figure: The created figure object.\n",
    "    \"\"\"\n",
    "    # Symmetrize matrix\n",
    "    mat = 0.5 * (mat + mat.T)\n",
    "\n",
    "    # Mask lower triangle\n",
    "    mat[np.tril_indices_from(mat, k=-1)] = np.nan\n",
    "    n = mat.shape[0]\n",
    "    X, Y = _get_45deg_mesh(mat)\n",
    "\n",
    "    # Coordinate normalization\n",
    "    C = 1\n",
    "    half_pixel_diag = 1 / (2*C)\n",
    "    pixel_side = 1 / (C * np.sqrt(2))\n",
    "    X = X * pixel_side + half_pixel_diag\n",
    "    Y = Y * pixel_side\n",
    "    Y = -Y  # Flip Y\n",
    "\n",
    "    # Clean up coordinates\n",
    "    X = np.nan_to_num(X, nan=0.0, posinf=1.0, neginf=-1.0)\n",
    "    Y = np.nan_to_num(Y, nan=0.0, posinf=1.0, neginf=-1.0)\n",
    "\n",
    "    # Set up normalization\n",
    "    vcenter = (vmin + vmax) / 2\n",
    "    norm = TwoSlopeNorm(vmin=vmin, vcenter=vcenter, vmax=vmax)\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)\n",
    "    im = ax.pcolormesh(X, Y, mat, cmap=cmap, norm=norm, shading=\"auto\")\n",
    "\n",
    "    # Set axis properties\n",
    "    ax.set_aspect(\"equal\")\n",
    "\n",
    "    # Calculate coordinate extents\n",
    "    finite_mask = np.isfinite(mat)\n",
    "    if finite_mask.any():\n",
    "        vertex_mask = np.zeros_like(X, dtype=bool)\n",
    "        vertex_mask[:-1, :-1] |= finite_mask\n",
    "        vertex_mask[:-1, 1:] |= finite_mask\n",
    "        vertex_mask[1:, :-1] |= finite_mask\n",
    "        vertex_mask[1:, 1:] |= finite_mask\n",
    "\n",
    "        valid_x = X[vertex_mask]\n",
    "        valid_y = Y[vertex_mask]\n",
    "        x_min, x_max = np.nanmin(valid_x), np.nanmax(valid_x)\n",
    "        y_min, y_max = np.nanmin(valid_y), np.nanmax(valid_y)\n",
    "    else:\n",
    "        x_min, x_max = np.nanmin(X), np.nanmax(X)\n",
    "        y_min, y_max = np.nanmin(Y), np.nanmax(Y)\n",
    "\n",
    "    triangle_bottom_y = y_min\n",
    "    triangle_top_y = y_max\n",
    "\n",
    "    xlim_pad = 0.1\n",
    "    ylim_pad = 0.1\n",
    "    y_range = (y_max - y_min) if y_max > y_min else 1.0\n",
    "    bottom_pad = 0.001 * y_range\n",
    "\n",
    "    ax.set_xlim(x_min - xlim_pad, x_max + xlim_pad)\n",
    "    ax.set_ylim(y_min - bottom_pad, y_max + ylim_pad)\n",
    "\n",
    "    # Add colorbar\n",
    "    if colorbar:\n",
    "        pyramid_y_min = y_min\n",
    "        pyramid_y_max = y_max\n",
    "\n",
    "        cb = plt.colorbar(im, ax=ax, shrink=1.0, aspect=20, pad=0.1)\n",
    "        cb.ax.tick_params(width=LINE_WIDTH*0.8, length=3, labelsize=7, direction=\"in\")\n",
    "        cb.outline.set_linewidth(LINE_WIDTH*0.8)\n",
    "\n",
    "        ax_pos = ax.get_position()\n",
    "        axes_y_min_data = y_min - bottom_pad\n",
    "        axes_y_max_data = y_max + ylim_pad\n",
    "        axes_height_data = axes_y_max_data - axes_y_min_data\n",
    "\n",
    "        if axes_height_data > 0:\n",
    "            pyramid_bottom_frac = (\n",
    "                (pyramid_y_min - axes_y_min_data) / axes_height_data\n",
    "            )\n",
    "            pyramid_top_frac = (\n",
    "                (pyramid_y_max - axes_y_min_data) / axes_height_data\n",
    "            )\n",
    "        else:\n",
    "            pyramid_bottom_frac = 0\n",
    "            pyramid_top_frac = 1\n",
    "\n",
    "        cb_pos = cb.ax.get_position()\n",
    "        new_y0 = ax_pos.y0 + ax_pos.height * pyramid_bottom_frac\n",
    "        new_height = ax_pos.height * (pyramid_top_frac - pyramid_bottom_frac)\n",
    "        cb.ax.set_position(\n",
    "            [cb_pos.x0, new_y0, cb_pos.width, new_height]\n",
    "        )\n",
    "\n",
    "    # Format axes\n",
    "    use_bp_display = positions.upper() == 'BP'\n",
    "\n",
    "    ax.tick_params(width=LINE_WIDTH, length=3, labelsize=7)\n",
    "    for sp in (\"top\", \"right\", \"left\"):\n",
    "        ax.spines[sp].set_visible(False)\n",
    "    ax.spines[\"bottom\"].set_linewidth(LINE_WIDTH)\n",
    "    ax.spines[\"bottom\"].set_position(('data', triangle_bottom_y))\n",
    "    ax.xaxis.set_ticks_position('bottom')\n",
    "    ax.xaxis.set_label_position('bottom')\n",
    "\n",
    "    if show_xticks and mat.size > 0:\n",
    "        n_tokens = mat.shape[0]\n",
    "        tick_positions = [0, n_tokens - 1]\n",
    "        if use_bp_display:\n",
    "            max_bp = (n_tokens - 1) * token_resolution\n",
    "            tick_labels = ['0', str(max_bp)]\n",
    "        else:\n",
    "            tick_labels = ['0', str(n_tokens - 1)]\n",
    "        ax.set_xticks(tick_positions)\n",
    "        ax.set_xticklabels(tick_labels)\n",
    "    else:\n",
    "        ax.set_xticks([])\n",
    "\n",
    "    if not show_yticks:\n",
    "        ax.set_yticks([])\n",
    "\n",
    "    if use_bp_display:\n",
    "        xlabel_text = \"Position (base pairs)\"\n",
    "    else:\n",
    "        xlabel_text = \"Token position\"\n",
    "    if show_xlabel:\n",
    "        ax.set_xlabel(xlabel_text, fontsize=8)\n",
    "\n",
    "    if show_titles and variant_id:\n",
    "        ax.set_title(variant_id, fontsize=9, pad=10)\n",
    "\n",
    "    ax.set_facecolor(\"white\")\n",
    "\n",
    "    if filename:\n",
    "        dirname = os.path.dirname(filename)\n",
    "        if dirname:\n",
    "            os.makedirs(dirname, exist_ok=True)\n",
    "        plt.savefig(\n",
    "            filename, dpi=600, bbox_inches=\"tight\", facecolor='white'\n",
    "        )\n",
    "        print(f\"✅ Saved attention plot: {filename}\")\n",
    "        if not show:\n",
    "            plt.close(fig)\n",
    "    elif show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close(fig)\n",
    "\n",
    "    return fig\n",
    "\n",
    "print(\"✅ Pyramid-style attention plotting function defined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qeziBpBmDK7o"
   },
   "outputs": [],
   "source": [
    "# Plot pyramid-style attention map (cropped to center window)\n",
    "LAYER_INDEX = layer_index\n",
    "TOKEN_RESOLUTION = token_resolution\n",
    "window_size = min(1024, len(attention_mean))  # tokens to show\n",
    "\n",
    "# Crop to center window\n",
    "n_tokens_full = len(attention_mean)\n",
    "start_crop = (n_tokens_full - window_size) // 2\n",
    "end_crop = start_crop + window_size\n",
    "attention_mean_cropped = attention_mean[\n",
    "    start_crop:end_crop, start_crop:end_crop\n",
    "].copy()\n",
    "\n",
    "print(f\"Plotting attention map:\")\n",
    "print(f\"  Layer: {LAYER_INDEX}\")\n",
    "print(\n",
    "    f\"  Window: {window_size} tokens \"\n",
    "    f\"(cropped from {n_tokens_full} total)\"\n",
    ")\n",
    "attn_crop_min = attention_mean_cropped.min()\n",
    "attn_crop_max = attention_mean_cropped.max()\n",
    "print(f\"  Value range: [{attn_crop_min:.6f}, {attn_crop_max:.6f}]\")\n",
    "\n",
    "# Plot with pyramid style using shared utility\n",
    "title = f\"Attention Map - Layer {LAYER_INDEX} (Head Average)\"\n",
    "fig = plot_attention_panel(\n",
    "    attention_mean_cropped,\n",
    "    cmap=\"Blues\",\n",
    "    vmin=0.0001,\n",
    "    vmax=0.005,\n",
    "    colorbar=True,\n",
    "    figsize=(10, 10),\n",
    "    dpi=150,\n",
    "    show_xlabel=True,\n",
    "    show_xticks=True,\n",
    "    positions='TOKEN',\n",
    "    token_resolution=TOKEN_RESOLUTION,\n",
    "    show=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **How to read this plot.** The x axis is token position across the input window, one token per `token_resolution` bp (128 bp), up to 1,024 tokens around the center, so it covers the full input rather than only the predicted span. Each pixel is the last-layer attention between two tokens, averaged over heads and symmetrized. The map is the upper triangle rotated by 45°: the diagonal runs along the base and holds each token with itself and its neighbours, the local context. A pixel sits above the midpoint of its two tokens, higher the farther apart they are, so dark blobs high above the base are long-range interactions. Darker means more attention, and the colour scale is clipped at `vmax`, so the strongest pairs saturate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J7jJHp0uDK7o"
   },
   "source": [
    "## 7) 🎯 Compute Saliency Maps\n",
    "\n",
    "Following the manuscript methods: Gradient-based attribution provides a first-order, nucleotide-resolved estimate of how NTv3 predictions depend on the underlying DNA sequence. We differentiate model outputs with respect to the one-hot encoded DNA input using a differentiable embedding override: a forward hook replaces the token embedding lookup with `one_hot @ embedding.weight`.\n",
    "\n",
    "For the selected functional track(s), we aggregate logits over sequence positions in a window and apply a log transform.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FhAzknM4DK7o"
   },
   "outputs": [],
   "source": [
    "def find_token_embedding(model, vocab_size):\n",
    "    \"\"\"Return the embedding layer that maps DNA tokens to vectors.\n",
    "\n",
    "    Args:\n",
    "        model: The model to search.\n",
    "        vocab_size: Size of the DNA token vocabulary.\n",
    "\n",
    "    Returns:\n",
    "        torch.nn.Embedding: The token embedding layer.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        embedding = model.get_input_embeddings()\n",
    "    except (AttributeError, NotImplementedError):\n",
    "        embedding = None\n",
    "    if isinstance(embedding, torch.nn.Embedding):\n",
    "        return embedding\n",
    "\n",
    "    # Fall back to the only embedding table sized to the DNA vocabulary\n",
    "    candidates = [\n",
    "        module for module in model.modules()\n",
    "        if isinstance(module, torch.nn.Embedding)\n",
    "        and module.num_embeddings == vocab_size\n",
    "    ]\n",
    "    if len(candidates) != 1:\n",
    "        raise ValueError(\n",
    "            f\"Expected one token embedding of size {vocab_size}, \"\n",
    "            f\"found {len(candidates)}. \"\n",
    "            \"Saliency computation requires access to embedding weights.\"\n",
    "        )\n",
    "    return candidates[0]\n",
    "\n",
    "\n",
    "class SaliencyComputer:\n",
    "    \"\"\"Compute gradient-based saliency maps for NTv3_tracks.\"\"\"\n",
    "\n",
    "    def __init__(self, model, tokenizer, track_indices, species=\"human\",\n",
    "                 device=None, target_window_bp=None):\n",
    "        \"\"\"Initialize saliency computer.\n",
    "\n",
    "        Args:\n",
    "            model: The model to compute saliency for.\n",
    "            tokenizer: Tokenizer for input sequences.\n",
    "            track_indices: List of track indices to compute saliency for.\n",
    "            species: Species name. Defaults to \"human\".\n",
    "            device: Device to use. If None, auto-detects. Defaults to None.\n",
    "            target_window_bp: Optional (start, end) tuple in bp along the\n",
    "                output axis to aggregate over. Indices are relative to the\n",
    "                output (0 = first predicted position). If None (default),\n",
    "                sums over all output positions. The output is at 1 bp\n",
    "                resolution.\n",
    "        \"\"\"\n",
    "        # Auto-detect device if not provided\n",
    "        if device is None:\n",
    "            device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "        self.model = model\n",
    "        self.tokenizer = tokenizer\n",
    "        # Only these channels of the track head are computed\n",
    "        self.track_indices = torch.tensor(track_indices, device=device)\n",
    "        self.species = species\n",
    "        self.device = device\n",
    "        self.target_window_bp = target_window_bp\n",
    "        self.embedding_layer = find_token_embedding(model, len(tokenizer))\n",
    "\n",
    "        # Species id from the vocabulary shipped with the model config\n",
    "        species_vocab = model.config.track_metadata[\"species_vocabulary\"]\n",
    "        self.species_ids = torch.tensor([species_vocab[species]], device=device)\n",
    "\n",
    "    def compute_saliency(self, sequence):\n",
    "        \"\"\"Compute saliency map for given sequence.\n",
    "\n",
    "        Args:\n",
    "            sequence: DNA sequence string to compute saliency for. Its\n",
    "                length must be a multiple of 128.\n",
    "\n",
    "        Returns:\n",
    "            tuple: (gradients, one_hot_sequence) where gradients is the\n",
    "                saliency scores and one_hot_sequence is the one-hot encoded\n",
    "                input sequence, both (positions, vocab_size).\n",
    "        \"\"\"\n",
    "        # Tokenize\n",
    "        batch = self.tokenizer(\n",
    "            [sequence], add_special_tokens=False, return_tensors=\"pt\"\n",
    "        )\n",
    "        token_ids = batch[\"input_ids\"].to(self.device)\n",
    "\n",
    "        # Create one-hot encoding\n",
    "        one_hot = F.one_hot(\n",
    "            token_ids[0], num_classes=self.embedding_layer.num_embeddings\n",
    "        ).float()\n",
    "        one_hot.requires_grad_(True)\n",
    "\n",
    "        # Replace the embedding lookup by one_hot @ weight, which is differentiable\n",
    "        def embed_one_hot(module, inputs, output):\n",
    "            return (one_hot @ module.weight).unsqueeze(0).to(output.dtype)\n",
    "\n",
    "        # Forward pass with gradient tracking; the token tensor is passed positionally\n",
    "        handle = self.embedding_layer.register_forward_hook(embed_one_hot)\n",
    "        try:\n",
    "            outputs = self.model(\n",
    "                token_ids,\n",
    "                conditions_tokens=[self.species_ids],\n",
    "                track_indices=self.track_indices,\n",
    "                skip_lm_head=True,\n",
    "            )\n",
    "        finally:\n",
    "            handle.remove()\n",
    "\n",
    "        # [batch, positions, len(track_indices)], read from the returned dict\n",
    "        logits = outputs[\"bigwig_tracks_logits\"]\n",
    "\n",
    "        # Optionally slice output positions (output is at 1 bp resolution)\n",
    "        if self.target_window_bp is not None:\n",
    "            start, end = self.target_window_bp\n",
    "            seq_len = logits.shape[1]\n",
    "            start = max(0, min(start, seq_len))\n",
    "            end = max(start, min(end, seq_len))\n",
    "            logits = logits[:, start:end, :]\n",
    "\n",
    "        score = torch.log(logits.sum() + 1e-6)\n",
    "\n",
    "        # Backward pass\n",
    "        score.backward()\n",
    "\n",
    "        gradients = one_hot.grad.cpu().detach().numpy()\n",
    "        one_hot_np = one_hot.detach().cpu().numpy()\n",
    "\n",
    "        return gradients, one_hot_np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_YwKWAj-DK7o"
   },
   "outputs": [],
   "source": [
    "# Select track for saliency computation\n",
    "# The track id must exist for this assembly\n",
    "saliency_track_id = \"ENCSR000EFT\"  # K562 GATA1 ChIP-seq\n",
    "\n",
    "# Resolve track index\n",
    "if saliency_track_id not in bigwig_names:\n",
    "    raise ValueError(f\"Track not available for {assembly}: {saliency_track_id!r}\")\n",
    "saliency_track_idx = bigwig_names.index(saliency_track_id)\n",
    "print(\n",
    "    f\"Computing saliency for track: {saliency_track_id} \"\n",
    "    f\"(index {saliency_track_idx})\"\n",
    ")\n",
    "\n",
    "# Use shorter sequence for saliency (faster computation)\n",
    "saliency_seq_length = len(seq) // 2\n",
    "seq_start = (len(seq) - saliency_seq_length) // 2\n",
    "seq_end = seq_start + saliency_seq_length\n",
    "saliency_seq = seq[seq_start:seq_end]\n",
    "\n",
    "print(\n",
    "    f\"Saliency sequence: {len(saliency_seq):,} bp \"\n",
    "    f\"(center crop of full sequence)\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G44_fw_bDK7o"
   },
   "outputs": [],
   "source": [
    "# Define target window: center +/- 1024 bp of the output axis\n",
    "promoter_window_bp = 1024\n",
    "output_len = int(len(saliency_seq) * 0.375)  # predicted center span at 1 bp resolution\n",
    "center = output_len // 2\n",
    "target_window_bp = (center - promoter_window_bp, center + promoter_window_bp + 1)\n",
    "\n",
    "# Initialize saliency computer\n",
    "saliency_computer = SaliencyComputer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    track_indices=[saliency_track_idx],\n",
    "    species=species,\n",
    "    device=device,\n",
    "    target_window_bp=target_window_bp,\n",
    ")\n",
    "\n",
    "# Compute saliency\n",
    "print(\"Computing saliency (this may take a moment)...\")\n",
    "saliency_scores, one_hot_sequence = saliency_computer.compute_saliency(\n",
    "    saliency_seq\n",
    ")\n",
    "\n",
    "print(f\"\\n✅ Saliency computed\")\n",
    "print(f\"   Shape: {saliency_scores.shape}\")\n",
    "sal_min = saliency_scores.min()\n",
    "sal_max = saliency_scores.max()\n",
    "print(f\"   Value range: [{sal_min:.6f}, {sal_max:.6f}]\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A-eR1tK9DK7o"
   },
   "outputs": [],
   "source": [
    "# Nucleotide channels in logo order (A, T, C, G), upper case and soft-masked\n",
    "upper_ids = tokenizer(\"ATCG\", add_special_tokens=False)[\"input_ids\"]\n",
    "lower_ids = tokenizer(\"atcg\", add_special_tokens=False)[\"input_ids\"]\n",
    "assert tokenizer.unk_token_id not in upper_ids + lower_ids\n",
    "\n",
    "# Read each position from the channels of its own case: (positions, 4)\n",
    "soft_masked = np.array([base.islower() for base in saliency_seq])\n",
    "channels = np.where(soft_masked[:, None], lower_ids, upper_ids)\n",
    "gradients_acgt = np.take_along_axis(saliency_scores, channels, axis=1)\n",
    "one_hot_acgt = np.take_along_axis(one_hot_sequence, channels, axis=1)\n",
    "\n",
    "# Apply gradient correction (subtract mean over the four nucleotides at each position)\n",
    "saliency_corrected = gradients_acgt - np.mean(\n",
    "    gradients_acgt, axis=1, keepdims=True\n",
    ")\n",
    "\n",
    "# Calculate genomic positions\n",
    "saliency_start = start + seq_start\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mU9TdySBDK7o"
   },
   "outputs": [],
   "source": [
    "# Plot saliency sequence logo\n",
    "from fast_logomaker import FastLogo\n",
    "\n",
    "# Nucleotide color scheme\n",
    "NUCLEOTIDE_COLORS = {\n",
    "    'A': '#00b050',  # green\n",
    "    'C': '#3973fc',  # blue\n",
    "    'G': '#fdb932',  # yellow/gold\n",
    "    'T': '#ff554d'   # red\n",
    "}\n",
    "\n",
    "# True = contribution mode (multiply by one-hot), False = hypothetical mode\n",
    "CONTRIBUTION_MODE = True\n",
    "LOGO_WINDOW_SIZE = 128  # Number of positions to show in logo\n",
    "\n",
    "# Saliency is at nucleotide resolution. Channels are already A, T, C, G.\n",
    "n_saliency = len(saliency_corrected)\n",
    "cropped_start = (n_saliency - LOGO_WINDOW_SIZE) // 2\n",
    "cropped_end = cropped_start + LOGO_WINDOW_SIZE\n",
    "window_start_bp = saliency_start + cropped_start\n",
    "window_end_bp = saliency_start + cropped_end\n",
    "\n",
    "print(f\"Plotting saliency sequence logo...\")\n",
    "print(f\"  Contribution mode: {CONTRIBUTION_MODE}\")\n",
    "print(f\"  Window: {window_start_bp:,} - {window_end_bp:,} bp\")\n",
    "\n",
    "saliency_acgt = saliency_corrected\n",
    "if CONTRIBUTION_MODE:\n",
    "    saliency_acgt = saliency_acgt * one_hot_acgt\n",
    "\n",
    "window_saliency = saliency_acgt[cropped_start:cropped_end]\n",
    "print(f\"  Window size: {len(window_saliency)} positions\")\n",
    "\n",
    "logo_processor = FastLogo(\n",
    "    window_saliency[np.newaxis, :, :],\n",
    "    alphabet=['A', 'T', 'C', 'G'],\n",
    "    font_name='DejaVu Sans',\n",
    "    font_weight='bold',\n",
    "    color_scheme=NUCLEOTIDE_COLORS,\n",
    "    fade_below=0.0,\n",
    "    shade_below=0.0,\n",
    "    width=0.9,\n",
    "    figsize=(14, 3),\n",
    "    center_values=not CONTRIBUTION_MODE,\n",
    "    border=False,\n",
    "    fixed_ylim=True,\n",
    "    batch_size=1,\n",
    ")\n",
    "logo_processor.process_all()\n",
    "fig, ax = logo_processor.draw_single(0, border=False)\n",
    "logo_title = f'{chrom}:{window_start_bp:,}-{window_end_bp:,}'\n",
    "ax.set_title(logo_title, fontsize=11)\n",
    "\n",
    "n_positions = len(window_saliency)\n",
    "tick_positions = [0, n_positions // 4, n_positions // 2,\n",
    "                  3 * n_positions // 4, n_positions - 1]\n",
    "tick_labels = [f'{window_start_bp + pos:,}' for pos in tick_positions]\n",
    "ax.set_xticks(tick_positions)\n",
    "ax.set_xticklabels(tick_labels, fontsize=8)\n",
    "ax.set_xlabel('Genomic Position', fontsize=9)\n",
    "\n",
    "plt.yticks([])\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "print(\"\\n✅ Saliency logo plotted\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **How to read this plot.** The x axis is the genomic coordinate on `chrom`, for the central logo window of the saliency sequence. The score is the log of the predicted `saliency_track_id` signal summed over `target_window_bp`, ±`promoter_window_bp` bp (1,024 bp) around the center of the predicted span, not over the whole span. The logo window sits inside it. In contribution mode, each letter is the reference base and its height is the gradient after subtracting the mean over the four bases. A tall letter above the axis means that base raises the predicted signal. A letter below the axis means it lowers it. Saliency ranks positions. It does not predict the effect of a mutation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JF7UOTNHDK7o"
   },
   "source": [
    "## 📝 Summary\n",
    "\n",
    "This notebook demonstrated three complementary interpretability modalities of the NTv3 tracks model:\n",
    "\n",
    "- ✅ **Functional tracks**: Predicted regulatory activity across thousands of assays\n",
    "- ✅ **Attention maps**: Long-range information integration patterns\n",
    "- ✅ **Saliency maps**: Gradient-based sequence attributions\n",
    "\n",
    "### 💡 Next Steps\n",
    "\n",
    "- Try different genomic regions by changing `chrom`, `center_position`, `sequence_length`\n",
    "- Try another species by changing `species`; `assembly` must be its entry in `SPECIES_TO_ASSEMBLY`\n",
    "- Explore different tracks by modifying `tracks_to_plot` and `saliency_track_id`; the ids must be in that assembly's `bigwig_names`\n",
    "- Compare attention patterns across different layers\n",
    "- Compute saliency for different track combinations\n",
    "\n",
    "### 📚 Additional Resources\n",
    "\n",
    "- For track prediction tutorials, see: [01_tracks_prediction.ipynb](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/01_tracks_prediction.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}