beamax.solvers.msgb_solvers.msgb_solver

Forward, time-reversal, and adjoint MSGB orchestration, including stage overrides and optional beam sharding.

beamax.solvers.msgb_solvers.msgb_solver

MSGBExperimentalConfig(coefficient_selection: CoefficientSelection = 'auto', boxes_per_chunk: int = DEFAULT_BOXES_PER_CHUNK, forward_kernel: ForwardKernel = 'auto', inverse_evaluator: InverseEvaluator = 'auto') dataclass

Experimental per-stage execution overrides.

"auto" selects supported defaults. Streaming top-n requires real, zero-velocity, fixed-size input. Generic Pallas supports GPU and TPU; hom_diag_3d_pallas requires a GPU.

ShardingStrategy(mesh: Mesh, beam_axis: str = 'x') dataclass

Strategy for sharding beam parameters across devices.

Attributes:

Name	Type	Description
mesh	Mesh	JAX device mesh for multi-device parallelization.
beam_axis	str	Mesh axis used to shard beams.

shard_beam_params(p0: jnp.ndarray, M0: jnp.ndarray, x0: jnp.ndarray, omega: jnp.ndarray, a0: jnp.ndarray, modes: jnp.ndarray) -> Tuple[jnp.ndarray, ...]

Shard forward beam parameters along the beam dimension.

Parameters:

Name	Type	Description	Default
p0	ndarray	Beam momenta.	required
M0	ndarray	Beam Hessian matrices.	required
x0	ndarray	Beam positions.	required
omega	ndarray	Beam frequencies.	required
a0	ndarray	Beam amplitudes.	required
modes	ndarray	Beam branch signs.	required

Returns:

Type	Description
Tuple[ndarray, ...]	Device-placed arrays with sharding specifications applied.

shard_tr_params(pts: jnp.ndarray, Mts: jnp.ndarray, xts: jnp.ndarray, omega_ts: jnp.ndarray, ats: jnp.ndarray, signum: jnp.ndarray, ts: jnp.ndarray) -> Tuple[jnp.ndarray, ...]

Shard time-reversal beam parameters along the beam dimension.

Parameters:

Name	Type	Description	Default
pts	ndarray	Beam momenta at the boundary.	required
Mts	ndarray	Beam Hessians at the boundary.	required
xts	ndarray	Boundary positions.	required
omega_ts	ndarray	Beam frequencies.	required
ats	ndarray	Beam amplitudes.	required
signum	ndarray	Beam branch signs.	required
ts	ndarray	Per-beam time intervals.	required

Returns:

Type	Description
Tuple[ndarray, ...]	Device-placed arrays with sharding specifications applied.

MSGBSolver(thr: Union[int, float], thr_strat: str, batch_size: int, input_type: str, ode_solver: SolverFn, sum_method: str, tr_ode_solver: Optional[SolverFn] = None, sharding: Optional[ShardingStrategy] = None, ode_config: Optional[SolverConfig] = None, adjoint_relative_guard: float = 0.05, pallas_config: Optional[PallasConfig] = None, experimental_config: Optional[MSGBExperimentalConfig] = None)

Bases: Module

Multiscale Gaussian Beam solver for the linear wave equation.

Implements forward, time-reversal, and adjoint operators by:

1. Decomposing the initial pressure into wave-packet coefficients via :class:beamax.transforms.MSWPT.
2. Thresholding to retain only significant coefficients.
3. Integrating a small Hamiltonian ODE per retained beam.
4. Summing (or scanning) the beam contributions at the sensor positions.

Parameters:

Name	Type	Description	Default
thr	int or float	Absolute magnitude threshold or top-k count, depending on thr_strat.	required
thr_strat	str	Thresholding strategy: "hard" or "top_n".	required
batch_size	int	Batch size along the beam axis for ODE integration. Tune to fit device memory; larger values amortise kernel launches.	required
input_type	(spatial, fourier)	Domain the caller provides p0 in.	"spatial"
ode_solver	SolverFn	Forward-time ODE integrator for beam dynamics (typically one of :mod:beamax.gb.gb_solvers).	required
sum_method	str	Method for summing beam contributions. One of "all_real", "scan_real", "all_complex", "scan_complex", or the experimental accelerator mode "pallas_real".	required
tr_ode_solver	SolverFn	ODE integrator for the time-reversal dynamics. Falls back to ode_solver when None.	None
sharding	ShardingStrategy	Multi-device sharding strategy. None runs on a single device.	None
ode_config	SolverConfig	Numerical configuration passed through to the ODE integrator. Falls back to SolverConfig.from_precision().	None
adjoint_relative_guard	float	Dimensionless near-grazing exclusion \(\Gamma/|\tau|\) used by the principal-symbol adjoint. Tune together with the ODE configuration.	5e-2
pallas_config	PallasConfig	Experimental static kernel configuration used by pallas_real and any explicitly selected Pallas stage.	None
experimental_config	MSGBExperimentalConfig	Per-stage pipeline overrides. Omitting it is equivalent to passing a configuration with every field set to "auto": streamed top-n selection and terminal-only inverse evaluation apply whenever they are output-equivalent, and Pallas kernels stay off.	None

