cluster

#!/usr/bin/env python2
from sklearn.cluster import KMeans
from sklearn.cluster import SpectralClustering
from sklearn.cluster import spectral_clustering
from sklearn.cluster import DBSCAN
from sklearn.cluster import AgglomerativeClustering
from sklearn.cluster import MeanShift
from sklearn.metrics.pairwise import cosine_similarity 
from scipy.sparse import csr_matrix
from sklearn.feature_selection import VarianceThreshold
from sklearn.decomposition import PCA
from multiprocessing import Pool
from functools import partial 
from collections import OrderedDict
import numpy as np
import argparse
import json
import sys
import csv
import glob
import os

def cluster_mk_results(args, r, index_to_file):
  labels = []
  for i,c in enumerate(r.labels_):
    labels.append((str(index_to_file[i]), int(c)))
  results = {}
  results['name'] = args.name
  results['labels'] = labels
  return results

def cluster_lifted(args, index_to_file, vecs, sel, c):
  na = np.array(vecs)
  if sel != None:
    before_dim = na.shape
    na = sel.fit_transform(na)
    after_dim = na.shape
    print "Reduced dimensionality from %s to %s" % (str(before_dim), str(after_dim))
  
  if args.pca != None:
    pca = PCA(n_components=args.pca, svd_solver='full')
    before_dim = na.shape
    na = pca.fit_transform(na)
    after_dim = na.shape
    print "PCA reduced dimensionality from %s to %s" % (str(before_dim), str(after_dim))

  r = c.fit(na)
  return cluster_mk_results(args, r, index_to_file)

def cluster_meanshift(args, index_to_file, vecs, sel=None):
  ms = MeanShift(n_jobs=-1)
  return cluster_lifted(args, index_to_file, vecs, sel, ms)

def cluster_agglomerative(args, index_to_file, vecs, sel=None):
  agg = AgglomerativeClustering(n_clusters=args.clusters)
  return cluster_lifted(args, index_to_file, vecs, sel, agg)

def cluster_dbscan(args, index_to_file, vecs, sel=None):
  dbscan = DBSCAN(n_jobs=-1)
  return cluster_lifted(args, index_to_file, vecs, sel, dbscan)

def cluster_kmeans(args, index_to_file, vecs, sel=None):
  kmeans = KMeans(n_clusters=args.clusters, n_jobs=-1)
  return cluster_lifted(args, index_to_file, vecs, sel, kmeans)

def cluster_spectral(args, index_to_file, vecs, sel=None):
  na = np.array(vecs)
  if sel != None:
    before_dim = na.shape
    na = sel.fit_transform(na)
    after_dim = na.shape
    print "Reduced dimensionality from %s to %s" % (str(before_dim), str(after_dim))
  
  if args.pca != None:
    pca = PCA(n_components=args.pca, svd_solver='full')
    before_dim = na.shape
    na = pca.fit_transform(na)
    after_dim = na.shape
    print "PCA reduced dimensionality from %s to %s" % (str(before_dim), str(after_dim))

  na = cosine_similarity(na)
  na = np.exp(-na / na.std())
  spec = SpectralClustering(n_clusters=args.clusters, n_jobs=-1, affinity='precomputed')
  r = spec.fit(na)

  return cluster_mk_results(args, r, index_to_file)

def loadfile(f):
  try:
    l = json.load(open(f, 'r'))
  except ValueError:
    print "Could not read {}".format(f)
    return (None,None)
  return (l['file'], l['tuples'])

def alignmap(inmap, total_key_set):
  fname,m = inmap 
  diff = total_key_set.difference(m.keys())
  for k in diff:
    m[k] = 0
  return (fname,m)

def createvec(inmap, total_keys):
  fname,m = inmap 
  vec = []
  for k in total_keys:
    vec.append(m[k])
  return (fname,vec)

def main(args):
  p = Pool(1)
  # Read in all the vector files as dictionaries. 
  print "reading in files" 
  files = []
  for i in args.files:
    if os.path.isdir(i):
      for j in glob.glob(i+os.path.sep+"*.json"):
        files.append(j)
    else:
      files.append(i)
  from_files = p.map(loadfile, files)

  print "aligning dictionaries"
  # Pad out all the dicts in map.dicts to have the same key set.
  tks = set()
  for (a,m) in from_files:
    tks = tks.union(m.keys())

  alignmappartial = partial(alignmap, total_key_set=tks)
  aligned_dicts = p.map(alignmappartial, from_files)

  print "creating vectors"
  tks_ordered = OrderedDict.fromkeys(tks)
  createvecpartial = partial(createvec, total_keys=tks_ordered)

  # Now, turn these into actual vectors, with an order controlled by
  # the set of all keys. 
  vecs_with_filename = p.map(createvecpartial, aligned_dicts)

  # Now, go through and flatten vecs_with_filename. 
  vecs = []
  index_to_file = {}
  index = 0
  for (f,v) in vecs_with_filename:
    vecs.append(v)
    index_to_file[index] = f
    index = index + 1
  
  if args.delog:
    for i in range(0, len(vecs)):
      for j in range(0, len(vecs[i])):
        vecs[i][j] = vecs[i][j] ** 2

  if args.dump_datavec:
    print "Writing out data vector to CSV only"
    t = vecs[0]
    header = ["filename"]
    header.extend([ "feature_%d" % i for i in range(0,len(t))])
    contents = [header]
    for i in range(0, len(vecs)):
      c = []
      c.append(index_to_file[i])
      c.extend(vecs[i])
      contents.append(c)
    outcsv = open(args.dump_datavec, 'w')
    writer = csv.writer(outcsv)
    for i in contents:
      writer.writerow(i)
    return 0

  # Now we can actually do some clustering. 
  print "clustering with %s" % args.cluster_method
  cluster_meth = None
  if args.cluster_method == 'kmeans':
    cluster_meth = cluster_kmeans
  elif args.cluster_method == 'spectral':
    cluster_meth = cluster_spectral
  elif args.cluster_method == 'dbscan':
    cluster_meth = cluster_dbscan
  elif args.cluster_method == 'agg':
    cluster_meth = cluster_agglomerative
  elif args.cluster_method == 'meanshift':
    cluster_meth = cluster_meanshift
  else:
    print "Unknown cluster method %s" % args.cluster_method
    return 1

  sel = None
  if args.variance_threshold != None:
    sel = VarianceThreshold(threshold=(args.variance_threshold * (1-args.variance_threshold)))
  results = cluster_meth(args, index_to_file, vecs, sel)
  # Write the results of this clustering out to a file. 
  out = open(args.outfile, 'w')
  json.dump(results, out)
  out.close()
  return 0

if __name__ == '__main__':
  parser = argparse.ArgumentParser("cluster")
  parser.add_argument('--name', type=str)
  parser.add_argument('--outfile', type=str, default='cluster.json')
  parser.add_argument('--dump-datavec', type=str, default=None)
  parser.add_argument('files', nargs='+')
  parser.add_argument('--pca', type=float, default=None)
  parser.add_argument('--clusters', type=int, default=8)
  parser.add_argument('--delog', default=False, action='store_true')
  parser.add_argument('--variance-threshold', default=None, type=float)
  parser.add_argument('--cluster-method', type=str, default='kmeans', choices=['kmeans', 'spectral', 'dbscan', 'agg', 'meanshift'])
  args = parser.parse_args()
  sys.exit(main(args))