visualize_model.py

"""
Visualize metric data from neural network.
"""
from __future__ import division
from math import ceil, floor, log10
from itertools import islice
from argparse import ArgumentParser

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button
from matplotlib.text import Annotation
from mpl_toolkits.mplot3d import Axes3D

from network.tower import TowerNetwork
from data import vec, NUM_BASES, generate_kmers, levenshtein,\
                 get_nucleotide_code

COLOR = 'blue'
SCATTER = 'scatter'
DISTRIBUTION = 'distribution'

################################################################################
#                  visualize embedding (scatter plot)
################################################################################

def plot_graph(model_file, dim):
    """
    Show plot from neural network model.
    :param dim: 1,2 or 3(D)
    :param model_file: path to model of neural network.
    """
    if dim not in [1, 2, 3]:
        raise Exception('Illegal dimension. Must be 1, 2 or 3.')

    # restore neural network
    sess, net = TowerNetwork.load(model_file)

    batch_size = 180
    num_kmers = NUM_BASES**net.kmer_length

    # 0: x values, 1: y values, 2: z values
    points = [[] for _ in range(dim)]
    if dim == 1:
        points.append([0]*num_kmers)

    # generate all possible kmers
    kmers = generate_kmers(net.kmer_length)

    for _ in range(int(ceil(num_kmers/batch_size))):
        # calculate embbeding
        batch = [vec(kmer) for kmer in islice(kmers, batch_size)]
        emb = sess.run(net.out, feed_dict={net.x: batch, net.keep_prob: 1.0})

        for i in range(dim):
            points[i].extend(emb[:, i])

    # plot datapoints
    title = 'Embedding of all possible {0}-mers'
    fig = plt.figure()
    fig.suptitle(title.format(net.kmer_length))
    if dim == 3:
        axe = fig.add_subplot(111, projection='3d')
    else:
        axe = fig.add_subplot(111)
    axe.scatter(*points, c=COLOR, marker='o', picker=True)


    def on_pick(event):
        """
        Show label for selected point.
        :param event: click event
        """
        # indices of selected dots (multiple, if they are close to each other)
        i = int(event.ind[0])
        # skip to right k-mer and create a readable label
        kmers = generate_kmers(net.kmer_length)
        next(islice(kmers, i, i), None)
        label = get_nucleotide_code(next(kmers))
        # show label
        pos = (event.mouseevent.xdata, event.mouseevent.ydata)
        axe.add_artist(Annotation(label, xy=pos, xycoords='data'))
        axe.figure.canvas.draw_idle()

    def on_clear_all(event):
        """
        Clear all labels.
        :param event: click event
        """
        axe.cla()
        axe.scatter(*points, c=COLOR, marker='o', picker=True)
        axe.figure.canvas.draw_idle()

    # add click handler
    fig.canvas.mpl_connect('pick_event', on_pick)

    # add clear button
    axe_clear_all = plt.axes([0.0, 0.0, 0.1, 0.05])
    bt_clear_all = Button(axe_clear_all, 'Clear all')
    bt_clear_all.on_clicked(on_clear_all)

    # show plot
    plt.show()


################################################################################
#             evaluate learned distance measure (bar plot)
################################################################################

def visualize_distribution(model_file, num_cols=2):
    """
    Visualize average distance between two kmers with edit distance specific
    edit distance.
    :param model_file: path to model of neural network
    :param num_cols: number of columns for the created figure
    """
    # restore neural network
    sess, net = TowerNetwork.load(model_file)

    # save embbeding points
    batch_size = 180

    # index i: all embedded distances for kmers with edit distance i+1
    values = [[] for _ in range(net.kmer_length)]

    # pairs: (NUM_BASES^kmer_length+2-1) choose 2 - NUM_BASES^kmer_length
    for i, k in enumerate(generate_kmers(net.kmer_length)):
        # embbeding for reference kmer
        emb = sess.run(net.out, feed_dict={net.x: [vec(k)], net.keep_prob: 1.0})
        num_kmers = NUM_BASES**net.kmer_length

        kmers = generate_kmers(net.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 _ in range(int(ceil(num_kmers/batch_size))):
            batch = list(islice(kmers, batch_size))
            if len(batch) == 0:
                continue
            input_x = list(map(vec, batch))
            emb_batch = sess.run(net.out, feed_dict={net.x: input_x,
                                                     net.keep_prob: 1.0})
            for j, k_2 in enumerate(batch):
                dis = levenshtein(k, k_2)
                # calculate euclidean distance for data points
                values[dis-1].append(np.linalg.norm(emb-emb_batch[j]))

    # show results as bar plot
    fig = plt.figure()
    title = 'Embbeded distances for {0}-mer pairs'
    fig.suptitle(title.format(net.kmer_length))

    step_size = 0.1
    # x: number to round, p: step to round to nearest
    nearest = lambda x, p: round(round(x / p) * p, -int(floor(log10(p))))
    nearest_num = lambda x: nearest(x, step_size)
    # xticks values
    x_max = np.amax([np.amax(val) for val in values])
    x_val = np.arange(0.0, nearest_num(x_max)+step_size, step_size)

    # configure sub plots
    ind = np.arange(len(x_val))
    ecl = ['black']*len(x_val)
    lwidth = [0.5]*len(x_val)
    num_rows = int(ceil(net.kmer_length/num_cols))
    title_str = 'edit distance: {0:d} median: {1:.2f} avg: {2:.2f} std: {3:.2f}'

    # generate subplots
    for i, val in enumerate(values):
        # mean and standard deviation
        mean = np.median(val)
        avg = np.average(val)
        std = np.std(val)

        # round values to easily count them
        val = list(map(nearest_num, val))
        y_val = [val.count(x) for x in x_val]

        # create bar plot
        axe = fig.add_subplot(int('{0}{1}{2}'.format(num_rows, num_cols, i+1)))
        axe.set_xlabel('embedded distance')
        axe.set_ylabel('amount of k-mers')
        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%2 else x for j, x in enumerate(x_val)])
        axe.set_title(title_str.format(i+1, mean, avg, std))

    # show plot
    plt.show()


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

if __name__ == '__main__':
    parser = ArgumentParser()
    # scatter only
    parser.add_argument('-d', '--dimension',
                        help="dimension of the scatter plot as integer"\
                        " 1, 2, 3(D) [scatter only]",
                        choices=[1, 2, 3], type=int, default=3)
    # distribution only
    parser.add_argument('-c', '--columns',
                        help="number of aligned subplots per row"\
                        " [distribution only]",
                        type=int, default=2)
    # general
    parser.add_argument('-p', '--plot_type',
                        help="type of plot to be drawn",
                        choices=[SCATTER, DISTRIBUTION],
                        default='scatter',
                        type=str)
    parser.add_argument('model_file',
                        help="path to the trained neural network model",
                        type=str)

    args = parser.parse_args()
    if args.plot_type == SCATTER:
        plot_graph(args.model_file, args.dimension)
    elif args.plot_type == DISTRIBUTION:
        visualize_distribution(args.model_file, args.columns)