qra/tests/test_logistic.py

import numpy as np
import pytest

from irae_risk.logistic import LogisticPolyRegression


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


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)
    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)
    regularized = model.goodness_of_fit(
        x,
        y,
        theta,
        regularization=True,
    )
    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)
    regularized = model.goodness_of_fit(
        x,
        y,
        theta,
        regularization=True,
    )

    assert regularized == pytest.approx(default)