dragon-ml-toolbox 12.11.0__py3-none-any.whl → 12.13.0__py3-none-any.whl

{dragon_ml_toolbox-12.11.0.dist-info → dragon_ml_toolbox-12.13.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragon-ml-toolbox
-Version: 12.11.0
+Version: 12.13.0
 Summary: A collection of tools for data science and machine learning projects.
 Author-email: "Karl L. Loza Vidaurre" <luigiloza@gmail.com>
 License-Expression: MIT

{dragon_ml_toolbox-12.11.0.dist-info → dragon_ml_toolbox-12.13.0.dist-info}/RECORD

@@ -1,19 +1,19 @@
-dragon_ml_toolbox-12.11.0.dist-info/licenses/LICENSE,sha256=L35WDmmLZNTlJvxF6Vy7Uy4SYNi6rCfWUqlTHpoRMoU,1081
-dragon_ml_toolbox-12.11.0.dist-info/licenses/LICENSE-THIRD-PARTY.md,sha256=iy2r_R7wjzsCbz_Q_jMsp_jfZ6oP8XW9QhwzRBH0mGY,1904
+dragon_ml_toolbox-12.13.0.dist-info/licenses/LICENSE,sha256=L35WDmmLZNTlJvxF6Vy7Uy4SYNi6rCfWUqlTHpoRMoU,1081
+dragon_ml_toolbox-12.13.0.dist-info/licenses/LICENSE-THIRD-PARTY.md,sha256=iy2r_R7wjzsCbz_Q_jMsp_jfZ6oP8XW9QhwzRBH0mGY,1904
 ml_tools/ETL_cleaning.py,sha256=2VBRllV8F-ZiPylPp8Az2gwn5ztgazN0BH5OKnRUhV0,20402
 ml_tools/ETL_engineering.py,sha256=KfYqgsxupAx6e_TxwO1LZXeu5mFkIhVXJrNjP3CzIZc,54927
 ml_tools/GUI_tools.py,sha256=Va6ig-dHULPVRwQYYtH3fvY5XPIoqRcJpRW8oXC55Hw,45413
 ml_tools/MICE_imputation.py,sha256=X273Qlgoqqg7KTmoKd75YDyAPB0UIbTzGP3xsCmRh3E,11717
-ml_tools/ML_callbacks.py,sha256=-XRIZEy3CPJWTHcoReyIw53FZlTs3pWcTVVnncTQQSc,13909
-ml_tools/ML_datasetmaster.py,sha256=t6q6mU9lz2rYKTVPKjA7yZ5ImV7_NykiciHaYnqIEpA,30822
-ml_tools/ML_evaluation.py,sha256=tLswOPgH4G1KExSMn0876YtNkbxPh-W3J4MYOjomMWA,16208
-ml_tools/ML_evaluation_multi.py,sha256=6OZyQ4SM9ALh38mOABmiHgIQDWcovsD_iOo7Bg9YZCE,12516
+ml_tools/ML_callbacks.py,sha256=2ZazJjlbClP-ALc8q0ru2oalkugbhO3TFwPg4RFZpck,14056
+ml_tools/ML_datasetmaster.py,sha256=kedCGneR3S2zui0_JFZN6TBL5e69XWkdpkE_QohyqSM,31433
+ml_tools/ML_evaluation.py,sha256=h7fAtk0lS4gTqQ46fiVjucTvFlX4rsufKnEtate6Nu0,18381
+ml_tools/ML_evaluation_multi.py,sha256=Kn9n5lfxo7A0TvgIDMx8UHZCvzTqv1ViezzwJBF-ypM,15970
 ml_tools/ML_inference.py,sha256=ymFvncFsU10PExq87xnEj541DKV5ck0nMuK8ToJHzVQ,23067
-ml_tools/ML_models.py,sha256=pSCV6KbmVnPZr49Kbyg7g25CYaWBWJr6IinBHKgVKGw,28042
+ml_tools/ML_models.py,sha256=G64NPhYZfYvHTIUwkIrMrNLgfDTKJwqdc8jwesPqB9E,28090
 ml_tools/ML_optimization.py,sha256=es3TlQbY7RYgJMZnznkjYGbUxFnAqzZxE_g3_qLK9Q8,22960
 ml_tools/ML_scaler.py,sha256=tw6onj9o8_kk3FQYb930HUzvv1zsFZe2YZJdF3LtHkU,7538
 ml_tools/ML_simple_optimization.py,sha256=W2mce1XFCuiOHTOjOsCNbETISHn5MwYlYsTIXH5hMMo,18177
-ml_tools/ML_trainer.py,sha256=_g48w5Ak-wQr5fGHdJqlcpnzv3gWyL1ghkOhy9VOZbo,23930
+ml_tools/ML_trainer.py,sha256=UmCuKr_GzQGYqhEZ-kaRv9Buj44DsNyuOzmOM7Fw8N0,24569
 ml_tools/ML_utilities.py,sha256=EnKpPTnJ2qjZmz7kvows4Uu5CfSA7ByRmI1v2-KarKw,9337
 ml_tools/PSO_optimization.py,sha256=fVHeemqilBS0zrGV25E5yKwDlGdd2ZKa18d8CZ6Q6Fk,22961
 ml_tools/RNN_forecast.py,sha256=Qa2KoZfdAvSjZ4yE78N4BFXtr3tTr0Gx7tQJZPotsh0,1967
@@ -35,7 +35,7 @@ ml_tools/optimization_tools.py,sha256=P074YCuZzkqkONnAsM-Zb9DTX_i8cRkkJLpwAWz6CR
 ml_tools/path_manager.py,sha256=CyDU16pOKmC82jPubqJPT6EBt-u-3rGVbxyPIZCvDDY,18432
 ml_tools/serde.py,sha256=ll2mVC0sO2jIEdG3K6xMcgEN13N4YSb8VjviGvw_ers,4949
 ml_tools/utilities.py,sha256=OcAyV1tEcYAfOWlGjRgopsjDLxU3DcI5EynzvWV4q3A,15754
-dragon_ml_toolbox-12.11.0.dist-info/METADATA,sha256=VOs19HzZ0j8xvEuKO9sIMDCGIPPQA22x3Lnh2H9Mw9c,6167
-dragon_ml_toolbox-12.11.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dragon_ml_toolbox-12.11.0.dist-info/top_level.txt,sha256=wm-oxax3ciyez6VoO4zsFd-gSok2VipYXnbg3TH9PtU,9
-dragon_ml_toolbox-12.11.0.dist-info/RECORD,,
+dragon_ml_toolbox-12.13.0.dist-info/METADATA,sha256=p3-oOSqq1hhJj13KjIXeFnwBu3UTfBJu5mTDL9MCpdU,6167
+dragon_ml_toolbox-12.13.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dragon_ml_toolbox-12.13.0.dist-info/top_level.txt,sha256=wm-oxax3ciyez6VoO4zsFd-gSok2VipYXnbg3TH9PtU,9
+dragon_ml_toolbox-12.13.0.dist-info/RECORD,,

ml_tools/ML_callbacks.py

@@ -113,18 +113,19 @@ class TqdmProgressBar(Callback):
 class EarlyStopping(Callback):
     """
     Stop training when a monitored metric has stopped improving.
-    
-    Args:
-        monitor (str): Quantity to be monitored. Defaults to 'val_loss'.
-        min_delta (float): Minimum change in the monitored quantity to qualify as an improvement.
-        patience (int): Number of epochs with no improvement after which training will be stopped.
-        mode (str): One of {'auto', 'min', 'max'}. In 'min' mode, training will stop when the quantity
-                    monitored has stopped decreasing; in 'max' mode it will stop when the quantity
-                    monitored has stopped increasing; in 'auto' mode, the direction is automatically
-                    inferred from the name of the monitored quantity.
-        verbose (int): Verbosity mode.
     """
     def __init__(self, monitor: str=PyTorchLogKeys.VAL_LOSS, min_delta: float=0.0, patience: int=5, mode: Literal['auto', 'min', 'max']='auto', verbose: int=1):
+        """
+        Args:
+            monitor (str): Quantity to be monitored. Defaults to 'val_loss'.
+            min_delta (float): Minimum change in the monitored quantity to qualify as an improvement.
+            patience (int): Number of epochs with no improvement after which training will be stopped.
+            mode (str): One of {'auto', 'min', 'max'}. In 'min' mode, training will stop when the quantity
+                        monitored has stopped decreasing; in 'max' mode it will stop when the quantity
+                        monitored has stopped increasing; in 'auto' mode, the direction is automatically
+                        inferred from the name of the monitored quantity.
+            verbose (int): Verbosity mode.
+        """
         super().__init__()
         self.monitor = monitor
         self.patience = patience
@@ -188,22 +189,23 @@ class EarlyStopping(Callback):
 
 class ModelCheckpoint(Callback):
     """
-    Saves the model to a directory with automated filename generation and rotation. The filename includes the epoch and score.
-
-    - If `save_best_only` is True, it saves the single best model, deleting the
-      previous best. 
-    - If `save_best_only` is False, it keeps the 3 most recent checkpoints,
-      deleting the oldest ones automatically.
-
-    Args:
-        save_dir (str): Directory where checkpoint files will be saved.
-        monitor (str): Metric to monitor for `save_best_only=True`.
-        save_best_only (bool): If true, save only the best model.
-        mode (str): One of {'auto', 'min', 'max'}.
-        verbose (int): Verbosity mode.
+    Saves the model weights to a directory with automated filename generation and rotation. 
     """
     def __init__(self, save_dir: Union[str,Path], checkpoint_name: Optional[str]=None, monitor: str = PyTorchLogKeys.VAL_LOSS,
                  save_best_only: bool = True, mode: Literal['auto', 'min', 'max']= 'auto', verbose: int = 0):
+        """
+        - If `save_best_only` is True, it saves the single best model, deleting the previous best. 
+        - If `save_best_only` is False, it keeps the 3 most recent checkpoints, deleting the oldest ones automatically.
+
+        Args:
+            save_dir (str): Directory where checkpoint files will be saved.
+            checkpoint_name (str| None): If None, the filename will include the epoch and score.
+            monitor (str): Metric to monitor for `save_best_only=True`.
+            save_best_only (bool): If true, save only the best model.
+            mode (str): One of {'auto', 'min', 'max'}.
+            verbose (int): Verbosity mode.
+        """
+        
         super().__init__()
         self.save_dir = make_fullpath(save_dir, make=True, enforce="directory")
         if not self.save_dir.is_dir():
@@ -306,17 +308,16 @@ class ModelCheckpoint(Callback):
 class LRScheduler(Callback):
     """
     Callback to manage a PyTorch learning rate scheduler.
-
-    This callback automatically calls the scheduler's `step()` method at the
-    end of each epoch. It also logs a message when the learning rate changes.
-
-    Args:
-        scheduler: An initialized PyTorch learning rate scheduler.
-        monitor (str, optional): The metric to monitor for schedulers that
-                                 require it, like `ReduceLROnPlateau`.
-                                 Should match a key in the logs (e.g., 'val_loss').
     """
     def __init__(self, scheduler, monitor: Optional[str] = None):
+        """
+        This callback automatically calls the scheduler's `step()` method at the
+        end of each epoch. It also logs a message when the learning rate changes.
+
+        Args:
+            scheduler: An initialized PyTorch learning rate scheduler.
+            monitor (str, optional): The metric to monitor for schedulers that require it, like `ReduceLROnPlateau`. Should match a key in the logs (e.g., 'val_loss').
+        """
         super().__init__()
         self.scheduler = scheduler
         self.monitor = monitor

ml_tools/ML_datasetmaster.py

@@ -81,8 +81,7 @@ class _PytorchDataset(Dataset):
             _LOGGER.error(f"Dataset {self.__class__} has not been initialized with any target names.")
 
 
-# --- Abstract Base Class (New) ---
-# --- Abstract Base Class (Corrected) ---
+# --- Abstract Base Class ---
 class _BaseDatasetMaker(ABC):
     """
     Abstract base class for dataset makers. Contains shared logic for
@@ -150,6 +149,14 @@ class _BaseDatasetMaker(ABC):
     @property
     def target_names(self) -> list[str]:
         return self._target_names
+    
+    @property
+    def number_of_features(self) -> int:
+        return len(self._feature_names)
+    
+    @property
+    def number_of_targets(self) -> int:
+        return len(self._target_names)
 
     @property
     def id(self) -> Optional[str]:
@@ -180,14 +187,14 @@ class _BaseDatasetMaker(ABC):
                           filename=DatasetKeys.TARGET_NAMES,
                           verbose=verbose)
 
-    def save_scaler(self, save_dir: Union[str, Path], verbose: bool=True) -> None:
+    def save_scaler(self, directory: Union[str, Path], verbose: bool=True) -> None:
         """
         Saves the fitted PytorchScaler's state to a .pth file.
 
         The filename is automatically generated based on the dataset id.
         
         Args:
-            save_dir (str | Path): The directory where the scaler will be saved.
+            directory (str | Path): The directory where the scaler will be saved.
         """
         if not self.scaler: 
             _LOGGER.error("No scaler was fitted or provided.")
@@ -195,7 +202,7 @@ class _BaseDatasetMaker(ABC):
         if not self.id: 
             _LOGGER.error("Must set the dataset `id` before saving scaler.")
             raise ValueError()
-        save_path = make_fullpath(save_dir, make=True, enforce="directory")
+        save_path = make_fullpath(directory, make=True, enforce="directory")
         sanitized_id = sanitize_filename(self.id)
         filename = f"{DatasetKeys.SCALER_PREFIX}{sanitized_id}.pth"
         filepath = save_path / filename
@@ -203,6 +210,15 @@ class _BaseDatasetMaker(ABC):
         if verbose:
             _LOGGER.info(f"Scaler for dataset '{self.id}' saved as '{filepath.name}'.")
 
+    def save_artifacts(self, directory: Union[str, Path], verbose: bool=True) -> None:
+        """
+        Convenience method to save feature names, target names, and the scaler (if a scaler was fitted)
+        """
+        self.save_feature_names(directory=directory, verbose=verbose)
+        self.save_target_names(directory=directory, verbose=verbose)
+        if self.scaler is not None:
+            self.save_scaler(directory=directory, verbose=verbose)
+
 
 # Single target dataset
 class DatasetMaker(_BaseDatasetMaker):

ml_tools/ML_evaluation.py

@@ -18,7 +18,7 @@ from sklearn.metrics import (
 import torch
 import shap
 from pathlib import Path
-from typing import Union, Optional, List
+from typing import Union, Optional, List, Literal
 
 from .path_manager import make_fullpath
 from ._logger import _LOGGER
@@ -249,13 +249,15 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[s
     plt.savefig(hist_path)
     _LOGGER.info(f"📊 Residuals histogram saved as '{hist_path.name}'")
     plt.close(fig_hist)
-
+    
 
 def shap_summary_plot(model, 
                       background_data: Union[torch.Tensor,np.ndarray], 
                       instances_to_explain: Union[torch.Tensor,np.ndarray], 
                       feature_names: Optional[list[str]], 
-                      save_dir: Union[str, Path]):
+                      save_dir: Union[str, Path],
+                      device: torch.device = torch.device('cpu'),
+                      explainer_type: Literal['deep', 'kernel'] = 'deep'):
     """
     Calculates SHAP values and saves summary plots and data.
 
@@ -265,48 +267,85 @@ def shap_summary_plot(model,
         instances_to_explain (torch.Tensor): The specific data instances to explain.
         feature_names (list of str | None): Names of the features for plot labeling.
         save_dir (str | Path): Directory to save SHAP artifacts.
+        device (torch.device): The torch device for SHAP calculations.
+        explainer_type (Literal['deep', 'kernel']): The explainer to use.
+            - 'deep': (Default) Uses shap.DeepExplainer. Fast and efficient for
+              PyTorch models.
+            - 'kernel': Uses shap.KernelExplainer. Model-agnostic but EXTREMELY
+              slow and memory-intensive.
     """
-    # everything to numpy
-    if isinstance(background_data, np.ndarray):
-        background_data_np = background_data
-    else:
-        background_data_np = background_data.numpy()
-        
-    if isinstance(instances_to_explain, np.ndarray):
-        instances_to_explain_np = instances_to_explain
-    else:
-        instances_to_explain_np = instances_to_explain.numpy()
     
-    # --- Data Validation Step ---
-    if np.isnan(background_data_np).any() or np.isnan(instances_to_explain_np).any():
-        _LOGGER.error("Input data for SHAP contains NaN values. Aborting explanation.")
-        return
-    
-    print("\n--- SHAP Value Explanation ---")
+    print(f"\n--- SHAP Value Explanation Using {explainer_type.upper()} Explainer ---")
     
     model.eval()
-    model.cpu()
-    
-    # 1. Summarize the background data.
-    # Summarize the background data using k-means. 10-50 clusters is a good starting point.
-    background_summary = shap.kmeans(background_data_np, 30) 
-    
-    # 2. Define a prediction function wrapper that SHAP can use. It must take a numpy array and return a numpy array.
-    def prediction_wrapper(x_np: np.ndarray) -> np.ndarray:
-        # Convert numpy data to torch tensor
-        x_torch = torch.from_numpy(x_np).float()
-        with torch.no_grad():
-            # Get model output
-            output = model(x_torch)
-        # Return as numpy array
-        return output.cpu().numpy().flatten()
-
-    # 3. Create the KernelExplainer
-    explainer = shap.KernelExplainer(prediction_wrapper, background_summary)
+    # model.cpu() # Run explanations on CPU
     
-    print("Calculating SHAP values with KernelExplainer...")
-    shap_values = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
+    shap_values = None
+    instances_to_explain_np = None
+
+    if explainer_type == 'deep':
+        # --- 1. Use DeepExplainer (Preferred) ---
+        
+        # Ensure data is torch.Tensor
+        if isinstance(background_data, np.ndarray):
+            background_data = torch.from_numpy(background_data).float()
+        if isinstance(instances_to_explain, np.ndarray):
+            instances_to_explain = torch.from_numpy(instances_to_explain).float()
+            
+        if torch.isnan(background_data).any() or torch.isnan(instances_to_explain).any():
+            _LOGGER.error("Input data for SHAP contains NaN values. Aborting explanation.")
+            return
+
+        background_data = background_data.to(device)
+        instances_to_explain = instances_to_explain.to(device)
+
+        explainer = shap.DeepExplainer(model, background_data)
+        # print("Calculating SHAP values with DeepExplainer...")
+        shap_values = explainer.shap_values(instances_to_explain)
+        instances_to_explain_np = instances_to_explain.cpu().numpy()
+
+    elif explainer_type == 'kernel':
+        # --- 2. Use KernelExplainer (Slow Fallback) ---
+        _LOGGER.warning(
+            "Using KernelExplainer. This is memory-intensive and slow. "
+            "Consider reducing 'n_samples' if the process terminates unexpectedly."
+        )
+
+        # Ensure data is np.ndarray
+        if isinstance(background_data, torch.Tensor):
+            background_data_np = background_data.cpu().numpy()
+        else:
+            background_data_np = background_data
+            
+        if isinstance(instances_to_explain, torch.Tensor):
+            instances_to_explain_np = instances_to_explain.cpu().numpy()
+        else:
+            instances_to_explain_np = instances_to_explain
+        
+        if np.isnan(background_data_np).any() or np.isnan(instances_to_explain_np).any():
+            _LOGGER.error("Input data for SHAP contains NaN values. Aborting explanation.")
+            return
+        
+        # Summarize background data
+        background_summary = shap.kmeans(background_data_np, 30) 
+        
+        def prediction_wrapper(x_np: np.ndarray) -> np.ndarray:
+            x_torch = torch.from_numpy(x_np).float().to(device)
+            with torch.no_grad():
+                output = model(x_torch)
+            # Return as numpy array
+            return output.cpu().numpy()
+
+        explainer = shap.KernelExplainer(prediction_wrapper, background_summary)
+        # print("Calculating SHAP values with KernelExplainer...")
+        shap_values = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
+        # instances_to_explain_np is already set
     
+    else:
+        _LOGGER.error(f"Invalid explainer_type: '{explainer_type}'. Must be 'deep' or 'kernel'.")
+        raise ValueError()
+
+    # --- 3. Plotting and Saving ---
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
     plt.ioff()
     
@@ -326,8 +365,9 @@ def shap_summary_plot(model,
     shap.summary_plot(shap_values, instances_to_explain_np, feature_names=feature_names, plot_type="dot", show=False)
     ax = plt.gca()
     ax.set_xlabel("SHAP Value Impact", labelpad=10)
-    cb = plt.gcf().axes[-1]
-    cb.set_ylabel("", size=1)
+    if plt.gcf().axes and len(plt.gcf().axes) > 1:
+        cb = plt.gcf().axes[-1]
+        cb.set_ylabel("", size=1)
     plt.title("SHAP Feature Importance")
     plt.tight_layout()
     plt.savefig(dot_path)
@@ -337,8 +377,14 @@ def shap_summary_plot(model,
     # Save Summary Data to CSV
     shap_summary_filename = SHAPKeys.SAVENAME + ".csv"
     summary_path = save_dir_path / shap_summary_filename
-    # Ensure the array is 1D before creating the DataFrame
-    mean_abs_shap = np.abs(shap_values).mean(axis=0).flatten()
+    
+    # Handle multi-class (list of arrays) vs. regression (single array)
+    if isinstance(shap_values, list):
+        mean_abs_shap = np.abs(np.stack(shap_values)).mean(axis=0).mean(axis=0)
+    else:
+        mean_abs_shap = np.abs(shap_values).mean(axis=0)
+
+    mean_abs_shap = mean_abs_shap.flatten()
     
     if feature_names is None:
         feature_names = [f'feature_{i}' for i in range(len(mean_abs_shap))]
@@ -351,7 +397,7 @@ def shap_summary_plot(model,
     summary_df.to_csv(summary_path, index=False)
     
     _LOGGER.info(f"📝 SHAP summary data saved as '{summary_path.name}'")
-    plt.ion()
+    plt.ion()    
 
 
 def plot_attention_importance(weights: List[torch.Tensor], feature_names: Optional[List[str]], save_dir: Union[str, Path], top_n: int = 10):

ml_tools/ML_evaluation_multi.py

@@ -19,11 +19,12 @@ from sklearn.metrics import (
     jaccard_score
 )
 from pathlib import Path
-from typing import Union, List
+from typing import Union, List, Literal
 
 from .path_manager import make_fullpath, sanitize_filename
 from ._logger import _LOGGER
 from ._script_info import _script_info
+from .keys import SHAPKeys
 
 
 __all__ = [
@@ -231,10 +232,12 @@ def multi_target_shap_summary_plot(
     instances_to_explain: Union[torch.Tensor, np.ndarray],
     feature_names: List[str],
     target_names: List[str],
-    save_dir: Union[str, Path]
+    save_dir: Union[str, Path],
+    device: torch.device = torch.device('cpu'),
+    explainer_type: Literal['deep', 'kernel'] = 'deep'
 ):
     """
-    Calculates SHAP values for a multi-target model and saves summary plots for each target.
+    Calculates SHAP values for a multi-target model and saves summary plots and data for each target.
 
     Args:
         model (torch.nn.Module): The trained PyTorch model.
@@ -243,40 +246,91 @@ def multi_target_shap_summary_plot(
         feature_names (List[str]): Names of the features for plot labeling.
         target_names (List[str]): Names of the output targets.
         save_dir (str | Path): Directory to save SHAP artifacts.
+        device (torch.device): The torch device for SHAP calculations.
+        explainer_type (Literal['deep', 'kernel']): The explainer to use.
+            - 'deep': (Default) Uses shap.DeepExplainer. Fast and efficient.
+            - 'kernel': Uses shap.KernelExplainer. Model-agnostic but slow and memory-intensive.
     """
-    # Convert all data to numpy
-    background_data_np = background_data.numpy() if isinstance(background_data, torch.Tensor) else background_data
-    instances_to_explain_np = instances_to_explain.numpy() if isinstance(instances_to_explain, torch.Tensor) else instances_to_explain
-
-    if np.isnan(background_data_np).any() or np.isnan(instances_to_explain_np).any():
-        _LOGGER.error("Input data for SHAP contains NaN values. Aborting explanation.")
-        return
-
-    _LOGGER.info("--- Multi-Target SHAP Value Explanation ---")
+    _LOGGER.info(f"--- Multi-Target SHAP Value Explanation (Using: {explainer_type.upper()}Explainer) ---")
     model.eval()
-    model.cpu()
-
-    # 1. Summarize the background data.
-    background_summary = shap.kmeans(background_data_np, 30)
-
-    # 2. Define a prediction function wrapper for the multi-target model.
-    def prediction_wrapper(x_np: np.ndarray) -> np.ndarray:
-        x_torch = torch.from_numpy(x_np).float()
-        with torch.no_grad():
-            output = model(x_torch)
-        return output.cpu().numpy()
-
-    # 3. Create the KernelExplainer.
-    explainer = shap.KernelExplainer(prediction_wrapper, background_summary)
-
-    print("Calculating SHAP values with KernelExplainer...")
-    # For multi-output models, shap_values is a list of arrays.
-    shap_values_list = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
+    # model.cpu()
+
+    shap_values_list = None
+    instances_to_explain_np = None
+
+    if explainer_type == 'deep':
+        # --- 1. Use DeepExplainer (Preferred) ---
+        
+        # Ensure data is torch.Tensor
+        if isinstance(background_data, np.ndarray):
+            background_data = torch.from_numpy(background_data).float()
+        if isinstance(instances_to_explain, np.ndarray):
+            instances_to_explain = torch.from_numpy(instances_to_explain).float()
+            
+        if torch.isnan(background_data).any() or torch.isnan(instances_to_explain).any():
+            _LOGGER.error("Input data for SHAP contains NaN values. Aborting explanation.")
+            return
+
+        background_data = background_data.to(device)
+        instances_to_explain = instances_to_explain.to(device)
+
+        explainer = shap.DeepExplainer(model, background_data)
+        print("Calculating SHAP values with DeepExplainer...")
+        # DeepExplainer returns a list of arrays for multi-output models
+        shap_values_list = explainer.shap_values(instances_to_explain)
+        instances_to_explain_np = instances_to_explain.cpu().numpy()
+
+    elif explainer_type == 'kernel':
+        # --- 2. Use KernelExplainer (Slow Fallback) ---
+        _LOGGER.warning(
+            "Using KernelExplainer. This is memory-intensive and slow. "
+            "Consider reducing 'n_samples' if the process terminates."
+        )
+        
+        # Convert all data to numpy
+        background_data_np = background_data.numpy() if isinstance(background_data, torch.Tensor) else background_data
+        instances_to_explain_np = instances_to_explain.numpy() if isinstance(instances_to_explain, torch.Tensor) else instances_to_explain
+
+        if np.isnan(background_data_np).any() or np.isnan(instances_to_explain_np).any():
+            _LOGGER.error("Input data for SHAP contains NaN values. Aborting explanation.")
+            return
+
+        background_summary = shap.kmeans(background_data_np, 30)
+
+        def prediction_wrapper(x_np: np.ndarray) -> np.ndarray:
+            x_torch = torch.from_numpy(x_np).float().to(device)
+            with torch.no_grad():
+                output = model(x_torch)
+            return output.cpu().numpy() # Return full multi-output array
+
+        explainer = shap.KernelExplainer(prediction_wrapper, background_summary)
+        print("Calculating SHAP values with KernelExplainer...")
+        # KernelExplainer also returns a list of arrays for multi-output models
+        shap_values_list = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
+        # instances_to_explain_np is already set
+        
+    else:
+        _LOGGER.error(f"Invalid explainer_type: '{explainer_type}'. Must be 'deep' or 'kernel'.")
+        raise ValueError("Invalid explainer_type")
+
+    # --- 3. Plotting and Saving (Common Logic) ---
+    
+    if shap_values_list is None or instances_to_explain_np is None:
+        _LOGGER.error("SHAP value calculation failed. Aborting plotting.")
+        return
+        
+    # Ensure number of SHAP value arrays matches number of target names
+    if len(shap_values_list) != len(target_names):
+        _LOGGER.error(
+            f"SHAP explanation mismatch: Model produced {len(shap_values_list)} "
+            f"outputs, but {len(target_names)} target_names were provided."
+        )
+        return
 
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
     plt.ioff()
 
-    # 4. Iterate through each target's SHAP values and generate plots.
+    # Iterate through each target's SHAP values and generate plots.
     for i, target_name in enumerate(target_names):
         print(f"  -> Generating SHAP plots for target: '{target_name}'")
         shap_values_for_target = shap_values_list[i]
@@ -293,11 +347,28 @@ def multi_target_shap_summary_plot(
         # Save Dot Plot for the target
         shap.summary_plot(shap_values_for_target, instances_to_explain_np, feature_names=feature_names, plot_type="dot", show=False)
         plt.title(f"SHAP Feature Importance for '{target_name}'")
+        if plt.gcf().axes and len(plt.gcf().axes) > 1:
+            cb = plt.gcf().axes[-1]
+            cb.set_ylabel("", size=1)
         plt.tight_layout()
         dot_path = save_dir_path / f"shap_dot_plot_{sanitized_target_name}.svg"
         plt.savefig(dot_path)
         plt.close()
-
+        
+        # --- Save Summary Data to CSV for this target ---
+        shap_summary_filename = f"{SHAPKeys.SAVENAME}_{sanitized_target_name}.csv"
+        summary_path = save_dir_path / shap_summary_filename
+        
+        # For a specific target, shap_values_for_target is just a 2D array
+        mean_abs_shap = np.abs(shap_values_for_target).mean(axis=0).flatten()
+        
+        summary_df = pd.DataFrame({
+            SHAPKeys.FEATURE_COLUMN: feature_names,
+            SHAPKeys.SHAP_VALUE_COLUMN: mean_abs_shap
+        }).sort_values(SHAPKeys.SHAP_VALUE_COLUMN, ascending=False)
+        
+        summary_df.to_csv(summary_path, index=False)
+        
     plt.ion()
     _LOGGER.info(f"All SHAP plots saved to '{save_dir_path.name}'")
 

ml_tools/ML_models.py

@@ -304,7 +304,7 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     def __init__(self, *,
                  in_features: int,
                  out_targets: int,
-                 categorical_map: Dict[int, int],
+                 categorical_index_map: Dict[int, int],
                  embedding_dim: int = 32,
                  num_heads: int = 8,
                  num_layers: int = 6,
@@ -313,7 +313,7 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
         Args:
             in_features (int): The total number of columns in the input data (features).
             out_targets (int): Number of output targets (1 for regression).
-            categorical_map (Dict[int, int]): Maps categorical column index to its cardinality (number of unique categories).
+            categorical_index_map (Dict[int, int]): Maps categorical column index to its cardinality (number of unique categories).
             embedding_dim (int): The dimension for all feature embeddings. Must be divisible by num_heads.
             num_heads (int): The number of heads in the multi-head attention mechanism.
             num_layers (int): The number of sub-encoder-layers in the transformer encoder.
@@ -340,20 +340,20 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
         super().__init__()
         
          # --- Validation ---
-        if categorical_map and max(categorical_map.keys()) >= in_features:
-            _LOGGER.error(f"A categorical index ({max(categorical_map.keys())}) is out of bounds for the provided input features ({in_features}).")
+        if categorical_index_map and max(categorical_index_map.keys()) >= in_features:
+            _LOGGER.error(f"A categorical index ({max(categorical_index_map.keys())}) is out of bounds for the provided input features ({in_features}).")
             raise ValueError()
         
         # --- Derive numerical indices ---
         all_indices = set(range(in_features))
-        categorical_indices_set = set(categorical_map.keys())
+        categorical_indices_set = set(categorical_index_map.keys())
         numerical_indices = sorted(list(all_indices - categorical_indices_set))
         
         # --- Save configuration ---
         self.in_features = in_features
         self.out_targets = out_targets
         self.numerical_indices = numerical_indices
-        self.categorical_map = categorical_map
+        self.categorical_map = categorical_index_map
         self.embedding_dim = embedding_dim
         self.num_heads = num_heads
         self.num_layers = num_layers
@@ -362,7 +362,7 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
         # --- 1. Feature Tokenizer ---
         self.tokenizer = _FeatureTokenizer(
             numerical_indices=numerical_indices,
-            categorical_map=categorical_map,
+            categorical_map=categorical_index_map,
             embedding_dim=embedding_dim
         )
         

ml_tools/ML_trainer.py

@@ -340,9 +340,10 @@ class MLTrainer:
     def explain(self,
                 save_dir: Union[str,Path], 
                 explain_dataset: Optional[Dataset] = None, 
-                n_samples: int = 1000,
+                n_samples: int = 300,
                 feature_names: Optional[List[str]] = None,
-                target_names: Optional[List[str]] = None):
+                target_names: Optional[List[str]] = None,
+                explainer_type: Literal['deep', 'kernel'] = 'deep'):
         """
         Explains model predictions using SHAP and saves all artifacts.
 
@@ -359,6 +360,9 @@ class MLTrainer:
             feature_names (list[str] | None): Feature names.
             target_names (list[str] | None): Target names for multi-target tasks.
             save_dir (str | Path): Directory to save all SHAP artifacts.
+            explainer_type (Literal['deep', 'kernel']): The explainer to use.
+                - 'deep': (Default) Uses shap.DeepExplainer. Fast and efficient for PyTorch models.
+                - 'kernel': Uses shap.KernelExplainer. Model-agnostic but EXTREMELY slow and memory-intensive. Use with a very low 'n_samples'< 100.
         """
         # Internal helper to create a dataloader and get a random sample
         def _get_random_sample(dataset: Dataset, num_samples: int):
@@ -410,6 +414,9 @@ class MLTrainer:
             else:
                 _LOGGER.error("Could not extract `feature_names` from the dataset. It must be provided if the dataset object does not have a `feature_names` attribute.")
                 raise ValueError()
+            
+        # move model to device
+        self.model.to(self.device)
 
         # 3. Call the plotting function
         if self.kind in ["regression", "classification"]:
@@ -418,7 +425,9 @@ class MLTrainer:
                 background_data=background_data,
                 instances_to_explain=instances_to_explain,
                 feature_names=feature_names,
-                save_dir=save_dir
+                save_dir=save_dir,
+                explainer_type=explainer_type,
+                device=self.device
             )
         elif self.kind in ["multi_target_regression", "multi_label_classification"]:
             # try to get target names
@@ -442,7 +451,9 @@ class MLTrainer:
                 instances_to_explain=instances_to_explain,
                 feature_names=feature_names, # type: ignore
                 target_names=target_names, # type: ignore
-                save_dir=save_dir
+                save_dir=save_dir,
+                explainer_type=explainer_type,
+                device=self.device
             )
 
     def _attention_helper(self, dataloader: DataLoader):