Coverage for ml_workbench/runner.py: 6%
748 statements
« prev ^ index » next coverage.py v7.11.0, created at 2026-01-06 16:09 +0200
1"""Runner for executing ML experiments."""
3# ToDo: Add support for estimator specific transformer: https://chatgpt.com/share/e/6923056a-7a44-8012-a36d-d822b913db60
5from __future__ import annotations
7from copy import deepcopy
8from pathlib import Path
9from typing import TYPE_CHECKING, Any
10import tempfile
11import os
13import matplotlib
15# matplotlib.use("Agg") # Use non-interactive backend
16import matplotlib.pyplot as plt
17import mlflow
18import numpy as np
19import pandas as pd
20import yaml
22try:
23 import seaborn as sns
24except ImportError:
25 sns = None
26from sklearn.compose import ColumnTransformer
27from sklearn.impute import SimpleImputer
28from sklearn.metrics import (
29 accuracy_score,
30 confusion_matrix,
31 f1_score,
32 mean_absolute_error,
33 mean_squared_error,
34 precision_score,
35 r2_score,
36 recall_score,
37)
38from sklearn.model_selection import (
39 GridSearchCV,
40 GroupShuffleSplit,
41 RandomizedSearchCV,
42 train_test_split,
43)
44from sklearn.pipeline import Pipeline
45from sklearn.preprocessing import OneHotEncoder, StandardScaler
47from .dataset import Dataset
48from .feature import Feature
49from .model import Model
50from .mlflow_conf import MlflowConf
52if TYPE_CHECKING: 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true
53 from .experiment import Experiment
56class ModelRunner:
57 """Handles training, evaluation, and MLflow logging for a specific model.
59 Responsibilities:
60 - Create model-specific preprocessing pipeline
61 - Create model-specific pipeline (preprocessor + model)
62 - Train the model
63 - Evaluate model performance
64 - Calculate feature weights/importances
65 - Log results to MLflow
67 Parameters
68 ----------
69 model_name : str
70 Name of the model to run
71 numerical_features : List[str]
72 List of numerical feature column names
73 categorical_features : List[str]
74 List of categorical feature column names
75 experiment : Experiment
76 Experiment specification
77 X_train : pd.DataFrame
78 Training features
79 y_train : pd.Series
80 Training target
81 X_holdout : Optional[pd.DataFrame]
82 Holdout features (optional)
83 y_holdout : Optional[pd.Series]
84 Holdout target (optional)
85 verbose : bool, optional
86 Whether to print progress information, by default True
87 """
89 def __init__(
90 self,
91 model_name: str,
92 numerical_features: list[str],
93 categorical_features: list[str],
94 experiment: Experiment,
95 X_train: pd.DataFrame, # noqa: N803
96 y_train: pd.Series,
97 X_holdout: pd.DataFrame | None = None, # noqa: N803
98 y_holdout: pd.Series | None = None,
99 verbose: bool = True,
100 ):
101 self.model_name = model_name
102 self.numerical_features = numerical_features
103 self.categorical_features = categorical_features
104 self.experiment = experiment
105 self.config = experiment.config
106 self.X_train = X_train
107 self.y_train = y_train
108 self.X_holdout = X_holdout
109 self.y_holdout = y_holdout
110 self.verbose = verbose
112 # Pipeline and model storage
113 self.full_pipeline: Pipeline | None = (
114 None # full pipeline with preprocessor and model
115 )
116 self.best_pipeline: Pipeline | None = (
117 None # best pipeline after cross-validation
118 )
119 self.is_cross_validation: bool = False
120 # cross validation results on train dataset (used to select the best model)
121 self.cv_results_: dict[str, Any] | None = None
122 self.cv_best_params_: dict[str, Any] | None = None
123 self.cv_best_score_: float | None = None
125 # metrics and selection score on holdout and train datasets (used to select the best model)
126 self.metrics: dict[str, float] | None = None
127 self.selection_score: float | None = (
128 None # the score used to select the best model if multiple metrics are specified
129 )
131 def _log(self, message: str) -> None:
132 """Print message if verbose is enabled."""
133 if self.verbose:
134 print(f"[ModelRunner:{self.model_name}] {message}") # noqa: T201
136 def _build_preprocessing_pipeline(self) -> ColumnTransformer:
137 """Build preprocessing pipeline for features.
139 Creates separate pipelines for numerical and categorical features:
140 - Numerical: integer-to-float conversion + imputation + standardization
141 - Categorical: imputation + one-hot encoding
143 The integer-to-float conversion ensures MLflow schema compatibility by
144 converting integer columns to float64, preventing schema enforcement
145 errors when missing values are present at inference time.
147 Returns
148 -------
149 ColumnTransformer
150 Preprocessing pipeline
151 """
152 self._log("Building preprocessing pipeline")
154 # Build transformers
155 transformers = []
157 if self.numerical_features:
158 numerical_transformer = Pipeline(
159 steps=[
160 ("imputer", SimpleImputer(strategy="median")),
161 ("scaler", StandardScaler()),
162 ]
163 )
164 transformers.append(("num", numerical_transformer, self.numerical_features))
166 if self.categorical_features:
167 categorical_transformer = Pipeline(
168 steps=[
169 (
170 "imputer",
171 SimpleImputer(strategy="constant", fill_value="missing"),
172 ),
173 (
174 "onehot",
175 OneHotEncoder(handle_unknown="ignore", sparse_output=False),
176 ),
177 ]
178 )
179 transformers.append((
180 "cat",
181 categorical_transformer,
182 self.categorical_features,
183 ))
185 if not transformers:
186 raise ValueError("No features to process") # noqa: TRY003
188 return ColumnTransformer(transformers=transformers, remainder="drop")
190 def _create_pipeline(self) -> Pipeline:
191 """Create full pipeline with preprocessor and model.
193 Returns
194 -------
195 Pipeline
196 Full pipeline with preprocessor and model steps
197 """
198 # Build preprocessing pipeline if not already built
199 preprocessor = self._build_preprocessing_pipeline()
201 # Load model configuration
202 model_obj = Model(self.model_name, self.config)
203 model_instance = model_obj.instantiate()
205 # Create full pipeline
206 pipeline = Pipeline(
207 steps=[("preprocessor", preprocessor), ("model", model_instance)]
208 )
210 # self.model = model_instance
211 self.full_pipeline = pipeline
213 return pipeline
215 def _calculate_metrics(
216 self,
217 y_true: pd.Series,
218 y_pred: np.ndarray,
219 metrics_list: list[str] | None = None,
220 ) -> dict[str, float]:
221 """Calculate evaluation metrics.
223 Parameters
224 ----------
225 y_true : pd.Series
226 True target values
227 y_pred : np.ndarray
228 Predicted values
229 metrics_list : Optional[List[str]]
230 List of metric names to calculate. If None, uses experiment metrics. The first metrics will be used to select the best model if multiple specified.
232 Returns
233 -------
234 Dict[str, float]
235 Dictionary of metric names to values
236 """
237 if metrics_list is None:
238 metrics_list = self.experiment.spec.metrics or ["r2", "mse"]
240 results = {}
242 for idx, metric_name in enumerate(metrics_list):
243 metric_name_lower = metric_name.lower()
244 metric_value = None
245 direction_flag = 1
247 try:
248 if metric_name_lower.startswith("r2"):
249 metric_value = r2_score(y_true, y_pred)
250 results["r2_score"] = metric_value
251 direction_flag = 1 # higher is better
252 elif metric_name_lower in ["mse", "mean_squared_error"]:
253 metric_value = mean_squared_error(y_true, y_pred)
254 results["mean_squared_error"] = metric_value
255 direction_flag = -1 # lower is better
256 elif metric_name_lower in ["rmse", "root_mean_squared_error"]:
257 metric_value = np.sqrt(mean_squared_error(y_true, y_pred))
258 results["root_mean_squared_error"] = metric_value
259 direction_flag = -1 # lower is better
260 elif metric_name_lower in ["mae", "mean_absolute_error"]:
261 metric_value = mean_absolute_error(y_true, y_pred)
262 results["mean_absolute_error"] = metric_value
263 direction_flag = -1 # lower is better
264 elif metric_name_lower.startswith("accuracy"):
265 metric_value = accuracy_score(y_true, y_pred)
266 results["accuracy_score"] = metric_value
267 direction_flag = 1 # higher is better
268 elif metric_name_lower.startswith("precision"):
269 metric_value = precision_score(y_true, y_pred, average="weighted")
270 results["precision_score"] = metric_value
271 direction_flag = 1 # higher is better
272 elif metric_name_lower.startswith("recall"):
273 metric_value = recall_score(y_true, y_pred, average="weighted")
274 results["recall_score"] = metric_value
275 direction_flag = 1 # higher is better
276 elif metric_name_lower.startswith("f1"):
277 metric_value = f1_score(y_true, y_pred, average="weighted")
278 results["f1_score"] = metric_value
279 direction_flag = 1 # higher is better
280 else:
281 self._log(f"Warning: Unknown metric '{metric_name}'")
283 # Always set selection_score from the *first* metric in the list
284 if idx == 0 and metric_value is not None:
285 results["selection_score"] = direction_flag * metric_value
287 except Exception as e:
288 self._log(f"Error calculating metric '{metric_name}': {e}")
290 return results
292 def calculate_feature_weights(self) -> pd.DataFrame:
293 """Calculate feature importances or coefficients.
295 Returns
296 -------
297 pd.DataFrame
298 DataFrame with feature names and their weights/importances
299 """
300 # Get the fitted pipeline from the search
301 best_pipeline = self.best_pipeline
303 # Find the column transformer and ridge model in the pipeline
304 # Assumes the pipeline steps are: ('preprocessor', ...), ('ridge', ...)
305 preprocessor = best_pipeline.named_steps["preprocessor"]
306 model = best_pipeline.named_steps["model"]
308 # Get feature names from preprocessor
309 feature_names = preprocessor.get_feature_names_out()
310 # Drop prefix before '__' in feature names, if present
311 feature_names = [
312 name.split("__", 1)[-1] if "__" in name else name for name in feature_names
313 ]
315 # Extract weights/importances from model
316 weights = None
317 weight_type = None
319 if hasattr(model, "coef_"):
320 # Linear models (coefficients)
321 weights = model.coef_
322 weight_type = "coefficient"
323 elif hasattr(model, "feature_importances_"):
324 # Tree-based models (feature importances)
325 weights = model.feature_importances_
326 weight_type = "importance"
327 else:
328 self._log(
329 "Warning: Model does not have coefficients or feature importances"
330 )
331 return pd.DataFrame()
333 # Handle multi-dimensional coefficients (e.g., multi-class classification)
334 if len(weights.shape) > 1:
335 weights = np.abs(weights).mean(axis=0)
337 # Create DataFrame
338 df = pd.DataFrame({
339 "feature": feature_names[: len(weights)],
340 "weight": weights,
341 "type": weight_type,
342 })
343 df = df.sort_values(by="weight", ascending=False, key=abs)
345 # self._log("Top 5 features by weight:")
346 # for _idx, row in df.head(5).iterrows():
347 # self._log(f" {row['feature']}: {row['weight']:.6f}")
349 return df
351 def plot_feature_weights(self) -> matplotlib.figure.Figure | None:
352 """
353 Plot feature importances or coefficients from the best estimator (tuned search).
354 Reuses the calculate_feature_weights method.
355 Returns a matplotlib Figure object for notebook display or file export.
357 Returns
358 -------
359 matplotlib.figure.Figure or None
360 The figure, or None if no feature weights are available.
361 """
363 try:
364 # Reuse calculate_feature_weights to get DataFrame
365 weights_df = self.calculate_feature_weights()
366 except Exception as e:
367 self._log(f"Error calculating feature weights: {e}")
368 return None
370 if (
371 weights_df is None
372 or weights_df.empty
373 or "feature" not in weights_df
374 or "weight" not in weights_df
375 ):
376 self._log("No feature weights available to plot.")
377 return None
379 # Sort features by absolute weight/importance (descending)
380 weights_df_sorted = weights_df.reindex(
381 weights_df.weight.abs().sort_values(ascending=False).index
382 )
384 # Plot
385 fig, ax = plt.subplots(
386 figsize=(10, max(4, min(0.5 * len(weights_df_sorted), 16)))
387 )
388 ax.barh(
389 weights_df_sorted["feature"],
390 weights_df_sorted["weight"],
391 color="tab:blue",
392 alpha=0.85,
393 )
395 ylabel = (
396 "Coefficient"
397 if "coefficient" in weights_df_sorted.columns.to_numpy().tolist()
398 or "coefficient" in weights_df_sorted.get("type", "")
399 else "Importance"
400 )
401 ax.set_xlabel(ylabel)
402 ax.set_title(f"Feature {ylabel}s for Best Estimator ({self.model_name})")
403 ax.invert_yaxis()
404 ax.grid(axis="x", linestyle="--", alpha=0.5)
405 fig.tight_layout()
407 return fig
409 def plot_cv_mean_score(self, figsize=(10, 10)) -> matplotlib.figure.Figure | None:
410 """
411 Plot mean test score across CV splits with confidence intervals (std error).
412 Returns
413 -------
414 matplotlib.figure.Figure or None
415 The figure, or None if no CV results are available.
416 """
418 if not self.cv_results_:
419 self._log("No cross-validation results found to plot.")
420 return None
422 cv_results = self.cv_results_
423 if "mean_test_score" not in cv_results or "std_test_score" not in cv_results:
424 self._log("CV results are missing 'mean_test_score' or 'std_test_score'.")
425 return None
427 mean_scores = np.array(cv_results["mean_test_score"])
428 std_scores = np.array(cv_results["std_test_score"])
429 x_labels = [f"{v}" for v in cv_results["params"]]
431 # Plot mean test score with confidence intervals
432 fig, ax = plt.subplots(figsize=figsize)
433 ax.tick_params(axis="x", rotation=90) # Make x labels vertical
434 ax.errorbar(
435 x_labels,
436 mean_scores,
437 yerr=std_scores,
438 fmt="o",
439 capsize=4,
440 label="Mean Test Score ±1 Std",
441 )
442 ax.set_xlabel("Hyperparameter Combination Index")
443 ax.set_ylabel("Mean Test Score")
444 ax.set_title(f"Hyperparameter Search Mean Test Score ±1 Std ({self.model_name})")
445 ax.legend()
446 ax.grid(True)
447 fig.tight_layout()
449 return fig
451 def plot_confusion_matrix(
452 self,
453 y_true: pd.Series | np.ndarray,
454 y_pred: pd.Series | np.ndarray,
455 title_prefix: str | None = None,
456 ) -> matplotlib.figure.Figure | None:
457 """Create confusion matrix plot for classification.
459 Parameters
460 ----------
461 y_true : pd.Series | np.ndarray
462 True target values
463 y_pred : pd.Series | np.ndarray
464 Predicted values
465 title_prefix : str, optional
466 Prefix to add to plot title (e.g., 'holdout data', 'test data', etc.)
467 If None, no prefix is added.
469 Returns
470 -------
471 matplotlib.figure.Figure | None
472 Matplotlib figure object with confusion matrix plot, or None if error occurs
473 """
474 try:
475 # Convert to numpy arrays if needed
476 y_true = y_true.to_numpy() if isinstance(y_true, pd.Series) else y_true
477 y_pred = y_pred.to_numpy() if isinstance(y_pred, pd.Series) else y_pred
479 # Get unique class labels
480 classes = sorted(set(np.unique(y_true)) | set(np.unique(y_pred)))
482 # Calculate confusion matrix with explicit labels
483 cm = confusion_matrix(y_true, y_pred, labels=classes)
485 # Create plot
486 fig, ax = plt.subplots(figsize=(8, 6))
488 # Use seaborn if available for better visualization
489 if sns is not None:
490 sns.heatmap(
491 cm,
492 annot=True,
493 fmt="d",
494 cmap="Blues",
495 ax=ax,
496 cbar_kws={"label": "Count"},
497 xticklabels=classes,
498 yticklabels=classes,
499 )
500 else:
501 # Fallback to matplotlib if seaborn not available
502 im = ax.imshow(cm, interpolation="nearest", cmap="Blues")
503 ax.figure.colorbar(im, ax=ax)
505 # Add text annotations
506 thresh = cm.max() / 2.0
507 for i in range(cm.shape[0]):
508 for j in range(cm.shape[1]):
509 ax.text(
510 j,
511 i,
512 format(cm[i, j], "d"),
513 ha="center",
514 va="center",
515 color="white" if cm[i, j] > thresh else "black",
516 )
518 # Build title suffix
519 title_suffix = f" ({title_prefix})" if title_prefix else ""
521 ax.set(
522 xlabel="Predicted Label",
523 ylabel="True Label",
524 title=f"Confusion Matrix{title_suffix}",
525 )
526 ax.set_xticks(np.arange(len(classes)))
527 ax.set_yticks(np.arange(len(classes)))
528 ax.set_xticklabels(classes)
529 ax.set_yticklabels(classes)
531 fig.tight_layout()
533 except Exception as e:
534 self._log(f"Error creating confusion matrix plot: {e}")
535 return None
536 else:
537 return fig
539 def plot_regression(
540 self,
541 y_true: pd.Series | np.ndarray,
542 y_pred: pd.Series | np.ndarray,
543 title_prefix: str | None = None,
544 ) -> matplotlib.figure.Figure | None:
545 """Create actual vs predicted and residuals vs predicted plots for regression.
547 Parameters
548 ----------
549 y_true : pd.Series | np.ndarray
550 True target values
551 y_pred : pd.Series | np.ndarray
552 Predicted values
553 title_prefix : str, optional
554 Prefix to add to plot titles (e.g., 'holdout data', 'test data', etc.)
555 If None, no prefix is added.
557 Returns
558 -------
559 matplotlib.figure.Figure | None
560 Matplotlib figure object with two subplots, or None if error occurs
561 """
562 try:
563 y_true = y_true.to_numpy() if isinstance(y_true, pd.Series) else y_true
564 y_pred = y_pred.to_numpy() if isinstance(y_pred, pd.Series) else y_pred
565 # Calculate residuals
566 residuals = y_true - y_pred
568 # Create figure with two subplots
569 fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
571 # Build title suffix
572 title_suffix = f" ({title_prefix})" if title_prefix else ""
574 # Plot 1: Actual vs Predicted
575 ax1.scatter(y_true, y_pred, alpha=0.6, s=20)
577 # Add diagonal line (perfect prediction)
578 min_val = min(y_true.min(), y_pred.min())
579 max_val = max(y_true.max(), y_pred.max())
580 ax1.plot(
581 [min_val, max_val],
582 [min_val, max_val],
583 "r--",
584 lw=2,
585 label="Perfect prediction",
586 )
588 ax1.set_xlabel("Actual Values")
589 ax1.set_ylabel("Predicted Values")
590 ax1.set_title(f"Actual vs Predicted{title_suffix}")
591 ax1.legend()
592 ax1.grid(True, alpha=0.3)
594 # Plot 2: Residuals vs Predicted
595 ax2.scatter(y_pred, residuals, alpha=0.6, s=20)
597 # Add horizontal line at y=0
598 ax2.axhline(y=0, color="r", linestyle="--", lw=2, label="Zero residual")
600 ax2.set_xlabel("Predicted Values")
601 ax2.set_ylabel("Residuals")
602 ax2.set_title(f"Residuals vs Predicted{title_suffix}")
603 ax2.legend()
604 ax2.grid(True, alpha=0.3)
606 fig.tight_layout()
608 except Exception as e:
609 self._log(f"Error creating regression plots: {e}")
610 return None
611 else:
612 return fig
614 def plot_distribution(
615 self,
616 y_true: pd.Series | np.ndarray,
617 y_pred: pd.Series | np.ndarray,
618 title_prefix: str | None = None,
619 ) -> matplotlib.figure.Figure | None:
620 """
621 Plot the distribution (KDE) of actual and predicted values.
623 Parameters
624 ----------
625 y_true : pd.Series | np.ndarray
626 True target values
627 y_pred : pd.Series | np.ndarray
628 Predicted values
629 title_prefix : str or None
630 Optional prefix for the plot title.
632 Returns
633 -------
634 matplotlib.figure.Figure or None
635 The figure, or None if an error occurs.
636 """
637 try:
638 y_true = y_true.to_numpy() if isinstance(y_true, pd.Series) else y_true
639 y_pred = y_pred.to_numpy() if isinstance(y_pred, pd.Series) else y_pred
641 fig, ax = plt.subplots(figsize=(8, 6))
642 title_suffix = f" ({title_prefix})" if title_prefix else ""
644 # Plot the KDE for true values
645 pd.Series(y_true).plot(kind="kde", ax=ax, label="Actual", color="tab:blue")
646 # Plot the KDE for predicted values
647 pd.Series(y_pred).plot(
648 kind="kde", ax=ax, label="Predicted", color="tab:orange"
649 )
651 ax.set_xlabel("Value")
652 ax.set_ylabel("Density")
653 ax.set_title(
654 f"Distribution (KDE) of Actual and Predicted Values{title_suffix}"
655 )
656 ax.legend()
657 ax.grid(True, alpha=0.3)
658 fig.tight_layout()
660 except Exception as e:
661 self._log(f"Error plotting distributions: {e}")
662 return None
663 else:
664 return fig
666 def fit_and_evaluate(self) -> float:
667 """Train the model and evaluate it.
669 This method encapsulates the full model lifecycle:
670 - Creates the pipeline
671 - If tuning is configured, performs cross-validation hyperparameter search
672 - Fits the model (or best tuned model) on training data
673 - Evaluates on holdout sets
674 - Calculates feature weights
676 Returns
677 -------
678 float
679 R2 score of the model - used to select the best model
680 """
681 self._log("Starting model training and evaluation")
683 # Load model configuration to check for tuning
684 model_obj = Model(self.model_name, self.config)
685 tuning_config = model_obj.spec.tuning
687 # Create initial pipeline (integer-to-float conversion is handled in the pipeline)
688 self._create_pipeline()
690 # Check if tuning is configured (non-empty dict)
691 if tuning_config and len(tuning_config) > 0:
692 self._log("Tuning configuration detected - performing cross-validation")
693 self.is_cross_validation = True
695 # Extract tuning parameters
696 method = tuning_config.get("method", "grid_search")
697 inner_cv = tuning_config.get("inner_cv", 5)
698 scoring = tuning_config.get("scoring", "neg_mean_squared_error")
699 param_grid = tuning_config.get("param_grid", {})
701 # Build parameter grid with proper prefix for pipeline
702 # Parameters need to be prefixed with "model__" since model is a step in the pipeline
703 pipeline_param_grid = {}
705 # Handle logspace if present as standalone key (overwrites alpha)
706 if "logspace" in param_grid:
707 logspace_config = param_grid["logspace"]
708 start = logspace_config.get("start", -2)
709 stop = logspace_config.get("stop", 2)
710 num = logspace_config.get("num", 50)
711 logspace_values = np.logspace(start, stop, num).tolist()
712 # Apply logspace to alpha (most common use case)
713 pipeline_param_grid["model__alpha"] = logspace_values
714 self._log(
715 f"Using logspace for alpha: {num} values from 10^{start} to 10^{stop}"
716 )
718 # Process regular parameters
719 for param_name, param_values in param_grid.items():
720 if (
721 param_name == "logspace"
722 ): # Skip standalone logspace key (already handled)
723 continue
725 # Check if this parameter has nested logspace structure
726 if isinstance(param_values, dict) and "logspace" in param_values:
727 logspace_config = param_values["logspace"]
728 start = logspace_config.get("start", -2)
729 stop = logspace_config.get("stop", 2)
730 num = logspace_config.get("num", 50)
731 logspace_values = np.logspace(start, stop, num).tolist()
732 pipeline_param_grid[f"model__{param_name}"] = logspace_values
733 self._log(
734 f"Using logspace for {param_name}: {num} values from 10^{start} to 10^{stop}"
735 )
736 elif isinstance(param_values, list):
737 # Regular list of values
738 pipeline_param_grid[f"model__{param_name}"] = param_values
740 # Choose search method
741 if method == "random_search":
742 n_iter = tuning_config.get("n_iter", 10)
743 search = RandomizedSearchCV(
744 self.full_pipeline,
745 pipeline_param_grid,
746 cv=inner_cv,
747 scoring=scoring,
748 n_iter=n_iter,
749 verbose=1 if self.verbose else 0,
750 n_jobs=-1,
751 )
752 else: # Default to grid_search
753 search = GridSearchCV(
754 self.full_pipeline,
755 pipeline_param_grid,
756 cv=inner_cv,
757 scoring=scoring,
758 verbose=1 if self.verbose else 0,
759 n_jobs=-1,
760 )
762 # Perform cross-validation search
763 self._log(f"Performing {method} with {inner_cv}-fold CV")
764 search.fit(self.X_train, self.y_train)
766 # Store CV results for future analysis
767 self.cv_results_ = search.cv_results_
768 self.cv_best_params_ = search.best_params_
769 self.cv_best_score_ = search.best_score_
771 # Set best_pipeline as the best tuned pipeline
772 self.best_pipeline = search.best_estimator_
773 self._log(f"Best parameters: {search.best_params_}")
774 self._log(f"Best CV score: {search.best_score_:.6f}")
776 # Fit best_pipeline on full training data
777 self._log("Fitting best pipeline on full training data")
778 self.best_pipeline.fit(self.X_train, self.y_train)
780 else:
781 # No tuning - standard fit
782 self.best_pipeline = self.full_pipeline
783 self._log("Training model")
784 self.best_pipeline.fit(self.X_train, self.y_train)
786 self._log("Model training completed")
788 # start evaluation part of the model (if holdout exists)
789 metrics = {}
790 selection_score = None
791 # evaluate on train set
792 if self.X_train is not None and self.y_train is not None:
793 self._log("Evaluating model on test set")
794 y_pred_train = self.best_pipeline.predict(self.X_train)
795 train_metrics = self._calculate_metrics(
796 self.y_train, y_pred_train, self.experiment.spec.metrics
797 )
798 for metric, value in train_metrics.items():
799 self._log(f" train_{metric}: {value:.6f}")
800 if "selection_score" in train_metrics:
801 selection_score = train_metrics["selection_score"]
802 self._log(f"Train Selection score: {selection_score:.6f}")
803 train_metrics = {f"train_{k}": v for k, v in train_metrics.items()}
804 metrics.update(train_metrics)
805 # evaluate on holdout set
806 if self.X_holdout is not None and self.y_holdout is not None:
807 self._log("Evaluating model on holdout set")
808 y_pred_holdout = self.best_pipeline.predict(self.X_holdout)
809 holdout_metrics = self._calculate_metrics(
810 self.y_holdout, y_pred_holdout, self.experiment.spec.metrics
811 )
812 for metric, value in holdout_metrics.items():
813 self._log(f" holdout_{metric}: {value:.6f}")
814 if "selection_score" in holdout_metrics:
815 selection_score = holdout_metrics["selection_score"]
816 self._log(f"Holdout Selection score: {selection_score:.6f}")
817 holdout_metrics = {f"holdout_{k}": v for k, v in holdout_metrics.items()}
818 metrics.update(holdout_metrics)
820 self.metrics = metrics
821 self.selection_score = selection_score
822 return self.selection_score
824 def predict(self, X): # noqa: N803
825 """
826 Make predictions using the best trained pipeline.
828 Parameters
829 ----------
830 X : pd.DataFrame or np.ndarray
831 Input features for prediction.
833 Returns
834 -------
835 np.ndarray
836 Predicted outputs.
837 """
838 if not hasattr(self, "best_pipeline") or self.best_pipeline is None:
839 raise RuntimeError("The model pipeline is not trained yet.") # noqa: TRY003
840 return self.best_pipeline.predict(X)
842 def mlflow_store(self) -> None:
843 """
844 Store the model in MLflow.
845 Assuming to run with active run
846 """
847 mlflow.log_param("model_name", self.model_name)
848 mlflow.log_param("model_pipeline", str(self.best_pipeline))
850 # Get the parameters from the 'model' step of the pipeline and log each as a parameter
851 model_params = self.best_pipeline.named_steps['model'].get_params()
852 for param_name, param_value in model_params.items():
853 mlflow.log_param(f"model__{param_name}", param_value)
855 # log metrics
856 for k, v in self.metrics.items():
857 mlflow.log_metric(k, v)
859 # log model with a signature
860 # Get a DataFrame containing the first 10 rows of X_train
861 sample_df = self.X_train.head(10).copy()
862 signature = mlflow.models.infer_signature(sample_df, self.best_pipeline.predict(sample_df))
863 mlflow.sklearn.log_model(self.best_pipeline, name="model", signature=signature)
865 # create temp directory to create artifacts and load them into MlFlow
866 with tempfile.TemporaryDirectory() as temp_dir:
867 # log cross validation results
868 if self.is_cross_validation and self.cv_results_ is not None:
869 mlflow.log_param("cross_validation",True)
870 fig = self.plot_cv_mean_score()
871 if fig:
872 fig.savefig(os.path.join(temp_dir, "cv_mean_score.png"))
873 mlflow.log_artifact(os.path.join(temp_dir, "cv_mean_score.png"))
874 else:
875 mlflow.log_param("cross_validation",False)
877 # log plot_feature_weights
878 weights_df = self.calculate_feature_weights()
879 if weights_df is not None and not weights_df.empty:
880 weights_df.to_csv(os.path.join(temp_dir, "feature_weights.csv"), index=False)
881 mlflow.log_artifact(os.path.join(temp_dir, "feature_weights.csv"))
882 fig = self.plot_feature_weights()
883 if fig:
884 fig.savefig(os.path.join(temp_dir, "feature_weights.png"))
885 mlflow.log_artifact(os.path.join(temp_dir, "feature_weights.png"))
887 # Check if it is regression or classification experiment
888 # INSERT_YOUR_CODE
889 # Precompute predictions for train and holdout to avoid redundant calls
890 y_pred_train = self.predict(self.X_train)
891 y_pred_holdout = self.predict(self.X_holdout) if self.X_holdout is not None else None
893 task_type = self.experiment.get_type()
894 if task_type == "classification":
895 # log plot_confusion_matrix
896 fig = self.plot_confusion_matrix(self.y_train, y_pred_train, "train")
897 if fig:
898 fig.savefig(os.path.join(temp_dir, "train_confusion_matrix.png"))
899 mlflow.log_artifact(os.path.join(temp_dir, "train_confusion_matrix.png"))
900 # log plot_confusion_matrix
901 if y_pred_holdout is not None:
902 fig = self.plot_confusion_matrix(self.y_holdout, y_pred_holdout, "holdout")
903 if fig:
904 fig.savefig(os.path.join(temp_dir, "holdout_confusion_matrix.png"))
905 mlflow.log_artifact(os.path.join(temp_dir, "holdout_confusion_matrix.png"))
906 else:
907 # log plot_regression
908 fig = self.plot_regression(self.y_train, y_pred_train, "train")
909 if fig:
910 fig.savefig(os.path.join(temp_dir, "train_regression.png"))
911 mlflow.log_artifact(os.path.join(temp_dir, "train_regression.png"))
913 # log plot_distribution
914 fig = self.plot_distribution(self.y_train, y_pred_train, "train")
915 if fig:
916 fig.savefig(os.path.join(temp_dir, "train_distribution.png"))
917 mlflow.log_artifact(os.path.join(temp_dir, "train_distribution.png"))
919 if y_pred_holdout is not None:
920 # log plot_regression
921 fig = self.plot_regression(self.y_holdout, y_pred_holdout, "holdout")
922 if fig:
923 fig.savefig(os.path.join(temp_dir, "holdout_regression.png"))
924 mlflow.log_artifact(os.path.join(temp_dir, "holdout_regression.png"))
926 # log plot_distribution
927 fig = self.plot_distribution(self.y_holdout, y_pred_holdout, "holdout")
928 if fig:
929 fig.savefig(os.path.join(temp_dir, "holdout_distribution.png"))
930 mlflow.log_artifact(os.path.join(temp_dir, "holdout_distribution.png"))
932 plt.close('all')
934class Runner:
935 """Execute ML experiments with full lifecycle management.
937 Responsibilities:
938 - Load datasets
939 - Build preprocessing pipelines
940 - Handle train/validation/test splits
941 - Coordinate ModelRunner instances for each model
943 Parameters
944 ----------
945 experiment : Experiment
946 Experiment specification to run
947 verbose : bool, optional
948 Whether to print progress information, by default True
949 """
951 def __init__(self, experiment: Experiment, verbose: bool = True):
952 self.experiment = experiment
953 self.verbose = verbose
954 self.config = experiment.config
956 # Data storage
957 self.dataset: pd.DataFrame | None = None
958 self.X_train: pd.DataFrame | None = None
959 self.y_train: pd.Series | None = None
960 self.X_holdout: pd.DataFrame | None = None
961 self.y_holdout: pd.Series | None = None
963 # Feature tracking
964 self.numerical_features: list[str] = []
965 self.categorical_features: list[str] = []
966 self.all_features: list[str] = []
968 # ModelRunner instances
969 self.model_runners: list[ModelRunner] = []
971 self.best_model: ModelRunner | None = None
972 self.best_model_score: str | None = None
974 def get_models(self) -> list[ModelRunner]:
975 return self.model_runners
977 def get_best_model(self) -> ModelRunner | None:
978 return self.best_model
980 def get_best_model_score(self) -> float:
981 return self.best_model_score
983 def _log(self, message: str) -> None:
984 """Print message if verbose is enabled."""
985 if self.verbose:
986 print(f"[Runner] {message}") # noqa: T201
988 def _load_dataset(self) -> pd.DataFrame:
989 """Load dataset specified in experiment.
991 Also infers experiment type from target column if not specified in configuration.
993 Returns
994 -------
995 pd.DataFrame
996 Loaded dataset
997 """
998 self._log(f"Loading dataset: {self.experiment.spec.dataset}")
999 dataset = Dataset(self.experiment.spec.dataset, self.config)
1000 self.dataset = dataset.read_pandas()
1001 self._log(
1002 f"Dataset loaded: {len(self.dataset)} rows, {len(self.dataset.columns)} columns"
1003 )
1005 # Infer experiment type from target if not specified
1006 if self.experiment.spec.type is None:
1007 inferred_type = self.experiment.infer_type_from_dataset(self.dataset)
1008 self._log(
1009 f"Inferred experiment type: {inferred_type} (from target column dtype)"
1010 )
1011 else:
1012 self._log(f"Using configured experiment type: {self.experiment.spec.type}")
1014 return self.dataset
1016 def _infer_column_types(self, df: pd.DataFrame) -> tuple[list[str], list[str]]:
1017 """Infer which columns are numerical vs categorical based on dtypes.
1019 Parameters
1020 ----------
1021 df : pd.DataFrame
1022 DataFrame to infer from
1023 columns : List[str]
1024 Columns to classify
1026 Returns
1027 -------
1028 Tuple[List[str], List[str]]
1029 Lists of (numerical_features, categorical_features)
1030 """
1031 # Fill numerical and categorical list from spec (if available)
1032 feature_name = getattr(self.experiment.spec, "features", None)
1033 if feature_name and hasattr(self, "config"):
1034 try:
1035 feature = Feature(feature_name, self.config)
1036 # Add columns from 'numerical' and 'categorical' in spec if present
1037 numerical = (
1038 list(feature.spec.numerical)
1039 if hasattr(feature.spec, "numerical") and feature.spec.numerical
1040 else []
1041 )
1042 categorical = (
1043 list(feature.spec.categorical)
1044 if hasattr(feature.spec, "categorical") and feature.spec.categorical
1045 else []
1046 )
1047 column = (
1048 list(feature.spec.column)
1049 if hasattr(feature.spec, "column") and feature.spec.column
1050 else []
1051 )
1052 except Exception:
1053 # Fall back to type inference if Feature cannot be constructed
1054 numerical = []
1055 categorical = []
1056 column = []
1057 else:
1058 numerical = []
1059 categorical = []
1060 column = []
1062 # handle a special keywork __all__ in the column list
1063 if "__all__" in column:
1064 column = [
1065 col
1066 for col in df.columns
1067 if col not in numerical and col not in categorical
1068 ]
1070 # just to be in safe side, remove "target" column"
1071 target = self.experiment.spec.target
1072 column = [col for col in column if col != target]
1073 numerical = [col for col in numerical if col not in column]
1074 categorical = [col for col in categorical if col not in column]
1076 # Check that all specified numerical and categorical columns exist in df
1077 missing_numerical = [col for col in numerical if col not in df.columns]
1078 missing_categorical = [col for col in categorical if col not in df.columns]
1079 if missing_numerical:
1080 self._log(
1081 f"Warning: The following numerical feature(s) are not in the dataset: {missing_numerical}"
1082 )
1083 numerical = [col for col in numerical if col in df.columns]
1084 if missing_categorical:
1085 self._log(
1086 f"Warning: The following categorical feature(s) are not in the dataset: {missing_categorical}"
1087 )
1088 categorical = [col for col in categorical if col in df.columns]
1090 # Process columns from the 'column' parameter (excluding those already specified)
1091 for col in column:
1092 if col not in df.columns:
1093 self._log(f"Warning: Column '{col}' not found in dataset")
1094 continue
1095 if col in numerical or col in categorical:
1096 continue
1098 dtype = df[col].dtype
1099 if pd.api.types.is_numeric_dtype(dtype):
1100 numerical.append(col)
1101 else:
1102 categorical.append(col)
1104 return numerical, categorical
1106 def _drop_outliers(self) -> None:
1107 """Detect outliers in numeric features and target, and mark them in the dataset.
1109 This method:
1110 - Detects outliers using Z-score method (default threshold: 3.0 standard deviations)
1111 - Sets 'is_outlier' boolean column in the dataset dataframe
1112 - Logs the number of records before and after outlier removal
1113 - Only processes numeric features and target column
1115 The 'is_outlier' column is always created, even if outlier removal is disabled
1116 (drop_outliers is None or 0.0). Rows with is_outlier==True will be excluded
1117 from train/holdout splits in _split_data() and data_load().
1118 """
1119 if self.dataset is None:
1120 raise RuntimeError("Dataset must be loaded before dropping outliers") # noqa: TRY003
1122 # Get drop_outliers threshold from experiment config
1123 threshold = self.experiment.spec.drop_outliers
1125 # Initialize is_outlier column to False
1126 self.dataset["is_outlier"] = False
1128 # If threshold is None or 0.0, outlier detection is disabled
1129 if threshold is None or threshold == 0.0:
1130 self._log("Outlier detection disabled (drop_outliers is None or 0.0)")
1131 return
1133 self._log(f"Detecting outliers using Z-score threshold: {threshold}")
1135 # Get numeric columns (features + target)
1136 target = self.experiment.spec.target
1137 numeric_columns = []
1139 # Add numeric features
1140 if hasattr(self, "numerical_features") and self.numerical_features:
1141 numeric_columns.extend(self.numerical_features)
1143 # Add target if it's numeric
1144 if target in self.dataset.columns and pd.api.types.is_numeric_dtype(
1145 self.dataset[target].dtype
1146 ):
1147 numeric_columns.append(target)
1149 # Remove duplicates and ensure columns exist
1150 numeric_columns = list({
1151 col for col in numeric_columns if col in self.dataset.columns
1152 })
1154 if not numeric_columns:
1155 self._log("No numeric columns found for outlier detection")
1156 return
1158 self._log(f"Checking outliers in columns: {numeric_columns}")
1160 # Record initial number of rows
1161 initial_count = len(self.dataset)
1163 # Calculate Z-scores for each numeric column
1164 outlier_mask = pd.Series(False, index=self.dataset.index)
1166 for col in numeric_columns:
1167 # Calculate Z-scores: (value - mean) / std
1168 col_mean = self.dataset[col].mean()
1169 col_std = self.dataset[col].std()
1171 # Skip if std is 0 (constant column) or NaN
1172 if col_std == 0 or pd.isna(col_std):
1173 self._log(f" {col}: skipped (constant or NaN std)")
1174 continue
1176 z_scores = np.abs((self.dataset[col] - col_mean) / col_std)
1177 # Mark as outlier if Z-score exceeds threshold
1178 col_outliers = z_scores > threshold
1179 outlier_mask |= col_outliers
1181 # Log column-specific outlier counts
1182 col_outlier_count = col_outliers.sum()
1183 if col_outlier_count > 0:
1184 self._log(f" {col}: {col_outlier_count} outliers detected")
1186 # Set is_outlier column
1187 self.dataset["is_outlier"] = outlier_mask
1189 # Count total outliers
1190 total_outliers = outlier_mask.sum()
1191 final_count = initial_count - total_outliers
1193 # Log results
1194 self._log("Outlier detection complete:")
1195 self._log(f" Initial records: {initial_count}")
1196 self._log(
1197 f" Outliers detected: {total_outliers} ({total_outliers / initial_count * 100:.2f}%)"
1198 )
1199 self._log(f" Records after removal: {final_count}")
1201 def _prepare_features(self) -> tuple[list[str], list[str]]:
1202 """Prepare and classify features from feature set.
1204 Collects features from feature specifications and classifies them
1205 as numerical or categorical. If feature specifications don't include
1206 types, infers from DataFrame dtypes.
1208 Returns
1209 -------
1210 Tuple[List[str], List[str]]
1211 Lists of (numerical_features, categorical_features)
1212 """
1213 if self.dataset is None:
1214 raise RuntimeError("Dataset must be loaded before preparing features") # noqa: TRY003
1216 self._log("Preparing features")
1218 # Collect features from feature specifications
1219 numerical_cols, categorical_cols = self._infer_column_types(self.dataset)
1221 # If no explicit type specifications, infer from all feature columns
1222 all_feature_cols = numerical_cols + categorical_cols
1224 # Remove target column if present
1225 target = self.experiment.spec.target
1226 # Check if target accidentally appears in feature columns (should not occur)
1227 if (
1228 target in all_feature_cols
1229 or target in numerical_cols
1230 or target in categorical_cols
1231 ):
1232 raise ValueError( # noqa: TRY003
1233 f"Target column '{target}' was included in the feature set"
1234 )
1236 if not (numerical_cols or categorical_cols):
1237 raise ValueError( # noqa: TRY003
1238 "No features were found in the feature set. At least one feature must be specified."
1239 )
1241 self.numerical_features = numerical_cols
1242 self.categorical_features = categorical_cols
1243 self.all_features = numerical_cols + categorical_cols
1245 # Convert column types: categorical to string, numerical to numeric
1246 self.dataset = self._convert_column_types(self.dataset)
1248 self._log(f"Numerical features ({len(numerical_cols)}): {numerical_cols}")
1249 self._log(f"Categorical features ({len(categorical_cols)}): {categorical_cols}")
1251 return numerical_cols, categorical_cols
1253 def _convert_column_types(self, df: pd.DataFrame) -> pd.DataFrame:
1254 """Convert categorical columns to string and numerical columns to numeric types.
1256 Parameters
1257 ----------
1258 df : pd.DataFrame
1259 DataFrame to convert column types for
1261 Returns
1262 -------
1263 pd.DataFrame
1264 DataFrame with converted column types
1265 """
1266 df = df.copy()
1268 # Convert categorical columns to string
1269 if self.categorical_features:
1270 for col in self.categorical_features:
1271 if col in df.columns:
1272 df[col] = df[col].astype(str)
1274 # Convert numerical columns to numeric (handles mixed types)
1275 if self.numerical_features:
1276 for col in self.numerical_features:
1277 if col in df.columns:
1278 df[col] = pd.to_numeric(df[col], errors="coerce")
1280 return df
1282 def _get_grouping_values(self, df: pd.DataFrame, name: str) -> pd.Series:
1283 """Get grouping values from column or index.
1285 Parameters
1286 ----------
1287 df : pd.DataFrame
1288 DataFrame to extract grouping values from
1289 name : str
1290 Name of the column or index to use for grouping
1292 Returns
1293 -------
1294 pd.Series
1295 Series containing grouping values aligned with df.index
1297 Raises
1298 ------
1299 ValueError
1300 If the name is not found in columns or index name
1301 """
1302 if name in df.columns:
1303 return df[name]
1304 if df.index.name is not None and df.index.name == name:
1305 # Convert index to Series with same index as df for alignment
1306 return pd.Series(df.index, index=df.index, name=name)
1307 index_info = (
1308 f"Index name: {df.index.name}" if df.index.name else "Index has no name"
1309 )
1310 raise ValueError( # noqa: TRY003
1311 f"Grouping column/index '{name}' not found in dataset. "
1312 f"Available columns: {list(df.columns)}, "
1313 f"{index_info}"
1314 )
1316 def _split_data(self) -> None:
1317 """Split data into training and hold-out sets.
1319 Respects:
1320 - hold_out configuration for creating a separate hold-out set
1321 - do_not_split_by for grouped splitting (prevents data leakage)
1323 Adds a binary column 'is_holdout' to self.dataset indicating hold-out samples.
1324 """
1325 if self.dataset is None:
1326 raise RuntimeError("Dataset must be loaded before splitting") # noqa: TRY003
1328 self._log("Splitting data")
1330 target = self.experiment.spec.target
1331 if target not in self.dataset.columns:
1332 raise ValueError(f"Target column '{target}' not found in dataset") # noqa: TRY003
1334 df = self.dataset.copy()
1335 # Initialize hold-out indicator column
1336 df["is_holdout"] = False
1338 # Filter out outliers before splitting
1339 if "is_outlier" in df.columns:
1340 outlier_mask = df["is_outlier"]
1341 outliers_count = outlier_mask.sum()
1342 if outliers_count > 0:
1343 self._log(
1344 f"Excluding {outliers_count} outlier rows from train/holdout splits"
1345 )
1346 df = df[~outlier_mask].copy()
1347 else:
1348 self._log("No outliers to exclude from splits")
1349 else:
1350 # If is_outlier column doesn't exist, create it as False
1351 df["is_outlier"] = False
1353 # Separate features and target
1354 y = df[target]
1355 X = df[self.all_features] # noqa: N806
1357 # Handle hold-out set if specified
1358 hold_out_config = self.experiment.spec.hold_out
1359 hold_out_fraction = (
1360 hold_out_config.get("fraction", 0.0) if hold_out_config else 0.0
1361 )
1363 do_not_split_by = self.experiment.spec.do_not_split_by
1365 if hold_out_fraction > 0:
1366 self._log(f"Creating hold-out set: {hold_out_fraction:.1%}")
1367 hold_out_random_state = hold_out_config.get("random_state", 42)
1369 # Check if we need grouped splitting for hold-out
1370 if do_not_split_by:
1371 # Create composite grouping key from multiple columns/index if needed
1372 # GroupShuffleSplit needs a single array-like, so we combine multiple columns
1373 if len(do_not_split_by) == 1:
1374 groups = self._get_grouping_values(df, do_not_split_by[0])
1375 else:
1376 # Create a composite group identifier from multiple columns/index
1377 # Get each grouping value (column or index)
1378 group_series = [
1379 self._get_grouping_values(df, name) for name in do_not_split_by
1380 ]
1381 # Combine into DataFrame for easy joining
1382 group_df = pd.concat(group_series, axis=1)
1383 # Convert to string tuples to create unique group identifiers
1384 groups = group_df.apply(
1385 lambda row: "_".join(str(val) for val in row), axis=1
1386 )
1387 splitter = GroupShuffleSplit(
1388 n_splits=1,
1389 test_size=hold_out_fraction,
1390 random_state=hold_out_random_state,
1391 )
1392 train_idx, holdout_idx = next(splitter.split(X, y, groups=groups))
1393 # Split data using positional indices
1394 self.X_train = X.iloc[train_idx].copy()
1395 self.y_train = y.iloc[train_idx].copy()
1396 self.X_holdout = X.iloc[holdout_idx].copy()
1397 self.y_holdout = y.iloc[holdout_idx].copy()
1398 # Mark hold-out samples using original DataFrame indices
1399 holdout_original_idx = X.index[holdout_idx]
1400 df.loc[holdout_original_idx, "is_holdout"] = True
1401 else:
1402 # Use train_test_split for non-grouped splitting
1403 X_train_temp, X_holdout_temp, y_train_temp, y_holdout_temp = ( # noqa: N806
1404 train_test_split(
1405 X,
1406 y,
1407 test_size=hold_out_fraction,
1408 random_state=hold_out_random_state,
1409 shuffle=True,
1410 )
1411 )
1412 # Assign split data
1413 self.X_train = X_train_temp.copy()
1414 self.y_train = y_train_temp.copy()
1415 self.X_holdout = X_holdout_temp.copy()
1416 self.y_holdout = y_holdout_temp.copy()
1417 # Mark hold-out samples in the dataset
1418 df.loc[X_holdout_temp.index, "is_holdout"] = True
1420 self._log(f"Hold-out set size: {len(self.X_holdout)}")
1421 else:
1422 # No hold-out split, all data is training
1423 self.X_train = X.copy()
1424 self.y_train = y.copy()
1425 self.X_holdout = None
1426 self.y_holdout = None
1428 # Update self.dataset with the hold-out indicator (but keep original rows including outliers)
1429 # Only update is_holdout column, don't filter out outliers from self.dataset
1430 # Use direct assignment to avoid deprecation warning
1431 # Initialize is_holdout column to False
1432 self.dataset["is_holdout"] = False
1433 # Then update based on df
1434 for idx in df.index:
1435 self.dataset.loc[idx, "is_holdout"] = df.loc[idx, "is_holdout"]
1437 self._log(f"Train set size: {len(self.X_train)}")
1438 if self.X_holdout is not None:
1439 self._log(f"Hold-out set size: {len(self.X_holdout)}")
1441 def data_save(self, filepath: str | None = None) -> str:
1442 """Save self.dataset to a parquet file.
1444 Parameters
1445 ----------
1446 filepath : str, optional
1447 Path where to save the parquet file. If None, generates a default
1448 filename based on experiment name.
1450 Returns
1451 -------
1452 str
1453 Path to the saved parquet file
1455 Raises
1456 ------
1457 RuntimeError
1458 If dataset is None (not loaded yet)
1459 """
1460 if self.dataset is None:
1461 raise RuntimeError("Dataset must be loaded before saving") # noqa: TRY003
1463 if filepath is None:
1464 # Generate default filename based on experiment name
1465 filepath = f"{self.experiment.name}_dataset.parquet"
1467 self._log(f"Saving dataset to {filepath}")
1468 self.dataset.to_parquet(filepath, index=False)
1469 self._log(
1470 f"Dataset saved successfully: {len(self.dataset)} rows, {len(self.dataset.columns)} columns"
1471 )
1473 return filepath
1475 def data_load(self, filepath: str) -> None:
1476 """Load dataset from parquet file, prepare features, and split based on is_holdout column.
1478 This method loads a previously saved dataset (from data_save), prepares features,
1479 and recreates the train/holdout splits based on the is_holdout column. The result
1480 is equivalent to calling data_preparation + data_save, then loading the saved file.
1482 Parameters
1483 ----------
1484 filepath : str
1485 Path to the parquet file to load
1487 Raises
1488 ------
1489 FileNotFoundError
1490 If the parquet file doesn't exist
1491 ValueError
1492 If required columns (target, is_holdout) are missing from the loaded dataset
1493 """
1494 self._log(f"Loading dataset from {filepath}")
1496 # Load dataset from parquet
1497 if not Path(filepath).exists():
1498 raise FileNotFoundError(f"Parquet file not found: {filepath}") # noqa: TRY003
1500 self.dataset = pd.read_parquet(filepath)
1501 self._log(
1502 f"Dataset loaded: {len(self.dataset)} rows, {len(self.dataset.columns)} columns"
1503 )
1505 # Verify is_holdout column exists
1506 if "is_holdout" not in self.dataset.columns:
1507 raise ValueError( # noqa: TRY003
1508 "Dataset must contain 'is_holdout' column. "
1509 "Please ensure the file was saved using data_save() after data_preparation()."
1510 )
1512 # Step 1: Prepare features (classify as numerical/categorical)
1513 self._prepare_features()
1515 # Step 2: Split data based on is_holdout column
1516 target = self.experiment.spec.target
1517 if target not in self.dataset.columns:
1518 raise ValueError(f"Target column '{target}' not found in dataset") # noqa: TRY003
1520 df = self.dataset.copy()
1522 # Filter out outliers before splitting
1523 if "is_outlier" in df.columns:
1524 outlier_mask = df["is_outlier"]
1525 outliers_count = outlier_mask.sum()
1526 if outliers_count > 0:
1527 self._log(
1528 f"Excluding {outliers_count} outlier rows from train/holdout splits"
1529 )
1530 df = df[~outlier_mask].copy()
1531 else:
1532 self._log("No outliers to exclude from splits")
1533 else:
1534 # If is_outlier column doesn't exist, create it as False
1535 df["is_outlier"] = False
1537 # Separate features and target
1538 y = df[target]
1539 X = df[self.all_features] # noqa: N806
1541 # Split based on is_holdout column
1542 is_holdout = df["is_holdout"]
1544 # Training set: rows where is_holdout is False
1545 train_mask = ~is_holdout
1546 self.X_train = X.loc[train_mask].copy()
1547 self.y_train = y.loc[train_mask].copy()
1549 # Holdout set: rows where is_holdout is True
1550 holdout_mask = is_holdout
1551 if holdout_mask.any():
1552 self.X_holdout = X.loc[holdout_mask].copy()
1553 self.y_holdout = y.loc[holdout_mask].copy()
1554 self._log(f"Hold-out set size: {len(self.X_holdout)}")
1555 else:
1556 self.X_holdout = None
1557 self.y_holdout = None
1558 self._log("No hold-out set found in loaded dataset")
1560 self._log(f"Train set size: {len(self.X_train)}")
1561 if self.X_holdout is not None:
1562 self._log(f"Hold-out set size: {len(self.X_holdout)}")
1564 def data_preparation(self) -> None:
1565 """Prepare data for the experiment.
1567 This method performs the initial data preparation steps:
1568 1. Load dataset
1569 2. Prepare features (classify as numerical/categorical)
1570 3. Split data into train/validation/test sets
1571 """
1572 self._log("Preparing data for experiment")
1574 # Step 1: Load dataset
1575 self._load_dataset()
1577 # Step 2: Prepare features (classify as numerical/categorical)
1578 self._prepare_features()
1580 # Step 3: Drop outliers
1581 self._drop_outliers()
1583 # Step 4: Split data
1584 self._split_data()
1586 def get_config(self) -> dict[str, Any]:
1587 """Return dictionary similar to original YAML config with all fields populated.
1589 Includes inferred feature types (numerical/categorical) and experiment type
1590 (regression/classification). The returned dictionary has the same structure
1591 as the original YAML configuration file.
1593 Returns
1594 -------
1595 Dict[str, Any]
1596 Dictionary with same structure as YAML config, with inferred fields populated.
1597 Includes:
1598 - All original config sections (datasets, features, models, experiments, etc.)
1599 - Inferred experiment type in experiments section
1600 - Inferred feature types (numerical/categorical) in features section
1602 Raises
1603 ------
1604 RuntimeError
1605 If features have not been prepared yet (need to call prepare_features() first)
1606 """
1607 # Start with a deep copy of the original config
1608 config_dict = deepcopy(self.config.to_dict())
1610 # Get experiment name
1611 exp_name = self.experiment.name
1613 # Update experiment section with inferred type
1614 if "experiments" in config_dict and exp_name in config_dict["experiments"]:
1615 exp_dict = config_dict["experiments"][exp_name]
1616 if isinstance(exp_dict, dict):
1617 # Add inferred type if available (either configured or inferred)
1618 inferred_type = self.experiment.get_type()
1619 if inferred_type is not None:
1620 exp_dict["type"] = inferred_type
1622 # Update features section with inferred types
1623 feature_name = self.experiment.spec.features
1624 if "features" in config_dict and feature_name in config_dict["features"]:
1625 feature_dict = config_dict["features"][feature_name]
1626 if isinstance(feature_dict, dict):
1627 # Check if features have been prepared
1628 if not hasattr(self, "numerical_features") or not hasattr(
1629 self, "categorical_features"
1630 ):
1631 raise RuntimeError( # noqa: TRY003
1632 "Features have not been prepared yet. "
1633 "Call prepare_features() or run() first before calling get_config()."
1634 )
1636 # If features have been prepared, update with inferred types
1637 # Remove 'columns' field if present (replaced by numerical/categorical)
1638 if "columns" in feature_dict:
1639 # Check if it was using __all__ - if so, we've expanded it
1640 # Remove columns key as we've inferred types from it
1641 feature_dict.pop("columns")
1642 # If it was just __all__, we can note that it was inferred
1643 # Otherwise, we've inferred types from the columns list
1645 # Update with inferred types
1646 feature_dict["numerical"] = self.numerical_features.copy()
1647 feature_dict["categorical"] = self.categorical_features.copy()
1649 # Ensure lists are not empty (remove empty lists if both are empty)
1650 # But keep them if they exist - empty lists are valid
1652 return config_dict
1654 def run(self, skip_mlflow: bool = False) -> dict[str, Any]:
1655 """Execute the complete experiment workflow.
1657 Parameters
1658 ----------
1659 skip_mlflow : bool, optional
1660 Whether to skip MLflow logging, by default False
1661 If True, the experiment will be run without logging to MLflow
1662 Returns
1663 -------
1664 Dict[str, Any]
1665 Results dictionary containing metrics and artifacts for all models
1666 """
1667 self._log(f"Starting experiment: {self.experiment.name}")
1669 # Check if data is already prepared
1670 if self.X_train is None:
1671 self.data_preparation()
1673 # Step 4: Create ModelRunner instances for each model
1674 if not self.numerical_features and not self.categorical_features:
1675 raise RuntimeError("No features prepared before creating ModelRunners") # noqa: TRY003
1676 if self.X_train is None or self.y_train is None:
1677 raise RuntimeError("Data must be split before creating ModelRunners") # noqa: TRY003
1679 self.model_runners = []
1680 for model_name in self.experiment.spec.models:
1681 model_runner = ModelRunner(
1682 model_name=model_name,
1683 numerical_features=self.numerical_features,
1684 categorical_features=self.categorical_features,
1685 experiment=self.experiment,
1686 X_train=self.X_train,
1687 y_train=self.y_train,
1688 X_holdout=self.X_holdout,
1689 y_holdout=self.y_holdout,
1690 verbose=self.verbose,
1691 )
1692 self.model_runners.append(model_runner)
1694 # Step 5: Run each ModelRunner
1695 results = {}
1696 results["models"] = {}
1697 results["best_model"] = None
1698 results["best_model_score"] = None
1700 best_model_score = float("-inf")
1701 for model_runner in self.model_runners:
1702 self._log(f"\n{'=' * 60}")
1703 self._log(f"Processing model: {model_runner.model_name}")
1704 self._log(f"{'=' * 60}")
1706 model_score = model_runner.fit_and_evaluate()
1707 self._log(f"Model {model_runner.model_name} score: {model_score:.6f}")
1708 if model_score > best_model_score:
1709 best_model_score = model_score
1710 self.best_model = model_runner
1711 results["models"][model_runner.model_name] = model_runner.metrics
1713 self.best_model_score = best_model_score
1714 results["best_model"] = self.best_model.model_name
1715 results["best_model_score"] = self.best_model_score
1717 self._log(f"\n{'=' * 60}")
1718 self._log("Experiment completed successfully!")
1719 self._log(f"Best model: {self.best_model.model_name}")
1720 self._log(f"Best model score: {self.best_model_score:.6f}")
1721 self._log(f"{'=' * 60}")
1723 if not skip_mlflow:
1724 self.mlflow_store()
1726 return results
1728 def mlflow_store(self) -> None:
1729 """
1730 Store the experiment in MLflow.
1731 """
1733 mlflow_config = MlflowConf(self.config)
1734 if not mlflow_config.is_enabled():
1735 return
1737 mlflow_experiment_name = mlflow_config.get_name(self.experiment.name)
1738 if mlflow_config.get_type() == "databricks":
1739 mlflow.set_tracking_uri("databricks")
1740 if not mlflow_experiment_name.startswith("/Shared/"):
1741 mlflow_experiment_name = f"/Shared/{mlflow_experiment_name}"
1742 mlflow.set_experiment(mlflow_experiment_name)
1743 else:
1744 # Only set tracking URI if not already set (e.g., by test fixtures)
1745 current_uri = mlflow.get_tracking_uri()
1746 if current_uri is None or current_uri == "":
1747 mlflow.set_tracking_uri("sqlite:///mlflow.db")
1748 mlflow.set_experiment(mlflow_experiment_name)
1750 # Get experiment description
1751 experiment_description = self.experiment.spec.description or ""
1753 with mlflow.start_run(description=experiment_description) as parent_run:
1754 parent_run_id = parent_run.info.run_id
1756 models = self.get_models()
1757 if len(models) > 1:
1758 # log multiple models as nested run
1759 for model in models:
1761 if self.best_model == model :
1762 name = f"BEST-{model.model_name}"
1763 else:
1764 name = model.model_name
1766 with mlflow.start_run(run_name=name, nested=True) as child_run:
1767 model.mlflow_store()
1769 # log the best model as parent run
1770 self.best_model.mlflow_store()
1772 # add addtional parameters to the parent run
1773 mlflow.log_param("experiment_name", self.experiment.name)
1774 mlflow.log_param("experiment_type", self.experiment.get_type())
1775 if experiment_description:
1776 mlflow.log_param("experiment_description", experiment_description)
1778 # add tags to the parent run
1779 mlflow.set_tags(mlflow_config.get_tags())
1781 # store experiment config into temp directory and upload it to mlflow
1782 with tempfile.TemporaryDirectory() as temp_dir:
1783 config_path = os.path.join(temp_dir, "experiment_config.yaml")
1784 config_dict = self.get_config()
1785 with open(config_path, 'w') as f:
1786 yaml.dump(config_dict, f, sort_keys=False, default_flow_style=False)
1787 mlflow.log_artifact(config_path)
1789__all__ = ["Runner", "ModelRunner"]