clustering/clustering.py

   1 #!/usr/bin/env python3
   2
   3 import pandas as pd
   4 import numpy as np
   5 from sklearn.cluster import AffinityPropagation
   6 import fire
   7
   8 def affinity_clustering(similarities, output, damping=0.9, max_iter=100000, convergence_iter=30, preference_quantile=0.5, random_state=1968, verbose=True):
   9     '''
  10     similarities: feather file with a dataframe of similarity scores
  11     preference_quantile: parameter controlling how many clusters to make. higher values = more clusters. 0.85 is a good value with 3000 subreddits.
  12     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.
  13     '''
  14
  15     df = pd.read_feather(similarities)
  16     n = df.shape[0]
  17     mat = np.array(df.drop('subreddit',1))
  18     mat[range(n),range(n)] = 1
  19
  20     preference = np.quantile(mat,preference_quantile)
  21
  22     print(f"preference is {preference}")
  23
  24     print("data loaded")
  25
  26     clustering = AffinityPropagation(damping=damping,
  27                                      max_iter=max_iter,
  28                                      convergence_iter=convergence_iter,
  29                                      copy=False,
  30                                      preference=preference,
  31                                      affinity='precomputed',
  32                                      verbose=verbose,
  33                                      random_state=random_state).fit(mat)
  34
  35
  36     print(f"clustering took {clustering.n_iter_} iterations")
  37     clusters = clustering.labels_
  38
  39     print(f"found {len(set(clusters))} clusters")
  40
  41     cluster_data = pd.DataFrame({'subreddit': df.subreddit,'cluster':clustering.labels_})
  42
  43     cluster_sizes = cluster_data.groupby("cluster").count()
  44     print(f"the largest cluster has {cluster_sizes.subreddit.max()} members")
  45
  46     print(f"the median cluster has {cluster_sizes.subreddit.median()} members")
  47
  48     print(f"{(cluster_sizes.subreddit==1).sum()} clusters have 1 member")
  49
  50     cluster_data.to_feather(output)
  51
  52 if __name__ == "__main__":
  53     fire.Fire(affinity_clustering)