sdtw_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SquiggleFilter Accuracy Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# manage imports \n",
    "from sklearn import metrics\n",
    "from itertools import repeat\n",
    "from numba import njit\n",
    "from glob import glob\n",
    "from scipy import stats\n",
    "import random, h5py, re, os\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.lines import Line2D  \n",
    "import multiprocessing as mp\n",
    "import seaborn as sns\n",
    "import pandas as pd         \n",
    "import ont_fast5_api\n",
    "\n",
    "# set matplotlib params\n",
    "%matplotlib inline\n",
    "mpl.rcParams['figure.dpi'] = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from pyguppy_client_lib.pyclient import PyGuppyClient\n",
    "# from pyguppy_client_lib.helper_functions import package_read, basecall_with_pyguppy\n",
    "# import mappy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user-defined datasets\n",
    "virus = \"lambda\"\n",
    "#virus = \"covid\"\n",
    "other = \"human\"\n",
    "\n",
    "# user-defined filepaths\n",
    "data_dir = \"data\"\n",
    "img_dir = \"img\"\n",
    "os.makedirs(img_dir, exist_ok=True)\n",
    "kmer_model_fn, k = f\"{data_dir}/dna_kmer_model.txt\", 6 # 6-mer model\n",
    "ref_fn = f\"{data_dir}/{virus}/reference.fasta\"\n",
    "virus_fast5_dir = f\"{data_dir}/{virus}/fast5\"\n",
    "other_fast5_dir = f\"{data_dir}/{other}/fast5\"\n",
    "results_dir = f\"scripts/results/artifact-eval\"\n",
    "\n",
    "# user-defined analysis params\n",
    "virus_max_reads = 1000\n",
    "other_max_reads = 1000\n",
    "prefix_lengths = np.array(range(1000,9001,1000))\n",
    "nprefixes = len(prefix_lengths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Expected Virus Reference Signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_fasta(fasta_fn):\n",
    "    ''' Get base sequence from FASTA filename. '''\n",
    "    with open(fasta_fn, 'r') as fasta:\n",
    "        return ''.join(fasta.read().split('\\n')[1:])\n",
    "\n",
    "    \n",
    "def rev_comp(bases):\n",
    "    ''' Get reverse complement of sequence. '''\n",
    "    return bases.replace('A','t').replace('T','a')\\\n",
    "        .replace('G','c').replace('C','g').upper()[::-1]\n",
    "\n",
    "\n",
    "def load_model(kmer_model_fn):\n",
    "    ''' Load k-mer model file into Python dict. '''\n",
    "    kmer_model = {}\n",
    "    with open(kmer_model_fn, 'r') as model_file:\n",
    "        for line in model_file:\n",
    "            kmer, current = line.split()\n",
    "            kmer_model[kmer] = float(current)\n",
    "    return kmer_model\n",
    "\n",
    "\n",
    "def discrete_normalize(seq, bits=8, minval=-4, maxval=4):\n",
    "    ''' Approximate normalization which converts signal to integer of desired precision. '''\n",
    "    mean = int(np.mean(seq))\n",
    "    mean_avg_dev = int(np.mean(np.abs(seq - mean)))\n",
    "    norm_seq = (seq - mean) / mean_avg_dev\n",
    "    norm_seq[norm_seq < minval] = minval # threshold\n",
    "    norm_seq[norm_seq > maxval] = maxval \n",
    "    norm_seq = ( (norm_seq - minval) * (2**(bits)/(maxval-minval)) ).astype(int)\n",
    "    return norm_seq\n",
    "\n",
    "\n",
    "def ref_signal(fasta, kmer_model):\n",
    "    ''' Convert reference FASTA to expected reference signal (approximate z-scores). '''\n",
    "    signal = np.zeros(len(fasta))\n",
    "    for kmer_start in range(len(fasta)-k):\n",
    "        signal[kmer_start] = kmer_model[fasta[kmer_start:kmer_start+k]]\n",
    "    return discrete_normalize(signal*100) # increase dist between floats before rounding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create COVID reference using (z-score normalized) \n",
    "# expected k-mer currents for forward/backward reference FASTA\n",
    "ref_fasta = get_fasta(ref_fn)\n",
    "kmer_model = load_model(kmer_model_fn)\n",
    "fwd_ref_sig = ref_signal(ref_fasta, kmer_model)\n",
    "rev_ref_sig = ref_signal(rev_comp(ref_fasta), kmer_model)\n",
    "ref = np.concatenate((fwd_ref_sig, rev_ref_sig))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_stall_end(signal, stall_threshold=3, \n",
    "                  stall_events=2, stall_event_len=3):\n",
    "    ''' Determine the end of the DNA stall region. '''\n",
    "    \n",
    "    # take average of a few samples to reduce variation\n",
    "    events = []\n",
    "    for event in range(0, len(signal), stall_event_len):\n",
    "        events.append(np.mean(signal[event:event+stall_event_len]))\n",
    "    \n",
    "    # find where we exceed threshold for a few consecutive events\n",
    "    above_threshold_count = 0\n",
    "    event_pos = 0\n",
    "    for event in events:\n",
    "        event_pos += 1\n",
    "        if event > stall_threshold:\n",
    "            above_threshold_count += 1\n",
    "        else:\n",
    "            above_threshold_count = 0\n",
    "        if above_threshold_count == stall_events:\n",
    "            break\n",
    "            \n",
    "    # find where we go back below threshold\n",
    "    below_threshold_count = 0\n",
    "    for event in events[event_pos:]:\n",
    "        event_pos += 1\n",
    "        if event < stall_threshold:\n",
    "            below_threshold_count += 1\n",
    "        else:\n",
    "            below_threshold_count = 0\n",
    "        if below_threshold_count == stall_events:\n",
    "            break\n",
    "            \n",
    "    return event_pos * stall_event_len\n",
    "\n",
    "\n",
    "def trim(signal):\n",
    "    ''' Trims signal by detecting stall (and eventually adapter). '''\n",
    "    stall_end = get_stall_end(stats.zscore(signal))\n",
    "    return signal[stall_end+1000 : stall_end+1000+max(prefix_lengths)], stall_end\n",
    "\n",
    "\n",
    "def filter_outliers(signal, minval=-4, maxval=4):\n",
    "    ''' Force FAST5 signal values into expected range. '''\n",
    "    \n",
    "    # return empty signals as-is\n",
    "    if not len(signal): return signal\n",
    "    \n",
    "    # upper threshold\n",
    "    for idx, x in enumerate(signal):\n",
    "        if x > maxval:\n",
    "            # other values above max -> threshold to max\n",
    "            if (idx+1 < len(signal) and signal[idx+1] > 3) or \\\n",
    "            (idx > 0 and signal[idx-1] > maxval):\n",
    "                signal[idx] = maxval\n",
    "            # otherwise, single outlier -> interpolate\n",
    "            elif idx == 0:\n",
    "                signal[idx] = signal[1]\n",
    "            elif idx+1 == len(signal):\n",
    "                signal[idx] = signal[idx-1]\n",
    "            else:\n",
    "                signal[idx] = (signal[idx-1] + signal[idx+1]) / 2\n",
    "                \n",
    "    # lower threshold\n",
    "    for idx, x in enumerate(signal):\n",
    "        if x < minval:\n",
    "            # other values below min -> threshold to min\n",
    "            if (idx+1 < len(signal) and signal[idx+1] < -3) or \\\n",
    "            (idx > 0 and signal[idx-1] < minval):\n",
    "                signal[idx] = minval\n",
    "            # otherwise, single outlier -> interpolate\n",
    "            elif idx == 0:\n",
    "                signal[idx] = signal[1]\n",
    "            elif idx+1 == len(signal):\n",
    "                signal[idx] = signal[idx-1]\n",
    "            else:\n",
    "                signal[idx] = (signal[idx-1] + signal[idx+1]) / 2\n",
    "                \n",
    "    return signal\n",
    "\n",
    "\n",
    "class Read():\n",
    "    ''' Store FAST5 read data. '''\n",
    "    def __init__(self, signal, read_id, offset=0, scaling=1.0):                                                                                                                                                                               \n",
    "        self.signal = signal                                                                                                                                                                                                                  \n",
    "        self.read_id = read_id                                                                                                                                                                                                                \n",
    "        self.total_samples = len(signal)                                                                                                                                                                                                      \n",
    "        self.daq_offset = offset                                                                                                                                                                                                              \n",
    "        self.daq_scaling = scaling                                                                                                                                                                                                            \n",
    "        self.read_tag = random.randint(0, int(2**32 - 1))  \n",
    "\n",
    "        \n",
    "def ba_preprocess_read(uuid, length):\n",
    "    ''' Extract read data from FAST5 file for basecalling. '''\n",
    "    readname = f\"read_{uuid}\"\n",
    "    fast5_file = h5py.File(full_index[uuid], 'r')\n",
    "    signal = np.array(fast5_file[readname]['Raw']['Signal'][:], dtype=np.int16)\n",
    "    signal, trimmed = trim(signal)\n",
    "    if len(signal) < max(prefix_lengths): return None\n",
    "    signal_dig = fast5_file[readname]['channel_id'].attrs['digitisation']\n",
    "    signal_offset = fast5_file[readname]['channel_id'].attrs['offset']\n",
    "    signal_range = fast5_file[readname]['channel_id'].attrs['range']\n",
    "    signal_scaling = signal_range / signal_dig\n",
    "    return Read(signal, readname, offset=signal_offset, scaling=signal_scaling)\n",
    "\n",
    "    \n",
    "def preprocess_read(uuid):\n",
    "    ''' Return preprocessed read from specified FAST5 file. '''\n",
    "    readname = f\"read_{uuid}\"\n",
    "    fast5_file = h5py.File(full_index[uuid], 'r')\n",
    "    signal = np.array(fast5_file[readname]['Raw']['Signal'][:], dtype=np.int16)\n",
    "    length = signal.shape[0]\n",
    "    signal, trimmed = trim(signal)\n",
    "    if len(signal) < max(prefix_lengths): return None\n",
    "    new_signal = np.array(signal, dtype=float)\n",
    "    for start in range(0, len(signal), 500):\n",
    "        new_signal[start:start+500] = \\\n",
    "            discrete_normalize(signal[:start+500])[start:start+500]\n",
    "    return new_signal, trimmed, length\n",
    "\n",
    "\n",
    "def get_index(index_filename):\n",
    "    ''' Read index data structure from file. '''\n",
    "    index_file = open(index_filename, 'r')\n",
    "    index = {}\n",
    "    for line in index_file:\n",
    "        uuid, fname = re.split(r'\\t+', line)\n",
    "        index[uuid] = fname.rstrip()\n",
    "    index_file.close()\n",
    "    return index\n",
    "\n",
    "\n",
    "def create_index(fast5_dir, force=False):\n",
    "    ''' Create file which stores read FAST5 to UUID mappings. '''\n",
    "\n",
    "    # return existing index if possible\n",
    "    index_fn = f'{fast5_dir}/index.db'\n",
    "    if not force and os.path.exists(index_fn):\n",
    "        return get_index(index_fn)\n",
    "\n",
    "    # remove existing index\n",
    "    if os.path.exists(index_fn):\n",
    "        os.remove(index_fn)\n",
    "\n",
    "    # create new index    \n",
    "    index_file = open(index_fn, 'w')\n",
    "\n",
    "    # iterate through all FAST5 files in directory\n",
    "    for subdir, dirs, files in os.walk(fast5_dir):\n",
    "        for filename in files:\n",
    "            ext = os.path.splitext(filename)[-1].lower()\n",
    "            if ext == \".fast5\":\n",
    "\n",
    "                # print read uuid and filename to index\n",
    "                fast5_file = h5py.File(os.path.join(subdir, filename), 'r')\n",
    "                if 'Raw' in fast5_file: # single-FAST5\n",
    "                    for readname in fast5_file['Raw']['Reads']:\n",
    "                        uuid = fast5_file['Raw']['Reads'][readname].attrs['read_id']\n",
    "                        print('{}\\t{}'.format(uuid.decode('utf-8'), \\\n",
    "                                os.path.join(subdir, filename)), file=index_file)\n",
    "                else: # multi-FAST5\n",
    "                    for readname in fast5_file:\n",
    "                        uuid = readname[5:] # remove 'read_' naming prefix\n",
    "                        print('{}\\t{}'.format(uuid, \\\n",
    "                                os.path.join(subdir, filename)), file=index_file)\n",
    "\n",
    "    # cleanup and return results\n",
    "    index_file.close()\n",
    "    return get_index(index_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create read UUID -> FAST5 filename mapping\n",
    "virus_index = create_index(virus_fast5_dir)\n",
    "other_index = create_index(other_fast5_dir)\n",
    "full_index = {**virus_index, **other_index}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select random subset of reads\n",
    "if virus == \"covid\":\n",
    "    # COVID reads were reverse-transcribed RNA->DNA, so most are short.\n",
    "    # To ensure we have enough long reads for evaluating accuracy at longer sDTW read lengths, use a lot.\n",
    "    random.seed(42)\n",
    "    virus_readnames = random.choices(list(virus_index.keys()), k=virus_max_reads*1000)\n",
    "else: \n",
    "    random.seed(7)\n",
    "    virus_readnames = random.choices(list(virus_index.keys()), k=virus_max_reads*2)\n",
    "other_readnames = random.choices(list(other_index.keys()), k=other_max_reads*2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Baseline: Basecalling and Alignment Read Until\n",
    "Running baseline adds lots of complicated dependencies, and isn't that useful to verify. Commenting for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # initialize minimap2 aligner\n",
    "# aligner = mappy.Aligner(\n",
    "#     fn_idx_in = ref_fn,\n",
    "#     preset = \"map-ont\",\n",
    "#     best_n = 1,\n",
    "#     k = 15\n",
    "# )\n",
    "\n",
    "# # initialize guppy-lite basecall server\n",
    "# basecaller = PyGuppyClient(\n",
    "#     address = \"127.0.0.1:1234\", \n",
    "#     config = \"dna_r9.4.1_450bps_fast.cfg\",\n",
    "#     server_file_load_timeout=10\n",
    "# )\n",
    "# basecaller.connect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def basecall(packets):\n",
    "#     calls = []                                                                                                                                                                                                                            \n",
    "#     sent, rcvd = 0, 0                                                                                                                                                                                                                         \n",
    "#     while sent < len(packets):                                                                                                                                                                                                            \n",
    "#         success = basecaller.pass_read(packets[sent])                                                                                                                                                                                     \n",
    "#         if not success:                                                                                                                                                                                                                       \n",
    "#             print('ERROR: Failed to basecall read.')                                                                                                                                                                                          \n",
    "#             break                                                                                                                                                                                                                             \n",
    "#         else:                                                                                                                                                                                                                                 \n",
    "#             sent += 1                                                                                                                                                                                                                         \n",
    "#     while rcvd < len(packets):                                                                                                                                                                                                            \n",
    "#         result = basecaller.get_completed_reads()                                                                                                                                                                                             \n",
    "#         rcvd += len(result)                                                                                                                                                                                                                   \n",
    "#         calls.extend(result)\n",
    "#     return calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ba_virus_scores = np.zeros((nprefixes, virus_max_reads))\n",
    "# ba_other_scores = np.zeros((nprefixes, other_max_reads))\n",
    "# with mp.Pool() as pool:\n",
    "#     for prefix_idx, length in enumerate(prefix_lengths):\n",
    "        \n",
    "#         # trim reads\n",
    "#         ba_virus_reads = list(filter(None, pool.starmap(\n",
    "#                     ba_preprocess_read, zip(virus_readnames, repeat(length)))))[:virus_max_reads]\n",
    "#         ba_other_reads = list(filter(None, pool.starmap(\n",
    "#                     ba_preprocess_read, zip(other_readnames, repeat(length)))))[:other_max_reads]\n",
    "                           \n",
    "#         # package read data\n",
    "#         virus_pkts = [package_read(\n",
    "#             read_tag = read.read_tag, \n",
    "#             read_id = read.read_id, \n",
    "#             raw_data = read.signal, \n",
    "#             daq_offset = float(read.daq_offset), \n",
    "#             daq_scaling = float(read.daq_scaling)\n",
    "#         ) for read in ba_virus_reads]\n",
    "#         other_pkts = [package_read(\n",
    "#             read_tag = read.read_tag, \n",
    "#             read_id = read.read_id, \n",
    "#             raw_data = read.signal, \n",
    "#             daq_offset = float(read.daq_offset), \n",
    "#             daq_scaling = float(read.daq_scaling)\n",
    "#         ) for read in ba_other_reads]\n",
    "\n",
    "#         # basecall\n",
    "#         virus_calls = basecall(virus_pkts)\n",
    "#         other_calls = basecall(other_pkts)\n",
    "        \n",
    "#         # align\n",
    "#         for call_idx, call in enumerate(virus_calls):\n",
    "#             try:\n",
    "#                 alignment = next(aligner.map(call['datasets']['sequence']))\n",
    "#                 ba_virus_scores[prefix_idx, call_idx] = alignment.mapq\n",
    "#             except(StopIteration):\n",
    "#                 pass # no alignment\n",
    "#         for call_idx, call in enumerate(other_calls):\n",
    "#             try:\n",
    "#                 alignment = next(aligner.map(call['datasets']['sequence']))\n",
    "#                 ba_other_scores[prefix_idx, call_idx] = alignment.mapq\n",
    "#             except(StopIteration):\n",
    "#                 pass # no alignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hardware Algorithm: sDTW Alignment Read Until"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trim all reads\n",
    "with mp.Pool() as pool:\n",
    "    virus_reads, virus_trims, virus_lengths = \\\n",
    "        list(map(list, zip(*filter(None, pool.map(preprocess_read, virus_readnames)))))\n",
    "    other_reads, other_trims, other_lengths = \\\n",
    "        list(map(list, zip(*filter(None, pool.map(preprocess_read, other_readnames)))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# warn user if not enough long reads for accuracy analyses\n",
    "if len(virus_reads) < virus_max_reads:\n",
    "    print(f'ERROR: only {len(virus_reads)} virus reads long enough, requested {virus_max_reads}')\n",
    "if len(other_reads) < other_max_reads:\n",
    "    print(f'ERROR: only {len(other_reads)} other reads long enough, requested {other_max_reads}')\n",
    "    \n",
    "# keep only 'max_reads' for further analysis\n",
    "virus_reads, virus_trims, virus_lengths = virus_reads[:virus_max_reads], \\\n",
    "    virus_trims[:virus_max_reads], virus_lengths[:virus_max_reads]\n",
    "other_reads, other_trims, other_lengths = other_reads[:other_max_reads], \\\n",
    "    other_trims[:other_max_reads], other_lengths[:other_max_reads]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit()\n",
    "def sdtw(seq):\n",
    "    ''' Returns minimum alignment score for subsequence DTW. '''\n",
    "    \n",
    "    # initialize column vectors\n",
    "    prev_consec = np.zeros((len(seq)))\n",
    "    curr_consec = np.zeros((len(seq)))\n",
    "    prev_cost = np.zeros((len(seq)))\n",
    "    curr_cost = np.zeros((len(seq)))\n",
    "    min_cost = np.zeros((len(seq)))\n",
    "    prev_cost[0] = abs(seq[0]-ref[0])\n",
    "    min_cost[0] = abs(seq[0]-ref[0])\n",
    "    for i in range(1, len(seq)):\n",
    "        prev_cost[i] = prev_cost[i-1] + abs(seq[i]-ref[0])\n",
    "        min_cost[i] = min_cost[i-1] + abs(seq[i]-ref[0])\n",
    "    \n",
    "    # compute entire cost matrix\n",
    "    for j in range(1, len(ref)):\n",
    "        bonus = 10\n",
    "        for i in range(1, len(seq)):\n",
    "            move = prev_cost[i-1] - prev_consec[i-1]*bonus < curr_cost[i-1]\n",
    "            if move:\n",
    "                curr_consec[i] = 0\n",
    "                curr_cost[i] = prev_cost[i-1] - prev_consec[i-1]*bonus + abs(seq[i]-ref[j])\n",
    "            else:\n",
    "                curr_consec[i] = min(10, prev_consec[i] + 1)\n",
    "                curr_cost[i] = curr_cost[i-1] + abs(seq[i]-ref[j])\n",
    "            min_cost[i] = min(min_cost[i], curr_cost[i])\n",
    "        for i in range(len(seq)):\n",
    "            prev_consec[i] = curr_consec[i]\n",
    "            curr_consec[i] = 0\n",
    "            prev_cost[i] = curr_cost[i]\n",
    "            curr_cost[i] = 0\n",
    "    \n",
    "    # return cost of optimal alignment\n",
    "    cost_mins = np.zeros((len(prefix_lengths),))\n",
    "    for i, l in enumerate(prefix_lengths):\n",
    "        if l <= len(seq):\n",
    "            cost_mins[i] = min_cost[l]\n",
    "    return cost_mins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with mp.Pool() as pool:\n",
    "    print(f'Aligning {virus} reads...', flush=True)\n",
    "    virus_scores_list = pool.map(sdtw, virus_reads)\n",
    "    print(f'Aligning {other} reads...', flush=True)\n",
    "    other_scores_list = pool.map(sdtw, other_reads)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# move data to numpy array for easy sorting/calculations\n",
    "virus_scores = np.zeros((nprefixes, len(virus_scores_list)))\n",
    "for idx, scores in enumerate(virus_scores_list):\n",
    "    for i in range(nprefixes):\n",
    "        virus_scores[i,idx] = scores[i]\n",
    "other_scores = np.zeros((nprefixes, len(other_scores_list)))\n",
    "for idx, scores in enumerate(other_scores_list):\n",
    "    for i in range(nprefixes):\n",
    "        other_scores[i,idx] = scores[i]\n",
    "        \n",
    "# save results\n",
    "os.makedirs(results_dir, exist_ok=True)\n",
    "np.save(f\"{results_dir}/prefix_lengths\", prefix_lengths)\n",
    "np.save(f\"{results_dir}/{virus}_trims\", virus_trims)\n",
    "np.save(f\"{results_dir}/{virus}_lengths\", virus_lengths)\n",
    "np.save(f\"{results_dir}/{virus}_scores\", virus_scores)\n",
    "np.save(f\"{results_dir}/{other}-{virus}_trims\", other_trims)\n",
    "np.save(f\"{results_dir}/{other}-{virus}_lengths\", other_lengths)\n",
    "np.save(f\"{results_dir}/{other}-{virus}_scores\", other_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # save basecall-align results as well\n",
    "# np.save(f\"{results_dir}/ba_{virus}_scores\", ba_virus_scores)\n",
    "# np.save(f\"{results_dir}/ba_{other}-{virus}_scores\", ba_other_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load cached results (allows skipping data generation above)\n",
    "prefix_lengths = np.load(f\"{results_dir}/prefix_lengths.npy\")\n",
    "virus_trims = np.load(f\"{results_dir}/{virus}_trims.npy\")\n",
    "virus_lengths = np.load(f\"{results_dir}/{virus}_lengths.npy\")\n",
    "virus_scores = np.load(f\"{results_dir}/{virus}_scores.npy\")\n",
    "other_trims = np.load(f\"{results_dir}/{other}-{virus}_trims.npy\")\n",
    "other_lengths = np.load(f\"{results_dir}/{other}-{virus}_lengths.npy\")\n",
    "other_scores = np.load(f\"{results_dir}/{other}-{virus}_scores.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load cached basecall-align results as well\n",
    "ba_virus_scores = np.load(f\"{results_dir}/ba_{virus}_scores.npy\")\n",
    "ba_other_scores = np.load(f\"{results_dir}/ba_{other}-{virus}_scores.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_stats(virus_scores, other_scores, thresh):\n",
    "    ''' Return F-scores (assumes sorted input). '''\n",
    "    fscores = np.zeros(nprefixes)\n",
    "    precs = np.zeros(nprefixes)\n",
    "    recalls = np.zeros(nprefixes)\n",
    "    for i in range(nprefixes):\n",
    "        # short reads don't receive a score, so ignore in accuracy metrics\n",
    "        long_virus = np.count_nonzero(virus_scores[i])\n",
    "        short_virus = virus_scores.shape[1]-long_virus\n",
    "        tp = np.searchsorted(virus_scores[i], thresh) - short_virus\n",
    "        fn = long_virus - tp\n",
    "        long_other = np.count_nonzero(other_scores[i])\n",
    "        short_other = other_scores.shape[1]-long_other\n",
    "        fp = np.searchsorted(other_scores[i], thresh) - short_other\n",
    "        precs[i] = 0 if not tp+fp else tp / (tp+fp)\n",
    "        recalls[i] = 0 if not tp+fn else tp / (tp+fn)  \n",
    "        fscores[i] = 0 if not tp+fp+fn else tp / (tp + 0.5*(fp + fn))\n",
    "    return fscores, precs, recalls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort arrays (for fast f-score calculation)\n",
    "virus_scores = np.sort(virus_scores)\n",
    "other_scores = np.sort(other_scores)\n",
    "min_score = min(np.min(virus_scores), np.min(other_scores))\n",
    "max_score = max(np.max(virus_scores), np.max(other_scores))\n",
    "\n",
    "# calculate all f-scores, and save the best thresholds\n",
    "best_threshs = np.zeros(nprefixes)\n",
    "best_fscores = np.zeros(nprefixes)\n",
    "best_precs = np.zeros(nprefixes)\n",
    "best_recalls = np.zeros(nprefixes)\n",
    "for thresh in np.linspace(min_score, max_score, num=100):\n",
    "    fscores, precs, recalls = get_stats(virus_scores, other_scores, thresh)\n",
    "    for i in range(nprefixes):\n",
    "        if fscores[i] > best_fscores[i]:\n",
    "            best_fscores[i] = fscores[i]\n",
    "            best_precs[i] = precs[i]\n",
    "            best_recalls[i] = recalls[i]\n",
    "            best_threshs[i] = thresh + 0.01\n",
    "            \n",
    "# save accuracy results\n",
    "np.save(f\"{results_dir}/{other}-{virus}_fscores\", best_fscores)\n",
    "np.save(f\"{results_dir}/{other}-{virus}_precisions\", best_precs)\n",
    "np.save(f\"{results_dir}/{other}-{virus}_recalls\", best_recalls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Accuracy Analysis (Figure 17a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtw_lengths = [1000, 2000, 3000, 5000]\n",
    "ba_lengths = [1000]\n",
    "\n",
    "# initialize plots\n",
    "mpl.rcParams.update({'font.size': 14})\n",
    "dtw_indices = [np.where(prefix_lengths == t)[0] for t in dtw_lengths]\n",
    "ba_indices = [np.where(prefix_lengths == t)[0] for t in ba_lengths]\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "# all DTW plots\n",
    "linestyles=['-', '--', '-.', ':', '-']\n",
    "dashes = [[10,0], [4,2], [6,1.5,1,1.5], [8,4], [8,2,2,2,2,2]]\n",
    "colors = ['C0', 'C1', 'C2', 'C4', 'C5']\n",
    "j = 0\n",
    "for i, l in zip(dtw_indices, dtw_lengths):\n",
    "    i = int(i)\n",
    "    minval = min(np.min(virus_scores[i]), np.min(other_scores[i]))-1\n",
    "    maxval = max(np.max(virus_scores[i]), np.max(other_scores[i]))+1\n",
    "    thresholds = np.linspace(minval, maxval, num=100)\n",
    "\n",
    "    other_discard_rate, virus_discard_rate = [], []\n",
    "    for t in thresholds:\n",
    "        virus_discard_rate.append(np.sum(virus_scores[i] > t) / len(virus_scores[i]))\n",
    "        other_discard_rate.append(np.sum(other_scores[i] > t) / len(other_scores[i]))\n",
    "    ax.plot(virus_discard_rate, other_discard_rate, linestyle=linestyles[j], linewidth=3, dashes=dashes[j], color=colors[j])\n",
    "    j+=1\n",
    "    \n",
    "# # all Guppy + Minimap2 plots\n",
    "# for i, l in zip(ba_indices, ba_lengths):\n",
    "#     i = int(i)\n",
    "#     minval = min(np.min(ba_virus_scores[i]), np.min(ba_other_scores[i]))-1\n",
    "#     maxval = max(np.max(ba_virus_scores[i]), np.max(ba_other_scores[i]))+1\n",
    "#     thresholds = np.linspace(minval, maxval, num=100)\n",
    "\n",
    "#     other_discard_rate, virus_discard_rate = [], []\n",
    "#     for t in thresholds:\n",
    "#         virus_discard_rate.append(sum(ba_virus_scores[i] < t) / len(ba_virus_scores[i]))\n",
    "#         other_discard_rate.append(sum(ba_other_scores[i] < t) / len(ba_other_scores[i]))\n",
    "#     ax.plot(virus_discard_rate, other_discard_rate, linestyle=':', linewidth=3, color='C3')\n",
    "    \n",
    "ax.set_xlabel('Lambda Phage Discard Rate', fontsize=18)\n",
    "ax.set_ylabel('Human Discard Rate', fontsize=18)\n",
    "ax.set_xlim((-0.1, 1.1))\n",
    "ax.set_ylim((-0.1, 1.1))\n",
    "ax.set_title('Accuracy')\n",
    "plt.savefig(f\"{img_dir}/fig17_discard-rates.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Until Runtime\n",
    "Analyze runtime as a function of accuracy and multi-stage thresholding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Reads():\n",
    "    def __init__(self, results_dir, virus_acc, other_acc, virus_eval=None):\n",
    "        if not virus_eval:\n",
    "            virus_eval = virus_acc\n",
    "        \n",
    "        self.prop_virus = 0.01        # proportion virus\n",
    "        self.prop_other = 1 - self.prop_virus\n",
    "        self.prefix_lengths = np.load(f'{results_dir}/prefix_lengths.npy')\n",
    "\n",
    "        self.ba_virus_scores = np.load(f'{results_dir}/ba_{virus_acc}_scores.npy')\n",
    "        self.virus_scores = np.load(f'{results_dir}/{virus_acc}_scores.npy')\n",
    "        self.virus_lengths = np.load(f'{results_dir}/{virus_eval}_lengths.npy')\n",
    "        if virus_eval == \"covid\":\n",
    "            # we've subselected reads long enough to do our accuracy analysis,\n",
    "            # but actual average read length of reverse transcribed DNA is much lower\n",
    "            self.avg_virus_length = 11814 # calculated from dataset\n",
    "            self.target_genome_size = 30_000.0 # bases\n",
    "        else:\n",
    "            self.avg_virus_length = np.mean(self.virus_lengths)\n",
    "            self.target_genome_size = 100_000.0 # bases\n",
    "        self.virus_trims = np.load(f'{results_dir}/{virus_eval}_trims.npy')\n",
    "        self.avg_virus_trim = np.mean(self.virus_trims)\n",
    "\n",
    "        self.ba_other_scores = np.load(f'{results_dir}/ba_{other}-{virus_acc}_scores.npy')\n",
    "        self.other_scores = np.load(f\"{results_dir}/{other}-{virus_acc}_scores.npy\")\n",
    "        self.other_lengths = np.load(f\"{results_dir}/{other}-{virus_eval}_lengths.npy\")\n",
    "        if virus_eval == \"covid\":\n",
    "            # we use a normal human DNA dataset, but take average read length from\n",
    "            # a reverse-transcribed DNA dataset here to match lab prep for COVID\n",
    "            self.avg_other_length = 12551 # calculated from rtDNA dataset\n",
    "        else:\n",
    "            self.avg_other_length = np.mean(self.other_lengths)\n",
    "        self.other_trims = np.load(f\"{results_dir}/{other}-{virus_eval}_trims.npy\")\n",
    "        self.avg_other_trim = np.mean(self.other_trims)\n",
    "    \n",
    "    \n",
    "class Flowcell():\n",
    "    def __init__(self, channels = 512):\n",
    "        self.chemistry = 'r9.4.1'\n",
    "        self.sampling_rate = 4000     # samples/sec, known\n",
    "        self.minknow_latency = 0.070  # sec, measured\n",
    "        self.channels = channels\n",
    "        \n",
    "        \n",
    "class Classifier():\n",
    "    def __init__(self, method='sf', watts=1):\n",
    "        self.method = method\n",
    "        if method == 'sf': # SquiggleFilter measured\n",
    "            self.throughput = 23_364 * 2000 * watts\n",
    "            self.latency = 0.0000408\n",
    "        elif method == 'ba': # BasecallAlign measured\n",
    "            self.throughput = 550 * 2000\n",
    "            self.latency = 0.149\n",
    "        else:\n",
    "            raise Exception(\"Unknown Read Until classifier type.\")\n",
    "    \n",
    "\n",
    "class Run():\n",
    "    def __init__(self, reads, clf='sf', flowcell=Flowcell(), watts=1):\n",
    "        self.flowcell = flowcell\n",
    "        self.clf = Classifier(clf, watts)\n",
    "        self.reads = reads\n",
    "        self.target_coverage = 30.0\n",
    "        self.coverage_bias = 10.0     # estimated\n",
    "        self.fwd_tr = 400.0           # bases / sec\n",
    "        self.rev_tr = 100_000.0       # bases / sec\n",
    "        self.capture_time = 1.0       # sec\n",
    "        self.sr = self.flowcell.sampling_rate\n",
    "    \n",
    "    \n",
    "    def get_simple_runtime(self):\n",
    "        max_throughput = self.flowcell.channels * self.sr\n",
    "        virus_time = self.reads.prop_virus * \\\n",
    "            (self.capture_time + self.reads.avg_virus_length/self.sr)\n",
    "        other_time = self.reads.prop_other * \\\n",
    "            (self.capture_time + self.reads.avg_other_length/self.sr)\n",
    "        useful_time = self.reads.prop_virus * \\\n",
    "            (self.reads.avg_virus_length-self.reads.avg_virus_trim) / self.sr\n",
    "        useful_throughput = max_throughput * useful_time / (virus_time + other_time)\n",
    "        duration = self.reads.target_genome_size * (self.sr/self.fwd_tr) * \\\n",
    "            self.target_coverage * self.coverage_bias / useful_throughput\n",
    "        return duration\n",
    "    \n",
    "    \n",
    "    def get_read_until_runtime(self, prefix_indices, thresholds):\n",
    "        \n",
    "        # get runtime without read until\n",
    "        max_throughput = self.flowcell.channels * self.sr\n",
    "        simple_virus_time = self.reads.prop_virus * \\\n",
    "            (self.capture_time + self.reads.avg_virus_length/self.sr)\n",
    "        simple_other_time = self.reads.prop_other * \\\n",
    "            (self.capture_time + self.reads.avg_other_length/self.sr)\n",
    "        simple_useful_time = self.reads.prop_virus * \\\n",
    "            (self.reads.avg_virus_length-self.reads.avg_virus_trim) / self.sr\n",
    "        simple_prop_useful_time = simple_useful_time / \\\n",
    "            (simple_virus_time + simple_other_time)\n",
    "        \n",
    "        # determine percentage of pores that can perform Read Until\n",
    "        # - estimate required basecall throughput from simple sequencing\n",
    "        bc_time = self.reads.prop_virus * self.reads.avg_virus_length/self.sr + \\\n",
    "            self.reads.prop_other * self.reads.avg_other_length/self.sr\n",
    "        bc_throughput = max_throughput * (bc_time/ (simple_virus_time + simple_other_time))\n",
    "        prop_ru = min(1.0, self.clf.throughput/bc_throughput)\n",
    "        prop_simple = 1 - prop_ru\n",
    "        \n",
    "        # calculate sequencing runtime for multiple thresholds\n",
    "        if self.clf.method == 'sf':\n",
    "            rem_virus_scores = self.reads.virus_scores.copy()\n",
    "            rem_other_scores = self.reads.other_scores.copy()\n",
    "        else:\n",
    "            rem_virus_scores = -self.reads.ba_virus_scores.copy()\n",
    "            rem_other_scores = -self.reads.ba_other_scores.copy()   \n",
    "            \n",
    "        eject_virus_time, eject_other_time = 0, 0\n",
    "        for i, thresh in zip(prefix_indices, thresholds):\n",
    "            \n",
    "            # device continues sequencing as we make a read-until decision\n",
    "            length = self.reads.prefix_lengths[i]\n",
    "            samples = length + self.sr * \\\n",
    "                (self.clf.latency + self.flowcell.minknow_latency)\n",
    "            reversal_latency = samples / (self.rev_tr * (self.sr/self.fwd_tr))\n",
    "            latency = self.clf.latency + self.flowcell.minknow_latency + reversal_latency\n",
    "            \n",
    "            # choose which reads to keep, count those ejected\n",
    "            keep_virus = rem_virus_scores[i] < thresh\n",
    "            rem_virus_scores = rem_virus_scores[:,keep_virus]\n",
    "            n_eject_virus = len(keep_virus) - sum(keep_virus)\n",
    "            keep_other = rem_other_scores[i] < thresh\n",
    "            rem_other_scores = rem_other_scores[:,keep_other]\n",
    "            n_eject_other = len(keep_other) - sum(keep_other)\n",
    "            \n",
    "            # update time spent sequencing ejected reads\n",
    "            eject_virus_time += n_eject_virus * self.reads.prop_virus * \\\n",
    "                (self.capture_time + length/self.sr + latency)\n",
    "            eject_other_time += n_eject_other * self.reads.prop_other * \\\n",
    "                (self.capture_time + length/self.sr + latency)\n",
    "            \n",
    "        # update total time spent sequencing each type of read\n",
    "        ru_useful_time = len(rem_virus_scores[0]) * self.reads.prop_virus * \\\n",
    "            (self.reads.avg_virus_length - self.reads.avg_virus_trim) / self.sr\n",
    "        ru_virus_time = eject_virus_time + len(rem_virus_scores[0]) * self.reads.prop_virus * \\\n",
    "            (self.capture_time + self.reads.avg_virus_length / self.sr)\n",
    "        ru_other_time = eject_other_time + len(rem_other_scores[0]) * self.reads.prop_other * \\\n",
    "            (self.capture_time + self.reads.avg_other_length / self.sr)\n",
    "        ru_prop_useful_time = ru_useful_time / (ru_virus_time + ru_other_time)\n",
    "        \n",
    "        # calculate duration based on simple/read until split\n",
    "        prop_useful_time = prop_ru * ru_prop_useful_time + \\\n",
    "            prop_simple * simple_prop_useful_time\n",
    "        useful_throughput = prop_useful_time * max_throughput + 0.0001\n",
    "        duration = self.reads.target_genome_size * (self.sr/self.fwd_tr) * \\\n",
    "            self.target_coverage * self.coverage_bias / useful_throughput\n",
    "        return duration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read Until Runtime (Figure 17b and 17c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the filter accuracy metrics from lambda phage and the known read length distribution from COVID to estimate the optimal read until thresholds for COVID. To do this, cached statistics must have been generated for both 'covid' and 'lambda'. First run all cells with `virus='lambda'` (in Cell 3), then do the same with `virus='covid'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_indices = [0,1,2,4,7]\n",
    "plot_prefixes = [1000,2000,3000,5000,8000]\n",
    "best_threshs = [0]*nprefixes\n",
    "\n",
    "if os.path.isfile(f'{results_dir}/lambda_scores.npy') and os.path.isfile(f'{results_dir}/covid_scores.npy'):\n",
    "    read_data = Reads(results_dir, virus_acc='lambda', other_acc=other, virus_eval='covid')\n",
    "    run = Run(read_data)\n",
    "    for i1, p1 in zip(plot_indices, plot_prefixes):\n",
    "        max_score = max(np.max(virus_scores[i1]), np.max(other_scores[i1]))\n",
    "        min_score = max(1, min(np.min(virus_scores[i1]), np.min(other_scores[i1])))\n",
    "        times = []\n",
    "        thresholds = np.linspace(min_score, max_score, 100)\n",
    "        for t in thresholds:\n",
    "            times.append(run.get_read_until_runtime([i1], [t]))\n",
    "\n",
    "        best = np.argmin(times)\n",
    "        best_threshs[i1] = thresholds[best]\n",
    "else:\n",
    "    print(f\"You must run data analysis on covid and lambda before determining optimal thresholds.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for v in ['lambda', 'covid']:\n",
    "    if os.path.isfile(f'{results_dir}/{v}_scores.npy'):\n",
    "        read_data = Reads(results_dir, v, other)\n",
    "        run = Run(read_data)\n",
    "        xmax = 80000\n",
    "        plt.plot(np.linspace(0,xmax,20), [run.get_simple_runtime()]*20, color='k', linestyle=':')\n",
    "\n",
    "        linestyles=['-', '--', '-.', ':', '-']\n",
    "        dashes = [[10,0], [4,2], [6,1.5,1,1.5], [8,4], [8,2,2,2,2,2]]\n",
    "        colors = ['C0', 'C1', 'C2', 'C4', 'C5']\n",
    "        j = 0\n",
    "        first=True\n",
    "        plot_indices = [0,1,2,4,7]\n",
    "        plot_prefixes = [1000,2000,3000,5000,8000]\n",
    "\n",
    "        for i1, p1 in zip(plot_indices, plot_prefixes):\n",
    "            max_score = max(np.max(virus_scores[i1]), np.max(other_scores[i1]))\n",
    "            min_score = max(1, min(np.min(virus_scores[i1]), np.min(other_scores[i1])))\n",
    "            times = []\n",
    "            thresholds = np.linspace(min_score, max_score, 100)\n",
    "            for t in thresholds:\n",
    "                times.append(run.get_read_until_runtime([i1], [t]))\n",
    "\n",
    "            if v == \"lambda\":\n",
    "                best = np.argmin(times)\n",
    "                if first:\n",
    "                    first=False\n",
    "                    plt.plot(thresholds[best], times[best], marker='*', markersize=9, color='k', linestyle='', zorder=2)\n",
    "                else:\n",
    "                    plt.plot(thresholds[best], times[best], marker='*', markersize=9, color='k', linestyle='', label='_nolegend_', zorder=2)\n",
    "            else:\n",
    "                if first:\n",
    "                    best = best_threshs[i1]\n",
    "                    first=False\n",
    "                    time = np.interp(best, thresholds, times)\n",
    "                    plt.plot(best, time, marker='*', markersize=9, color='k', linestyle='', zorder=2)\n",
    "                else:\n",
    "                    best = best_threshs[i1]\n",
    "                    time = np.interp(best, thresholds, times)\n",
    "                    plt.plot(best, time, marker='*', markersize=9, color='k', linestyle='', label='_nolegend_', zorder=2)\n",
    "\n",
    "            plt.plot(thresholds, times, dashes=dashes[j], color=colors[j], linewidth=3, zorder=1)\n",
    "            plt.ylabel('Read Until runtime (s)', fontsize=18)\n",
    "            plt.xlabel('sDTW Alignment Score Threshold', fontsize=18)\n",
    "            j+= 1\n",
    "\n",
    "        plt.xlim(0,xmax)\n",
    "        plt.ylim(3000, run.get_simple_runtime()*1.3)\n",
    "        plt.yscale('log')\n",
    "        if v == 'lambda':\n",
    "            plt.title('Lambda Phage')\n",
    "        else:\n",
    "            plt.title('SARS-CoV-2')\n",
    "            plt.plot([], [], linestyle=':', linewidth=3, color='C3')\n",
    "            leg = plt.legend(['No Read Until', 'Selected Threshold'] + [f'SquiggleFilter {l} samples' for l in plot_prefixes] + ['Guppy-lite 1000 samples'], \n",
    "                             loc=(1.04,0.1), handlelength=4)\n",
    "            leg.legendHandles[2]._legmarker.set_markersize(9)\n",
    "        plt.savefig(f\"{img_dir}/fig17_{virus}-read-until-time.png\")\n",
    "        plt.show()\n",
    "    else:\n",
    "        print(f\"You must run data analysis on '{v}' before generating corresponding plot.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Score Distribution (Figure 11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_prefixes = [1000, 2000, 3000, 5000, 8000]\n",
    "fig, ax = plt.subplots(1,len(plot_prefixes), figsize=(12,2))\n",
    "j = 0\n",
    "for i, l in enumerate(prefix_lengths):\n",
    "    if l in plot_prefixes:\n",
    "        minval = min(np.min(virus_scores[i]), np.min(other_scores[i]))\n",
    "        maxval = max(np.max(virus_scores[i]), np.max(other_scores[i]))\n",
    "        ax[j].hist(virus_scores[i], bins=np.linspace(minval, maxval, num=100), facecolor='#1a4099', alpha=0.7)\n",
    "        ax[j].hist(other_scores[i], bins=np.linspace(minval, maxval, num=100), facecolor='#ebc100', alpha=0.5)\n",
    "        ax[j].axvline(best_threshs[i], color='k', linestyle='--')\n",
    "        ax[j].set_title('{} Samples'.format(l), fontsize=16)\n",
    "        ax[j].set_ylim((0,130))\n",
    "        ax[j].set_xlim(0, best_threshs[i]*2)\n",
    "        j += 1\n",
    "\n",
    "    ax[0].legend(['threshold', virus, other], loc='upper left', fontsize=15)\n",
    "    ax[2].set_xlabel('sDTW Alignment Cost', fontsize=20)\n",
    "    ax[1].set_yticklabels([])\n",
    "    ax[2].set_yticklabels([])\n",
    "    ax[3].set_yticklabels([])\n",
    "    ax[4].set_yticklabels([])\n",
    "    ax[0].set_ylabel('Read Count', fontsize=20)\n",
    "\n",
    "plt.savefig(f\"{img_dir}/fig11_cost-distribution.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other Figures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Progression of SARS-CoV-2 Testing (Figure 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 10})   \n",
    "df = pd.read_csv(f'{data_dir}/owid-covid-data.csv', header=None) \n",
    "df[1] = df[1].diff()                                                             \n",
    "                                                                                 \n",
    "ax = df.plot(color=\"#6eaaef\", figsize=(4.5,2), zorder=1)                         \n",
    "ax.set_yticks([100_000, 500_000, 1_000_000, 2_000_000])                          \n",
    "ax.set_yticklabels([\"100,000\", \"500,000\", \"1 million\", \"2 million\"])             \n",
    "ax.set_ylim(0, 2_500_000)                                                        \n",
    "                                                                                 \n",
    "ax.set_xticks([0, 196, 366])                                                     \n",
    "ax.set_xticklabels([\"Jan 1 2020\", \"July 15\", \"Jan 1 2021\"])                      \n",
    "ax.legend().remove()                                                             \n",
    "                                                                                \n",
    "plt.vlines(x=13, linestyle=\"--\", ymin=0, ymax=2_500_000, color=\"k\")              \n",
    "plt.text(14, 1_400_000, 'SARS-CoV-2\\nSequenced')                                 \n",
    "                                                                                 \n",
    "plt.arrow(14, 200_000, 181, 0, fc='k', ec='k', zorder=2)                         \n",
    "plt.plot(18, 200_000, marker='<', color='k')                                     \n",
    "plt.plot(189, 200_000, marker='>', color='k')                                    \n",
    "plt.text(23, 260_000, 'Insufficient Testing')                                    \n",
    "                                                                                 \n",
    "plt.vlines(x=196, linestyle='--', ymin=0, ymax=2_500_000, color='k')             \n",
    "plt.text(197, 1_400_000, '1 Million\\nDaily Tests')                               \n",
    "                                                                                 \n",
    "lines = [Line2D([0], [0], color=\"#6eaaef\", lw=3)]                                \n",
    "ax.legend(lines, [\"Daily US COVID-19 Tests\"], loc='upper right')                 \n",
    "                                                                               \n",
    "plt.tight_layout()                                                               \n",
    "plt.savefig(f\"{img_dir}/fig2_covid-tests.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pipeline Runtime Proportions (Figure 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 14})                                    \n",
    "                                                                                 \n",
    "fig, axes = plt.subplots(1,2)                                                    \n",
    "tools = ['Basecalling', 'Alignment', 'Variant Calling']                          \n",
    "covid_prop100 = [87.72, 2.75, 9.53]                                              \n",
    "covid_prop1000 = [95.95, 3.01, 1.04]                                             \n",
    "colors = [\"#1a4099\", \"#ebc100\", \"#9c9c9c\"]                                       \n",
    "labels0 = [f'{tool} ({pct}%)' for tool, pct in zip(tools, covid_prop100)]        \n",
    "labels1 = [f'{tool} ({pct}%)' for tool, pct in zip(tools, covid_prop1000)]          \n",
    "                                                                                 \n",
    "patches0, _ = axes[0].pie(covid_prop100, startangle=90, colors=colors)           \n",
    "patches1, _ = axes[1].pie(covid_prop1000, startangle=90, colors=colors)          \n",
    "axes[0].legend(patches0, labels0, loc=(0,-0.3))                                  \n",
    "axes[1].legend(patches1, labels1, loc=(0,-0.3))                                  \n",
    "plt.tight_layout()                                                               \n",
    "fig.savefig(f'{img_dir}/fig5_pipeline-runtimes.png') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Virus Genome Sizes (Figure 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 17})                                                                                 \n",
    "                                                                                 \n",
    "def add_datapoint(ax, data):                                                     \n",
    "    name = data[0]                                                               \n",
    "    size = data[1]                                                               \n",
    "    na = data[2]                                                                 \n",
    "    strands = data[3]                                                            \n",
    "    height = data[4]                                                             \n",
    "    ax.text(size, height+0.2, name, {'color': \"#ebc100\" if na == \"DNA\" else \"#1a4099\"}, ha=\"left\", va=\"center\")\n",
    "    ax.plot(size, height, color=\"#ebc100\" if na == \"DNA\" else \"#1a4099\", markersize=12, marker=\"P\" if strands==\"ds\" else \".\")\n",
    "                                                                                 \n",
    "virus_genomes = [                                                                \n",
    "    [\"Coxsackie\", 7200, \"RNA\", \"ss\", 0.7],                                       \n",
    "    [\"Polio\", 7500, \"RNA\", \"ss\", 1.4],                                           \n",
    "    [\"HIV\", 9800, \"RNA\", \"ss\", 2.1],                                             \n",
    "    [\"Zika\", 10000, \"RNA\", \"ss\", 2.8],          \n",
    "    [\"Dengue\", 11000, \"RNA\", \"ss\", 3.5],                                         \n",
    "    [\"Yellow Fever\", 11000, \"RNA\", \"ss\", 4.2],                                   \n",
    "    [\"Borna Disease\", 8900, \"RNA\", \"ss\", 4.9],                                   \n",
    "    [\"West Nile\", 11500, \"RNA\", \"ss\", 5.6],                                      \n",
    "    [\"Influenza A\", 13500, \"RNA\", \"ss\", 0.7],                                    \n",
    "    [\"Influenza B\", 14500, \"RNA\", \"ss\", 1.4],                                    \n",
    "    [\"Mumps\", 15400, \"RNA\", \"ss\", 2.1],                                          \n",
    "    [\"Measles\", 15900, \"RNA\", \"ss\", 2.8],                                        \n",
    "    [\"Ebola\", 19000, \"RNA\", \"ss\", 3.5],                                          \n",
    "    [\"Reovirus\", 29200, \"RNA\", \"ds\", 0.7],                                       \n",
    "    [\"SARS\", 29700, \"RNA\", \"ss\", 1.4],                                           \n",
    "    [\"Coronavirus\", 29900, \"RNA\", \"ss\", 2.1],                                    \n",
    "    [\"Herpes Simplex\", 152000, \"DNA\", \"ds\", 0.7],                                \n",
    "    [\"Smallpox\", 186000, \"DNA\", \"ds\", 1.4],                                      \n",
    "]                                                                                \n",
    "                                                                                 \n",
    "fig, ax = plt.subplots(figsize=(10,4))                                           \n",
    "ax.set_ylim(0, 6.5)                                                              \n",
    "ax.set_yticks([])                                                                \n",
    "for v in virus_genomes:                                                          \n",
    "    add_datapoint(ax, v)                                                         \n",
    "ax.set_xscale('log')                                                             \n",
    "ax.set_xlim(5000, 500000)                                                        \n",
    "ax.set_xticks([5_000, 10_000, 50_000, 100_000, 250_000])                         \n",
    "ax.set_xticklabels([\"5,000\",\"10,000\",\"50,000\",\"100,000\", \"250,000\"])             \n",
    "ax.set_xlabel(\"Genome length (bases)\")                                           \n",
    "                                                                                 \n",
    "# legend                                                                         \n",
    "ax.add_patch(mpl.patches.Rectangle((110_000, 4), 375_000, 2.3, fill=None, alpha=1))          \n",
    "ax.plot(120000, 5.8, color=\"#1a4099\", marker=\".\", markersize=15)                 \n",
    "ax.text(130000, 5.8, \"single stranded RNA\", color=\"#1a4099\", va=\"center\")        \n",
    "ax.plot(120000, 5.1, color=\"#1a4099\", marker=\"P\", markersize=15)                 \n",
    "ax.text(130000, 5.1, \"double stranded RNA\", color=\"#1a4099\", va=\"center\")        \n",
    "ax.plot(120000, 4.4, color=\"#ebc100\", marker=\"P\", markersize=15)                 \n",
    "ax.text(130000, 4.4, \"double stranded DNA\", color=\"#ebc100\", va=\"center\")        \n",
    "                                                                                 \n",
    "# lines                                                                          \n",
    "for x in np.arange(0, 6.5, 0.4):                                                 \n",
    "    ax.plot(50_000, x, marker='P', markersize=5, color='k')                      \n",
    "    ax.plot(100_000, x, marker='.', markersize=5, color='k')                     \n",
    "                                                                                 \n",
    "plt.tight_layout()                                                               \n",
    "plt.savefig(f'{img_dir}/fig10_virus-genomes.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Latency Results (Figure 16a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 18})                                     \n",
    "fig, ax = plt.subplots()                                                         \n",
    "                                                                                 \n",
    "tools = ['Guppy', 'Guppy\\nLite', 'Squiggle\\nFilter']                             \n",
    "colors = ['C3', 'C1', 'C2']                                                      \n",
    "latencies = [ 1.304, 0.149, (107_000) / (2.5 * 10**9) ]                          \n",
    "                                                                                 \n",
    "ax.bar(tools, latencies, color=colors)                                           \n",
    "                                                                                 \n",
    "base_lines = [1, 10, 100, 500]                                                   \n",
    "base_latencies = [.0025, .025, .25, 1.25]                                        \n",
    "for latency, base in zip(base_latencies, base_lines):                            \n",
    "    plt.hlines(y=latency, xmin=-0.5, xmax=1.5, color='k', linestyle=':', linewidth=3)\n",
    "    label = \" 1 base\" if base == 1 else f' {base} bases'                         \n",
    "    plt.text(1.5, latency, label, ha='left', va='center')                        \n",
    "                                                                                 \n",
    "ax.set_ylabel('Latency (s)')                                                     \n",
    "ax.set_yscale('log')                                                             \n",
    "plt.tight_layout()                                                               \n",
    "fig.savefig(f'{img_dir}/fig16_latency.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Throughput Results (Figure 16b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 20})                                         \n",
    "fig, ax = plt.subplots(figsize=(15,5))                                                \n",
    "                                                                                      \n",
    "systems = ['Guppy', 'Guppy\\nLite', ' Guppy', ' Guppy\\nLite', 'Squiggle\\nFilter', ' Squiggle\\nFilter']\n",
    "colors = ['C3', 'C1', 'C3', 'C1', 'C2', 'C2']                                         \n",
    "                                                                                      \n",
    "titan_thru_hac = 387                                                                  \n",
    "titan_thru_fast = 1524                                                                \n",
    "compute_ratio_hac = 480941/222959                                                     \n",
    "compute_ratio_fast = 109888/39656                                                     \n",
    "                                                                                      \n",
    "throughputs = [                                                                       \n",
    "        titan_thru_hac/compute_ratio_hac,                                             \n",
    "        titan_thru_fast/compute_ratio_fast,                                           \n",
    "        titan_thru_hac,                                                               \n",
    "        titan_thru_fast,                                                              \n",
    "        2.5*10**9 / 107_000,                                                          \n",
    "        5 * 2.5*10**9 / 107_000 ]                                                     \n",
    "                                                                                      \n",
    "ax.bar(systems, throughputs, color=colors)                                            \n",
    "                                                                                      \n",
    "ax.axhline(y=512*4000/2000, color='k', linestyle='--', linewidth=3)                   \n",
    "ax.axhline(y=5*512*4000/2000, color='k', linestyle=':', linewidth=3)                  \n",
    "ax.axhline(y=24*3000*4000/2000, color='k', linestyle='-.', linewidth=3)               \n",
    "ax.axhline(y=100*512*4000/2000, color='k', linestyle='-', linewidth=3)                \n",
    "ax.legend(['MinION output', 'GridION output', 'PromethION output', 'Future MinION output'], labelspacing=0.2, loc=(1.04,0.5))\n",
    "                                                                                      \n",
    "ax.set_ylabel('Throughput (reads/s)')                                                 \n",
    "ax.set_yscale('log')                                                                  \n",
    "ax.set_xticklabels(systems)                                                           \n",
    "ax.set_xlim(-0.7, 5.7)                                                                \n",
    "                                                                                      \n",
    "                                                                                      \n",
    "ax.text(0.5, 6000, \"Jetson\\nXavier\\n\" + r\"$350\\mathrm{mm}^2$\" + \"\\n30W\", ha='center')\n",
    "ax.text(2.5, 6000, \"Titan\\nXP\\n\" + r\"$471\\mathrm{mm}^2$\" + \"\\n250W\", ha='center')\n",
    "ax.text(4, 6000, \"1-Tile\\nASIC\\n\" + r\"$3\\mathrm{mm}^2$\" + \"\\n3W\", ha='center')   \n",
    "ax.text(5, 6000, \"5-Tile\\nASIC\\n\" + r\"$13\\mathrm{mm}^2$\" + \"\\n14W\", ha='center') \n",
    "                                                                                      \n",
    "plt.tight_layout()                                                                    \n",
    "fig.savefig(f'{img_dir}/fig16_throughput.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Algorithm Modifications (Figure 18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_dirs = [\"final_baseline\", \"final_no_ref_skip\", \"final_discrete_norm\", \"final_abs_value\", \"final_all\", \"final_all_bonus\"]\n",
    "linestyles=['-', '-.', '--', ':', '-', ':']\n",
    "markers=['','','','','*','*']\n",
    "mpl.rcParams.update({'font.size': 16})\n",
    "fig, axs = plt.subplots(1,1, figsize=(8,4))\n",
    "for i, d in enumerate(results_dirs):\n",
    "    fscores = np.load(f\"scripts/results/{d}/fscores.npy\")\n",
    "    axs.plot(np.arange(500,5001,500), fscores, marker=markers[i], linestyle=linestyles[i])\n",
    "axs.set_ylim(0.5,1)\n",
    "axs.set_xlim(0,7000)\n",
    "axs.set_xlabel('Samples')\n",
    "axs.set_ylabel('F-Score')\n",
    "axs.legend([\"Baseline\", \"No Reference Deletions\", \"Integer Normalization\", \"Absolute Difference\", \"All Efficiency Optimizations\", \"All Eff Opts + Match Bonus\"],)\n",
    "plt.savefig(f\"{img_dir}/fig18_alg-mods.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mutation-Dependent Accuracy (Figure 19)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams.update({'font.size': 18})                                    \n",
    "                                                                                 \n",
    "fig, ax = plt.subplots(1, figsize=(8,3))                                         \n",
    "                                                                                 \n",
    "n_mut = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000]        \n",
    "f_sub = [0.901, 0.9, 0.902, 0.9, 0.9, 0.902, 0.899, 0.896, 0.891, 0.878, 0.836, 0.662, 0.209, 0.027]\n",
    "f_ins = [0.901, 0.903, 0.901, 0.901, 0.9, 0.902, 0.9, 0.9, 0.893, 0.883, 0.877, 0.807, 0.604, 0.093]\n",
    "f_del = [0.901, 0.901, 0.903, 0.9, 0.9, 0.898, 0.9, 0.897, 0.884, 0.867, 0.806, 0.438, 0.031, 0.021]\n",
    "                                                                                 \n",
    "ax.plot(n_mut, f_sub, linewidth=4, linestyle='-', color='r', alpha=0.5)          \n",
    "ax.plot(n_mut, f_ins, linewidth=4, linestyle='--', color='g', alpha=0.5)         \n",
    "ax.plot(n_mut, f_del, linewidth=4, linestyle=':', color='b', alpha=0.5)          \n",
    "ax.set_ylim(0,1)                                                                 \n",
    "ax.set_xscale('log')                                                             \n",
    "ax.set_ylabel('F-score')                                                         \n",
    "ax.set_xlabel('Number of Mutations')                                             \n",
    "ax.legend(['Substitutions', 'Insertions', 'Deletions'])                          \n",
    "                                                                                 \n",
    "plt.tight_layout()                                                               \n",
    "plt.savefig(f'{img_dir}/fig19_mutation-dependent-accuracy.png', dpi=300)   \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Projected Future SquiggleFilter Read Until Benefits (Figure 21)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "read_data = Reads(results_dir, virus, other)\n",
    "mpl.rcParams.update({'font.size': 22})\n",
    "\n",
    "simple_runtimes = []\n",
    "sf_runtimes = [[], []]\n",
    "ba_runtimes = []\n",
    "channels = np.logspace(2, 9, num=30)\n",
    "for channel_count in channels:\n",
    "    flowcell = Flowcell(channel_count)\n",
    "    for i, w in enumerate([1,5]):\n",
    "        sf_run = Run(read_data, 'sf', flowcell, w)\n",
    "        sf_runtimes[i].append(sf_run.get_read_until_runtime([1], [12935]))\n",
    "    ba_run = Run(read_data, 'ba', flowcell)\n",
    "    simple_runtimes.append(sf_run.get_simple_runtime())\n",
    "    ba_runtimes.append(ba_run.get_read_until_runtime([0], [-1]))\n",
    "fig = plt.figure(figsize=(12,4))\n",
    "ax = fig.add_subplot(1,1,1)\n",
    "ax.set_yscale('log')\n",
    "ax.set_xscale('log')\n",
    "ax.plot(channels, simple_runtimes, color='k', linestyle=':', linewidth=3)\n",
    "ax.plot(channels, ba_runtimes, linestyle=':', color='C3', linewidth=4)\n",
    "ax.plot(channels, sf_runtimes[0], linestyle='--', color='C2', linewidth=4, alpha=0.7)\n",
    "ax.plot(channels, sf_runtimes[1], linestyle='-.', color='C2', linewidth=4, alpha=0.7)\n",
    "ax.axvline(512, color='k', linestyle='--', linewidth=3)\n",
    "ax.axvline(5*512, color='k', linestyle='--', alpha=0.5, linewidth=3)\n",
    "ax.axvline(3000*24, color='k', linestyle='--', alpha=0.5, linewidth=3)\n",
    "ax.axvline(100*512, color='k', linestyle='--', linewidth=3)\n",
    "ax.text(512, 10, 'MinION', fontsize=22)\n",
    "ax.text(5*512, 3*10**5, 'GridION', fontsize=22, alpha=0.5)\n",
    "ax.text(3000*24, 3*10**5, 'PromethION', fontsize=22, alpha=0.5)\n",
    "ax.text(100*512, 10, 'Future MinION', fontsize=22)\n",
    "ax.set_ylim(5, 10**6)\n",
    "ax.set_xlabel('Sequencer Channels', fontsize=22)\n",
    "ax.set_ylabel('Sequencing Time', fontsize=22)\n",
    "leg = ax.legend(['No Read Until', 'Guppy-lite (30W)', '1-Tile SquiggleFilter (3W)', '5-Tile SquiggleFilter (14W)'], \n",
    "                loc=(0.45,0.4), fontsize=22, labelspacing=0)\n",
    "leg.legendHandles[0]._legmarker.set_markersize(9)\n",
    "plt.savefig(f\"{img_dir}/fig21_future-benefits.png\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sf-venv",
   "language": "python",
   "name": "sf-venv"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}