Notebook - Post Inference Analysis 1 - Analysis of fractional distance.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis of fractional distance\n",
    "\n",
    "To determine whether more data equates to reaching a canonical form of the reconstructions we performed ASR on a number of varying sized alignments for two datasets, DHAD and KARI. \n",
    "\n",
    "The reconstructions were generated by using GRASP in November, 2019. All reconstructions are using joint reconstructions and only the consensus sequence is considered.\n",
    "\n",
    "Method to determine the fractional distance:\n",
    "\n",
    "```\n",
    "    1. Load reconstructed tree (output of GRASP) \n",
    "    2. Choose the smallest tree as the reference\n",
    "    3. For every tree, with respect to the reference tree from (2), calculate equivalent nodes in the larger tree\n",
    "    4. Pick several nodes to inspect from (2)\n",
    "    5. Using the downloaded consensus sequences from GRASP, load the alignment and pull out the consensus sequence for equivalent nodes in each of the alignments using the dictionary created in (3)\n",
    "    6. Remove all the gaps for each set of sequences from 5\n",
    "    7. Sort the sequences from (6) so that they are in order from smallest alignment to largest and align sequences from (6) using MAFFT\n",
    "    8. Reload the alignment and calculate the fractional distances using binfpy\n",
    "    9. Plot heatmaps\n",
    "    \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Useful functions below (used for the analysis and visualisation)\n",
    "\n",
    "The functions below are used for each analysis, don't need to be read in detail (except for write_heatmap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import sys\n",
    "import os\n",
    "import seaborn as sns\n",
    "sys.path.append('./src/binfpy')\n",
    "sys.path.append('./src/SimilarNodeFinder/src')\n",
    "\n",
    "import sequence\n",
    "import webservice\n",
    "import numpy as np\n",
    "import pickle\n",
    "import csv\n",
    "import scipy.stats as stats\n",
    "import matplotlib.pyplot as plt\n",
    "from tree_object import TreeObject\n",
    "from tree_controller import TreeController\n",
    "from scipy.stats import pearsonr\n",
    "\n",
    "\n",
    "cyp2_str = 'CYP2'\n",
    "kari_str = 'KARI1'\n",
    "dhad_str = 'DHAD'\n",
    "\n",
    "\n",
    "\n",
    "# Choose your colour\n",
    "colour_map = \"Greens\"\n",
    "colour_map = \"BuPu\"\n",
    "colour_map = \"YlGnBu\"\n",
    "colour_map = \"Blues\"\n",
    "colour_map = \"mako\"\n",
    "colour_map = \"PiYG\"\n",
    "colour_map = \"gist_rainbow\"\n",
    "# Would recommend the below\n",
    "colour_map = \"BuPu\"\n",
    "colour_map = \"viridis\"\n",
    "colour_map = \"BuPu_r\"\n",
    "# Set the number of ticks between each tick, one means every label is showing\n",
    "numerics = ['int16', 'int32', 'int64', 'float16', 'float32', 'float64']\n",
    "\n",
    "\n",
    "n_ticks = 1\n",
    "\n",
    "def save_plt(fig, name, dpi=300):\n",
    "    # Old version of save fig doesn't make it high res.\n",
    "    fig.savefig(name, format = \"png\", bbox_inches = 'tight', pad_inches = 0, dpi=dpi)\n",
    "\n",
    "\n",
    "def generate_alignment(filename):\n",
    "    filename_fasta = filename + '.fa'\n",
    "    filename_aln = filename + '.aln'\n",
    "    ! mafft --auto --quiet --thread 4 $filename_fasta > $filename_aln\n",
    "\n",
    "def writeToCsv(writer, num_seqs, node_name, seq_name, dists):\n",
    "    \"\"\"\n",
    "    Helper to write to the csv file\n",
    "    \"\"\"\n",
    "    to_write = [str(num_seqs), str(node_name), str(seq_name),]\n",
    "    \n",
    "    for d in dists:\n",
    "        to_write.append(str(d))\n",
    "\n",
    "    writer.writerow(to_write)\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "               \n",
    "def plotFracHeatmap(df, node, labels, vmax=0.6):\n",
    "    fig, ax = plt.subplots()\n",
    "    \n",
    "    im, cbar = heatmap(df[fact_names], labels, labels, node + \" fractional distance\", ax=ax, cmap=\"BuPu\", cbarlabel=\"Distance\")#, vmax=vmax)\n",
    "    plt.show()\n",
    "    return fig\n",
    "\n",
    "def heatmap(data, row_labels, col_labels, title, ax=None,\n",
    "            cbar_kw={}, cbarlabel=\"\", **kwargs):\n",
    "    \"\"\"\n",
    "    Create a heatmap from a numpy array and two lists of labels.\n",
    "\n",
    "    Arguments:\n",
    "        data       : A 2D numpy array of shape (N,M)\n",
    "        row_labels : A list or array of length N with the labels\n",
    "                     for the rows\n",
    "        col_labels : A list or array of length M with the labels\n",
    "                     for the columns\n",
    "    Optional arguments:\n",
    "        ax         : A matplotlib.axes.Axes instance to which the heatmap\n",
    "                     is plotted. If not provided, use current axes or\n",
    "                     create a new one.\n",
    "        cbar_kw    : A dictionary with arguments to\n",
    "                     :meth:`matplotlib.Figure.colorbar`.\n",
    "        cbarlabel  : The label for the colorbar\n",
    "    All other arguments are directly passed on to the imshow call.\n",
    "    \"\"\"\n",
    "\n",
    "    if not ax:\n",
    "        ax = plt.gca()\n",
    "\n",
    "    # Plot the heatmap\n",
    "    im = ax.imshow(data, **kwargs)\n",
    "\n",
    "    # Create colorbar\n",
    "    cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw)\n",
    "    cbar.ax.set_ylabel(cbarlabel, rotation=-90, va=\"bottom\")\n",
    "\n",
    "    # We want to show all ticks...\n",
    "    ax.set_xticks(np.arange(data.shape[1]))\n",
    "    ax.set_yticks(np.arange(data.shape[0]))\n",
    "    # ... and label them with the respective list entries.\n",
    "    ax.set_xticklabels(col_labels)\n",
    "    ax.set_yticklabels(row_labels)\n",
    "\n",
    "    # Let the horizontal axes labelling appear on top.\n",
    "\n",
    "    # Rotate the tick labels and set their alignment.\n",
    "    plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "    # Turn spines off and create white grid.\n",
    "    for edge, spine in ax.spines.items():\n",
    "        spine.set_visible(False)\n",
    "\n",
    "    ax.set_xticks(np.arange(data.shape[1]+1)-.5, minor=True)\n",
    "    ax.set_yticks(np.arange(data.shape[0]+1)-.5, minor=True)\n",
    "    # ax.grid(which=\"minor\", color=\"w\", linestyle='-', linewidth=3)\n",
    "    ax.tick_params(which=\"minor\", bottom=False, left=False)\n",
    "    ax.set_title(title)\n",
    "    \n",
    "    return im, cbar\n",
    "\n",
    "def write_heatmap(output_dir_str, sequences_of_interest, node_name_dict):\n",
    "    \"\"\"\n",
    "    \n",
    "    Calculates the file for the heatmap creation. This involves aligning all the sequences \n",
    "    (from the various reconstructions) for each node of interest.\n",
    "    Then re-reading this in and calculating the fractional distance.\n",
    "    \n",
    "    We also keep track of the average length of sequences for a given node (just out of interest)\n",
    "    and print these results\n",
    "    \n",
    "    \"\"\"\n",
    "    with open(output_dir_str, 'w+') as output_writer:\n",
    "        writer = csv.writer(output_writer, delimiter=',', quotechar='\"', quoting=csv.QUOTE_MINIMAL)\n",
    "        row = ['num_seqs', 'aln_name', 'seq_name']\n",
    "        # Write the first row\n",
    "        for s in sequences_of_interest:\n",
    "            row.append(str(s))\n",
    "        # Write our header row\n",
    "        writer.writerow(row)\n",
    "        \n",
    "        # Here we want to iterate through our node dictionary and get all the consensus sequences\n",
    "        # we're interested in. These are by default gappy.\n",
    "        for node in node_name_dict:\n",
    "            # Get the consensus sequences from all the reconstructions for that node (e.g. N1 would have sequences\n",
    "            # from recon1, recon2 etc.)\n",
    "            # Write this to a file\n",
    "            outfilename = output_dir_str.split('.')[0] + node \n",
    "            seqs_all = node_name_dict[node]['seqs']\n",
    "            seqs_dict = {}\n",
    "            seqs_names = []\n",
    "            # Now we want to build up a dictionary of the names to the sequences so we can order them easily\n",
    "            for s in seqs_all:\n",
    "                seqs_names.append(s.name)\n",
    "                seqs_dict[s.name] = s\n",
    "            seqs = []\n",
    "            # Order the names\n",
    "            seqs_names.sort()\n",
    "            # Recreate a list of sequences now in a sorted order\n",
    "            for s in seqs_names:\n",
    "                seqs.append(seqs_dict[s])\n",
    "\n",
    "            non_gappy_seqs = []\n",
    "            seq_lens = []\n",
    "            # Now we want to save this as an alignment, but we want to remove the gaps\n",
    "            for s in seqs:\n",
    "                count_no_gaps = 0\n",
    "                no_gap = ''\n",
    "                for c in s:\n",
    "                    if c != '-':\n",
    "                        count_no_gaps += 1\n",
    "                        no_gap += c\n",
    "                # Create a new sequence object just with the gaps removed\n",
    "                non_gappy_seqs.append(sequence.Sequence(no_gap, sequence.Protein_Alphabet, name=s.name))\n",
    "                seq_lens.append(count_no_gaps)\n",
    "\n",
    "            # Calculate various metrics for the length\n",
    "            var_lens = np.var(seq_lens)\n",
    "            sd_lens = np.std(seq_lens)\n",
    "            mean_len = np.mean(seq_lens)\n",
    "            median_len = np.median(seq_lens)\n",
    "            print(\"-------------- \" + node + '-------------')\n",
    "            print('var: ', var_lens, ', stdev: ', sd_lens, ', mean: ', mean_len, ', median: ', median_len)\n",
    "\n",
    "            # Write this to a fasta file and generate an alignemnt from this\n",
    "            sequence.writeFastaFile(outfilename + '.fa', non_gappy_seqs)\n",
    "\n",
    "            generate_alignment(outfilename)\n",
    "\n",
    "            # Read in the aligned sequences and from this determine the consensus etc.\n",
    "            seqs_aligned = sequence.readFastaFile(outfilename + '.aln', sequence.Protein_Alphabet, gappy=True)\n",
    "\n",
    "            alignment = sequence.Alignment(seqs_aligned)\n",
    "\n",
    "            # For each sequence in the alignment calculate the divergence from the consensus using fractional\n",
    "            distances = alignment.calcDistances(\"fractional\")\n",
    "            seq_dists = {}\n",
    "            # Now we want to assign the distances to a given sequence (so we can again ensure the order is maintained)\n",
    "            for s_idx in range(0, len(seqs_aligned)):\n",
    "                seq = seqs_aligned[s_idx]\n",
    "                dists = distances[s_idx]\n",
    "                seq_dists[seq.name] = dists\n",
    "            # Finally for each sequence, write it to a file\n",
    "            for seq in seqs:     \n",
    "                writeToCsv(writer, seq.name.split('_')[0], node, seq.name, seq_dists[seq.name])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DHAD\n",
    "\n",
    "First we will analyse the DHADs. These were made using a base dataset of 1612. However, we will calculate the nodes with respect to an original tree of 585 sequences, titled sp_cured_3_01112018_reconstructed-tree_GRASP.nwk\n",
    "\n",
    "There are 16 datasets, each 500 sequences more than the preceding until we reach 9112 sequences. \n",
    "\n",
    "Note this will take a few minutes to run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (1 - 3) Load the trees and calculate the matching nodes to our reference\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tree_object import TreeObject\n",
    "from tree_controller import TreeController\n",
    "\n",
    "input_dir_str = 'Files/PIA_1_Files/GRASP_recons/' + dhad_str + '/'\n",
    "output_dir_str = 'Files/PIA_1_Files/output/'\n",
    "\n",
    "\n",
    "trees = []\n",
    "files = os.listdir(input_dir_str)\n",
    "for f in files:\n",
    "    if '.nwk' in f:\n",
    "        trees.append(f)\n",
    "\n",
    "# Choose the base tree\n",
    "base_tree = 'sp_cured_3_01112018_reconstructed-tree_GRASP.nwk'\n",
    "node_matching_dict = {}\n",
    "\n",
    "# Create our node matching dictionary, print out if we get any errors.\n",
    "for t in trees:\n",
    "    try:\n",
    "        t1 = TreeObject(input_dir_str + base_tree) \n",
    "        t2 = TreeObject(input_dir_str + t)\n",
    "        tc = TreeController()\n",
    "        node_matching = tc.get_similar_nodes_between_trees(t1, t2, True)\n",
    "        node_matching_dict[t] = node_matching\n",
    "        print(\"Complete: \" + t)\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (4) Find the matching nodes in the larger trees for our chosen nodes of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we make our main dictionary using the node identifiers for our base/reference tree\n",
    "node_name_dict = {'N1': {'seqs': []}, 'N423': {'seqs': []}, 'N560': {'seqs': []}}\n",
    "\n",
    "nodes = ['N1', 'N423', 'N560']\n",
    "\n",
    "for t, tree in node_matching_dict.items():\n",
    "    for n in nodes:\n",
    "        try:\n",
    "            \"\"\"\n",
    "            Here we're using t_num as our method for matching our trees and our alignments, this\n",
    "            will be specific to each dataset. Our trees are: r_1000_2612_dhad_18032019.nwk_recon.nwk\n",
    "            So by spliting on _ and taking the 1st element, we get 1000 for the tree above.          \n",
    "            \"\"\"\n",
    "            t_num = t.split('_')[1]\n",
    "            node_name_dict[n][t_num] = tree.get(n).original_label\n",
    "            print(\"Matched: \\t\" + n + \"\\t -> \\t\" + tree.get(n).get_label())\n",
    "        except Exception as e:\n",
    "            print(n, t, e)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (5) Load in the alignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dhad_files = []\n",
    "files = os.listdir(input_dir_str)\n",
    "for f in files:\n",
    "    if '.aln' in f:\n",
    "        dhad_files.append(f)\n",
    "\n",
    "print(dhad_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (6) Pick out the consensus sequences from each alignment for our nodes of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Here we want to iterate through our alignment files (dhad_files) and pick out the consensus\n",
    "sequences that match with our base/reference tree nodes of interest\n",
    "\"\"\"\n",
    "for f in dhad_files:\n",
    "    lengths = []\n",
    "    num_seqs = []\n",
    "    # Read in the fasta file. This is the output of GRASP (for the joint consensus ancestors)\n",
    "    seqs_aligned = sequence.readFastaFile(input_dir_str + f, sequence.Protein_Alphabet, gappy=True)\n",
    "    seq_dict = {}\n",
    "    for seq in seqs_aligned:\n",
    "        seq_dict[seq.name] = seq\n",
    "    i = 0\n",
    "    \"\"\"\n",
    "    Here we want a string to match between our reconstructed trees and our alignments from GRASP\n",
    "    For here our trees are annotated as: r_5500_7112_dhad_18032019.nwk_recon.nwk\n",
    "    and our alignments as: recon_5500.aln\n",
    "    Thus we want to match the 5500 from the tree to the 5500 from the alignment\n",
    "\n",
    "    \"\"\"\n",
    "    match_string = f.split('_')[1].split('.')[0]\n",
    "    # Check if this is a sequence of interest\n",
    "    for node_name, node_dict in node_name_dict.items():\n",
    "        for seq_nums, matching_node_name in node_dict.items():\n",
    "            if seq_nums == match_string:\n",
    "                s = seq_dict[matching_node_name]\n",
    "                no_gap = ''\n",
    "                for c in s:\n",
    "                    if c != '-':\n",
    "                        no_gap += c\n",
    "                seq = sequence.Sequence(no_gap, sequence.Protein_Alphabet, name=seq_nums + '_' + s.name)\n",
    "                node_dict['seqs'].append(seq)\n",
    "    print(\"Finished: \\t\" + f + \" \\t Matched: \" + str(len(node_dict['seqs'])) + \" nodes\")\n",
    "    i += 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (7 - 8) Calculate the fractional distances\n",
    "\n",
    "Now we have the sequences for each of the nodes, we want to:\n",
    "```\n",
    "1. Remove all the gaps\n",
    "2. Create an alignment of the sequences\n",
    "3. Calculate the fractional distances from these\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write out which sequences were interested in (this is really the alignment names, so number of sequences)\n",
    "sequences_of_interest = [1612, 2112, 2612, 3112, 3612, 4112, 4612, 5112, 5612, 6112, 6612, 7112, 7612, 8112, 8612, 9112]\n",
    "\n",
    "# Write heatmap calculates the fractional distances\n",
    "write_heatmap(output_dir_str + 'DHAD_heatmap.csv', sequences_of_interest, node_name_dict)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (9) Plot the heatmaps using the file that we just wrote"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "csv_file_name = output_dir_str + 'DHAD_heatmap.csv'\n",
    "nodes_of_interest = ['N1', 'N423', 'N560']\n",
    "sequences_of_interest = [1612, 2112, 2612, 3112, 3612, 4112, 4612, 5112, 5612, 6112, 6612, 7112, 7612, 8112, 8612, 9112]\n",
    "                         \n",
    "# Read in the csv\n",
    "df = pd.read_csv(csv_file_name)\n",
    "\n",
    "# Ensure our names are strings\n",
    "fact_names = []\n",
    "for n in sequences_of_interest:\n",
    "    fact_names.append(str(n))\n",
    "\n",
    "# Group the data\n",
    "df_N1 = df[df['aln_name'] == 'N1']\n",
    "df_N423 = df[df['aln_name'] == 'N423']\n",
    "df_N560 = df[df['aln_name'] == 'N560']\n",
    "\n",
    "# Plot the heatmaps\n",
    "fig = plotFracHeatmap(df_N1, \"N1\", sequences_of_interest, vmax=0.4)\n",
    "save_plt(fig, output_dir_str + 'DHAD_N1_heatmap_frac_dist.png')\n",
    "\n",
    "fig = plotFracHeatmap(df_N423, \"N423\", sequences_of_interest, vmax=0.4)\n",
    "save_plt(fig, output_dir_str + 'DHAD_N423_heatmap_frac_dist.png')\n",
    "\n",
    "fig = plotFracHeatmap(df_N560, \"N560\", sequences_of_interest, vmax=0.4)\n",
    "save_plt(fig, output_dir_str + 'DHAD_N560_heatmap_frac_dist.png')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Repeat the analysis for CYP2\n",
    "\n",
    "Here we want to run the same analysis as above, but on the CYP2 dataset. Here rather than randomly adding sequences\n",
    "we're looking at adding families of proteins. \n",
    "\n",
    "We have the CYP2U1, CYP2R, and CYP2D.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (1 - 3) Load the trees and calculate the matching nodes to our reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tree_object import TreeObject\n",
    "from tree_controller import TreeController\n",
    "\n",
    "input_dir_str = 'Files/PIA_1_Files/GRASP_recons/' + cyp2_str + '/'\n",
    "output_dir_str = 'Files/PIA_1_Files/output/'\n",
    "\n",
    "trees = []\n",
    "files = os.listdir(input_dir_str)\n",
    "for f in files:\n",
    "    if '.nwk' in f:\n",
    "        trees.append(f)\n",
    "\n",
    "print(trees)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (4) Find the matching nodes in the larger trees for our chosen nodes of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_tree = '20191216_CYP2U_Paper_Version_reconstructed-tree_GRASP.nwk'\n",
    "nodes = ['N0', 'N1', 'N4']\n",
    "node_name_dict = {'N0': {'seqs': []}, 'N1': {'seqs': []}, 'N4': {'seqs': []}}\n",
    "\n",
    "nodes = ['N0', 'N1', 'N4']\n",
    "node_matching_dict = {}\n",
    "for t in trees:\n",
    "    t1 = TreeObject(input_dir_str + base_tree) \n",
    "    t2 = TreeObject(input_dir_str + t)\n",
    "    tc = TreeController()\n",
    "    t_num = t.split('_')[1]\n",
    "    node_matching = tc.get_similar_nodes_between_trees(t1, t2, True)\n",
    "    node_matching_dict[t] = node_matching\n",
    "    \n",
    "    for label_original, label_match in node_matching.items():\n",
    "        if label_original == 'N0':\n",
    "            print(label_original, label_match.get_label())\n",
    "            \n",
    "    for n in nodes:\n",
    "        node_name_dict[n][t_num] = node_matching_dict[t][n].get_label()\n",
    "        print(n, node_matching_dict[t][n].get_label())\n",
    "    print(\"Complete: \" + t)\n",
    "print(node_name_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (5) Load in the alignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cyp2_files = []\n",
    "sequences_of_interest = []\n",
    "\n",
    "files = os.listdir(input_dir_str)\n",
    "for f in files:\n",
    "    if '.fasta' in f:\n",
    "        cyp2_files.append(f)\n",
    "        # Use the 2U-2R-2D as the order\n",
    "        sequences_of_interest.append(f.split('_')[1])\n",
    "        \n",
    "print(cyp2_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (6) Pick out the consensus sequences from each alignment for our nodes of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Again build up our node_name_dict using our sequences of interest\n",
    "\n",
    "\n",
    "for t, tree in node_matching_dict.items():\n",
    "    for n in nodes:\n",
    "        try:\n",
    "            \"\"\"\n",
    "            Here we're using t_num as our method for matching our trees and our alignments, this\n",
    "            will be specific to each dataset. Our trees are: 20191216_CYP2U-CYP2R_Paper_Version_reconstructed-tree_GRASP.nwk\n",
    "            So by splitting on _ and taking the 1st element, we get 2 for the tree above, \n",
    "            20191216_CYP2U-CYP2R_Paper_Version_reconstructed-tree_GRASP.nwk = CYP2U-CYP2R etc.        \n",
    "            \"\"\"\n",
    "            t_num = t.split('_')[1]\n",
    "            node_name_dict[n][t_num] = tree.get(n).original_label\n",
    "            print(\"Matched: \\t\" + n + \"\\t -> \\t\" + tree.get(n).get_label())\n",
    "        except Exception as e:\n",
    "            print(n, t, e)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Here we want to iterate through our alignment files and pick out the consensus\n",
    "sequences that match with our base/reference tree nodes of interest\n",
    "\"\"\"\n",
    "for f in cyp2_files:\n",
    "    lengths = []\n",
    "    num_seqs = []\n",
    "    seqs_aligned = sequence.readFastaFile(input_dir_str + f, sequence.Protein_Alphabet, gappy=True)\n",
    "    seq_dict = {}\n",
    "    for seq in seqs_aligned:\n",
    "        seq_dict[seq.name] = seq\n",
    "    i = 0\n",
    "    \"\"\"\n",
    "    UNIQUE TO EACH DATASET:\n",
    "        \n",
    "    Here we want a string to match between our reconstructed trees and our alignments from GRASP\n",
    "    For here our trees are annotated as: 20191216_CYP2U-CYP2R_Paper_Version_reconstructed-tree_GRASP.nwk\n",
    "    and our alignments as: 20191216_CYP2U-CYP2R-CYP2D_Paper_Version_joint-ancestors_GRASP.fasta \n",
    "    Thus we want to match the [2] from the tree with the [2] from the alignment\n",
    "\n",
    "    \"\"\"\n",
    "    num_seqs = f.split('_')[1]\n",
    "    \n",
    "    # Check if this is a sequence of interest\n",
    "    for node_name, node_dict in node_name_dict.items():\n",
    "        for seq_nums, matching_node_name in node_dict.items():\n",
    "            if seq_nums == num_seqs:\n",
    "                s = seq_dict[matching_node_name]\n",
    "                no_gap = ''\n",
    "                for c in s:\n",
    "                    if c != '-':\n",
    "                        no_gap += c\n",
    "                seq = sequence.Sequence(no_gap, sequence.Protein_Alphabet, name=seq_nums + '_' + s.name)\n",
    "                node_dict['seqs'].append(seq)\n",
    "                \n",
    "    print(\"Finished: \\t\" + f + \" \\t Matched: \" + str(len(node_dict['seqs'])) + \" nodes\")\n",
    "    i += 1\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (7 - 8) Calculate the fractional distances\n",
    "\n",
    "Now we have the sequences for each of the nodes, we want to:\n",
    "```\n",
    "1. Remove all the gaps\n",
    "2. Create an alignment of the sequences\n",
    "3. Calculate the fractional distances from these\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write out which sequecnes were interested in (this is really the alignment names, so number of sequences)\n",
    "print(sequences_of_interest, node_name_dict)\n",
    "write_heatmap(output_dir_str + 'CYP2_heatmap.csv', sequences_of_interest, node_name_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate the size and variation of the sequences\n",
    "\n",
    "For each node of interest calculate the size and differences for each of these\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variances = []\n",
    "medians = []\n",
    "means = []\n",
    "corelations = []\n",
    "\n",
    "for node_name, node in node_name_dict.items():\n",
    "    seqs = node['seqs']\n",
    "    seq_lens = []\n",
    "    # Now for each sequence we want to get the length (no gaps)\n",
    "    for seq in seqs:\n",
    "        seq_len = 0\n",
    "        for i in seq:\n",
    "            if i != '-':\n",
    "                seq_len += 1\n",
    "        seq_lens.append(seq_len)\n",
    "    print(node_name, np.var(seq_lens), np.mean(seq_lens), np.median(seq_lens))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (9) Plot heatmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "csv_file_name = output_dir_str + 'CYP2_heatmap.csv'\n",
    "df = pd.read_csv(csv_file_name)\n",
    "fact_names = []\n",
    "for n in sequences_of_interest:\n",
    "    fact_names.append(str(n))\n",
    "\n",
    "print(fact_names)\n",
    "df_N0 = df[df['aln_name'] == 'N0']\n",
    "df_N1 = df[df['aln_name'] == 'N1']\n",
    "df_N4 = df[df['aln_name'] == 'N4']\n",
    "\n",
    "\n",
    "fig = plotFracHeatmap(df_N0, \"N0\", sequences_of_interest, vmax=0.6)\n",
    "save_plt(fig, output_dir_str + 'CYP2_N0_heatmap_frac_dist.png')\n",
    "\n",
    "fig = plotFracHeatmap(df_N1, \"N1\", sequences_of_interest, vmax=0.6)\n",
    "save_plt(fig, output_dir_str + 'CYP2_N1_heatmap_frac_dist.png')\n",
    "\n",
    "fig = plotFracHeatmap(df_N4, \"N4\", sequences_of_interest, vmax=0.6)\n",
    "save_plt(fig, output_dir_str + 'CYP2_N4_heatmap_frac_dist.png')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}