`beamax.solvers.kwave_solver`

Wrapper around k-wave-python for forward, time reversal, and adjoint workflows.

Notes

Optional dependency: install with pip install .[kwave].
Supports 2D/3D k-Wave backends.

API Reference

`beamax.solvers.kwave_solver`

`KWaveSolver(simulation_options=None, execution_options=None, **kwargs)`

Bases: Solver

:mod:k-wave-python wrapper (2D/3D) for forward, time-reversal, and adjoint solves.

Provides a :class:beamax.solvers.Solver-compatible interface backed by the k-Wave pseudo-spectral time-domain solver. Used as the "reference" in examples and as the low-frequency leg of :class:HybridSolver.

Parameters:

Name	Type	Description	Default
`simulation_options`	`SimulationOptions`	Legacy simulation-options object. If provided with `execution_options`, they are converted to the unified kwargs dict.	`None`
`execution_options`	`SimulationExecutionOptions`	Legacy execution-options object. See `simulation_options`.	`None`
`**kwargs`		Unified k-Wave arguments forwarded to :func:`kwave.kspaceFirstOrder`. If neither legacy options nor kwargs are supplied, sensible CPU defaults are used.	`{}`

Notes

Requires the [kwave] extra (pip install 'beamax[kwave]').
Boundary handling uses a PML; :attr:Domain.periodic flags control PML-inside vs. PML-outside placement via the k-Wave options.
Forward solves run on the C++ backend by default; time-reversal and adjoint currently fall back to the pure-Python backend until the upstream CppSimulation path ships source preprocessing for time-varying pressure sources.
binary_path can be supplied directly or via BEAMAX_KWAVE_BINARY_PATH. Direct kwargs take precedence.

Examples:

>>> import jax.numpy as jnp
>>> from beamax import Domain, Sensor
>>> from beamax.solvers import KWaveSolver
>>> domain = Domain(N=(64, 64), dx=(1e-3, 1e-3), c=1500.0, periodic=(False, False))
>>> solver = KWaveSolver(pml_size=10, device="cpu")

Initialize k-Wave solver options.

Parameters:

Name	Type	Description	Default
`simulation_options`	`SimulationOptions`	Legacy simulation options object.	`None`
`execution_options`	`SimulationExecutionOptions`	Legacy execution options object.	`None`
`**kwargs`		Unified keyword options forwarded to :func:`kwave.kspaceFirstOrder`.	`{}`

`forward(p0: Union[np.ndarray, jnp.ndarray], domain: Domain, sensors: Union[np.ndarray, jnp.ndarray], ts: Union[np.ndarray, jnp.ndarray], *, record: str = 'p') -> np.ndarray`

Forward k-Wave simulation for linear wave equation.

Parameters:

Name	Type	Description	Default
`p0`	`array`	Initial pressure.	required
`domain`	`Domain`		required
`sensors`	`array`	Sensor mask or positions (solver expects mask).	required
`ts`	`(array, shape(Nt))`		required
`record`	`('p', ...)`	Quantity to record.	`"p"`

Returns:

Type	Description
`ndarray`	Sensor time series `(Nt, Ns)` or solver-specific shape.

`time_reversal(data: Union[np.ndarray, jnp.ndarray], domain, sensors, sources, ts: Union[np.ndarray, jnp.ndarray], *, record: str = 'p_final', data_layout: str = 'auto') -> np.ndarray`

Run classic k-Wave time reversal.

Parameters:

Name	Type	Description	Default
`data`	`(ndarray or ndarray, shape(Ns, Nt) or (Nt, Ns))`	Sensor measurements.	required
`domain`	`Domain`	Reconstruction domain.	required
`sensors`	`array - like`	Sensor mask.	required
`sources`	`array - like`	Source mask where the reversed data are injected.	required
`ts`	`(ndarray or ndarray, shape(Nt))`	Time grid.	required
`record`	`str`	k-Wave field to return.	`"p_final"`
`data_layout`	`('auto', 'ns_nt', 'nt_ns')`	Sensor-data layout interpretation.	`"auto"`

Returns:

Type	Description
`ndarray`	Requested k-Wave record.

Notes

Enforces p(x_s, t) = sensor_data(t, x_s) as a Dirichlet source.

`adjoint(data: Union[np.ndarray, jnp.ndarray], domain, sensors, sources, ts: Union[np.ndarray, jnp.ndarray], *, record: str = 'p_final', data_layout: str = 'auto') -> np.ndarray`

