qra/tests/test_logistic.py

import inspect

import numpy as np
import pytest

from irae_risk.logistic import DEFAULT_RANDOM_SEED, LogisticPolyRegression


@pytest.fixture
def regression_data():
    x = np.linspace(-2.0, 2.0, 20)
    y = np.array([0, 1] * 10)
    theta = np.array([-0.2, 0.8])
    return x, y, theta


def test_stochastic_methods_share_a_fixed_default_seed():
    methods_and_seed_parameters = (
        (LogisticPolyRegression.fit, "seed"),
        (LogisticPolyRegression.goodness_of_fit, "bootstrap_seed"),
        (LogisticPolyRegression.get_model_quantiles_normal, "seed"),
        (LogisticPolyRegression.get_normal_theta, "seed"),
        (LogisticPolyRegression.get_nonparam_boots_theta, "seed"),
        (LogisticPolyRegression.get_nonparam_stratified_boots_theta, "seed"),
        (LogisticPolyRegression.get_parametric_boots_theta, "seed"),
    )

    for method, parameter in methods_and_seed_parameters:
        default = inspect.signature(method).parameters[parameter].default
        assert default == DEFAULT_RANDOM_SEED


def test_differential_evolution_fit_is_reproducible_by_default(
    regression_data,
):
    x, y, _ = regression_data
    model = LogisticPolyRegression(degree=1)

    first = model.fit(x, y, method="diff_evol")
    second = model.fit(x, y, method="diff_evol")

    np.testing.assert_array_equal(first["theta"], second["theta"])
    assert first["cost"] == second["cost"]


def test_x50_and_s50_for_linear_logistic_model():
    model = LogisticPolyRegression(degree=1)
    theta = np.array([-2.0, 0.5])

    x50 = model.get_x50(theta)

    assert x50 == pytest.approx(4.0)
    assert model.model(np.array([x50]), theta)[0] == pytest.approx(0.5)
    assert model.get_s50(theta) == pytest.approx(0.5 / 4.0)


def test_x50_selects_first_crossing_of_unconstrained_polynomial():
    model = LogisticPolyRegression(degree=3)
    # F(x) = (x + 2) (x - 1) (x - 3)
    theta = np.array([6.0, -5.0, -2.0, 1.0])

    assert model.get_x50(theta) == pytest.approx(-2.0)
    assert model.get_s50(theta) == pytest.approx(15.0 / 4.0)


def test_x50_and_s50_are_invariant_to_monotonic_theta_symmetries():
    model = LogisticPolyRegression(degree=3, mono=True)
    theta = np.array([-0.5, 0.4, 1.2, -0.3])
    equivalent_theta = np.array([-0.5, -0.4, -1.2, 0.3])

    x50 = model.get_x50(theta)

    assert model.model(np.array([x50]), theta)[0] == pytest.approx(0.5)
    assert model.get_x50(equivalent_theta) == pytest.approx(x50)
    assert model.get_s50(equivalent_theta) == pytest.approx(
        model.get_s50(theta)
    )


def test_x50_supports_flat_monotonic_midpoint():
    model = LogisticPolyRegression(degree=3, mono=True)
    # F(x) = (x - 1)**3 / 3, so p'(x50) = 0 at x50 = 1.
    theta = np.array([-1.0 / 3.0, 0.0, 1.0, -1.0])

    x50 = model.get_x50(theta)

    assert x50 == pytest.approx(1.0, abs=1e-5)
    assert model.model(np.array([x50]), theta)[0] == pytest.approx(0.5)
    assert model.get_s50(theta) == pytest.approx(0.0, abs=1e-9)


def test_x50_rejects_models_without_a_unique_real_midpoint():
    quadratic = LogisticPolyRegression(degree=2)

    with pytest.raises(ValueError, match="does not reach 0.5"):
        quadratic.get_x50(np.array([1.0, 0.0, 1.0]))

    with pytest.raises(ValueError, match="not uniquely defined"):
        quadratic.get_x50(np.zeros(3))


