Optimization API#

The optimization namespace separates problem descriptions, globalization strategies, and solver classes. All public solvers return (solution, final_cost, history).

Problem classes#

class geojax.optimization.Minimize(*, M, cost, x0=None, solver=None, key=None, grad=None, egrad=None, precon=None, ehess_vec=None, rhess_vec=None)#

Bases: object

Riemannian minimization problem.

Parameters:

  • M (Any) – Manifold object.

  • cost (CostFn) – Scalar objective cost(x) to minimize.

  • grad (Optional[GradFn]) – Optional Riemannian gradient grad(x). If supplied, it takes priority over egrad and autodiff.

  • egrad (Optional[GradFn]) – Optional ambient Euclidean gradient egrad(x). It is converted to a Riemannian gradient using M.egrad_to_rgrad.

  • precon (Optional[PreconFn]) – Optional preconditioner precon(x, grad). If omitted, the identity preconditioner is used.

  • ehess_vec (Optional[HessVecFn]) – Optional ambient Euclidean Hessian-vector product. Automatic conversion requires the geometry to advertise exact ehess_to_rhess support.

  • rhess_vec (Optional[HessVecFn]) – Optional Riemannian Hessian-vector product. Supply this for second-order methods whenever the geometry does not advertise an exact automatic conversion.

  • x0 (Optional[Array])

  • solver (Optional[Any])

  • key (Optional[Array | int])

split_key()#

Return a fresh JAX PRNG key and advance the problem key.

Return type:

Any

solve()#

Solve the problem using the configured class-style solver.

Return type:

tuple[Any, float, List[InfoEntry]]

ehess_vec(x, u)#

Ambient Euclidean Hessian-vector product.

When only a Riemannian grad callback is available, this method returns its directional derivative instead. rhess_vec handles that path separately through the geometry’s advertised connection support.

Parameters:

  • x (Any)

  • u (Any)

Return type:

Any

rhess_vec(x, u)#

Riemannian Hessian-vector product.

Parameters:

  • x (Any)

  • u (Any)

Return type:

Any

hessian_operator(x)#

Return u -> rhess_vec(x, u).

Parameters:

x (Any)

Return type:

Callable[[Any], Any]

class geojax.optimization.LeastSquares(*, M, residual, x0=None, solver=None, key=None, jacobian_vec=None, adjoint_jacobian=None, precon=None)#

Bases: Minimize

Nonlinear least-squares problem on a manifold.

The objective is 0.5 * ||residual(x)||_2**2. Residuals may be arrays or arbitrary pytrees of arrays. Jacobian-vector and adjoint-Jacobian products are obtained with JAX autodiff unless callbacks are supplied.

Parameters:

  • M (Any)

  • residual (ResidualFn)

  • x0 (Optional[Array])

  • solver (Optional[Any])

  • key (Optional[Array | int])

  • jacobian_vec (Optional[JacobianVecFn])

  • adjoint_jacobian (Optional[AdjointJacobianFn])

  • precon (Optional[Callable[[Array, Array], Array]])

residual_value(x)#

Evaluate the residual pytree.

Parameters:

x (Any)

Return type:

Any

residual_norm(x)#

Return the Euclidean norm of the residual.

Parameters:

x (Any)

Return type:

Any

jacobian_vec(x, u)#

Apply the residual Jacobian to tangent vector u.

Parameters:

  • x (Any)

  • u (Any)

Return type:

Any

adjoint_jacobian(x, z)#

Apply the adjoint residual Jacobian and return a tangent vector.

Parameters:

  • x (Any)

  • z (Any)

Return type:

Any

normal_operator(x, u, damping=0.0)#

Apply J(x)^* J(x) + damping * I to u.

Parameters:

  • x (Any)

  • u (Any)

  • damping (float)

Return type:

Any

class geojax.optimization.FiniteSum(*, M, loss, num_terms, x0=None, solver=None, key=None, precon=None)#

Bases: Minimize

Finite-sum problem mean_i loss(x, i) for stochastic solvers.

