Source code for sparho.search

"""Top-level outer loops: ``grad_search`` (plain GD) and ``hoag_search`` (HOAG).

Both are pure imperative ``for``-loops that thread Solver + Criterion +
Hypergrad into a single call and operate in **log space**: ``hp0`` is
interpreted as a positive ``α``, and the search steps in ``θ = log α``. This
keeps ``α > 0`` without projection and matches how ``α`` typically varies —
across orders of magnitude. The chain rule ``dC/dθ = dC/dα · α`` is applied
inside the loop.

After the search, the inner solver is run once more on the **full** problem
at the best ``α`` seen (``best_value`` minimum) and the resulting ``coef``
goes into ``SearchResult.best_coef``. For ``CrossVal`` this matters — the
per-fold ``coef`` reported by the criterion is the last-fold fit, not the
full-data fit the user wants.

:func:`grad_search` is classical bilevel approximate-gradient descent: plain
``θ ← θ - lr · dC/dθ`` with a fixed learning rate. A baseline for
comparison and for problems where step-size choice is well-understood.

:func:`hoag_search` is the HOAG (Pedregosa 2016) port and the recommended
default for non-smooth bilevel HP optimization. It owns its outer loop: step
size adapted from a Lipschitz proxy ``L``, acceptance test has a ``+C·tol``
slack term that tolerates noise from inner-solver inaccuracy, and inner
tolerance is part of the loop state (optionally decreased exponentially
across iterations).
"""

from __future__ import annotations

import warnings
from collections.abc import Callable

import numpy as np

from .core.types import Hyperparam
from .criteria import Criterion
from .hypergrad import implicit_forward
from .problem import Problem
from .solver import Solver
from .state import IterationRecord, SearchResult


def _cg_status_from_warnings(captured: list[warnings.WarningMessage]) -> str:
    """Map captured ``implicit_forward`` warnings to a single ``cg_status``.

    Returns ``"nonfinite"`` if any captured warning indicates a non-finite
    CG output (the strictly worse failure mode), ``"nonconvergence"`` if CG
    returned a finite but unconverged result, otherwise ``"ok"``.
    """
    status = "ok"
    for w in captured:
        if not issubclass(w.category, RuntimeWarning):
            continue
        msg = str(w.message)
        if "non-finite hypergradient" in msg or "non-finite gradient" in msg:
            return "nonfinite"
        if "CG failed" in msg:
            if "finite=False" in msg:
                return "nonfinite"
            status = "nonconvergence"
    return status


HypergradFn = Callable[..., Hyperparam]






# ---------------------------------------------------------------- HOAG






# ---------------------------------------------------------------- helpers


def _as_positive(hp: Hyperparam, n_features: int) -> Hyperparam:
    """Validate ``hp > 0`` and (for vector-α) length match against ``n_features``."""
    if isinstance(hp, np.ndarray):
        arr = np.asarray(hp, dtype=np.float64)
        if arr.ndim != 1:
            raise ValueError(f"hp0 must be a 1-D vector for per-feature α, got ndim={arr.ndim}")
        if arr.shape[0] != n_features:
            raise ValueError(
                f"hp0 length ({arr.shape[0]}) must equal problem.n_features ({n_features})"
            )
        if np.any(arr <= 0):
            raise ValueError("hp0 must be strictly positive componentwise")
        return arr
    val = float(hp)
    if val <= 0:
        raise ValueError("hp0 must be strictly positive")
    return val


def _log(hp: Hyperparam) -> Hyperparam:
    if isinstance(hp, np.ndarray):
        return np.log(hp)
    return float(np.log(hp))


def _exp(theta: Hyperparam) -> Hyperparam:
    if isinstance(theta, np.ndarray):
        return np.exp(theta)
    return float(np.exp(theta))


def _elementwise_mul(a: Hyperparam, b: Hyperparam) -> Hyperparam:
    if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
        return np.asarray(a, dtype=np.float64) * np.asarray(b, dtype=np.float64)
    return float(a) * float(b)


def _norm(x: Hyperparam) -> float:
    if isinstance(x, np.ndarray):
        return float(np.linalg.norm(x))
    return abs(float(x))


def _sub(a: Hyperparam, b: Hyperparam) -> Hyperparam:
    if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
        return np.asarray(a, dtype=np.float64) - np.asarray(b, dtype=np.float64)
    return float(a) - float(b)


def _scale(a: Hyperparam, c: float) -> Hyperparam:
    if isinstance(a, np.ndarray):
        return c * np.asarray(a, dtype=np.float64)
    return c * float(a)


def _copy(a: Hyperparam) -> Hyperparam:
    if isinstance(a, np.ndarray):
        return np.asarray(a, dtype=np.float64).copy()
    return float(a)