`beamax.utils`

General utilities shared across decomposition, transforms, and solvers.

Typical Use Cases

FFT/interpolation helpers for resampling and signal manipulation.
Batching/index utilities for coefficient and beam pipelines.
Miscellaneous helpers for device checks, synthetic data generation, and numeric convenience operations.

API Reference

`beamax.utils`

Utilities public API (explicit). Heavy things must be imported inside call-sites.

`Interpolator(grid_points: List[jnp.ndarray], values: jnp.ndarray, *, method: str = 'linear', boundary: str = '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 `∇c(x)`.
`hessian`	Hessian `∇²c(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`.

`get_devices()`

Inspect available JAX devices.

Returns:

Type	Description
`Tuple[bool, bool, bool]`	Presence flags for (CPU, GPU, TPU). Prints the detected devices.

`memory_estimate(dims: jnp.ndarray, dtype: jnp.dtype) -> str`

Estimate memory footprint for an array shape/dtype.

Parameters:

Name	Type	Description	Default
`dims`	`ndarray`	Shape tuple or array of dimensions.	required
`dtype`	`dtype`		required

Returns:

Type	Description
`str`	Estimated memory usage.

`memory_str(x: jnp.ndarray) -> str`

Memory usage of an existing array.

Parameters:

Name	Type	Description	Default
`x`	`ndarray`		required

Returns:

Type	Description
`str`	Human-friendly memory string.

`array_str(x: Union[jnp.ndarray, None]) -> Union[str, None]`

Short descriptor for an array (shape/dtype/memory).

Parameters:

Name	Type	Description	Default
`x`	`ndarray \| None`		required

Returns:

Type	Description
`str \| None`

`detect_root() -> Path`

Locate the repository root used by examples for output files.

Priority

1) BEAMAX_ROOT environment variable 2) Current working directory upward search 3) Package source location upward search 4) Current working directory

Returns:

Type	Description
`Path`

`example_plot_dir(example_file: str | os.PathLike[str]) -> Path`

Return the plot output directory for a public example script.

Examples mirror their first directory under examples/:

examples/forward/custom_lf_spectral_backend.py -> <root>/plots/forward
examples/rays/2d_ray_bending.py -> <root>/plots/rays

If example_file is outside the detected checkout's examples tree, the file's immediate parent directory name is used as the category.

Parameters:

Name	Type	Description	Default
`example_file`	`str or PathLike`	Usually `__file__` from an example script.	required

Returns:

Type	Description
`Path`	Existing output directory for plots from that example category.

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

`make_c_function_from_grid(c_map: jnp.ndarray, spacing: Optional[Tuple[float, ...]] = None, origin: Optional[Tuple[float, ...]] = None, *, method: str = 'linear', boundary: str = '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`	`('clamp', 'nearest', 'reflect', 'mirror', 'wrap', 'periodic')`	Boundary extension used by the interpolant. `"clamp"` is an alias for nearest-edge extension; `"periodic"` is an alias for `"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 a.e.
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.

`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	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	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` (0..Ni-1).	required

Returns:

Type	Description
jnp.ndarray, shape `box_shape_tuple`

`rel_l2(a, b)`

Compute relative L2 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`	`\|\|a - b\|\|_2 / (\|\|a\|\|_2 + 1e-30)`.

`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 global box index → dyadic level.

Parameters:

Name	Type	Description	Default
`dyadic_decomp`	`DyadicDecomposition`		required
`box_num`	`int`		required

Returns:

Type	Description
`ndarray`	Scalar (0-D) array with the level `ℓ` such that cumulative boxes up to `ℓ` exceed `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 `vmap`ped over `level`.	required

Returns:

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