notebooks/pre-processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dc975da1-4566-4d6c-a671-9371a3ea4300",
   "metadata": {},
   "source": [
    "# Pre-processing read alignments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a58d239b-ffe1-461e-80f3-6128bb89be95",
   "metadata": {},
   "source": [
    "**Summary:**\n",
    "\n",
    "This notebook compiles the scripts that are used as part of the pre-processing step of our analysis in order to produce a data table (referred in other notebooks as **dataset**) used accross all downstream analysis. \n",
    "\n",
    "It is organized in the following sections:<br>\n",
    "\n",
    "- **1. Extraction of genomic informations:**<br>\n",
    "Scripts used for extracting genomic features, namely genomic coordinates of the read alignment and chromosome.  \n",
    "\n",
    "- **2. Extraction of transcriptomic informations:**<br>\n",
    "Scripts used for extracting transcriptomics features such as cDNA start & end positions, gene and isoform on which the read is mapped, read orientation, length of the different regions (5' and 3' soft-clips, aligned portion of the sequence, etc). \n",
    "\n",
    "- **3. Generation of a summary dataframe:**<br>\n",
    "Regroupment of genomic and transcriptomic informations inside a single dataframe.<br> \n",
    "Also features a script for correcting genomic start & end positions based on transcriptomic alignments.\n",
    "\n",
    "- **4. Strand Switching Primer (SSP) search:**<br>\n",
    "Method used for identifying SSP sequences in reads\n",
    "\n",
    "- **5. Spliced Leader (SL) search:**<br>\n",
    "Method used for identifying splice leader sequences in reads (see notebooks for **Figure 2 and Figure 3**)\n",
    "\n",
    "- **6. Endogenous Hairpin search:**<br>\n",
    "Method used for identifying hairpin structure (see notebook for **Figure 4**)\n",
    "\n",
    "- **7. Result of Sequence Search:**<br>\n",
    "Method used for accepting and measuring SL, hairpin, Unidentified and SSP reads propensity at each gene's position\n",
    "(see notebooks for **Figure 5 and Figure 6**)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27f08e42-2a67-45c9-8856-61e52db11e88",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-01-31T17:34:03.500891Z",
     "iopub.status.busy": "2022-01-31T17:34:03.488799Z",
     "iopub.status.idle": "2022-01-31T17:34:03.725465Z",
     "shell.execute_reply": "2022-01-31T17:34:03.715320Z",
     "shell.execute_reply.started": "2022-01-31T17:34:03.497774Z"
    }
   },
   "source": [
    "---\n",
    "<br>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69a3f0d2-d651-42e0-9b1c-adb9664ae655",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa72ea86-8549-4c69-ac09-4112cf9835b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# datasets & scientific libraries\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "# json format library\n",
    "import json\n",
    "\n",
    "# regex library\n",
    "import re\n",
    "from re import search\n",
    "\n",
    "# plotting libraries\n",
    "import seaborn as sns\n",
    "import matplotlib\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "from matplotlib.ticker import MaxNLocator, MultipleLocator, AutoMinorLocator\n",
    "\n",
    "# sequence manipulation libraries\n",
    "import pysam\n",
    "import pyfastx\n",
    "import parasail\n",
    "import pyranges as pr\n",
    "from Bio.Seq import Seq\n",
    "from Bio import SeqIO\n",
    "\n",
    "# multiprocessing on all CPU cores\n",
    "from multiprocessing import Pool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0b09358-7520-4d20-8c57-999b83798e50",
   "metadata": {},
   "source": [
    "## Set environmental constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b27169d-07fd-4f3d-8524-bc385ab6dc63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path where files are stored\n",
    "path = '/Volumes/elegans/rna_sequencing'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f988d5-47c9-4813-b5ff-5610e375dd50",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set runs ID\n",
    "runs = ['SSP_1', 'SSP_2', 'SSP_3', 'SSP_4', 'SSP_5', 'SSP_6', 'SL1_1', 'NP_1', 'NP_2', 'NP_3', 'NP_4', 'NP_5']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab8f580c-0a35-4987-948e-68c7bd2a2f5f",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "043625c4-253b-4e10-b143-5f766368d1d9",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82a672a6-ac4a-4660-b55a-9dfe7fb9625c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-01-31T17:38:21.576699Z",
     "iopub.status.busy": "2022-01-31T17:38:21.575283Z",
     "iopub.status.idle": "2022-01-31T17:38:21.587175Z",
     "shell.execute_reply": "2022-01-31T17:38:21.585003Z",
     "shell.execute_reply.started": "2022-01-31T17:38:21.576636Z"
    }
   },
   "source": [
    "# 1. Extract informations from genomic alignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86c666a5-7736-43ab-9ef3-a9e655705b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract features from genomic alignments (read name, chromosome, genomic start & end positions)\n",
    "\n",
    "def genomic_stats(input_file, output_file):\n",
    "\n",
    "    name = []\n",
    "    start = []\n",
    "    end = []\n",
    "    chrom = []\n",
    "\n",
    "    alignment = pysam.AlignmentFile(input_file, 'rb')\n",
    "\n",
    "    for read in alignment:\n",
    "\n",
    "        # Only take into account primary reads\n",
    "        if not read.is_supplementary and not read.is_secondary and not read.is_unmapped and read.seq is not None:\n",
    "\n",
    "            name.append(read.query_name)\n",
    "            chrom.append(alignment.get_reference_name(read.reference_id))\n",
    "            start.append(read.reference_start+1) # 0-based coordinates - see doc\n",
    "            end.append(read.reference_end) # Points to one past last coordinate (so 1-based ?)\n",
    "\n",
    "    # Build dataframe with lists\n",
    "    table = pd.DataFrame(dict(read=name, chromosome=chrom, genomic_start=start, genomic_end=end))\n",
    "\n",
    "    # Save dataframe as a .tsv file\n",
    "    table.to_csv(output_file, sep='\\t', index=None)\n",
    "    \n",
    "    alignment.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf240c14-701f-4645-99b4-351849822f76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Handler function\n",
    "def genomic_stats_handler(ID):\n",
    "    \n",
    "    input_file = f'{path}/{ID}/{ID}-genome_sorted.bam'\n",
    "    output_name = f'{path}/{ID}-genomic_stats.tsv'\n",
    "    \n",
    "    # run script\n",
    "    genomic_stats(input_file, output_name)\n",
    "    \n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd6c001-3188-406d-9878-a2a423221f1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(genomic_stats_handler, runs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e433b4c7-f724-42d9-9642-1b24b4154dc5",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b4d30e-b4d9-492e-b596-5d48cdea0fc1",
   "metadata": {},
   "source": [
    "# 2. Extract informations from transcriptomic alignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed9dc3b-af2d-4da0-a9b2-80eefcb52a6d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Convert isoform name into gene name\n",
    "# ex: MTCE.35.1 -> MTCE.35\n",
    "\n",
    "def isoform_to_gene(isoform):\n",
    "    \n",
    "    match = re.search(r\"\\w+.\\d+\", isoform)\n",
    "\n",
    "    if match is not None:\n",
    "        return match.group(0)\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae635ffd-b520-460a-8514-ff9db76f3b92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract features from transcriptomic alignments (read name, isoform, gene, transcriptomic start & end positions, etc.)\n",
    "\n",
    "def transcriptomic_stats(input_file, output_file):\n",
    "\n",
    "    rows = []\n",
    "    \n",
    "    transcriptome = pysam.AlignmentFile(input_file, 'rb')\n",
    "\n",
    "    for read in transcriptome:\n",
    "\n",
    "        # Only take into account primary reads\n",
    "        if not read.is_unmapped and not read.is_secondary and not read.is_supplementary and read.seq is not None:\n",
    "\n",
    "            name = read.query_name\n",
    "            isoform = transcriptome.get_reference_name(read.reference_id)\n",
    "            \n",
    "            # start & end positions relative to the transcript of reference\n",
    "            transcript_start = read.reference_start+1\n",
    "            transcript_end = read.reference_end\n",
    "\n",
    "            # Find read orientation\n",
    "            orientation = 'antisense' if read.is_reverse else 'sense'\n",
    "\n",
    "            # sequence length\n",
    "            seq = read.seq\n",
    "            seq_length = len(str(seq))\n",
    "            \n",
    "            # alignment length\n",
    "            aln_start = read.query_alignment_start\n",
    "            aln_end = read.query_alignment_end\n",
    "            aln_length = len(seq[aln_start:aln_end])\n",
    "\n",
    "            # start is 0 so soft-clip is rannging from [-n to 0] with n being the soft-clip length\n",
    "            # which equals to aln start position\n",
    "            \n",
    "            # softclips length\n",
    "            # start is 0 so soft-clip is rannging from [-n to 0] with n being the soft-clip length\n",
    "            # which equals to aln start position\n",
    "            five_prime_sc = -aln_start\n",
    "            three_prime_sc = len(seq[aln_start:])\n",
    "\n",
    "            # Determine if long or short 5' soft-clip\n",
    "            softclip = 1 if aln_start > 80 else 0\n",
    "\n",
    "            # format as dict\n",
    "            read_row = {'read':name, 'isoform':isoform, 'read_orientation':orientation, 'softclip':softclip,\n",
    "                        'transcriptomic_start':transcript_start, 'transcriptomic_end':transcript_end,\n",
    "                        'sequence_length':seq_length, 'alignment_length':aln_length,\n",
    "                        'SC5':five_prime_sc, 'alignment_start':aln_start, 'alignment_end':aln_end, 'SC3':three_prime_sc}\n",
    "\n",
    "            rows.append(read_row)\n",
    "            \n",
    "    # Generate dataframe\n",
    "    dataframe = pd.DataFrame(rows)\n",
    "\n",
    "    # Find gene_ID with isoform value\n",
    "    dataframe['gene'] = dataframe['isoform'].apply(isoform_to_gene)\n",
    "\n",
    "    # Set columns order\n",
    "    dataframe = dataframe[['read', 'gene', 'isoform', 'read_orientation', 'softclip', 'transcriptomic_start', 'transcriptomic_end',\n",
    "                           'sequence_length', 'alignment_length', 'SC5', 'alignment_start', 'alignment_end', 'SC3']]\n",
    "    \n",
    "    # Save dataframe\n",
    "    dataframe.to_csv(output_file, sep='\\t', index=None)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe0e821c-0fa1-4302-96b1-29ec3bf6e24c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Handler function\n",
    "def transcriptomic_stats_handler(ID):\n",
    "\n",
    "    input_file = f'{path}/{ID}/{ID}-transcriptome_sorted.bam'\n",
    "    output_file = f'{path}/{ID}-transcriptome_stats.tsv'\n",
    "\n",
    "    # run script\n",
    "    transcriptomic_stats(input_file, output_file)\n",
    "    \n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7501090b-a1a9-46ad-bea8-2fd540609fa5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(transcriptomic_stats_handler, runs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdd8f7ea-4074-41d2-b9e9-101f2214b3ea",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c81e5896-ef63-4b33-83bb-5a478e84f8ab",
   "metadata": {},
   "source": [
    "# 3. Generation of a summary dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b46e9158-a410-4e0c-b56a-9ffa0c99d83c",
   "metadata": {},
   "source": [
    "## 3.a. Parsing genomic and transcriptomic informations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db187254-5e03-43b5-b37c-4c0948b7514e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set empty list to store intermediary dataframes\n",
    "dflist = []\n",
    "\n",
    "# loop over runs\n",
    "runs = ['SSP_1', 'SSP_2', 'SSP_3', 'SSP_4', 'SSP_5', 'SSP_6', 'SL1_1', 'NP_1', 'NP_2', 'NP_3', 'NP_4', 'NP_5']\n",
    "\n",
    "for ID in runs:\n",
    "\n",
    "    genome = f'{path}/{ID}-genomic_stats.tsv'\n",
    "    transcriptome = f'{path}/{ID}-transcriptome_stats.tsv'\n",
    "\n",
    "    # open genomic stats\n",
    "    genome_stats = pd.read_csv(genome, sep='\\t')\n",
    "    genome_stats = genome_stats.set_index('read')\n",
    "\n",
    "    # open transcriptomics stats\n",
    "    transcriptome_stats = pd.read_csv(transcriptome, sep='\\t')\n",
    "    transcriptome_stats = transcriptome_stats.set_index('read')\n",
    "\n",
    "    # concat tables together\n",
    "    final = pd.concat([genome_stats, transcriptome_stats], join='outer', axis=1)\n",
    "    \n",
    "    # only keep reads that are mapped against both genome and transcriptome files\n",
    "    final = final.loc[(final['genomic_start'].notnull()) & (final['transcriptomic_start'].notnull())]\n",
    "\n",
    "    final.index.name = 'read'\n",
    "    final = final.reset_index()\n",
    "\n",
    "    final['run'] = ID\n",
    "\n",
    "    final = final[['read', 'gene', 'isoform', 'chromosome', 'read_orientation', 'softclip', 'run',\n",
    "                   'genomic_start', 'genomic_end', 'transcriptomic_start', 'transcriptomic_end', \n",
    "                   'sequence_length', 'alignment_length', 'SC5', 'alignment_start', 'alignment_end', 'SC3']]\n",
    "    \n",
    "    dflist.append(final)\n",
    "\n",
    "stats = pd.concat(dflist, ignore_index=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f477dc6-77d7-4753-9756-095c3fff81c8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Build dictionnary of gene orientation \n",
    "strand = pd.read_csv(f'{path}/ref/gene&strand.tsv', sep='\\t')\n",
    "strands = {x['gene']: x['strand'] for idx, x in strand.iterrows()}\n",
    "\n",
    "# set gene orientation\n",
    "stats['gene_orientation'] = stats.apply(lambda x: strands[x['gene']], axis=1)\n",
    "\n",
    "# Set Start and end correctly\n",
    "stats['start'] = stats.apply(lambda x: x['genomic_start'] if strands[x['gene']] == '+' else x['genomic_end'], axis=1)\n",
    "stats['end'] = stats.apply(lambda x: x['genomic_end'] if strands[x['gene']] == '+' else x['genomic_start'], axis=1)\n",
    "\n",
    "stats['genomic_start'] = stats['start']\n",
    "stats['genomic_end'] = stats['end']\n",
    "\n",
    "# reorder and drop 'start' and 'end' columns\n",
    "stats = stats[['read', 'gene', 'isoform', 'chromosome', 'gene_orientation', 'read_orientation', 'softclip', 'run',\n",
    "               'genomic_start', 'genomic_end', 'transcriptomic_start', 'transcriptomic_end', \n",
    "               'sequence_length', 'alignment_length', 'SC5', 'alignment_start', 'alignment_end', 'SC3']]\n",
    "\n",
    "# save result\n",
    "stats.to_csv(f'{path}/dataset.tsv', sep='\\t', index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ec48528-6adc-41b7-933f-1f6de5336645",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1708310-5ae8-4445-95a6-be393b68a0b9",
   "metadata": {},
   "source": [
    "## 3.b. Correction of genomic coordinates based on transcriptome alignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae26c0d-0768-462b-85af-4e0529ffbdb1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Open GTF file as PyRanges-object\n",
    "gtf = pr.read_gtf(f'{path}/ref/c_elegans.PRJNA13758.WS270.canonical_geneset.gtf')\n",
    "gtf = gtf.df\n",
    "\n",
    "strand = gtf.loc[gtf['Feature'] == 'transcript'].set_index('transcript_id')['Strand'].to_dict()\n",
    "\n",
    "gtf = gtf.loc[~gtf['exon_id'].isna()][['exon_id', 'Start', 'End']].set_index('exon_id')[['Start', 'End']].to_dict()\n",
    "\n",
    "# Open transcriptome_relative_coordinates\n",
    "with open(f'{path}/ref/transcript_exons.json', 'r') as dico:\n",
    "    exonics_positions = json.loads(dico.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9fd1762-749a-4f10-8ca0-900b6a0754eb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# This script converts a transcriptomic position into a genomic position based on a given isoform\n",
    "def transcriptome_based_correction(transcript, transcriptome_position):\n",
    "    \n",
    "    exons = exonics_positions[transcript]\n",
    "\n",
    "    for exon, i in exons.items():\n",
    "\n",
    "        if i[0] <= transcriptome_position <= i[1]:\n",
    "\n",
    "            delta = transcriptome_position - i[0]\n",
    "\n",
    "            if strand[transcript] == '+':\n",
    "                exon_start = gtf['Start'][f'{transcript}.e{exon}']\n",
    "                genomic_position = exon_start + delta\n",
    "            else:\n",
    "                exon_start = gtf['End'][f'{transcript}.e{exon}']\n",
    "                genomic_position = exon_start - delta\n",
    "\n",
    "            return genomic_position\n",
    "\n",
    "        else:\n",
    "            if int(exon) < len(exons):\n",
    "                continue\n",
    "            else:\n",
    "                return np.NaN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c83bd16-460b-4901-8695-665fe27f18e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Transcriptomic alignment positions are used to refine genomic positions\n",
    "def correct_genomic_position(input_file):\n",
    "\n",
    "    table = pd.read_csv(input_file, sep='\\t')\n",
    "    \n",
    "    # correct start and end positons\n",
    "    table['corrected_genomic_start'] = table.apply(lambda x: transcriptome_based_correction(x['isoform'], x['transcriptomic_start']), axis=1)\n",
    "    table['corrected_genomic_end'] = table.apply(lambda x: transcriptome_based_correction(x['isoform'], x['transcriptomic_end']), axis=1)\n",
    "    \n",
    "    # useful ?\n",
    "    table = table.loc[table['gene'].notnull()]\n",
    "    \n",
    "    # reorder columns\n",
    "    final = table[['read', 'gene', 'isoform', 'chromosome', 'gene_orientation', 'read_orientation', 'softclip', 'run',\n",
    "                   'corrected_genomic_start', 'corrected_genomic_end', 'genomic_start', 'genomic_end', \n",
    "                   'transcriptomic_start', 'transcriptomic_end', \n",
    "                   'sequence_length', 'alignment_length', 'SC5', 'alignment_start', 'alignment_end', 'SC3']]\n",
    "    \n",
    "    # return table with corrected positions\n",
    "    return final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c00d9441-5d35-46be-b49a-e4793014534e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Run script\n",
    "corrected_dataset = correct_genomic_position(input_file=f'dataset.tsv')\n",
    "\n",
    "# clean up columns type\n",
    "col_types = {'read':object, 'gene':object, 'isoform':object, 'chromosome':object,\n",
    "             'genomic_start':float, 'genomic_end':float, 'corrected_genomic_start':float, 'corrected_genomic_end':float,\n",
    "             'read_orientation':object, 'softclip':int, 'run':object}\n",
    "\n",
    "corrected_dataset = corrected_dataset.astype(col_types)\n",
    "\n",
    "# save table\n",
    "corrected_dataset.to_csv(f'{path}/dataset.tsv', sep='\\t', index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "910f5787-4345-434a-a334-8486cfb8b1a4",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c61da2ef-1187-4d34-af8c-2f326ad2ccc2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-01-31T18:41:29.263946Z",
     "iopub.status.busy": "2022-01-31T18:41:29.260207Z",
     "iopub.status.idle": "2022-01-31T18:41:29.291939Z",
     "shell.execute_reply": "2022-01-31T18:41:29.288436Z",
     "shell.execute_reply.started": "2022-01-31T18:41:29.263805Z"
    }
   },
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ecaf7a0-db70-44fb-b423-f129abac43e0",
   "metadata": {},
   "source": [
    "# 4. Sequence search: Strand Switching Primer (SSP)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12993e9a-9900-4d05-973f-5f5e310d32ab",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 4.a. Extraction of the full 5'soft-clip sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a75be8-134c-4ac9-b93a-566ee193b319",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_full_softclip(input_file, output_file, ALN=0):\n",
    "    \n",
    "    # create output file\n",
    "    with open(output_file, 'w+') as file:\n",
    "    \n",
    "        # alignments\n",
    "        alignments = pysam.AlignmentFile(input_file, 'rb')\n",
    "\n",
    "        for read in alignments:\n",
    "\n",
    "            if not read.is_unmapped and not read.is_supplementary and not read.is_secondary and read.seq is not None:\n",
    "\n",
    "                name = read.query_name\n",
    "\n",
    "                start = read.query_alignment_start\n",
    "                seq = str(read.seq[0:(start+ALN)])\n",
    "\n",
    "                # write to file\n",
    "                file.write(f'>{name}\\n{seq}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8538cad-3d46-4804-a899-2bf4325af17a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Handler function\n",
    "def extract_full_softclip_handler(ID):\n",
    "    \n",
    "    input_file = f'{path}/{ID}/{ID}-transcriptome_sorted.bam'\n",
    "    output_file= f'{path}/{ID}/{ID}-five_prime_softclip[SC=FULL|ALN=80].fa'\n",
    "\n",
    "    # run script\n",
    "    extract_full_softclip(input_file, output_file, ALN=80)\n",
    "\n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7827b588-90d2-4a4d-aa7b-c4aaf334600d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(extract_full_softclip_handler, runs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9d613f5-b0fe-4a9a-9a34-43c0319bf7dd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 4.b. SSP sequence search on 5' soft-clip sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f6b2c75-54e3-41a8-97d9-604274d99c92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SSP_ALIGN_PARAMS = {'match': 1,\n",
    "                    'mismatch': 0,\n",
    "                    'gap_open': 2,    # penalty to create a gap\n",
    "                    'gap_extend': 1}  # penalty to extend a gap (must have created before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f27f5027-a5df-41a3-a4c2-90ac6e9044ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def semi_global_alignment(reference, query, params):\n",
    "\n",
    "    subs_mat = parasail.matrix_create(\"ACGT\", params['match'], params['mismatch'])\n",
    "    alignment = parasail.sg_trace_striped_32(reference, query, params['gap_open'], params['gap_extend'], subs_mat)\n",
    "\n",
    "    return alignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35576d66-1951-4bd5-bb5f-6a7464a6c079",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def evaluate_alignment(sequence, motif, seqtype, sensitivity=0.7):\n",
    "    \n",
    "    ref = len(motif)\n",
    "    \n",
    "    # initialize\n",
    "    aln_score = None\n",
    "    aln_size = None\n",
    "    aln_dist = None\n",
    "    \n",
    "    # initial alignment (full seq)\n",
    "    aln = semi_global_alignment(sequence, motif, params=SSP_ALIGN_PARAMS)\n",
    "    score = aln.score\n",
    "    \n",
    "    # shorter alignments\n",
    "    if score < len(motif) * sensitivity:\n",
    "        \n",
    "        for pos in reversed(range(15, len(motif), 1)):\n",
    "            \n",
    "            if seqtype == \"5SC\":\n",
    "                aln = semi_global_alignment(sequence, motif[-pos:], params=SSP_ALIGN_PARAMS)\n",
    "            elif seqtype == \"3SC\":\n",
    "                aln = semi_global_alignment(sequence, motif[:pos], params=SSP_ALIGN_PARAMS)\n",
    "            else:\n",
    "                break\n",
    "                \n",
    "            score = aln.score\n",
    "\n",
    "            if score < sensitivity * pos:\n",
    "                continue\n",
    "            else:\n",
    "                aln_score = score\n",
    "                aln_size = pos\n",
    "                aln_dist = len(sequence) - (int(aln.end_query) + 1)\n",
    "                break\n",
    "\n",
    "    else:\n",
    "        aln_score = score\n",
    "        aln_size = len(motif)\n",
    "        aln_dist = len(sequence) - (int(aln.end_query) + 1)\n",
    "        \n",
    "    return aln_score, aln_size, aln_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59f41c95-f193-444a-9537-cf243e8d7b81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def search_SSP(input_file, output_file, offset=0):\n",
    "    \n",
    "    reads = []\n",
    "    SSP_scores = []\n",
    "    SSP_sizes = []\n",
    "    SSP_distances = []\n",
    "    \n",
    "    for read in SeqIO.parse(input_file, \"fasta\"):\n",
    "        \n",
    "        name = read.id\n",
    "        SC5_sequence = str(read.seq)\n",
    "        \n",
    "        score, size, dist = evaluate_alignment(SC5_sequence, 'TTTCTGTTGGTGCTGATATTGCTGGG', '5SC')\n",
    "\n",
    "        if score is None:\n",
    "            reads.append(name)\n",
    "            SSP_scores.append(0)\n",
    "            SSP_sizes.append(0)\n",
    "            SSP_distances.append(None)\n",
    "        else:\n",
    "            reads.append(name)\n",
    "            SSP_scores.append(score)\n",
    "            SSP_sizes.append(size)\n",
    "            SSP_distances.append(-(dist-offset))\n",
    "\n",
    "    data = pd.DataFrame(dict(read=reads, SSP_score=SSP_scores, SSP_size=SSP_sizes, SSP_dist=SSP_distances))\n",
    "    data = data.set_index('read')\n",
    "    data.to_csv(output_file, sep='\\t', index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d2ea85b-65a5-4e18-a758-9ff979a154eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_SSP_handler(ID):\n",
    "    \n",
    "    input_file = f'{path}/{ID}/{ID}-five_prime_softclip[SC=FULL|ALN=80].fa'\n",
    "    output_file= f'{path}/{ID}/{ID}-SSP_search.tsv'\n",
    "\n",
    "    # run script\n",
    "    search_SSP(input_file, output_file, offset=80)\n",
    "\n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46f929d2-9550-4015-a7d4-402d7598f70d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(search_SSP_handler, runs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be95a11-43f3-41be-9c7c-3ff71ebc9150",
   "metadata": {},
   "source": [
    "## 4.c. Processing SSP search result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7b01e4d-e76a-4dae-8d05-c99326edc887",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_list = []\n",
    "\n",
    "for ID in runs:\n",
    "    \n",
    "    # open SL search result as dataframe\n",
    "    df = pd.read_csv(f'{path}/{ID}/{ID}-SSP_search.tsv', sep='\\t')\n",
    "    \n",
    "    # append to new dataframe\n",
    "    df_list.append(df)\n",
    "    \n",
    "# concat dataframes\n",
    "ssp_result = pd.concat(df_list, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d90a6923-e7b9-4df7-afcf-43cfc550b6b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge SL result dataframe with dataset\n",
    "dataset_SSP = corrected_dataset.merge(ssp_result, on='read', how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e7a31fa-c28b-43b6-a260-07e7ebb56064",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Determine sequences found based on scores\n",
    "dataset_SSP['SSP_FOUND'] = dataset_SSP.apply(lambda x: 'FOUND' if x['SSP_dist']==x['SSP_dist'] and x['SSP_dist']-x['SC5']<=80 else 'NOT FOUND', axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e19c18ca-2890-4bd1-a1b3-1735f120a9c7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save result\n",
    "dataset_SSP.to_csv(f'{path}/dataset_+SSP.tsv', sep='\\t', index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb2d14dc-1022-40e3-9e3e-97ae0551fdb4",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b150643f-3907-4348-8e65-067c98597311",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d463f090-3a82-44d2-acdc-62a623079f18",
   "metadata": {},
   "source": [
    "# 5. Sequence search: Splice Leader (SL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7729de9e-25c2-4271-8a6a-ac669db67936",
   "metadata": {},
   "source": [
    "This section of the notebook corresponds to the code used for searching splice leader (SL) sequences in the different reads obtained from the sequencing of *C. elegans* transcriptome.\n",
    "\n",
    "It is composed of three sub-sections:<br>\n",
    "**a.** Extraction of the last 100bp of the 5' soft-clip sequence.<br>\n",
    "**b.** Definition of the algorithm used for identifying a SL sequence.<br>\n",
    "**c.** Inclusion of the SL search result to the dataset table."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44c483eb-24f4-4309-a911-b1c28d8913d0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 5.a. Extracting 5' soft-clip sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf3ded64-12ff-47ac-b8b9-bec94b29e2cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract last 100bp of 5' soft-clip sequence and first 20bp of primary alignment from transcriptomic alignments\n",
    "\n",
    "def extract_softclip_end(input_file, output_file, SC, ALN=0):\n",
    "\n",
    "    alignments = pysam.AlignmentFile(input_file, 'rb')\n",
    "\n",
    "    with open(output_file, 'w+') as fasta:\n",
    "\n",
    "        for read in alignments:\n",
    "\n",
    "            if not read.is_secondary and not read.is_supplementary and not read.is_unmapped and read.seq is not None:\n",
    "\n",
    "                name = read.query_name\n",
    "                # equivalent to soft-clip length\n",
    "                start = read.query_alignment_start\n",
    "\n",
    "                # softclip longer than 100bp\n",
    "                if start > SC:\n",
    "                    seq = read.seq[(start-SC) :(start+ALN)]\n",
    "\n",
    "                # softclip smaller than 100bp\n",
    "                elif start <= SC:\n",
    "                    seq = read.seq[:(start+ALN)]\n",
    "\n",
    "                fasta.write(f'>{name}\\n{seq}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e01d2949-611c-4fcf-a33e-7fbe2f1cd2cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Handler function\n",
    "def extract_softclip_end_handler(ID):\n",
    "    \n",
    "    input_file = f'{path}/{ID}/{ID}-transcriptome_sorted.bam'\n",
    "    output_file= f'{path}/{ID}/{ID}-five_prime_softclip[SC=100|ALN=20].fa'\n",
    "\n",
    "    # run script\n",
    "    extract_softclip_end(input_file, output_file, SC=100, ALN=20)\n",
    "\n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e85ddf9-1315-48bc-aaf9-02ff8e9336f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(extract_softclip_end_handler, runs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dc45748-82e6-447c-b9ea-b6fc5bc796e5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-05T09:06:57.197274Z",
     "iopub.status.busy": "2021-10-05T09:06:57.196988Z",
     "iopub.status.idle": "2021-10-05T09:06:57.202826Z",
     "shell.execute_reply": "2021-10-05T09:06:57.201398Z",
     "shell.execute_reply.started": "2021-10-05T09:06:57.197244Z"
    },
    "tags": []
   },
   "source": [
    "## 5.b. SL Search (1st Pass - on 5' soft-clip end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cacf19e-f5df-4dcc-808a-6aa8089f0083",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SL_ALIGN_PARAMS = {'match': 1,\n",
    "                   'mismatch': -1,\n",
    "                   'gap_open': 2,    # penalty to create a gap\n",
    "                   'gap_extend': 1}  # penalty to extend a gap (must have created before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbbe0954-9a18-4389-a673-ce0022293932",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# For each read, all known SL1 and SL2 sequences are matched against their 5' soft-clip sequence.\n",
    "# A match is returned if the alignment score is higher or equal to 70% (sensitivity parameter) of the sequence length\n",
    "# ex: a 10bp sequence will be accepted if the returned score is higher of equal at 7.\n",
    "\n",
    "# If no match is returned, the evaluated SL sequence is shortened (on the 5' side) and the match is performed again.\n",
    "# ex: sequence 'AAATTGTGTGTGT' returned no match, we then search 'ATTGTGTGTGT'\n",
    "\n",
    "# The process is repeated until we reach a minimal sequence of 7bp. If no match is found, the SL is considered not found.\n",
    "# If various SL sequences return the same score they are equally accepted\n",
    "\n",
    "def search_spliced_leaders(input_file, output_file, sensitivity=0.7, offset=0):\n",
    "\n",
    "    seq_file = open(f'{path}/ref/SL_sequences.fasta')\n",
    "    SPLICELEADERS = {record.id: str(record.seq) for record in SeqIO.parse(seq_file, \"fasta\")}\n",
    "\n",
    "    reads = {}\n",
    "\n",
    "    for record in SeqIO.parse(input_file, \"fasta\"):\n",
    "\n",
    "        ref = len(record.seq)\n",
    "        \n",
    "        aln_scores = {}\n",
    "        position = {}\n",
    "        distance = {}\n",
    "\n",
    "        for sl_name in SPLICELEADERS:\n",
    "\n",
    "            sl_seq = SPLICELEADERS[sl_name]\n",
    "\n",
    "            sl_length = len(sl_seq)\n",
    "\n",
    "            aln = semi_global_alignment(str(record.seq), sl_seq, params=SL_ALIGN_PARAMS)\n",
    "            score = aln.score\n",
    "\n",
    "            if score < sl_length * sensitivity:\n",
    "\n",
    "                for pos in reversed(range(7, sl_length, 1)):\n",
    "\n",
    "                    aln = semi_global_alignment(str(record.seq), sl_seq[-pos:], params=SL_ALIGN_PARAMS)\n",
    "                    score = aln.score\n",
    "\n",
    "                    if score < sensitivity * pos:\n",
    "                        continue\n",
    "                    else:\n",
    "                        aln_scores[sl_name] = score\n",
    "                        position[sl_name] = pos\n",
    "                        distance[sl_name] = ref - (int(aln.end_query) + 1)\n",
    "                        break\n",
    "            else:\n",
    "                aln_scores[sl_name] = score\n",
    "                position[sl_name] = sl_length\n",
    "                distance[sl_name] = ref - (int(aln.end_query) + 1)\n",
    "\n",
    "        if len(aln_scores) > 0:\n",
    "\n",
    "            #### Get best % match SL\n",
    "            top_score = max(aln_scores.values())\n",
    "            best_matches = [sl for sl, value in aln_scores.items() if value == top_score]\n",
    "\n",
    "            if len(best_matches) == 1:\n",
    "\n",
    "                sl_found = best_matches[0]\n",
    "                distance = distance[sl_found]\n",
    "                reads[record.id] = (sl_found, top_score, distance)\n",
    "\n",
    "            else:\n",
    "\n",
    "                distance = [dist for sl, dist in distance.items() if sl in best_matches]\n",
    "                small_dist = min(distance)\n",
    "                ix = [n for n, dist in enumerate(distance) if dist == small_dist]\n",
    "                \n",
    "                closest_match = [sl for n, sl in enumerate(best_matches) if n in ix]\n",
    "\n",
    "                sl_found = ' / '.join(closest_match)\n",
    "\n",
    "                reads[record.id] = (sl_found, top_score, small_dist)\n",
    "\n",
    "\n",
    "    final = pd.DataFrame.from_dict(reads, orient='index')\n",
    "    final.columns = ['SL','score','distance_to_start']\n",
    "    \n",
    "    # account for bases aligned at the beginning\n",
    "    final['distance_to_start'] = -(final['distance_to_start'] - offset)\n",
    "    \n",
    "    final.index.name = 'read'\n",
    "    final.to_csv(output_file, sep='\\t', index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf2c9a7-3483-4484-93e4-d92b81cb2778",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Handler function\n",
    "def search_spliced_leaders_handler(ID):\n",
    "    \n",
    "    input_file = f'{path}/{ID}-five_prime_softclip[SC=100|ALN=20].fa'\n",
    "    output_file= f'{path}/{ID}/{ID}-SL_search.tsv'\n",
    "\n",
    "    # run script\n",
    "    search_spliced_leaders(input_file, output_file, sensitivity=0.7, offset=20)\n",
    "\n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0ca205-d89a-4fc1-97e4-f5368ba23645",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(search_spliced_leaders_handler, runs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c55d5c-cdab-4571-a164-33d1a1475440",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set new dataframe\n",
    "sl_result = pd.DataFrame()\n",
    "\n",
    "for ID in runs:\n",
    "    \n",
    "    # open SL search result as dataframe\n",
    "    df = pd.read_csv(f'{path}/{ID}/{ID}-SL_search.tsv', sep='\\t')\n",
    "    \n",
    "    # append to new dataframe\n",
    "    sl_result = pd.concat([sl_result, df], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a8fe08e-ae7f-4bde-9639-83ad96a5dc8e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Returns a binary result (FOUND / NOT FOUND) based on the detection of a SL with high confidence or not\n",
    "# Robust SL are defined based on their alignment score (score of 10 or higher)\n",
    "# Or based on how close they are found to the start of the alignment\n",
    "\n",
    "def robust_SL(SL, score, dist):\n",
    "    \n",
    "    if SL:\n",
    "        if score > 9 or dist > -3:\n",
    "            return 'FOUND'\n",
    "        else:\n",
    "            return 'NOT FOUND'\n",
    "    else:\n",
    "        return 'NOT FOUND'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e96442b-0db5-44e1-886e-3e4e636ec5b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def robust_variant(SL, score, dist):\n",
    "\n",
    "    if SL:\n",
    "        \n",
    "        # any match which scored 10 or plus\n",
    "        if score > 9 :\n",
    "            return 'FOUND'\n",
    "        \n",
    "        # matchs which score 8 or 9 IF immediatily near the ATG\n",
    "        else:\n",
    "            if score > 7 and dist > -3 :\n",
    "                return 'FOUND'\n",
    "            else:\n",
    "                return 'NOT FOUND'\n",
    "    else:\n",
    "        return 'NOT FOUND'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5098bf30-1f52-4db7-a6c6-f7a5f9f33042",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Find robust SL match\n",
    "sl_result['ROBUST_SL'] = sl_result.apply(lambda x: robust_SL(SL=x['SL'], score=x['score'], dist=x['distance_to_start']), axis=1)\n",
    "\n",
    "# Find robust SL variant\n",
    "sl_result['ROBUST_VARIANT'] = sl_result.apply(lambda x: robust_variant(SL=x['SL'], score=x['score'], dist=x['distance_to_start']), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f844fe-6915-4845-9d6e-34881d9b2718",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# filter for non-SL reads\n",
    "reads_sl = sl_result[sl_result['ROBUST_SL']=='FOUND']\n",
    "reads_not_sl = sl_result[sl_result['ROBUST_SL']=='NOT FOUND']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "226b3a27-30ff-4fe8-9dbe-b0bb6c6a0855",
   "metadata": {},
   "source": [
    "## 5.c. SL Search (2nd pass - on whole 5' soft-clip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "853e1f7e-1956-4266-b5ce-83ac3593e9ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Handler function\n",
    "def search_spliced_leaders_2nd_pass(ID):\n",
    "    \n",
    "    input_file = f'{path}/{ID}/{ID}-five_prime_softclip[SC=FULL|ALN=80].fa'\n",
    "    output_file= f'{path}/{ID}/{ID}-SL_search_(2nd_pass).tsv'\n",
    "\n",
    "    # run script\n",
    "    search_spliced_leaders(input_file, output_file, sensitivity=0.7, offset=80)\n",
    "\n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b575e9c-aef9-453f-99c6-e8892a31fa81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# processing on all CPU cores\n",
    "with Pool(4) as p:\n",
    "    p.map(search_spliced_leaders_2nd_pass, runs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b488cc-4bd0-4020-8dff-c7f99db44446",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set new dataframe\n",
    "sl_result_2ndpass = pd.DataFrame()\n",
    "\n",
    "for ID in runs:\n",
    "    \n",
    "    # open SL search result as dataframe\n",
    "    df = pd.read_csv(f'{path}/{ID}/{ID}-SL_search_(2nd_pass).tsv', sep='\\t')\n",
    "    \n",
    "    # append to new dataframe\n",
    "    sl_result_2ndpass = pd.concat([sl_result_2ndpass, df], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5be081b-c7d9-455c-85fc-cce2b444c995",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Find robust SL match\n",
    "sl_result_2ndpass['ROBUST_SL'] = sl_result_2ndpass.apply(lambda x: 'FOUND' if x['SL']==x['SL'] and x['score'] >= 12 else 'NOT FOUND', axis=1)\n",
    "\n",
    "# Find robust SL variant\n",
    "sl_result_2ndpass['ROBUST_VARIANT'] = 'NOT FOUND'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab24e13f-b65f-4bdb-8379-93ae0b1f4a23",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# take new SL reads\n",
    "reads_no_sl = list(reads_not_sl['read'])\n",
    "sl_found_2ndpass = sl_result_2ndpass[(sl_result_2ndpass['ROBUST_SL']=='FOUND') & (sl_result_2ndpass['read'].isin(reads_no_sl))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a46514c-a247-49a6-b723-de4d9d2d3108",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# then merge back\n",
    "sl_result.set_index('read', inplace=True)\n",
    "sl_result.update(sl_found_2ndpass.set_index('read'))\n",
    "sl_result = sl_result.reset_index()\n",
    "\n",
    "sl_result.columns = ['read','SL', 'SL_score', 'SL_distance', 'ROBUST_SL_FOUND', 'VARIANT_SL_FOUND'] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d64398e-6e13-4c3d-94a0-138d1c26d817",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge SL result dataframe with dataset\n",
    "dataset_SSP_SL = dataset_SSP.merge(sl_result, on='read', how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3514a78f-4063-43e1-a8f7-4a7b2aec9f40",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save result\n",
    "dataset_SSP_SL.to_csv(f'{path}/dataset_+SSP+SL.tsv', sep='\\t', index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d922fb6-9e17-4087-a187-69cf34f3dbc8",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb0f6f3d-67d0-4cb2-9a34-21130ecb063f",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32b21bc3-37fb-4604-a053-a197074de29a",
   "metadata": {},
   "source": [
    "# 6. Sequence search: Endogenous hairpins (Mimic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b778b38-6e20-420c-b9a6-886a27a9106e",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2022-02-03T10:14:35.906837Z",
     "iopub.status.idle": "2022-02-03T10:14:35.907244Z",
     "shell.execute_reply": "2022-02-03T10:14:35.907097Z"
    }
   },
   "source": [
    "This section of the notebook corresponds to the code used for searching hairpin sequences (referred as **hairpin mimic**) sequences in non-SL reads.\n",
    "\n",
    "\n",
    "**a.** Definition of the algorithm used for identifying hairpin mimics.<br>\n",
    "**b.** Processing result of the search algorithm and integration to the dataset table.<br>\n",
    "**c.** Determination of SL and hairpin mimic events in each gene.<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99ce74ac-2e76-40f8-b5da-fa259c7cd98b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 6.a hairpin search algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6d50314-3207-4c83-80fe-f20f1944f87f",
   "metadata": {},
   "outputs": [],
   "source": [
    "HAIRPIN_ALIGN_PARAMS = {'match': 1,\n",
    "                        'mismatch': 0,\n",
    "                        'gap_open': 2,    # penalty to create a gap\n",
    "                        'gap_extend': 1}  # penalty to extend a gap (must have created before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4e9274d-9202-458a-b7e0-12831ac4b62a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def hairpin_search(input_file, output_file, length_s1=12, length_s2=12, offset_value=0):\n",
    "    \n",
    "    read = []\n",
    "    scores = []\n",
    "    first_stem = []\n",
    "    second_stem = []\n",
    "    loop_scores = []\n",
    "    \n",
    "    for record in SeqIO.parse(input_file, \"fasta\"):\n",
    "        \n",
    "        # set read's constant \n",
    "        seq = record.seq\n",
    "        seq_length = len(seq)\n",
    "        min_index = 50\n",
    "        offset = offset_value - seq_length \n",
    "\n",
    "        # initialize score & distance & stem regions\n",
    "        hairpin_found = False\n",
    "        best_score = 0\n",
    "        last_best_n = 0\n",
    "        s1 = None\n",
    "        s2 = None\n",
    "        loop_size = None\n",
    "\n",
    "        # start recursive search\n",
    "        # from last position to first positions\n",
    "        for n in reversed(range(min_index, seq_length+1, 2)):\n",
    "\n",
    "            # determine loop space available for computing loop\n",
    "            available_loop = n - min_index\n",
    "            max_loop = 10 if available_loop > 10 else available_loop\n",
    "            \n",
    "            for loop in reversed(range(0, max_loop, 1)):\n",
    "                \n",
    "                # generate 2 regions\n",
    "                end_s2 = n\n",
    "                start_s2 = end_s2 - length_s2\n",
    "                stem2 = str(seq[start_s2:end_s2])\n",
    "\n",
    "                end_s1 = start_s2 - 1 - loop\n",
    "                start_s1 = end_s1 - length_s1\n",
    "                stem1 = str(Seq(seq[start_s1:end_s1]).reverse_complement())\n",
    "                \n",
    "                # perform semi-global alignment and retrieve score\n",
    "                aln = semi_global_alignment(stem1, stem2, params=HAIRPIN_ALIGN_PARAMS)\n",
    "                score = aln.score\n",
    "\n",
    "                if score > best_score:\n",
    "\n",
    "                    best_score = score\n",
    "                    s1 = (start_s1+offset, end_s1+offset)\n",
    "                    s2 = (start_s2+offset, end_s2+offset)\n",
    "                    loop_size = loop\n",
    "\n",
    "                    if score >= 10:\n",
    "                        hairpin_found = True\n",
    "                        continue # we give a chance to find a better loop\n",
    "\n",
    "            if hairpin_found:\n",
    "                break # we do not test further regions\n",
    "\n",
    "        # after recursive search\n",
    "        read.append(record.id)\n",
    "        scores.append(best_score)\n",
    "        first_stem.append(s1)\n",
    "        second_stem.append(s2)\n",
    "        loop_scores.append(loop_size)     \n",
    "    \n",
    "    result = pd.DataFrame(dict(read=read, HAIRPIN_score=scores, HAIRPIN_stem1=first_stem, HAIRPIN_stem2=second_stem, HAIRPIN_loop=loop_scores))\n",
    "\n",
    "    result.to_csv(output_file, sep='\\t', index=None)\n",
    "                       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60bf379a-4f6d-47ca-afac-477428bdfd68",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def hairpin_search_handler(ID):\n",
    "    \n",
    "    input_file=f'{path}/{ID}/{ID}-five_prime_softclip[SC=FULL|ALN=80].fa'\n",
    "    output_file=f'{path}/{ID}/{ID}-HAIRPIN_search.tsv'\n",
    "                                  \n",
    "    # search hairpins\n",
    "    hairpin_search(input_file, output_file, length_s1=12, length_s2=12, offset_value=80)\n",
    "    \n",
    "    print(f'Completed run {ID}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1abb739e-727a-48bd-ba92-0248b92e43eb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with Pool(4) as p:\n",
    "    p.map(hairpin_search_handler, runs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b1b23a0-2764-4a80-a5b8-65fc38df335b",
   "metadata": {},
   "source": [
    "## 6.b. Processing the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d84885f-1628-4f8f-9dfc-14ada38ce9be",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_list = []\n",
    "for ID in runs:\n",
    "    df = pd.read_csv(f'{path}/{ID}/{ID}-HAIRPIN_search.tsv', sep='\\t')\n",
    "    df_list.append(df)\n",
    "    \n",
    "hairpin_result = pd.concat(df_list)\n",
    "\n",
    "# Merge SL result dataframe with dataset\n",
    "dataset_SSP_SL_hairpin = dataset_SSP_SL.merge(hairpin_result, on='read', how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d5debe8-23a3-4ae4-a653-3ac5a8ea90b8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def hairpin_found(read_orientation, sl_found, hairpin_score):\n",
    "    \n",
    "    # hairpin accepted only on antisense reads for which we could not detect an SL sequence\n",
    "    if read_orientation == 'antisense' and sl_found != 'FOUND' and hairpin_score >= 10:\n",
    "        return 'FOUND' \n",
    "    else:\n",
    "        return 'NOT FOUND'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a58625f7-05ea-4d51-9d9a-fcba1edf728a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Determine if there is an hairpin found or not\n",
    "dataset_SSP_SL_hairpin['HAIRPIN_FOUND'] = dataset_SSP_SL_hairpin.apply(lambda x: hairpin_found(x['read_orientation'], x['ROBUST_SL_FOUND'], x['HAIRPIN_score']), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8de7e8b6-5369-4b6f-97ce-da9b0457c14d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# save\n",
    "dataset_SSP_SL_hairpin.to_csv(f'{path}/dataset_+SSP+SL+HAIRPIN.tsv', sep='\\t', index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "223e8ef4-f9f8-42b7-ae39-f81bb42a8150",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15562e2d-fe5a-4bf0-b6d5-7495b1f69cde",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4dab5a5-f487-4141-9611-a7fbf730a679",
   "metadata": {},
   "source": [
    "# 7. Processing sequence search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20b1901-fcfc-4276-819b-032b4c320168",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute stats per gene per start positions (%SL / %hairpin / %unidentified)\n",
    "\n",
    "rows = []\n",
    "\n",
    "for (gene, position), reads in dataset_SSP_SL_hairpin.groupby(['gene','corrected_genomic_start']):\n",
    "\n",
    "    # Total of reads at a given position\n",
    "    total = len(reads)\n",
    "\n",
    "    position = int(position)\n",
    "    \n",
    "    ###########################\n",
    "    # Measure number of SSP  at the position \n",
    "    # (this is measured independantly from SL/hairpin/unidentified)\n",
    "    ###########################\n",
    "    ssp_reads = len(reads[reads['SSP_FOUND'] == 'FOUND'])\n",
    "    ssp_percent = round(ssp_reads/total*100, 2)\n",
    "    \n",
    "    ###########################\n",
    "    # Measure number of SL / hairpin / unidentified at the position\n",
    "    ###########################\n",
    "    sl_reads = len(reads[reads['ROBUST_SL_FOUND'] == 'FOUND'])\n",
    "    sl_percent = round(sl_reads/total*100, 2)\n",
    "    \n",
    "    hairpin_reads = len(reads[reads['HAIRPIN_FOUND'] == 'FOUND'])\n",
    "    hairpin_percent = round(hairpin_reads/total*100, 2)\n",
    "    \n",
    "    unidentified_reads = total - (sl_reads+hairpin_reads)\n",
    "    unidentified_percent = round(unidentified_reads/total*100, 2)\n",
    "    \n",
    "\n",
    "    ###########################\n",
    "    # Measure ratio SL1 / SL2\n",
    "    ###########################\n",
    "    \n",
    "    # Sense reads from SL1_1 experiment are not used for counting SL1/SL2/Hairpin percentages\n",
    "    evaluated_reads = reads[~((reads['run'] == 'SL1_1') & (reads['read_orientation'] == 'antisense'))]\n",
    "    evaluated = len(evaluated_reads)\n",
    "    \n",
    "    if evaluated > 0:\n",
    "    \n",
    "        # robust variants are used to measure SL1/SL2 ratio \n",
    "        # (-> subgroup of ROBUST SL for which we have a ROBUST VARIANT)\n",
    "        robust_variant = evaluated_reads[evaluated_reads['VARIANT_SL_FOUND'] == 'FOUND']\n",
    "        \n",
    "        # count SL1 and SL2 variants\n",
    "        sl1 = len(robust_variant[robust_variant['SL'].str.contains('SL1')])\n",
    "        sl2 = len(robust_variant[~(robust_variant['SL'].str.contains('SL1'))])\n",
    "        \n",
    "        # measure SL2/SL1 ratio\n",
    "        sl2_ratio = sl2/(sl1+sl2) if (sl1+sl2)>0 else None\n",
    "        \n",
    "    else:\n",
    "        sl1 = 0\n",
    "        sl2 = 0\n",
    "        sl2_ratio = None\n",
    "        \n",
    "    # Put numbers into a dictionnary\n",
    "    row = {'gene':gene, 'position':position, 'total':total, 'SSP':ssp_reads, '%SSP':ssp_percent,\n",
    "           'SL':sl_reads, 'HAIRPIN':hairpin_reads, 'UNIDENTIFIED':unidentified_reads,\n",
    "           '%SL':sl_percent, '%HAIRPIN':hairpin_percent, '%UNIDENTIFIED':unidentified_percent,\n",
    "           'SL1_variants':sl1, 'SL2_variants':sl2, 'SL2_ratio':sl2_ratio }\n",
    "    \n",
    "    rows.append(row)\n",
    "    \n",
    "# create dataframe\n",
    "result = pd.DataFrame(rows)\n",
    "result = result.sort_values(['gene','position'])\n",
    "\n",
    "# save table\n",
    "result.to_csv(f'{path}/start_positions_stats.tsv', sep='\\t', index=None)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "elegans_transplicing",
   "language": "python",
   "name": "elegans_transplicing"
  },
  "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.7.13"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}