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 falls back to the pure-Python backend until the upstream CppSimulation path ships source preprocessing for time-varying pressure sources. The adjoint also uses the Python backend and applies the Appendix-B source and terminal-field scalings needed for the discrete transpose.
• 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). Sensor channels follow NumPy C mask order for both the Python and standalone C++/CUDA backends.

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

Discrete k-Wave adjoint following Arridge et al., Appendix B.

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	'p_final'	Terminal pressure field. Other records cannot be converted to the initial-pressure transpose and are therefore rejected.	"p_final"

Returns:

Type	Description
ndarray	Euclidean discrete adjoint image A.T @ data with shape domain.N.

Notes

The returned field is the algebraic transpose under unweighted discrete sums. To represent a one-cell planar detector under the thesis's continuous surface and volume rectangle rules, multiply this result by dt / dx_normal.

The source is Eq. (B.2) of Arridge et al. (2016). k-Wave's additive pressure-source preprocessing multiplies a user source by 2*dt/(d*c*dx) before adding it to each split density field. Consequently the user source must be

rho_source*c_source*dx/(4*dt) * beta,

and the terminal pressure must be divided pointwise by c**2*rho. additive-no-correction is intentional: the optional cosine k-space filter on the injected source is not part of Eq. (B.2). This mode still applies k-Wave's additive-source amplitude scaling and does not disable the sinc k-space correction used by the propagation operators.

The current k-wave-python Python backend uses a different pressure scaling for each spatial axis on anisotropic grids. A single scalar pressure signal cannot then create the equal split-density increments required by Eq. (B.2), so this implementation requires isotropic grid spacing.

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.