beamax.plotter

Plotting helpers for wavefields, coefficients, and solver diagnostics.

Key Objects

• PlotHelper: convenience wrapper around common matplotlib visualizations.
• Free plotting functions: animations and comparison-style diagnostics for examples/analysis.

API Reference

beamax.plotter

Plotting helpers for beamax workflows (matplotlib/pyvista).

This module collects visualization utilities built on the optional viz dependencies: matplotlib for standard figures and pyvista for selected 3D views. Most functions accept NumPy/JAX arrays and are intended for diagnostics rather than production rendering.

PlotHelper(figsize=(12, 5), font_size=12, title_font_size=14, colormap='viridis')

Convenience wrapper around common plotting patterns.

Parameters:

Name	Type	Description	Default
figsize	Tuple[int, int]	Figure size passed to matplotlib.	(12, 5)
font_size	int	Base font size for axes labels.	12
title_font_size	int	Title font size.	14
colormap	str	Name of a matplotlib colormap to use for heatmaps/surfaces.	"viridis"

Store global style defaults for plots.

Parameters:

Name	Type	Description	Default
figsize	Tuple[int, int]	Figure size passed to matplotlib.	(12, 5)
font_size	int	Base font size for labels.	12
title_font_size	int	Title font size.	14
colormap	str	Matplotlib colormap name.	"viridis"

setup_figure(title=None, projection=None)

Create a figure/axes pair with optional projection and title.

Parameters:

Name	Type	Description	Default
title	str | None		None
projection	str | None	e.g. "3d" for surface plots.	None

Returns:

Type	Description
(Figure, Axes)

set_labels(ax, xlabel=None, ylabel=None, zlabel=None)

Apply axis labels if provided.

Parameters:

Name	Type	Description	Default
ax	Axes	Target axes.	required
xlabel	str	X-axis label.	None
ylabel	str	Y-axis label.	None
zlabel	str	Z-axis label for 3D axes.	None

save_plot(filename)

Save the current matplotlib figure and close it.

Parameters:

Name	Type	Description	Default
filename	str or path - like	Output path passed to :func:matplotlib.pyplot.savefig.	required

plot_coeffs_with_boxes(coeffs, dyadic_decomp, *, ax=None, plane: str = 'xy', slice_index: int | None = None, cmap: str | None = None, title: str | None = None, show_colorbar: bool = True, extent=None)

Plot transform coefficients with dyadic box overlays (1D/2D/3D).

Parameters:

Name	Type	Description	Default
coeffs	array - like	Coefficients arranged on a regular grid (1D/2D/3D).	required
dyadic_decomp	DyadicDecomposition	Decomposition whose boxes will be overlaid.	required
ax	matplotlib Axes | None	Axes to draw on; if None a new figure/axes is created.	None
plane	(xy, xz, yz)	Projection plane for 3D coefficients. Ignored for 1D/2D.	"xy","xz","yz"
slice_index	int | None	Slice to show for 3D; defaults to central slice along the axis orthogonal to plane.	None
cmap	str | None	Matplotlib colormap name.	None
title	str | None	Title for the axes.	None
show_colorbar	bool	Whether to attach a colorbar (if a new figure is created).	True
extent	tuple | None	Optional imshow extent to use for 2D/3D projections.	None

Returns:

Type	Description
(fig, ax)

plot_wavefield(data, complex=False, *args, **kwargs)

Dispatch to the appropriate real/complex wavefield plotter.

Parameters:

Name	Type	Description	Default
data	ndarray	Wavefield with shape (Nx, ...) (1D/2D/3D).	required
complex	bool	If True, splits into real/imaginary panels.	False
*args		Passed through to the underlying plotting helper (e.g. titles, axes labels, sensor markers).	()
**kwargs		Passed through to the underlying plotting helper (e.g. titles, axes labels, sensor markers).	()

plot_pressure_time(p, t, title=None, filename=None)

Plot time traces or space–time image of pressure data.

Parameters:

Name	Type	Description	Default
p	ndarray	Pressure array shaped (Nt,) or (Nt, Ns).	required
t	ndarray	Time samples aligned with the first dimension of p.	required
title	str | None		None
filename	str | None	If provided, save instead of displaying interactively.	None

plot_centers(centers, title=None, filename=None)

Visualize dyadic box centres in 1D/2D/3D.

Parameters:

