{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Predicting Mean Ribosome Load with a 1D Convolutional Neural Network\n",
        "\n",
        "CSI 4106 — Introduction to Artificial Intelligence\n",
        "\n",
        "Marcel Turcotte  \n",
        "Version: Aug 31, 2026 14:55\n",
        "\n",
        "# Introduction\n",
        "\n",
        "Convolutional neural networks are most often introduced using images.\n",
        "Images provide compelling visual examples, but training image\n",
        "classifiers can require substantial computational resources. In this\n",
        "notebook, we study a smaller one-dimensional problem for which\n",
        "convolution is equally natural.\n",
        "\n",
        "Our machine-learning task can be stated without biological\n",
        "prerequisites:\n",
        "\n",
        "> Given a string of length 50 over the alphabet `{A, C, G, T}`, predict\n",
        "> a continuous value that usually lies between 0 and 13.\n",
        "\n",
        "The strings are 5′ untranslated regions (5′ UTRs), and the target is the\n",
        "mean ribosome load (MRL). A 5′ UTR is a segment of an RNA molecule that\n",
        "precedes the protein-coding region. MRL summarizes how many ribosomes\n",
        "are associated with the RNA and therefore provides an experimental\n",
        "measure related to translation.\n",
        "\n",
        "The experiment is adapted from Sample et al. and incorporates\n",
        "architectural ideas used in later work, including Tang et al. Our\n",
        "partitioning strategy is deliberately different. After removing\n",
        "low-coverage measurements, we randomly partition the data so that the\n",
        "test set is drawn from the same filtered population as the training set.\n",
        "\n",
        "> **Why a one-dimensional convolution?**\n",
        ">\n",
        "> A short convolutional kernel examines a local substring at every\n",
        "> position. The same learned weights are reused across the entire\n",
        "> sequence. This gives a 1D CNN the same two important inductive biases\n",
        "> as an image CNN: **local connectivity** and **parameter sharing**.\n",
        "\n",
        "## Learning objectives\n",
        "\n",
        "By the end of this notebook, you should be able to:\n",
        "\n",
        "- represent fixed-length strings as numeric tensors using one-hot\n",
        "  encoding;\n",
        "- explain how a `Conv1D` layer processes a sequence;\n",
        "- distinguish training, validation, and test data;\n",
        "- use early stopping to control neural-network training;\n",
        "- interpret loss, Pearson correlation, and $R^2$ learning curves; and\n",
        "- explain how predictions can be highly correlated with observations\n",
        "  while remaining poorly calibrated.\n",
        "\n",
        "# Preparation\n",
        "\n",
        "We use Keras with TensorFlow as its backend. NumPy and pandas support\n",
        "data processing, while Matplotlib and seaborn support visualization.\n",
        "Scikit-learn and SciPy provide familiar preprocessing and evaluation\n",
        "functions."
      ],
      "id": "dfb81364-f962-4316-998c-c886716b57dd"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "from contextlib import redirect_stderr, redirect_stdout\n",
        "from io import StringIO\n",
        "from pathlib import Path\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "from scipy.stats import pearsonr, spearmanr\n",
        "from sklearn.metrics import mean_squared_error, r2_score\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import callbacks, layers"
      ],
      "id": "imports"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "All important choices are collected in one configuration cell. Keeping\n",
        "these values together makes the experiment easier to inspect and\n",
        "reproduce."
      ],
      "id": "31dfe05c-a41d-472e-9576-8e3a5f8ef657"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "TensorFlow 2.21.0; training device: CPU"
          ]
        }
      ],
      "source": [
        "SEED = 42\n",
        "SEQUENCE_LENGTH = 50\n",
        "\n",
        "N_RETAINED = 275_000\n",
        "N_TRAIN = 220_000\n",
        "N_VALIDATION = 27_500\n",
        "N_TEST = 27_500\n",
        "\n",
        "BATCH_SIZE = 512\n",
        "MAX_EPOCHS = 50\n",
        "PATIENCE = 5\n",
        "\n",
        "# Shared limits make the two prediction plots directly comparable.\n",
        "MRL_PLOT_LIMITS = (-0.5, 13.5)\n",
        "\n",
        "DATA_DIR = Path(\"data\")\n",
        "MODEL_DIR = Path(\"models\")\n",
        "DATA_DIR.mkdir(exist_ok=True)\n",
        "MODEL_DIR.mkdir(exist_ok=True)\n",
        "\n",
        "DATASETS = {\n",
        "    \"GSM3130435_egfp_unmod_1.csv.gz\": (\n",
        "        \"https://ftp.ncbi.nlm.nih.gov/geo/samples/GSM3130nnn/\"\n",
        "        \"GSM3130435/suppl/GSM3130435_egfp_unmod_1.csv.gz\"\n",
        "    ),\n",
        "    \"GSE232927_processed_defined_end_tcell_r1.csv.gz\": (\n",
        "        \"https://ftp.ncbi.nlm.nih.gov/geo/series/GSE232nnn/\"\n",
        "        \"GSE232927/suppl/\"\n",
        "        \"GSE232927_processed_defined_end_tcell_r1.csv.gz\"\n",
        "    ),\n",
        "    \"GSE232927_processed_defined_end_tcell_r2.csv.gz\": (\n",
        "        \"https://ftp.ncbi.nlm.nih.gov/geo/series/GSE232nnn/\"\n",
        "        \"GSE232927/suppl/\"\n",
        "        \"GSE232927_processed_defined_end_tcell_r2.csv.gz\"\n",
        "    ),\n",
        "}\n",
        "\n",
        "assert N_TRAIN + N_VALIDATION + N_TEST == N_RETAINED\n",
        "\n",
        "keras.utils.set_random_seed(SEED)\n",
        "sns.set_theme(style=\"ticks\", context=\"notebook\")\n",
        "\n",
        "accelerators = tf.config.list_physical_devices(\"GPU\")\n",
        "device_description = \"GPU\" if accelerators else \"CPU\"\n",
        "print(f\"TensorFlow {tf.__version__}; training device: {device_description}\")"
      ],
      "id": "configuration"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The helper below downloads a dataset only when it is not already\n",
        "available. `keras.utils.get_file` also returns the local path, which we\n",
        "can pass directly to pandas. Its progress display is suppressed by\n",
        "default because terminal control characters can produce long outputs in\n",
        "a rendered notebook. Passing `show_progress=True` restores the\n",
        "interactive progress bar."
      ],
      "id": "3095707a-7197-43ab-bf7f-5e3ff740b404"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "def download_dataset(filename, show_progress=False):\n",
        "    \"\"\"Download one of the configured datasets and return its local path.\"\"\"\n",
        "    if filename not in DATASETS:\n",
        "        raise KeyError(f\"Unknown dataset: {filename}\")\n",
        "\n",
        "    arguments = {\n",
        "        \"fname\": filename,\n",
        "        \"origin\": DATASETS[filename],\n",
        "        \"cache_dir\": str(Path.cwd()),\n",
        "        \"cache_subdir\": str(DATA_DIR),\n",
        "    }\n",
        "\n",
        "    if show_progress:\n",
        "        path = keras.utils.get_file(**arguments)\n",
        "    else:\n",
        "        with redirect_stdout(StringIO()), redirect_stderr(StringIO()):\n",
        "            path = keras.utils.get_file(**arguments)\n",
        "\n",
        "    return Path(path)"
      ],
      "id": "download-helper"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Loading and validating the main dataset\n",
        "\n",
        "The source files use slightly different column names. We standardize\n",
        "them at the boundary of our program so that all subsequent functions can\n",
        "use the same simple vocabulary: `utr`, `mrl`, and `read_count`."
      ],
      "id": "74500398-ae52-4440-8933-ce492c912bfb"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "def load_utr_dataset(path, expected_length=SEQUENCE_LENGTH):\n",
        "    \"\"\"Load and validate a processed UTR dataset.\n",
        "\n",
        "    The returned DataFrame always contains the columns ``utr``, ``mrl``, and\n",
        "    ``read_count``.\n",
        "    \"\"\"\n",
        "    data = pd.read_csv(path)\n",
        "\n",
        "    if \"utr\" not in data.columns and \"UTR\" in data.columns:\n",
        "        data = data.rename(columns={\"UTR\": \"utr\"})\n",
        "\n",
        "    if \"total_reads\" in data.columns:\n",
        "        read_column = \"total_reads\"\n",
        "    elif \"total\" in data.columns:\n",
        "        read_column = \"total\"\n",
        "    else:\n",
        "        raise ValueError(\"The dataset has no total_reads or total column.\")\n",
        "\n",
        "    required = {\"utr\", \"rl\", read_column}\n",
        "    missing = required.difference(data.columns)\n",
        "    if missing:\n",
        "        raise ValueError(f\"Missing required columns: {sorted(missing)}\")\n",
        "\n",
        "    data = (\n",
        "        data.loc[:, [\"utr\", \"rl\", read_column]]\n",
        "        .rename(columns={\"rl\": \"mrl\", read_column: \"read_count\"})\n",
        "        .copy()\n",
        "    )\n",
        "\n",
        "    data[\"utr\"] = data[\"utr\"].astype(\"string\").str.upper()\n",
        "    data[\"mrl\"] = pd.to_numeric(data[\"mrl\"], errors=\"coerce\")\n",
        "    data[\"read_count\"] = pd.to_numeric(\n",
        "        data[\"read_count\"], errors=\"coerce\"\n",
        "    )\n",
        "\n",
        "    if data[[\"utr\", \"mrl\", \"read_count\"]].isna().any().any():\n",
        "        raise ValueError(\"Missing or non-numeric values were found.\")\n",
        "\n",
        "    valid_sequence = data[\"utr\"].str.fullmatch(\n",
        "        rf\"[ACGTN]{{{expected_length}}}\"\n",
        "    )\n",
        "    if not valid_sequence.all():\n",
        "        n_invalid = int((~valid_sequence).sum())\n",
        "        raise ValueError(\n",
        "            f\"Found {n_invalid} sequences that are not length \"\n",
        "            f\"{expected_length} strings over A, C, G, T, and N.\"\n",
        "        )\n",
        "\n",
        "    if (data[\"read_count\"] < 0).any():\n",
        "        raise ValueError(\"Read counts must be non-negative.\")\n",
        "\n",
        "    if data[\"utr\"].duplicated().any():\n",
        "        raise ValueError(\"The dataset contains duplicate UTR sequences.\")\n",
        "\n",
        "    return data.reset_index(drop=True)"
      ],
      "id": "load-helper"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loaded 326,033 sequences."
          ]
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "main_path = download_dataset(\"GSM3130435_egfp_unmod_1.csv.gz\")\n",
        "main_data = load_utr_dataset(main_path)\n",
        "\n",
        "print(f\"Loaded {len(main_data):,} sequences.\")\n",
        "main_data.head()"
      ],
      "id": "cell-load-main-data"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The target and read counts are measured on very different scales. A\n",
        "compact summary gives us an initial view of both variables."
      ],
      "id": "42fbd258-eafb-457f-8760-c749fe88ed86"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "main_data[[\"mrl\", \"read_count\"]].describe(\n",
        "    percentiles=[0.01, 0.25, 0.50, 0.75, 0.99]\n",
        ").round(3)"
      ],
      "id": "cell-summarize-main-data"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Filtering measurements by read count\n",
        "\n",
        "Measurements based on very few reads are less precise. We therefore\n",
        "retain exactly the 275,000 sequences with the highest read counts."
      ],
      "id": "ec302700-0b23-4e78-a07d-fe28a5451331"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Read-count boundary: 294\n",
            "Retained: 275,000 sequences (84.3%)"
          ]
        }
      ],
      "source": [
        "def select_top_by_reads(data, n):\n",
        "    \"\"\"Return the n highest-read rows and the read count at the boundary.\"\"\"\n",
        "    if not 0 < n <= len(data):\n",
        "        raise ValueError(\"n must be between 1 and the number of rows.\")\n",
        "\n",
        "    ranked = data.sort_values(\n",
        "        \"read_count\", ascending=False, kind=\"stable\"\n",
        "    ).reset_index(drop=True)\n",
        "    cutoff = float(ranked.loc[n - 1, \"read_count\"])\n",
        "    return ranked.iloc[:n].copy(), cutoff\n",
        "\n",
        "\n",
        "retained_data, read_cutoff = select_top_by_reads(main_data, N_RETAINED)\n",
        "retained_fraction = N_RETAINED / len(main_data)\n",
        "\n",
        "print(f\"Read-count boundary: {read_cutoff:,.0f}\")\n",
        "print(f\"Retained: {N_RETAINED:,} sequences ({retained_fraction:.1%})\")"
      ],
      "id": "select-top-helper"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read counts are strongly right-skewed. We plot\n",
        "$log_{10}(\\text{reads}+1)$ so that both low- and high-coverage\n",
        "measurements remain visible. The red line is the read count of the\n",
        "275,000th ranked sequence."
      ],
      "id": "4a5ca067-ff04-417e-a453-459796ce5186"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "log_read_counts = np.log10(main_data[\"read_count\"].to_numpy() + 1)\n",
        "log_cutoff = np.log10(read_cutoff + 1)\n",
        "\n",
        "fig, ax = plt.subplots(figsize=(9, 4.5))\n",
        "sns.histplot(log_read_counts, bins=80, color=\"steelblue\", ax=ax)\n",
        "ax.axvline(log_cutoff, color=\"red\", linewidth=2)\n",
        "ax.text(\n",
        "    log_cutoff,\n",
        "    ax.get_ylim()[1] * 0.92,\n",
        "    f\"  cutoff = {read_cutoff:,.0f} reads\",\n",
        "    color=\"red\",\n",
        "    va=\"top\",\n",
        ")\n",
        "ax.set(\n",
        "    xlabel=r\"$\\log_{10}(\\mathrm{read\\ count}+1)$\",\n",
        "    ylabel=\"Number of sequences\",\n",
        "    title=\"Read coverage before filtering\",\n",
        ")\n",
        "sns.despine()\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ],
      "id": "cell-plot-read-histogram"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Note**\n",
        ">\n",
        "> Because read counts are integers, multiple sequences can be tied at\n",
        "> the boundary. We use a stable ranking and retain exactly 275,000 rows.\n",
        "> The red line therefore represents the boundary value, not a rule that\n",
        "> resolves ties.\n",
        "\n",
        "# Representing nucleotide strings\n",
        "\n",
        "A neural network operates on numbers rather than characters. We\n",
        "represent each nucleotide using four binary values:\n",
        "\n",
        "| Nucleotide  |   A |   C |   G |   T |\n",
        "|:------------|----:|----:|----:|----:|\n",
        "| A           |   1 |   0 |   0 |   0 |\n",
        "| C           |   0 |   1 |   0 |   0 |\n",
        "| G           |   0 |   0 |   1 |   0 |\n",
        "| T           |   0 |   0 |   0 |   1 |\n",
        "| N (unknown) |   0 |   0 |   0 |   0 |\n",
        "\n",
        "A string of length 50 consequently becomes a $50 \\times 4$ matrix."
      ],
      "id": "20b0bc7e-b597-43c2-8347-3010dc45a783"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "def one_hot_encode(sequences, sequence_length=SEQUENCE_LENGTH):\n",
        "    \"\"\"Encode nucleotide strings as an (n, sequence_length, 4) array.\"\"\"\n",
        "    sequences = pd.Series(sequences, dtype=\"string\").str.upper()\n",
        "\n",
        "    valid = sequences.str.fullmatch(rf\"[ACGTN]{{{sequence_length}}}\")\n",
        "    if not valid.all():\n",
        "        raise ValueError(\"Every sequence must have the expected format.\")\n",
        "\n",
        "    # Convert all ASCII characters at once and look up their channel indices.\n",
        "    characters = np.frombuffer(\n",
        "        \"\".join(sequences).encode(\"ascii\"), dtype=np.uint8\n",
        "    ).reshape(-1, sequence_length)\n",
        "\n",
        "    channel_lookup = np.full(256, -1, dtype=np.int8)\n",
        "    for channel, nucleotide in enumerate(\"ACGT\"):\n",
        "        channel_lookup[ord(nucleotide)] = channel\n",
        "\n",
        "    channels = channel_lookup[characters]\n",
        "    encoded = np.zeros(\n",
        "        (len(sequences), sequence_length, 4), dtype=np.float32\n",
        "    )\n",
        "\n",
        "    rows, positions = np.nonzero(channels >= 0)\n",
        "    encoded[rows, positions, channels[rows, positions]] = 1.0\n",
        "    return encoded"
      ],
      "id": "one-hot-helper"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We visualize the beginning of one sequence before encoding the complete\n",
        "dataset."
      ],
      "id": "9d6f5d65-c207-4f2d-a3a6-dfb8f5000ca4"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Complete sequence length: 50\n",
            "Encoded shape: (50, 4)"
          ]
        }
      ],
      "source": [
        "example_sequence = retained_data.loc[0, \"utr\"]\n",
        "example_encoding = one_hot_encode([example_sequence])[0]\n",
        "positions_to_show = 12\n",
        "\n",
        "fig, ax = plt.subplots(figsize=(9, 2.8))\n",
        "sns.heatmap(\n",
        "    example_encoding[:positions_to_show].T,\n",
        "    cmap=\"Blues\",\n",
        "    vmin=0,\n",
        "    vmax=1,\n",
        "    cbar=False,\n",
        "    linewidths=0.5,\n",
        "    linecolor=\"lightgray\",\n",
        "    xticklabels=np.arange(1, positions_to_show + 1),\n",
        "    yticklabels=list(\"ACGT\"),\n",
        "    ax=ax,\n",
        ")\n",
        "ax.set(\n",
        "    xlabel=\"Position in the sequence\",\n",
        "    ylabel=\"Channel\",\n",
        "    title=f\"Sequence prefix: {example_sequence[:positions_to_show]}\",\n",
        ")\n",
        "plt.tight_layout()\n",
        "plt.show()\n",
        "\n",
        "print(\"Complete sequence length:\", len(example_sequence))\n",
        "print(\"Encoded shape:\", example_encoding.shape)"
      ],
      "id": "cell-visualize-one-hot"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Experiment 1: random partitioning\n",
        "\n",
        "## Creating training, validation, and test sets\n",
        "\n",
        "We randomly shuffle the retained sequences and then create three\n",
        "disjoint partitions:\n",
        "\n",
        "- 220,000 sequences for fitting the model;\n",
        "- 27,500 sequences for early stopping and model selection; and\n",
        "- 27,500 sequences for the final evaluation.\n",
        "\n",
        "The test set remains untouched until training is complete."
      ],
      "id": "1c4bbef0-a67a-44ec-aa5a-7ed275362876"
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [
        "def random_partition(data, seed=SEED):\n",
        "    \"\"\"Create the agreed random training, validation, and test partitions.\"\"\"\n",
        "    if len(data) != N_RETAINED:\n",
        "        raise ValueError(f\"Expected exactly {N_RETAINED:,} retained rows.\")\n",
        "\n",
        "    shuffled = data.sample(frac=1, random_state=seed).reset_index(drop=True)\n",
        "\n",
        "    train = shuffled.iloc[:N_TRAIN].copy()\n",
        "    validation = shuffled.iloc[\n",
        "        N_TRAIN:N_TRAIN + N_VALIDATION\n",
        "    ].copy()\n",
        "    test = shuffled.iloc[N_TRAIN + N_VALIDATION:].copy()\n",
        "\n",
        "    if (len(train), len(validation), len(test)) != (\n",
        "        N_TRAIN,\n",
        "        N_VALIDATION,\n",
        "        N_TEST,\n",
        "    ):\n",
        "        raise AssertionError(\"Unexpected partition sizes.\")\n",
        "\n",
        "    train_sequences = set(train[\"utr\"])\n",
        "    validation_sequences = set(validation[\"utr\"])\n",
        "    test_sequences = set(test[\"utr\"])\n",
        "\n",
        "    if not train_sequences.isdisjoint(validation_sequences):\n",
        "        raise AssertionError(\"Training and validation sets overlap.\")\n",
        "    if not train_sequences.isdisjoint(test_sequences):\n",
        "        raise AssertionError(\"Training and test sets overlap.\")\n",
        "    if not validation_sequences.isdisjoint(test_sequences):\n",
        "        raise AssertionError(\"Validation and test sets overlap.\")\n",
        "\n",
        "    return train, validation, test\n",
        "\n",
        "\n",
        "train_data, validation_data, test_data = random_partition(retained_data)"
      ],
      "id": "random-partition-helper"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The three partitions should have similar MRL and read-count\n",
        "distributions. This small table is a useful check on the randomization."
      ],
      "id": "0b5b0b22-f05e-4d6a-8567-9225ab286c5a"
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "def summarize_partitions(partitions):\n",
        "    \"\"\"Summarize the size and target distribution of named partitions.\"\"\"\n",
        "    rows = []\n",
        "    for name, data in partitions.items():\n",
        "        rows.append(\n",
        "            {\n",
        "                \"partition\": name,\n",
        "                \"n\": len(data),\n",
        "                \"mean MRL\": data[\"mrl\"].mean(),\n",
        "                \"SD MRL\": data[\"mrl\"].std(),\n",
        "                \"median reads\": data[\"read_count\"].median(),\n",
        "            }\n",
        "        )\n",
        "    return pd.DataFrame(rows).set_index(\"partition\")\n",
        "\n",
        "\n",
        "random_partition_summary = summarize_partitions(\n",
        "    {\n",
        "        \"Training\": train_data,\n",
        "        \"Validation\": validation_data,\n",
        "        \"Test\": test_data,\n",
        "    }\n",
        ")\n",
        "random_partition_summary.round(3)"
      ],
      "id": "cell-summarize-random-partitions"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We standardize the training targets to have mean zero and standard\n",
        "deviation one. The same transformation is then applied to validation and\n",
        "test targets. Fitting this transformation on training data only prevents\n",
        "information from the validation and test sets from leaking into model\n",
        "development."
      ],
      "id": "87097a94-bd32-4857-8c95-5b8d55119871"
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training features: (220000, 50, 4) float32\n",
            "Validation features: (27500, 50, 4) float32\n",
            "Test features: (27500, 50, 4) float32\n",
            "Training targets: (220000, 1) float32"
          ]
        }
      ],
      "source": [
        "target_scaler = StandardScaler()\n",
        "\n",
        "y_train = target_scaler.fit_transform(\n",
        "    train_data[[\"mrl\"]]\n",
        ").astype(np.float32)\n",
        "y_validation = target_scaler.transform(\n",
        "    validation_data[[\"mrl\"]]\n",
        ").astype(np.float32)\n",
        "y_test = target_scaler.transform(\n",
        "    test_data[[\"mrl\"]]\n",
        ").astype(np.float32)\n",
        "\n",
        "x_train = one_hot_encode(train_data[\"utr\"])\n",
        "x_validation = one_hot_encode(validation_data[\"utr\"])\n",
        "x_test = one_hot_encode(test_data[\"utr\"])\n",
        "\n",
        "print(\"Training features:\", x_train.shape, x_train.dtype)\n",
        "print(\"Validation features:\", x_validation.shape, x_validation.dtype)\n",
        "print(\"Test features:\", x_test.shape, x_test.dtype)\n",
        "print(\"Training targets:\", y_train.shape, y_train.dtype)"
      ],
      "id": "prepare-arrays"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Building the convolutional neural network\n",
        "\n",
        "The input shape is `(50, 4)`: 50 positions and four nucleotide channels.\n",
        "Keras uses this channels-last representation directly, so no tensor\n",
        "permutation is required."
      ],
      "id": "9207f29f-c522-4454-b8a1-b722689b7b8b"
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"mrl_cnn\"</span>\n",
              "</pre>"
            ]
          }
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓\n",
              "┃<span style=\"font-weight: bold\"> Layer (type)                    </span>┃<span style=\"font-weight: bold\"> Output Shape           </span>┃<span style=\"font-weight: bold\">       Param # </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩\n",
              "│ convolution_1 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Conv1D</span>)          │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │         <span style=\"color: #00af00; text-decoration-color: #00af00\">5,280</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ convolution_2 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Conv1D</span>)          │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │       <span style=\"color: #00af00; text-decoration-color: #00af00\">204,960</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ batch_normalization_1           │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │           <span style=\"color: #00af00; text-decoration-color: #00af00\">640</span> │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">BatchNormalization</span>)            │                        │               │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ dropout_1 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Dropout</span>)             │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │             <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ convolution_3 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Conv1D</span>)          │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │       <span style=\"color: #00af00; text-decoration-color: #00af00\">204,960</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ batch_normalization_2           │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │           <span style=\"color: #00af00; text-decoration-color: #00af00\">640</span> │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">BatchNormalization</span>)            │                        │               │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ dropout_2 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Dropout</span>)             │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">160</span>)        │             <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ convolution_4 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Conv1D</span>)          │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">80</span>)         │       <span style=\"color: #00af00; text-decoration-color: #00af00\">102,480</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ batch_normalization_3           │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">80</span>)         │           <span style=\"color: #00af00; text-decoration-color: #00af00\">320</span> │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">BatchNormalization</span>)            │                        │               │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ dropout_3 (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Dropout</span>)             │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">50</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">80</span>)         │             <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ flatten (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Flatten</span>)               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">4000</span>)           │             <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ latent_representation (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Dense</span>)   │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">80</span>)             │       <span style=\"color: #00af00; text-decoration-color: #00af00\">320,080</span> │\n",
              "├─────────────────────────────────┼────────────────────────┼───────────────┤\n",
              "│ scaled_mrl (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">Dense</span>)              │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">1</span>)              │            <span style=\"color: #00af00; text-decoration-color: #00af00\">81</span> │\n",
              "└─────────────────────────────────┴────────────────────────┴───────────────┘\n",
              "</pre>"
            ]
          }
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">839,441</span> (3.20 MB)\n",
              "</pre>"
            ]
          }
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">838,641</span> (3.20 MB)\n",
              "</pre>"
            ]
          }
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">800</span> (3.12 KB)\n",
              "</pre>"
            ]
          }
        }
      ],
      "source": [
        "def build_mrl_cnn(\n",
        "    sequence_length=SEQUENCE_LENGTH,\n",
        "    n_filters=160,\n",
        "    kernel_size=8,\n",
        "    latent_dimension=80,\n",
        "):\n",
        "    \"\"\"Build the 1D CNN used to predict standardized MRL.\"\"\"\n",
        "    model = keras.Sequential(\n",
        "        [\n",
        "            keras.Input(shape=(sequence_length, 4), name=\"one_hot_utr\"),\n",
        "            layers.Conv1D(\n",
        "                n_filters,\n",
        "                kernel_size,\n",
        "                padding=\"same\",\n",
        "                activation=\"relu\",\n",
        "                name=\"convolution_1\",\n",
        "            ),\n",
        "            layers.Conv1D(\n",
        "                n_filters,\n",
        "                kernel_size,\n",
        "                padding=\"same\",\n",
        "                activation=\"relu\",\n",
        "                name=\"convolution_2\",\n",
        "            ),\n",
        "            layers.BatchNormalization(name=\"batch_normalization_1\"),\n",
        "            layers.Dropout(0.2, name=\"dropout_1\"),\n",
        "            layers.Conv1D(\n",
        "                n_filters,\n",
        "                kernel_size,\n",
        "                padding=\"same\",\n",
        "                activation=\"relu\",\n",
        "                name=\"convolution_3\",\n",
        "            ),\n",
        "            layers.BatchNormalization(name=\"batch_normalization_2\"),\n",
        "            layers.Dropout(0.4, name=\"dropout_2\"),\n",
        "            layers.Conv1D(\n",
        "                n_filters // 2,\n",
        "                kernel_size,\n",
        "                padding=\"same\",\n",
        "                activation=\"relu\",\n",
        "                name=\"convolution_4\",\n",
        "            ),\n",
        "            layers.BatchNormalization(name=\"batch_normalization_3\"),\n",
        "            layers.Dropout(0.2, name=\"dropout_3\"),\n",
        "            layers.Flatten(name=\"flatten\"),\n",
        "            layers.Dense(\n",
        "                latent_dimension,\n",
        "                activation=\"relu\",\n",
        "                name=\"latent_representation\",\n",
        "            ),\n",
        "            layers.Dense(1, name=\"scaled_mrl\"),\n",
        "        ],\n",
        "        name=\"mrl_cnn\",\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "model = build_mrl_cnn()\n",
        "model.summary()"
      ],
      "id": "model-helper"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Every convolution uses a kernel width of eight and `padding=\"same\"`.\n",
        "Consequently, the spatial dimension remains 50 while the network\n",
        "constructs increasingly abstract representations of local sequence\n",
        "patterns. The final linear unit outputs one standardized MRL prediction.\n",
        "\n",
        "## Training the model\n",
        "\n",
        "Keras provides a streaming implementation of $R^2$. We implement Pearson\n",
        "correlation in the same way by accumulating the sufficient statistics\n",
        "over an entire epoch. Calculating correlation separately in each\n",
        "mini-batch and then averaging those values would not equal the\n",
        "correlation over the full dataset."
      ],
      "id": "932ec9a4-a047-493e-93e3-8f1f9c291298"
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [],
      "source": [
        "@keras.utils.register_keras_serializable(package=\"CSI4106\")\n",
        "class PearsonCorrelation(keras.metrics.Metric):\n",
        "    \"\"\"Pearson correlation accumulated over all examples in an epoch.\"\"\"\n",
        "\n",
        "    def __init__(self, name=\"pearson_r\", **kwargs):\n",
        "        super().__init__(name=name, **kwargs)\n",
        "        self.count = self.add_weight(name=\"count\", initializer=\"zeros\")\n",
        "        self.sum_true = self.add_weight(name=\"sum_true\", initializer=\"zeros\")\n",
        "        self.sum_pred = self.add_weight(name=\"sum_pred\", initializer=\"zeros\")\n",
        "        self.sum_true_squared = self.add_weight(\n",
        "            name=\"sum_true_squared\", initializer=\"zeros\"\n",
        "        )\n",
        "        self.sum_pred_squared = self.add_weight(\n",
        "            name=\"sum_pred_squared\", initializer=\"zeros\"\n",
        "        )\n",
        "        self.sum_products = self.add_weight(\n",
        "            name=\"sum_products\", initializer=\"zeros\"\n",
        "        )\n",
        "\n",
        "    def update_state(self, y_true, y_pred, sample_weight=None):\n",
        "        y_true = tf.cast(tf.reshape(y_true, [-1]), self.dtype)\n",
        "        y_pred = tf.cast(tf.reshape(y_pred, [-1]), self.dtype)\n",
        "\n",
        "        self.count.assign_add(tf.cast(tf.size(y_true), self.dtype))\n",
        "        self.sum_true.assign_add(tf.reduce_sum(y_true))\n",
        "        self.sum_pred.assign_add(tf.reduce_sum(y_pred))\n",
        "        self.sum_true_squared.assign_add(tf.reduce_sum(tf.square(y_true)))\n",
        "        self.sum_pred_squared.assign_add(tf.reduce_sum(tf.square(y_pred)))\n",
        "        self.sum_products.assign_add(tf.reduce_sum(y_true * y_pred))\n",
        "\n",
        "    def result(self):\n",
        "        covariance = (\n",
        "            self.sum_products\n",
        "            - self.sum_true * self.sum_pred / self.count\n",
        "        )\n",
        "        true_variation = (\n",
        "            self.sum_true_squared\n",
        "            - tf.square(self.sum_true) / self.count\n",
        "        )\n",
        "        pred_variation = (\n",
        "            self.sum_pred_squared\n",
        "            - tf.square(self.sum_pred) / self.count\n",
        "        )\n",
        "        denominator = tf.sqrt(true_variation * pred_variation)\n",
        "        return tf.math.divide_no_nan(covariance, denominator)\n",
        "\n",
        "    def reset_state(self):\n",
        "        for variable in self.variables:\n",
        "            variable.assign(0)"
      ],
      "id": "pearson-metric"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Mean squared error is both our optimization objective and the loss\n",
        "displayed during training. Pearson correlation and $R^2$ are monitoring\n",
        "metrics; they do not change the gradient updates."
      ],
      "id": "bcc1f39e-b6f9-45f4-8909-8a31abcde87f"
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [],
      "source": [
        "def compile_mrl_cnn(model):\n",
        "    \"\"\"Configure the optimizer, loss, and epoch-level metrics.\"\"\"\n",
        "    model.compile(\n",
        "        optimizer=keras.optimizers.Adam(learning_rate=1e-3),\n",
        "        loss=\"mean_squared_error\",\n",
        "        metrics=[\n",
        "            PearsonCorrelation(),\n",
        "            keras.metrics.R2Score(name=\"r2\"),\n",
        "        ],\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "model = compile_mrl_cnn(model)"
      ],
      "id": "compile-model"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Early stopping monitors validation loss. When validation loss has not\n",
        "improved for five epochs, training stops and Keras restores the weights\n",
        "from the best epoch. A checkpoint is also written to disk so that the\n",
        "model can be reused in later notebooks.\n",
        "\n",
        "> **Expected runtime**\n",
        ">\n",
        "> This is the most computationally expensive cell in the notebook.\n",
        "> Runtime depends strongly on the available TensorFlow accelerator. The\n",
        "> cell trains on all 220,000 examples and may take several minutes."
      ],
      "id": "43e63f36-49ff-4a04-b2cd-a7461aae65ff"
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/50\n",
            "430/430 - 129s - 301ms/step - loss: 0.4637 - pearson_r: 0.7361 - r2: 0.5363 - val_loss: 1.5230 - val_pearson_r: 0.7143 - val_r2: -5.2182e-01\n",
            "Epoch 2/50\n",
            "430/430 - 127s - 296ms/step - loss: 0.2081 - pearson_r: 0.8899 - r2: 0.7919 - val_loss: 0.1739 - val_pearson_r: 0.9196 - val_r2: 0.8263\n",
            "Epoch 3/50\n",
            "430/430 - 127s - 296ms/step - loss: 0.1468 - pearson_r: 0.9237 - r2: 0.8532 - val_loss: 0.1460 - val_pearson_r: 0.9374 - val_r2: 0.8541\n",
            "Epoch 4/50\n",
            "430/430 - 127s - 295ms/step - loss: 0.1285 - pearson_r: 0.9336 - r2: 0.8715 - val_loss: 0.1218 - val_pearson_r: 0.9440 - val_r2: 0.8783\n",
            "Epoch 5/50\n",
            "430/430 - 126s - 293ms/step - loss: 0.1202 - pearson_r: 0.9380 - r2: 0.8798 - val_loss: 0.1220 - val_pearson_r: 0.9462 - val_r2: 0.8781\n",
            "Epoch 6/50\n",
            "430/430 - 130s - 302ms/step - loss: 0.1141 - pearson_r: 0.9412 - r2: 0.8859 - val_loss: 0.1196 - val_pearson_r: 0.9486 - val_r2: 0.8805\n",
            "Epoch 7/50\n",
            "430/430 - 127s - 296ms/step - loss: 0.1106 - pearson_r: 0.9431 - r2: 0.8894 - val_loss: 0.1036 - val_pearson_r: 0.9501 - val_r2: 0.8965\n",
            "Epoch 8/50\n",
            "430/430 - 128s - 297ms/step - loss: 0.1069 - pearson_r: 0.9450 - r2: 0.8931 - val_loss: 0.1129 - val_pearson_r: 0.9510 - val_r2: 0.8872\n",
            "Epoch 9/50\n",
            "430/430 - 129s - 301ms/step - loss: 0.1045 - pearson_r: 0.9463 - r2: 0.8955 - val_loss: 0.0982 - val_pearson_r: 0.9521 - val_r2: 0.9019\n",
            "Epoch 10/50\n",
            "430/430 - 130s - 302ms/step - loss: 0.1023 - pearson_r: 0.9475 - r2: 0.8977 - val_loss: 0.1040 - val_pearson_r: 0.9519 - val_r2: 0.8961\n",
            "Epoch 11/50\n",
            "430/430 - 128s - 299ms/step - loss: 0.1001 - pearson_r: 0.9486 - r2: 0.8999 - val_loss: 0.0949 - val_pearson_r: 0.9525 - val_r2: 0.9052\n",
            "Epoch 12/50\n",
            "430/430 - 97s - 226ms/step - loss: 0.0988 - pearson_r: 0.9493 - r2: 0.9012 - val_loss: 0.0953 - val_pearson_r: 0.9532 - val_r2: 0.9048\n",
            "Epoch 13/50\n",
            "430/430 - 98s - 229ms/step - loss: 0.0973 - pearson_r: 0.9501 - r2: 0.9027 - val_loss: 0.0998 - val_pearson_r: 0.9530 - val_r2: 0.9003\n",
            "Epoch 14/50\n",
            "430/430 - 97s - 226ms/step - loss: 0.0956 - pearson_r: 0.9510 - r2: 0.9044 - val_loss: 0.0928 - val_pearson_r: 0.9537 - val_r2: 0.9073\n",
            "Epoch 15/50\n",
            "430/430 - 98s - 229ms/step - loss: 0.0944 - pearson_r: 0.9516 - r2: 0.9056 - val_loss: 0.0960 - val_pearson_r: 0.9536 - val_r2: 0.9041\n",
            "Epoch 16/50\n",
            "430/430 - 103s - 239ms/step - loss: 0.0931 - pearson_r: 0.9523 - r2: 0.9069 - val_loss: 0.0936 - val_pearson_r: 0.9539 - val_r2: 0.9065\n",
            "Epoch 17/50\n",
            "430/430 - 100s - 232ms/step - loss: 0.0915 - pearson_r: 0.9531 - r2: 0.9085 - val_loss: 0.1032 - val_pearson_r: 0.9542 - val_r2: 0.8969\n",
            "Epoch 18/50\n",
            "430/430 - 115s - 268ms/step - loss: 0.0902 - pearson_r: 0.9538 - r2: 0.9098 - val_loss: 0.1068 - val_pearson_r: 0.9531 - val_r2: 0.8933\n",
            "Epoch 19/50\n",
            "430/430 - 141s - 328ms/step - loss: 0.0893 - pearson_r: 0.9543 - r2: 0.9107 - val_loss: 0.0974 - val_pearson_r: 0.9537 - val_r2: 0.9027\n",
            "Epoch 19: early stopping\n",
            "Restoring model weights from the end of the best epoch: 14."
          ]
        }
      ],
      "source": [
        "random_split_model_path = MODEL_DIR / \"mrl_cnn_random_split.keras\"\n",
        "\n",
        "training_callbacks = [\n",
        "    callbacks.EarlyStopping(\n",
        "        monitor=\"val_loss\",\n",
        "        mode=\"min\",\n",
        "        patience=PATIENCE,\n",
        "        restore_best_weights=True,\n",
        "        verbose=1,\n",
        "    ),\n",
        "    callbacks.ModelCheckpoint(\n",
        "        filepath=str(random_split_model_path),\n",
        "        monitor=\"val_loss\",\n",
        "        mode=\"min\",\n",
        "        save_best_only=True,\n",
        "        verbose=0,\n",
        "    ),\n",
        "]\n",
        "\n",
        "history = model.fit(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    validation_data=(x_validation, y_validation),\n",
        "    epochs=MAX_EPOCHS,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    shuffle=True,\n",
        "    callbacks=training_callbacks,\n",
        "    verbose=2,\n",
        ")"
      ],
      "id": "train-model"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualizing learning dynamics\n",
        "\n",
        "The Keras `History` object records every loss and metric named during\n",
        "model compilation. We display three complementary views of learning."
      ],
      "id": "898249e6-2d94-4b4e-813e-5d09dd3340a2"
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [],
      "source": [
        "def plot_training_history(history):\n",
        "    \"\"\"Plot training and validation loss, Pearson r, and R-squared.\"\"\"\n",
        "    values = history.history\n",
        "    epochs = np.arange(1, len(values[\"loss\"]) + 1)\n",
        "    best_epoch = int(np.argmin(values[\"val_loss\"])) + 1\n",
        "\n",
        "    panels = [\n",
        "        (\"loss\", \"Mean squared error\", \"Loss\"),\n",
        "        (\"pearson_r\", \"Pearson correlation\", r\"Pearson $r$\"),\n",
        "        (\"r2\", \"Coefficient of determination\", r\"$R^2$\"),\n",
        "    ]\n",
        "\n",
        "    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))\n",
        "    for ax, (metric, title, ylabel) in zip(axes, panels):\n",
        "        ax.plot(epochs, values[metric], label=\"Training\", linewidth=2)\n",
        "        ax.plot(\n",
        "            epochs,\n",
        "            values[f\"val_{metric}\"],\n",
        "            label=\"Validation\",\n",
        "            linewidth=2,\n",
        "        )\n",
        "        ax.axvline(\n",
        "            best_epoch,\n",
        "            color=\"black\",\n",
        "            linestyle=\":\",\n",
        "            linewidth=1.5,\n",
        "            label=\"Best validation loss\" if metric == \"loss\" else None,\n",
        "        )\n",
        "        ax.set(\n",
        "            xlabel=\"Epoch\",\n",
        "            ylabel=ylabel,\n",
        "            title=title,\n",
        "        )\n",
        "        ax.legend()\n",
        "        sns.despine(ax=ax)\n",
        "\n",
        "    fig.suptitle(\"Learning dynamics\", fontsize=15)\n",
        "    fig.tight_layout()\n",
        "    return fig, axes"
      ],
      "id": "history-plot-helper"
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "plot_training_history(history)\n",
        "plt.show()"
      ],
      "id": "cell-plot-training-history"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluating the test set\n",
        "\n",
        "The network predicts standardized MRL. We transform its output back to\n",
        "the original experimental scale before calculating final metrics or\n",
        "plotting the results."
      ],
      "id": "c5634520-911c-4c0a-8363-3cc09c71f4e5"
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {},
      "outputs": [],
      "source": [
        "def predict_mrl(model, encoded_sequences, scaler, batch_size=BATCH_SIZE):\n",
        "    \"\"\"Predict MRL and return values on the original experimental scale.\"\"\"\n",
        "    scaled_predictions = model.predict(\n",
        "        encoded_sequences,\n",
        "        batch_size=batch_size,\n",
        "        verbose=0,\n",
        "    )\n",
        "    return scaler.inverse_transform(scaled_predictions).ravel()\n",
        "\n",
        "\n",
        "def regression_metrics(observed, predicted):\n",
        "    \"\"\"Return the principal regression metrics for one evaluation set.\"\"\"\n",
        "    observed = np.asarray(observed).ravel()\n",
        "    predicted = np.asarray(predicted).ravel()\n",
        "\n",
        "    return {\n",
        "        \"n\": len(observed),\n",
        "        \"Pearson r\": pearsonr(observed, predicted)[0],\n",
        "        \"Spearman rho\": spearmanr(observed, predicted)[0],\n",
        "        \"R2\": r2_score(observed, predicted),\n",
        "        \"RMSE\": np.sqrt(mean_squared_error(observed, predicted)),\n",
        "    }"
      ],
      "id": "evaluation-helpers"
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "test_observed = test_data[\"mrl\"].to_numpy()\n",
        "test_predicted = predict_mrl(model, x_test, target_scaler)\n",
        "\n",
        "test_metrics = regression_metrics(test_observed, test_predicted)\n",
        "pd.DataFrame([test_metrics], index=[\"Random test set\"]).round(4)"
      ],
      "id": "cell-evaluate-test-set"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "With tens of thousands of points, an ordinary scatter plot suffers from\n",
        "severe overplotting. A hexagonal density plot shows where observations\n",
        "concentrate, while marginal histograms show the distribution along each\n",
        "axis."
      ],
      "id": "cdbb5347-2f2b-4d23-a458-f36bb7be8bc4"
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {},
      "outputs": [],
      "source": [
        "def plot_observed_vs_predicted(\n",
        "    observed,\n",
        "    predicted,\n",
        "    title,\n",
        "    limits=MRL_PLOT_LIMITS,\n",
        "):\n",
        "    \"\"\"Plot prediction density, marginal distributions, and calibration lines.\"\"\"\n",
        "    observed = np.asarray(observed).ravel()\n",
        "    predicted = np.asarray(predicted).ravel()\n",
        "    metrics = regression_metrics(observed, predicted)\n",
        "\n",
        "    grid = sns.JointGrid(\n",
        "        x=observed,\n",
        "        y=predicted,\n",
        "        height=7,\n",
        "        ratio=5,\n",
        "        space=0.05,\n",
        "        xlim=limits,\n",
        "        ylim=limits,\n",
        "    )\n",
        "\n",
        "    grid.ax_joint.hexbin(\n",
        "        observed,\n",
        "        predicted,\n",
        "        gridsize=60,\n",
        "        mincnt=1,\n",
        "        bins=\"log\",\n",
        "        cmap=\"viridis\",\n",
        "    )\n",
        "    sns.histplot(x=observed, bins=50, color=\"steelblue\", ax=grid.ax_marg_x)\n",
        "    sns.histplot(y=predicted, bins=50, color=\"darkorange\", ax=grid.ax_marg_y)\n",
        "\n",
        "    line_x = np.asarray(limits)\n",
        "    grid.ax_joint.plot(\n",
        "        line_x,\n",
        "        line_x,\n",
        "        color=\"red\",\n",
        "        linestyle=\"--\",\n",
        "        linewidth=2,\n",
        "        label=\"Perfect calibration\",\n",
        "    )\n",
        "\n",
        "    slope, intercept = np.polyfit(observed, predicted, deg=1)\n",
        "    grid.ax_joint.plot(\n",
        "        line_x,\n",
        "        intercept + slope * line_x,\n",
        "        color=\"white\",\n",
        "        linewidth=2,\n",
        "        label=\"Fitted relationship\",\n",
        "    )\n",
        "\n",
        "    annotation = (\n",
        "        f\"n = {metrics['n']:,}\\n\"\n",
        "        f\"Pearson r = {metrics['Pearson r']:.3f}\\n\"\n",
        "        f\"R² = {metrics['R2']:.3f}\\n\"\n",
        "        f\"fit: y = {intercept:.2f} + {slope:.2f}x\"\n",
        "    )\n",
        "    grid.ax_joint.text(\n",
        "        0.03,\n",
        "        0.97,\n",
        "        annotation,\n",
        "        transform=grid.ax_joint.transAxes,\n",
        "        ha=\"left\",\n",
        "        va=\"top\",\n",
        "        color=\"white\",\n",
        "        bbox={\"boxstyle\": \"round\", \"facecolor\": \"black\", \"alpha\": 0.65},\n",
        "    )\n",
        "    grid.ax_joint.set(\n",
        "        xlabel=\"Experimental MRL\",\n",
        "        ylabel=\"Predicted MRL\",\n",
        "    )\n",
        "    grid.ax_joint.legend(loc=\"lower right\")\n",
        "    grid.fig.suptitle(title, y=1.02, fontsize=15)\n",
        "    return grid"
      ],
      "id": "prediction-plot-helper"
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "plot_observed_vs_predicted(\n",
        "    test_observed,\n",
        "    test_predicted,\n",
        "    title=\"Random test set\",\n",
        ")\n",
        "plt.show()"
      ],
      "id": "cell-plot-test-predictions"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluation on an independent T-cell experiment\n",
        "\n",
        "The preceding test set came from the same experiment and filtered\n",
        "population as the training data. We now ask a more difficult question:\n",
        "does the model transfer to measurements made in a different cellular\n",
        "context?\n",
        "\n",
        "Two biological replicates are available. A sequence measured more deeply\n",
        "in one replicate should receive more influence from that replicate. We\n",
        "therefore combine replicate MRLs using their read counts as weights:\n",
        "\n",
        "$$\n",
        "\\operatorname{MRL}_{\\mathrm{merged}}\n",
        "=\n",
        "\\frac{\n",
        "  \\operatorname{MRL}_1 n_1 + \\operatorname{MRL}_2 n_2\n",
        "}{n_1+n_2},\n",
        "$$\n",
        "\n",
        "where $n_1$ and $n_2$ are the replicate read counts."
      ],
      "id": "9f1b18b8-6b27-4ad7-8167-423456aabdb4"
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {},
      "outputs": [],
      "source": [
        "def merge_replicates(replicate_1, replicate_2):\n",
        "    \"\"\"Outer-join two replicates and calculate read-weighted MRL.\"\"\"\n",
        "    first = replicate_1.rename(\n",
        "        columns={\n",
        "            \"mrl\": \"mrl_1\",\n",
        "            \"read_count\": \"read_count_1\",\n",
        "        }\n",
        "    )\n",
        "    second = replicate_2.rename(\n",
        "        columns={\n",
        "            \"mrl\": \"mrl_2\",\n",
        "            \"read_count\": \"read_count_2\",\n",
        "        }\n",
        "    )\n",
        "\n",
        "    merged = first.merge(second, how=\"outer\", on=\"utr\", validate=\"one_to_one\")\n",
        "\n",
        "    for replicate in (1, 2):\n",
        "        read_column = f\"read_count_{replicate}\"\n",
        "        mrl_column = f\"mrl_{replicate}\"\n",
        "\n",
        "        merged[read_column] = merged[read_column].fillna(0)\n",
        "        inconsistent = (\n",
        "            (merged[read_column] > 0) & merged[mrl_column].isna()\n",
        "        )\n",
        "        if inconsistent.any():\n",
        "            raise ValueError(\n",
        "                f\"Replicate {replicate} has positive reads but missing MRL.\"\n",
        "            )\n",
        "        merged[mrl_column] = merged[mrl_column].fillna(0)\n",
        "\n",
        "    merged[\"read_count\"] = (\n",
        "        merged[\"read_count_1\"] + merged[\"read_count_2\"]\n",
        "    )\n",
        "\n",
        "    positive_reads = merged[\"read_count\"] > 0\n",
        "    merged = merged.loc[positive_reads].copy()\n",
        "    merged[\"mrl\"] = (\n",
        "        merged[\"mrl_1\"] * merged[\"read_count_1\"]\n",
        "        + merged[\"mrl_2\"] * merged[\"read_count_2\"]\n",
        "    ) / merged[\"read_count\"]\n",
        "\n",
        "    return merged.loc[:, [\"utr\", \"mrl\", \"read_count\"]].reset_index(\n",
        "        drop=True\n",
        "    )"
      ],
      "id": "merge-replicates-helper"
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Replicate 1: 349,633 sequences\n",
            "Replicate 2: 352,724 sequences\n",
            "Merged union: 387,794 sequences\n",
            "Retained: 100,000 sequences\n",
            "Read-count boundary: 888"
          ]
        }
      ],
      "source": [
        "tcell_r1_path = download_dataset(\n",
        "    \"GSE232927_processed_defined_end_tcell_r1.csv.gz\"\n",
        ")\n",
        "tcell_r2_path = download_dataset(\n",
        "    \"GSE232927_processed_defined_end_tcell_r2.csv.gz\"\n",
        ")\n",
        "\n",
        "tcell_r1 = load_utr_dataset(tcell_r1_path)\n",
        "tcell_r2 = load_utr_dataset(tcell_r2_path)\n",
        "\n",
        "tcell_merged = merge_replicates(tcell_r1, tcell_r2)\n",
        "tcell_data, tcell_read_cutoff = select_top_by_reads(tcell_merged, 100_000)\n",
        "\n",
        "print(f\"Replicate 1: {len(tcell_r1):,} sequences\")\n",
        "print(f\"Replicate 2: {len(tcell_r2):,} sequences\")\n",
        "print(f\"Merged union: {len(tcell_merged):,} sequences\")\n",
        "print(f\"Retained: {len(tcell_data):,} sequences\")\n",
        "print(f\"Read-count boundary: {tcell_read_cutoff:,.0f}\")"
      ],
      "id": "load-and-merge-tcell"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The independent data must not influence preprocessing. We reuse both the\n",
        "one-hot encoder and the target scaler fitted during the original\n",
        "experiment."
      ],
      "id": "51fa859b-9862-4945-9d4a-51f32d157b70"
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "x_tcell = one_hot_encode(tcell_data[\"utr\"])\n",
        "tcell_observed = tcell_data[\"mrl\"].to_numpy()\n",
        "tcell_predicted = predict_mrl(model, x_tcell, target_scaler)\n",
        "\n",
        "tcell_metrics = regression_metrics(tcell_observed, tcell_predicted)\n",
        "\n",
        "evaluation_summary = pd.DataFrame(\n",
        "    [test_metrics, tcell_metrics],\n",
        "    index=[\"Random test set\", \"Independent T-cell set\"],\n",
        ")\n",
        "evaluation_summary.round(4)"
      ],
      "id": "cell-evaluate-tcell"
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "plot_observed_vs_predicted(\n",
        "    tcell_observed,\n",
        "    tcell_predicted,\n",
        "    title=\"Independent T-cell experiment\",\n",
        ")\n",
        "plt.show()"
      ],
      "id": "cell-plot-tcell-predictions"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Correlation, ranking, and calibration\n",
        "\n",
        "The two red and white reference lines answer different questions:\n",
        "\n",
        "- The dashed red identity line represents perfect numerical predictions.\n",
        "- The fitted white line represents the linear relationship actually\n",
        "  present in the data.\n",
        "\n",
        "Pearson correlation measures the strength of a linear relationship. It\n",
        "is unchanged when every prediction receives the same offset and positive\n",
        "scaling. Predictive $R^2$ is stricter:\n",
        "\n",
        "$$\n",
        "R^2\n",
        "=\n",
        "1-\n",
        "\\frac{\\sum_i(y_i-\\hat{y}_i)^2}\n",
        "     {\\sum_i(y_i-\\bar{y})^2}.\n",
        "$$\n",
        "\n",
        "It therefore penalizes offsets and scaling errors. The independent data\n",
        "can exhibit a high Pearson correlation while obtaining a lower $R^2$\n",
        "when the white fitted line differs from the red identity line.\n",
        "\n",
        "Spearman correlation, included in the results table, measures monotonic\n",
        "rank-order agreement more directly. This is useful for our eventual\n",
        "design application: a genetic algorithm primarily needs the model to\n",
        "rank promising candidate sequences effectively.\n",
        "\n",
        "> **$R^2$ is not generally the same as $r^2$**\n",
        ">\n",
        "> Some literature reports the square of Pearson correlation and labels\n",
        "> it “R-squared.” Here, `R2` always denotes the predictive coefficient\n",
        "> of determination calculated from errors relative to the identity line.\n",
        "\n",
        "## Saving the reusable predictor\n",
        "\n",
        "The best model was saved during training. We also save the two numbers\n",
        "required to transform predictions back to the original MRL scale. These\n",
        "artifacts will allow a later genetic-algorithm notebook to score new\n",
        "strings without loading the experimental training data."
      ],
      "id": "7dfbe676-8df2-4e5d-a1d3-9a7f47c673e9"
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saved model: models/mrl_cnn_random_split.keras\n",
            "Saved target transformation: models/mrl_target_scaler.npz"
          ]
        }
      ],
      "source": [
        "scaler_path = MODEL_DIR / \"mrl_target_scaler.npz\"\n",
        "np.savez(\n",
        "    scaler_path,\n",
        "    mean=target_scaler.mean_.astype(np.float32),\n",
        "    scale=target_scaler.scale_.astype(np.float32),\n",
        ")\n",
        "\n",
        "print(\"Saved model:\", random_split_model_path)\n",
        "print(\"Saved target transformation:\", scaler_path)"
      ],
      "id": "save-artifacts"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The complete prediction interface is now: validate strings, one-hot\n",
        "encode them, apply the model, and undo target standardization."
      ],
      "id": "f53886fd-0fc2-419c-8d56-d1ec0c7cd234"
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "example_sequences = pd.Series(\n",
        "    [\n",
        "        \"A\" * SEQUENCE_LENGTH,\n",
        "        \"C\" * SEQUENCE_LENGTH,\n",
        "        (\"ACGT\" * 13)[:SEQUENCE_LENGTH],\n",
        "    ],\n",
        "    name=\"utr\",\n",
        ")\n",
        "\n",
        "example_predictions = predict_mrl(\n",
        "    model,\n",
        "    one_hot_encode(example_sequences),\n",
        "    target_scaler,\n",
        ")\n",
        "\n",
        "pd.DataFrame(\n",
        "    {\n",
        "        \"utr\": example_sequences,\n",
        "        \"predicted_mrl\": example_predictions,\n",
        "    }\n",
        ").round({\"predicted_mrl\": 3})"
      ],
      "id": "cell-predict-literal-sequences"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Experiment 2: partitioning by read count\n",
        "\n",
        "Our primary experiment uses a random test set because it asks how well\n",
        "the model generalizes to new sequences from the same filtered\n",
        "population. Sample et al. asked a different question and reserved the\n",
        "most deeply measured sequences for evaluation.\n",
        "\n",
        "To isolate the effect of this design choice, we keep exactly the same\n",
        "partition sizes as before:\n",
        "\n",
        "- the 27,500 highest-read sequences form the test set;\n",
        "- 220,000 of the remaining sequences form the training set; and\n",
        "- the final 27,500 sequences form the validation set.\n",
        "\n",
        "Only the rule used to select the test set changes."
      ],
      "id": "b4b23c02-9e3c-4f99-aee6-233e4d9c3f9a"
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "def read_count_partition(data, seed=SEED):\n",
        "    \"\"\"Reserve the highest-read rows for testing, then split the remainder.\"\"\"\n",
        "    if len(data) != N_RETAINED:\n",
        "        raise ValueError(f\"Expected exactly {N_RETAINED:,} retained rows.\")\n",
        "\n",
        "    ranked = data.sort_values(\n",
        "        \"read_count\", ascending=False, kind=\"stable\"\n",
        "    ).reset_index(drop=True)\n",
        "\n",
        "    test = ranked.iloc[:N_TEST].copy()\n",
        "    remaining = ranked.iloc[N_TEST:].sample(\n",
        "        frac=1,\n",
        "        random_state=seed,\n",
        "    ).reset_index(drop=True)\n",
        "\n",
        "    train = remaining.iloc[:N_TRAIN].copy()\n",
        "    validation = remaining.iloc[N_TRAIN:].copy()\n",
        "\n",
        "    if (len(train), len(validation), len(test)) != (\n",
        "        N_TRAIN,\n",
        "        N_VALIDATION,\n",
        "        N_TEST,\n",
        "    ):\n",
        "        raise AssertionError(\"Unexpected partition sizes.\")\n",
        "\n",
        "    return train, validation, test\n",
        "\n",
        "\n",
        "read_train_data, read_validation_data, read_test_data = (\n",
        "    read_count_partition(retained_data)\n",
        ")\n",
        "\n",
        "read_partition_summary = summarize_partitions(\n",
        "    {\n",
        "        \"Training\": read_train_data,\n",
        "        \"Validation\": read_validation_data,\n",
        "        \"Test (highest reads)\": read_test_data,\n",
        "    }\n",
        ")\n",
        "read_partition_summary.round(3)"
      ],
      "id": "cell-read-count-partition-helper"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The distributions below show that choosing the test set by read count\n",
        "changes both measurement coverage and, potentially, the target\n",
        "distribution."
      ],
      "id": "b89ec2e9-e928-4848-9b2e-d20a1b4017ad"
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "partition_plot_data = pd.concat(\n",
        "    [\n",
        "        read_train_data[[\"mrl\", \"read_count\"]].assign(\n",
        "            partition=\"Training\"\n",
        "        ),\n",
        "        read_validation_data[[\"mrl\", \"read_count\"]].assign(\n",
        "            partition=\"Validation\"\n",
        "        ),\n",
        "        read_test_data[[\"mrl\", \"read_count\"]].assign(\n",
        "            partition=\"Test (highest reads)\"\n",
        "        ),\n",
        "    ],\n",
        "    ignore_index=True,\n",
        ")\n",
        "partition_plot_data[\"log_read_count\"] = np.log10(\n",
        "    partition_plot_data[\"read_count\"] + 1\n",
        ")\n",
        "\n",
        "fig, axes = plt.subplots(1, 2, figsize=(13, 4.2))\n",
        "sns.histplot(\n",
        "    data=partition_plot_data,\n",
        "    x=\"log_read_count\",\n",
        "    hue=\"partition\",\n",
        "    bins=60,\n",
        "    stat=\"density\",\n",
        "    common_norm=False,\n",
        "    element=\"step\",\n",
        "    fill=False,\n",
        "    ax=axes[0],\n",
        ")\n",
        "axes[0].set(\n",
        "    xlabel=r\"$\\log_{10}(\\mathrm{read\\ count}+1)$\",\n",
        "    ylabel=\"Density\",\n",
        "    title=\"Read-count distribution\",\n",
        ")\n",
        "\n",
        "sns.histplot(\n",
        "    data=partition_plot_data,\n",
        "    x=\"mrl\",\n",
        "    hue=\"partition\",\n",
        "    bins=60,\n",
        "    stat=\"density\",\n",
        "    common_norm=False,\n",
        "    element=\"step\",\n",
        "    fill=False,\n",
        "    ax=axes[1],\n",
        "    legend=False,\n",
        ")\n",
        "axes[1].set(\n",
        "    xlabel=\"Experimental MRL\",\n",
        "    ylabel=\"Density\",\n",
        "    title=\"Target distribution\",\n",
        ")\n",
        "\n",
        "for ax in axes:\n",
        "    sns.despine(ax=ax)\n",
        "fig.tight_layout()\n",
        "plt.show()"
      ],
      "id": "cell-compare-read-partitions"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now retrain the same architecture from scratch. The target\n",
        "transformation is again fitted using training data only."
      ],
      "id": "d4fadda1-b446-44f4-9cc7-c7c270fe59cc"
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training features: (220000, 50, 4) float32\n",
            "Validation features: (27500, 50, 4) float32\n",
            "Test features: (27500, 50, 4) float32"
          ]
        }
      ],
      "source": [
        "read_target_scaler = StandardScaler()\n",
        "\n",
        "read_y_train = read_target_scaler.fit_transform(\n",
        "    read_train_data[[\"mrl\"]]\n",
        ").astype(np.float32)\n",
        "read_y_validation = read_target_scaler.transform(\n",
        "    read_validation_data[[\"mrl\"]]\n",
        ").astype(np.float32)\n",
        "\n",
        "read_x_train = one_hot_encode(read_train_data[\"utr\"])\n",
        "read_x_validation = one_hot_encode(read_validation_data[\"utr\"])\n",
        "read_x_test = one_hot_encode(read_test_data[\"utr\"])\n",
        "\n",
        "print(\"Training features:\", read_x_train.shape, read_x_train.dtype)\n",
        "print(\"Validation features:\", read_x_validation.shape, read_x_validation.dtype)\n",
        "print(\"Test features:\", read_x_test.shape, read_x_test.dtype)"
      ],
      "id": "prepare-read-count-arrays"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The seed, architecture, optimizer, batch size, early-stopping rule, and\n",
        "maximum number of epochs are identical to Experiment 1. Consequently,\n",
        "the data partitioning strategy is the only intentional experimental\n",
        "difference.\n",
        "\n",
        "> **Expected runtime**\n",
        ">\n",
        "> This cell performs a second complete model fit and therefore\n",
        "> approximately doubles the total training time of the notebook."
      ],
      "id": "295cceee-fffc-42f7-bd30-c88858e03484"
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/50\n",
            "430/430 - 147s - 343ms/step - loss: 0.4517 - pearson_r: 0.7438 - r2: 0.5483 - val_loss: 1.7051 - val_pearson_r: 0.7071 - val_r2: -7.0563e-01\n",
            "Epoch 2/50\n",
            "430/430 - 149s - 347ms/step - loss: 0.2059 - pearson_r: 0.8911 - r2: 0.7941 - val_loss: 0.1677 - val_pearson_r: 0.9242 - val_r2: 0.8322\n",
            "Epoch 3/50\n",
            "430/430 - 149s - 346ms/step - loss: 0.1493 - pearson_r: 0.9223 - r2: 0.8507 - val_loss: 0.1373 - val_pearson_r: 0.9363 - val_r2: 0.8626\n",
            "Epoch 4/50\n",
            "430/430 - 148s - 344ms/step - loss: 0.1333 - pearson_r: 0.9310 - r2: 0.8667 - val_loss: 0.1372 - val_pearson_r: 0.9396 - val_r2: 0.8628\n",
            "Epoch 5/50\n",
            "430/430 - 147s - 343ms/step - loss: 0.1255 - pearson_r: 0.9352 - r2: 0.8745 - val_loss: 0.1241 - val_pearson_r: 0.9418 - val_r2: 0.8759\n",
            "Epoch 6/50\n",
            "430/430 - 148s - 344ms/step - loss: 0.1201 - pearson_r: 0.9380 - r2: 0.8799 - val_loss: 0.1189 - val_pearson_r: 0.9437 - val_r2: 0.8811\n",
            "Epoch 7/50\n",
            "430/430 - 149s - 346ms/step - loss: 0.1162 - pearson_r: 0.9401 - r2: 0.8838 - val_loss: 0.1162 - val_pearson_r: 0.9440 - val_r2: 0.8838\n",
            "Epoch 8/50\n",
            "430/430 - 140s - 326ms/step - loss: 0.1134 - pearson_r: 0.9416 - r2: 0.8866 - val_loss: 0.1140 - val_pearson_r: 0.9455 - val_r2: 0.8860\n",
            "Epoch 9/50\n",
            "430/430 - 141s - 328ms/step - loss: 0.1107 - pearson_r: 0.9430 - r2: 0.8893 - val_loss: 0.1141 - val_pearson_r: 0.9458 - val_r2: 0.8859\n",
            "Epoch 10/50\n",
            "430/430 - 138s - 320ms/step - loss: 0.1089 - pearson_r: 0.9440 - r2: 0.8911 - val_loss: 0.1114 - val_pearson_r: 0.9469 - val_r2: 0.8886\n",
            "Epoch 11/50\n",
            "430/430 - 137s - 319ms/step - loss: 0.1066 - pearson_r: 0.9452 - r2: 0.8934 - val_loss: 0.1058 - val_pearson_r: 0.9476 - val_r2: 0.8942\n",
            "Epoch 12/50\n",
            "430/430 - 137s - 318ms/step - loss: 0.1052 - pearson_r: 0.9459 - r2: 0.8948 - val_loss: 0.1094 - val_pearson_r: 0.9478 - val_r2: 0.8906\n",
            "Epoch 13/50\n",
            "430/430 - 136s - 317ms/step - loss: 0.1034 - pearson_r: 0.9469 - r2: 0.8966 - val_loss: 0.1084 - val_pearson_r: 0.9474 - val_r2: 0.8915\n",
            "Epoch 14/50\n",
            "430/430 - 137s - 319ms/step - loss: 0.1022 - pearson_r: 0.9475 - r2: 0.8978 - val_loss: 0.1028 - val_pearson_r: 0.9484 - val_r2: 0.8972\n",
            "Epoch 15/50\n",
            "430/430 - 137s - 319ms/step - loss: 0.1011 - pearson_r: 0.9481 - r2: 0.8989 - val_loss: 0.1044 - val_pearson_r: 0.9485 - val_r2: 0.8956\n",
            "Epoch 16/50\n",
            "430/430 - 135s - 315ms/step - loss: 0.0998 - pearson_r: 0.9488 - r2: 0.9002 - val_loss: 0.1061 - val_pearson_r: 0.9488 - val_r2: 0.8938\n",
            "Epoch 17/50\n",
            "430/430 - 138s - 320ms/step - loss: 0.0983 - pearson_r: 0.9496 - r2: 0.9017 - val_loss: 0.1015 - val_pearson_r: 0.9488 - val_r2: 0.8984\n",
            "Epoch 18/50\n",
            "430/430 - 134s - 311ms/step - loss: 0.0971 - pearson_r: 0.9502 - r2: 0.9029 - val_loss: 0.1037 - val_pearson_r: 0.9488 - val_r2: 0.8963\n",
            "Epoch 19/50\n",
            "430/430 - 132s - 308ms/step - loss: 0.0954 - pearson_r: 0.9511 - r2: 0.9046 - val_loss: 0.1018 - val_pearson_r: 0.9490 - val_r2: 0.8982\n",
            "Epoch 20/50\n",
            "430/430 - 132s - 308ms/step - loss: 0.0943 - pearson_r: 0.9517 - r2: 0.9057 - val_loss: 0.1016 - val_pearson_r: 0.9491 - val_r2: 0.8984\n",
            "Epoch 21/50\n",
            "430/430 - 132s - 307ms/step - loss: 0.0933 - pearson_r: 0.9522 - r2: 0.9067 - val_loss: 0.1031 - val_pearson_r: 0.9492 - val_r2: 0.8969\n",
            "Epoch 22/50\n",
            "430/430 - 133s - 310ms/step - loss: 0.0922 - pearson_r: 0.9528 - r2: 0.9078 - val_loss: 0.1020 - val_pearson_r: 0.9485 - val_r2: 0.8980\n",
            "Epoch 22: early stopping\n",
            "Restoring model weights from the end of the best epoch: 17."
          ]
        }
      ],
      "source": [
        "# Reset the random state before constructing the second model.\n",
        "keras.utils.set_random_seed(SEED)\n",
        "read_count_model = compile_mrl_cnn(build_mrl_cnn())\n",
        "read_count_model_path = MODEL_DIR / \"mrl_cnn_read_count_split.keras\"\n",
        "\n",
        "read_history = read_count_model.fit(\n",
        "    read_x_train,\n",
        "    read_y_train,\n",
        "    validation_data=(read_x_validation, read_y_validation),\n",
        "    epochs=MAX_EPOCHS,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    shuffle=True,\n",
        "    callbacks=[\n",
        "        callbacks.EarlyStopping(\n",
        "            monitor=\"val_loss\",\n",
        "            mode=\"min\",\n",
        "            patience=PATIENCE,\n",
        "            restore_best_weights=True,\n",
        "            verbose=1,\n",
        "        ),\n",
        "        callbacks.ModelCheckpoint(\n",
        "            filepath=str(read_count_model_path),\n",
        "            monitor=\"val_loss\",\n",
        "            mode=\"min\",\n",
        "            save_best_only=True,\n",
        "        ),\n",
        "    ],\n",
        "    verbose=2,\n",
        ")"
      ],
      "id": "train-read-count-model"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We first inspect the learning dynamics and verify that early stopping\n",
        "selected a model with stable validation performance."
      ],
      "id": "1d0605b3-dd8a-40f4-bc85-fb21e30e62e0"
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "plot_training_history(read_history)\n",
        "plt.show()"
      ],
      "id": "cell-plot-read-count-training-history"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The final evaluation uses the 27,500 sequences with the highest read\n",
        "counts. The comparison table places its results beside those from the\n",
        "random test set."
      ],
      "id": "97702ebc-5893-4b42-81f3-2eadf29bc8ba"
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "read_test_observed = read_test_data[\"mrl\"].to_numpy()\n",
        "read_test_predicted = predict_mrl(\n",
        "    read_count_model,\n",
        "    read_x_test,\n",
        "    read_target_scaler,\n",
        ")\n",
        "\n",
        "split_comparison = pd.DataFrame(\n",
        "    [\n",
        "        test_metrics,\n",
        "        regression_metrics(read_test_observed, read_test_predicted),\n",
        "    ],\n",
        "    index=[\"Random test selection\", \"Highest-read test selection\"],\n",
        ")\n",
        "split_comparison.round(4)"
      ],
      "id": "cell-evaluate-read-count-test-set"
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "plot_observed_vs_predicted(\n",
        "    read_test_observed,\n",
        "    read_test_predicted,\n",
        "    title=\"Test set selected by read count\",\n",
        ")\n",
        "plt.show()"
      ],
      "id": "cell-plot-read-count-test-predictions"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluation on the independent T-cell experiment\n",
        "\n",
        "We now apply the second model to the same 100,000 merged T-cell\n",
        "sequences used to evaluate Experiment 1. Reusing an identical external\n",
        "evaluation set lets us compare the effect of the training-partition\n",
        "strategy directly. Predictions from the second model must be returned to\n",
        "the experimental MRL scale using the target transformation fitted on its\n",
        "own training partition."
      ],
      "id": "09b70049-aa8d-4ca4-b4e1-1090fef7254e"
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "\n",
              "</div>"
            ]
          }
        }
      ],
      "source": [
        "read_tcell_predicted = predict_mrl(\n",
        "    read_count_model,\n",
        "    x_tcell,\n",
        "    read_target_scaler,\n",
        ")\n",
        "read_tcell_metrics = regression_metrics(\n",
        "    tcell_observed,\n",
        "    read_tcell_predicted,\n",
        ")\n",
        "\n",
        "cross_cell_comparison = pd.DataFrame(\n",
        "    [\n",
        "        tcell_metrics,\n",
        "        read_tcell_metrics,\n",
        "    ],\n",
        "    index=[\n",
        "        \"Experiment 1: random partition\",\n",
        "        \"Experiment 2: read-count partition\",\n",
        "    ],\n",
        ")\n",
        "cross_cell_comparison.round(4)"
      ],
      "id": "cell-evaluate-read-count-model-on-tcell"
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [
        "plot_observed_vs_predicted(\n",
        "    tcell_observed,\n",
        "    read_tcell_predicted,\n",
        "    title=\"Independent T-cell experiment — read-count-partitioned model\",\n",
        ")\n",
        "plt.show()"
      ],
      "id": "cell-plot-read-count-model-tcell-predictions"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we save the target transformation associated with the second\n",
        "model. The best model itself was saved by `ModelCheckpoint` during\n",
        "training."
      ],
      "id": "7b536f82-46d1-40fb-aa35-e21c9a89ec58"
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saved model: models/mrl_cnn_read_count_split.keras\n",
            "Saved target transformation: models/mrl_target_scaler_read_count_split.npz"
          ]
        }
      ],
      "source": [
        "read_scaler_path = MODEL_DIR / \"mrl_target_scaler_read_count_split.npz\"\n",
        "np.savez(\n",
        "    read_scaler_path,\n",
        "    mean=read_target_scaler.mean_.astype(np.float32),\n",
        "    scale=read_target_scaler.scale_.astype(np.float32),\n",
        ")\n",
        "\n",
        "print(\"Saved model:\", read_count_model_path)\n",
        "print(\"Saved target transformation:\", read_scaler_path)"
      ],
      "id": "save-read-count-artifacts"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The two experiments estimate different forms of generalization. Their\n",
        "scores should therefore be interpreted alongside the displayed target\n",
        "distributions. In particular, $R^2$ depends on the variance of the\n",
        "observed values in its test set, so a change in $R^2$ cannot\n",
        "automatically be attributed solely to a change in model quality.\n",
        "\n",
        "# References\n",
        "\n",
        "- Sample, P. J., Wang, B., Reid, D. W., Presnyak, V., McFadyen, I. J.,\n",
        "  Morris, D. R., & Seelig, G. (2019). Human 5′ UTR design and variant\n",
        "  effect prediction from a massively parallel translation assay. *Nature\n",
        "  Biotechnology*, *37*(7), 803–809.\n",
        "  <https://doi.org/10.1038/s41587-019-0164-5>\n",
        "- Castillo-Hair, S., Fedak, S., Wang, B., Linder, J., Havens, K., Certo,\n",
        "  M., & Seelig, G. (2024). Optimizing 5′ UTRs for mRNA-delivered gene\n",
        "  editing using deep learning. *Nature Communications*, *15*(1), 5284.\n",
        "  <https://doi.org/10.1038/s41467-024-49508-2>\n",
        "- Tang, X., Huo, M., Chen, Y., et al. (2024). A novel deep generative\n",
        "  model for mRNA vaccine development: Designing 5′ UTRs with\n",
        "  N1-methyl-pseudouridine modification. *Acta Pharmaceutica Sinica B*,\n",
        "  *14*(4), 1814–1826. <https://doi.org/10.1016/j.apsb.2023.11.003>"
      ],
      "id": "2c81c488-1f49-47cb-9236-2482cc585b13"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "path": "/Users/turcotte/.virtualenvs/csi4106/share/jupyter/kernels/python3"
    }
  }
}