Optimization solvers

ds.inference.optsolve(problem, optimizer=...) takes one of these solvers. See optimization solvers.

Recommended defaults

NewtonTrustRegion() is the default for fewer than about 100 parameters, and TrustRegion() for a lowered hierarchical study. For large problems such as embedded neural networks, run Optax with optax.adamw. Bound every run with max_steps.

diff_soup.inference.NewtonTrustRegion

The Newton trust-region method, steps on the Gauss-Newton matrix or the exact Hessian.

It steps on the Gauss-Newton matrix of a least-squares objective and keeps box bounds. hessian="true" restores the exact Hessian of previous releases. It is the default optimizer of ds.inference.optsolve.

Parameters:

Name	Type	Description	Default
max_steps		The step budget.	required
hessian		"auto" steps on J^T J when the problem has residual rows and on the exact Hessian otherwise, "gauss-newton" or "true" forces one.	required
subproblem		"exact", "cg", or "auto", which picks CG only for an exact Hessian above dense_cutoff parameters without HessianScale.	required
linear_solver		"eigh", "cholesky" or "qr" for the exact subproblem, or "auto" for "eigh" on the exact Hessian and Cholesky or QR otherwise.	required
dense_cutoff		The parameter count above which subproblem="auto" may pick CG and hessian_tile="auto" tiles the exact Hessian.	required
radius_init		The initial trust radius, in scaled parameter units.	required
radius_max		The largest trust radius.	required
eta		The ratio of achieved to predicted decrease a step must reach to be accepted.	required
grad_tol		The gradient-norm tolerance, or None for the dtype's default.	required
step_tol		The step-norm tolerance relative to the parameter norm, or None for the dtype's default.	required
cost_tol		The tolerance on the change in the objective relative to the objective, or None for the dtype's default.	required
cost_atol		An absolute floor under the cost-change test, zero when None.	required
x_scale		The parameter scaling, HessianScale() by default, with JacScale refused under hessian="true" and HessianScale under CG.	required
composition		"forward-over-reverse" or "forward-over-forward" for the exact Hessian, or "auto" to pick one the time-stepping adjoint supports.	required
hessian_refresh		How often the exact Hessian is recomputed, with SR1 updates between: "adaptive", "every", or every k steps for an integer k.	required
hessian_tile		How many directions of the exact Hessian or Gauss-Newton Jacobian are formed at once: 8 to 128, "full", or "auto" to fit the free memory.	required
cg_max_iter		The conjugate-gradient iteration budget, the parameter count when None.	required
callbacks		Callbacks run together with those passed to ds.inference.optsolve.	required

diff_soup.inference.TrustRegion

The trust-region Gauss-Newton method for a hierarchical study, with each step solved cluster by cluster.

It needs a study lowered with ds.lower(study). Box bounds are kept. Pass checkpoint= to ds.inference.optsolve to save the fit and resume it later.

Parameters:

Name	Type	Description	Default
hessian	Any	"gauss-newton" steps on the Gauss-Newton curvature, "true" on the exact Hessian.	'gauss-newton'
nested_at	Any	A level at which the fit splits into an outer step and a refit of each member's own parameters.	None
inner	Any	A TrustRegion without nested_at whose budget and tolerances the member refits use, or None for two iterations per member.	None
radius_init	Any	The initial trust radius, in scaled parameter units.	1.0
radius_max	Any	The largest trust radius.	1000.0
eta	Any	The ratio of achieved to predicted decrease a step must reach to be accepted.	0.0001
grad_tol	Any	The gradient-norm tolerance, or None to test no gradient.	None
step_tol	Any	The step-norm tolerance relative to the parameter norm, or None for the dtype's default.	None
cost_tol	Any	The tolerance on the change in the objective relative to the objective, or None for the dtype's default.	None
cost_atol	Any	An absolute floor under the cost-change test, zero when None.	None
stationarity_tol	Any	The relative predicted decrease below which the fit ends, or None for the larger of cost_tol and three times the measured noise.	None
max_consecutive_invalid_trials	Any	The number of trials in a row with failed solves that ends the fit as STALLED.	5
noise_trial_limit	Any	A rejected trial counts as noise only when its predicted decrease is within this multiple of the study's evaluation noise.	None
noise_relaxed_ratio	Any	If True, the ratio of achieved to predicted decrease allows for the measured evaluation noise.	False
gradient_floor	Any	The relative accuracy of the derivatives, below which the fit ends DERIVATIVE_ACCURACY, "auto" to derive it from the solvers' rtol, or None.	'auto'
max_steps	Any	The step budget, counted across resumes.	30
max_evaluations	Any	The budget of objective evaluations, which ends the fit as MAX_STEPS.	None
inner_start_iterations	Any	The number of member refit iterations at the start point, before the first outer step.	0
inner_fraction	Any	A member refit ends when its predicted decrease falls to this fraction of the outer step's.	0.01
inner_retreats	Any	How many times a member whose solve fails halves its step before the attempt fails.	3
target_accepted_steps	Any	Ends the fit as ACCEPTED_STEPS after this many accepted steps.	None
x_scale	Any	The parameter scaling, HessianScale() when None, JacScale() or a ConstantScale, with JacScale refused under hessian="true".	None
weights	Any	One weight per member of the first level, or a (level, weights) pair, multiplying each member's share of the objective.	None
weight_priors	Any	If True, the weights also multiply each member's prior and regularization terms.	True
constraints	Any	Linear equality constraints on the joint parameter vector, as LinearConstraint objects.	()
gauges	Any	Gauge constraints that pin directions the data leave free.	()
fixed_indices	Any	The indices of the joint parameter vector held at their starting values.	()
callbacks	Any	Callbacks run together with those passed to ds.inference.optsolve.	()

diff_soup.inference.LevenbergMarquardt

The Levenberg-Marquardt method, damped Gauss-Newton steps on a least-squares problem.

It needs the residual and its Jacobian, which a LeastSquaresProblem provides. Box bounds are kept.

Parameters:

Name	Type	Description	Default
max_steps		The step budget.	required
lam_init		The initial damping.	required
lam_max		The damping above which it restarts from lam_init.	required
nu_init		The factor the damping grows by on a rejection after an accepted step, itself doubled on each further rejection.	required
nu_start		The growth factor before the first accepted step, nu_init when None.	required
accept_cutoff		The ratio of achieved to predicted decrease a step must exceed to be accepted.	required
grad_tol		The gradient-norm tolerance, or None for the dtype's default.	required
step_tol		The step-norm tolerance relative to the parameter norm, or None for the dtype's default.	required
cost_tol		The tolerance on the change in the objective relative to the objective, or None for the dtype's default.	required
cost_atol		An absolute floor under the cost-change test, zero when None.	required
max_step_norm		The longest step allowed, or 0 for no limit.	required
linear_solver		"cholesky", "qr", "solve", or "auto" for QR in single precision with a full Jacobian and Cholesky otherwise.	required
x_scale		The parameter scaling of the damping, JacScale() by default.	required
jacobian_tile		How many Jacobian directions are formed at once: 8 to 128, "full", or "auto", whole when that fits the free memory.	required
callbacks		Callbacks run together with those passed to ds.inference.optsolve.	required

diff_soup.inference.BFGS

The limited-memory BFGS method, quasi-Newton steps with a backtracking line search.

It needs the objective and its gradient. Box bounds are kept by projecting each line-search trial. A failed line search at a non-finite objective or gradient ends the run with DIVERGED.

Parameters:

Name	Type	Description	Default
max_steps		The step budget.	required
grad_tol		The gradient-norm tolerance, or None for the dtype's default.	required
step_tol		The step-norm tolerance relative to the parameter norm, or None for the dtype's default.	required
c1		The sufficient-decrease parameter of the line search.	required
max_linesearch		The number of step halvings the line search tries per step.	required
memory		The number of past steps the curvature estimate is built from.	required
ad		"reverse" takes the gradient from the reverse adjoint, "forward" from the forward-mode Jacobian.	required
callbacks		Callbacks run together with those passed to ds.inference.optsolve.	required

