furax.linalg

def cg(
    A: AbstractLinearOperator,
    b: PyTree[Num[Array, '...']],
    x0: PyTree[Num[Array, '...']] | None = None,
    *,
    preconditioner: AbstractLinearOperator | None = None,
    max_steps: int = 500,
    atol: float = 0.0,
    rtol: float = 1e-5,
    stabilise_every: int = 10,
    negative_curvature: Literal['ignore', 'error', 'truncate'] = 'ignore',
    loop_kind: Literal['lax', 'checkpointed', 'bounded'] = 'lax',
    iteration_callback: Callable[[Array, Array], None] | None = None,
) -> CGResult:
    """Conjugate Gradient solver for symmetric positive definite systems Ax = b.

    The residual norm is recorded at every iteration (see ``CGResult.residuals``)
    so convergence can be monitored. Inputs may be sharded along their contracting
    dimensions. The solve is forward-mode differentiable by default; reverse-mode
    (``grad``/``vjp``) requires ``loop_kind='bounded'`` or ``'checkpointed'``.

    Convergence is declared when ``||r|| <= atol + rtol * ||b||``. With ``atol``
    and ``rtol`` both 0 the criterion is never met and the solver runs for exactly
    ``max_steps`` steps; otherwise it stops early once the residual is small enough,
    and ``max_steps`` is only a ceiling. ``A`` is assumed positive definite; negative
    curvature is handled per ``negative_curvature``.

    Args:
        A: A symmetric positive definite linear operator.
        b: Right-hand side of the system.
        x0: Initial guess. Defaults to zeros.
        preconditioner: Optional preconditioner M such that MA is better
            conditioned. M must be symmetric positive definite.
        max_steps: Maximum number of iteration steps.
        atol: Absolute tolerance on the residual norm.
        rtol: Relative tolerance on the residual norm (scaled by ``||b||``).
        stabilise_every: If set, replace the recursively-updated residual with
            the true residual ``b - A x`` every N iterations (after steps
            N, 2N, 3N, ...). This counters floating-point drift at the cost of
            one extra matvec per stabilisation.
        negative_curvature: How to handle negative curvature ``p^T A p < 0``, which a
            positive definite ``A`` never produces. One of:

            - ``'ignore'`` (default): assume ``A`` is positive definite and do not
                check. Fastest; the other modes add a per-iteration check.
            - ``'error'``: raise as soon as negative curvature is encountered (a
                debugging aid for verifying ``A``; configurable via Equinox's
                ``EQX_ON_ERROR``, see ``equinox.error_if``).
            - ``'truncate'``: stop and return the last iterate before the bad
                direction (truncated CG, as used by Newton-CG on indefinite Hessians).
        loop_kind: Lowering for the iteration loop (``equinox.internal.while_loop``).
            ``'lax'`` (default) is fastest and forward-mode differentiable only.
            ``'bounded'`` adds reverse-mode AD but its cost scales with ``max_steps``
            (the checkpoint structure runs to the ceiling regardless of early exit),
            so keep ``max_steps`` tight. ``'checkpointed'`` is reverse-mode only.
        iteration_callback: Optional host callback called after each step with
            ``(step, r_norm)`` as 0-d JAX arrays.  Runs via
            ``jax.debug.callback`` so it is JIT-compatible and ordered.

    Returns:
        CGResult with the solution, per-iteration residual norms, and iteration count.

    Examples:
        >>> import jax.numpy as jnp
        >>> from furax import DiagonalOperator
        >>> from furax.tree import as_structure
        >>> from furax.linalg import cg
        >>> d = jnp.array([1., 2., 3., 4., 5.])
        >>> A = DiagonalOperator(d, in_structure=as_structure(d))
        >>> b = jnp.ones(5)
        >>> result = cg(A, b, max_steps=20)
        >>> expected = jnp.array([1.0, 0.5, 1 / 3, 0.25, 0.2])
        >>> bool(jnp.allclose(result.solution, expected, atol=1e-4))
        True
    """
    if negative_curvature not in ('ignore', 'error', 'truncate'):
        raise ValueError(
            f'negative_curvature must be ignore/error/truncate, got {negative_curvature!r}'
        )
    truncate = negative_curvature == 'truncate'
    check_curvature = negative_curvature == 'error'

    if x0 is None:
        x0 = tree.zeros_like(b)

    has_scale = atol > 0 or rtol > 0
    norm_b = tree.norm(b)
    abs_tol = atol + rtol * norm_b

    def _converged(r_norm: Array) -> Array:
        return has_scale & (r_norm <= abs_tol)

    def _stable_r(x: PyTree) -> PyTree:
        return tree.sub(b, A(x))

    def _cheap_r(r: PyTree, alpha: Array, Ap: PyTree) -> PyTree:
        return tree.sub(r, tree.mul(alpha, Ap))

    # Initial residual r = b - A @ x0
    r = _stable_r(x0)

    # Apply preconditioner: z = M r
    M = preconditioner or (lambda _: _)
    z = M(r)

    p = z
    rz = tree.dot(r, z)  # r^T z (or r^T M^{-1} r)
    eps = jnp.finfo(rz.real.dtype).eps
    # `rz` is a residual energy, so square the relative norm floor. The 100*eps
    # margin catches roundoff-level residuals before stabilisation can amplify them.
    rz_floor = (100 * eps) ** 2 * jnp.abs(rz)

    r0_norm = tree.norm(r)
    # residuals[0] = initial residual; residuals[i+1] = residual after step i.
    # If i+1 >= max_steps (loop ran to completion), the last write is dropped by JAX.
    residuals = jnp.zeros(max_steps).at[0].set(r0_norm)

    def cond_fn(c: _CGCarry) -> Array:
        return ~c.converged & ~c.truncated & (c.step < max_steps)

    def body_fn(c: _CGCarry) -> _CGCarry:
        x, r, p, rz, i = c.x, c.r, c.p, c.rz, c.step

        Ap = A(p)
        pAp = tree.dot(p, Ap)
        active = jnp.abs(rz) > rz_floor

        # `pAp == 0` only at convergence (p -> 0), handled by `safe_pAp`
        # `pAp < 0` is genuine negative curvature, which a positive definite A should never produce
        if truncate:
            # Negative curvature: take no step this iteration and flag the loop to stop.
            truncated = active & (pAp < 0)
            safe_pAp = jnp.where(pAp == 0, 1.0, pAp)
            alpha = jnp.where(active & ~truncated, rz / safe_pAp, 0.0)
        else:
            truncated = c.truncated
            if check_curvature:
                pAp = eqx.error_if(
                    pAp, active & (pAp < 0), 'cg: negative curvature detected p^T A p < 0'
                )
            safe_pAp = jnp.where(pAp == 0, 1.0, pAp)
            alpha = jnp.where(active, rz / safe_pAp, 0.0)

        x = tree.add(x, tree.mul(alpha, p))

        # Choose between stable (true) and cheap (recursive) residual
        if stabilise_every > 0:
            r = jax.lax.cond(
                i % stabilise_every == stabilise_every - 1,
                lambda: _stable_r(x),
                lambda: _cheap_r(r, alpha, Ap),
            )
        else:
            r = _cheap_r(r, alpha, Ap)

        r_norm = tree.norm(r)

        if iteration_callback is not None:
            # `ordered = True` would error in distributed mode
            jax.debug.callback(iteration_callback, i, r_norm)

        z = M(r)
        rz_new = tree.dot(r, z)
        # Once the preconditioned residual energy reaches the floating-point floor,
        # fixed-iteration CG should become inert. This avoids dividing true-residual
        # roundoff by true-residual roundoff after stabilisation.
        safe_rz = jnp.where(active, rz, 1.0)
        beta = jnp.where(jnp.abs(rz_new) > rz_floor, rz_new / safe_rz, 0.0)

        p = tree.add(z, tree.mul(beta, p))

        converged = _converged(r_norm)
        residuals = c.residuals.at[i + 1].set(r_norm)

        return _CGCarry(x, r, p, rz_new, converged, truncated, i + 1, residuals)

    init_carry = _CGCarry(
        x=x0,
        r=r,
        p=p,
        rz=rz,
        converged=_converged(r0_norm),
        truncated=jnp.array(False),
        step=jnp.int32(0),
        residuals=residuals,
    )
    # `residuals` is only ever written (`.at[i].set`) inside the loop, never read, so mark it a
    # buffer for an efficient in-place scatter. `loop_kind` selects the lowering: 'lax' is fastest
    # (forward-mode AD only), 'bounded'/'checkpointed' add reverse-mode AD (see `loop_kind` arg).
    out = eqxi.while_loop(
        cond_fn,
        body_fn,
        init_carry,
        max_steps=max_steps,
        buffers=lambda c: c.residuals,
        kind=loop_kind,
    )

    return CGResult(solution=out.x, residuals=out.residuals, num_steps=out.step)

