analyze_clusters

#!/usr/bin/env python2.7
from sklearn import metrics
from functools import partial
from multiprocessing import Pool
import itertools as it
import argparse
import json
import sys

def mk_vectors(A, B):
# make A and B into maps.
  Aks = set()
  Bks = set()
  for (elem,q) in A:
    Aks.add(elem)
  for (elem,q) in B:
    Bks.add(elem)

  A_B = Aks.intersection(Bks)

  Amap = {}
  for (elem,label) in A:
    if elem in A_B:
      Amap[elem] = label
  Bmap = {}
  for (elem,label) in B:
    if elem in A_B:
      Bmap[elem] = label

# Put the labels into the vectors with the same
# ordering. Now, the ith entry in Afvec/Bfvec corresponds
# to a particular element in A/B.
  assert frozenset(Amap.keys()) == frozenset(Bmap.keys())
  Afvec = []
  Bfvec = []
  for i in Amap.keys():
    Afvec.append(Amap[i])
    Bfvec.append(Bmap[i])

  print len(Afvec)
  print len(Bfvec)
  return (Afvec,Bfvec)

def mk_sets(A, B):
  Amap = {}
  for (elem,label) in A:
    Amap[label] = Amap.get(label, [])
    Amap[label].append(elem)
  Bmap = {}
  for (elem,label) in B:
    Bmap[label] = Bmap.get(label, [])
    Bmap[label].append(elem)

  ASet = frozenset([frozenset(i) for i in Amap.values()])
  BSet = frozenset([frozenset(i) for i in Bmap.values()])

  return (ASet,BSet)

def f_calculate(param, t, g):
  a,b = param
  tp = fp = fn = 0
  truth_a = t[a]
  truth_b = t[b]
  guess_a = g[a]
  guess_b = g[b]

  if truth_a==truth_b:
    if guess_a==guess_b:
      tp += 1
    else:
      fn += 1
  elif guess_a==guess_b:
    fp += 1
  # True negatives not used
  return (tp, fp, fn)

def f_measure(truth, guess, beta=1.0):
  t_map = create_mapping(truth)
  g_map = create_mapping(guess)

  g_collapsed = set().union(*guess) # g_collapsed = guess[0] U guess[1] U ... U guess[n]
 
  p = Pool()
  mapfun = partial(f_calculate, t=t_map, g=g_map)

  parresults = p.map(mapfun, it.combinations(g_collapsed, 2))
  tp,fp,fn = reduce(lambda x,y:(x[0]+y[0],x[1]+y[1],x[2]+y[2]), parresults)

  precision = float(tp)/(tp+fp)
  recall = float(tp)/(tp+fn)

  return (beta**2 + 1)*precision*recall/((beta**2)*precision+recall)

def create_mapping(s):
  s_collapsed = set().union(*s)

  mapping = {}

  for ele in s_collapsed:
    set_of_ele = None
    for x in s:
      if ele in x:
        set_of_ele = x
        break
    mapping[ele] = x

  return mapping

def do_fmi(A, B):
  AVec, BVec = mk_vectors(A, B)
  return metrics.fowlkes_mallows_score(AVec, BVec)

def do_fmeasure(A, B):
  AVec, BVec = mk_vectors(A, B)
  return metrics.f1_score(AVec, BVec, average='micro')

def do_ari(A, B):
  AVec, BVec = mk_vectors(A, B)
  return metrics.adjusted_rand_score(AVec, BVec)

def do_ami(A, B):
  AVec, BVec = mk_vectors(A, B)
  return metrics.adjusted_mutual_info_score(AVec, BVec)

def main(args):
  clusters = []
  truth = json.load(open(args.groundtruth, 'r'))
  cmp_meth = None
  if args.method == "fmi":
    cmp_meth = do_fmi
  elif args.method == "f":
    cmp_meth = do_fmeasure
  elif args.method == "ari":
    cmp_meth = do_ari
  elif args.method == "ami":
    cmp_meth = do_ami
  else:
    print "Invalid clustering algorithm"
    return 1

  for i in args.clusters:
    clusters.append(json.load(open(i, 'r')))

  for B in clusters:
    print "%s,%s,%.4f" % (truth["name"], B["name"], cmp_meth(truth["labels"],B["labels"]))

  return 0

if __name__ == '__main__':
  parser = argparse.ArgumentParser("analyze_clusters")
  parser.add_argument("groundtruth", help="Ground truth data", type=str)
  parser.add_argument("clusters", help="Clusters in JSON format", nargs='+')
  parser.add_argument("-m", "--method", help="Clustering method", type=str, default="fmi", choices=["fmi", "f", "ari", "ami"])
  args = parser.parse_args()
  sys.exit(main(args))