sklearn/utils/tests/test_optimize.py - main - githubmirror/scikit-learn

sklearn/utils/tests/test_optimize.py
239 строк8 KB
Christian Lorentzen
ENH add Array API to newton-cg in LogisticRegression (#34412)
03 июл 2026, 09:44
03 июл 2026, 09:44add0e83
import warnings

import numpy as np
import pytest
from scipy.optimize import fmin_ncg

from sklearn import config_context
from sklearn.exceptions import ConvergenceWarning
from sklearn.utils._array_api import move_to, yield_namespace_device_dtype_combinations
from sklearn.utils._bunch import Bunch
from sklearn.utils._testing import _array_api_for_tests, assert_allclose
from sklearn.utils.optimize import _check_optimize_result, _newton_cg


def test_newton_cg(global_random_seed):
    # Test that newton_cg gives same result as scipy's fmin_ncg

    rng = np.random.RandomState(global_random_seed)
    A = rng.normal(size=(10, 10))
    x0 = np.ones(10)

    def func(x):
        Ax = A.dot(x)
        return 0.5 * (Ax).dot(Ax)

    def grad(x):
        return A.T.dot(A.dot(x))

    def hess(x, p):
        return p.dot(A.T.dot(A.dot(x.all())))

    def grad_hess(x):
        return grad(x), lambda x: A.T.dot(A.dot(x))

    # func is a definite positive quadratic form, so the minimum is at x = 0
    # hence the use of absolute tolerance.
    assert np.all(np.abs(_newton_cg(grad_hess, func, grad, x0, tol=1e-10)[0]) <= 1e-7)
    assert_allclose(
        _newton_cg(grad_hess, func, grad, x0, tol=1e-7)[0],
        fmin_ncg(f=func, x0=x0, fprime=grad, fhess_p=hess),
        atol=1e-5,
    )


@pytest.mark.parametrize(
    "array_namespace, device_name, dtype_name",
    yield_namespace_device_dtype_combinations(),
)
def test_newton_cg_array_api_compliance(array_namespace, device_name, dtype_name):
    """Test that newton_cg works with Array API input."""
    xp, device = _array_api_for_tests(array_namespace, device_name)
    A = xp.asarray(np.array([[3, -1], [-1, 1]]).astype(dtype_name), device=device)
    y = xp.asarray(np.arange(2).astype(dtype_name), device=device)
    x0 = xp.asarray(np.ones(2).astype(dtype_name), device=device)

    def func(x):
        return 0.5 * (y - A @ x) @ (y - A @ x)

    def grad(x):
        return A.T @ (A @ x - y)

    def hess(x, p):
        return A.T @ (A @ p)

    def grad_hess(x):
        return grad(x), lambda p: hess(x, p)

    with config_context(array_api_dispatch=True):
        res = _newton_cg(grad_hess, func, grad, x0, tol=1e-10)

    assert_allclose(
        move_to(res[0], xp=np, device="cpu"),
        [1 / 2, 3 / 2],
        atol=1e-10,
    )


@pytest.mark.parametrize("verbose", [0, 1, 2])
def test_newton_cg_verbosity(capsys, verbose):
    """Test the std output of verbose newton_cg solver."""
    A = np.eye(2)
    b = np.array([1, 2], dtype=float)

    _newton_cg(
        grad_hess=lambda x: (A @ x - b, lambda z: A @ z),
        func=lambda x: 0.5 * x @ A @ x - b @ x,
        grad=lambda x: A @ x - b,
        x0=np.zeros(A.shape[0]),
        verbose=verbose,
    )  # returns array([1., 2])
    captured = capsys.readouterr()

    if verbose == 0:
        assert captured.out == ""
    else:
        msg = [
            "Newton-CG iter = 1",
            "Check Convergence",
            "max |gradient|",
            "Solver did converge at loss = ",
        ]
        for m in msg:
            assert m in captured.out

    if verbose >= 2:
        msg = [
            "Inner CG solver iteration 1 stopped with",
            "sum(|residuals|) <= tol",
            "Line Search",
            "try line search wolfe1",
            "wolfe1 line search was successful",
        ]
        for m in msg:
            assert m in captured.out

    if verbose >= 2:
        # Set up a badly scaled singular Hessian with a completely wrong starting
        # position. This should trigger 2nd line search check
        A = np.array([[1.0, 2], [2, 4]]) * 1e30  # collinear columns
        b = np.array([1.0, 2.0])
        # Note that scipy.optimize._linesearch LineSearchWarning inherits from
        # RuntimeWarning, but we do not want to import from non public APIs.
        with pytest.warns((RuntimeWarning, UserWarning)):
            _newton_cg(
                grad_hess=lambda x: (A @ x - b, lambda z: A @ z),
                func=lambda x: 0.5 * x @ A @ x - b @ x,
                grad=lambda x: A @ x - b,
                x0=np.array([-2.0, 1]),  # null space of hessian
                verbose=verbose,
            )
        captured = capsys.readouterr()
        msg = [
            "wolfe1 line search was not successful",
            "check loss |improvement| <= eps * |loss_old|:",
            "check sum(|gradient|) < sum(|gradient_old|):",
            "last resort: try line search wolfe2",
        ]
        for m in msg:
            assert m in captured.out

        # Function with locally negative curvature at x0=b.
        # f   = (x - 2) * (x - 1) * (x + 1) * (x + 4) + 2x
        # f'  = 2x (2x^2 + 3x - 9)
        # f'' = 12x^2 + 12x - 18
        # global min at x=-3, local max (saddlepoint) at x=0
        # negative curvature between (-1-sqrt(7))/2 and (-1+sqrt(7))/2
        b = np.array([(-1 + np.sqrt(7)) / 2 - 1e-2])  # point of negative curvature
        with pytest.warns(ConvergenceWarning):
            _newton_cg(
                grad_hess=lambda x: (
                    2 * x * (2 * x**2 + 3 * x - 9),
                    lambda z: (12 * x**2 + 12 * x - 18)[:, None] @ z,
                ),
                func=lambda x: ((x - 2) * (x - 1) * (x + 1) * (x + 4) + 2 * x)[0],
                grad=lambda x: 2 * x * (2 * x**2 + 3 * x - 9),
                x0=b,
                verbose=verbose,
                maxiter=2,
            )
        captured = capsys.readouterr()
        msg = [
            "Inner CG solver iteration 0 detected a negative curvature",
            "curvature at p < -eps * ||p||^2",
        ]
        for m in msg:
            assert m in captured.out

        # Successful inner CG.
        A = np.diag([1e-3, 1, 1e3])
        b = np.array([-2.0, 1, 2.0])
        with pytest.warns(ConvergenceWarning):
            _newton_cg(
                grad_hess=lambda x: (A @ x - b, lambda z: A @ z),
                func=lambda x: 0.5 * x @ A @ x - b @ x,
                grad=lambda x: A @ x - b,
                x0=np.ones_like(b),
                verbose=verbose,
                maxiter=2,
                maxinner=1,
            )
        captured = capsys.readouterr()
        msg = ["Inner CG solver stopped reaching maxiter=1"]
        for m in msg:
            assert m in captured.out


def test_check_optimize():
    # Mock some lbfgs output using a Bunch instance:
    result = Bunch()

    # First case: no warnings
    result.nit = 1
    result.status = 0
    result.message = "OK"

    with warnings.catch_warnings():
        warnings.simplefilter("error")
        _check_optimize_result("lbfgs", result)

    # Second case: warning about implicit `max_iter`: do not recommend the user
    # to increase `max_iter` this is not a user settable parameter.
    result.status = 1
    result.message = "STOP: TOTAL NO. OF ITERATIONS REACHED LIMIT"
    with pytest.warns(ConvergenceWarning) as record:
        _check_optimize_result("lbfgs", result)

    assert len(record) == 1
    warn_msg = record[0].message.args[0]
    assert "lbfgs failed to converge after 1 iteration(s)" in warn_msg
    assert result.message in warn_msg
    assert "Increase the number of iterations" not in warn_msg
    assert "scale the data" in warn_msg

    # Third case: warning about explicit `max_iter`: recommend user to increase
    # `max_iter`.
    with pytest.warns(ConvergenceWarning) as record:
        _check_optimize_result("lbfgs", result, max_iter=1)

    assert len(record) == 1
    warn_msg = record[0].message.args[0]
    assert "lbfgs failed to converge after 1 iteration(s)" in warn_msg
    assert result.message in warn_msg
    assert "Increase the number of iterations" in warn_msg
    assert "scale the data" in warn_msg

    # Fourth case: other convergence problem before reaching `max_iter`: do not
    # recommend increasing `max_iter`.
    result.nit = 2
    result.status = 2
    result.message = "ABNORMAL"
    with pytest.warns(ConvergenceWarning) as record:
        _check_optimize_result("lbfgs", result, max_iter=10)

    assert len(record) == 1
    warn_msg = record[0].message.args[0]
    assert "lbfgs failed to converge after 2 iteration(s)" in warn_msg
    assert result.message in warn_msg
    assert "Increase the number of iterations" not in warn_msg
    assert "scale the data" in warn_msg