Skip to content

← JAX Numerical Computing step 12 of 25

Hard Primitives

2-D FFT DC Component

Jump to the task

Find the zero-frequency part of a matrix

A two-dimensional Fourier transform describes variation along both axes of a matrix. A frequency index selects a pattern along each axis. The [0, 0] coefficient measures the component that is constant along both: the zero-frequency, or DC, component.

With JAX's default forward normalization, every input contributes directly to this coefficient, so it equals the sum of all entries. This exercise returns its magnitude, which is abs(sum(x)) for real input.

Worked example

import jax.numpy as jnp

x = jnp.array([[-1.0, -2.0],
               [0.0, 1.0]], dtype=jnp.float32)
spectrum = jnp.fft.fft2(x)
dc_coefficient = spectrum[0, 0]    # -2+0j
magnitude = jnp.abs(dc_coefficient) # 2.0

The signed sum is -2, but the requested magnitude is 2. Summing absolute values would give 4, which is a different quantity. The matrix mean is -0.5, obtained by dividing the signed DC coefficient by the four entries.

Problem

Implement fft_2d_dc(x):

  1. Compute jnp.fft.fft2(x) with its default options.
  2. Select coefficient [0, 0].
  3. Return its magnitude with jnp.abs.

x is a two-dimensional real float32 JAX array with both dimensions nonempty. Return a real nonnegative scalar, not the full spectrum.

The identity abs(sum(x)) is a useful check. Use the FFT in your implementation to practice locating the coefficient.

Shape, axes and normalization

For this two-dimensional input, fft2 transforms both axes and returns a complex array with the same shape. More generally, its default axes are the last two; a leading batch axis is not transformed automatically. The transform is separable, so applying a one-dimensional FFT along each of the two axes gives the same mathematical result, up to numerical error.

The default forward transform leaves the sum unscaled. Passing norm="forward" would instead make the DC coefficient equal to the mean; norm="ortho" would divide the sum by the square root of the element count. Use the default in this task.

Avoid mixing up representations

Do not apply fftshift before reading [0, 0]: shifting moves the DC bin toward the center. Take the magnitude after selecting the coefficient; abs(sum(x)) is generally not sum(abs(x)).

For real inputs, rfft2 can store fewer coefficients by using conjugate symmetry. Its default output shape is (m, n // 2 + 1) for input (m, n). That is useful for other tasks, but this exercise introduces fft2 and its unshifted DC coefficient. A zero DC value means the entries sum to zero; it does not mean the entire input is zero.