test_utils.py

import pytest
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from msdbook.utils import fit_logit, plot_contour_map
import warnings
from statsmodels.tools.sm_exceptions import HessianInversionWarning

warnings.simplefilter("ignore", HessianInversionWarning)


@pytest.fixture
def sample_data():
    """Fixture to provide sample data for testing."""
    np.random.seed(0)  # For reproducibility

    # Number of samples
    n = 100

    # Generate some random data
    df = pd.DataFrame(
        {
            "Success": np.random.randint(0, 2, size=n),  # Binary outcome variable (0 or 1)
            "Predictor1": np.random.randn(n),  # Random values for Predictor1
            "Predictor2": np.random.randn(n),  # Random values for Predictor2
            "Interaction": np.random.randn(n),  # Random values for Interaction term
        }
    )

    return df


predictor1 = "Predictor1"
predictor2 = "Predictor2"
interaction = "Interaction"
intercept = "Intercept"
success = "Success"


# @pytest.mark.parametrize("predictors, expected_params, min_coeff, max_coeff", [
#     (['Predictor1', 'Predictor2'], np.array([0.34060709, -0.26968773, 0.31551482, 0.45824332]), 1e-5, 10),  # Adjusted expected params
# ])
@pytest.mark.filterwarnings("ignore:Inverting hessian failed, no bse or cov_params available")
@pytest.mark.parametrize(
    "sample_data, df_resid, df_model, llf",
    [
        (
            pd.DataFrame(
                {
                    predictor1: [1.0, 2.0, 3.0],
                    predictor2: [3.0, 4.0, 5.0],
                    interaction: [2.0, 4.0, 6.0],
                    intercept: [1.0, 1.0, 1.0],
                    success: [1.0, 1.0, 0.0],
                }
            ),
            0.0,
            2.0,
            -6.691275315650184e-06,
        ),
        (
            pd.DataFrame(
                {
                    predictor1: [5.0, 6.0, 7.0],
                    predictor2: [7.0, 8.0, 9.0],
                    interaction: [3.0, 6.0, 9.0],
                    intercept: [1.0, 1.0, 1.0],
                    success: [1.0, 0.0, 1.0],
                }
            ),
            0.0,
            2.0,
            -2.4002923915238235e-06,
        ),
        (
            pd.DataFrame(
                {
                    predictor1: [0.5, 1.5, 2.5],
                    predictor2: [1.0, 2.0, 3.0],
                    interaction: [0.2, 0.4, 0.6],
                    intercept: [1.0, 1.0, 1.0],
                    success: [0.0, 1.0, 1.0],
                }
            ),
            0.0,
            2.0,
            -1.7925479970021486e-05,
        ),
    ],
)
def test_fit_logit(sample_data, df_resid, df_model, llf):
    predictors = [predictor1, predictor2]
    result = fit_logit(sample_data, predictors)
    assert result.df_resid == df_resid
    assert result.df_model == df_model
    assert result.llf == llf


def test_plot_contour_map(sample_data):
    """Test the plot_contour_map function."""
    fig, ax = plt.subplots()

    # Fit a logit model for the purpose of plotting
    result = fit_logit(sample_data, ["Predictor1", "Predictor2"])

    # Dynamically generate grid and levels based on input data to reflect the data range
    xgrid_min, xgrid_max = sample_data["Predictor1"].min(), sample_data["Predictor1"].max()
    ygrid_min, ygrid_max = sample_data["Predictor2"].min(), sample_data["Predictor2"].max()
    xgrid = np.linspace(xgrid_min - 1, xgrid_max + 1, 50)
    ygrid = np.linspace(ygrid_min - 1, ygrid_max + 1, 50)
    levels = np.linspace(0, 1, 10)

    contour_cmap = "viridis"
    dot_cmap = "coolwarm"

    # Call the plot function
    contourset = plot_contour_map(
        ax,
        result,
        sample_data,
        contour_cmap,
        dot_cmap,
        levels,
        xgrid,
        ygrid,
        "Predictor1",
        "Predictor2",
    )

    # Verify the plot and axis limits/labels are correct
    assert contourset is not None
    assert ax.get_xlim() == (xgrid.min(), xgrid.max())
    assert ax.get_ylim() == (ygrid.min(), ygrid.max())
    assert ax.get_xlabel() == "Predictor1"
    assert ax.get_ylabel() == "Predictor2"

    # Verify that scatter plot is present by checking the number of points
    assert len(ax.collections) > 0
    plt.close(fig)


def test_empty_data():
    """Test with empty data to ensure no errors."""
    empty_df = pd.DataFrame({"Success": [], "Predictor1": [], "Predictor2": [], "Interaction": []})

    # Test if fitting with empty data raises an error
    with pytest.raises(ValueError):
        fit_logit(empty_df, ["Predictor1", "Predictor2"])

    # Close any created figures
    plt.close("all")


def test_invalid_predictors(sample_data):
    """Test with invalid predictors."""
    invalid_predictors = ["InvalidPredictor1", "InvalidPredictor2"]

    with pytest.raises(KeyError):
        fit_logit(sample_data, invalid_predictors)


def test_logit_with_interaction(sample_data):
    """Test logistic regression with interaction term."""
    data = sample_data.copy()
    data["Interaction"] = data["Predictor1"] * data["Predictor2"]
    result = fit_logit(data, ["Predictor1", "Predictor2"])

    # Ensure the interaction term is included in the result
    assert "Interaction" in result.params.index


def test_fit_logit_comprehensive(sample_data):
    """Comprehensive test for fit_logit checking various aspects."""
    # Check valid predictors
    result = fit_logit(sample_data, ["Predictor1", "Predictor2"])

    # Validate coefficients are reasonable (not exceeding expected ranges)
    assert np.all(np.abs(result.params) < 10)  # Coefficients should not exceed 10

    # Check if all expected predictors are present in the result
    for predictor in ["Intercept", "Predictor1", "Predictor2", "Interaction"]:
        assert predictor in result.params.index

    # Check p-values are valid
    assert np.all(np.isfinite(result.pvalues))  # P-values should be finite numbers