Parameters:

  • M (Any)

  • loss (LossFn)

  • num_terms (int)

  • x0 (Optional[Array])

  • solver (Optional[Any])

  • key (Optional[Array | int])

  • precon (Optional[Callable[[Array, Array], Array]])

batch_cost_and_grad(x, indices)#

Return mini-batch cost and Riemannian gradient.

Parameters:

  • x (Any)

  • indices (Any)

Return type:

tuple[Any, Any]

sample_batch(key, batch_size, *, replace=True)#

Draw uniformly distributed term indices.

Parameters:

  • key (Any)

  • batch_size (int)

  • replace (bool)

Return type:

Any

Line searches#

class geojax.optimization.LineSearchProtocol(*args, **kwargs)#

Bases: Protocol

Protocol implemented by all public line-search strategies.

class geojax.optimization.ConstantStep(stepsize=1.0, normalize_step=False)#

Bases: object

Take a fixed multiplier or fixed Riemannian-length step.

Parameters:

  • stepsize (float)

  • normalize_step (bool)

class geojax.optimization.BacktrackingArmijo(contraction_factor=0.5, sufficient_decrease=0.0001, max_steps=25, initial_stepsize=1.0, normalize_step=True)#

Bases: object

Monotone Armijo search along a retraction curve.

Parameters:

  • contraction_factor (float)

  • sufficient_decrease (float)

  • max_steps (int)

  • initial_stepsize (float)

  • normalize_step (bool)

class geojax.optimization.AdaptiveArmijo(contraction_factor=0.5, sufficient_decrease=0.0001, max_steps=25, initial_stepsize=1.0, normalize_step=True, optimism=2.0)#

Bases: BacktrackingArmijo

Armijo search initialized from progress in the preceding iteration.

Parameters:

  • contraction_factor (float)

  • sufficient_decrease (float)

  • max_steps (int)

  • initial_stepsize (float)

  • normalize_step (bool)

  • optimism (float)

class geojax.optimization.StrongWolfe(sufficient_decrease=0.0001, curvature=0.9, initial_stepsize=1.0, expansion=2.0, max_stepsize=50.0, max_steps=20, max_zoom_steps=25, normalize_step=True)#

Bases: object

Strong-Wolfe search using transported directional derivatives.

The derivative at a trial point is evaluated by pairing its Riemannian gradient with the transported initial direction. This is exact for a geodesic paired with parallel transport and is the standard vector- transport proxy for a general retraction.

Parameters:

  • sufficient_decrease (float)

  • curvature (float)

  • initial_stepsize (float)

  • expansion (float)

  • max_stepsize (float)

  • max_steps (int)

  • max_zoom_steps (int)

  • normalize_step (bool)

class geojax.optimization.LineSearchState(previous_cost=None, previous_directional_derivative=None, previous_alpha=None, previous_stepsize=None)#

Bases: object

Information carried between consecutive line searches.

Parameters:

  • previous_cost (float | None)

  • previous_directional_derivative (float | None)

  • previous_alpha (float | None)

  • previous_stepsize (float | None)

class geojax.optimization.LineSearchResult(point, cost, gradient, stepsize, alpha, stats, state)#

Bases: object

Point, values, diagnostics, and reusable state from a line search.

Parameters:

Smooth solvers#

class geojax.optimization.SteepestDescent(requires_gradient=True, tolgradnorm=1e-06, maxiter=1000, maxtime=inf, minstepsize=1e-10, verbosity=2, line_search=<factory>, statsfun=None, stopfun=None, key=None)#

Bases: object

Minimize a smooth objective along its negative Riemannian gradient.

Parameters:

  • requires_gradient (bool)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • line_search (LineSearchProtocol)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

  • key (Any | None)

solve(problem)#

Solve problem and return (solution, final_cost, history).

Parameters:

problem (Any)

Return type:

tuple[Any, float, List[InfoEntry]]

