"""Includes utility functions for the project. Generic functions to generate data, seeds, etc."""
from typing import Callable, Tuple
import jax
from jax import Array
import jax.random as jr
import jax.numpy as jnp
[docs]
def keygen(key, nkeys):
"""
Generate randomness that JAX can use by splitting the JAX keys.
Parameters
----------
key : jax.random.PRNGKey
The random key for JAX.
nkeys : int
Number of keys to generate.
Returns
-------
tuple
A tuple containing the new key for further generators and the key generator.
"""
keys = jr.split(key, nkeys + 1)
return keys[0], (k for k in keys[1:])
[docs]
def initialise_stable_dynamics(key: Tuple[int, int], n_dim: int, T: int, radii: float = 0.6) -> Array:
"""
Generate a state matrix with stable dynamics (all eigenvalues inside the unit circle) discrete dynamics.
Parameters
----------
key : tuple of int
Random key.
n_dim : int
State dimensions.
T : int
Time steps.
radii : float, optional
Spectral radius, by default 0.6.
Returns
-------
Array
Matrix A with stable dynamics.
"""
mat = jr.normal(key, (n_dim, n_dim)) * radii
mat /= jnp.sqrt(n_dim)
return jnp.tile(mat, (T, 1, 1))
[docs]
def initialise_stable_time_varying_dynamics(key: Tuple[int, int], n_dim: int, T: int, radii: float = 0.6) -> Array:
"""
Generate a state matrix with stable dynamics (all eigenvalues inside the unit circle).
Parameters
----------
key : tuple of int
Random key.
n_dim : int
State dimensions.
T : int
Time steps.
radii : float, optional
Spectral radius, by default 0.6.
Returns
-------
Array
Matrix A with stable dynamics.
"""
mat = jr.normal(key, (T, n_dim, n_dim)) * radii
mat /= jnp.sqrt(n_dim)
# mat -= jnp.eye(n_dim)
return mat
[docs]
def linearise(fun: Callable) -> Callable:
"""
Find Jacobian with respect to x and u inputs.
Parameters
----------
fun : Callable
Function with arguments (t, x, u, params).
Returns
-------
Callable
Jacobian tuple evaluated at args 1 and 2.
"""
return jax.jacrev(fun, argnums=(1, 2))
[docs]
def quadratise(fun: Callable) -> Callable:
"""
Find Hessian with respect to x and u inputs.
Parameters
----------
fun : Callable
Function with arguments (t, x, u, params).
Returns
-------
tuple
Hessian tuple cross evaluated with args 1 and 2.
"""
return jax.jacfwd(jax.jacrev(fun, argnums=(1, 2)), argnums=(1, 2))
[docs]
def time_map(fun: Callable) -> Callable:
"""
Vectorise function in time. Assumes 0th-axis is time for x and u args of fun, the last
arg (theta) of Callable function assumed to be time-invariant.
Parameters
----------
fun : Callable
Function that takes args (t, x[Txn], u[Txm], theta).
Returns
-------
Callable
Vectorised function along args 1 and 2 0th-axis.
"""
return jax.vmap(fun, in_axes=(0, 0, 0, None))
[docs]
def smooth_l1(x: Array, eps: float = 1e-6) -> Array:
"""
Smooth approximation of the L1 norm, sum(sqrt(x^2 + eps) - sqrt(eps)).
Twice differentiable, so it can be used as an iLQR sparsity cost. Approaches ``sum(|x|)`` as
``eps -> 0``; anneal ``eps`` by re-solving with smaller values, warm-starting each solve.
Parameters
----------
x : Array
Input array.
eps : float, optional
Smoothing parameter, by default 1e-6.
Returns
-------
Array
Smoothed L1 norm of x.
"""
return jnp.sum(jnp.sqrt(x**2 + eps) - jnp.sqrt(eps))
