collocation.py

import os
import re
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from collections import Counter

# CONFIGURATION
CSV_DIR = "epoch_csvs"
# List of your 5-year epoch labels (must match your CSV filenames)
EPOCHS = ["2007_2011", "2012_2016", "2017_2021", "2022_2025"]

# How many tokens apart to consider a "collocation"
WINDOW_SIZE = 3

# How many top pairs (edges) to keep in the graph
TOP_EDGES = 60

# Output folder for graphs
OUT_DIR = "collocation_graphs"
os.makedirs(OUT_DIR, exist_ok=True)

# HELPER: TOKENIZE A STRING
TOKEN_PATTERN = re.compile(r"[a-z']+")

def tokenize(text):
    """
    Lowercase + keep only a–z and apostrophes, return a list of tokens.
    """
    return TOKEN_PATTERN.findall(text.lower())

# FUNCTION: BUILD GRAPH FOR ONE EPOCH
def build_epoch_graph(epoch_label, window=WINDOW_SIZE, top_n=TOP_EDGES):
    """
    1. Reads the CSV for the given epoch.
    2. Concatenates all lyrics into one long string.
    3. Tokenizes into a list of words.
    4. For each position i in tokens, looks at tokens[j] for j in [i+1, i+window].
    5. Counts each pair (token_i, token_j) in a Counter.
    6. Builds a NetworkX graph from the top_n most common pairs (as weighted edges).
    """
    csv_path = os.path.join(CSV_DIR, f"odd_future_{epoch_label}.csv")
    df = pd.read_csv(csv_path)
    # Concatenate all lyrics in this epoch into one string
    all_lyrics = " ".join(df["lyrics"].fillna("")).lower()
    tokens = tokenize(all_lyrics)

    # Count co-occurrences within the sliding window
    pair_counter = Counter()
    N = len(tokens)
    for i in range(N):
        w1 = tokens[i]
        # Look ahead up to `window` tokens
        for j in range(i + 1, min(i + window + 1, N)):
            w2 = tokens[j]
            # We can avoid counting identical tokens with itself
            if w1 != w2:
                # Sort the pair for undirected edge consistency (w1,w2) same as (w2,w1)
                pair = tuple(sorted((w1, w2)))
                pair_counter[pair] += 1

    # Take the top_n most common pairs to create edges
    top_pairs = pair_counter.most_common(top_n)

    # Build a NetworkX graph
    G = nx.Graph()
    for (w1, w2), weight in top_pairs:
        G.add_edge(w1, w2, weight=weight)

    return G

# LOOP OVER EPOCHS AND SAVE GRAPHS
for epoch in EPOCHS:
    print(f"Building collocation graph for epoch: {epoch}")
    G = build_epoch_graph(epoch_label=epoch,
                          window=WINDOW_SIZE,
                          top_n=TOP_EDGES)

    # Draw the graph
    plt.figure(figsize=(6, 6))
    # spring_layout for node positions (gives a “force‐directed” look)
    pos = nx.spring_layout(G, seed=42, k=0.25)

    # Draw nodes and edges.  
    # Node size fixed at 30, edge width ~ scaled by weight (divided by a constant for visibility)
    edge_weights = [G[u][v]["weight"] for u, v in G.edges()]
    # Normalize edge width: dividing by max weight to keep widths small
    max_w = max(edge_weights) if edge_weights else 1
    widths = [3 * (w / max_w) for w in edge_weights]

    nx.draw_networkx_nodes(G, pos, node_size=30, node_color="skyblue", alpha=0.8)
    nx.draw_networkx_edges(G, pos, width=widths, edge_color="gray", alpha=0.6)
    nx.draw_networkx_labels(G, pos, font_size=6, font_color="black")

    plt.title(f"Collocation Network: {epoch}", fontsize=10)
    plt.axis("off")
    plt.tight_layout()

    out_path = os.path.join(OUT_DIR, f"net_{epoch}.png")
    plt.savefig(out_path, dpi=150)
    plt.close()
    print(f"[✓] Saved graph for {epoch} → {out_path}")

print("All collocation graphs generated.")