siamese_generate_graphs.py

"""
Generate and analyze semantic networks from word embeddings.

This script constructs semantic networks from word embeddings and calculates
network metrics to analyze exploration patterns. It supports processing multiple
participant trajectories in parallel and outputs detailed network statistics.

Functionality:
    1. Load word embedding models.
    2. Construct semantic networks from participant trajectories.
    3. Calculate network metrics including clustering, path lengths, and edge weights.
    4. Output metrics for analyzing exploration styles (Hunter, Busybody, Dancer).

Example:
    $ python siamese_generate_graphs.py -v 3

Attributes:
    logger (logging.Logger): Configured logging instance for tracking script execution.
    log_filename (str): Path to log file with timestamp.

Authors:
    Stephen Steinle
    My Duong
    Amber Nugen
    Kareem Sinan

Created:
    2025
"""

import argparse
import logging
import multiprocessing as mp
import os
from datetime import datetime

import bct
import gensim.downloader as api
import networkx as nx
import numpy as np
import pandas as pd
from gensim.models import KeyedVectors
from tqdm import tqdm

import siamese_equations as eq

# Create logs directory if it doesn't exist
os.makedirs("logs", exist_ok=True)

# Configure logging
log_filename = (
    f'logs/siamese_generate_graphs_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'
)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[
        logging.FileHandler(log_filename),
        logging.StreamHandler(),  # This will also logger.info to console
    ],
)
logger = logging.getLogger(__name__)
logger.info(f"Logging to file: {log_filename}")


def create_word_graph(words, threshold, W2V):
    """
    Creates a 2D graph with words as nodes and edges based on cosine distance between embeddings.

    Parameters:
    words (list of str): A list of words.
    threshold (float): A float threshold for edge inclusion based on cosine distance.

    Returns:
    G (networkx.Graph): The generated graph with words as nodes and edges based on the distance threshold.
    """
    g = nx.Graph()
    # Add nodes to the graph
    for word in words:
        g.add_node(word)

    # Calculate cosine distances and add edges
    for i, word1 in enumerate(words):
        for j, word2 in enumerate(words):
            if i == j or word1 == word2:
                continue
            # Calculate cosine distance between word embeddings
            distance = 1 - W2V.similarity(word1, word2)

            if distance == 0:
                logger.warning(
                    f"Zero distance between {word1} and {word2} during networkx graph creation"
                )

            # Add edge if adjacent or meets threshold
            if distance < threshold:
                g.add_edge(word1, word2, weight=1 / distance)
            elif j - i == 1:
                g.add_edge(word1, word2, weight=1 / distance)
    return g


def calculate_network_metrics_dict_weighted(g):
    """
    Calculate the suite of network metrics for graph g
    g: graph object output from the function utils_networkx.make_graph_links

    N: Number of nodes in the network N
    E: Number of edges in the network
    rho: Density
    C: Clustering coefficient
    k: average degree
    spl: Characteristic path length
    eff: gloabal efficiency
    cp: core-periphery structure: Borgatti-Everett method
    mdl: minimum description length of the hierarchical degree-corrected stochastic block model
    B: Number of blocks in the SBM on the lowest level of the hierarchy
    """
    # if disconnected, take largest component
    g_conmax = g.subgraph(max(nx.connected_components(g), key=len))

    N = nx.number_of_nodes(g)  # nodes
    E = nx.number_of_edges(g)  # edges

    k = np.mean([v for k, v in nx.degree(g, weight="weight")])  # degree
    spl = nx.average_shortest_path_length(
        g_conmax, weight="weight", method="dijkstra"
    )  # shortest path length
    eff = nx.global_efficiency(g)  # global efficiency

    if len(g_conmax) > 1:
        cp = bct.core_periphery_dir(nx.convert_matrix.to_numpy_array(g))
    else:
        cp = [np.nan, np.nan]

    return {
        "N": N,
        "E": E,
        "average_degree": k,
        "Characteristic_path_length": spl,
        "global_efficiency": eff,
        "core_periphery_structure": np.around(cp[1], 8),
    }


