gauss_newton_cg

jaxns.experimental.solvers.gauss_newton_cg

Module Contents

DomainType

ObjectiveRet

CT

convert_to_real(x)

Return a real-valued twin of x and a merge-back function.

Parameters:

x (CT)

Return type:

Tuple[_CT, Callable[[_CT], CT]]

class NewtonDiagnostic

Bases: NamedTuple

iteration: jaxns.internals.types.IntArray

g_norm: jaxns.internals.types.FloatArray

mu: jaxns.internals.types.FloatArray

damping: jaxns.internals.types.FloatArray

cg_iters: jaxns.internals.types.IntArray

f: jaxns.internals.types.FloatArray

f_prop: jaxns.internals.types.FloatArray

f_quad: jaxns.internals.types.FloatArray

delta_f_pred: jaxns.internals.types.FloatArray

delta_f_actual: jaxns.internals.types.FloatArray

gain_ratio: jaxns.internals.types.FloatArray

accepted: jaxns.internals.types.BoolArray

in_trust_region: jaxns.internals.types.BoolArray

delta_x_norm: jaxns.internals.types.FloatArray

ddelta_x_norm: jaxns.internals.types.FloatArray

newton_cg_solver(obj_fn, x0, args=(), maxiter=100, maxiter_cg=100, gtol=3e-05, p_accept=0.01, p_lower=0.25, p_upper=1.1, mu_init=1.0, mu_min=1e-06, mu_in_factor=5, mu_out_factor=0.1, approx_hvp=False, verbose=False)

Trust-region Newton-CG minimiser.

Identical call signature and adaptive-μ logic as lm_solver, but uses a scalar objective instead of residuals, and solves

(H + damping·I) δx = -∇f

by CG with Hessian–vector products.

Returns:

• x_final (pytree matching x0 (merged back to complex if needed))

• diagnostics (NewtonDiagnostic[…] array with length = maxiter)

Parameters:

• obj_fn (Callable[Ellipsis, ObjectiveRet])

• x0 (DomainType)

• args (tuple)

• maxiter (int)

• maxiter_cg (int)

• gtol (float)

• p_accept (float)

• p_lower (float)

• p_upper (float)

• mu_init (float)

• mu_min (float)

• mu_in_factor (float)

• mu_out_factor (float)

• approx_hvp (bool)

• verbose (bool)

Return type:

Tuple[DomainType, NewtonDiagnostic]