beamax.gb.core

Low level Gaussian beam field evaluation kernels.

Scope

• Beam synthesis in spatial domain.
• Real and complex beam formulations used by higher level MSGB solvers.

API Reference

beamax.gb.core

wrap_position(position: Float[Array, 'b t d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d']) -> Float[Array, 'b t d']

Apply periodic wrapping selectively per axis.

Parameters:

Name	Type	Description	Default
position	(ndarray, shape(b, t, d))		required
domain_size	(ndarray, shape(d))		required
periodic	(ndarray, shape(d), bool)		required

Returns:

Type	Description
(ndarray, shape(b, t, d))	Wrapped where periodic is True; unchanged elsewhere.

compute_phase(xt: Float[Array, 'b t d'], pt: Float[Array, 'b t d'], mt: Complex[Array, 'b t d d'], sensors: Float[Array, '*S d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d']) -> Complex[Array, 'b t *S']

GB phase at sensors: (p · Δx) + 0.5 Δxᵀ M Δx.

Parameters:

Name	Type	Description	Default
xt	(ndarray, shape(b, t, d))		required
pt	(ndarray, shape(b, t, d))		required
mt	(ndarray, shape(b, t, d, d), complex)		required
sensors	(ndarray, shape(*S, d))		required
domain_size	(ndarray, shape(d))		required
periodic	(ndarray, shape(d), bool)		required

Returns:

Type	Description
(ndarray, shape(b, t, *S), complex)	Phase values (real part used in oscillatory term).

compute_diff(xt: Float[Array, 'b t d'], sensors: Float[Array, '*S d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d']) -> Float[Array, 'b t *S d']

Sensor–ray displacement with periodic wrap.

Parameters:

Name	Type	Description	Default
xt	(ndarray, shape(b, t, d))		required
sensors	(ndarray, shape(*S, d))		required
domain_size	(ndarray, shape(d))		required
periodic	(ndarray, shape(d), bool)		required

Returns:

Type	Description
(ndarray, shape(b, t, *S, d))	Δx = sensors - xt (broadcast), wrapped as needed.

compute_gaussian_beam(x0: Float[Array, 'b d'], p0: Float[Array, 'b d'], M0: Complex[Array, 'b d d'], a0: Complex[Array, ' b'], omega0: Float[Array, ' b'], mode: Num[Array, ' b'], c: Callable[[Float[Array, '... d']], Float[Array, ...]], lam: float, ts: Float[Array, ' Nt'], sensors: Float[Array, '*S d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d'], ode_solver: SolverFn, solver_config: Optional[SolverConfig] = None) -> Complex[Array, 'Nt *S b']

Complex GB field at sensors, keeping beam axis.

Parameters:

Name	Type	Description	Default
x0	(ndarray, shape(b, d))		required
p0	(ndarray, shape(b, d))		required
M0	(ndarray, shape(b, d, d), complex)		required
a0	(ndarray, shape(b), complex)		required
omega0	(ndarray, shape(b))	Angular frequencies (|p| scaled).	required
mode	(ndarray, shape(b))	±1 per beam.	required
c	Callable[[ndarray], ndarray]	Speed of sound.	required
lam	float	Absorption parameter (affects amplitude ODE if enabled).	required
ts	(ndarray, shape(Nt))		required
sensors	(ndarray, shape(*S, d))		required
domain_size	(ndarray, shape(d))		required
periodic	(ndarray, shape(d), bool)		required
ode_solver	SolverFn	Integrator returning (xt, pt, Mt, At).	required
solver_config	Optional[SolverConfig]		None

Returns:

Type	Description
(ndarray, shape(Nt, *S, b), complex)	Field contributions per beam (not summed).

Notes

• Calls ode_solver once; wraps positions; phases from xt, pt, Mt.
• Overall factor At * exp(i omega0 * phase).

safe_angle_eps(z, eps=1e-12)

Phase angle with zero-safe branch for (0+0j).

Parameters:

Name	Type	Description	Default
z	ndarray	Complex array.	required
eps	float	Substitute real part when both real/imag are exactly zero.	1e-12

