rm3100/tests/test_analysis.py

"""Spectral and statistical helpers, checked against signals of known answer.

Each test feeds in something whose spectrum or deviation is known analytically,
so a normalisation slip -- the easy mistake in Welch and Allan code, and an
invisible one on real data -- shows up as a factor rather than a wobble.
"""

import math
from pathlib import Path

import numpy as np
import pytest

import capture
import characterize
import compare
import logger
import plot
from synthetic import drifting_times, field_counts, write_capture


# --------------------------------------------------------------------------
# welch_asd
# --------------------------------------------------------------------------

def test_welch_asd_recovers_the_level_of_white_noise():
    """A white signal of sd s at rate fs sits at s/sqrt(fs/2) per root hertz."""
    fs, sd = 250.0, 20.0
    v = np.random.default_rng(1).normal(0, sd, 200_000)
    freqs, asd = characterize.welch_asd(v, fs)
    assert np.median(asd) == pytest.approx(sd / np.sqrt(fs / 2), rel=0.05)


def test_welch_asd_scales_with_amplitude_not_length():
    fs = 250.0
    rng = np.random.default_rng(2)
    short = characterize.welch_asd(rng.normal(0, 10, 50_000), fs)[1]
    long = characterize.welch_asd(rng.normal(0, 10, 200_000), fs)[1]
    assert np.median(short) == pytest.approx(np.median(long), rel=0.1)
    louder = characterize.welch_asd(rng.normal(0, 20, 50_000), fs)[1]
    assert np.median(louder) == pytest.approx(2 * np.median(short), rel=0.1)


def test_welch_asd_puts_a_tone_in_the_right_bin():
    fs, tone = 250.0, 60.0
    t = np.arange(100_000) / fs
    v = np.sin(2 * np.pi * tone * t)
    freqs, asd = characterize.welch_asd(v, fs)
    assert freqs[np.argmax(asd)] == pytest.approx(tone, abs=fs / 4096)


def test_welch_asd_drops_the_dc_bin():
    """A large DC offset must not appear as signal; detrending removes it."""
    fs = 250.0
    v = 50_000 + np.random.default_rng(3).normal(0, 1, 20_000)
    freqs, asd = characterize.welch_asd(v, fs)
    assert freqs[0] > 0
    assert asd.max() < 10


def test_welch_asd_removes_a_linear_ramp():
    """A slow drift would otherwise smear energy across the low bins."""
    fs = 250.0
    n = 40_000
    rng = np.random.default_rng(4)
    noise = rng.normal(0, 5, n)
    ramped = noise + np.linspace(0, 5000, n)
    flat_asd = characterize.welch_asd(noise, fs)[1]
    ramp_asd = characterize.welch_asd(ramped, fs)[1]
    assert np.median(ramp_asd) == pytest.approx(np.median(flat_asd), rel=0.05)


def test_welch_asd_frequencies_stop_at_nyquist():
    fs = 250.0
    freqs, _ = characterize.welch_asd(
        np.random.default_rng(5).normal(0, 1, 20_000), fs)
    assert freqs[-1] == pytest.approx(fs / 2)
    assert len(freqs) == len(characterize.welch_asd(
        np.random.default_rng(5).normal(0, 1, 20_000), fs)[1])


def test_welch_asd_handles_a_capture_barely_long_enough():
    """capture.py's floor is 64 samples, so the spectrum code must survive it."""
    freqs, asd = characterize.welch_asd(
        np.random.default_rng(6).normal(0, 1, 64), 250.0)
    assert len(freqs) == len(asd) > 0
    assert np.all(np.isfinite(asd))


# --------------------------------------------------------------------------
# allan_deviation
# --------------------------------------------------------------------------

def test_allan_deviation_of_white_noise_falls_as_root_tau():
    """White noise gives slope -1/2 on a log-log ADEV plot."""
    fs, sd = 100.0, 10.0
    v = np.random.default_rng(7).normal(0, sd, 100_000)
    taus, devs = characterize.allan_deviation(v, fs)
    # Fit the log-log slope over the well-averaged decades.
    keep = (taus > 10 / fs) & (taus < 1000 / fs)
    slope = np.polyfit(np.log(taus[keep]), np.log(devs[keep]), 1)[0]
    assert slope == pytest.approx(-0.5, abs=0.05)


