beamax.gb.core

Real and complex Gaussian beam field evaluation.

beamax.gb.core

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 scaled by \(\lVert p\rVert\).	required
mode	(ndarray, shape(b))	Branch value \(\pm1\) 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 \(A_t\exp(i\omega_0\Phi)\).

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 \(\mathcal{O}(N_tS)\) memory instead of materializing the beam axis with \(\mathcal{O}(bN_tS)\) storage.

compute_gaussian_beam_real_pallas(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, pallas_config: PallasConfig | None = None, initial_field: jax.Array | None = None, return_padded: bool = False) -> Float[Array, 'Nt *S']

Evaluate a real GB field with XLA trajectories and Pallas reduction.

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+\overline{u} =2|A|\cos\!\left( \omega\left[p\mathbin{\cdot}\Delta x +\frac{1}{2}\Delta x^{\mathsf T}M_R\Delta x\right] +\arg A \right) \exp\!\left( -\frac{\omega}{2}\Delta x^{\mathsf T}M_I\Delta x \right). \]

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.

compute_gaussian_beam_real_TR_xla_terminal(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, '*S']

Evaluate the terminal TR field with output-sized XLA accumulation.

compute_gaussian_beam_real_TR_pallas_terminal(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, pallas_config: PallasConfig | None = None, initial_field: jax.Array | None = None, return_padded: bool = False) -> Float[Array, '*S']

Evaluate the terminal TR field with Pallas beam reduction.

sum_gaussian_beam_real_trajectories_xla(xt: Float[Array, 'b Nt d'], pt: Float[Array, 'b Nt d'], Mt: Complex[Array, 'b Nt d d'], At: Complex[Array, 'b Nt 1'], omega0: Float[Array, ' b'], sensors: Float[Array, '*S d'], domain_size: Float[Array, ' d'], periodic: Bool[Array, ' d'], *, precision=jax.lax.Precision.HIGHEST) -> Float[Array, 'Nt *S']

Sum trajectory fields with output-sized XLA accumulation.

A singleton time axis provides terminal-only evaluation.