beamax.solvers.hybrid_solver

Low-/high-frequency solver composition. HybridBackend adapts operation callables; HybridContext supplies their domain and reconstruction state.

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. Larger values sharpen the transition.	12.0
order	int	Spline order for zoom interpolation; one of \(0,1,\ldots,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(domain: Domain, component_domain: Domain, component_sensor_mask: jnp.ndarray, ts: jnp.ndarray, target_shape: Tuple[int, ...], sources: Any = None) dataclass

Runtime inputs passed to a low-frequency hybrid backend.

:class:HybridSolver owns splitting, optional downsampling, windowing, interpolation, and merging. The context contains only the geometry and shapes needed by the backend adapter.

Attributes:

Name	Type	Description
domain	Domain	Full-resolution output domain. For forward this is the physical simulation domain; for inverse operations this is the reconstruction domain.
component_domain	Domain	Domain for the LF component passed to the backend.
component_sensor_mask	ndarray	Sensor mask aligned to component_domain.
ts	ndarray	Time grid passed to the LF operation.
target_shape	tuple[int, ...]	Shape the LF result will be interpolated/truncated to before merging.
sources	(object, optional)	Source geometry for inverse operations.

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.

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 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.	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