NewtonKrylov

class cubie.integrators.matrix_free_solvers.NewtonKrylov(precision: type[float16] | type[float32] | type[float64] | dtype[float16] | dtype[float32] | dtype[float64], solver_width: int, linear_solver: LinearSolverBase, norm: CorrectionNorm | None = None, **kwargs)[source]

Bases: MatrixFreeSolver

Factory for Newton–Krylov solver device functions.

Uses a matrix-free linear solver for each Newton correction.

Parameters:
  • precision (PrecisionDType) – Numerical precision for computations.

  • solver_width (int) – Solver vector length.

  • linear_solver (LinearSolverBase) – LinearSolverBase instance for solving linear systems.

  • **kwargs – Forwarded to NewtonKrylovConfig and the norm factory. Includes prefixed tolerance parameters (newton_atol, newton_rtol).

__init__(precision: type[float16] | type[float32] | type[float64] | dtype[float16] | dtype[float32] | dtype[float64], solver_width: int, linear_solver: LinearSolverBase, norm: CorrectionNorm | None = None, **kwargs) None[source]

Initialize NewtonKrylov with parameters.

Parameters:
  • precision (PrecisionDType) – Numerical precision for computations.

  • solver_width (int) – Solver vector length.

  • linear_solver (LinearSolverBase) – Inner linear solver, constructed with zero_initial_guess=True.

  • norm (CorrectionNorm, optional) – Correction norm. Defaults to DIRK scaling.

  • **kwargs – Newton solver settings.

_require_child_zero_guess() None[source]

Reject a linear solver not built for a zero guess.

build() NewtonKrylovCache[source]

Compile the Newton solver.

property device_function: Callable

Return cached Newton-Krylov solver device function.

property krylov_atol: ndarray

Return the Krylov absolute tolerance array from nested linear solver.

property krylov_max_iters: int

Return max linear iterations from nested linear solver.

property krylov_residual_floor: float

Return the linear solver’s weighted-residual floor.

property krylov_residual_reduction: float

Return the linear solver’s relative stopping factor.

property krylov_rtol: ndarray

Return the Krylov relative tolerance array from nested linear solver.

property linear_correction_type: str

Return correction type from nested linear solver.

property newton_atol: ndarray

Return absolute tolerance array.

property newton_max_iters: int

Return maximum Newton iterations.

property newton_rtol: ndarray

Return relative tolerance array.

register_buffers() None[source]

Register buffers according to locations in compile settings.

update(updates_dict: Dict[str, Any] | None = None, silent: bool = False, **kwargs) Set[str][source]

Update compile settings and invalidate cache if changed.

Parameters:
  • updates_dict (dict, optional) – Dictionary of settings to update.

  • silent (bool, default False) – If True, suppress warnings about unrecognized keys.

  • **kwargs – Additional settings as keyword arguments.

Returns:

Set of recognized parameter names that were updated.

Return type:

set

class cubie.integrators.matrix_free_solvers.NewtonKrylovConfig(precision: type[float16] | type[float32] | type[float64] | dtype[float16] | dtype[float32] | dtype[float64], instance_label: str = '', solver_width: int = 0, norm_fn: Callable | None = None, max_iters: int = 8, use_cached_auxiliaries: bool = False, residual_fn: Callable | None = None, krylov_linear_solver_fn: Callable | None = None, delta_location: str = 'local', residual_location: str = 'local', krylov_iters_local_location: str = 'local', prev_theta_location: str = 'local', *, jit_flags: JITFlags = NOTHING, unroll: UnrollFlags = NOTHING)[source]

Bases: MatrixFreeSolverConfig

Configuration for NewtonKrylov solver compilation.

precision

Numerical precision for computations.

Type:

PrecisionDType

solver_width

Solver vector length.

Type:

int

max_iters

Maximum Newton iterations permitted, defaulting to eight.

Type:

int

norm_fn

Compiled correction norm for convergence checks.

Type:

Optional[Callable]

residual_fn

Device function evaluating residuals.

Type:

Optional[Callable]

krylov_linear_solver_fn

Device function for solving linear systems.

Type:

Optional[Callable]

delta_location

Memory location for delta buffer.

Type:

str

residual_location

Memory location for residual buffer.

Type:

str

krylov_iters_local_location

Memory location for the single-element Krylov iteration counter buffer.

Type:

str

prev_theta_location

Memory location for the persistent contraction-estimate buffer.

Type:

str

Notes

Tolerance arrays (newton_atol, newton_rtol) are managed by the solver’s norm factory and accessed via NewtonKrylov.newton_atol/newton_rtol properties.

delta_location: str
krylov_iters_local_location: str
krylov_linear_solver_fn: Callable | None
max_iters: int
prev_theta_location: str
residual_fn: Callable | None
residual_location: str
use_cached_auxiliaries: bool
class cubie.integrators.matrix_free_solvers.NewtonKrylovCache(nonlinear_solver_fn: Callable)[source]

Bases: CUDADispatcherCache

Cache container for NewtonKrylov outputs.

nonlinear_solver_fn

Compiled CUDA device function for Newton-Krylov solving.

Type:

Callable

nonlinear_solver_fn: Callable