FFT Spectrum Analyzer (python, written by Codex)

Codex wrote this python project from the task below. It installed and ran on a clean Ubuntu 22.04 machine as written.

Task given to the agent:

TASK: FFT Spectrum Analyzer

Write a program that performs Fast Fourier Transform (FFT) analysis on time-domain signal data, identifying dominant frequencies, computing power spectral density, and supporting windowing functions.

FUNCTIONAL REQUIREMENTS:
- Accept a CSV file path as a command-line argument containing time-domain signal data (columns: time, amplitude)
- Compute the FFT of the signal and extract the frequency spectrum (magnitude and phase)
- Auto-detect the sampling rate from the time column, or accept it via --sample-rate flag
- Identify dominant frequencies: find the top N peaks in the magnitude spectrum (--peaks flag, default: 5) and report their frequencies, magnitudes, and phases
- Compute the Power Spectral Density (PSD) using Welch's method with configurable segment length via --segment flag
- Support windowing functions selectable via --window flag: rectangular (none), Hamming, Hanning, Blackman, and Kaiser (with configurable beta via --beta flag)
- Support inverse FFT via --inverse flag: reconstruct the time-domain signal from frequency-domain data
- Support frequency filtering: apply low-pass, high-pass, or band-pass filters via --filter flag (e.g., --filter low:1000 for 1kHz low-pass) and output the filtered signal
- Export the frequency spectrum data as CSV via --export flag
- Print analysis summary to console: sampling rate, number of samples, frequency resolution, dominant frequencies with magnitudes, and total signal power
- Save the full analysis as JSON with --output flag (default: fft_analysis.json)
- If no input is given, generate a sample signal composed of three sine waves at known frequencies (100Hz, 250Hz, 800Hz) with added white noise, sampled at 4000Hz for 1 second, analyze it, and show that the FFT correctly identifies the three component frequencies
- Handle errors: non-uniform sampling, insufficient data points, and signals with DC offset

Create a complete Python project for a clean Ubuntu 22.04 machine with only Python 3.10+ installed. Include:
- Source code
- requirements.txt with all dependencies (direct and transitive) pinned to exact versions
- README.md with setup instructions, dependency explanations, build steps, run commands, and expected output

[]

# FFT Spectrum Analyzer (Python)

## Requirements
- Ubuntu 22.04
- Python 3.10+

## Install
```bash
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
```

## Run
```bash
python src/main.py signal.csv --peaks 5 --window hamming
python src/main.py signal.csv --sample-rate 4000 --segment 512 --filter low:1000 --export spectrum.csv
python src/main.py signal.csv --inverse --output fft_analysis.json
```

If no input is given, a synthetic signal is generated (100Hz, 250Hz, 800Hz + noise).

numpy==2.1.1
scipy==1.14.1

#!/usr/bin/env python3
import argparse
import json
import sys
from pathlib import Path

import numpy as np
from scipy import signal


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="FFT Spectrum Analyzer")
    p.add_argument("file", nargs="?")
    p.add_argument("--sample-rate", type=float)
    p.add_argument("--peaks", type=int, default=5)
    p.add_argument("--segment", type=int, default=256)
    p.add_argument("--window", choices=["rectangular", "hamming", "hanning", "blackman", "kaiser"], default="rectangular")
    p.add_argument("--beta", type=float, default=8.0)
    p.add_argument("--inverse", action="store_true")
    p.add_argument("--filter")
    p.add_argument("--export")
    p.add_argument("--output", default="fft_analysis.json")
    return p.parse_args()


def read_signal(path: Path):
    data = np.loadtxt(path, delimiter=",", skiprows=1)
    t = data[:, 0]
    x = data[:, 1]
    return t, x


def write_sample(path: Path):
    sr = 4000
    t = np.arange(0, 1, 1 / sr)
    x = (
        1.2 * np.sin(2 * np.pi * 100 * t)
        + 0.8 * np.sin(2 * np.pi * 250 * t + 0.7)
        + 0.5 * np.sin(2 * np.pi * 800 * t + 1.2)
        + np.random.default_rng(42).normal(0, 0.15, len(t))
    )
    arr = np.column_stack([t, x])
    np.savetxt(path, arr, delimiter=",", header="time,amplitude", comments="")


