5_user_communities_global.py

#!/usr/bin/env python3
"""Create user communities on a global level.

For each user we compute a user embedding based on the average document
embeddings within all topics minus the topics centroid.
We then cluster users based on these embeddings.
"""


import gzip
import json
from collections import defaultdict
from typing import Tuple, Dict

import hdbscan
import numpy as np
from umap import UMAP
from tqdm.auto import tqdm

from util import (
    NumpyArrayEncoder,
    get_users,
    get_embeddings,
    get_doc_topics,
    get_topics,
    get_topic_events,
    EVENT_USER_COMMUNITIES_PATH,
    USER_EMBEDDINGS_PATH,
)

umap_model = UMAP(
    n_components=5,
    n_neighbors=20,
    min_dist=0.0,
    metric="cosine",
    random_state=42,
    n_jobs=1,
)
umap_model_2d = UMAP(
    n_components=2,
    n_neighbors=umap_model.n_neighbors,
    min_dist=umap_model.min_dist,
    metric=umap_model.metric,
    random_state=umap_model.random_state,
    n_jobs=1,
)
hdbscan_model = hdbscan.HDBSCAN(
    min_cluster_size=30,
    metric="euclidean",
    prediction_data=True,
    min_samples=10,
    core_dist_n_jobs=1,
)

doc_topics = get_doc_topics()
users = get_users()
embeddings = get_embeddings()


def get_user_communities() -> Tuple[Dict[int, int], Dict[int, np.ndarray]]:
    # compute user embeddings based on topic embeddings
    user_topicwise_embeddings = defaultdict(list)
    for topic in tqdm(get_topics()[1:], desc="Topics"):
        topic_doc_indices = np.where(doc_topics == topic)[0]
        topic_users = users[topic_doc_indices]
        topic_embeddings = embeddings[topic_doc_indices]
        topic_embedding_centroid = np.mean(topic_embeddings, axis=0)
        # get embeddings for each user in the topic and compute bias
        for user in np.unique(topic_users):
            user_indices = np.where(users == user)[0]
            user_topic_embeddings = embeddings[user_indices]
            # skip user if less than 10 documents
            if user_topic_embeddings.shape[0] < 10:
                continue
            user_topic_embedding_centroid = np.mean(user_topic_embeddings, axis=0)
            user_topic_embedding_bias = (
                user_topic_embedding_centroid - topic_embedding_centroid
            )
            user_topicwise_embeddings[user].append(user_topic_embedding_bias)
    # compute global user embeddings based on bias in topic embeddings
    user_global_embeddings = {
        user: np.mean(embeddings, axis=0)
        for user, embeddings in user_topicwise_embeddings.items()
    }

    # 2D user embeddings for visualization
    umap_embeddings_2d = umap_model_2d.fit_transform(
        list(user_global_embeddings.values())
    )
    user_embeddings_map_2d = {
        user: embedding
        for user, embedding in zip(user_global_embeddings.keys(), umap_embeddings_2d)
    }

    # cluster users
    # and remove outliers with soft clustering
    umap_embeddings = umap_model.fit_transform(list(user_global_embeddings.values()))
    soft_clusters = hdbscan_model.fit_predict(umap_embeddings)
    cluster_membership_vecs = hdbscan.all_points_membership_vectors(hdbscan_model)
    hard_clusters = np.argmax(cluster_membership_vecs, axis=1)
    selected_clusters = hard_clusters
    # mapping of user to community
    user_community_map = {
        user: cluster
        for user, cluster in zip(user_global_embeddings.keys(), selected_clusters)
    }
    return user_community_map, user_embeddings_map_2d


if __name__ == "__main__":
    with gzip.open(EVENT_USER_COMMUNITIES_PATH, "wt") as f_event_user_comms, gzip.open(
        USER_EMBEDDINGS_PATH, "wt"
    ) as f_user_emb:
        (
            user_community_map,
            user_embeddings_map_2d,
        ) = get_user_communities()
        # save user embeddings
        for user, embedding in user_embeddings_map_2d.items():
            f_user_emb.write(
                json.dumps(
                    {
                        "user": user,
                        "embedding": embedding,
                    },
                    cls=NumpyArrayEncoder,
                )
                + "\n"
            )
        for topic_id in get_topics()[1:]:
            for event in get_topic_events(topic_id).values():
                for doc_id in event.doc_indices:
                    f_event_user_comms.write(
                        json.dumps(
                            {
                                "topic": topic_id,
                                "event": event.id,
                                "doc": doc_id,
                                "user": users[doc_id],
                                "community": user_community_map.get(users[doc_id], -1),
                            },
                            cls=NumpyArrayEncoder,
                        )
                        + "\n"
                    )