Nonlinear solvers

Newton and BoxConstrainedActiveSet are standalone solvers for a bordered system R(u, lam) = 0, g(u, lam) = 0. The engines of ds.pdesolve run their own corrector, described in nonlinear solvers.

Warning

BoxConstrainedActiveSet is experimental.

diff_soup.Newton

Newton's method for a bordered system R(u, lam) = 0, g(u, lam) = 0.

It converges on a weighted RMS norm of the update, scaled by the observed contraction rate. A contraction rate above 0.9 stops it as stalled, which terminate reports as NONLINEAR_DIVERGED.

diff_soup.BoxConstrainedActiveSet

Newton's method with box bounds on selected state entries, by an active-set iteration.

Entries held at a bound are pinned, and iterates are clipped into the box. It stops when the max-norm of the pinned residual falls to config.atol. Its derivative uses the forward implicit rule under adjoint=ForwardAdjoint() and the reverse rule otherwise.

Attributes:

Name	Type	Description
config	BoxConstrainedUnifiedConfig	Iteration and bound settings, with the uniform lower_scalar, upper_scalar, and constrained_indices.
bounds	Optional[BoxBounds]	Per-entry BoxBounds, used only when lower or upper is a tuple, otherwise the bounds come from config.

config class-attribute instance-attribute

config: BoxConstrainedUnifiedConfig = eqx.field(static=True, default_factory=BoxConstrainedUnifiedConfig)

bounds class-attribute instance-attribute

bounds: Optional[BoxBounds] = eqx.field(static=True, default=None)

diff_soup.nonlinear.BoxBounds dataclass

Lower and upper bounds on selected entries of the state.

Select the entries with exactly one of mask or indices. Bounds are Python floats or tuples with one value per selected entry, never JAX arrays.

Parameters:

Name	Type	Description	Default
lower	_FloatLike	Lower bound, a float or one value per selected entry.	0.0
upper	_FloatLike	Upper bound, strictly above lower at every selected entry.	1.0
mask	Optional[tuple[bool, ...]]	One boolean per state entry, True where bounded.	None
indices	Optional[tuple[int, ...]]	Indices of the bounded state entries.	None

resolve

resolve(*, n_state: int, dtype) -> '_ResolvedBounds'

Expand the bounds to per-entry arrays for a state of n_state entries.

Parameters:

Name	Type	Description	Default
n_state	int	Length of the state.	required
dtype		Dtype of the bound arrays.	required

Returns:

Type	Description
'_ResolvedBounds'	The selection mask and full-length lower and upper bound arrays, zero outside the selection.

Raises:

Type	Description
ValueError	If the selection or the bounds do not fit n_state, or lower is not below upper.

---

diff_soup.nonlinear.NonlinearSolver

The details matter only when writing a corrector of your own.

diff_soup.nonlinear.NonlinearSolver

Base class of the Newton correctors for a bordered system R(u, lam) = 0, g(u, lam) = 0.

solve runs the whole iteration, and init, step, and terminate drive it one iteration at a time. A solve never raises on non-convergence, so read the returned status.

Attributes:

Name	Type	Description
config	NonlinearConfig	Iteration settings, a NonlinearConfig.
linear	Optional[LinearSolver]	Linear solver stored on the corrector and not read by its methods, which take backend.
adjoint	AbstractAdjoint	Differentiation strategy, used only by BoxConstrainedActiveSet.

config class-attribute instance-attribute

config: NonlinearConfig = eqx.field(static=True, default_factory=NonlinearConfig)

linear class-attribute instance-attribute

linear: Optional[LinearSolver] = None

adjoint class-attribute instance-attribute

adjoint: AbstractAdjoint = eqx.field(default_factory=ImplicitAdjoint)

init

init(operator: Any, u0: Array, lam0: Array, theta: Optional[Array], ctx: Any) -> NonlinearState

Return the iteration state at (u0, lam0) before the first step.

Parameters:

Name	Type	Description	Default
operator	Any	Unused by the base class.	required
u0	Array	Initial state.	required
lam0	Array	Initial global unknowns.	required
theta	Optional[Array]	Unused by the base class.	required
ctx	Any	Unused by the base class.	required