def test_allan_deviation_starts_near_the_sample_sd():
    """At tau = one sample the deviation is the sd of the differences."""
    fs, sd = 100.0, 10.0
    v = np.random.default_rng(8).normal(0, sd, 50_000)
    taus, devs = characterize.allan_deviation(v, fs)
    assert taus[0] == pytest.approx(1 / fs)
    assert devs[0] == pytest.approx(sd, rel=0.05)


def test_allan_deviation_turns_up_on_a_ramp():
    """Drift is what an upturn means; a pure ramp must produce one."""
    fs = 100.0
    n = 50_000
    v = np.random.default_rng(9).normal(0, 1, n) + np.linspace(0, 500, n)
    taus, devs = characterize.allan_deviation(v, fs)
    assert devs[-1] > devs[np.argmin(devs)] * 5


def test_allan_deviation_of_a_constant_is_zero():
    taus, devs = characterize.allan_deviation(np.full(10_000, 42.0), 100.0)
    assert np.allclose(devs, 0.0, atol=1e-9)


def test_allan_deviation_taus_increase_and_stay_in_range():
    taus, devs = characterize.allan_deviation(
        np.random.default_rng(10).normal(0, 1, 10_000), 100.0)
    assert np.all(np.diff(taus) > 0)
    assert taus[-1] <= 10_000 / 4 / 100.0
    assert len(taus) == len(devs)


# --------------------------------------------------------------------------
# rolling_mean
# --------------------------------------------------------------------------

