Manuscript Stemma Reconstruction

def make_distance_matrix(noise_scale=0.0, seed=42):
    """Return (names, D) using exact tree-additive distances.

    When noise_scale=0.0 (default) the distances are exact and NJ recovers
    the true topology and branch lengths to machine precision.
    A non-zero noise_scale adds symmetric Gaussian noise to simulate the
    imperfect additivity of real manuscript data.
    """
    n = len(MANUSCRIPTS)
    D = np.zeros((n, n))
    for i, a in enumerate(MANUSCRIPTS):
        for j, b in enumerate(MANUSCRIPTS):
            if i < j:
                key = (a, b) if (a, b) in _BRANCH else (b, a)
                D[i, j] = D[j, i] = _BRANCH[key]
    if noise_scale > 0.0:
        rng = np.random.default_rng(seed)
        noise = rng.normal(0, noise_scale, (n, n))
        noise = (noise + noise.T) / 2
        np.fill_diagonal(noise, 0)
        D = np.maximum(D + noise, 0.0)
    return list(MANUSCRIPTS), D

def upgma(names, dist_matrix):
    """Run the UPGMA algorithm and return a list of (node_a, node_b, length) edges.

    `names` is a list of manuscript labels.
    `dist_matrix` is a symmetric distance matrix with zero diagonal.

    UPGMA (Unweighted Pair-Group Method with Arithmetic means) builds a rooted,
    ultrametric tree.  At each step it merges the pair with the smallest pairwise
    distance, assigns each branch half that distance, then recomputes distances to
    the new node as the average of the merged nodes' distances.

    Returns a list of (node_a, node_b, branch_length) triples.  For each internal
    node u created by merging i and j, two edges are added: (u, i, D(i,j)/2) and
    (u, j, D(i,j)/2).  The final two nodes are joined by one last edge.
    """
    names = list(names)
    D = dist_matrix.copy().astype(float)
    edges = []
    # node_counter starts after all leaf indices so generated names never
    # collide with manuscript labels such as "A", "B", etc.
    node_counter = len(names)

    while len(names) > 2:
        n = len(names)

        # Find the pair (i, j) with i < j that has the smallest distance.
        np.fill_diagonal(D, np.inf)
        idx = np.argmin(D)
        np.fill_diagonal(D, 0.0)
        i, j = divmod(int(idx), n)
        if i > j:
            i, j = j, i

        d_ij = D[i, j]
        # Each branch from the new node u to i and j has length D(i,j)/2.
        half = d_ij / 2.0
        u_name = f"node{node_counter}"
        node_counter += 1
        edges.append((u_name, names[i], half))
        edges.append((u_name, names[j], half))

        # Distance from u to every remaining taxon k is the average of the
        # distances from i and from j to k (UPGMA averaging rule).
        new_row = np.array(
            [0.5 * (D[i, k] + D[j, k]) for k in range(n)]
        )

        # Remove rows/columns for i and j; append row/column for u.
        keep = [k for k in range(n) if k != i and k != j]
        D_new = D[np.ix_(keep, keep)]
        extra = new_row[keep]
        size = len(keep) + 1
        D2 = np.zeros((size, size))
        D2[: len(keep), : len(keep)] = D_new
        D2[: len(keep), -1] = extra
        D2[-1, : len(keep)] = extra
        D = D2
        names = [names[k] for k in keep] + [u_name]

    # Two nodes remain; connect them with the remaining distance.
    edges.append((names[0], names[1], D[0, 1]))
    return edges

def hamming_distance(seq_a, seq_b):
    """Return the fraction of loci where seq_a and seq_b differ.

    Both sequences must be 1-D integer arrays of the same length.
    A value of 0.0 means identical manuscripts; 1.0 means they differ at
    every locus.  This is the Hamming distance normalized by sequence length.
    """
    return float(np.sum(seq_a != seq_b)) / len(seq_a)

