wrapshap 0.1__tar.gz → 0.1.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (28) hide show
  1. {wrapshap-0.1/wrapshap.egg-info → wrapshap-0.1.2}/PKG-INFO +2 -2
  2. {wrapshap-0.1 → wrapshap-0.1.2}/README.md +4 -0
  3. {wrapshap-0.1 → wrapshap-0.1.2}/setup.py +2 -2
  4. wrapshap-0.1.2/wrapshap/__init__.py +1 -0
  5. wrapshap-0.1.2/wrapshap/_internal/__init__.py +0 -0
  6. wrapshap-0.1.2/wrapshap/_internal/_already_wrapped.py +531 -0
  7. wrapshap-0.1.2/wrapshap/_internal/_already_wrapped_LGB.py +574 -0
  8. wrapshap-0.1.2/wrapshap/_internal/_wrapshap.py +525 -0
  9. wrapshap-0.1.2/wrapshap/already_wrapped.py +2 -0
  10. wrapshap-0.1.2/wrapshap/gpu_implementation/__init__.py +1 -0
  11. wrapshap-0.1.2/wrapshap/gpu_implementation/_internal/__init__.py +0 -0
  12. wrapshap-0.1.2/wrapshap/gpu_implementation/_internal/_already_wrapped_gpu.py +213 -0
  13. wrapshap-0.1.2/wrapshap/gpu_implementation/_internal/_already_wrapped_lgb_gpu.py +217 -0
  14. wrapshap-0.1.2/wrapshap/gpu_implementation/_internal/_wrapshap_gpu.py +774 -0
  15. wrapshap-0.1.2/wrapshap/gpu_implementation/already_wrapped_gpu.py +2 -0
  16. wrapshap-0.1.2/wrapshap/gpu_implementation/wrapshap_gpu.py +157 -0
  17. wrapshap-0.1.2/wrapshap/wrapshap.py +215 -0
  18. {wrapshap-0.1 → wrapshap-0.1.2/wrapshap.egg-info}/PKG-INFO +2 -2
  19. wrapshap-0.1.2/wrapshap.egg-info/SOURCES.txt +24 -0
  20. {wrapshap-0.1 → wrapshap-0.1.2}/wrapshap.egg-info/requires.txt +1 -1
  21. wrapshap-0.1.2/wrapshap.egg-info/top_level.txt +1 -0
  22. wrapshap-0.1/wrapshap.egg-info/SOURCES.txt +0 -10
  23. wrapshap-0.1/wrapshap.egg-info/top_level.txt +0 -1
  24. {wrapshap-0.1 → wrapshap-0.1.2}/LICENSE +0 -0
  25. {wrapshap-0.1 → wrapshap-0.1.2}/README_PyPI.md +0 -0
  26. {wrapshap-0.1 → wrapshap-0.1.2}/pyproject.toml +0 -0
  27. {wrapshap-0.1 → wrapshap-0.1.2}/setup.cfg +0 -0
  28. {wrapshap-0.1 → wrapshap-0.1.2}/wrapshap.egg-info/dependency_links.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: wrapshap
3
- Version: 0.1
3
+ Version: 0.1.2
4
4
  Summary: SHAP-based feature selection and exploration for high-dimensional data
5
5
  Home-page: https://github.com/fefurger/wrapshap
6
6
  Author: Félix Furger
@@ -22,7 +22,7 @@ Requires-Dist: lightgbm>=4.6.0
22
22
  Requires-Dist: seaborn>=0.13.2
23
23
  Requires-Dist: matplotlib>=3.8.4
24
24
  Provides-Extra: gpu
25
- Requires-Dist: tensorflow>=2.16.1; extra == "gpu"
25
+ Requires-Dist: tensorflow<=2.20,>=2.16.1; extra == "gpu"
26
26
  Dynamic: author
27
27
  Dynamic: author-email
28
28
  Dynamic: classifier
