beamax.gb.gb_utils

Hamiltonian, geometric, and matrix utilities for Gaussian beam dynamics.

beamax.gb.gb_utils

G(x: Float[Array, '... d'], p: Float[Array, '... d'], mode: Num[Array, ...], c: Callable[[Float[Array, '... d']], Float[Array, ...]]) -> Float[Array, ...]

Hamiltonian for acoustics: \(G(x,p)=\sigma c(x)\lVert p\rVert\), where \(\sigma=\mathtt{mode}\in\{-1,1\}\).

Parameters:

Name	Type	Description	Default
x	(ndarray, shape(..., d))	Positions.	required
p	(ndarray, shape(..., d))	Momenta.	required
mode	(ndarray, shape(...))	Branch selector \(\sigma\in\{-1,1\}\).	required
c	Callable[[ndarray], ndarray]	Speed of sound \(c(x)\).	required

Returns:

Type	Description
(ndarray, shape(...))	Hamiltonian values. dtype follows c(x)/inputs.

Notes

JIT-compiled; pure; vectorize with vmap as needed.

Gx(x: Float[Array, '... d'], p: Float[Array, '... d'], mode: Num[Array, ...], c: Callable[[Float[Array, '... d']], Float[Array, ...]]) -> Float[Array, '... d']

Gradient \(\partial G/\partial x\) for \(G(x,p)=\sigma c(x)\lVert p\rVert\).

Parameters:

Name	Type	Description	Default
x	(ndarray, shape(..., d))		required
p	(ndarray, shape(..., d))		required
mode	(ndarray, shape(...))		required
c	Callable		required

Returns:

Type	Description
(ndarray, shape(..., d))	\(\sigma\nabla c(x)\lVert p\rVert\).

Gp(x: Float[Array, '... d'], p: Float[Array, '... d'], mode: Num[Array, ...], c: Callable[[Float[Array, '... d']], Float[Array, ...]]) -> Float[Array, '... d']

Gradient \(\partial G/\partial p\) for \(G(x,p)=\sigma c(x)\lVert p\rVert\).

Parameters:

Name	Type	Description	Default
x	(ndarray, shape(..., d))		required
p	(ndarray, shape(..., d))		required
mode	(ndarray, shape(...))		required
c	Callable		required

Returns:

Type	Description
(ndarray, shape(..., d))	\(\sigma c(x)p/\lVert p\rVert\).

check_M0(M0: Complex[Array, 'b d d']) -> None

Validate initial Hessian M0: symmetric with \(\operatorname{Im}(M_0)\succ0\).

Parameters:

Name	Type	Description	Default
M0	(ndarray, shape(b, d, d), complex)		required

Raises:

Type	Description
ValueError	If symmetry fails or \(\operatorname{Im}(M_0)\) is not positive definite (per batch).

prepare_M0(alpha0: Optional[Complex[Array, 'b d']], M0: Optional[Complex[Array, 'b d d']]) -> Complex[Array, 'b d d']

Construct an initial Hessian or pass through an explicit one.

Parameters:

Name	Type	Description	Default
alpha0	(ndarray | None, shape(b, d), complex)	If given, produces \(M_0=\operatorname{diag}(\alpha_0)\). \(\operatorname{Im}(\alpha_0)\) should be positive.	required
M0	(ndarray | None, shape(b, d, d), complex)	Explicit Hessian. Mathematical callers must ensure symmetry and \(\operatorname{Im}(M_0)\succ0\); call :func:check_M0 in eager validation code.	required

Returns:

Type	Description
(ndarray, shape(b, d, d), complex)	Constructed or supplied Hessian.

Raises:

Type	Description
ValueError	If both or neither of alpha0 and M0 are provided.

Notes

This helper is used inside JIT-compiled solver paths, so it cannot raise a Python exception from data-dependent positivity checks. Use :func:check_M0 before entering a compiled solve when accepting untrusted user-supplied Hessians.

is_diagonal(M0: Num[Array, 'b d d']) -> Num[Array, '']

Test whether M0 is diagonal (per batch).

Parameters:

Name	Type	Description	Default
M0	(ndarray, shape(b, d, d))		required

Returns:

Type	Description
ndarray	0-D boolean array — True when M0 is diagonal.