def draw_stemma(edges, manuscripts, filename):
    """Save an unrooted stemma tree diagram to `filename`.

    Manuscript nodes (leaves) are drawn in steelblue; reconstructed
    intermediate ancestor nodes are drawn in lightgrey.
    Edge labels show the branch length (fraction of differing loci)
    rounded to two decimal places.
    """
    G = nx.Graph()
    for a, b, length in edges:
        G.add_edge(a, b, weight=round(length, 4))

    pos = nx.spring_layout(G, seed=3)

    leaf_nodes = [n for n in G.nodes() if n in manuscripts]
    internal_nodes = [n for n in G.nodes() if n not in manuscripts]

    fig, ax = plt.subplots(figsize=(6, 5))
    nx.draw_networkx_nodes(
        G, pos, nodelist=leaf_nodes, node_color="steelblue", node_size=700, ax=ax
    )
    nx.draw_networkx_nodes(
        G, pos, nodelist=internal_nodes, node_color="lightgrey", node_size=500, ax=ax
    )
    nx.draw_networkx_labels(
        G,
        pos,
        ax=ax,
        font_size=10,
        font_color="white",
        labels={n: n for n in leaf_nodes},
    )
    nx.draw_networkx_labels(
        G, pos, ax=ax, font_size=8, labels={n: n for n in internal_nodes}
    )
    nx.draw_networkx_edges(G, pos, ax=ax)
    edge_labels = {(a, b): f"{d:.2f}" for a, b, d in edges}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=7, ax=ax)

    ax.set_title("Reconstructed manuscript stemma")
    ax.axis("off")
    plt.tight_layout()
    plt.savefig(filename)
    plt.close(fig)

def make_manuscripts(
    n_loci=N_LOCI,
    mutation_prob=MUTATION_PROB,
    seed=SEED,
):
    """Return a dict mapping each manuscript name to its variant sequence.

    The copying tree is:

        archetype
        |-- alpha  (intermediate ancestor of A and B)
        |   |-- A
        |   `-- B
        `-- beta   (intermediate ancestor of C and D)
            |-- C
            `-- D

    Each branch introduces random mutations independently.  The archetype
    starts as all-zero; mutations flip a locus from 0 to 1 (or back).
    Only the four surviving manuscripts are returned; the archetype and
    intermediate ancestors are not observed.
    """
    rng = np.random.default_rng(seed)
    archetype = np.zeros(n_loci, dtype=np.int8)
    alpha = _copy(archetype, mutation_prob, rng)
    beta = _copy(archetype, mutation_prob, rng)
    mss = {
        "A": _copy(alpha, mutation_prob, rng),
        "B": _copy(alpha, mutation_prob, rng),
        "C": _copy(beta, mutation_prob, rng),
        "D": _copy(beta, mutation_prob, rng),
    }
    return mss

import pytest
import numpy as np
from generate_stemma import make_manuscripts, MANUSCRIPTS
from stemma import (
    hamming_distance,
    make_distance_matrix,
    upgma,
)


# ---------------------------------------------------------------------------
# hamming_distance
# ---------------------------------------------------------------------------


def test_hamming_identical():
    # Two identical sequences have distance 0.0.
    seq = np.array([0, 1, 0, 1, 1], dtype=np.int8)
    assert hamming_distance(seq, seq.copy()) == pytest.approx(0.0)


def test_hamming_all_different():
    # Sequences that differ at every locus have distance 1.0.
    seq_a = np.zeros(10, dtype=np.int8)
    seq_b = np.ones(10, dtype=np.int8)
    assert hamming_distance(seq_a, seq_b) == pytest.approx(1.0)


def test_hamming_half():
    # Sequences differing at exactly half the loci have distance 0.5.
    seq_a = np.zeros(8, dtype=np.int8)
    seq_b = np.array([1, 1, 1, 1, 0, 0, 0, 0], dtype=np.int8)
    assert hamming_distance(seq_a, seq_b) == pytest.approx(0.5)


# ---------------------------------------------------------------------------
# Distance matrix
# ---------------------------------------------------------------------------


def test_distance_matrix_symmetry():
    # The distance matrix must be symmetric with zero diagonal.
    names, D = make_distance_matrix()
    n = len(names)
    assert D.shape == (n, n)
    np.testing.assert_allclose(D, D.T, atol=1e-12)
    np.testing.assert_allclose(np.diag(D), 0.0, atol=1e-12)


