1import numpy as np2from pamle.likelihoods import Exponential, Logistic, gather_blocks3from pamle.algorithms import (local_split_fits, worker_corrected_estimates, pilot_estimator, global_newton,4 aee_lin_xi, csl, analytic_bias_correction, estimate_all, quantize_relative)567def _exp_design(seed=0, N=2 ** 16, K=2 ** 8):8 rng = np.random.default_rng(seed); ex = Exponential()9 d = ex.simulate(rng, N, 1.0)10 return ex, d, gather_blocks(d, np.arange(N).reshape(K, N // K)), N, K111213def test_newton_identity_on_data():14 ex, d, db, N, K = _exp_design()15 th_c = 1 / d['x'].mean()16 for x in (0.7, 1.0, 1.4):17 path, diag = global_newton(ex, db, np.array([x]), q=1)18 assert abs(path[1][0] - (2 * x - x * x * d['x'].mean())) < 1e-1219 assert abs((path[1][0] - th_c) + (x - th_c) ** 2 / th_c) < 1e-12202122def test_pilot_bias_orders_exponential():23 """Average over many samples: raw bias ~ theta0/m, jackknife bias ~ -2 theta0/m^2 (exact formulas)."""24 rng = np.random.default_rng(1); ex = Exponential(); m = 32; K = 2000; reps = 20025 raw, j2, j3 = [], [], []26 for _ in range(reps):27 d = ex.simulate(rng, m * K, 1.0); db = gather_blocks(d, np.arange(m * K).reshape(K, m))28 fits = local_split_fits(ex, db, (1, 2, 4))29 raw.append(pilot_estimator(fits[1][0][:, 0]))30 j2.append(pilot_estimator(worker_corrected_estimates(fits, (1, 2))[0]))31 j3.append(pilot_estimator(worker_corrected_estimates(fits, (1, 2, 4))[0]))32 raw, j2, j3 = np.array(raw).ravel(), np.array(j2).ravel(), np.array(j3).ravel()33 se = 1 / np.sqrt(m * K * reps)34 assert abs(raw.mean() - m / (m - 1)) < 4 * se35 assert abs(j2.mean() - (1 - 2 / ((m - 1) * (m - 2)))) < 4 * se36 assert abs(j3.mean() - (1 + 8 / ((m - 1) * (m - 2) * (m - 4)))) < 4 * se373839def test_estimate_all_agrees_with_central_when_predicted():40 rng = np.random.default_rng(2); m = Logistic(3); N = 2 ** 15; K = 6441 d = m.simulate(rng, N, [0.5, 1.0, -1.0]); db = gather_blocks(d, np.arange(N).reshape(K, N // K))42 out, diag = estimate_all(m, db, splits_max=3, q_max=2)43 c = out['central']44 gap = {k: np.sqrt(N) * np.linalg.norm(v - c) for k, v in out.items()}45 assert gap['r1_q2'] < gap['r1_q1'] < gap['r1_q0']46 assert gap['r2_q1'] < gap['r1_q1']47 assert gap['r2_q2'] < 1e-6 and gap['r3_q2'] < 1e-548 assert gap['csl2'] < gap['csl1']49 assert diag['central_status'] == 0505152def test_aee_and_csl_basic():53 ex, d, db, N, K = _exp_design(seed=3)54 fits = local_split_fits(ex, db, (1,))55 th = fits[1][0][:, 0]; H = fits[1][2][:, 0]56 a = aee_lin_xi(th, H)57 assert np.isfinite(a).all()58 path = csl(ex, db, th.mean(axis=0), rounds=1)59 # exponential model: L_0 and L_N differ by a linear term, so the full CSL surrogate maximizer IS the centralized MLE60 assert abs(path[1][0] - 1 / d['x'].mean()) < 1e-1061 # the quadratic (one-step) CSL variant with machine 0's Hessian equals the exact Newton step here (data-free Hessian)62 single, _ = global_newton(ex, db, th.mean(axis=0), q=1, hessian_mode="single")63 exact, _ = global_newton(ex, db, th.mean(axis=0), q=1)64 assert abs(single[1][0] - exact[1][0]) < 1e-12656667def test_analytic_bias_coefficient_exponential():68 ex, d, db, N, K = _exp_design(seed=4)69 th = local_split_fits(ex, db, (1,))[1][0][:, 0]70 thc, b = analytic_bias_correction(ex, db, th)71 assert abs(b.mean() - 1.0) < 0.05 # b_1 = theta0 = 1727374def test_quantize_relative():75 A = np.array([[1.234567, -0.001234], [3.0, 1e-9]])76 Q = quantize_relative(A, 8)77 assert np.all(np.abs(Q - A) <= np.abs(A) * 2.0 ** -8 + 1e-300)78 assert np.allclose(quantize_relative(A, 52), A)79