Cla. Decision Tree
Train a Decision Tree classification model.
Classification Decision Tree
Processing
This brick trains a Decision Tree classification model using your provided data. Think of a Decision Tree as a flowchart: it asks a sequence of yes/no questions based on your data features (e.g., "Is the value greater than 5?") to categorize items into specific classes.
The brick handles the complete machine learning workflow:
- Data Splitting: It automatically separates your data into training and testing sets to ensure the model is evaluated fairly.
- Training: It builds the tree logic to best separate your data classes.
- Optimization (Optional): It can automatically run multiple trials to find the best settings (hyperparameters) for your specific dataset.
- Evaluation: It calculates key performance metrics (like Accuracy and F1-Score) so you know how well the model performs.
Inputs
- X
- The features data (independent variables). This contains the information used to make predictions. It must contain only numerical or boolean data.
- y
- The target data (labels/dependent variable). This is the column you want the model to learn how to predict.
Inputs Types
| Input | Types |
|---|---|
X |
DataFrame |
y |
DataSeries, NDArray, List, |
You can check the list of supported types here: Available Type Hints.
Outputs
- Model
- The trained Decision Tree Classifier object (Scikit-Learn estimator). This object can be saved or used in subsequent bricks to make predictions on new data.
- Model Classes
- A reference table mapping the model's internal numerical indices to the actual class labels (names).
- SHAP
- The SHAP TreeExplainer object. This is used to generate explanations for why the model made specific predictions (requires the "SHAP Explainer" option to be enabled).
- Metrics
- A summary of the model's performance on the test set, including Accuracy, Precision, Recall, F1-Score, and ROC AUC.
- CV Metrics
- The results from Cross-Validation (if enabled). This shows the mean and standard deviation of performance metrics across multiple folds, providing a more robust estimate of model reliability.
- Features Importance
- A list of your input columns ranked by how useful they were in making predictions.
- Prediction Set
- A DataFrame containing the test data, the model's predictions, the actual true values, and the prediction probabilities. Useful for auditing errors.
- HPO Trials
- A history of all the settings tried during Hyperparameter Optimization (if enabled).
- HPO Best
- A dictionary containing the best combination of parameters found during optimization.
The Prediction Set output contains the following specific data fields:
- feature_{name}: The original input feature columns.
- proba (or proba_{class}): The probability score assigned by the model for the prediction.
- y_true: The actual label/class from your original data.
- y_pred: The label/class predicted by the model.
- is_false_prediction: A boolean (True/False) indicating if the model got it wrong.
The Features Importance output contains:
- feature: The name of the column.
- importance: A score indicating how much this feature contributed to the model's decisions (higher is better).
Outputs Types
| Output | Types |
|---|---|
Model |
Any |
Model Classes |
DataFrame |
SHAP |
Any |
Metrics |
DataFrame, Dict |
CV Metrics |
DataFrame |
Features Importance |
DataFrame |
Prediction Set |
DataFrame |
HPO Trials |
DataFrame |
HPO Best |
Dict |
You can check the list of supported types here: Available Type Hints.
Options
The Classification Decision Tree brick contains some changeable options:
- Criterion
- The function to measure the quality of a split.
- Gini: Measures impurity (how mixed the classes are). Generally faster.
- Entropy: Measures information gain. Can create more balanced trees but is slightly slower.
- Split Strategy
- How the model chooses to split a node.
- Best: Always chooses the optimal split (standard behavior).
- Random: Chooses the best random split. Can help prevent overfitting and is faster.
- Max Depth (0 = Unlimited)
- Controls how "deep" the tree can grow. A deeper tree can capture more complex patterns but is more likely to memorize the data (overfitting). 0 means the depth is unlimited.
- Feature Sampling
- The number of features to consider when looking for the best split.
- Automatic 30% / 50%: Uses a fixed percentage of total features.
- Square root: Uses the square root of the total feature count (Standard for classification).
- Logarithm: Uses the base-2 logarithm of the feature count.
- None: Considers every single feature at every split.
- Min Samples per Leaf
- The minimum number of data points required to be at a leaf node (the end of a branch). Increasing this smooths the model (reduces overfitting).
- Auto Split Data
- If enabled, automatically splits your data into Training and Testing sets. If disabled, you can manually set the percentage.
- Test/Validation Set %
- The percentage of data to hold back for testing/validation (ignored if Auto Split is on).
- Retrain On Full Data
- If enabled, the model will be re-trained on 100% of the provided data after the metrics are calculated. Use this when preparing a final model for production.
- Average Strategy
- How metrics are calculated for multiclass problems.
- auto: Automatically selects based on class balance.
- binary: Only for two-class problems.
- micro: Calculate metrics globally by counting total true positives, false negatives and false positives.
- macro: Calculate metrics for each label, and find their unweighted mean.
- weighted: Calculate metrics for each label, and find their average weighted by support (the number of true instances for each label).
- Enable Cross-Validation
- If enabled, splits the data into "folds" (groups) and trains/tests multiple times to ensure the scores are reliable.
- Hyperparameter Optim.
- If enabled, the brick will run multiple experiments to automatically find the best settings (like Depth and Criterion) for your data.
- Optimization Metric
- The score the optimization process tries to maximize (e.g., maximize "F1 Score" to balance precision and recall).
- Optimization Method
- The algorithm used to search for the best parameters. "Tree-structured Parzen" is generally the most efficient.
- Tree-structured Parzen: A Bayesian optimization method that models good vs bad parameter regions using probability distributions and prioritizes sampling where success is statistically more likely.
- Gaussian Process: ses a probabilistic regression model (Gaussian Process) to estimate performance uncertainty and selects new trials using acquisition functions.
- CMA-ES: An evolutionary strategy that adapts the covariance matrix of a multivariate normal distribution to efficiently search complex, non-linear, non-convex spaces.
- Random Sobol Search: Uses low-discrepancy quasi-random sequences to ensure uniform coverage of the parameter space, avoiding clustering and gaps.
- Random Search: Uniform random sampling of parameter configurations without learning or feedback between iterations.
- Optimization Iterations
- The number of different parameter combinations to try during optimization.
- Positive Label (Binary Only)
- Explicitly state which class is "Positive" (e.g., "1", "True", "Churned"). Essential for correct Precision/Recall scores in binary classification.
- SHAP Explainer
- If enabled, generates a SHAP explainer object to visualize feature contribution.
- SHAP Sampler
- If enabled, uses a smaller, representative sample of your data (via K-Means or Random Sampling) to calculate SHAP values. This significantly speeds up the explanation process for large datasets with minimal loss in accuracy.
- SHAP Feature Perturbation
- Defines how the explainer handles correlated features when calculating importance.
- Interventional: Breaks the dependencies between features. This is better for understanding the causal effect of a specific feature change, assuming you could change it independently.
- Tree Path Dependent: Follows the natural correlations in the data (how features usually change together). This is better for understanding the model's behavior on realistic data distributions.
- Number of Jobs
- Controls how many CPU cores are used for training and cross-validation.
- 1: Sequential processing (slower, uses less system resources).
- 2, 4, 8: Specific number of cores.
- All: Uses all available cores on the machine for maximum speed.
- Random State
- A seed number (integer) that ensures reproducibility. Using the same seed with the same data ensures you get the exact same split and model results every time you run the brick.
- Brick Caching
- If enabled, the brick saves the results to a temporary cache based on the inputs. If you run the flow again with the exact same inputs and options, it loads the results immediately instead of retraining, saving significant time.
- Verbose Logging
- If enabled, prints detailed progress updates, metric calculations, and optimization steps to the execution logs. Useful for debugging or monitoring long-running tasks.
import logging
import warnings
import shap
import json
import xxhash
import hashlib
import tempfile
import sklearn
import scipy
import joblib
import numpy as np
import pandas as pd
import polars as pl
from pathlib import Path
from scipy import sparse
from optuna.samplers import (
TPESampler,
RandomSampler,
GPSampler,
CmaEsSampler,
QMCSampler,
)
import optuna
from optuna import Study
from optuna.trial import FrozenTrial
from optuna.pruners import HyperbandPruner
from optuna import create_study
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split, cross_validate, StratifiedKFold
from sklearn.metrics import (
accuracy_score,
precision_score,
recall_score,
f1_score,
roc_auc_score,
make_scorer,
)
from dataclasses import dataclass
from datetime import datetime
from coded_flows.types import (
Union,
Dict,
List,
Tuple,
NDArray,
DataFrame,
DataSeries,
Any,
Tuple,
)
from coded_flows.utils import CodedFlowsLogger
logger = CodedFlowsLogger(name="Cla. Decision Tree", level=logging.INFO)
optuna.logging.set_verbosity(optuna.logging.ERROR)
warnings.filterwarnings("ignore", category=optuna.exceptions.ExperimentalWarning)
DataType = Union[
pd.DataFrame, pl.DataFrame, np.ndarray, sparse.spmatrix, pd.Series, pl.Series
]
@dataclass
class _DatasetFingerprint:
"""Lightweight fingerprint of a dataset."""
hash: str
shape: tuple
computed_at: str
data_type: str
method: str
class _UniversalDatasetHasher:
"""
High-performance dataset hasher optimizing for zero-copy operations
and native backend execution (C/Rust).
"""
def __init__(
self,
data_size: int,
method: str = "auto",
sample_size: int = 100000,
verbose: bool = False,
):
self.method = method
self.sample_size = sample_size
self.data_size = data_size
self.verbose = verbose
def hash_data(self, data: DataType) -> _DatasetFingerprint:
"""
Main entry point: hash any supported data format.
Auto-detects format and applies optimal strategy.
"""
if isinstance(data, pd.DataFrame):
return self._hash_pandas(data)
elif isinstance(data, pl.DataFrame):
return self._hash_polars(data)
elif isinstance(data, pd.Series):
return self._hash_pandas_series(data)
elif isinstance(data, pl.Series):
return self._hash_polars_series(data)
elif isinstance(data, np.ndarray):
return self._hash_numpy(data)
elif sparse.issparse(data):
return self._hash_sparse(data)
else:
raise TypeError(f"Unsupported data type: {type(data)}")
def _hash_pandas(self, df: pd.DataFrame) -> _DatasetFingerprint:
"""
Optimized Pandas hashing using pd.util.hash_pandas_object.
Avoids object-to-string conversion overhead.
"""
method = self._determine_method(self.data_size, self.method)
self.verbose and logger.info(
f"Hashing Pandas: {self.data_size:,} rows - {method}"
)
target_df = df
if method == "sampled":
target_df = self._get_pandas_sample(df)
hasher = xxhash.xxh128()
self._hash_schema(
hasher,
{
"columns": df.columns.tolist(),
"dtypes": {k: str(v) for (k, v) in df.dtypes.items()},
"shape": df.shape,
},
)
try:
row_hashes = pd.util.hash_pandas_object(target_df, index=False)
hasher.update(memoryview(row_hashes.values))
except Exception as e:
self.verbose and logger.warning(
f"Fast hash failed, falling back to slow hash: {e}"
)
self._hash_pandas_fallback(hasher, target_df)
return _DatasetFingerprint(
hash=hasher.hexdigest(),
shape=df.shape,
computed_at=datetime.now().isoformat(),
data_type="pandas",
method=method,
)
def _get_pandas_sample(self, df: pd.DataFrame) -> pd.DataFrame:
"""Deterministic slicing for sampling (Zero randomness)."""
if self.data_size <= self.sample_size:
return df
chunk = self.sample_size // 3
head = df.iloc[:chunk]
mid_idx = self.data_size // 2
mid = df.iloc[mid_idx : mid_idx + chunk]
tail = df.iloc[-chunk:]
return pd.concat([head, mid, tail])
def _hash_pandas_fallback(self, hasher, df: pd.DataFrame):
"""Legacy fallback for complex object types."""
for col in df.columns:
val = df[col].astype(str).values
hasher.update(val.astype(np.bytes_).tobytes())
def _hash_polars(self, df: pl.DataFrame) -> _DatasetFingerprint:
"""
Optimized Polars hashing using native Rust execution.
"""
method = self._determine_method(self.data_size, self.method)
self.verbose and logger.info(
f"Hashing Polars: {self.data_size:,} rows - {method}"
)
target_df = df
if method == "sampled" and self.data_size > self.sample_size:
indices = self._get_sample_indices(self.data_size, self.sample_size)
target_df = df.gather(indices)
hasher = xxhash.xxh128()
self._hash_schema(
hasher,
{
"columns": df.columns,
"dtypes": [str(t) for t in df.dtypes],
"shape": df.shape,
},
)
row_hashes = target_df.hash_rows()
hasher.update(memoryview(row_hashes.to_numpy()))
return _DatasetFingerprint(
hash=hasher.hexdigest(),
shape=df.shape,
computed_at=datetime.now().isoformat(),
data_type="polars",
method=method,
)
def _hash_pandas_series(self, series: pd.Series) -> _DatasetFingerprint:
"""Hash Pandas Series using the fastest vectorized method."""
self.verbose and logger.info(f"Hashing Pandas Series: {self.data_size:,} rows")
hasher = xxhash.xxh128()
self._hash_schema(
hasher,
{
"name": series.name if series.name else "None",
"dtype": str(series.dtype),
"shape": series.shape,
},
)
try:
row_hashes = pd.util.hash_pandas_object(series, index=False)
hasher.update(memoryview(row_hashes.values))
except Exception as e:
self.verbose and logger.warning(f"Series hash failed, falling back: {e}")
hasher.update(memoryview(series.astype(str).values.tobytes()))
return _DatasetFingerprint(
hash=hasher.hexdigest(),
shape=series.shape,
computed_at=datetime.now().isoformat(),
data_type="pandas_series",
method="full",
)
def _hash_polars_series(self, series: pl.Series) -> _DatasetFingerprint:
"""Hash Polars Series using native Polars expressions."""
self.verbose and logger.info(f"Hashing Polars Series: {self.data_size:,} rows")
hasher = xxhash.xxh128()
self._hash_schema(
hasher,
{"name": series.name, "dtype": str(series.dtype), "shape": series.shape},
)
try:
row_hashes = series.hash()
hasher.update(memoryview(row_hashes.to_numpy()))
except Exception as e:
self.verbose and logger.warning(
f"Polars series native hash failed. Falling back."
)
hasher.update(str(series.to_list()).encode())
return _DatasetFingerprint(
hash=hasher.hexdigest(),
shape=series.shape,
computed_at=datetime.now().isoformat(),
data_type="polars_series",
method="full",
)
def _hash_numpy(self, arr: np.ndarray) -> _DatasetFingerprint:
"""
Optimized NumPy hashing using Buffer Protocol (Zero-Copy).
"""
hasher = xxhash.xxh128()
self._hash_schema(
hasher,
{"shape": arr.shape, "dtype": str(arr.dtype), "strides": arr.strides},
)
if arr.flags["C_CONTIGUOUS"] or arr.flags["F_CONTIGUOUS"]:
hasher.update(memoryview(arr))
else:
hasher.update(memoryview(np.ascontiguousarray(arr)))
return _DatasetFingerprint(
hash=hasher.hexdigest(),
shape=arr.shape,
computed_at=datetime.now().isoformat(),
data_type="numpy",
method="full",
)
def _hash_sparse(self, matrix: sparse.spmatrix) -> _DatasetFingerprint:
"""
Optimized sparse hashing. Hashes underlying data arrays directly.
"""
if not (sparse.isspmatrix_csr(matrix) or sparse.isspmatrix_csc(matrix)):
matrix = matrix.tocsr()
hasher = xxhash.xxh128()
self._hash_schema(
hasher, {"shape": matrix.shape, "format": matrix.format, "nnz": matrix.nnz}
)
hasher.update(memoryview(matrix.data))
hasher.update(memoryview(matrix.indices))
hasher.update(memoryview(matrix.indptr))
return _DatasetFingerprint(
hash=hasher.hexdigest(),
shape=matrix.shape,
computed_at=datetime.now().isoformat(),
data_type=f"sparse_{matrix.format}",
method="sparse",
)
def _determine_method(self, rows: int, requested: str) -> str:
if requested != "auto":
return requested
if rows < 5000000:
return "full"
return "sampled"
def _hash_schema(self, hasher, schema: Dict[str, Any]):
"""Compact schema hashing."""
hasher.update(
json.dumps(schema, sort_keys=True, separators=(",", ":")).encode()
)
def _get_sample_indices(self, total_rows: int, sample_size: int) -> list:
"""Calculate indices for sampling without generating full range lists."""
chunk = sample_size // 3
indices = list(range(min(chunk, total_rows)))
mid_start = max(0, total_rows // 2 - chunk // 2)
mid_end = min(mid_start + chunk, total_rows)
indices.extend(range(mid_start, mid_end))
last_start = max(0, total_rows - chunk)
indices.extend(range(last_start, total_rows))
return sorted(list(set(indices)))
def _validate_numerical_data(data):
"""
Validates if the input data (NumPy array, Pandas DataFrame/Series,
Polars DataFrame/Series, or SciPy sparse matrix) contains only
numerical (integer, float) or boolean values.
Args:
data: The input data structure to check.
Raises:
TypeError: If the input data contains non-numerical and non-boolean types.
ValueError: If the input data is of an unsupported type.
"""
if sparse.issparse(data):
if not (
np.issubdtype(data.dtype, np.number) or np.issubdtype(data.dtype, np.bool_)
):
raise TypeError(
f"Sparse matrix contains unsupported data type: {data.dtype}. Only numerical or boolean types are allowed."
)
return
elif isinstance(data, np.ndarray):
if not (
np.issubdtype(data.dtype, np.number) or np.issubdtype(data.dtype, np.bool_)
):
raise TypeError(
f"NumPy array contains unsupported data type: {data.dtype}. Only numerical or boolean types are allowed."
)
return
elif isinstance(data, (pd.DataFrame, pd.Series)):
d_types = data.dtypes.apply(lambda x: x.kind)
non_numerical_mask = ~d_types.isin(["i", "f", "b"])
if non_numerical_mask.any():
non_numerical_columns = (
data.columns[non_numerical_mask].tolist()
if isinstance(data, pd.DataFrame)
else [data.name]
)
raise TypeError(
f"Pandas {('DataFrame' if isinstance(data, pd.DataFrame) else 'Series')} contains non-numerical/boolean data. Offending column(s) and types: {data.dtypes[non_numerical_mask].to_dict()}"
)
return
elif isinstance(data, (pl.DataFrame, pl.Series)):
pl_numerical_types = [
pl.Int8,
pl.Int16,
pl.Int32,
pl.Int64,
pl.UInt8,
pl.UInt16,
pl.UInt32,
pl.UInt64,
pl.Float32,
pl.Float64,
pl.Boolean,
]
if isinstance(data, pl.DataFrame):
for col, dtype in data.schema.items():
if dtype not in pl_numerical_types:
raise TypeError(
f"Polars DataFrame column '{col}' has unsupported data type: {dtype}. Only numerical or boolean types are allowed."
)
elif isinstance(data, pl.Series):
if data.dtype not in pl_numerical_types:
raise TypeError(
f"Polars Series has unsupported data type: {data.dtype}. Only numerical or boolean types are allowed."
)
return
else:
raise ValueError(
f"Unsupported data type provided: {type(data)}. Function supports NumPy, Pandas, Polars, and SciPy sparse matrices."
)
def _normalize_hpo_df(df):
df = df.copy()
param_cols = [c for c in df.columns if c.startswith("params_")]
df[param_cols] = df[param_cols].astype("string[pyarrow]")
return df
def _get_shape_and_sparsity(X: Any) -> Tuple[int, int, float, bool]:
"""
Efficiently extracts shape and estimates sparsity without converting
the entire dataset to numpy.
"""
(n_samples, n_features) = (0, 0)
is_sparse = False
sparsity = 0.0
if hasattr(X, "nnz") and hasattr(X, "shape"):
(n_samples, n_features) = X.shape
is_sparse = True
sparsity = 1.0 - X.nnz / (n_samples * n_features)
return (n_samples, n_features, sparsity, is_sparse)
if hasattr(X, "height") and hasattr(X, "width"):
(n_samples, n_features) = (X.height, X.width)
return (n_samples, n_features, 0.0, False)
if hasattr(X, "shape") and hasattr(X, "iloc"):
(n_samples, n_features) = X.shape
return (n_samples, n_features, 0.0, False)
if isinstance(X, list):
X = np.array(X)
if hasattr(X, "shape"):
(n_samples, n_features) = X.shape
return (n_samples, n_features, 0.0, False)
raise ValueError("Unsupported data type")
def _choose_class_weight(y):
from collections import Counter
counts = np.array(list(Counter(y).values()))
n_min = counts.min()
r = counts.max() / n_min
if r <= 1.5:
return None
if r <= 10:
return "balanced"
if n_min >= 50:
return "balanced_subsample"
return "balanced"
def _smart_split(
n_samples,
X,
y,
*,
random_state=42,
shuffle=True,
stratify=None,
fixed_test_split=None,
verbose=True,
):
"""
Parameters
----------
n_samples : int
Number of samples in the dataset (len(X) or len(y))
X : array-like
Features
y : array-like
Target
random_state : int
shuffle : bool
stratify : array-like or None
For stratified splitting (recommended for classification)
Returns
-------
If return_val=True → X_train, X_val, X_test, y_train, y_val, y_test
If return_val=False → X_train, X_test, y_train, y_test
"""
if fixed_test_split:
test_ratio = fixed_test_split
val_ratio = fixed_test_split
elif n_samples <= 1000:
test_ratio = 0.2
val_ratio = 0.1
elif n_samples < 10000:
test_ratio = 0.15
val_ratio = 0.15
elif n_samples < 100000:
test_ratio = 0.1
val_ratio = 0.1
elif n_samples < 1000000:
test_ratio = 0.05
val_ratio = 0.05
else:
test_ratio = 0.01
val_ratio = 0.01
(X_train, X_test, y_train, y_test) = train_test_split(
X,
y,
test_size=test_ratio,
random_state=random_state,
shuffle=shuffle,
stratify=stratify,
)
val_size_in_train = val_ratio / (1 - test_ratio)
verbose and logger.info(
f"Split → Train: {1 - test_ratio:.2%} | Test: {test_ratio:.2%} (no validation set)"
)
return (X_train, X_test, y_train, y_test, val_size_in_train)
def _get_best_metric_average_strategy(y_true, balance_threshold: float = 0.5) -> str:
"""
Analyzes y_true to determine the best averaging strategy.
Args:
y_true: Input array (Numpy array, Pandas Series, or Polars Series).
balance_threshold: Float (0 to 1). If min_class_count / max_class_count
is below this, the data is considered imbalanced.
Returns:
str: 'binary', 'weighted', or 'macro'
"""
counts = None
if hasattr(y_true, "value_counts") and hasattr(y_true, "values"):
counts = y_true.value_counts().values
elif hasattr(y_true, "value_counts") and hasattr(y_true, "to_numpy"):
vc = y_true.value_counts()
if "count" in vc.columns:
counts = vc["count"].to_numpy()
else:
counts = vc[:, 1].to_numpy()
elif isinstance(y_true, np.ndarray):
(_, counts) = np.unique(y_true, return_counts=True)
else:
(_, counts) = np.unique(np.array(y_true), return_counts=True)
if counts is None or len(counts) == 0:
raise ValueError("Input y_true appears to be empty.")
n_classes = len(counts)
if n_classes <= 2:
return "binary"
min_c = np.min(counts)
max_c = np.max(counts)
ratio = min_c / max_c
if ratio < balance_threshold:
return "weighted"
else:
return "macro"
def _ensure_feature_names(X, feature_names=None):
if isinstance(X, pd.DataFrame):
return list(X.columns)
if isinstance(X, np.ndarray):
if feature_names is None:
feature_names = [f"feature_{i}" for i in range(X.shape[1])]
return feature_names
raise TypeError("X must be a pandas DataFrame or numpy ndarray")
def _perform_cross_validation(
model,
X,
y,
cv_folds,
average_strategy,
shuffle,
random_state,
n_jobs,
verbose,
pos_label=None,
) -> dict[str, Any]:
"""Perform cross-validation on the model."""
verbose and logger.info(f"Performing {cv_folds}-fold cross-validation...")
cv = StratifiedKFold(n_splits=cv_folds, shuffle=shuffle, random_state=random_state)
if average_strategy == "binary":
scoring = {
"accuracy": "accuracy",
"precision": make_scorer(
precision_score, average="binary", pos_label=pos_label
),
"recall": make_scorer(recall_score, average="binary", pos_label=pos_label),
"f1": make_scorer(f1_score, average="binary", pos_label=pos_label),
"roc_auc": "roc_auc",
}
else:
average_strategy_suffix = f"_{average_strategy}"
roc_average_strategy_suffix = (
f"_{average_strategy}" if average_strategy == "weighted" else ""
)
roc_auc_ovr_suffix = "_ovr"
scoring = (
f"f1{average_strategy_suffix}",
"accuracy",
f"precision{average_strategy_suffix}",
f"recall{average_strategy_suffix}",
f"roc_auc{roc_auc_ovr_suffix}{roc_average_strategy_suffix}",
)
cv_results = cross_validate(
model, X, y, cv=cv, scoring=scoring, return_train_score=True, n_jobs=n_jobs
)
def get_score_mean_std(metric_key):
if metric_key in cv_results:
return (cv_results[metric_key].mean(), cv_results[metric_key].std())
return (0.0, 0.0)
if average_strategy == "binary":
(accuracy_mean, accuracy_std) = get_score_mean_std("test_accuracy")
(precision_mean, precision_std) = get_score_mean_std("test_precision")
(recall_mean, recall_std) = get_score_mean_std("test_recall")
(f1_mean, f1_std) = get_score_mean_std("test_f1")
(roc_auc_mean, roc_auc_std) = get_score_mean_std("test_roc_auc")
else:
(accuracy_mean, accuracy_std) = get_score_mean_std("test_accuracy")
(precision_mean, precision_std) = get_score_mean_std(
f"test_precision{average_strategy_suffix}"
)
(recall_mean, recall_std) = get_score_mean_std(
f"test_recall{average_strategy_suffix}"
)
(f1_mean, f1_std) = get_score_mean_std(f"test_f1{average_strategy_suffix}")
roc_key = f"test_roc_auc{roc_auc_ovr_suffix}{roc_average_strategy_suffix}"
(roc_auc_mean, roc_auc_std) = get_score_mean_std(roc_key)
verbose and logger.info(
f"CV Accuracy : {accuracy_mean:.4f} (+/- {accuracy_std:.4f})"
)
verbose and logger.info(
f"CV Precision : {precision_mean:.4f} (+/- {precision_std:.4f})"
)
verbose and logger.info(f"CV Recall : {recall_mean:.4f} (+/- {recall_std:.4f})")
verbose and logger.info(f"CV F1 Score : {f1_mean:.4f} (+/- {f1_std:.4f})")
verbose and logger.info(
f"CV ROC-AUC : {roc_auc_mean:.4f} (+/- {roc_auc_std:.4f})"
)
CV_metrics = pd.DataFrame(
{
"Metric": ["Accuracy", "Precision", "Recall", "F1-Score", "ROC AUC"],
"Mean": [accuracy_mean, precision_mean, recall_mean, f1_mean, roc_auc_mean],
"Std": [accuracy_std, precision_std, recall_std, f1_std, roc_auc_std],
}
)
return CV_metrics
def _compute_score(model, X, y, metric, average_strategy, pos_label=None):
score_params = {"average": average_strategy, "zero_division": 0}
y_pred = model.predict(X)
if average_strategy != "binary":
y_score = model.predict_proba(X)
else:
score_params["pos_label"] = pos_label
if pos_label is not None:
classes = list(model.classes_)
try:
pos_idx = classes.index(pos_label)
except ValueError:
pos_idx = 1 if len(classes) > 1 else 0
y_score = model.predict_proba(X)[:, pos_idx]
else:
y_score = model.predict_proba(X)[:, 1]
if metric == "Accuracy":
score = accuracy_score(y, y_pred)
elif metric == "Precision":
score = precision_score(y, y_pred, **score_params)
elif metric == "Recall":
score = recall_score(y, y_pred, **score_params)
elif metric == "F1 Score":
score = f1_score(y, y_pred, **score_params)
elif metric == "ROC-AUC":
if average_strategy != "binary":
score = roc_auc_score(
y, y_score, multi_class="ovr", average=average_strategy
)
else:
score = roc_auc_score(y, y_score)
return score
def _get_cv_scoring_object(metric, average_strategy, pos_label=None):
"""
Returns a scoring object (string or callable) suitable for cross_validate.
Used during HPO.
"""
if average_strategy == "binary":
if metric == "F1 Score":
return make_scorer(f1_score, average="binary", pos_label=pos_label)
elif metric == "Accuracy":
return "accuracy"
elif metric == "Precision":
return make_scorer(precision_score, average="binary", pos_label=pos_label)
elif metric == "Recall":
return make_scorer(recall_score, average="binary", pos_label=pos_label)
elif metric == "ROC-AUC":
return "roc_auc"
else:
average_strategy_suffix = f"_{average_strategy}"
roc_auc_ovr_suffix = "_ovr"
roc_average_strategy_suffix = (
f"_{average_strategy}" if average_strategy == "weighted" else ""
)
if metric == "F1 Score":
return f"f1{average_strategy_suffix}"
elif metric == "Accuracy":
return "accuracy"
elif metric == "Precision":
return f"precision{average_strategy_suffix}"
elif metric == "Recall":
return f"recall{average_strategy_suffix}"
elif metric == "ROC-AUC":
return f"roc_auc{roc_auc_ovr_suffix}{roc_average_strategy_suffix}"
def _hyperparameters_optimization(
X,
y,
constant_hyperparameters,
optimization_metric,
metric_average_strategy,
val_ratio,
shuffle_split,
stratify_split,
use_cross_val,
cv_folds,
n_trials=50,
strategy="maximize",
sampler="Tree-structured Parzen",
seed=None,
n_jobs=-1,
verbose=False,
pos_label=None,
):
direction = "maximize" if strategy.lower() == "maximize" else "minimize"
sampler_map = {
"Tree-structured Parzen": TPESampler(seed=seed),
"Gaussian Process": GPSampler(seed=seed),
"CMA-ES": CmaEsSampler(seed=seed),
"Random Search": RandomSampler(seed=seed),
"Random Sobol Search": QMCSampler(seed=seed),
}
if sampler in sampler_map:
chosen_sampler = sampler_map[sampler]
else:
logger.warning(f"Sampler '{sampler}' not recognized → falling back to TPE")
chosen_sampler = TPESampler(seed=seed)
chosen_pruner = HyperbandPruner()
if use_cross_val:
cv = StratifiedKFold(
n_splits=cv_folds, shuffle=shuffle_split, random_state=seed
)
cv_score_obj = _get_cv_scoring_object(
optimization_metric, metric_average_strategy, pos_label
)
else:
(X_train, X_val, y_train, y_val) = train_test_split(
X,
y,
test_size=val_ratio,
random_state=seed,
stratify=y if stratify_split else None,
shuffle=shuffle_split,
)
def logging_callback(study: Study, trial: FrozenTrial):
"""Callback function to log trial progress"""
verbose and logger.info(
f"Trial {trial.number} finished with value: {trial.value} and parameters: {trial.params}"
)
try:
verbose and logger.info(f"Best value so far: {study.best_value}")
verbose and logger.info(f"Best parameters so far: {study.best_params}")
except ValueError:
verbose and logger.info(f"No successful trials completed yet")
verbose and logger.info(f"" + "-" * 50)
def objective(trial):
try:
criterion = trial.suggest_categorical("criterion", ["gini", "entropy"])
splitter = trial.suggest_categorical("splitter", ["best", "random"])
max_depth = trial.suggest_int("max_depth", 10, 120, log=True)
min_samples_split = trial.suggest_int("min_samples_split", 2, 32)
min_samples_leaf = trial.suggest_int("min_samples_leaf", 1, 32)
max_features = trial.suggest_categorical(
"max_features", ["sqrt", "log2", None, 0.5, 0.8]
)
class_weight = constant_hyperparameters.get("class_weight")
model = DecisionTreeClassifier(
criterion=criterion,
splitter=splitter,
min_samples_leaf=min_samples_leaf,
max_features=max_features,
max_depth=max_depth,
min_samples_split=min_samples_split,
class_weight=class_weight,
random_state=seed,
)
if use_cross_val:
scores = cross_validate(
model, X, y, cv=cv, n_jobs=n_jobs, scoring=cv_score_obj
)
return scores["test_score"].mean()
else:
model.fit(X_train, y_train)
score = _compute_score(
model,
X_val,
y_val,
optimization_metric,
metric_average_strategy,
pos_label,
)
return score
except Exception as e:
verbose and logger.error(
f"Trial {trial.number} failed with error: {str(e)}"
)
raise
study = create_study(
direction=direction, sampler=chosen_sampler, pruner=chosen_pruner
)
study.optimize(
objective,
n_trials=n_trials,
catch=(Exception,),
n_jobs=n_jobs,
callbacks=[logging_callback],
)
verbose and logger.info(f"Optimization completed!")
verbose and logger.info(
f" Best Criterion : {study.best_params['criterion']}"
)
verbose and logger.info(
f" Best Split Strategy : {study.best_params['splitter']}"
)
verbose and logger.info(
f" Best Max Depth : {study.best_params['max_depth']}"
)
verbose and logger.info(
f" Best Feature Sampling : {study.best_params['max_features']}"
)
verbose and logger.info(
f" Best Min Samples per Leaf : {study.best_params['min_samples_leaf']}"
)
verbose and logger.info(
f" Best Min Samples Split : {study.best_params['min_samples_split']}"
)
verbose and logger.info(
f" Best {optimization_metric:<22}: {study.best_value:.4f}"
)
verbose and logger.info(f" Sampler used : {sampler}")
verbose and logger.info(f" Direction : {direction}")
if use_cross_val:
verbose and logger.info(f" Cross-validation : {cv_folds}-fold")
else:
verbose and logger.info(
f" Validation : single train/val split"
)
trials = study.trials_dataframe()
trials["best_value"] = trials["value"].cummax()
cols = list(trials.columns)
value_idx = cols.index("value")
cols = [c for c in cols if c != "best_value"]
new_order = cols[: value_idx + 1] + ["best_value"] + cols[value_idx + 1 :]
trials = trials[new_order]
return (study.best_params, trials)
def _combine_test_data(
X_test, y_true, y_pred, y_proba, class_names, features_names=None
):
"""
Combine X_test, y_true, y_pred, and y_proba into a single DataFrame.
Parameters:
-----------
X_test : pandas/polars DataFrame, numpy array, or scipy sparse matrix
Test features
y_true : pandas/polars Series, numpy array, or list
True labels
y_pred : pandas/polars Series, numpy array, or list
Predicted labels
y_proba : pandas/polars Series/DataFrame, numpy array (1D or 2D), or list
Prediction probabilities - can be:
- 1D array for binary classification (probability of positive class)
- 2D array for multiclass (probabilities for each class)
class_names : list or array-like
Names of the classes in order.
For binary classification with 1D y_proba, only the positive class name is needed.
Returns:
--------
pandas.DataFrame
Combined DataFrame with features, probabilities, y_true, and y_pred
"""
if sparse.issparse(X_test):
X_df = pd.DataFrame(X_test.toarray())
elif isinstance(X_test, np.ndarray):
X_df = pd.DataFrame(X_test)
elif hasattr(X_test, "to_pandas"):
X_df = X_test.to_pandas()
elif isinstance(X_test, pd.DataFrame):
X_df = X_test.copy()
else:
raise TypeError(f"Unsupported type for X_test: {type(X_test)}")
if X_df.columns.tolist() == list(range(len(X_df.columns))):
X_df.columns = (
[f"feature_{i}" for i in range(len(X_df.columns))]
if features_names is None
else features_names
)
if isinstance(y_true, list):
y_true_series = pd.Series(y_true, name="y_true")
elif isinstance(y_true, np.ndarray):
y_true_series = pd.Series(y_true, name="y_true")
elif hasattr(y_true, "to_pandas"):
y_true_series = y_true.to_pandas()
y_true_series.name = "y_true"
elif isinstance(y_true, pd.Series):
y_true_series = y_true.copy()
y_true_series.name = "y_true"
else:
raise TypeError(f"Unsupported type for y_true: {type(y_true)}")
if isinstance(y_pred, list):
y_pred_series = pd.Series(y_pred, name="y_pred")
elif isinstance(y_pred, np.ndarray):
y_pred_series = pd.Series(y_pred, name="y_pred")
elif hasattr(y_pred, "to_pandas"):
y_pred_series = y_pred.to_pandas()
y_pred_series.name = "y_pred"
elif isinstance(y_pred, pd.Series):
y_pred_series = y_pred.copy()
y_pred_series.name = "y_pred"
else:
raise TypeError(f"Unsupported type for y_pred: {type(y_pred)}")
if isinstance(y_proba, list):
y_proba_array = np.array(y_proba)
elif isinstance(y_proba, np.ndarray):
y_proba_array = y_proba
elif hasattr(y_proba, "to_pandas"):
y_proba_pd = y_proba.to_pandas()
if isinstance(y_proba_pd, pd.Series):
y_proba_array = y_proba_pd.values
else:
y_proba_array = y_proba_pd.values
elif isinstance(y_proba, pd.Series):
y_proba_array = y_proba.values
elif isinstance(y_proba, pd.DataFrame):
y_proba_array = y_proba.values
else:
raise TypeError(f"Unsupported type for y_proba: {type(y_proba)}")
def sanitize_class_name(class_name):
"""Convert class name to valid column name by replacing spaces and special chars"""
return str(class_name).replace(" ", "_").replace("-", "_")
if y_proba_array.ndim == 1:
y_proba_df = pd.DataFrame({"proba": y_proba_array})
else:
n_classes = y_proba_array.shape[1]
if len(class_names) == n_classes:
proba_columns = [f"proba_{sanitize_class_name(cls)}" for cls in class_names]
else:
proba_columns = [f"proba_{i}" for i in range(n_classes)]
y_proba_df = pd.DataFrame(y_proba_array, columns=proba_columns)
y_proba_df = y_proba_df.reset_index(drop=True)
X_df = X_df.reset_index(drop=True)
y_true_series = y_true_series.reset_index(drop=True)
y_pred_series = y_pred_series.reset_index(drop=True)
is_false_prediction = pd.Series(
y_true_series != y_pred_series, name="is_false_prediction"
).reset_index(drop=True)
result_df = pd.concat(
[X_df, y_proba_df, y_true_series, y_pred_series, is_false_prediction], axis=1
)
return result_df
def _get_feature_importance(model, feature_names=None, sort=True, top_n=None):
"""
Extract feature importance from model.
Parameters:
-----------
model : Fitted scikit-learn model
feature_names : list or array-like, optional
Names of features. If None, uses generic names like 'feature_0', 'feature_1', etc.
sort : bool, default=True
Whether to sort features by importance (descending)
top_n : int, optional
If specified, returns only the top N most important features
Returns:
--------
pd.DataFrame
DataFrame with columns: 'feature', 'importance'
Importance values represent the mean decrease in impurity (Gini importance)
"""
importances = model.feature_importances_
if feature_names is None:
feature_names = [f"feature_{i}" for i in range(len(importances))]
importance_df = pd.DataFrame({"feature": feature_names, "importance": importances})
if sort:
importance_df = importance_df.sort_values("importance", ascending=False)
importance_df = importance_df.reset_index(drop=True)
if top_n is not None:
importance_df = importance_df.head(top_n)
return importance_df
def _smart_shap_background(
X: Union[np.ndarray, pd.DataFrame],
model_type: str = "tree",
seed: int = 42,
verbose: bool = True,
) -> Union[np.ndarray, pd.DataFrame, object]:
"""
Intelligently prepares a background dataset for SHAP based on model type.
Strategies:
- Tree: Higher sample cap (1000), uses Random Sampling (preserves data structure).
- Other: Lower sample cap (100), uses K-Means (maximizes info density).
"""
(n_rows, n_features) = X.shape
if model_type == "tree":
max_samples = 1000
use_kmeans = False
else:
max_samples = 100
use_kmeans = True
if n_rows <= max_samples:
verbose and logger.info(
f"✓ Dataset small ({n_rows} <= {max_samples}). Using full data."
)
return X
verbose and logger.info(
f"⚡ Large dataset detected ({n_rows} rows). Optimization Strategy: {('K-Means' if use_kmeans else 'Random Sampling')}"
)
if use_kmeans:
try:
verbose and logger.info(
f" Summarizing to {max_samples} weighted centroids..."
)
return shap.kmeans(X, max_samples)
except Exception as e:
logger.warning(
f" K-Means failed ({str(e)}). Falling back to random sampling."
)
return shap.sample(X, max_samples, random_state=seed)
else:
verbose and logger.info(f" Sampling {max_samples} random rows...")
return shap.sample(X, max_samples, random_state=seed)
def _class_index_df(model):
columns = {"index": pd.Series(dtype="int64"), "class": pd.Series(dtype="object")}
if model is None:
return pd.DataFrame(columns)
classes = getattr(model, "classes_", None)
if classes is None:
return pd.DataFrame(columns)
return pd.DataFrame({"index": range(len(classes)), "class": classes})
def train_cla_decision_tree(
X: DataFrame, y: Union[DataSeries, NDArray, List], options=None
) -> Tuple[
Any,
DataFrame,
Any,
Union[DataFrame, Dict],
DataFrame,
DataFrame,
DataFrame,
DataFrame,
Dict,
]:
options = options or {}
criterion = options.get("criterion", "gini").lower()
splitter = options.get("splitter", "best").lower()
feature_strategy = options.get("feature_strategy", "Square root")
max_depth_input = options.get("max_depth", 0)
min_samples_leaf = options.get("min_samples_leaf", 1)
if feature_strategy == "Automatic 30%":
max_features = 0.3
elif feature_strategy == "Automatic 60%":
max_features = 0.5
elif feature_strategy == "Square root":
max_features = "sqrt"
elif feature_strategy == "Logarithm":
max_features = "log2"
elif feature_strategy == "None":
max_features = None
else:
max_features = "sqrt"
max_depth = None if max_depth_input == 0 else max_depth_input
auto_split = options.get("auto_split", True)
test_val_size = options.get("test_val_size", 15) / 100
shuffle_split = options.get("shuffle_split", True)
stratify_split = options.get("stratify_split", True)
retrain_on_full = options.get("retrain_on_full", False)
custom_average_strategy = options.get("custom_average_strategy", "auto")
use_cross_validation = options.get("use_cross_validation", False)
cv_folds = options.get("cv_folds", 5)
use_hpo = options.get("use_hyperparameter_optimization", False)
optimization_metric = options.get("optimization_metric", "F1 Score")
optimization_method = options.get("optimization_method", "Tree-structured Parzen")
optimization_iterations = options.get("optimization_iterations", 50)
pos_label_option = options.get("pos_label", "").strip()
if pos_label_option == "":
pos_label_option = None
return_shap_explainer = options.get("return_shap_explainer", False)
use_shap_sampler = options.get("use_shap_sampler", False)
shap_feature_perturbation = options.get(
"shap_feature_perturbation", "Interventional"
)
metrics_as = options.get("metrics_as", "Dataframe")
n_jobs_str = options.get("n_jobs", "1")
random_state = options.get("random_state", 42)
activate_caching = options.get("activate_caching", False)
verbose = options.get("verbose", True)
n_jobs_int = -1 if n_jobs_str == "All" else int(n_jobs_str)
skip_computation = False
Model = None
Metrics = pd.DataFrame()
CV_Metrics = pd.DataFrame()
Features_Importance = pd.DataFrame()
SHAP = None
HPO_Trials = pd.DataFrame()
HPO_Best = None
accuracy = None
precision = None
recall = None
f1 = None
roc_auc = None
(n_samples, n_features, sparsity, is_sparse) = _get_shape_and_sparsity(X)
shap_feature_names = _ensure_feature_names(X)
if activate_caching:
verbose and logger.info(f"Caching is activate")
data_hasher = _UniversalDatasetHasher(n_samples, verbose=verbose)
X_hash = data_hasher.hash_data(X).hash
y_hash = data_hasher.hash_data(y).hash
all_hash_base_text = f"HASH BASE TEXTPandas Version {pd.__version__}POLARS Version {pl.__version__}Numpy Version {np.__version__}Scikit Learn Version {sklearn.__version__}Scipy Version {scipy.__version__}{('SHAP Version ' + shap.__version__ if return_shap_explainer else 'NO SHAP Version')}{X_hash}{y_hash}{criterion}{splitter}{feature_strategy}{max_depth_input}{min_samples_leaf}{('Use HPO' if use_hpo else 'No HPO')}{(optimization_metric if use_hpo else 'No HPO Metric')}{(optimization_method if use_hpo else 'No HPO Method')}{(optimization_iterations if use_hpo else 'No HPO Iter')}{(cv_folds if use_cross_validation else 'No CV')}{('Auto Split' if auto_split else test_val_size)}{shuffle_split}{stratify_split}{return_shap_explainer}{shap_feature_perturbation}{use_shap_sampler}{random_state}{(pos_label_option if pos_label_option else 'default_pos')}{custom_average_strategy}"
all_hash = hashlib.sha256(all_hash_base_text.encode("utf-8")).hexdigest()
verbose and logger.info(f"Hash was computed: {all_hash}")
temp_folder = Path(tempfile.gettempdir())
cache_folder = temp_folder / "coded-flows-cache"
cache_folder.mkdir(parents=True, exist_ok=True)
model_path = cache_folder / f"{all_hash}.model"
metrics_dict_path = cache_folder / f"metrics_{all_hash}.json"
metrics_df_path = cache_folder / f"metrics_{all_hash}.parquet"
cv_metrics_path = cache_folder / f"cv_metrics_{all_hash}.parquet"
hpo_trials_path = cache_folder / f"hpo_trials_{all_hash}.parquet"
hpo_best_params_path = cache_folder / f"hpo_best_params_{all_hash}.json"
features_importance_path = (
cache_folder / f"features_importance_{all_hash}.parquet"
)
prediction_set_path = cache_folder / f"prediction_set_{all_hash}.parquet"
shap_path = cache_folder / f"{all_hash}.shap"
skip_computation = model_path.is_file()
if not skip_computation:
try:
_validate_numerical_data(X)
except Exception as e:
verbose and logger.error(
f"Only numerical or boolean types are allowed for 'X' input!"
)
raise
if hasattr(y, "nunique"):
n_classes = y.nunique()
elif hasattr(y, "n_unique"):
n_classes = y.n_unique()
else:
n_classes = len(np.unique(y))
is_multiclass = n_classes > 2
features_names = X.columns if hasattr(X, "columns") else None
min_samples_split = 2 * min_samples_leaf
class_weight = _choose_class_weight(y)
if class_weight in ["balanced", "balanced_subsample"]:
verbose and logger.info(
f"Class imbalance, using '{class_weight}' mode for class weight"
)
fixed_test_split = None if auto_split else test_val_size
(X_train, X_test, y_train, y_test, val_ratio) = _smart_split(
n_samples,
X,
y,
random_state=random_state,
shuffle=shuffle_split,
stratify=y if stratify_split else None,
fixed_test_split=fixed_test_split,
verbose=verbose,
)
if custom_average_strategy == "auto":
metric_average_strategy = _get_best_metric_average_strategy(y_test)
else:
metric_average_strategy = custom_average_strategy
effective_pos_label = None
if metric_average_strategy == "binary":
unique_classes = np.unique(y_train)
if pos_label_option is not None:
if pos_label_option in unique_classes:
effective_pos_label = pos_label_option
else:
try:
as_int = int(pos_label_option)
if as_int in unique_classes:
effective_pos_label = as_int
except ValueError:
pass
elif 1 in unique_classes:
effective_pos_label = 1
elif "1" in unique_classes:
effective_pos_label = "1"
if effective_pos_label is not None:
verbose and logger.info(f"Using positive label: {effective_pos_label}")
elif effective_pos_label is None and metric_average_strategy == "binary":
error_message = 'The target appears to be binary, but no positive label was provided and no "1" class exists in the label set.'
verbose and logger.error(error_message)
raise ValueError(error_message)
if use_hpo:
verbose and logger.info(f"Performing Hyperparameters Optimization")
constant_hyperparameters = {"class_weight": class_weight}
(HPO_Best, HPO_Trials) = _hyperparameters_optimization(
X_train,
y_train,
constant_hyperparameters,
optimization_metric,
metric_average_strategy,
val_ratio,
shuffle_split,
stratify_split,
use_cross_validation,
cv_folds,
optimization_iterations,
"maximize",
optimization_method,
random_state,
n_jobs_int,
verbose=verbose,
pos_label=effective_pos_label,
)
HPO_Trials = _normalize_hpo_df(HPO_Trials)
criterion = HPO_Best["criterion"]
splitter = HPO_Best["splitter"]
max_depth = HPO_Best["max_depth"]
min_samples_split = HPO_Best["min_samples_split"]
min_samples_leaf = HPO_Best["min_samples_leaf"]
max_features = HPO_Best["max_features"]
Model = DecisionTreeClassifier(
criterion=criterion,
splitter=splitter,
min_samples_leaf=min_samples_leaf,
max_features=max_features,
max_depth=max_depth,
min_samples_split=min_samples_split,
class_weight=class_weight,
random_state=random_state,
)
if use_cross_validation and (not use_hpo):
verbose and logger.info(
f"Using Cross-Validation to measure performance metrics"
)
CV_Metrics = _perform_cross_validation(
Model,
X_train,
y_train,
cv_folds,
metric_average_strategy,
shuffle_split,
random_state,
n_jobs_int,
verbose,
pos_label=effective_pos_label,
)
Model.fit(X_train, y_train)
y_pred = Model.predict(X_test)
if is_multiclass:
y_score = Model.predict_proba(X_test)
elif effective_pos_label is not None:
try:
pos_idx = list(Model.classes_).index(effective_pos_label)
y_score = Model.predict_proba(X_test)[:, pos_idx]
except ValueError:
y_score = Model.predict_proba(X_test)[:, 1]
else:
y_score = Model.predict_proba(X_test)[:, 1]
score_params = {"average": metric_average_strategy, "zero_division": 0}
if effective_pos_label:
score_params["pos_label"] = effective_pos_label
accuracy = accuracy_score(y_test, y_pred)
precision = precision_score(y_test, y_pred, **score_params)
recall = recall_score(y_test, y_pred, **score_params)
f1 = f1_score(y_test, y_pred, **score_params)
if is_multiclass:
roc_auc = roc_auc_score(
y_test, y_score, multi_class="ovr", average=metric_average_strategy
)
else:
roc_auc = roc_auc_score(y_test, y_score)
if metrics_as == "Dataframe":
Metrics = pd.DataFrame(
{
"Metric": [
"Accuracy",
"Precision",
"Recall",
"F1-Score",
"ROC AUC",
],
"Value": [accuracy, precision, recall, f1, roc_auc],
}
)
else:
Metrics = {
"Accuracy": accuracy,
"Precision": precision,
"Recall": recall,
"F1-Score": f1,
"ROC AUC": roc_auc,
}
verbose and logger.info(f"Accuracy : {accuracy:.4f}")
verbose and logger.info(f"Precision : {precision:.4f}")
verbose and logger.info(f"Recall : {recall:.4f}")
verbose and logger.info(f"F1-Score : {f1:.4f}")
verbose and logger.info(f"ROC-AUC : {roc_auc:.4f}")
Prediction_Set = _combine_test_data(
X_test, y_test, y_pred, y_score, Model.classes_, features_names
)
verbose and logger.info(f"Prediction Set created")
if retrain_on_full:
verbose and logger.info(
"Retraining model on full dataset for production deployment"
)
Model.fit(X, y)
verbose and logger.info(
"Model successfully retrained on full dataset. Reported metrics remain from original held-out test set."
)
Features_Importance = _get_feature_importance(Model, features_names)
verbose and logger.info(f"Features Importance computed")
if return_shap_explainer:
if shap_feature_perturbation == "Interventional":
SHAP = shap.TreeExplainer(
Model,
(
_smart_shap_background(
X if retrain_on_full else X_train,
model_type="tree",
seed=random_state,
verbose=verbose,
)
if use_shap_sampler
else X if retrain_on_full else X_train
),
feature_names=shap_feature_names,
)
else:
SHAP = shap.TreeExplainer(
Model,
feature_names=shap_feature_names,
feature_perturbation="tree_path_dependent",
)
verbose and logger.info(f"SHAP explainer generated")
if activate_caching:
verbose and logger.info(f"Caching output elements")
joblib.dump(Model, model_path)
if isinstance(Metrics, dict):
with metrics_dict_path.open("w", encoding="utf-8") as f:
json.dump(Metrics, f, ensure_ascii=False, indent=4)
else:
Metrics.to_parquet(metrics_df_path)
if use_cross_validation and (not use_hpo):
CV_Metrics.to_parquet(cv_metrics_path)
if use_hpo:
HPO_Trials.to_parquet(hpo_trials_path)
with hpo_best_params_path.open("w", encoding="utf-8") as f:
json.dump(HPO_Best, f, ensure_ascii=False, indent=4)
Features_Importance.to_parquet(features_importance_path)
Prediction_Set.to_parquet(prediction_set_path)
if return_shap_explainer:
with shap_path.open("wb") as f:
joblib.dump(SHAP, f)
verbose and logger.info(f"Caching done")
else:
verbose and logger.info(f"Skipping computations and loading cached elements")
Model = joblib.load(model_path)
verbose and logger.info(f"Model loaded")
if metrics_dict_path.is_file():
with metrics_dict_path.open("r", encoding="utf-8") as f:
Metrics = json.load(f)
else:
Metrics = pd.read_parquet(metrics_df_path)
verbose and logger.info(f"Metrics loaded")
if use_cross_validation and (not use_hpo):
CV_Metrics = pd.read_parquet(cv_metrics_path)
verbose and logger.info(f"Cross Validation metrics loaded")
if use_hpo:
HPO_Trials = pd.read_parquet(hpo_trials_path)
with hpo_best_params_path.open("r", encoding="utf-8") as f:
HPO_Best = json.load(f)
verbose and logger.info(
f"Hyperparameters Optimization trials and best params loaded"
)
Features_Importance = pd.read_parquet(features_importance_path)
verbose and logger.info(f"Features Importance loaded")
Prediction_Set = pd.read_parquet(prediction_set_path)
verbose and logger.info(f"Prediction Set loaded")
if return_shap_explainer:
with shap_path.open("rb") as f:
SHAP = joblib.load(f)
verbose and logger.info(f"SHAP Explainer loaded")
Model_Classes = _class_index_df(Model)
return (
Model,
Model_Classes,
SHAP,
Metrics,
CV_Metrics,
Features_Importance,
Prediction_Set,
HPO_Trials,
HPO_Best,
)
Brick Info
- shap>=0.47.0
- scikit-learn
- pandas
- numpy
- torch
- numba>=0.56.0
- shap
- cmaes
- optuna
- scipy
- polars
- xxhash