multistart_minimize

class braintools.optim.multistart_minimize(fun, x0, *, method='lbfgs', max_iter=500, tol=1e-06, bounds=None, learning_rate=0.01)

Minimise a scalar function from many starting points in parallel.

Every start runs the same local optimiser inside a lax.while_loop; the loop is vmap-ped over the starts and the whole run is compiled into one XLA program. A start stops when it converges, when its loss becomes non-finite, or after max_iter iterations, while the other starts keep going.

Parameters:

• fun (Callable[[Any], Any]) – fun(x) -> scalar. A pure, differentiable function of one start x (a pytree shaped like a single slice of x0). Its value must be a plain or dimensionless scalar.

• x0 (Any) – Starting points: every leaf has a leading axis of size S, the number of starts. Leaves may be brainunit Quantities; fun then receives Quantities in the same units.

• method (str | GradientTransformation) – Local optimiser. 'lbfgs' is optax.lbfgs() with its default zoom line search (learning_rate is ignored); 'adam' and 'gd' are optax.adam() and optax.sgd() with step size learning_rate. Any other optax transformation can be passed directly.

• max_iter (int) – Maximum number of iterations per start (static).

• tol (float) – Convergence tolerance. A start has converged when the largest absolute gradient entry (in the unconstrained space) is below tol, or when the loss changes by a relative amount \(|f_k - f_{k-1}| \le \mathrm{tol}\,|f_{k-1}|\) in one iteration (so a step that leaves the loss unchanged also counts).

• bounds (Any) – Box bounds mirroring the structure of x0: a (low, high) pair or None in place of every leaf (a bare pair for a single-leaf x0). Bounds broadcast against one start of the leaf and may be Quantities in the leaf’s unit. low or high may be None for a one-sided bound. A bounded leaf is optimised in an unconstrained space through brainstate.nn.SigmoidT (SoftplusT / NegSoftplusT for one-sided bounds) and mapped back, so iterates never leave the box.

• learning_rate (Any) – Step size for 'adam' and 'gd'.

Returns:

(x, fun, converged, n_iter, best_x, best_fun, best_index, history); see MultiStartResult.

Return type:

MultiStartResult

Raises:

• ValueError – If method is unknown, max_iter is not a positive integer, x0 has no leaves or leaves with different leading sizes, or a bound is malformed (wrong structure, non-finite, low >= high).

• TypeError – If fun returns a value with a physical unit, or a bound’s unit does not match its leaf.

See also

multistart_init

Uniform starting points within bounds.

braintools.fit.Fitter

Model fitting built on this minimiser.

Notes

Convergence test. The gradient test uses the gradient with respect to the constrained value (the chain-rule gradient divided by the slope of the bijector), so a start near a bound, where the sigmoid is flat, is not mistaken for a stationary point. At an optimum on a bound that gradient stays non-zero, and only the loss-change test can stop the start. The loss-change test is used for 'lbfgs' only: its line search makes a negligible decrease a sign of convergence, whereas the per-step change of a fixed-step method ('adam', 'gd', a custom optax transformation) reflects the step size rather than the distance to the optimum (with a large constant offset in the loss, a relative test would stop such a method far from the minimum). These methods therefore stop on the gradient test or after max_iter. Conversely, for 'lbfgs' on a loss with a large constant offset the relative test can stop early; lower tol (or remove the offset) then.

Starts on a bound. Starting values are moved 1e-3 of the interval width inside the bounds (1e-3 * max(1, |bound|) for a one-sided bound) before the bijector is inverted, so starts on (or outside) a bound are finite and still move.

Non-finite losses. A start whose loss at x0 is not finite stops immediately; a start whose loss becomes non-finite during the run (it diverged) stops and keeps its last finite point. Either way it is reported with converged=False and n_iter < max_iter, and it is never selected as the best while another start has a finite loss and did not diverge. If every start diverged, the lowest finite loss among them is selected; if no loss is finite, best_index is 0 and best_fun is not finite. A start that already satisfies the convergence test at x0 runs zero iterations.

The function contains no host synchronisation, so it can itself be called inside jax.jit (with the keyword options static) or jax.vmap. It is not meant to be differentiated.

Examples

>>> import jax.numpy as jnp
>>> from braintools.optim import multistart_minimize
>>> def rosenbrock(p):
...     return (1. - p['x']) ** 2 + 100. * (p['y'] - p['x'] ** 2) ** 2
>>> x0 = {'x': jnp.array([-1.5, 0., 2.]), 'y': jnp.array([2., -1., 1.])}
>>> res = multistart_minimize(rosenbrock, x0, max_iter=200)
>>> res.converged
Array([ True,  True,  True], dtype=bool)
>>> print(jnp.round(res.best_x['x'], 3), jnp.round(res.best_x['y'], 3))
1.0 1.0
>>>
>>> res = multistart_minimize(lambda x: jnp.sin(3. * x) + 0.1 * x ** 2,
...                           jnp.linspace(-2.5, 2.5, 5), bounds=(-3., 3.))
>>> print(jnp.round(res.best_x, 3))
-0.512