MCM-ICM_Toolkit/tests/test_models.py

"""models 包测试。"""

from __future__ import annotations

import numpy as np
import pytest

_NOT_IMPLEMENTED = "未实现：见 docs/03-roadmap.md L1"


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_fit_linear_recovers_known_coefficients(linear_data):
    """真值 y = 2x + 1，拟合系数应在 1% 内。"""
    from mcm_toolkit.models import fit_linear

    x, y = linear_data
    r = fit_linear(x.reshape(-1, 1), y)
    assert np.isclose(r.params["coef"][0], 2.0, rtol=0.01)
    assert np.isclose(r.params["intercept"], 1.0, atol=0.05)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_fit_polynomial_matches_numpy(linear_data):
    from mcm_toolkit.models import fit_polynomial

    x, y = linear_data
    r = fit_polynomial(x.reshape(-1, 1), y, degree=2)
    assert np.allclose(r.params["coef"], np.polyfit(x, y, 2), rtol=1e-6)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_metrics_known_values():
    from mcm_toolkit.models import mae, r2_score, rmse

    y_true = np.array([1.0, 2.0, 3.0])
    y_pred = np.array([1.0, 2.0, 5.0])
    assert np.isclose(rmse(y_true, y_pred), np.sqrt(4 / 3))
    assert np.isclose(mae(y_true, y_pred), 2 / 3)
    assert r2_score(y_true, y_pred) < 1.0


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_r2_is_one_for_perfect_fit():
    from mcm_toolkit.models import r2_score

    y = np.array([1.0, 2.0, 3.0, 4.0])
    assert np.isclose(r2_score(y, y), 1.0)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_compare_models_returns_table_sorted_by_first_metric(linear_data):
    from mcm_toolkit.models import compare_models, fit_linear, fit_polynomial

    x, y = linear_data
    X = x.reshape(-1, 1)
    tbl = compare_models(
        {"linear": fit_linear(X, y), "poly2": fit_polynomial(X, y, degree=2)},
        X,
        y,
        cv=0,
    )
    assert set(tbl.index) >= {"linear", "poly2"}
    assert "rmse" in tbl.columns