{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ee89a6ec-b731-44cc-89af-260af70ce678",
   "metadata": {},
   "source": [
    "# Context swap test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0e520279-ba5f-4872-a38d-27487ed26c56",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-10 15:32:02.506610: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE4.1 SSE4.2 AVX AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from creme import creme\n",
    "from creme import utils\n",
    "from creme import custom_model\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3083d60-240d-4c05-afc2-02bded899c81",
   "metadata": {},
   "source": [
    "Load Enformer and example sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5e3979fd-a872-41fa-afc4-109ee830b141",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-10 15:32:03.877407: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE4.1 SSE4.2 AVX AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-06-10 15:32:05.259255: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 69489 MB memory:  -> device: 0, name: NVIDIA A100 80GB PCIe, pci bus id: 0000:85:00.0, compute capability: 8.0\n"
     ]
    }
   ],
   "source": [
    "data_dir = \"../../../data\"\n",
    "track_index = [5111]\n",
    "model = custom_model.Enformer(track_index=track_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b84bb66-1c29-448f-b78f-3ec690e39d22",
   "metadata": {},
   "outputs": [],
   "source": [
    "fasta_path = f'{data_dir}/GRCh38.primary_assembly.genome.fa'\n",
    "seq_parser = utils.SequenceParser(fasta_path)\n",
    "\n",
    "genes = ['ABCA8_chr17_68955392_-', 'NFKBIZ_chr3_101849513_+']\n",
    "gene_seqs = {}\n",
    "\n",
    "for gene in genes:\n",
    "    gene_name, chrom, start, strand = gene.split('_')\n",
    "    seq = seq_parser.extract_seq_centered(chrom, int(start), strand, model.seq_length)\n",
    "    gene_seqs[gene_name] = seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61c69d44-b67f-4f71-b10b-5dc6c7b867f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TSS bin indeces\n",
    "bins = [447, 448]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "502ab1fc-bcb5-444d-ab7b-8351db4f7463",
   "metadata": {},
   "outputs": [],
   "source": [
    "abca8_wt = model.predict(gene_seqs['ABCA8'])[0,:,0]\n",
    "nfkbiz_wt = model.predict(gene_seqs['NFKBIZ'])[0,:,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34adb807-c057-4993-aa1a-dc2bf161d288",
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_track([abca8_wt], color='green', zoom=[0, 896], marks=bins)\n",
    "utils.plot_track([nfkbiz_wt], color='red', zoom=[0, 896], marks=bins)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "727193db-efab-48b9-9fe8-7fd4ed36b89a",
   "metadata": {},
   "source": [
    "**Context swap test**\n",
    "In this example we will swap the TSS tiles (5Kb) of an enhancing context sequence with a silencing one and vice versa. \n",
    "\n",
    "To run context swap test we need:\n",
    "- a loaded model\n",
    "- onehot encoded sequence (WT) of the source from which the TSS is taken\n",
    "- onehot encoded sequence (WT) of the target into which the TSS is embedded\n",
    "- a coordinate interval where the TSS is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa851a1f-f266-4142-821b-333d08aa684b",
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_halflen = model.seq_length // 2\n",
    "half_window_size = 2500\n",
    "N_shuffles = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0193ee8-b6d3-4e89-ad06-7448ae91ef62",
   "metadata": {},
   "source": [
    "TSS (from an enhancing sequence) embedded in a silencing context\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52086ff5-48a3-45b2-99ad-818cd054e529",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_mut = creme.context_swap_test(model, gene_seqs['ABCA8'], gene_seqs['NFKBIZ'],\n",
    "                                   [seq_halflen - half_window_size, seq_halflen + half_window_size])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f025f4-2a2b-4f68-8967-5195c346326d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax=utils.plot_track([abca8_wt], color='green', zoom=[400, 500], alpha=1, label='WT')\n",
    "utils.plot_track(pred_mut[:,:,0], color='purple', zoom=[400, 500], ax=ax, alpha=1, label='WT TSS + \\ncontext from \\nsilencing sequence')\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba882f12-f6aa-4661-8ddf-d259947d7230",
   "metadata": {},
   "source": [
    "TSS (from a silencing sequence) embedded in an enhancing context\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6497499c-9174-4b2b-af68-7f318e468d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_mut = creme.context_swap_test(model, gene_seqs['NFKBIZ'], gene_seqs['ABCA8'], \n",
    "                                   [seq_halflen - half_window_size, seq_halflen + half_window_size])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "421f16c9-2200-4ba4-aed6-b6caaa5bfb83",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax=utils.plot_track([nfkbiz_wt], color='red', zoom=[400, 500], alpha=1, label='WT')\n",
    "utils.plot_track(pred_mut[:,:,0], color='purple', zoom=[400, 500], ax=ax, alpha=1, label='WT TSS + \\ncontext from \\nenhancing sequence')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa56a3b-5005-4077-b535-9c6d14b8bd91",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}