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 numpy as np
import pytest

import characterize
import compare
import logger
import plot


# --------------------------------------------------------------------------
# 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