diff_soup.inference.GradientDescent

Gradient descent, steps along the negative gradient with a fixed learning rate and optional momentum.

It needs the objective and its gradient. Box bounds are kept by projecting each iterate.

Parameters:

Name	Type	Description	Default
lr		The learning rate.	required
momentum		The heavy-ball momentum coefficient, where 0 is plain gradient descent.	required
max_steps		The step budget.	required
grad_tol		The gradient-norm tolerance, or None for the dtype's default.	required
ad		"reverse" takes the gradient from the reverse adjoint, "forward" from the forward-mode Jacobian.	required
callbacks		Callbacks run together with those passed to ds.inference.optsolve.	required

diff_soup.inference.Optax

An Optax optimizer, first-order steps from a gradient transformation such as optax.adamw.

It needs the objective and its gradient, and the optional optax dependency. Box bounds are kept by projection after each update. The result reports the evaluated point of lowest objective.

Parameters:

Name	Type	Description	Default
optimizer		The Optax gradient transformation, optax.adam(1e-3) when None.	required
max_steps		The step budget.	required
grad_tol		The gradient-norm tolerance, or None for the dtype's default.	required
log_every		Ignored.	required
ad		"reverse" takes the gradient from the reverse adjoint, "forward" from the forward-mode Jacobian.	required
callbacks		Callbacks run together with those passed to ds.inference.optsolve.	required

Example

import optax

solver = ds.inference.Optax(optimizer=optax.adamw(1e-2), max_steps=200)
result = ds.inference.optsolve(infproblem, solver)

---

Scales

The x_scale= setting of a solver takes one of these.

diff_soup.inference.HessianScale

Scales each parameter by the magnitude of its curvature diagonal, the default of the trust-region solvers.

NewtonTrustRegion refuses it with a CG subproblem.

diff_soup.inference.JacScale

Scales each parameter by the squared norm of its Jacobian column, the default of LevenbergMarquardt.

The trust-region solvers refuse it with hessian="true".

diff_soup.inference.ConstantScale

Scales each parameter by a fixed typical magnitude, with the damping diagonal 1 / scale**2.

Parameters:

Name	Type	Description	Default
scale		The typical magnitude of the parameters, one scalar or one entry per parameter.	required

---

Constrained fits

diff_soup.inference.constrained_optsolve

constrained_optsolve(problem: Any, optimizer: Any = None, *, constraints: InequalityConstraints, outer_max: int = 15, feas_tol: float = 1e-06, stat_tol: float = 1e-06, rho0: float = 1000000.0, callbacks: tuple = ()) -> ConstrainedOptResult

Fit a problem subject to inequality constraints with an augmented-Lagrangian outer loop.

Each outer iteration runs ds.inference.optsolve on the objective plus a penalty on the constraints, then updates the multipliers. It needs JAX 64-bit mode.

Parameters:

Name	Type	Description	Default
problem	Any	The problem to fit, as for ds.inference.optsolve.	required
optimizer	Any	The inner optimizer, NewtonTrustRegion() when None.	None
constraints	InequalityConstraints	The InequalityConstraints to satisfy.	required
outer_max	int	The largest number of outer iterations.	15
feas_tol	float	The largest constraint violation accepted at convergence.	1e-06
stat_tol	float	The tolerance on the norm of the Lagrangian gradient, relative to the gradient scale.	1e-06
rho0	float	The initial penalty weight.	1000000.0
callbacks	tuple	Callbacks passed to every inner fit.	()

Returns:

Type	Description
ConstrainedOptResult	A ConstrainedOptResult with the last inner fit as result, the multipliers, the kkt measures and converged.

diff_soup.inference.InequalityConstraints dataclass

Inequality constraints g(params) <= 0 for ds.inference.constrained_optsolve.

Attributes:

Name	Type	Description
g	Callable	A function of the fitted parameters returning one value per constraint, satisfied where every value is at most zero.
n_constraints	int	The length of the array g returns.