def test_x50_validates_theta():
    model = LogisticPolyRegression(degree=1)

    with pytest.raises(ValueError, match="theta must have shape"):
        model.get_x50(np.array([1.0]))

    with pytest.raises(ValueError, match="finite"):
        model.get_x50(np.array([0.0, np.nan]))


def test_goodness_of_fit_uses_unpenalized_likelihood_by_default(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1, lam=(0.0, 0.5))

    result = model.goodness_of_fit(x, y, theta, bootstrap_samples=20)
    expected_llf = -model.get_nllf(x, y, theta)
    k = len(theta)
    n = len(x)

    assert result["LLF"] == pytest.approx(expected_llf)
    assert result["AIC"] == pytest.approx(2 * k - 2 * expected_llf)
    assert result["BIC"] == pytest.approx(k * np.log(n) - 2 * expected_llf)


def test_goodness_of_fit_can_include_regularization(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1, lam=(0.0, 0.5))

    unregularized = model.goodness_of_fit(
        x, y, theta, bootstrap_samples=20
    )
    regularized = model.goodness_of_fit(
        x,
        y,
        theta,
        regularization=True,
        bootstrap_samples=20,
    )
    penalty = model.penalty(theta)

    assert regularized["LLF"] == pytest.approx(
        -model.get_cost(x, y, theta)
    )
    assert unregularized["LLF"] - regularized["LLF"] == pytest.approx(penalty)
    assert regularized["AIC"] - unregularized["AIC"] == pytest.approx(
        2 * penalty
    )
    assert regularized["BIC"] - unregularized["BIC"] == pytest.approx(
        2 * penalty
    )


def test_regularization_switch_has_no_effect_without_penalty(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1)

    default = model.goodness_of_fit(x, y, theta, bootstrap_samples=20)
    regularized = model.goodness_of_fit(
        x,
        y,
        theta,
        regularization=True,
        bootstrap_samples=20,
    )

    assert regularized == pytest.approx(default)


def test_goodness_of_fit_reports_reproducible_bootstrap_deviance(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1)

    first = model.goodness_of_fit(
        x,
        y,
        theta,
        bootstrap_samples=25,
        bootstrap_seed=123,
    )
    second = model.goodness_of_fit(
        x,
        y,
        theta,
        bootstrap_samples=25,
        bootstrap_seed=123,
    )

    assert "chi2" not in first
    assert "p-value(chi2)" not in first
    assert first["deviance"] == pytest.approx(2 * model.get_nllf(x, y, theta))
    assert first["p-value(deviance_bootstrap)"] == second[
        "p-value(deviance_bootstrap)"
    ]
    assert 0 < first["p-value(deviance_bootstrap)"] <= 1
    assert first["deviance_bootstrap_samples"] == 25


def test_covariance_methods_match_sandwich_formulas(regression_data):
    x, y, theta = regression_data
    ridge = 0.5
    model = LogisticPolyRegression(degree=1, lam=(0.0, ridge))

    design = np.column_stack([np.ones_like(x), x])
    probabilities = model.model(x, theta)
    weights = probabilities * (1 - probabilities)
    information = (design.T * weights) @ design
    bread = information + 2 * ridge * np.eye(2)
    bread_inv = np.linalg.pinv(bread, hermitian=True)

    scores = (y - probabilities)[:, None] * design
    robust_meat = scores.T @ scores

    expected_model = bread_inv @ information @ bread_inv
    expected_robust = bread_inv @ robust_meat @ bread_inv

    np.testing.assert_allclose(
        model.get_cov(x, y, theta),
        expected_model,
    )
    np.testing.assert_allclose(
        model.get_cov(x, y, theta, method="robust_sandwich"),
        expected_robust,
    )
    np.testing.assert_allclose(
        model.get_cov(x, y, theta, method="inverse_hessian"),
        bread_inv,
    )


def test_covariance_rejects_unknown_method_and_l1_penalty(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1)

    with pytest.raises(ValueError, match="Unknown covariance method"):
        model.get_cov(x, y, theta, method="not-a-method")

    l1_model = LogisticPolyRegression(degree=1, lam=(0.1, 0.0))
    with pytest.raises(ValueError, match="nonzero L1"):
        l1_model.get_cov(x, y, theta)