Co-occurrence Matrices › Code Implementation

# Sample corpus
corpus = [
    "the cat sat on the mat",
    "the cat ate the rat",
    "the dog sat on the log",
    "the dog ran in the park",
    "the cat chased the rat in the park",
    "dogs and cats are common pets",
    "a cat or a dog makes a good pet",
    "the rat ran from the cat",
    "cats and dogs play in the garden",
    "the mat sat under the cat",
]

# Tokenize and build vocabulary
all_tokens = []
sentences = []
for sent in corpus:
    tokens = sent.lower().split()
    sentences.append(tokens)
    all_tokens.extend(tokens)

# Count frequencies and build vocabulary (keep words with freq >= 2)
from collections import Counter

freq = Counter(all_tokens)
vocab = sorted(w for w, c in freq.items() if c >= 2)
word2idx = {w: i for i, w in enumerate(vocab)}
idx2word = {i: w for w, i in word2idx.items()}
V = len(vocab)

def build_cooccurrence_matrix(
    sentences, word2idx, window_size=2, weighted=False
):
    """Build a word-word co-occurrence matrix.

    Args:
        sentences: List of tokenized sentences
        word2idx: Mapping from word to index
        window_size: Context window half-size
        weighted: If True, weight co-occurrences by 1/distance

    Returns:
        Dense numpy array of shape (V, V)
    """
    V = len(word2idx)
    C = np.zeros((V, V), dtype=float)

    for sent in sentences:
        # Map tokens to indices, skipping OOV words
        indices = [word2idx[w] for w in sent if w in word2idx]

        for i, target_idx in enumerate(indices):
            # Context window bounds
            left = max(0, i - window_size)
            right = min(len(indices), i + window_size + 1)

            for j in range(left, right):
                if i == j:
                    continue
                context_idx = indices[j]
                distance = abs(i - j)
                weight = 1.0 / distance if weighted else 1.0
                C[target_idx, context_idx] += weight

    return C


# Build both versions
C_uniform = build_cooccurrence_matrix(
    sentences, word2idx, window_size=2, weighted=False
)
C_weighted = build_cooccurrence_matrix(
    sentences, word2idx, window_size=2, weighted=True
)

total_entries = V * V
nonzero_count = np.count_nonzero(C_uniform)
sparsity = 1.0 - (nonzero_count / total_entries)
density = nonzero_count / total_entries

def build_word_document_matrix(corpus, word2idx):
    """Build a word-document co-occurrence matrix.

    Returns:
        Dense array of shape (V, D)
    """
    V = len(word2idx)
    D = len(corpus)
    M = np.zeros((V, D), dtype=float)

    for doc_idx, sent in enumerate(corpus):
        tokens = sent.lower().split()
        for word in tokens:
            if word in word2idx:
                M[word2idx[word], doc_idx] += 1

    return M


M_doc = build_word_document_matrix(corpus, word2idx)

# Compare word similarity from word-word vs word-document matrices
from numpy.linalg import norm


def cosine_similarity(u, v):
    """Compute cosine similarity between two vectors."""
    u_norm = norm(u)
    v_norm = norm(v)
    if u_norm == 0 or v_norm == 0:
        return 0.0
    return np.dot(u, v) / (u_norm * v_norm)


# Find most similar words to 'cat' in each matrix
target = "cat"
target_idx = word2idx[target]

# Word-word similarities
ww_sims = {}
for w, idx in word2idx.items():
    if w != target:
        ww_sims[w] = cosine_similarity(C_uniform[target_idx], C_uniform[idx])

# Word-document similarities
wd_sims = {}
for w, idx in word2idx.items():
    if w != target:
        wd_sims[w] = cosine_similarity(M_doc[target_idx], M_doc[idx])

# Build matrices with different window sizes to compare
C_window1 = build_cooccurrence_matrix(
    sentences, word2idx, window_size=1, weighted=False
)
C_window4 = build_cooccurrence_matrix(
    sentences, word2idx, window_size=4, weighted=False
)

# Build matrices at several window sizes
window_sizes = [1, 2, 3, 4, 5]
target_word = "dog"
target_idx = word2idx[target_word]
compare_word = "cat"
compare_idx = word2idx[compare_word]

# Track similarity between 'dog' and 'cat' across window sizes
dog_cat_sims = []
for ws in window_sizes:
    C = build_cooccurrence_matrix(
        sentences, word2idx, window_size=ws, weighted=False
    )
    sim = cosine_similarity(C[target_idx], C[compare_idx])
    dog_cat_sims.append(sim)

from scipy.sparse import lil_matrix


def build_sparse_cooccurrence(
    sentences, word2idx, window_size=2, weighted=False
):
    """Build co-occurrence matrix in sparse LIL format, then convert to CSR."""
    V = len(word2idx)
    # LIL (List of Lists) format is efficient for incremental construction
    C = lil_matrix((V, V), dtype=float)

    for sent in sentences:
        indices = [word2idx[w] for w in sent if w in word2idx]
        for i, target_idx in enumerate(indices):
            left = max(0, i - window_size)
            right = min(len(indices), i + window_size + 1)
            for j in range(left, right):
                if i == j:
                    continue
                context_idx = indices[j]
                distance = abs(i - j)
                weight = 1.0 / distance if weighted else 1.0
                C[target_idx, context_idx] += weight

    # Convert to CSR for efficient row-wise operations
    return C.tocsr()


C_sparse = build_sparse_cooccurrence(
    sentences, word2idx, window_size=2, weighted=True
)

# Get all nonzero counts from the uniform matrix
nonzero_counts = C_uniform[C_uniform > 0].flatten()
nonzero_counts_sorted = np.sort(nonzero_counts)[::-1]