Name	Type	Description	Default
centers	ndarray	Array of integer centre coordinates with shape (B, d).	required
title	str | None		None
filename	str | None	Optional path to save the figure instead of showing.	None

plot_centers_slice_grid(centers, axis='z', slice_values=None, cols=3, title=None, font_size=None, figsize=None, save_path=None, show=True)

Plot labelled centers as 2D scatter slices along a chosen axis.

Parameters:

Name	Type	Description	Default
centers	(array - like, shape(n_points, 3))	3D coordinates.	required
axis	(x, y, z, 0, 1, 2)	Which axis to slice along.	"x","y","z",0,1,2
slice_values	iterable | None	Specific coordinate values to plot. If None, uses all unique values.	None
cols	int	Number of subplot columns.	3
title	str | None	Figure title.	None
font_size	int | None	Base font size for labels (defaults to self.font_size).	None
figsize	tuple | None	Optional matplotlib figsize. Defaults to (cols*4, rows*4).	None
save_path	str | Path | None	If provided, path to save the PNG.	None
show	bool	Whether to call plt.show().	True

plot_coefficients(coeffs: jnp.ndarray, wpt, sort_data=False, log_scale=False, filename=None)

Scatter-plot MSWPT coefficients coloured by dyadic level.

Parameters:

Name	Type	Description	Default
coeffs	ndarray	Flattened coefficient vector.	required
wpt	MSWPT	Transform object (used to infer level layout).	required
sort_data	bool	If True, sort by magnitude descending.	False
log_scale	bool	Plot log10(|c|) instead of |c|.	False
filename	str | None	Optional path to save the figure.	None

plot_speed_of_sound(c: jnp.ndarray, domain: Domain, title: str = 'Speed of Sound Map', filename: str = None, num_contours: int = 20)

Plot a speed-of-sound field in 2D/3D.

Parameters:

Name	Type	Description	Default
c	ndarray	Scalar field sampled on domain.N.	required
domain	Domain	Provides grid spacing and dimensionality.	required
title	str		"Speed of Sound Map"
filename	str | None	Save path (optional).	None
num_contours	int	Number of contour levels for 2D plots.	20

Raises:

Type	Description
ValueError	If domain.ndim is not 2 or 3.

plot_coeff_surf(coeffs, N, title=None, filename=None)

Surface plot for a flattened 4N² coefficient vector.

Parameters:

Name	Type	Description	Default
coeffs	ndarray	1D array expected to be length 4 * N**2.	required
N	int	Reshape factor (resulting grid is (2N, 2N)).	required
title	str | None		None
filename	str | None		None

use_beamax_style() -> None

Apply a lightweight matplotlib style for example figures.

Notes

This function updates a small set of rcParams directly so examples remain portable in editable checkouts, wheels, and notebooks.

animate_wavefield_2d(data, skip_frames=1, cmap='viridis', label='', timesteps=None)

Animate a time-series of 2D fields with matplotlib.

Parameters:

Name	Type	Description	Default
data	ndarray	Array shaped (Nt, Ny, Nx).	required
skip_frames	int	Step between frames to thin out long sequences.	1
cmap	str	Colormap name.	"viridis"
label	str	Colorbar label.	""
timesteps	ndarray | None	Optional time stamps used in the frame title.	None

animate_wavefield_1d(data, skip_frames=1)

Animate 1D traces over time.

Parameters:

Name	Type	Description	Default
data	ndarray	Array shaped (Nt, Nx).	required
skip_frames	int	Step between frames to thin out long sequences.	1

animate_wavefield(data, skip_frames=1, cmap='viridis')

Dispatch animation to 1D/2D helpers based on data dimensionality.

Parameters:

Name	Type	Description	Default
data	ndarray	Array shaped (Nt, Nx) or (Nt, Ny, Nx).	required
skip_frames	int		1
cmap	str		"viridis"

plot_wavefields_interactive(reference_all, p0_msgb_all, ts, domain: Domain, XY, plot_type='imshow', fixed_diff_colorbar=False)

Interactive comparison between Reference and MSGB wavefields.

Parameters:

Name	Type	Description	Default
reference_all	ndarray	Wavefield sequence from the Reference solver (Nt, ...).	required
p0_msgb_all	ndarray	Wavefield sequence from the MSGB solver (Nt, ...).	required
ts	ndarray	Time points associated with frames.	required
domain	Domain	Provides spatial coordinates.	required
XY	ndarray	Meshgrid used for plotting.	required
plot_type	(imshow, surf)	Choose between image and 3D surface rendering for 2D data.	"imshow"
fixed_diff_colorbar	bool	If True, keep the difference plot color limits fixed across frames.	False

