g_math 0.6.6

#!/usr/bin/env python3
"""mpmath references for the ODE compute-tier state gate
(tests/data/ode_compute_state_refs.rs).

Each reference is the SAME discrete scheme the library runs (RK4,
Dormand-Prince 5th-order solution at a fixed step, Stormer-Verlet, the RK4
geodesic integrator, the Euler parallel transport with optional
re-orthogonalization) evaluated in mpmath at 120 digits: the exact-arithmetic
result of the scheme, so the difference the test measures is rounding alone,
not discretization error.

Problems have right-hand sides that are exact in fixed point (the harmonic
oscillator, y' = t - y, the flat metric) or cheap (the upper half-plane
H^2 with closed-form Christoffel symbols 1/y). Inputs and step sizes are
dyadic with at most 8 fraction bits, stored exactly at every gated split
(8 to 24 fraction bits) and on every wider profile; the geodesic steps are
T/N with N = 96 (not dyadic), which is what the library forms at the compute
tier. The test parses each literal with the exact literal parser, which rounds
it to nearest at the build's split, so one table serves every build.

Usage: python3 scripts/generate_ode_compute_state_refs.py
"""
from fractions import Fraction

from mpmath import mp, mpf

mp.dps = 140
DIGITS = 120

OUT = "tests/data/ode_compute_state_refs.rs"


def lit(v):
    """A decimal literal: exact for a dyadic Fraction, else 120 digits."""
    if isinstance(v, Fraction) and (v.denominator & (v.denominator - 1)) == 0:
        sign = "-" if v < 0 else ""
        v = abs(v)
        d = 0
        while (v * 10 ** d).denominator != 1:
            d += 1
        n = int(v * 10 ** d)
        if d == 0:
            return f"{sign}{n}"
        s = str(n).rjust(d + 1, "0")
        return f"{sign}{s[:-d]}.{s[-d:]}"
    if isinstance(v, Fraction):
        v = mpf(v.numerator) / v.denominator
    return mp.nstr(v, DIGITS, strip_zeros=False, min_fixed=-10**9, max_fixed=10**9)


def m(fr):
    return mpf(fr.numerator) / fr.denominator


def axpy(x, h, terms):
    """x + h * sum_j c_j k_j per component."""
    return [x[i] + h * sum(c * k[i] for c, k in terms) for i in range(len(x))]


def rk4_step(f, t, x, h):
    k1 = f(t, x)
    k2 = f(t + h / 2, axpy(x, h, [(mpf(1) / 2, k1)]))
    k3 = f(t + h / 2, axpy(x, h, [(mpf(1) / 2, k2)]))
    k4 = f(t + h, axpy(x, h, [(mpf(1), k3)]))
    return axpy(x, h, [(mpf(1) / 6, k1), (mpf(2) / 6, k2), (mpf(2) / 6, k3), (mpf(1) / 6, k4)])


def q(n, d):
    return mpf(n) / d


def dp45_x5(f, t, x, h):
    k1 = f(t, x)
    k2 = f(t + h * q(1, 5), axpy(x, h, [(q(1, 5), k1)]))
    k3 = f(t + h * q(3, 10), axpy(x, h, [(q(3, 40), k1), (q(9, 40), k2)]))
    k4 = f(t + h * q(4, 5), axpy(x, h, [(q(44, 45), k1), (q(-56, 15), k2), (q(32, 9), k3)]))
    k5 = f(t + h * q(8, 9), axpy(x, h, [
        (q(19372, 6561), k1), (q(-25360, 2187), k2), (q(64448, 6561), k3), (q(-212, 729), k4)]))
    k6 = f(t + h, axpy(x, h, [
        (q(9017, 3168), k1), (q(-355, 33), k2), (q(46732, 5247), k3), (q(49, 176), k4),
        (q(-5103, 18656), k5)]))
    return axpy(x, h, [
        (q(35, 384), k1), (q(500, 1113), k3), (q(125, 192), k4), (q(-2187, 6784), k5), (q(11, 84), k6)])


def osc(_t, x):
    return [x[1], -x[0]]


def forced(t, x):
    return [t - x[0]]


def h2_gamma(p):
    """Gamma[k][i][j] for the upper half-plane metric (1/y^2) I."""
    iy = 1 / p[1]
    g = [[[mpf(0)] * 2 for _ in range(2)] for _ in range(2)]
    g[0][0][1] = g[0][1][0] = -iy
    g[1][0][0] = iy
    g[1][1][1] = -iy
    return g


def geodesic_rhs(gamma):
    def f(_t, s):
        x, v = s[:2], s[2:]
        g = gamma(x)
        dv = [-sum(g[k][i][j] * v[i] * v[j] for i in range(2) for j in range(2)) for k in range(2)]
        return v + dv
    return f


def flat_gamma(_p):
    return [[[mpf(0)] * 2 for _ in range(2)] for _ in range(2)]


def every(points, stride):
    last = len(points) - 1
    return [(k, points[k]) for k in range(0, last + 1) if k % stride == 0 or k == last]


