reproduction-bundle/src/run_completeness.py

#!/usr/bin/env python3
"""Independent finite certificates for both cases of Proposition C.2.

Case 1 exhausts every nonconstant event for several independent categorical
product spaces.  Case 2 exhausts every nonconstant event over an eight-leaf
taxonomy whose predicates are pairwise nested or mutually exclusive.  Each
event is compiled into CI conjunctions and ME disjunctions, then compared with
an independently differentiated direct probability sum.
"""

from __future__ import annotations

import argparse
import itertools
import json
from pathlib import Path

import numpy as np


def softmax(x: np.ndarray) -> np.ndarray:
    z = x - x.max()
    e = np.exp(z)
    return e / e.sum()


def categorical_atoms(logits: list[np.ndarray]) -> tuple[list[np.ndarray], list[list[tuple[float, np.ndarray]]]]:
    probabilities = [softmax(x.astype(float)) for x in logits]
    dimension = sum(map(len, probabilities))
    atoms: list[list[tuple[float, np.ndarray]]] = []
    offset = 0
    for probs in probabilities:
        group = []
        for index, probability in enumerate(probs):
            score = np.zeros(dimension)
            score[offset : offset + len(probs)] = -probs
            score[offset + index] += 1.0
            group.append((float(probability), score))
        atoms.append(group)
        offset += len(probs)
    return probabilities, atoms


def assignment_term(atoms: list[list[tuple[float, np.ndarray]]], assignment: tuple[int, ...]) -> tuple[float, np.ndarray]:
    probability = 1.0
    score = np.zeros_like(atoms[0][0][1])
    for group, value in enumerate(assignment):
        p, s = atoms[group][value]
        probability *= p
        score += s
    return probability, score


def compile_me(terms: list[tuple[float, np.ndarray]]) -> tuple[float, np.ndarray]:
    probability = sum(p for p, _ in terms)
    gradient = sum((p * s for p, s in terms), np.zeros_like(terms[0][1]))
    return probability, gradient / probability


def direct_product_event(
    probabilities: list[np.ndarray],
    atoms: list[list[tuple[float, np.ndarray]]],
    selected: tuple[tuple[int, ...], ...],
) -> tuple[float, np.ndarray]:
    probability = 0.0
    gradient = np.zeros_like(atoms[0][0][1])
    for assignment in selected:
        p = float(np.prod([probabilities[g][v] for g, v in enumerate(assignment)]))
        s = sum((atoms[g][v][1] for g, v in enumerate(assignment)), np.zeros_like(gradient))
        probability += p
        gradient += p * s
    return probability, gradient / probability


def case1_certificate(seeds: int) -> dict:
    domain_shapes = [(2, 2, 2), (2, 3), (3, 3)]
    cases = 0
    max_probability_error = 0.0
    max_score_error = 0.0
    for seed in range(seeds):
        rng = np.random.default_rng(seed)
        for shape in domain_shapes:
            probabilities, atoms = categorical_atoms([rng.normal(size=n) for n in shape])
            assignments = tuple(itertools.product(*(range(n) for n in shape)))
            for mask in range(1, 2 ** len(assignments) - 1):
                selected = tuple(a for i, a in enumerate(assignments) if mask & (1 << i))
                compiled = compile_me([assignment_term(atoms, a) for a in selected])
                direct = direct_product_event(probabilities, atoms, selected)
                max_probability_error = max(max_probability_error, abs(compiled[0] - direct[0]))
                max_score_error = max(max_score_error, float(np.max(np.abs(compiled[1] - direct[1]))))
                cases += 1
    return {
        "domain_shapes": [list(x) for x in domain_shapes],
        "seeds": seeds,
        "all_nonconstant_events_checked": cases,
        "max_probability_error": max_probability_error,
        "max_score_error": max_score_error,
        "passed": max_probability_error < 1e-14 and max_score_error < 1e-14,
    }


def valid_taxonomy_sets() -> list[frozenset[int]]:
    # Complete binary taxonomy over eight terminal leaves.
    return [
        frozenset(range(8)),
        frozenset(range(4)),
        frozenset(range(4, 8)),
        frozenset((0, 1)),
        frozenset((2, 3)),
        frozenset((4, 5)),
        frozenset((6, 7)),
        *(frozenset((i,)) for i in range(8)),
    ]


def pairwise_nested_or_me(sets: list[frozenset[int]]) -> bool:
    for i, left in enumerate(sets):
        for right in sets[i + 1 :]:
            if left & right and not (left <= right or right <= left):
                return False
    return True


def case2_certificate(seeds: int) -> dict:
    taxonomy = valid_taxonomy_sets()
    assert pairwise_nested_or_me(taxonomy)
    cases = 0
    max_probability_error = 0.0
    max_score_error = 0.0
    for seed in range(seeds):
        probabilities, atoms = categorical_atoms([np.random.default_rng(10_000 + seed).normal(size=8)])
        probs = probabilities[0]
        for mask in range(1, 2**8 - 1):
            selected = tuple(i for i in range(8) if mask & (1 << i))
            compiled = compile_me([atoms[0][i] for i in selected])
            probability = float(probs[list(selected)].sum())
            gradient = sum((probs[i] * atoms[0][i][1] for i in selected), np.zeros(8))
            direct = (probability, gradient / probability)
            max_probability_error = max(max_probability_error, abs(compiled[0] - direct[0]))
            max_score_error = max(max_score_error, float(np.max(np.abs(compiled[1] - direct[1]))))
            cases += 1
    invalid_overlap = [frozenset((0, 1)), frozenset((1, 2))]
    return {
        "taxonomy_predicates": len(taxonomy),
        "taxonomy_leaves": 8,
        "pairwise_nested_or_mutually_exclusive": True,
        "seeds": seeds,
        "all_nonconstant_semantic_events_checked": cases,
        "max_probability_error": max_probability_error,
        "max_score_error": max_score_error,
        "invalid_overlap_control_rejected": not pairwise_nested_or_me(invalid_overlap),
        "passed": (
            max_probability_error < 1e-14
            and max_score_error < 1e-14
            and not pairwise_nested_or_me(invalid_overlap)
        ),
    }


def run(output_dir: Path, seeds: int) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    report = {
        "paper": "OAM1jJsMGp",
        "claim": "Anchored claim 5 / Proposition C.2 completeness",
        "case_1_independent_categorical_groups": case1_certificate(seeds),
        "case_2_nested_or_me_taxonomy": case2_certificate(seeds),
    }
    report["all_checks_pass"] = all(section["passed"] for section in report.values() if isinstance(section, dict))
    (output_dir / "completeness_report.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))
    return report


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=Path("outputs/completeness"))
    parser.add_argument("--seeds", type=int, default=25)
    args = parser.parse_args()
    run(args.output_dir, args.seeds)


if __name__ == "__main__":
    main()