multistart_init

Contents

multistart_init#

class braintools.optim.multistart_init(key, bounds, n_starts)#

Draw uniform starting points within box bounds.

Parameters:

  • key (Any) – A jax.random key.

  • bounds (Any) – A pytree whose leaves are (low, high) pairs (scalars, arrays or brainunit Quantities of a common unit). A bare pair describes a single leaf.

  • n_starts (int) – Number of starting points S.

Returns:

The structure of bounds with every pair replaced by an array of shape (S, *broadcast(low, high).shape) drawn uniformly from [low, high); a Quantity in the unit of low when the bounds carry one.

Return type:

Any

Raises:

ValueError – If n_starts is not a positive integer.

See also

multistart_minimize

Minimise from the drawn starts.

Notes

n_starts fixes the output shape, so it must be a Python int. The function works under jax.jit (with n_starts static) and jax.vmap over key.

Examples

>>> import jax
>>> from braintools.optim import multistart_init
>>> x0 = multistart_init(jax.random.PRNGKey(0), {'a': (0., 1.), 'b': (-2., 2.)}, 4)
>>> x0['a'].shape, x0['b'].shape
((4,), (4,))
>>> bool(((x0['b'] >= -2.) & (x0['b'] < 2.)).all())
True
Contents