"""
Main PipelineOptimizer class implementation.
"""
import numpy as np
import numpy.typing as npt
# Import Bayesian Optimization
from bayes_opt import BayesianOptimization
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import GroupKFold, GroupShuffleSplit, LeavePGroupsOut
from sklearn.pipeline import Pipeline
from spectoprep.logging import configure_logging, get_logger
from spectoprep.modelling.ridge import OptimizedRidgeCV
from spectoprep.pipeline.builder import build_preprocessor_from_bayes
# Import from package modules
from spectoprep.pipeline.config import AVAILABLE_STEPS, DEFAULT_PARAM_BOUNDS, INCOMPATIBLE_SETS
from spectoprep.pipeline.utils import generate_pipeline_configurations
# Valid cross-validation strategies for the objective.
_CV_METHODS = ("group_shuffle_split", "group_kfold", "leave_p_group_out")
# Alpha grid searched internally by the downstream RidgeCV at each evaluation.
_RIDGE_ALPHAS = np.logspace(-3, 3, 25)
[docs]
class PipelineOptimizer:
"""
A class for optimizing machine learning pipelines using Bayesian optimization.
It precomputes possible pipeline configurations and then searches over both
the pipeline configuration (encoded as an index) and the hyperparameters.
"""
[docs]
def __init__(
self,
X_train: npt.NDArray,
y_train: npt.NDArray,
preprocessing_steps: list[str] | str | None = None,
X_test: npt.NDArray | None = None,
y_test: npt.NDArray | None = None,
cv_method: str = "group_shuffle_split",
n_splits: int = 3,
test_size: float = 0.3,
n_groups_out: int = 2,
random_state: int = 42,
groups: npt.NDArray | None = None,
max_pipeline_length: int = 5,
n_jobs: int = -1,
allowed_preprocess_combinations: int | list[int] | tuple[int, ...] | None = None,
ridge_alphas: npt.ArrayLike | None = None,
log_level: str = "INFO",
):
"""
Initialize the PipelineOptimizer.
Args:
X_train: Training features.
y_train: Training targets.
preprocessing_steps: List of preprocessing steps to use.
X_test: Test features (optional).
y_test: Test targets (optional).
cv_method: One of "group_shuffle_split", "group_kfold", or
"leave_p_group_out".
n_splits: Number of CV splits.
test_size: Test set fraction (if using GroupShuffleSplit).
n_groups_out: Number of groups left out (if using LeavePGroupsOut).
random_state: Random seed.
groups: Optional group labels; if None, one group per sample is used.
max_pipeline_length: Maximum number of steps in pipeline.
n_jobs: Number of parallel jobs for compatible estimators.
allowed_preprocess_combinations: Allowed lengths for preprocessing combinations.
ridge_alphas: Alpha grid searched internally by the downstream
``RidgeCV`` at each evaluation. Defaults to
``np.logspace(-3, 3, 25)``. A single value fixes the penalty,
which is useful for two-stage (fixed-model) comparisons.
log_level: Logging level (INFO, DEBUG, WARNING, ERROR).
"""
# Set up structured logging. configure_logging validates the level and
# is idempotent; a host application may also call it first.
configure_logging(level=log_level)
self.log = get_logger("PipelineOptimizer")
# Store data attributes. Validate the feature matrix type early, before
# any shape access below, so a non-array input fails with a clear error.
if not isinstance(X_train, np.ndarray):
raise ValueError("X_train and y_train must be numpy arrays")
self.X_train = X_train
self.y_train = np.ravel(y_train)
self.X_test = X_test
self.y_test = np.ravel(y_test) if y_test is not None else None
# Store cross-validation parameters
self.cv_method = cv_method
self.n_splits = n_splits
self.test_size = test_size
self.n_groups_out = n_groups_out
# Cross-validated R² of every evaluated parameter set (the objective
# itself only tracks RMSE), keyed by the proposed parameters so that
# reporting can recover it for each Bayesian-optimization trial.
self._trial_r2: dict[tuple, float] = {}
# Store other configuration
self.ridge_alphas = (
_RIDGE_ALPHAS
if ridge_alphas is None
else np.atleast_1d(np.asarray(ridge_alphas, dtype=float))
)
self.random_state = random_state
self.max_pipeline_length = max_pipeline_length
self.n_jobs = n_jobs
self.allowed_preprocess_combinations = (
[1, 2] if allowed_preprocess_combinations is None else allowed_preprocess_combinations
)
# Score assigned to a candidate that cannot be evaluated (a transform fails on a
# fold, ...): twice the target's standard deviation, i.e. worse than predicting the
# mean. A scale-aware value keeps the surrogate model of Bayesian optimisation
# well-conditioned, unlike an arbitrary huge constant.
y_sd = float(np.std(self.y_train))
self.failure_score = -2.0 * y_sd if y_sd > 0 else -1.0
self._failed_keys: set[tuple] = set()
# Handle groups
if groups is not None:
self.groups = np.ravel(groups)
else:
self.groups = np.arange(self.X_train.shape[0])
# Validate and set preprocessing steps
self.preprocessing_steps = self._validate_preprocessing_steps(preprocessing_steps)
# Validate inputs
self._validate_inputs()
# Generate all valid pipeline configurations
self.all_pipelines = generate_pipeline_configurations(
self.preprocessing_steps,
INCOMPATIBLE_SETS,
self.max_pipeline_length,
self.allowed_preprocess_combinations,
)
# Build the Bayesian search bounds. The downstream RidgeCV selects its
# own penalty by cross-validation, so ``ridge_alpha`` is not searched.
self.param_bounds = DEFAULT_PARAM_BOUNDS.copy()
self.param_bounds.pop("ridge_alpha", None)
self.param_bounds["pipeline_config"] = (0, len(self.all_pipelines) - 1)
# Log initialization information as a single structured event
self.log.info(
"optimizer_initialized",
n_preprocessing_steps=len(self.preprocessing_steps),
n_pipeline_configs=len(self.all_pipelines),
cv_method=cv_method,
)
def _validate_preprocessing_steps(self, steps: list[str] | str | None = None) -> list[str]:
"""
Validate and standardize preprocessing steps.
Args:
steps: List of preprocessing step names
Returns:
Validated list of preprocessing step names
"""
if steps is None:
# Default steps
steps = ["scaler", "pca", "robust_scaler", "select_k_best"]
if isinstance(steps, str):
steps = [steps]
invalid_steps = set(steps) - set(AVAILABLE_STEPS.keys())
if invalid_steps:
raise ValueError(
f"Invalid preprocessing steps: {invalid_steps}. Available: {list(AVAILABLE_STEPS.keys())}"
)
return list(steps)
def _validate_inputs(self) -> None:
"""
Validate input data and parameters.
Raises:
ValueError: If inputs are invalid
"""
if not isinstance(self.X_train, np.ndarray) or not isinstance(self.y_train, np.ndarray):
raise ValueError("X_train and y_train must be numpy arrays")
if self.X_train.shape[0] != self.y_train.shape[0]:
raise ValueError(
f"X_train and y_train must have the same number of samples. Got {self.X_train.shape[0]} and {self.y_train.shape[0]}"
)
if (self.X_test is None) != (self.y_test is None):
raise ValueError("Both X_test and y_test must be provided together")
if self.X_test is not None and self.X_test.shape[1] != self.X_train.shape[1]:
raise ValueError(
f"X_test and X_train must have the same number of features. Got {self.X_test.shape[1]} and {self.X_train.shape[1]}"
)
if self.cv_method not in _CV_METHODS:
raise ValueError(f"cv_method must be one of {_CV_METHODS}")
if len(self.groups) != self.X_train.shape[0]:
raise ValueError("Groups must have the same number of samples as X_train.")
def _build_cv(self):
"""Construct the group-aware cross-validation splitter for scoring."""
if self.cv_method == "group_shuffle_split":
return GroupShuffleSplit(
n_splits=self.n_splits,
test_size=self.test_size,
random_state=self.random_state,
)
if self.cv_method == "group_kfold":
return GroupKFold(n_splits=self.n_splits)
return LeavePGroupsOut(n_groups=self.n_groups_out)
@staticmethod
def _fit_pipeline(pipeline: Pipeline, X, y, groups) -> None:
"""Fit ``pipeline``, keeping the ridge penalty selection replicate-aware.
The optimizer scores candidates with group-aware folds, but the downstream
``OptimizedRidgeCV`` also runs its own internal cross-validation to choose
the penalty. When replicates exist (some group has more than one spectrum)
those internal folds must respect the groups too, otherwise replicates leak
into the penalty selection. When every spectrum is its own group nothing is
grouped and the ridge keeps its usual shuffled k-fold.
"""
ridge = pipeline.named_steps.get("ridge") if hasattr(pipeline, "named_steps") else None
replicated = groups is not None and len(np.unique(groups)) < len(groups)
if ridge is not None and replicated and hasattr(ridge, "groups"):
ridge.groups = np.asarray(groups)
try:
pipeline.fit(X, y)
finally:
ridge.groups = None
else:
pipeline.fit(X, y)
@staticmethod
def _trial_key(params: dict) -> tuple:
"""Hashable key identifying one evaluated parameter set."""
return tuple(sorted((name, float(value)) for name, value in params.items()))
[docs]
def bayes_objective(self, **params) -> float:
"""Objective function for Bayesian optimization.
The score is the negative mean RMSE under **group-aware
cross-validation on the training data only**. Any supplied test set is
deliberately never touched here: using it to select preprocessing or
hyperparameters would leak the held-out data and inflate reported
performance. The test set is used only for final reporting in
:meth:`get_best_pipeline_predictions`.
Args:
**params: Parameters proposed by the Bayesian optimizer.
Returns:
float: Negative cross-validated RMSE, or a large penalty
(``failure_score``, twice the target's standard deviation) if no valid
predictions could be produced.
"""
try:
# Extract pipeline configuration index and limit to valid range
pipeline_config_index = int(round(params["pipeline_config"]))
pipeline_config_index = max(0, min(pipeline_config_index, len(self.all_pipelines) - 1))
pipeline_config = self.all_pipelines[pipeline_config_index]
# Build pipeline steps
steps = []
for step in pipeline_config:
transformer = build_preprocessor_from_bayes(
step, params, self.X_train.shape, self.random_state, self.n_jobs
)
steps.append((step, transformer))
# Add a cross-validated Ridge as the final estimator; its penalty is
# selected internally by RidgeCV rather than by the Bayesian search.
steps.append(("ridge", OptimizedRidgeCV(alphas=self.ridge_alphas)))
pipeline = Pipeline(steps)
# Configure group-aware cross-validation on the training data.
cv = self._build_cv()
# Perform cross-validation
all_predictions: list[float] = []
all_actuals: list[float] = []
for train_idx, val_idx in cv.split(self.X_train, self.y_train, groups=self.groups):
try:
X_train_fold = self.X_train[train_idx]
X_val_fold = self.X_train[val_idx]
y_train_fold = self.y_train[train_idx]
y_val_fold = self.y_train[val_idx]
# Check condition number to avoid numerical instability
if np.linalg.cond(X_train_fold) > 1e10:
self.log.warning("fold_skipped_high_condition", config=pipeline_config)
continue
self._fit_pipeline(pipeline, X_train_fold, y_train_fold, self.groups[train_idx])
preds = pipeline.predict(X_val_fold)
preds = np.ravel(preds)
all_predictions.extend(preds)
all_actuals.extend(y_val_fold)
except np.linalg.LinAlgError:
self.log.warning("fold_linalg_error", config=pipeline_config, exc_info=True)
continue
except Exception:
self.log.warning("fold_error", config=pipeline_config, exc_info=True)
continue
# Check if we have valid predictions
if not all_predictions:
self.log.warning("trial_no_valid_predictions", config=pipeline_config)
self._failed_keys.add(self._trial_key(params))
return self.failure_score # Penalty score
# Calculate metrics
actual = np.array(all_actuals)
rmse = np.sqrt(mean_squared_error(actual, np.array(all_predictions)))
ss_tot = float(np.sum((actual - actual.mean()) ** 2))
if ss_tot > 0:
ss_res = float(np.sum((actual - np.array(all_predictions)) ** 2))
self._trial_r2[self._trial_key(params)] = 1.0 - ss_res / ss_tot
score = -rmse
self.log.info("cv_evaluated", config=pipeline_config, rmse=float(rmse))
return score
except Exception:
self.log.error("bayes_objective_failed", exc_info=True)
self._failed_keys.add(self._trial_key(params))
return self.failure_score # Penalty score for failed configurations
def _acquisition(self, name: str):
"""Build the acquisition object for ``name`` (``None`` keeps the library default).
``"ucb"`` is the default of the ``bayesian-optimization`` library (upper
confidence bound, kappa 2.576) and is returned as ``None`` so that default
behaviour is unchanged. ``"ei"`` and ``"poi"`` need the acquisition classes of
``bayesian-optimization`` >= 2.
"""
if name == "ucb":
return None
try:
from bayes_opt import acquisition
except ImportError: # bayesian-optimization < 2 passes ``acq=`` to ``maximize`` instead
return None
cls = (
acquisition.ExpectedImprovement
if name == "ei"
else acquisition.ProbabilityOfImprovement
)
return cls(xi=0.01, random_state=self.random_state)
[docs]
def bayesian_optimize(
self, init_points: int = 10, n_iter: int = 50, acquisition_function: str = "ucb"
) -> tuple[dict, Pipeline]:
"""
Run Bayesian optimization to find the best pipeline configuration and hyperparameters.
Args:
init_points: Number of random initial points
n_iter: Number of Bayesian optimization iterations
acquisition_function: ``"ucb"`` (upper confidence bound, the library
default and the default here), ``"ei"`` (expected improvement) or
``"poi"`` (probability of improvement).
Returns:
Tuple containing:
- Dict of best parameters
- Fitted Pipeline with best configuration
"""
# Validate the acquisition function
if acquisition_function not in ["ucb", "ei", "poi"]:
self.log.warning(
"unknown_acquisition_function", requested=acquisition_function, fallback="ucb"
)
acquisition_function = "ucb"
acquisition = self._acquisition(acquisition_function)
# Create optimizer (``acquisition=None`` keeps the library default, UCB)
optimizer = BayesianOptimization(
f=self.bayes_objective,
pbounds=self.param_bounds,
random_state=self.random_state,
verbose=2,
**({"acquisition_function": acquisition} if acquisition is not None else {}),
)
# Run optimization
if acquisition is None and acquisition_function != "ucb":
# bayesian-optimization < 2: select the acquisition function in maximize()
optimizer.maximize( # type: ignore[call-arg]
init_points=init_points, n_iter=n_iter, acq=acquisition_function, xi=0.01
)
else:
optimizer.maximize(init_points=init_points, n_iter=n_iter)
# Store optimizer for later analysis
self.optimizer = optimizer
# Extract best parameters
best_result = optimizer.max
if best_result is None:
raise RuntimeError("Optimization produced no evaluated points.")
best_params = best_result["params"]
# Build the best pipeline from the best parameters
pipeline_config_index = int(round(best_params["pipeline_config"]))
pipeline_config_index = max(0, min(pipeline_config_index, len(self.all_pipelines) - 1))
best_pipeline_config = self.all_pipelines[pipeline_config_index]
# Create and fit the best pipeline
steps = []
for step in best_pipeline_config:
transformer = build_preprocessor_from_bayes(
step, best_params, self.X_train.shape, self.random_state, self.n_jobs
)
steps.append((step, transformer))
steps.append(("ridge", OptimizedRidgeCV(alphas=self.ridge_alphas)))
best_pipeline = Pipeline(steps)
self._fit_pipeline(best_pipeline, self.X_train, self.y_train, self.groups)
# Log best pipeline details
self.log.info(
"best_pipeline_selected",
config=best_pipeline_config,
score=float(best_result["target"]),
steps=[name for name, _ in steps],
)
return best_params, best_pipeline
[docs]
def get_best_pipeline_predictions(
self, best_pipeline: Pipeline
) -> tuple[npt.NDArray, float, float]:
"""
Get predictions using the best pipeline.
Args:
best_pipeline: Fitted pipeline object
Returns:
Tuple containing:
- Predictions array
- RMSE score
- R² score
"""
# Fit the pipeline to training data
self._fit_pipeline(best_pipeline, self.X_train, self.y_train, self.groups)
# If test data is available, use it for evaluation
if self.X_test is not None:
y_test = self.y_test
assert y_test is not None # guaranteed by _validate_inputs
preds = best_pipeline.predict(self.X_test)
preds = np.ravel(preds)
rmse = np.sqrt(mean_squared_error(y_test, preds))
r2 = 1 - np.sum((y_test - preds) ** 2) / np.sum((y_test - np.mean(y_test)) ** 2)
else:
# Out-of-fold predictions placed back into the original sample
# order so callers can plot ``y_train`` against ``preds`` without
# fold-concatenation mismatch.
cv = self._build_cv()
preds = np.full(self.y_train.shape[0], np.nan, dtype=float)
for train_idx, val_idx in cv.split(self.X_train, self.y_train, groups=self.groups):
X_train_fold = self.X_train[train_idx]
X_val_fold = self.X_train[val_idx]
y_train_fold = self.y_train[train_idx]
self._fit_pipeline(
best_pipeline, X_train_fold, y_train_fold, self.groups[train_idx]
)
preds[val_idx] = np.ravel(best_pipeline.predict(X_val_fold))
if np.isnan(preds).any():
raise RuntimeError(
"Cross-validation did not predict every training sample; "
"check groups / cv_method / n_splits."
)
rmse = np.sqrt(mean_squared_error(self.y_train, preds))
r2 = 1 - np.sum((self.y_train - preds) ** 2) / np.sum(
(self.y_train - np.mean(self.y_train)) ** 2
)
return preds, rmse, r2
[docs]
def get_all_tested_pipelines(self) -> list[dict]:
"""
Get details of all tested pipeline configurations.
Returns:
List of dictionaries with pipeline details
"""
if not hasattr(self, "optimizer"):
raise AttributeError("No optimizer found. Please run bayesian_optimize() first.")
results = []
for i, res in enumerate(self.optimizer.res):
params = res["params"]
pipeline_index = int(round(params["pipeline_config"]))
pipeline_index = max(0, min(pipeline_index, len(self.all_pipelines) - 1))
pipeline_config = self.all_pipelines[pipeline_index]
# Calculate metrics from the saved objective score (negative RMSE)
failed = self._trial_key(params) in self._failed_keys
rmse = float("inf") if failed else -res["target"]
r2 = self._trial_r2.get(self._trial_key(params), float("nan"))
# Create a dictionary with trial information. R² is the pooled
# cross-validated R² recorded when the trial was scored; it is NaN
# (never None) if unavailable, e.g. for penalised trials, so that
# downstream numeric code (sorting, plotting) does not raise.
result_dict = {
"trial": i,
"pipeline_config": pipeline_config,
"params": {k: v for k, v in params.items() if k != "pipeline_config"},
"rmse": rmse,
"r2": r2,
}
results.append(result_dict)
return results
[docs]
def print_evaluated_pipelines(self) -> None:
"""
Print details for all evaluated pipelines from the Bayesian optimizer.
This method assumes that bayesian_optimize() has been run and that
self.optimizer exists.
"""
if not hasattr(self, "optimizer"):
self.log.warning("no_optimizer", hint="run bayesian_optimize() first")
return
print("Evaluated pipelines:")
for i, res in enumerate(self.optimizer.res):
params = res["params"]
# Convert the continuous pipeline_config parameter to an integer index
pipeline_index = int(round(params["pipeline_config"]))
# Clamp the index to the valid range
pipeline_index = max(0, min(pipeline_index, len(self.all_pipelines) - 1))
pipeline_config = self.all_pipelines[pipeline_index]
target = res["target"]
print(f"Trial {i}:")
print(f" Pipeline configuration: {pipeline_config}")
print(f" Hyperparameters: {params}")
print(f" Objective (score): {target:.4f}")
[docs]
def export_best_pipeline(self, file_path: str) -> None:
"""
Export the best pipeline configuration and hyperparameters to a file.
Args:
file_path: Path to save the export file
Raises:
AttributeError: If optimizer hasn't been run yet
"""
if not hasattr(self, "optimizer"):
raise AttributeError("No optimizer found. Please run bayesian_optimize() first.")
import json
best_result = self.optimizer.max
if best_result is None:
raise AttributeError("No optimization results available.")
best_params = best_result["params"]
pipeline_index = int(round(best_params["pipeline_config"]))
pipeline_index = max(0, min(pipeline_index, len(self.all_pipelines) - 1))
best_pipeline_config = self.all_pipelines[pipeline_index]
export_data = {
"best_score": best_result["target"],
"pipeline_config": list(best_pipeline_config),
"hyperparameters": {k: v for k, v in best_params.items() if k != "pipeline_config"},
}
with open(file_path, "w") as f:
json.dump(export_data, f, indent=2)
self.log.info("best_pipeline_exported", path=file_path)
[docs]
def summarize_optimization(self) -> dict:
"""
Generate a summary of the optimization results.
Returns:
Dictionary containing optimization summary metrics
"""
if not hasattr(self, "optimizer"):
raise AttributeError("No optimizer found. Please run bayesian_optimize() first.")
results = self.optimizer.res
targets = [r["target"] for r in results]
# Extract best pipeline configuration
best_result = self.optimizer.max
if best_result is None:
raise AttributeError("No optimization results available.")
best_params = best_result["params"]
pipeline_index = int(round(best_params["pipeline_config"]))
best_pipeline_config = self.all_pipelines[pipeline_index]
# Calculate improvement and convergence metrics
initial_performance = min(targets[:5]) if len(targets) >= 5 else min(targets)
final_performance = best_result["target"]
improvement = final_performance - initial_performance
# Check for convergence by looking at the last few iterations
n_last = min(5, len(targets))
recent_targets = targets[-n_last:]
converged = (max(recent_targets) - min(recent_targets)) < 0.001
# Count unique pipeline configurations evaluated
unique_configs = set()
for res in results:
idx = int(round(res["params"]["pipeline_config"]))
idx = max(0, min(idx, len(self.all_pipelines) - 1))
unique_configs.add(idx)
# Create summary
summary = {
"n_trials": len(results),
"n_unique_configs": len(unique_configs),
"best_score": final_performance,
"best_pipeline": list(best_pipeline_config),
"improvement": improvement,
"converged": converged,
"best_rmse": (
float("inf")
if self._trial_key(best_params) in self._failed_keys
else -final_performance
),
}
return summary