Newton Krylov

newton_krylov

Explicit eager NumPy Newton-GMRES with checked residuals and backtracking.

This implementation has data-dependent host decisions and supplies no automatic differentiation rule. It never constructs a dense Jacobian or probes a basis.

NewtonKrylovConfig(max_iterations=8, rtol=1e-08, atol=1e-10, damping=1.0, linear_options=(lambda: LinearSolveOptions(rtol=1e-06, atol=0.0))(), max_backtracks=12, backtrack_factor=0.5, armijo=0.0001) dataclass

Bases: NewtonConfig

Outer RMS tolerances, inner true-residual gate and Armijo backtracking.

Inherited max_iterations is a maximum number of applied Newton updates, not an exact unroll count. The outer gate is atol + rtol * initial_rms. Inner options independently use the checked linear solver's L2 gate max(atol, rtol * ||rhs||) and its documented iteration budget.

Attributes:

Name	Type	Description
linear_options	LinearSolveOptions	Explicit inner tolerance and iteration budget.
max_backtracks	int	Maximum reductions after the initial damped trial.
backtrack_factor	float	Multiplicative step-length reduction, strictly in (0,1).
armijo	float	Required fractional RMS decrease per unit step, strictly in (0,1).

__post_init__()

Validate outer, inner and globalization settings.

Raises:

Type	Description
TypeError	If inner options do not use the existing checked contract.
ValueError	If backtracking settings are invalid.

Source code in src/op_engine/newton_krylov.py

def __post_init__(self) -> None:
    """Validate outer, inner and globalization settings.

    Raises:
        TypeError: If inner options do not use the existing checked contract.
        ValueError: If backtracking settings are invalid.
    """
    NewtonConfig.__post_init__(self)
    if not isinstance(self.linear_options, LinearSolveOptions):
        msg = "linear_options must be LinearSolveOptions"
        raise TypeError(msg)
    if (
        not isinstance(self.max_backtracks, Integral)
        or isinstance(self.max_backtracks, bool)
        or self.max_backtracks < 0
    ):
        msg = "max_backtracks must be a nonnegative integer"
        raise ValueError(msg)
    if any(
        not np.isfinite(value) or not 0.0 < value < 1.0
        for value in (self.backtrack_factor, self.armijo)
    ):
        msg = "backtrack_factor and armijo must be finite and strictly in (0,1)"
        raise ValueError(msg)

NewtonKrylovSolver(config=NewtonKrylovConfig(), linear_solver=ScipyGmresSolver(), preconditioner_factory=None) dataclass

Stateless eager NumPy Newton-GMRES with a per-linearization inverse factory.

The default inner solver is :class:ScipyGmresSolver (left inverse preconditioning). A failed checked inner solve returns nonlinear nonconvergence without applying its candidate; malformed callbacks and arbitrary linear-solver exceptions remain fatal. The integration boundary applies its existing fixed/adaptive/replay failure policy.

Attributes:

Name	Type	Description
config	NewtonKrylovConfig	Outer/inner tolerances and bounded line search.
linear_solver	LinearSolver	Existing checked eager linear-solver callback.
preconditioner_factory	InversePreconditionerFactory | None	Optional factory receiving the exact current request, operator and :class:NewtonLinearization context.

__post_init__()

Validate explicit configuration and callbacks.

Raises:

Type	Description
TypeError	If configuration or callbacks are unsupported.

Source code in src/op_engine/newton_krylov.py

def __post_init__(self) -> None:
    """Validate explicit configuration and callbacks.

    Raises:
        TypeError: If configuration or callbacks are unsupported.
    """
    if not isinstance(self.config, NewtonKrylovConfig):
        msg = "config must be NewtonKrylovConfig"
        raise TypeError(msg)
    if not callable(self.linear_solver) or (
        self.preconditioner_factory is not None
        and not callable(self.preconditioner_factory)
    ):
        msg = "linear solver and preconditioner factory must be callable"
        raise TypeError(msg)

solve(problem, initial_guess)

Return a root candidate and explicit eager termination evidence.

Returns:

