neml2.solvers#

Python-native linear and nonlinear solvers.

_result – RetCode, NonlinearResult.
_exceptions – ConvergenceError (recoverable solve failure).
dense_lu – DenseLU.
schur_complement – SchurComplement.
gmres – GMRES (matrix-free iterative).
bicgstab – BiCGStab (matrix-free iterative).
newton – Newton.
newton_linesearch – NewtonWithLineSearch.
_helpers – private reduction / scaling helpers.

class neml2.solvers.BiCGStab(*, restart=40, max_its=1000, abs_tol=0.0, rel_tol=0.0001, preconditioner=None, cache_strategy='none', cache_max_its=10)[source]#

Bases: _KrylovSolver

Matrix-free BiCGStab linear solver for the implicit Newton step.

A short-recurrence alternative to GMRES (no growing Krylov basis, so a bounded per-iteration memory), two matvecs per iteration. Like GMRES it is matrix-free (applies J.v via Matvec) and configured as a Newton linear_solver. Has no restart width (the max_its budget bounds the iteration count).

Parameters:

restart (int)
max_its (int)
abs_tol (float)
rel_tol (float)
preconditioner (Preconditioner | None)
cache_strategy (str)
cache_max_its (int)

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

BiCGStab

hit = <neml2.schema.HitSchema object>#

method = 'bicgstab'#: Krylov method tag ("gmres" / "bicgstab"); overridden per subclass.

class neml2.solvers.BlockJacobiPreconditioner[source]#

Bases: Preconditioner

Block-Jacobi: M = blockdiag(A) over the per-variable on-diagonal blocks.

apply_module(system)[source]#

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

BlockJacobiPreconditioner

hit = <neml2.schema.HitSchema object>#

kind = 'block_jacobi'#: Metadata tag (informational; the behavior is defined by the graphs).

setup_module(system)[source]#

The setup module, or None for the identity (no) preconditioner.

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module

exception neml2.solvers.ConvergenceError#: Bases: RuntimeError

class neml2.solvers.DenseLU[source]#

Bases: object

Dense LU linear solver. Assembles the (possibly) sparse matrix into a dense one and uses a standard LU decomposition.

The dense-vs-block-diagonal dispatch is implicit in the wrapped Tensor operand: Tensor.solve forwards to torch.linalg.solve on the trailing matrix dims, so leading sub-batch axes (if any) batch naturally as independent LUs per site – no special-case code at this layer.

SECTION = 'Solvers'#

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

DenseLU

hit = <neml2.schema.HitSchema object>#

solve(A: AssembledMatrix, b: AssembledVector) → AssembledVector[source]#
solve(A: AssembledMatrix, b: AssembledMatrix) → AssembledMatrix

class neml2.solvers.FullPreconditioner[source]#

Bases: Preconditioner

Full (LU) preconditioner: M = A (exact per-solve; with a caching cache_strategy this is modified Newton – factor once, reuse).

apply_module(system)[source]#

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

FullPreconditioner

hit = <neml2.schema.HitSchema object>#

kind = 'full'#: Metadata tag (informational; the behavior is defined by the graphs).

setup_module(system)[source]#

The setup module, or None for the identity (no) preconditioner.

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module

class neml2.solvers.GMRES(*, restart=40, max_its=1000, abs_tol=0.0, rel_tol=0.0001, preconditioner=None, cache_strategy='none', cache_max_its=10)[source]#

Bases: _KrylovSolver

Matrix-free restarted GMRES(m) linear solver for the implicit Newton step.

Never assembles the Jacobian: each inner iteration applies J.v (the compiled/eager Matvec). Best for the well-conditioned near-identity systems of implicit backward-Euler updates, where a handful of matvecs beat an O(N^3) dense factorization. Configure it as a Newton linear_solver.

Parameters:

restart (int)
max_its (int)
abs_tol (float)
rel_tol (float)
preconditioner (Preconditioner | None)
cache_strategy (str)
cache_max_its (int)

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

GMRES

hit = <neml2.schema.HitSchema object>#

method = 'gmres'#: Krylov method tag ("gmres" / "bicgstab"); overridden per subclass.

class neml2.solvers.JacobiPreconditioner[source]#

Bases: Preconditioner

Point-Jacobi: M = diag(A) (elementwise scaling).

apply_module(system)[source]#

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

JacobiPreconditioner

hit = <neml2.schema.HitSchema object>#

kind = 'jacobi'#: Metadata tag (informational; the behavior is defined by the graphs).

setup_module(system)[source]#

The setup module, or None for the identity (no) preconditioner.

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module

