beamax.utils

Shared array, FFT, interpolation, indexing, batching, device, and memory-planning utilities.

Arrays

Array shape and resampling utilities.

Centered pad / crop, Fourier-based resampling, nearest-neighbour resampling, and a wrap-around centered-box extractor. These helpers operate on raw arrays and have no MSWPT/MSGB-specific dependencies.

interpolate_nearest(array: jnp.ndarray, new_shape: Tuple) -> jnp.ndarray

Nearest-neighbour resampling to new_shape.

Parameters:

Name	Type	Description	Default
array	ndarray		required
new_shape	Tuple[int, ...]		required

Returns:

Type	Description
ndarray	Resampled array.

pad_array(array: jnp.ndarray, desired_size: Tuple[int, ...], mode: str = 'constant') -> jnp.ndarray

Centered pad per axis to reach desired_size.

Parameters:

Name	Type	Description	Default
array	ndarray		required
desired_size	Tuple[int, ...]		required
mode	('constant', 'edge')		"constant"

Returns:

Type	Description
ndarray

pad_zero(array: jnp.ndarray, desired_size: Tuple) -> jnp.ndarray

Zero-pad to centered target size.

Parameters:

Name	Type	Description	Default
array	ndarray		required
desired_size	Tuple[int, ...]		required

Returns:

Type	Description
ndarray

pad_edge(array: jnp.ndarray, desired_size: Tuple) -> jnp.ndarray

Edge-pad (replicate border) to centered target size.

Parameters:

Name	Type	Description	Default
array	ndarray		required
desired_size	Tuple[int, ...]		required

Returns:

Type	Description
ndarray

crop_centered(array: jnp.ndarray, desired_size: Tuple[int, ...]) -> jnp.ndarray

Centered crop per axis. No-op if any target dim exceeds current.

Parameters:

Name	Type	Description	Default
array	ndarray		required
desired_size	Tuple[int, ...]		required

Returns:

Type	Description
ndarray

interpolate_fourier(array: jnp.ndarray, desired_size: Tuple[int, ...], input_type: str, output_type: str) -> jnp.ndarray

Unitary FFT-based resampling (pad when upsampling, crop when downsampling).

Parameters:

Name	Type	Description	Default
array	ndarray		required
desired_size	Tuple[int, ...]		required
input_type	('spatial', 'fourier')		"spatial"
output_type	('spatial', 'fourier')		"spatial"

Returns:

Type	Description
ndarray	Resampled array in output_type domain.

Notes

Assumes periodic boundaries. Uses pad_array then crop_centered in frequency.

extract_centered_box(arr, box_shape_tuple, center)

Wrap-around extraction of a centered N-D box (JIT/static-friendly).

Parameters:

Name	Type	Description	Default
arr	jnp.ndarray, shape S = (N1, ..., Nd)		required
box_shape_tuple	Tuple[int, ...]	Static Python tuple of ints (box sizes per axis).	required
center	(ndarray, shape(d))	Center index in the same index space as arr, with entries \(0, \ldots, N_i - 1\).	required

Returns:

Type	Description
jnp.ndarray, shape `box_shape_tuple`

rel_l2(a, b)

Compute the relative \(\ell^2\) error between two arrays.

Parameters:

Name	Type	Description	Default
a	array - like	Reference array.	required
b	array - like	Comparison array.	required

Returns:

Type	Description
float	\(\lVert a - b\rVert_2 / (\lVert a\rVert_2 + 10^{-30})\).

Fourier transforms

unitary_fft(arr: jnp.ndarray) -> jnp.ndarray

Unitary N-D FFT with centred zero-frequency component.

Parameters:

Name	Type	Description	Default
arr	ndarray	Real or complex array in the spatial domain.	required

Returns:

Type	Description
ndarray	Fourier transform with norm="ortho" and fftshift applied.

unitary_ifft(arr: jnp.ndarray) -> jnp.ndarray

Unitary inverse FFT matching :func:unitary_fft.

Parameters:

Name	Type	Description	Default
arr	ndarray	Fourier-domain array (already shifted).	required

Returns:

Type	Description
ndarray	Spatial-domain array with unitary scaling.

