differt2d.optimize module

Optimization toolbox.

All the functions present in this toolbox support Just-in-time compilation with jax.jit.

Examples

>>> from differt2d.optimize import minimize
>>> import chex
>>> import jax
>>> import jax.numpy as jnp
>>> @jax.jit
... def parabola_min(a, b, c):
...     def f(x):
...         x = a * (x + b) + c
...         return jnp.dot(x, x)
...
...     return minimize(f, jnp.array(0.0))
>>>
>>> x, y = parabola_min(2.0, 1.0, 1.0)
>>> chex.assert_trees_all_close(x, -1.5, rtol=1e-2)
>>> chex.assert_trees_all_close(y, +0.0, atol=1e-3)

minimize(fun, x0, args=(), steps=100, optimizer=None)

Minimizes a scalar function of one or more variables.

Parameters:

• fun (Callable[[Array, 'n'], Unpack[TypeVarTuple]], Array, '']]) – The objective function to be minimized.

• x0 (Array, 'n']) – The initial guess.

• args (tuple[Unpack[TypeVarTuple]], default: ()) – Positional arguments passed to fun.

• steps (int, default: 100) – The number of steps to perform.

• optimizer (Optional[GradientTransformation], default: None) – The optimizer to use. If not provided, uses optax.adam with a learning rate of 0.1.

Return type:

tuple[Array, 'n'], Array, '']]

Returns:

The solution array and the corresponding loss.

Examples:

>>> from differt2d.optimize import minimize
>>> import chex
>>> import jax.numpy as jnp
>>> def f(x, offset=1.0):
...     x = x - offset
...     return jnp.dot(x, x)
>>>
>>> x, y = minimize(f, jnp.zeros(10))
>>> chex.assert_trees_all_close(x, jnp.ones(10), rtol=1e-2)
>>> chex.assert_trees_all_close(y, 0.0, atol=1e-4)
>>>
>>> # It is also possible to pass positional arguments
>>> x, y = minimize(f, jnp.zeros(10), args=(2.0,))
>>> chex.assert_trees_all_close(x, 2.0 * jnp.ones(10), rtol=1e-2)
>>> chex.assert_trees_all_close(y, 0.0, atol=1e-3)

minimize_random_uniform(fun, key, n, **kwargs)

Minimizes a scalar function of one or more variables, with initial guess drawn randomly from a uniform distribution.

Parameters:

• fun (Callable[[Array, '{n}'], Unpack[TypeVarTuple]], Array, '']]) – The objective function to be minimized.

• key (Union[Array, ''], Array, '2']]) – The random key used to generate the initial guess.

• n (int) – The size of the random vector to generate.

• kwargs (Any) – Keyword arguments passed to minimize.

Return type:

tuple[Array, '{n}'], Array, '']]

Returns:

The solution array and the corresponding loss.

Examples:

>>> from differt2d.optimize import minimize_random_uniform
>>> import chex
>>> import jax
>>> import jax.numpy as jnp
>>> def f(x):
...     x = x - 1.0
...     return jnp.dot(x, x)
>>>
>>> x, y = minimize_random_uniform(f, jax.random.PRNGKey(1234), 10)
>>> chex.assert_trees_all_close(x, jnp.ones(10), rtol=1e-2)
>>> chex.assert_trees_all_close(y, 0.0, atol=1e-3)

minimize_many_random_uniform(fun, key, n, many=10, **kwargs)

Minimizes many times a scalar function of one or more variables, with initial guess drawn randomly from a uniform distribution, and returns the best minimum out of the many trials.

Parameters:

• fun (Callable[[Array, '{n}'], Unpack[TypeVarTuple]], Array, '']]) – The objective function to be minimized.

• key (Union[Array, ''], Array, '2']]) – The random key used to generate the initial guesses.

• n (int) – The size of the random vector to generate.

• many (int, default: 10) – How many times the minimization should be performed.

• kwargs (Any) – Keyword arguments passed to minimize_random_uniform.

Return type:

tuple[Array, '{n}'], Array, '']]

Returns:

The solution array and the corresponding loss.

Examples:

>>> from differt2d.optimize import minimize_many_random_uniform
>>> import chex
>>> import jax
>>> import jax.numpy as jnp
>>> def f(x):
...     x = x - 1.0
...     return jnp.dot(x, x)
>>>
>>> x, y = minimize_many_random_uniform(f, jax.random.PRNGKey(1234), 10)
>>> chex.assert_trees_all_close(x, jnp.ones(10), rtol=1e-2)
>>> chex.assert_trees_all_close(y, 0.0, atol=1e-4)