beamax.solvers.hybrid_solver

High-level hybrid solver combining low-frequency and high-frequency backends.

Key Objects

• HybridSolver: orchestrates LF/HF split, optional downsampling, and merge logic.
• HybridBackend: adapter for arbitrary LF solver operations. It accepts any subset of forward, time_reversal, and adjoint callables, as long as at least one is present.
• HybridContext: stable payload passed to LF backend callables. It contains domains, masks, sensors, time grids, WPT objects, target shape, and split configuration.
• HybridSolverConfig: typed configuration for interpolation, cutoff behavior, and runtime options.

LF backends do not need to subclass Solver. Each operation callable receives (component_array, context) and may adapt that context into whatever argument order the underlying LF solver expects. Missing operations fail only when that operation is called.

For existing beamax-style solvers, use HybridBackend.from_beamax_solver(existing_solver).

API Reference

beamax.solvers.hybrid_solver

Hybrid solver for combining low- and high-frequency propagation backends.

The high-frequency component is typically handled by MSGB, while the low-frequency component is supplied through a small adapter callable API.

HybridSolverConfig(box_corners: Optional[jnp.ndarray] = None, cutoff_freq: Optional[float] = None, downsample: bool = True, use_pow2: bool = True, input_type: str = 'spatial', interp_method: str = 'fourier', dt_oversample: int = 30, beta: float = 12.0, order: int = 3, window_type: str = 'kaiser', use_windowing: bool = True, use_time_extension: bool = True) dataclass

Configuration for hybrid solver.

Parameters:

Name	Type	Description	Default
box_corners	ndarray	Indices defining LF/HF split region in Fourier space	None
cutoff_freq	float	Alternative to box_corners - split by frequency magnitude	None
downsample	bool	Whether to solve LF on downsampled grid	True
use_pow2	bool	Round downsampled size to power of 2	True
input_type	str	"spatial" or "fourier" - domain of input data	'spatial'
interp_method	str	"fourier" (periodic domains) or "zoom" (non-periodic)	'fourier'
dt_oversample	int	Number of extra time steps for windowing (forward only)	30
beta	float	Kaiser window shape parameter (higher = sharper transition)	12.0
order	int	Spline order for zoom interpolation (0-5)	3
window_type	str	"kaiser" or "tukey" - windowing function type	'kaiser'
use_windowing	bool	Whether to apply windowing (typically True for forward, False for TR)	True
use_time_extension	bool	Whether to extend time array (typically True for forward, False for TR)	True

HybridContext(operation: HybridOperationName, config: HybridSolverConfig, domain: Domain, input_domain: Domain, component_domain: Domain, full_sensors: Any, component_sensors: Any, full_sensor_mask: jnp.ndarray, component_sensor_mask: jnp.ndarray, ts: jnp.ndarray, original_ts: jnp.ndarray, target_shape: Tuple[int, ...], sources: Any = None, wpt: Optional[MSWPT] = None, data_wpt: Optional[MSWPT] = None, img_wpt: Optional[MSWPT] = None, data_domain: Optional[Domain] = None) dataclass

Runtime context passed to a low-frequency hybrid backend operation.

The backend receives the already split low-frequency component plus this object. It may use any of the fields it understands and ignore the rest; the :class:HybridSolver still owns splitting, optional downsampling, time extension/windowing, interpolation, and HF/LF merging.

Attributes:

