Vocabulary Richness in Historical Texts

def make_corpus(
    authors=AUTHORS,
    vocab_sizes=VOCAB_SIZES,
    texts_per_author=TEXTS_PER_AUTHOR,
    text_length=TEXT_LENGTH,
    zipf_exp=ZIPF_EXP,
    seed=SEED,
):
    """Return a Polars DataFrame with columns text_id, author, word.

    Each author has a distinct vocabulary size; words are sampled from
    a Zipfian frequency distribution (frequency proportional to 1/rank^zipf_exp).
    A larger vocabulary size produces more distinct words per token and thus
    a higher type-token ratio.
    """
    rng = np.random.default_rng(seed)
    records = []
    text_id = 0
    for author, vocab_size in zip(authors, vocab_sizes):
        ranks = np.arange(1, vocab_size + 1, dtype=float)
        probs = ranks ** (-zipf_exp)
        probs /= probs.sum()
        word_forms = [f"w{i:04d}" for i in range(vocab_size)]
        for _ in range(texts_per_author):
            sampled = rng.choice(word_forms, size=text_length, p=probs)
            for w in sampled:
                records.append({"text_id": text_id, "author": author, "word": w})
            text_id += 1
    return pl.DataFrame(records)

def type_token_ratio(words):
    """Return unique words / total words (type-token ratio, TTR).

    TTR falls as text length increases even when vocabulary richness is
    constant, because longer texts inevitably repeat high-frequency words.
    Use MATTR for length-fair comparisons across texts of different sizes.
    """
    words = list(words)
    if not words:
        return 0.0
    return len(set(words)) / len(words)


def mattr(words, window=MATTR_WINDOW):
    """Return the Moving-Average Type-Token Ratio.

    Computes TTR over each consecutive window of width `window`, then
    returns the mean of those window TTRs.  Because every window has the
    same length, the result is length-independent and can fairly compare
    texts of different sizes.  Falls back to global TTR when the text is
    shorter than the window.
    """
    words = list(words)
    n = len(words)
    if n <= window:
        return type_token_ratio(words)
    window_ttrs = [
        type_token_ratio(words[i : i + window]) for i in range(n - window + 1)
    ]
    return float(np.mean(window_ttrs))

def compute_richness(df, window=MATTR_WINDOW):
    """Return a DataFrame with TTR and MATTR for each text.

    Input df must have columns text_id, author, word.
    Output has columns text_id, author, ttr, mattr.
    """
    rows = []
    for text_id in df["text_id"].unique().sort():
        subset = df.filter(pl.col("text_id") == text_id)
        author = subset["author"][0]
        words = subset["word"].to_list()
        rows.append(
            {
                "text_id": int(text_id),
                "author": author,
                "ttr": type_token_ratio(words),
                "mattr": mattr(words, window),
            }
        )
    return pl.DataFrame(rows)

def plot_richness(richness_df, filename):
    """Save a bar chart of MATTR per text, coloured by author."""
    df = richness_df.with_columns(
        pl.concat_str(
            [
                pl.col("author"),
                pl.lit(" \u2013 text "),
                (pl.col("text_id") + 1).cast(pl.String),
            ]
        ).alias("label")
    )
    chart = (
        alt.Chart(df)
        .mark_bar()
        .encode(
            x=alt.X("label:N", title="Text", sort=None),
            y=alt.Y("mattr:Q", title="MATTR", scale=alt.Scale(zero=False)),
            color=alt.Color("author:N", title="Author"),
        )
        .properties(
            width=320,
            height=250,
            title="Moving-Average Type-Token Ratio by Text",
        )
    )
    chart.save(filename)

import pytest
from generate_vocab import make_corpus
from vocab import type_token_ratio, mattr, compute_richness
import polars as pl


def test_ttr_all_unique():
    # When every word is distinct the TTR is exactly 1.0.
    assert type_token_ratio(["a", "b", "c", "d"]) == pytest.approx(1.0)


def test_ttr_all_same():
    # When the same word is repeated n times the TTR is 1/n.
    assert type_token_ratio(["x"] * 10) == pytest.approx(0.1)


def test_ttr_empty():
    # Empty input returns 0.0 without raising an exception.
    assert type_token_ratio([]) == pytest.approx(0.0)


def test_mattr_short_text_equals_ttr():
    # When the text is shorter than the window, MATTR falls back to global TTR.
    words = ["a", "b", "c"]
    assert mattr(words, window=10) == pytest.approx(type_token_ratio(words))


def test_mattr_window_equals_length():
    # When text length equals the window width there is exactly one window,
    # so MATTR equals TTR exactly.
    words = ["a", "b", "c", "a"]
    assert mattr(words, window=4) == pytest.approx(type_token_ratio(words))


def test_richness_order_by_vocab_size():
    # Mean MATTR must increase with vocabulary size: Author A (200) <
    # Author B (400) < Author C (800).  The Zipfian generator guarantees
    # that a larger vocabulary produces more distinct tokens per window.
    df = make_corpus()
    richness = compute_richness(df)
    means = richness.group_by("author").agg(pl.col("mattr").mean()).sort("author")
    mattr_vals = means["mattr"].to_list()
    # Authors sort alphabetically: A, B, C -- matching ascending vocab sizes.
    assert mattr_vals[0] < mattr_vals[1] < mattr_vals[2]