rqutils.math module

Math utility functions

Math API

rqutils.math.matrix_ufunc(operator: Callable, mat: ArrayLike, hermitian: int | bool = 0, with_diagonals: bool = False, npmod: ModuleType = <module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/rqutils/envs/v0.2.0/lib/python3.12/site-packages/numpy/__init__.py'>, save_errors: bool = False) → NDArray | tuple[NDArray, NDArray]

Apply a unitary-invariant unary matrix operator to an array of normal matrices.

The argument mat must be an array of normal (i.e. square diagonalizable) matrices in the last two dimensions. This function unitary-diagonalizes the matrices, applies operator to the diagonals, and inverts the diagonalization.

Diagonalization and gradient

When using this function with an autodiff library (e.g. JAX), the gradient diverges when an input parameter controls off-diagonal elements of mat but mat is diagonal. Use an alternative function (that is hopefully available) in such cases:

# Reshape the matrix to gather all off-diagonal elements to a block ([:, 1:])
mat_dim = mat.shape[-1]
diag_checker = mat.reshape(-1, mat_dim ** 2)
# The very last element is a part of diagonal -> can ignore for this purpose
diag_checker = diag_checker[:, :-1].reshape(-1, mat_dim - 1, mat_dim + 1)
is_diagonal = ~jnp.any(diag_checker[:, :, 1:], axis=(1, 2))
has_diagonal = jnp.any(is_diagonal)

result = jax.lax.cond(has_diagonal,
                      alternative_X,
                      functools.partial(matrix_ufunc, X),
                      mat)

Parameters:

operator – Unary operator to be applied to the diagonals of mat.
mat – Array of normal matrices (shape (…, n, n)). No check on normality is performed.
hermitian – 1 or True -> mat is Hermitian, -1 -> mat is anti-hermitian, 0 or False -> otherwise
with_diagonals – If True, also return the array operator(eigenvalues).

Returns:

An array corresponding to operator(mat). If diagonals==True, another array corresponding to operator(eigvals).

rqutils.math.matrix_exp(mat: ArrayLike, hermitian: int | bool = 0, with_diagonals: bool = False, npmod: ModuleType = <module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/rqutils/envs/v0.2.0/lib/python3.12/site-packages/numpy/__init__.py'>, save_errors: bool = False) → NDArray: matrix_ufunc(exp, …)

rqutils.math.matrix_angle(mat: ArrayLike, hermitian: int | bool = 0, with_diagonals: bool = False, npmod: ModuleType = <module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/rqutils/envs/v0.2.0/lib/python3.12/site-packages/numpy/__init__.py'>, save_errors=False) → NDArray: matrix_ufunc(angle, …)