Returns:

Type	Description
NonlinearState	A NonlinearState with the iteration count at zero.

step

step(operator: Any, backend: LinearSolver, analysis: Any, state: NonlinearState, theta: Optional[Array], ctx: Any) -> tuple[NonlinearState, dict]

Take one Newton iteration with a freshly evaluated Jacobian.

The caller stops on terminate and bounds the number of iterations.

Parameters:

Name	Type	Description	Default
operator	Any	Residual and Jacobian provider, as in solve.	required
backend	LinearSolver	Linear solver for the state block.	required
analysis	Any	Analysis of the state block's sparsity.	required
state	NonlinearState	Current NonlinearState.	required
theta	Optional[Array]	Parameters passed to operator.	required
ctx	Any	Context passed to operator.	required

Returns:

Type	Description
tuple[NonlinearState, dict]	The next NonlinearState and a dict with "update_norm" and "iteration".

terminate

terminate(operator: Any, state: NonlinearState) -> tuple[Array, RESULT]

Report whether to stop, and the status of the iteration so far.

Parameters:

Name	Type	Description	Default
operator	Any	Unused by the base class.	required
state	NonlinearState	Current NonlinearState.	required

Returns:

Type	Description
tuple[Array, RESULT]	A boolean stop flag and a RESULT, which is SUCCESS when converged, NONLINEAR_DIVERGED when convergence stalled, and UNKNOWN otherwise.

solve

solve(operator: Any, backend: LinearSolver, analysis: Any, u0: Array, lam0: Array, theta: Optional[Array], ctx: Any, *, frozen_blocks: Optional[tuple] = None, factor_refit: bool = True, mass_refresh_fn: Optional[Any] = None, early_exit: bool = False) -> tuple[Array, Array, NewtonInfo]

Run Newton's method from (u0, lam0) and return the last iterate with its diagnostics.

Parameters:

Name	Type	Description	Default
operator	Any	Object with residual(u, lam, theta, ctx) returning (R, g) and jacobian_blocks(u, lam, theta, ctx).	required
backend	LinearSolver	Linear solver for the state block.	required
analysis	Any	Analysis of the state block's sparsity, from diff_soup.linear.analyze.	required
u0	Array	Initial state.	required
lam0	Array	Initial global unknowns.	required
theta	Optional[Array]	Parameters passed to operator.	required
ctx	Any	Context passed to operator.	required
frozen_blocks	Optional[tuple]	Jacobian blocks to factor in place of evaluating them at u0, used when config.jacobian_cache is "reuse".	None
factor_refit	bool	Accepted and ignored.	True
mass_refresh_fn	Optional[Any]	Function of (u, lam) returning ctx.mass_diag for each iterate.	None
early_exit	bool	If True, stop at convergence with a loop that is not reverse-differentiable.	False

Returns:

Type	Description
tuple[Array, Array, NewtonInfo]	The final u, the final lam, and a NewtonInfo.

Raises:

Type	Description
TypeError	If early_exit is set while u0 is traced and config.safe_mode is off.

diff_soup.nonlinear.NewtonInfo

Diagnostics of one Newton solve, returned by NonlinearSolver.solve.

Attributes:

Name	Type	Description
converged	Array	Whether the convergence test passed.
iterations	Array	Newton iterations taken.
residual_norm	Array	Max-norm of the residual [R; g] at the returned iterate.
tooslow	Array	Whether the contraction rate stalled before convergence.
ndecomps	Array	Factorizations, equal to factorizations.
nsolves	Array	Linear solves.
nfevals	Array	Residual evaluations, including the final one.
npds	Array	Jacobian evaluations, equal to jacobian_evals.
jacobian_evals	Array	Jacobian evaluations.
factorizations	Array	Factorizations of the state block.
solve_only	Array	Solves against an existing factorization.
cached_retries	Array	1 when the solve retried with fresh Jacobians after a stall, else 0.
rho	Array	Last observed contraction rate.
residual	Array	Residual [R; g] at the returned iterate when config.return_residual is set, else an empty array.