mu_kernel_arb.py

"""The immutable proof kernel for the mu upper bound (beta <= 2||R_m||_inf).

EVERYTHING that participates in producing a certified bound lives in this one
file: the imports that bind the interval-arithmetic provider, the working
precision and rounding constants, the exact conversions, Bernstein's F, the
witness context, the Lemma-M cell bound, the branch-and-bound, AND the domain driver (witness validation, the Lemma-T tail check, exact interval and split-piece construction, job dispatch, whole-run aggregation and outward emission -- pass-6 external review: a correct cell kernel proves nothing if the driver feeds it the wrong domains, so the driver lives inside the immutable unit).  The proof
package hashes THIS FILE'S COMPLETE BYTES as the kernel identity going
forward, and, for certificates whose shards were produced by the earlier
monolithic mu_certify_arb.py, the verifier proves every top-level definition,
constant, and import statement of this module identical to the producer's
source (pass-5 external review: a symbol-list closure guard that tracked only
assignments missed helper functions and import bindings; a whole-module hash
cannot).

Aggregation, sharding, manifests, and the CLI live in mu_certify_arb.py and
may evolve freely; this module changes only when the MATHEMATICS changes, and
any such change forces re-derivation of every certificate bound to it.
"""

from __future__ import annotations

import json
from decimal import Decimal
from fractions import Fraction

from flint import arb, ctx

import os
from multiprocessing import Pool

__all__ = ["arb", "ctx", "PREC", "CEIL_PLACES", "WIDTH_FLOOR", "_CTX",
           "_frac_to_arb", "dec_ceil_str", "_a", "_ball", "F_arb", "Ctx", "_init", "_ceil_frac", "_ukey", "_ukey_or_none", "_dyadic_below", "_bound", "_do_interval", "certify"]

PREC = 160
CEIL_PLACES = 40
WIDTH_FLOOR = Fraction(1, 10 ** 12)


def _frac_to_arb(fr: Fraction) -> arb:
    return arb(int(fr.numerator)) / arb(int(fr.denominator))


