liblaf.peach.optim.pncg ¶

def init(self) -> ConvergenceState:
    """Create an empty convergence state shaped like `params`."""
    return ConvergenceState()

def primary_success(self, state: ConvergenceState) -> Bool[Tensor, ""]:
    """Check the primary absolute-relative gradient tolerance."""
    return state.grad_norm <= max(
        self.atol_primary, self.rtol_primary * state.grad_norm_first
    )

def secondary_success(self, state: ConvergenceState) -> Bool[Tensor, ""]:
    """Check the looser secondary gradient tolerance."""
    return state.grad_norm <= max(
        self.atol_secondary, self.rtol_secondary * state.grad_norm_first
    )

def terminate(self, state: ConvergenceState) -> tuple[bool, Result]:
    """Return the convergence decision and result code."""
    if state.step < self.max_steps:
        if self.primary_success(state):
            return True, Result.PRIMARY_SUCCESS
        if state.stagnation_count > self.patience:
            return True, Result.STAGNATION
        return False, Result.INTERRUPT
    if self.secondary_success(state):
        return True, Result.SECONDARY_SUCCESS
    return True, Result.MAX_STEPS_REACHED

def update(
    self, state: ConvergenceState, g: Vector, *, line_search_ok: bool
) -> ConvergenceState:
    """Record the current gradient and gradient norms."""
    grad_norm: Scalar = torch.linalg.vector_norm(g)
    if state.step == 0:
        state.grad_norm_first = grad_norm
    if line_search_ok:
        state.stagnation_count = 0
    else:
        state.stagnation_count += 1
    state.grad_norm = grad_norm
    state.step += 1
    return state

def __call__(
    self,
    g: Vector,
    g_prev: Vector,
    P: Vector,
    p_prev: Vector,
    *,
    restart: bool = False,
) -> Vector:
    """Compute a descent direction, restarting when requested."""
    if restart:
        return -P * g
    beta: Scalar = dai_kou_plus(g=g, g_prev=g_prev, P=P, p_prev=p_prev)
    Pg: Vector = -P * g
    p: Vector = Pg + beta * p_prev
    return torch.where(torch.dot(p, g) < 0.0, p, Pg)

def hess_diag(self, state: HessianDampingState, H_diag: Vector) -> Vector:
    """Return `abs(H_diag) + factor * mean(abs(H_diag))`."""
    H_diag: Vector = torch.abs(H_diag)
    H_diag_mean: Scalar = torch.mean(H_diag)
    state.hess_diag_mean = H_diag_mean
    return H_diag + state.factor * H_diag_mean

def hess_quad(self, state: HessianDampingState, p: Vector, pHp: Scalar) -> Scalar:
    """Add the damping contribution to a Hessian quadratic form."""
    return pHp + state.factor * state.hess_diag_mean * torch.dot(p, p)

def init(self) -> HessianDampingState:
    """Create damping state with the initial factor."""
    return HessianDampingState(factor=self.initial)

def update(
    self,
    state: HessianDampingState,
    *,
    actual_decrease: Scalar,
    line_search_step: int,
    predicted_decrease: Scalar,
) -> None:
    """Adapt the damping factor from line-search behavior."""
    factor: float = state.factor
    if line_search_step == 0 and actual_decrease > predicted_decrease:
        factor *= 0.5
    else:
        factor *= (line_search_step + 1) ** 2
    state.factor = min(max(factor, self.factor_min), self.factor_max)

def __call__[X](
    self,
    state: LineSearchState,
    problem: Problem[X],
    model_state: X,
    m: Scalar,
    p: Vector,
    params: Vector,
    pHp: Scalar,
) -> None:
    """Run line search along direction `p` from `params`.

    The initial step length is the smaller of the Newton proposal and the
    configured infinity-norm bound. If the problem implements
    [`Problem.max_step_size`][liblaf.peach.optim.base.Problem.max_step_size],
    that hook receives the proposed displacement `alpha * p` and returns a
    safe fraction of it. Every accepted or rejected trial is materialized
    through [`Problem.update`][liblaf.peach.optim.base.Problem.update].
    """
    alpha_upper: Scalar = self.line_search_upper(
        p=p, max_step_norm=self.max_step_norm
    )
    alpha_newton: Scalar = self.line_search_newton(m=m, pHp=pHp)
    alpha: Scalar = torch.minimum(alpha_upper, alpha_newton)
    if is_implemented(problem, Problem.max_step_size):
        step_fraction: Scalar = problem.max_step_size(model_state, alpha * p)
        step_fraction: Scalar = torch.as_tensor(step_fraction)
        alpha *= torch.clamp(step_fraction, 0.0, 1.0)

    f0: Scalar = state.f_alpha
    for step in range(self.max_steps + 1):
        if step > 0:
            alpha *= 0.5
        problem.update(model_state, params + alpha * p)
        f_alpha: Scalar = problem.fun(model_state)
        if self.armijo_condition(f_alpha=f_alpha, f0=f0, alpha=alpha, m=m):
            state.ok = True
            break
    else:
        state.ok = False
    state.alpha = alpha
    state.f_alpha = f_alpha
    state.f0 = f0
    state.step = step

def armijo_condition(
    self, f_alpha: Scalar, f0: Scalar, alpha: Scalar, m: Scalar
) -> Bool[Tensor, ""]:
    """Evaluate the Armijo sufficient-decrease condition."""
    return f_alpha <= f0 + alpha * self.armijo * m