plot_wavefields_interactive_wpt(reference_all, p0_msgb_all, ts, domain: Domain, XY, wpt, dyadic_decomp, plot_type='imshow', fixed_diff_colorbar=False)

Interactive wavefield viewer that also shows MSWPT coefficients.

Parameters:

Name	Type	Description	Default
reference_all	ndarray	Wavefield sequence from the Reference solver (Nt, ...).	required
p0_msgb_all	ndarray	Wavefield sequence from the MSGB solver (Nt, ...).	required
ts	ndarray	Time points associated with frames.	required
domain	Domain		required
XY	ndarray	Meshgrid used for plotting.	required
wpt	MSWPT	Transform used to compute coefficients.	required
dyadic_decomp	DyadicDecomposition	Associated dyadic tiling used for labelling levels.	required
plot_type	(imshow, surf)		"imshow"
fixed_diff_colorbar	bool	If True, keep the difference plot color limits fixed across frames.	False

plot_GB_ellipse(x0, p0, a0, m0, figsize=(10, 10), scale_factor=1.0)

Visualize Gaussian beams as ellipses with arrows.

Parameters:

Name	Type	Description	Default
x0	ndarray	Beam positions, shape (b, d).	required
p0	ndarray	Propagation directions, shape (b, d).	required
a0	ndarray	Complex amplitudes, shape (b,).	required
m0	ndarray	Hessian matrices, shape (b, d, d).	required
figsize	Tuple[int, int]		(10, 10)
scale_factor	float	Scales ellipse size for legibility.	1.0

Returns:

Type	Description
(Figure, Axes)

GB_point_line_plot(X, Y, p0, p0_b, x0_b)

Overlay beam start points and directions on an initial pressure field.

Parameters:

Name	Type	Description	Default
X	ndarray	Meshgrid arrays for plotting the pressure field.	required
Y	ndarray	Meshgrid arrays for plotting the pressure field.	required
p0	ndarray	Initial pressure map (same shape as X/Y).	required
p0_b	ndarray	Beam directions, shape (b, d).	required
x0_b	ndarray	Beam start positions, shape (b, d).	required

plot_coeffs(data, dyadic_decomp, wpt)

Quick visualizer for MSWPT coefficients and Fourier-space centres.

Parameters:

Name	Type	Description	Default
data	ndarray	Spatial-domain input field.	required
dyadic_decomp	DyadicDecomposition		required
wpt	MSWPT	Transform instance; uses forward with spatial input.	required

Notes

Only tested for 2D layouts; higher dimensions will need refinements.

plot_mswpt_coeffs(ax: plt.Axes, coeffs_array: jnp.ndarray, dyadic_decomp: DyadicDecomposition, cutoff_freq: float | None = None, box_corners: np.ndarray | jnp.ndarray | tuple[int, int] | None = None, asymptote: bool = False, log_scale: bool = False)

Plot MSWPT coefficients with dyadic boxes, highlights, and asymptotes.

Parameters:

Name	Type	Description	Default
ax	Axes	Target axes.	required
coeffs_array	ndarray	Two-dimensional coefficient image.	required
dyadic_decomp	DyadicDecomposition	Decomposition providing box centres and sizes.	required
cutoff_freq	float	Highlight boxes whose centre norm is below this threshold.	None
box_corners	array - like	Pair of opposing global box indices defining the highlighted low-frequency region.	None
asymptote	bool	Whether to overlay diagonal asymptote guides.	False
log_scale	bool	Whether to display coefficient magnitudes with logarithmic normalization.	False

Returns:

Type	Description
AxesImage	Image handle returned by imshow.

plot_mswpt_coeffs_3d(ax: plt.Axes, coeffs_array: jnp.ndarray, dyadic_decomp: DyadicDecomposition, cutoff_freq: float | None = None, box_corners: np.ndarray | jnp.ndarray | tuple[int, int] | None = None, asymptote: bool = False)

Plot a 2D projection (e.g., MIP) of 3D MSWPT coefficients with projected dyadic boxes.

The input coeffs_array is expected to be 2D (already projected), and boxes are projected onto the (x, y) plane.

Returns:

Name	Type	Description
im	AxesImage	The imshow image handle (useful for colorbars).