def eigvalsh(A: Float[Array, '... n n'], batch_size: int = 10_000) -> Float[Array, '... n']:
    """Eigenvalues of batched symmetric matrices, sorted ascending.

    Uses an analytic closed-form for n<=3, and jax.numpy.linalg.eigvalsh for larger matrices.

    Args:
        A: Batched symmetric matrix array with shape (..., n, n).
        batch_size: Number of matrices to process per XLA kernel launch for n>3.
            When batch_size >= total number of matrices the full batch is processed
            in one launch (maximum parallelism). Reduce to limit peak memory usage.
    """
    if A.ndim < 2 or A.shape[-1] != A.shape[-2]:
        raise ValueError(f'eigvalsh requires (..., n, n) input, got shape {A.shape}')
    n = A.shape[-1]
    if n == 1:
        return A[..., :1, 0]
    elif n == 2:
        return _eigvalsh_2x2(A)  # type: ignore[no-any-return]
    elif n == 3:
        return _eigvalsh_3x3(A)  # type: ignore[no-any-return]
    else:
        leading = A.shape[:-2]
        flat = A.reshape(-1, n, n)
        result = jax.lax.map(jnp.linalg.eigvalsh, flat, batch_size=batch_size)
        return result.reshape(*leading, n)  # type: ignore[no-any-return]

