"""Tests for the estimation-basics module.""" import numpy as np import pytest from estimation import (bias_variance_mse, weighted_least_squares, crlb_linear_gaussian, polynomial_design) # --------------------------------------------------------------------------- # bias_variance_mse # --------------------------------------------------------------------------- class TestBiasVarianceMSE: def test_decomposition_is_exact(self): """MSE = bias^2 + variance must hold exactly on the sample.""" rng = np.random.default_rng(0) estimates = rng.normal(2.5, 0.8, 5000) bias, var, mse = bias_variance_mse(estimates, 2.0) assert mse == pytest.approx(bias**2 + var, rel=1e-12) def test_known_values(self): estimates = np.array([1.0, 2.0, 3.0, 4.0]) # mean 2.5, var (1/n) = 1.25 bias, var, mse = bias_variance_mse(estimates, 2.0) assert bias == pytest.approx(0.5) assert var == pytest.approx(1.25) assert mse == pytest.approx(1.5) # mean of 1,0,1,4 def test_unbiased_estimator_has_near_zero_bias(self): rng = np.random.default_rng(1) means = np.array([rng.normal(3.0, 1.0, 40).mean() for _ in range(4000)]) bias, _, _ = bias_variance_mse(means, 3.0) assert abs(bias) < 0.01 def test_zero_variance_when_all_equal(self): bias, var, mse = bias_variance_mse(np.full(10, 5.0), 4.0) assert var == pytest.approx(0.0) assert bias == pytest.approx(1.0) assert mse == pytest.approx(1.0) def test_rejects_bad_input(self): with pytest.raises(ValueError): bias_variance_mse(np.ones((3, 3)), 1.0) with pytest.raises(ValueError): bias_variance_mse(np.array([1.0]), 1.0) # --------------------------------------------------------------------------- # weighted_least_squares # --------------------------------------------------------------------------- class TestWeightedLeastSquares: def test_recovers_exact_fit_without_noise(self): x = np.arange(1, 11, dtype=float) X = polynomial_design(x, (1, 2)) theta_true = np.array([2.0, -0.5]) theta, _ = weighted_least_squares(X, X @ theta_true, 0.3 * x) assert theta == pytest.approx(theta_true, abs=1e-10) def test_unweighted_matches_numpy_lstsq(self): rng = np.random.default_rng(2) X = rng.normal(size=(30, 3)) y = rng.normal(size=30) theta, _ = weighted_least_squares(X, y, None) expected = np.linalg.lstsq(X, y, rcond=None)[0] assert theta == pytest.approx(expected, rel=1e-10) def test_scalar_noise_var_equals_unweighted(self): rng = np.random.default_rng(3) X = rng.normal(size=(20, 2)) y = rng.normal(size=20) theta_scalar, _ = weighted_least_squares(X, y, 4.0) theta_none, _ = weighted_least_squares(X, y, None) assert theta_scalar == pytest.approx(theta_none, rel=1e-10) def test_weights_shift_the_fit_toward_precise_samples(self): """A corrupted sample declared noisy should barely move the fit.""" x = np.arange(1, 11, dtype=float) X = polynomial_design(x, (0, 1)) y = 1.0 + 2.0 * x y[-1] += 50.0 # one wild sample var = np.ones_like(x) var[-1] = 1e6 # ... declared very noisy theta_w, _ = weighted_least_squares(X, y, var) theta_u, _ = weighted_least_squares(X, y, None) assert theta_w == pytest.approx([1.0, 2.0], abs=1e-3) assert abs(theta_u[1] - 2.0) > 0.5 # unweighted is dragged off def test_covariance_matches_monte_carlo(self): """The returned covariance must describe the real spread of the fit.""" rng = np.random.default_rng(4) x = np.arange(1, 13, dtype=float) X = polynomial_design(x, (1, 3)) theta_true = np.array([1.0, 0.05]) noise_var = 0.04 * x _, cov = weighted_least_squares(X, X @ theta_true, noise_var) estimates = np.array([ weighted_least_squares(X, X @ theta_true + rng.normal(0, np.sqrt(noise_var)), noise_var)[0] for _ in range(4000) ]) empirical = np.cov(estimates.T) # 4000 trials give roughly a 2% standard error on a variance. assert np.diag(empirical) == pytest.approx(np.diag(cov), rel=0.10) def test_rejects_bad_input(self): X = np.ones((5, 2)) with pytest.raises(ValueError): weighted_least_squares(X, np.ones(4), None) # shape mismatch with pytest.raises(ValueError): weighted_least_squares(X, np.ones(5), -np.ones(5)) # negative variance with pytest.raises(ValueError): weighted_least_squares(np.ones((2, 5)), np.ones(2), None) # underdetermined # --------------------------------------------------------------------------- # crlb_linear_gaussian # --------------------------------------------------------------------------- class TestCRLB: def test_sample_mean_bound_is_sigma_squared_over_n(self): """The canonical worked example on the page.""" N, sigma = 50, 2.0 crlb = crlb_linear_gaussian(np.ones((N, 1)), sigma**2) assert crlb[0, 0] == pytest.approx(sigma**2 / N) def test_matches_wls_covariance(self): """WLS attains the bound: the two expressions must agree exactly.""" x = np.arange(1, 13, dtype=float) X = polynomial_design(x, (1, 3, 4)) noise_var = 0.04 * x _, cov = weighted_least_squares(X, np.zeros_like(x), noise_var) assert crlb_linear_gaussian(X, noise_var) == pytest.approx(cov, rel=1e-12) def test_wls_variance_attains_the_bound(self): """Monte Carlo: no unbiased estimator beats it, and WLS meets it.""" rng = np.random.default_rng(5) x = np.arange(1, 13, dtype=float) X = polynomial_design(x, (1, 3, 4)) theta_true = np.array([1.0, 0.1, -0.008]) noise_var = 0.04 * x crlb = crlb_linear_gaussian(X, noise_var) estimates = np.array([ weighted_least_squares(X, X @ theta_true + rng.normal(0, np.sqrt(noise_var)), noise_var)[0] for _ in range(4000) ]) for j in range(3): assert estimates[:, j].var() == pytest.approx(crlb[j, j], rel=0.10) def test_more_data_lowers_the_bound_as_one_over_n(self): sigma2 = 1.5 b_100 = crlb_linear_gaussian(np.ones((100, 1)), sigma2)[0, 0] b_400 = crlb_linear_gaussian(np.ones((400, 1)), sigma2)[0, 0] assert b_400 == pytest.approx(b_100 / 4) # variance ~ 1/N def test_noisier_data_raises_the_bound(self): X = np.ones((20, 1)) assert (crlb_linear_gaussian(X, 4.0)[0, 0] > crlb_linear_gaussian(X, 1.0)[0, 0]) def test_needs_no_data(self): """The bound depends on the design only, not on any observation.""" X = polynomial_design(np.arange(1, 9, dtype=float), (0, 1)) assert crlb_linear_gaussian(X, 2.0).shape == (2, 2) def test_rejects_nonpositive_variance(self): with pytest.raises(ValueError): crlb_linear_gaussian(np.ones((4, 1)), 0.0) # --------------------------------------------------------------------------- # polynomial_design # --------------------------------------------------------------------------- class TestPolynomialDesign: def test_columns_are_the_requested_powers(self): x = np.array([1.0, 2.0, 3.0]) X = polynomial_design(x, (0, 1, 3)) assert X.shape == (3, 3) assert X[:, 0] == pytest.approx([1, 1, 1]) assert X[:, 1] == pytest.approx([1, 2, 3]) assert X[:, 2] == pytest.approx([1, 8, 27]) def test_skipping_powers_changes_the_model(self): """The page's point: omitting terms is prior knowledge, not a typo.""" x = np.arange(1, 13, dtype=float) assert polynomial_design(x, (1, 3, 4)).shape == (12, 3) assert polynomial_design(x, (0, 1, 2, 3, 4)).shape == (12, 5) def test_rejects_empty_powers(self): with pytest.raises(ValueError): polynomial_design(np.arange(5.0), ())