Skip to content

← JAX Numerical Computing step 10 of 25

Medium Primitives

1-D FFT Magnitude

Jump to the task

Measure a signal's frequency components

A discrete Fourier transform represents a sequence using complex frequency coefficients. jnp.fft.fft computes that transform with an FFT algorithm. Taking the absolute value of each coefficient gives its magnitude.

The magnitude tells you the size of a coefficient, but discards its phase. It is not a power spectrum: power calculations use squared magnitudes, with normalization chosen for the application.

Worked example

import jax.numpy as jnp

x = jnp.array([0.0, 1.0, 0.0, -1.0], dtype=jnp.float32)
coefficients = jnp.fft.fft(x)
# [0+0j, 0-2j, 0+0j, 0+2j]
magnitudes = jnp.abs(coefficients)
# [0.0, 2.0, 0.0, 2.0]

This sequence samples one cycle of a sine wave. Its positive and negative frequency coefficients have equal magnitude. The input amplitude is one, but those magnitudes are two: the default forward FFT is not normalized by the sequence length, and a real sine wave has two frequency components. Do not interpret raw bin magnitude as the signal's amplitude without accounting for the normalization and how the spectrum is represented.

Problem

Implement fft_magnitude(x) using jnp.fft.fft(x) and jnp.abs.

  • x is a nonempty one-dimensional real float32 JAX array of length n.
  • Return a real, nonnegative array of length n.
  • Use the default FFT normalization and original bin order.
  • Do not square the magnitudes, divide by n, shift the bins, or return only the nonnegative-frequency half.

Read the first bin correctly

With the default normalization, coefficient zero is sum(x). Its magnitude is abs(sum(x)), which may differ from both the signed sum and the sum of absolute values. A constant sequence [-2, -2, -2, -2] has coefficient zero equal to -8 and first-bin magnitude equal to 8; its other coefficients vanish. The mean is the signed zero-frequency coefficient divided by n.

Bin order and symmetry

Bin zero is the zero-frequency, or DC, component. Positive-frequency bins come next, followed by negative-frequency bins at the end. For even n, the middle bin is the Nyquist frequency, at half the sampling rate. Use jnp.fft.fftfreq(n, d=sample_spacing) when you need frequency labels.

For real input, coefficient n-k is the complex conjugate of coefficient k, so their magnitudes match. rfft uses this symmetry to return only n // 2 + 1 coefficients, but this exercise requests all n bins. Its speed advantage depends on the input size and backend, not a fixed ratio.