def auto_sample_rate(t: np.ndarray) -> float:
    if len(t) < 2:
        raise ValueError("Insufficient data points")
    dt = np.diff(t)
    m = float(np.mean(dt))
    if np.std(dt) > 0.05 * m:
        raise ValueError("Non-uniform sampling detected")
    return 1.0 / m


def make_window(name: str, n: int, beta: float):
    if name == "rectangular":
        return np.ones(n)
    if name == "hamming":
        return np.hamming(n)
    if name == "hanning":
        return np.hanning(n)
    if name == "blackman":
        return np.blackman(n)
    if name == "kaiser":
        return np.kaiser(n, beta)
    return np.ones(n)


def apply_filter_fft(X: np.ndarray, freqs: np.ndarray, spec: str | None):
    if not spec:
        return X
    mask = np.ones_like(freqs, dtype=bool)
    if spec.startswith("low:"):
        cutoff = float(spec.split(":", 1)[1])
        mask = np.abs(freqs) <= cutoff
    elif spec.startswith("high:"):
        cutoff = float(spec.split(":", 1)[1])
        mask = np.abs(freqs) >= cutoff
    elif spec.startswith("band:"):
        low, high = map(float, spec.split(":", 1)[1].split("-"))
        af = np.abs(freqs)
        mask = (af >= low) & (af <= high)
    Xf = X.copy()
    Xf[~mask] = 0
    return Xf


def main() -> int:
    args = parse_args()
    file_path = Path(args.file) if args.file else Path("sample_signal.csv")

    if not args.file:
        write_sample(file_path)

    try:
        t, x_raw = read_signal(file_path)
        sr = args.sample_rate or auto_sample_rate(t)

        w = make_window(args.window, len(x_raw), args.beta)
        x = x_raw * w

        X = np.fft.fft(x)
        freqs = np.fft.fftfreq(len(x), d=1.0 / sr)

        pos = freqs >= 0
        f_pos = freqs[pos]
        mag = np.abs(X[pos])
        phase = np.angle(X[pos])

        # peak detection
        pk_idx, _ = signal.find_peaks(mag)
        top = pk_idx[np.argsort(mag[pk_idx])[::-1][: args.peaks]] if len(pk_idx) else np.array([], dtype=int)
        peaks = [
            {"frequency": float(f_pos[i]), "magnitude": float(mag[i]), "phase": float(phase[i])}
            for i in top
        ]

        # Welch PSD
        win = "boxcar" if args.window == "rectangular" else ("hann" if args.window == "hanning" else args.window)
        f_psd, pxx = signal.welch(x, fs=sr, nperseg=min(args.segment, len(x)), window=win)

        reconstructed = np.fft.ifft(X).real.tolist() if args.inverse else None

        filtered_signal = None
        if args.filter:
            Xf = apply_filter_fft(X, freqs, args.filter)
            filtered_signal = np.fft.ifft(Xf).real.tolist()

        if args.export:
            out = np.column_stack([f_pos, mag, phase])
            np.savetxt(args.export, out, delimiter=",", header="frequency,magnitude,phase", comments="")

        report = {
            "config": vars(args),
            "samplingRate": float(sr),
            "samples": int(len(x_raw)),
            "frequencyResolution": float(sr / len(x_raw)),
            "dominantFrequencies": peaks,
            "totalSignalPower": float(np.mean(x_raw**2)),
            "spectrum": {
                "frequency": f_pos.tolist(),
                "magnitude": mag.tolist(),
                "phase": phase.tolist(),
            },
            "psd": {
                "frequency": f_psd.tolist(),
                "power": pxx.tolist(),
            },
            "reconstructed": reconstructed,
            "filteredSignal": filtered_signal,
        }

        print(f"Sampling rate: {sr:.3f} Hz")
        print(f"Samples: {len(x_raw)}")
        print(f"Resolution: {sr / len(x_raw):.6f} Hz")
        print("Top peaks:")
        for i, p in enumerate(peaks, start=1):
            print(f"{i}. f={p['frequency']:.3f}Hz mag={p['magnitude']:.5f} phase={p['phase']:.4f}")

        Path(args.output).write_text(json.dumps(report, indent=2), encoding="utf-8")
        return 0
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())