01_reference_processing.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "output-file: reference_processing.html\n",
    "title: reference_processing\n",
    "\n",
    "---\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp reference_processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "from collections import defaultdict\n",
    "from interlap import InterLap\n",
    "from joblib import Parallel, delayed\n",
    "from LAFITE.utils import Vividict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def gtf_line_split(entry):\n",
    "    \"\"\"split gtf line\n",
    "    \"\"\"\n",
    "    lst = entry.rstrip().split('\\t')\n",
    "    chrom, source, feature, start, end, score, strand, frame, attribute = lst\n",
    "    start = int(start)\n",
    "    end = int(end)\n",
    "\n",
    "    return chrom, source, feature, start, end, score, strand, frame, attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def gtf2splicing(gtf, transcript_ref='transcript_id', keepAttribute=False, no_transcript=False):\n",
    "    \"\"\"\n",
    "    preprocess the gtf file to gene and isoform level\n",
    "    \"\"\"\n",
    "    isoform_structure_dict = Vividict()\n",
    "    with open(gtf) as f:\n",
    "        for line in f:\n",
    "            if not line.startswith('#'):\n",
    "                chrom, source, feature, start, end, score, strand, frame, attributes = gtf_line_split(\n",
    "                    line)\n",
    "\n",
    "                try:\n",
    "                    attributes = [a.strip().replace('\"', '') for a in attributes.strip(\n",
    "                        ';').split('\"; ') if len(a) > 0]\n",
    "                    attributes = dict([a.split(' ', 1)[0:2]\n",
    "                                      for a in attributes])\n",
    "\n",
    "                except:\n",
    "                    raise ValueError(\"Fatal: please check input GTF format\\n\")\n",
    "\n",
    "                if keepAttribute:\n",
    "                    if no_transcript:\n",
    "                        # this is only for uncompleted gtf file\n",
    "                        transcript_id = attributes[transcript_ref]\n",
    "                        if transcript_id not in isoform_structure_dict[(chrom, strand)]:\n",
    "                            isoform_structure_dict[(chrom, strand)][transcript_id] = [\n",
    "                                attributes]\n",
    "\n",
    "                    elif feature in ('transcript', 'mRNA'):\n",
    "                        transcript_id = attributes[transcript_ref]\n",
    "                        isoform_structure_dict[(chrom, strand)][transcript_id] = [\n",
    "                            attributes]\n",
    "\n",
    "                if feature == \"exon\":\n",
    "                    transcript_id = attributes[transcript_ref]\n",
    "                    if transcript_id in isoform_structure_dict[(chrom, strand)]:\n",
    "                        isoform_structure_dict[(chrom, strand)][transcript_id].extend(\n",
    "                            [start, end])\n",
    "                    else:\n",
    "                        isoform_structure_dict[(chrom, strand)][transcript_id] = [\n",
    "                            start, end]\n",
    "\n",
    "    return isoform_structure_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "class RefAnnotationExtraction:\n",
    "    \"\"\"\n",
    "    extraction the splicing/exon information from the reference annotation\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, chrom, strand, chrand_ref_trans_structure_dict):\n",
    "        self.chrom = chrom\n",
    "        self.strand = strand\n",
    "        self.chrand_ref_trans_structure_dict = chrand_ref_trans_structure_dict\n",
    "\n",
    "    def annotation_extraction(self):\n",
    "\n",
    "        # reference exon start and end sites, the key will be chromosome, strand\n",
    "        chrand_ref_exon = set()\n",
    "        chrand_ref_junction = set()  # splicing junction from reference annotation\n",
    "        # splicing list for every chromosome, e.g. {'chr1,+':[[1,2,3,4,5],[111,223,44]], 'chr1,-':[[xxx,xxx,xxx],[xxx,xx,xxxx]]}\n",
    "        chrand_ref_mutple_exon_trans = defaultdict(dict)\n",
    "        # splicing list for every chromosome, but only for single exon transcript\n",
    "        chrand_ref_single_exon_trans = set()\n",
    "        chrand_left_sj_set = set()\n",
    "        chrand_right_sj_set = set()\n",
    "        chrand_tss_dict = set()  # transcript start sites set\n",
    "\n",
    "        for isoform, full_block in self.chrand_ref_trans_structure_dict.items():\n",
    "            full_block.sort()\n",
    "            iso_splicing = full_block[1:-1]\n",
    "            start, end = full_block[0], full_block[-1]\n",
    "\n",
    "            # record the TSS\n",
    "            chrand_tss_dict.add(\n",
    "                start) if self.strand == '+' else chrand_tss_dict.add(end)\n",
    "\n",
    "            if not iso_splicing:\n",
    "                # record single exon trans start and end site\n",
    "                chrand_ref_single_exon_trans.add((start, end, isoform))\n",
    "            else:\n",
    "                # record the splicing site, splicing junction, exon for multi exon trans\n",
    "                full_block = iter(full_block)\n",
    "                for idx, exon in enumerate(zip(full_block, full_block)):\n",
    "                    chrand_ref_exon.add(exon)\n",
    "                    if idx > 0:\n",
    "                        right_sj = exon[0]\n",
    "                        chrand_left_sj_set.add(left_sj)\n",
    "                        chrand_right_sj_set.add(right_sj)\n",
    "                        chrand_ref_junction.add((left_sj, right_sj))\n",
    "                    left_sj = exon[1]\n",
    "\n",
    "                # record the splicing structure and the tss and tes for multi exon trans\n",
    "                if self.strand == '-':\n",
    "                    se_site = [end, start, isoform]\n",
    "                else:\n",
    "                    se_site = [start, end, isoform]\n",
    "                chrand_ref_mutple_exon_trans[tuple(iso_splicing)] = se_site\n",
    "\n",
    "        return chrand_ref_exon, chrand_ref_junction, chrand_ref_single_exon_trans, chrand_ref_mutple_exon_trans, chrand_left_sj_set, chrand_right_sj_set, chrand_tss_dict\n",
    "\n",
    "    def annotation_sorting(self):\n",
    "\n",
    "        chrand_ref_exon, chrand_ref_junction, chrand_ref_single_exon_trans, chrand_ref_mutple_exon_trans, chrand_left_sj_set, chrand_right_sj_set, chrand_tss_dict = self.annotation_extraction()\n",
    "\n",
    "        chrand_tss_dict = sorted(chrand_tss_dict)\n",
    "        # sort multi exon transcript by splicing junction number\n",
    "        chrand_ref_mutple_exon_trans = dict(sorted(\n",
    "            chrand_ref_mutple_exon_trans.items(), key=lambda d: len(d[0]), reverse=True))\n",
    "\n",
    "        chrand_left_sj_set = sorted(chrand_left_sj_set)\n",
    "        chrand_right_sj_set = sorted(chrand_right_sj_set)\n",
    "\n",
    "        # covert to interlap data format\n",
    "        if chrand_ref_single_exon_trans:\n",
    "            t = list(chrand_ref_single_exon_trans)\n",
    "            chrand_ref_single_exon_trans = InterLap()\n",
    "            chrand_ref_single_exon_trans.update(t)\n",
    "\n",
    "        if chrand_ref_exon:\n",
    "            t = list(chrand_ref_exon)\n",
    "            chrand_ref_exon = InterLap()\n",
    "            chrand_ref_exon.update(t)\n",
    "\n",
    "        return self.chrom, self.strand, chrand_ref_exon, chrand_ref_junction, chrand_ref_single_exon_trans, chrand_ref_mutple_exon_trans, chrand_left_sj_set, chrand_right_sj_set, chrand_tss_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "class RefProcessWrapper:\n",
    "    def __init__(self, ref_gtf, thread):\n",
    "        self.ref_gtf = ref_gtf\n",
    "        self.thread = thread\n",
    "\n",
    "    def process(self):\n",
    "        preprocess_lst = []\n",
    "        ref_trans_structure_dict = gtf2splicing(self.ref_gtf)\n",
    "        for (chrom, strand), chrand_ref_trans_structure_dict in ref_trans_structure_dict.items():\n",
    "            preprocess_lst.append(RefAnnotationExtraction(\n",
    "                chrom, strand, chrand_ref_trans_structure_dict))\n",
    "        with Parallel(n_jobs=self.thread) as parallel:\n",
    "            results = parallel(delayed(lambda x: x.annotation_sorting())(\n",
    "                job) for job in preprocess_lst)\n",
    "\n",
    "        return results\n",
    "\n",
    "    def result_collection(self):\n",
    "        results = self.process()\n",
    "        # reference exon start and end sites, the key will be chromosome, strand\n",
    "        ref_exon = defaultdict(dict)\n",
    "        # splicing junction from reference annotation\n",
    "        ref_junction = defaultdict(dict)\n",
    "        # splicing list for every chromosome, e.g. {'chr1,+':[[1,2,3,4,5],[111,223,44]], 'chr1,-':[[xxx,xxx,xxx],[xxx,xx,xxxx]]}\n",
    "        ref_mutple_exon_trans = Vividict()\n",
    "        # splicing list for every chromosome, but only for single exon transcript\n",
    "        ref_single_exon_trans = defaultdict(dict)\n",
    "        left_sj_set = defaultdict(dict)\n",
    "        right_sj_set = defaultdict(dict)\n",
    "        tss_dict = defaultdict(dict)  # transcript start sites set\n",
    "\n",
    "        for result in results:\n",
    "            chrom, strand, chrand_ref_exon, chrand_ref_junction, chrand_ref_single_exon_trans, chrand_ref_mutple_exon_trans, chrand_left_sj_set, chrand_right_sj_set, chrand_tss_dict = result\n",
    "\n",
    "            tss_dict[(chrom, strand)] = chrand_tss_dict\n",
    "            if chrand_ref_exon:\n",
    "                ref_exon[(chrom, strand)] = chrand_ref_exon\n",
    "            if chrand_ref_junction:\n",
    "                ref_junction[(chrom, strand)] = chrand_ref_junction\n",
    "            if chrand_ref_mutple_exon_trans:\n",
    "                ref_mutple_exon_trans[(chrom, strand)\n",
    "                                      ] = chrand_ref_mutple_exon_trans\n",
    "            if chrand_ref_single_exon_trans:\n",
    "                ref_single_exon_trans[(chrom, strand)\n",
    "                                      ] = chrand_ref_single_exon_trans\n",
    "            if chrand_left_sj_set:\n",
    "                left_sj_set[(chrom, strand)] = chrand_left_sj_set\n",
    "            if chrand_right_sj_set:\n",
    "                right_sj_set[(chrom, strand)] = chrand_right_sj_set\n",
    "\n",
    "        return ref_exon, ref_junction, ref_single_exon_trans, ref_mutple_exon_trans, left_sj_set, right_sj_set, tss_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def short_reads_sj_import(sj_tab, left_sj_set, right_sj_set):\n",
    "    \"\"\"import the splicing junctions detected from short reads data (STAR SJ_tab)\n",
    "    \"\"\"\n",
    "    for file in sj_tab:\n",
    "        with open(file) as f:\n",
    "            for line in f:\n",
    "                line = line.strip('\\n').split(\"\\t\")\n",
    "                chrom = line[0]\n",
    "                if line[3] in ['+', '1']:\n",
    "                    strand = '+'\n",
    "                else:\n",
    "                    strand = '-'\n",
    "                left_sj = int(line[1])-1\n",
    "                right_sj = int(line[2])+1\n",
    "                left_sj_set[(chrom, strand)].add(left_sj)\n",
    "                right_sj_set[(chrom, strand)].add(right_sj)\n",
    "\n",
    "    for i in left_sj_set:\n",
    "        left_sj_set[i] = sorted(left_sj_set[i])\n",
    "    for i in right_sj_set:\n",
    "        right_sj_set[i] = sorted(right_sj_set[i])\n",
    "\n",
    "    return left_sj_set, right_sj_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def cage_tss_import(cage_tss, tss_dict):\n",
    "    \"\"\"import the splicing junctions detected from short reads data (STAR SJ_tab)\n",
    "    \"\"\"\n",
    "    for file in cage_tss:\n",
    "        with open(file) as f:\n",
    "            for line in f:\n",
    "                line = line.strip('\\n').split(\"\\t\")\n",
    "                chrom, strand = line[0], line[3]\n",
    "                tss_center = round(abs(int(line[2])-int(line[1]))/2 + 1)\n",
    "                if tss_center not in tss_dict[(chrom, strand)]:\n",
    "                    tss_dict[(chrom, strand)].add(tss_center)\n",
    "\n",
    "    for i in tss_dict:\n",
    "        tss_dict[i] = sorted(tss_dict[i])\n",
    "\n",
    "    return tss_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "\n",
    "def annotation_reshape(isoform_structure_dict):\n",
    "    reshaped_multi_exon_isoform_dict = Vividict()\n",
    "    reshaped_single_exon_isoform_dict = Vividict()\n",
    "    single_exon_isoform_interlap = defaultdict(set)\n",
    "\n",
    "    for chr_str, transcript in isoform_structure_dict.items():\n",
    "        chrom, strand = chr_str\n",
    "        for transcript_id, isoform_detail in transcript.items():\n",
    "            attributes, *isoform_structure = isoform_detail\n",
    "            isoform_structure.sort()\n",
    "            splicing = isoform_structure[1:-1]\n",
    "            if splicing:\n",
    "                reshaped_multi_exon_isoform_dict[(chrom, strand)][tuple(splicing)] = [\n",
    "                    transcript_id, isoform_structure[0], isoform_structure[1], attributes]\n",
    "            else:\n",
    "                single_exon_isoform_interlap[(chrom, strand)].add(\n",
    "                    tuple(isoform_structure))\n",
    "                reshaped_single_exon_isoform_dict[(chrom, strand)][tuple(isoform_structure)] = [\n",
    "                    transcript_id, isoform_structure[0], isoform_structure[1], attributes]\n",
    "\n",
    "    # convert to interlap\n",
    "    for i in single_exon_isoform_interlap:\n",
    "        t = list(single_exon_isoform_interlap[i])\n",
    "        single_exon_isoform_interlap[i] = InterLap()\n",
    "        single_exon_isoform_interlap[i].update(t)\n",
    "    return reshaped_multi_exon_isoform_dict, reshaped_single_exon_isoform_dict, single_exon_isoform_interlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def split_bed_line(bed_line, extends=False):\n",
    "    \"\"\"\n",
    "    split bed line\n",
    "    \"\"\"\n",
    "    cells = bed_line.strip('\\n').split(\"\\t\")\n",
    "    chrom = start = end = name = score = strand = thick_start = thick_end = item_rgb = block_count = block_sizes = block_starts = others = None\n",
    "    if len(cells) >= 3:\n",
    "        chrom = cells[0]\n",
    "        start = int(cells[1])\n",
    "        end = int(cells[2])\n",
    "    if len(cells) >= 6:\n",
    "        name = cells[3]\n",
    "        score = cells[4]\n",
    "        strand = cells[5]\n",
    "    if len(cells) >= 12:\n",
    "        thick_start = int(cells[6])\n",
    "        thick_end = int(cells[7])\n",
    "        item_rgb = cells[8]\n",
    "        block_count = int(cells[9])\n",
    "        block_sizes = [int(i) for i in cells[10].rstrip(',').split(',')]\n",
    "        block_starts = [int(i) for i in cells[11].rstrip(',').split(',')]\n",
    "        if extends:\n",
    "            others = cells[12:]\n",
    "            return chrom, start, end, name, score, strand, thick_start, thick_end, item_rgb, block_count, block_sizes, block_starts, others\n",
    "    return chrom, start, end, name, score, strand, thick_start, thick_end, item_rgb, block_count, block_sizes, block_starts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def bed_block_to_splicing(start, block_count, block_starts, block_sizes):\n",
    "    read_splicing = []\n",
    "    if block_count > 1:\n",
    "        for i in range(block_count-1):\n",
    "            left_sj = start + block_starts[i] + block_sizes[i]\n",
    "            right_sj = start + block_starts[i+1] + 1\n",
    "            read_splicing.extend([left_sj, right_sj])\n",
    "    return read_splicing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "def read_assignment(bed, reshaped_multi_exon_isoform_dict, reshaped_single_exon_isoform_dict, single_exon_isoform_interlap, only_known=True, ratio=0.8):\n",
    "    cp_read_dict = defaultdict(list)\n",
    "    with open(bed) as f:\n",
    "        for line in f:\n",
    "            chrom, start, end, name, score, strand, thick_start, thick_end, item_rgb, block_count, block_sizes, block_starts = split_bed_line(\n",
    "                line)\n",
    "            if block_count > 1:\n",
    "                read_splicing = tuple(bed_block_to_splicing(\n",
    "                    start, block_count, block_starts, block_sizes))\n",
    "                if read_splicing in reshaped_multi_exon_isoform_dict[(chrom, strand)]:\n",
    "                    attributes = reshaped_multi_exon_isoform_dict[(\n",
    "                        chrom, strand)][read_splicing][3]\n",
    "                    if attributes['isoform_class'] == 'Reference':\n",
    "                        if 'reference_id' in attributes:\n",
    "                            transcript = attributes['reference_id']\n",
    "                            cp_read_dict[transcript].append(name)\n",
    "                        # else:\n",
    "                        # \ttranscript = reshaped_multi_exon_isoform_dict[(chrom, strand)][read_splicing][0]\n",
    "                        # \tprint(f'ERROR {transcript}')\n",
    "\n",
    "                    elif not only_known:\n",
    "                        transcript = reshaped_multi_exon_isoform_dict[(\n",
    "                            chrom, strand)][read_splicing][0]\n",
    "                        cp_read_dict[transcript].append(name)\n",
    "\n",
    "            elif single_exon_isoform_interlap[(chrom, strand)]:\n",
    "                overlap_gene = list(single_exon_isoform_interlap[(\n",
    "                    chrom, strand)].find((start+1, end)))\n",
    "                if overlap_gene:\n",
    "                    for i in overlap_gene:\n",
    "\n",
    "                        attributes = reshaped_single_exon_isoform_dict[(\n",
    "                            chrom, strand)][(i[0], i[1])][3]\n",
    "                        # transcript = reshaped_single_exon_isoform_dict[(chrom, strand)][(i[0],i[1])][0]\n",
    "                        inter = list(\n",
    "                            set(range(start, end+1)).intersection(range(i[0], i[1]+1)))\n",
    "                        inter.sort()\n",
    "                        if (inter[-1]-inter[0]+1)/(i[1]-i[0]+1) > ratio:\n",
    "                            if attributes['isoform_class'] == 'Reference':\n",
    "                                if 'reference_id' in attributes:\n",
    "                                    transcript = attributes['reference_id']\n",
    "                                    cp_read_dict[transcript].append(name)\n",
    "                            elif not only_known:\n",
    "\n",
    "                                transcript = reshaped_single_exon_isoform_dict[(\n",
    "                                    chrom, strand)][(i[0], i[1])][0]\n",
    "                                cp_read_dict[transcript].append(name)\n",
    "\n",
    "    return cp_read_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | hide\n",
    "\n",
    "from nbdev import nbdev_export\n",
    "nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}