Authorship Attribution by N-gram Profiles

def make_corpus(
    authors=AUTHORS,
    chars=CHARS,
    preferred_weight=PREFERRED_WEIGHT,
    other_weight=OTHER_WEIGHT,
    chars_per_author=CHARS_PER_AUTHOR,
    word_length_min=WORD_LENGTH_MIN,
    word_length_max=WORD_LENGTH_MAX,
    words_per_text=WORDS_PER_TEXT,
    texts_per_author=TEXTS_PER_AUTHOR,
    seed=SEED,
):
    """Return a Polars DataFrame with columns author, text_id, role, text.

    role is 'train' for the first texts_per_author-1 texts per author and
    'test' for the last.  Each text is a space-separated sequence of synthetic
    words; each author's words are biased toward their preferred character set
    so that character n-gram profiles are clearly distinct between authors.
    """
    rng = np.random.default_rng(seed)
    n_chars = len(chars)
    records = []
    text_id = 0

    for a_idx, author in enumerate(authors):
        probs = np.full(n_chars, other_weight)
        start = a_idx * chars_per_author
        end = start + chars_per_author
        probs[start:end] = preferred_weight
        probs /= probs.sum()

        for t_idx in range(texts_per_author):
            words = []
            for _ in range(words_per_text):
                length = int(rng.integers(word_length_min, word_length_max + 1))
                word = "".join(
                    chars[int(rng.choice(n_chars, p=probs))] for _ in range(length)
                )
                words.append(word)
            role = "test" if t_idx == texts_per_author - 1 else "train"
            records.append(
                {
                    "author": author,
                    "text_id": text_id,
                    "role": role,
                    "text": " ".join(words),
                }
            )
            text_id += 1

    return pl.DataFrame(records)

def char_ngrams(text, n):
    """Return a Counter of character n-grams in text.

    Spaces are included so that word-boundary patterns (e.g., the bigram
    formed by the last character of one word and the space before the next)
    contribute to the profile alongside within-word patterns.
    """
    return Counter(text[i : i + n] for i in range(len(text) - n + 1))

def build_profile(texts, n=NGRAM_SIZE):
    """Build a normalised character n-gram frequency profile from a list of texts.

    Counts are pooled across all texts, then divided by the total count so that
    the profile is a probability distribution over observed n-grams.
    Returns a dict mapping n-gram string to relative frequency.
    """
    counts = Counter()
    for text in texts:
        counts += char_ngrams(text, n)
    total = sum(counts.values())
    return {ng: c / total for ng, c in counts.items()}

def cosine_similarity(profile_a, profile_b):
    """Return the cosine similarity between two n-gram frequency profiles.

    Profiles are dicts of {ngram: frequency}.  The similarity is computed as:
      dot(a, b) / (norm(a) * norm(b))
    where only n-grams present in both profiles contribute to the dot product,
    and each norm is taken over all n-grams in that profile.
    Returns a value in [0, 1]: 1 means identical profiles, 0 means no shared n-grams.
    """
    shared = set(profile_a) & set(profile_b)
    dot = sum(profile_a[ng] * profile_b[ng] for ng in shared)
    norm_a = sum(v * v for v in profile_a.values()) ** 0.5
    norm_b = sum(v * v for v in profile_b.values()) ** 0.5
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)

def attribute(unknown_profile, candidate_profiles):
    """Return candidates ranked by cosine similarity to the unknown profile.

    unknown_profile is a dict produced by build_profile for one unknown text.
    candidate_profiles is a dict mapping author name to its profile dict.
    Returns a list of (author, similarity) pairs sorted from highest to lowest.
    """
    scores = [
        (author, cosine_similarity(unknown_profile, profile))
        for author, profile in candidate_profiles.items()
    ]
    return sorted(scores, key=lambda pair: pair[1], reverse=True)

def plot_similarity(ranked, filename):
    """Save a horizontal bar chart of cosine similarity scores."""
    df = pl.DataFrame([{"author": author, "similarity": sim} for author, sim in ranked])
    chart = (
        alt.Chart(df)
        .mark_bar()
        .encode(
            y=alt.Y("author:N", title="Candidate author", sort=None),
            x=alt.X(
                "similarity:Q",
                title="Cosine similarity",
                scale=alt.Scale(domain=[0.0, 1.0]),
            ),
            color=alt.Color("author:N", legend=None),
        )
        .properties(
            width=320,
            height=160,
            title="Authorship attribution by character bigram similarity",
        )
    )
    chart.save(filename)

import polars as pl
import pytest
from generate_authorship import make_corpus, AUTHORS
from authorship import char_ngrams, build_profile, cosine_similarity, attribute


def test_char_ngrams_basic():
    # "abc" with n=2 yields exactly the bigrams "ab" and "bc".
    counts = char_ngrams("abc", 2)
    assert counts["ab"] == 1
    assert counts["bc"] == 1
    assert len(counts) == 2


def test_char_ngrams_repeated():
    # "aaa" with n=2 yields two overlapping "aa" bigrams.
    counts = char_ngrams("aaa", 2)
    assert counts["aa"] == 2


def test_profile_sums_to_one():
    # A profile built from any non-empty text must sum to 1.
    profile = build_profile(["abcabc"], n=2)
    assert sum(profile.values()) == pytest.approx(1.0)


def test_cosine_identical_profiles():
    # Identical profiles have cosine similarity exactly 1.0.
    p = {"ab": 0.5, "cd": 0.5}
    assert cosine_similarity(p, p) == pytest.approx(1.0)


def test_cosine_disjoint_profiles():
    # Profiles with no shared n-grams have cosine similarity 0.0.
    p1 = {"ab": 1.0}
    p2 = {"cd": 1.0}
    assert cosine_similarity(p1, p2) == pytest.approx(0.0)


def test_attribution_correct_author():
    # Each test text must be attributed to its true author.
    df = make_corpus()
    train_df = df.filter(pl.col("role") == "train")
    test_df = df.filter(pl.col("role") == "test")

    candidate_profiles = {
        author: build_profile(
            train_df.filter(pl.col("author") == author)["text"].to_list()
        )
        for author in AUTHORS
    }

    for row in test_df.iter_rows(named=True):
        unknown_profile = build_profile([row["text"]])
        ranked = attribute(unknown_profile, candidate_profiles)
        predicted = ranked[0][0]
        assert predicted == row["author"], (
            f"Misattributed: true={row['author']}, predicted={predicted}"
        )