similarity

#!/usr/bin/python

import os
import sys
import uuid
import argparse
import numpy as np
import pandas as pd
import subprocess as sub
import matplotlib.pyplot as plt


#########################################
# Arguments
#########################################

def parse_args():
    parser = argparse.ArgumentParser(description='Compute ETA2 similarity between volumes.')

    # input
    parser.add_argument('--volumes', '-v', metavar='file',
                           type=argparse.FileType('r'), help='Input volume(s)', nargs='+',
                           required=True)
    parser.add_argument('--names', '-n', metavar='name', 
                           help='volume name(s); match order of -v', nargs='+',
                           required=True)

    # outputs
    parser.add_argument('--csv', metavar='filename', type=argparse.FileType('w'),
                           help='output csv file of matrix; otherwise stdout', required=False)
    parser.add_argument('--heatmap', metavar='matrix_image.png', type=argparse.FileType('w'),
                           help='output matrix heatmap of eta2 values', required=False)

    # options
    parser.add_argument('--mask', metavar='file',
                           type=argparse.FileType('r'),
                           help='Mask volume (highly suggested!)', required=False)

    # metric
    parser.add_argument('--metric', '-m', metavar='eta2/pearson',
                           choices=['eta2','pearson'], default='eta2',
                           help='Metric of similarity (only use pearson on 4d files)', required=False)

    args = parser.parse_args()

    if len(args.names) > 0 and len(args.names) != len(args.volumes):
        print("Error: when implementing names, the number of names must match the number of input volumes")
        sys.exit()

    if len(args.volumes) < 2:
        print("Error: you need at least 2 volumes to make a comparision, silly.")
        sys.exit()

    return args

#########################################
# Heatmap
#########################################

# creates square heatmap (row/col labels are same)
# * current setup highlights positive-only relationships
def heatmap(matrix, labels=None, limits=[0,1], cm=plt.cm.YlGn_r):
    # limits
    mn,mx = limits

    # Plot it out
    fig, ax = plt.subplots()
    heatmap = ax.pcolor(matrix, cmap=cm, alpha=0.95, vmin=mn, vmax=mx)

    # Format
    fig = plt.gcf()
    fig.set_size_inches(12, 10)

    # turn off the frame
    ax.set_frame_on(False)

    # put the major ticks at the middle of each cell
    ax.set_yticks(np.arange(matrix.shape[0]) + 0.5, minor=False)
    ax.set_xticks(np.arange(matrix.shape[1]) + 0.5, minor=False)

    # want a more natural, table-like display
    ax.invert_yaxis()
    ax.xaxis.tick_top()

    # Set the labels
    if labels is None: labels = matrix.index

    # note I could have used nba_sort.columns but made "labels" instead
    ax.set_xticklabels(labels, minor=False)
    ax.set_yticklabels(labels, minor=False)

    # rotate the
    plt.xticks(rotation=90)

    ax.grid(False)

    # Turn off all the ticks
    ax = plt.gca()

    # insert color bar
    plt.colorbar(heatmap)

    for t in ax.xaxis.get_major_ticks():
        t.tick1On = False
        t.tick2On = False
    for t in ax.yaxis.get_major_ticks():
        t.tick1On = False
        t.tick2On = False

    return fig


#########################################
# Similarity & Clustering
#########################################

# returns symmetric matrix representing the
# eta2 values between each variable-pairs.
def eta2(volumes,func=float,mask=None):
    maskstr = '-mask {}'.format(mask.name) if mask is not None else ''
    cmdstr  = '3ddot -full {} -doeta2 {}'.format(maskstr, ' '.join(volumes))

    p = sub.Popen(cmdstr, shell=True, stdout=sub.PIPE)
    out, err = p.communicate()

    if err:
        print "Error processing volumes:"
        print err
        sys.exit(1)

    values = [[func(y) for y in x.split()]
                        for x in out.split('\n')
                        if  len(x) > 1]
    return values


def tmpvol():
    """ returns prefix path for temp volume file """
    return os.path.join('/tmp',str(uuid.uuid4()))


def pearson(volumes, func=float, mask=None):
    corr={}
    for v in volumes:
        for u in volumes:
            if v not in corr:
                corr[v] = {}
            if u not in corr:
                corr[u] = {}
            if v == u:
                corr[v][u] = 1
                continue
            if v in corr and u in corr[v]: continue

            print "correlating {} with {}".format(v,u)

            vucorr=tmpvol()
            # compute corr
            cmdstr='3dTcorrelate -prefix {} {} {}'.format(vucorr, v, u)
            p = sub.Popen(cmdstr, shell=True, stdout=sub.PIPE)
            out,err = p.communicate()

            # compute mean corr
            maskstr = ' -mask {}'.format(mask.name) if mask is not None else ''
            cmdstr='3dmaskave {} {}'.format(maskstr, vucorr+'+????.BRIK')
            p = sub.Popen(cmdstr, shell=True, stdout=sub.PIPE)
            out,err = p.communicate()

            if err:
                print "Error processing volumes:"
                print err
                sys.exit(1)

            vucorrval = func(out.split()[0])

            # store correlation
            corr[v][u] = vucorrval
            corr[u][v] = vucorrval

    # return matrix
    return [[corr[v][u] for v in volumes] 
                        for u in volumes]


#########################################
# Main Application
#########################################

if __name__ == '__main__':
    args = parse_args()

    if args.metric == 'eta2':
        matrix = eta2([x.name for x in args.volumes], func=float, mask=args.mask)
    elif args.metric == 'pearson':
        matrix = pearson([x.name for x in args.volumes], func=float, mask=args.mask)

    df = pd.DataFrame.from_dict({n:vals for n,vals in zip(args.names,matrix)})
    df.index = args.names

    # output csv format
    if args.csv:
        df.to_csv(args.csv)
        print "CSV file saved to {}".format(args.csv.name)
    else:
        print " showing {} values: ".format(args.metric)
        print "-" * 30
        print df
        print "-" * 30

    # heatmap
    if args.heatmap:
        fig = heatmap(np.array(matrix), labels=args.names)
        plt.tight_layout()
        plt.savefig(args.heatmap)
        print "heatmap saved to {}".format(args.heatmap)