cooc.py

#!/usr/bin/env python
# by @mazieres
from itertools import permutations
from collections import defaultdict

import pandas as pd
import numpy as np
from matplotlib import pyplot as plt


def mat_cooc(d):
    res = defaultdict(lambda: defaultdict(int))
    for k, v in d.iteritems():
        for a, b in permutations(v, 2):
            res[a][b] += 1
    df = pd.DataFrame(res)
    return df.apply(lambda x: x/df.sum(0))

def draw_mat(df, title='Cooc', cmap=plt.cm.Blues):
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    plt.yticks(np.arange(len(df.index)), df.index)
    plt.xticks(np.arange(len(df.columns)), df.columns, rotation=90)
    plt.tick_params(axis='both', labelsize=15)
    plt.xlabel('', fontsize=15)
    plt.ylabel('')