1import numpy as np2import pytest3from pamle.likelihoods import Exponential, Logistic, Poisson, gather_blocks, STATUS_OK, STATUS_DIVERGED456def fd_check(model, theta, db, eps=1e-6):7 p = theta.shape[1]8 g, H = model.score_hessian(theta, db)9 gfd = np.array([(model.loglik(theta + eps * np.eye(p)[i], db) - model.loglik(theta - eps * np.eye(p)[i], db))[0] / (2 * eps) for i in range(p)])10 Hfd = np.array([(model.score_hessian(theta + eps * np.eye(p)[i], db)[0] - model.score_hessian(theta - eps * np.eye(p)[i], db)[0])[0] / (2 * eps) for i in range(p)])11 S = np.linalg.inv(-H)12 T = model.third_contraction(theta, db, S)13 Tfd = np.array([np.sum((model.score_hessian(theta + eps * np.eye(p)[i], db)[1] - model.score_hessian(theta - eps * np.eye(p)[i], db)[1])[0] / (2 * eps) * S[0]) for i in range(p)])14 return np.max(np.abs(g[0] - gfd)), np.max(np.abs(H[0] - Hfd)), np.max(np.abs(T[0] - Tfd))151617@pytest.mark.parametrize("model,theta0", [(Logistic(4), [0.3, 0.5, -0.7, 0.2]), (Poisson(3), [0.4, 0.3, -0.2]), (Exponential(), [1.7])])18def test_derivatives_against_finite_differences(model, theta0):19 rng = np.random.default_rng(11)20 d = model.simulate(rng, 4000, theta0)21 db = {k: v[None] for k, v in d.items()}22 th = np.asarray(theta0, float)[None] * 0.8 + 0.0523 eg, eh, et = fd_check(model, th, db)24 assert eg < 1e-8 and eh < 1e-7 and et < 1e-6252627def test_fit_matches_closed_form_exponential():28 rng = np.random.default_rng(3); ex = Exponential()29 d = ex.simulate(rng, 500, 2.0)30 th, info = ex.fit(d)31 assert abs(th[0] - 1 / d['x'].mean()) < 1e-12 and info['status'] == STATUS_OK323334def test_batched_fit_equals_individual_fits():35 rng = np.random.default_rng(4); m = Logistic(3)36 d = m.simulate(rng, 2048, [0.2, 0.8, -0.5])37 db = gather_blocks(d, np.arange(2048).reshape(8, 256))38 th, info = m.fit_batched(db, tol=1e-11)39 assert np.all(info['status'] == STATUS_OK)40 for k in range(8):41 thk, _ = m.fit({kk: v[k] for kk, v in db.items()}, tol=1e-11)42 assert np.allclose(th[k], thk, atol=1e-8)434445def test_score_zero_at_mle_tight():46 rng = np.random.default_rng(5); m = Poisson(3)47 d = m.simulate(rng, 5000, [0.5, 0.3, -0.2])48 th, info = m.fit(d, tol=1e-12)49 assert info['grad_max'] < 1e-12505152def test_separation_is_flagged():53 """Perfectly separated logistic sample: the MLE does not exist; status must not be OK."""54 m = Logistic(2)55 X = np.hstack([np.ones((40, 1)), np.linspace(-1, 1, 40)[:, None]])56 y = (X[:, 1] > 0).astype(float)57 th, info = m.fit({'X': X, 'y': y}, maxit=200)58 assert info['status'] != STATUS_OK596061def test_random_blocks_cover_everything():62 from pamle.partitioning import random_blocks, equal_blocks, skewed_sizes, unequal_blocks63 rng = np.random.default_rng(0)64 b = random_blocks(1000, 10, rng); assert sorted(b.ravel().tolist()) == list(range(1000))65 assert equal_blocks(12, 3).shape == (3, 4)66 s = skewed_sizes(1000, 7, 4.0); assert s.sum() == 1000 and s.max() / s.min() > 367 ub = unequal_blocks(s); assert sum(len(u) for u in ub) == 100068