SDR-Rover/experiments/lab024b_pluto_welch_spectrum.py

"""
Lab024B. Усреднённый спектр Pluto+ методом Уэлча.

Схема подключения:
    антенна 40–860 МГц -> RX1

Передатчики TX1 и TX2 не используются.
"""

from pathlib import Path
import csv

import adi
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import find_peaks, welch


# ---------------------------------------------------------------------
# Настройки Pluto+
# ---------------------------------------------------------------------

PLUTO_URI = "ip:192.168.2.1"

CENTER_FREQUENCY_HZ = 100_000_000
SAMPLE_RATE_HZ = 2_400_000
RX_BANDWIDTH_HZ = 2_000_000

RX_BUFFER_SIZE = 65_536
BUFFER_COUNT = 16
DISCARD_BUFFER_COUNT = 2


# ---------------------------------------------------------------------
# Настройки спектрального анализа
# ---------------------------------------------------------------------

WELCH_SEGMENT_LENGTH = 8_192
WELCH_OVERLAP_LENGTH = 4_096

SMOOTHING_BANDWIDTH_HZ = 40_000

MINIMUM_PEAK_DISTANCE_HZ = 120_000
MINIMUM_PEAK_PROMINENCE_DB = 4.0
MINIMUM_PEAK_HEIGHT_DB = -25.0

DC_EXCLUSION_HALF_WIDTH_HZ = 40_000

MAXIMUM_CANDIDATE_COUNT = 12


# ---------------------------------------------------------------------
# Выходные файлы
# ---------------------------------------------------------------------

OUTPUT_DIRECTORY = Path("data/processed/lab024b")

SPECTRUM_CSV_PATH = OUTPUT_DIRECTORY / "lab024b_welch_spectrum.csv"
CANDIDATES_CSV_PATH = OUTPUT_DIRECTORY / "lab024b_fm_candidates.csv"
GRAPH_PATH = OUTPUT_DIRECTORY / "lab024b_welch_spectrum.png"
REPORT_PATH = OUTPUT_DIRECTORY / "lab024b_report.txt"


def configure_receiver() -> adi.Pluto:
    """
    Подключается к Pluto+ и настраивает приёмник RX1.
    """

    print("Подключение к Pluto+...")

    sdr = adi.Pluto(uri=PLUTO_URI)

    # Используем первый приёмный канал.
    sdr.rx_enabled_channels = [0]

    sdr.sample_rate = SAMPLE_RATE_HZ
    sdr.rx_lo = CENTER_FREQUENCY_HZ
    sdr.rx_rf_bandwidth = RX_BANDWIDTH_HZ

    # Автоматическая регулировка усиления.
    sdr.gain_control_mode_chan0 = "slow_attack"

    sdr.rx_buffer_size = RX_BUFFER_SIZE

    return sdr


def receive_samples(sdr: adi.Pluto) -> np.ndarray:
    """
    Получает несколько последовательных IQ-буферов.

    Первые буферы отбрасываются, поскольку сразу после настройки
    приёмника могут наблюдаться переходные процессы АРУ и фильтров.
    """

    print()
    print("Отбрасывание переходных буферов...")

    for _ in range(DISCARD_BUFFER_COUNT):
        _ = sdr.rx()

    received_buffers: list[np.ndarray] = []

    print("Получение рабочих IQ-буферов...")

    for buffer_number in range(1, BUFFER_COUNT + 1):
        samples = np.asarray(sdr.rx(), dtype=np.complex64)
        received_buffers.append(samples)

        print(
            f"  Буфер {buffer_number:02d}/{BUFFER_COUNT}: "
            f"{len(samples)} отсчётов"
        )

    combined_samples = np.concatenate(received_buffers)

    if len(combined_samples) == 0:
        raise RuntimeError("Pluto+ не вернул IQ-сэмплы.")

    if not np.all(np.isfinite(combined_samples)):
        raise RuntimeError("В IQ-буфере обнаружены NaN или Inf.")

    return combined_samples