class geojax.optimization.ConjugateGradient(tolgradnorm=1e-06, maxiter=1000, maxtime=inf, minstepsize=1e-10, verbosity=2, beta_type='H-S', orth_value=inf, line_search=<factory>, statsfun=None, stopfun=None, key=None)#

Bases: object

Nonlinear Riemannian conjugate-gradient solver.

beta_type may be 'S-D'/'steep', 'F-R', 'P-R', 'H-S', 'H-Z', 'L-S', 'P-R-SATO' or 'H-S-SATO'. The default is 'H-S', following Manopt.

Parameters:

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • beta_type (str)

  • orth_value (float)

  • line_search (LineSearchProtocol)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

  • key (Any | None)

solve(problem)#

Solve problem and return (solution, final_cost, history).

Parameters:

problem (Any)

Return type:

tuple[Any, float, List[InfoEntry]]

class geojax.optimization.BarzilaiBorwein(requires_gradient: 'bool' = True, bb_type: 'str' = 'alternate', initial_stepsize: 'float' = 1.0, min_stepsize: 'float' = 1e-12, max_stepsize: 'float' = 1000000000000.0, line_search: 'LineSearchProtocol' = <factory>, tolgradnorm: 'float' = 1e-06, maxiter: 'int' = 1000, maxtime: 'float' = inf, minstepsize: 'float' = 1e-14, verbosity: 'int' = 2, statsfun: 'Optional[StatsFn]' = None, stopfun: 'Optional[StopFn]' = None)#

Bases: object

Parameters:

  • requires_gradient (bool)

  • bb_type (str)

  • initial_stepsize (float)

  • min_stepsize (float)

  • max_stepsize (float)

  • line_search (LineSearchProtocol)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.LBFGS(requires_gradient: 'bool' = True, memory: 'int' = 10, tolgradnorm: 'float' = 1e-06, maxiter: 'int' = 1000, maxtime: 'float' = inf, minstepsize: 'float' = 1e-10, verbosity: 'int' = 2, line_search: 'LineSearchProtocol' = <factory>, cautious_update: 'bool' = True, cautious_threshold: 'float' = 1e-10, statsfun: 'Optional[StatsFn]' = None, stopfun: 'Optional[StopFn]' = None)#

Bases: object

Parameters:

  • requires_gradient (bool)

  • memory (int)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • line_search (LineSearchProtocol)

  • cautious_update (bool)

  • cautious_threshold (float)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.NewtonCG(requires_gradient=True, tolgradnorm=1e-06, maxiter=200, maxtime=inf, minstepsize=1e-10, verbosity=2, maxinner=100, cg_relative_tolerance=None, cg_absolute_tolerance=1e-10, curvature_tolerance=0.0, line_search=<factory>, statsfun=None, stopfun=None)#

Bases: object

Riemannian Newton method with an inexact Hessian solve.

Hessian-vector products come from problem.rhess_vec. The inner conjugate-gradient solve is truncated on non-positive curvature; a non-descent Newton direction falls back to preconditioned steepest descent.

Parameters:

  • requires_gradient (bool)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • maxinner (int)

  • cg_relative_tolerance (float | None)

  • cg_absolute_tolerance (float)

  • curvature_tolerance (float)

  • line_search (LineSearchProtocol)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.TrustRegions(requires_gradient=True, tolgradnorm=1e-06, maxiter=200, maxtime=inf, minstepsize=0.0, verbosity=2, initial_radius=1.0, max_radius=100.0, rho_prime=0.1, kappa=0.1, theta=1.0, maxinner=250, statsfun=None, stopfun=None)#

Bases: object

Approximate Riemannian trust-regions method.

Uses truncated conjugate gradient to approximately minimize the quadratic model. A geometry must advertise an exact automatic Hessian conversion or the problem must supply rhess_vec explicitly.