Type	Description
NonlinearSolveResult	Owned candidate and residual with checked inner/outer diagnostics.

Raises:

Type	Description
TypeError	If input or callback arrays are not native real NumPy float32/float64 with matching shape/dtype/namespace.
ValueError	If the state is empty or no residual JVP is supplied.

Source code in src/op_engine/newton_krylov.py

def solve(
    self, problem: NonlinearProblem, initial_guess: Array, /
) -> NonlinearSolveResult:
    """Return a root candidate and explicit eager termination evidence.

    Returns:
        Owned candidate and residual with checked inner/outer diagnostics.

    Raises:
        TypeError: If input or callback arrays are not native real NumPy
            float32/float64 with matching shape/dtype/namespace.
        ValueError: If the state is empty or no residual JVP is supplied.
    """
    real_dtypes = (np.float32, np.float64)
    if (
        not isinstance(initial_guess, np.ndarray)
        or initial_guess.dtype not in real_dtypes
    ):
        msg = "NewtonKrylovSolver requires eager NumPy float32/float64 arrays"
        raise TypeError(msg)
    if initial_guess.size < 1 or problem.jvp is None:
        msg = (
            "NewtonKrylovSolver requires a nonempty state and NonlinearProblem.jvp"
        )
        raise ValueError(msg)
    if not np.isfinite(initial_guess).all():
        msg = "initial_guess must be finite"
        raise ValueError(msg)
    iterate = initial_guess.copy()
    iterate.flags.writeable = False
    work = _Work()
    residual = self._residual(problem, iterate, work)
    initial_norm = _rms(residual)
    tolerance = self.config.atol + self.config.rtol * initial_norm
    step_norm = 0.0
    termination = NewtonKrylovTermination.MAX_ITERATIONS
    for _ in range(self.config.max_iterations):
        norm = _rms(residual)
        if not np.isfinite(norm):
            termination = NewtonKrylovTermination.NONFINITE_RESIDUAL
            break
        if norm <= tolerance:
            termination = NewtonKrylovTermination.CONVERGED
            break
        direction = self._direction(problem, iterate, residual, work)
        if direction is None:
            termination = NewtonKrylovTermination.LINEAR_FAILURE
            break
        trial = self._trial(
            problem, iterate, direction, norm=norm, tolerance=tolerance, work=work
        )
        if trial is None:
            termination = NewtonKrylovTermination.LINE_SEARCH_FAILURE
            break
        candidate, residual = trial
        step_norm = _rms(candidate - iterate)
        iterate = candidate
        work.iterations += 1
    norm = _rms(residual)
    converged = np.isfinite(norm) and norm <= tolerance
    if converged:
        termination = NewtonKrylovTermination.CONVERGED
    diagnostics = NonlinearSolveDiagnostics(
        converged=np.asarray(converged),
        iterations=work.iterations,
        residual_evaluations=work.residuals,
        jacobian_evaluations=0,
        initial_residual_norm=np.asarray(initial_norm, dtype=iterate.dtype),
        residual_norm=np.asarray(norm, dtype=iterate.dtype),
        step_norm=np.asarray(step_norm, dtype=iterate.dtype),
        linear_solves=tuple(work.linear),
        termination=int(termination),
        jvp_evaluations=work.actions,
        backtracks=work.backtracks,
    )
    return NonlinearSolveResult(iterate, residual, diagnostics)

NewtonKrylovTermination

Bases: IntEnum

Eager termination codes, retained in stage/integration diagnostics.

NewtonLinearization(iteration, iterate, problem_context) dataclass

Unique per-linearization preconditioner context, with an owned iterate.

iterate is a read-only snapshot. Its full shape, dtype and coefficients, problem_context (including SDIRK scale/time when applicable) and this record's identity define the preparation lifetime. A factory is invoked afresh for every linearization; a caller retaining a cache owns invalidation. No iterate history is retained in solve diagnostics.

NonlinearStageContext(time, step_size, diagonal_step, stage_index) dataclass

Explicit SDIRK stage metadata, freshly bound on every stage and retry.