`beamax.transforms`

Multiscale wave-packet transforms and frame construction.

`beamax.transforms`

`MSWPT(dyadic_decomp, redundancy, windowing)`

Bases: Module

Multiscale Wave-Packet Transform.

Parameters:

Name	Type	Description	Default
`dyadic_decomp`	`DyadicDecomposition`	Frequency tiling (centres, per-level box lengths).	required
`redundancy`	`int`	Supported translation-lattice redundancy (1 or 2). Static under JIT.	required
`windowing`	`(rectangular, rectangular_mirror, none)`	Windowing for tile filters. `"none"` is available for explicit synthesis/beam-comparison atoms only; the fast packed analysis rejects it because an untruncated Gaussian is not supported on the packed local patch. Static under JIT.	`"rectangular"`

Attributes:

Name	Type	Description
`dyadic_decomp`	`DyadicDecomposition`
`redundancy`	`int`
`windowing`	`str`
`complex_dtype`	`dtype`	`complex64` unless JAX x64 is enabled, in which case `complex128`.
`sum_gsquare`	`(ndarray, shape(*N))`	\(\sum_b g_b^2\), precomputed.
`boxes_cumsum`	`Tuple[int, ...]`	Cumulative number of boxes per level (static).
`coeff_shapes`	`Tuple[Tuple[int, ...], ...]`	Shape per level: (n_boxes_level, *support_shape).
`coeffs_cumsum`	`Tuple[int, ...]`	Flat coefficient offsets per level (static).
`total_coeffs`	`int`	Total number of flat coefficients (static).
`gfilts_packed`	`List[ndarray]`	Per-level tile `g` packed to minimal support boxes.
`_support_shapes`	`List[Tuple[int, ...]]`	Per-level support shapes (static).
`_box_shapes`	`List[Tuple[int, ...]]`	Per-level "box lengths" in each axis (static).
`half_mask`	`ndarray`	Read-only mask selecting one spatial-frequency representative from each conjugate pair.

Notes

Constructor performs shape bookkeeping to keep runtime kernels static-JIT friendly.
All heavy transforms are pure JAX functions; no side effects.

Build a transform instance with precomputed static metadata.

Parameters:

Name	Type	Description	Default
`dyadic_decomp`	`DyadicDecomposition`	Provides centres, per-level box lengths, and meshgrid.	required
`redundancy`	`int`	1 or 2. Controls per-level support sizes and total coeff count.	required
`windowing`	`str`	Window type for Gaussian tiles.	required

Notes

Precomputes cumulative box/coeff offsets and per-level support shapes.
Chooses complex dtype from global JAX precision flag.
Packs representative per-level g filters for later slicing/rolls.

`half_mask: jnp.ndarray` `property`

Mask selecting one representative from each conjugate frame pair.

`forward(data: Num[Array, '*N'], input_type: str) -> Num[Array, ' total_coeffs']`

Forward MSWPT.

Parameters:

Name	Type	Description	Default
`data`	`(ndarray, shape(*N), real or complex)`	Input field in spatial or Fourier domain.	required
`input_type`	`(spatial, fourier)`	Declares the domain of `data`.	`"spatial"`

Returns:

Type	Description
`(ndarray, shape(total_coeffs), complex)`	Flat coefficient vector.

Notes

Converts to Fourier with unitary_fft if needed.
Divides by \(\sum_b g_b^2\) to apply the pointwise canonical-dual analysis filters for this painless frame construction.
JIT-compiled while preserving the reusable transform state. The input data buffer may still be donated for memory efficiency.

`inverse(coeffs: Num[Array, ' total_coeffs'], output_type: str) -> Num[Array, '*N']`

Fast synthesis MSWPT; the exact inverse of :meth:forward for its supported window configurations.

Parameters:

Name	Type	Description	Default
`coeffs`	`(ndarray, shape(total_coeffs), complex)`	Flat coefficient vector produced by :meth:`forward`.	required
`output_type`	`(spatial, fourier)`	Domain of the returned array.	`"spatial"`

Returns:

Type	Description
`(ndarray, shape(*N), complex)`	Reconstructed field in the requested domain.

Notes

With windowing="none" this remains a synthesis-only map for the Gaussian restricted to the transform's nominal packed support (the same packed atom as windowing="rectangular"). It is not the global unwindowed atom returned by :func:compute_frames.

The synthesis mirrors the analysis steps in :meth:forward:

\[ \begin{aligned} P_b &= \operatorname{extract}\left(F / \sum_j g_j^2, x_b\right), & c_b &= \operatorname{IFFT}\left( \operatorname{roll}(g_b P_b, +r_b) \right), \\ T_b &= \operatorname{FFT}(c_b), & F &\leftarrow F + g_b\operatorname{roll}(T_b, -r_b). \end{aligned} \]

The final update is a periodic scatter-add at the true box centre \(x_b\).

`compute_windowed_gaussian(centre: Int[Array, ' d'], meshgrid: Int[Array, 'N d'], box_length, box_aspect_ratio: Union[Int[Array, ' d'], Tuple[int, ...]], domain_length: Union[Int[Array, ' d'], Tuple[int, ...]], redundancy: int, windowing: str) -> Float[Array, 'N']`

Windowed N-D Gaussian in Fourier index space.

Parameters:

Name	Type	Description	Default
`centre`	`(ndarray, shape(d))`	Tile centre in Fourier index units.	required
`meshgrid`	`(ndarray, shape(*N, d))`	Integer-centred Fourier grid from `DyadicDecomposition.fourier_meshgrid`.	required
`box_length`	`int`	Smallest-axis tile length for this level.	required
`box_aspect_ratio`	`(ndarray, shape(d))`	Per-axis aspect multipliers. Every component satisfies \(a_i \ge 1\), with at least one satisfying \(a_i = 1\).	required
`domain_length`	`(array - like, shape(d))`	Per-axis grid lengths. These define the periodic wrap independently on each Fourier axis.	required
`redundancy`	`int`	Supported translation-lattice redundancy (1 or 2). Redundancy 2 does not by itself make the Gaussian window frame tight.	required
`windowing`	`(none, rectangular, rectangular_mirror)`	Windowing function applied to the Gaussian.	`"none"`

Returns:

Type	Description
`(ndarray, shape(*N))`	Windowed Gaussian weights. dtype = float32/float64.

Notes

Pure JAX, JIT-safe. Periodisation handled by modulo arithmetic.

`single_filter_idx(centre_idx, meshgrid: Int[Array, 'N d'], dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Float[Array, 'N']`

Filter for a single tile (by global box index).

Parameters:

Name	Type	Description	Default
`centre_idx`	`int`	Global box index in `dyadic_decomp.centres_ndim`.	required
`meshgrid`	`(ndarray, shape(*N, d))`	Fourier meshgrid.	required
`dyadic_decomp`	`DyadicDecomposition`	Dyadic parameters providing centres and per-level box lengths.	required
`redundancy`	`int`	Supported translation-lattice redundancy (1 or 2).	required
`windowing`	`str`	Windowing function (see `compute_windowed_gaussian`).	`'rectangular'`

Returns:

Type	Description
`(ndarray, shape(*N))`	Filter values.

`single_filter_coord(centre: Int[Array, ' d'], level: int, meshgrid: Int[Array, 'N d'], dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Float[Array, 'N']`

Filter for a single tile (by (centre, level) pair).

Parameters:

Name	Type	Description	Default
`centre`	`(ndarray, shape(d))`	Tile centre in Fourier indices.	required
`level`	`int`	Dyadic level \(\ell \in \{0, \ldots, L-1\}\).	required
`meshgrid`	`(ndarray, shape(*N, d))`	Fourier meshgrid.	required
`dyadic_decomp`	`DyadicDecomposition`	Decomposition parameters.	required
`redundancy`	`int`	1 or 2.	required
`windowing`	`str`	Windowing function.	`'rectangular'`

Returns:

Type	Description
`(ndarray, shape(*N))`	Filter values.

`compute_frame_phase(dyadic_decomp: DyadicDecomposition, boxidx: Union[int, Int[Array, ...]], k: Int[Array, '... d'], redundancy: int) -> Num[Array, ...]`

Unit-modulus phase induced by local-patch parity recentering.

The fast transform parity-rolls each local Fourier patch before applying its FFT. This factor converts the clean global modulation \(\exp\left(-2\pi i\sum_{j=1}^{d}m_jk_j/S_j\right)\) into the exact coefficient convention used by the implemented transform, including rectangular supports and both supported redundancies.

`compute_gh_filters(dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Tuple[Float[Array, B * N], Float[Array, B * N]]`

Compute the \(g_b\) tiles and their duals \(h_b = g_b / \sum_{b'} g_{b'}^2\).

Parameters:

Name	Type	Description	Default
`dyadic_decomp`	`DyadicDecomposition`	Frequency tiling.	required
`redundancy`	`int`	1 or 2.	required
`windowing`	`str`	Windowing function.	`'rectangular'`

Returns:

Type	Description
`(ndarray, ndarray)`	`(gfilt, hfilt)` each with shape (num_boxes, *N).

`compute_frames(dyadic_decomp: DyadicDecomposition, boxidx: int, k: Int[Array, ' d'], fourier_space: Num[Array, 'N d'], redundancy: int, windowing: str = 'rectangular') -> Num[Array, 'N']`

Frame atom for box boxidx with plane-wave modulation.

Parameters:

Name	Type	Description	Default
`dyadic_decomp`	`DyadicDecomposition`	Frequency tiling.	required
`boxidx`	`int`	Box index.	required
`k`	`(ndarray, shape(d))`	Local spatial-translation multi-index.	required
`fourier_space`	`(ndarray, shape(*N, d))`	Integer Fourier-index coordinates.	required
`redundancy`	`int`	1 or 2.	required
`windowing`	`str`	Windowing function.	`'rectangular'`

Returns:

Type	Description
`(ndarray, shape(*N))`	Complex atom, dtype = complex64/complex128.

beamax.transforms

beamax.transforms

MSWPT(dyadic_decomp, redundancy, windowing)

half_mask: jnp.ndarray property

forward(data: Num[Array, '*N'], input_type: str) -> Num[Array, ' total_coeffs']

inverse(coeffs: Num[Array, ' total_coeffs'], output_type: str) -> Num[Array, '*N']

compute_windowed_gaussian(centre: Int[Array, ' d'], meshgrid: Int[Array, '*N d'], box_length, box_aspect_ratio: Union[Int[Array, ' d'], Tuple[int, ...]], domain_length: Union[Int[Array, ' d'], Tuple[int, ...]], redundancy: int, windowing: str) -> Float[Array, '*N']

single_filter_idx(centre_idx, meshgrid: Int[Array, '*N d'], dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Float[Array, '*N']

single_filter_coord(centre: Int[Array, ' d'], level: int, meshgrid: Int[Array, '*N d'], dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Float[Array, '*N']

compute_frame_phase(dyadic_decomp: DyadicDecomposition, boxidx: Union[int, Int[Array, ...]], k: Int[Array, '... d'], redundancy: int) -> Num[Array, ...]

compute_gh_filters(dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Tuple[Float[Array, B * N], Float[Array, B * N]]

compute_frames(dyadic_decomp: DyadicDecomposition, boxidx: int, k: Int[Array, ' d'], fourier_space: Num[Array, '*N d'], redundancy: int, windowing: str = 'rectangular') -> Num[Array, '*N']

`beamax.transforms`

`beamax.transforms`

`MSWPT(dyadic_decomp, redundancy, windowing)`

`half_mask: jnp.ndarray` `property`

`forward(data: Num[Array, '*N'], input_type: str) -> Num[Array, ' total_coeffs']`

`inverse(coeffs: Num[Array, ' total_coeffs'], output_type: str) -> Num[Array, '*N']`

`compute_windowed_gaussian(centre: Int[Array, ' d'], meshgrid: Int[Array, 'N d'], box_length, box_aspect_ratio: Union[Int[Array, ' d'], Tuple[int, ...]], domain_length: Union[Int[Array, ' d'], Tuple[int, ...]], redundancy: int, windowing: str) -> Float[Array, 'N']`

`single_filter_idx(centre_idx, meshgrid: Int[Array, 'N d'], dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Float[Array, 'N']`

`single_filter_coord(centre: Int[Array, ' d'], level: int, meshgrid: Int[Array, 'N d'], dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Float[Array, 'N']`

`compute_frame_phase(dyadic_decomp: DyadicDecomposition, boxidx: Union[int, Int[Array, ...]], k: Int[Array, '... d'], redundancy: int) -> Num[Array, ...]`

`compute_gh_filters(dyadic_decomp: DyadicDecomposition, redundancy: int, windowing: str = 'rectangular') -> Tuple[Float[Array, B * N], Float[Array, B * N]]`

`compute_frames(dyadic_decomp: DyadicDecomposition, boxidx: int, k: Int[Array, ' d'], fourier_space: Num[Array, 'N d'], redundancy: int, windowing: str = 'rectangular') -> Num[Array, 'N']`