def lanczos_eigh(
    A: AbstractLinearOperator,
    v0: PyTree[Num[Array, '...']],
    *,
    k: int = 20,
    m: int | None = None,
) -> LanczosResult:
    """Lanczos algorithm for computing k eigenvalues via an m-dimensional Krylov subspace.

    Builds an m-dimensional Krylov subspace (m >= k) and computes all m Ritz pairs.
    When m == k the method returns all m Ritz pairs sorted ascending.  When m > k,
    only the k Ritz pairs with the smallest residual norms are returned.  Selecting
    by residual norm (rather than by eigenvalue magnitude) picks the pairs that have
    converged most reliably within the subspace, which need not be the extremal ones.

    Note:
        There is no ``which`` parameter (unlike ``lanczos_tr``).  Targeted selection
        of smallest/largest pairs would require either restarts or post-filtering of
        unconverged Ritz values, neither of which is meaningful for a single-shot
        m-step factorization.  Use ``lanczos_tr`` if you need extremal eigenpairs.

    Note:
        Early breakdown (``beta_j == 0``, i.e. invariant subspace reached) is not
        detected.  The corresponding Lanczos vector becomes zero and the remaining
        iterations produce zero contributions; affected Ritz pairs will have zero
        residual norms but their eigenvectors should not be trusted.

    The cheap Lanczos residual bound is used:
        ||A y_i - θ_i y_i|| ≈ |β_m| |s_i[m-1]|
    where s_i is the i-th eigenvector of the m×m tridiagonal T_m.

    Args:
        A: A Hermitian linear operator.
        v0: Initial vector for the Krylov subspace.
        k: Number of eigenpairs to return.
        m: Size of the Krylov subspace.  Must be >= k.  Defaults to min(2k, n).
            Larger m builds a richer subspace and can yield more accurate Ritz
            pairs, at the cost of m matrix-vector products and O(m) vector storage.

    Returns:
        LanczosResult containing the k best eigenvalues, eigenvectors, and their
        residual norms, sorted by eigenvalue ascending.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> from furax import DiagonalOperator
        >>> from furax.tree import as_structure, normal_like
        >>> d = jnp.array([1., 2., 3., 4., 5.])
        >>> A = DiagonalOperator(d, in_structure=as_structure(d))
        >>> v0 = normal_like(as_structure(d), jax.random.PRNGKey(0))
        >>> result = lanczos_eigh(A, v0, k=5)
        >>> result.eigenvalues
        Array([1., 2., 3., 4., 5.], dtype=float32)
    """
    m = m or _default_m(A, k)
    if m < k:
        raise ValueError(f'm ({m}) must be >= k ({k})')

    # Run Lanczos to build tridiagonal matrix in m-dimensional Krylov subspace
    alpha, beta, V, beta_last, _ = lanczos_tridiag(A, v0, m)
    ritz_values, ritz_vectors = jax.scipy.linalg.eigh_tridiagonal(alpha, beta, eigvals_only=False)

    eigenvectors = _vecmat(V, ritz_vectors)  # y_i = V s_i

    # ||A y_i - θ_i y_i|| ≈ |β_m| |s_i[-1]|
    residual_norms = jnp.abs(beta_last) * jnp.abs(ritz_vectors[-1, :])

    # Select the k best Ritz pairs by residual norm, then sort by eigenvalue ascending
    best_idx = jnp.argsort(residual_norms)[:k]
    best_idx = best_idx[jnp.argsort(ritz_values[best_idx])]
    return LanczosResult(
        eigenvalues=ritz_values[best_idx],
        eigenvectors=jax.tree.map(lambda leaf: leaf[best_idx], eigenvectors),
        residual_norms=residual_norms[best_idx],
    )

