02_characterize-genomic-context.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "species = [\"bambooshark\",\n",
    "           \"homo\",\n",
    "           \"gallus\",\n",
    "           \"xenopus\",\n",
    "           \"lepisosteus\",\n",
    "           \"bonytongue\",\n",
    "           \"ictalurus\",\n",
    "           \"danio\",\n",
    "           \"esox\",\n",
    "           \"gadus\",\n",
    "           \"takifugu\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load modules \n",
    "\n",
    "import phylopandas as phy\n",
    "from phylogenetics import tools\n",
    "import phylogenetics\n",
    "\n",
    "%matplotlib inline \n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import re, pickle, glob, os, random, string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# ---------------------------------------------------\n",
    "# Functions for plotting genomic context \n",
    "# ---------------------------------------------------\n",
    "\n",
    "def _smooth_hits(hits_df,window_size,global_bounds,\n",
    "                 weight_by_evalue=True,scale_min_feature_to=3):\n",
    "    \"\"\"\n",
    "    Generate running average of hits along the block.\n",
    "    \n",
    "    hits_df: dataframe with all hits on this block\n",
    "    window_size: size for smoothing window\n",
    "    global_bounds: tuple holding total size of block: (start,end)\n",
    "    weight_by_evalue: whether or not to weight average by the -log(evalue) of\n",
    "                      the hits.\n",
    "    scale_min_feature_to: to make faster, take the smallest feature in the data\n",
    "                          and rescale so this is scale_min_feature_to units long.\n",
    "    \"\"\"\n",
    "    \n",
    "    # copy df\n",
    "    df = hits_df.copy()\n",
    "    \n",
    "    # Go through every gene, looking for minimum gene length\n",
    "    for i in range(len(df.subject_start)):\n",
    "                \n",
    "        # Flip all genes so they go s -> e\n",
    "        if df.subject_start.iloc[i] > df.subject_end.iloc[i]:\n",
    "            s = df.subject_start.iloc[i]\n",
    "            e = df.subject_end.iloc[i]\n",
    "            \n",
    "            df.subject_start.iloc[i] = e\n",
    "            df.subject_end.iloc[i] = s\n",
    "\n",
    "    # See if we've seen a smaller gene\n",
    "    min_feature_size = window_size\n",
    "    min_seen = np.min(df.subject_end - df.subject_start)\n",
    "    if min_seen < min_feature_size:\n",
    "        min_feature_size = min_seen\n",
    "        \n",
    "    # Make minimum feature size scale_min_feature_to units long\n",
    "    scale = 10**np.floor(np.log10(min_feature_size/scale_min_feature_to))\n",
    "    \n",
    "    # Rescale indexes to be smaller and more managable\n",
    "    df.subject_start = np.array(np.round(df.subject_start/scale),dtype=np.int)\n",
    "    df.subject_end = np.array(np.round(df.subject_end/scale),dtype=np.int)\n",
    "    window_size = np.int(np.round(window_size/scale))\n",
    "    global_bounds = list(global_bounds[:])\n",
    "    global_bounds[0] = np.int(np.round(global_bounds[0]/scale))\n",
    "    global_bounds[1] = np.int(np.round(global_bounds[1]/scale))\n",
    "    \n",
    "    # Sort values from small to large\n",
    "    df = df.sort_values(by=\"subject_start\")\n",
    "    \n",
    "    # Edit e_values to be -log(e_value\n",
    "    e_values = np.copy(df.e_value)\n",
    "    e_values[np.array([e == 0 for e in e_values])] = 5.5e-126\n",
    "    df.e_value = -np.log(e_values)\n",
    "    \n",
    "    N = np.max(df.subject_end) - np.min(df.subject_start)\n",
    "    \n",
    "    # Window size can't be bigger than the total length of the fragment\n",
    "    if N < window_size:\n",
    "        window_size = N\n",
    "\n",
    "    \"\"\"\n",
    "    offset = df.subject_start.iloc[0]\n",
    "    to_convolve = np.zeros(N,dtype=np.float)\n",
    "    for i in range(len(df.subject_start)):\n",
    "        s = df.subject_start.iloc[i] - offset\n",
    "        e = df.subject_end.iloc[i] - offset\n",
    "        to_convolve[s:(e + 1)] += df.e_value[i]\n",
    "    \n",
    "    indexes = np.arange(len(to_convolve)) + offset\n",
    "    smoothed = np.convolve(to_convolve,\n",
    "                           np.ones((window_size,))/window_size,\n",
    "                           mode='same')\n",
    "    \"\"\"\n",
    "\n",
    "    # Hit sets are hits that need to be convolved given hit size.  Stick the\n",
    "    # first hit into the first hit set\n",
    "    hit_sets = [[]]\n",
    "    hit_sets[-1].append([df.subject_start.iloc[i],\n",
    "                         df.subject_end.iloc[i],\n",
    "                         df.e_value.iloc[i]])\n",
    "    \n",
    "    # Go over remaining hits\n",
    "    for i in range(1,len(df.subject_start)):\n",
    "        \n",
    "        # No overlap, so create new set in list\n",
    "        if df.subject_start.iloc[i] > (df.subject_end.iloc[i-1] + window_size):\n",
    "            hit_sets.append([])\n",
    "        \n",
    "        # Append hit to last hit set in the list\n",
    "        hit_sets[-1].append([df.subject_start.iloc[i],\n",
    "                             df.subject_end.iloc[i],\n",
    "                             df.e_value.iloc[i]])\n",
    "    \n",
    "    all_indexes = []\n",
    "    all_smoothed = []\n",
    "    for hit_set in hit_sets:\n",
    "        \n",
    "        # All hit boundaries in this set\n",
    "        all_positions = [hit[0] for hit in hit_set]\n",
    "        all_positions.extend([hit[1] for hit in hit_set])\n",
    "        \n",
    "        # Get outer edges of hit set\n",
    "        total_min = np.min(all_positions) - window_size\n",
    "        total_max = np.max(all_positions) + window_size\n",
    "            \n",
    "        # Populate array for doing convolution\n",
    "        to_convolve = np.zeros(total_max - total_min,dtype=np.float)\n",
    "        for hit in hit_set:\n",
    "            s = hit[0] - total_min\n",
    "            e = hit[1] - total_min\n",
    "            \n",
    "            if weight_by_evalue:\n",
    "                v = hit[2]\n",
    "            else:\n",
    "                v = 1.0\n",
    "                \n",
    "            to_convolve[s:(e+1)] += v\n",
    "        \n",
    "        # Do convolution\n",
    "        smoothed = np.convolve(to_convolve, \n",
    "                               np.ones((window_size,))/window_size,\n",
    "                               mode='same')\n",
    "        \n",
    "        indexes = np.arange(len(smoothed),dtype=np.float) + total_min\n",
    "        \n",
    "        all_indexes.append(np.copy(indexes))\n",
    "        all_smoothed.append(smoothed)\n",
    "\n",
    "    if all_indexes[0][-1] > global_bounds[0]:\n",
    "        all_indexes.insert(0,np.array([global_bounds[0]],dtype=np.float))\n",
    "        all_smoothed.insert(0,np.zeros(1,dtype=np.float))\n",
    "\n",
    "    if all_indexes[-1][-1] < global_bounds[0]:\n",
    "        all_indexes.append(np.array([global_bounds[1]],dtype=np.float))\n",
    "        all_smoothed.append(np.zeros(1,dtype=np.float))\n",
    "        \n",
    "    indexes = np.concatenate(all_indexes)\n",
    "    smoothed = np.concatenate(all_smoothed)\n",
    "    \n",
    "    outside_mask = np.logical_and(indexes >= global_bounds[0],\n",
    "                                  indexes <= global_bounds[1])\n",
    "\n",
    "    return indexes*scale, smoothed\n",
    "\n",
    "\n",
    "def plot_hits_on_assembly(hit_df,\n",
    "                          subjects_to_plot=None,\n",
    "                          window_size=0.0001,\n",
    "                          focal_query=None,\n",
    "                          focal_locations=None,\n",
    "                          block_lengths=None,\n",
    "                          use_evalue_hit_height=True,\n",
    "                          min_evalue=5.5e-126,\n",
    "                          num_blocks_to_plot=5,\n",
    "                          min_hits_cutoff=4,\n",
    "                          figsize=(10,5),\n",
    "                          dpi=90,\n",
    "                          only_plot_focal=False,\n",
    "                          manual_title=None):\n",
    "    \"\"\"\n",
    "    hit_df\n",
    "    subjects_to_plot: either string for single subject species to make\n",
    "                      plot for or a list of strings containing subject\n",
    "                      species.  if None, plot all subject species in \n",
    "                      the dataframe.\n",
    "    window_size: size of window for making bar plots of hits/per genomic\n",
    "                 region. if < 1, use as a fraction of total size.  if > 1\n",
    "                 use as an actual window. \n",
    "    focal_query: focal query used to generate the regions of interest. \n",
    "                 If defined, the focal_query is *not* included in the\n",
    "                 bar height, and will be highlighted with its own color.\n",
    "                 If None, ignore. \n",
    "    focal_locations: dictionary containing locations of known focal genes.\n",
    "                     key should be subject_species, value should be a \n",
    "                     tuple like (block,location_on_block).  If None,\n",
    "                     ignore. \n",
    "    block_lengths: dictionary of block lengths for this species\n",
    "    use_evalue_hit_height: weight the height of bars by sum(-log(evalue)) \n",
    "                           for hits in the window\n",
    "    min_evalue: assign any evalue below min_evalue to min_evalue. \n",
    "                (needed to deal with evalue = 0)\n",
    "    num_blocks_to_plot: number of blocks to plot per subject\n",
    "    min_hits_cutoff: don't plot a block if it has fewer than\n",
    "                     min_hits_cutoff hits\n",
    "    figsize: figure size (passed directly to matplotlib.pyplot.figure)\n",
    "    dpi: figure dpi (passed directly to matplotlib.pyplot.figure)\n",
    "    only_plot_focal: only plot block if it has a focal hit\n",
    "    \"\"\"\n",
    "    \n",
    "    plot_y_min = -0.2\n",
    "    plot_y_max =  1.0\n",
    "    focal_color = \"green\"\n",
    "    gene_vertical = -0.1\n",
    "    species_colors = {\"homo\":\"indigo\",\"danio\":\"peru\"}\n",
    "    \n",
    "    # -------------------------------------------------------------- \n",
    "    # Sanity checking\n",
    "    # --------------------------------------------------------------\n",
    "\n",
    "    if focal_locations is not None and focal_query is None:\n",
    "        err = \"focal_locations can only be specified if a focal_query\"\n",
    "        err += \" is specified.\\n\"\n",
    "        raise ValueError(err)\n",
    "    \n",
    "    # -------------------------------------------------------------- \n",
    "    # Figure out what subject species we are going to make plots for\n",
    "    # --------------------------------------------------------------\n",
    "        \n",
    "    if subjects_to_plot is None:\n",
    "        subjects_to_plot = list(set(hit_df.subject_species))\n",
    "        subjects_to_plot.sort()\n",
    "    else:\n",
    "        if type(subjects_to_plot) is str:\n",
    "            subjects_to_plot = [subjects_to_plot]\n",
    "\n",
    "    # --------------------------------------------------------------\n",
    "    # X-axis and which blocks to plot\n",
    "    #\n",
    "    # Find the longest chunk of genome we'll have to plot.  Put\n",
    "    # every assembly/chromosome etc. on this x-length scale so plots\n",
    "    # are comprable.  Main output is \"plot_bounds\" dictionary, which\n",
    "    # gives us the bounds for any plot that involves a subject species\n",
    "    # and block pair. \n",
    "    #\n",
    "    # Also record the number of hits on each block, creating a \n",
    "    # dictionary that keys each subject species to a list that holds\n",
    "    # the blocks ranked from most to fewest hits. \n",
    "    # --------------------------------------------------------------\n",
    "    \n",
    "    # Do pass through all pairs, finding maximum span of each block,\n",
    "    # as well as the maxium span seen overall. \n",
    "    biggest_span = -1\n",
    "    all_min_max = {}\n",
    "    block_hit_numbers = {}\n",
    "    \n",
    "    subject_species = set(hit_df.subject_species)\n",
    "    for ss in subject_species:\n",
    "        \n",
    "        all_min_max[ss] = {}\n",
    "        block_hit_numbers[ss] = []\n",
    "        \n",
    "        ss_df = hit_df[hit_df.subject_species == ss]\n",
    "        blocks = set(ss_df.block)\n",
    "        for b in blocks:\n",
    "            \n",
    "            block_df = ss_df[ss_df.block == b]\n",
    "            all_positions = block_df.subject_start.tolist()\n",
    "            all_positions.extend(block_df.subject_end.tolist())\n",
    "            min_seen = np.min(all_positions)\n",
    "            max_seen = np.max(all_positions)\n",
    "    \n",
    "            all_min_max[ss][b] = (min_seen,max_seen)\n",
    "        \n",
    "            if (max_seen - min_seen) > biggest_span:\n",
    "                biggest_span = (max_seen - min_seen)\n",
    "                \n",
    "            num_hits_on_block = len(block_df)\n",
    "            block_hit_numbers[ss].append((num_hits_on_block,b))\n",
    "        \n",
    "        # Sort blocks so they go from highest number of hits to \n",
    "        # lowest number, keyed to each species\n",
    "        block_hit_numbers[ss].sort(reverse=True)\n",
    "    \n",
    "    biggest_span = biggest_span*1.1\n",
    "    \n",
    "    # Go back through pairs, creating plot bounds that add offsets to\n",
    "    # left and right so region of interest is always centered. \n",
    "    x_plot_bounds = {}\n",
    "    for ss in all_min_max.keys():\n",
    "        x_plot_bounds[ss] = {}\n",
    "        for b in all_min_max[ss].keys():\n",
    "            min_seen = all_min_max[ss][b][0]\n",
    "            max_seen = all_min_max[ss][b][1]\n",
    "            span = (max_seen - min_seen)\n",
    "            \n",
    "            leftover = (biggest_span - span)/2. \n",
    "            \n",
    "            x_plot_bounds[ss][b] = (min_seen,max_seen,leftover)\n",
    "    \n",
    "    # Figure out window size to use\n",
    "    if window_size < 1:\n",
    "        window_size = np.int(np.round(window_size*biggest_span))\n",
    "    \n",
    "    # --------------------------------------------------------------\n",
    "    # Y-axis\n",
    "    #\n",
    "    # Find the maximum height we could get on our output plot\n",
    "    # if all of the query species \n",
    "    # --------------------------------------------------------------\n",
    "    \n",
    "    # Get the maximum number of query genes that come from a single \n",
    "    # query species\n",
    "    num_queries = []\n",
    "    query_species = set(hit_df.query_species)\n",
    "    for qs in query_species:\n",
    "        qs_df = hit_df[hit_df.query_species == qs]\n",
    "        query_names = set(qs_df.query_name)\n",
    "        num_queries.append(len(query_names))\n",
    "    \n",
    "    # Maximum possible height is total number of queries less the focal\n",
    "    # sequence, which we ignore in the score. \n",
    "    height_norm_factor = np.max(num_queries)\n",
    "    if focal_query is not None:\n",
    "        height_norm_factor = height_norm_factor - 1\n",
    "        \n",
    "    if use_evalue_hit_height:\n",
    "        height_norm_factor *= -np.log(min_evalue)\n",
    "    \n",
    "    # HACK HACK HACK\n",
    "    height_norm_factor = 53\n",
    "    \n",
    "    # --------------------------------------------------------------\n",
    "    # Get down to plotting business\n",
    "    # --------------------------------------------------------------\n",
    "    \n",
    "    for ss in subjects_to_plot:\n",
    "        \n",
    "        df = all_hits[all_hits.subject_species == ss]\n",
    "        \n",
    "        # Construct dictionary of focal locations, with start keyed\n",
    "        # to block name.\n",
    "        if focal_locations is not None:\n",
    "            \n",
    "            try:\n",
    "                tmp = focal_locations[ss]\n",
    "                focal_for_species = dict([(t[0],t[1]) for t in tmp])\n",
    "            except KeyError:\n",
    "                focal_for_species = {}\n",
    "        else:\n",
    "            focal_for_species = {}\n",
    "    \n",
    "        num_blocks = len(block_hit_numbers[ss])\n",
    "        if num_blocks > num_blocks_to_plot:\n",
    "            num_blocks = num_blocks_to_plot\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        fig, ax = plt.subplots(num_blocks,1,\n",
    "                               figsize=(figsize[0],figsize[1]*(num_blocks)),\n",
    "                               dpi=dpi,frameon=False)\n",
    "        \n",
    "        for i in range(num_blocks):\n",
    "\n",
    "            \n",
    "            num_hits = block_hit_numbers[ss][i][0] \n",
    "            block = block_hit_numbers[ss][i][1]\n",
    "            \n",
    "            # If too few hits, we're done plotting for this subject\n",
    "            # because the block hits are ordered from most to fewest\n",
    "            # hits. \n",
    "            if num_hits < min_hits_cutoff:\n",
    "                break\n",
    "            \n",
    "            # Find the focal position on this block if it exists\n",
    "            focal_position = None\n",
    "            try:\n",
    "                focal_position = focal_for_species[block]\n",
    "            except KeyError:\n",
    "                pass\n",
    "            \n",
    "            if only_plot_focal and focal_position is None:\n",
    "                continue\n",
    "            \n",
    "            x_bounds = x_plot_bounds[ss][block][:2]\n",
    "            x_offset = x_plot_bounds[ss][block][2]\n",
    "            \n",
    "            #fig = figure(1, figsize=figsize, dpi=dpi)\n",
    "            #ax[i] = fig.add_subplot(111,frameon=False)\n",
    "            ax[i].set_xlim((x_bounds[0]-x_offset,x_bounds[1]+x_offset))\n",
    "            ax[i].set_ylim((plot_y_min,plot_y_max))\n",
    "            \n",
    "            # Try to add line for focal position, if this is in this\n",
    "            # species on this block.\n",
    "            if focal_position is not None:\n",
    "                ax[i].plot((focal_position,focal_position),\n",
    "                             (plot_y_min,gene_vertical),\"--\",color=focal_color)\n",
    "\n",
    "            block_df = df[df.block == block]\n",
    "            \n",
    "            for qs in query_species:\n",
    "                \n",
    "                qs_df = block_df[block_df.query_species == qs]\n",
    "                \n",
    "                ax[i].plot(qs_df.subject_start,\n",
    "                             np.ones(len(qs_df.subject_start))*gene_vertical,\n",
    "                             \"^\",color=species_colors[qs])\n",
    "                \n",
    "                tmp_df = qs_df.copy()\n",
    "                tmp_df = tmp_df[tmp_df.query_name != focal_query]\n",
    "                \n",
    "                if len(tmp_df.e_value) > 0:\n",
    "                    indexes, smoothed = _smooth_hits(tmp_df,\n",
    "                                                     window_size=window_size,\n",
    "                                                     global_bounds=x_bounds,\n",
    "                                                     weight_by_evalue=use_evalue_hit_height)\n",
    "                \n",
    "                    smoothed = smoothed/height_norm_factor\n",
    "                    ax[i].fill_between(indexes,smoothed,color=species_colors[qs],alpha=0.5)\n",
    "                    ax[i].plot(indexes,smoothed,'-',color=species_colors[qs],linewidth=0.5)\n",
    "                \n",
    "                num_hits_drawn = len(qs_df)\n",
    "                for j in range(num_hits_drawn):\n",
    "                    if qs == \"homo\":\n",
    "                        x_value = x_bounds[0] - x_offset*0.9\n",
    "                    else:\n",
    "                        x_value = x_bounds[0] + x_offset*0.9\n",
    "                        \n",
    "                    \n",
    "                    #ax[i].text(qs_df.iloc[j].subject_start,\n",
    "                    #             0.8,qs_df.iloc[j].query_name,\n",
    "                    #             color=species_colors[qs],ha=\"right\",\n",
    "                    #             rotation=90)\n",
    "                    \n",
    "                    \n",
    "                    num = qs_df.iloc[j].subject_start\n",
    "                    name = qs_df.iloc[j].query_name\n",
    "                    out = \"{}: {}\".format(num,name)\n",
    "                    ax[i].text(x_value,\n",
    "                             (num_hits_drawn - j )/18, #num_hits_drawn,\n",
    "                             out,\n",
    "                             color=species_colors[qs],ha=\"left\")\n",
    "                            \n",
    "                               \n",
    "                \n",
    "\n",
    "            # Plot genome line\n",
    "            if block_lengths is not None:\n",
    "                \n",
    "                block_length = block_lengths[block]\n",
    "\n",
    "                block_start_location = 0\n",
    "                block_end_location = block_length \n",
    "\n",
    "                if block_start_location < (x_bounds[0]-x_offset):\n",
    "                    block_start_location = x_bounds[0]-x_offset\n",
    "                    ax[i].plot((block_start_location,block_start_location),(-0.2,1),\"b--\")\n",
    "\n",
    "                if block_end_location > (x_bounds[1]+x_offset):\n",
    "                    block_end_location = x_bounds[1] + x_offset\n",
    "                    ax[i].plot((block_end_location,block_end_location),(-0.2,1),\"b--\")\n",
    "                    \n",
    "            else:\n",
    "                block_start_location = x_bounds[0]\n",
    "                block_end_location = x_bounds[1]\n",
    "            \n",
    "            # Plot query hits on top\n",
    "            focal_query_df = block_df[block_df.query_name == focal_query]\n",
    "            ax[i].plot(focal_query_df.subject_start,\n",
    "                         np.ones(len(focal_query_df.subject_start))*gene_vertical,\n",
    "                          \"^\",color=focal_color)\n",
    "\n",
    "            # Plot genome\n",
    "            ax[i].plot((block_start_location,block_end_location),(0,0),\"k-\",linewidth=2) #,zorder=1000000)\n",
    "            \n",
    "                        \n",
    "            if i == 0 and manual_title is not None:\n",
    "                title = \"{}\\n\\n{}\".format(manual_title,block)\n",
    "                ax[i].set_title(title) #manual_title,pad=-5)\n",
    "            else:\n",
    "                ax[i].set_title(\" \\n \\n{}\".format(block)) #\"{}, {}\".format(ss,block))\n",
    "                \n",
    "            ax[i].set_xlim((x_bounds[0]-x_offset,x_bounds[1]+x_offset))\n",
    "            ax[i].set_ylim((plot_y_min,plot_y_max))\n",
    "            ax[i].spines['top'].set_visible(False)\n",
    "            ax[i].spines['right'].set_visible(False)\n",
    "            #ax[i].spines['bottom'].set_visible(False)\n",
    "            #ax[i].spines['left'].set_visible(False)\n",
    "            \n",
    "            ax[i].set_xlabel(\"genome position (Mb)\")\n",
    "            ax[i].set_ylabel(\"relative evalue/bp\")\n",
    "    \n",
    "        plt.tight_layout()\n",
    "            #plt.savefig(\"new_figs/{}-{}-{}.pdf\".format(ss,i,block))\n",
    "            #plt.show()\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data written out from 01_blast-genomes.ipynby\n",
    "block_lengths = {}\n",
    "for s in species:\n",
    "    block_lengths[s] = pickle.load(open(\"{}_length.pickle\".format(s),\"rb\"))\n",
    "all_hits = pd.read_csv(\"all-hits.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct ortholog locations by reverse BLAST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "# Reverse blast hits from TLR4 against human and danio\n",
    "# This is slow.  It saves out homo_rblast_hits.pickle and\n",
    "# danio_rblast_hits.pickle. These can be loaded in the next\n",
    "# cell, meaning you only need to run this cell once, even if \n",
    "# you tweak the downstream analysis. \n",
    "\n",
    "homo_rblast_hits = {}\n",
    "danio_rblast_hits = {}\n",
    "for s in species:\n",
    "    s_df = all_hits[np.logical_and(all_hits.subject_species == s,all_hits.query_name==\"TLR4\")]\n",
    "\n",
    "    print(s)\n",
    "    \n",
    "    homo_rblast_hits[s] = {}\n",
    "    danio_rblast_hits[s] = {}\n",
    "    for i in range(len(s_df.sequence)):\n",
    "        uid = s_df.iloc[i].uid\n",
    "        homo_rblast_hits[s][uid] = tools.blast.local_blast(s_df.iloc[i],db=\"reference/homo\",hitlist_size=10)\n",
    "        danio_rblast_hits[s][uid] = tools.blast.local_blast(s_df.iloc[i],db=\"reference/danio\",hitlist_size=10)\n",
    "\n",
    "pickle.dump(homo_rblast_hits,open(\"homo_rblast_hits.pickle\",\"wb\"))\n",
    "pickle.dump(danio_rblast_hits,open(\"danio_rblast_hits.pickle\",\"wb\"))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dictionary keying species to TLR4 ortholog locations \n",
    "\n",
    "homo_rblast_hits = pickle.load(open(\"homo_rblast_hits.pickle\",\"rb\"))\n",
    "danio_rblast_hits = pickle.load(open(\"danio_rblast_hits.pickle\",\"rb\"))\n",
    "pattern = re.compile(\"TLR4|toll-like receptor 4|toll like receptor 4\",flags=re.IGNORECASE)\n",
    "worst_allowed = 0\n",
    "\n",
    "ortholog_locations = {}\n",
    "for s in species:\n",
    "    s_df = all_hits[np.logical_and(all_hits.subject_species == s,all_hits.query_name==\"TLR4\")]\n",
    "\n",
    "    for i in range(len(s_df.sequence)):\n",
    "        uid = s_df.iloc[i].uid\n",
    "        homo_hits = homo_rblast_hits[s][uid]\n",
    "        for j in range(len(homo_hits.hit_def)):\n",
    "            \n",
    "            if pattern.search(homo_hits.hit_def.iloc[j]):\n",
    "                if j <= worst_allowed:\n",
    "                    \n",
    "                    block = s_df[s_df.uid == uid].block.iloc[0]\n",
    "                    start = s_df[s_df.uid == uid].subject_start.iloc[0]\n",
    "                    end = s_df[s_df.uid == uid].subject_end.iloc[0]\n",
    "                    out_tuple = (block,start,end)\n",
    "                    try:\n",
    "                        ortholog_locations[s].append(out_tuple)\n",
    "                    except KeyError:\n",
    "                        ortholog_locations[s] = [out_tuple]\n",
    "                    \n",
    "tlr4_ortholog_locations = ortholog_locations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate graphs plotting hit e-value/bp along genomes of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for i, s in enumerate(species):\n",
    "\n",
    "    print(s)\n",
    "    plot_hits_on_assembly(all_hits,\n",
    "                          subjects_to_plot=s,\n",
    "                          window_size=0.0001,\n",
    "                          focal_query=\"TLR4\",\n",
    "                          focal_locations=tlr4_ortholog_locations,\n",
    "                          block_lengths=block_lengths[s],\n",
    "                          num_blocks_to_plot=10,\n",
    "                          min_hits_cutoff=1,\n",
    "                          figsize=(10,4),\n",
    "                          use_evalue_hit_height=True,\n",
    "                          manual_title=\"Fig S{}: hits on {} chromosomes\".format(i+2,s),\n",
    "                          dpi=300)\n",
    "    \n",
    "    plt.savefig(\"Fig-S{}.pdf\".format(i+2))\n"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}