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

Community Data Science Collective || Want to submit a patch?