beamax.gb.gb_solvers
Gaussian beam ODE solvers and integration configuration.
beamax.gb.gb_solvers
SolverFn
Bases: Protocol
Protocol for Gaussian beam ODE integrators.
Implementations integrate beam initial data and return beam positions, momenta, Hessians, and amplitudes over time.
Notes
Expected signature is (x0, p0, M0, a0, mode, ts, c, *args, **kwargs)
returning (xt, pt, Mt, At). Standard shapes are:
x0, p0:(b, d)M0:(b, d, d)a0, mode:(b,)ts:(Nt,)xt, pt:(b, Nt, d)Mt:(b, Nt, d, d)At:(b, Nt, 1)
SolverConfig(solver: diffrax.AbstractSolver = diffrax.Tsit5(), max_steps: int = 4096, rtol: float = 1e-07, atol: float = 1e-09, pcoeff: float = 0.0, icoeff: float = 1.0, dcoeff: float = 0.0, dt0: float | None = None)
dataclass
Configuration for ODE solver settings.
from_precision(use_x64: Optional[bool] = None, solver: Optional[diffrax.AbstractSolver] = None, **overrides)
classmethod
Create config with precision-appropriate tolerances.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
use_x64
|
bool | None
|
If None, auto-detects from jax.config.x64_enabled |
None
|
solver
|
AbstractSolver | None
|
Override default solver (Tsit5) |
None
|
**overrides
|
Override any other config fields (max_steps, rtol, etc.) |
{}
|
Examples:
Auto-detect precision, use defaults
config = SolverConfig.from_precision()
Force float32 tolerances, but increase max_steps
config = SolverConfig.from_precision(use_x64=False, max_steps=8192)
Auto precision, custom solver and tolerances
config = SolverConfig.from_precision( solver=diffrax.Dopri5(), rtol=1e-5, max_steps=10000 )
Everything custom
config = SolverConfig.from_precision( use_x64=False, solver=diffrax.Dopri8(), max_steps=2048, rtol=1e-3, atol=1e-5, pcoeff=0.3 )
solve_hom_diag(x0: jnp.ndarray, p0: jnp.ndarray, M0: jnp.ndarray, a0: jnp.ndarray, mode: jnp.ndarray, ts: jnp.ndarray, c: Callable, lam=None, config=None) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
Solver for homogeneous media with simplified equations.
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 Hessian matrices. Only diagonal entries are used. |
required |
a0
|
(ndarray, shape(b))
|
Initial amplitudes. |
required |
mode
|
(ndarray, shape(b))
|
Hamiltonian branch signs. |
required |
ts
|
(ndarray, shape(Nt))
|
Time grid. |
required |
c
|
Callable
|
Homogeneous sound-speed function. |
required |
lam
|
Any
|
Ignored compatibility argument. |
None
|
config
|
Any
|
Ignored compatibility argument. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
xt |
(ndarray, shape(b, Nt, d))
|
Beam positions over time. |
pt |
(ndarray, shape(b, Nt, d))
|
Beam momenta over time. |
Mt |
(ndarray, shape(b, Nt, d, d))
|
Beam Hessians over time. |
At |
(ndarray, shape(b, Nt, 1))
|
Beam amplitudes over time. |
Notes
Assumes \(c(x)\) is homogeneous, M0 is diagonal, and
\(d\in\{1,2,3\}\). The diagonal and dimensionality assumptions are relaxed by
:func:solve_hom_general.
solve_hom_general(x0: jnp.ndarray, p0: jnp.ndarray, m0: jnp.ndarray, a0: jnp.ndarray, mode: jnp.ndarray, ts: jnp.ndarray, c: Callable, lam=None, config=None) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
Solver for homogeneous media with simplified equations.
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 Hessian matrices. |
required |
a0
|
(ndarray, shape(b))
|
Initial amplitudes. |
required |
mode
|
(ndarray, shape(b))
|
Hamiltonian branch signs. |
required |
ts
|
(ndarray, shape(Nt))
|
Time grid. |
required |
c
|
Callable
|
Homogeneous sound-speed function. |
required |
lam
|
Any
|
Ignored compatibility argument. |
None
|
config
|
Any
|
Ignored compatibility argument. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
xt |
(ndarray, shape(b, Nt, d))
|
Beam positions over time. |
pt |
(ndarray, shape(b, Nt, d))
|
Beam momenta over time. |
Mt |
(ndarray, shape(b, Nt, d, d))
|
Beam Hessians over time. |
At |
(ndarray, shape(b, Nt, 1))
|
Beam amplitudes over time. |
solve_hom_TR(xT: jnp.ndarray, pT: jnp.ndarray, mT: jnp.ndarray, aT: jnp.ndarray, mode: jnp.ndarray, ts: jnp.ndarray, c: Callable, lam=None, config=None) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
Time-reversal solver for homogeneous media.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
xT
|
(ndarray, shape(b, d))
|
Beam positions at the reference final time. |
required |
pT
|
(ndarray, shape(b, d))
|
Beam momenta at the reference final time. |
required |
mT
|
(ndarray, shape(b, d, d))
|
Beam Hessians at the reference final time. |
required |
aT
|
(ndarray, shape(b) or (b, 1))
|
Beam amplitudes at the reference final time. |
required |
mode
|
(ndarray, shape(b) or (b, 1))
|
Hamiltonian branch signs. |
required |
ts
|
ndarray
|
Per-beam or shared time grid. |
required |
c
|
Callable
|
Homogeneous sound-speed function. |
required |
lam
|
Any
|
Ignored compatibility argument. |
None
|
config
|
Any
|
Ignored compatibility argument. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
x0 |
(ndarray, shape(b, Nt, d))
|
Time-reversed beam positions. |
p0_time |
(ndarray, shape(b, Nt, d))
|
Time-reversed beam momenta. |
m0 |
(ndarray, shape(b, Nt, d, d))
|
Time-reversed Hessians. |
a0 |
(ndarray, shape(b, Nt, 1))
|
Time-reversed amplitudes. |
solve_ODE_base(x0: jnp.ndarray, p0: jnp.ndarray, M0: jnp.ndarray, a0: jnp.ndarray, mode: jnp.ndarray, ts: jnp.ndarray, c: Callable, lam: float = 0.0, solver_config: Optional[SolverConfig] = None) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
Solve the coupled system of ODEs for the GB with configurable solver settings.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x0
|
(ndarray, shape(d))
|
Initial beam position for one vmapped beam. |
required |
p0
|
(ndarray, shape(d))
|
Initial momentum for one vmapped beam. |
required |
M0
|
(ndarray, shape(d, d))
|
Initial Hessian for one vmapped beam. |
required |
a0
|
(ndarray, shape(1) or scalar)
|
Initial amplitude. |
required |
mode
|
ndarray
|
Hamiltonian branch sign. |
required |
ts
|
(ndarray, shape(Nt))
|
Time grid. |
required |
c
|
Callable
|
Sound-speed function. |
required |
lam
|
float
|
Absorption coefficient. |
0.0
|
solver_config
|
SolverConfig
|
Numerical solver configuration. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
xt |
(ndarray, shape(Nt, d))
|
Beam positions. |
pt |
(ndarray, shape(Nt, d))
|
Beam momenta. |
Mt |
(ndarray, shape(Nt, d, d))
|
Beam Hessians. |
At |
(ndarray, shape(Nt, 1))
|
Beam amplitudes. |
solve_ODE_batch_t(x0: jnp.ndarray, p0: jnp.ndarray, M0: jnp.ndarray, A0: jnp.ndarray, mode: jnp.ndarray, ts: jnp.ndarray, c: Callable, lam: Optional[float] = None, solver_config: Optional[SolverConfig] = None) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
Solve the coupled system of ODEs for the GB motion with per-batch time points.
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 Hessian matrices. |
required |
A0
|
(ndarray, shape(b))
|
Initial amplitudes. |
required |
mode
|
(ndarray, shape(b))
|
Hamiltonian branch signs. |
required |
ts
|
(ndarray, shape(b, Nt))
|
Per-beam time grids, commonly |
required |
c
|
Callable
|
Sound-speed function. |
required |
lam
|
float
|
Absorption coefficient. Currently unused by this no-absorption RHS. |
None
|
solver_config
|
SolverConfig
|
Numerical solver configuration. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
xt |
(ndarray, shape(b, Nt, d))
|
Beam positions. |
pt |
(ndarray, shape(b, Nt, d))
|
Beam momenta. |
Mt |
(ndarray, shape(b, Nt, d, d))
|
Beam Hessians. |
At |
(ndarray, shape(b, Nt, 1))
|
Beam amplitudes. |
solve_ODE_batch_t_terminal(x0: jnp.ndarray, p0: jnp.ndarray, M0: jnp.ndarray, A0: jnp.ndarray, mode: jnp.ndarray, ts: jnp.ndarray, c: Callable, lam: Optional[float] = None, solver_config: Optional[SolverConfig] = None) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
Solve per-beam intervals and save each terminal state.
The return values retain a singleton time axis.
Returns:
| Name | Type | Description |
|---|---|---|
xt, pt : jnp.ndarray, shape (b, 1, d)
|
Terminal beam positions and momenta. |
|
Mt |
(ndarray, shape(b, 1, d, d))
|
Terminal complex Hessians. |
At |
(ndarray, shape(b, 1, 1))
|
Terminal complex amplitudes. |