Parameters:

  • requires_gradient (bool)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • initial_radius (float)

  • max_radius (float)

  • rho_prime (float)

  • kappa (float)

  • theta (float)

  • maxinner (int)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.AdaptiveRegularizationCubics(requires_gradient=True, tolgradnorm=1e-06, maxiter=200, maxtime=inf, minstepsize=0.0, verbosity=2, initial_sigma=1.0, min_sigma=1e-12, max_sigma=1000000000000.0, acceptance_threshold=0.1, very_successful_threshold=0.9, decrease_factor=0.5, increase_factor=2.0, subproblem_iterations=10, subproblem_tolerance=0.1, subproblem_backtracks=20, statsfun=None, stopfun=None)#

Bases: object

Adaptive cubic regularization on a Riemannian manifold.

At x the solver approximately minimizes

<g, eta> + 0.5 <eta, Hess f(x)[eta]> + sigma/3 ||eta||^3

in the tangent space. The subproblem begins at its exact Cauchy point and may be refined by model-gradient steps. The ratio of actual to predicted decrease controls acceptance and the next regularization parameter.

Parameters:

  • requires_gradient (bool)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • initial_sigma (float)

  • min_sigma (float)

  • max_sigma (float)

  • acceptance_threshold (float)

  • very_successful_threshold (float)

  • decrease_factor (float)

  • increase_factor (float)

  • subproblem_iterations (int)

  • subproblem_tolerance (float)

  • subproblem_backtracks (int)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

Structured objectives#

class geojax.optimization.GaussNewton(requires_gradient=True, tolgradnorm=1e-06, maxiter=200, maxtime=inf, minstepsize=1e-10, verbosity=2, maxinner=100, cg_relative_tolerance=0.001, cg_absolute_tolerance=1e-10, line_search=<factory>, statsfun=None, stopfun=None)#

Bases: object

Gauss-Newton using matrix-free normal equations in each tangent space.

Parameters:

  • requires_gradient (bool)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • maxinner (int)

  • cg_relative_tolerance (float)

  • cg_absolute_tolerance (float)

  • line_search (LineSearchProtocol)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.LevenbergMarquardt(requires_gradient=True, tolgradnorm=1e-06, maxiter=200, maxtime=inf, minstepsize=0.0, verbosity=2, initial_damping=0.001, min_damping=1e-12, max_damping=1000000000000.0, damping_increase=2.0, damping_decrease=3.0, acceptance_threshold=0.0001, maxinner=100, cg_relative_tolerance=0.001, cg_absolute_tolerance=1e-10, statsfun=None, stopfun=None)#

Bases: object

Damped Gauss-Newton method with gain-ratio damping updates.

Parameters:

  • requires_gradient (bool)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • initial_damping (float)

  • min_damping (float)

  • max_damping (float)

  • damping_increase (float)

  • damping_decrease (float)

  • acceptance_threshold (float)

  • maxinner (int)

  • cg_relative_tolerance (float)

  • cg_absolute_tolerance (float)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.StochasticGradient(requires_gradient=True, batch_size=1, step_schedule=<factory>, momentum=0.0, clip_norm=None, replace=True, evaluation_period=10, tolgradnorm=0.0, maxiter=1000, maxtime=inf, minstepsize=0.0, verbosity=2, key=None, statsfun=None, stopfun=None)#

Bases: object

Mini-batch Riemannian stochastic gradient with optional momentum.

Parameters:

  • requires_gradient (bool)

  • batch_size (int)

  • step_schedule (StepScheduleProtocol)

  • momentum (float)

  • clip_norm (float | None)

  • replace (bool)

  • evaluation_period (int)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • key (Any | int | None)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.StepScheduleProtocol(*args, **kwargs)#

Bases: Protocol

Callable schedule returning the multiplier for iteration k.

class geojax.optimization.ConstantSchedule(stepsize=0.01)#

Bases: object

Constant stochastic-gradient step multiplier.

Parameters:

stepsize (float)

class geojax.optimization.PolynomialDecay(initial_stepsize=0.1, decay_rate=0.01, power=0.5, minimum_stepsize=0.0)#

Bases: object

Schedule initial_stepsize / (1 + decay_rate * k)**power.

Parameters:

  • initial_stepsize (float)

  • decay_rate (float)

  • power (float)

  • minimum_stepsize (float)

