Skip to content

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/validate initial Hessian.

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)

If given, must be symmetric with Im(M0) ≻ 0.

required

Returns:

Type Description
(ndarray, shape(b, d, d), complex)

Validated/constructed Hessian.

Raises:

Type Description
ValueError

If both or neither of alpha0 and M0 are provided.

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.