def test_rolling_mean_matches_a_naive_partial_window_mean():
    """The point of the count normalisation is that the ends do not taper."""
    v = np.arange(50, dtype=float)
    window = 7
    got = plot.rolling_mean(v, window)
    for i in (0, 1, 25, 48, 49):
        lo = max(0, i - window // 2)
        hi = min(len(v), i + window // 2 + 1)
        assert got[i] == pytest.approx(v[lo:hi].mean())


def test_rolling_mean_preserves_a_constant_including_the_ends():
    v = np.full(100, 7.0)
    assert np.allclose(plot.rolling_mean(v, 21), 7.0)


def test_rolling_mean_reduces_noise_by_root_window():
    v = np.random.default_rng(11).normal(0, 10, 100_000)
    smoothed = plot.rolling_mean(v, 25)
    assert smoothed.std() == pytest.approx(10 / np.sqrt(25), rel=0.1)


@pytest.mark.parametrize("window", [0, 1])
def test_rolling_mean_is_a_no_op_below_two(window):
    v = np.arange(10, dtype=float)
    assert plot.rolling_mean(v, window) is v


# --------------------------------------------------------------------------
# i2c_bus_time
# --------------------------------------------------------------------------

def test_i2c_bus_time_matches_the_hand_calculation():
    """A poll (1 byte) and a results read (9), each n+3 bytes of 9 bits + 3."""
    bits = (1 + 3) * 9 + 3 + (9 + 3) * 9 + 3
    assert bits == 150
    assert logger.i2c_bus_time(400) == pytest.approx(bits / 400_000.0)


def test_i2c_bus_time_is_inversely_proportional_to_speed():
    assert logger.i2c_bus_time(100) == pytest.approx(
        4 * logger.i2c_bus_time(400))
    assert logger.i2c_bus_time(750) == pytest.approx(0.200e-3, abs=5e-6)


def test_i2c_bus_time_at_the_default_is_a_small_share_of_the_period():
    """The documented 6% of a cc=100 period at 750 kHz."""
    import rm3100
    share = logger.i2c_bus_time(750) / rm3100.sample_period(100)
    assert share == pytest.approx(0.06, abs=0.005)


# --------------------------------------------------------------------------
# compare.axis_ratio_spread
# --------------------------------------------------------------------------

def _record(mean):
    mean = np.array(mean, dtype=float)
    return {"mean": mean, "field": float(np.linalg.norm(mean))}


def test_axis_ratio_spread_is_zero_for_a_pure_gain_change():
    a = _record([10_000, 20_000, -15_000])
    b = _record([10_600, 21_200, -15_900])          # every axis x1.06
    spread = compare.axis_ratio_spread(a, b, b["mean"] / a["mean"])
    assert spread == pytest.approx(0.0, abs=1e-9)


def test_axis_ratio_spread_detects_movement():
    a = _record([10_000, 20_000, -15_000])
    b = _record([11_400, 21_000, -13_000])          # each axis moved differently
    spread = compare.axis_ratio_spread(a, b, b["mean"] / a["mean"])
    assert spread > compare.RATIO_SPREAD_OK


def test_axis_ratio_spread_ignores_an_axis_carrying_no_field():
    """A near-zero mean makes its ratio noise, which used to read as movement."""
    a = _record([10_000, 20_000, 5])
    b = _record([10_600, 21_200, -30])              # X and Y are a clean x1.06
    spread = compare.axis_ratio_spread(a, b, b["mean"] / a["mean"])
    assert spread == pytest.approx(0.0, abs=1e-9)


def test_axis_ratio_spread_abstains_with_too_few_usable_axes():
    a = _record([10_000, 3, 5])
    b = _record([10_600, -8, 2])
    assert compare.axis_ratio_spread(a, b, b["mean"] / a["mean"]) is None


# --------------------------------------------------------------------------
# coherent_amplitude, host_time_base, and the time-base test they support
# --------------------------------------------------------------------------

TB_DT = 1.0 / 250.0
TB_ROWS = 25_500
TB_FIELD = (-40000.0, 16000.0, 12000.0)


def _drifted(path, ppm_per_second=0.0, jitter_s=0.0, tones=(), noise_nt=5.0,
             rows=TB_ROWS, seed=0):
    counts = field_counts(rows, TB_FIELD, 100, noise_nt, seed, tones)
    return capture.load(write_capture(
        path, rows=rows, dt=TB_DT, counts=counts,
        times=drifting_times(rows, TB_DT, ppm_per_second, jitter_s, seed)))


def test_coherent_amplitude_recovers_a_planted_tone():
    """2*|mean(v exp(-i w t))| is the tone's peak amplitude, not its rms."""
    fs, n = 250.0, 20_000
    t = np.arange(n) / fs
    v = 7.0 * np.cos(2 * np.pi * 31.0 * t + 0.4)
    assert characterize.coherent_amplitude(v, t, 31.0)[0] == pytest.approx(7.0,
                                                                          rel=1e-3)


def test_coherent_amplitude_is_near_zero_away_from_the_tone():
    fs, n = 250.0, 20_000
    t = np.arange(n) / fs
    v = 7.0 * np.cos(2 * np.pi * 31.0 * t)
    assert characterize.coherent_amplitude(v, t, 40.0)[0] < 0.05


def test_coherent_amplitude_accepts_many_frequencies_at_once():
    t = np.arange(20_000) / 250.0
    out = characterize.coherent_amplitude(np.cos(2 * np.pi * 31.0 * t), t,
                                          [20.0, 31.0, 40.0])
    assert len(out) == 3
    assert out[1] > 10 * max(out[0], out[2])


def test_coherent_amplitude_ignores_a_dc_offset():
    t = np.arange(20_000) / 250.0
    v = 3.0 * np.cos(2 * np.pi * 31.0 * t)
    assert characterize.coherent_amplitude(v + 5000.0, t, 31.0)[0] == \
        pytest.approx(characterize.coherent_amplitude(v, t, 31.0)[0], rel=1e-6)


def test_coherent_amplitude_reads_an_unevenly_sampled_tone():
    """The reason this exists rather than a call to welch_asd.

    Samples taken at irregular times still carry the tone; a spectrum, which can
    only assume they were evenly spaced, puts the energy somewhere else.
    """
    rng = np.random.default_rng(0)
    t = np.sort(rng.uniform(0, 80.0, 20_000))
    v = 7.0 * np.cos(2 * np.pi * 31.0 * t)
    assert characterize.coherent_amplitude(v, t, 31.0)[0] == pytest.approx(7.0,
                                                                          rel=0.02)
    # The same samples read as if they were uniform: the tone is gone.
    uniform = np.arange(len(t)) * (t[-1] / len(t))
    assert characterize.coherent_amplitude(v, uniform, 31.0)[0] < 1.0


def test_host_time_base_at_degree_one_is_the_uniform_grid(tmp_path):
    """Degree 1 is the grid already in use, so it is the family's null case."""
    cap = _drifted(tmp_path / "d.csv", ppm_per_second=8.0)
    linear = characterize.host_time_base(cap, degree=1)
    grid = (cap.sample_index - cap.sample_index[0]) * cap.dt_true
    assert np.allclose(linear, grid, atol=1e-9)


def test_host_time_base_follows_a_planted_drift(tmp_path):
    """Degree 3 tracks the ramp the uniform grid cannot absorb."""
    cap = _drifted(tmp_path / "d.csv", ppm_per_second=8.0)
    truth = cap.system_time - cap.system_time[0]
    grid = (cap.sample_index - cap.sample_index[0]) * cap.dt_true
    fitted = characterize.host_time_base(cap)
    assert np.abs(fitted - truth).max() < np.abs(grid - truth).max() / 100


def test_host_time_base_rejects_read_jitter(tmp_path):
    """Scheduling noise says nothing about when the chip sampled, so it goes."""
    cap = _drifted(tmp_path / "j.csv", jitter_s=2e-3)
    fitted = characterize.host_time_base(cap)
    raw = cap.system_time - cap.system_time[0]
    assert (fitted - raw).std() == pytest.approx(2e-3, rel=0.1)


def test_host_time_base_degrades_on_a_capture_shorter_than_its_degree(tmp_path):
    cap = _drifted(tmp_path / "s.csv", rows=200)
    assert len(characterize.host_time_base(cap, degree=500)) == 200


def test_a_wall_clock_tone_prefers_the_host_time_base(tmp_path):
    """The discriminator's claim, on data where the answer is planted.

    A tone at a fixed frequency in wall time is smeared by a drifting grid, so
    placing the samples by the host clock recovers amplitude the index grid lost.
    """
    cap = _drifted(tmp_path / "w.csv", ppm_per_second=40.0, noise_nt=1.0)
    # 31 Hz in wall time, written against the true sample times.
    t = cap.system_time - cap.system_time[0]
    v = 40.0 * np.cos(2 * np.pi * 31.0 * t)
    grid = (cap.sample_index - cap.sample_index[0]) * cap.dt_true
    on_grid = characterize.coherent_amplitude(v, grid, 31.0)[0]
    on_host = characterize.coherent_amplitude(v, characterize.host_time_base(cap),
                                              31.0)[0]
    assert on_host > 1.5 * on_grid


def test_a_sample_locked_tone_prefers_the_index_grid(tmp_path):
    """The control. A tone locked to the sampling is wrecked by the host clock."""
    cap = _drifted(tmp_path / "l.csv", ppm_per_second=40.0, noise_nt=1.0)
    index = cap.sample_index - cap.sample_index[0]
    v = 40.0 * np.cos(2 * np.pi * 0.124 * index)     # a fixed fraction of fs
    grid = index * cap.dt_true
    on_grid = characterize.coherent_amplitude(v, grid, 0.124 / cap.dt_true)[0]
    on_host = characterize.coherent_amplitude(v, characterize.host_time_base(cap),
                                              0.124 / cap.dt_true)[0]
    assert on_host < 0.75 * on_grid


# --------------------------------------------------------------------------
# sibling_path and the figures
# --------------------------------------------------------------------------

@pytest.mark.parametrize("path, suffix, expected", [
    ("run_noise.png", "drift", "run_drift.png"),
    ("run.png", "drift", "run_drift.png"),
    ("/a/b/run_noise.png", "timebase", "/a/b/run_timebase.png"),
    ("run_noise.pdf", "spectrogram", "run_spectrogram.png"),
    ("noise_run_noise.png", "drift", "noise_run_drift.png"),
    ("a.b.c_noise.png", "drift", "a.b.c_drift.png"),
])
def test_sibling_path_names_a_companion_figure(path, suffix, expected):
    assert characterize.sibling_path(path, suffix) == expected


def test_make_drift_writes_a_figure(tmp_path):
    cap = _drifted(tmp_path / "d.csv", ppm_per_second=8.0)
    out = characterize.make_drift(cap, str(tmp_path / "d_drift.png"))
    assert Path(out).stat().st_size > 0


def test_make_drift_declines_a_capture_too_short_to_window(tmp_path):
    """A short capture is a reason to say so, not to abort the whole run."""
    cap = _drifted(tmp_path / "s.csv", rows=2_000)
    with pytest.raises(ValueError, match="too few|too short"):
        characterize.make_drift(cap, str(tmp_path / "s_drift.png"))


def test_make_time_base_writes_a_figure(tmp_path):
    cap = _drifted(tmp_path / "t.csv", ppm_per_second=8.0,
                   tones=[(0, 0.25, 60.0)])
    out = characterize.make_time_base(cap, str(tmp_path / "t_timebase.png"))
    assert Path(out).stat().st_size > 0


def test_make_time_base_declines_a_constant_field(tmp_path):
    cap = _drifted(tmp_path / "c.csv", noise_nt=0.0)
    with pytest.raises(ValueError, match="constant|no lines"):
        characterize.make_time_base(cap, str(tmp_path / "c_timebase.png"))


# --------------------------------------------------------------------------
# Measuring the oscillator's drift
# --------------------------------------------------------------------------

def test_window_rates_recovers_a_planted_ramp(tmp_path):
    """The rate falls at the planted rate, in fractional terms per second."""
    ppm_per_second = 8.0
    cap = _drifted(tmp_path / "d.csv", ppm_per_second=ppm_per_second)
    centres, rates = characterize.window_rates(cap, 10.0)
    assert len(rates) == 10
    # A period ramping up as (1 + r*t) is a rate falling as (1 - r*t).
    slope = np.polyfit(centres, rates / rates.mean(), 1)[0]
    assert slope == pytest.approx(-ppm_per_second * 1e-6, rel=0.05)


def test_window_rates_is_flat_without_drift(tmp_path):
    cap = _drifted(tmp_path / "f.csv")
    _, rates = characterize.window_rates(cap, 10.0)
    assert np.ptp(rates) / rates.mean() < 1e-9


def test_window_rates_centres_span_the_capture(tmp_path):
    cap = _drifted(tmp_path / "c.csv")
    centres, rates = characterize.window_rates(cap, 10.0)
    assert len(centres) == len(rates)
    assert centres[0] == pytest.approx(5.0)
    assert centres[-1] == pytest.approx(95.0)


@pytest.mark.parametrize("window", [0.0, -1.0, 60.0, 1e6])
def test_window_rates_declines_impossible_windows(tmp_path, window):
    """Non-positive, or too long to fit two windows in the capture."""
    cap = _drifted(tmp_path / "s.csv")
    centres, rates = characterize.window_rates(cap, window)
    assert len(centres) == 0 and len(rates) == 0


def test_rate_stability_recovers_a_planted_ramp(tmp_path):
    """Peak to peak across eight windows, against the ramp that produced it."""
    cap = _drifted(tmp_path / "r.csv", ppm_per_second=8.0)
    # Eight windows of 12.5 s: centres 6.25 s and 93.75 s apart, so the spread
    # is the ramp over 87.5 s, not over the full 100 s.
    expected = 8.0e-6 * cap.duration * (1 - 1 / 8)
    assert characterize.rate_stability(cap) == pytest.approx(expected, rel=0.05)


def test_rate_stability_is_zero_on_a_uniform_grid(tmp_path):
    assert characterize.rate_stability(_drifted(tmp_path / "u.csv")) < 1e-9


def test_rate_stability_declines_a_capture_too_short_to_window(tmp_path):
    cap = _drifted(tmp_path / "t.csv", rows=200)
    assert characterize.rate_stability(cap) == 0.0


def test_rate_allan_of_a_linear_drift_rises_with_tau(tmp_path):
    """sigma_y(tau) = |D| * tau / sqrt(2) for a deterministic frequency ramp.

    This is the shape the real captures show above a few seconds, and it is what
    separates a warming oscillator from a random walk.
    """
    ppm_per_second = 8.0
    cap = _drifted(tmp_path / "a.csv", ppm_per_second=ppm_per_second)
    taus, devs = characterize.rate_allan(cap, [4.0, 8.0, 16.0])
    assert list(taus) == [4.0, 8.0, 16.0]
    for tau, dev in zip(taus, devs):
        assert dev == pytest.approx(ppm_per_second * 1e-6 * tau / math.sqrt(2),
                                    rel=0.05)


def test_rate_allan_falls_with_tau_when_only_jitter_is_present(tmp_path):
    """Independent read jitter averages down; drift does not. Opposite slopes."""
    cap = _drifted(tmp_path / "j.csv", jitter_s=2e-3)
    taus, devs = characterize.rate_allan(cap, [4.0, 16.0])
    assert devs[0] > devs[1] * 2


def test_rate_allan_is_negligible_on_a_uniform_grid(tmp_path):
    cap = _drifted(tmp_path / "z.csv")
    _, devs = characterize.rate_allan(cap, [4.0, 8.0])
    assert np.all(devs < 1e-9)


def test_rate_allan_drops_taus_that_yield_too_few_windows(tmp_path):
    """Two windows give one difference, which is not an estimate of anything."""
    cap = _drifted(tmp_path / "d2.csv")
    taus, devs = characterize.rate_allan(cap, [10.0, 40.0, 200.0])
    assert list(taus) == [10.0]      # 40 s gives two windows, 200 s none
    assert len(devs) == 1


def test_rate_allan_accepts_an_empty_request(tmp_path):
    taus, devs = characterize.rate_allan(_drifted(tmp_path / "e.csv"), [])
    assert len(taus) == 0 and len(devs) == 0