Listing: Listing

from fractions import Fraction as F
from random import Random
import sys

ZERO, ONE = F(0), F(1)

def clip01(x):
    return max(ZERO, min(ONE, x))

def bin_ge(n, k, p):
    """Exact P[Bin(n,p) >= k] for rational p."""
    assert isinstance(n, int) and n >= 1
    assert isinstance(k, int)
    assert isinstance(p, F) and 0 <= p <= 1
    if k <= 0:
        return ONE
    if k > n:
        return ZERO
    if p == 0:
        return ZERO
    if p == 1:
        return ONE
    if k == n:
        return p ** n
    a, b = p.numerator, p.denominator
    c = b - a
    term = c ** n
    total = 0
    for i in range(n + 1):
        if i >= k:
            total += term
        if i < n:
            num = term * (n - i) * a
            den = (i + 1) * c
            assert num % den == 0
            term = num // den
    return F(total, b ** n)

def cp_outward(n, k, tail, steps=32):
    """Conservative lower and upper CP endpoints."""
    assert isinstance(n, int) and n >= 1
    assert isinstance(k, int) and 0 <= k <= n
    assert isinstance(tail, F) and 0 < tail < 1
    assert isinstance(steps, int) and steps >= 1

    if k == 0:
        lower = ZERO
    else:
        lo, hi = ZERO, ONE
        for _ in range(steps):
            mid = (lo + hi) / 2
            if bin_ge(n, k, mid) <= tail:
                lo = mid
            else:
                hi = mid
        lower = lo

    if k == n:
        upper = ONE
    else:
        lo, hi = ZERO, ONE
        for _ in range(steps):
            mid = (lo + hi) / 2
            cdf = ONE - bin_ge(n, k + 1, mid)
            if cdf >= tail:
                lo = mid
            else:
                hi = mid
        upper = hi
    return lower, upper

def sqrt_upper(x, steps=40):
    assert isinstance(x, F) and 0 <= x <= 1
    assert isinstance(steps, int) and steps >= 1
    if x == 0:
        return ZERO
    lo, hi = ZERO, ONE
    for _ in range(steps):
        mid = (lo + hi) / 2
        if mid * mid >= x:
            hi = mid
        else:
            lo = mid
    return hi

def certificate_details(counts, alpha_s=F(1, 20),
                        bcal=ZERO, b=ZERO, steps=32,
                        sqrt_steps=40):
    assert len(counts) == 3
    assert isinstance(alpha_s, F) and 0 < alpha_s < 1
    assert isinstance(bcal, F) and bcal >= 0
    assert isinstance(b, F) and b >= 0
    endpoints = [
        cp_outward(n, k, alpha_s / 6, steps)
        for n, k in counts
    ]
    pl = sum((v[0] for v in endpoints), ZERO)
    pu = sum((v[1] for v in endpoints), ZERO)
    fl = clip01((pl - 1) / 2 - bcal)
    fu = clip01((pu - 1) / 2 + bcal)
    elo = max(ZERO, ONE - fu - b)
    ehi = min(ONE, b + 2 * sqrt_upper(ONE - fl, sqrt_steps))
    return endpoints, (fl, fu, elo, ehi)

def certify(counts, alpha_s=F(1, 20),
            bcal=ZERO, b=ZERO, steps=32):
    """Version-1-compatible four-endpoint interface."""
    return certificate_details(
        counts, alpha_s, bcal, b, steps
    )[1]

def accepts(fl, epsilon, b=ZERO):
    """Exact test of the mathematical certificate expression."""
    assert isinstance(fl, F) and 0 <= fl <= 1
    assert isinstance(epsilon, F) and 0 <= epsilon < 1
    assert isinstance(b, F) and b >= 0
    if epsilon < b:
        return False
    return 4 * (ONE - fl) <= (epsilon - b) ** 2

def bernoulli(rng, p):
    """Rational threshold; the PRNG execution is diagnostic."""
    assert isinstance(p, F) and 0 <= p <= 1
    return int(rng.randrange(p.denominator) < p.numerator)

