Skip to content

← JAX Numerical Computing step 11 of 25

Medium Primitives

Real FFT (rfft)

Jump to the task

Keep the nonnegative-frequency coefficients

For a real signal, the negative-frequency Fourier coefficients are complex conjugates of the positive-frequency coefficients. They contain redundant information. jnp.fft.rfft returns only the nonnegative-frequency side: n // 2 + 1 complex coefficients for a real input of length n.

Taking their absolute values gives magnitudes. The result has fewer entries than a full FFT for most input lengths, but the coefficients themselves use the same default normalization. Do not double or rescale them in this task.

Worked example

import jax.numpy as jnp

x = jnp.array([0.0, 1.0, 0.0, 0.0, 0.0], dtype=jnp.float32)
coefficients = jnp.fft.rfft(x)
magnitudes = jnp.abs(coefficients)
# approximately [1.0, 1.0, 1.0], length 5 // 2 + 1 = 3

restored = jnp.fft.irfft(coefficients, n=x.size)
# approximately [0.0, 1.0, 0.0, 0.0, 0.0]

This shifted impulse has unit magnitude in each bin, but its coefficients also carry phase. The inverse uses the complex coefficients, not the magnitudes. Feeding the magnitudes to irfft would lose the shift.

Problem

Implement rfft_magnitude(x) using jnp.fft.rfft(x) followed by jnp.abs.

  • x is a nonempty one-dimensional real float32 JAX array of length n.
  • Return a real nonnegative array of shape (n // 2 + 1,).
  • Keep the default normalization and bin order.
  • Return magnitudes, not squared magnitudes or complex coefficients.

The first coefficient is sum(x), so the first returned magnitude is abs(sum(x)). Do not replace it with the signed sum.

Odd and even lengths

Bin zero is DC. For even n, the final bin is the Nyquist frequency, exactly half the sampling rate, and its coefficient is real. For odd n, there is no Nyquist bin; the final coefficient can have an imaginary part. Use jnp.fft.rfftfreq(n, d=sample_spacing) to label the frequencies.

The output length does not uniquely identify the input length. Both a four-sample and a five-sample signal produce three coefficients. Without n, irfft assumes the even length 2 * (number_of_coefficients - 1). Pass the original length explicitly to reconstruct odd-length signals and to handle a single-sample signal correctly.

Input type and performance

JAX's rfft requires real input and rejects complex arrays. Use fft if the signal itself is complex; taking only its real part would change the problem.

The stored spectrum has n // 2 + 1 coefficients instead of n, which is about half as many for large n. Runtime savings depend on the size and backend; this is not a guarantee of a two-times speedup. Magnitudes alone also omit phase, so they cannot generally reconstruct the original signal.