`beamax.solvers.msgb_solvers.tr_solver_utils`

Internal utilities for MSGB time-reversal reconstruction.

Scope

Time-reversal parameter assembly from data-domain coefficients.
TR beam aggregation and geometric/matrix helper routines.

API Reference

`beamax.solvers.msgb_solvers.tr_solver_utils`

`compute_mT_linear_system(xT: jnp.ndarray, pT: jnp.ndarray, mT_spc: Union[None, jnp.ndarray], mT_spc_time: Union[None, jnp.ndarray], mode: jnp.ndarray, c: Callable) -> jnp.ndarray`

Compute the linear system relating spatial and spacetime Hessians.

Parameters:

Name	Type	Description	Default
`xT`	`(ndarray, shape(b, d))`	Beam positions at the final time.	required
`pT`	`(ndarray, shape(b, d))`	Beam momenta at the final time.	required
`mT_spc`	`ndarray or None`	Spatial Hessian representation.	required
`mT_spc_time`	`ndarray or None`	Spacetime Hessian representation.	required
`mode`	`ndarray`	Beam branch signs.	required
`c`	`Callable`	Sound-speed function.	required

Returns:

Type	Description
`ndarray`	Converted Hessian representation.

Raises:

Type	Description
`ValueError`	If both Hessian representations are `None`.

`mT_forward(xT: jnp.ndarray, pT: jnp.ndarray, mT_spc: jnp.ndarray, mode: jnp.ndarray, c: Callable) -> jnp.ndarray`

Convert final-time spatial Hessians to spacetime Hessians.

Parameters:

Name	Type	Description	Default
`xT`	`(ndarray, shape(b, d))`	Beam positions at the final time.	required
`pT`	`(ndarray, shape(b, d))`	Beam momenta at the final time.	required
`mT_spc`	`(ndarray, shape(b, d, d))`	Spatial Hessian matrices.	required
`mode`	`ndarray`	Beam branch signs.	required
`c`	`Callable`	Sound-speed function.	required

Returns:

Type	Description
`(ndarray, shape(b, d, d))`	Spacetime Hessian matrices.

`mT_inverse(xT, pT, mT_spc_time, mode, c)`

Convert final-time spacetime Hessians to spatial Hessians.

Parameters:

Name	Type	Description	Default
`xT`	`(ndarray, shape(b, d))`	Beam positions at the final time.	required
`pT`	`(ndarray, shape(b, d))`	Beam momenta at the final time.	required
`mT_spc_time`	`(ndarray, shape(b, d, d))`	Spacetime Hessian matrices.	required
`mode`	`ndarray`	Beam branch signs.	required
`c`	`Callable`	Sound-speed function.	required

Returns:

Type	Description
`(ndarray, shape(b, d, d))`	Spatial Hessian matrices.

`find_constant_columns(arr, max_constant_axes=None)`

Find columns where all values are the same (e.g., the boundary normal axis).

Parameters:

Name	Type	Description	Default
`arr`	`(ndarray, shape(N, d))`	Array to check.	required
`max_constant_axes`	`int`	Maximum number of constant axes to return. For TR this is typically 1.	`None`

Returns:

Name	Type	Description
`constant_mask`	`(ndarray, shape(d))`	Boolean mask of constant columns.
`constant_values`	`ndarray`	Values of constant columns, padded with zeros.
`constant_axes`	`ndarray`	Indices of constant columns, padded with `-1`.

`compute_TR_parameters(significant_coeffs: jnp.ndarray, domain_data: Domain, wpt_data: MSWPT, sources: Sensor) -> Tuple[jnp.ndarray, ...]`

Compute the components for the time reversal problem.

Parameters:

Name	Type	Description	Default
`significant_coeffs`	`ndarray`	Flattened significant coefficient indices in the boundary-data WPT.	required
`domain_data`	`Domain`	Domain describing the boundary data coordinates.	required
`wpt_data`	`MSWPT`	Boundary-data wave-packet transform.	required
`sources`	`Sensor`	Source/sensor geometry on the acquisition boundary.	required

Returns:

Name	Type	Description
`pts`	`(ndarray, shape(B, d_spatial))`	Momentum at the final time.
`Mts`	`(ndarray, shape(B, d_spatial, d_spatial))`	Hessian matrix at the final time.
`xts`	`(ndarray, shape(B, d_spatial))`	Position at the final time.
	`ωs : jnp.ndarray, shape (B,)`	Beam frequency scale at the final time.
`ats`	`(ndarray, shape(B, 1))`	Beam amplitude scale at the final time.
`signum`	`(ndarray, shape(B, 1))`	Gaussian beam mode sign.
`ts`	`(ndarray, shape(B, 2))`	Per-beam time interval.

`compute_TR_result(params: Tuple[jnp.ndarray, ...], c: Callable, lam: float, sensors: jnp.ndarray, domain_size: jnp.ndarray, periodic: jnp.ndarray, ode_solver: SolverFn = None, aggregate_method: str = 'scan', solver_config: Optional[SolverConfig] = None, dt0: float | None = None) -> jnp.ndarray`

Compute time-reversal solution using Gaussian beams.

Unified interface similar to compute_forward_result.

Parameters:

Name	Type	Description	Default
`params`	`Tuple[ndarray, ...]`	Beam parameters `(p0, M0, x0, omega, a0, mode, ts)`.	required
`c`	`Callable`	Sound-speed function.	required
`lam`	`float`	Absorption parameter.	required
`sensors`	`ndarray`	Sensor positions.	required
`domain_size`	`ndarray`	Domain size.	required
`periodic`	`ndarray`	Boundary periodicity flags.	required
`ode_solver`	`SolverFn`	ODE solver. If `None`, uses :func:`solve_ODE_batch_t`.	`None`
`aggregate_method`	`(scan, vmap, all)`	Beam aggregation method.	`"scan"`
`solver_config`	`SolverConfig`	Numerical ODE configuration.	`None`
`dt0`	`float`	Optional initial time step passed to `solve_ODE_batch_t`.	`None`

Returns:

Type	Description
`ndarray`	Time-reversed field at sensor locations at the final time point.

Notes

TR requires per-beam time intervals (shape (b, 2)), so the ODE solver
must support this. Default is solve_ODE_batch_t. If passing a custom
solver, ensure it handles per-beam time arrays correctly.