@@ -1,3 +1,7 @@
1
+ This package allows you to perform SHAP-based feature selection and exploration for high-dimensional data, with both CPU and GPU support.
2
+
3
+ **Paper:** [Furger et al. (2026), *Wrapshap: a high dimension feature selection framework to explore biological systems through machine learning explanations*](https://www.nature.com/articles/s41540-026-00779-z).
4
+
1
5
  ## How to install Wrapshap
2
6
 
3
7
  Both CPU and GPU versions of Wrapshap are available.
@@ -6,7 +6,7 @@ README = Path(__file__).parent / "README_PyPI.md"
6
6
 
7
7
  setup(
8
8
  name='wrapshap',
9
- version='0.1',
9
+ version='0.1.2',
10
10
  author='Félix Furger',
11
11
  author_email='fefurger@hotmail.com',
12
12
  description='SHAP-based feature selection and exploration for high-dimensional data',
@@ -27,7 +27,7 @@ setup(
27
27
  ],
28
28
  extras_require={
29
29
  'gpu': [
30
- 'tensorflow>=2.16.1'
30
+ 'tensorflow>=2.16.1, <=2.20'
31
31
  ]
32
32
  },
33
33
  license="Apache-2.0",
@@ -0,0 +1 @@
1
+ from .wrapshap import Wrapshap
File without changes
@@ -0,0 +1,531 @@
1
+ """
2
+ This module provides utility functions and convenient classes of XGBoost regressors and classifiers already wrapped
3
+ with Wrapshap.
4
+ """
5
+
6
+ # Imports
7
+
8
+ import numpy as np
9
+ import pandas as pd
10
+ import xgboost as xgb
11
+ import seaborn as sns
12
+ import matplotlib.pyplot as plt
13
+ import matplotlib.lines as mlines
14
+
15
+ from functools import wraps
16
+ from abc import ABC, abstractmethod
17
+ from sklearn.preprocessing import LabelEncoder
18
+ from sklearn.model_selection import train_test_split
19
+ from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
20
+ from sklearn.metrics import auc, roc_curve
21
+ from sklearn.metrics import r2_score, mean_absolute_error
22
+
23
+ from ..wrapshap import Wrapshap
24
+
25
+ # ----------------------------------------------------------------------------------------------------------------------
26
+ # Public API - These utilities are used by other modules and re-exported.
27
+ # ----------------------------------------------------------------------------------------------------------------------
28
+
29
+
30
+ def wrapshap_required(*attrs):
31
+ def decorator(func):
32
+ @wraps(func)
33
+ def wrapper(self, *args, **kwargs):
34
+ for attr in attrs:
35
+ if getattr(self, attr, None) is None:
36
+ if self.model is None:
37
+ print('No model found.')
38
+ print('Fitting model with default parameters..')
39
+ self.fit()
40
+ print('Done.')
41
+ print("Computing SHAP values and predictions..")
42
+ self.compute_shap_and_predictions()
43
+ print('Done.')
44
+ break
45
+ return func(self, *args, **kwargs)
46
+ return wrapper
47
+ return decorator
48
+
49
+
50
+ class WrappedXGB(ABC):
51
+ """
52
+ This base class provides an XGBoost already conveniently wrapped with Wrapshap. This allows you to fit a model and
53
+ compute its SHAP values in a straigthforward way. If you want to have more control over the model and the SHAP
54
+ computation, please train your model separately and directly use the Wrapshap class itself.
55
+
56
+ Attributes:
57
+ X (np.array): The original data as a 2D array.
58
+ y (np.array): The original target as a 1D array.
59
+ X_train (np.array): The original train set as a 2D array.
60
+ X_test (np.array): The original test set as a 2D array.
61
+ y_train (np.array): The original train target as a 1D array.
62
+ y_test (np.array): The original test target as a 1D array.
63
+ shap_train (np.array): The SHAP values for the train set as a 2D array.
64
+ shap_test (np.array): The SHAP values for the test set as a 2D array.
65
+ y_pred_train (np.array): The complex model predictions for the train set as a 1D array.
66
+ y_pred_test (np.array): The complex model predictions for the test set as a 1D array.
67
+ feature_names (list): The feature names.
68
+ random_state (int): The chosen random state for reproductibility.
69
+ model (xgb.Booster): The XGBoost model.
70
+ xgb_params (dict): The XGBoost model parameters.
71
+ wrapshap (Wrapshap): The underlying Wrapshap object being called.
72
+ """
73
+
74
+ def __init__(self, X: np.ndarray, y: np.ndarray, feature_names: list = None, random_state: int = 42, **kwargs):
75
+ """
76
+ Parameters:
77
+ X (np.array): The original data as a 2D array.
78
+ y (np.array): The original target as a 1D array.
79
+ feature_names (list): The feature names.
80
+ random_state (int): The chosen random state for reproductibility.
81
+ **kwargs: Any parameters to add to the XGBoost configuration.
82
+ """
83
+
84
+ if type(self) is WrappedXGB:
85
+ raise TypeError(f"{type(self)} is an abstract class and cannot be instanciated directly.")
86
+
87
+ if not isinstance(X, np.ndarray) or X.ndim != 2:
88
+ raise ValueError("'X' should be a 2D numpy array.")
89
+
90
+ if not isinstance(y, np.ndarray) or y.ndim != 1:
91
+ raise ValueError("'y' should be a 1D numpy array.")
92
+
93
+ if feature_names is None:
94
+ self.feature_names = [f"feature_{i}" for i in range(1, X.shape[1]+1)]
95
+ else:
96
+ if isinstance(feature_names, (list, pd.Series)):
97
+ feature_names = [str(f) for f in feature_names]
98
+ else:
99
+ raise TypeError("'feature_names' should be a list of strings.")
100
+ self.feature_names = feature_names
101
+
102
+ self.xgb_params = {
103
+ 'eta': 0.1,
104
+ 'max_depth': 3,
105
+ 'seed': random_state
106
+ }
107
+
108
+ self.xgb_params.update(kwargs)
109
+
110
+ self.X = X
111
+ self.y = y
112
+ self.random_state = random_state
113
+
114
+ self.model = None
115
+ self.shap_train = None
116
+ self.shap_test = None
117
+ self.y_pred_train = None
118
+ self.y_pred_test = None
119
+ self.wrapshap = None
120
+
121
+
122
+ def compute_shap_and_predictions(self):
123
+ """
124
+ Computes the XGBoost model's SHAP values and predictions.
125
+ """
126
+
127
+ if self.model is None:
128
+ print("No XGBoost model was found.")
129
+ return
130
+
131
+ self.shap_train = self.model.predict(xgb.DMatrix(self.X_train, self.y_train), pred_contribs=True)[:,:-1]
132
+ self.shap_test = self.model.predict(xgb.DMatrix(self.X_test, self.y_test), pred_contribs=True)[:,:-1]
133
+
134
+ @wrapshap_required('wrapshap')
135
+ def shapRFE(self, surrogate_metrics: bool = False, top: int = None, plot:bool = True,time_info:bool = False):
136
+ """
137
+ Simply calls Wrapshap implementation of Recursive Feature Elimination (RFE).
138
+
139
+ Parameters:
140
+ surrogate_metrics (bool): Whether the compute and return the surrogate metrics or not.
141
+ top (int): When different than None, instead of performing RFE on all features, perfoms it only on the top
142
+ features (the ones with biggest mean absolute SHAP values)
143
+
144
+ Returns:
145
+ tuple[list, dict]: A tuple containing the list of ranked features (from most to least important) and a
146
+ dictionary with the corresponding performance metrics.
147
+ """
148
+
149
+ return self.wrapshap.shapRFE(
150
+ surrogate_metrics=surrogate_metrics,
151
+ top=top,
152
+ plot = plot,
153
+ time_info=time_info
154
+ )
155
+
156
+
157
+ @wrapshap_required('wrapshap')
158
+ def shapFFS(self, surrogate_metrics: bool = False, top: int = None, plot:bool = True,time_info:bool = True):
159
+ """
160
+ Simply calls Wrapshap implementation of Forward Feature Selection (FFS).
161
+
162
+ Parameters:
163
+ surrogate_metrics (bool): Whether the compute and return the surrogate metrics or not.
164
+ top (int): When different than None, instead of performing FFS on all features, perfoms it only on the top
165
+ features (the ones with biggest mean absolute SHAP values)
166
+
167
+ Returns:
168
+ tuple[list, dict]: A tuple containing the list of ranked features (from most to least important) and a
169
+ dictionary with the corresponding performance metrics.
170
+ """
171
+
172
+ return self.wrapshap.shapFFS(
173
+ surrogate_metrics=surrogate_metrics,
174
+ top=top,
175
+ plot = plot,
176
+ time_info=time_info
177
+ )
178
+
179
+
180
+ @wrapshap_required('wrapshap')
181
+ def shapBFE(self, surrogate_metrics: bool = False, top: int = None, plot:bool = True, time_info:bool=False):
182
+ """
183
+ Simply calls Wrapshap implementation of Forward Feature Selection (BFE).
184
+
185
+ Parameters:
186
+ surrogate_metrics (bool): Whether the compute and return the surrogate metrics or not.
187
+ top (int): When different than None, instead of performing BFE on all features, perfoms it only on the top
188
+ features (the ones with biggest mean absolute SHAP values)
189
+
190
+ Returns:
191
+ tuple[list, dict]: A tuple containing the list of ranked features (from most to least important) and a
192
+ dictionary with the corresponding performance metrics.
193
+ """
194
+
195
+ return self.wrapshap.shapBFE(
196
+ surrogate_metrics=surrogate_metrics,
197
+ top=top,
198
+ plot = plot,
199
+ time_info = time_info
200
+ )
201
+
202
+
203
+ @wrapshap_required('wrapshap')
204
+ def get_surrogate_metrics(self, feature_ranking: list, metrics_dict: dict = {}):
205
+ """
206
+ Iteratively computes the surrogate metrics for subsets from the provided feature ranking.
207
+
208
+ Parameters:
209
+ feature_ranking (list): The ranked features.
210
+ metrics_dict (dict): The metrics dictionary to which surrogate metrics will be added.
211
+
212
+ Returns:
213
+ dict: A dictionary containing the computed surrogate performance metrics.
214
+ """
215
+
216
+ return self.wrapshap.get_surrogate_metrics(
217
+ feature_ranking=feature_ranking,
218
+ metrics_dict=metrics_dict,
219
+ )
220
+
221
+
222
+ @wrapshap_required('wrapshap')
223
+ def shap_summary_plot(self, return_ranking: bool = False, max_display: int = None):
224
+ """
225
+ Plots a SHAP Summary Plot.
226
+
227
+ Parameters:
228
+ return_ranking (bool): Whether or not to return the features ranked by mean absolute SHAP values (ShapAbs).
229
+ max_display (int): How many top features to include in the plot.
230
+
231
+ Returns:
232
+ list: If return_ranking is True, the list of ranked features.
233
+ """
234
+
235
+ return self.wrapshap.shap_summary_plot(
236
+ return_ranking=return_ranking,
237
+ max_display=max_display
238
+ )
239
+
240
+ @abstractmethod
241
+ def plot_performance(self):
242
+ pass
243
+
244
+
245
+ @abstractmethod
246
+ def fit(self):
247
+ pass
248
+
249
+
250
+ def _fit(self, eval_size, num_boost_round, early_stopping_rounds, verbose_eval):
251
+ X_train, X_val, y_train, y_val = train_test_split(self.X_train, self.y_train,
252
+ test_size=eval_size, random_state=self.random_state)
253
+
254
+ dtrain = xgb.DMatrix(X_train, label=y_train)
255
+ dvalid = xgb.DMatrix(X_val, label=y_val)
256
+
257
+ self.model = xgb.train(
258
+ params=self.xgb_params,
259
+ dtrain=dtrain,
260
+ num_boost_round=num_boost_round,
261
+ evals=[(dvalid, 'eval')],
262
+ early_stopping_rounds=early_stopping_rounds,
263
+ verbose_eval=verbose_eval
264
+ )
265
+
266
+
267
+ class WrappedBinaryClassifierXGB(WrappedXGB):
268
+ """
269
+ This class extends the base WrappedXGB class and provides an XGBoost Binary Classifier already conveniently wrapped
270
+ with Wrapshap. Objective will be binary logistic. Please see WrappedXGB for more information.
271
+ """
272
+
273
+ def __init__(
274
+ self, X: np.ndarray, y: np.ndarray, feature_names: list = None,
275
+ test_size: float = 0.2, random_state: int = 42, encode_label: bool = True,
276
+ **kwargs
277
+ ):
278
+ """
279
+ Parameters:
280
+ X (np.array): The original data as a 2D array.
281
+ y (np.array): The original target as a 1D array.
282
+ feature_names (list): The feature names.
283
+ test_size (float): Portion of the data being used as the test set.
284
+ random_state (int): The chosen random state for reproductibility.
285
+ encode_label (bool): Whether or not to encode the labels before fitting.
286
+ **kwargs: Any parameters to add to the XGBoost configuration.
287
+ """
288
+
289
+ super().__init__(X, y, feature_names, random_state, **kwargs)
290
+
291
+ self.xgb_params['objective'] = 'binary:logistic'
292
+
293
+ if encode_label:
294
+ label_encoder = LabelEncoder()
295
+ self.y = label_encoder.fit_transform(self.y)
296
+
297
+ self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.X, self.y,
298
+ test_size=test_size,
299
+ random_state=random_state)
300
+
301
+ def fit(
302
+ self, scale_pos_weight: bool = False, eval_size: float = 0.2, num_boost_round: int = 10000,
303
+ early_stopping_rounds: int = 10, verbose_eval: bool = False, threshold_prob: float = 0.5
304
+ ):
305
+ """
306
+ Fits the classifier on training data.
307
+
308
+ Parameters:
309
+ scale_pos_weight (bool): Whether or not to apply weights proportional to the class imbalance to mitigate it.
310
+ eval_size (float): Portion of the train set being used as the validation set.
311
+ num_boost_round (int): Maximum number of boosting rounds.
312
+ early_stopping_rounds: (int): Number of rounds with no improvement on validation data before training stops.
313
+ verbose_eval (bool): Whether or not to display validation metrics during training.
314
+ threshold_prob (float): Discriminatory threshold used for classification metrics computation.
315
+ """
316
+
317
+ if not scale_pos_weight:
318
+ _ = self.xgb_params.pop('scale_pos_weight', None)
319
+ elif 'scale_pos_weight' not in self.xgb_params:
320
+ self.xgb_params['scale_pos_weight'] = float(np.sum(self.y_train == 0)) / np.sum(self.y_train == 1)
321
+
322
+ super()._fit(eval_size, num_boost_round, early_stopping_rounds, verbose_eval)
323
+
324
+ self.plot_performance(threshold_prob=threshold_prob)
325
+
326
+
327
+ def plot_performance(self, threshold_prob: float = 0.5):
328
+ """
329
+ Displays model performance.
330
+
331
+ Parameters:
332
+ threshold_prob (float): Discriminatory threshold used for classification metrics computation.
333
+ """
334
+
335
+ if self.model is None:
336
+ print("No XGBoost model was found. Please call .fit()")
337
+ return
338
+
339
+ dtrain = xgb.DMatrix(self.X_train, label=self.y_train)
340
+ dtest = xgb.DMatrix(self.X_test, label=self.y_test)
341
+
342
+ y_train_pred_proba = self.model.predict(dtrain)
343
+ y_test_pred_proba = self.model.predict(dtest)
344
+
345
+ _plot_classification_metrics(self.y_train, self.y_test, y_train_pred_proba, y_test_pred_proba, threshold_prob)
346
+
347
+
348
+ def compute_shap_and_predictions(self):
349
+ """
350
+ Computes model's SHAP values and predictions and instantiates the underlying Wrapshap object.
351
+ """
352
+ super().compute_shap_and_predictions()
353
+
354
+ self.y_pred_train = self.model.predict(
355
+ xgb.DMatrix(self.X_train, label=self.y_train), output_margin=True).reshape(-1,1).squeeze()
356
+ self.y_pred_test = self.model.predict(
357
+ xgb.DMatrix(self.X_test, label=self.y_test), output_margin=True).reshape(-1,1).squeeze()
358
+
359
+ self.wrapshap = Wrapshap(
360
+ X_train=self.X_train, X_test=self.X_test, y_train=self.y_train, y_test=self.y_test,
361
+ shap_train=self.shap_train, shap_test=self.shap_test,
362
+ y_pred_train=self.y_pred_train, y_pred_test=self.y_pred_test,
363
+ feature_names=self.feature_names,
364
+ task_type='classification'
365
+ )
366
+
367
+
368
+ class WrappedRegressorXGB(WrappedXGB):
369
+ """
370
+ This class extends the base WrappedXGB class and provides an XGBoost Regressor already conveniently wrapped
371
+ with Wrapshap. Objective will be squared error. Please see WrappedXGB for more information.
372
+ """
373
+
374
+ def __init__(
375
+ self, X: np.ndarray, y: np.ndarray, feature_names: list = None,
376
+ test_size: float = 0.2, random_state: int = 42,
377
+ **kwargs
378
+ ):
379
+ """
380
+ Parameters:
381
+ X (np.array): The original data as a 2D array.
382
+ y (np.array): The original target as a 1D array.
383
+ feature_names (list): The feature names.
384
+ test_size (float): Portion of the data being used as the test set.
385
+ random_state (int): The chosen random state for reproductibility.
386
+ **kwargs: Any parameters to add to the XGBoost configuration.
387
+ """
388
+
389
+ super().__init__(X, y, feature_names, random_state, **kwargs)
390
+
391
+ self.xgb_params['objective'] = 'reg:squarederror'
392
+
393
+ self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.X, self.y,
394
+ test_size=test_size,
395
+ random_state=random_state)
396
+
397
+ def fit(
398
+ self, eval_size: float = 0.2, num_boost_round: int = 10000,
399
+ early_stopping_rounds: int = 10, verbose_eval: bool = False
400
+ ):
401
+ """
402
+ Fits the classifier on training data.
403
+
404
+ Parameters:
405
+ eval_size (float): Portion of the train set being used as the validation set.
406
+ num_boost_round (int): Maximum number of boosting rounds.
407
+ early_stopping_rounds: (int): Number of rounds with no improvement on validation data before training stops.
408
+ verbose_eval (bool): Whether or not to display validation metrics during training.
409
+ threshold_prob (float): Discriminatory threshold used for classification metrics computation.
410
+ """
411
+
412
+ super()._fit(eval_size, num_boost_round, early_stopping_rounds, verbose_eval)
413
+ self.plot_performance()
414
+
415
+
416
+ def plot_performance(self):
417
+ """
418
+ Displays model performance.
419
+ """
420
+
421
+ if self.model is None:
422
+ print("No XGBoost model was found. Please call .fit()")
423
+ return
424
+
425
+ dtrain = xgb.DMatrix(self.X_train, label=self.y_train)
426
+ dtest = xgb.DMatrix(self.X_test, label=self.y_test)
427
+
428
+ y_train_pred = self.model.predict(dtrain)
429
+ y_test_pred = self.model.predict(dtest)
430
+
431
+ _plot_regression_metrics(self.y_train, self.y_test, y_train_pred, y_test_pred)
432
+
433
+
434
+ def compute_shap_and_predictions(self):
435
+ """
436
+ Computes model's SHAP values and predictions and instantiates the underlying Wrapshap object.
437
+ """
438
+
439
+ super().compute_shap_and_predictions()
440
+
441
+ self.y_pred_train = self.model.predict(
442
+ xgb.DMatrix(self.X_train, label=self.y_train)).reshape(-1,1).squeeze()
443
+ self.y_pred_test = self.model.predict(
444
+ xgb.DMatrix(self.X_test, label=self.y_test)).reshape(-1,1).squeeze()
445
+
446
+ self.wrapshap = Wrapshap(
447
+ X_train=self.X_train, X_test=self.X_test, y_train=self.y_train, y_test=self.y_test,
448
+ shap_train=self.shap_train, shap_test=self.shap_test,
449
+ y_pred_train=self.y_pred_train, y_pred_test=self.y_pred_test,
450
+ feature_names=self.feature_names,
451
+ task_type='regression'
452
+ )
453
+
454
+
455
+ # ---------------------------------------------------------------------------------------------------------------------
456
+ # Internal Functions - Below are utilities intended only for internal use.
457
+ # ---------------------------------------------------------------------------------------------------------------------
458
+
459
+
460
+ def _plot_classification_metrics(y_train, y_test, y_train_pred_proba, y_test_pred_proba, threshold_prob=0.5):
461
+ fig = plt.figure(figsize=(15, 15))
462
+
463
+ grid = plt.GridSpec(4, 2, wspace=0.4, hspace=0.3)
464
+
465
+ y_train_pred = (y_train_pred_proba >= threshold_prob).astype(int)
466
+ y_test_pred = (y_test_pred_proba >= threshold_prob).astype(int)
467
+
468
+ scores = {
469
+ 'Accuracy': (accuracy_score(y_train, y_train_pred), accuracy_score(y_test, y_test_pred)),
470
+ 'Precision': (precision_score(y_train, y_train_pred), precision_score(y_test, y_test_pred)),
471
+ 'Recall': (recall_score(y_train, y_train_pred), recall_score(y_test, y_test_pred)),
472
+ 'F1': (f1_score(y_train, y_train_pred), f1_score(y_test, y_test_pred))
473
+ }
474
+
475
+ ax1 = fig.add_subplot(grid[0:2, :])
476
+ fpr_train, tpr_train, _ = roc_curve(y_train, y_train_pred_proba)
477
+ roc_auc_train = auc(fpr_train, tpr_train)
478
+ fpr_test, tpr_test, _ = roc_curve(y_test, y_test_pred_proba)
479
+ roc_auc_test = auc(fpr_test, tpr_test)
480
+ ax1.plot(fpr_train, tpr_train, label=f'Train ROC (AUC = {roc_auc_train:.2f})')
481
+ ax1.plot(fpr_test, tpr_test, label=f'Test ROC (AUC = {roc_auc_test:.2f})')
482
+ ax1.plot([0, 1], [0, 1], color='navy', linestyle='--')
483
+ ax1.set_xlim([0.0, 1.0])
484
+ ax1.set_ylim([0.0, 1.05])
485
+ ax1.set_xlabel('False Positive Rate')
486
+ ax1.set_ylabel('True Positive Rate')
487
+ ax1.set_title(f'ROC curve (XGBoost)')
488
+
489
+ legend_handles = [mlines.Line2D([], [], color='blue', marker='_', linestyle='-', markersize=10, label=f'Train ROC (AUC = {roc_auc_train:.2f})'),
490
+ mlines.Line2D([], [], color='orange', marker='_', linestyle='-', markersize=10, label=f'Test ROC (AUC = {roc_auc_test:.2f})')]
491
+ for score_name, (train_score, test_score) in scores.items():
492
+ dummy_line = mlines.Line2D([], [], color='none', marker='_', linestyle='-', markersize=0,
493
+ label=f'{score_name}: train {train_score:.2f} / test {test_score:.2f}')
494
+ legend_handles.append(dummy_line)
495
+ ax1.legend(handles=legend_handles, loc="lower right")
496
+
497
+ ax2 = fig.add_subplot(grid[2, 0])
498
+ cf_matrix_train = confusion_matrix(y_train, y_train_pred)
499
+ sns.heatmap(cf_matrix_train, annot=True, cmap='Blues', fmt='g', ax=ax2)
500
+ ax2.set_xlabel('Predicted')
501
+ ax2.set_ylabel('True')
502
+ ax2.set_title('Confusion Matrix (Train)')
503
+
504
+ ax3 = fig.add_subplot(grid[2, 1])
505
+ cf_matrix_test = confusion_matrix(y_test, y_test_pred)
506
+ sns.heatmap(cf_matrix_test, annot=True, cmap='Blues', fmt='g', ax=ax3)
507
+ ax3.set_xlabel('Predicted')
508
+ ax3.set_ylabel('True')
509
+ ax3.set_title('Confusion Matrix (Test)')
510
+ plt.show()
511
+
512
+
513
+ def _plot_regression_metrics(y_train, y_test, y_train_pred, y_test_pred):
514
+ train_r2 = r2_score(y_train, y_train_pred)
515
+ test_r2 = r2_score(y_test, y_test_pred)
516
+ train_mae = mean_absolute_error(y_train, y_train_pred)
517
+ test_mae = mean_absolute_error(y_test, y_test_pred)
518
+
519
+ _, ax = plt.subplots(1, 2, figsize=(12, 6))
520
+
521
+ ax[0].scatter(y_train, y_train_pred, alpha=0.3)
522
+ ax[0].set_xlabel('True Values')
523
+ ax[0].set_ylabel('Predictions')
524
+ ax[0].set_title(f'Training Set: R2 = {train_r2:.2f}, MAE = {train_mae:.2f}')
525
+
526
+ ax[1].scatter(y_test, y_test_pred, alpha=0.3)
527
+ ax[1].set_xlabel('True Values')
528
+ ax[1].set_ylabel('Predictions')
529
+ ax[1].set_title(f'Test Set: R2 = {test_r2:.2f}, MAE = {test_mae:.2f}')
530
+
531
+ plt.show()