dragon-ml-toolbox 6.4.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

ml_tools/ML_scaler.py

@@ -0,0 +1,197 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+from typing import Union, List, Optional
+from ._logger import _LOGGER
+from ._script_info import _script_info
+from .path_manager import make_fullpath
+
+__all__ = [
+    "PytorchScaler"
+]
+
+class PytorchScaler:
+    """
+    Standardizes continuous features in a PyTorch dataset by subtracting the
+    mean and dividing by the standard deviation.
+
+    The scaler is fitted on a training dataset and can then be saved and
+    loaded for consistent transformation during inference.
+    """
+    def __init__(self,
+                 mean: Optional[torch.Tensor] = None,
+                 std: Optional[torch.Tensor] = None,
+                 continuous_feature_indices: Optional[List[int]] = None):
+        """
+        Initializes the scaler.
+
+        Args:
+            mean (torch.Tensor, optional): The mean of the features to scale.
+            std (torch.Tensor, optional): The standard deviation of the features.
+            continuous_feature_indices (List[int], optional): The column indices of the features to standardize.
+        """
+        self.mean_ = mean
+        self.std_ = std
+        self.continuous_feature_indices = continuous_feature_indices
+
+    @classmethod
+    def fit(cls, dataset: Dataset, continuous_feature_indices: List[int], batch_size: int = 64) -> 'PytorchScaler':
+        """
+        Fits the scaler by computing the mean and std dev from a dataset using a
+        fast, single-pass, vectorized algorithm.
+
+        Args:
+            dataset (Dataset): The PyTorch Dataset to fit on.
+            continuous_feature_indices (List[int]): The column indices of the
+                features to standardize.
+            batch_size (int): The batch size for iterating through the dataset.
+
+        Returns:
+            PytorchScaler: A new, fitted instance of the scaler.
+        """
+        if not continuous_feature_indices:
+            _LOGGER.warning("⚠️ No continuous feature indices provided. Scaler will not be fitted.")
+            return cls()
+
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+        
+        running_sum, running_sum_sq = None, None
+        count = 0
+        num_continuous_features = len(continuous_feature_indices)
+
+        for features, _ in loader:
+            if running_sum is None:
+                device = features.device
+                running_sum = torch.zeros(num_continuous_features, device=device)
+                running_sum_sq = torch.zeros(num_continuous_features, device=device)
+
+            continuous_features = features[:, continuous_feature_indices].to(device)
+            
+            running_sum += torch.sum(continuous_features, dim=0)
+            running_sum_sq += torch.sum(continuous_features**2, dim=0) # type: ignore
+            count += continuous_features.size(0)
+
+        if count == 0:
+             _LOGGER.warning("⚠️ Dataset is empty. Scaler cannot be fitted.")
+             return cls(continuous_feature_indices=continuous_feature_indices)
+
+        # Calculate mean
+        mean = running_sum / count
+
+        # Calculate standard deviation
+        if count < 2:
+            _LOGGER.warning(f"⚠️ Only one sample found. Standard deviation cannot be calculated and is set to 1.")
+            std = torch.ones_like(mean)
+        else:
+            # var = E[X^2] - (E[X])^2
+            var = (running_sum_sq / count) - mean**2
+            std = torch.sqrt(torch.clamp(var, min=1e-8)) # Clamp for numerical stability
+
+        _LOGGER.info(f"Scaler fitted on {count} samples for {num_continuous_features} continuous features.")
+        return cls(mean=mean, std=std, continuous_feature_indices=continuous_feature_indices)
+
+    def transform(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Applies standardization to the specified continuous features.
+
+        Args:
+            data (torch.Tensor): The input data tensor.
+
+        Returns:
+            torch.Tensor: The transformed data tensor.
+        """
+        if self.mean_ is None or self.std_ is None or self.continuous_feature_indices is None:
+            _LOGGER.warning("⚠️ Scaler has not been fitted. Returning original data.")
+            return data
+
+        data_clone = data.clone()
+        
+        # Ensure mean and std are on the same device as the data
+        mean = self.mean_.to(data.device)
+        std = self.std_.to(data.device)
+        
+        # Extract the columns to be scaled
+        features_to_scale = data_clone[:, self.continuous_feature_indices]
+        
+        # Apply scaling, adding epsilon to std to prevent division by zero
+        scaled_features = (features_to_scale - mean) / (std + 1e-8)
+        
+        # Place the scaled features back into the cloned tensor
+        data_clone[:, self.continuous_feature_indices] = scaled_features
+        
+        return data_clone
+
+    def inverse_transform(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Applies the inverse of the standardization transformation.
+
+        Args:
+            data (torch.Tensor): The scaled data tensor.
+
+        Returns:
+            torch.Tensor: The original-scale data tensor.
+        """
+        if self.mean_ is None or self.std_ is None or self.continuous_feature_indices is None:
+            _LOGGER.warning("⚠️ Scaler has not been fitted. Returning original data.")
+            return data
+            
+        data_clone = data.clone()
+        
+        mean = self.mean_.to(data.device)
+        std = self.std_.to(data.device)
+        
+        features_to_inverse = data_clone[:, self.continuous_feature_indices]
+        
+        # Apply inverse scaling
+        original_scale_features = (features_to_inverse * (std + 1e-8)) + mean
+        
+        data_clone[:, self.continuous_feature_indices] = original_scale_features
+        
+        return data_clone
+
+    def save(self, filepath: Union[str, Path]):
+        """
+        Saves the scaler's state (mean, std, indices) to a .pth file.
+
+        Args:
+            filepath (str | Path): The path to save the file.
+        """
+        path_obj = make_fullpath(filepath)
+        state = {
+            'mean': self.mean_,
+            'std': self.std_,
+            'continuous_feature_indices': self.continuous_feature_indices
+        }
+        torch.save(state, path_obj)
+        _LOGGER.info(f"✅ PytorchScaler state saved to '{path_obj.name}'.")
+
+    @staticmethod
+    def load(filepath: Union[str, Path]) -> 'PytorchScaler':
+        """
+        Loads a scaler's state from a .pth file.
+
+        Args:
+            filepath (str | Path): The path to the saved scaler file.
+
+        Returns:
+            PytorchScaler: An instance of the scaler with the loaded state.
+        """
+        path_obj = make_fullpath(filepath, enforce="file")
+        state = torch.load(path_obj)
+        _LOGGER.info(f"✅ PytorchScaler state loaded from '{path_obj.name}'.")
+        return PytorchScaler(
+            mean=state['mean'],
+            std=state['std'],
+            continuous_feature_indices=state['continuous_feature_indices']
+        )
+    
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the scaler."""
+        if self.continuous_feature_indices:
+            num_features = len(self.continuous_feature_indices)
+            return f"PytorchScaler(fitted for {num_features} features)"
+        return "PytorchScaler(not fitted)"
+
+
+def info():
+    _script_info(__all__)

ml_tools/ML_trainer.py

@@ -6,7 +6,7 @@ from torch import nn
 import numpy as np
 
 from .ML_callbacks import Callback, History, TqdmProgressBar
-from .ML_evaluation import classification_metrics, regression_metrics, plot_losses, shap_summary_plot
+from .ML_evaluation import classification_metrics, regression_metrics, plot_losses, shap_summary_plot, plot_attention_importance
 from ._script_info import _script_info
 from .keys import PyTorchLogKeys
 from ._logger import _LOGGER
@@ -282,8 +282,11 @@ class MLTrainer:
         print("\n--- Training History ---")
         plot_losses(self.history, save_dir=save_dir)
     
-    def explain(self, explain_dataset: Optional[Dataset] = None, n_samples: int = 1000, 
-                feature_names: Optional[List[str]] = None, save_dir: Optional[Union[str,Path]] = None):
+    def explain(self,
+                feature_names: Optional[List[str]], 
+                save_dir: Union[str,Path], 
+                explain_dataset: Optional[Dataset] = None, 
+                n_samples: int = 1000):
         """
         Explains model predictions using SHAP and saves all artifacts.
 
@@ -328,14 +331,14 @@ class MLTrainer:
         # 1. Get background data from the trainer's train_dataset
         background_data = _get_random_sample(self.train_dataset, n_samples)
         if background_data is None:
-            print("Warning: Trainer's train_dataset is empty or invalid. Skipping SHAP analysis.")
+            _LOGGER.error("❌ Trainer's train_dataset is empty or invalid. Skipping SHAP analysis.")
             return
 
         # 2. Determine target dataset and get explanation instances
         target_dataset = explain_dataset if explain_dataset is not None else self.test_dataset
         instances_to_explain = _get_random_sample(target_dataset, n_samples)
         if instances_to_explain is None:
-            print("Warning: Explanation dataset is empty or invalid. Skipping SHAP analysis.")
+            _LOGGER.error("❌ Explanation dataset is empty or invalid. Skipping SHAP analysis.")
             return
 
         # 3. Call the plotting function
@@ -347,7 +350,80 @@ class MLTrainer:
             save_dir=save_dir
         )
     
