plot_dist_graph.py

from matplotlib import pyplot as plt
import numpy as np

def dist_sort(dists, ext=None, done=None, verbose=False):
    """
    sort keys of an input dictionary of dictionaries (half matrix of distances) 
    put closest together
    
    :param None ext: used by the recursion, first and last value of ordered list
    :param None done: used by the recursion, ordered list of keys
    :param False verbose: prints pair of values found to be close at each recurtion
       and the order list (done).
    
    :returns: list of ordered keys
    """
    if not done:
        done = list()
    a, b = max([(dists[k][l], k, l) for k in dists for l in dists[k] 
                if (k != l) and (not ext or (k in ext or l in ext) and not (k in done and l in done))])[1:]
    if not ext:
        new, old = done = [a, b]
    else:
        new, old = (a, b) if b in done else (b, a)
        done.insert(bool(done.index(old)) * len(done), new)
    if verbose:
        print new, old, done
    if len(done) == len(dists):
        return done
    return dist_sort(dists, set([done[0], done[-1]]), done, verbose)


def awesome_plot(pw_dists, keys=None, annot=2, exclusive=None, verbose=False, 
                 scale_width=5, text_size=8, pwr=1):
    """
    Plots a circular graph of distances between elements given as a dictionary of 
    dictionaries (half matrix of distances)
    
    :param None keys: list of keys to be used as order. If None, dist_sort function 
       will be used.
    :param 2 annot: number of time per edge, the distance is written
    :param None exclusive: number of excusive elements by category (dictionnary).
    :param 5 scale_width: modifier for line widths and areas
    :param 1 pwr: to control the transformation from values to widths
    :param False verbose: prints pair of values found to be close at each recurtion
       and the order list (done).    
    """
    mod = max([pw_dists[k][l] for k in pw_dists for l in pw_dists[k]] + 
              (exclusive.values() if exclusive else []))
    mod *= float(scale_width)
    if not keys:
        keys = dist_sort(pw_dists, verbose=verbose)
    if not exclusive:
        exclusive = {}
    frac = 2 * np.pi / len(pw_dists)
    coords = {}
    for i, k in enumerate(keys):
        coords[k] = np.cos(frac * i + np.pi / 2), np.sin(frac * i + np.pi / 2)
    patches = []
    cnt = 0
    # draw exclusive elements
    dists = []
    for i, k in enumerate(exclusive):
        dist = (0.98 + (((pw_dists[k].get(k, mod * 0.15) / mod)/np.pi)**.5)**pwr + 
                ((exclusive[k] / mod / np.pi)**.5)**pwr)
        dists.append(dist + ((exclusive[k] / mod / np.pi)**.5)**pwr + 0.02)
        c2 = plt.Circle((coords[k][0] * dist, coords[k][1] * dist),
                        (((exclusive[k] / mod)/np.pi)**.5)**pwr, color='lightgrey', ec='lightgrey')
        plt.gcf().gca().add_artist(c2)
        val = (int(exclusive[k]) if float(exclusive[k]).is_integer() 
               else round(exclusive[k], 2))
        plt.text(coords[k][0] * dist, coords[k][1] * dist, val, size=text_size ,
                 verticalalignment='center'   , horizontalalignment='center')
    plt.xlim((-max(dists), max(dists)))
    plt.ylim((-max(dists), max(dists)))
    plt.axis('off')
    # draw edges
    count = 0
    for i, k in enumerate(pw_dists):
        for j, l in enumerate(pw_dists[k]):
            if k == l:
                continue
            # size of each side of the triangle formed by the two elements and 
            # their projections
            a = abs(coords[k][1] - coords[l][1])
            c = abs(coords[k][0] - coords[l][0])
            b = (c**2 + a**2)**.5
            angle = np.rad2deg(np.arccos((a**2 + b**2 - c**2) / (2 * a * b)))
            # correct angles for rectangle direction
            if coords[k][0] < coords[l][0] and coords[k][1] < coords[l][1]:
                angle = -angle
            elif coords[k][0] < coords[l][0] and coords[k][1] > coords[l][1]:
                angle = angle - 180
            elif coords[k][0] > coords[l][0] and coords[k][1] > coords[l][1]:
                angle = -angle - 180
            elif coords[k][0] < coords[l][0] and coords[k][1] > coords[l][1]:
                angle = -angle - 180
            d = (pw_dists[k][l] / mod)**pwr
            # correctors in x and y because rectangle are pin at one corner
            x = np.cos(np.deg2rad(angle)) * d / 2
            y = np.sin(np.deg2rad(angle)) * d / 2
            r = plt.Rectangle((coords[k][0] - x, coords[k][1] - y), d, b, 
                              angle=angle, color='lightgrey')
            cnt+=1
            plt.gcf().gca().add_artist(r)
            if annot == 1:
                if len(keys)%2:
                    divs = [2.]
                else:
                    divs = [2.25 - 0.5 * (count % 2)]
            elif annot == 2:
                divs = [4, 1.33]
            else:
                divs = []
            for div in divs:
                val = (int(pw_dists[k][l]) 
                       if float(pw_dists[k][l]).is_integer() 
                       else round(pw_dists[k][l], 2))
                plt.text(coords[k][0] + (coords[l][0] - coords[k][0]) / div, 
                         coords[k][1] + (coords[l][1] - coords[k][1]) / div,
                         val, verticalalignment='center',
                         horizontalalignment='center', size=text_size)
            count += 1
    # draw total elements and names
    for k in keys:
        c1 = plt.Circle(coords[k], (((pw_dists[k].get(k, mod * 0.15) / mod)/np.pi)**.5)**pwr, color='w', ec='lightgrey')
        plt.gcf().gca().add_artist(c1)
        plt.text(coords[k][0], coords[k][1], k, size=text_size+2,
                 verticalalignment='bottom'   , horizontalalignment='center')
        if annot:
            val = (int(pw_dists[k].get(k, float('NAN')))
                   if float(pw_dists[k].get(k, float('NAN'))).is_integer()
                   else round(pw_dists[k].get(k, float('NAN')), 2))
            plt.text(coords[k][0], coords[k][1], '(' + str(val) + ')', 
                     size=text_size,
                     verticalalignment='top', horizontalalignment='center')
    plt.show()

from random import random

aa = {}
size = 6
for i in range(size):
    aa.setdefault(str(i), {})
    for j in range(i, size):
        aa[str(i)][str(j)] = random()

plt.figure(figsize=(10, 10))
awesome_plot(aa, keys=None, annot=1, exclusive=dict([(k, random()) for k in aa]), 
             text_size=9, scale_width=5)