def init(self, fun: Scalar) -> LineSearchState:
    """Create an empty line-search state."""
    return LineSearchState(f_alpha=fun)

@staticmethod
def line_search_newton(m: Scalar, pHp: Scalar) -> Scalar:
    """Return the Newton step length or `1.0` when curvature is unsuitable."""
    alpha: Scalar = -m / pHp
    return torch.where((pHp <= 0.0) | (m >= 0.0), 1.0, alpha)

@staticmethod
def line_search_upper(p: Vector, max_step_norm: float) -> Scalar:
    """Return the largest step satisfying the infinity-norm bound."""
    if not math.isfinite(max_step_norm):
        return torch.as_tensor(torch.inf)
    p_norm: Scalar = torch.linalg.vector_norm(p, ord=torch.inf)
    return torch.where(p_norm == 0.0, 0.0, max_step_norm / p_norm)

@override
def init[X](self, problem: BaseProblem[X], model_state: X, params: Vector) -> State:
    """Initialize `Pncg` state.

    The stored parameter vector is cloned from `params`; subsequent steps
    mutate the clone instead of the caller-owned input tensor.
    """
    problem: Problem[X] = cast("Problem[X]", problem)
    fun: Scalar = problem.fun(model_state)
    return self.State(
        fun=fun,
        params=params.clone(),
        convergence_state=self.criteria.init(),
        hess_damping_state=self.hess_damping.init(),
        line_search_state=self.line_search.init(fun=fun),
    )

def minimize[X](
    self, problem: BaseProblem[X], model_state: X, params: Vector
) -> Solution[S, T]:
    """Run optimization until the configured termination rule stops.

    After every step, [`Problem.callback`][liblaf.peach.optim.base.Problem.callback]
    is called when the concrete problem implements it. The callback receives
    the current model state and the same mutable optimizer state that will be
    stored on the returned [`Solution`][liblaf.peach.optim.base.Solution].

    Args:
        problem: Optimization problem that supplies objective hooks.
        model_state: Initial model state used by objective hooks.
        params: Initial optimizer parameter vector.

    Returns:
        The final solution.
    """
    problem: Problem[X] = cast("Problem[X]", problem)
    opt_state: S = self.init(problem, model_state, params)
    while True:
        self.step(problem, model_state, opt_state)
        if is_implemented(problem, Problem.callback):
            problem.callback(model_state, opt_state)
        ok, result = self.terminate(problem, model_state, opt_state)
        if ok:
            break
    solution: Solution[S, T] = self.postprocess(
        problem, model_state, opt_state, result
    )
    return solution

@override
def postprocess[X](
    self, problem: BaseProblem[X], model_state: X, opt_state: State, result: Result
) -> Solution:
    """Build a `Pncg` solution object."""
    stats: Pncg.Stats = self.Stats()
    return Solution(result=result, state=opt_state, stats=stats)

@override
def step[X](
    self, problem: BaseProblem[X], model_state: X, opt_state: State
) -> None:
    """Run one `Pncg` step.

    The method evaluates derivatives at `model_state`, writes the accepted
    direction, slope, Hessian diagnostics, line-search diagnostics, and
    parameter update into `opt_state`, then mutates `model_state` to the
    accepted trial.
    """
    problem: Problem[X] = cast("Problem[X]", problem)

    g: Vector = problem.grad(model_state)
    H_diag: Vector = problem.hess_diag(model_state)
    H_diag_damp: Vector = self.hess_damping.hess_diag(
        state=opt_state.hess_damping_state, H_diag=H_diag
    )
    P: Vector = torch.reciprocal(H_diag_damp)

    p: Vector = self.direction(
        g=g,
        g_prev=opt_state.grad,
        P=P,
        p_prev=opt_state.direction,
        restart=opt_state.step == 0 or not opt_state.line_search_state.ok,
    )
    opt_state.direction = p
    opt_state.grad = g
    opt_state.hess_diag = H_diag
    slope: Scalar = torch.dot(g, p)
    opt_state.slope = slope

    pHp: Scalar = problem.hess_quad(model_state, p)
    opt_state.hess_quad = pHp
    pHp_damp: Scalar = self.hess_damping.hess_quad(
        state=opt_state.hess_damping_state, p=p, pHp=pHp
    )
    self.line_search(
        opt_state.line_search_state,
        problem=problem,
        model_state=model_state,
        m=slope,
        p=p,
        params=opt_state.params,
        pHp=pHp_damp,
    )
    ls_state: Pncg.LineSearchState = opt_state.line_search_state
    opt_state.fun = ls_state.f_alpha

    alpha: Scalar = ls_state.alpha
    actual_decrease: Scalar = ls_state.f0 - ls_state.f_alpha
    predicted_decrease: Scalar = -alpha * slope - 0.5 * alpha**2 * pHp_damp
    self.hess_damping.update(
        opt_state.hess_damping_state,
        actual_decrease=actual_decrease,
        line_search_step=ls_state.step,
        predicted_decrease=predicted_decrease,
    )

    self.criteria.update(
        opt_state.convergence_state, g=g, line_search_ok=ls_state.ok
    )

    opt_state.params += alpha * p

@override
def terminate[X](
    self, problem: BaseProblem[X], model_state: X, opt_state: State
) -> tuple[bool, Result]:
    """Delegate stopping to the configured convergence criteria."""
    return self.criteria.terminate(opt_state.convergence_state)