_calc_kernel

tit._calc_kernel

numba kernel for the K>=2 sphere maximum of :func:tit.calc.envelope_max.

One element at a time, parallel over elements: track the three eigen-branches of cos(t) C+ - sin(t) C- over tit.calc._THETA_GRID, root-find every sign change of g = cos^2(t) u - sin^2(t) v along a branch, add safeguard candidates, then polish the best candidates with Newton ascent. Each step is load-bearing: removing any of them, or lowering tit.calc._ROOT_ITERS / _NEWTON_ITERS, fails the dense-search tests in tests/test_calc.py. The 3x3 eigensolver is cyclic Jacobi warm-started from the previous eigenvectors (LAPACK's per-call overhead was most of the cost of a NumPy version).

numba is a hard dependency: the image ships it and container/blueprint/verify_runtime.py fails the build without it. Accuracy is pinned against independent oracles in tests/test_calc.py (host) and tests/numerical/test_calc_kernel.py (image).

max_sphere

max_sphere(E, grid, extra, root_iters, newton_iters)

K>=2 sphere maximum of E (2K, n, 3): (amp, n_hat, ok), parallel over elements.

ok is False where an eigen-solve did not converge; tit.calc raises on any such row.

Source code in tit/_calc_kernel.py

@njit(cache=True, error_model="numpy", parallel=True)
def max_sphere(E, grid, extra, root_iters, newton_iters):
    """K>=2 sphere maximum of ``E`` ``(2K, n, 3)``: ``(amp, n_hat, ok)``, parallel over elements.

    ``ok`` is False where an eigen-solve did not converge; ``tit.calc`` raises on any such row.
    """
    # grid, extra (n, K, 3: each beating pair's K=1 optimum), root_iters, newton_iters come from
    # tit.calc (_THETA_GRID, _pair_optima, _ROOT_ITERS, _NEWTON_ITERS). A null element gives 0;
    # tit.calc sets its direction.
    K = E.shape[0] // 2
    n_el = E.shape[1]
    G = grid.shape[0]
    amp = np.empty(n_el)
    n_hat = np.empty((n_el, 3))
    ok = np.empty(n_el, dtype=np.bool_)
    block = 64  # elements per scratch allocation
    cg, sg = np.cos(grid), np.sin(grid)
    for blk in prange((n_el + block - 1) // block):
        S = np.empty((K, 3))
        D = np.empty((K, 3))
        Cp = np.empty((3, 3))
        Cm = np.empty((3, 3))
        Cw = np.empty((3, 3))
        work = (
            np.empty((G, 3, 3)),
            np.empty((G, 3)),
            np.empty(3),
            np.empty((3, 3)),
            np.empty((3, 3)),
            np.empty((3, 3)),
            np.empty(3),
            np.empty(3),
        )
        for e in range(blk * block, min(n_el, (blk + 1) * block)):
            for k in range(K):
                for x in range(3):
                    S[k, x] = E[2 * k, e, x] + E[2 * k + 1, e, x]
                    D[k, x] = E[2 * k, e, x] - E[2 * k + 1, e, x]
            # C+ = S^T S / 2, C- = D^T D / 2, Cw = sum_k (E_ka E_kb^T + E_kb E_ka^T) = C+ - C-
            Cp[:] = 0.0
            Cm[:] = 0.0
            Cw[:] = 0.0
            for k in range(K):
                for x in range(3):
                    for y in range(3):
                        Cp[x, y] += S[k, x] * S[k, y]
                        Cm[x, y] += D[k, x] * D[k, y]
                        Cw[x, y] += E[2 * k, e, x] * E[2 * k + 1, e, y]
            Cp *= 0.5
            Cm *= 0.5
            for x in range(3):  # Cw + Cw^T
                for y in range(x, 3):
                    Cw[x, y] = Cw[y, x] = Cw[x, y] + Cw[y, x]
            amp[e], ok[e] = _element(
                S,
                D,
                Cp,
                Cm,
                Cw,
                grid,
                cg,
                sg,
                extra[e],
                root_iters,
                newton_iters,
                n_hat[e],
                work,
            )
    return amp, n_hat, ok