dragon-ml-toolbox 14.3.1__py3-none-any.whl → 16.0.0__py3-none-any.whl

ml_tools/ML_evaluation.py

@@ -21,10 +21,11 @@ from pathlib import Path
 from typing import Union, Optional, List, Literal
 import warnings
 
-from .path_manager import make_fullpath
+from .path_manager import make_fullpath, sanitize_filename
 from ._logger import _LOGGER
 from ._script_info import _script_info
-from .keys import SHAPKeys, PyTorchLogKeys
+from ._keys import SHAPKeys, PyTorchLogKeys
+from .ML_configuration import RegressionMetricsFormat, ClassificationMetricsFormat
 
 
 __all__ = [
@@ -35,10 +36,13 @@ __all__ = [
     "plot_attention_importance"
 ]
 
+DPI_value = 250
+
 
 def plot_losses(history: dict, save_dir: Union[str, Path]):
     """
     Plots training & validation loss curves from a history object.
+    Also plots the learning rate if available in the history.
 
     Args:
         history (dict): A dictionary containing 'train_loss' and 'val_loss'.
@@ -46,29 +50,52 @@ def plot_losses(history: dict, save_dir: Union[str, Path]):
     """
     train_loss = history.get(PyTorchLogKeys.TRAIN_LOSS, [])
     val_loss = history.get(PyTorchLogKeys.VAL_LOSS, [])
+    lr_history = history.get(PyTorchLogKeys.LEARNING_RATE, [])
     
     if not train_loss and not val_loss:
-        print("Warning: Loss history is empty or incomplete. Cannot plot.")
+        _LOGGER.warning("Loss history is empty or incomplete. Cannot plot.")
         return
 
-    fig, ax = plt.subplots(figsize=(10, 5), dpi=100)
+    fig, ax = plt.subplots(figsize=(10, 5), dpi=DPI_value)
+    
+    # --- Plot Losses (Left Y-axis) ---
+    line_handles = [] # To store line objects for the legend
     
     # Plot training loss only if data for it exists
     if train_loss:
         epochs = range(1, len(train_loss) + 1)
-        ax.plot(epochs, train_loss, 'o-', label='Training Loss')
+        line1, = ax.plot(epochs, train_loss, 'o-', label='Training Loss', color='tab:blue')
+        line_handles.append(line1)
     
     # Plot validation loss only if data for it exists
     if val_loss:
         epochs = range(1, len(val_loss) + 1)
-        ax.plot(epochs, val_loss, 'o-', label='Validation Loss')
+        line2, = ax.plot(epochs, val_loss, 'o-', label='Validation Loss', color='tab:orange')
+        line_handles.append(line2)
     
     ax.set_title('Training and Validation Loss')
     ax.set_xlabel('Epochs')
-    ax.set_ylabel('Loss')
-    ax.legend()
-    ax.grid(True)
-    plt.tight_layout()
+    ax.set_ylabel('Loss', color='tab:blue')
+    ax.tick_params(axis='y', labelcolor='tab:blue')
+    ax.grid(True, linestyle='--')
+    
+    # --- Plot Learning Rate (Right Y-axis) ---
+    if lr_history:
+        ax2 = ax.twinx() # Create a second y-axis
+        epochs = range(1, len(lr_history) + 1)
+        line3, = ax2.plot(epochs, lr_history, 'g--', label='Learning Rate')
+        line_handles.append(line3)
+        
+        ax2.set_ylabel('Learning Rate', color='g')
+        ax2.tick_params(axis='y', labelcolor='g')
+        # Use scientific notation if the LR is very small
+        ax2.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
+    
+    # Combine legends from both axes
+    ax.legend(handles=line_handles, loc='best')
+    
+    # ax.grid(True)
+    plt.tight_layout()    
     
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
     save_path = save_dir_path / "loss_plot.svg"
@@ -78,23 +105,49 @@ def plot_losses(history: dict, save_dir: Union[str, Path]):
     plt.close(fig)
 
 
-def classification_metrics(save_dir: Union[str, Path], y_true: np.ndarray, y_pred: np.ndarray, y_prob: Optional[np.ndarray] = None, 
-                           cmap: str = "Blues"):
+def classification_metrics(save_dir: Union[str, Path], 
+                           y_true: np.ndarray, 
+                           y_pred: np.ndarray, 
+                           y_prob: Optional[np.ndarray] = None, 
+                           config: Optional[ClassificationMetricsFormat] = None):
     """
     Saves classification metrics and plots.
 
     Args:
         y_true (np.ndarray): Ground truth labels.
         y_pred (np.ndarray): Predicted labels.
-        y_prob (np.ndarray, optional): Predicted probabilities for ROC curve.
-        cmap (str): Colormap for the confusion matrix.
+        y_prob (np.ndarray): Predicted probabilities for ROC curve.
+        config (ClassificationMetricsFormat): Formatting configuration object.
         save_dir (str | Path): Directory to save plots.
     """
-    print("--- Classification Report ---")
+    # --- Parse Config or use defaults ---
+    if config is None:
+        # Create a default config if one wasn't provided
+        config = ClassificationMetricsFormat()
+    
+    original_rc_params = plt.rcParams.copy()
+    plt.rcParams.update({'font.size': config.font_size})
+    
+    # print("--- Classification Report ---")
+    
+    # --- Parse class_map ---
+    map_labels = None
+    map_display_labels = None
+    if config.class_map:
+        # Sort the map by its values (the indices) to ensure correct order
+        try:
+            sorted_items = sorted(config.class_map.items(), key=lambda item: item[1])
+            map_labels = [item[1] for item in sorted_items]
+            map_display_labels = [item[0] for item in sorted_items]
+        except Exception as e:
+            _LOGGER.warning(f"Could not parse 'class_map': {e}")
+            map_labels = None
+            map_display_labels = None
+    
     # Generate report as both text and dictionary
-    report_text: str = classification_report(y_true, y_pred) # type: ignore
-    report_dict: dict = classification_report(y_true, y_pred, output_dict=True) # type: ignore
-    print(report_text)
+    report_text: str = classification_report(y_true, y_pred, labels=map_labels, target_names=map_display_labels) # type: ignore
+    report_dict: dict = classification_report(y_true, y_pred, output_dict=True, labels=map_labels, target_names=map_display_labels) # type: ignore
+    # print(report_text)
     
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
     # Save text report
@@ -104,8 +157,15 @@ def classification_metrics(save_dir: Union[str, Path], y_true: np.ndarray, y_pre
 
     # --- Save Classification Report Heatmap ---
     try:
-        plt.figure(figsize=(8, 6), dpi=100)
-        sns.heatmap(pd.DataFrame(report_dict).iloc[:-1, :].T, annot=True, cmap='viridis', fmt='.2f')
+        plt.figure(figsize=(8, 6), dpi=DPI_value)
+        sns.set_theme(font_scale=1.2) # Scale seaborn font
+        sns.heatmap(pd.DataFrame(report_dict).iloc[:-1, :].T, 
+                    annot=True, 
+                    cmap=config.cmap, 
+                    fmt='.2f',
+                    vmin=0.0,
+                    vmax=1.0)
+        sns.set_theme(font_scale=1.0) # Reset seaborn scale
         plt.title("Classification Report")
         plt.tight_layout()
         heatmap_path = save_dir_path / "classification_report_heatmap.svg"
@@ -114,72 +174,224 @@ def classification_metrics(save_dir: Union[str, Path], y_true: np.ndarray, y_pre
         plt.close()
     except Exception as e:
         _LOGGER.error(f"Could not generate classification report heatmap: {e}")
-
+        
+    # --- labels for Confusion Matrix ---
+    plot_labels = map_labels
+    plot_display_labels = map_display_labels
+    
     # Save Confusion Matrix
-    fig_cm, ax_cm = plt.subplots(figsize=(6, 6), dpi=100)
-    ConfusionMatrixDisplay.from_predictions(y_true, y_pred, cmap=cmap, ax=ax_cm)
+    fig_cm, ax_cm = plt.subplots(figsize=(6, 6), dpi=DPI_value)
+    disp_ = ConfusionMatrixDisplay.from_predictions(y_true, 
+                                            y_pred, 
+                                            cmap=config.cmap, 
+                                            ax=ax_cm, 
+                                            normalize='true',
+                                            labels=plot_labels,
+                                            display_labels=plot_display_labels)
+    
+    disp_.im_.set_clim(vmin=0.0, vmax=1.0)
+    
+    # Turn off gridlines
+    ax_cm.grid(False)
+    
+    # Manually update font size of cell texts
+    for text in ax_cm.texts:
+        text.set_fontsize(config.font_size)
+
+    fig_cm.tight_layout()
+    
     ax_cm.set_title("Confusion Matrix")
     cm_path = save_dir_path / "confusion_matrix.svg"
     plt.savefig(cm_path)
     _LOGGER.info(f"❇️ Confusion matrix saved as '{cm_path.name}'")
     plt.close(fig_cm)
 
-    # Plotting logic for ROC and PR Curves
-    if y_prob is not None and y_prob.ndim > 1 and y_prob.shape[1] >= 2:
-        # Use probabilities of the positive class
-        y_score = y_prob[:, 1]
+
+    # Plotting logic for ROC, PR, and Calibration Curves
+    if y_prob is not None and y_prob.ndim == 2:
+        num_classes = y_prob.shape[1]
         
-        # --- Save ROC Curve ---
-        fpr, tpr, _ = roc_curve(y_true, y_score)
-        auc = roc_auc_score(y_true, y_score)
-        fig_roc, ax_roc = plt.subplots(figsize=(6, 6), dpi=100)
-        ax_roc.plot(fpr, tpr, label=f'AUC = {auc:.2f}')
-        ax_roc.plot([0, 1], [0, 1], 'k--')
-        ax_roc.set_title('Receiver Operating Characteristic (ROC) Curve')
-        ax_roc.set_xlabel('False Positive Rate')
-        ax_roc.set_ylabel('True Positive Rate')
-        ax_roc.legend(loc='lower right')
-        ax_roc.grid(True)
-        roc_path = save_dir_path / "roc_curve.svg"
-        plt.savefig(roc_path)
-        _LOGGER.info(f"📈 ROC curve saved as '{roc_path.name}'")
-        plt.close(fig_roc)
-
-        # --- Save Precision-Recall Curve ---
-        precision, recall, _ = precision_recall_curve(y_true, y_score)
-        ap_score = average_precision_score(y_true, y_score)
-        fig_pr, ax_pr = plt.subplots(figsize=(6, 6), dpi=100)
-        ax_pr.plot(recall, precision, label=f'AP = {ap_score:.2f}')
-        ax_pr.set_title('Precision-Recall Curve')
-        ax_pr.set_xlabel('Recall')
-        ax_pr.set_ylabel('Precision')
-        ax_pr.legend(loc='lower left')
-        ax_pr.grid(True)
-        pr_path = save_dir_path / "pr_curve.svg"
-        plt.savefig(pr_path)
-        _LOGGER.info(f"📈 PR curve saved as '{pr_path.name}'")
-        plt.close(fig_pr)
+        # --- Determine which classes to loop over ---
+        class_indices_to_plot = []
+        plot_titles = []
+        save_suffixes = []
+
+        if num_classes == 2:
+            # Binary case: Only plot for the positive class (index 1)
+            class_indices_to_plot = [1]
+            plot_titles = [""] # No extra title
+            save_suffixes = [""] # No extra suffix
+            _LOGGER.debug("Generating binary classification plots (ROC, PR, Calibration).")
         
-        # --- Save Calibration Plot ---
-        if y_prob.ndim > 1 and y_prob.shape[1] >= 2:
-            y_score = y_prob[:, 1] # Use probabilities of the positive class
+        elif num_classes > 2:
+            _LOGGER.debug(f"Generating One-vs-Rest plots for {num_classes} classes.")
+            # Multiclass case: Plot for every class (One-vs-Rest)
+            class_indices_to_plot = list(range(num_classes))
             
-            fig_cal, ax_cal = plt.subplots(figsize=(8, 8), dpi=100)
-            CalibrationDisplay.from_predictions(y_true, y_score, n_bins=15, ax=ax_cal)
+            # --- Use class_map names if available ---
+            use_generic_names = True
+            if map_display_labels and len(map_display_labels) == num_classes:
+                try:
+                    # Ensure labels are safe for filenames
+                    safe_names = [sanitize_filename(name) for name in map_display_labels]
+                    plot_titles = [f" ({name} vs. Rest)" for name in map_display_labels]
+                    save_suffixes = [f"_{safe_names[i]}" for i in class_indices_to_plot]
+                    use_generic_names = False
+                except Exception as e:
+                    _LOGGER.warning(f"Failed to use 'class_map' for plot titles: {e}. Reverting to generic names.")
+                    use_generic_names = True
             
-            ax_cal.set_title('Reliability Curve')
+            if use_generic_names:
+                plot_titles = [f" (Class {i} vs. Rest)" for i in class_indices_to_plot]
+                save_suffixes = [f"_class_{i}" for i in class_indices_to_plot]
+        
+        else:
+            # Should not happen, but good to check
+            _LOGGER.warning(f"Probability array has invalid shape {y_prob.shape}. Skipping ROC/PR/Calibration plots.")
+
+        # --- Loop and generate plots ---
+        for i, class_index in enumerate(class_indices_to_plot):
+            plot_title = plot_titles[i]
+            save_suffix = save_suffixes[i]
+
+            # Get scores for the current class
+            y_score = y_prob[:, class_index]
+            
+            # Binarize y_true for the current class
+            y_true_binary = (y_true == class_index).astype(int)
+            
+            # --- Save ROC Curve ---
+            fpr, tpr, thresholds = roc_curve(y_true_binary, y_score)
+            
+            try:
+                # Calculate Youden's J statistic (tpr - fpr)
+                J = tpr - fpr
+                # Find the index of the best threshold
+                best_index = np.argmax(J)
+                optimal_threshold = thresholds[best_index]
+                
+                # Define the filename
+                threshold_filename = f"best_threshold{save_suffix}.txt"
+                threshold_path = save_dir_path / threshold_filename
+                
+                # Get the class name for the report
+                class_name = ""
+                # Check if we have display labels and the current index is valid
+                if map_display_labels and class_index < len(map_display_labels):
+                    class_name = map_display_labels[class_index]
+                    if num_classes > 2:
+                        # Add 'vs. Rest' for multiclass one-vs-rest plots
+                        class_name += " (vs. Rest)"
+                else:
+                    # Fallback to the generic title or default binary name
+                    class_name = plot_title.strip() or "Binary Positive Class"
+                
+                # Create content for the file
+                file_content = (
+                    f"Optimal Classification Threshold (Youden's J Statistic)\n"
+                    f"Class: {class_name}\n"
+                    f"--------------------------------------------------\n"
+                    f"Threshold: {optimal_threshold:.6f}\n"
+                    f"True Positive Rate (TPR): {tpr[best_index]:.6f}\n"
+                    f"False Positive Rate (FPR): {fpr[best_index]:.6f}\n"
+                )
+                
+                threshold_path.write_text(file_content, encoding="utf-8")
+                _LOGGER.info(f"💾 Optimal threshold saved as '{threshold_path.name}'")
+
+            except Exception as e:
+                _LOGGER.warning(f"Could not calculate or save optimal threshold: {e}")
+            
+            # Calculate AUC. 
+            auc = roc_auc_score(y_true_binary, y_score) 
+            
+            fig_roc, ax_roc = plt.subplots(figsize=(6, 6), dpi=DPI_value)
+            ax_roc.plot(fpr, tpr, label=f'AUC = {auc:.2f}', color=config.ROC_PR_line)
+            ax_roc.plot([0, 1], [0, 1], 'k--')
+            ax_roc.set_title(f'Receiver Operating Characteristic{plot_title}')
+            ax_roc.set_xlabel('False Positive Rate')
+            ax_roc.set_ylabel('True Positive Rate')
+            ax_roc.legend(loc='lower right')
+            ax_roc.grid(True)
+            roc_path = save_dir_path / f"roc_curve{save_suffix}.svg"
+            plt.savefig(roc_path)
+            plt.close(fig_roc)
+
+            # --- Save Precision-Recall Curve ---
+            precision, recall, _ = precision_recall_curve(y_true_binary, y_score)
+            ap_score = average_precision_score(y_true_binary, y_score)
+            fig_pr, ax_pr = plt.subplots(figsize=(6, 6), dpi=DPI_value)
+            ax_pr.plot(recall, precision, label=f'Avg Precision = {ap_score:.2f}', color=config.ROC_PR_line)
+            ax_pr.set_title(f'Precision-Recall Curve{plot_title}')
+            ax_pr.set_xlabel('Recall')
+            ax_pr.set_ylabel('Precision')
+            ax_pr.legend(loc='lower left')
+            ax_pr.grid(True)
+            pr_path = save_dir_path / f"pr_curve{save_suffix}.svg"
+            plt.savefig(pr_path)
+            plt.close(fig_pr)
+            
+            # --- Save Calibration Plot ---
+            fig_cal, ax_cal = plt.subplots(figsize=(8, 8), dpi=DPI_value)
+
+            # --- Step 1: Get binned data *without* plotting ---
+            with plt.ioff(): # Suppress showing the temporary plot
+                fig_temp, ax_temp = plt.subplots()
+                cal_display_temp = CalibrationDisplay.from_predictions(
+                    y_true_binary, # Use binarized labels
+                    y_score, 
+                    n_bins=config.calibration_bins, 
+                    ax=ax_temp,
+                    name="temp" # Add a name to suppress potential warnings
+                )
+                # Get the x, y coordinates of the binned data
+                line_x, line_y = cal_display_temp.line_.get_data() # type: ignore
+                plt.close(fig_temp) # Close the temporary plot
+
+            # --- Step 2: Build the plot from scratch ---
+            ax_cal.plot([0, 1], [0, 1], 'k--', label='Perfectly calibrated')
+            
+            sns.regplot(
+                x=line_x, 
+                y=line_y,
+                ax=ax_cal,
+                scatter=False, 
+                label=f"Calibration Curve ({config.calibration_bins} bins)",
+                line_kws={
+                    'color': config.ROC_PR_line, 
+                    'linestyle': '--', 
+                    'linewidth': 2,
+                    }
+            )
+            
+            ax_cal.set_title(f'Reliability Curve{plot_title}')
             ax_cal.set_xlabel('Mean Predicted Probability')
             ax_cal.set_ylabel('Fraction of Positives')
+            
+            # --- Step 3: Set final limits *after* plotting ---
+            ax_cal.set_ylim(0.0, 1.0) 
+            ax_cal.set_xlim(0.0, 1.0)
+            
+            ax_cal.legend(loc='lower right')
             ax_cal.grid(True)
             plt.tight_layout()
             
-            cal_path = save_dir_path / "calibration_plot.svg"
+            cal_path = save_dir_path / f"calibration_plot{save_suffix}.svg"
             plt.savefig(cal_path)
-            _LOGGER.info(f"📈 Calibration plot saved as '{cal_path.name}'")
             plt.close(fig_cal)
+            
+        _LOGGER.info(f"📈 Saved {len(class_indices_to_plot)} sets of ROC, Precision-Recall, and Calibration plots.")
+            
+    # restore RC params
+    plt.rcParams.update(original_rc_params)
 
 
-def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[str, Path]):
+def regression_metrics(
+    y_true: np.ndarray, 
+    y_pred: np.ndarray, 
+    save_dir: Union[str, Path],
+    config: Optional[RegressionMetricsFormat] = None
+):
     """
     Saves regression metrics and plots.
 
@@ -187,7 +399,19 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[s
         y_true (np.ndarray): Ground truth values.
         y_pred (np.ndarray): Predicted values.
         save_dir (str | Path): Directory to save plots and report.
+        config (RegressionMetricsFormat, optional): Formatting configuration object.
     """
+    
+    # --- Parse Config or use defaults ---
+    if config is None:
+        # Create a default config if one wasn't provided
+        config = RegressionMetricsFormat()
+        
+    # --- Set Matplotlib font size ---
+    original_rc_params = plt.rcParams.copy()
+    plt.rcParams.update({'font.size': config.font_size})
+    
+    # --- Calculate Metrics ---
     rmse = np.sqrt(mean_squared_error(y_true, y_pred))
     mae = mean_absolute_error(y_true, y_pred)
     r2 = r2_score(y_true, y_pred)
@@ -209,11 +433,13 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[s
     report_path.write_text(report_string)
     _LOGGER.info(f"📝 Regression report saved as '{report_path.name}'")
 
-    # Save residual plot
+    # --- Save residual plot ---
     residuals = y_true - y_pred
-    fig_res, ax_res = plt.subplots(figsize=(8, 6), dpi=100)
-    ax_res.scatter(y_pred, residuals, alpha=0.6)
-    ax_res.axhline(0, color='red', linestyle='--')
+    fig_res, ax_res = plt.subplots(figsize=(8, 6), dpi=DPI_value)
+    ax_res.scatter(y_pred, residuals, 
+                   alpha=config.scatter_alpha, 
+                   color=config.scatter_color)
+    ax_res.axhline(0, color=config.residual_line_color, linestyle='--')
     ax_res.set_xlabel("Predicted Values")
     ax_res.set_ylabel("Residuals")
     ax_res.set_title("Residual Plot")
@@ -224,10 +450,15 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[s
     _LOGGER.info(f"📈 Residual plot saved as '{res_path.name}'")
     plt.close(fig_res)
 
-    # Save true vs predicted plot
-    fig_tvp, ax_tvp = plt.subplots(figsize=(8, 6), dpi=100)
-    ax_tvp.scatter(y_true, y_pred, alpha=0.6)
-    ax_tvp.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 'k--', lw=2)
+    # --- Save true vs predicted plot ---
+    fig_tvp, ax_tvp = plt.subplots(figsize=(8, 6), dpi=DPI_value)
+    ax_tvp.scatter(y_true, y_pred, 
+                   alpha=config.scatter_alpha, 
+                   color=config.scatter_color)
+    ax_tvp.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 
+                linestyle='--', 
+                lw=2,
+                color=config.ideal_line_color)
     ax_tvp.set_xlabel('True Values')
     ax_tvp.set_ylabel('Predictions')
     ax_tvp.set_title('True vs. Predicted Values')
@@ -238,9 +469,11 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[s
     _LOGGER.info(f"📉 True vs. Predicted plot saved as '{tvp_path.name}'")
     plt.close(fig_tvp)
     
-    # Save Histogram of Residuals
-    fig_hist, ax_hist = plt.subplots(figsize=(8, 6), dpi=100)
-    sns.histplot(residuals, kde=True, ax=ax_hist)
+    # --- Save Histogram of Residuals ---
+    fig_hist, ax_hist = plt.subplots(figsize=(8, 6), dpi=DPI_value)
+    sns.histplot(residuals, kde=True, ax=ax_hist, 
+                 bins=config.hist_bins, 
+                 color=config.scatter_color)
     ax_hist.set_xlabel("Residual Value")
     ax_hist.set_ylabel("Frequency")
     ax_hist.set_title("Distribution of Residuals")
@@ -251,6 +484,9 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Union[s
     _LOGGER.info(f"📊 Residuals histogram saved as '{hist_path.name}'")
     plt.close(fig_hist)
     
+    # --- Restore RC params ---
+    plt.rcParams.update(original_rc_params)
+    
 
 def shap_summary_plot(model, 
                       background_data: Union[torch.Tensor,np.ndarray], 
@@ -276,7 +512,7 @@ def shap_summary_plot(model,
               slow and memory-intensive.
     """
     
-    print(f"\n--- SHAP Value Explanation Using {explainer_type.upper()} Explainer ---")
+    _LOGGER.info(f"📊 Running SHAP Value Explanation Using {explainer_type.upper()} Explainer")
     
     model.eval()
     # model.cpu() # Run explanations on CPU
@@ -348,9 +584,9 @@ def shap_summary_plot(model,
         _LOGGER.error(f"Invalid explainer_type: '{explainer_type}'. Must be 'deep' or 'kernel'.")
         raise ValueError()
     
-    if not isinstance(shap_values, list) and shap_values.ndim == 3 and shap_values.shape[2] == 1:
+    if not isinstance(shap_values, list) and shap_values.ndim == 3 and shap_values.shape[2] == 1: # type: ignore
         # _LOGGER.info("Squeezing SHAP values from (N, F, 1) to (N, F) for regression plot.")
-        shap_values = shap_values.squeeze(-1)
+        shap_values = shap_values.squeeze(-1) # type: ignore
 
     # --- 3. Plotting and Saving ---
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
@@ -455,7 +691,7 @@ def plot_attention_importance(weights: List[torch.Tensor], feature_names: Option
     # --- Step 3: Create and save the plot for top N features ---
     plot_df = summary_df.head(top_n).sort_values('mean_attention', ascending=True)
     
-    plt.figure(figsize=(10, 8), dpi=100)
+    plt.figure(figsize=(10, 8), dpi=DPI_value)
 
     # Create horizontal bar plot with error bars
     plt.barh(