+    def _attention_helper(self, dataloader: DataLoader):
+        """
+        Private method to yield model attention weights batch by batch for evaluation.
 
+        Args:
+            dataloader (DataLoader): The dataloader to predict on.
+
+        Yields:
+            (torch.Tensor): Attention weights
+        """
+        self.model.eval()
+        self.model.to(self.device)
+        
+        with torch.no_grad():
+            for features, target in dataloader:
+                features = features.to(self.device)
+                attention_weights = None
+                
+                # Get model output
+                # Unpack logits and weights from the special forward method
+                _output, attention_weights = self.model.forward_attention(features) # type: ignore
+                
+                if attention_weights is not None:
+                    attention_weights = attention_weights.cpu()
+
+                yield attention_weights
+    
+    def explain_attention(self, save_dir: Union[str, Path], feature_names: Optional[List[str]], explain_dataset: Optional[Dataset] = None):
+        """
+        Generates and saves a feature importance plot based on attention weights.
+
+        This method only works for models with a `forward_attention` method.
+
+        Args:
+            save_dir (str | Path): Directory to save the plot and summary data.
+            feature_names (List[str] | None): Names for the features for plot labeling.
+            explain_dataset (Dataset, optional): A specific dataset to explain. If None, the trainer's test dataset is used.
+        """
+        
+        print("\n--- Attention Analysis ---")
+        
+        # --- Step 1: Check if the model supports this explanation ---
+        if not hasattr(self.model, 'forward_attention'):
+            _LOGGER.error("❌ Model does not have a `forward_attention` method. Skipping attention explanation.")
+            return
+
+        # --- Step 2: Set up the dataloader ---
+        dataset_to_use = explain_dataset if explain_dataset is not None else self.test_dataset
+        if not isinstance(dataset_to_use, Dataset):
+            _LOGGER.error("❌ The explanation dataset is empty or invalid. Skipping attention analysis.")
+            return
+        
+        explain_loader = DataLoader(
+            dataset=dataset_to_use, batch_size=32, shuffle=False,
+            num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
+            pin_memory=("cuda" in self.device.type)
+        )
+        
+        # --- Step 3: Collect weights ---
+        all_weights = []
+        for att_weights_b in self._attention_helper(explain_loader):
+            if att_weights_b is not None:
+                all_weights.append(att_weights_b)
+
+        # --- Step 4: Call the plotting function ---
+        if all_weights:
+            plot_attention_importance(
+                weights=all_weights,
+                feature_names=feature_names,
+                save_dir=save_dir
+            )
+        else:
+            _LOGGER.error("❌ No attention weights were collected from the model.")
+    
     def callbacks_hook(self, method_name: str, *args, **kwargs):
         """Calls the specified method on all callbacks."""
         for callback in self.callbacks: