Skip to content

beamax.decomposition

Dyadic decomposition utilities for partitioning Fourier space into level-wise boxes.

Key Object

  • DyadicDecomposition: stores box sizes, level counts, and center coordinates used by MSWPT and beam parameterization.

API Reference

beamax.decomposition

DyadicDecomposition(num_levels: int, N: Tuple[int, ...], num_boxes_levels: Tuple[int, ...], box_aspect_ratio: Tuple[int, ...], boxes_per_dim_levels: Optional[Tuple[Tuple[int, ...], ...]] = None)

Bases: Module

Multi-level dyadic tiling of Fourier space on a rectangular grid.

The decomposition partitions Fourier space into frequency boxes organised across num_levels scales. Each coarser level covers roughly half the resolution of the next, and the outer ring of boxes at each non-final level is retained while an inner cutout (covered by the finer level) is removed — this gives the "dyadic" structure.

Parameters:

Name Type Description Default
num_levels int

Number of scale levels. Must satisfy 1 <= num_levels <= floor(log2(N_ref / num_boxes_levels[0])) + 1 where N_ref = min(N).

 required
N Tuple[int, ...]

Grid shape of the underlying domain per spatial axis.

 required
num_boxes_levels Tuple[int, ...]

Boxes per side along the smallest spatial axis, one entry per level, from coarsest to finest. Typically a dyadic progression such as (4, 8) or (2, 4, 8).

 required
box_aspect_ratio Tuple[int, ...]

Integer aspect ratio of each box per axis, applied at every level.

 required
boxes_per_dim_levels Tuple[Tuple[int, ...], ...]

Explicit per-axis box count per level. If None (default), the per-axis count is derived from num_boxes_levels and the domain aspect ratio.

 None

Attributes:

Name Type Description
num_boxes_ndim (ndarray, shape(num_levels))

Number of boxes retained at each level (outer boxes minus inner cutout).
centres_ndim (ndarray, shape(total_num_boxes, ndim))

Integer Fourier-index centres of every box across all levels.
total_num_boxes int

Sum of num_boxes_ndim.

Examples:

>>> from beamax import DyadicDecomposition
>>> decomp = DyadicDecomposition(
...     num_levels=2,
...     N=(64, 64),
...     num_boxes_levels=(4, 8),
...     box_aspect_ratio=(1, 1),
... )
>>> int(decomp.total_num_boxes)
76
Notes
  • Raises ValueError during construction if the parameter combination cannot tile the grid cleanly (e.g. box lengths do not divide the smallest side).
  • num_levels, N, num_boxes_levels, box_aspect_ratio, and boxes_per_dim_levels are marked static for JAX pytree flattening.

Construct a dyadic decomposition and precompute box metadata.

Parameters:

Name Type Description Default
num_levels int

Number of dyadic levels.

 required
N Tuple[int, ...]

Grid shape per spatial axis.

 required
num_boxes_levels Tuple[int, ...]

Number of boxes along the smallest axis at each level.

 required
box_aspect_ratio Tuple[int, ...]

Per-axis integer box aspect ratio.

 required
boxes_per_dim_levels Tuple[Tuple[int, ...], ...]

Explicit per-axis box counts per level.

 None

Raises:

Type Description
ValueError

If the requested tiling is invalid.
ndim: int property

Number of spatial dimensions (len(N)).

num_boxes_ndim_cumsum: Int[Array, ' num_levels'] property

Cumulative num_boxes_ndim; useful for indexing per-level slices of centres_ndim.

fourier_meshgrid: Int[Array, '*N d'] property

Zero-centred Fourier-index meshgrid.

Returns:

Type Description
(ndarray, shape(*N, ndim))

Stacked meshgrid of axes arange(-Ni//2, Ni//2).
scaling: Float[Array, ' d'] property

Per-axis scaling factor that maps the domain aspect ratio onto the box aspect ratio.

Returns:

Type Description
(ndarray, shape(ndim))

(N // min(N)) / box_aspect_ratio.
box_lengths: Int[Array, ' num_levels'] property

Box side length (in Fourier indices, along the smallest axis) at each level.

Returns:

Type Description
(ndarray, shape(num_levels), int32)

Length at level L is min(N) // (num_boxes_levels[-1] * 2**(num_levels-1-L)); halves moving coarse → fine.
validate_params(num_levels: int, N: Tuple[int, ...], num_boxes_levels: Tuple[int, ...], box_aspect_ratio: Tuple[int, ...], boxes_per_dim_levels: Optional[Tuple[Tuple[int, ...], ...]] = None) -> None

Validate parameters for :class:DyadicDecomposition.

Parameters:

Name Type Description Default
num_levels int

Number of dyadic levels.

 required
N Tuple[int, ...]

Grid shape per spatial axis.

 required
num_boxes_levels Tuple[int, ...]

Number of boxes along the smallest axis at each level.

 required
box_aspect_ratio Tuple[int, ...]

Per-axis integer box aspect ratio.

 required
boxes_per_dim_levels Tuple[Tuple[int, ...], ...]

Explicit per-axis box counts per level.

 None

Raises:

Type Description
ValueError

If the parameters cannot define a valid dyadic tiling.