k-Wave adjoint following MATLAB demo convention.

Parameters:

Name	Type	Description	Default
`data`	`(array, shape(Ns, Nt) or (Nt, Ns))`	Sensor measurements.	required
`domain`	`Domain`		required
`sensors`	`array`	Sensor mask.	required
`sources`	`array`	Source mask.	required
`ts`	`(array, shape(Nt))`		required
`record`	`str`	What to record at the end (e.g., "p_final").	`'p_final'`

Returns:

Type	Description
`ndarray`	Adjoint image `(N...)`.
`Based off the original MATLAB k-Wave adjoint example: https://github.com/ucl-bug/k-wave/blob/main/k-Wave/examples/example_pr_2D_adjoint.m`

`TimedKWaveSolver(*args, mode: str = 'stdout', **kwargs)`

Bases: KWaveSolver

k-Wave wrapper that also returns a timing.

Timing modes

mode="stdout" (default): Parse k-Wave's own "Total execution time: " line from stdout. Intended to reflect simulation kernel time, not wrapper overhead. Fallback to wall-clock if parsing fails.

mode="wall": Wall-clock around the internal _run() call (includes whatever k-Wave does internally, excludes your pre/post Python work).

Initialize timed k-Wave solver.

Parameters:

Name	Type	Description	Default
`*args`		Positional arguments forwarded to :class:`KWaveSolver`.	`()`
`mode`	`('stdout', 'wall')`	Timing mode. `"stdout"` parses k-Wave output; `"wall"` uses Python wall-clock timing.	`"stdout"`
`**kwargs`		Keyword arguments forwarded to :class:`KWaveSolver`.	`{}`

Raises:

Type	Description
`ValueError`	If `mode` is not `"stdout"` or `"wall"`.

`forward(p0, domain, sensors, ts, *, record: str = 'p', **solver_kwargs) -> Tuple[np.ndarray, float]`

Same as KWaveSolver.forward, but returns (result, seconds).

Parameters:

Name	Type	Description	Default
`p0`	`array - like`	Initial pressure.	required
`domain`	`Domain`	Computational domain.	required
`sensors`	`array - like`	Sensor mask or positions.	required
`ts`	`(array - like, shape(Nt))`	Time grid.	required
`record`	`str`	k-Wave field to record.	`"p"`
`**solver_kwargs`		Additional solver keyword arguments.	`{}`

Returns:

Name	Type	Description
`result`	`ndarray`	Forward simulation result.
`seconds`	`float`	Execution time.

`time_reversal(data, domain, sensors, sources, ts, *, record: str = 'p_final', data_layout: str = 'auto', **solver_kwargs) -> Tuple[np.ndarray, float]`

Same as KWaveSolver.time_reversal, but returns (result, seconds).

Parameters:

Name	Type	Description	Default
`data`	`array - like`	Sensor measurements.	required
`domain`	`Domain`	Reconstruction domain.	required
`sensors`	`array - like`	Sensor mask.	required
`sources`	`array - like`	Source mask.	required
`ts`	`(array - like, shape(Nt))`	Time grid.	required
`record`	`str`	k-Wave field to return.	`"p_final"`
`data_layout`	`('auto', 'ns_nt', 'nt_ns')`	Sensor-data layout interpretation.	`"auto"`
`**solver_kwargs`		Additional solver keyword arguments.	`{}`

Returns:

Name	Type	Description
`result`	`ndarray`	Time-reversal result.
`seconds`	`float`	Execution time.

`adjoint(data, domain, sensors, sources, ts, *, record: str = 'p_final', data_layout: str = 'auto', **solver_kwargs) -> Tuple[np.ndarray, float]`

Same as KWaveSolver.adjoint, but returns (result, seconds).

Parameters:

Name	Type	Description	Default
`data`	`array - like`	Sensor measurements.	required
`domain`	`Domain`	Reconstruction domain.	required
`sensors`	`array - like`	Sensor mask.	required
`sources`	`array - like`	Source mask.	required
`ts`	`(array - like, shape(Nt))`	Time grid.	required
`record`	`str`	k-Wave field to return.	`"p_final"`
`data_layout`	`('auto', 'ns_nt', 'nt_ns')`	Sensor-data layout interpretation.	`"auto"`
`**solver_kwargs`		Additional solver keyword arguments.	`{}`

Returns:

Name	Type	Description
`result`	`ndarray`	Adjoint result.
`seconds`	`float`	Execution time.