def calculate_welch_spectrum(
    samples: np.ndarray,
    sample_rate_hz: float,
    center_frequency_hz: float,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Рассчитывает спектральную плотность мощности методом Уэлча.

    Возвращает:
        absolute_frequencies_hz — абсолютные радиочастоты;
        power_density           — мощность в линейном масштабе.
    """

    baseband_frequencies_hz, power_density = welch(
        samples,
        fs=sample_rate_hz,
        window="hann",
        nperseg=WELCH_SEGMENT_LENGTH,
        noverlap=WELCH_OVERLAP_LENGTH,
        detrend=False,
        return_onesided=False,
        scaling="density",
    )

    baseband_frequencies_hz = np.fft.fftshift(
        baseband_frequencies_hz
    )

    power_density = np.fft.fftshift(power_density)

    absolute_frequencies_hz = (
        center_frequency_hz + baseband_frequencies_hz
    )

    return absolute_frequencies_hz, power_density


def smooth_power_spectrum(
    power_density: np.ndarray,
    frequency_step_hz: float,
) -> np.ndarray:
    """
    Сглаживает спектр скользящим средним в линейном масштабе.

    Усреднять следует мощность, а не значения в децибелах.
    """

    smoothing_bin_count = int(
        round(SMOOTHING_BANDWIDTH_HZ / frequency_step_hz)
    )

    smoothing_bin_count = max(1, smoothing_bin_count)

    # Нечётная длина ядра обеспечивает симметрию относительно центра.
    if smoothing_bin_count % 2 == 0:
        smoothing_bin_count += 1

    kernel = (
        np.ones(smoothing_bin_count, dtype=np.float64)
        / smoothing_bin_count
    )

    smoothed_power = np.convolve(
        power_density,
        kernel,
        mode="same",
    )

    return smoothed_power


def convert_power_to_relative_db(
    raw_power: np.ndarray,
    smoothed_power: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Переводит мощность в относительные децибелы.

    0 дБ соответствует максимальному значению сглаженного спектра.
    """

    minimum_positive_value = np.finfo(np.float64).tiny

    raw_power_db = 10.0 * np.log10(
        raw_power + minimum_positive_value
    )

    smoothed_power_db = 10.0 * np.log10(
        smoothed_power + minimum_positive_value
    )

    reference_power_db = np.max(smoothed_power_db)

    raw_relative_db = raw_power_db - reference_power_db
    smoothed_relative_db = smoothed_power_db - reference_power_db

    return raw_relative_db, smoothed_relative_db


def detect_signal_candidates(
    frequencies_hz: np.ndarray,
    smoothed_power_db: np.ndarray,
    center_frequency_hz: float,
) -> list[dict[str, float]]:
    """
    Ищет локальные максимумы сглаженного спектра.

    Максимумы сначала рассчитываются по исходному спектру.
    После этого отбрасываются точки из DC-зоны и с краёв диапазона.

    Такой порядок не создаёт искусственных провалов -200 дБ
    и не искажает значение prominence.
    """

    frequency_step_hz = float(
        np.mean(np.diff(frequencies_hz))
    )

    minimum_distance_bins = max(
        1,
        int(
            round(
                MINIMUM_PEAK_DISTANCE_HZ
                / abs(frequency_step_hz)
            )
        ),
    )

    edge_bin_count = max(
        1,
        int(
            round(
                SMOOTHING_BANDWIDTH_HZ
                / abs(frequency_step_hz)
            )
        ),
    )

    # Важно: find_peaks получает настоящий спектр без вставок -200 дБ.
    peak_indices, peak_properties = find_peaks(
        smoothed_power_db,
        height=MINIMUM_PEAK_HEIGHT_DB,
        prominence=MINIMUM_PEAK_PROMINENCE_DB,
        distance=minimum_distance_bins,
    )

    candidates: list[dict[str, float]] = []

    for peak_number, peak_index in enumerate(peak_indices):
        frequency_hz = float(frequencies_hz[peak_index])

        # Отбрасываем возможный аппаратный DC-пик.
        if (
            abs(frequency_hz - center_frequency_hz)
            <= DC_EXCLUSION_HALF_WIDTH_HZ
        ):
            continue

        # Отбрасываем максимумы возле краёв наблюдаемого диапазона.
        if peak_index < edge_bin_count:
            continue

        if peak_index >= len(smoothed_power_db) - edge_bin_count:
            continue

        candidates.append(
            {
                "frequency_hz": frequency_hz,
                "relative_power_db": float(
                    smoothed_power_db[peak_index]
                ),
                "prominence_db": float(
                    peak_properties["prominences"][peak_number]
                ),
            }
        )

    # Оставляем наиболее выраженные сигналы.
    candidates.sort(
        key=lambda item: item["prominence_db"],
        reverse=True,
    )

    candidates = candidates[:MAXIMUM_CANDIDATE_COUNT]

    # В итоговой таблице располагаем сигналы по частоте.
    candidates.sort(
        key=lambda item: item["frequency_hz"]
    )

    return candidates


def save_spectrum_csv(
    frequencies_hz: np.ndarray,
    raw_power_db: np.ndarray,
    smoothed_power_db: np.ndarray,
) -> None:
    """
    Сохраняет полный рассчитанный спектр в CSV-файл.

    Для каждой частотной точки сохраняются:
    - частота в герцах;
    - частота в мегагерцах;
    - исходная относительная мощность;
    - сглаженная относительная мощность.
    """

    with SPECTRUM_CSV_PATH.open(
        "w",
        encoding="utf-8",
        newline="",
    ) as csv_file:
        writer = csv.writer(csv_file)

        writer.writerow(
            [
                "frequency_hz",
                "frequency_mhz",
                "raw_relative_power_db",
                "smoothed_relative_power_db",
            ]
        )

        for frequency_hz, raw_db, smoothed_db in zip(
            frequencies_hz,
            raw_power_db,
            smoothed_power_db,
        ):
            writer.writerow(
                [
                    f"{frequency_hz:.3f}",
                    f"{frequency_hz / 1e6:.6f}",
                    f"{raw_db:.6f}",
                    f"{smoothed_db:.6f}",
                ]
            )


def save_candidates_csv(
    candidates: list[dict[str, float]],
) -> None:
    """
    Сохраняет найденные кандидаты на радиосигналы.
    """

    with CANDIDATES_CSV_PATH.open(
        "w",
        encoding="utf-8",
        newline="",
    ) as csv_file:
        writer = csv.writer(csv_file)

        writer.writerow(
            [
                "frequency_hz",
                "frequency_mhz",
                "relative_power_db",
                "prominence_db",
            ]
        )

        for candidate in candidates:
            writer.writerow(
                [
                    f"{candidate['frequency_hz']:.3f}",
                    f"{candidate['frequency_hz'] / 1e6:.6f}",
                    f"{candidate['relative_power_db']:.3f}",
                    f"{candidate['prominence_db']:.3f}",
                ]
            )


def save_report(
    sample_count: int,
    candidates: list[dict[str, float]],
) -> None:
    """
    Создаёт текстовый отчёт лабораторной.
    """

    report_lines = [
        "Lab024B. Усреднённый спектр Pluto+ методом Уэлча",
        "",
        f"URI: {PLUTO_URI}",
        f"Центральная частота: {CENTER_FREQUENCY_HZ / 1e6:.3f} МГц",
        f"Частота дискретизации: {SAMPLE_RATE_HZ / 1e6:.3f} Мвыб/с",
        f"Полоса RX: {RX_BANDWIDTH_HZ / 1e6:.3f} МГц",
        f"Количество IQ-сэмплов: {sample_count}",
        f"Размер сегмента Уэлча: {WELCH_SEGMENT_LENGTH}",
        f"Перекрытие сегментов: {WELCH_OVERLAP_LENGTH}",
        "",
        f"Найдено кандидатов: {len(candidates)}",
        "",
    ]

    for candidate_number, candidate in enumerate(
        candidates,
        start=1,
    ):
        report_lines.append(
            f"{candidate_number:02d}. "
            f"{candidate['frequency_hz'] / 1e6:.6f} МГц, "
            f"уровень {candidate['relative_power_db']:.2f} дБ, "
            f"выраженность {candidate['prominence_db']:.2f} дБ"
        )

    REPORT_PATH.write_text(
        "\n".join(report_lines),
        encoding="utf-8",
    )


def create_graph(
    frequencies_hz: np.ndarray,
    raw_power_db: np.ndarray,
    smoothed_power_db: np.ndarray,
    candidates: list[dict[str, float]],
) -> None:
    """
    Строит исходный и сглаженный спектры.
    """

    plt.figure(figsize=(13, 7))

    plt.plot(
        frequencies_hz / 1e6,
        raw_power_db,
        linewidth=0.6,
        alpha=0.45,
        label="Спектр Уэлча",
    )

    plt.plot(
        frequencies_hz / 1e6,
        smoothed_power_db,
        linewidth=1.5,
        label="Сглаженный спектр",
    )

    for candidate in candidates:
        frequency_mhz = candidate["frequency_hz"] / 1e6
        power_db = candidate["relative_power_db"]

        plt.scatter(
            [frequency_mhz],
            [power_db],
            marker="o",
        )

        plt.annotate(
            f"{frequency_mhz:.3f}",
            xy=(frequency_mhz, power_db),
            xytext=(0, 10),
            textcoords="offset points",
            ha="center",
            fontsize=8,
            rotation=45,
        )

    plt.axvspan(
        (CENTER_FREQUENCY_HZ - DC_EXCLUSION_HALF_WIDTH_HZ) / 1e6,
        (CENTER_FREQUENCY_HZ + DC_EXCLUSION_HALF_WIDTH_HZ) / 1e6,
        alpha=0.15,
        label="Исключённая DC-зона",
    )

    plt.title(
        "Lab024B. Усреднённый спектр Pluto+ и кандидаты на FM-сигналы"
    )
    plt.xlabel("Частота, МГц")
    plt.ylabel("Относительная мощность, дБ")
    plt.grid(True)
    plt.ylim(-60, 5)
    plt.legend()
    plt.tight_layout()

    plt.savefig(GRAPH_PATH, dpi=160)
    plt.show()


def main() -> None:
    OUTPUT_DIRECTORY.mkdir(parents=True, exist_ok=True)

    sdr = configure_receiver()

    print()
    print("Параметры приёмника:")
    print(f"  URI:                   {PLUTO_URI}")
    print(
        f"  Центральная частота:   "
        f"{sdr.rx_lo / 1e6:.3f} МГц"
    )
    print(
        f"  Частота дискретизации: "
        f"{sdr.sample_rate / 1e6:.3f} Мвыб/с"
    )
    print(
        f"  Полоса RX:             "
        f"{sdr.rx_rf_bandwidth / 1e6:.3f} МГц"
    )
    print(
        f"  Режим усиления:        "
        f"{sdr.gain_control_mode_chan0}"
    )

    samples = receive_samples(sdr)

    print()
    print(f"Всего получено: {len(samples)} IQ-сэмплов")

    frequencies_hz, power_density = calculate_welch_spectrum(
        samples=samples,
        sample_rate_hz=float(sdr.sample_rate),
        center_frequency_hz=float(sdr.rx_lo),
    )

    frequency_step_hz = float(
        np.mean(np.diff(frequencies_hz))
    )

    smoothed_power = smooth_power_spectrum(
        power_density=power_density,
        frequency_step_hz=frequency_step_hz,
    )

    raw_power_db, smoothed_power_db = (
        convert_power_to_relative_db(
            raw_power=power_density,
            smoothed_power=smoothed_power,
        )
    )

    candidates = detect_signal_candidates(
        frequencies_hz=frequencies_hz,
        smoothed_power_db=smoothed_power_db,
        center_frequency_hz=float(sdr.rx_lo),
    )

    save_spectrum_csv(
        frequencies_hz=frequencies_hz,
        raw_power_db=raw_power_db,
        smoothed_power_db=smoothed_power_db,
    )

    save_candidates_csv(candidates)
    save_report(len(samples), candidates)

    print()
    print("Найденные кандидаты на радиосигналы:")
    print()

    if candidates:
        print(
            "  №   Частота, МГц   Уровень, дБ   "
            "Выраженность, дБ"
        )
        print(
            "  --  ------------   -----------   "
            "-----------------"
        )

        for candidate_number, candidate in enumerate(
            candidates,
            start=1,
        ):
            print(
                f"  {candidate_number:2d}  "
                f"{candidate['frequency_hz'] / 1e6:12.6f}   "
                f"{candidate['relative_power_db']:11.2f}   "
                f"{candidate['prominence_db']:17.2f}"
            )
    else:
        print("  Кандидаты не найдены.")

    create_graph(
        frequencies_hz=frequencies_hz,
        raw_power_db=raw_power_db,
        smoothed_power_db=smoothed_power_db,
        candidates=candidates,
    )

    print()
    print("Созданы файлы:")
    print(f"  Спектр:      {SPECTRUM_CSV_PATH}")
    print(f"  Кандидаты:   {CANDIDATES_CSV_PATH}")
    print(f"  График:      {GRAPH_PATH}")
    print(f"  Отчёт:       {REPORT_PATH}")

    print()
    print("Lab024B выполнена успешно.")


if __name__ == "__main__":
    main()