geope.jax.dexpm

geope.jax.dexpm.Ui(x, basis)

Compute a unitary from a linear combination of Hermitian basis matrices.

Constructs \(U = \exp(i \sum_k x_k B_k)\).

Parameters:

Name	Type	Description	Default
x	Array	Coefficient vector of shape (K,).	required
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required

Returns:

Type	Description
Array	A unitary matrix of shape (d, d).

geope.jax.dexpm.get_Ui_fn(basis)

Create a partial unitary function with a fixed basis.

Parameters:

Name	Type	Description	Default
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required

Returns:

Type	Description
Callable[[Array], Array]	A callable that accepts a coefficient vector and returns
Callable[[Array], Array]	the corresponding unitary matrix.

geope.jax.dexpm.dexpm_block(A, x)

Compute the derivative of the matrix exponential via the block method.

Implements the block-matrix approach of Al-Mohy & Higham (2009) <https://arxiv.org/pdf/1506.00628>_, Eq. (31), extracting \(d\exp(iA)/dA \cdot x\) from the upper-right block of \(\exp(i[[A, x], [0, A]])\).

Parameters:

Name	Type	Description	Default
A	Array	The Hamiltonian matrix of shape (d, d).	required
x	Array	The direction matrix of shape (d, d).	required

Returns:

Type	Description
Array	The directional derivative matrix of shape (d, d).

geope.jax.dexpm.dexpm(x, basis)

Compute the derivative of the exponential map for all basis directions.

For each basis element \(B_k\), computes \(\partial \exp(i \sum_j x_j B_j) / \partial x_k\).

Parameters:

Name	Type	Description	Default
x	Array	Coefficient vector of shape (K,).	required
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required

Returns:

Type	Description
Array	An array of shape (d, d, K) whose last axis indexes the
Array	partial derivatives with respect to each coefficient.

geope.jax.dexpm.dexpm_batched(x, basis, batch_size)

Batched derivative of the exponential map.

Same as dexpm but uses jax.lax.map with a configurable batch_size to limit peak memory usage.

Parameters:

Name	Type	Description	Default
x	Array	Coefficient vector of shape (K,).	required
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required
batch_size	int	Number of basis elements to process per batch.	required

Returns:

Type	Description
Array	An array of shape (d, d, K).

geope.jax.dexpm.get_dexpm(basis, batch_size=None)

Create a JIT-compiled exponential-map derivative function.

Parameters:

Name	Type	Description	Default
basis	Array	Array of Hermitian matrices of shape (K, d, d).	required
batch_size	int | None	Optional batch size. If None, the full vmap variant is used; otherwise the batched variant.	None

Returns:

Type	Description
Callable[[Array], Array]	A callable that accepts a coefficient vector and returns
Callable[[Array], Array]	the derivative array of shape (d, d, K).