doc tests made more robust

Author: Julien Roussel

qolmat/imputations/softimpute.py

@@ -59,14 +59,13 @@ class SoftImpute(BaseEstimator, TransformerMixin):
     >>> from qolmat.imputations.softimpute import SoftImpute
     >>> D = np.array([[1, 2, np.nan, 4], [1, 5, 3, np.nan], [4, 2, 3, 2], [1, 1, 5, 4]])
     >>> Omega = ~np.isnan(D)
-    >>> M, A = SoftImpute(random_state=11, tau=1).decompose(D, Omega)
-    >>> print(M + A)
-    [[1.         2.         3.04868607 4.        ]
-     [1.         5.         3.         3.37501463]
-     [4.         2.         3.         2.        ]
-     [1.         1.         5.         4.        ]]
-    >>> print(SoftImpute.cost_function(D, M, A, Omega, tau=1))
-    18.520174964466026
+    >>> M, A = SoftImpute(random_state=10, tau=1).decompose(D, Omega)
+    >>> naive_cost = SoftImpute.cost_function(
+    ...     D, np.where(Omega, M, 0), np.zeros_like(M), Omega, tau=1
+    ... )
+    >>> minimal_cost = SoftImpute.cost_function(D, M, A, Omega, tau=1)
+    >>> minimal_cost < naive_cost
+    np.True_
 
     """
 

tests/imputations/test_softimpute.py

@@ -2,59 +2,74 @@
 import pytest
 from numpy.typing import NDArray
 
-from qolmat.imputations import softimpute
+from qolmat.imputations.softimpute import SoftImpute
 
-X = np.random.rand(100, 100)
-X[np.random.choice(100, 10), np.random.choice(100, 10)] = np.nan
-X_non_regression_test = np.array(
-    [[1, 2, np.nan, 4], [1, 5, 3, np.nan], [4, 2, 3, 2], [1, 1, 5, 4]]
-)
-X_expected = np.array([[1, 2, 2.9066, 4], [1, 5, 3, 2.1478], [4, 2, 3, 2], [1, 1, 5, 4]])
-tau = 1
-max_iterations = 30
-random_state = 50
+
+@pytest.fixture
+def X_random() -> NDArray:
+    """Generate random matrix with missing values."""
+    rng = np.random.RandomState(42)
+    X = rng.rand(100, 100)
+    X[rng.choice(100, 10), rng.choice(100, 10)] = np.nan
+    return X
+
+
+@pytest.fixture
+def X_non_regression() -> NDArray:
+    """Get small test matrix for non-regression tests."""
+    return np.array([[1, 2, np.nan, 4], [1, 5, 3, np.nan], [4, 2, 3, 2], [1, 1, 5, 4]])
+
+
+@pytest.fixture
+def X_expected() -> NDArray:
+    """Get expected imputed values for non-regression test."""
+    return np.array([[1, 2, 2.9066, 4], [1, 5, 3, 2.1478], [4, 2, 3, 2], [1, 1, 5, 4]])
+
+
+@pytest.fixture
+def default_params() -> dict:
+    """Get default parameters for SoftImpute."""
+    return {"tau": 1, "max_iterations": 30, "random_state": 50}
 
 
 def test_initialized_default() -> None:
     """Test that initialization does not crash and has default parameters."""
-    model = softimpute.SoftImpute()
+    model = SoftImpute()
     assert model.period == 1
     assert model.rank is None
     assert model.tolerance == 1e-05
 
 
 def test_initialized_custom() -> None:
     """Test that initialization does not crash and has custom parameters."""
-    model = softimpute.SoftImpute(period=2, rank=10)
+    model = SoftImpute(period=2, rank=10)
     assert model.period == 2
     assert model.rank == 10
     assert model.tau is None
 
 
-@pytest.mark.parametrize("X", [X])
-def test_soft_impute_decompose(X: NDArray) -> None:
+def test_soft_impute_decompose(X_random: NDArray, default_params: dict) -> None:
     """Test fit instance and decomposition is computed."""
-    tau = 1
-    model = softimpute.SoftImpute(tau=tau)
-    Omega = ~np.isnan(X)
-    X_imputed = np.where(Omega, X, 0)
-    cost_all_in_M = model.cost_function(X, X_imputed, np.full_like(X, 0), Omega, tau)
-    cost_all_in_A = model.cost_function(X, np.full_like(X, 0), X_imputed, Omega, tau)
-    M, A = model.decompose(X, Omega)
-    cost_final = model.cost_function(X, M, A, Omega, tau)
-    assert isinstance(model, softimpute.SoftImpute)
-    assert M.shape == X.shape
-    assert A.shape == X.shape
+    tau = default_params["tau"]
+    model = SoftImpute(tau=tau)
+    Omega = ~np.isnan(X_random)
+    X_imputed = np.where(Omega, X_random, 0)
+    cost_all_in_M = model.cost_function(X_random, X_imputed, np.full_like(X_random, 0), Omega, tau)
+    cost_all_in_A = model.cost_function(X_random, np.full_like(X_random, 0), X_imputed, Omega, tau)
+    M, A = model.decompose(X_random, Omega)
+    cost_final = model.cost_function(X_random, M, A, Omega, tau)
+    assert isinstance(model, SoftImpute)
+    assert M.shape == X_random.shape
+    assert A.shape == X_random.shape
     assert not np.any(np.isnan(M))
     assert not np.any(np.isnan(A))
     assert cost_final < cost_all_in_M
     assert cost_final < cost_all_in_A
 
 
-@pytest.mark.parametrize("X", [X])
-def test_soft_impute_convergence(X: NDArray) -> None:
+def test_soft_impute_convergence() -> None:
     """Test type of the check convergence."""
-    model = softimpute.SoftImpute()
+    model = SoftImpute()
     M = model.random_state.uniform(size=(10, 20))
     U, D, V = np.linalg.svd(M, full_matrices=False)
     ratio = model._check_convergence(U, D, V.T, U, D, V.T)
@@ -63,7 +78,7 @@ def test_soft_impute_convergence(X: NDArray) -> None:
 
 def test_soft_impute_convergence_with_none() -> None:
     """Test check type None and raise error."""
-    model = softimpute.SoftImpute()
+    model = SoftImpute()
     with pytest.raises(ValueError):
         _ = model._check_convergence(
             np.array([1]),
@@ -73,3 +88,22 @@ def test_soft_impute_convergence_with_none() -> None:
             np.array([1]),
             np.array([1]),
         )
+
+
+def test_decompose_loss_minimized(X_random: NDArray, default_params: dict) -> None:
+    """Test that the loss function is at a local minimum."""
+    tau = default_params["tau"]
+    imputer = SoftImpute(random_state=123, tau=tau)
+    Omega = ~np.isnan(X_random)
+    M, A = imputer.decompose(X_random, Omega)
+    X_imputed = M + A
+    cost_imputed = SoftImpute.cost_function(X_imputed, M, A, Omega, tau)
+    for i in range(10):
+        Delta = 1.1 ** (i - 9) * imputer.random_state.uniform(0, 1, size=X_random.shape)
+        X_perturbed = X_imputed + Delta
+        cost_perturbed = SoftImpute.cost_function(X_perturbed, M, A, Omega, tau)
+        assert cost_perturbed > cost_imputed
+    M = np.zeros(X_random.shape)
+    A = X_random.copy()
+    cost_perturbed = SoftImpute.cost_function(X_random, M, A, Omega, tau)
+    assert cost_perturbed > cost_imputed