Real FFT (rfft)

Why this matters

jnp.fft.rfft(x) is the FFT specialised for real-valued input. Because the DFT of a real signal has Hermitian symmetry (the upper half of bins are conjugates of the lower half), rfft only computes and returns the non-redundant half: n // 2 + 1 complex bins instead of n.

This matters because:

• ~2× faster than full fft for the same real input.
• ~2× less memory for the frequency representation.
• Audio, images, and sensor readings are almost always real — rfft is the default in practice.

Key facts:

• Output shape: (n // 2 + 1,) for input of length n.
• Inverse: jnp.fft.irfft(X, n) — pass the original length n explicitly.
• Not interchangeable with fft for complex inputs.

Worked mini-example

import jax.numpy as jnp

x = jnp.array([1.0, 0.0, 0.0, 0.0])   # length 4
mags = jnp.abs(jnp.fft.rfft(x))
# output has length 4//2 + 1 = 3
# mags = [1.0, 1.0, 1.0]

Common pitfalls

• Output length is n // 2 + 1, not n — don’t assume same-length output.
• Inverse needs n — call irfft(X, n) with the original signal length; otherwise even-length signals reconstruct fine, odd-length signals may not.
• Complex input is undefined — rfft silently produces wrong results for complex input; use fft instead.

Problem

Implement rfft_magnitude(x) that returns the magnitude of each rfft bin.

• x: 1-D jax array (real).
• Returns: 1-D array (n // 2 + 1,) — magnitudes of the half-spectrum.