Initialize the MSGB solver.

Parameters:

Name	Type	Description	Default
thr	int or float	Threshold value for coefficient selection.	required
thr_strat	str	Thresholding strategy name.	required
batch_size	int	Number of beams per batch for scanned aggregation.	required
input_type	(spatial, fourier)	Domain of inputs supplied to the solver.	"spatial"
ode_solver	SolverFn	Forward ODE solver.	required
sum_method	str	Aggregation mode string. Must be one of the values listed in the class-level parameter documentation.	required
tr_ode_solver	SolverFn	ODE solver for time reversal. Defaults to ode_solver.	None
sharding	ShardingStrategy	Multi-device sharding strategy.	None
ode_config	SolverConfig	Numerical ODE solver configuration. Defaults to SolverConfig.from_precision().	None
adjoint_relative_guard	float	Dimensionless near-grazing exclusion for the adjoint. Must lie in \([0,1)\).	5e-2
pallas_config	PallasConfig	Experimental static Pallas layout and launch configuration.	None
experimental_config	MSGBExperimentalConfig	Per-stage overrides: force the streamed selector, choose the inverse evaluator, or opt in to the experimental Pallas forward/inverse kernels. Omitting it selects the supported defaults for every stage.	None

forward(p0: jnp.ndarray, domain: Domain, sensors: Union[Sensor, jnp.ndarray], ts: jnp.ndarray, wpt: MSWPT, *, dpdt: Optional[jnp.ndarray] = None) -> jnp.ndarray

Solve the forward wave equation with MSGB:

\[ \partial_t^2u-c(\mathbf{x})^2\Delta u=0, \qquad u(0,\mathbf{x})=p_0(\mathbf{x}), \qquad \partial_tu(0,\mathbf{x})=\dot p_0(\mathbf{x}). \]

The dpdt argument supplies \(\dot p_0\) and is zero for standard photoacoustic tomography.

Parameters:

Name	Type	Description	Default
p0	(ndarray, shape(*N))	Initial pressure field. Real or complex; dtype selects the underlying beam formulation.	required
domain	Domain	Computational domain and medium.	required
sensors	Sensor or ndarray	Sensor geometry. Either a :class:Sensor or an array of positions in physical units, shape (Ns, ndim).	required
ts	(ndarray, shape(Nt))	Time grid.	required
wpt	MSWPT	Wave-packet transform used to build the beam decomposition.	required
dpdt	ndarray	Initial time derivative. Defaults to zeros (standard PAT).	None

Returns:

Type	Description
(ndarray, shape(Nt, Ns))	Pressure at each sensor over time.

forward_with_params(p0: jnp.ndarray, domain: Domain, sensors: Union[Sensor, jnp.ndarray], ts: jnp.ndarray, wpt: MSWPT, *, dpdt: Optional[jnp.ndarray] = None) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, ...]]

Forward MSGB solve plus diagnostic beam parameters.

This is the explicit diagnostic variant of :meth:forward. Most users should call :meth:forward, which returns only sensor data.

Returns:

Name	Type	Description
sensor_data	(ndarray, shape(Nt, Ns))	Pressure at each sensor over time.
params	tuple of jnp.ndarray	Beam parameters used in the solve: (p0s, M0s, x0s, omegas, a0s, modes).

time_reversal(data: jnp.ndarray, domain: Domain, sensors: Sensor, sources: Sensor, ts, data_domain: Domain, data_wpt: MSWPT) -> jnp.ndarray

MSGB time-reversal reconstruction.

Parameters:

Name	Type	Description	Default
data	(ndarray, shape(Nt, Ns))	Sensor time series to time-reverse.	required
domain	Domain	Reconstruction domain. Must have periodic all False — time-reversal here assumes free-space boundaries.	required
sensors	Sensor	Sensor geometry corresponding to data.	required
sources	Sensor	Source positions used to seed the TR beams (often the same boundary as sensors).	required
ts	(ndarray, shape(Nt))	Time grid corresponding to data.	required
data_domain	Domain	Domain on which data was acquired (may differ from domain under downsampling).	required
data_wpt	MSWPT	Wave-packet transform on data_domain used to analyse data.	required

