visualize_data.py

"""
Visualize distribution of generated batches
"""
from __future__ import division
from argparse import ArgumentParser
from os.path import join
from glob import glob
from itertools import islice

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from data import NUM_BASES, generate_kmers, levenshtein, get_nucleotide_code, \
                 generate_triplets, generate_tuplets, generate_tuplets_fair, \
                 generate_triplets_fair

COLOR = 'blue'

################################################################################
#              visualize distribution of k-mers in a batch
################################################################################

def visualize_kmer_distribution(kmer_length, samples, generator, skip=5):
    """
    Visualize distribution of k-mers.
    :param kmer_length: length of the k-mers to generate
    :param samples: amount of triplets to generate
    :pram gen: generator to use for k-mer creation
    :param skip: amount of labels to skip before next label is displayed
    """
    # count how often each k-mer occurs
    y_val = [0]*NUM_BASES**kmer_length

    # convert k-mer to decimal number
    k = kmer_length-1
    to_decimal = lambda a: np.sum([c*NUM_BASES**(k-i) for i, c in enumerate(a)])

    # count in kmers list how often each kmer occurs
    for i, tri in enumerate(generator):
        if i == samples:
            break
        for kmer in tri:
            if isinstance(kmer, tuple):
                y_val[to_decimal(kmer)] += 1

    # labels for each k-mer
    x_val = list(map(get_nucleotide_code, generate_kmers(kmer_length)))

    # show results as bar plot
    fig = plt.figure()
    title = 'Frequency of each k-mer with length {0} | samples: {1}'
    # helper
    ind = np.arange(len(x_val))
    ecl = ['black']*len(x_val)
    lwidth = [0.5]*len(x_val)

    # create bar plot
    axe = fig.add_subplot(111)
    axe.set_xlabel('k-mer')
    axe.set_ylabel('frequency')
    axe.bar(ind, y_val, width=0.9, edgecolor=ecl, linewidth=lwidth,
            color=COLOR, align='center')
    axe.set_xticks(ind)
    axe.set_xticklabels(['' if j%skip else x for j, x in enumerate(x_val)])
    axe.set_title(title.format(kmer_length, np.sum(y_val)))
    fig.autofmt_xdate(bottom=0.2, rotation=90, ha='center')
    plt.show()


def visualize_edit_distances_tuplets(kmer_length, samples):
    """
    Count how often each edit distance occurs for a given k-mer-length and
    how often a pair with each edit distance is generated by the tuplets
    generator.
    :param samples: amount of tuplets to generate
    """
    # actual distribution
    y_val_1 = [0]*kmer_length
    for i, k_1 in enumerate(generate_kmers(kmer_length)):
        kmers = generate_kmers(kmer_length)
        # skip pairs we already calculated
        # e.g (0,0,0), (1,0,0) is the same as (1,0,0), (0,0,0)
        next(islice(kmers, i+1, i+1), None)
        for k_2 in kmers:
            edit_dis = levenshtein(k_1, k_2)
            y_val_1[edit_dis-1] += 1
    print(y_val_1)

    # distribution generated by tuplet generator
    y_val_2 = [0]*kmer_length
    for i, tuplet in enumerate(generate_tuplets(kmer_length, 1, False)):
        if i == samples:
            break
        y_val_2[tuplet[-1]-1] += 1
    print(y_val_2)

    # distribution generated by tuplet generator
    y_val_3 = [0]*kmer_length
    for i, tuplet in enumerate(generate_tuplets_fair(kmer_length, 1, False)):
        if i == samples:
            break
        y_val_3[tuplet[-1]-1] += 1
    print(y_val_3)

    # show results as bar plot
    fig = plt.figure()
    titles = ['real distribution of pairs', 'generated distribution of pairs',\
              'fair generated distribution of pairs']
    for i, y_val in enumerate([y_val_1, y_val_2, y_val_3]):
        # labels for each k-mer
        x_val = list(range(1, kmer_length+1))
        # helper
        ind = np.arange(len(x_val))
        ecl = ['black']*len(x_val)
        lwidth = [0.5]*len(x_val)

        # create bar plot
        axe = fig.add_subplot(int('1{0}{1}'.format(len(titles), i+1)))
        axe.set_xlabel('edit distance')
        axe.set_ylabel('frequency')
        axe.bar(ind, y_val, width=0.9, edgecolor=ecl, linewidth=lwidth,
                color=COLOR, align='center')
        axe.set_xticks(ind)
        axe.set_xticklabels(x_val)
        axe.set_title(titles[i])
    plt.show()


