dragon-ml-toolbox 10.2.0__py3-none-any.whl → 14.2.0__py3-none-any.whl

ml_tools/ML_evaluation.py

@@ -18,10 +18,13 @@ from sklearn.metrics import (
 import torch
 import shap
 from pathlib import Path
+from typing import Union, Optional, List, Literal
+import warnings
+
 from .path_manager import make_fullpath
 from ._logger import _LOGGER
-from typing import Union, Optional, List
 from ._script_info import _script_info
+from .keys import SHAPKeys, PyTorchLogKeys
 
 
 __all__ = [
@@ -41,8 +44,8 @@ def plot_losses(history: dict, save_dir: Union[str, Path]):
         history (dict): A dictionary containing 'train_loss' and 'val_loss'.
         save_dir (str | Path): Directory to save the plot image.
     """
-    train_loss = history.get('train_loss', [])
-    val_loss = history.get('val_loss', [])
+    train_loss = history.get(PyTorchLogKeys.TRAIN_LOSS, [])
+    val_loss = history.get(PyTorchLogKeys.VAL_LOSS, [])
     
     if not train_loss and not val_loss:
         print("Warning: Loss history is empty or incomplete. Cannot plot.")
@@ -247,13 +250,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'] = 'kernel'):
     """
     Calculates SHAP values and saves summary plots and data.
 
@@ -263,54 +268,105 @@ 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': 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()
+    # model.cpu() # Run explanations on 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)
+    shap_values = None
+    instances_to_explain_np = None
+
+    if explainer_type == 'deep':
+        # --- 1. Use DeepExplainer  ---
+        
+        # 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)
+        
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning)
+            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 ---
+        _LOGGER.warning(
+            "KernelExplainer is memory-intensive and slow. Consider reducing the number of instances to explain 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
     
-    print("Calculating SHAP values with KernelExplainer...")
-    shap_values = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
+    else:
+        _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:
+        # _LOGGER.info("Squeezing SHAP values from (N, F, 1) to (N, F) for regression plot.")
+        shap_values = shap_values.squeeze(-1)
+
+    # --- 3. Plotting and Saving ---
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
     plt.ioff()
     
+    # Convert instances to a DataFrame. robust way to ensure SHAP correctly maps values to feature names.
+    if feature_names is None:
+        # Create generic names if none were provided
+        num_features = instances_to_explain_np.shape[1]
+        feature_names = [f'feature_{i}' for i in range(num_features)]
+        
+    instances_df = pd.DataFrame(instances_to_explain_np, columns=feature_names)
+    
     # Save Bar Plot
     bar_path = save_dir_path / "shap_bar_plot.svg"
-    shap.summary_plot(shap_values, instances_to_explain_np, feature_names=feature_names, plot_type="bar", show=False)
+    shap.summary_plot(shap_values, instances_df, plot_type="bar", show=False)
     ax = plt.gca()
     ax.set_xlabel("SHAP Value Impact", labelpad=10)
     plt.title("SHAP Feature Importance")
@@ -321,11 +377,12 @@ def shap_summary_plot(model,
 
     # Save Dot Plot
     dot_path = save_dir_path / "shap_dot_plot.svg"
-    shap.summary_plot(shap_values, instances_to_explain_np, feature_names=feature_names, plot_type="dot", show=False)
+    shap.summary_plot(shap_values, instances_df, 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)
@@ -333,17 +390,21 @@ def shap_summary_plot(model,
     plt.close()
 
     # Save Summary Data to CSV
-    summary_path = save_dir_path / "shap_summary.csv"
-    # Ensure the array is 1D before creating the DataFrame
-    mean_abs_shap = np.abs(shap_values).mean(axis=0).flatten()
+    shap_summary_filename = SHAPKeys.SAVENAME + ".csv"
+    summary_path = save_dir_path / shap_summary_filename
     
-    if feature_names is None:
-        feature_names = [f'feature_{i}' for i in range(len(mean_abs_shap))]
+    # 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()
         
     summary_df = pd.DataFrame({
-        'feature': feature_names,
-        'mean_abs_shap_value': mean_abs_shap
-    }).sort_values('mean_abs_shap_value', ascending=False)
+        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)
     
@@ -351,7 +412,7 @@ def shap_summary_plot(model,
     plt.ion()
 
 
-def plot_attention_importance(weights: List[torch.Tensor], feature_names: Optional[List[str]], save_dir: Union[str, Path]):
+def plot_attention_importance(weights: List[torch.Tensor], feature_names: Optional[List[str]], save_dir: Union[str, Path], top_n: int = 10):
     """
     Aggregates attention weights and plots global feature importance.
 
@@ -362,6 +423,7 @@ def plot_attention_importance(weights: List[torch.Tensor], feature_names: Option
         weights (List[torch.Tensor]): A list of attention weight tensors from each batch.
         feature_names (List[str] | None): Names of the features for plot labeling.
         save_dir (str | Path): Directory to save the plot and summary CSV.
+        top_n (int): The number of top features to display in the plot.
     """
     if not weights:
         _LOGGER.error("Attention weights list is empty. Skipping importance plot.")
@@ -390,11 +452,10 @@ def plot_attention_importance(weights: List[torch.Tensor], feature_names: Option
     summary_df.to_csv(summary_path, index=False)
     _LOGGER.info(f"📝 Attention summary data saved as '{summary_path.name}'")
 
-    # --- Step 3: Create and save the plot ---
-    plt.figure(figsize=(10, 8), dpi=100)
+    # --- Step 3: Create and save the plot for top N features ---
+    plot_df = summary_df.head(top_n).sort_values('mean_attention', ascending=True)
     
-    # Sort for plotting
-    plot_df = summary_df.sort_values('mean_attention', ascending=True)
+    plt.figure(figsize=(10, 8), dpi=100)
 
     # Create horizontal bar plot with error bars
     plt.barh(
@@ -408,7 +469,7 @@ def plot_attention_importance(weights: List[torch.Tensor], feature_names: Option
         color='cornflowerblue'
     )
     
-    plt.title('Global Feature Importance')
+    plt.title('Top Features by Attention')
     plt.xlabel('Average Attention Weight')
     plt.ylabel('Feature')
     plt.grid(axis='x', linestyle='--', alpha=0.6)

ml_tools/ML_evaluation_multi.py

@@ -19,11 +19,14 @@ from sklearn.metrics import (
     jaccard_score
 )
 from pathlib import Path
-from typing import Union, List
+from typing import Union, List, Literal
+import warnings
 
 from .path_manager import make_fullpath, sanitize_filename
 from ._logger import _LOGGER
 from ._script_info import _script_info
+from .keys import SHAPKeys
+
 
 __all__ = [
     "multi_target_regression_metrics",
@@ -230,10 +233,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'] = 'kernel'
 ):
     """
-    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.
@@ -242,40 +247,93 @@ 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': 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 ---
+        
+        # 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)
+        
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning)
+            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  ---
+        _LOGGER.warning(
+            "KernelExplainer is memory-intensive and slow. Consider reducing the number of instances to explain if the process terminates unexpectedly."
+        )
+        
+        # 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]
@@ -292,11 +350,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_inference.py

@@ -9,7 +9,8 @@ from .ML_scaler import PytorchScaler
 from ._script_info import _script_info
 from ._logger import _LOGGER
 from .path_manager import make_fullpath
-from .keys import PyTorchInferenceKeys
+from .keys import PyTorchInferenceKeys, PyTorchCheckpointKeys
+
 
 __all__ = [
     "PyTorchInferenceHandler",
@@ -55,11 +56,21 @@ class _BaseInferenceHandler(ABC):
         model_p = make_fullpath(state_dict, enforce="file")
 
         try:
-            # Load the state dictionary and apply it to the model structure
-            self.model.load_state_dict(torch.load(model_p, map_location=self.device))
+            # Load whatever is in the file
+            loaded_data = torch.load(model_p, map_location=self.device)
+
+            # Check if it's the new checkpoint dictionary or an old weights-only file
+            if isinstance(loaded_data, dict) and PyTorchCheckpointKeys.MODEL_STATE in loaded_data:
+                # It's a new training checkpoint, extract the weights
+                self.model.load_state_dict(loaded_data[PyTorchCheckpointKeys.MODEL_STATE])
+            else:
+                # It's an old-style file (or just a state_dict), load it directly
+                self.model.load_state_dict(loaded_data)
+            
+            _LOGGER.info(f"Model state loaded from '{model_p.name}'.")
+                
             self.model.to(self.device)
             self.model.eval()  # Set the model to evaluation mode
-            _LOGGER.info(f"Model state loaded from '{model_p.name}' and set to evaluation mode.")
         except Exception as e:
             _LOGGER.error(f"Failed to load model state from '{model_p}': {e}")
             raise
@@ -71,7 +82,6 @@ class _BaseInferenceHandler(ABC):
             _LOGGER.warning("CUDA not available, switching to CPU.")
             device_lower = "cpu"
         elif device_lower == "mps" and not torch.backends.mps.is_available():
-            # Your M-series Mac will appreciate this check!
             _LOGGER.warning("Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
             device_lower = "cpu"
         return torch.device(device_lower)