flu_reference_mapping_template.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Project: <project_name>\n",
    "## sample subset: <primer_set>\n",
    "## analysis date: <analysis_date>\n",
    "## user: <name>\n",
    "\n",
    "### Description\n",
    "<High-level description of samples, primers, and references used>\n",
    "\n",
    "### General steps\n",
    "- prep fastq files\n",
    "- fastqc on raw fastq files\n",
    "- clean fastq files using seqyclean\n",
    "- fastqc on filtered fastq files\n",
    "- align using bwa\n",
    "- trim primers using samtools ampliconclip\n",
    "- generate consensus sequence with ivar consensus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import re\n",
    "import subprocess\n",
    "\n",
    "import pandas as pd\n",
    "from Bio import SeqIO\n",
    "\n",
    "pd.set_option('display.max_rows', 1000)\n",
    "\n",
    "DOCKER_BASE = 'docker run --rm -v ${PWD}:/data -u $(id -u):$(id -g)'\n",
    "FASTQC_DOCKER = f'{DOCKER_BASE} staphb/fastqc:0.12.1'\n",
    "MULTIQC_DOCKER = f'{DOCKER_BASE} staphb/multiqc:1.8'\n",
    "SEQYCLEAN_DOCKER = f'{DOCKER_BASE} staphb/seqyclean:1.10.09'\n",
    "BWA_DOCKER = f'{DOCKER_BASE} staphb/bwa:0.7.17'\n",
    "SAMTOOLS_DOCKER = f'{DOCKER_BASE} staphb/samtools:1.20'\n",
    "IVAR_DOCKER = f'{DOCKER_BASE} staphb/ivar:1.4.2'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General Set Up\n",
    "\n",
    "Edit the below input variables as needed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# working directory absolute path \n",
    "wd = ''  # e.g. /home/sam_baird/flu/h5/test_h5_0005_nextseq/avrl_h5n1_250bp/\n",
    "\n",
    "project_name = '' # e.g. test_h5_0005_nextseq\n",
    "\n",
    "### INPUT FILES ###\n",
    "# just provide the filenames, no paths\n",
    "\n",
    "# if multiple primer sets in single project, use subworkbook split by primer_set\n",
    "workbook = '' # e.g. test_h5_0005_nextseq_workbook_avrl_h5n1_250bp.tsv\n",
    "\n",
    "seqyclean_contaminants_fasta = ''\n",
    "\n",
    "primer_bed = ''\n",
    "\n",
    "# keys: short name (used for naming subdirectories and files)\n",
    "# values: filename\n",
    "# references can be either single gene segments or multifastas\n",
    "reference_sequences = {  # e.g. 'h5n1_bovine': 'A_Bovine_texas_029328-01_UtoT.fasta'\n",
    "    '' : '',\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up directories and paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "os.chdir(wd)\n",
    "refs_dir = 'references'\n",
    "if not os.path.exists(refs_dir):\n",
    "    os.mkdir(refs_dir)\n",
    "reference_sequence_paths = {k: os.path.join(refs_dir, v) for k, v in reference_sequences.items()}\n",
    "\n",
    "fastq_dir = 'fastq_files'\n",
    "if not os.path.exists(fastq_dir):\n",
    "    os.mkdir(fastq_dir)\n",
    "\n",
    "raw_alignment_dir = 'bwa_alignment'\n",
    "trimmed_alignment_dir = 'bwa_alignment_trimmed'\n",
    "primer_bed_path = os.path.join(refs_dir, primer_bed)\n",
    "fastqc_raw_dir = 'fastqc_raw'\n",
    "multiqc_raw_dir = 'multiqc_raw'\n",
    "multiqc_clean_dir = 'multiqc_clean'\n",
    "seqyclean_dir = 'seqyclean'\n",
    "fastqc_clean_dir = 'fastqc_clean'\n",
    "alignment_dir = 'bwa_alignment'\n",
    "aln_metrics_dir = 'alignment_metrics'\n",
    "consensus_dir = 'consensus_sequences'\n",
    "amplicon_dir = 'amplicon_metrics'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Manual steps needed**: In the working directory, add the:\n",
    "- workbook\n",
    "- FASTQ files to the `fastq_files` directory\n",
    "- reference genomes to the `references` directory\n",
    "- seqyclean contaminants FASTA to the `references` directory\n",
    "- primer BED file to the `references` directory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read workbook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "path_wb = os.path.join(wd, workbook)\n",
    "wb = pd.read_csv(path_wb, sep = '\\t')\n",
    "\n",
    "sample_names = wb.sample_name.tolist()\n",
    "wb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FastQC on raw FASTQ files\n",
    "Assess the quality of the raw FASTQ files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(fastqc_raw_dir):\n",
    "    os.mkdir(fastqc_raw_dir)\n",
    "\n",
    "for sample in sample_names:\n",
    "    print(sample)\n",
    "    \n",
    "    fastq_1 = f'{sample}_R1.fastq.gz'\n",
    "    fastq_2 = f'{sample}_R2.fastq.gz'\n",
    "\n",
    "    outdir = os.path.join(fastqc_raw_dir, sample)\n",
    "    if not os.path.exists(outdir):\n",
    "        os.mkdir(outdir)\n",
    "        \n",
    "    !{FASTQC_DOCKER} fastqc --outdir fastqc_raw/{sample} fastq_files/{fastq_1}\n",
    "    !{FASTQC_DOCKER} fastqc --outdir fastqc_raw/{sample} fastq_files/{fastq_2}\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MultiQC on raw FASTQ files\n",
    "Combine FastQC results into a single report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(multiqc_raw_dir):\n",
    "    os.mkdir(multiqc_raw_dir)\n",
    "\n",
    "!{MULTIQC_DOCKER} multiqc fastqc_raw --outdir {multiqc_raw_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CSV report on raw FASTQ files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_reads(sample_names, fastqc_dir):\n",
    "    for sample in sample_names:\n",
    "        print(sample)\n",
    "        # get fastqc path for zip files\n",
    "        zip1 = glob.glob(os.path.join(fastqc_dir, sample, f'{sample}*1_fastqc.zip'))[0]\n",
    "        zip2 = glob.glob(os.path.join(fastqc_dir, sample, f'{sample}*2_fastqc.zip'))[0]\n",
    "        \n",
    "        # set unzip outpath\n",
    "        unzip_outpath = os.path.join(fastqc_dir, sample)\n",
    "        \n",
    "        # unzip\n",
    "        !unzip -q -o {zip1} -d {unzip_outpath}\n",
    "        !unzip -q -o {zip2} -d {unzip_outpath}\n",
    "        \n",
    "        # read in file and create summary out file\n",
    "        # set up dataframe\n",
    "        df = pd.DataFrame()\n",
    "        df['sample_name'] = [sample]\n",
    "        \n",
    "        # R1\n",
    "        path = glob.glob(os.path.join(unzip_outpath, f'{sample}*_fastqc', 'fastqc_data.txt'))[0]\n",
    "        \n",
    "        command = f'cat {path} | grep \"Total Sequences\" | cut -f 2'\n",
    "        total_seqs = subprocess.check_output(command, shell=True, text=True).rstrip('\\n')\n",
    "        df['r1_total_reads'] = [total_seqs]\n",
    "\n",
    "        command = f'cat {path} | grep \"Sequences flagged as poor quality\" | cut -f 2'\n",
    "        flagged_seqs = subprocess.check_output(command, shell=True, text=True).rstrip('\\n')\n",
    "        df['r1_flagged_reads_as_poor_quality'] = [flagged_seqs]\n",
    "\n",
    "        command = f'cat {path} | grep \"Sequence length\" | cut -f 2'\n",
    "        seq_len = subprocess.check_output(command, shell=True, text=True).rstrip('\\n')\n",
    "        df['r1_read_len'] = [seq_len]\n",
    "        \n",
    "        # R2\n",
    "        path = glob.glob(os.path.join(unzip_outpath, f'{sample}*2_fastqc', 'fastqc_data.txt'))[0]\n",
    "        \n",
    "        command = f'cat {path} | grep \"Total Sequences\" | cut -f 2'\n",
    "        total_seqs = subprocess.check_output(command, shell = True, text = True).rstrip('\\n')\n",
    "        df['r2_total_reads'] = [total_seqs]\n",
    "\n",
    "        command = f'cat {path} | grep \"Sequences flagged as poor quality\" | cut -f 2'\n",
    "        flagged_seqs = subprocess.check_output(command, shell = True, text = True).rstrip('\\n')\n",
    "        df['r2_flagged_reads_as_poor_quality'] = [flagged_seqs]\n",
    "\n",
    "        command = f'cat {path} | grep \"Sequence length\" | cut -f 2'\n",
    "        seq_len = subprocess.check_output(command, shell = True, text = True).rstrip('\\n')\n",
    "        df['r2_read_len'] = [seq_len]\n",
    "        \n",
    "        outfile = os.path.join(unzip_outpath, f'{sample}_summary_metrics.tsv')\n",
    "        df.to_csv(outfile, sep = '\\t', index = False)\n",
    "\n",
    "summarize_reads(sample_names, fastqc_raw_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seqyclean on raw FASTQ files\n",
    "Filter by quality and remove known contaminants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(seqyclean_dir):\n",
    "    os.mkdir(seqyclean_dir)\n",
    "\n",
    "seqyclean_commands = []\n",
    "for sample in sample_names:    \n",
    "    fastq_1 = f'fastq_files/{sample}_R1.fastq.gz'\n",
    "    fastq_2 = f'fastq_files/{sample}_R2.fastq.gz'         \n",
    "    out_name = f'seqyclean/{sample}_clean'\n",
    "\n",
    "    !{SEQYCLEAN_DOCKER} seqyclean -minlen 25 -qual 30 30 -gz -1 {fastq_1} -2 {fastq_2} -c references/Adapters_plus_PhiX_174.fasta -o {out_name}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FastQC on cleaned FASTQ files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(fastqc_clean_dir):\n",
    "    os.mkdir(fastqc_clean_dir)\n",
    "\n",
    "for sample in sample_names:\n",
    "    print(sample)\n",
    "    \n",
    "    fastq_1 = f'{sample}_clean_PE1.fastq.gz'\n",
    "    fastq_2 = f'{sample}_clean_PE2.fastq.gz'\n",
    "\n",
    "    outdir = os.path.join(fastqc_clean_dir, sample)\n",
    "    if not os.path.exists(outdir):\n",
    "        os.mkdir(outdir)\n",
    "                \n",
    "    !{FASTQC_DOCKER} fastqc --outdir fastqc_clean/{sample} seqyclean/{fastq_1}\n",
    "    !{FASTQC_DOCKER} fastqc --outdir fastqc_clean/{sample} seqyclean/{fastq_2}\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MultiQC on cleaned FASTQ files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(multiqc_clean_dir):\n",
    "    os.mkdir(multiqc_clean_dir)\n",
    "\n",
    "!{MULTIQC_DOCKER} multiqc fastqc_clean --outdir {multiqc_clean_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CSV report on raw FASTQ files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_reads_summary = summarize_reads(sample_names, fastqc_clean_dir)\n",
    "clean_reads_summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get summary metrics on reads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# combine fastqc sumary from raw and filtered data\n",
    "df_raw_list = []\n",
    "df_filtered_list = []\n",
    "    \n",
    "    \n",
    "for sample in sample_names:    \n",
    "    # first read in raw summary file\n",
    "    path = os.path.join(wd, 'fastqc_raw', sample, f'{sample}_summary_metrics.tsv')\n",
    "    df = pd.read_csv(path, sep = '\\t')\n",
    "    \n",
    "    df_raw_list.append(df)\n",
    "    \n",
    "    # now read in filtered summary file\n",
    "    path = os.path.join(wd, 'fastqc_clean', sample, f'{sample}_summary_metrics.tsv')\n",
    "    df = pd.read_csv(path, sep = '\\t')\n",
    "    \n",
    "    df_filtered_list.append(df)\n",
    "    \n",
    "    \n",
    "# concatenate dataframes and join\n",
    "raw_df = pd.concat(df_raw_list).reset_index(drop = True)\n",
    "raw_df = raw_df.set_index('sample_name')\n",
    "\n",
    "filtered_df = pd.concat(df_filtered_list).reset_index(drop = True)\n",
    "filtered_df = filtered_df.set_index('sample_name')\n",
    "\n",
    "\n",
    "combined_df = raw_df.join(filtered_df, how = 'left', lsuffix = '_raw', rsuffix = '_filtered')\n",
    "combined_df = combined_df.reset_index()\n",
    "\n",
    "# calculate total reads diff\n",
    "combined_df['total_reads_diff'] = combined_df['r1_total_reads_raw'] - combined_df['r1_total_reads_filtered']\n",
    "combined_df['raw_total_reads_paired'] = combined_df['r1_total_reads_raw']\n",
    "combined_df['filtered_total_reads_paired'] = combined_df['r1_total_reads_filtered']\n",
    "\n",
    "# add primer set\n",
    "for row in range(combined_df.shape[0]):\n",
    "    sample = combined_df.sample_name[row]\n",
    "\n",
    "# print(j.columns)\n",
    "\n",
    "col_order = ['sample_name', 'raw_total_reads_paired', 'filtered_total_reads_paired',\n",
    "             'total_reads_diff', 'r1_total_reads_raw', 'r1_flagged_reads_as_poor_quality_raw',\n",
    "             'r1_read_len_raw', 'r2_total_reads_raw', 'r2_flagged_reads_as_poor_quality_raw',\n",
    "             'r2_read_len_raw', 'r1_total_reads_filtered', 'r1_flagged_reads_as_poor_quality_filtered',\n",
    "             'r1_read_len_filtered', 'r2_total_reads_filtered',\n",
    "             'r2_flagged_reads_as_poor_quality_filtered', 'r2_read_len_filtered']\n",
    "\n",
    "combined_df = combined_df[col_order]\n",
    "combined_df['project_name'] = project_name\n",
    "\n",
    "outfile = os.path.join(wd, f'{project_name}_reads_QC_summary.csv')\n",
    "combined_df.to_csv(outfile, index = False)\n",
    "combined_df\n",
    "\n",
    "combined_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Align to each reference using BWA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(alignment_dir):\n",
    "    os.mkdir(alignment_dir)\n",
    "\n",
    "for sample in sample_names:\n",
    "    # align each sample to each reference\n",
    "    for ref in reference_sequence_paths:\n",
    "\n",
    "        fastq_1 = os.path.join(seqyclean_dir, f'{sample}_clean_PE1.fastq.gz')\n",
    "        fastq_2 = os.path.join(seqyclean_dir, f'{sample}_clean_PE2.fastq.gz')\n",
    "\n",
    "        ref_base_name = ref.split('/')[-1]\n",
    "\n",
    "        # each reference gets its own subdirectory\n",
    "        out_subdir = os.path.join(alignment_dir, ref_base_name)\n",
    "        if not os.path.exists(out_subdir):\n",
    "            os.mkdir(out_subdir)\n",
    "\n",
    "        outsam = os.path.join(out_subdir, f'{sample}.sam')\n",
    "        outbam = os.path.join(out_subdir, f'{sample}.aln.sorted.bam')\n",
    "\n",
    "        !{BWA_DOCKER} bwa index -p {ref_base_name} -a is {reference_sequence_paths[ref]}\n",
    "        !{BWA_DOCKER} bwa mem -t 6 {ref_base_name} {fastq_1} {fastq_2} -f {outsam}\n",
    "        !{SAMTOOLS_DOCKER} samtools view -bS {outsam} | samtools sort -o {outbam}\n",
    "        !rm {outsam}  # save storage space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a sorted primer BED file for samtools ampliconclip/ampliconstats "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "bed_df = pd.read_csv(primer_bed_path, sep='\\t', header=None)\n",
    "\n",
    "def extract_amplicon_name(primer_name):\n",
    "    last_separator = max(primer_name.rfind('-'), primer_name.rfind('_'))\n",
    "    if last_separator != -1:\n",
    "        amplicon = primer_name[:last_separator]\n",
    "    else:\n",
    "        raise ValueError(f'Could not parse amplicon name of primer {primer_name}')\n",
    "    return amplicon\n",
    "\n",
    "bed_df['amplicon'] = bed_df[3].apply(extract_amplicon_name)\n",
    "\n",
    "# amplicon number assumed to be last string of digits in amplicon_name\n",
    "bed_df['amplicon_number'] = bed_df[3].str.extract(r'(\\d+)(?!.*\\d)', expand=False)\n",
    "bed_df['amplicon_number'] = bed_df['amplicon_number'].astype(int)\n",
    "\n",
    "bed_df = bed_df.sort_values([0, 'amplicon_number', 5])\n",
    "\n",
    "# check that the BAM reference name(s) can be matched to reference(s) in the BED\n",
    "# In the WDL, we probably want to either error out or not output a BAM in trim step\n",
    "bed_references = set(bed_df[0])\n",
    "for bam in glob.glob(f'{alignment_dir}/*/*'):\n",
    "    bam_sq_header = !{SAMTOOLS_DOCKER} samtools view -H {bam} | grep '^@SQ'\n",
    "    bam_references = set()\n",
    "    for line in bam_sq_header:\n",
    "        bam_references.add(line.split('\\t')[1].removeprefix('SN:'))\n",
    "    missing = bam_references - bed_references\n",
    "    if missing:\n",
    "        print(f'WARNING: The reference(s) {missing} from BAM could not be matched to a corresponding reference in BED for {bam}.\\n')\n",
    "\n",
    "sorted_bed_path = os.path.join(refs_dir, f'sorted_{primer_bed}')\n",
    "bed_df.drop(columns=['amplicon', 'amplicon_number']).to_csv(sorted_bed_path, sep='\\t', header=False, index=False)\n",
    "bed_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trim primers from alignment using samtools ampliconclip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample in sample_names:\n",
    "    for ref in reference_sequence_paths:\n",
    "\n",
    "        out_subdir = os.path.join(trimmed_alignment_dir, ref)\n",
    "        os.makedirs(out_subdir, exist_ok=True)\n",
    "\n",
    "        in_subdir = os.path.join(raw_alignment_dir, ref)\n",
    "        aln_bam = os.path.join(in_subdir, f'{sample}.aln.sorted.bam')\n",
    "\n",
    "        trimmed_bam = os.path.join(out_subdir, f'{sample}_trimmed.bam')\n",
    "        trimmed_sorted_bam = os.path.join(out_subdir, f'{sample}_trimmed.sorted.bam')\n",
    "\n",
    "        # we use --both-ends since small enough fragments will sequence into the opposite primer\n",
    "        # first column in BED must match reference name(s) in BAM for primers to be trimmed\n",
    "        !{SAMTOOLS_DOCKER} samtools ampliconclip --both-ends -b {sorted_bed_path} -o {trimmed_bam} {aln_bam}\n",
    "\n",
    "        !{SAMTOOLS_DOCKER} samtools sort {trimmed_bam} -o {trimmed_sorted_bam}\n",
    "        !{SAMTOOLS_DOCKER} samtools index {trimmed_sorted_bam}\n",
    "        !rm {trimmed_bam}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate alignment metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(aln_metrics_dir):\n",
    "    os.mkdir(aln_metrics_dir)\n",
    "    \n",
    "for sample in sample_names:\n",
    "    for ref in reference_sequence_paths:\n",
    "        aln_dir = os.path.join('bwa_alignment_trimmed', ref)\n",
    "        bam_file = os.path.join(aln_dir, f'{sample}_trimmed.sorted.bam')\n",
    "\n",
    "        out_subdir_aln_metrics = os.path.join(aln_metrics_dir, ref)\n",
    "        if not os.path.exists(out_subdir_aln_metrics):\n",
    "            os.mkdir(out_subdir_aln_metrics)\n",
    "\n",
    "        out_bam_coverage = os.path.join(out_subdir_aln_metrics, f'{sample}_{ref}_coverage.txt')\n",
    "        out_bam_stats = os.path.join(out_subdir_aln_metrics, f'{sample}_{ref}_stats.txt')\n",
    "\n",
    "        !{SAMTOOLS_DOCKER} samtools coverage -o {out_bam_coverage} {bam_file}\n",
    "        !{SAMTOOLS_DOCKER} sh -c \"samtools stats {bam_file} > {out_bam_stats}\"\n",
    "\n",
    "        # Output per-segment stats for multi-segment references\n",
    "        segments = !{SAMTOOLS_DOCKER} samtools idxstats {bam_file} | cut -f1 | grep -v '*'\n",
    "        if len(segments) > 1:\n",
    "            for segment in segments:\n",
    "                out_bam_coverage = os.path.join(out_subdir_aln_metrics,\n",
    "                                                segment,\n",
    "                                                f'{sample}_{ref}_{segment}_coverage.txt')\n",
    "\n",
    "                out_bam_stats = os.path.join(out_subdir_aln_metrics,\n",
    "                                            segment,\n",
    "                                            f'{sample}_{ref}_{segment}_stats.txt')\n",
    "\n",
    "                !{SAMTOOLS_DOCKER} samtools coverage --region {segment} -o {out_bam_coverage} {bam_file}\n",
    "                !{SAMTOOLS_DOCKER} sh -c \"samtools stats {bam_file} {segment} > {out_bam_stats}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Determine consensus sequence using ivar consensus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(consensus_dir):\n",
    "    os.mkdir(consensus_dir)\n",
    "\n",
    "for sample in sample_names:\n",
    "    for ref in reference_sequence_paths:\n",
    "        out_subdir_consensus = os.path.join(consensus_dir, ref)\n",
    "        if not os.path.exists(out_subdir_consensus):\n",
    "            os.mkdir(out_subdir_consensus)\n",
    "\n",
    "        trimmed_sorted_bam = os.path.join(trimmed_alignment_dir, ref, f'{sample}_trimmed.sorted.bam')\n",
    "        pileup_txt = os.path.join(trimmed_alignment_dir, ref, f'{sample}_pileup.txt')\n",
    "\n",
    "        # generate pileup\n",
    "        !{SAMTOOLS_DOCKER} samtools faidx {reference_sequence_paths[ref]}\n",
    "        !{SAMTOOLS_DOCKER} samtools mpileup -A -aa -d 600000 -B -Q 20 -q 20 -f {reference_sequence_paths[ref]} {trimmed_sorted_bam} -o {pileup_txt}\n",
    "\n",
    "        records = SeqIO.parse(reference_sequence_paths[ref], format='fasta')\n",
    "        num_records = len(list(records))\n",
    "        # if reference is a multifasta, write separate consensus for each gene segment\n",
    "        if num_records > 1:\n",
    "            segment_fastas = []\n",
    "            for segment in SeqIO.parse(reference_sequence_paths[ref], format='fasta'):\n",
    "                out_subdir_consensus_segment = os.path.join(consensus_dir, ref, segment.id)\n",
    "                os.makedirs(out_subdir_consensus_segment, exist_ok=True)\n",
    "                consensus_prefix = os.path.join(out_subdir_consensus_segment, f'{sample}_{ref}_{segment.id}_consensus')\n",
    "\n",
    "                command = f\"cat {pileup_txt} | grep {segment.id} | ivar consensus -p {consensus_prefix} -q 20 -t 0.6 -m 10\"\n",
    "                print(command)\n",
    "                !{IVAR_DOCKER} sh -c \"{command}\"\n",
    "                \n",
    "                segment_fasta = f'{consensus_prefix}.fa'\n",
    "                segment_fastas.append(segment_fasta)\n",
    "                \n",
    "            # also output a fasta with all the gene segments\n",
    "            segments_fastas_str = ' '.join(segment_fastas)\n",
    "            !cat {segments_fastas_str} > {out_subdir_consensus}/{sample}_{ref}_consensus.fa\n",
    "        else:\n",
    "            out_subdir_consensus = os.path.join(consensus_dir, ref)\n",
    "            consensus_prefix = os.path.join(out_subdir_consensus, f'{sample}_{ref}_consensus')\n",
    "            !{IVAR_DOCKER} sh -c \"cat {pileup_txt} | ivar consensus -p {consensus_prefix} -q 20 -t 0.6 -m 10\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate percent coverage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reference length with and without primers\n",
    "reference_lengths = {}\n",
    "for ref in reference_sequence_paths:\n",
    "    records = list(SeqIO.parse(reference_sequence_paths[ref], format='fasta'))\n",
    "    num_records = len(records)\n",
    "    if num_records > 1:\n",
    "        segment_lengths = {}\n",
    "        for segment in records:\n",
    "            total_length = len(segment.seq)\n",
    "            segment_bed_df = bed_df[bed_df[0] == segment.id]\n",
    "\n",
    "#                       Primer0                      PrimerN\n",
    "#                       ======>                    <=========\n",
    "#             sequence: A  T  C  G  G  C  G  A  T  T  A  A  A \n",
    "#             0-based:  0  1  2  3  4  5  6  7  8  9  10 11 12\n",
    "#             1-based:  1  2  3  4  5  6  7  8  9  10 11 12 13\n",
    "\n",
    "#             Primer1 BED coordinates: [0,3)\n",
    "#             PrimerN BED coordinates: [9,13)\n",
    "#             Insert length: 9 - 3 = 6\n",
    "\n",
    "            primer_insert_length = segment_bed_df[1].max() - segment_bed_df[2].min()\n",
    "\n",
    "            segment_lengths[segment.id] = (total_length, primer_insert_length)\n",
    "        reference_lengths[ref] = segment_lengths\n",
    "    else:\n",
    "        record = records[0]\n",
    "        total_length = len(record.seq)\n",
    "        segment_bed_df = bed_df[bed_df[0] == record.id]\n",
    "        if segment_bed_df.empty:\n",
    "            raise ValueError(f'Reference {record.id} not found in BED')\n",
    "        primer_insert_length = segment_bed_df[1].max() - segment_bed_df[2].min()\n",
    "        reference_lengths[ref] = (total_length, primer_insert_length)\n",
    "\n",
    "reference_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_percent_coverage(sample, consensus, total_reference_length, primer_insert_length):\n",
    "    record = SeqIO.read(consensus, 'fasta')\n",
    "\n",
    "    seq_len = len(record.seq)\n",
    "    number_ns = record.seq.count('N')\n",
    "    total_seq_len = seq_len - number_ns\n",
    "\n",
    "    # count N's in front of sequence and N's at back of sequence\n",
    "    seq_str = str(record.seq)\n",
    "\n",
    "    if re.search( '^N+', seq_str):\n",
    "        string = re.findall('^N+', seq_str)[0]\n",
    "        front_ns = len(string)\n",
    "    else:\n",
    "        front_ns = 0\n",
    "\n",
    "    if re.search( 'N+$', seq_str):\n",
    "        string = re.findall('N+$', seq_str)[0]\n",
    "        back_ns = len(string)\n",
    "    else:\n",
    "        back_ns = 0\n",
    "\n",
    "    percent_coverage_total_length = round(total_seq_len/total_reference_length * 100,1)\n",
    "    percent_coverage_primer_insert_length = round(total_seq_len/primer_insert_length * 100,1)\n",
    "    \n",
    "    consensus_path_split = consensus.split('/')\n",
    "    if len(consensus_path_split) == 4:\n",
    "        reference = f'{consensus_path_split[1]}: {consensus_path_split[2]}'\n",
    "        subdir = os.path.join(consensus_path_split[1], consensus_path_split[2])\n",
    "    else:\n",
    "        reference = consensus_path_split[1]\n",
    "        subdir = consensus_path_split[1]\n",
    "\n",
    "    df = pd.DataFrame()\n",
    "    df['sample_name'] = [sample]\n",
    "    df['reference'] = reference\n",
    "    df['percent_coverage_total_length'] = percent_coverage_total_length\n",
    "    df['percent_coverage_primer_insert_length'] = percent_coverage_primer_insert_length\n",
    "    df['number_aln_bases'] = seq_len\n",
    "    df['ambiguous_bases'] = number_ns\n",
    "    df['number_aln_nonambiguous_bases'] = total_seq_len\n",
    "    df['front_ns'] = front_ns\n",
    "    df['back_ns'] = back_ns\n",
    "    \n",
    "    os.makedirs(os.path.join(aln_metrics_dir, subdir), exist_ok=True)\n",
    "    fname = f'{consensus_path_split[-1].removesuffix(\"_consensus.fa\")}_percent_coverage.csv'\n",
    "    outfile = os.path.join(aln_metrics_dir, subdir, fname)\n",
    "    df.to_csv(outfile, index=False)\n",
    "\n",
    "    return df\n",
    "\n",
    "percent_coverage_dfs = []\n",
    "for sample in sample_names:\n",
    "    for ref in reference_lengths:\n",
    "        # if reference is multi-segment\n",
    "        if isinstance(reference_lengths[ref], dict):\n",
    "            for segment in reference_lengths[ref]:\n",
    "                consensus = os.path.join(consensus_dir, ref, segment, f'{sample}_{ref}_{segment}_consensus.fa')\n",
    "                percent_coverage_df = calculate_percent_coverage(\n",
    "                    sample,\n",
    "                    consensus,\n",
    "                    reference_lengths[ref][segment][0],\n",
    "                    reference_lengths[ref][segment][1],\n",
    "                )\n",
    "                percent_coverage_dfs.append(percent_coverage_df)\n",
    "        \n",
    "        else:\n",
    "            consensus = os.path.join(consensus_dir, ref, f'{sample}_{ref}_consensus.fa')\n",
    "            percent_coverage_df = calculate_percent_coverage(\n",
    "                sample,\n",
    "                consensus,\n",
    "                reference_lengths[ref][0],\n",
    "                reference_lengths[ref][1],\n",
    "            )\n",
    "            percent_coverage_dfs.append(percent_coverage_df)\n",
    "            \n",
    "concat_percent_coverage_df = pd.concat(percent_coverage_dfs).reset_index(drop=True)\n",
    "concat_percent_coverage_outfile = os.path.join(aln_metrics_dir, f'{project_name}_percent_coverage.csv')\n",
    "concat_percent_coverage_df.to_csv(concat_percent_coverage_outfile, index=False)\n",
    "concat_percent_coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def summarize_coverage(sample, stats_dir):\n",
    "    stats_dir_split = stats_dir.split('/')\n",
    "    if len(stats_dir_split) == 3:\n",
    "        # ref_segment\n",
    "        ref_str = f'{stats_dir_split[1]}_{stats_dir_split[2]}'\n",
    "    else:\n",
    "        # just ref\n",
    "        ref_str = f'{stats_dir_split[1]}'\n",
    "    \n",
    "    bam_coverage = os.path.join(stats_dir, f'{sample}_{ref_str}_coverage.txt')\n",
    "    bam_cov_df = pd.read_csv(bam_coverage, sep = '\\t')\n",
    "\n",
    "    # pull out number of reads mapped and use with fastqc to calculate \n",
    "    # the number of unmapped reads and the percent of reads mapped\n",
    "    bam_stats = os.path.join(stats_dir, f'{sample}_{ref_str}_stats.txt')\n",
    "    with open (bam_stats, 'r') as file:\n",
    "        for line in file:\n",
    "            if re.search('reads mapped:', line):\n",
    "                reads_mapped = float(line.split('\\t')[2].strip())\n",
    "\n",
    "\n",
    "    # pull out number of reads from fastqc output\n",
    "    fastqc_metrics_file = os.path.join(wd, \"fastqc_clean\", sample, f'{sample}_summary_metrics.tsv')\n",
    "    fastqc_df = pd.read_csv(fastqc_metrics_file, sep = '\\t')\n",
    "    total_raw_reads = fastqc_df['r1_total_reads'][0] + fastqc_df['r2_total_reads'][0]\n",
    "\n",
    "    # calcuate the number of unmapped reads and the percent reads mapped\n",
    "    if total_raw_reads != 0:\n",
    "        percent_reads_mapped = round(reads_mapped/total_raw_reads * 100, 1)\n",
    "    else:\n",
    "        percent_reads_mapped = 0\n",
    "\n",
    "    reads_unmapped = total_raw_reads - reads_mapped\n",
    "\n",
    "    # pull out percent coverage 1 and 2 from percent coverage\n",
    "    percent_coverage_file = os.path.join(stats_dir, f'{sample}_{ref_str}_percent_coverage.csv')\n",
    "    percent_df = pd.read_csv(percent_coverage_file)\n",
    "\n",
    "\n",
    "    df = pd.DataFrame()\n",
    "    df['sample_name'] = [sample]\n",
    "    df['reference'] = [percent_df['reference'].iloc[0]]\n",
    "\n",
    "    df['percent_reads_mapped'] = percent_reads_mapped\n",
    "    df['total_raw_reads'] = total_raw_reads\n",
    "    df['reads_mapped'] = reads_mapped\n",
    "    df['reads_unmapped'] = reads_unmapped\n",
    "    df['av_depth'] = [bam_cov_df['meandepth'].iloc[0]]\n",
    "    df['meanmapq'] = [bam_cov_df['meanmapq'].iloc[0]]\n",
    "\n",
    "    df['percent_coverage_total_length'] = [percent_df['percent_coverage_total_length'].iloc[0]]\n",
    "    df['percent_coverage_primer_insert_length'] = [percent_df['percent_coverage_primer_insert_length'].iloc[0]]\n",
    "    df['front_ns'] = [percent_df['front_ns'].iloc[0]]\n",
    "    df['back_ns'] = [percent_df['back_ns'].iloc[0]]\n",
    "    \n",
    "    outfile = os.path.join(stats_dir, f'{sample}_{ref_str}_aln_metrics_summary.csv')\n",
    "\n",
    "    df.to_csv(outfile, index=False)\n",
    "\n",
    "    return df\n",
    "\n",
    "coverage_summary_dfs = []\n",
    "for sample in sample_names:\n",
    "    for ref in reference_sequence_paths:\n",
    "        records = SeqIO.parse(reference_sequence_paths[ref], format='fasta')\n",
    "        num_records = len(list(records))\n",
    "        if num_records > 1:\n",
    "            for segment in SeqIO.parse(reference_sequence_paths[ref], format='fasta'):\n",
    "                stats_dir = os.path.join(aln_metrics_dir, ref, segment.id)\n",
    "                coverage_summary_df = summarize_coverage(sample, stats_dir)\n",
    "                coverage_summary_dfs.append(coverage_summary_df)\n",
    "        else:\n",
    "            stats_dir = os.path.join(aln_metrics_dir, ref)\n",
    "            coverage_summary_df = summarize_coverage(sample, stats_dir)\n",
    "            coverage_summary_dfs.append(coverage_summary_df)\n",
    "\n",
    "concat_coverage_summary_df = pd.concat(coverage_summary_dfs).reset_index(drop=True)\n",
    "concat_coverage_summary_outfile = os.path.join(aln_metrics_dir, f'{project_name}_aln_metrics_summary.csv')\n",
    "concat_coverage_summary_df.to_csv(concat_coverage_summary_outfile, index=False)\n",
    "concat_coverage_summary_df       "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate per-amplicon coverage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(amplicon_dir):\n",
    "    os.mkdir(amplicon_dir)\n",
    "\n",
    "num_primers = bed_df[3].nunique()\n",
    "amplicons = list(bed_df['amplicon'].unique())\n",
    "if num_primers / len(amplicons) != 2:\n",
    "    raise ValueError('Either more than two primers per amplicon or cannot properly parse primer names in BED')\n",
    "\n",
    "amplicon_min_depths_dfs = []\n",
    "for sample in sample_names:\n",
    "    for ref in reference_sequence_paths:\n",
    "        trimmed_sorted_bam = os.path.join(trimmed_alignment_dir, ref, f'{sample}_trimmed.sorted.bam')\n",
    "        amplicon_subdir = os.path.join(amplicon_dir, ref)\n",
    "        if not os.path.exists(amplicon_subdir):\n",
    "            os.mkdir(amplicon_subdir)\n",
    "        ampliconstats_out = os.path.join(amplicon_subdir, f'{sample}_{ref}_ampliconstats.txt')\n",
    "        ampliconstats_plots_dir = os.path.join(amplicon_subdir, f'{sample}_{ref}_ampliconstats_plots')\n",
    "        os.makedirs(ampliconstats_plots_dir, exist_ok=True)\n",
    "        plot_ampliconstats_prefix = os.path.join(ampliconstats_plots_dir, f'{sample}_{ref}_ampliconstats')\n",
    "        \n",
    "        # calculate coverage stats per amplicon\n",
    "        !{SAMTOOLS_DOCKER} samtools ampliconstats {sorted_bed_path} {trimmed_sorted_bam} -o {ampliconstats_out}\n",
    "        !{SAMTOOLS_DOCKER} plot-ampliconstats {plot_ampliconstats_prefix} {ampliconstats_out}\n",
    "\n",
    "        # only select amplicons relevant to BAM   \n",
    "        bam_sq_header = !{SAMTOOLS_DOCKER} samtools view -H {trimmed_sorted_bam} | grep '^@SQ'\n",
    "        bam_references = set()\n",
    "        for line in bam_sq_header:\n",
    "            bam_references.add(line.split('\\t')[1].removeprefix('SN:'))   \n",
    "        subset_bed_df = bed_df[bed_df[0].isin(bam_references)]\n",
    "        amplicons = subset_bed_df['amplicon'].unique().tolist()\n",
    "\n",
    "        # get depth at every position\n",
    "        # use -J so that deletions don't get counted as 0 depth\n",
    "        depths_out = os.path.join(amplicon_subdir, f'{sample}_{ref}_depths.txt')\n",
    "        !{SAMTOOLS_DOCKER} samtools depth -aJ {trimmed_sorted_bam} -o {depths_out}\n",
    "\n",
    "        depths_df = pd.read_csv(depths_out, header=None, sep='\\t',\n",
    "                                names=['reference', 'position', 'depth'])\n",
    "\n",
    "        os.remove(depths_out)\n",
    "\n",
    "        # determine minimum depth per amplicon\n",
    "        if not depths_df.empty:\n",
    "            amplicon_min_depths = []\n",
    "            for amplicon in amplicons:\n",
    "                amplicon_record = {}\n",
    "                amplicon_start = subset_bed_df[(subset_bed_df['amplicon'] == amplicon)\n",
    "                                               & (subset_bed_df[5] == '+')][2].values[0] + 1\n",
    "\n",
    "                amplicon_end = subset_bed_df[(subset_bed_df['amplicon'] == amplicon)\n",
    "                                             & (subset_bed_df[5] == '-')][1].values[0]\n",
    "                \n",
    "                current_reference = subset_bed_df[subset_bed_df['amplicon'] == amplicon][0].iloc[0]\n",
    "\n",
    "                subset_depths_df = depths_df[(depths_df['reference'] == current_reference)\n",
    "                                             & (depths_df['position'] >= amplicon_start)\n",
    "                                             & (depths_df['position'] <= amplicon_end)]\n",
    "                                \n",
    "                amplicon_min_depth = subset_depths_df['depth'].min()\n",
    "\n",
    "                amplicon_record['sample_name'] = sample\n",
    "                amplicon_record['reference'] = ref\n",
    "                amplicon_record['amplicon'] = amplicon\n",
    "                amplicon_record['start'] = amplicon_start\n",
    "                amplicon_record['end'] = amplicon_end\n",
    "                amplicon_record['min_depth'] = amplicon_min_depth if pd.notna(amplicon_min_depth) else 0\n",
    "                amplicon_min_depths.append(amplicon_record)\n",
    "\n",
    "            amplicon_depths_df = pd.DataFrame.from_records(amplicon_min_depths)\n",
    "            amplicon_min_depths_out = os.path.join(amplicon_subdir, f'{sample}_{ref}_amplicon_min_depths.tsv')\n",
    "            amplicon_depths_df.to_csv(amplicon_min_depths_out, sep='\\t', index=False)\n",
    "            amplicon_min_depths_dfs.append(amplicon_depths_df)\n",
    "\n",
    "concat_amplicon_min_depths_df = pd.concat(amplicon_min_depths_dfs).reset_index(drop=True)\n",
    "concat_amplicon_min_depths_out = os.path.join(amplicon_dir, f'{project_name}_amplicon_min_depths.tsv')\n",
    "concat_amplicon_min_depths_df.to_csv(concat_amplicon_min_depths_out, sep='\\t', index=False)\n",
    "concat_amplicon_min_depths_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}