class geojax.optimization.CosineDecay(initial_stepsize=0.1, final_stepsize=0.0, decay_steps=1000)#

Bases: object

Cosine interpolation from initial_stepsize to final_stepsize.

Parameters:

  • initial_stepsize (float)

  • final_stepsize (float)

  • decay_steps (int)

class geojax.optimization.AlternatingGradient(requires_gradient=True, block_order=None, tolgradnorm=1e-06, maxiter=1000, maxtime=inf, minstepsize=1e-10, verbosity=2, line_search=<factory>, statsfun=None, stopfun=None)#

Bases: object

Cycle through Product factors using one gradient block at a time.

block_order refers to the leaves of the Product factor pytree in JAX’s deterministic flattening order. When omitted, every leaf is visited once per outer iteration.

Parameters:

  • requires_gradient (bool)

  • block_order (Sequence[int] | None)

  • tolgradnorm (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • verbosity (int)

  • line_search (LineSearchProtocol)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

Derivative-free solvers#

class geojax.optimization.ParticleSwarm(requires_gradient: 'bool' = False, swarm_size: 'int' = 30, inertia: 'float' = 0.5, cognitive: 'float' = 1.5, social: 'float' = 1.5, initial_velocity_scale: 'float' = 0.1, maxiter: 'int' = 200, maxtime: 'float' = inf, minstepsize: 'float' = 0.0, tolgradnorm: 'float' = -inf, verbosity: 'int' = 2, statsfun: 'Optional[StatsFn]' = None, stopfun: 'Optional[StopFn]' = None)#

Bases: object

Parameters:

  • requires_gradient (bool)

  • swarm_size (int)

  • inertia (float)

  • cognitive (float)

  • social (float)

  • initial_velocity_scale (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • tolgradnorm (float)

  • verbosity (int)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

class geojax.optimization.NelderMead(requires_gradient: 'bool' = False, initial_scale: 'float' = 0.1, reflection: 'float' = 1.0, expansion: 'float' = 2.0, contraction: 'float' = 0.5, shrink: 'float' = 0.5, tolcostspread: 'float' = 1e-10, maxiter: 'int' = 1000, maxtime: 'float' = inf, minstepsize: 'float' = 0.0, tolgradnorm: 'float' = -inf, verbosity: 'int' = 2, statsfun: 'Optional[StatsFn]' = None, stopfun: 'Optional[StopFn]' = None)#

Bases: object

Parameters:

  • requires_gradient (bool)

  • initial_scale (float)

  • reflection (float)

  • expansion (float)

  • contraction (float)

  • shrink (float)

  • tolcostspread (float)

  • maxiter (int)

  • maxtime (float)

  • minstepsize (float)

  • tolgradnorm (float)

  • verbosity (int)

  • statsfun (Callable[[Any, Any, InfoEntry], Dict[str, Any]] | None)

  • stopfun (Callable[[Any, Any, InfoEntry], Tuple[bool, str]] | None)

Iteration records#

class geojax.optimization.InfoEntry(iter, cost, gradnorm, stepsize, time, linesearch=None, beta=None, reason='', extra=<factory>)#

Bases: object

Per-iteration optimization statistics.

This mirrors the useful parts of Manopt’s info struct-array in a Python dataclass. beta is used by conjugate-gradient methods and left as None by algorithms that do not define a beta parameter.

Parameters:

  • iter (int)

  • cost (float)

  • gradnorm (float)

  • stepsize (float)

  • time (float)

  • linesearch (LineSearchStats | None)

  • beta (float | None)

  • reason (str)

  • extra (Dict[str, Any])

class geojax.optimization.LineSearchStats(costevals, stepsize, alpha, accepted, gradevals=0, method='', reason='')#

Bases: object

Common diagnostics returned by a line-search strategy.

Parameters:

  • costevals (int)

  • stepsize (float)

  • alpha (float)

  • accepted (bool)

  • gradevals (int)

  • method (str)

  • reason (str)