diff_soup.inference.Gauge

A linear equality that fixes a direction the data leave free, for gauges= of TrustRegion.

Build one with Gauge.pin, Gauge.zero_mean, Gauge.zero_mean_across or Gauge.declared.

pin classmethod

pin(name: str, value: Any) -> 'Gauge'

Hold one coordinate of the fitted vector at a value.

Parameters:

Name	Type	Description	Default
name	str	The name of the coordinate, matching exactly one of the fit's parameter names.	required
value	Any	The finite value it is held at.	required

Returns:

Type	Description
'Gauge'	The gauge.

zero_mean classmethod

zero_mean(name: str, weights: Any = None) -> 'Gauge'

Hold the weighted mean of one parameter group at zero.

Parameters:

Name	Type	Description	Default
name	str	The name of the parameter group.	required
weights	Any	One weight per scalar of the group, or a basis object with mean_weights(), or None for equal weights.	None

Returns:

Type	Description
'Gauge'	The gauge.

zero_mean_across classmethod

zero_mean_across(owners: Any, weights: Any = None) -> 'Gauge'

Hold at zero the weighted mean, across the members of a one-level study, of one number per member.

Each member's number is a weighted sum of one of its parameter groups, and at least two members are needed.

Parameters:

Name	Type	Description	Default
owners	Any	A mapping from each member's group name to the weights that sum the group to one number, or a sequence of scalar group names.	required
weights	Any	One nonnegative weight per entry of owners, multiplied by the solver's cluster weights, or None for equal weights.	None

Returns:

Type	Description
'Gauge'	The gauge.

declared classmethod

declared(name: str) -> 'Gauge'

Record a known free direction without constraining it.

Parameters:

Name	Type	Description	Default
name	str	A label for the free direction.	required

Returns:

Type	Description
'Gauge'	The gauge, which constrains nothing.

---

Custom solvers and problems

A solver of your own implements OptSolver and receives an OptimizationProblem, or a LeastSquaresProblem when the objective is half a squared residual norm.

diff_soup.inference.OptSolver

Base class for an optimizer that ds.inference.optsolve runs.

A subclass implements solve and sets requires. ds.inference.optsolve refuses a problem that cannot provide what the solver requires.

Attributes:

Name	Type	Description
requires	frozenset	The names of what solve needs from the problem, such as "value_gradient", plus "resume" when solve takes state=.

requires class-attribute

requires: frozenset = frozenset()

solve abstractmethod

solve(problem: Any, *, callbacks: tuple = ()) -> 'OptResult'

Minimize the problem and return the result.

Parameters:

Name	Type	Description	Default
problem	Any	The OptimizationProblem to minimize.	required
callbacks	tuple	The Callback objects to call during the run.	()

Returns:

Type	Description
'OptResult'	The OptResult of the run.

diff_soup.inference.OptimizationProblem

The objective an optimizer minimizes, with its starting point, bounds and derivatives.

ds.inference.optsolve builds one from the problem it is given and passes it to OptSolver.solve. structure names the functions it provides, and asking for any other raises ConfigurationError.

Attributes:

Name	Type	Description
x0	Any	The starting point, a 1-D array.
structure	frozenset	The names of the functions this problem provides, such as "value_gradient", "hessian" or "gauss_newton".
bounds	Any	The (lower, upper) arrays aligned with x0, or None when unbounded.

x0 instance-attribute

x0: Any

value property

value: Callable

The objective, a function value(x) returning a scalar.

value_and_gradient property

value_and_gradient: Callable

A function value_and_gradient(x) returning the objective and its gradient.

gauss_newton property

gauss_newton: Callable

A function gauss_newton(x) returning the objective, its gradient, the Gauss-Newton matrix and a failure count.

Available when structure holds "gauss_newton".

from_value classmethod

from_value(x0: Any, value: Callable, *, regularization: Any = None, bounds: Any = None, is_valid: Callable | None = None, unflatten: Callable | None = None, static: Any = None, hessian: bool = False) -> OptimizationProblem

