version of weekly_cosine_similarities.py from klone

[cdsc_reddit.git] / similarities / weekly_cosine_similarities.py

diff --git a/similarities/weekly_cosine_similarities.py b/similarities/weekly_cosine_similarities.py
index 2b3c90be18cd08d7fbed3ae6e5c62bd77966499f..aeafe743ca1f0ecfe832741d2081fcbee5ae3a14 100644 (file)
--- a/similarities/weekly_cosine_similarities.py
+++ b/similarities/weekly_cosine_similarities.py
@@ -8,7 +8,22 @@ import fire
 from itertools import islice
 from pathlib import Path
 from similarities_helper import *
+from multiprocessing import pool
 
+def _week_similarities(tempdir, term_colname, week):
+        print(f"loading matrix: {week}")
+        mat = read_tfidf_matrix_weekly(tempdir.name, term_colname, week)
+        print('computing similarities')
+        sims = column_similarities(mat)
+        del mat
+
+        names = subreddit_names.loc[subreddit_names.week == week]
+        sims = pd.DataFrame(sims.todense())
+
+        sims = sims.rename({i: sr for i, sr in enumerate(names.subreddit.values)}, axis=1)
+        sims['_subreddit'] = names.subreddit.values
+
+        write_weekly_similarities(outfile, sims, week, names)
 
 #tfidf = spark.read.parquet('/gscratch/comdata/users/nathante/subreddit_tfidf_weekly.parquet')
 def cosine_similarities_weekly(tfidf_path, outfile, term_colname, min_df = None, included_subreddits = None, topN = 500):
@@ -35,25 +50,18 @@ def cosine_similarities_weekly(tfidf_path, outfile, term_colname, min_df = None,
     subreddit_names['subreddit_id_new'] = subreddit_names['subreddit_id_new'] - 1
     spark.stop()
 
-    weeks = list(subreddit_names.week.drop_duplicates())
-    for week in weeks:
-        print(f"loading matrix: {week}")
-        mat = read_tfidf_matrix_weekly(tempdir.name, term_colname, week)
-        print('computing similarities')
-        sims = column_similarities(mat)
-        del mat
+    weeks = sorted(list(subreddit_names.week.drop_duplicates()))
+    # do this step in parallel if we have the memory for it.
+    # should be doable with pool.map
 
-        names = subreddit_names.loc[subreddit_names.week == week]
-        sims = pd.DataFrame(sims.todense())
-
-        sims = sims.rename({i: sr for i, sr in enumerate(names.subreddit.values)}, axis=1)
-        sims['subreddit'] = names.subreddit.values
-
-        write_weekly_similarities(outfile, sims, week, names)
+    def week_similarities_helper(week):
+        _week_similarities(tempdir, term_colname, week)
 
+    with Pool(40) as pool: # maybe it can be done with 40 cores on the huge machine?
+        list(pool.map(weeks,week_similarities_helper))
 
-def author_cosine_similarities_weekly(outfile, min_df=None , included_subreddits=None, topN=500):
-    return cosine_similarities_weekly('/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
+def author_cosine_similarities_weekly(outfile, min_df=2 , included_subreddits=None, topN=500):
+    return cosine_similarities_weekly('/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors_30k.parquet',
                                       outfile,
                                       'author',
                                       min_df,
@@ -61,7 +69,7 @@ def author_cosine_similarities_weekly(outfile, min_df=None , included_subreddits
                                       topN)
 
 def term_cosine_similarities_weekly(outfile, min_df=None, included_subreddits=None, topN=500):
-    return cosine_similarities_weekly('/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
+    return cosine_similarities_weekly('/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms_30k.parquet',
                                       outfile,
                                       'term',
                                       min_df,
@@ -69,5 +77,5 @@ def term_cosine_similarities_weekly(outfile, min_df=None, included_subreddits=No
                                       topN)
 
 if __name__ == "__main__":
-    fire.Fire({'author':author_cosine_similarities_weekly,
-               'term':term_cosine_similarities_weekly})
+    fire.Fire({'authors':author_cosine_similarities_weekly,
+               'terms':term_cosine_similarities_weekly})