def test_distance_matrix_within_family_smaller():
    # Within-family distances (A–B and C–D) must be smaller than cross-family
    # distances (A–C, A–D, B–C, B–D), reflecting the true copying tree.
    names, D = make_distance_matrix()
    idx = {n: i for i, n in enumerate(names)}
    d_ab = D[idx["A"], idx["B"]]
    d_cd = D[idx["C"], idx["D"]]
    cross = [D[idx[a], idx[b]] for a in "AB" for b in "CD"]
    assert d_ab < min(cross)
    assert d_cd < min(cross)


# ---------------------------------------------------------------------------
# UPGMA topology and branch lengths
# ---------------------------------------------------------------------------


def test_upgma_topology_recovery():
    # A and B must share a private internal ancestor, and C and D must share
    # a different one.  UPGMA merges the pair with the smallest distance first:
    # D(A,B) = 0.08 is the global minimum, so A and B are joined first into
    # node4; then D(C,D) = 0.10 < D(node4,C) = D(node4,D), so C and D are
    # joined into node5; finally node4 and node5 are joined.
    names, D = make_distance_matrix()
    edges = upgma(names, D)

    # Build an adjacency map.
    adj = {}
    for a, b, _ in edges:
        adj.setdefault(a, []).append(b)
        adj.setdefault(b, []).append(a)

    # Find the internal node connected to each manuscript.
    parent_a = [n for n in adj["A"] if n not in MANUSCRIPTS][0]
    parent_b = [n for n in adj["B"] if n not in MANUSCRIPTS][0]
    parent_c = [n for n in adj["C"] if n not in MANUSCRIPTS][0]
    parent_d = [n for n in adj["D"] if n not in MANUSCRIPTS][0]

    assert parent_a == parent_b, "A and B should share an internal ancestor"
    assert parent_c == parent_d, "C and D should share an internal ancestor"
    assert parent_a != parent_c, (
        "AB family and CD family should have different ancestors"
    )


def test_upgma_branch_lengths():
    # UPGMA branch lengths for the reference distance matrix.
    #
    # Step 1: merge A and B (D=0.08).  Each branch length = 0.08/2 = 0.04.
    #   D(node4, C) = (0.23 + 0.21) / 2 = 0.22
    #   D(node4, D) = (0.21 + 0.19) / 2 = 0.20
    #
    # Step 2: merge C and D (D=0.10 < 0.20 < 0.22).  Each branch length = 0.05.
    #   D(node4, node5) = (0.22 + 0.20) / 2 = 0.21
    #
    # Step 3: join node4 and node5 with the remaining distance = 0.21.
    #
    # Tolerance 1e-10 covers only accumulated floating-point rounding.
    names, D = make_distance_matrix()
    edges = upgma(names, D)

    expected = {
        frozenset(["node4", "A"]): 0.04,
        frozenset(["node4", "B"]): 0.04,
        frozenset(["node5", "C"]): 0.05,
        frozenset(["node5", "D"]): 0.05,
        frozenset(["node4", "node5"]): 0.21,
    }
    recovered = {frozenset([a, b]): length for a, b, length in edges}
    for key, true_len in expected.items():
        assert key in recovered, f"Missing edge {key}"
        assert recovered[key] == pytest.approx(true_len, abs=1e-10)


def test_upgma_all_lengths_positive():
    # Every branch length in the UPGMA tree must be strictly positive.
    names, D = make_distance_matrix()
    edges = upgma(names, D)
    for a, b, length in edges:
        assert length > 0, f"Non-positive branch length on edge {a}--{b}: {length}"


# ---------------------------------------------------------------------------
# Stochastic generator topology
# ---------------------------------------------------------------------------


def test_generator_topology():
    # Manuscripts generated stochastically must also produce a distance matrix
    # whose minimum within-family distance is smaller than the minimum
    # cross-family distance, so UPGMA can still find the correct grouping.
    mss = make_manuscripts()
    within = [
        hamming_distance(mss["A"], mss["B"]),
        hamming_distance(mss["C"], mss["D"]),
    ]
    cross = [hamming_distance(mss[a], mss[b]) for a in "AB" for b in "CD"]
    assert max(within) < min(cross), (
        "Within-family distances should be smaller than cross-family distances"
    )