Returns:

Type	Description
ndarray	atan2(Im(z), Re(z_safe)).

compute_gaussian_beam_real(x0: Float[Array, 'b d'], p0: Float[Array, 'b d'], M0: Complex[Array, 'b d d'], a0: Complex[Array, ' b'], omega0: Float[Array, ' b'], mode: Num[Array, ' b'], c: Callable[[Float[Array, '... d']], Float[Array, ...]], lam: float, ts: Float[Array, ' Nt'], sensors: Float[Array, '*S d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d'], ode_solver: SolverFn, solver_config: Optional[SolverConfig] = None) -> Float[Array, Nt * S]

Real-valued streaming GB: scan over beams, vmap over time.

Parameters:

Name	Type	Description	Default
x0	(ndarray, shape(b, d))	Initial beam positions.	required
p0	(ndarray, shape(b, d))	Initial momenta.	required
M0	(ndarray, shape(b, d, d))	Initial complex Hessian matrices.	required
a0	(ndarray, shape(b))	Initial amplitudes.	required
omega0	(ndarray, shape(b))	Beam angular frequencies.	required
mode	(ndarray, shape(b))	Hamiltonian branch signs.	required
c	Callable	Sound-speed function.	required
lam	float	Absorption parameter passed to ode_solver.	required
ts	(ndarray, shape(Nt))	Time grid.	required
sensors	(ndarray, shape(*S, d))	Sensor positions.	required
domain_size	(ndarray, shape(d))	Physical domain size.	required
periodic	(ndarray, shape(d))	Per-axis periodicity flags.	required
ode_solver	SolverFn	ODE integrator returning (xt, pt, Mt, At).	required
solver_config	SolverConfig	Optional numerical configuration for ode_solver.	None

Returns:

Type	Description
(ndarray, shape(Nt, *S))	Real-valued summed field at the sensor positions.

Notes

Uses O(Nt * S) memory instead of materializing the beam axis with O(b * Nt * S) storage.

compute_gaussian_beam_real_TR(x0: Float[Array, 'b d'], p0: Float[Array, 'b d'], M0: Complex[Array, 'b d d'], a0: Complex[Array, ' b'], omega0: Float[Array, ' b'], mode: Num[Array, ' b'], c: Callable[[Float[Array, '... d']], Float[Array, ...]], lam: float, ts: Float[Array, 'b Nt'], sensors: Float[Array, '*S d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d'], ode_solver: SolverFn, solver_config: Optional[SolverConfig] = None) -> Float[Array, 'Nt *S b']

Compute a collection of Gaussian Beams in n-dimensions, assuming the resulting field is real.

u + bar(u) = 2|A|cos(ω(x p + 0.5 x Mr x) + angle(A))exp(-ω/2 x Mi x).

This should require 2 times less memory than the complex version, and should be faster.

Parameters:

Name	Type	Description	Default
x0	(ndarray, shape(b, d))	Initial beam positions.	required
p0	(ndarray, shape(b, d))	Initial momenta.	required
M0	(ndarray, shape(b, d, d))	Initial complex Hessian matrices.	required
a0	(ndarray, shape(b))	Initial amplitudes.	required
omega0	(ndarray, shape(b))	Beam angular frequencies.	required
mode	(ndarray, shape(b))	Hamiltonian branch signs.	required
c	Callable	Sound-speed function.	required
lam	float	Absorption parameter passed to ode_solver.	required
ts	(ndarray, shape(Nt))	Time grid.	required
sensors	(ndarray, shape(*S, d))	Sensor positions.	required
domain_size	(ndarray, shape(d))	Physical domain size.	required
periodic	(ndarray, shape(d))	Per-axis periodicity flags.	required
ode_solver	SolverFn	ODE integrator returning (xt, pt, Mt, At).	required
solver_config	SolverConfig	Optional numerical configuration for ode_solver.	None

Returns:

Type	Description
(ndarray, shape(Nt, *S, b))	Real-valued beam field with the beam axis retained.