notebooks/2024-coding-noncoding-prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "648a1842-3b73-4d49-bd59-dea29a8f62ad",
   "metadata": {},
   "source": [
    "## Comparing ESM-based models and RNASamba models for predicting coding and noncoding transcripts\n",
    "__Keith Cheveralls__<br>\n",
    "__March 2024__<br>\n",
    "\n",
    "This notebook documents the visualizations that were used to compare the performance of ESM-based models and RNASamba models trained to predict whether transcripts are coding or noncoding. This was motivated by developing an approach that used ESM embeddings to identifying sORFs for the [peptigate pipeline](https://github.com/Arcadia-Science/peptigate).\n",
    "\n",
    "The predictions from ESM-based models and RNASamba models on which this notebook depends were generated outside of this notebook. Predictions from ESM-based models were generated using the commands namespaced under the `plmutils orf-classification` CLI. Predictions from RNASamba models were generated using the script found in the `/scripts/rnasamba` subdirectory of this repo. The CLI commands that were used are briefly documented in the sections below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1299118b-1c19-4bc6-898f-6a6a0dad81e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import pathlib\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "from plmutils.models import calc_metrics\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a6da10-8436-46cf-b67f-755648570f18",
   "metadata": {},
   "source": [
    "### Dataset metadata\n",
    "The metadata associated with the 16 species used for these comparisons is included below for completeness. Note that the plots in this notebook label species using the `species_id` defined in this metadata (rather than the full species name)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f52ad9cb-1306-451b-8f51-8f5c449ba43a",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_csv_content = \"\"\"\n",
    "species_id\tspecies_common_name\troot_url\tgenome_name\tcdna_endpoint\tncrna_endpoint\tgenome_abbreviation\n",
    "hsap\thuman\thttps://ftp.ensembl.org/pub/release-111/fasta/homo_sapiens/\tHomo_sapiens.GRCh38\tcdna/Homo_sapiens.GRCh38.cdna.all.fa.gz\tncrna/Homo_sapiens.GRCh38.ncrna.fa.gz\tGRCh38\n",
    "scer\tyeast\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/fungi/release-58/fasta/saccharomyces_cerevisiae/\tSaccharomyces_cerevisiae.R64-1-1\tcdna/Saccharomyces_cerevisiae.R64-1-1.cdna.all.fa.gz\tncrna/Saccharomyces_cerevisiae.R64-1-1.ncrna.fa.gz\tR64-1-1\n",
    "cele\tworm\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/metazoa/release-58/fasta/caenorhabditis_elegans/\tCaenorhabditis_elegans.WBcel235\tcdna/Caenorhabditis_elegans.WBcel235.cdna.all.fa.gz\tncrna/Caenorhabditis_elegans.WBcel235.ncrna.fa.gz\tWBcel235\n",
    "atha\tarabadopsis\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/plants/release-58/fasta/arabidopsis_thaliana/\tArabidopsis_thaliana.TAIR10\tcdna/Arabidopsis_thaliana.TAIR10.cdna.all.fa.gz\tncrna/Arabidopsis_thaliana.TAIR10.ncrna.fa.gz\tTAIR10\n",
    "dmel\tdrosophila\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/metazoa/release-58/fasta/drosophila_melanogaster/\tDrosophila_melanogaster.BDGP6.46\tcdna/Drosophila_melanogaster.BDGP6.46.cdna.all.fa.gz\tncrna/Drosophila_melanogaster.BDGP6.46.ncrna.fa.gz\tBDGP6.46\n",
    "ddis\tdictyostelium_discoideum\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/protists/release-58/fasta/dictyostelium_discoideum/\tDictyostelium_discoideum.dicty_2.7\tcdna/Dictyostelium_discoideum.dicty_2.7.cdna.all.fa.gz\tncrna/Dictyostelium_discoideum.dicty_2.7.ncrna.fa.gz\tdicty_2.7\n",
    "mmus\tmouse\thttps://ftp.ensembl.org/pub/release-111/fasta/mus_musculus/\tMus_musculus.GRCm39\tcdna/Mus_musculus.GRCm39.cdna.all.fa.gz\tncrna/Mus_musculus.GRCm39.ncrna.fa.gz\tGRCm39\n",
    "drer\tzebrafish\thttps://ftp.ensembl.org/pub/release-111/fasta/danio_rerio/\tDanio_rerio.GRCz11\tcdna/Danio_rerio.GRCz11.cdna.all.fa.gz\tncrna/Danio_rerio.GRCz11.ncrna.fa.gz\tGRCz11\n",
    "ggal\tchicken\thttps://ftp.ensembl.org/pub/release-111/fasta/gallus_gallus/\tGallus_gallus.bGalGal1.mat.broiler.GRCg7b\tcdna/Gallus_gallus.bGalGal1.mat.broiler.GRCg7b.cdna.all.fa.gz\tncrna/Gallus_gallus.bGalGal1.mat.broiler.GRCg7b.ncrna.fa.gz\tbGalGal1.mat.broiler.GRCg7b\n",
    "oind\trice\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/plants/release-58/fasta/oryza_indica/\tOryza_indica.ASM465v1\tcdna/Oryza_indica.ASM465v1.cdna.all.fa.gz\tncrna/Oryza_indica.ASM465v1.ncrna.fa.gz\tASM465v1\n",
    "zmay\tmaize\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/plants/release-58/fasta/zea_mays/\tZea_mays.Zm-B73-REFERENCE-NAM-5.0\tcdna/Zea_mays.Zm-B73-REFERENCE-NAM-5.0.cdna.all.fa.gz\tncrna/Zea_mays.Zm-B73-REFERENCE-NAM-5.0.ncrna.fa.gz\tZm-B73-REFERENCE-NAM-5.0\n",
    "xtro\tfrog\thttps://ftp.ensembl.org/pub/release-111/fasta/xenopus_tropicalis/\tXenopus_tropicalis.UCB_Xtro_10.0\tcdna/Xenopus_tropicalis.UCB_Xtro_10.0.cdna.all.fa.gz\tncrna/Xenopus_tropicalis.UCB_Xtro_10.0.ncrna.fa.gz\tUCB_Xtro_10.0\n",
    "rnor\trat\thttps://ftp.ensembl.org/pub/release-111/fasta/rattus_norvegicus/\tRattus_norvegicus.mRatBN7.2\tcdna/Rattus_norvegicus.mRatBN7.2.cdna.all.fa.gz\tncrna/Rattus_norvegicus.mRatBN7.2.ncrna.fa.gz\tmRatBN7\n",
    "amel\thoneybee\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/metazoa/release-58/fasta/apis_mellifera/\tApis_mellifera.Amel_HAv3.1\tcdna/Apis_mellifera.Amel_HAv3.1.cdna.all.fa.gz\tncrna/Apis_mellifera.Amel_HAv3.1.ncrna.fa.gz\tAmel_HAv3.1\n",
    "spom\tfission_yeast\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/fungi/release-58/fasta/schizosaccharomyces_pombe/\tSchizosaccharomyces_pombe.ASM294v2\tcdna/Schizosaccharomyces_pombe.ASM294v2.cdna.all.fa.gz\tncrna/Schizosaccharomyces_pombe.ASM294v2.ncrna.fa.gz\tASM294v2\n",
    "tthe\ttetrahymena\thttps://ftp.ensemblgenomes.ebi.ac.uk/pub/protists/release-58/fasta/tetrahymena_thermophila/\tTetrahymena_thermophila.JCVI-TTA1-2.2\tcdna/Tetrahymena_thermophila.JCVI-TTA1-2.2.cdna.all.fa.gz\tncrna/Tetrahymena_thermophila.JCVI-TTA1-2.2.ncrna.fa.gz\tJCVI-TTA1-2.2\n",
    "\"\"\"  # noqa: E501\n",
    "\n",
    "metadata = pd.read_csv(io.StringIO(metadata_csv_content), sep=\"\\t\")\n",
    "metadata.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dba1886-9654-4508-bb45-4627d1bbcbc3",
   "metadata": {},
   "source": [
    "### Heatmap plotting functions\n",
    "These are functions used later in the notebook to generate heatmap visualizations of the matrices of model performance metrics for all pairs of training and test species. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9393913d-3b59-4308-bc79-6c9e5abf1107",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_heatmap(df, column=\"accuracy\", model_name=\"unknown\", ax=None, **heatmap_kwargs):\n",
    "    \"\"\"\n",
    "    Plot the values in the given column as a square heatmap of training vs test species\n",
    "    (with training species on the x-axis and test species on the y-axis).\n",
    "\n",
    "    Note: \"training species\" is the species used to train the model and \"test species\"\n",
    "    is the species used to test each trained model.\n",
    "    \"\"\"\n",
    "    df = df.pivot(index=\"test_species_id\", columns=\"training_species_id\", values=column)\n",
    "\n",
    "    if ax is None:\n",
    "        plt.figure(figsize=(8, 6))\n",
    "        ax = plt.gca()\n",
    "\n",
    "    sns.heatmap(\n",
    "        df,\n",
    "        cmap=\"coolwarm\",\n",
    "        annot=True,\n",
    "        annot_kws={\"size\": 6},\n",
    "        fmt=\".1f\",\n",
    "        square=True,\n",
    "        ax=ax,\n",
    "        **heatmap_kwargs,\n",
    "    )\n",
    "\n",
    "    name = column.replace(\"_\", \" \")\n",
    "    if name.lower() == \"mcc\":\n",
    "        name = name.upper()\n",
    "    else:\n",
    "        name = name[0].upper() + name[1:]\n",
    "\n",
    "    ax.set_xlabel(\"Training species\")\n",
    "    ax.set_ylabel(\"Test species\")\n",
    "    ax.set_title(f\"{name} | {model_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2501a48f-026e-470a-aecc-af14b388df11",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_heatmaps(df_left, df_right, column, model_names):\n",
    "    \"\"\"\n",
    "    Plot a row of three heatmaps: one for the left dataframe, one for the right dataframe,\n",
    "    and the third (the rightmost) for the difference between the two (right minus left).\n",
    "    \"\"\"\n",
    "    fig, axs = plt.subplots(nrows=1, ncols=3, figsize=(18, 5))\n",
    "\n",
    "    df_merged = pd.merge(df_left, df_right, on=(\"training_species_id\", \"test_species_id\"))\n",
    "    df_merged[column] = df_merged[f\"{column}_y\"] - df_merged[f\"{column}_x\"]\n",
    "\n",
    "    plot_heatmap(df_left, column=column, model_name=model_names[0], ax=axs[0])\n",
    "    plot_heatmap(df_right, column=column, model_name=model_names[1], ax=axs[1])\n",
    "    plot_heatmap(df_merged, column=column, model_name=\"difference\", ax=axs[2], vmin=-1, vmax=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1220a1c-bebd-406a-8a3b-1c237f0d4e86",
   "metadata": {},
   "source": [
    "### ESM-based model predictions\n",
    "\n",
    "These predictions were generated using the `plmutils orf-prediction` CLI. \n",
    "\n",
    "First, download the Ensembl datasets listed in the user-provided metadata CSV file (see above for the file used with this notebook):\n",
    "```\n",
    "plmutils orf-prediction download-data \\\n",
    "    output/data/ensembl-dataset-metadata.tsv \\\n",
    "    output/data/\n",
    "```\n",
    "\n",
    "Next, construct deduplicated sets of coding and noncoding transcripts. Deduplication is achieved by clustering transcripts by sequence identity and retaining only one representative sequence from each cluster.\n",
    "```\n",
    "plmutils orf-prediction construct-data \\\n",
    "    output/data/ensembl-dataset-metadata.tsv \\\n",
    "    output/data/ \\\n",
    "    --subsample-factor 1\n",
    "```\n",
    "\n",
    "Next, find putative ORFs from coding and noncoding transcripts, retain only the longest putative ORF from each transcript, and generate the embedding of the protein sequence for which it codes:\n",
    "```\n",
    "plmutils orf-prediction translate-and-embed \\\n",
    "    output/data/processed/final/coding-dedup-ssx1/transcripts\n",
    "\n",
    "plmutils orf-prediction translate-and-embed \\\n",
    "    output/data/processed/final/noncoding-dedup-ssx1/transcripts   \n",
    "```\n",
    "\n",
    "Finally, train models using these embeddings to predict whether a given ORF orginated from a coding or noncoding transcript. Separate models are trained on, and used to make predictions for, each species. This results in a matrix of model performance metrics for all pairs of species (one used to train the model, the other to evaluate it). The `--output-dirpath` in the command below corresponds to the directories passed to the `calc_metrics_from_smallesm_results` function defined below. (This command was run manually with and without `--max-length 100` to train models on all ORFs and only sORFs, respectively).\n",
    "```\n",
    "plmutils orf-prediction train-and-evaluate \\\n",
    "    --coding-dirpath output/data/processed/final/coding-dedup-ssx1/embeddings/esm2_t6_8M_UR50D \\\n",
    "    --noncoding-dirpath output/data/processed/final/noncoding-dedup-ssx1/embeddings/esm2_t6_8M_UR50D \\\n",
    "    --output-dirpath output/data/esm-model-results-ssx1-all\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66bd922d-0d56-43d9-92c2-14fa846a2521",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_metrics_from_smallesm_results(results_dirpath, max_length=None):\n",
    "    \"\"\"\n",
    "    Calculate classification metrics from ESM-based model results.\n",
    "    \"\"\"\n",
    "    all_metrics = []\n",
    "    prediction_filepaths = pathlib.Path(results_dirpath).glob(\"*.csv\")\n",
    "    for prediction_filepath in prediction_filepaths:\n",
    "        df = pd.read_csv(prediction_filepath)\n",
    "\n",
    "        if max_length is not None:\n",
    "            df = df.loc[df.sequence_length < max_length]\n",
    "\n",
    "        metrics = calc_metrics(\n",
    "            y_true=(df.true_label == \"coding\"),\n",
    "            y_pred_proba=df.predicted_probability.values,\n",
    "        )\n",
    "        metrics[\"training_species_id\"] = df.iloc[0].training_species_id\n",
    "        metrics[\"test_species_id\"] = df.iloc[0].testing_species_id\n",
    "        metrics[\"num_coding\"] = (df.true_label == \"coding\").sum()\n",
    "        metrics[\"num_noncoding\"] = (df.true_label != \"coding\").sum()\n",
    "\n",
    "        all_metrics.append(metrics)\n",
    "    df = pd.DataFrame(all_metrics)\n",
    "    df[\"true_negative_rate\"] = df.num_true_negative / df.num_noncoding\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "288daad0-42e1-4d85-af2c-3bab2cf97dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_esm_trained_all_eval_all = calc_metrics_from_smallesm_results(\n",
    "    \"../output/results/2024-03-01-esm-model-results-ssx1-all/\",\n",
    "    max_length=None,\n",
    ")\n",
    "metrics_esm_trained_all_eval_short = calc_metrics_from_smallesm_results(\n",
    "    \"../output/results/2024-03-01-esm-model-results-ssx1-all/\",\n",
    "    max_length=100,\n",
    ")\n",
    "metrics_esm_trained_short_eval_all = calc_metrics_from_smallesm_results(\n",
    "    \"../output/results/2024-02-29-esm-model-results-ssx1-max-length-100/\",\n",
    "    max_length=None,\n",
    ")\n",
    "metrics_esm_trained_short_eval_short = calc_metrics_from_smallesm_results(\n",
    "    \"../output/results/2024-02-29-esm-model-results-ssx1-max-length-100/\",\n",
    "    max_length=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d07fdfa-3229-45a1-90ac-dd8105d75de7",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_esm_trained_all_eval_all.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605d8b14-f5f2-42f3-87a0-615f710fde6c",
   "metadata": {},
   "source": [
    "#### Compare ESM-based models trained on all ORFs and only sORFs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d00df1d9-d7ed-4c05-8b1f-cc1237d22926",
   "metadata": {},
   "outputs": [],
   "source": [
    "# models trained on either all ORFs or only sORFs and evaluated on only sORFs.\n",
    "plot_heatmaps(\n",
    "    metrics_esm_trained_all_eval_short,\n",
    "    metrics_esm_trained_short_eval_short,\n",
    "    column=\"mcc\",\n",
    "    model_names=(\"ESM-based (trained all, eval short)\", \"ESM-based (trained short, eval short)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb07a2e1-f409-4349-b2d1-3311239481e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# models trained only on sORFs and evaluated on all or only sORFs.\n",
    "plot_heatmaps(\n",
    "    metrics_esm_trained_short_eval_all,\n",
    "    metrics_esm_trained_short_eval_short,\n",
    "    column=\"mcc\",\n",
    "    model_names=(\"ESM-based (trained short, eval all)\", \"ESM-based (trained short, eval short)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68e8ed9-24f4-469c-8b18-53b861ed69ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# models trained on all ORFs or only sORFs, but evaluated on all sequences.\n",
    "plot_heatmaps(\n",
    "    metrics_esm_trained_all_eval_all,\n",
    "    metrics_esm_trained_short_eval_all,\n",
    "    column=\"mcc\",\n",
    "    model_names=(\"ESM-based (trained all, eval all)\", \"ESM-based (trained short, eval all)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce5e80f0-a89b-48de-8bd1-11ee7538b9f5",
   "metadata": {},
   "source": [
    "### RNASamba predictions\n",
    "\n",
    "These predictions were generated by the script `plm-utils/scripts/rnasamba/train_and_evaluate.py` using the same datasets of deduplicated coding and noncoding transcripts generated by the `plmutils orf-prediction construct-data` command describe above. \n",
    "\n",
    "To train RNASamba models on all sequences:\n",
    "```\n",
    "python scripts/rnasamba-comparison/train_and_evaluate.py \\\n",
    "--coding-dirpath output/data/processed/final/coding-dedup-ssx1/transcripts \\\n",
    "--noncoding-dirpath output/data/processed/final/noncoding-dedup-ssx1/transcripts \\\n",
    "--output-dirpath 2024-02-28-rnasamba-results-ssx1-all \\\n",
    "```\n",
    "\n",
    "To train RNASamba models on transcripts corresponding to sORFs:\n",
    "```\n",
    "python scripts/rnasamba-comparison/train_and_evaluate.py \\\n",
    "--coding-dirpath output/data/processed/final/coding-dedup-ssx1/transcripts \\\n",
    "--noncoding-dirpath output/data/processed/final/noncoding-dedup-ssx1/transcripts \\\n",
    "--output-dirpath output/data/2024-02-28-rnasamba-results-ssx1-min-peptide-length-100 \\\n",
    "--max-length 100\n",
    "```\n",
    "The `--output-dirpath` above corresponds to the directory passed to the `calc_metrics_from_rnasamba_results` function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d08c522-1cb4-4406-ae1f-548192593506",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_metrics_from_rnasamba_results(rnasamba_results_dirpath):\n",
    "    \"\"\"\n",
    "    Aggregate the results from RNASamba models trained in the script\n",
    "    `scripts/rnasamba-comparison/train_and_evaluate.py`.\n",
    "    \"\"\"\n",
    "    all_metrics = []\n",
    "    dirpaths = [p for p in rnasamba_results_dirpath.glob(\"trained-on*\") if p.is_dir()]\n",
    "    for dirpath in dirpaths:\n",
    "        # dirnames are of the form 'trained-on-{species_id}-filtered'.\n",
    "        training_species_id = dirpath.stem.split(\"-\")[2]\n",
    "\n",
    "        prediction_filepaths = dirpath.glob(\"*.tsv\")\n",
    "        for prediction_filepath in prediction_filepaths:\n",
    "            # filenames are of the form '{species_id}-preds.csv'.\n",
    "            test_species_id = prediction_filepath.stem.split(\"-\")[0]\n",
    "\n",
    "            df = pd.read_csv(prediction_filepath, sep=\",\")\n",
    "            metrics = calc_metrics(\n",
    "                y_true=(df.true_label == \"coding\"), y_pred_proba=df.coding_score.values\n",
    "            )\n",
    "            metrics[\"training_species_id\"] = training_species_id\n",
    "            metrics[\"test_species_id\"] = test_species_id\n",
    "            metrics[\"num_coding\"] = (df.true_label == \"coding\").sum()\n",
    "            metrics[\"num_noncoding\"] = (df.true_label != \"coding\").sum()\n",
    "\n",
    "            all_metrics.append(metrics)\n",
    "\n",
    "    df = pd.DataFrame(all_metrics)\n",
    "    df[\"true_negative_rate\"] = df.num_true_negative / df.num_noncoding\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af4d3fd7-13f3-4be4-a311-789b27b503ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# models trained and tested on all transcripts.\n",
    "rnasamba_results_dirpath_all = pathlib.Path(\n",
    "    \"../output/results/2024-02-23-rnasamba-models-clustered-ssx3/\"\n",
    ")\n",
    "\n",
    "# models trained and tested only on transcripts whose longest ORFs are sORFs.\n",
    "rnasamba_results_dirpath_short = pathlib.Path(\n",
    "    \"../output/results/2024-02-28-rnasamba-results-ssx1-max-peptide-length-100/\"\n",
    ")\n",
    "\n",
    "metrics_rs_trained_all_eval_all = calc_metrics_from_rnasamba_results(rnasamba_results_dirpath_all)\n",
    "metrics_rs_trained_short_eval_short = calc_metrics_from_rnasamba_results(\n",
    "    rnasamba_results_dirpath_short\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c3109a1-6cb7-483a-a9f7-cf56c50a9b34",
   "metadata": {},
   "source": [
    "#### Compare RNASamba models trained on all or only sORFs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06119a03-ee3c-43e5-9351-a7a214aee265",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_heatmaps(\n",
    "    metrics_rs_trained_all_eval_all,\n",
    "    metrics_rs_trained_short_eval_short,\n",
    "    column=\"mcc\",\n",
    "    model_names=(\"RNASamba (all)\", \"RNASamba (short)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4f49649-2f71-4bd0-b68c-0535616ab27c",
   "metadata": {},
   "source": [
    "### Compare RNASamba and ESM-based models\n",
    "\n",
    "These are the most important plots in this notebook. They compare the performance of ESM-based models to RNASamba models by plotting the heatmap of performance metrics side by side."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3164f191-667a-4331-a7e2-260b1c6e8140",
   "metadata": {},
   "source": [
    "#### Models trained and evaluated on all transcripts (for RNASamba) or ORFs (for ESM-based)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a66f2fa1-b5f4-43b6-a7f5-4f5e2c42bdce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# overall performance (MCC metric)\n",
    "plot_heatmaps(\n",
    "    metrics_rs_trained_all_eval_all,\n",
    "    metrics_esm_trained_all_eval_all,\n",
    "    column=\"mcc\",\n",
    "    model_names=(\"RNASamba (all)\", \"ESM-based (all)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "360f6d2f-e1f1-4268-ab69-7d3964b5ee74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# recall (also the true positive rate, or num_true_positive / num_coding)\n",
    "plot_heatmaps(\n",
    "    metrics_rs_trained_all_eval_all,\n",
    "    metrics_esm_trained_all_eval_all,\n",
    "    column=\"recall\",\n",
    "    model_names=(\"RNASamba (all)\", \"ESM-based (all)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1075eca7-c526-431d-8c69-8d3f83958452",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the true negative rate.\n",
    "plot_heatmaps(\n",
    "    metrics_rs_trained_all_eval_all,\n",
    "    metrics_esm_trained_all_eval_all,\n",
    "    column=\"true_negative_rate\",\n",
    "    model_names=(\"RNASamba (all)\", \"ESM-based (all)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "239c44e1-df73-40cc-8866-bd99cef97769",
   "metadata": {},
   "source": [
    "#### Models trained only on short sequences (< 100aa)\n",
    "\n",
    "For RNASamba, this means the models were trained only on transcripts whose longest ORF was an sORF (less than 100aa long). \n",
    "\n",
    "Note that the class imbalance in this case is severe (most species do not have many coding transcripts whose longest ORF is an sORF) and this likely at least partly explains why the RNASamba models perform so poorly, as we do not compensate for the class imbalance during training (while we do compensate for it when training the ESM-based models). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e9806c2-8a27-4578-82b8-2ad01adda71f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_heatmaps(\n",
    "    metrics_rs_trained_short_eval_short,\n",
    "    metrics_esm_trained_short_eval_short,\n",
    "    column=\"mcc\",\n",
    "    model_names=(\"RNASamba (short)\", \"ESM-based (short)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49169892-fb85-48a5-8c0b-31c756be038a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_heatmaps(\n",
    "    metrics_rs_trained_short_eval_short,\n",
    "    metrics_esm_trained_short_eval_short,\n",
    "    column=\"recall\",\n",
    "    model_names=(\"RNASamba (short)\", \"ESM-based (short)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea2a582-10a2-41bf-a53c-210f46a6d579",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_heatmaps(\n",
    "    metrics_rs_trained_short_eval_short,\n",
    "    metrics_esm_trained_short_eval_short,\n",
    "    column=\"true_negative_rate\",\n",
    "    model_names=(\"RNASamba (short)\", \"ESM-based (short)\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f86f37-90f7-4f4a-89f0-c1662a7660f5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "18bc48c8-b690-4b35-90e1-934f070d432d",
   "metadata": {},
   "source": [
    "### Aside: blasting against peptipedia\n",
    "\n",
    "We were curious whether some of the false positives from ESM-based models represented genuine sORFs from lncRNAs (which are annotated as noncoding). As a way to examine this, we blasted all of the putative ORFs against peptipedia, and plotted the distribution of max evalues from putative sORFs for which the ESM-based model made either true and false positive predictions. If the model correctly identifies genuine sORFs from lncRNAs, we'd expect to see an enrichment of low evalues among the false positives.\n",
    "\n",
    "The command `plmutils orf-classification blast-peptipedia` was used to generate the directory of blast results that are loaded and concatenated by `concat_smallesm_results` function below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f531f10-63c7-4cdf-b5cd-0d015f175034",
   "metadata": {},
   "outputs": [],
   "source": [
    "def concat_smallesm_results(results_dirpath):\n",
    "    \"\"\"\n",
    "    Load and concatenate the predictions from esm-based models.\n",
    "    \"\"\"\n",
    "    dfs = []\n",
    "    prediction_filepaths = pathlib.Path(results_dirpath).glob(\"*.csv\")\n",
    "    for prediction_filepath in prediction_filepaths:\n",
    "        dfs.append(pd.read_csv(prediction_filepath))\n",
    "\n",
    "    return pd.concat(dfs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bebaeef2-50fc-4a79-ad08-4774d4dabee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# predictions from models trained on all putative ORFs.\n",
    "esm_trained_all_preds = concat_smallesm_results(\n",
    "    \"../output/results/2024-03-01-esm-model-results-ssx1-all/\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4de0e81-e278-445a-8028-5445b79dd063",
   "metadata": {},
   "outputs": [],
   "source": [
    "# predictions from models trained on short peptides (< 100aa).\n",
    "esm_trained_short_preds = concat_smallesm_results(\n",
    "    \"../output/results/2024-02-29-esm-model-results-ssx1-max-length-100/\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6774f99-9863-4db4-8aca-c95e3b4ddc01",
   "metadata": {},
   "outputs": [],
   "source": [
    "esm_trained_all_preds.shape, esm_trained_short_preds.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7cc1e18-071c-46a3-90b5-807f71300c1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "esm_trained_short_preds.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4835d5db-44bf-4bb1-84bd-3fa15410617b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# count the number of peptides from coding and noncoding transcripts to make sure\n",
    "# that the class imbalance between coding and noncoding is not too severe.\n",
    "# (we only need to look at preds from one model, since each model is tested with all species).\n",
    "hsap_preds = esm_trained_all_preds.loc[esm_trained_all_preds.training_species_id == \"hsap\"].copy()\n",
    "pd.merge(\n",
    "    hsap_preds.groupby([\"testing_species_id\", \"true_label\"]).count().sequence_id,\n",
    "    (\n",
    "        hsap_preds.loc[hsap_preds.sequence_length < 100]\n",
    "        .groupby([\"testing_species_id\", \"true_label\"])\n",
    "        .count()\n",
    "        .sequence_id\n",
    "    ),\n",
    "    left_index=True,\n",
    "    right_index=True,\n",
    "    suffixes=(\"_all\", \"_short\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efca0b8e-bc37-4342-bc10-519c6582dc07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def concat_blast_results(dirpaths):\n",
    "    \"\"\"\n",
    "    Aggregate the blast results generated by `plmutils orf-classification blast-peptipedia`.\n",
    "    \"\"\"\n",
    "    blast_results_columns = (\n",
    "        \"qseqid sseqid full_sseq pident length qlen slen mismatch gapopen qstart qend sstart send evalue bitscore\"  # noqa: E501\n",
    "    ).split(\" \")\n",
    "\n",
    "    dfs = []\n",
    "    for dirpath in dirpaths:\n",
    "        filepaths = pathlib.Path(dirpath).glob(\"*.tsv\")\n",
    "        for filepath in filepaths:\n",
    "            try:\n",
    "                df = pd.read_csv(filepath, sep=\"\\t\")\n",
    "            except Exception:\n",
    "                continue\n",
    "            df.columns = blast_results_columns\n",
    "            dfs.append(df)\n",
    "    return pd.concat(dfs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef7f7d8-64c2-45ef-8d24-93aca614a7af",
   "metadata": {},
   "outputs": [],
   "source": [
    "blast_results = concat_blast_results(\n",
    "    [\n",
    "        \"../output/data/processed/final/coding-dedup-ssx1/blast-peptipedia-results/\",\n",
    "        \"../output/data/processed/final/noncoding-dedup-ssx1/blast-peptipedia-results/\",\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ae380b-9230-490f-a231-c67363c227d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use the log of the evalue for readability.\n",
    "blast_results[\"evalue\"] = np.log(blast_results.evalue)\n",
    "\n",
    "# we only need to examine the minimum evalue for all hits to each peptide.\n",
    "min_evalues = blast_results.groupby(\"qseqid\").evalue.min().reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5192dbbd-78cd-4e98-8c9f-c47ac466e56b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# merge the minimum evalues with the model predictions.\n",
    "esm_trained_short_preds_w_evalues = pd.merge(\n",
    "    esm_trained_short_preds, min_evalues, left_on=\"sequence_id\", right_on=\"qseqid\", how=\"inner\"\n",
    ")\n",
    "\n",
    "esm_trained_all_preds_w_evalues = pd.merge(\n",
    "    esm_trained_all_preds, min_evalues, left_on=\"sequence_id\", right_on=\"qseqid\", how=\"inner\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a3fad9c-79e0-43e8-80eb-6f8750e0ad7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "esm_trained_short_preds_w_evalues_short_only = esm_trained_short_preds_w_evalues.loc[\n",
    "    esm_trained_short_preds_w_evalues.sequence_length < 100\n",
    "].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f44f49f-3376-431b-81ee-c5850ba9d44a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity-check: count the number of peptides that had hits in peptipedia.\n",
    "(\n",
    "    esm_trained_short_preds_w_evalues_short_only\n",
    "    # we only need to look at one model\n",
    "    .loc[esm_trained_short_preds_w_evalues_short_only.training_species_id == \"hsap\"]\n",
    "    .groupby([\"testing_species_id\", \"true_label\"])\n",
    "    .count()[[\"sequence_id\"]]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a947cac-ee4c-4d27-8306-d274d5cad195",
   "metadata": {},
   "source": [
    "#### Histograms of evalues for coding and noncoding transcripts\n",
    "\n",
    "This was to determine whether the false positives were enriched for peptides that had hits in peptipedia, which would suggest that they correspond to genuine sORFs from lncRNAs (and are therefore not actually false positives)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6582d0a9-d027-44f0-a7d1-fdffbe9ac336",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we only look at preds for short peptides from the human dataset\n",
    "# because it is one of the only that has a decent number of short peptides\n",
    "# with peptipedia hits and are from noncoding transcripts.\n",
    "preds = esm_trained_all_preds_w_evalues.loc[\n",
    "    (esm_trained_all_preds_w_evalues.training_species_id == \"hsap\")\n",
    "    & (esm_trained_all_preds_w_evalues.testing_species_id == \"hsap\")\n",
    "    & (esm_trained_all_preds_w_evalues.sequence_length < 100)\n",
    "]\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(16, 6))\n",
    "\n",
    "min_min_evalue = -150\n",
    "bins = np.arange(min_min_evalue, 0, -min_min_evalue / 30)\n",
    "kwargs = dict(bins=bins, density=False, alpha=0.5)\n",
    "\n",
    "# left axis: coding transcripts\n",
    "ax = axs[0]\n",
    "ax.hist(\n",
    "    preds[(preds.true_label == \"coding\") & (preds.predicted_probability > 0.5)].evalue,\n",
    "    label=\"True positives\",\n",
    "    color=\"blue\",\n",
    "    **kwargs,\n",
    ")\n",
    "ax.hist(\n",
    "    preds[(preds.true_label == \"coding\") & (preds.predicted_probability < 0.5)].evalue,\n",
    "    label=\"False negatives\",\n",
    "    color=\"red\",\n",
    "    **kwargs,\n",
    ")\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"Minimum log evalue\")\n",
    "ax.set_ylabel(\"Density\")\n",
    "ax.set_title(\"Coding transcripts\")\n",
    "\n",
    "# right axis: noncoding transcripts\n",
    "ax = axs[1]\n",
    "ax.hist(\n",
    "    preds[(preds.true_label == \"noncoding\") & (preds.predicted_probability < 0.5)].evalue,\n",
    "    label=\"True negatives\",\n",
    "    color=\"blue\",\n",
    "    **kwargs,\n",
    ")\n",
    "_ = ax.hist(\n",
    "    preds[(preds.true_label == \"noncoding\") & (preds.predicted_probability > 0.5)].evalue,\n",
    "    label=\"False positives\",\n",
    "    color=\"red\",\n",
    "    **kwargs,\n",
    ")\n",
    "ax.legend()\n",
    "ax.set_title(\"Noncoding transcripts\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe9235c-c9df-4020-a8ed-7a15a561b6ec",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "esm-py311-env",
   "language": "python",
   "name": "esm-py311-env"
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}