def process_user_list_weighted(
    pair: tuple[str, int, list[str]],
    threshold: float,
    data: pd.DataFrame,
    dataset: str,
    W2V: KeyedVectors,
) -> dict[str, float | str]:
    """Process a user's word list to generate weighted graph metrics and descriptions.

    Args:
        pair (tuple[str, int, list[str]]): Tuple containing (user_id, list_number, word_list)
        threshold (float): Threshold value for creating word graph edges
        data (pd.DataFrame): DataFrame containing metadata about the lists
        dataset (str): Name of dataset being processed ("SNAFU", "AMHR", or "FF")
        W2V (KeyedVectors): Word2Vec model for calculating word similarities

    Returns:
        dict[str, float | str]: Dictionary containing graph metrics and descriptions including:
            - Basic metadata (id, listnum, category, etc.)
            - Network metrics from calculate_network_metrics_dict_weighted()
            - Edge weight descriptions
            - Node weight descriptions
            - Node gravity metrics
            - Forward flow and other derived metrics
    """
    if dataset == "SNAFU":
        fal_dict = {"id": pair[0], "listnum": pair[1]}
        fal_dict["category"] = data[
            (data["id"] == pair[0]) & (data["listnum"] == pair[1])
        ]["category"].values[0]
        fal_dict["group"] = data[
            (data["id"] == pair[0]) & (data["listnum"] == pair[1])
        ]["group"].values[0]
        clean_list = pair[2]
    elif dataset == "AMHR":
        fal_dict = {
            "id": pair[0],
            "listnum": pair[1],
            "category": data[(data["User_ID"] == pair[0])]["Cue"].values[0],
        }
        clean_list = pair[2]
    elif dataset == "FF":
        fal_dict = {"id": pair[0], "listnum": pair[1]}

        clean_list = pair[2]

        fal_dict["category"] = data[
            (data["id"] == pair[0]) & (data["listnum"] == pair[1])
        ]["item_FF"].values[0]
    else:
        logger.info("Invalid dataset")
        return

    fal_dict["threshold"] = threshold

    graph_nx = create_word_graph(clean_list, threshold, W2V)

    net_mets = calculate_network_metrics_dict_weighted(graph_nx)

    edge_descriptions = eq.describe_edge_weights(graph_nx)

    node_weights = eq.describe_node_weights(clean_list)

    node_gravity_1 = eq.describe_node_gravity_1(clean_list, W2V)

    node_gravity_2 = eq.describe_node_gravity_2(clean_list, W2V)

    description_dict = {
        "weights_min": edge_descriptions.minmax[0],
        "weights_max": edge_descriptions.minmax[1],
        "weights_mean": edge_descriptions.mean,
        "weights_variance": edge_descriptions.variance,
        "weights_skewness": edge_descriptions.skewness,
        "weights_kurtosis": edge_descriptions.kurtosis,
    }

    fal_dict |= {f"{k}": v for k, v in net_mets.items()}
    fal_dict |= {f"{k}": v for k, v in description_dict.items()}

    new_weighted = eq.compute_relevant_weighted_graph_metrics(graph_nx)
    fal_dict |= {f"{k}": v for k, v in new_weighted.items()}

    fal_dict["bb_h"] = eq.weighted_hunter_busybody(fal_dict)
    fal_dict["ff"] = sum(eq.forward_flow(clean_list, W2V)) / (len(clean_list) - 1)

    fal_dict |= {f"{k}": v for k, v in node_weights.items()}
    fal_dict |= {f"{k}": v for k, v in node_gravity_1.items()}
    fal_dict |= {f"{k}": v for k, v in node_gravity_2.items()}

    return fal_dict


def clean_fals(fals: list[str], W2V: KeyedVectors) -> list[str]:
    """Clean a list of words by filtering out words not in Word2Vec vocabulary.

    Args:
        fals (list[str]): List of words to clean
        W2V (KeyedVectors): Word2Vec model containing vocabulary to check against

    Returns:
        list[str]: Cleaned list with unknown words replaced by "_", avoiding consecutive "_"s
    """
    clean_list = []

    for word in fals:
        if word in W2V:
            clean_list.append(word)
        elif clean_list and clean_list[-1] == "_":
            continue
        else:
            clean_list.append("_")

    return clean_list


