copilot fixes

Author: anwurl

octopus/modules/octo/training.py

@@ -217,7 +217,15 @@ def _relabel_processed_output(
             # In this case, column order is preserved (no ColumnTransformer reordering)
             output_cols = list(self.feature_cols)
 
+        n_cols = processed_data.shape[1]
         if set(output_cols) != set(self.feature_cols):
+            # If column count also mismatches, raise a clear error
+            if n_cols != len(self.feature_cols):
+                raise ValueError(
+                    f"Pipeline output has {n_cols} columns but expected {len(self.feature_cols)}. "
+                    f"Pipeline columns: {output_cols}, expected: {list(self.feature_cols)}. "
+                    f"This may indicate extra/unexpected columns were passed to the transformer."
+                )
             logger.warning(
                 "Pipeline output columns %s do not match feature_cols %s. Falling back to positional labeling.",
                 output_cols,
@@ -773,7 +781,7 @@ def calculate_fi_featuresused_shap(self, partition="dev", bg_max=200):
             else:
                 feature_names = [f"f{i}" for i in range(n_features)]
 
-        # Build predict function that converts numpy to DataFrame for sklearn compatibility
+        # Build predict function as fallback for sklearn compatibility
         _feature_cols = self.feature_cols
 
         if getattr(self, "ml_type", None) in (MLType.BINARY, MLType.MULTICLASS) and hasattr(
@@ -787,14 +795,18 @@ def predict_fn(X):
             def predict_fn(X):
                 return np.asarray(self.model.predict(pd.DataFrame(np.asarray(X), columns=_feature_cols)))
 
-        # Build explainer
+        # Build explainer: try model directly first for fast Tree/Linear explainers,
+        # fall back to callable wrapper if that fails
+        X_bg_df = pd.DataFrame(X_bg, columns=_feature_cols)
+        X_eval_df = pd.DataFrame(X_eval, columns=_feature_cols)
+
         try:
-            # Let SHAP auto-select the best explainer (Tree for tree models, Kernel otherwise)
-            explainer = shap.Explainer(predict_fn, X_bg)
-            sv = explainer(X_eval)
+            # Try model directly — SHAP can auto-detect Tree/Linear explainers for speed
+            explainer = shap.Explainer(self.model, X_bg_df)
+            sv = explainer(X_eval_df)
         except Exception as e1:
-            logger.debug(f"SHAP auto explainer failed: {e1}. Falling back to callable + Kernel.")
-            # Use the generic constructor so SHAP picks Kernel with the given background
+            logger.debug(f"SHAP auto explainer with model failed: {e1}. Falling back to callable wrapper.")
+            # Fall back to callable approach (always works, but uses slower KernelExplainer)
             explainer = shap.Explainer(predict_fn, X_bg)
             sv = explainer(X_eval)
 
@@ -911,8 +923,8 @@ def predict(self, x: pd.DataFrame) -> np.ndarray:
         if isinstance(x, np.ndarray):
             x = pd.DataFrame(x, columns=self.feature_cols)
         elif isinstance(x, pd.DataFrame):
-            # Reset index to avoid sklearn ColumnTransformer issues
-            x = x.reset_index(drop=True)
+            # Subset to feature_cols to prevent extra columns flowing through ColumnTransformer
+            x = x[self.feature_cols].reset_index(drop=True)
 
         # Apply the same preprocessing pipeline used during training
         x_processed = self._transform_to_dataframe(x)
@@ -931,8 +943,8 @@ def predict_proba(self, x: pd.DataFrame) -> np.ndarray:
         if isinstance(x, np.ndarray):
             x = pd.DataFrame(x, columns=self.feature_cols)
         elif isinstance(x, pd.DataFrame):
-            # Reset index to avoid sklearn ColumnTransformer issues
-            x = x.reset_index(drop=True)
+            # Subset to feature_cols to prevent extra columns flowing through ColumnTransformer
+            x = x[self.feature_cols].reset_index(drop=True)
 
         # Apply the same preprocessing pipeline used during training
         x_processed = self._transform_to_dataframe(x)

octopus/predict/notebook_utils.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import re
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -57,8 +58,9 @@ def find_latest_study(studies_root: str | UPath, prefix: str) -> str:
     """
     root = UPath(studies_root)
     # Match timestamped directories: prefix-YYYYMMDD_HHMMSS
+    _timestamp_pattern = re.compile(re.escape(prefix) + r"-\d{8}_\d{6}$")
     candidates = sorted(
-        [d for d in root.glob(f"{prefix}-*") if d.is_dir()],
+        [d for d in root.glob(f"{prefix}-*") if d.is_dir() and _timestamp_pattern.search(d.name)],
         key=lambda p: p.name,
         reverse=True,
     )

tests/modules/octo/test_column_ordering.py

@@ -8,10 +8,9 @@
     pytest tests/modules/octo/test_column_ordering.py -v
 """
 
-import warnings
-
 import numpy as np
 import pandas as pd
+import pytest
 
 from octopus.models import Models
 from octopus.models.hyperparameter import (
@@ -171,12 +170,12 @@ def _create_training(
     )
 
 
+@pytest.mark.filterwarnings("ignore")
 class TestColumnOrdering:
     """Tests for ColumnTransformer column ordering with mixed types."""
 
     def test_x_train_processed_columns_match_feature_cols(self):
         """Verify x_train_processed has columns in feature_cols order after fit."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -190,7 +189,6 @@ def test_x_train_processed_columns_match_feature_cols(self):
 
     def test_x_dev_processed_columns_match_feature_cols(self):
         """Verify x_dev_processed has columns in feature_cols order after fit."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -201,7 +199,6 @@ def test_x_dev_processed_columns_match_feature_cols(self):
 
     def test_x_test_processed_columns_match_feature_cols(self):
         """Verify x_test_processed has columns in feature_cols order after fit."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -212,7 +209,6 @@ def test_x_test_processed_columns_match_feature_cols(self):
 
     def test_numerical_data_in_numerical_column(self):
         """Verify that numerical columns in x_train_processed contain actual numerical data."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -228,7 +224,6 @@ def test_numerical_data_in_numerical_column(self):
 
     def test_categorical_data_in_categorical_column(self):
         """Verify that categorical columns in x_train_processed contain actual categorical data."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -247,7 +242,6 @@ def test_internal_fi_labels_correct_with_mixed_types(self):
 
         Target is strongly correlated with num1, so num1 should have highest importance.
         """
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data(n_samples=500)
         data_train, data_dev, data_test = _split_data(data)
 
@@ -267,7 +261,6 @@ def test_internal_fi_labels_correct_with_mixed_types(self):
 
     def test_permutation_fi_labels_correct_with_mixed_types(self):
         """Verify permutation FI labels are correct when mixed column types exist."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data(n_samples=500)
         data_train, data_dev, data_test = _split_data(data)
 
@@ -287,7 +280,6 @@ def test_permutation_fi_labels_correct_with_mixed_types(self):
 
     def test_all_numerical_columns_no_regression(self):
         """Verify all-numerical columns still work correctly (regression test)."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_numerical_only_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -304,7 +296,6 @@ def test_all_numerical_columns_no_regression(self):
 
     def test_predict_works_with_mixed_types(self):
         """Verify predict() works correctly with mixed column types."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -317,7 +308,6 @@ def test_predict_works_with_mixed_types(self):
 
     def test_predict_classification_with_mixed_types(self):
         """Verify predict_proba() works correctly with mixed column types for classification."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_mixed_type_data()
         data_train, data_dev, data_test = _split_data(data)
 
@@ -341,7 +331,6 @@ def test_predict_classification_with_mixed_types(self):
 
     def test_relabel_fallback_when_get_feature_names_out_fails(self):
         """Verify fallback when get_feature_names_out() is not available."""
-        warnings.filterwarnings("ignore")
         data, feature_cols, feature_groups = _create_numerical_only_data()
         data_train, data_dev, data_test = _split_data(data)