class neml2.solvers.Newton(*, linear_solver=None, atol=1e-10, rtol=1e-08, miters=25, substep_del_tol=1e-06)[source]#

Bases: object

The standard Newton-Raphson solver which always takes the ‘full’ Newton step.

Parameters:

linear_solver (_LinearSolver | None)
atol (float)
rtol (float)
miters (int)
substep_del_tol (float)

SECTION = 'Solvers'#: Inherited by NewtonWithLineSearch.

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

Newton

hit = <neml2.schema.HitSchema object>#

solve(system)[source]#

Solve the nonlinear system via the shared C++ Newton solver.

The iteration control (convergence test, line search) lives in C++ and is shared with the AOTI runtime; this wrapper supplies the residual / Newton-step callbacks (the RHS residual + the Jacobian operator chained into LinearSolve, the same modules the AOTI runtime compiles, run eagerly here) plus the per-group layouts, then commits the converged iterate back into the system.

A matrix-free iterative linear_solver (GMRES / BiCGStab) takes the sibling _solve_iterative() path – the same shared C++ Newton loop, but each step’s linear solve is a Krylov iteration over the Matvec callback rather than a baked direct solve.

Parameters:: system (ModelNonlinearSystem)
Return type:: NonlinearResult

class neml2.solvers.NewtonWithLineSearch(*, linear_solver=None, atol=1e-10, rtol=1e-08, miters=25, substep_del_tol=1e-06, linesearch_type='BACKTRACKING', max_linesearch_iterations=10, linesearch_cutback=2.0, linesearch_stopping_criteria=0.001, check_negative_criterion=False)[source]#

Bases: Newton

The Newton-Raphson solver with line search.

Parameters:

linear_solver (DenseLU | SchurComplement | None)
atol (float)
rtol (float)
miters (int)
substep_del_tol (float)
linesearch_type (str)
max_linesearch_iterations (int)
linesearch_cutback (float)
linesearch_stopping_criteria (float)
check_negative_criterion (bool)

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

NewtonWithLineSearch

hit = <neml2.schema.HitSchema object>#

class neml2.solvers.NoPreconditioner[source]#

Bases: Preconditioner

Identity preconditioner (unpreconditioned Krylov).

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

NoPreconditioner

hit = <neml2.schema.HitSchema object>#

kind = 'none'#: Metadata tag (informational; the behavior is defined by the graphs).

class neml2.solvers.NonlinearResult(ret, iterations, log=())[source]#

Bases: object

Result metadata returned by a nonlinear solve.

Parameters:

ret (RetCode)
iterations (int)
log (tuple[str, ...])

iterations: int#

log: tuple[str, ...] = ()#: Per-iteration convergence history (ITERATION ... / LS ITERATION ... lines, max-norm over the batch). Opt-in data path: populated only when the C++ solver’s collect_log is set; empty otherwise. Console verbosity is separate – see the newton channel in neml2.log.

ret: RetCode#

class neml2.solvers.Preconditioner[source]#

Bases: object

Base preconditioner. Concrete subclasses set kind and return their setup/apply modules; the base is the none (identity) case.

SECTION = 'Solvers'#

apply_module(system)[source]#

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module | None

is_preconditioner = True#: Marks a preconditioner object (GMRES/BiCGStab + the exporter dispatch on it).

kind = 'none'#: Metadata tag (informational; the behavior is defined by the graphs).

setup_module(system)[source]#

The setup module, or None for the identity (no) preconditioner.

Parameters:: system (ModelNonlinearSystem)
Return type:: nn.Module | None

class neml2.solvers.RetCode(value)[source]#

Bases: Enum

Nonlinear solver return code.

FAILURE = 2#

MAXITER = 1#

SUCCESS = 0#

class neml2.solvers.SchurComplement(*, residual_primary_group=0, unknown_primary_group=0, primary_solver=None, schur_solver=None)[source]#

Bases: object

Schur complement linear solver. Solves a block-partitioned system A x = b by forming and solving the Schur complement of the primary block.

The six-step factorisation reads naturally on top of the Tensor-backed AssembledMatrix / AssembledVector arithmetic – A_pp.solve(...), A_sp @ Y, A_ss - A_sp @ Y all forward to the typed primitive.

Parameters:

residual_primary_group (int)
unknown_primary_group (int)

SECTION = 'Solvers'#

classmethod from_hit(node, factory)[source]#

Parameters:

node (nmhit.Node)
factory (_NativeInputFile)

Return type:

SchurComplement

hit = <neml2.schema.HitSchema object>#

solve(A: AssembledMatrix, b: AssembledVector) → AssembledVector[source]#
solve(A: AssembledMatrix, b: AssembledMatrix) → AssembledMatrix