def generate_snafu_data(
    snafu_df: pd.DataFrame,
    nproc: int = 30,
    threshold: int = 3,
    W2V: KeyedVectors | None = None,
) -> pd.DataFrame:
    """Generate graph descriptions for SNAFU data by processing the provided dataframe.

    Args:
        snafu_df (pd.DataFrame): DataFrame containing SNAFU task data
        nproc (int): Number of processes to use for parallel processing. Defaults to 30.
        threshold (int): Threshold value for processing. Defaults to 3.
        W2V (KeyedVectors | None): Word2Vec model for word embeddings. Defaults to None.

    Returns:
        pd.DataFrame: Graph descriptions for SNAFU data with unique id/listnum combinations
    """
    t_lists = 0

    ul_list = []

    for user in snafu_df["id"].unique():
        t_lists += len(snafu_df[snafu_df["id"] == user]["listnum"].unique())
        for lid in snafu_df[snafu_df["id"] == user]["listnum"].unique():
            lid = int(lid)
            ulist = snafu_df[(snafu_df["id"] == user) & (snafu_df["listnum"] == lid)][
                "item"
            ].values.tolist()
            clean_list = clean_fals(ulist, W2V)
            if len(set(clean_list)) >= 5:
                ul_list.append([user, lid, clean_list])
            ul_list.append([user, lid, clean_list])

    chunks = [(pair, threshold, snafu_df, "SNAFU", W2V) for pair in ul_list]

    with mp.Pool(nproc) as pool:
        results = pool.starmap(
            process_user_list_weighted, tqdm(chunks, total=len(chunks))
        )

    snafu_graph_descriptions_weighted = pd.DataFrame(results)

    snafu_graph_descriptions_weighted = (
        snafu_graph_descriptions_weighted.drop_duplicates(subset=["id", "listnum"])
    )

    return snafu_graph_descriptions_weighted


def generate_ff_data(
    ff_df: pd.DataFrame | None = None,
    nproc: int = 30,
    threshold: int = 3,
    W2V: KeyedVectors | None = None,
) -> pd.DataFrame:
    """Generate graph descriptions for FF (Fluency Finder) data by processing the provided dataframe.

    Args:
        ff_df (pd.DataFrame | None): DataFrame containing FF task data. Defaults to None.
        nproc (int): Number of processes to use for parallel processing. Defaults to 30.
        threshold (int): Threshold value for processing. Defaults to 3.
        W2V (KeyedVectors | None): Word2Vec model for word embeddings. Defaults to None.

    Returns:
        pd.DataFrame: Graph descriptions for FF data
    """
    ff_df["listnum"] = ff_df.groupby("id").cumcount()

    t_lists = 0

    ul_list = []

    for user in ff_df["id"].unique():
        t_lists += len(ff_df[ff_df["id"] == user]["listnum"].unique())
        for lid in ff_df[ff_df["id"] == user]["listnum"].unique():
            lid = int(lid)
            ulist = (
                ff_df[(ff_df["id"] == user) & (ff_df["listnum"] == lid)]
                .apply(lambda row: list(row.iloc[1:-1].dropna()), axis=1)
                .values[0]
            )
            clean_list = clean_fals(ulist, W2V)
            if len(set(clean_list)) >= 5:
                ul_list.append([user, lid, clean_list])

    chunks = [(pair, threshold, ff_df, "FF", W2V) for pair in ul_list]

    with mp.Pool(nproc) as pool:
        results = pool.starmap(
            process_user_list_weighted, tqdm(chunks, total=len(chunks))
        )

    return pd.DataFrame(results)


def generate_amhr_data(
    cfa_df: pd.DataFrame,
    sft_df: pd.DataFrame,
    nproc: int = 30,
    threshold: int = 3,
    W2V: KeyedVectors = None,
) -> pd.DataFrame:
    """Generate graph descriptions for AMHR data by processing CFA and SFT dataframes.

    Args:
        cfa_df (pd.DataFrame): DataFrame containing CFA task data
        sft_df (pd.DataFrame): DataFrame containing SFT task data
        nproc (int, optional): Number of processes to use for parallel processing. Defaults to 30.
        threshold (int, optional): Threshold value for processing. Defaults to 3.
        W2V (KeyedVectors, optional): Word2Vec model for word embeddings. Defaults to None.

    Returns:
        pd.DataFrame: Combined graph descriptions for CFA and SFT data
    """
    cfa_users = set(cfa_df["User_ID"].unique())
    sft_users = set(sft_df["User_ID"].unique())
    fa_users = cfa_users.intersection(sft_users)

    ul_list = []

    for user in fa_users:
        ulist = (
            sft_df[(sft_df["User_ID"] == user) & (sft_df["Task"] == "SFT")]
            .apply(lambda row: list(row.iloc[2:].dropna()), axis=1)
            .values[0]
        )
        clean_list = clean_fals(ulist, W2V)
        if len(set(clean_list)) >= 5:
            ul_list.append([user, 0, clean_list])

    chunks = [(pair, threshold, sft_df, "AMHR", W2V) for pair in ul_list]

    with mp.Pool(nproc) as pool:
        results = pool.starmap(
            process_user_list_weighted, tqdm(chunks, total=len(chunks))
        )

    sft_graph_descriptions_weighted = pd.DataFrame(results)

    ul_list = []

    for user in fa_users:
        ulist = (
            cfa_df[(cfa_df["User_ID"] == user) & (cfa_df["Task"] == "CFA")]
            .apply(lambda row: list(row.iloc[2:].dropna()), axis=1)
            .values[0]
        )
        clean_list = clean_fals(ulist, W2V)
        if len(set(clean_list)) >= 5:
            ul_list.append([user, 0, clean_list])

    chunks = [(pair, threshold, cfa_df, "AMHR", W2V) for pair in ul_list]

    with mp.Pool(nproc) as pool:
        results = pool.starmap(
            process_user_list_weighted, tqdm(chunks, total=len(chunks))
        )

    cfa_graph_descriptions_weighted = pd.DataFrame(results)

    sft_cfa_graph_descriptions_wegihted = pd.concat(
        [cfa_graph_descriptions_weighted, sft_graph_descriptions_weighted]
    )

    sft_cfa_graph_descriptions_wegihted["listnum"] = 0

    return sft_cfa_graph_descriptions_wegihted


