path_laplacian_retrieval_lab.py

#!/usr/bin/env python3
"""A deterministic retrieval experiment for the Loadbear path-Laplacian series.

The graph is intentionally synthetic.  It has a planted relevant chain and a
nearby, lexical-looking but irrelevant support cluster.  That lets us ask a
useful question with a known answer: as indirect paths get stronger, do we find
the distant relevant nodes before we start spending context on the wrong topic?

The paper's operator is used for transport:

    L_hat(alpha) = sum_d d**(-alpha) L^(d)
    dx/dt = -L_hat(alpha) x

The small query prior below is *not* from the paper.  It stands in for whatever
lexical/semantic scorer a real retrieval system already has.  Keeping it
separate from transport is the point of the experiment.

Run:
    python3 scripts/path_laplacian_retrieval_lab.py
    python3 scripts/path_laplacian_retrieval_lab.py --json
"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict, dataclass
from math import inf


@dataclass(frozen=True)
class Node:
    label: str
    relevant: bool
    tokens: int
    query_prior: float


@dataclass(frozen=True)
class Result:
    alpha: float
    selected: tuple[str, ...]
    recall: float
    leakage: int
    tokens: int


# The first four candidates are the answer-bearing chain.  The last four make
# up a second topic cluster reached through one sparse bridge.  "renewal FAQ"
# is a deliberate lexical decoy: it has a decent query prior, but is not part
# of the answer to an approval-follow-up query.
NODES = (
    Node("customer request", False, 24, 1.00),  # query anchor; never selected
    Node("pricing owner", True, 30, 0.58),
    Node("approval policy", True, 42, 0.94),
    Node("follow-up task", True, 34, 0.82),
    Node("contract exception", True, 46, 0.66),
    Node("support bridge", False, 28, 0.14),
    Node("renewal FAQ", False, 38, 0.72),
    Node("incident timeline", False, 44, 0.31),
    Node("on-call rota", False, 30, 0.17),
)

# A relevant chain 0--1--2--3--4, followed by a sparse bridge 2--5 into the
# unrelated support cluster 5--6--7--8.
EDGES = ((0, 1), (1, 2), (2, 3), (3, 4), (2, 5), (5, 6), (6, 7), (7, 8))
QUERY_ANCHOR = 0
# The four answer-bearing nodes occupy exactly 152 tokens.  A lexical decoy
# costs 38 tokens, so it can crowd out the fourth answer-bearing node when the
# ranking leaks into the support cluster.
TOKEN_BUDGET = 152
TRANSPORT_WEIGHT = 0.82
DIFFUSION_TIME = 0.90
EULER_STEP = 0.01


def shortest_distances() -> list[list[float]]:
    """Floyd-Warshall is plenty for this nine-node fixture."""
    size = len(NODES)
    distances = [[0.0 if row == column else inf for column in range(size)] for row in range(size)]
    for left, right in EDGES:
        distances[left][right] = distances[right][left] = 1.0
    for middle in range(size):
        for left in range(size):
            for right in range(size):
                distances[left][right] = min(distances[left][right], distances[left][middle] + distances[middle][right])
    return distances


DISTANCES = shortest_distances()


def path_laplacian(alpha: float) -> list[list[float]]:
    """Build L_hat from shortest-path layers, not from invented graph edges."""
    size = len(NODES)
    laplacian = [[0.0 for _ in range(size)] for _ in range(size)]
    for left in range(size):
        for right in range(left + 1, size):
            distance = DISTANCES[left][right]
            if distance == inf:
                continue
            weight = distance ** (-alpha)
            laplacian[left][left] += weight
            laplacian[right][right] += weight
            laplacian[left][right] -= weight
            laplacian[right][left] -= weight
    return laplacian


def matvec(matrix: list[list[float]], vector: list[float]) -> list[float]:
    return [sum(entry * vector[index] for index, entry in enumerate(row)) for row in matrix]


def diffuse(alpha: float) -> list[float]:
    """Explicit Euler approximation of x' = -L_hat x from the query anchor."""
    laplacian = path_laplacian(alpha)
    values = [0.0 for _ in NODES]
    values[QUERY_ANCHOR] = 1.0
    for _ in range(round(DIFFUSION_TIME / EULER_STEP)):
        disagreement = matvec(laplacian, values)
        values = [value - EULER_STEP * delta for value, delta in zip(values, disagreement)]
    return values


def rank(alpha: float) -> list[tuple[float, int]]:
    """Combine fixed query evidence with transport, then rank candidates.

    The combination is deliberately ordinary weighted addition.  It is here so
    the lab tests a *retrieval policy*, rather than quietly treating diffusion
    time as a relevance label.
    """
    transported = diffuse(alpha)
    transport_max = max(transported[index] for index in range(len(NODES)) if index != QUERY_ANCHOR)
    scored = []
    for index, node in enumerate(NODES):
        if index == QUERY_ANCHOR:
            continue
        transport = transported[index] / transport_max if transport_max else 0.0
        score = TRANSPORT_WEIGHT * transport + (1.0 - TRANSPORT_WEIGHT) * node.query_prior
        scored.append((score, index))
    return sorted(scored, reverse=True)


def select(alpha: float) -> Result:
    """Greedily pack ranked context until the fixed token budget is exhausted."""
    selected: list[str] = []
    spent = 0
    relevant = 0
    leakage = 0
    for _, index in rank(alpha):
        node = NODES[index]
        if spent + node.tokens > TOKEN_BUDGET:
            continue
        selected.append(node.label)
        spent += node.tokens
        if node.relevant:
            relevant += 1
        else:
            leakage += 1
    total_relevant = sum(node.relevant for node in NODES)
    return Result(alpha, tuple(selected), relevant / total_relevant, leakage, spent)


def run() -> list[Result]:
    return [select(alpha) for alpha in (0.0, 0.5, 1.0, 2.0, 3.0, 4.0, 8.0)]


def verify(results: list[Result]) -> None:
    """Keep the deliberately planted tradeoff from silently drifting away."""
    by_alpha = {result.alpha: result for result in results}
    for alpha in (1.0, 2.0, 3.0):
        assert by_alpha[alpha].recall == 1.0 and by_alpha[alpha].leakage == 0
    for alpha in (0.0, 4.0, 8.0):
        assert by_alpha[alpha].recall == 0.75 and by_alpha[alpha].leakage == 1


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--json", action="store_true", help="emit machine-readable results")
    parser.add_argument("--verify", action="store_true", help="assert the fixture's intended tradeoff")
    args = parser.parse_args()
    results = run()
    if args.verify:
        verify(results)
    if args.json:
        print(json.dumps([asdict(result) for result in results], indent=2))
        return

    print("alpha  recall  leak  tokens  selected")
    for result in results:
        print(
            f"{result.alpha:>4.1f}   {result.recall:>4.2f}    {result.leakage:>1}    "
            f"{result.tokens:>3}   {', '.join(result.selected)}"
        )


if __name__ == "__main__":
    main()