10_wrapper.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "output-file: wrapper.html\n",
    "title: lafite_wrapper\n",
    "---\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp lafite_wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "import sys\n",
    "import os\n",
    "import subprocess\n",
    "import argparse\n",
    "import pickle\n",
    "\n",
    "from numpy import percentile\n",
    "from time import strftime\n",
    "\n",
    "from LAFITE.reference_processing import RefProcessWrapper, short_reads_sj_import, cage_tss_import, annotation_reshape, gtf2splicing, split_bed_line, bed_block_to_splicing, read_assignment\n",
    "from LAFITE.preprocessing import read_grouping, polya_signal_import, PolyAFinder\n",
    "from LAFITE.utils import temp_dir_creation, bam2bed, keep_tmp_file\n",
    "from LAFITE.read_collapsing import CoCoWrapper\n",
    "from LAFITE.tailFinder import TailFinderWrapper\n",
    "from LAFITE.refine import RefineWrapper\n",
    "from LAFITE.output import OutputAssembly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exports\n",
    "\n",
    "class LafiteWrapper:\n",
    "    def __init__(self, bam, bedtools, full_cleanup, gtf, genome, min_count_tss_tes, mis_intron_length, min_novel_trans_count,\n",
    "                 min_single_exon_coverage, min_single_exon_len, label, output, polya, polyA_motif_file, relative_abundance_threshold,\n",
    "                 sj_correction_window, short_sj_tab, thread, tss_cutoff, tss_peak, read_assign, assign_known):\n",
    "        self.bam = bam\n",
    "        self.bedtools = bedtools\n",
    "        self.full_cleanup = full_cleanup\n",
    "        self.gtf = gtf\n",
    "        self.genome = genome\n",
    "        self.min_count_tss_tes = min_count_tss_tes\n",
    "        self.mis_intron_length = mis_intron_length\n",
    "        self.min_novel_trans_count = min_novel_trans_count\n",
    "        self.min_single_exon_coverage = min_single_exon_coverage\n",
    "        self.min_single_exon_len = min_single_exon_len\n",
    "        self.label = label\n",
    "        self.output = output\n",
    "        self.polya = polya\n",
    "        self.polyA_motif_file = polyA_motif_file\n",
    "        self.relative_abundance_threshold = relative_abundance_threshold\n",
    "        self.sj_correction_window = sj_correction_window\n",
    "        self.short_sj_tab = short_sj_tab\n",
    "        self.thread = thread\n",
    "        self.tss_cutoff = tss_cutoff\n",
    "        self.tss_peak = tss_peak\n",
    "        self.read_assign = read_assign\n",
    "        self.assign_known = assign_known\n",
    "\n",
    "    def revisit_parameter(self):\n",
    "        \"\"\"\n",
    "        revisit input parameters\"\"\"\n",
    "\n",
    "        sys.stdout.write(f'\\nInput parameters:\\n')\n",
    "        sys.stdout.write(f'Read alignment file: {self.bam}\\n')\n",
    "        sys.stdout.write(f'Reference gene annotation: {self.gtf}\\n')\n",
    "        sys.stdout.write(f'Reference genome annotation: {self.genome}\\n')\n",
    "        if self.polya:\n",
    "            sys.stdout.write(f'Reads Polyadenylation events: {self.polya}\\n')\n",
    "        elif self.polyA_motif_file:\n",
    "            sys.stdout.write(\n",
    "                f'PolyA motif file for read Polyadenylation event estimation: {self.polyA_motif_file}\\n')\n",
    "        else:\n",
    "            raise ValueError(\n",
    "                'Fatal: Please provide either polyA motif file or reads polyadenylation event result\\n')\n",
    "        sys.stdout.write(\n",
    "            f'Edit distance to reference splicing site allowed for splicing correction: {self.sj_correction_window}\\n')\n",
    "        sys.stdout.write(f'Output assembly file: {self.output}\\n')\n",
    "\n",
    "    def run_lafite(self):\n",
    "        \"\"\"\n",
    "        LAFITE wrapper\"\"\"\n",
    "\n",
    "        self.revisit_parameter()\n",
    "        # create temp directory for log and intermediate files\n",
    "        try:\n",
    "            tmp_folder = temp_dir_creation(os.path.dirname(self.output))\n",
    "            tmp_dir = tmp_folder.name\n",
    "        except:\n",
    "            raise ValueError(\n",
    "                'Fatal: Please provide a valid path for output files\\n')\n",
    "\n",
    "        # reference gene annotation processing\n",
    "        sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") +\n",
    "                         ': Preprocessing reference gene annotation\\n')\n",
    "        ref_exon, ref_junction, ref_single_exon_trans, ref_mutple_exon_trans, left_sj_set, right_sj_set, tss_dict = RefProcessWrapper(\n",
    "            self.gtf, self.thread).result_collection()\n",
    "\n",
    "        if self.short_sj_tab:\n",
    "            left_sj_set, right_sj_set = short_reads_sj_import(\n",
    "                self.short_sj_tab, left_sj_set, right_sj_set)\n",
    "\n",
    "        if self.tss_peak:\n",
    "            tss_dict = cage_tss_import(self.tss_peak, tss_dict)\n",
    "\n",
    "        # processing alignment bam file\n",
    "        # convert alignment bam file to bed12 format\n",
    "        sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") +\n",
    "                         ': Preprocessing alignment file\\n')\n",
    "        try:\n",
    "            bam2bed_cmd = bam2bed(self.bam, tmp_dir, self.bedtools)\n",
    "        except:\n",
    "            raise ValueError(\n",
    "                'Fatal: Please provide a valid path for bam file and bedtools\\n')\n",
    "        p = subprocess.run(bam2bed_cmd, shell=True)\n",
    "        if p.returncode == 0:\n",
    "            pass\n",
    "        else:\n",
    "            raise ValueError('Fatal: Error in bam conversion\\n')\n",
    "\n",
    "        # read grouping according to chromosome and strand\n",
    "        outbed = os.path.join(tmp_dir, 'bam.bed')\n",
    "        junction_dict, processed_read = read_grouping(outbed, self.genome)\n",
    "\n",
    "        # polyA info import\n",
    "        if self.polya:\n",
    "            sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") +\n",
    "                             ': Loading polyA information\\n')\n",
    "            polya_dict = polya_signal_import(self.polya)\n",
    "        else:\n",
    "            sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") +\n",
    "                             ': No reads Polyadenylation event provided, detecting from sequence\\n')\n",
    "            polya_dict = PolyAFinder(\n",
    "                processed_read, self.genome, self.polyA_motif_file).polya_estimation()\n",
    "\n",
    "        # read correction and collapsing\n",
    "        sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") +\n",
    "                         ': Collapssing corrected reads\\n')\n",
    "        collected_single_exon_read, collected_multi_exon_read, collected_rss, collected_res = CoCoWrapper(\n",
    "            self.thread, processed_read, ref_exon, ref_junction, ref_single_exon_trans, ref_mutple_exon_trans, left_sj_set, right_sj_set, junction_dict, self.sj_correction_window, polya_dict, self.mis_intron_length, tmp_dir).result_collection()\n",
    "\n",
    "        # identify putative TSS and TES for collapsed reads\n",
    "        processed_collected_multi_exon_read, three_prime_exon = TailFinderWrapper(\n",
    "            collected_multi_exon_read, self.min_count_tss_tes, self.thread).result_collection()\n",
    "\n",
    "        # calculating the tss_cutoff and tes_cutoff:\n",
    "        if not self.tss_cutoff:\n",
    "            self.tss_cutoff = percentile(collected_rss, 75)\n",
    "        tes_cutoff = percentile(collected_res, 70)\n",
    "        print(f'TSS cutoff: {self.tss_cutoff}')\n",
    "        print(f'TES cutoff: {tes_cutoff}')\n",
    "\n",
    "        # identify high quality isoforms from collapsed reads\n",
    "        sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") +\n",
    "                         ': Revisiting the collapsed reads to get high-concensus full-length isoforms\\n')\n",
    "        collected_refined_isoforms = RefineWrapper(processed_collected_multi_exon_read, collected_single_exon_read, ref_mutple_exon_trans, ref_single_exon_trans, three_prime_exon,\n",
    "                                                   tss_dict, self.tss_cutoff, tes_cutoff, self.min_novel_trans_count, self.min_single_exon_coverage, self.min_single_exon_len, self.thread, tmp_dir).result_collection()\n",
    "\n",
    "        # output refined isoforms\n",
    "        OutputAssembly(collected_refined_isoforms, self.output,\n",
    "                       self.label, self.relative_abundance_threshold).write_out()\n",
    "\n",
    "        if self.read_assign:\n",
    "            reshaped_multi_exon_isoform_dict, reshaped_single_exon_isoform_dict, single_exon_isoform_interlap = annotation_reshape(\n",
    "                gtf2splicing(self.output, keepAttribute=True, no_transcript=True))\n",
    "            read_assign_res = read_assignment(os.path.join(tmp_dir, 'Corrected_reads.bed'), reshaped_multi_exon_isoform_dict,\n",
    "                                              reshaped_single_exon_isoform_dict, single_exon_isoform_interlap, only_known=self.assign_known)\n",
    "            with open(self.output.replace('.gtf', 'read_assignment.pkl'), 'wb') as output_pkl:\n",
    "                pickle.dump(read_assign_res, output_pkl)\n",
    "\n",
    "        # keep intermediate results\n",
    "        if self.full_cleanup:\n",
    "            os.system(keep_tmp_file(self.output, tmp_dir))\n",
    "\n",
    "        sys.stdout.write(strftime(\"%Y-%m-%d %H:%M:%S\") + ': All done!\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "\n",
    "def main():\n",
    "\n",
    "    parser = argparse.ArgumentParser(description='Low-abundance Aware Full-length Isoform clusTEr')\n",
    "    parser.add_argument('-b', dest='bam', help='path to the alignment file in bam format', required=True)\n",
    "    parser.add_argument('-B', dest='bedtools', type=str, default='bedtools', help='path to the executable bedtools')\n",
    "    parser.add_argument('-g', dest='gtf', help='path to the reference gene annotation in GTF format', required=True)\n",
    "    parser.add_argument('-f', dest='genome', help='path to the reference genome fasta', required=True)\n",
    "    parser.add_argument('-o', dest='output', help='path to the output file', required=True)\n",
    "    parser.add_argument('-n', dest='min_count_tss_tes', type=int, default=3, help='minimum number of reads supporting a alternative TSS or TES, default: 3')\n",
    "    parser.add_argument('-i', dest='mis_intron_length', type=int, default=150, help='length cutoff for correcting unexpected small intron, default: 150')\n",
    "    parser.add_argument('-c', dest='min_novel_trans_count', type=int, default=3, help='minimum occurrences required for a isoform from novel loci, default: 3')\n",
    "    parser.add_argument('-s', dest='min_single_exon_coverage', type=int, default=4, help='minimum read coverage required for a novel single-exon transcript, default: 4')\n",
    "    parser.add_argument('-l', dest='min_single_exon_len', type=int, default=100, help='minimum length for single-exon transcript, default: 100')\n",
    "    parser.add_argument('-L', dest='label', type=str, default='LAFT', help='name prefix for output transcripts, default: LAFT')\n",
    "    parser.add_argument('-p', dest='polya', help='path to the file contains read Polyadenylation event')\n",
    "    parser.add_argument('-m', dest='polyA_motif_file', help='path to the polya motif file')\n",
    "    parser.add_argument('-r', dest='relative_abundance_threshold', type=int, default=0.01, help='minimum abundance of the predicted multi-exon transcripts as a fraction of the total transcript assembled at a given locus, default: 0.01')\n",
    "    parser.add_argument('-j', dest='short_sj_tab', default=None, help='path to the short read splice junction file')\n",
    "    parser.add_argument('-w', dest='sj_correction_window', type=int, default=40, help='edit distance to reference splicing site for splicing correction, default: 40')\n",
    "    parser.add_argument('--no_full_cleanup', dest='full_cleanup', action='store_true', help='keep all intermediate files')\n",
    "    parser.add_argument('-t', dest='thread', type=int, default=4, help='number of the threads, default: 4')\n",
    "    parser.add_argument('-T', dest='tss_peak', default=None, help='path to the TSS peak file')\n",
    "    parser.add_argument('-d', dest='tss_cutoff', type=int, default=None, help='minimum TSS distance for a transcript to be considered as a novel transcript')\n",
    "    parser.add_argument('--read-assignment', dest='read_assignment',action='store_true',help='output the read assignment')\n",
    "    parser.add_argument('--assign-known', dest='assign_known',action='store_true',help='only assign reads to known transcript')\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    LafiteWrapper(args.bam, args.bedtools, args.full_cleanup, args.gtf, args.genome, args.min_count_tss_tes, \n",
    "                  args.mis_intron_length, args.min_novel_trans_count, args.min_single_exon_coverage,\n",
    "                  args.min_single_exon_len, args.label, args.output, args.polya, args.polyA_motif_file, args.relative_abundance_threshold, args.sj_correction_window, args.short_sj_tab, args.thread, args.tss_cutoff, \n",
    "                  args.tss_peak,args.read_assignment,args.assign_known).run_lafite()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Input parameters:\n",
      "Read alignment file: test/test.bam\n",
      "Reference gene annotation: test/ref.gtf\n",
      "Reference genome annotation: test/ref.fasta\n",
      "PolyA motif file for read Polyadenylation event estimation: test/human_mouse_polyA_motif.txt\n",
      "Edit distance to reference splicing site allowed for splicing correction: 40\n",
      "Output assembly file: test/tmp_reuslt.gtf\n",
      "2023-08-24 14:15:06: Preprocessing reference gene annotation\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-08-24 14:15:06: Preprocessing alignment file\n",
      "2023-08-24 14:15:07: No reads Polyadenylation event provided, detecting from sequence\n",
      "2023-08-24 14:15:08: Collapssing corrected reads\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV3 -:   0%|          | 0/2164 [00:00<?, ?it/s]7032.95it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV3 -: 100%|██████████| 2164/2164 [00:00<00:00, 44250.87it/s]\n",
      "\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV5 -: 100%|██████████| 61/61 [00:00<00:00, 6155.63it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV3 +: 100%|██████████| 4048/4048 [00:00<00:00, 65133.24it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV1 -:  69%|██████▊   | 4367/6362 [00:00<00:00, 43660.78it/s]\n",
      "\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV4 -: 100%|██████████| 3829/3829 [00:00<00:00, 78338.36it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV6 -: 100%|██████████| 466/466 [00:00<00:00, 13518.13it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV7 -:   0%|          | 0/1396 [00:00<?, ?it/s] 40953.20it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV7 -: 100%|██████████| 1396/1396 [00:00<00:00, 51852.61it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV5 +: 100%|██████████| 4483/4483 [00:00<00:00, 35857.09it/s]\n",
      "2023-08-24 14:15:08: Collapsing raw reads from SIRV6 +: 100%|██████████| 11292/11292 [00:00<00:00, 58401.33it/s]\n",
      "2023-08-24 14:15:09: Calculating pupative TSS and TES for collapsed read: 100%|██████████| 1082/1082 [00:06<00:00, 156.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TSS cutoff: 31.0\n",
      "TES cutoff: 16.0\n",
      "2023-08-24 14:15:17: Revisiting the collapsed reads to get high-concensus full-length isoforms\n",
      "2023-08-24 14:15:17: All done!\n"
     ]
    }
   ],
   "source": [
    "output = 'test/tmp_reuslt.gtf'\n",
    "gtf = 'test/ref.gtf'\n",
    "genome = 'test/ref.fasta'\n",
    "bam = 'test/test.bam'\n",
    "full_cleanup=True\n",
    "polyA_motif_file = 'test/human_mouse_polyA_motif.txt'\n",
    "thread = 10\n",
    "sj_correction_window = 40\n",
    "mis_intron_length = 150\n",
    "min_count_tss_tes = 3\n",
    "min_single_exon_coverage = 4\n",
    "min_novel_trans_count = 3\n",
    "min_single_exon_len = 100\n",
    "label = 'LFT'\n",
    "relative_abundance_threshold = 0.01\n",
    "\n",
    "LafiteWrapper(bam, 'bedtools', full_cleanup, gtf, genome, min_count_tss_tes, \n",
    "                  mis_intron_length, min_novel_trans_count, min_single_exon_coverage,\n",
    "                  min_single_exon_len, label, output, None, polyA_motif_file, relative_abundance_threshold, sj_correction_window, None, thread, None, \n",
    "                  None,True,True).run_lafite()"
   ]
  },
  {
   "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": [
    "!nbdev_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}