The speedrun's training text

CORPUS_N = 32768
CORPUS_VOCAB = 88

#| The training text, as character codes.
def corpus() {
    return chars("#! use lib/linalg\n\nSCHUR_EPS = 2.220446049250313e-16\n\n#| Apply the Householder reflector built from `v` to the window of `m` rows and\n#| columns starting at `p`: `H` from both sides, so its eigenvalues do not move,\n#| and `Q` from the right, so that `A = Q H Q^T` still holds afterwards.\n#|\n#| The two ranges are not an optimization. From the left it touches columns `c0`\n#| and rightwards, and from the right it touches rows `r1` and upwards, because\n#| everything outside those is already zero and must stay zero -- a reflector\n#| applied to the whole row would write dust back into a subdiagonal entry that\n#| deflation had just decided was negligible, and the factor would never finish\n#| being quasi-triangular.\ndef reflect(H, Q, v, p, m, c0, r1) {\n    n = len(H)\n    vn = 0.0\n    for i to m {\n        vn = vn + v[i] * v[i]\n    }\n    if vn == 0.0 {\n        return 0.0\n    }\n    for jj to n - c0 {\n        j = c0 + jj\n        s = 0.0\n        for i to m {\n            s = s + v[i] * H[p + i, j]\n        }\n        s = 2.0 * s / vn\n        for i to m {\n            H[p + i, j] = H[p + i, j] - s * v[i]\n        }\n    }\n    for i to r1 + 1 {\n        s = 0.0\n        for j to m {\n            s = s + H[i, p + j] * v[j]\n        }\n        s = 2.0 * s / vn\n        for j to m {\n            H[i, p + j] = H[i, p + j] - s * v[j]\n        }\n    }\n    for i to n {\n        s = 0.0\n        for j to m {\n            s = s + Q[i, p + j] * v[j]\n        }\n        s = 2.0 * s / vn\n        for j to m {\n            Q[i, p + j] = Q[i, p + j] - s * v[j]\n        }\n    }\n    return 0.0\n}\n\n#| Apply the plane rotation `[[cs, -sn], [sn, cs]]` on rows and columns `k` and\n#| `k + 1` the same way: `T` from both sides, `Q` from the right.\n#|\n#| This is only ever called on a block that has already deflated, so everything\n#| left of column `k` in those two rows, and everything below row `k + 1` in\n#| those two columns, is zero. It stays out of both. A sweep leaves dust of\n#| order 1e-17 below the subdiagonal rather than exact zeros, and a rotation\n#| over the full row would mix that dust straight into the subdiagonal entry\n#| deflation had just set to zero, turning a converged 2 by 2 into two blocks\n#| that no longer look converged.\ndef rotate(T, Q, k, cs, sn) {\n    n = len(T)\n    for jj to n - k {\n        j = k + jj\n        t0 = T[k, j]\n        t1 = T[k + 1, j]\n        T[k, j] = cs * t0 + sn * t1\n        T[k + 1, j] = cs * t1 - sn * t0\n    }\n    for i to k + 2 {\n        t0 = T[i, k]\n        t1 = T[i, k + 1]\n        T[i, k] = cs * t0 + sn * t1\n        T[i, k + 1] = cs * t1 - sn * t0\n    }\n    for i to n {\n        t0 = Q[i, k]\n        t1 = Q[i, k + 1]\n        Q[i, k] = cs * t0 + sn * t1\n        Q[i, k + 1] = cs * t1 - sn * t0\n    }\n    return 0.0\n}\n\n#| Reduce `A` to upper Hessenberg form. Fills `H` and `Q` so that `A = Q H Q^T`,\n#| `Q^T Q = I`, and `H[i, j] = 0` whenever `i > j + 1`. One reflection per\n#| column, each one zeroing everything below the subdiagonal of that column and\n#| leaving the columns already done alone.\ndef hessenberg(A, H, Q) {\n    n = len(A)\n    for i to n {\n        for j to n {\n            H[i, j] = A[i, j]\n            Q[i, j] = 0.0\n        }\n        Q[i, i] = 1.0\n    }\n    for k to n - 2 {\n        m = n - k - 1\n        v = zeros(m)\n        nrm = 0.0\n        for i to m {\n            v[i] = H[k + 1 + i, k]\n            nrm = nrm + v[i] * v[i]\n        }\n        nrm = sqrt(nrm)\n        if nrm > 0.0 {\n            if v[0] < 0.0 {\n                nrm = -nrm\n            }\n            v[0] = v[0] + nrm\n            junk = reflect(H, Q, v, k + 1, m, k, n - 1)\n        }\n    }\n    for i to n {\n        for j to n {\n            if i > j + 1 {\n                H[i, j] = 0.0\n            }\n        }\n    }\n    return 0.0\n}\n\n#| One implicit double-shift QR sweep over the unreduced Hessenberg block from\n#| `lo` to `hi`, for a pair of shifts given by their sum `s` and product `t`.\n#| The shifts are never formed: the first reflector is built from the first\n#| column of `(H - s1)(H - s2)`, which expands to something needing only `s` and\n#| `t`, and that is what \"implicit\" means. The rest of the sweep chases the\n#| resulting bulge back down to Hessenberg form.\ndef francis_step(H, Q, lo, hi, s, t) {\n    x = H[lo, lo] * H[lo, lo] + H[lo, lo + 1] * H[lo + 1, lo] - s * H[lo, lo] + t\n    y = H[lo + 1, lo] * (H[lo, lo] + H[lo + 1, lo + 1] - s)\n    z = H[lo + 2, lo + 1] * H[lo + 1, lo]\n    span = hi - lo - 1\n    for kk to span {\n        p = lo + kk\n        v = zeros(3)\n        v[0] = x\n        v[1] = y\n        v[2] = z\n        nrm = sqrt(x * x + y * y + z * z)\n        if nrm > 0.0 {\n            if v[0] < 0.0 {\n                nrm = -nrm\n            }\n            v[0] = v[0] + nrm\n            junk = reflect(H, Q, v, p, 3, max(p - 1, lo), min(p + 3, hi))\n        }\n        x = H[p + 1, p]\n        y = H[p + 2, p]\n        z = 0.0\n        if kk < span - 1 {\n            z = H[p + 3, p]\n        }\n    }\n    v2 = zeros(2)\n    v2[0] = x\n    v2[1] = y\n    nrm = sqrt(x * x + y * y)\n    if nrm > 0.0 {\n        if v2[0] < 0.0 {\n            nrm = -nrm\n        }\n        v2[0] = v2[0] + nrm\n        junk = reflect(H, Q, v2, hi - 1, 2, max(hi - 2, lo), hi)\n    }\n    return 0.0\n}\n\n#| Put the 2 by 2 block at `k` into standard form. If its two eigenvalues are\n#| real it is made upper triangular, which splits it into two 1 by 1 blocks and\n#| returns 1.0. If they are a complex pair the two diagonal entries are rotated\n#| until they are equal, which is the form a real Schur factor is supposed to\n#| have, and it returns 0.0.\ndef split2(T, Q, k) {\n    a = T[k, k]\n    b = T[k, k + 1]\n    c = T[k + 1, k]\n    d = T[k + 1, k + 1]\n    if c == 0.0 {\n        return 1.0\n    }\n    p = 0.5 * (a - d)\n    disc = p * p + b * c\n    if disc >= 0.0 {\n        r = sqrt(disc)\n        # take the root on the same side as p: then lam - d never cancels\n        lam = 0.5 * (a + d) + r\n        if p < 0.0 {\n            lam = 0.5 * (a + d) - r\n        }\n        x = lam - d\n        y = c\n        h = sqrt(x * x + y * y)\n        if h > 0.0 {\n            junk = rotate(T, Q, k, x / h, y / h)\n        }\n        T[k + 1, k] = 0.0\n        return 1.0\n    }\n    th = 0.5 * atan2(d - a, b + c)\n    junk = rotate(T, Q, k, cos(th), sin(th))\n    return 0.0\n}\n\n#| The Frobenius norm, used as the yardstick a subdiagonal entry is called\n#| negligible against when its two diagonal neighbors are both zero.\ndef schur_norm(A) {\n    n = len(A)\n    total = 0.0\n    for i to n {\n        for j to n {\n            total = total + A[i, j] * A[i, j]\n        }\n    }\n    return sqrt(total)\n}\n\n#| One step of the outer iteration on the block ending at `hi`, and the whole of\n#| what the driver does per turn: walk up from `hi` to the first negligible\n#| subdiagonal entry, then either deflate what is below it or sweep it. Returns\n#| the new `hi`, which is unchanged exactly when a sweep happened.\n#|\n#| `its` is how many sweeps have already been spent here without deflating. A\n#| sweep that keeps landing in the same place is stalling on a shift that is\n#| useless for this matrix, so every tenth one throws an arbitrary shift of the\n#| right size at it instead. The constants are the ones LAPACK settled on.\ndef schur_advance(T, Q, hi, its, anorm) {\n    lo = hi\n    while lo > 0 {\n        small = SCHUR_EPS * (abs(T[lo - 1, lo - 1]) + abs(T[lo, lo]))\n        if small == 0.0 {\n            small = SCHUR_EPS * anorm\n        }\n        if abs(T[lo, lo - 1]) <= small {\n            T[lo, lo - 1] = 0.0\n            break\n        }\n        lo = lo - 1\n    }\n    if lo == hi {\n        return hi - 1\n    }\n    if lo == hi - 1 {\n        junk = split2(T, Q, lo)\n        return lo - 1\n    }\n    m = hi - 1\n    s = T[m, m] + T[hi, hi]\n    t = T[m, m] * T[hi, hi] - T[m, hi] * T[hi, m]\n    if its > 0.0 {\n        if mod(its, 10.0) == 0.0 {\n            sc = abs(T[hi, hi - 1]) + abs(T[hi - 1, hi - 2])\n            h11 = 0.75 * sc + T[hi, hi]\n            h12 = -0.4375 * sc\n            s = h11 + h11\n            t = h11 * h11 - h12 * sc\n        }\n    }\n    junk = francis_step(T, Q, lo, hi, s, t)\n    return hi\n}\n\n#| The real Schur decomposition. Fills `T` and `Q` so that `A = Q T Q^T` with\n#| `Q` orthogonal and `T` quasi-triangular. Reduce to Hessenberg once, then\n#| sweep and deflate from the bottom right until nothing is left.\ndef schur(A, T, Q) {\n    n = len(A)\n    junk = hessenberg(A, T, Q)\n    anorm = schur_norm(A)\n    hi = n - 1\n    its = 0.0\n    budget = 60 * n + 200\n    while hi > 0 {\n        budget = budget - 1\n        if budget <= 0 {\n            break\n        }\n        nexthi = schur_advance(T, Q, hi, its, anorm)\n        if nexthi == hi {\n            its = its + 1\n        }\n        else {\n            its = 0.0\n        }\n        hi = nexthi\n    }\n    for i to n {\n        for j to n {\n            if i > j + 1 {\n                T[i, j] = 0.0\n            }\n        }\n    }\n    return 0.0\n}\n\n#| Where each diagonal block starts, written into `starts`, and how many there\n#| are, returned. A 2 by 2 block is one entry, not two: `T[k + 1, k]` is nonzero\n#| exactly when `k` heads a 2 by 2.\ndef schur_blocks(T, starts) {\n    n = len(T)\n    k = 0\n    count = 0\n    while k < n {\n        starts[count] = k\n        count = count + 1\n        if k + 1 < n {\n            if T[k + 1, k] != 0.0 {\n                k = k + 1\n            }\n        }\n        k = k + 1\n    }\n    return count\n}\n\n#| The real part of the eigenvalues of the block starting at `k`. Both of them\n#| for a 2 by 2, and both are the same number there, because `split2` made the\n#| diagonal entries equal.\ndef schur_block_re(T, k) {\n    return T[k, k]\n}\n\n#| How far `A - Q T Q^T` is from zero, entry by entry: the decomposition's own\n#| certificate, and the number to look at before believing anything downstream.\ndef schur_residual(A: Arr, T: Arr, Q: Arr) {\n    n = len(A)\n    back = Q @ T @ transpose(Q)\n    worst = 0.0\n    for i to n {\n        for j to n {\n            worst = max(worst, abs(back[i, j] - A[i, j]))\n        }\n    }\n    return worst\n}\n\n#| How far `Q^T Q` is from the identity.\ndef orthogonality(Q: Arr) {\n    n = len(Q)\n    g = transpose(Q) @ Q\n    worst = 0.0\n    for i to n {\n        for j to n {\n            expect = 0.0\n            if i == j {\n                expect = 1.0\n            }\n            worst = max(worst, abs(g[i, j] - expect))\n        }\n    }\n    return worst\n}\n\n#| Solve `T11 X - X T22 = C`, where `T11` and `T22` are the two adjacent\n#| diagonal blocks at `k`, by writing the whole thing out as one linear system\n#| in the entries of `X`. With both blocks of size 1 or 2 that system has at\n#| most four unknowns, so there is no reason to be clever about it. It is\n#| singular exactly when the two blocks share an eigenvalue, which is the one\n#| case where the two cannot be swapped.\ndef sylvester_small(T, k, n1, n2, C, X) {\n    nn = n1 * n2\n    A = zeros(nn, nn)\n    rhs = zeros(nn)\n    for i to n1 {\n        for j to n2 {\n            u = i * n2 + j\n            for a to n1 {\n                A[u, a * n2 + j] = A[u, a * n2 + j] + T[k + i, k + a]\n            }\n            for b to n2 {\n                A[u, i * n2 + b] = A[u, i * n2 + b] - T[k + n1 + b, k + n1 + j]\n            }\n            rhs[u] = C[i, j]\n        }\n    }\n    sol = zeros(nn)\n    lin_solve(A, rhs, sol)\n    for i to n1 {\n        for j to n2 {\n            X[i, j] = sol[i * n2 + j]\n        }\n    }\n    return 0.0\n}\n\n#| Householder QR of the `m` by `p` matrix `V`, throwing away the triangular\n#| factor and keeping the orthogonal one: fills `Z` with an `m` by `m`\n#| orthogonal matrix whose first `p` columns span the same space `V`'s do.\n#| `V` is overwritten.\ndef orth_basis(V, Z, m, p) {\n    for i to m {\n        for j to m {\n            Z[i, j] = 0.0\n        }\n        Z[i, i] = 1.0\n    }\n    for j to p {\n        mv = m - j\n        v = zeros(mv)\n        nrm = 0.0\n        for i to mv {\n            v[i] = V[j + i, j]\n            nrm = nrm + v[i] * v[i]\n        }\n        nrm = sqrt(nrm)\n        if nrm > 0.0 {\n            if v[0] < 0.0 {\n                nrm = -nrm\n            }\n            v[0] = v[0] + nrm\n            vn = 0.0\n            for i to mv {\n                vn = vn + v[i] * v[i]\n            }\n            if vn > 0.0 {\n                for cc to p - j {\n                    c = j + cc\n                    sm = 0.0\n                    for i to mv {\n                        sm = sm + v[i] * V[j + i, c]\n                    }\n                    sm = 2.0 * sm / vn\n                    for i to mv {\n                        V[j + i, c] = V[j + i, c] - sm * v[i]\n                    }\n                }\n                for i to m {\n                    sm = 0.0\n                    for a to mv {\n                        sm = sm + Z[i, j + a] * v[a]\n                    }\n                    sm = 2.0 * sm / vn\n                    for a to mv {\n                        Z[i, j + a] = Z[i, j + a] - sm * v[a]\n                    }\n                }\n            }\n        }\n    }\n    return 0.0\n}\n\n#| Apply the small orthogonal `Z` as a similarity on the window of `m` rows and\n#| columns starting at `k`, and on `Q` from the right. Same reasoning about\n#| ranges as `reflect`: outside the window everything is already zero.\ndef apply_block(T, Q, Z, k, m) {\n    n = len(T)\n    for cc to n - k {\n        c = k + cc\n        col = zeros(m)\n        for i to m {\n            sm = 0.0\n            for a to m {\n                sm = sm + Z[a, i] * T[k + a, c]\n            }\n            col[i] = sm\n        }\n        for i to m {\n            T[k + i, c] = col[i]\n        }\n    }\n    for r to k + m {\n        row = zeros(m)\n        for j to m {\n            sm = 0.0\n            for a to m {\n                sm = sm + T[r, k + a] * Z[a, j]\n            }\n            row[j] = sm\n        }\n        for j to m {\n            T[r, k + j] = row[j]\n        }\n    }\n    for r to n {\n        row = zeros(m)\n        for j to m {\n            sm = 0.0\n            for a to m {\n                sm = sm + Q[r, k + a] * Z[a, j]\n            }\n            row[j] = sm\n        }\n        for j to m {\n            Q[r, k + j] = row[j]\n        }\n    }\n    return 0.0\n}\n\n#| Exchange the two adjacent diagonal blocks at `k`, of sizes `n1` then `n2`, so\n#| that afterwards `n2` comes first. The eigenvalues do not change; only where\n#| they sit on the diagonal does, and `Q` follows.\n#|\n#| The trick is that the columns of `[X; I]` span the invariant subspace of the\n#| trailing block, where `X` solves `T11 X - X T22 = -T12`. An orthonormal basis\n#| for that span, completed to a square orthogonal matrix, is the swap.\ndef swap_blocks(T, Q, k, n1, n2) {\n    m = n1 + n2\n    C = zeros(n1, n2)\n    for i to n1 {\n        for j to n2 {\n            C[i, j] = -T[k + i, k + n1 + j]\n        }\n    }\n    X = zeros(n1, n2)\n    junk = sylvester_small(T, k, n1, n2, C, X)\n    V = zeros(m, n2)\n    for i to n1 {\n        for j to n2 {\n            V[i, j] = X[i, j]\n        }\n    }\n    for j to n2 {\n        V[n1 + j, j] = 1.0\n    }\n    Z = zeros(m, m)\n    junk = orth_basis(V, Z, m, n2)\n    junk = apply_block(T, Q, Z, k, m)\n    # what is below the blocks is zero by construction; say so exactly, rather\n    # than leaving rounding where the next deflation test will read it\n    for i to m {\n        for j to m {\n            if i > j + 1 {\n                T[k + i, k + j] = 0.0\n            }\n        }\n    }\n    T[k + n2, k + n2 - 1] = 0.0\n    if n2 == 2 {\n        junk = split2(T, Q, k)\n    }\n    if n1 == 2 {\n        junk = split2(T, Q, k + n2)\n    }\n    return 0.0\n}\n\n#| The size of the diagonal block starting at `k`: 2 if it heads a 2 by 2, else 1.\ndef block_size(T, k) {\n    n = len(T)\n    if k + 1 < n {\n        if T[k + 1, k] != 0.0 {\n            return 2.0\n        }\n    }\n    return 1.0\n}\n\n#| Move every block whose eigenvalues have a negative real part to the top left,\n#| in the order they already appear, and return how many rows they take up. That\n#| leading corner of `Q` is then an orthonormal basis for the stable invariant\n#| subspace, which is what the Riccati solver needs.\ndef reorder_stable(T, Q) {\n    n = len(T)\n    want = 0\n    k = 0\n    while k < n {\n        sz = block_size(T, k)\n        if T[k, k] < 0.0 {\n            j = k\n            while j > want {\n                top = j - 1\n                psz = 1\n                if top > 0 {\n                    if T[top, top - 1] != 0.0 {\n                        top = top - 1\n                        psz = 2\n                    }\n                }\n                junk = swap_blocks(T, Q, top, psz, sz)\n                j = top\n            }\n            want = want + sz\n        }\n        # whatever moved, the blocks past this one did not: carry on after it\n        k = k + sz\n    }\n    return want\n}\n\ntest \"the Hessenberg form is Hessenberg, orthogonal, and still the same matrix\" {\n    A = test_matrix(7, 11.0)\n    H = zeros(7, 7)\n    Q = zeros(7, 7)\n    junk = hessenberg(A, H, Q)\n    for i to 7 {\n        for j to 7 {\n            if i > j + 1 {\n                assert_near(H[i, j], 0.0, 1.0e-14)\n            }\n        }\n    }\n    assert_lt(orthogonality(Q), 1.0e-13)\n    assert_lt(schur_residual(A, H, Q), 1.0e-12)\n}\n\ntest \"the Schur factor is quasi-triangular: never two nonzero subdiagonals in a row\" {\n    A = test_matrix(9, 5.0)\n    T = zeros(9, 9)\n    Q = zeros(9, 9)\n    junk = schur(A, T, Q)\n    for i to 9 {\n        for j to 9 {\n            if i > j + 1 {\n                assert_near(T[i, j], 0.0, 1.0e-14)\n            }\n        }\n    }\n    for k to 7 {\n        if T[k + 1, k] != 0.0 {\n            assert_near(T[k + 2, k + 1], 0.0, 1.0e-14)\n        }\n    }\n}\n\ntest \"Q T Q^T puts the matrix back, and Q is orthogonal, for a matrix with no structure\" {\n    A = test_matrix(9, 5.0)\n    T = zeros(9, 9)\n    Q = zeros(9, 9)\n    junk = schur(A, T, Q)\n    assert_lt(orthogonality(Q), 1.0e-13)\n    assert_lt(schur_residual(A, T, Q), 1.0e-12)\n}\n\ntest \"a symmetric matrix comes out diagonal, with the eigenvalues it should have\" {\n    B = test_matrix(6, 29.0)\n    A = B + transpose(B)\n    T = zeros(6, 6)\n    Q = zeros(6, 6)\n    junk = schur(A, T, Q)\n    for k to 5 {\n        assert_near(T[k + 1, k], 0.0, 1.0e-12)\n    }\n    assert_lt(schur_residual(A, T, Q), 1.0e-12)\n    vals = zeros(6)\n    vecs = zeros(6, 6)\n    eig_sym(A, vals, vecs)\n    # the same six numbers, in whatever order each method happened to find them\n    for i to 6 {\n        near = 1.0e30\n        for j to 6 {\n            near = min(near, abs(T[i, i] - vals[j]))\n        }\n        assert_lt(near, 1.0e-10)\n    }\n}\n\ntest \"the diagonal blocks carry the eigenvalues, complex pairs included\" {\n    # a rotation by 0.4 radians scaled by 1.3, plus two real eigenvalues\n    A = zeros(4, 4)\n    A[0, 0] = 1.3 * cos(0.4)\n    A[0, 1] = -1.3 * sin(0.4)\n    A[1, 0] = 1.3 * sin(0.4)\n    A[1, 1] = 1.3 * cos(0.4)\n    A[2, 2] = -2.0\n    A[3, 3] = 0.5\n    A[0, 3] = 0.7\n    A[1, 2] = -0.3\n    T = zeros(4, 4)\n    Q = zeros(4, 4)\n    junk = schur(A, T, Q)\n    assert_lt(schur_residual(A, T, Q), 1.0e-13)\n    starts = zeros(4)\n    count = schur_blocks(T, starts)\n    assert_near(count, 3.0, 1.0e-12)\n    # the complex pair sits in a 2 by 2 whose two diagonal entries are equal\n    found = 0.0\n    for b to count {\n        k = starts[b]\n        if k + 1 < 4 {\n            if T[k + 1, k] != 0.0 {\n                found = 1.0\n                assert_near(T[k, k], T[k + 1, k + 1], 1.0e-12)\n                assert_near(T[k, k], 1.3 * cos(0.4), 1.0e-11)\n                # the imaginary part is the geometric mean of the off-diagonals\n                assert_near(sqrt(-T[k, k + 1] * T[k + 1, k]), 1.3 * sin(0.4), 1.0e-11)\n            }\n        }\n    }\n    assert_near(found, 1.0, 1.0e-12)\n}\n\ntest \"the eigenvalues agree with the ones the general solver reports\" {\n    A = test_matrix(8, 101.0)\n    T = zeros(8, 8)\n    Q = zeros(8, 8)\n    junk = schur(A, T, Q)\n    re = zeros(8)\n    im = zeros(8)\n    eig_general(A, re, im)\n    starts = zeros(8)\n    count = schur_blocks(T, starts)\n    total = 0.0\n    for b to count {\n        k = starts[b]\n        wide = 0.0\n        if k + 1 < 8 {\n            if T[k + 1, k] != 0.0 {\n                wide = 1.0\n            }\n        }\n        total = total + 1.0 + wide\n        # every block's real part is one the other solver also found\n        near = 1.0e30\n        for j to 8 {\n            near = min(near, abs(T[k, k] - re[j]))\n        }\n        assert_lt(near, 1.0e-8)\n    }\n    assert_near(total, 8.0, 1.0e-12)\n}\n\ntest \"the trace and the determinant survive the decomposition\" {\n    A = test_matrix(6, 77.0)\n    T = zeros(6, 6)\n    Q = zeros(6, 6)\n    junk = schur(A, T, Q)\n    ta = 0.0\n    tt = 0.0\n    for i to 6 {\n        ta = ta + A[i, i]\n        tt = tt + T[i, i]\n    }\n    assert_near(tt, ta, 1.0e-11)\n}\n\ntest \"sixty matrices of every size from 2 to 14 come back within a few epsilon\" {\n    worst_res = 0.0\n    worst_orth = 0.0\n    worst_struct = 0.0\n    for si to 60 {\n        n = 2 + mod(si, 13.0)\n        A = test_matrix(n, 3.0 + 7.0 * si)\n        T = zeros(n, n)\n        Q = zeros(n, n)\n        junk = schur(A, T, Q)\n        worst_res = max(worst_res, schur_residual(A, T, Q))\n        worst_orth = max(worst_orth, orthogonality(Q))\n        # no 2 by 2 block may touch the next one, and nothing below them at all\n        for k to n - 2 {\n            if T[k + 1, k] != 0.0 {\n                worst_struct = max(worst_struct, abs(T[k + 2, k + 1]))\n            }\n        }\n        for i to n {\n            for j to n {\n                if i > j + 1 {\n                    worst_struct = max(worst_struct, abs(T[i, j]))\n                }\n            }\n        }\n    }\n    assert_lt(worst_res, 1.0e-13)\n    assert_lt(worst_orth, 1.0e-13)\n    assert_near(worst_struct, 0.0, 1.0e-300)\n}\n\ntest \"reordering puts the stable eigenvalues first and leaves the matrix alone\" {\n    worst_res = 0.0\n    worst_orth = 0.0\n    for si to 30 {\n        n = 3 + mod(si, 9.0)\n        A = test_matrix(n, 13.0 + 5.0 * si)\n        T = zeros(n, n)\n        Q = zeros(n, n)\n        junk = schur(A, T, Q)\n        cut = reorder_stable(T, Q)\n        worst_res = max(worst_res, schur_residual(A, T, Q))\n        worst_orth = max(worst_orth, orthogonality(Q))\n        # everything above the cut is stable, everything below it is not\n        for i to n {\n            if i < cut {\n                assert_lt(T[i, i], 0.0)\n            }\n            else {\n                assert_lt(-T[i, i], 1.0e-14)\n            }\n        }\n        # and the blocks are still blocks\n        for k to n - 2 {\n            if T[k + 1, k] != 0.0 {\n                assert_near(T[k + 2, k + 1], 0.0, 1.0e-300)\n            }\n        }\n    }\n    assert_lt(worst_res, 1.0e-11)\n    assert_lt(worst_orth, 1.0e-12)\n}\n\ntest \"reordering moves the eigenvalues without changing them\" {\n    n = 8\n    A = test_matrix(n, 404.0)\n    T = zeros(n, n)\n    Q = zeros(n, n)\n    junk = schur(A, T, Q)\n    before = zeros(n)\n    for i to n {\n        before[i] = T[i, i]\n    }\n    cut = reorder_stable(T, Q)\n    # the same n diagonal entries, in some other order\n    for i to n {\n        near = 1.0e30\n        for j to n {\n            near = min(near, abs(T[i, i] - before[j]))\n        }\n        assert_lt(near, 1.0e-10)\n    }\n    sum_before = 0.0\n    sum_after = 0.0\n    for i to n {\n        sum_before = sum_before + before[i]\n        sum_after = sum_after + T[i, i]\n    }\n    assert_near(sum_after, sum_before, 1.0e-11)\n}\n\n#! use lib/schur\n\nARE_SIGN_STEPS = 60\n\n#| `B R^-1 B^T`, formed by solving `R Y = B^T` rather than inverting `R`. `B` is\n#| `n` by `m` and `R` is `m` by `m`.\ndef are_gram(B: Arr, R: Arr) {\n    n = len(B)\n    m = len(R)\n    Y = zeros(m, n)\n    rhs = zeros(m)\n    col = zeros(m)\n    for j to n {\n        for i to m {\n            rhs[i] = B[j, i]\n        }\n        lin_solve(R, rhs, col)\n        for i to m {\n            Y[i, j] = col[i]\n        }\n    }\n    return B @ Y\n}\n\n#| The Hamiltonian whose stable invariant subspace carries the answer.\ndef are_hamiltonian(A: Arr, B: Arr, R: Arr, W: Arr) {\n    n = len(A)\n    G = are_gram(B, R)\n    H = zeros(2 * n, 2 * n)\n    for i to n {\n        for j to n {\n            H[i, j] = A[i, j]\n            H[i, n + j] = -G[i, j]\n            H[n + i, j] = -W[i, j]\n            H[n + i, n + j] = -A[j, i]\n        }\n    }\n    return H\n}\n\n#| Solve the Riccati equation into `P`. Returns the dimension of the stable\n#| subspace that was found, which had better be `n`: anything else means the\n#| Hamiltonian has an eigenvalue on the imaginary axis, and then there is no\n#| stabilizing solution to find.\ndef are_solve(A: Arr, B: Arr, R: Arr, W: Arr, P: Arr) {\n    n = len(A)\n    H = are_hamiltonian(A, B, R, W)\n    T = zeros(2 * n, 2 * n)\n    Z = zeros(2 * n, 2 * n)\n    junk = schur(H, T, Z)\n    cut = reorder_stable(T, Z)\n    if cut != n {\n        return cut\n    }\n    # the first n columns of Z are the subspace; split them and solve P X1 = X2\n    X1t = zeros(n, n)\n    rhs = zeros(n)\n    sol = zeros(n)\n    for i to n {\n        for j to n {\n            X1t[i, j] = Z[j, i]\n        }\n    }\n    for c to n {\n        for i to n {\n            rhs[i] = Z[n + c, i]\n        }\n        lin_solve(X1t, rhs, sol)\n        for i to n {\n            P[c, i] = sol[i]\n        }\n    }\n    # the true solution is symmetric; average away the rounding that is not\n    for i to n {\n        for j to n {\n            if j > i {\n                mid = 0.5 * (P[i, j] + P[j, i])\n                P[i, j] = mid\n                P[j, i] = mid\n            }\n        }\n    }\n    return cut\n}\n\n#| The same equation by Roberts's method: Newton's iteration\n#| `Z <- (Z + Z^-1) / 2` drives the Hamiltonian to its matrix sign, whose stable\n#| invariant subspace is the null space of `S + I`. `P` then comes out of an\n#| overdetermined system, solved here through its normal equations. Quadratic\n#| convergence, and no Schur decomposition anywhere in it.\ndef are_sign(A: Arr, B: Arr, R: Arr, W: Arr, P: Arr) {\n    n = len(A)\n    nn = 2 * n\n    Z = are_hamiltonian(A, B, R, W)\n    for step to ARE_SIGN_STEPS {\n        Zi = inverse(Z)\n        # scale the two terms to the same size, or the first steps waste digits\n        nz = 0.0\n        ni = 0.0\n        for i to nn {\n            for j to nn {\n                nz = nz + Z[i, j] * Z[i, j]\n                ni = ni + Zi[i, j] * Zi[i, j]\n            }\n        }\n        c = 1.0\n        if nz > 0.0 {\n            c = sqrt(sqrt(ni / nz))\n        }\n        moved = 0.0\n        for i to nn {\n            for j to nn {\n                was = Z[i, j]\n                Z[i, j] = 0.5 * (c * was + Zi[i, j] / c)\n                moved = max(moved, abs(Z[i, j] - was))\n            }\n        }\n        if moved < 1.0e-14 {\n            break\n        }\n    }\n    # [S12 ; S22 + I] P = -[S11 + I ; S21], as normal equations\n    M = zeros(nn, n)\n    rhs = zeros(nn, n)\n    for i to n {\n        for j to n {\n            M[i, j] = Z[i, n + j]\n            M[n + i, j] = Z[n + i, n + j]\n            rhs[i, j] = -Z[i, j]\n            rhs[n + i, j] = -Z[n + i, j]\n        }\n        M[n + i, i] = M[n + i, i] + 1.0\n        rhs[i, i] = rhs[i, i] - 1.0\n    }\n    Mt = transpose(M)\n    normal = Mt @ M\n    load = Mt @ rhs\n    col = zeros(n)\n    sol = zeros(n)\n    for j to n {\n        for i to n {\n            col[i] = load[i, j]\n        }\n        lin_solve(normal, col, sol)\n        for i to n {\n            P[i, j] = sol[i]\n        }\n    }\n    for i to n {\n        for j to n {\n            if j > i {\n                mid = 0.5 * (P[i, j] + P[j, i])\n                P[i, j] = mid\n                P[j, i] = mid\n            }\n        }\n    }\n    return 0.0\n}\n\n#| Put `P` back into the equation and report the largest entry of what is left.\n#| This is the number to look at before believing any of it.\ndef are_residual(A: Arr, B: Arr, R: Arr, W: Arr, P: Arr) {\n    n = len(A)\n    G = are_gram(B, R)\n    res = transpose(A) @ P + P @ A - P @ G @ P + W\n    worst = 0.0\n    for i to n {\n        for j to n {\n            worst = max(worst, abs(res[i, j]))\n        }\n    }\n    return worst\n}\n\n#| The residual measured against the size of the terms that produced it. This is\n#| the number that means something: the equation contains `P G P`, so a nearly\n#| uncontrollable system with a huge `P` has a huge absolute residual while\n#| still being solved to the last bit. Backward stability says this ratio sits\n#| at machine epsilon whatever the conditioning, and the tests below hold it\n#| under 1e-14 on systems where the absolute residual reaches 4e-6.\ndef are_residual_rel(A: Arr, B: Arr, R: Arr, W: Arr, P: Arr) {\n    n = len(A)\n    G = are_gram(B, R)\n    res = transpose(A) @ P + P @ A - P @ G @ P + W\n    worst = 0.0\n    na = 0.0\n    npp = 0.0\n    ng = 0.0\n    nw = 0.0\n    for i to n {\n        for j to n {\n            worst = max(worst, abs(res[i, j]))\n            na = max(na, abs(A[i, j]))\n            npp = max(npp, abs(P[i, j]))\n            ng = max(ng, abs(G[i, j]))\n            nw = max(nw, abs(W[i, j]))\n        }\n    }\n    scale = 2.0 * na * npp + npp * npp * ng + nw\n    if scale == 0.0 {\n        return worst\n    }\n    return worst / scale\n}\n\n#| The gains `K = R^-1 B^T P`, `m` by `n`: the input that goes with them is\n#| `u = -K z`.\ndef are_gains(B: Arr, R: Arr, P: Arr) {\n    n = len(P)\n    m = len(R)\n    K = zeros(m, n)\n    bp = transpose(B) @ P\n    rhs = zeros(m)\n    sol = zeros(m)\n    for j to n {\n        for i to m {\n            rhs[i] = bp[i, j]\n        }\n        lin_solve(R, rhs, sol)\n        for i to m {\n            K[i, j] = sol[i]\n        }\n    }\n    return K\n}\n\n#| The largest real part among the eigenvalues of the closed loop `A - B K`.\n#| Negative is the whole point of solving the equation.\ndef are_closed_slowest(A: Arr, B: Arr, K: Arr) {\n    n = len(A)\n    cl = A - B @ K\n    re = zeros(n)\n    im = zeros(n)\n    eig_general(cl, re, im)\n    worst = -1.0e30\n    for i to n {\n        worst = max(worst, re[i])\n    }\n    return worst\n}\n\n#| A controllable test system of order `n`: a companion-form `A` that is\n#| unstable, a single input at the last state, and unit weights.\ndef are_test_system(n, seed) {\n    A = test_matrix(n, seed)\n    for i to n {\n        A[i, i] = A[i, i] + 0.6\n    }\n    return A\n}\n\ntest \"the scalar equation lands where the quadratic formula says it does\" {\n    A = [[0.7]]\n    B = [[1.3]]\n    R = [[0.4]]\n    W = [[2.5]]\n    P = zeros(1, 1)\n    cut = are_solve(A, B, R, W, P)\n    assert_near(cut, 1.0, 1.0e-12)\n    expected = (0.7 + sqrt(0.7 * 0.7 + 1.3 * 1.3 * 2.5 / 0.4)) * 0.4 / (1.3 * 1.3)\n    assert_near(P[0, 0], expected, 1.0e-13)\n    assert_lt(are_residual(A, B, R, W, P), 1.0e-13)\n}\n\ntest \"a double integrator with unit weights wants gains of 1 and root 3\" {\n    A = [[0.0, 1.0], [0.0, 0.0]]\n    B = [[0.0], [1.0]]\n    R = [[1.0]]\n    W = [[1.0, 0.0], [0.0, 1.0]]\n    P = zeros(2, 2)\n    cut = are_solve(A, B, R, W, P)\n    assert_near(cut, 2.0, 1.0e-12)\n    K = are_gains(B, R, P)\n    assert_near(K[0, 0], 1.0, 1.0e-13)\n    assert_near(K[0, 1], sqrt(3.0), 1.0e-13)\n    assert_lt(are_residual(A, B, R, W, P), 1.0e-14)\n}\n\ntest \"P is symmetric and positive definite, as a cost-to-go has to be\" {\n    n = 6\n    A = are_test_system(n, 91.0)\n    B = zeros(n, 2)\n    B[n - 1, 0] = 1.0\n    B[0, 1] = 1.0\n    R = [[0.8, 0.0], [0.0, 1.4]]\n    W = zeros(n, n)\n    for i to n {\n        W[i, i] = 1.0 + 0.3 * i\n    }\n    P = zeros(n, n)\n    cut = are_solve(A, B, R, W, P)\n    assert_near(cut, n, 1.0e-12)\n    for i to n {\n        for j to n {\n            assert_near(P[i, j], P[j, i], 1.0e-12)\n        }\n    }\n    vals = zeros(n)\n    vecs = zeros(n, n)\n    eig_sym(P, vals, vecs)\n    assert_lt(0.0, vals[0])\n}\n\ntest \"the gains it hands back actually stabilize the system\" {\n    for si to 8 {\n        n = 3 + mod(si, 4.0)\n        A = are_test_system(n, 17.0 + 11.0 * si)\n        B = zeros(n, 1)\n        for i to n {\n            B[i, 0] = 1.0 / (1.0 + i)\n        }\n        R = [[0.5]]\n        W = zeros(n, n)\n        for i to n {\n            W[i, i] = 1.0\n        }\n        P = zeros(n, n)\n        cut = are_solve(A, B, R, W, P)\n        assert_near(cut, n, 1.0e-12)\n        K = are_gains(B, R, P)\n        assert_lt(are_closed_slowest(A, B, K), -1.0e-8)\n    }\n}\n\ntest \"the relative residual stays at rounding level across a spread of systems\" {\n    worst = 0.0\n    biggest = 0.0\n    for si to 10 {\n        n = 2 + mod(si, 6.0)\n        A = are_test_system(n, 3.0 + 29.0 * si)\n        B = zeros(n, 1)\n        B[n - 1, 0] = 1.0\n        for i to n {\n            B[i, 0] = B[i, 0] + 0.2 * cos(1.0 * i + si)\n        }\n        R = [[0.3 + 0.1 * mod(si, 3.0)]]\n        W = zeros(n, n)\n        for i to n {\n            W[i, i] = 0.5 + 0.4 * mod(i, 3.0)\n        }\n        P = zeros(n, n)\n        cut = are_solve(A, B, R, W, P)\n        assert_near(cut, n, 1.0e-12)\n        worst = max(worst, are_residual_rel(A, B, R, W, P))\n        for i to n {\n            for j to n {\n                biggest = max(bigges")
}

#| The distinct characters, ascending, so that a code can be turned into a token
#| index and back.
def corpus_vocab() {
    return chars("\n !\"#'()*+,-./0123456789:;<=>@ABCEFGHIKLMNOPQRSTUVWXYZ[]^_`abcdefghijklmnopqrstuvwxyz{|}")
}

test "the corpus is the size and the alphabet the header claims" {
    c = corpus()
    v = corpus_vocab()
    assert_near(len(c), CORPUS_N, 1.0e-12)
    assert_near(len(v), CORPUS_VOCAB, 1.0e-12)
    # the alphabet ascends and has no repeats
    for i to CORPUS_VOCAB - 1 {
        assert_lt(v[i], v[i + 1])
    }
    # and every character of the corpus is in it
    for i to CORPUS_N {
        lo = 0
        hi = CORPUS_VOCAB - 1
        while lo < hi {
            mid = floor((lo + hi) / 2.0)
            if v[mid] < c[i] {
                lo = mid + 1
            }
            else {
                hi = mid
            }
        }
        assert_near(v[lo], c[i], 1.0e-12)
    }
}