Name	Type	Description
operation	{forward, time_reversal, adjoint}	Operation currently being dispatched.
config	HybridSolverConfig	Hybrid split/downsampling/windowing configuration.
domain	Domain	Full-resolution output domain. For forward this is the physical simulation domain; for inverse operations this is the reconstruction domain.
input_domain	Domain	Full-resolution domain of the data being split. This is usually the same as domain for forward and data_domain for inverse operations.
component_domain	Domain	Domain for the LF component actually passed to the backend. It may be downsampled when config.downsample is true.
full_sensors	object	Original sensor object or mask supplied by the caller.
component_sensors	object	Sensor representation aligned to component_domain.
full_sensor_mask	ndarray	Full-resolution sensor mask.
component_sensor_mask	ndarray	Sensor mask aligned to component_domain.
ts	ndarray	Time grid passed to the LF operation. This may be extended for forward solves.
original_ts	ndarray	Time grid supplied by the caller before any hybrid extension.
target_shape	tuple[int, ...]	Shape the LF result will be interpolated/truncated to before merging.
sources	(object, optional)	Source geometry for inverse operations.
wpt, data_wpt, img_wpt	(MSWPT, optional)	Wave-packet transforms relevant to the operation.
data_domain	(Domain, optional)	Full-resolution data domain for inverse operations.

split_config: HybridSolverConfig property

Alias for config for adapters that name it by responsibility.

HybridBackend(forward: Optional[HybridOperation] = None, time_reversal: Optional[HybridOperation] = None, adjoint: Optional[HybridOperation] = None, name: str = 'low-frequency backend') dataclass

Adapter for a low-frequency backend used by :class:HybridSolver.

Parameters:

Name	Type	Description	Default
forward	callable	Operation callables with signature callable(component_array, context) -> array. At least one operation must be provided.	None
time_reversal	callable	Operation callables with signature callable(component_array, context) -> array. At least one operation must be provided.	None
adjoint	callable	Operation callables with signature callable(component_array, context) -> array. At least one operation must be provided.	None
name	str	Human-readable backend name used in error messages.	'low-frequency backend'

operations: Tuple[str, ...] property

Operation names implemented by this backend.

supports(operation: HybridOperationName) -> bool

Return whether this backend implements operation.

require(operation: HybridOperationName) -> HybridOperation

Return an operation callable or raise a clear missing-operation error.

Parameters:

Name	Type	Description	Default
operation	(forward, time_reversal, adjoint)	Operation required by the hybrid solve.	"forward"

from_beamax_solver(solver: Any, *, name: Optional[str] = None) -> HybridBackend classmethod

Wrap a beamax-style solver object as a low-frequency backend.

The wrapped solver may implement any subset of forward, time_reversal, and adjoint. Each method is called with the current component-domain objects from :class:HybridContext.

InterpolationStrategy

Bases: ABC

Abstract base for spatial interpolation strategies.

interpolate(data: jnp.ndarray, target_shape: Tuple) -> jnp.ndarray abstractmethod

Interpolate data to a target shape.

Parameters:

Name	Type	Description	Default
data	ndarray	Input array.	required
target_shape	Tuple[int, ...]	Desired output shape.	required

Returns:

Type	Description
ndarray	Interpolated array.

FourierInterpolation

Bases: InterpolationStrategy

Fourier-based interpolation (periodic boundaries assumed).

interpolate(data: jnp.ndarray, target_shape: Tuple) -> jnp.ndarray

Interpolate by Fourier-domain padding or cropping.

Parameters:

Name	Type	Description	Default
data	ndarray	Spatial-domain array to resize.	required
target_shape	Tuple[int, ...]	Desired output shape.	required

Returns:

Type	Description
ndarray	Real-valued resized array with sample-value normalization.

ZoomInterpolation(order: int = 3)

Bases: InterpolationStrategy

Spline-based interpolation (handles non-periodic domains).

Initialize spline interpolation order.

Parameters:

Name	Type	Description	Default
order	int	Spline order passed to :func:scipy.ndimage.zoom.	3

interpolate(data: jnp.ndarray, target_shape: Tuple) -> jnp.ndarray

Interpolate by spline zoom factors.

Parameters:

Name	Type	Description	Default
data	ndarray	Input array.	required
target_shape	Tuple[int, ...]	Desired output shape.	required

Returns:

Type	Description
ndarray	Resized array from :func:scipy.ndimage.zoom.