convert_space(array: jnp.ndarray, input_space: str, target_space: str) -> jnp.ndarray

Convert an array between spatial and Fourier domains.

Parameters:

Name	Type	Description	Default
array	ndarray		required
input_space	(spatial, fourier)	Declares the domain of array.	"spatial"
target_space	(spatial, fourier)	Desired output domain.	"spatial"

Returns:

Type	Description
ndarray	Array in the requested domain.

Raises:

Type	Description
ValueError	If an unsupported conversion is requested.

Interpolation

Interpolator(grid_points: List[jnp.ndarray], values: jnp.ndarray, *, method: str = 'linear', boundary: BoundarySpec = 'clamp', smooth_sigma: Optional[float | Sequence[float]] = None)

Thin wrapper around make_c_function_from_grid using axis vectors.

Parameters:

Name	Type	Description	Default
grid_points	List[ndarray]	1D axis arrays (length d), each strictly increasing.	required
values	ndarray	Grid values shaped to match grid_points.	required

Methods:

Name	Description
__call__	Evaluate interpolant at x (shape (..., d)).
grad	Gradient \(\nabla c(\mathbf{x})\).
hessian	Hessian \(\nabla^2 c(\mathbf{x})\).

Construct an interpolator from axis vectors and grid values.

Parameters:

Name	Type	Description	Default
grid_points	List[ndarray]	One strictly increasing one-dimensional coordinate vector per axis.	required
values	ndarray	Grid values with dimensionality matching grid_points.	required
method	str	Passed through to :func:make_c_function_from_grid.	'linear'
boundary	str	Passed through to :func:make_c_function_from_grid.	'linear'
smooth_sigma	str	Passed through to :func:make_c_function_from_grid.	'linear'

Raises:

Type	Description
ValueError	If axis count and value dimensionality disagree, or if any axis is not one-dimensional with at least two points.

grad(x: jnp.ndarray) -> jnp.ndarray

Evaluate the gradient of the interpolant.

Parameters:

Name	Type	Description	Default
x	(ndarray, shape(d))	Physical query coordinate.	required

Returns:

Type	Description
(ndarray, shape(d))	Gradient at x.

hessian(x: jnp.ndarray) -> jnp.ndarray

Evaluate the Hessian of the interpolant.

Parameters:

Name	Type	Description	Default
x	(ndarray, shape(d))	Physical query coordinate.	required

Returns:

Type	Description
(ndarray, shape(d, d))	Hessian at x.

make_c_function_from_grid(c_map: jnp.ndarray, spacing: Optional[Tuple[float, ...]] = None, origin: Optional[Tuple[float, ...]] = None, *, method: str = 'linear', boundary: BoundarySpec = 'clamp', smooth_sigma: Optional[float | Sequence[float]] = None)

Build a JAX-differentiable interpolant over a rectilinear grid.

Parameters:

Name	Type	Description	Default
c_map	(ndarray, shape(*N))	Grid values.	required
spacing	Tuple[float, ...] | None	Per-axis spacing. Defaults to 1.0.	None
origin	Tuple[float, ...] | None	Per-axis origin. Defaults to 0.0.	None
method	('linear', 'bspline3')	Interpolation method. "linear" preserves the original n-linear behaviour. "bspline3" builds a tensor-product cubic B-spline interpolant whose value, gradient, and Hessian are continuous away from boundary handling.	"linear"
boundary	str | sequence of str	Boundary extension for every axis, or one mode per axis. Supported modes are "clamp", "nearest", "reflect", "mirror", "wrap", and "periodic". "clamp" aliases nearest-edge extension and "periodic" aliases "wrap".	"clamp"
smooth_sigma	float | sequence of float | None	Optional Gaussian pre-smoothing of the grid values, with standard deviation measured in grid cells. Smoothing is applied once when the interpolant is constructed.	None

Returns:

Type	Description
Callable[[ndarray], ndarray]	Function c_fun(x) with x shape (..., d) in physical units.

Notes

• The default method="linear" path is piecewise-linear and differentiable almost everywhere.
• The method="bspline3" path precomputes cubic spline coefficients using SciPy, then evaluates the spline using JAX operations so jax.grad and jax.hessian can differentiate with respect to query coordinates.