def table(name, rows, width, doc):
    out = [f"/// {doc}", f"pub const {name}: &[(usize, [&str; {width}])] = &["]
    for k, vals in rows:
        out.append(f"    ({k}, [" + ", ".join(f'"{lit(v)}"' for v in vals) + "]),")
    out.append("];\n")
    return "\n".join(out)


def main():
    parts = ["// Generated by scripts/generate_ode_compute_state_refs.py. Do not edit.\n"
             "// (step index, state) pairs: the exact-arithmetic result of the same scheme.\n"]

    # RK4 harmonic oscillator, h = 1/64, 256 steps
    h = mpf(1) / 64
    x, pts = [mpf(1), mpf(0)], []
    pts.append(x)
    for k in range(256):
        x = rk4_step(osc, k * h, x, h)
        pts.append(x)
    parts.append(table("RK4_OSC", every(pts, 16), 2,
                       "rk4_integrate x0' = x1, x1' = -x0 from (1, 0), h = 1/64, 256 steps."))

    # RK4 y' = t - y, h = 1/64, 256 steps
    x, pts = [mpf(1)], []
    pts.append(x)
    for k in range(256):
        x = rk4_step(forced, k * h, x, h)
        pts.append(x)
    parts.append(table("RK4_FORCED", every(pts, 16), 1,
                       "rk4_integrate y' = t - y from 1, h = 1/64, 256 steps."))

    # Dormand-Prince 5th-order solution, fixed h = 1/32, 128 steps
    h32 = mpf(1) / 32
    x, pts = [mpf(1), mpf(0)], []
    pts.append(x)
    for k in range(128):
        x = dp45_x5(osc, k * h32, x, h32)
        pts.append(x)
    parts.append(table("DP45_OSC", every(pts, 8), 2,
                       "rk45_integrate harmonic oscillator at a fixed h = 1/32 (no rejections), 128 steps."))

    # Stormer-Verlet, H = (p^2 + q^2) / 2, h = 1/64, 256 steps; state (q, p)
    qv, pv, pts = mpf(1), mpf(0), []
    pts.append([qv, pv])
    for _ in range(256):
        ph = pv + h / 2 * (-qv)
        qv = qv + h * ph
        pv = ph + h / 2 * (-qv)
        pts.append([qv, pv])
    parts.append(table("VERLET_OSC", every(pts, 16), 2,
                       "verlet_integrate H = (p^2 + q^2) / 2 from (q, p) = (1, 0), h = 1/64, 256 steps."))

    # Geodesics, T = 1, N = 96 (h = 1/96, not dyadic)
    n_steps = 96
    hg = mpf(1) / n_steps
    for name, gamma, x0, v0, doc in [
        ("GEODESIC_FLAT", flat_gamma, [q(1, 2), q(-1, 4)], [q(3, 4), q(5, 4)],
         "geodesic_integrate EuclideanMetric (dim 2) from (0.5, -0.25), v = (0.75, 1.25), T = 1, 96 steps."),
        ("GEODESIC_H2", h2_gamma, [mpf(0), mpf(1)], [q(3, 4), q(1, 2)],
         "geodesic_integrate HyperbolicMetric from (0, 1), v = (0.75, 0.5), T = 1, 96 steps."),
    ]:
        f = geodesic_rhs(gamma)
        s, pts = x0 + v0, [x0]
        for k in range(n_steps):
            s = rk4_step(f, k * hg, s, hg)
            pts.append(s[:2])
        parts.append(table(name, every(pts, 8), 2, doc))

    # Parallel transport (Euler along a given curve) on H^2, 64 segments
    def transport(curve, v, reorth):
        for step in range(len(curve) - 1):
            p, pn = curve[step], curve[step + 1]
            dx = [pn[i] - p[i] for i in range(2)]
            g = h2_gamma(p)
            v = [v[k] - sum(g[k][i][j] * v[i] * dx[j] for i in range(2) for j in range(2)) for k in range(2)]
            if reorth > 0 and (step + 1) % reorth == 0:
                nn = dx[0] * dx[0] + dx[1] * dx[1]
                if nn != 0:
                    c = (v[0] * dx[0] + v[1] * dx[1]) / nn
                    v = [v[i] - dx[i] * c for i in range(2)]
        return v

    horiz = [[q(k, 64), mpf(1)] for k in range(65)]
    diag = [[q(k, 64), 1 + q(k, 128)] for k in range(65)]
    v0 = [q(1, 2), q(3, 4)]
    rows = [
        (0, transport(horiz, v0, 0)),
        (1, transport(diag, v0, 0)),
        (2, transport(diag, v0, 16)),
    ]
    parts.append(table("TRANSPORT_H2", rows, 2,
                       "parallel_transport_ode HyperbolicMetric, V0 = (0.5, 0.75), 64 segments: "
                       "0 = y = 1, x = k/64; 1 = x = k/64, y = 1 + k/128; 2 = case 1 re-orthogonalized every 16."))

    with open(OUT, "w") as fh:
        fh.write("\n".join(parts))
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()