Word Frequency Change Over Time

def make_corpus(
    decades=DECADES,
    tokens_per_decade=TOKENS_PER_DECADE,
    bg_vocab_size=BACKGROUND_VOCAB_SIZE,
    target_words=TARGET_WORDS,
    target_base_freq=TARGET_BASE_FREQ,
    target_slope=TARGET_SLOPE,
    seed=SEED,
):
    """Return a Polars DataFrame with columns decade, word, count.

    For each decade, target word frequencies follow the injected linear trend
    (base_freq + slope * decade_index).  The remaining probability mass is
    shared equally among bg_vocab_size background words and sampled with a
    multinomial draw, so each run with the same seed produces identical counts.
    """
    rng = np.random.default_rng(seed)
    records = []
    for decade_index, decade in enumerate(decades):
        # Target word frequencies for this decade.
        target_freqs = [
            base + slope * decade_index
            for base, slope in zip(target_base_freq, target_slope)
        ]
        remaining = 1.0 - sum(target_freqs)
        bg_freq = remaining / bg_vocab_size

        # Build probability vector: target words first, then background words.
        bg_words = [f"word{i:03d}" for i in range(bg_vocab_size)]
        all_words = target_words + bg_words
        all_probs = target_freqs + [bg_freq] * bg_vocab_size

        counts = rng.multinomial(tokens_per_decade, all_probs)
        for word, count in zip(all_words, counts):
            records.append({"decade": decade, "word": word, "count": int(count)})

    return pl.DataFrame(records)

def normalize(df):
    """Return df with a 'freq' column: each word's share of tokens that decade.

    freq = count / sum(count) for the same decade.
    The result is a probability distribution over words for each decade,
    so values sum to 1.0 within each decade.
    """
    totals = df.group_by("decade").agg(pl.col("count").sum().alias("total"))
    return df.join(totals, on="decade").with_columns(
        (pl.col("count") / pl.col("total")).alias("freq")
    )

def linear_trend(decade_indices, freqs):
    """Return the OLS slope: change in normalized frequency per decade.

    Uses the closed-form ordinary least-squares formula:

        slope = sum((x - x_mean) * (y - y_mean)) / sum((x - x_mean)^2)

    where x is the decade index (0, 1, ..., n-1) and y is the normalized
    frequency.  The result is in units of frequency change per 10-year period.
    """
    x = np.asarray(decade_indices, dtype=float)
    y = np.asarray(freqs, dtype=float)
    x_c = x - x.mean()
    denom = np.dot(x_c, x_c)
    if denom == 0.0:
        return 0.0
    return float(np.dot(x_c, y) / denom)


def compute_trends(freq_df, words):
    """Return a dict mapping each word to its OLS slope (per decade).

    freq_df must have columns decade, word, freq, sorted by decade.
    Decade index 0 corresponds to the earliest decade in the data.
    """
    min_decade = freq_df["decade"].min()
    trends = {}
    for word in words:
        subset = freq_df.filter(pl.col("word") == word).sort("decade")
        decade_indices = [(d - min_decade) // 10 for d in subset["decade"].to_list()]
        freqs = subset["freq"].to_list()
        trends[word] = linear_trend(decade_indices, freqs)
    return trends

def plot_trajectories(freq_df, words, filename):
    """Save an Altair line chart of normalized frequency over time for each word.

    Each word appears as a separate line; the x-axis is the decade and the
    y-axis is the normalized frequency (proportion of all tokens that decade).
    """
    chart_df = freq_df.filter(pl.col("word").is_in(words)).sort(["word", "decade"])
    chart = (
        alt.Chart(chart_df)
        .mark_line(point=True)
        .encode(
            x=alt.X("decade:O", title="Decade"),
            y=alt.Y("freq:Q", title="Normalized frequency"),
            color=alt.Color("word:N", title="Word"),
        )
        .properties(
            width=400,
            height=260,
            title="Word frequency trajectories by decade",
        )
    )
    chart.save(filename)

import pytest
import polars as pl
from generate_wordshift import make_corpus, TARGET_WORDS, TARGET_SLOPE
from wordshift import normalize, linear_trend, compute_trends


def test_normalize_sums_to_one():
    # After normalization, frequencies within each decade must sum to 1.0.
    df = make_corpus()
    freq_df = normalize(df)
    decade_totals = freq_df.group_by("decade").agg(
        pl.col("freq").sum().alias("total_freq")
    )
    for total in decade_totals["total_freq"].to_list():
        assert total == pytest.approx(1.0, abs=1e-9)


def test_normalize_non_negative():
    # All normalized frequencies must be non-negative.
    df = make_corpus()
    freq_df = normalize(df)
    assert freq_df["freq"].min() >= 0.0


def test_linear_trend_flat():
    # A constant frequency sequence has slope exactly 0.
    indices = list(range(11))
    freqs = [0.015] * 11
    assert linear_trend(indices, freqs) == pytest.approx(0.0)


def test_linear_trend_known_slope():
    # A perfectly linear sequence has the exact injected slope.
    # slope = 0.002, starting at 0.005, over 11 decades.
    indices = list(range(11))
    freqs = [0.005 + 0.002 * i for i in indices]
    assert linear_trend(indices, freqs) == pytest.approx(0.002, abs=1e-12)


def test_linear_trend_single_point():
    # A single-point series has an undefined denominator; the function
    # must return 0.0 without raising an exception.
    assert linear_trend([0], [0.5]) == pytest.approx(0.0)


def test_slope_signs_correct():
    # telegraph must have a positive slope, candle a negative slope,
    # and steam a slope near zero.  Tolerance of 0.001 is 5x the
    # expected standard error given 5000 tokens per decade and 11 decades.
    df = make_corpus()
    freq_df = normalize(df)
    trends = compute_trends(freq_df, TARGET_WORDS)
    assert trends["telegraph"] > 0.0
    assert trends["candle"] < 0.0
    assert trends["steam"] == pytest.approx(0.0, abs=0.001)


def test_slope_magnitude():
    # The recovered slopes should be within 0.001 of the true injected values.
    # Standard error of the OLS slope is < 0.0002 for these parameters,
    # so 0.001 allows for more than 5 standard errors of sampling variation.
    df = make_corpus()
    freq_df = normalize(df)
    trends = compute_trends(freq_df, TARGET_WORDS)
    true_slopes = dict(zip(TARGET_WORDS, TARGET_SLOPE))
    for word in TARGET_WORDS:
        assert trends[word] == pytest.approx(true_slopes[word], abs=0.001), (
            f"{word}: expected slope {true_slopes[word]}, got {trends[word]:.5f}"
        )