Coefficient indexing

Coefficient indexing utilities used across the MSWPT transform and solvers.

These helpers translate between flat coefficient indices and (level, multi-index) coordinates, and batch leading-axis-batched beam parameter arrays for scan/vmap aggregation.

batch_data(*args, batch_size, zero_padded_args=())

Batch leading beam axis into (num_batches, batch_size, ...).

Parameters:

Name	Type	Description	Default
*args	Tuple[ndarray, ...]	Arrays with leading dimension b.	()
batch_size	int		required
zero_padded_args	Tuple[int, ...]	Indices into args that should be zero-padded in the last batch; others repeat their last entry.	()

Returns:

Type	Description
Tuple[ndarray, ...]	Batched arrays (with padding if needed).

find_level(dyadic_decomp: DyadicDecomposition, box_num: int | Int[Array, '...']) -> Int[Array, '']

Map a global box index to its dyadic level.

Parameters:

Name	Type	Description	Default
dyadic_decomp	DyadicDecomposition		required
box_num	int		required

Returns:

Type	Description
ndarray	Scalar (0-D) array containing the level \(\ell\) such that the cumulative box count through level \(\ell\) exceeds box_num. JAX traces it as an integer scalar inside JIT.

find_tensor_and_multiindex(flat_indices: jnp.ndarray, shapes: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]

Decode flat indices over concatenated tensors.

Parameters:

Name	Type	Description	Default
flat_indices	(ndarray, shape(m))		required
shapes	(ndarray, shape(L, k))	Shapes of each tensor (per level).	required

Returns:

Name	Type	Description
array_indices	(ndarray, shape(m))	Which tensor each flat index belongs to.
multidimensional_indices	(ndarray, shape(m, k))	Unravelled indices within that tensor.

compute_coeff_shapes(dyadic_decomp: DyadicDecomposition, redundancy: int, level) -> jnp.ndarray

Per-level coefficient tensor shapes.

Parameters:

Name	Type	Description	Default
dyadic_decomp	DyadicDecomposition		required
redundancy	int	1 (basis) or 2 (frame).	required
level	int	Vectorized: function is vmapped over level.	required

Returns:

Type	Description
(ndarray, shape(k))	(num_boxes_level, *support_shape).

Memory planning

Device capability and memory planning for MSGB solves.

DeviceCapabilities(backend: str, device_kind: str, device_count: int, memory_limit_bytes: Optional[int], x64_enabled: bool, pallas_triton_available: bool, pallas_tpu_available: bool, jax_version: str) dataclass

Snapshot of the local JAX environment relevant to MSGB planning.

MemoryEstimate(operation: Operation, policy: str, config_summary: str, items: Tuple[Tuple[str, int], ...], lower_bound_bytes: int, safety_factor: float, budget_bytes: Optional[int], notes: Tuple[str, ...] = tuple()) dataclass

Itemized lower bound, planning margin, and budget verdict.

planning_bytes: int property

Return \(\lceil sL\rceil\) for safety factor \(s\) and lower bound \(L\).

device_capabilities() -> DeviceCapabilities

Inspect JAX and read the default device's memory limit when available.

estimate_msgb_memory(operation: Operation, *, grid_shape: Sequence[int], top_n: int, batch_size: int, num_times: int = 0, num_detectors: int = 0, data_shape: Optional[Sequence[int]] = None, selection: Selection = 'streaming_top_n', forward_kernel: ForwardKernel = 'xla', redundancy: int = 2, total_coeffs: Optional[int] = None, real_dtype: Any | None = None, safety_factor: float = 2.0, budget_bytes: Optional[int] = None, use_device_budget: bool = True) -> MemoryEstimate

Estimate device memory needs for one MSGB operation.

Pass total_coeffs for an exact materialized-vector size. real_dtype defaults to float64 when JAX x64 is enabled and float32 otherwise. The planning estimate is \(sL\), where \(L\) is the allocation lower bound and \(s\) is the explicit safety_factor. The fused homogeneous 3-D Pallas kernel requires float32. This estimate is not a measured allocator peak.