# name, correlations, exact Bell overlap, implemented error
CASES = [
    ("identity", (ONE, -ONE, ONE), ONE, ZERO),
    ("pure replacement", (ZERO, ZERO, ZERO), F(1,4), ONE),
    ("mixed replacement", (ZERO, ZERO, ZERO), F(1,4), F(3,4)),
    ("dephasing", (F(4,5), -F(4,5), ONE), F(9,10), F(1,10)),
    ("rotation", (F(7,25), -F(7,25), ONE), F(16,25), F(3,5)),
    ("amplitude damping", (F(1,2), -F(1,2), F(1,4)),
     F(9,16), F(3,4)),
    ("completed loss", (F(4,5), -F(4,5), F(4,5)),
     F(17,20), F(3,20)),
]

def suite(seed=265, repetitions=20, shots=40,
          alpha_s=F(1,20), epsilon=F(1,10)):
    assert repetitions >= 1 and shots >= 1
    rng = Random(seed)
    signs = (1, -1, 1)
    total_covered = 0
    total_accepted = 0
    total_invalid = 0

    print("Python:", sys.version.replace("\n", " "))
    print("seed", seed, "repetitions", repetitions,
          "shots", shots, "alpha_s", alpha_s)
    print("bcal=0 b=0; CP steps=32; sqrt steps=40")
    print("acceptance uses exact squared threshold")

    for name, correlations, fidelity, error in CASES:
        inside = 0
        accepted_count = 0
        invalid_count = 0

        for rep in range(repetitions):
            counts = []
            for s, c in zip(signs, correlations):
                p = (ONE + s * c) / 2
                k = sum(bernoulli(rng, p) for _ in range(shots))
                counts.append((shots, k))

            endpoints, result = certificate_details(
                counts, alpha_s=alpha_s
            )
            fl, fu, elo, ehi = result
            covered = (fl <= fidelity <= fu and
                       elo <= error <= ehi)
            passed = accepts(fl, epsilon)

            inside += int(covered)
            accepted_count += int(passed)
            invalid_count += int(passed and error > epsilon)

            if rep == 0:
                print("\nCASE:", name)
                print("true F_e", fidelity, "true e", error)
                print("counts", counts)
                print("CP endpoints", endpoints)
                print("F_L", fl, "F_U", fu)
                print("e lower", elo, "e upper", ehi)
                print("covered", covered, "accepted", passed)

        total_covered += inside
        total_accepted += accepted_count
        total_invalid += invalid_count
        print("coverage", inside, "/", repetitions,
              "accepted", accepted_count,
              "invalid", invalid_count)

    datasets = len(CASES) * repetitions
    print("\nTOTAL datasets", datasets,
          "covered", total_covered,
          "accepted", total_accepted,
          "invalid", total_invalid)

def optional_checks():
    """New checks; execution is not claimed in the manuscript."""
    from math import comb

    # Verify the integer recurrence against direct binomial sums.
    for n in range(1, 9):
        for p in (ZERO, F(1,7), F(1,2), F(6,7), ONE):
            for k in range(n + 2):
                direct = sum(
                    (F(comb(n, i)) * p**i * (ONE-p)**(n-i)
                     for i in range(k, n + 1)),
                    ZERO
                )
                assert bin_ge(n, k, p) == direct

    # Check the reported all-pass arithmetic.
    endpoints, result = certificate_details([(40,40)] * 3)
    assert endpoints == [(F(1905242823,2147483648), ONE)] * 3
    assert result == (
        F(3568244821,4294967296), ONE, ZERO,
        F(28267292255,34359738368)
    )
    assert not accepts(result[0], F(1,10))

    # Exact integer check of the ideal all-pass sample boundary.
    assert 120 * 599**2870 > 600**2870
    assert 120 * 599**2871 < 600**2871

    # Exercise an informative acceptance regime using outward bounds.
    below = certify([(2870,2870)] * 3)
    above = certify([(2871,2871)] * 3)
    assert not accepts(below[0], F(1,10))
    assert accepts(above[0], F(1,10))
    assert above[3] <= F(1,10)

    # Check the reported square-root bracket directly.
    u = F(28267292255,34359738368)
    x = 4 * (ONE - F(3568244821,4294967296))
    assert (u - F(1,2**39))**2 < x <= u**2

    print("Optional arithmetic/threshold checks completed.")

if __name__ == "__main__":
    if "--checks" in sys.argv[1:]:
        optional_checks()
    suite()