def visualize_edit_distances_triplets(kmer_length, samples):
    """
    """
    def flatten(arr):
        res = []
        for i in y:
            for j in i:
                res.append(j)
        return res

    x_val = [str((i, j)) for i in range(1, kmer_length)
             for j in range(i+1, kmer_length+1)]
    print(x_val)

    # real distribution
    y = np.array([np.zeros(i) for i in range(kmer_length-1, 0, -1)])
    for i, k_1 in enumerate(generate_kmers(kmer_length, False)):
        kmers = generate_kmers(kmer_length, False)
        next(islice(kmers, i+1, i+1), None)
        for j, k_2 in enumerate(kmers):
            kmers2 = generate_kmers(kmer_length, False)
            next(islice(kmers2, i+j+2, i+2), None)
            for _, k_3 in enumerate(kmers2):
                l_1 = levenshtein(k_1, k_2)
                l_2 = levenshtein(k_1, k_3)
                if l_1 == l_2:
                    continue
                if l_1 > l_2:
                    l_2, l_1 = l_1, l_2
                y[l_1-1][l_2-(l_1+1)] += 1

    # flatten
    y_val_1 = flatten(y)
    print(y_val_1)

    # real unfair distribution
    y = np.array([np.zeros(i) for i in range(kmer_length-1, 0, -1)])
    for i, triplet in enumerate(generate_triplets(kmer_length, 1, False)):
        if i == samples:
            break
        x, x_pos, x_neg = triplet
        l_1, l_2 = levenshtein(x, x_pos), levenshtein(x, x_neg)
        y[l_1-1][l_2-(l_1+1)] += 1

    y_val_2 = flatten(y)
    print(y_val_2)

    # real fair distribution
    y = np.array([np.zeros(i) for i in range(kmer_length-1, 0, -1)])
    for i, triplet in enumerate(generate_triplets_fair(kmer_length, 1, False)):
        if i == samples:
            break
        x, x_pos, x_neg = triplet
        l_1, l_2 = levenshtein(x, x_pos), levenshtein(x, x_neg)
        y[l_1-1][l_2-(l_1+1)] += 1

    y_val_3 = flatten(y)
    print(y_val_3)

    # show results as bar plot
    fig = plt.figure()
    titles = ['real distribution of pairs', 'generated distribution of pairs',\
              'fair generated distribution of pairs']
    for i, y_val in enumerate([y_val_1, y_val_2, y_val_3]):
        # helper
        ind = np.arange(len(x_val))
        ecl = ['black']*len(x_val)
        lwidth = [0.5]*len(x_val)

        # create bar plot
        axe = fig.add_subplot(int('1{0}{1}'.format(len(titles), i+1)))
        axe.set_xlabel('edit distance')
        axe.set_ylabel('frequency')
        axe.bar(ind, y_val, width=0.9, edgecolor=ecl, linewidth=lwidth,
                color=COLOR, align='center')
        axe.set_xticks(ind)
        axe.set_xticklabels(x_val)
        axe.set_title(titles[i])
    plt.show()

################################################################################
#                           main function
################################################################################

if __name__ == '__main__':
    parser = ArgumentParser()
    # general
    parser.add_argument('-k', '--kmer_length',
                        help="length of the k-mers",
                        default=3,
                        type=int)
    plot_types = ['batch_distribution', 'edit_distances_tuplets', \
                  'edit_distances_triplets']
    parser.add_argument('-p', '--plot_type',
                        help="Type of plot to generate",
                        default=plot_types[0],
                        choices=plot_types)
    parser.add_argument('-s', '--samples',
                        help="amount of samples to generate",
                        default=1000000,
                        type=int)


    args = parser.parse_args()

    # create generator
    if args.plot_type == plot_types[0]:
        gen = generate_triplets(args.kmer_length, encode=False)
        visualize_kmer_distribution(args.kmer_length, args.samples, gen)
    elif args.plot_type == plot_types[1]:
        visualize_edit_distances_tuplets(args.kmer_length, args.samples)
    elif args.plot_type == plot_types[2]:
        visualize_edit_distances_triplets(args.kmer_length, args.samples)