def lanczos_tr(
    A: AbstractLinearOperator,
    v0: PyTree[Num[Array, '...']],
    *,
    k: int = 20,
    m: int | None = None,
    which: LanczosWhich = 'LM',
    max_restarts: int = 300,
    tol: float = 1e-10,
) -> LanczosResult:
    r"""Thick-restart Lanczos for computing k eigenpairs of a Hermitian operator.

    Each restart cycle maintains an m-step Lanczos factorization:

    $$
    A V_m = V_m H + \beta_\text{last}\, v_\text{last}\, e_{m-1}^T
    $$

    where H is a bordered tridiagonal inner matrix.

    A pair is converged when the cheap Lanczos residual bound (ARPACK criterion)

    $$
    |\beta_\text{last}| \cdot |S[m-1, i]| \le \text{tol} \cdot \max(|\theta_i|, \epsilon \|A\|)
    $$

    where ``S`` is the eigenvector matrix of the m×m inner matrix H, so
    ``S[m-1, i]`` is the last component of the i-th eigenvector of H, and
    ``max|theta|`` is used as an estimate of ``||A||``.

    Uses full reorthogonalization throughout.  No locking: all k pairs are
    recomputed at every restart regardless of convergence status.

    Note:
        ``residual_norms`` is the cheap bound, not the true residual.  Exactly
        ``0`` means "converged below the detectable coupling", not zero error.

    Note:
        Lanczos converges fastest to *extremal* eigenvalues; interior ones
        (near the middle of the spectrum) converge slowly.  This makes ``which``
        targets that select interior pairs hard: ``'SM'`` for an indefinite
        operator picks eigenvalues closest to zero, which are interior, and
        restarts alone will not converge them unless ``m`` is a large fraction
        of ``n`` (no shift-invert is implemented).  ``'LM'``, ``'LA'``, ``'SA'``
        and the two ends of ``'BE'`` are extremal and converge with small ``m``.

    Args:
        A: A Hermitian linear operator.
        v0: Initial vector for the Krylov subspace.
        k: Number of eigenpairs to compute.
        m: Size of the Krylov subspace.  Must be > k.
            Defaults to min(2*k, n).
        which: Which k eigenpairs to target.  One of:
            - 'LM' (default): k largest magnitude (|λ|).
            - 'SM': k smallest magnitude (|λ|).
            - 'LA': k largest algebraic.
            - 'SA': k smallest algebraic.
            - 'BE': half (k//2) from each end of the spectrum; for odd k the
              extra pair comes from the high (largest algebraic) end.
        max_restarts: Maximum number of restart cycles.
        tol: Convergence tolerance; see criterion above.

    Returns:
        LanczosResult containing eigenvalues, eigenvectors, and residual norms,
        sorted by eigenvalue ascending.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> from furax import DiagonalOperator
        >>> from furax.tree import as_structure, normal_like
        >>> d = jnp.array([1., 2., 3., 4., 5.])
        >>> A = DiagonalOperator(d, in_structure=as_structure(d))
        >>> v0 = normal_like(as_structure(d), jax.random.PRNGKey(0))
        >>> result = lanczos_tr(A, v0, k=2, which='SA')
        >>> result.eigenvalues  # Should be approximately [1, 2]
        Array([1., 2.], dtype=float32)
    """
    if which not in get_args(LanczosWhich):
        raise ValueError(f'which must be one of {get_args(LanczosWhich)}, got {which!r}')
    m = m or _default_m(A, k)
    if m <= k:
        raise ValueError(f'm ({m}) must be > k ({k})')

    def _select_wanted(theta):  # type: ignore[no-untyped-def]
        if which == 'LM':  # largest magnitude
            sorted_idx = jnp.argsort(-jnp.abs(theta))
        elif which == 'SM':  # smallest magnitude
            sorted_idx = jnp.argsort(jnp.abs(theta))
        elif which == 'LA':  # largest algebraic
            sorted_idx = jnp.argsort(-theta)
        elif which == 'SA':  # smallest algebraic
            sorted_idx = jnp.argsort(theta)
        else:  # 'BE': half from each end of the spectrum
            sorted_idx = jnp.argsort(theta)
            n_low = k // 2
            n_high = k - n_low
            return jnp.concatenate([sorted_idx[:n_low], sorted_idx[-n_high:]])
        return sorted_idx[:k]

    def _check_converged(theta, beta_last, S, wanted_idx):  # type: ignore[no-untyped-def]
        # ARPACK criterion: |β_m| |s_i[-1]| ≤ tol * max(|θ_i|, eps*||A||)
        # Floor prevents stall when θ_i ≈ 0; max|θ| estimates ||A||.
        ritz_res = jnp.abs(beta_last) * jnp.abs(S[-1, wanted_idx])  # |β_m| |s_i[-1]|
        eps = jnp.finfo(theta.dtype).eps
        scale = jnp.maximum(jnp.abs(theta[wanted_idx]), eps * jnp.max(jnp.abs(theta)))
        return jnp.all(ritz_res <= tol * scale)

    # Initial m-step factorization
    alpha, beta, V, beta_last, v_last = lanczos_tridiag(A, v0, m)
    theta, S = jax.scipy.linalg.eigh_tridiagonal(alpha, beta, eigvals_only=False)
    wanted_idx = _select_wanted(theta)
    init_converged = _check_converged(theta, beta_last, S, wanted_idx)

    def cond_fn(state):  # type: ignore[no-untyped-def]
        *_, iteration, converged, _theta, _S, _wanted_idx = state
        return jnp.logical_and(iteration < max_restarts, ~converged)

    def body_fn(state):  # type: ignore[no-untyped-def]
        V, beta_last, v_last, iteration, _converged, theta, S, wanted_idx = state

        V_k = _vecmat(V, S[:, wanted_idx])  # U_k = V S[:,wanted]  (Ritz vectors)
        theta_k = theta[wanted_idx]  # θ_k
        h = beta_last * S[-1, wanted_idx]  # h_i = β_m s_i[-1]  (coupling)

        alpha_ext, beta_ext, V, beta_last, v_last = _tr_extend(A, V_k, v_last, k, m)

        H = _build_bordered_tridiag(theta_k, h, alpha_ext, beta_ext, k, m)
        theta, S = jnp.linalg.eigh(H)  # H S = S diag(θ)

        wanted_idx = _select_wanted(theta)
        converged = _check_converged(theta, beta_last, S, wanted_idx)

        return V, beta_last, v_last, iteration + 1, converged, theta, S, wanted_idx

    init_state = (V, beta_last, v_last, jnp.array(0), init_converged, theta, S, wanted_idx)
    V, beta_last, _v_last, _iters, _conv, theta, S, wanted_idx = jax.lax.while_loop(
        cond_fn, body_fn, init_state
    )

    # Sort selected pairs by eigenvalue ascending
    wanted_idx = wanted_idx[jnp.argsort(theta[wanted_idx])]
    eigenvalues = theta[wanted_idx]
    eigenvectors = _vecmat(V, S[:, wanted_idx])  # y_i = V s_i
    residual_norms = jnp.abs(beta_last) * jnp.abs(S[-1, wanted_idx])  # |β_m| |s_i[-1]|

    return LanczosResult(
        eigenvalues=eigenvalues,
        eigenvectors=eigenvectors,
        residual_norms=residual_norms,
    )