beamax.gb.gb_utils

Mathematical helper routines for Gaussian beam dynamics.

Scope

• Hamiltonian derivatives and geometric quantities.
• Matrix helpers and vectorized utilities used by ODE solvers.

API Reference

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) = mode * c(x) * ||p||.

Parameters:

Name	Type	Description	Default
x	(ndarray, shape(..., d))	Positions.	required
p	(ndarray, shape(..., d))	Momenta.	required
mode	(ndarray, shape(...))	±1 branch selector.	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']

∂G/∂x for G(x,p) = mode * c(x) * ||p||.

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))	mode * ∇c(x) * ||p||.

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

∂G/∂p for G(x,p) = mode * c(x) * ||p||.

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))	mode * c(x) * p / ||p||.

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

Validate initial Hessian M0: symmetric and Im(M0) ≻ 0.

Parameters:

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

Raises:

Type	Description
ValueError	If symmetry fails or Im(M0) 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 diagonal M0 = diag(alpha0). Im(alpha0) should be > 0.	required
M0	(ndarray | None, shape(b, d, d), complex)	Explicit Hessian. Mathematical callers must ensure symmetry and Im(M0) ≻ 0; 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.