Pose a problem from a scalar objective of a 1-D array.

The gradient comes from reverse-mode autodiff, and every function is compiled with jax.jit.

Parameters:

Name	Type	Description	Default
x0	Any	The starting point, a 1-D array.	required
value	Callable	A JAX-traceable function value(x) returning the scalar objective.	required
regularization	Any	A penalty the solver adds to the objective, such as Tikhonov, or None.	None
bounds	Any	A (lower, upper) pair aligned with x0, where None on either side leaves it unbounded.	None
is_valid	Callable | None	A function is_valid(x) returning whether a trial point may be evaluated.	None
unflatten	Callable | None	A function rebuilding the fitted model's arrays from x, used with static.	None
static	Any	The static part of the fitted model, combined with unflatten(x).	None
hessian	bool	If True, also provide the exact Hessian by autodiff.	False

Returns:

Type	Description
OptimizationProblem	An OptimizationProblem with the "value_gradient" structure, and "hessian" when requested.

diff_soup.inference.LeastSquaresProblem

An optimization problem whose objective is half the squared norm of a residual vector.

ds.inference.optsolve builds one from an InferenceProblem, and from_residual poses one from any JAX-traceable residual function. ds.inference.LevenbergMarquardt needs this form.

residual property

residual: Callable

The residual, a function residual(x) returning a 1-D array.

residual_and_jacobian property

residual_and_jacobian: Callable

A function residual_and_jacobian(x) returning the residual and its Jacobian.

from_residual classmethod

from_residual(x0: Any, residual: Callable, *, jacobian: Callable | None = None, regularization: Any = None, bounds: Any = None, is_valid: Callable | None = None, unflatten: Callable | None = None, static: Any = None, segment_fns: tuple = (), hessian: bool = True) -> LeastSquaresProblem

Pose a least-squares problem from a residual function of a 1-D array.

Parameters:

Name	Type	Description	Default
x0	Any	The starting point, a 1-D array.	required
residual	Callable	A JAX-traceable function residual(x) returning the residual vector.	required
jacobian	Callable | None	A function jacobian(x) returning the Jacobian of residual, or None to compute it with jax.jacfwd.	None
regularization	Any	A penalty the solver adds to the objective, such as Tikhonov, or None.	None
bounds	Any	A (lower, upper) pair aligned with x0, where None on either side leaves it unbounded.	None
is_valid	Callable | None	A function is_valid(x) returning whether a trial point may be evaluated.	None
unflatten	Callable | None	A function rebuilding the fitted model's arrays from x, used with static.	None
static	Any	The static part of the fitted model, combined with unflatten(x).	None
segment_fns	tuple	Residual functions whose outputs concatenate to residual(x), with () meaning (residual,).	()
hessian	bool	If True, also provide the exact Hessian of the objective by autodiff.	True

Returns:

Type	Description
LeastSquaresProblem	A LeastSquaresProblem.

Example

t = jnp.linspace(0.0, 2.0, 20)
y = 3.0 * jnp.exp(-1.3 * t)
problem = ds.inference.LeastSquaresProblem.from_residual(
    jnp.array([1.0, 1.0]), lambda x: x[0] * jnp.exp(-x[1] * t) - y)
result = ds.inference.optsolve(problem, ds.inference.LevenbergMarquardt())

diff_soup.inference.XScale

Base class for the parameter scaling of the optimizers' damping and trust regions.

A subclass implements damping_diagonal. Levenberg-Marquardt damps by that diagonal, and the trust-region solvers measure steps in its square root.

damping_diagonal abstractmethod

damping_diagonal(J: ndarray, x: ndarray) -> jnp.ndarray

Return the damping diagonal, one positive entry per parameter.

Parameters:

Name	Type	Description	Default
J	ndarray	The residual Jacobian, or the curvature matrix of a trust-region solver.	required
x	ndarray	The current parameters.	required

Returns:

Type	Description
ndarray	A 1-D array with the length of x.