HybridSolver(*, hf_solver: Any, lf_backend: HybridBackend, config: Optional[HybridSolverConfig] = None, **config_kwargs)

Bases: Module

Hybrid MSGB / low-frequency backend solver.

Splits the input pressure in frequency: a fast high-frequency solver (typically :class:beamax.solvers.MSGBSolver) handles the sparse, highly oscillatory content, while a configurable low-frequency backend handles the smooth residual, optionally on a coarser grid. The two results are added back together on the target grid.

Parameters:

Name	Type	Description	Default
hf_solver	MSGBSolver or object	Solver used for the high-frequency component.	required
lf_backend	HybridBackend	Low-frequency backend adapter. Each operation callable receives the LF component and a :class:HybridContext.	required
config	HybridSolverConfig	Frequency-split and windowing configuration.	None

Notes

The LF backend does not need to subclass :class:beamax.solvers.Solver. It only needs to provide at least one operation through :class:HybridBackend. Missing operations fail when called, not at hybrid construction time.

Initialize hybrid solver.

Parameters:

Name	Type	Description	Default
hf_solver	MSGBSolver or object	High-frequency solver.	required
lf_backend	HybridBackend	Low-frequency backend adapter.	required
config	HybridSolverConfig	Configuration object.	None
**config_kwargs		Alternative: pass config parameters as kwargs.	{}

create_with_domain(*, hf_solver: Any, lf_backend: HybridBackend, domain: Domain, config: Optional[HybridSolverConfig] = None, **config_kwargs) -> HybridSolver staticmethod

Factory method with automatic interpolation selection based on domain.

Uses Fourier interpolation for periodic domains, spline interpolation for non-periodic domains.

Parameters:

Name	Type	Description	Default
hf_solver	MSGBSolver or object	High-frequency solver.	required
lf_backend	HybridBackend	Low-frequency backend adapter.	required
domain	Domain	Domain whose periodic flags determine interpolation choice.	required
config	HybridSolverConfig	Explicit configuration. If provided, config_kwargs are ignored.	None
**config_kwargs		Configuration options used when config is None.	{}

Returns:

Type	Description
HybridSolver	Configured hybrid solver.

forward(p0: jnp.ndarray, domain: Domain, sensors: Sensor, ts: jnp.ndarray, wpt: MSWPT) -> jnp.ndarray

Forward solve with hybrid approach.

Parameters:

Name	Type	Description	Default
p0	ndarray	Initial pressure field	required
domain	Domain	Computational domain	required
sensors	Sensor	Sensor geometry	required
ts	ndarray	Time points	required
wpt	MSWPT	Wavelet packet transform for frequency splitting	required

Returns:

Type	Description
ndarray	Sensor data, shape (Nt, Ns)

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

Time reversal with hybrid approach.

Parameters:

Name	Type	Description	Default
data	ndarray	Recorded sensor data	required
domain	Domain	Reconstruction domain	required
sensors	Sensor	Sensor geometry	required
sources	Sensor	Source positions for data injection	required
ts	ndarray	Time points	required
data_domain	Domain	Domain where data was recorded	required
data_wpt	MSWPT	Wavelet transform for data domain	required
img_wpt	MSWPT	Wavelet transform for reconstruction domain	required

Returns:

Type	Description
ndarray	Reconstructed initial pressure field

Notes

Time extension/windowing not used for time reversal (applies only to forward).

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

Adjoint solve with hybrid approach.

Similar to time_reversal but uses adjoint method.

Parameters:

Name	Type	Description	Default
data	ndarray	Recorded sensor data	required
domain	Domain	Reconstruction domain	required
sensors	Sensor	Sensor geometry	required
sources	Sensor	Source positions for data injection	required
ts	ndarray	Time points	required
data_domain	Domain	Domain where data was recorded	required
data_wpt	MSWPT	Wavelet transform for data domain	required
img_wpt	MSWPT	Wavelet transform for reconstruction domain	required

Returns:

Type	Description
ndarray	Adjoint reconstruction