"""Tests for the lock-in detection module. The important ones are the bound checks: the module's central claim, var(A_hat) = S(f0)/T, is pinned against a numerically-built Fisher information matrix for both white and AR(1) noise, per the arc's standing rule that no Cramer-Rao bound ships on a closed form alone. """ import numpy as np import pytest from scipy.linalg import toeplitz from scipy.signal import freqz, welch from lockin import (front_end_noise, lockin_block, lockin_stream, enbw, amplitude_noise_std) class TestLockinBlockCleanTone: def test_sine_reference_recovers_amplitude_and_phase(self): n = np.arange(1024) f0 = 64 / 1024 x = 1.7 * np.cos(2 * np.pi * f0 * n + 0.9) a, p = lockin_block(x, f0) assert abs(a - 1.7) < 1e-12 assert abs(p - 0.9) < 1e-12 def test_square_reference_is_calibrated(self): # 16 samples per cycle: the sampled-square residual is ~0.65%. n = np.arange(16 * 200) f0 = 1.0 / 16 x = 0.8 * np.cos(2 * np.pi * f0 * n + 0.4) a, p = lockin_block(x, f0, reference='square') assert abs(a - 0.8) < 0.008 assert abs(p - 0.4) < 0.001 def test_square_reference_residual_shrinks_as_M_squared(self): # The docstring's ~1.7/M^2 claim: quadrupling the samples per # cycle must shrink the calibration residual ~16x. errs = {} for M in (8, 16, 64): n = np.arange(M * 400) a, _ = lockin_block(np.cos(2 * np.pi * n / M), 1.0 / M, reference='square') errs[M] = abs(a - 1.0) assert 0.02 < errs[8] < 0.035 assert errs[64] < errs[16] / 10 < errs[8] / 30 def test_fs_units_are_consistent(self): fs = 8000.0 t = np.arange(4000) / fs x = 0.5 * np.cos(2 * np.pi * 1000.0 * t - 0.3) a, p = lockin_block(x, 1000.0, fs=fs) assert abs(a - 0.5) < 1e-12 assert abs(p + 0.3) < 1e-12 def test_rejects_other_frequencies(self): # A tone on a different bin contributes nothing over whole cycles. n = np.arange(1024) x = 3.0 * np.cos(2 * np.pi * (100 / 1024) * n + 1.0) a, _ = lockin_block(x, 64 / 1024) assert a < 1e-12 def test_empty_input_raises(self): with pytest.raises(ValueError): lockin_block(np.array([]), 0.1) def test_unknown_reference_raises(self): with pytest.raises(ValueError): lockin_block(np.ones(64), 0.1, reference='triangle') class TestBoundsAgainstNumericFisher: """The standing-rule checks: closed forms vs the exact FIM.""" def test_white_noise_bound_matches_numeric_fisher(self): # var(A_hat) >= 2 sigma^2 / N for known f0 and unknown (A, phi). # Build the 2x2 FIM from the model derivatives, invert, compare. A, phi, sigma, N = 1.3, 0.7, 0.2, 512 f0 = 60 / N psi = 2 * np.pi * f0 * np.arange(N) + phi G = np.vstack([np.cos(psi), -A * np.sin(psi)]) cov = np.linalg.inv(G @ G.T / sigma**2) assert abs(cov[0, 0] / (2 * sigma**2 / N) - 1) < 1e-3 def test_colored_noise_bound_is_psd_at_f0_over_T(self): # In AR(1) noise the closed form S(f0)/T must match the exact # FIM built with the true covariance. This is the page's # central claim. for a_coef, f0 in [(0.9, 0.25), (0.9, 0.1), (0.5, 0.2)]: N, fs = 512, 1.0 se2 = 1.0 acf = se2 / (1 - a_coef**2) * a_coef**np.arange(N) Ci = np.linalg.inv(toeplitz(acf)) psi = 2 * np.pi * f0 * np.arange(N) + 0.7 A = 1.0 G = np.vstack([np.cos(psi), -A * np.sin(psi)]) cov = np.linalg.inv(G @ Ci @ G.T) S1 = 2 * se2 / abs(1 - a_coef * np.exp(-2j * np.pi * f0))**2 / fs T = N / fs assert abs(cov[0, 0] / (S1 / T) - 1) < 0.01, \ f"S(f0)/T should match the exact FIM at a={a_coef}, f0={f0}" def test_lockin_attains_the_white_noise_bound(self): rng = np.random.default_rng(2) A, sigma, N = 1.0, 0.3, 512 f0 = 60 / N n = np.arange(N) tone = A * np.cos(2 * np.pi * f0 * n + 0.5) est = np.array([lockin_block(tone + rng.normal(0, sigma, N), f0)[0] for _ in range(3000)]) # 3000 trials give ~2.6% standard error on a variance; allow 10%. assert abs(est.var() / (2 * sigma**2 / N) - 1) < 0.10 def test_amplitude_noise_std_equals_crlb_rewritten(self): # S/T with S = 2 sigma^2/fs and T = N/fs is exactly 2 sigma^2/N. sigma, fs, N = 0.3, 8000.0, 4096 S1 = 2 * sigma**2 / fs assert abs(amplitude_noise_std(S1, N / fs)**2 - 2 * sigma**2 / N) < 1e-15 class TestSquareReferencePenalties: def test_variance_penalty_is_pi_squared_over_8(self): rng = np.random.default_rng(5) N, sigma = 4096, 1.0 f0 = 256 / N theta = 2 * np.pi * f0 * np.arange(N) tone = np.cos(theta + 0.7) a_sin, a_sq = [], [] for _ in range(4000): x = tone + rng.normal(0, sigma, N) a_sin.append(lockin_block(x, f0)[0]) a_sq.append(lockin_block(x, f0, reference='square')[0]) ratio = np.var(a_sq) / np.var(a_sin) assert abs(ratio / (np.pi**2 / 8) - 1) < 0.08 def test_square_reference_hears_odd_harmonics(self): # An interferer at 3 f0 reads as 1/3 of its amplitude, at 5 f0 # as 1/5; even harmonics and the sine reference read nothing. N = 64 * 400 f0 = 1.0 / 64 n = np.arange(N) for k, weight in [(3, 1 / 3), (5, 1 / 5)]: xk = np.cos(2 * np.pi * k * f0 * n + 1.1) a_sq, _ = lockin_block(xk, f0, reference='square') a_sin, _ = lockin_block(xk, f0) assert abs(a_sq - weight) < 0.05 * weight, f"harmonic {k}" assert a_sin < 1e-10, f"sine reference must reject {k} f0" for k in (2, 4): xk = np.cos(2 * np.pi * k * f0 * n + 0.3) a_even, _ = lockin_block(xk, f0, reference='square') assert a_even < 1e-6, f"even harmonic {k} must be rejected" class TestFrontEndNoise: def test_white_only(self): rng = np.random.default_rng(0) x = front_end_noise(100000, 0.5, corner=0.0, rng=rng) assert abs(x.std() - 0.5) < 0.01 def test_psd_calibration(self): # White floor S_w above the corner, S_w * corner/f below it. rng = np.random.default_rng(1) fs, sigma_w, corner = 2000.0, 1.0, 20.0 x = front_end_noise(2**20, sigma_w, corner, fs=fs, rng=rng) f, P = welch(x, fs=fs, nperseg=8192) S_w = 2 * sigma_w**2 / fs model = S_w * (1 + corner / f[1:]) hi = f[1:] > 500 lo = (f[1:] > 4) & (f[1:] < 10) assert abs(np.mean(P[1:][hi] / model[hi]) - 1) < 0.05 assert abs(np.mean(P[1:][lo] / model[lo]) - 1) < 0.15 def test_input_validation(self): with pytest.raises(ValueError): front_end_noise(1, 1.0, 10.0) with pytest.raises(ValueError): front_end_noise(100, -1.0, 10.0) with pytest.raises(ValueError): front_end_noise(100, 1.0, -1.0) class TestLockinStream: def test_tracks_a_clean_tone(self): fs, f0, tau = 2000.0, 250.0, 0.05 t = np.arange(int(2 * fs)) / fs x = 0.7 * np.cos(2 * np.pi * f0 * t + 0.2) r, phi = lockin_stream(x, f0, fs=fs, tau=tau) settle = int(8 * tau * fs) assert abs(np.mean(r[settle:]) - 0.7) < 0.01 assert abs(np.mean(phi[settle:]) - 0.2) < 0.02 def test_settles_with_time_constant_tau(self): # Amplitude step at t=0: the envelope estimate must reach # 1 - 1/e of the step one time constant later. fs, f0, tau = 2000.0, 250.0, 0.1 t = np.arange(int(2 * fs)) / fs x = np.cos(2 * np.pi * f0 * t) r, _ = lockin_stream(x, f0, fs=fs, tau=tau) k = int(tau * fs) assert abs(r[k] - (1 - np.exp(-1))) < 0.03 def test_output_noise_variance_is_2_S_Bn(self): # var(I) = 2 S(f0) B_n with B_n = 1/(4 tau). rng = np.random.default_rng(3) fs, f0, tau, sigma = 2000.0, 250.0, 0.05, 1.0 n = 400000 t = np.arange(n) / fs x = np.cos(2 * np.pi * f0 * t + 0.4) + sigma * rng.standard_normal(n) r, _ = lockin_stream(x, f0, fs=fs, tau=tau) settle = int(10 * tau * fs) S1 = 2 * sigma**2 / fs pred = 2 * S1 * (1 / (4 * tau)) # r fluctuates like the in-phase component at this high SNR. assert abs(np.var(r[settle:]) / pred - 1) < 0.10 def test_too_short_tau_raises(self): with pytest.raises(ValueError): lockin_stream(np.ones(100), 0.1, fs=1.0, tau=1.0) def test_empty_input_raises(self): with pytest.raises(ValueError): lockin_stream(np.array([]), 0.1, tau=100.0) class TestEnbw: def test_rectangular_is_half_over_T(self): fs, N = 2000.0, 1000 assert abs(enbw(np.ones(N), fs=fs) - fs / (2 * N)) < 1e-12 def test_scale_invariant(self): w = np.hanning(500) assert abs(enbw(w) - enbw(7.3 * w)) < 1e-15 def test_hann_is_1_5_times_rectangular(self): N = 4096 ratio = enbw(np.hanning(N)) / enbw(np.ones(N)) assert abs(ratio - 1.5) < 0.01 def test_one_pole_enbw_is_quarter_over_tau(self): # Not a window, so checked directly against the |H|^2 integral. fs, tau = 2000.0, 0.1 b = np.exp(-1 / (fs * tau)) w, h = freqz([1 - b], [1, -b], worN=2**16, fs=fs) bn = np.trapezoid(np.abs(h)**2, w) / np.abs(h[0])**2 assert abs(bn / (1 / (4 * tau)) - 1) < 0.01 def test_input_validation(self): with pytest.raises(ValueError): enbw(np.array([])) with pytest.raises(ValueError): enbw(np.array([1.0, -1.0])) class TestRayleighFloor: def test_noise_only_magnitude_has_a_floor(self): # With no signal, each quadrature is N(0, 2 sigma^2/N), so the # magnitude is Rayleigh with mean sqrt(pi/2) * sqrt(2 sigma^2/N): # averaging magnitudes converges to that floor, not to zero, # while averaging the complex outputs does go to zero. rng = np.random.default_rng(7) sigma, N = 1.0, 1024 f0 = 64 / N theta = 2 * np.pi * f0 * np.arange(N) mags, zs = [], [] for _ in range(2000): x = rng.normal(0, sigma, N) z = 2 * np.mean(x * np.exp(-1j * theta)) mags.append(np.abs(z)) zs.append(z) floor = np.sqrt(np.pi / 2) * np.sqrt(2 * sigma**2 / N) assert abs(np.mean(mags) / floor - 1) < 0.05 assert np.abs(np.mean(zs)) < floor / 10 class TestAmplitudeNoiseStd: def test_input_validation(self): with pytest.raises(ValueError): amplitude_noise_std(-1.0, 1.0) with pytest.raises(ValueError): amplitude_noise_std(1.0, 0.0)