def dec_ceil_str(fr: Fraction, places: int) -> str:
    """Decimal string, rounded UP, so the printed endpoint is itself proved."""
    scaled = fr * 10 ** places
    n = -((-scaled.numerator) // scaled.denominator)      # ceiling division
    sign = "-" if n < 0 else ""
    n = abs(n)
    d = str(n).rjust(places + 1, "0")
    return f"{sign}{d[:-places]}.{d[-places:]}" if places else f"{sign}{d}"


def _a(fr: Fraction) -> arb:
    return arb(int(fr.numerator)) / arb(int(fr.denominator))


def _ball(p: Fraction, q: Fraction) -> arb:
    """A ball covering the whole interval [p, q].

    Built with ``union`` of the two exact endpoints, so the enclosure is
    rigorous; constructing it from a float radius could round the radius *down*
    and lose points of the cell.
    """
    return _a(p).union(_a(q))


def F_arb(t: arb) -> arb:
    """Bernstein's F(t) = (t/2)[psi(t/2+3/4) - psi(t/2+1/4)] (Arb: rigorous)."""
    h = t / 2
    return h * ((h + arb(3) / 4).digamma() - (h + arb(1) / 4).digamma())


class Ctx:
    """Witness data, prepared once per process."""

    def __init__(self, path):
        with open(path) as fh:
            w = json.load(fh)
        self.m = int(w["m"])
        self.a = [Fraction(int(p[0]), int(p[1])) for p in w["a"]]
        self.c = [Fraction(2 * k - 1, 2) for k in range(1, self.m + 1)]
        self.lam = Fraction(1, 2) - self.a[0]
        self.a_arb = [_a(x) for x in self.a]
        self.c2_arb = [_a(x * x) for x in self.c]
        self.c_arb = [_a(x) for x in self.c]
        self.all_pos = all(x > 0 for x in self.a[1:])

    def H(self, p: Fraction, skip=None) -> arb:
        t = _a(p)
        t2 = t * t
        s = F_arb(t) - self.a_arb[0]
        for i in range(self.m):
            if i == skip:
                continue
            s = s - self.a_arb[i + 1] / (t2 - self.c2_arb[i])
        return s


_CTX = None

def _init(path, prec):
    global _CTX
    ctx.prec = prec
    _CTX = Ctx(path)


def _ceil_frac(x: arb, places: int = CEIL_PLACES, maxstep: int = 64) -> Fraction:
    """Smallest ``n/10**places`` *certified* to be >= every point of the ball ``x``.

    The Arb -> decimal boundary is where a computer-assisted proof is easiest to
    break: ``float(x)`` rounds to NEAREST and can land below the ball's upper
    endpoint, and ``repr`` of that float is the shortest round-tripping decimal,
    which can be lower still.  Either one silently converts "we proved U" into
    "we claim something smaller than U".

    This uses only Arb *comparisons*, which are conservative -- ``a >= b`` is
    true only when it holds for every point of both balls -- so the returned
    rational is a proved upper bound, never an approximation of one.
    """
    P = 10 ** places
    Pa = arb(P)
    # Seed from Arb's own decimal string.  The seed only has to be close: the
    # loop below is what certifies, via a conservative Arb comparison.  (A float
    # seed cannot work -- binary64 carries ~16 digits and `places` may exceed
    # that by decades.)
    seed = Fraction(Decimal(x.str(places + 25, radius=False)))
    n = (seed * P).__floor__() - 2
    for _ in range(maxstep):
        if arb(n) / Pa >= x:
            return Fraction(n, P)
        n += 1
    raise RuntimeError(f"no certified ceiling for {x!r} at {places} places")


def _ukey(x: arb) -> arb:
    """Exact (radius-0), FINITE upper endpoint of a ball -- the ONLY thing
    allowed into heapq / min / max / target comparisons.

    Arb comparisons are a PARTIAL order: overlapping non-exact balls return
    False for both `<` directions, so heapq's root need not be the frontier
    maximum and Python's max() can return a non-maximal ball (found by external
    review; the P0 defect).  Exact zero-radius values restore a total order.

    Radius zero is NOT sufficient on its own -- a second review point.  In
    python-flint 0.6.0 both ``arb('nan').upper()`` and ``arb('inf').upper()``
    report ``rad() == 0``, yet NaN compares False against everything and would
    silently destroy the ordering exactly the way overlapping balls did.  The
    finiteness test is what actually makes the key totally ordered.
    """
    u = x.upper()
    if not u.rad() == 0:
        raise AssertionError("upper() returned a non-exact value; cannot order")
    if not u.is_finite():
        raise AssertionError(f"non-finite ordering key {u!r} (NaN/Inf pass rad()==0)")
    return u


def _ukey_or_none(x: arb):
    """``_ukey(x)`` when x is an exact finite key, else None (meaning +infinity).

    ``H`` is genuinely infinite at a pole, so "no finite key" is a legitimate
    answer here and must be represented, not swallowed.
    """
    u = x.upper()
    if not u.rad() == 0 or not u.is_finite():
        return None
    return u


def _dyadic_below(fr: Fraction, bits: int = 128) -> arb:
    """Exact (radius-0) dyadic ``d`` with ``d <= fr``.

    The branch-and-bound target is a decimal rational, which is generally NOT
    dyadic, so ``_frac_to_arb(target)`` returns a BALL.  Comparing an exact cell
    key against a ball reintroduces the partial order the ``_ukey`` fix removed:
    when they overlap, ``best > target`` and ``-negb <= target`` are BOTH False.
    The old behaviour was conservative (an undecided comparison means "keep
    refining" and "also try the pole-corrected branch"), so no invalid bound
    could be emitted -- but the stated invariant "all ordered comparisons are
    exact" was false, and the G7 gates did not cover it.

    Rounding the target DOWN to a dyadic restores exactness in the safe
    direction: ``best <= d`` implies ``best <= target``, and a smaller target
    only ever causes more refinement.
    """
    n = (fr * (1 << bits)).__floor__()
    d = arb(int(n)) / arb(1 << bits)
    if not d.rad() == 0 or not d.is_finite():
        raise AssertionError("dyadic target is not an exact finite value")
    if not Fraction(int(n), 1 << bits) <= fr:
        raise AssertionError("dyadic target is not below the requested target")
    return d


def _bound(p, q, kl, kr, target, cache):
    """EXACT upper bound on |R_m| over [p,q] (inside [c_kl, c_kr]); Lemma M.

    Returns ``None`` to mean "+infinity": no finite endpoint bound is available
    for this cell and it MUST be bisected before it can be bounded at all.

    That case is real and was previously mishandled.  Every pole-to-pole root
    cell [c_{k}, c_{k+1}] has a pole at BOTH endpoints, so whichever pole the
    correction removes, ``H`` is still infinite at the other end.  Arb returns
    NaN there, and NaN compares False against everything -- so the old
    ``max(|H(p)|, |H(q)|)`` silently returned the finite endpoint and produced a
    number that bounds nothing.  (It did not corrupt the published m=8000 run:
    those bogus root values all landed ~0.15, well above the 0.14008474995
    target, so every root cell was popped and bisected, and a cell can only
    survive in the final frontier if its key is <= target.  But the guard was a
    property of the numbers, not of the code.)  Now such a cell is ranked ahead
    of every finite cell and forced to bisect; its children touch at most one
    pole each and bound normally.
    """
    W = _CTX
    cm = (arb.pi() * _ball(p, q)).cos().abs_upper()   # |cos(pi t)| on the cell

    def Hv(x, skip):
        key = (x, skip)
        v = cache.get(key)
        if v is None:
            v = W.H(x, skip=skip)
            cache[key] = v
        return v

    def endmax(skip):
        """Exact max of |H| over the two endpoints with pole ``skip`` removed,
        or None if either endpoint is still a pole of the remaining sum."""
        best = None
        for x in (p, q):
            u = _ukey_or_none(Hv(x, skip).abs_upper())
            if u is None:
                return None
            if best is None or u > best:
                best = u
        return best

    cands = []
    at_pole = (kl is not None and p == W.c[kl]) or (kr is not None and q == W.c[kr])
    if not at_pole:
        mh = endmax(None)
        if mh is not None:
            cands.append(_ukey(cm * mh))
    if not cands or cands[0] > target:
        for k in (kl, kr):
            if k is None:
                continue
            mh = endmax(k)
            if mh is None:                 # the other endpoint is a pole
                continue
            extra = (W.a_arb[k + 1] * arb.pi() / (_a(p) + W.c_arb[k])).abs_upper()
            cands.append(_ukey(cm * mh + extra))
    if not cands:
        return None
    best = cands[0]                        # min over independently valid bounds
    for c in cands[1:]:
        if c < best:
            best = c
    return best


def _do_interval(job):
    """Branch and bound one pole interval; returns its certified sup bound.

    Heap entries are ``(rank, key, n, p, q, kl, kr)``.  ``rank`` is 0 for a cell
    with no finite bound yet and 1 otherwise, so an unbounded cell always sorts
    ahead of every bounded one and can never be mistaken for the frontier
    maximum.  Within a rank the keys are exact finite arbs (or the constant 0),
    so every comparison the heap makes is a total order.
    """
    import heapq
    p0, q0, kl, kr, target, max_cells = job
    # exact dyadic <= target, so EVERY comparison below is exact-vs-exact
    target = _dyadic_below(target) if isinstance(target, Fraction) else target
    cache = {}
    heap = []
    n = 0

    def push(b, p, q, l, r):
        nonlocal n
        n += 1
        heapq.heappush(heap, (0, 0, n, p, q, l, r) if b is None
                       else (1, -b, n, p, q, l, r))

    push(_bound(p0, q0, kl, kr, target, cache), p0, q0, kl, kr)
    cells = 1
    hit_floor = False
    while heap and cells < max_cells:
        rank, key, _, p, q, a_kl, a_kr = heap[0]
        if rank == 1 and -key <= target:
            break
        if q - p < WIDTH_FLOOR:
            # Stopped refining because the cell got too narrow, NOT because the
            # frontier reached the target.  The returned bound is still valid,
            # but the interval did not converge and the caller must say so:
            # `converged` used to mean only "did not hit the cell cap", which
            # silently blessed this case (external review).
            hit_floor = True
            break
        heapq.heappop(heap)
        mid = (p + q) / 2
        for (pp, qq, l, r) in ((p, mid, a_kl, None), (mid, q, None, a_kr)):
            push(_bound(pp, qq, l, r, target, cache), pp, qq, l, r)
            cells += 1
    unbounded = [h for h in heap if h[0] == 0]
    if unbounded:
        # Refusing to return a number here is the point: a cell with no finite
        # bound makes the interval's sup unproved, so there is nothing to report.
        raise RuntimeError(
            f"interval [{p0},{q0}]: {len(unbounded)} cell(s) still have no finite "
            f"bound after {cells} cells (width floor {WIDTH_FLOOR}, cap {max_cells}); "
            f"the sup is NOT certified")
    worst = None
    for h in heap:                          # exact finite keys: a total order
        v = -h[1]
        if worst is None or v > worst:
            worst = v
    # Hand back an exact rational that is CERTIFIED >= worst.  Returning
    # float(worst) would round to nearest and can land below the Arb bound,
    # turning a proof of "sup <= U" into a claim of something smaller than U.
    return _ceil_frac(worst), cells, str(worst), hit_floor



def certify(path, target, T0=None, max_cells=400_000, procs=None, prec=PREC,
            verbose=True, out_places=20, interval_range=None, split=None,
            partial=False):
    """Prove ||R_m||_inf <= U on [0,inf).

    Returns ``(two_u_str, info)`` where ``two_u_str`` is a decimal rounded
    **up**, so the printed endpoint is itself a proved bound on ``beta``.  Every
    stage from the Arb balls to that string is exact or outward: worker results
    come back as certified rationals, the max and the doubling are exact, and
    the final decimal is a ceiling.  ``info["endpoint_verified"]`` records the
    regression check that re-parsing the emitted string still dominates the
    exact bound.
    """
    ctx.prec = prec
    # a partial run must be DECLARED as such: it returns a shard payload and
    # never a global certificate (the aggregation wrapper sets `partial` when
    # given a shard path; the message keeps the wording the gates assert on)
    if (interval_range is not None or split is not None) and not partial:
        raise ValueError(
            "partial certification (interval_range/split) must write a shard via "
            "shard_out; a partial run cannot return a global certificate")
    path = os.path.abspath(path)   # one canonical path: opened AND hashed (review)
    target = Fraction(Decimal(str(target))) if not isinstance(target, Fraction) else target
    W = Ctx(path)
    if not W.all_pos:
        raise ValueError("witness fails Lemma T/M hypothesis: some a_k <= 0")
    if W.lam < 0:
        raise ValueError("witness fails Lemma T hypothesis: a_0 > 1/2")
    if len(W.a) != W.m + 1:
        raise ValueError(f"witness has {len(W.a)} coefficients, need m+1 = {W.m + 1}")
    m = W.m
    if T0 is None:
        T0 = Fraction(m)
    if not T0 > W.c[-1]:
        raise ValueError(f"T0 = {T0} must exceed the last pole c_m = {W.c[-1]}")
    # ---- Lemma T ----
    HT0 = W.H(T0)
    lam_hi = _a(W.lam).abs_upper()
    if not HT0.lower() >= -lam_hi:
        raise ValueError("Lemma T check failed at T0")
    if verbose:
        print(f"[Lemma T] H({T0}) >= {HT0.lower()} >= -lambda; tail <= {lam_hi}",
              flush=True)
    # ---- finite part, one job per pole interval (independent by Lemma M) ----
    edges = [Fraction(0)] + list(W.c) + [T0]
    jobs = []
    for j in range(len(edges) - 1):
        kl = j - 1 if j - 1 >= 0 else None
        kr = j if j < m else None
        jobs.append((edges[j], edges[j + 1], kl, kr, target, max_cells))
    n_intervals = len(jobs)
    subrange = None
    if split is not None:
        # One pole interval, sliced into N equal rational pieces, piece I.
        # Lemma M gives monotonicity BETWEEN poles, so any exact sub-division of
        # a pole interval is legitimate: sup over the interval is the max of the
        # sups over the pieces, exactly as it is across intervals.  This matters
        # for scheduling, not just for restarts -- the first gap [0, c_1) carries
        # two alternation points and needs far more cells than any other, so as
        # a single job it becomes the serial critical path of the whole run.
        k, i, n = split
        if not (0 <= k < n_intervals and 0 <= i < n and n >= 1):
            raise ValueError(f"bad split {split} for {n_intervals} intervals")
        p_lo, p_hi, kl, kr, _, _ = jobs[k]
        step = (p_hi - p_lo) / n
        p = p_lo + step * i
        q = p_hi if i == n - 1 else p_lo + step * (i + 1)
        jobs = [(p, q, kl if i == 0 else None, kr if i == n - 1 else None,
                 target, max_cells)]
        subrange = {"interval": k, "piece": i, "pieces": n,
                    "p": str(p), "q": str(q)}
        interval_range = (k, k + 1)
    elif interval_range is not None:
        lo, hi = interval_range
        if not (0 <= lo < hi <= n_intervals):
            raise ValueError(f"interval_range {interval_range} outside [0,{n_intervals}]")
        jobs = jobs[lo:hi]
    procs = procs or min(os.cpu_count() or 4, 10)
    with Pool(procs, initializer=_init, initargs=(path, prec)) as pool:
        out = pool.map(_do_interval, jobs, chunksize=4)
    if partial:
        lo, hi = interval_range if interval_range is not None else (0, n_intervals)
        shard = {
            "kind": "mu_certify_arb shard",
            "m": m, "T0": str(T0), "prec": prec, "max_cells": max_cells,
            "target_num": str(target.numerator), "target_den": str(target.denominator),
            "interval_lo": lo, "interval_hi": hi, "n_intervals_total": n_intervals,
            "subrange": subrange,
            "bounds": [[str(o[0].numerator), str(o[0].denominator)] for o in out],
            "cells": [o[1] for o in out],
            "floored": [bool(o[3]) for o in out],
        }
        return None, shard
    if interval_range is not None or split is not None:
        # unreachable given the entry guard; kept so the invariant is enforced
        # where the damage would happen, not only where the argument arrives
        raise AssertionError("internal: partial run reached the global aggregation path")
    worst = max(o[0] for o in out)                 # exact Fractions, each >= its ball
    cells = sum(o[1] for o in out)
    capped = [j for j, o in enumerate(out) if o[1] >= max_cells]
    floored = [j for j, o in enumerate(out) if o[3]]
    tail_f = _ceil_frac(lam_hi)
    U = max(worst, tail_f)                         # exact
    two_u = 2 * U                                  # exact (dyadic-free)
    two_u_str = dec_ceil_str(two_u, out_places)
    if Fraction(two_u_str) < two_u:                # must never fire
        raise AssertionError(f"emitted decimal {two_u_str} is below the proved bound")
    if verbose:
        print(f"[finite]  sup on [0,{T0}] <= {dec_ceil_str(worst, out_places)}  "
              f"({cells} cells, {procs} procs)", flush=True)
        if capped:
            print(f"[WARNING] {len(capped)} interval(s) hit max_cells={max_cells}: "
                  f"{capped[:8]}{'...' if len(capped) > 8 else ''} -- the bound is "
                  f"valid but NOT converged; raise max_cells", flush=True)
    return two_u_str, {
        "m": m, "cells": cells, "procs": procs, "prec": prec, "T0": str(T0),
        "max_cells": max_cells,
        "target": dec_ceil_str(target, out_places),
        "two_U_exact_num": str(two_u.numerator), "two_U_exact_den": str(two_u.denominator),
        "U_finite": dec_ceil_str(worst, out_places),
        "U_tail": dec_ceil_str(tail_f, out_places),
        "intervals_hitting_max_cells": capped,
        "intervals_hitting_width_floor": floored,
        "width_floor": str(WIDTH_FLOOR),
        # `converged` means the branch-and-bound frontier actually reached the
        # target on EVERY interval: neither the cell cap nor the width floor
        # cut it short.  (It previously meant only "no cell cap", which blessed
        # width-floor stops -- external review.)
        "converged": not capped and not floored,
        "endpoint_verified": True,
    }