harden O2PLS: guard permutation_test against multi-output y, simplify _stack_columns, cache _ssq in orthogonal extraction, add all-scale predict and multi-output score tests

Author: jsture

src/scikit_opls/_o2pls_core.py

@@ -267,7 +267,8 @@ def _extract_one_orthogonal_component(
     # sequential filter that will later be replayed on new samples.
     score = X @ weight
     score_ssq = float(score @ score)
-    block_ssq = max(_ssq(X), 1.0)
+    x_ssq = _ssq(X)
+    block_ssq = max(x_ssq, 1.0)
     if score_ssq <= tol * block_ssq:
         return None
 
@@ -279,7 +280,7 @@ def _extract_one_orthogonal_component(
     if not np.all(np.isfinite(filtered)):
         return None
     # Refuse components that do not measurably reduce the block sum of squares.
-    if _ssq(filtered) >= _ssq(X) - tol * max(_ssq(X), 1.0):
+    if _ssq(filtered) >= x_ssq - tol * block_ssq:
         return None
 
     return OrthogonalBlockComponent(
@@ -348,12 +349,12 @@ def _replay_orthogonal_filter(
 
 
 def _stack_columns(
-    values: list[NDArray[np.float64]], n_rows: int, n_columns: int
+    values: list[NDArray[np.float64]], n_rows: int
 ) -> NDArray[np.float64]:
     """Column-stack values or return a correctly shaped empty matrix."""
     if values:
         return np.column_stack(values)
-    return np.zeros((n_rows, n_columns), dtype=np.float64)
+    return np.zeros((n_rows, 0), dtype=np.float64)
 
 
 def _reconstruction_r2(
@@ -505,12 +506,12 @@ def o2pls_fit(
     B_T = _lstsq_map(T, U)
     B_U = _lstsq_map(U, T)
 
-    X_orth_weights = _stack_columns(x_weights, n_x_features, 0)
-    X_orth_scores = _stack_columns(x_scores, n_samples, 0)
-    X_orth_loadings = _stack_columns(x_loadings, n_x_features, 0)
-    Y_orth_weights = _stack_columns(y_weights, n_y_features, 0)
-    Y_orth_scores = _stack_columns(y_scores, n_samples, 0)
-    Y_orth_loadings = _stack_columns(y_loadings, n_y_features, 0)
+    X_orth_weights = _stack_columns(x_weights, n_x_features)
+    X_orth_scores = _stack_columns(x_scores, n_samples)
+    X_orth_loadings = _stack_columns(x_loadings, n_x_features)
+    Y_orth_weights = _stack_columns(y_weights, n_y_features)
+    Y_orth_scores = _stack_columns(y_scores, n_samples)
+    Y_orth_loadings = _stack_columns(y_loadings, n_y_features)
 
     # Reconstruct nominal joint, orthogonal and residual parts in the original
     # preprocessed coordinates for diagnostics and estimator attributes.

src/scikit_opls/validation.py

@@ -180,7 +180,9 @@ def permutation_test(
         raise ValueError(f"n_permutations must be >= 1, got {n_permutations}")
 
     X = check_array(X, dtype=np.float64)
-    y = np.asarray(y, dtype=np.float64).ravel()
+    y = np.asarray(y, dtype=np.float64)
+    if y.ndim != 1:
+        raise ValueError("permutation_test currently requires a 1D response.")
     check_consistent_length(X, y)
     if not np.all(np.isfinite(y)):
         raise ValueError("y must contain only finite values.")

tests/test_o2pls.py

@@ -125,6 +125,13 @@ def test_score_matches_r2_score_for_1d_y():
     assert model.score(X, y) == pytest.approx(r2_score(y, model.predict(X)))
 
 
+def test_score_matches_r2_score_uniform_average_for_multi_output_y():
+    X, Y = _make_o2pls_data(seed=7)
+    model = O2PLS(n_components=2).fit(X, Y)
+
+    assert model.score(X, Y) == pytest.approx(r2_score(Y, model.predict(X)))
+
+
 @pytest.mark.parametrize(
     ("method", "arg_name"),
     [
@@ -246,6 +253,20 @@ def test_o2pls_predict_unscales_y():
     assert_allclose(model.predict(X_raw), expected)
 
 
+@pytest.mark.parametrize("scale", ["none", "center", "pareto", "standard"])
+def test_o2pls_predict_matches_scaled_filtered_coefficient_path_all_scales(scale):
+    X, y = _regression_data()
+    rng = np.random.default_rng(0)
+    Y = np.column_stack([y, y + 0.1 * rng.normal(size=y.shape)])
+    X_raw = 10.0 + 3.0 * X
+    Y_raw = -5.0 + 2.0 * Y
+    model = O2PLS(scale=scale).fit(X_raw, Y_raw)
+    expected = (
+        model.filter_transform_x(X_raw) @ model.coef_filtered_
+    ) * model.y_std_ + model.y_mean_
+    assert_allclose(model.predict(X_raw), expected)
+
+
 def test_o2pls_predict_x_unscales_x():
     X, y = _regression_data()
     rng = np.random.default_rng(0)