clustering/clustering.py

   1 #!/usr/bin/env python3
   2 # TODO: replace prints with logging.
   3 import sys
   4 import pandas as pd
   5 import numpy as np
   6 from sklearn.cluster import AffinityPropagation
   7 import fire
   8 from pathlib import Path
   9
  10 def read_similarity_mat(similarities, use_threads=True):
  11     df = pd.read_feather(similarities, use_threads=use_threads)
  12     mat = np.array(df.drop('_subreddit',1))
  13     n = mat.shape[0]
  14     mat[range(n),range(n)] = 1
  15     return (df._subreddit,mat)
  16
  17 def affinity_clustering(similarities, *args, **kwargs):
  18     subreddits, mat = read_similarity_mat(similarities)
  19     return _affinity_clustering(mat, subreddits, *args, **kwargs)
  20
  21 def _affinity_clustering(mat, subreddits, output, damping=0.9, max_iter=100000, convergence_iter=30, preference_quantile=0.5, random_state=1968, verbose=True):
  22     '''
  23     similarities: feather file with a dataframe of similarity scores
  24     preference_quantile: parameter controlling how many clusters to make. higher values = more clusters. 0.85 is a good value with 3000 subreddits.
  25     damping: parameter controlling how iterations are merged. Higher values make convergence faster and more dependable. 0.85 is a good value for the 10000 subreddits by author.
  26     '''
  27     print(f"damping:{damping}; convergenceIter:{convergence_iter}; preferenceQuantile:{preference_quantile}")
  28
  29     preference = np.quantile(mat,preference_quantile)
  30
  31     print(f"preference is {preference}")
  32     print("data loaded")
  33     sys.stdout.flush()
  34     clustering = AffinityPropagation(damping=damping,
  35                                      max_iter=max_iter,
  36                                      convergence_iter=convergence_iter,
  37                                      copy=False,
  38                                      preference=preference,
  39                                      affinity='precomputed',
  40                                      verbose=verbose,
  41                                      random_state=random_state).fit(mat)
  42
  43
  44     print(f"clustering took {clustering.n_iter_} iterations")
  45     clusters = clustering.labels_
  46
  47     print(f"found {len(set(clusters))} clusters")
  48
  49     cluster_data = pd.DataFrame({'subreddit': subreddits,'cluster':clustering.labels_})
  50
  51     cluster_sizes = cluster_data.groupby("cluster").count()
  52     print(f"the largest cluster has {cluster_sizes.subreddit.max()} members")
  53
  54     print(f"the median cluster has {cluster_sizes.subreddit.median()} members")
  55
  56     print(f"{(cluster_sizes.subreddit==1).sum()} clusters have 1 member")
  57
  58     sys.stdout.flush()
  59     cluster_data.to_feather(output)
  60     print(f"saved {output}")
  61     return clustering
  62
  63 if __name__ == "__main__":
  64     fire.Fire(affinity_clustering)