crykey_main_figures_clean.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import glob\n",
    "import re\n",
    "import os\n",
    "import copy\n",
    "import subprocess\n",
    "import time\n",
    "import datetime as dt\n",
    "from datetime import datetime, timedelta, date\n",
    "from collections import defaultdict, Counter, OrderedDict\n",
    "from functools import reduce\n",
    "\n",
    "import vcf\n",
    "import pysam\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import dask.dataframe as dd\n",
    "from Bio import SeqIO\n",
    "from Bio.SeqUtils import seq1\n",
    "from Bio.Seq import Seq\n",
    "\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import matplotlib.ticker as mtick\n",
    "from matplotlib.patches import Patch\n",
    "from matplotlib.lines import Line2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "from matplotlib.colors import LinearSegmentedColormap, ListedColormap\n",
    "import matplotlib.colors as mcolors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.gridspec as gridspec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.dates as mdates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vcf\n",
    "from operator import add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import fisher_exact"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove Problematic Sites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_crykey_output(population_df):\n",
    "    \"\"\"\n",
    "    This function filters the output of crykey by\n",
    "    1. removing cryptic lineage found in control samples\n",
    "    2. removing cryptic lineage containing variants of problematic sites\n",
    "    3. removing cryptic lineage located within 100 bp of the SARS-CoV-2 reference genome\n",
    "    The function returns a filtered dataframe.\n",
    "    The filtered dataframe is stored at: fixed_results/crykey_wastewater_houston.csv\n",
    "    \"\"\"\n",
    "    output = \"/home/Users/yl181/wastewater/quaid/output_incl_recombinant\"\n",
    "    dbs = '/home/Users/yl181/wastewater/quaid/quarc_dbs_01102023_incl_recombinant'\n",
    "    \n",
    "    sites = population_df['WWTP'].unique()\n",
    "    merged_df = pd.read_csv(os.path.join(output, 'merged_df.csv'))\n",
    "    \n",
    "    merged_df['Site'] = merged_df['Site'].str.upper()\n",
    "    merged_df = merged_df[merged_df.Site.isin(sites)]\n",
    "    \n",
    "    merged_df['Date'] = merged_df['Date'].apply(str)\n",
    "    merged_df['Date'] = merged_df['Date'].apply(lambda x: x.zfill(8))\n",
    "    merged_df['Date'] = merged_df['Date'].apply(lambda x: f'{x[0:2]}/{x[2:4]}/{x[4:]}')\n",
    "    merged_df['Date'] = pd.to_datetime(merged_df['Date'])\n",
    "    \n",
    "    exclude_end_bp = 100\n",
    "    exclude_sites = [187, 1059, 2094, 3037, 3130, 6990, 8022, 10323, 10741, 11074, 13408,\n",
    "                     14786, 19684, 20148, 21137, 24034, 24378, 25563, 26144, 26461, 26681,\n",
    "                     28077, 28826, 28854, 29700]\n",
    "    exclude_sites = exclude_sites + list(np.arange(1,1+exclude_end_bp)) + list(np.arange(29903-exclude_end_bp+1, 29903+1))\n",
    "    \n",
    "    keep_records = []\n",
    "    for idx, row in merged_df.iterrows():\n",
    "        nt_mutations = row['Nt Mutations'].split(\";\")\n",
    "        keep = True\n",
    "        for mut in nt_mutations:\n",
    "            pos = int(mut[1:-1])\n",
    "            if pos in exclude_sites:\n",
    "                keep = False\n",
    "                break\n",
    "\n",
    "        if keep:\n",
    "            keep_records.append(row)\n",
    "            \n",
    "    merged_df = pd.DataFrame(keep_records)\n",
    "    return merged_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coverage Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coverage_date_dataframe(coverage_dir):\n",
    "    \"\"\"\n",
    "    Sample collection dates to week starts\n",
    "    Check whether we have duplicate samples for the same week\n",
    "    Returns a mapping between sample collection date and week start\n",
    "    \"\"\"\n",
    "    coverage_dates = []\n",
    "    for coverage_date in os.listdir(coverage_dir):\n",
    "        date = pd.to_datetime(coverage_date[3:], format='%m%d%Y')\n",
    "        coverage_dates.append(date)\n",
    "    coverage_dates.sort()\n",
    "    coverage_df = pd.DataFrame(coverage_dates, columns=['Collection Date'])\n",
    "    coverage_df['Week Start'] = coverage_df['Collection Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "    try:\n",
    "        assert len(coverage_df['Collection Date'].unique()) == len(coverage_df['Week Start'].unique())\n",
    "    except AssertionError:\n",
    "        weekstart2collection = defaultdict(list)\n",
    "        for idx, row in coverage_df.iterrows():\n",
    "            weekstart2collection[row['Week Start']].append(row['Collection Date'])\n",
    "        for i in weekstart2collection:\n",
    "            if len(weekstart2collection[i]) != 1:\n",
    "                print(i, weekstart2collection[i])\n",
    "    return coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_missing_weeks(coverage_dir):\n",
    "    \"\"\"\n",
    "    Calculate missing weeks\n",
    "    Returns a mapping between sample collection date and week start\n",
    "    \"\"\"\n",
    "    coverage_df = coverage_date_dataframe(coverage_dir)\n",
    "    \n",
    "    total_week_count = int((coverage_df['Week Start'].max() - coverage_df['Week Start'].min()).days/7) + 1\n",
    "    missing_week_count = total_week_count - len(coverage_df['Week Start'].unique())\n",
    "    print(\"Total Weeks:\\t\", total_week_count)\n",
    "    print(\"Missing Weeks:\\t\", missing_week_count)\n",
    "    \n",
    "    start_date = coverage_df['Week Start'].min()\n",
    "    days_in_a_week = timedelta(days = 7)\n",
    "    missing_weeks = []\n",
    "    for i in range(total_week_count):\n",
    "        if not start_date + i*days_in_a_week in coverage_df['Week Start'].unique():\n",
    "            missing_weeks.append(start_date + i*days_in_a_week)\n",
    "    \n",
    "    return missing_weeks, coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_coverage_dataframe(coverage_date_df, population_df, missing_weeks):\n",
    "    sites = list(population_df['WWTP'].unique())\n",
    "    record_list = []\n",
    "    index_list = []\n",
    "    count = 0\n",
    "    for idx, row in coverage_date_df.iterrows():\n",
    "        date_dir = \"HHD\"+row['Collection Date'].strftime('%m%d%Y')\n",
    "        for site in sites:\n",
    "            coverage_f = os.path.join(coverage_dir, date_dir, f\"{site}-1\", f\"{site}-1.clean.sorted.coverage.txt\")\n",
    "            if not os.path.exists(coverage_f):\n",
    "                count += 1\n",
    "            else:\n",
    "                try:\n",
    "                    temp_df = pd.read_csv(coverage_f, header=None, sep='\\t', usecols = [1,2], names=['POS', 'DEP'])\n",
    "                    if not temp_df.empty:\n",
    "                        temp_df = pd.pivot_table(temp_df, index=['POS'], values=['DEP']).transpose()\n",
    "                        index_list.append((row['Collection Date'], site))\n",
    "                        # temp_df['Collection Date'] = row['Collection Date']\n",
    "                        # temp_df['Site'] = site\n",
    "                        record_list.append(temp_df.iloc[0].to_dict())\n",
    "                    else:\n",
    "                        count += 1\n",
    "                except:\n",
    "                    print(coverage_f)\n",
    "    print(\"Missing Samples:\", count + len(missing_weeks)*39)\n",
    "    \n",
    "    for missing_week in missing_weeks:\n",
    "        record_list.append({0:0})\n",
    "        index_list.append((missing_week, 'SS'))\n",
    "    \n",
    "    index = pd.MultiIndex.from_tuples(index_list, names=[\"Collection Date\", \"Site\"])\n",
    "    coverage_df = pd.DataFrame(record_list, index=index, dtype=pd.Int64Dtype())\n",
    "    columns = coverage_df.columns.to_list()\n",
    "    positions = []\n",
    "    for element in columns:\n",
    "        if isinstance(element, int):\n",
    "            positions.append(element)\n",
    "    positions.sort()\n",
    "    coverage_df = coverage_df[positions].fillna(0)\n",
    "    coverage_df = coverage_df.reset_index()\n",
    "    coverage_df = coverage_df.sort_values(['Collection Date', 'Site'], axis=0)\n",
    "    return coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_has_sample_dataframe(coverage_date_df, population_df, missing_weeks):\n",
    "    sites = list(population_df['WWTP'].unique())\n",
    "    record_list = []\n",
    "    index_list = []\n",
    "    for idx, row in coverage_date_df.iterrows():\n",
    "        date_dir = \"HHD\"+row['Collection Date'].strftime('%m%d%Y')\n",
    "        for site in sites:\n",
    "            coverage_f = os.path.join(coverage_dir, date_dir, f\"{site}-1\", f\"{site}-1.clean.sorted.coverage.txt\")\n",
    "            index_list.append((row['Collection Date'], site))\n",
    "            if (os.path.exists(coverage_f)) and (os.path.getsize(coverage_f) > 0):\n",
    "                record_list.append({'Has Sample': True})\n",
    "            else:\n",
    "                record_list.append({'Has Sample': False})\n",
    "\n",
    "    \n",
    "    for missing_week in missing_weeks:\n",
    "        for site in sites: \n",
    "            index_list.append((missing_week, site))\n",
    "            record_list.append({'Has Sample': False})\n",
    "    \n",
    "    index = pd.MultiIndex.from_tuples(index_list, names=[\"Collection Date\", \"Site\"])\n",
    "    has_sample_df = pd.DataFrame(record_list, index=index, dtype='boolean')\n",
    "    has_sample_df = has_sample_df.reset_index()\n",
    "    has_sample_df = has_sample_df.sort_values(['Collection Date', 'Site'], axis=0)\n",
    "    \n",
    "    return has_sample_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_build_coverage_dataframe():\n",
    "    coverage_dir = '/home/Users/yl181/wastewater/processed_data/Coverage'\n",
    "    missing_weeks, coverage_date_df = get_missing_weeks(coverage_dir)\n",
    "    coverage_df = build_coverage_dataframe(coverage_date_df, population_df, missing_weeks)\n",
    "    has_sample_df = build_has_sample_dataframe(coverage_date_df, vcf_dir, population_df, missing_weeks)\n",
    "    coverage_df.to_csv(os.path.join(fixed_results_dir, 'houston_coverage.csv'))\n",
    "    has_sample_df.to_csv(os.path.join(fixed_results_dir, 'houston_has_sample_df.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gisaid_total_count = 12988494"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_gisaid_occurance = 0.0001 * gisaid_total_count\n",
    "max_gisaid_occurance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Crykey Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_cryptic_lineage_df(mutation_rarity_df, merged_df, gisaid_total_count):\n",
    "    min_site = 2\n",
    "    max_gisaid_occurance = 0.0001 * gisaid_total_count\n",
    "    print('Minumum Detected Site:', \"\\t\", min_site)\n",
    "    print('Max GISAID Occurance:', \"\\t\", max_gisaid_occurance)\n",
    "\n",
    "    selected_data = mutation_rarity_df[(mutation_rarity_df['Site'] >= min_site) & (mutation_rarity_df['GISAID Count'] <= max_gisaid_occurance)].copy()\n",
    "    selected_data = selected_data.rename({'Site': 'Mutation Occurance'}, axis=1)\n",
    "    print('# the cryptic lineage after filtering:', \"\\t\", len(selected_data))\n",
    "    \n",
    "    present_weeks = []\n",
    "    present_sites = []\n",
    "    mean_freqs = []\n",
    "    max_freqs = []\n",
    "    max_durations = []\n",
    "    aa_mutations = []\n",
    "\n",
    "    for idx, row in selected_data.iterrows():\n",
    "        mutation_df = merged_df[merged_df['Nt Mutations'] == idx]\n",
    "        present_weeks.append(len(mutation_df['Date'].unique()))    \n",
    "        present_sites.append(len(mutation_df['Site'].unique()))\n",
    "        mean_freqs.append(mutation_df['Combined Freq'].mean())\n",
    "        max_freqs.append(mutation_df['Combined Freq'].max())\n",
    "        max_durations.append(int((mutation_df['Date'].unique().max()-mutation_df['Date'].unique().min()).astype('timedelta64[D]')/np.timedelta64(1, 'D') + 1))\n",
    "        aa_mutations.append(mutation_df['AA Mutations'].unique()[0])\n",
    "\n",
    "    selected_data['AA_Mutation'] = aa_mutations\n",
    "    selected_data['Present Weeks'] = present_weeks\n",
    "    selected_data['Present Sites'] = present_sites\n",
    "    selected_data['Mean Site Occurance'] = selected_data['Mutation Occurance']/selected_data['Present Weeks']\n",
    "    selected_data['Mean Allele Freq'] = mean_freqs\n",
    "    selected_data['Max Allele Freq'] = max_freqs\n",
    "    selected_data['Max Durations'] = max_durations\n",
    "    \n",
    "    selected_data = selected_data.sort_values(by=['Mutation Occurance', 'GISAID Count'], ascending=False)\n",
    "    \n",
    "    return selected_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_gisaid_rarity_dataframe(fixed_results_dir):\n",
    "    file_to_read = open(os.path.join(fixed_results_dir, \"query_result.pkl\"), \"rb\")\n",
    "    query_result = pickle.load(file_to_read)\n",
    "    \n",
    "    count_rarity_dict = dict()\n",
    "    for key in query_result:\n",
    "        count_rarity_dict[key] = sum(query_result[key].values())\n",
    "    gisaid_count_df = pd.DataFrame.from_dict(count_rarity_dict, columns=['GISAID Count'], orient='index')\n",
    "    mutation_rarity_df = site_count_df.merge(gisaid_count_df, left_index=True, right_index=True, validate='one_to_one')\n",
    "    \n",
    "    return mutation_rarity_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fixed_results_dir = \"/home/Users/yl181/wastewater/quarc_figures/fixed_results\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Samtools Coverage DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "coverage_df = pd.read_csv(os.path.join(fixed_results_dir, 'houston_coverage.csv'), index_col=False)\n",
    "coverage_df['Collection Date'] = coverage_df['Collection Date'].apply(str)\n",
    "coverage_df['Collection Date'] = pd.to_datetime(coverage_df['Collection Date'], format='%Y-%m-%d')\n",
    "coverage_df = coverage_df.drop(['Unnamed: 0'], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Has Sample DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_has_sample_dataframe(coverage_date_df, vcf_dir, population_df, missing_weeks):\n",
    "    sites = list(population_df['WWTP'].unique())\n",
    "    record_list = []\n",
    "    index_list = []\n",
    "    for idx, row in coverage_date_df.iterrows():\n",
    "        date_dir = \"HHD\"+row['Collection Date'].strftime('%m%d%Y')\n",
    "        for site in sites:\n",
    "            coverage_f = os.path.join(coverage_dir, date_dir, f\"{site}-1\", f\"{site}-1.clean.sorted.coverage.txt\")\n",
    "            index_list.append((row['Collection Date'], site))\n",
    "            if (os.path.exists(coverage_f)) and (os.path.getsize(coverage_f) > 0):\n",
    "                record_list.append({'Has Sample': True})\n",
    "            else:\n",
    "                record_list.append({'Has Sample': False})\n",
    "\n",
    "    \n",
    "    for missing_week in missing_weeks:\n",
    "        for site in sites: \n",
    "            index_list.append((missing_week, site))\n",
    "            record_list.append({'Has Sample': False})\n",
    "    \n",
    "    index = pd.MultiIndex.from_tuples(index_list, names=[\"Collection Date\", \"Site\"])\n",
    "    has_sample_df = pd.DataFrame(record_list, index=index, dtype='boolean')\n",
    "    has_sample_df = has_sample_df.reset_index()\n",
    "    has_sample_df = has_sample_df.sort_values(['Collection Date', 'Site'], axis=0)\n",
    "    \n",
    "    return has_sample_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "has_sample_df = pd.read_csv(os.path.join(fixed_results_dir, 'houston_has_sample_df.csv'), index_col=False)\n",
    "has_sample_df['Collection Date'] = has_sample_df['Collection Date'].apply(str)\n",
    "has_sample_df['Collection Date'] = pd.to_datetime(has_sample_df['Collection Date'], format='%Y-%m-%d')\n",
    "has_sample_df = has_sample_df.drop(['Unnamed: 0'], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Merged DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_df = pd.read_csv(os.path.join(fixed_results_dir, 'crykey_wastewater_houston.csv'), index_col=0)\n",
    "merged_df['Date'] = merged_df['Date'].apply(str)\n",
    "merged_df['Date'] = pd.to_datetime(merged_df['Date'], format='%Y-%m-%d')\n",
    "merged_df['Week Start'] = merged_df['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "print(\"# the cryptic lineage before filtering:\", len(merged_df['Nt Mutations'].unique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "site_count_df = pd.DataFrame(merged_df.groupby('Nt Mutations')['Site'].count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Used by\n",
    "    Supplementary Figure 1\n",
    "    Supplementary Figure 5\n",
    "\"\"\"\n",
    "\n",
    "mutation_rarity_df = build_gisaid_rarity_dataframe(fixed_results_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutation_rarity_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_cryptic_df = filter_cryptic_lineage_df(mutation_rarity_df, merged_df, gisaid_total_count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Viral Load and Population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viral_load_df = pd.read_csv(os.path.join(fixed_results_dir, 'houston_viral_load.csv')).dropna()\n",
    "viral_load_df['date']= pd.to_datetime(viral_load_df['date'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "population_df = pd.read_csv(os.path.join(fixed_results_dir, 'population.csv')).dropna()\n",
    "population_df = population_df[['WWTP', 'pop']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_WWTPs = list(population_df['WWTP'].values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 1 - Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'https://drive.google.com/file/d/1zxmsrMhb2vXBJgy-yQiQkBlpO-9FjK8J/view?usp=sharing'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 3 - Genomic Distribution of CLs in Wastewater Samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_colorbar(mappable):\n",
    "    last_axes = plt.gca()\n",
    "    ax = mappable.axes\n",
    "    fig = ax.figure\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"right\", size=\"3%\", pad=0.05)\n",
    "    cbar = fig.colorbar(mappable, cax=cax)\n",
    "    cbar.ax.get_yaxis().labelpad = 5\n",
    "    cbar.ax.set_ylabel('Rarity in GISAID', rotation=90, fontsize=fontsize)\n",
    "    plt.sca(last_axes)\n",
    "    return cbar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def synonymous_mutations(filtered_cryptic_df):\n",
    "    selected_data = filtered_cryptic_df.copy().rename({'AA_Mutation': 'AA Mutation'}, axis=1)\n",
    "    #selected_data = selected_data.reset_index().rename({'index': 'Nt Mutation'}, axis=1)\n",
    "    non_synonyms_exclusive_list = []\n",
    "    for idx, row in selected_data.iterrows():\n",
    "        aa_mutation_set = row['AA Mutation']\n",
    "        ns_count = 0\n",
    "        for aa_muts in aa_mutation_set.split(';'):\n",
    "            temp = aa_muts.split(\":\")[-1]\n",
    "            if temp[0] == temp[-1]:\n",
    "                ns_count += 1\n",
    "        if ns_count > 0:\n",
    "            non_synonyms_exclusive_list.append(False)\n",
    "        else:\n",
    "            non_synonyms_exclusive_list.append(True)\n",
    "\n",
    "    selected_data['Non-synonyms Only'] = non_synonyms_exclusive_list\n",
    "    \n",
    "    return selected_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gene_label(selected_data):\n",
    "    gene_list = []\n",
    "    snp_locations = []\n",
    "    s_gene_count = 0\n",
    "    n_gene_count = 0\n",
    "    orf1a_count = 0\n",
    "    orf1b_count = 0\n",
    "\n",
    "    for idx, row in selected_data.iterrows():\n",
    "        snps = idx.split(\";\")\n",
    "        snp_locations.append(int(np.array([int(i[1:-1]) for i in snps]).mean()))\n",
    "        aa_muts = row['AA Mutation'].split(';')\n",
    "        gene_set = set()\n",
    "        for aa_mut in aa_muts:\n",
    "            gene = aa_mut.split(':')[0]\n",
    "            gene_set.add(gene)\n",
    "\n",
    "        gene_list.append(\";\".join(list(gene_set)))\n",
    "\n",
    "        if 'S' in gene_set:\n",
    "            s_gene_count += 1\n",
    "        if 'N' in gene_set:\n",
    "            n_gene_count += 1\n",
    "        if 'ORF1a' in gene_set:\n",
    "            orf1a_count += 1\n",
    "        if 'ORF1b' in gene_set:\n",
    "            orf1b_count += 1      \n",
    "\n",
    "    selected_data['Gene'] = gene_list\n",
    "    selected_data['Nt Location'] = snp_locations\n",
    "    \n",
    "    print('Total CL #:', '\\t', selected_data.shape[0])\n",
    "    \n",
    "    print(\"S Gene CL #:\", '\\t', s_gene_count)\n",
    "    print(\"S Gene CL %:\", '\\t', s_gene_count/selected_data.shape[0])\n",
    "    print(\"N Gene CL #:\", '\\t', n_gene_count)\n",
    "    print(\"N Gene CL %:\", '\\t', n_gene_count/selected_data.shape[0])\n",
    "    \n",
    "    print(\"Minor AF (<0.5) CL #:\", '\\t', selected_data[selected_data['Mean Allele Freq']<0.5].shape[0])\n",
    "    print(\"Minor AF (<0.5) CL %:\", '\\t', selected_data[selected_data['Mean Allele Freq']<0.5].shape[0]/selected_data.shape[0])\n",
    "    \n",
    "    print(\"Consensus AF (>=0.5) CL %:\", '\\t', selected_data[selected_data['Mean Allele Freq']>=0.5].shape[0]/selected_data.shape[0])\n",
    "    \n",
    "    return selected_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference = SeqIO.read(\"SARS-CoV-2-reference.gb\", \"genbank\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_data = synonymous_mutations(filtered_cryptic_df)\n",
    "selected_data = get_gene_label(selected_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(selected_data[selected_data['Mean Allele Freq'] < 0.2])/selected_data.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes  = plt.subplots(3, 1, figsize=(14, 10),\n",
    "                          sharex=True,\n",
    "                          gridspec_kw={'height_ratios': [1, 2.5, 1], 'hspace':0.1})\n",
    "\n",
    "fontsize = 12\n",
    "alpha = 0.5\n",
    "\n",
    "cmap = mpl.cm.inferno_r\n",
    "\n",
    "bin_width = 400\n",
    "\n",
    "selected_data['Bin Index'] = selected_data['Nt Location']/bin_width\n",
    "selected_data['Bin Index'] = selected_data['Bin Index'].astype(int)\n",
    "selected_data['Bin Index'] = selected_data['Bin Index'] * bin_width + bin_width/2\n",
    "\n",
    "hist_df = pd.DataFrame(selected_data.groupby(['Bin Index']).count())\n",
    "\n",
    "bar_x = hist_df.index.to_list()\n",
    "bar_y1 = hist_df['AA Mutation'].values\n",
    "\n",
    "hist_df = pd.DataFrame(selected_data.groupby(['Bin Index'])['Non-synonyms Only'].sum())\n",
    "bar_y2 = hist_df['Non-synonyms Only'].values\n",
    "\n",
    "ax = axes[1]\n",
    "\n",
    "rarity = -np.log10((selected_data['GISAID Count'].values+1)/gisaid_total_count)\n",
    "main_plot = ax.scatter(x=selected_data['Nt Location'].values, \n",
    "                       y=selected_data['Mean Allele Freq'].values, \n",
    "                       c=rarity, \n",
    "                       s=selected_data['Present Weeks'].values*10,\n",
    "                       cmap=cmap,\n",
    "                       alpha=alpha)\n",
    "\n",
    "ax.set_xlim(0, 29903+1)\n",
    "#ax.set_xticks(np.arange(0, 29903+1, 5000))\n",
    "\n",
    "ax.set_ylim(0,1.1)\n",
    "ax.set_yticks(np.arange(0, 11, 2)/10)\n",
    "ax.set_ylabel('Mean AF of CRs', fontsize=fontsize)\n",
    "\n",
    "ax.xaxis.set_tick_params(labelsize=fontsize)\n",
    "ax.yaxis.set_tick_params(labelsize=fontsize)\n",
    "\n",
    "add_colorbar(main_plot)\n",
    "\n",
    "handles, labels = main_plot.legend_elements(prop=\"sizes\", alpha=alpha, num=4, func = lambda x: x/10)\n",
    "legend = ax.legend(handles, labels, loc=\"upper left\", title=\"Weeks of detection\", ncol=4, framealpha=1, borderpad=0.6)#, bbox_to_anchor=(0.5,1))\n",
    "legend.get_frame().set_facecolor('none')\n",
    "\n",
    "ax.text(-0.07, 1.05, 'a', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "ax = axes[0]\n",
    "colors = plt.get_cmap('Set3').colors\n",
    "\n",
    "gene_label_h = 0.5\n",
    "gene_label_y = -gene_label_h/2\n",
    "label_top = gene_label_h/2*1.5\n",
    "counter = 0\n",
    "\n",
    "for seq_feature in reference.features:\n",
    "    if seq_feature.type == 'gene':\n",
    "        \n",
    "        gene_name = seq_feature.qualifiers['gene'][0]\n",
    "        \n",
    "        rectangle = Rectangle((seq_feature.location.start, gene_label_y), seq_feature.location.end-seq_feature.location.start, gene_label_h,\n",
    "                              fill=True,\n",
    "                              facecolor=colors[counter],\n",
    "                              edgecolor='black')\n",
    "        ax.add_patch(rectangle)\n",
    "        \n",
    "        if counter != 0:\n",
    "            rx, ry = rectangle.get_xy()\n",
    "            cx = rx + rectangle.get_width()/2.0\n",
    "            cy = ry + rectangle.get_height()/2.0 + label_top\n",
    "            if label_top > 0:\n",
    "                va = 'top'\n",
    "            else:\n",
    "                va = 'bottom'\n",
    "            ax.annotate(gene_name, (cx, cy), color='black', weight='normal', fontsize=fontsize-2, ha='center', va=va, rotation=90)\n",
    "        else:\n",
    "            rx, ry = rectangle.get_xy()\n",
    "            cx = rx + rectangle.get_width()/2.0\n",
    "            cy = ry + rectangle.get_height()/2.0 - label_top\n",
    "            ax.annotate(gene_name, (cx, cy), color='black', weight='normal', fontsize=fontsize-2, ha='center', va='bottom', rotation=90)\n",
    "        \n",
    "        counter += 1\n",
    "        label_top = label_top * -1\n",
    "        \n",
    "\n",
    "ax.set_ylim(-label_top*3,label_top*3)\n",
    "comb_pos = -2.5\n",
    "\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"3%\", pad=0.05)\n",
    "\n",
    "ax.axis('off')\n",
    "cax.axis('off')\n",
    "\n",
    "ax = axes[2]\n",
    "\n",
    "ax.bar(bar_x, bar_y1, width=bin_width*0.85, color=\"black\", alpha=0.6, label='Including Synonymous Mutations')\n",
    "ax.bar(bar_x, bar_y2, width=bin_width*0.85, color=\"C1\", alpha=1, label='Non-synonymous Mutations Only')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"3%\", pad=0.05)\n",
    "\n",
    "#ax.axis('off')\n",
    "cax.axis('off')\n",
    "ax.set_xticks(np.arange(0, 29903+1, 5000))\n",
    "ax.set_xlabel('Position on Reference Genome', fontsize=fontsize)\n",
    "ax.set_ylim(0,100)\n",
    "ax.set_ylabel('CR Count', fontsize=fontsize)\n",
    "\n",
    "ax.xaxis.set_tick_params(labelsize=fontsize)\n",
    "ax.yaxis.set_tick_params(labelsize=fontsize)\n",
    "\n",
    "ax.legend(loc='upper left', borderpad=0.6)\n",
    "ax.text(-0.07, 1.05, 'b', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "#fig.show()\n",
    "\n",
    "fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/figure_3.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 4 - Longitude Distribution of CLs in Wastewater Samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_first_occ_df(merged_df, filtered_cryptic_df):\n",
    "    first_occ_df = filtered_cryptic_df.merge(merged_df.sort_values('Date')[['Date', 'Nt Mutations']].drop_duplicates('Nt Mutations'), left_index=True, right_on='Nt Mutations').set_index('Nt Mutations')\n",
    "    return first_occ_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_breadth_coverage_df():\n",
    "    coverage_df = pd.read_csv(os.path.join(fixed_results_dir, 'houston_breadth_coverage_df_v2.csv'))\n",
    "    coverage_df['Date'] = pd.to_datetime(coverage_df['Date'])\n",
    "    min_coverage = coverage_df[coverage_df['Date'] >= pd.Timestamp('2021-05-01')]['Breadth_Coverage'].min()\n",
    "    max_coverage = coverage_df[coverage_df['Date'] >= pd.Timestamp('2021-05-01')]['Breadth_Coverage'].max()\n",
    "    print(min_coverage, max_coverage)\n",
    "    coverage_df['Width'] = (coverage_df['Breadth_Coverage']-min_coverage)/(max_coverage-min_coverage)*0.9 + 0.1\n",
    "    #coverage_df['Width'] = coverage_df['Breadth_Coverage'] * 1.6\n",
    "    coverage_df[coverage_df['Date'] >= pd.Timestamp('2021-05-01')]['Width'].max()\n",
    "    \n",
    "    return coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tx_voc_df = pd.read_csv(os.path.join(fixed_results_dir, 'texas_gisaid_voc.csv'), index_col=False)\n",
    "tx_voc_df['Week'] = pd.to_datetime(tx_voc_df['Week'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_occ_df = get_first_occ_df(merged_df, filtered_cryptic_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_occ_df = synonymous_mutations(first_occ_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "breadth_coverage_df = load_breadth_coverage_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert breadth_coverage_df[breadth_coverage_df['Date'] >= pd.Timestamp('2021-05-01')]['Width'].max() == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_df = pd.DataFrame(first_occ_df.groupby(['Date'])['AA Mutation'].count())\n",
    "\n",
    "bar_x = time_df.index.to_list()\n",
    "bar_y1 = time_df['AA Mutation'].values\n",
    "time_min = time_df.index.min()\n",
    "time_max = time_df.index.max()\n",
    "\n",
    "time_df = pd.DataFrame(first_occ_df.groupby(['Date'])['Non-synonyms Only'].sum())\n",
    "time_df = time_df.merge(breadth_coverage_df, left_index=True, right_on=['Date'])\n",
    "bar_y2 = time_df['Non-synonyms Only'].values\n",
    "bar_w = time_df['Width'].values\n",
    "\n",
    "fontsize = 14\n",
    "fig, axes  = plt.subplots(2, 1, figsize=(14, 8), sharex=True,\n",
    "                         gridspec_kw={'height_ratios': [2, 1], 'hspace':0.15})\n",
    "\n",
    "ax = axes[0]\n",
    "\n",
    "ax.bar(bar_x, bar_y1, width=6*bar_w, label='Including Synonymous Mutations', color=\"black\", alpha=0.7)\n",
    "ax.bar(bar_x, bar_y2, width=6*bar_w, label='Non-synonymous Mutations Only', color='C1')\n",
    "\n",
    "ax.set_xlim(pd.Timestamp('2021-05-01')-timedelta(days=3), time_max+timedelta(days=3))\n",
    "ax.set_ylim(0,70)\n",
    "ax_load = ax.twinx()\n",
    "ax_load.plot(viral_load_df['date'], viral_load_df['Spline_WW_Percent_10'], 'k--')\n",
    "ax_load.set_ylim(0,1000)\n",
    "\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "\n",
    "ax.set_ylabel('Count of CRs', fontsize=fontsize)\n",
    "ax_load.yaxis.set_major_formatter(mtick.PercentFormatter())\n",
    "ax_load.set_ylabel('Viral Load', fontsize=fontsize)\n",
    "ax_load.tick_params(axis='y', which='major', labelsize=fontsize)\n",
    "ax.legend(fontsize=fontsize)\n",
    "ax.text(-0.12, 1.05, 'a', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "ax.set_title('Newly detected CRs in Houston wastewater', fontsize=fontsize)\n",
    "\n",
    "ax = axes[1]\n",
    "\n",
    "voc_df = pd.pivot_table(tx_voc_df, values='Accession ID', index='Week', columns=['VOC'], aggfunc='count', fill_value=0)\n",
    "ax.stackplot(voc_df.index, voc_df['Delta'], voc_df['Omicron'], voc_df['BA.2'], voc_df['BA.5'], voc_df['Other'],\n",
    "            labels = ['Delta', 'Omicron', 'BA.2', 'BA.5', 'Others'])\n",
    "\n",
    "\n",
    "ax.legend()\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax.set_title('Texas sequences in GISAID EpiCoV', fontsize=fontsize)\n",
    "\n",
    "ax.set_ylabel('Count of sequences', fontsize=fontsize)\n",
    "ax.set_ylim(0,6000)\n",
    "ax.legend(fontsize=fontsize, ncol=1, loc=2, bbox_to_anchor=(1,1))\n",
    "ax.text(-0.12, 1.05, 'b', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))\n",
    "\n",
    "fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/figure_4.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 6 - Long lasting CR12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_single_cryptic(target_mutation, merged_df):\n",
    "    selected_data = merged_df[merged_df['Nt Mutations'] == target_mutation].copy()\n",
    "    selected_data = selected_data.merge(population_df, left_on='Site', right_on='WWTP')\n",
    "    date_mean_freqs = selected_data.groupby(['Date'])['Combined Freq'].mean()\n",
    "    date_site_count = selected_data.groupby(['Date'])['Site'].count()\n",
    "    date_pop_sum = selected_data.groupby(['Date'])['pop'].sum()\n",
    "    aa_mut_label = selected_data['AA Mutations'].unique()[0]\n",
    "    \n",
    "    selected_data['Week Start'] = selected_data['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "    assert len(selected_data['Date'].unique()) == len(selected_data['Week Start'].unique())\n",
    "    \n",
    "    return selected_data, date_site_count, aa_mut_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_select_coverage_df(coverage_df, target_mutation, start_date, end_date, min_dp=10):\n",
    "    positions = []\n",
    "    for mut in target_mutation.split(\";\"):\n",
    "        positions.append(mut[1:-1])\n",
    "        \n",
    "    selected_coverage_df = coverage_df[['Collection Date', 'Site'] + positions].copy()\n",
    "    selected_coverage_df['Week Start'] = selected_coverage_df['Collection Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "    \n",
    "    days_in_a_week = timedelta(days = 7)\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    selected_dates = []\n",
    "    for i in range(total_week_count):\n",
    "        selected_dates.append(str(start_date + i*days_in_a_week))\n",
    "    selected_coverage_df = selected_coverage_df[selected_coverage_df['Week Start'].isin(selected_dates)]\n",
    "    \n",
    "    valid_coverage = []\n",
    "    for idx, row in selected_coverage_df.iterrows():\n",
    "        is_valid = True\n",
    "        for pos in positions:\n",
    "            is_valid = is_valid and (row[pos] >= min_dp)\n",
    "        valid_coverage.append(is_valid)\n",
    "    selected_coverage_df['Not Valid'] = np.invert(valid_coverage)\n",
    "        \n",
    "    return selected_coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_select_missing_sample_df(has_sample_df, start_date, end_date):  \n",
    "    selected_has_sample_df = has_sample_df[['Collection Date', 'Site', 'Has Sample']].copy()\n",
    "    selected_has_sample_df['Missing Sample'] = np.invert(selected_has_sample_df['Has Sample'])\n",
    "    selected_has_sample_df['Week Start'] = selected_has_sample_df['Collection Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "    \n",
    "    days_in_a_week = timedelta(days = 7)\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    selected_dates = []\n",
    "    for i in range(total_week_count):\n",
    "        selected_dates.append(str(start_date + i*days_in_a_week))\n",
    "    selected_has_sample_df = selected_has_sample_df[selected_has_sample_df['Week Start'].isin(selected_dates)]\n",
    "    \n",
    "    selected_has_sample_df = selected_has_sample_df.pivot_table(values=['Missing Sample'], index=['Site'], columns=['Week Start'])\n",
    "    selected_has_sample_df.columns = selected_has_sample_df.columns.droplevel()\n",
    "    return selected_has_sample_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_figure5_data(target_mutation, merged_df, coverage_df, has_sample_df, start_date=date(2021, 4, 15), end_date=date(2022, 10, 1)):\n",
    "    selected_cr_df, date_site_count, aa_mut_label = select_single_cryptic(target_mutation, merged_df)\n",
    "\n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "\n",
    "    selected_coverage_df = build_select_coverage_df(coverage_df, target_mutation, start_date, end_date)\n",
    "    coverage_mask = selected_coverage_df.pivot_table(values=['Not Valid'], index=['Site'], columns=['Week Start']).fillna(True)\n",
    "    coverage_mask.columns = coverage_mask.columns.droplevel()\n",
    "    selected_cr_heatmap_array = selected_cr_df.pivot_table(values=['Combined Freq'], index=['Site'], columns=['Week Start'])\n",
    "    selected_cr_heatmap_array.columns = selected_cr_heatmap_array.columns.droplevel()\n",
    "    for i in list(set(coverage_mask.columns) - set(selected_cr_heatmap_array.columns)):\n",
    "        selected_cr_heatmap_array[i] = 0\n",
    "    selected_cr_heatmap_array = selected_cr_heatmap_array[coverage_mask.columns].fillna(0)\n",
    "    sorted_index = selected_cr_heatmap_array.astype(bool).sum(axis=1).sort_values(ascending=False).index\n",
    "    sorted_index = selected_cr_heatmap_array.loc[sorted_index].ne(0).idxmax(axis=1).sort_values(ascending=False).index\n",
    "    bar_plot_df = pd.DataFrame(selected_cr_heatmap_array.astype(bool).sum(axis=0), \n",
    "                                   columns=['Count'])\n",
    "    \n",
    "    missing_sample_df = build_select_missing_sample_df(has_sample_df, start_date, end_date)\n",
    "    \n",
    "    missing_sample_df = missing_sample_df.loc[sorted_index]\n",
    "    sorted_heatmap_array = selected_cr_heatmap_array.loc[sorted_index]\n",
    "    sorted_mask = coverage_mask.loc[sorted_index]\n",
    "\n",
    "    try:\n",
    "        assert np.logical_and(sorted_mask.values, sorted_heatmap_array.astype(bool).values).sum() == False\n",
    "    except:\n",
    "        print(\"Inconsistency between coverage data between cryptic lineage data.\")\n",
    "        \n",
    "    return selected_cr_df, date_site_count, aa_mut_label, sorted_heatmap_array, sorted_mask, missing_sample_df, bar_plot_df, selected_coverage_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_cryptic_mutations = ['A26530G;C26577G;G26634A',\n",
    "                              'C6402T;G6456A',\n",
    "                              'A26530G;C26533T;C26542A;T26545G',\n",
    "                              'A26530G;T26545G',\n",
    "                              'A27259C;C27335T;A27344T;A27345T',\n",
    "                              'T29029C;A29039T',\n",
    "                              'A26530G;C26577G;C26625A',\n",
    "                              'C10449A;T10459C',\n",
    "                              'T15682A;T15685A',\n",
    "                              'A24966T;C25000T',\n",
    "                              'A27344T;A27345T;A27354G',\n",
    "                              'A29039T;G29049A']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_cryptic_df[filtered_cryptic_df['Present Weeks'] > 10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_mutation = selected_cryptic_mutations[12-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_cr_df, date_site_count, aa_mut_label, selected_cr_heatmap_array, sorted_mask, missing_sample_df, bar_plot_df, selected_coverage_df = \\\n",
    "build_figure5_data(target_mutation, merged_df, coverage_df, has_sample_df)\n",
    "depth_df = selected_coverage_df[selected_coverage_df['Not Valid'] == False].groupby('Week Start')[['29039', '29049']].mean()\n",
    "depth_df = depth_df.merge(bar_plot_df, how='right', left_index=True, right_index=True)\n",
    "depth_df['depth_total'] = (depth_df['29039'] + depth_df['29049'])/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customized Colormap, where 0-0.1 is light blue.\n",
    "vmin = 0\n",
    "vmax = 0.15\n",
    "v_lower_threshold = 0.01\n",
    "total_colors = 300\n",
    "v_lower_num_colors = int((v_lower_threshold/vmax)*total_colors)\n",
    "skip_min_colors = 0.1\n",
    "\n",
    "cmap_reds = plt.get_cmap('YlOrRd')\n",
    "colors = ['lightblue']*v_lower_num_colors +\\\n",
    "[cmap_reds((i+total_colors/(1-skip_min_colors)*skip_min_colors+v_lower_num_colors) / total_colors/(1-skip_min_colors)) for i in range(0, total_colors-v_lower_num_colors)]\n",
    "cmap = LinearSegmentedColormap.from_list('', colors, total_colors)\n",
    "\n",
    "# Master Layout\n",
    "fontsize = 12\n",
    "fig, axes  = plt.subplots(2, 2, figsize=(15, 10),\n",
    "                          sharex=False,\n",
    "                          gridspec_kw={'height_ratios': [4, 1],\n",
    "                                       'width_ratios': [37, 1],\n",
    "                                       'wspace': 0.04,\n",
    "                                       'hspace': 0.15})\n",
    "\n",
    "\n",
    "# Heatmap\n",
    "ax = axes[0][0]\n",
    "ax1 = sns.heatmap(selected_cr_heatmap_array, \n",
    "                  mask=sorted_mask,\n",
    "                  ax=axes[0,0],\n",
    "                  cbar_ax=axes[0,1],\n",
    "                  linewidths=0.2,\n",
    "                  cmap=cmap,\n",
    "                  cbar_kws={'label': 'Allele Frequency'\n",
    "                            },\n",
    "                  square=False,\n",
    "                  vmin=vmin,\n",
    "                  vmax=vmax)\n",
    "\n",
    "sns.heatmap(missing_sample_df, \n",
    "            cmap=ListedColormap(['white']), \n",
    "            linecolor='white', \n",
    "            linewidths=0.2,\n",
    "            cbar=False, mask=(missing_sample_df != 1),\n",
    "            ax=axes[0,0])\n",
    "\n",
    "ax1.axhline(y = 0, color='k',linewidth = 2)\n",
    "ax1.axhline(y = selected_cr_heatmap_array.shape[0], color = 'k', linewidth = 2)\n",
    "ax1.axvline(x = 0, color = 'k',linewidth = 2)\n",
    "ax1.axvline(x = selected_cr_heatmap_array.shape[1], color = 'k', linewidth = 2)\n",
    "\n",
    "ax.set_title(f'{target_mutation} ({aa_mut_label}) Detection in Houston Wastewater', fontsize=fontsize)\n",
    "\n",
    "axes[0,1].set_ylabel(\"Allele Frequency\",size=12);\n",
    "\n",
    "axes[0,0].set_facecolor(mcolors.CSS4_COLORS['lightgrey'])\n",
    "axes[0,0].set_ylabel(\"WWTP\")\n",
    "axes[0,0].set_xlabel(\"\")\n",
    "axes[0,0].set(xticklabels=[],  xlabel=None)\n",
    "ax1.tick_params(left=False, bottom=False)\n",
    "\n",
    "# Colormap for the heatmap\n",
    "axes[0,1].set_facecolor(\"w\")\n",
    "axes[0,1].xaxis.label.set_size(fontsize)\n",
    "\n",
    "ax.text(-0.07, 1.05, 'a', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "# Customized Legend for the heatmap\n",
    "legend_elements = [Line2D([0], [0], marker='s', color='w', label='Not Detected',\n",
    "                          markerfacecolor=cmap(0), markersize=8, markeredgecolor=cmap(0)),\n",
    "                   Line2D([0], [0], marker='s', color='w', label='No Coverage',\n",
    "                          markerfacecolor=mcolors.CSS4_COLORS['lightgrey'], markersize=8, markeredgecolor=mcolors.CSS4_COLORS['lightgrey']),\n",
    "                   Line2D([0], [0], marker='s', color='w', label='Missing Sample',\n",
    "                          markerfacecolor='w', markersize=8, markeredgecolor='black')]\n",
    "legend = axes[0,0].legend(handles=legend_elements, fontsize=fontsize, loc='upper center', bbox_to_anchor=(0.5, 0), ncol=len(legend_elements), frameon=False)\n",
    "frame = legend.get_frame()\n",
    "frame.set_facecolor('w')\n",
    "\n",
    "# Bar plot\n",
    "ax = axes[1][0]\n",
    "\n",
    "x_ticklabels = bar_plot_df.index.strftime('%Y-%m-%d').to_list()\n",
    "x_tick_pos = [i + 0.5 for i in range(len(bar_plot_df.index))]\n",
    "\n",
    "ax.bar(x_tick_pos, bar_plot_df['Count'].values, width=0.8, color='g')\n",
    "\n",
    "ax_twinx = ax.twinx()\n",
    "ax_twinx.plot(x_tick_pos, depth_df['depth_total'].fillna(0).values, color='b', linestyle='dashed')\n",
    "ax_twinx.set_ylim(0, depth_df['depth_total'].fillna(0).values.max()*1.2)\n",
    "ax_twinx.set_ylabel('Mean Coverage', fontsize=fontsize)\n",
    "\n",
    "ax.set_xlim([0, len(bar_plot_df.index)])\n",
    "ax.set_xticks(x_tick_pos[::4])\n",
    "ax.set_xticklabels(x_ticklabels[::4], rotation=90, fontsize=fontsize)\n",
    "\n",
    "ax.set_yticks(np.arange(0,20,5))\n",
    "ax.set_yticklabels(['   0', '   5', '  10', '  15'], fontsize=fontsize)\n",
    "ax.set_ylabel('Site Count', fontsize=fontsize)\n",
    "ax.set_ylim(0, bar_plot_df['Count'].values.max()*1.2)\n",
    "\n",
    "ax.text(-0.07, 1.15, 'b', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "# Remove extract plot\n",
    "axes[1,1].set_visible(False)\n",
    "\n",
    "fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/figure_6.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 7 - CR3 & CR 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def padding_date(date_string):\n",
    "    return str(date_string).zfill(8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## VCFs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_wastewater_vcf(has_sample_df, valid_WWTPs):\n",
    "    vcf_metadata = pd.read_csv('/home/Users/yl181/crykey_bu/input_metadata.tsv', sep='\\t')\n",
    "    vcf_metadata['Date'] = vcf_metadata['Sample_Collection_Date'].map(padding_date)\n",
    "    vcf_metadata['Date'] = pd.to_datetime(vcf_metadata['Date'], format='%m%d%Y')\n",
    "    vcf_metadata['Week Start'] = vcf_metadata['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "    \n",
    "    #print(vcf_metadata.columns)\n",
    "    \n",
    "    records = []\n",
    "    for idx, row in vcf_metadata.iterrows():\n",
    "        try:\n",
    "            if row['WWTP'] in valid_WWTPs and has_sample_df[(has_sample_df['Collection Date'] == row['Date']) & (has_sample_df['Site'] == row['WWTP'])]['Has Sample'].values[0]:\n",
    "                vcf_path = row['VCF']\n",
    "                if os.path.exists(vcf_path):\n",
    "                    for record in vcf.Reader(open(vcf_path, 'r')):\n",
    "                        ref = str(record.REF)\n",
    "                        pos = str(record.POS)\n",
    "                        alt = str(record.ALT[0])\n",
    "\n",
    "                        if len(ref) == 1 and len(alt) == 1:\n",
    "                            mut = ref+pos+alt\n",
    "                            records.append({'Week Start': row['Week Start'],\n",
    "                                            'WWTP': row['WWTP'],\n",
    "                                'Nt Mutation': mut,\n",
    "                                'Depth': int(record.INFO['DP']),\n",
    "                                'AF': float(record.INFO['AF'])})\n",
    "                    records.append({'Week Start': row['Week Start'],\n",
    "                                'WWTP': row['WWTP'],\n",
    "                                'Nt Mutation': 'Valid_VCF',\n",
    "                                'Depth': 0,\n",
    "                                'AF': 0})\n",
    "        except:\n",
    "            continue\n",
    "    vcf_df = pd.DataFrame(records)\n",
    "    return vcf_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "wastewater_vcf_df = parse_wastewater_vcf(has_sample_df, valid_WWTPs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_clinical_vcf():\n",
    "    clinical_metadata_df = pd.read_csv('/home/Users/yl181/wastewater/quarc_clinical_sampling/PRJNA764181.filtered.csv')\n",
    "    clinical_metadata_df['Date'] = pd.to_datetime(clinical_metadata_df['Collection_Date'], format='%Y-%m-%d')\n",
    "    clinical_metadata_df['Week Start'] = clinical_metadata_df['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "    \n",
    "    valid_samples = []\n",
    "    records = []\n",
    "    for idx, row in clinical_metadata_df.iterrows():\n",
    "        run_id = row['Run']\n",
    "        vcf_path = f'/home/Users/yl181/wastewater/quarc_clinical_sampling/Harvest_Variant_Outputs_Houston/{idx}_out/vcf_files_filtered/{run_id}.vcf'\n",
    "        if os.path.exists(vcf_path):\n",
    "            valid_samples.append(True)\n",
    "            for record in vcf.Reader(open(vcf_path, 'r')):\n",
    "                ref = str(record.REF)\n",
    "                pos = str(record.POS)\n",
    "                alt = str(record.ALT[0])\n",
    "\n",
    "                if len(ref) == 1 and len(alt) == 1:\n",
    "                    mut = ref+pos+alt\n",
    "                    records.append({'Week Start': row['Week Start'],\n",
    "                                    'WWTP': run_id,\n",
    "                        'Nt Mutation': mut,\n",
    "                        'Depth': int(record.INFO['DP']),\n",
    "                        'AF': float(record.INFO['AF'])})\n",
    "            records.append({'Week Start': row['Week Start'],\n",
    "                    'WWTP': run_id,\n",
    "                    'Nt Mutation': 'Valid_VCF',\n",
    "                    'Depth': 0,\n",
    "                    'AF': 0})\n",
    "        else:\n",
    "            valid_samples.append(False)\n",
    "            #print(vcf_path)\n",
    "    vcf_df = pd.DataFrame(records)\n",
    "    clinical_metadata_df['Valid_VCF'] = valid_samples\n",
    "    return vcf_df, clinical_metadata_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "clinical_vcf_df, clinical_metadata_df = parse_clinical_vcf()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Individual SNV DataFrame CR12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filtering_clinical_samples(clinical_cryptic_df, strand_bias=False):\n",
    "    crykey_calls = []\n",
    "    if strand_bias:\n",
    "        p_values = []\n",
    "        sb_values = []\n",
    "    \n",
    "    for idx, row in clinical_cryptic_df.iterrows():\n",
    "        supp_dp = row['Support DP']\n",
    "        total_dp = row['Total DP']\n",
    "            \n",
    "        if strand_bias:\n",
    "            dp1, dp2, dp3, dp4 = row['DP1'], row['DP2'], row['DP3'], row['DP4']\n",
    "            if dp3+dp4 > 0 and (dp3/(dp3+dp4) > 0.85 or dp3/(dp3+dp4) < 1-0.85):\n",
    "                p_value = fisher_exact([[dp1, dp2], [dp3, dp4]])[1]\n",
    "                try:\n",
    "                    sb_value = abs((dp3/(dp1+dp3)) - (dp4/(dp2+dp4)))/((dp3+dp4)/(dp1+dp2+dp3+dp4))\n",
    "                except ZeroDivisionError:\n",
    "                    sb_value = 0\n",
    "            else:\n",
    "                p_value = 1\n",
    "                sb_value = 0  \n",
    "            p_values.append(p_value)\n",
    "            sb_values.append(sb_value)\n",
    "        \n",
    "        if row['Crykey'] == True:\n",
    "            crykey_call = (supp_dp >= 5) and (supp_dp/total_dp >= 0.02)\n",
    "            if strand_bias:\n",
    "                crykey_call = crykey_call and (sb_value < 1) and (p_value > 0.05)\n",
    "                \n",
    "        else:\n",
    "            crykey_call = False\n",
    "        crykey_calls.append(crykey_call)\n",
    "            \n",
    "    clinical_cryptic_df['Crykey'] = crykey_calls\n",
    "    if strand_bias:\n",
    "        clinical_cryptic_df['SB_p_value'] = p_values\n",
    "        clinical_cryptic_df['SB'] = sb_values\n",
    "    \n",
    "    return clinical_cryptic_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_wastewater_independent_snv_df(merged_df, valid_WWTPs, from_source=False):\n",
    "    if from_source:\n",
    "        independent_snv_df = pd.read_csv('/home/Users/yl181/wastewater/crykey_wastewater_split/ww_20_split_af_all_dates.csv')\n",
    "        independent_snv_df['Date'] = independent_snv_df['Date'].map(padding_date)\n",
    "        independent_snv_df['Date'] = pd.to_datetime(independent_snv_df['Date'], format='%m%d%Y')\n",
    "        independent_snv_df['Week Start'] = independent_snv_df['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "\n",
    "        independent_snv_df = independent_snv_df[independent_snv_df.Site.isin(valid_WWTPs)]\n",
    "\n",
    "        crykey_calls = []\n",
    "        for idx, row in independent_snv_df.iterrows():\n",
    "            if merged_df[(merged_df['Nt Mutations'] == row['Nt Mutations']) & (merged_df['Week Start'] == row['Week Start']) & (merged_df['Site'] == row['Site'])].shape[0] == 1:\n",
    "                crykey_calls.append(True)\n",
    "            else:\n",
    "                crykey_calls.append(False)\n",
    "        independent_snv_df['Crykey'] = crykey_calls\n",
    "    else:\n",
    "        independent_snv_df = pd.read_csv(os.path.join(fixed_results_dir, 'wastewater_12crs_independent_snvs.csv'), index_col=False)\n",
    "        independent_snv_df['Date'] = pd.to_datetime(independent_snv_df['Date'], format='%Y-%m-%d')\n",
    "        independent_snv_df['Week Start'] = pd.to_datetime(independent_snv_df['Week Start'], format='%Y-%m-%d')\n",
    "        independent_snv_df = independent_snv_df.drop(['Unnamed: 0'], axis=1)\n",
    "    return independent_snv_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_clinical_independent_snv_df(from_source=False):\n",
    "    if from_source:\n",
    "        clinical_metadata_df = pd.read_csv('/home/Users/yl181/wastewater/quarc_clinical_sampling/PRJNA764181.filtered.csv')\n",
    "        clinical_metadata_df = clinical_metadata_df[['Run', 'Collection_Date']].copy().set_index('Run')\n",
    "\n",
    "        clinical_independent_snv_df = pd.read_csv('/home/Users/yl181/wastewater/crykey_wastewater_split/houston_20_split_af.csv')\n",
    "        clinical_independent_snv_df['Date'] = clinical_independent_snv_df['Site'].map(clinical_metadata_df['Collection_Date'].to_dict())\n",
    "        clinical_independent_snv_df['Date'] = pd.to_datetime(clinical_independent_snv_df['Date'], format='%Y-%m-%d')\n",
    "        clinical_independent_snv_df['Week Start'] = clinical_independent_snv_df['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "\n",
    "        crykey_calls = []\n",
    "        for idx, row in clinical_independent_snv_df.iterrows():\n",
    "            if row['Support DP'] >= 5 and row['Combined Freq'] >= 0.01:   \n",
    "                crykey_filter = True\n",
    "            else:\n",
    "                crykey_filter = False\n",
    "\n",
    "            if crykey_filter:\n",
    "                vcf_all_called = True\n",
    "                for nt in row['Nt Mutations'].split(';'):\n",
    "                    vcf_all_called = clinical_vcf_df[(clinical_vcf_df['Run'] == row['Site']) & (clinical_vcf_df['Nt Mutation'] == nt)].shape[0] * vcf_all_called\n",
    "                if vcf_all_called:\n",
    "                    crykey_calls.append(True)\n",
    "                else:\n",
    "                    crykey_calls.append(False)\n",
    "            else:\n",
    "                crykey_calls.append(False)\n",
    "\n",
    "        clinical_independent_snv_df['Crykey'] = crykey_calls\n",
    "    else:\n",
    "        clinical_independent_snv_df = pd.read_csv(os.path.join(fixed_results_dir, 'clinical_12crs_independent_snvs.csv'), index_col=False)\n",
    "        clinical_independent_snv_df['Date'] = pd.to_datetime(clinical_independent_snv_df['Date'], format='%Y-%m-%d')\n",
    "        clinical_independent_snv_df['Week Start'] = pd.to_datetime(clinical_independent_snv_df['Week Start'], format='%Y-%m-%d')\n",
    "        clinical_independent_snv_df = clinical_independent_snv_df.drop(['Unnamed: 0'], axis=1)\n",
    "        \n",
    "    return clinical_independent_snv_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wastewater_isnv_df = parse_wastewater_independent_snv_df(merged_df, valid_WWTPs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clinical_isnv_df = parse_clinical_independent_snv_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "clinical_isnv_df = filtering_clinical_samples(clinical_isnv_df, strand_bias=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lineage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def group_lineage(idx):\n",
    "    if idx.startswith('AY.'):\n",
    "        return 'Delta'\n",
    "    elif idx == 'BA.1.15' or idx.startswith('BA.1.15.'):\n",
    "        return 'BA.1.15'\n",
    "    elif idx == 'BA.1.17' or idx.startswith('BA.1.17.'):\n",
    "        return 'BA.1.17'\n",
    "    elif idx == 'BA.1.18' or idx.startswith('BA.1.18.'):\n",
    "        return 'BA.1.18'\n",
    "    elif idx == 'BA.1.20' or idx.startswith('BA.1.20.'):\n",
    "        return 'BA.1.20'\n",
    "    elif idx == 'BA.1.1' or idx.startswith('BA.1.1.'):\n",
    "        return 'BA.1.1'\n",
    "    else:\n",
    "        if idx.startswith('BA') or idx == 'B.1.1.529':\n",
    "            return 'Omicron'\n",
    "        else:\n",
    "            print('Error', idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "houston_lineage_df = pd.read_csv(os.path.join(fixed_results_dir, 'houston_lineage_report.csv'))\n",
    "houston_lineage_df = houston_lineage_df[['taxon', 'lineage']]\n",
    "houston_lineage_dict = houston_lineage_df.set_index('taxon')['lineage'].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lineage_plot(clinical_isnv_df, target_mutation, ax):\n",
    "    category_names = ['Delta', 'BA.1.1', 'BA.1.15', 'BA.1.17', 'BA.1.18', 'BA.1.20', 'Omicron']\n",
    "    category_colors = plt.get_cmap('RdYlGn')(np.linspace(0.15, 0.85, len(category_names)))\n",
    "\n",
    "    clinical_support_df = clinical_isnv_df[clinical_isnv_df['Crykey'] == True]\n",
    "    clinical_support_df['Lineage'] = clinical_support_df['Site'].map(houston_lineage_dict)\n",
    "    clinical_support_df['Lineage Label'] = clinical_support_df['Lineage'].map(group_lineage)\n",
    "    lineage_counter = Counter(clinical_support_df[clinical_support_df['Nt Mutations'] == target_mutation]['Lineage Label'].values)\n",
    "    \n",
    "    ax.invert_yaxis()\n",
    "    ax.xaxis.set_visible(False)\n",
    "    ax.yaxis.set_visible(False)\n",
    "    ax.set_xlim(0, sum(lineage_counter.values()))\n",
    "    starts = 0\n",
    "    for i, (colname, color) in enumerate(zip(category_names, category_colors)):\n",
    "        widths = lineage_counter[colname]\n",
    "        rects = ax.barh(0, widths, left=starts, height=0.5,\n",
    "                        label=colname, color=color)\n",
    "        starts = starts+widths\n",
    "\n",
    "    ax.legend(title='Lineages of CR supported clinical samples', ncol=len(category_names), bbox_to_anchor=(0.5, 1),\n",
    "              loc='lower center')\n",
    "    ax.set_ylim(-0.2,0.2)\n",
    "    return lineage_counter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Plot Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_isnv_plot_data(target_mutation, vcf_df, valid_samples, min_depth=10, start_date=date(2021, 3, 1), end_date=date(2022, 11, 3)):\n",
    "    \"\"\"\n",
    "    For individual SNV plot\n",
    "    \"\"\"\n",
    "    def get_weekly_stats(isnv, weekly_vcf_df, valid_samples, af_min=0.02):\n",
    "        afs = np.array(weekly_vcf_df[(weekly_vcf_df['Nt Mutation'] == isnv) & (weekly_vcf_df['AF'] >= af_min) & (weekly_vcf_df['WWTP'].isin(valid_samples))]['AF'].values)\n",
    "        valid_sample_count = len(valid_samples)\n",
    "        try:\n",
    "            prevalence = len(afs)/valid_sample_count\n",
    "            if prevalence>1:\n",
    "                print(isnv, len(afs), valid_sample_count)\n",
    "        except ZeroDivisionError:\n",
    "            prevalence = np.nan\n",
    "            \n",
    "        if len(afs) > 0:\n",
    "            min_af = abs(np.nanmin(afs) - afs.mean())\n",
    "            max_af = abs(np.nanmax(afs) - afs.mean())\n",
    "        else:\n",
    "            min_af = np.nan\n",
    "            max_af = np.nan\n",
    "            \n",
    "        return afs.mean(), min_af, max_af, afs.std(), prevalence\n",
    "    \n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    week_offset = timedelta(days = 7)\n",
    "    \n",
    "    isnv_mean_dict = defaultdict(list)\n",
    "    isnv_min_dict = defaultdict(list)\n",
    "    isnv_max_dict = defaultdict(list)\n",
    "    isnv_std_dict = defaultdict(list)\n",
    "    isnv_prev_dict = defaultdict(list)\n",
    "    for i in range(0, total_week_count):\n",
    "        week_start = start_date + i*week_offset        \n",
    "        weekly_vcf_df = vcf_df[vcf_df['Week Start'] == pd.to_datetime(week_start)]\n",
    "        for isnv in target_mutation.split(\";\"):\n",
    "            mean, min_af, max_af, std, prevalence = get_weekly_stats(isnv, weekly_vcf_df, valid_samples[i], af_min=0.02)\n",
    "            isnv_mean_dict[isnv].append(mean)\n",
    "            isnv_min_dict[isnv].append(min_af)\n",
    "            isnv_max_dict[isnv].append(max_af)\n",
    "            isnv_std_dict[isnv].append(std)\n",
    "            isnv_prev_dict[isnv].append(prevalence)\n",
    "    return isnv_mean_dict, isnv_min_dict, isnv_max_dict, isnv_std_dict, isnv_prev_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cr_plot_data(target_mutation, isnv_df, min_depth=10, start_date=date(2021, 3, 1), end_date=date(2022, 11, 3)):\n",
    "    \"\"\"\n",
    "    For CR plots\n",
    "    \"\"\"\n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    week_offset = timedelta(days = 7)\n",
    "    \n",
    "    selected_isnv_df = isnv_df[isnv_df['Nt Mutations'] == target_mutation].copy()\n",
    "\n",
    "    means = []\n",
    "    mins = []\n",
    "    maxs = []\n",
    "    stds = []\n",
    "    coverages = []\n",
    "    prevalences = []\n",
    "    \n",
    "    total_sample_counts = []\n",
    "    valid_sample_counts = []\n",
    "    valid_samples = []\n",
    "    detected_sample_counts = []\n",
    "    \n",
    "    for i in range(0, total_week_count):\n",
    "        week_start = start_date + i*week_offset \n",
    "        \n",
    "        weekly_selected_isnv_df = selected_isnv_df[selected_isnv_df['Week Start'] == pd.to_datetime(week_start)]\n",
    "        total_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        weekly_selected_isnv_df = weekly_selected_isnv_df[isnv_df['Total DP'] >= min_depth]\n",
    "        valid_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        valid_sample = weekly_selected_isnv_df['Site'].values\n",
    "        \n",
    "        afs = list(weekly_selected_isnv_df[(weekly_selected_isnv_df['Crykey'] == True)]['Combined Freq'].values)\n",
    "        mean = np.nanmean(afs)\n",
    "        std = np.nanstd(afs)\n",
    "        valid_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        try:\n",
    "            prevalence = len(afs)/valid_sample_count\n",
    "        except ZeroDivisionError:\n",
    "            prevalence = np.nan\n",
    "        coverage = weekly_selected_isnv_df['Total DP'].mean()\n",
    "        \n",
    "        means.append(mean)\n",
    "        try:\n",
    "            mins.append(abs(np.nanmin(afs)-mean))\n",
    "            maxs.append(abs(np.nanmax(afs)-mean))\n",
    "        except ValueError:\n",
    "            mins.append(np.nan)\n",
    "            maxs.append(np.nan)\n",
    "        stds.append(std)\n",
    "        prevalences.append(prevalence)\n",
    "        coverages.append(coverage)\n",
    "        \n",
    "        total_sample_counts.append(total_sample_count)\n",
    "        valid_sample_counts.append(valid_sample_count)\n",
    "        valid_samples.append(valid_sample)\n",
    "        detected_sample_counts.append(len(afs))\n",
    "        \n",
    "    return means, mins, maxs, stds, prevalences, coverages, total_sample_counts, valid_sample_counts, detected_sample_counts, valid_samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def date_labels(start_date=date(2021, 3, 1), end_date=date(2022, 11, 3)):\n",
    "    date_labels = []\n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    week_offset = timedelta(days = 7)\n",
    "    \n",
    "    for i in range(0, total_week_count):\n",
    "        start = start_date + i*week_offset\n",
    "        date_labels.append(start.strftime('%Y-%m-%d'))\n",
    "        \n",
    "    return date_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "date_labels = date_labels()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_mutation = selected_cryptic_mutations[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot3axes(axes, cr_no, vcf_df, isnv_df, data_type, width, total_weeks):\n",
    "    total_width = 0.8\n",
    "    means, mins, maxs, stds, prevalences, coverages, total_sample_counts, valid_sample_counts, detected_sample_counts, valid_samples = get_cr_plot_data(selected_cryptic_mutations[cr_no-1], isnv_df)\n",
    "    \n",
    "    ax = axes[1]\n",
    "    positions = np.arange(total_weeks)\n",
    "    ax.errorbar(positions, \n",
    "                means, \n",
    "                [mins, maxs], \n",
    "                marker='o', markersize=4,\n",
    "                linestyle='dotted', \n",
    "                linewidth=2, alpha=1, color=f'C0', \n",
    "                label=f'CR{cr_no} AF')\n",
    "    ax.legend(loc='upper left', framealpha=0.4)\n",
    "    ax.set_ylim(0,1.05)\n",
    "    ax.set_ylabel(f'AF\\n({data_type})')\n",
    "    \n",
    "    isnv_mean_dict, isnv_min_dict, isnv_max_dict, isnv_std_dict, isnv_prev_dict = get_isnv_plot_data(selected_cryptic_mutations[cr_no-1], vcf_df, valid_samples)\n",
    "    \n",
    "    ax = ax.twinx()\n",
    "    prevelence_barplot = ax.bar(x=positions,\n",
    "                                height=prevalences,\n",
    "                                alpha=0.3, color=f'C0',\n",
    "                                label=f'CR{cr_no} Prevalence')\n",
    "    ax.legend(loc='upper right', framealpha=0.4)\n",
    "    ax.set_ylim(0,1.05)\n",
    "    ax.set_ylabel(f'Prevalence\\n({data_type})')\n",
    "    \n",
    "    ax = axes[0]\n",
    "    \n",
    "    for i, isnv in enumerate(isnv_prev_dict):\n",
    "        positions = np.arange(total_weeks)-total_width/2+(i)*width+0.5*width\n",
    "        ax.errorbar(positions, \n",
    "                    isnv_mean_dict[isnv], \n",
    "                    [isnv_min_dict[isnv], isnv_max_dict[isnv]], \n",
    "                    marker='o', markersize=4,\n",
    "                    linestyle='dotted', \n",
    "                    linewidth=2, alpha=1, color=f'C{i+1}',\n",
    "                    label=f'{isnv} AF')\n",
    "    ax.legend(loc='upper left', framealpha=0.4)\n",
    "    ax.set_ylim(0,1.05)\n",
    "    ax.set_ylabel(f'AF\\n({data_type})')\n",
    "\n",
    "    ax = ax.twinx()\n",
    "    for i, isnv in enumerate(isnv_prev_dict):    \n",
    "        # Prevelence Bars\n",
    "        positions = np.arange(total_weeks)-total_width/2+(i)*width+0.5*width\n",
    "        prevelence_barplot = ax.bar(x=positions,\n",
    "                             height=isnv_prev_dict[isnv],\n",
    "                             alpha=0.5,\n",
    "                             width=width, color=f'C{i+1}', \n",
    "                             label=f\"{isnv} Prevalence\")\n",
    "    ax.legend(loc='upper right', framealpha=0.4)\n",
    "    ax.set_ylim(0,1.05)\n",
    "    ax.set_ylabel(f'Prevalence\\n({data_type})')\n",
    "\n",
    "    ax = axes[2]\n",
    "    coverage_barplot = ax.bar(x=np.arange(total_weeks),\n",
    "                              height=coverages,\n",
    "                              alpha=0.8, color='black',\n",
    "                              label=f'Mean Coverage at CR{cr_no} Positions')\n",
    "    \n",
    "    ax.legend(loc='upper right', framealpha=0.4)\n",
    "    ax.set_ylabel(f'Coverage\\n({data_type})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cr_plot(cr_no, cr_label, save_fig=False):\n",
    "    fontsize = 12\n",
    "    fig, axes  = plt.subplots(7, 1, figsize=(15, 10), sharex=True, constrained_layout=True, gridspec_kw={'height_ratios': [2, 2, 1, 2, 2, 1, 1], 'wspace':0.2})\n",
    "\n",
    "    total_weeks = len(date_labels)\n",
    "    total_width = 0.8\n",
    "    width = total_width*1/(len(selected_cryptic_mutations[cr_no-1].split(\";\")))\n",
    "    linewidth = 2\n",
    "\n",
    "    plot3axes(axes[0:3], cr_no, wastewater_vcf_df, wastewater_isnv_df, \"Wastewater\", width, total_weeks)\n",
    "    plot3axes(axes[3:6], cr_no, clinical_vcf_df, clinical_isnv_df, \"Clinical\", width, total_weeks)\n",
    "\n",
    "    ax = axes[-1]\n",
    "    positions = np.arange(total_weeks)\n",
    "\n",
    "    _, _, _, _, _, _, total_sample_counts, valid_sample_counts, detected_sample_counts, _ = get_cr_plot_data(selected_cryptic_mutations[cr_no-1], wastewater_isnv_df)\n",
    "\n",
    "    ax.bar(positions, total_sample_counts, color=mcolors.CSS4_COLORS['lightgrey'], alpha=1, label='Insufficient Coverage')\n",
    "    ax.bar(positions, valid_sample_counts, color=cmap(0), label=f'CR{cr_no} Not Detected')\n",
    "    ax.bar(positions, detected_sample_counts, color=cmap(100), label=f'CR{cr_no} Detected')\n",
    "    ax.legend()\n",
    "    xticks = []\n",
    "    xticklabels = []\n",
    "    for i in range(total_weeks):\n",
    "        if i%4 == 0:\n",
    "            xticks.append(i)\n",
    "            xticklabels.append(date_labels[i])\n",
    "\n",
    "    ax.set_xticks(xticks)\n",
    "    ax.set_xticklabels(xticklabels, rotation=90)\n",
    "\n",
    "    ax.set_xlim(-0.6, total_weeks-0.4)\n",
    "    ax.set_ylim(0, 39)\n",
    "    ax.set_ylabel('WWTP\\nCount')\n",
    "\n",
    "    for ax in axes[:-1]:\n",
    "        ax.xaxis.set_tick_params(labelbottom=False)\n",
    "        ax.tick_params(bottom=False)\n",
    "\n",
    "    if save_fig:\n",
    "        fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/high_dpi/{cr_label}_{cr_no}.png', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")\n",
    "    else:\n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inner(cr_no, subfigure, fig, fig_label):\n",
    "    inner = gridspec.GridSpecFromSubplotSpec(10, 1,\n",
    "                                             height_ratios=[2, 2, 1.5, 1.5, 1.5, 2, 2, 1.5, 0.6, 0.5], \n",
    "                                             wspace=0.2,\n",
    "                                             subplot_spec=subfigure)\n",
    "    axes = []\n",
    "    ax0 = fig.add_subplot(inner[0])\n",
    "    axes.append(ax0)\n",
    "    for j in range(1,10,1):\n",
    "        if j != 9:\n",
    "            ax = fig.add_subplot(inner[j], sharex=ax0)\n",
    "        else:\n",
    "            ax = fig.add_subplot(inner[j])\n",
    "        axes.append(ax)\n",
    "    \n",
    "    axes[4].set_visible(False)\n",
    "    axes[-2].set_visible(False)\n",
    "\n",
    "    total_weeks = len(date_labels)\n",
    "    total_width = 0.8\n",
    "    width = total_width*1/(len(selected_cryptic_mutations[cr_no-1].split(\";\")))\n",
    "    linewidth = 2\n",
    "\n",
    "    plot3axes(axes[0:3], cr_no, wastewater_vcf_df, wastewater_isnv_df, \"Wastewater\", width, total_weeks)\n",
    "    plot3axes(axes[5:8], cr_no, clinical_vcf_df, clinical_isnv_df, \"Clinical\", width, total_weeks)\n",
    "    \n",
    "    lineage_plot(clinical_isnv_df, selected_cryptic_mutations[cr_no-1], axes[-1])\n",
    "    \n",
    "    ax = axes[3]\n",
    "    positions = np.arange(total_weeks)\n",
    "\n",
    "    _, _, _, _, _, _, total_sample_counts, valid_sample_counts, detected_sample_counts, _ = get_cr_plot_data(selected_cryptic_mutations[cr_no-1], wastewater_isnv_df)\n",
    "\n",
    "    ax.bar(positions, total_sample_counts, color=mcolors.CSS4_COLORS['lightgrey'], alpha=1, label='Insufficient Coverage')\n",
    "    ax.bar(positions, valid_sample_counts, color=cmap(0), label=f'CR{cr_no} Not Detected')\n",
    "    ax.bar(positions, detected_sample_counts, color=cmap(100), label=f'CR{cr_no} Detected')\n",
    "    ax.legend(loc='upper left', framealpha=0.4)\n",
    "    xticks = []\n",
    "    xticklabels = []\n",
    "    for i in range(total_weeks):\n",
    "        if i%4 == 0:\n",
    "            xticks.append(i)\n",
    "            xticklabels.append(date_labels[i])\n",
    "\n",
    "    ax.set_xticks(xticks)\n",
    "    ax.set_xticklabels(xticklabels, rotation=90)\n",
    "\n",
    "    ax.set_xlim(-0.6, total_weeks-0.4)\n",
    "    ax.set_ylim(0, 39)\n",
    "    ax.set_ylabel('WWTP\\nCount')\n",
    "\n",
    "    for idx, ax in enumerate(axes):\n",
    "        if idx != 3:\n",
    "            ax.xaxis.set_tick_params(labelbottom=False, bottom=False)\n",
    "        if idx == 5:\n",
    "            ax.xaxis.set_tick_params(labelbottom=False, bottom=False, top=True)\n",
    "            ax.set_xticks(xticks)\n",
    "            \n",
    "    for idx, ax in enumerate(axes):\n",
    "        ax.yaxis.set_label_coords(-0.032, 0.5)\n",
    "    \n",
    "    if fig_label != 'x':\n",
    "        ax0.text(-0.06, 1.05, fig_label, transform=ax0.transAxes,\n",
    "                fontsize=20, fontweight='bold', va='top')\n",
    "    else:\n",
    "        labels = ['a', 'b', 'c', 'd', '', 'e', 'f', 'g', '', 'h']\n",
    "        for idx, ax in enumerate(axes):\n",
    "            ax.text(-0.065, 1.05, labels[idx], transform=ax.transAxes,\n",
    "                fontsize=16, fontweight='bold', va='top')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, 12))\n",
    "outer = gridspec.GridSpec(1, 1, wspace=0.15, hspace=0.05)\n",
    "\n",
    "cr_nos = [8]\n",
    "fig_labels = ['x']\n",
    "for i in range(len(cr_nos)):\n",
    "    inner(cr_nos[i], outer[i], fig, fig_labels[i])\n",
    "fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/pdf/Figure_8.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplemenatary 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, 12))\n",
    "outer = gridspec.GridSpec(1, 1, wspace=0.15, hspace=0.05)\n",
    "\n",
    "cr_nos = [3]\n",
    "fig_labels = ['x']\n",
    "for i in range(len(cr_nos)):\n",
    "    inner(cr_nos[i], outer[i], fig, fig_labels[i])\n",
    "fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/pdf/supp_figure_3.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplemenatary 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, 12))\n",
    "outer = gridspec.GridSpec(1, 1, wspace=0.15, hspace=0.05)\n",
    "\n",
    "cr_nos = [5]\n",
    "fig_labels = ['x']\n",
    "for i in range(len(cr_nos)):\n",
    "    inner(cr_nos[i], outer[i], fig, fig_labels[i])\n",
    "fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/pdf/supp_figure_4.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplemenatary 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, 12))\n",
    "outer = gridspec.GridSpec(1, 1, wspace=0.15, hspace=0.05)\n",
    "\n",
    "cr_nos = [12]\n",
    "fig_labels = ['x']\n",
    "for i in range(len(cr_nos)):\n",
    "    inner(cr_nos[i], outer[i], fig, fig_labels[i])\n",
    "fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/pdf/supp_figure_5.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cr_plot(1, '_', save_fig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 7 - Clinical Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_us_data(from_source=False):\n",
    "    if from_source:\n",
    "        clinical_metadata_df = pd.read_csv('/home/Users/yl181/wastewater/quarc_clinical_sampling/filtered_us_df.csv', index_col=0)\n",
    "        clinical_metadata_df['Date'] = pd.to_datetime(clinical_metadata_df['Collection_Date'], format='%Y-%m-%d')\n",
    "        clinical_metadata_df['Week Start'] = clinical_metadata_df['Date'].dt.to_period('W-SUN').apply(lambda r: r.start_time)\n",
    "\n",
    "        runidx2weekstart = clinical_metadata_df[['Run Index', 'Week Start']].drop_duplicates().set_index('Run Index')['Week Start'].to_dict()\n",
    "\n",
    "        runidx2accession = clinical_metadata_df[['Run Index', 'SRA Accession ID']].drop_duplicates().set_index('Run Index')['SRA Accession ID'].to_dict()\n",
    "\n",
    "        valid_sample_dict = dict()\n",
    "        records = []\n",
    "        for run_id in clinical_metadata_df['Run Index'].unique():\n",
    "            vcf_path = f'/home/Users/yl181/wastewater/quarc_clinical_sampling/Harvest_Variant_Outputs_US/{run_id}_out/vcf_files_filtered/{runidx2accession[run_id]}.vcf'\n",
    "            if os.path.exists(vcf_path):\n",
    "                valid_sample_dict[run_id] = True\n",
    "                for record in vcf.Reader(open(vcf_path, 'r')):\n",
    "                    ref = str(record.REF)\n",
    "                    pos = str(record.POS)\n",
    "                    alt = str(record.ALT[0])\n",
    "\n",
    "                    if len(ref) == 1 and len(alt) == 1:\n",
    "                        mut = ref+pos+alt\n",
    "                        records.append({'Week Start': runidx2weekstart[run_id],\n",
    "                                        'WWTP': runidx2accession[run_id],\n",
    "                            'Nt Mutation': mut,\n",
    "                            'Depth': int(record.INFO['DP']),\n",
    "                            'AF': float(record.INFO['AF'])})\n",
    "                records.append({'Week Start': runidx2weekstart[run_id],\n",
    "                        'WWTP': runidx2accession[run_id],\n",
    "                        'Nt Mutation': 'Valid_VCF',\n",
    "                        'Depth': 0,\n",
    "                        'AF': 0})\n",
    "            else:\n",
    "                valid_sample_dict[run_id] = False\n",
    "        vcf_df = pd.DataFrame(records)\n",
    "\n",
    "        assert sum(valid_sample_dict.values()) == len(runidx2accession)\n",
    "\n",
    "        crykey_calls = []\n",
    "        for idx, row in clinical_metadata_df.iterrows():\n",
    "            if row['Support DP'] >= 5 and row['Combined Freq'] >= 0.01:   \n",
    "                crykey_filter = True\n",
    "            else:\n",
    "                crykey_filter = False\n",
    "\n",
    "            if crykey_filter:\n",
    "                vcf_all_called = True\n",
    "                for nt in row['Nt Mutations'].split(';'):\n",
    "                    vcf_all_called = vcf_df[(vcf_df['WWTP'] == row['SRA Accession ID']) & (vcf_df['Nt Mutation'] == nt)].shape[0] * vcf_all_called\n",
    "                if vcf_all_called:\n",
    "                    crykey_calls.append(True)\n",
    "                else:\n",
    "                    crykey_calls.append(False)\n",
    "            else:\n",
    "                crykey_calls.append(False)\n",
    "\n",
    "        clinical_metadata_df['Crykey'] = crykey_calls\n",
    "\n",
    "        return clinical_metadata_df\n",
    "    else:\n",
    "        clinical_metadata_df = pd.read_csv(os.path.join(fixed_results_dir, 'us_crykey.csv'), index_col=0)\n",
    "        clinical_metadata_df['Collection_Date'] = pd.to_datetime(clinical_metadata_df['Date'], format='%Y-%m-%d')\n",
    "        clinical_metadata_df['Week Start'] = pd.to_datetime(clinical_metadata_df['Week Start'], format='%Y-%m-%d')\n",
    "        clinical_metadata_df = clinical_metadata_df.rename({'SRA Accession ID': 'Site'}, axis=1)\n",
    "        return clinical_metadata_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_heatmap_data(target_mutation, isnv_df, regions, min_depth=10, start_date=date(2021, 12, 6), end_date=date(2022, 1, 30), exclude=False):\n",
    "    \"\"\"\n",
    "    For CR plots\n",
    "    \"\"\"\n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    week_offset = timedelta(days = 7)\n",
    "    \n",
    "    selected_isnv_df = isnv_df[isnv_df['Nt Mutations'] == target_mutation].copy()\n",
    "    \n",
    "    if not exclude:\n",
    "        for region in regions:\n",
    "            selected_isnv_df = selected_isnv_df[selected_isnv_df['Region'] == region]\n",
    "    else:\n",
    "        for region in regions:\n",
    "            selected_isnv_df = selected_isnv_df[selected_isnv_df['Region'] != region]\n",
    "\n",
    "    means = []\n",
    "    mins = []\n",
    "    maxs = []\n",
    "    stds = []\n",
    "    coverages = []\n",
    "    prevalences = []\n",
    "    \n",
    "    total_sample_counts = []\n",
    "    valid_sample_counts = []\n",
    "    valid_samples = []\n",
    "    detected_sample_counts = []\n",
    "    \n",
    "    for i in range(0, total_week_count):\n",
    "        week_start = start_date + i*week_offset \n",
    "        \n",
    "        weekly_selected_isnv_df = selected_isnv_df[selected_isnv_df['Week Start'] == pd.to_datetime(week_start)]\n",
    "        total_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        weekly_selected_isnv_df = weekly_selected_isnv_df[isnv_df['Total DP'] >= min_depth]\n",
    "        valid_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        valid_sample = weekly_selected_isnv_df['Site'].values\n",
    "        \n",
    "        afs = list(weekly_selected_isnv_df[(weekly_selected_isnv_df['Crykey'] == True)]['Combined Freq'].values)\n",
    "        mean = np.nanmean(afs)\n",
    "        std = np.nanstd(afs)\n",
    "        valid_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        try:\n",
    "            prevalence = len(afs)/total_sample_count\n",
    "        except ZeroDivisionError:\n",
    "            prevalence = np.nan\n",
    "        coverage = weekly_selected_isnv_df['Total DP'].mean()\n",
    "        \n",
    "        means.append(mean)\n",
    "        try:\n",
    "            mins.append(abs(np.nanmin(afs)-mean))\n",
    "            maxs.append(abs(np.nanmax(afs)-mean))\n",
    "        except ValueError:\n",
    "            mins.append(np.nan)\n",
    "            maxs.append(np.nan)\n",
    "        stds.append(std)\n",
    "        prevalences.append(prevalence)\n",
    "        coverages.append(coverage)\n",
    "        \n",
    "        total_sample_counts.append(total_sample_count)\n",
    "        valid_sample_counts.append(valid_sample_count)\n",
    "        valid_samples.append(valid_sample)\n",
    "        detected_sample_counts.append(len(afs))\n",
    "        \n",
    "    return means, mins, maxs, stds, prevalences, coverages, total_sample_counts, valid_sample_counts, detected_sample_counts, valid_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_heatmap_array(mut_label, us_clinical_isnv_df):\n",
    "    states = ['Maryland', 'Massachusetts', 'California', 'Colorado', 'Utah']\n",
    "    heatmap_array = []\n",
    "\n",
    "    for state in states:\n",
    "        means, mins, maxs, stds, prevalences, coverages, total_sample_counts, valid_sample_counts, detected_sample_counts, valid_samples = get_heatmap_data(mut_label, us_clinical_isnv_df, [state])\n",
    "        heatmap_array.append(prevalences)\n",
    "\n",
    "    means, mins, maxs, stds, prevalences, coverages, total_sample_counts, valid_sample_counts, detected_sample_counts, valid_samples = get_heatmap_data(mut_label, us_clinical_isnv_df, states, exclude=True)\n",
    "    heatmap_array.append(prevalences)\n",
    "    \n",
    "    return heatmap_array, states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_boxplot_data(target_mutation, isnv_df, min_depth=10, start_date=date(2021, 12, 6), end_date=date(2022, 1, 30)):\n",
    "    \"\"\"\n",
    "    For CR plots\n",
    "    \"\"\"\n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    week_offset = timedelta(days = 7)\n",
    "    \n",
    "    selected_isnv_df = isnv_df[isnv_df['Nt Mutations'] == target_mutation].copy()\n",
    "\n",
    "    means = []\n",
    "    mins = []\n",
    "    maxs = []\n",
    "    stds = []\n",
    "    coverages = []\n",
    "    prevalences = []\n",
    "    weekly_afs = []\n",
    "    \n",
    "    total_sample_counts = []\n",
    "    valid_sample_counts = []\n",
    "    valid_samples = []\n",
    "    detected_sample_counts = []\n",
    "    \n",
    "    for i in range(0, total_week_count):\n",
    "        week_start = start_date + i*week_offset \n",
    "        \n",
    "        weekly_selected_isnv_df = selected_isnv_df[selected_isnv_df['Week Start'] == pd.to_datetime(week_start)]\n",
    "        total_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        weekly_selected_isnv_df = weekly_selected_isnv_df[isnv_df['Total DP'] >= min_depth]\n",
    "        valid_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        valid_sample = weekly_selected_isnv_df['Site'].values\n",
    "        \n",
    "        afs = list(weekly_selected_isnv_df[(weekly_selected_isnv_df['Crykey'] == True)]['Combined Freq'].values)\n",
    "        weekly_afs.append(afs)\n",
    "        mean = np.nanmean(afs)\n",
    "        std = np.nanstd(afs)\n",
    "        valid_sample_count = weekly_selected_isnv_df.shape[0]\n",
    "        try:\n",
    "            prevalence = len(afs)/valid_sample_count\n",
    "        except ZeroDivisionError:\n",
    "            prevalence = np.nan\n",
    "        coverage = weekly_selected_isnv_df['Total DP'].mean()\n",
    "        \n",
    "        means.append(mean)\n",
    "        try:\n",
    "            mins.append(abs(np.nanmin(afs)-mean))\n",
    "            maxs.append(abs(np.nanmax(afs)-mean))\n",
    "        except ValueError:\n",
    "            mins.append(np.nan)\n",
    "            maxs.append(np.nan)\n",
    "        stds.append(std)\n",
    "        prevalences.append(prevalence)\n",
    "        coverages.append(coverage)\n",
    "        \n",
    "        total_sample_counts.append(total_sample_count)\n",
    "        valid_sample_counts.append(valid_sample_count)\n",
    "        valid_samples.append(valid_sample)\n",
    "        detected_sample_counts.append(len(afs))\n",
    "        \n",
    "    return means, mins, maxs, stds, prevalences, coverages, total_sample_counts, valid_sample_counts, detected_sample_counts, valid_samples, weekly_afs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_xtick_wastewater_data(target_mutation, merged_df, start_date=date(2021, 12, 6), end_date=date(2022, 1, 30)):\n",
    "    start_date = start_date - timedelta(days=start_date.weekday())\n",
    "    end_date = end_date - timedelta(days=end_date.weekday())\n",
    "    total_week_count = int((end_date - start_date).days/7) + 1\n",
    "    week_offset = timedelta(days = 7)\n",
    "    \n",
    "    ww_df = merged_df[merged_df['Nt Mutations'] == target_mutation]\n",
    "    \n",
    "    x_ticklabels = []\n",
    "    ww_counts = []\n",
    "\n",
    "    for i in range(0, total_week_count):\n",
    "        start = start_date + i*week_offset\n",
    "        end = start_date + (i+1)*week_offset\n",
    "\n",
    "        ww_count = ww_df[ww_df['Week Start'] == pd.to_datetime(start)].shape[0]\n",
    "        time_label = start.strftime('%Y-%m-%d')\n",
    "        x_ticklabels.append(time_label)\n",
    "        if ww_count > 0:\n",
    "            ww_counts.append(2)\n",
    "        else:\n",
    "            ww_counts.append(np.nan)\n",
    "    \n",
    "    x_ticklabels.append(end.strftime('%Y-%m-%d'))\n",
    "    \n",
    "    return x_ticklabels, ww_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lineage_values(df, target_mutation, state):\n",
    "    lineage_mutation_df = pd.pivot_table(df, values='Site', index='lineage', columns='Nt Mutations', \n",
    "                   aggfunc='count').fillna(0)\n",
    "    lineage_mutation_df.index.unique()\n",
    "\n",
    "    group_labels = []\n",
    "    for idx, _ in lineage_mutation_df.iterrows():\n",
    "        if idx.startswith('AY.'):\n",
    "            group_labels.append('Delta')\n",
    "        elif idx == 'BA.1.15' or idx.startswith('BA.1.15.'):\n",
    "            group_labels.append('BA.1.15')\n",
    "        elif idx == 'BA.1.17' or idx.startswith('BA.1.17.'):\n",
    "            group_labels.append('BA.1.17')\n",
    "        elif idx == 'BA.1.18' or idx.startswith('BA.1.18.'):\n",
    "            group_labels.append('BA.1.18')\n",
    "        elif idx == 'BA.1.20' or idx.startswith('BA.1.20.'):\n",
    "            group_labels.append('BA.1.20')\n",
    "        elif idx == 'BA.1.1' or idx.startswith('BA.1.1.'):\n",
    "            group_labels.append('BA.1.1')\n",
    "        else:\n",
    "            group_labels.append('Omicron')\n",
    "    lineage_mutation_df['Label'] = group_labels\n",
    "    \n",
    "    lineage_mutation_df = lineage_mutation_df.groupby(lineage_mutation_df['Label']).sum()\n",
    "    \n",
    "    lineage_mutation_df = lineage_mutation_df.rename({target_mutation: state}, axis=1).transpose()\n",
    "    return lineage_mutation_df\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lineage_plot_data(target_mutation, houston_clinical_isnv_df, us_clinical_isnv_df):\n",
    "    states = ['Houston', 'Maryland', 'Massachusetts', 'California', 'Colorado', 'Utah'][::-1]\n",
    "    lineage_array = []\n",
    "    for state in states:\n",
    "        if state == 'Houston':\n",
    "            houston_clinical_isnv_df['lineage'] = houston_clinical_isnv_df['Site'].map(houston_lineage_dict)\n",
    "            houston_df = houston_clinical_isnv_df[(houston_clinical_isnv_df['Crykey'] == True) & (houston_clinical_isnv_df['Nt Mutations'] == target_mutation)].copy()\n",
    "            lineage_mutation_df = get_lineage_values(houston_df, target_mutation, state).copy()\n",
    "        else:\n",
    "            temp_df = us_clinical_isnv_df[(us_clinical_isnv_df['Region'] == state) & (us_clinical_isnv_df['Crykey'] == True) & (us_clinical_isnv_df['Nt Mutations'] == target_mutation)]\n",
    "            lineage_mutation_df = get_lineage_values(temp_df, target_mutation, state).copy()\n",
    "        \n",
    "        if not lineage_mutation_df.empty:\n",
    "            lineage_array.append(lineage_mutation_df)\n",
    "    \n",
    "    return lineage_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clinical_plot(cr_no, cr_label, save_fig=False):\n",
    "    fig, axes  = plt.subplots(3, 1, figsize=(8, 10), sharex=False, constrained_layout=True,\n",
    "                             gridspec_kw={'height_ratios': [5, 3, 3], 'hspace':0.35})\n",
    "    target_mutation = selected_cryptic_mutations[cr_no-1]\n",
    "\n",
    "    fontsize = 12\n",
    "    # Barplot Prevalence Rate in Houston\n",
    "\n",
    "    x_ticklabels, ww_counts = get_xtick_wastewater_data(target_mutation, merged_df)\n",
    "    _, _, _, _, p_rates, _, _, valid_sample_counts, _, _, freqs = get_boxplot_data(target_mutation, houston_clinical_isnv_df)\n",
    "\n",
    "    ax = axes[0]\n",
    "    x = np.arange(8)\n",
    "    ax.bar(x, ww_counts, color='gray', alpha=0.3, width=1)\n",
    "    bar = ax.bar(x, p_rates, color='c')\n",
    "\n",
    "    ax.set_ylim(0,1)\n",
    "    ax.set_ylabel('Prevalence Rate in\\nHouston Samples', fontsize=fontsize)\n",
    "    ax.set_title(\", \".join(target_mutation.split(';')))\n",
    "\n",
    "    # Boxplot of AF in Houston\n",
    "    ax_twin = ax.twinx()\n",
    "    ax_twin.boxplot(freqs, \n",
    "                    positions=x,\n",
    "                    showmeans=False)\n",
    "    ax_twin.set_ylim(0, 0.6)\n",
    "    ax_twin.set_ylabel('Intra-host\\nAllele Frequency', fontsize=fontsize)\n",
    "\n",
    "    ax.set_xlim(-0.5,7.5)\n",
    "    ax.set_xticks(np.arange(9)-0.5)\n",
    "    ax.set_xticklabels(x_ticklabels, rotation=90, fontsize=12)\n",
    "    ax.xaxis.set_ticks_position('none')\n",
    "    ax.set_xlabel(' ', fontsize=fontsize)\n",
    "\n",
    "    ax.text(-0.42, 0.98, f'CR{cr_no}',\n",
    "                fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "    ax.text(-0.25, 1.1, f'{cr_label}', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "    # Heatmap\n",
    "\n",
    "    heatmap_array, state_labels = get_heatmap_array(target_mutation, us_clinical_isnv_df)\n",
    "    state_labels.append('Other States')\n",
    "\n",
    "    sns.heatmap(heatmap_array, linewidth=2, \n",
    "                ax=axes[1],\n",
    "                vmin=0,\n",
    "                vmax=1,\n",
    "                cmap='flare',\n",
    "                yticklabels=state_labels,\n",
    "                xticklabels=[],\n",
    "                annot=True, fmt=\".3f\",\n",
    "                cbar_kws={'label': 'Prevalence Rate\\nin non-Texas Samples'},)\n",
    "\n",
    "    axes[1].figure.axes[-1].yaxis.label.set_size(fontsize)\n",
    "\n",
    "    axes[1].xaxis.set_ticks_position('none')\n",
    "    axes[1].yaxis.set_ticks_position('none')\n",
    "    axes[1].set_yticklabels(state_labels, fontsize=fontsize)\n",
    "    axes[1].set_xlabel(' ', fontsize=fontsize)\n",
    "\n",
    "    # bar\n",
    "    df = pd.concat(get_lineage_plot_data(target_mutation, houston_clinical_isnv_df, us_clinical_isnv_df), axis=0).fillna(0)\n",
    "    df = df.div(df.sum(axis=1), axis=0)\n",
    "    df = df[['BA.1.1', 'BA.1.15', 'BA.1.17', 'BA.1.18', 'BA.1.20', 'Omicron']]\n",
    "\n",
    "    ax = axes[2]\n",
    "    y = df.index.to_list()\n",
    "    lineages = df.columns.to_list()\n",
    "\n",
    "    category_names = ['Delta', 'BA.1.1', 'BA.1.15', 'BA.1.17', 'BA.1.18', 'BA.1.20', 'Omicron']\n",
    "    category_colors = plt.get_cmap('RdYlGn')(np.linspace(0.15, 0.85, len(category_names)))\n",
    "\n",
    "    for idx, lineage in enumerate(lineages):\n",
    "        color_idx = category_names.index(lineage)\n",
    "        if idx == 0: \n",
    "            ax.barh(y, df[lineage].to_list(), align='center', height=.5, label=lineage, edgecolor='black', color=category_colors[color_idx])\n",
    "            sum_array = df[lineage].to_list()\n",
    "        else:\n",
    "            ax.barh(y, df[lineage].to_list(), align='center', height=.5, label=lineage, left=sum_array, edgecolor='black', color=category_colors[color_idx])\n",
    "            sum_array = list(map(add, sum_array, df[lineage].to_list())) \n",
    "    ax.set_yticks(y)\n",
    "    ax.set_yticklabels(y, fontsize=fontsize)\n",
    "    ax.set_xlabel('Consensus Pango Lineage of CR Supported Samples', fontsize=fontsize)\n",
    "\n",
    "    ax.set_xticks([])\n",
    "\n",
    "    ax.spines[\"top\"].set_visible(False)\n",
    "    ax.spines[\"bottom\"].set_visible(False)\n",
    "\n",
    "    ax.set_xlim(0,1)\n",
    "    ax.legend(loc=4, bbox_to_anchor=(1.25, 0, 0, 0))\n",
    "\n",
    "    if save_fig:\n",
    "        fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/pdf/supp_figure_6{cr_label}.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")\n",
    "    else:\n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "houston_clinical_isnv_df = clinical_isnv_df.copy()\n",
    "us_clinical_isnv_df = parse_us_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "us_clinical_isnv_df = filtering_clinical_samples(us_clinical_isnv_df, strand_bias=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "category_names = ['Delta', 'BA.1.1', 'BA.1.15', 'BA.1.17', 'BA.1.18', 'BA.1.20', 'Omicron']\n",
    "category_colors = plt.get_cmap('RdYlGn')(np.linspace(0.15, 0.85, len(category_names)))\n",
    "    \n",
    "fontsize = 14\n",
    "num = int(len(selected_cryptic_mutations)/2)\n",
    "fig, axes  = plt.subplots(num, 4, figsize=(12, num*2), sharex=False, constrained_layout=True, gridspec_kw={'width_ratios': [5, 1, 5, 1], 'wspace':0.35})\n",
    "\n",
    "x = np.arange(8)\n",
    "\n",
    "plotted_mut = []\n",
    "\n",
    "for ax_idx, mut_label in enumerate(selected_cryptic_mutations[-num:]):\n",
    "    ax = axes[ax_idx][2]\n",
    "    target_mutation = mut_label\n",
    "    x_ticklabels, ww_counts = get_xtick_wastewater_data(target_mutation, merged_df)\n",
    "    means, mins, maxs, stds, p_rates, _, _, valid_sample_counts, _, _, freqs = get_boxplot_data(target_mutation, houston_clinical_isnv_df)\n",
    "\n",
    "    ax.bar(x, ww_counts, color='gray', alpha=0.2, width=1)\n",
    "    bar = ax.bar(x, p_rates, color='c')\n",
    "    ax.scatter(x, ww_counts, marker=\"P\", s=60, color='black')\n",
    "    ax.set_ylim(0,1.3)\n",
    "    ax.set_title(\", \".join(mut_label.split(';')))\n",
    "    ax.xaxis.set_ticks_position('none')\n",
    "    plotted_mut.append(mut_label)\n",
    "\n",
    "    ax.text(0.01, 0.95, f'CR{ax_idx+7}', transform=ax.transAxes,\n",
    "        fontsize=14, fontweight='bold', va='top')\n",
    "    ax_twin = ax.twinx()\n",
    "    ax_twin.errorbar(x, means, [mins, maxs], linestyle='dotted', marker='o', markersize=4)\n",
    "    ax_twin.set_ylim(0, 0.6)\n",
    "    ax_twin.set_yticks([0,0.25,0.5])\n",
    "    ax_twin.set_yticklabels(['0','.25','.50'])\n",
    "    ax.set_xlim(-0.5,7.5)\n",
    "    ax.set_xticks(np.arange(9)-0.5)\n",
    "    ax.set_xticklabels([])\n",
    "    \n",
    "    ax2 = axes[ax_idx][3]\n",
    "    bottom = 0\n",
    "    df = pd.concat(get_lineage_plot_data(target_mutation, houston_clinical_isnv_df, us_clinical_isnv_df), axis=0).fillna(0)\n",
    "    df = df.div(df.sum(axis=1), axis=0)\n",
    "    temp_lineage_dict = df.loc['Houston'].to_dict()\n",
    "    for idx in temp_lineage_dict:\n",
    "        color_idx = category_names.index(idx)\n",
    "        p = ax2.bar(0, temp_lineage_dict[idx], 1, label=idx, bottom=bottom, edgecolor='black', color=category_colors[color_idx])\n",
    "        bottom += temp_lineage_dict[idx]\n",
    "        \n",
    "    \n",
    "    ax2.set_ylim(0,1)\n",
    "    ax2.set_xlim(-0.2,0.2)\n",
    "    ax2.xaxis.set_ticks_position('none')\n",
    "    ax2.set_yticks([])\n",
    "    ax2.set_xticklabels([])\n",
    "    \n",
    "ax2.set_xticks([0])\n",
    "ax2.set_xticklabels(['Lineage'], rotation=90, fontsize=12)    \n",
    "#x_ticklabels.append((sampling_date_start+8*week_offset).strftime('%Y-%m-%d'))\n",
    "ax.set_xlim(-0.5,7.5)\n",
    "ax.set_xticks(np.arange(9)-0.5)\n",
    "ax.set_xticklabels(x_ticklabels, rotation=90, fontsize=12)\n",
    "\n",
    "patches = []\n",
    "for j, facecolor in enumerate(category_colors[::-1]):\n",
    "    patches.append(Patch(facecolor=facecolor, edgecolor='black',\n",
    "                         label=category_names[::-1][j]))\n",
    "\n",
    "ax2.legend(handles=patches, loc=4, bbox_to_anchor=(2.75, 0, 0, 0))\n",
    "\n",
    "for ax_idx, mut_label in enumerate(selected_cryptic_mutations[:num]):\n",
    "    ax = axes[ax_idx][0]\n",
    "    target_mutation = mut_label\n",
    "    x_ticklabels, ww_counts = get_xtick_wastewater_data(target_mutation, merged_df)\n",
    "    means, mins, maxs, stds, p_rates, _, _, valid_sample_counts, _, _, freqs = get_boxplot_data(target_mutation, houston_clinical_isnv_df)\n",
    "\n",
    "    ax.bar(x, ww_counts, color='gray', alpha=0.2, width=1)\n",
    "    bar = ax.bar(x, p_rates, color='c')\n",
    "    ax.scatter(x, ww_counts, marker=\"P\", s=60, color='black')\n",
    "    ax.set_ylim(0,1.3)\n",
    "    ax.set_title(\", \".join(mut_label.split(';')))\n",
    "    ax.xaxis.set_ticks_position('none')\n",
    "    plotted_mut.append(mut_label)\n",
    "\n",
    "    ax.text(0.01, 0.95, f'CR{ax_idx+1}', transform=ax.transAxes,\n",
    "        fontsize=14, fontweight='bold', va='top')\n",
    "    ax_twin = ax.twinx()\n",
    "    ax_twin.errorbar(x, means, [mins, maxs], linestyle='dotted', marker='o', markersize=4,)\n",
    "    ax_twin.set_ylim(0, 0.6)\n",
    "    ax_twin.set_yticks([0,0.25,0.5])\n",
    "    ax_twin.set_yticklabels(['0','.25','.50'])\n",
    "    ax.set_xlim(-0.5,7.5)\n",
    "    ax.set_xticks(np.arange(9)-0.5)\n",
    "    ax.set_xticklabels([])\n",
    "    \n",
    "    ax2 = axes[ax_idx][1]\n",
    "    bottom = 0\n",
    "    df = pd.concat(get_lineage_plot_data(target_mutation, houston_clinical_isnv_df, us_clinical_isnv_df), axis=0).fillna(0)\n",
    "    df = df.div(df.sum(axis=1), axis=0)\n",
    "    temp_lineage_dict = df.loc['Houston'].to_dict()\n",
    "    for idx in temp_lineage_dict:\n",
    "        color_idx = category_names.index(idx)\n",
    "        p = ax2.bar(0, temp_lineage_dict[idx], 1, label=idx, bottom=bottom, edgecolor='black', color=category_colors[color_idx])\n",
    "        bottom += temp_lineage_dict[idx]\n",
    "    \n",
    "    ax2.set_ylim(0,1)\n",
    "    ax2.set_xlim(-0.2,0.2)\n",
    "    ax2.xaxis.set_ticks_position('none')\n",
    "    ax2.set_yticks([])\n",
    "    ax2.set_xticklabels([])\n",
    "    \n",
    "ax2.set_xticks([0])\n",
    "ax2.set_xticklabels(['Lineage'], rotation=90, fontsize=12)\n",
    "ax.set_xlim(-0.5,7.5)\n",
    "ax.set_xticks(np.arange(9)-0.5)\n",
    "ax.set_xticklabels(x_ticklabels, rotation=90, fontsize=12)\n",
    "\n",
    "fig.text(-0.01, 0.5, 'Prevalence Rate in Clinical Samples', va='center', rotation='vertical', fontsize=12)\n",
    "fig.text(.815, 0.5, 'Mean Allele Frequency', va='center', rotation='vertical', fontsize=12)\n",
    "\n",
    "fig.savefig(f'/home/Users/yl181/wastewater/quarc_figures/pdf/Figure_7.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Supplementary Figure 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def supp_figure_clinical_plot(cr_no, cr_label, fig):\n",
    "    axes  = fig.subplots(3, 1, sharex=False, constrained_layout=True,\n",
    "                             gridspec_kw={'height_ratios': [5, 3, 3], 'hspace':0.35})\n",
    "    target_mutation = selected_cryptic_mutations[cr_no-1]\n",
    "\n",
    "    fontsize = 12\n",
    "    # Barplot Prevalence Rate in Houston\n",
    "\n",
    "    x_ticklabels, ww_counts = get_xtick_wastewater_data(target_mutation, merged_df)\n",
    "    _, _, _, _, p_rates, _, _, valid_sample_counts, _, _, freqs = get_boxplot_data(target_mutation, houston_clinical_isnv_df)\n",
    "\n",
    "    ax = axes[0]\n",
    "    x = np.arange(8)\n",
    "    ax.bar(x, ww_counts, color='gray', alpha=0.3, width=1)\n",
    "    bar = ax.bar(x, p_rates, color='c')\n",
    "\n",
    "    ax.set_ylim(0,1)\n",
    "    ax.set_ylabel('Prevalence Rate in\\nHouston Samples', fontsize=fontsize)\n",
    "    ax.set_title(\", \".join(target_mutation.split(';')))\n",
    "\n",
    "    # Boxplot of AF in Houston\n",
    "    ax_twin = ax.twinx()\n",
    "    ax_twin.boxplot(freqs, \n",
    "                    positions=x,\n",
    "                    showmeans=False)\n",
    "    ax_twin.set_ylim(0, 0.6)\n",
    "    ax_twin.set_ylabel('Intra-host\\nAllele Frequency', fontsize=fontsize)\n",
    "\n",
    "    ax.set_xlim(-0.5,7.5)\n",
    "    ax.set_xticks(np.arange(9)-0.5)\n",
    "    ax.set_xticklabels(x_ticklabels, rotation=90, fontsize=12)\n",
    "    ax.xaxis.set_ticks_position('none')\n",
    "    ax.set_xlabel(' ', fontsize=fontsize)\n",
    "\n",
    "    ax.text(-0.42, 0.98, f'CR{cr_no}',\n",
    "                fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "    ax.text(-0.25, 1.1, f'{cr_label}', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "    # Heatmap\n",
    "\n",
    "    heatmap_array, state_labels = get_heatmap_array(target_mutation, us_clinical_isnv_df)\n",
    "    state_labels.append('Other States')\n",
    "\n",
    "    sns.heatmap(heatmap_array, linewidth=2, \n",
    "                ax=axes[1],\n",
    "                vmin=0,\n",
    "                vmax=1,\n",
    "                cmap='flare',\n",
    "                yticklabels=state_labels,\n",
    "                xticklabels=[],\n",
    "                annot=True, fmt=\".3f\",\n",
    "                cbar_kws={'label': 'Prevalence Rate\\nin non-Texas Samples'},)\n",
    "\n",
    "    axes[1].figure.axes[-1].yaxis.label.set_size(fontsize)\n",
    "\n",
    "    axes[1].xaxis.set_ticks_position('none')\n",
    "    axes[1].yaxis.set_ticks_position('none')\n",
    "    axes[1].set_yticklabels(state_labels, fontsize=fontsize)\n",
    "    axes[1].set_xlabel(' ', fontsize=fontsize)\n",
    "\n",
    "    # bar\n",
    "    df = pd.concat(get_lineage_plot_data(target_mutation, houston_clinical_isnv_df, us_clinical_isnv_df), axis=0).fillna(0)\n",
    "    df = df.div(df.sum(axis=1), axis=0)\n",
    "    df = df[['BA.1.1', 'BA.1.15', 'BA.1.17', 'BA.1.18', 'BA.1.20', 'Omicron']]\n",
    "\n",
    "    ax = axes[2]\n",
    "    y = df.index.to_list()\n",
    "    lineages = df.columns.to_list()\n",
    "\n",
    "    category_names = ['Delta', 'BA.1.1', 'BA.1.15', 'BA.1.17', 'BA.1.18', 'BA.1.20', 'Omicron']\n",
    "    category_colors = plt.get_cmap('RdYlGn')(np.linspace(0.15, 0.85, len(category_names)))\n",
    "\n",
    "    for idx, lineage in enumerate(lineages):\n",
    "        color_idx = category_names.index(lineage)\n",
    "        if idx == 0: \n",
    "            ax.barh(y, df[lineage].to_list(), align='center', height=.5, label=lineage, edgecolor='black', color=category_colors[color_idx])\n",
    "            sum_array = df[lineage].to_list()\n",
    "        else:\n",
    "            ax.barh(y, df[lineage].to_list(), align='center', height=.5, label=lineage, left=sum_array, edgecolor='black', color=category_colors[color_idx])\n",
    "            sum_array = list(map(add, sum_array, df[lineage].to_list())) \n",
    "    ax.set_yticks(y)\n",
    "    ax.set_yticklabels(y, fontsize=fontsize)\n",
    "    ax.set_xlabel('Consensus Pango Lineage of CR Supported Samples', fontsize=fontsize)\n",
    "\n",
    "    ax.set_xticks([])\n",
    "\n",
    "    ax.spines[\"top\"].set_visible(False)\n",
    "    ax.spines[\"bottom\"].set_visible(False)\n",
    "\n",
    "    ax.set_xlim(0,1)\n",
    "    ax.legend(loc=4, bbox_to_anchor=(1.25, 0, 0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplemenatary 6a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clinical_plot(5, 'a', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplemenatary 6b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clinical_plot(8, 'b', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 5 - Duration and Rarity of Cryptic Lineages Found in Houston Wastewater"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_supp_figure1_data(selected_data, step=100):\n",
    "    week_count = sorted(selected_data['Present Weeks'].unique())\n",
    "    site_count_per_week = []\n",
    "    week_count_height = []\n",
    "\n",
    "    for week_idx in week_count:\n",
    "        site_count_per_week.append(selected_data[selected_data['Present Weeks'] == week_idx]['Mean Site Occurance'].values)\n",
    "        week_count_height.append(len(selected_data[selected_data['Present Weeks'] == week_idx]['Mean Site Occurance'].values))\n",
    "        \n",
    "    occ_list = np.arange(0,14,1) * step\n",
    "    occ_bar_height = []\n",
    "    for occ_idx in occ_list:\n",
    "        occ_bar_height.append(selected_data[(selected_data['GISAID Count'] >= occ_idx) & (selected_data['GISAID Count'] < occ_idx+step)].shape[0])\n",
    "\n",
    "    return week_count, site_count_per_week, week_count_height, occ_list, occ_bar_height"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "step = 100\n",
    "week_count, site_count_per_week, week_count_height, occ_list, occ_bar_height = create_supp_figure1_data(filtered_cryptic_df, step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_cryptic_df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize = 12\n",
    "fig, axes  = plt.subplots(2, 2, figsize=(9, 9),\n",
    "                          gridspec_kw={'height_ratios': [3, 1], 'width_ratios': [3, 1], 'hspace':0.08, 'wspace':0.08})\n",
    "\n",
    "ax = axes[0][0]\n",
    "scatter = ax.scatter(filtered_cryptic_df['Present Weeks'], filtered_cryptic_df['GISAID Count'],\n",
    "           s=filtered_cryptic_df['Mean Site Occurance']*40,\n",
    "           c=filtered_cryptic_df['Mean Allele Freq'],\n",
    "           cmap='viridis',\n",
    "           alpha=0.4)\n",
    "\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax.set_ylabel('Occurrence of CR in GISAID EpiCoV', fontsize=fontsize)\n",
    "ax.set_ylim(-10,1400)\n",
    "ax.set_xlim(-2,38)\n",
    "ax.set_xticklabels([])\n",
    "\n",
    "legend1 = ax.legend(*scatter.legend_elements(num=7),\n",
    "                    loc=\"upper right\", title=\"Mean Allele Freq\")\n",
    "ax.add_artist(legend1)\n",
    "\n",
    "legend_elements = [Line2D([0], [0], marker='o', color='w', label='2',\n",
    "                          markerfacecolor='gray', markersize=np.sqrt(80)),\n",
    "                   Line2D([0], [0], marker='o', color='w', label='4',\n",
    "                          markerfacecolor='gray', markersize=np.sqrt(160)),\n",
    "                   Line2D([0], [0], marker='o', color='w', label='6',\n",
    "                          markerfacecolor='gray', markersize=np.sqrt(240))]\n",
    "ax.legend(handles=legend_elements, loc='upper center', title=\"Site Count\\nper Week\", bbox_to_anchor=(0.64, 1))\n",
    "                   \n",
    "ax = axes[1][0]\n",
    "bars = ax.bar(week_count, week_count_height, log=False)\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax.set_ylabel('CR Count', fontsize=fontsize)\n",
    "ax.set_xlabel('Number of Weeks being Detected', fontsize=fontsize)\n",
    "ax.set_xlim(-2,38)\n",
    "\n",
    "ax = axes[0][1]\n",
    "hbars = ax.barh(occ_list, occ_bar_height, align='edge', height=step*0.9, log=False)\n",
    "ax.set_ylim(-10,1400)\n",
    "ax.set_xlim(0,500)\n",
    "ax.set_yticklabels([])\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax.invert_xaxis()\n",
    "ax.set_xlabel('CR Count', fontsize=fontsize)\n",
    "\n",
    "axes[1][1].set_visible(False)\n",
    "\n",
    "fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/figure_5.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 2 - Performance Benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_supp_figure5_boxplot_data(fixed_results_dir):\n",
    "    \"\"\"\n",
    "    Load benchmark results and generate boxplot data\n",
    "    Benchmark is done by sampling 50 queries from each bin, \n",
    "    and bins are defined by the occurance of the query in GISAID\n",
    "    bins are =0, 1-5, 6-10, 11-25, 26-50, 51-100, 101-250, 251-500, >500\n",
    "    \"\"\"\n",
    "    process_time = []\n",
    "    with open(os.path.join(fixed_results_dir, 'benchmark_result_50.txt'), 'r') as performance_benchmark_f:\n",
    "        for line in performance_benchmark_f.readlines():\n",
    "            process_time_bin = [float(i) for i in line.strip().split(\",\")]\n",
    "            process_time.append(process_time_bin)\n",
    "\n",
    "    mean_query_times = []\n",
    "    for time_bin in process_time:\n",
    "        mean_query_times.append(np.mean(time_bin))\n",
    "    \n",
    "    return process_time, mean_query_times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_supp_figure5_hist_data(gisaid_count_range, mutation_rarity_df):\n",
    "    \"\"\"\n",
    "    Generate Histogram Data by creating bins defined by the occurance of the query in GISAID\n",
    "    \"\"\"\n",
    "    hist_bars = []\n",
    "    labels = []\n",
    "    for idx, upper_range in enumerate(gisaid_count_range):\n",
    "        lower_range = gisaid_count_range[idx-1]\n",
    "        if lower_range > upper_range:\n",
    "            lower_range = -1\n",
    "            labels.append(f\"0\")\n",
    "        else:\n",
    "            labels.append(f\"{lower_range+1}-{upper_range}\")\n",
    "        hist_bars.append(len(mutation_rarity_df[(mutation_rarity_df['GISAID Count'] > lower_range)&(mutation_rarity_df['GISAID Count'] <= upper_range)]))\n",
    "    hist_bars.append(len(mutation_rarity_df[(mutation_rarity_df['GISAID Count'] > gisaid_count_range[-1])]))\n",
    "    labels.append(f\"> {gisaid_count_range[-1]}\")\n",
    "\n",
    "    # cumulative precentage\n",
    "    total_set_count = sum(hist_bars)\n",
    "\n",
    "    cum_precentages = []\n",
    "    for idx, item in enumerate(hist_bars):\n",
    "        cum_precentages.append(sum(hist_bars[:idx+1])/total_set_count*100)\n",
    "        \n",
    "    return hist_bars, labels, total_set_count, cum_precentages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_size = 50\n",
    "gisaid_count_range = [0, 5, 10, 25, 50, 100, 250, 500]\n",
    "\n",
    "hist_bars, labels, total_set_count, cum_precentages = create_supp_figure5_hist_data(gisaid_count_range, mutation_rarity_df)\n",
    "process_time, mean_query_times = create_supp_figure5_boxplot_data(fixed_results_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hist_bars[0]/total_set_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(hist_bars[0:5])/total_set_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize = 14\n",
    "fig, axes  = plt.subplots(2, 1, figsize=(12, 9),\n",
    "                          sharex=False,\n",
    "                          gridspec_kw={'height_ratios': [3, 2], 'hspace':0.6})\n",
    "\n",
    "ax = axes[0]\n",
    "ax.bar(np.arange(len(hist_bars)), hist_bars,\n",
    "      linewidth = 2,\n",
    "      edgecolor = 'black')\n",
    "\n",
    "ax.set_xticks(np.arange(len(hist_bars)))\n",
    "ax.set_xticklabels(labels, rotation=90)\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax.set_ylabel('Count of CR candidates', fontsize=fontsize)\n",
    "\n",
    "\n",
    "ax_twinx = ax.twinx()\n",
    "ax_twinx.plot(np.arange(len(hist_bars)), cum_precentages, \n",
    "              'o-c',\n",
    "              linewidth=2, markersize=10)\n",
    "ax_twinx.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax_twinx.set_ylabel('Cumulative percentage (%)', fontsize=fontsize)\n",
    "ax_twinx.set_ylim(0,110)\n",
    "\n",
    "ax.set_title(f'Occurrence distribution in GISAID EpiCoV for CR candidates', fontsize=fontsize)\n",
    "\n",
    "ax.text(-0.07, 1.08, 'a', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "ax = axes[1]\n",
    "bp = ax.boxplot(process_time,\n",
    "                positions=np.arange(len(hist_bars)),\n",
    "                showfliers=True,\n",
    "                meanline=False,\n",
    "                widths=0.8)\n",
    "\n",
    "ax.set_xticks(np.arange(len(hist_bars)))\n",
    "ax.set_xticklabels(labels, rotation=90)\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "ax.set_ylabel('Process time (s)', fontsize=fontsize)\n",
    "ax.set_xlabel(f'Occurrence in GISAID EpiCoV', fontsize=fontsize)\n",
    "ax.set_title(f'Process time per queries in each bin (sampling size n={sampling_size})', fontsize=fontsize)\n",
    "ax.set_ylim(0,4)\n",
    "\n",
    "ax.plot(np.arange(len(hist_bars)), mean_query_times, \n",
    "              'o-b',\n",
    "              linewidth=2, markersize=10)\n",
    "\n",
    "ax.text(-0.07, 1.08, 'b', transform=ax.transAxes,\n",
    "            fontsize=20, fontweight='bold', va='top')\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/figure_2.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplementary Figure 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "houston_clinical_metadata = pd.read_csv(os.path.join(fixed_results_dir, 'PRJNA764181.filtered.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "houston_clinical_metadata['Collection_Date'] = pd.to_datetime(houston_clinical_metadata['Collection_Date'], format='%Y-%m-%d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_date_start = pd.to_datetime('12/06/2021')\n",
    "week_offset = timedelta(days = 7)\n",
    "sample_counts = []\n",
    "x_ticklabels = []\n",
    "\n",
    "for i in range(0, 8):\n",
    "    start = sampling_date_start + i*week_offset\n",
    "    end = sampling_date_start + (i+1)*week_offset\n",
    "\n",
    "    weekly_clinical_df = houston_clinical_metadata[(houston_clinical_metadata['Collection_Date'] >= start) & (houston_clinical_metadata['Collection_Date'] < end)]\n",
    "    weekly_sample_count = weekly_clinical_df.shape[0]\n",
    "    sample_counts.append(weekly_sample_count)\n",
    "    time_label = start.strftime('%Y-%m-%d')\n",
    "    x_ticklabels.append(time_label)\n",
    "    \n",
    "x_ticklabels.append((sampling_date_start+8*week_offset).strftime('%Y-%m-%d'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize = 14\n",
    "fig, ax  = plt.subplots(1, 1, figsize=(6, 3), sharex=True, constrained_layout=True)\n",
    "\n",
    "x = np.arange(8)\n",
    "\n",
    "ax.bar(x, sample_counts, color='black', alpha=1, width=0.8)\n",
    "\n",
    "ax.set_ylim(0,2000)\n",
    "ax.set_xlim(-0.5,7.5)\n",
    "ax.set_xticks(np.arange(9)-0.5)\n",
    "ax.set_xticklabels(x_ticklabels, rotation=90)\n",
    "\n",
    "ax.set_title('Houston Clinical Sample Count')\n",
    "\n",
    "fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/supp_figure_1.pdf', transparent=False, facecolor='white', dpi=300, bbox_inches = \"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supplementary Figure 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_supp_figure2():\n",
    "    non_tx_metadata = pd.read_csv(os.path.join(fixed_results_dir, 'meta_non_tx.tsv'), sep='\\t', low_memory=False)\n",
    "    non_tx_metadata = non_tx_metadata.fillna('None')\n",
    "    non_tx_metadata = non_tx_metadata[(non_tx_metadata['Region'] != 'Texas') & ~((non_tx_metadata['Country'] == 'USA') & (non_tx_metadata['Region'] == 'None'))]\n",
    "    non_tx_metadata = non_tx_metadata[(non_tx_metadata['Country'] == 'USA')]\n",
    "\n",
    "    region_df = pd.DataFrame(pd.pivot_table(non_tx_metadata, values='Repository', index='Region', aggfunc='count'))\n",
    "    region_df = region_df.rename({'Repository': 'Count'}, axis=1)\n",
    "    other_count = region_df['Count'].sum() - region_df[region_df['Count'] >= 70]['Count'].sum()\n",
    "    selected_region_df = region_df[region_df['Count'] >= 70]\n",
    "    \n",
    "    fontsize = 14\n",
    "    fig, ax  = plt.subplots(1, 1, figsize=(6, 4), sharex=True, constrained_layout=True)\n",
    "\n",
    "    selected_region_df = region_df[region_df['Count'] >= 70]\n",
    "    x = np.arange(selected_region_df.shape[0] + 1)\n",
    "\n",
    "    sample_counts = selected_region_df['Count'].to_list() + [other_count]\n",
    "    x_ticklabels = selected_region_df.index.to_list() + [f'Other {36-12} States']\n",
    "    ax.bar(x, sample_counts, color='black', alpha=1, width=0.8)\n",
    "\n",
    "    ax.set_xlim(-0.5,len(sample_counts)-0.5)\n",
    "    ax.set_xticks(np.arange(len(sample_counts)))\n",
    "    ax.set_xticklabels(x_ticklabels, rotation=90, fontsize=12)\n",
    "\n",
    "    fig.savefig('/home/Users/yl181/wastewater/quarc_figures/pdf/supp_figure_2.pdf', dpi=300, facecolor='white', transparent=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_supp_figure2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.12"
  },
  "toc-autonumbering": false,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}