if __name__ == "__main__":
    """Main script execution for generating graph descriptions from multiple datasets.

    This script processes three different datasets (SNAFU, FF, and AMHR) to generate graph descriptions
    using word embeddings from Word2Vec. For each dataset, it:
    1. Loads the data from specified file paths
    2. Processes the data to generate graph descriptions
    3. Saves the resulting graph descriptions to pickle files

    Command line arguments:
        -v (str): Version of the data to save. Default is '0'.
        --snafu_path (str): Path to the SNAFU data CSV file. Default is 'sources/snafu_sample_cleaned.csv'
        --ff_path (str): Path to the FF data TSV file. Default is 'sources/FF_Input_Multi_Task.tsv'
        --amhr_path (str): Path to the AMHR data pickle file. Default is 'sources/amhr_data.pkl'
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-v",
        type=str,
        help="Version of the data to save. Default is '0'.",
        default="0",
    )
    parser.add_argument(
        "--snafu_path",
        type=str,
        help="Path to the SNAFU data CSV file.",
        default=f"sources/snafu_sample_cleaned.csv",
    )
    parser.add_argument(
        "--ff_path",
        type=str,
        help="Path to the FF data TSV file.",
        default=f"sources/FF_Input_Multi_Task.tsv",
    )
    parser.add_argument(
        "--amhr_path",
        type=str,
        help="Path to the CFA pickle file.",
        default=f"sources/amhr_data.pkl",
    )
    args = parser.parse_args()

    logger.info("Loading Word2Vec model...")
    w2v_path = api.load("word2vec-google-news-300", return_path=True)
    W2V = KeyedVectors.load_word2vec_format(w2v_path, binary=True)
    logger.info("Word2Vec model loaded.")

    logger.info("Generating SNAFU data...")

    logger.info(f"Loading SNAFU data from {args.snafu_path}...")
    snafu_df = pd.read_csv(args.snafu_path)
    snafu_df = generate_snafu_data(snafu_df, W2V=W2V)
    logger.info("Saving SNAFU data to disk...")
    snafu_df.to_pickle(
        f"graph_similarity/snafu_graph_descriptions_weighted_{args.v}.pkl"
    )

    logger.info("Generating FF data...")

    logger.info(f"Loading FF data from {args.ff_path}...")
    ff_df = pd.read_csv(args.ff_path, sep="\t")
    ff_df = generate_ff_data(ff_df, W2V=W2V)
    logger.info("Saving FF data to disk...")
    ff_df.to_pickle(f"graph_similarity/ff_graph_descriptions_weighted_{args.v}.pkl")

    logger.info("Generating AMHR data...")

    logger.info(f"Loading CFA data from {args.amhr_path}...")
    cfa_df = pd.read_pickle(args.amhr_path)
    cfa_df = cfa_df[cfa_df["Task"] == "CFA"]
    cfa_df.reset_index(drop=True, inplace=True)

    logger.info(f"Loading SFT data from {args.amhr_path}...")
    sft_df = pd.read_pickle(args.amhr_path)
    sft_df = sft_df[sft_df["Task"] == "SFT"]
    sft_df.reset_index(drop=True, inplace=True)

    amhr_df = generate_amhr_data(cfa_df, sft_df, W2V=W2V)

    logger.info("Saving AMHR data to disk...")
    amhr_df.to_pickle(f"graph_similarity/amhr_graph_descriptions_weighted_{args.v}.pkl")

    logger.info("Done.")