08-signal-processing

Analyze frequencies with FFT — decompose signals and filter noise. This example builds a signal from two sine waves plus noise, identifies the component frequencies, and filters the noise in the frequency domain.

dune exec nx/examples/08-signal-processing/main.exe

What You'll Learn

• Constructing synthetic signals from sine waves and noise
• Transforming to the frequency domain with rfft
• Mapping frequency bins to Hz with rfftfreq
• Identifying dominant frequency components by magnitude
• Filtering noise by zeroing small-magnitude frequency bins
• Reconstructing a clean signal with irfft

Key Functions

Output Walkthrough

Signal construction

A 256-sample signal at 256 Hz composed of two sine waves plus noise:

Signal: 256 samples at 256 Hz
Components: 5 Hz (amplitude 1.0) + 20 Hz (amplitude 0.5) + noise

Frequency analysis

After rfft, compute magnitudes scaled by 2/N to get amplitudes:

Dominant frequencies:
  5.0 Hz  (magnitude 1.002)
  20.0 Hz  (magnitude 0.501)

The FFT correctly recovers both sine components. Noise spreads across many bins with small magnitudes.

Noise filtering

Zero all frequency bins below a threshold, then reconstruct with irfft:

After filtering (threshold=0.2):
  Original first 8:  [1.29, 1.16, 0.81, 0.44, ...]
  Filtered first 8:  [1.00, 1.16, 0.79, 0.36, ...]

The filtered signal retains the two sine waves while removing noise.

Frequency bins

For an N-sample real signal, rfft produces N/2 + 1 complex bins from 0 Hz to the Nyquist frequency (sample_rate / 2):

Total bins: 129 (for 256-sample signal)

Try It

1. Add a third sine wave at 40 Hz and verify it appears in the dominant frequencies.
2. Raise the filter threshold to 0.5 and observe how the 20 Hz component gets removed (its amplitude is only 0.5).
3. Double the sample rate to 512 Hz and check how the frequency resolution changes.

Next Steps

Continue to 09-image-processing to apply convolution and pooling to 2D image data.

(** Analyze frequencies with FFT — decompose signals and filter noise.

    Build a signal from two sine waves plus noise. Use the real FFT to identify
    component frequencies, then filter the noise and reconstruct a clean signal.
*)

open Nx
open Nx.Infix

let () =
  (* Signal parameters. *)
  let n = 256 in
  let sample_rate = 256.0 in
  let dt = 1.0 /. sample_rate in

  (* Time axis: n samples at the given rate. *)
  let t = linspace float64 0.0 (Float.of_int n *. dt) n ~endpoint:false in

  (* Build signal: 5 Hz sine + 20 Hz sine + noise. *)
  let noise = randn Float64 [| n |] *$ 0.3 in
  let pi2 = 2.0 *. Float.pi in
  let signal_5hz = sin (t *$ (pi2 *. 5.0)) in
  let signal_20hz = sin (t *$ (pi2 *. 20.0)) *$ 0.5 in
  let signal = signal_5hz + signal_20hz + noise in

  Printf.printf "Signal: %d samples at %.0f Hz\n" n sample_rate;
  Printf.printf
    "Components: 5 Hz (amplitude 1.0) + 20 Hz (amplitude 0.5) + noise\n\n";

  (* Show first 8 samples. *)
  Printf.printf "First 8 samples: %s\n\n"
    (data_to_string (slice [ R (0, 8) ] signal));

  (* --- Real FFT: transform to frequency domain --- *)
  let spectrum = rfft signal in
  let freqs = rfftfreq ~d:dt n in

  (* Magnitudes (scaled by 2/N for single-sided spectrum). Extract real and
     imaginary parts to compute |z| = sqrt(re² + im²). *)
  let spectrum_arr = to_array spectrum in
  let re =
    create float64 (shape spectrum)
      (Array.map (fun c -> c.Complex.re) spectrum_arr)
  in
  let im =
    create float64 (shape spectrum)
      (Array.map (fun c -> c.Complex.im) spectrum_arr)
  in
  let magnitudes = sqrt ((re * re) + (im * im)) *$ (2.0 /. Float.of_int n) in

  (* Find the dominant frequencies (magnitude > 0.3). *)
  Printf.printf "Dominant frequencies:\n";
  let n_freqs = (shape magnitudes).(0) in
  for i = 0 to pred n_freqs do
    let mag = item [ i ] magnitudes in
    if Stdlib.( > ) mag 0.3 then
      Printf.printf "  %.1f Hz  (magnitude %.3f)\n" (item [ i ] freqs) mag
  done;
  print_newline ();

  (* --- Filter: zero out small frequency components --- *)
  let threshold = 0.2 in
  let mag_arr = to_array magnitudes in
  let filtered =
    Array.mapi
      (fun i c ->
        if Stdlib.( < ) mag_arr.(i) threshold then Complex.zero else c)
      (to_array spectrum)
  in
  let clean_spectrum = create Complex128 (shape spectrum) filtered in

  (* Inverse FFT back to time domain. *)
  let clean_signal = irfft ~n clean_spectrum in

  Printf.printf "After filtering (threshold=%.1f):\n" threshold;
  Printf.printf "  Original first 8:  %s\n"
    (data_to_string (slice [ R (0, 8) ] signal));
  Printf.printf "  Filtered first 8:  %s\n\n"
    (data_to_string (slice [ R (0, 8) ] clean_signal));

  (* --- Frequency bins explained --- *)
  Printf.printf "Frequency bins (first 10): %s\n"
    (data_to_string (slice [ R (0, 10) ] freqs));
  Printf.printf "Total bins: %d (for %d-sample signal)\n" (shape freqs).(0) n