Returns:

Type	Description
(ndarray, shape(*N))	Reconstructed initial pressure. Scaled by 2 to match the standard full-field time-reversal convention.

Raises:

Type	Description
ValueError	If any axis of domain is periodic.

Notes

Time reversal here uses per-beam time intervals (via :func:beamax.gb.solve_ODE_batch_t) regardless of the forward integrator, to accommodate the variable emission time of each beam.

time_reversal_with_params(data: jnp.ndarray, domain: Domain, sensors: Sensor, sources: Sensor, ts, data_domain: Domain, data_wpt: MSWPT) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, ...]]

Time-reversal MSGB solve plus diagnostic beam parameters.

This is the explicit diagnostic variant of :meth:time_reversal. Most users should call :meth:time_reversal, which returns only the reconstructed field.

Returns:

Name	Type	Description
p0_recon	(ndarray, shape(*N))	Reconstructed initial pressure.
params	tuple of jnp.ndarray	Beam parameters used in the solve.

adjoint(data: jnp.ndarray, domain: Domain, sensors: Union[Sensor, jnp.ndarray], sources: Sensor, ts: jnp.ndarray, data_domain: Domain, data_wpt: MSWPT, *, window: Optional[jnp.ndarray] = None) -> jnp.ndarray

Principal-symbol MSGB approximation of the continuous PAT adjoint.

Parameters:

Name	Type	Description	Default
data	ndarray	Boundary residual \(r(t,\mathbf{x}_s)\) on \(\Gamma\). Shape (Nt, Ns) or (Nt,), with time along axis 0.	required
domain	Domain	Reconstruction domain for \(q(T,\mathbf{x})\).	required
sensors	Sensor or ndarray	Adjoint evaluation locations. Image reconstruction typically uses domain.grid to evaluate \(q_T\) on the full grid.	required
sources	Sensor	Source geometry on \(\Gamma\), used to construct the boundary beam parameters (same role as in time reversal).	required
ts	ndarray	Time grid, shape (Nt,). Reserved for interface consistency.	required
data_domain	Domain	Domain describing the \((t,\mathbf{x}_s)\) grid of the boundary data. Its dx[0] is used as the time step \(\Delta t\).	required
data_wpt	MSWPT	MSWPT instance for analysing the boundary data / source.	required
window	ndarray	Sampled acquisition window. It may have the same shape as data or shape (Nt,), in which case it is broadcast over sensors. The default is one on the sampled acquisition array and assumes the residual is negligible at the temporal endpoints; otherwise pass a taper that vanishes there.	None

Returns:

Type	Description
ndarray	Principal-symbol approximation to \(P^*r\) for unweighted image and data \(L^2\) pairings. This is not the exact transpose of the thresholded discrete MSGB forward solver.

adjoint_with_params(data: jnp.ndarray, domain: Domain, sensors: Union[Sensor, jnp.ndarray], sources: Sensor, ts: jnp.ndarray, data_domain: Domain, data_wpt: MSWPT, *, window: Optional[jnp.ndarray] = None) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, ...]]

Adjoint MSGB solve plus diagnostic beam parameters.

This is the explicit diagnostic variant of :meth:adjoint. Most users should call :meth:adjoint, which returns only the adjoint field.

Returns:

Name	Type	Description
q_T	ndarray	Principal-symbol approximation to \(P^*r\) on the reconstruction domain under unweighted \(L^2\) pairings.
params	tuple of jnp.ndarray	Beam parameters used internally.

form_adjoint_source(data: jnp.ndarray, dt: float, c_at_sources: jnp.ndarray, window: Optional[jnp.ndarray] = None) -> jnp.ndarray

Form the acquisition-time source for the unweighted wave equation.

For detector residual \(r(s,\mathbf{x}_s)\), this returns

\[ -c(\mathbf{x}_s)^2\,\partial_s \left[w(s,\mathbf{x}_s)r(s,\mathbf{x}_s)\right]. \]

It is the acquisition-time representation of the time-reversed source in the continuous PAT adjoint. A one-dimensional window is interpreted as a time window and broadcast over detector axes.

apply_adjoint_image_weight(terminal_field: jnp.ndarray, c_at_image: jnp.ndarray) -> jnp.ndarray

Apply the \(c^{-2}\) weight for an unweighted image-space pairing.