dragon-ml-toolbox 14.7.0__py3-none-any.whl → 16.2.0__py3-none-any.whl

ml_tools/ML_models.py

@@ -7,16 +7,15 @@ import json
 from ._logger import _LOGGER
 from .path_manager import make_fullpath
 from ._script_info import _script_info
-from .keys import PytorchModelArchitectureKeys
+from ._keys import PytorchModelArchitectureKeys
 from ._schema import FeatureSchema
 
 
 __all__ = [
-    "MultilayerPerceptron",
-    "AttentionMLP",
-    "MultiHeadAttentionMLP",
-    "TabularTransformer",
-    "SequencePredictorLSTM",
+    "DragonMLP",
+    "DragonAttentionMLP",
+    "DragonMultiHeadAttentionNet",
+    "DragonTabularTransformer"
 ]
 
 
@@ -174,7 +173,7 @@ class _BaseAttention(_BaseMLP):
         return logits, attention_weights
 
 
-class MultilayerPerceptron(_BaseMLP):
+class DragonMLP(_BaseMLP):
     """
     Creates a versatile Multilayer Perceptron (MLP) for regression or classification tasks.
     """
@@ -208,10 +207,10 @@ class MultilayerPerceptron(_BaseMLP):
         # Extracts the number of neurons from each nn.Linear layer
         layer_sizes = [str(layer.in_features) for layer in self.mlp if isinstance(layer, nn.Linear)]
         
-        return self._repr_helper(name="MultilayerPerceptron", mlp_layers=layer_sizes)
+        return self._repr_helper(name="DragonMLP", mlp_layers=layer_sizes)
 
 
-class AttentionMLP(_BaseAttention):
+class DragonAttentionMLP(_BaseAttention):
     """
     A Multilayer Perceptron (MLP) that incorporates an Attention layer to dynamically weigh input features.
     
@@ -244,10 +243,10 @@ class AttentionMLP(_BaseAttention):
             if isinstance(layer, nn.Linear):
                 arch.append(str(layer.in_features))
         
-        return self._repr_helper(name="AttentionMLP", mlp_layers=arch)
+        return self._repr_helper(name="DragonAttentionMLP", mlp_layers=arch)
 
 
-class MultiHeadAttentionMLP(_BaseAttention):
+class DragonMultiHeadAttentionNet(_BaseAttention):
     """
     An MLP that incorporates a standard `nn.MultiheadAttention` layer to process
     the input features.
@@ -292,10 +291,10 @@ class MultiHeadAttentionMLP(_BaseAttention):
         )
         arch_str = f"{self.in_features} -> [MultiHead(h={self.num_heads})] -> {mlp_part}"
         
-        return f"MultiHeadAttentionMLP(arch: {arch_str})"
+        return f"DragonMultiHeadAttentionNet(arch: {arch_str})"
 
 
-class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
+class DragonTabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     """
     A Transformer-based model for tabular data tasks.
     
@@ -502,7 +501,7 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
         
         arch_str = " -> ".join(parts)
         
-        return f"TabularTransformer(arch: {arch_str})"
+        return f"DragonTabularTransformer(arch: {arch_str})"
 
 
 class _FeatureTokenizer(nn.Module):
@@ -662,93 +661,5 @@ class _MultiHeadAttentionLayer(nn.Module):
         return out, attn_weights.squeeze()
 
 
-class SequencePredictorLSTM(nn.Module, _ArchitectureHandlerMixin):
-    """
-    A simple LSTM-based network for sequence-to-sequence prediction tasks.
-
-    This model is designed for datasets where each input sequence maps to an
-    output sequence of the same length. It's suitable for forecasting problems
-    prepared by the `SequenceMaker` class.
-
-    The expected input shape is `(batch_size, sequence_length, features)`.
-
-    Args:
-        features (int): The number of features in the input sequence. Defaults to 1.
-        hidden_size (int): The number of features in the LSTM's hidden state.
-                           Defaults to 100.
-        recurrent_layers (int): The number of recurrent LSTM layers. Defaults to 1.
-        dropout (float): The dropout probability for all but the last LSTM layer.
-                         Defaults to 0.
-    """
-    def __init__(self, features: int = 1, hidden_size: int = 100,
-                 recurrent_layers: int = 1, dropout: float = 0):
-        super().__init__()
-
-        # --- Validation ---
-        if not isinstance(features, int) or features < 1:
-            raise ValueError("features must be a positive integer.")
-        if not isinstance(hidden_size, int) or hidden_size < 1:
-            raise ValueError("hidden_size must be a positive integer.")
-        if not isinstance(recurrent_layers, int) or recurrent_layers < 1:
-            raise ValueError("recurrent_layers must be a positive integer.")
-        if not (0.0 <= dropout < 1.0):
-            raise ValueError("dropout must be a float between 0.0 and 1.0.")
-        
-        # --- Save configuration ---
-        self.features = features
-        self.hidden_size = hidden_size
-        self.recurrent_layers = recurrent_layers
-        self.dropout = dropout
-        
-        # Build model
-        self.lstm = nn.LSTM(
-            input_size=features,
-            hidden_size=hidden_size,
-            num_layers=recurrent_layers,
-            dropout=dropout,
-            batch_first=True  # This is crucial for (batch, seq, feature) input
-        )
-        self.linear = nn.Linear(in_features=hidden_size, out_features=features)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Defines the forward pass.
-
-        Args:
-            x (torch.Tensor): The input tensor with shape
-                              (batch_size, sequence_length, features).
-
-        Returns:
-            torch.Tensor: The output tensor with shape
-                          (batch_size, sequence_length, features).
-        """
-        # The LSTM returns the full output sequence and the final hidden/cell states
-        lstm_out, _ = self.lstm(x)
-        
-        # Pass the LSTM's output sequence to the linear layer
-        predictions = self.linear(lstm_out)
-        
-        return predictions
-    
-    def get_architecture_config(self) -> dict:
-        """Returns the configuration of the model."""
-        return {
-            'features': self.features,
-            'hidden_size': self.hidden_size,
-            'recurrent_layers': self.recurrent_layers,
-            'dropout': self.dropout
-        }
-    
-    def __repr__(self) -> str:
-        """Returns the developer-friendly string representation of the model."""
-        return (
-            f"SequencePredictorLSTM(features={self.lstm.input_size}, "
-            f"hidden_size={self.lstm.hidden_size}, "
-            f"recurrent_layers={self.lstm.num_layers})"
-        )
-
-
-# ---- PyTorch models ---
-
 def info():
     _script_info(__all__)

ml_tools/ML_models_advanced.py

@@ -6,7 +6,7 @@ import json
 
 from ._logger import _LOGGER
 from .path_manager import make_fullpath
-from .keys import PytorchModelArchitectureKeys
+from ._keys import PytorchModelArchitectureKeys
 from ._schema import FeatureSchema
 from ._script_info import _script_info
 from .ML_models import _ArchitectureHandlerMixin

ml_tools/ML_optimization.py

@@ -14,9 +14,9 @@ from functools import partial
 from .path_manager import make_fullpath, sanitize_filename
 from ._logger import _LOGGER
 from ._script_info import _script_info
-from .ML_inference import PyTorchInferenceHandler
-from .keys import PyTorchInferenceKeys
-from .SQL import DatabaseManager
+from .ML_inference import DragonInferenceHandler
+from ._keys import PyTorchInferenceKeys
+from .SQL import DragonSQL
 from .optimization_tools import _save_result, create_optimization_bounds
 from .utilities import save_dataframe_filename
 from .math_utilities import discretize_categorical_values
@@ -24,14 +24,14 @@ from ._schema import FeatureSchema
 
 
 __all__ = [
-    "MLOptimizer",
+    "DragonOptimizer",
     "FitnessEvaluator",
     "create_pytorch_problem",
     "run_optimization"
 ]
 
 
-class MLOptimizer:
+class DragonOptimizer:
     """
     A wrapper class for setting up and running EvoTorch optimization tasks.
 
@@ -47,7 +47,7 @@ class MLOptimizer:
         >>> cont_bounds = {'feature_A': (0, 100), 'feature_B': (-10, 10)}
         >>>
         >>> # 3. Initialize the optimizer
-        >>> optimizer = MLOptimizer(
+        >>> optimizer = DragonOptimizer(
         ...     inference_handler=my_handler,
         ...     schema=schema,
         ...     continuous_bounds_map=cont_bounds,
@@ -63,7 +63,7 @@ class MLOptimizer:
         ... )
     """
     def __init__(self,
-                 inference_handler: PyTorchInferenceHandler,
+                 inference_handler: DragonInferenceHandler,
                  schema: FeatureSchema,
                  continuous_bounds_map: Dict[str, Tuple[float, float]],
                  task: Literal["min", "max"],
@@ -75,7 +75,7 @@ class MLOptimizer:
         Initializes the optimizer by creating the EvoTorch problem and searcher.
 
         Args:
-            inference_handler (PyTorchInferenceHandler): 
+            inference_handler (DragonInferenceHandler): 
                 An initialized inference handler containing the model.
             schema (FeatureSchema): 
                 The definitive schema object from data_exploration.
@@ -172,18 +172,18 @@ class FitnessEvaluator:
     A callable class that wraps the PyTorch model inference handler and performs
     on-the-fly discretization for the EvoTorch fitness function.
 
-    This class is automatically instantiated by MLOptimizer and passed to
+    This class is automatically instantiated by DragonOptimizer and passed to
     create_pytorch_problem, encapsulating the evaluation logic.
     """
     def __init__(self,
-                 inference_handler: PyTorchInferenceHandler,
+                 inference_handler: DragonInferenceHandler,
                  categorical_index_map: Optional[Dict[int, int]] = None,
                  discretize_start_at_zero: bool = True):
         """
         Initializes the fitness evaluator.
 
         Args:
-            inference_handler (PyTorchInferenceHandler): 
+            inference_handler (DragonInferenceHandler): 
                 An initialized inference handler containing the model.
             categorical_index_map (Dict[int, int] | None): 
                 Maps {column_index: cardinality} for discretization.
@@ -426,7 +426,7 @@ def run_optimization(
         _LOGGER.info(f"🏁 Starting optimal solution space analysis with {repetitions} repetitions...")
         
         first_run_logger = None # To store the logger from the first rep
-        db_context = DatabaseManager(db_path) if save_format in ['sqlite', 'both'] else nullcontext()
+        db_context = DragonSQL(db_path) if save_format in ['sqlite', 'both'] else nullcontext()
         
         with db_context as db_manager:
             # --- Setup Database Schema (if applicable) ---

ml_tools/ML_scaler.py

@@ -9,11 +9,11 @@ from .path_manager import make_fullpath
 
 
 __all__ = [
-    "PytorchScaler"
+    "DragonScaler"
 ]
 
 
-class PytorchScaler:
+class DragonScaler:
     """
     Standardizes continuous features in a PyTorch dataset by subtracting the
     mean and dividing by the standard deviation.
@@ -38,7 +38,7 @@ class PytorchScaler:
         self.continuous_feature_indices = continuous_feature_indices
 
     @classmethod
-    def fit(cls, dataset: Dataset, continuous_feature_indices: List[int], batch_size: int = 64) -> 'PytorchScaler':
+    def fit(cls, dataset: Dataset, continuous_feature_indices: List[int], batch_size: int = 64) -> 'DragonScaler':
         """
         Fits the scaler by computing the mean and std dev from a dataset using a
         fast, single-pass, vectorized algorithm.
@@ -50,7 +50,7 @@ class PytorchScaler:
             batch_size (int): The batch size for iterating through the dataset.
 
         Returns:
-            PytorchScaler: A new, fitted instance of the scaler.
+            DragonScaler: A new, fitted instance of the scaler.
         """
         if not continuous_feature_indices:
             _LOGGER.error("No continuous feature indices provided. Scaler will not be fitted.")
@@ -167,10 +167,10 @@ class PytorchScaler:
         }
         torch.save(state, path_obj)
         if verbose:
-            _LOGGER.info(f"PytorchScaler state saved as '{path_obj.name}'.")
+            _LOGGER.info(f"DragonScaler state saved as '{path_obj.name}'.")
 
     @staticmethod
-    def load(filepath: Union[str, Path], verbose: bool=True) -> 'PytorchScaler':
+    def load(filepath: Union[str, Path], verbose: bool=True) -> 'DragonScaler':
         """
         Loads a scaler's state from a .pth file.
 
@@ -178,13 +178,13 @@ class PytorchScaler:
             filepath (str | Path): The path to the saved scaler file.
 
         Returns:
-            PytorchScaler: An instance of the scaler with the loaded state.
+            DragonScaler: An instance of the scaler with the loaded state.
         """
         path_obj = make_fullpath(filepath, enforce="file")
         state = torch.load(path_obj)
         if verbose:
-            _LOGGER.info(f"PytorchScaler state loaded from '{path_obj.name}'.")
-        return PytorchScaler(
+            _LOGGER.info(f"DragonScaler state loaded from '{path_obj.name}'.")
+        return DragonScaler(
             mean=state['mean'],
             std=state['std'],
             continuous_feature_indices=state['continuous_feature_indices']
@@ -194,8 +194,8 @@ class PytorchScaler:
         """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)"
+            return f"DragonScaler(fitted for {num_features} features)"
+        return "DragonScaler(not fitted)"
 
 
 def info():

ml_tools/ML_sequence_datasetmaster.py

@@ -0,0 +1,341 @@
+import torch
+from torch.utils.data import Dataset
+import pandas
+import numpy
+from typing import Literal, Union, Tuple
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+from .path_manager import make_fullpath
+from ._logger import _LOGGER
+from ._script_info import _script_info
+from .ML_scaler import DragonScaler
+from .ML_datasetmaster import _PytorchDataset
+from ._keys import DatasetKeys, MLTaskKeys
+
+
+__all__ = [
+    "DragonDatasetSequence"
+]
+
+
+# --- SequenceMaker ---
+class DragonDatasetSequence:
+    """
+    Creates windowed PyTorch datasets from a univariate (one feature) sequential data.
+    
+    Automatic Pipeline:
+    
+    1. Split Data: Separate data into training, validation, and testing portions.
+    2. Normalize Data: Normalize the data. The scaler will be fitted on the training portion.
+    3. Generate Windows: Create the windowed sequences from the normalized splits.
+    """
+    def __init__(self, 
+                 prediction_mode: Literal["sequence-to-sequence", "sequence-to-value"],
+                 data: Union[pandas.DataFrame, pandas.Series, numpy.ndarray], 
+                 sequence_length: int,
+                 validation_size: float = 0.2,
+                 test_size: float = 0.1):
+        """
+        Initializes the dataset manager and automatically processes the data.
+        
+        The constructor runs the full pipeline:
+        1. Splits the data chronologically (train, validation, test).
+        2. Fits a DragonScaler on the training split.
+        3. Normalizes all splits using the fitted scaler.
+        4. Generates windowed datasets for training, validation, and testing.
+
+        Args:
+            prediction_mode: The type of sequence task.
+            data: The input univariate time-series data.
+                - If pandas.DataFrame: The index is used for the time axis
+                  and the *first column* is used as the sequence.
+                - If pandas.Series: The index is used for the time axis.
+                - If numpy.ndarray: A simple integer range is used for the time axis.
+            sequence_length (int): The number of time steps in each input window (X).
+            validation_size (float): The fraction of data to hold out for validation.
+            test_size (float): The fraction of data to hold out for testing.
+        """
+        self._train_dataset = None
+        self._test_dataset = None
+        self._val_dataset = None
+        self.sequence_length = sequence_length
+        self.scaler = None
+        
+        if not prediction_mode in [MLTaskKeys.SEQUENCE_SEQUENCE, MLTaskKeys.SEQUENCE_VALUE]:
+            _LOGGER.error(f"Unrecognized prediction mode: '{prediction_mode}'.")
+            raise ValueError()
+        else:
+            self.prediction_mode = prediction_mode
+        
+        if isinstance(data, pandas.DataFrame):
+            self.time_axis = data.index.values
+            self.sequence = data.iloc[:, 0].values.astype(numpy.float32)
+        elif isinstance(data, pandas.Series):
+            self.time_axis = data.index.values
+            self.sequence = data.values.astype(numpy.float32)
+        elif isinstance(data, numpy.ndarray):
+            self.time_axis = numpy.arange(len(data))
+            self.sequence = data.astype(numpy.float32)
+        else:
+            _LOGGER.error("Data must be a pandas DataFrame/Series or a numpy array.")
+            raise TypeError()
+            
+        self.train_sequence = None
+        self.val_sequence = None
+        self.test_sequence = None
+        
+        self.train_time_axis = None
+        self.val_time_axis = None 
+        self.test_time_axis = None 
+        
+        self._is_split = False
+        self._is_normalized = False
+        self._are_windows_generated = False
+        
+        # Automation
+        self._split_data(validation_size=validation_size, test_size=test_size)
+        self._normalize_data()
+        self._generate_windows()
+    
+    def _split_data(self, validation_size: float = 0.2, test_size: float = 0.1) -> None:
+        """
+        Splits the sequence chronologically into training, validation, and testing portions.
+        
+        To prevent windowing errors, the validation and test sets include an overlap of `sequence_length` from the preceding data.
+        """
+        if self._is_split:
+            _LOGGER.warning("Data has already been split.")
+            return
+            
+        if (validation_size + test_size) >= 1.0:
+            _LOGGER.error(f"The sum of validation_size ({validation_size}) and test_size ({test_size}) must be less than 1.0.")
+            raise ValueError("validation_size and test_size sum must be < 1.0")
+
+        total_size = len(self.sequence)
+        
+        # Calculate split indices
+        test_split_idx = int(total_size * (1 - test_size))
+        val_split_idx = int(total_size * (1 - test_size - validation_size))
+        
+        # --- Create sequences ---
+        # Train sequence is from the beginning to the validation index
+        self.train_sequence = self.sequence[:val_split_idx]
+        
+        # Validation sequence starts `sequence_length` before its split index for windowing
+        self.val_sequence = self.sequence[val_split_idx - self.sequence_length : test_split_idx]
+        
+        # Test sequence starts `sequence_length` before its split index for windowing
+        self.test_sequence = self.sequence[test_split_idx - self.sequence_length:]
+        
+        # --- Create time axes ---
+        self.train_time_axis = self.time_axis[:val_split_idx]
+        # The "plottable" validation/test time axes start from their respective split indices
+        self.val_time_axis = self.time_axis[val_split_idx : test_split_idx]
+        self.test_time_axis = self.time_axis[test_split_idx:]
+
+        self._is_split = True
+        _LOGGER.info(f"Sequence split into training ({len(self.train_sequence)}), validation ({len(self.val_sequence)}), and testing ({len(self.test_sequence)}) points.")
+    
+    def _normalize_data(self) -> None:
+        """
+        Normalizes the sequence data using DragonScaler. Must be called AFTER splitting to prevent data leakage from the test set.
+        """
+        if not self._is_split:
+            _LOGGER.error("Data must be split BEFORE normalizing.")
+            raise RuntimeError()
+
+        if self.scaler:
+            _LOGGER.warning("Data has already been normalized.")
+            return
+
+        # 1. DragonScaler requires a Dataset to fit. Create a temporary one.
+        # The scaler expects 2D data [n_samples, n_features].
+        train_features = self.train_sequence.reshape(-1, 1) # type: ignore
+
+        # _PytorchDataset needs labels, so we create dummy ones.
+        dummy_labels = numpy.zeros(len(train_features))
+        temp_train_ds = _PytorchDataset(train_features, dummy_labels, labels_dtype=torch.float32)
+
+        # 2. Fit the DragonScaler on the temporary training dataset.
+        # The sequence is a single feature, so its index is [0].
+        _LOGGER.info("Fitting DragonScaler on the training data...")
+        self.scaler = DragonScaler.fit(temp_train_ds, continuous_feature_indices=[0])
+
+        # 3. Transform sequences using the fitted scaler.
+        # The transform method requires a tensor, so we convert, transform, and convert back.
+        train_tensor = torch.tensor(self.train_sequence.reshape(-1, 1), dtype=torch.float32) # type: ignore
+        val_tensor = torch.tensor(self.val_sequence.reshape(-1, 1), dtype=torch.float32) # type: ignore
+        test_tensor = torch.tensor(self.test_sequence.reshape(-1, 1), dtype=torch.float32) # type: ignore
+
+        self.train_sequence = self.scaler.transform(train_tensor).numpy().flatten()
+        self.val_sequence = self.scaler.transform(val_tensor).numpy().flatten()
+        self.test_sequence = self.scaler.transform(test_tensor).numpy().flatten()
+
+        self._is_normalized = True
+        _LOGGER.info("Sequence data normalized using DragonScaler.")
+
+    def _generate_windows(self) -> None:
+        """
+        Generates overlapping windows for features and labels.
+        """
+        if not self._is_split:
+            _LOGGER.error("Cannot generate windows before splitting data.")
+            raise RuntimeError()
+        
+        if not self._is_normalized:
+            _LOGGER.error("Cannot generate windows before normalizing data.")
+            raise RuntimeError()
+        
+        if self._are_windows_generated:
+            _LOGGER.warning("Windows have already been generated.")
+            return
+
+        self._train_dataset = self._create_windowed_dataset(self.train_sequence) # type: ignore
+        self._val_dataset = self._create_windowed_dataset(self.val_sequence) # type: ignore
+        self._test_dataset = self._create_windowed_dataset(self.test_sequence) # type: ignore
+        
+        self._are_windows_generated = True
+        _LOGGER.info("Feature and label windows generated for train, validation, and test sets.")
+
+    def _create_windowed_dataset(self, data: numpy.ndarray) -> Dataset:
+        """Efficiently creates windowed features and labels using numpy."""
+        if len(data) <= self.sequence_length:
+            # Validation/Test sets of size 0 might be passed
+            _LOGGER.warning(f"Data length ({len(data)}) is not greater than sequence_length ({self.sequence_length}). Cannot create windows. Returning empty dataset.")
+            return _PytorchDataset(numpy.array([]), numpy.array([]), labels_dtype=torch.float32)
+            
+        if self.prediction_mode == MLTaskKeys.SEQUENCE_VALUE:
+            # sequence-to-value
+            features = data[:-1]
+            labels = data[self.sequence_length:]
+            
+            n_windows = len(features) - self.sequence_length + 1
+            bytes_per_item = features.strides[0]
+            strided_features = numpy.lib.stride_tricks.as_strided(
+                features, shape=(n_windows, self.sequence_length), strides=(bytes_per_item, bytes_per_item)
+            )
+            # Ensure labels align with the end of each feature window
+            aligned_labels = labels[:n_windows]
+            return _PytorchDataset(strided_features, aligned_labels, labels_dtype=torch.float32)
+        
+        else:
+            # Sequence-to-sequence
+            x_data = data[:-1]
+            y_data = data[1:]
+            
+            n_windows = len(x_data) - self.sequence_length + 1
+            bytes_per_item = x_data.strides[0]
+            
+            strided_x = numpy.lib.stride_tricks.as_strided(x_data, shape=(n_windows, self.sequence_length), strides=(bytes_per_item, bytes_per_item))
+            strided_y = numpy.lib.stride_tricks.as_strided(y_data, shape=(n_windows, self.sequence_length), strides=(bytes_per_item, bytes_per_item))
+            
+            return _PytorchDataset(strided_x, strided_y, labels_dtype=torch.float32)
+
+    def plot_splits(self, save_dir: Union[str, Path]):
+        """Plots the training, validation and testing data."""
+        if not self._is_split:
+            _LOGGER.error("Cannot plot before splitting data.")
+            raise RuntimeError()
+        
+        if self.scaler is None:
+            _LOGGER.error("Cannot plot: data has not been normalized, or scaler is missing.")
+            return
+        
+        save_path = make_fullpath(save_dir, make=True, enforce="directory")
+        full_path = save_path / "SequenceSplits.svg"
+
+        plt.figure(figsize=(15, 6))
+        plt.title("Sequential Data")
+        plt.grid(True)
+        plt.xlabel("Sequence")
+        plt.ylabel("Value")
+        
+        # Plot denormalized training data
+        plt.plot(self.train_time_axis, self.scaler.inverse_transform(self.train_sequence.reshape(-1, 1)), label='Train Data') # type: ignore
+        
+        # Plot denormalized validation data
+        # We must skip the overlapping 'sequence_length' part for plotting
+        val_plot_data = self.val_sequence[self.sequence_length:] # type: ignore
+        plt.plot(self.val_time_axis, self.scaler.inverse_transform(val_plot_data.reshape(-1, 1)), label='Validation Data', c='orange') # type: ignore
+
+        # Plot denormalized test data
+        # We must skip the overlapping 'sequence_length' part for plotting
+        test_plot_data = self.test_sequence[self.sequence_length:] # type: ignore
+        plt.plot(self.test_time_axis, self.scaler.inverse_transform(test_plot_data.reshape(-1, 1)), label='Test Data', c='green') # type: ignore
+
+        plt.legend()
+        
+        plt.tight_layout()
+        plt.savefig(full_path)
+        _LOGGER.info(f"📈 Sequence data splits saved as '{full_path.name}'.")
+        plt.close()
+
+    def get_datasets(self) -> Tuple[Dataset, Dataset, Dataset]:
+        """Returns the final train, validation, and test datasets."""
+        if not self._are_windows_generated:
+            _LOGGER.error("Windows have not been generated. Call .generate_windows() first.")
+            raise RuntimeError()
+        return self._train_dataset, self._val_dataset, self._test_dataset # type: ignore
+    
+    def save_scaler(self, directory: Union[str, Path], verbose: bool=True) -> None:
+        """
+        Saves the fitted DragonScaler's state to a .pth file.
+        
+        Args:
+            directory (str | Path): The directory where the scaler will be saved.
+        """
+        if not self.scaler: 
+            _LOGGER.error("No scaler was fitted or provided.")
+            raise RuntimeError()
+
+        save_path = make_fullpath(directory, make=True, enforce="directory")
+
+        filename = f"{DatasetKeys.SCALER_PREFIX}{self.prediction_mode}.pth"
+        filepath = save_path / filename
+        self.scaler.save(filepath, verbose=False)
+        if verbose:
+            _LOGGER.info(f"Scaler saved as '{filepath.name}'.")
+    
+    def get_last_training_sequence(self) -> numpy.ndarray:
+        """
+        Returns the final, un-scaled sequence from the training data.
+        """
+        if not self._is_split:
+            _LOGGER.error("Data has not been split. Cannot get last training sequence.")
+            raise RuntimeError()
+        
+        # The length of train_time_axis is our validation split index
+        val_split_idx = len(self.train_time_axis) # type: ignore
+        
+        if val_split_idx < self.sequence_length:
+            _LOGGER.error(f"Training data length ({val_split_idx}) is less than sequence_length ({self.sequence_length}).")
+            raise ValueError()
+
+        # Get the slice from the *original* sequence
+        start_idx = val_split_idx - self.sequence_length
+        end_idx = val_split_idx
+        
+        return self.sequence[start_idx:end_idx]
+    
+    def __repr__(self) -> str:
+        s = f"<{self.__class__.__name__}>:\n"
+        s += f"  Prediction Mode: {self.prediction_mode}\n"
+        s += f"  Sequence Length (Window): {self.sequence_length}\n"
+        s += f"  Total Data Points: {len(self.sequence)}\n"
+        s += "  --- Status ---\n"
+        s += f"  Split: {self._is_split}\n"
+        s += f"  Normalized: {self._is_normalized}\n"
+        s += f"  Windows Generated: {self._are_windows_generated}\n"
+        
+        if self._are_windows_generated:
+            train_len = len(self._train_dataset) if self._train_dataset else 0 # type: ignore
+            val_len = len(self._val_dataset) if self._val_dataset else 0 # type: ignore
+            test_len = len(self._test_dataset) if self._test_dataset else 0 # type: ignore
+            s += f"  Datasets (Train | Validation | Test): {train_len} | {val_len} | {test_len} windows\n"
+            
+        return s
+
+
+def info():
+    _script_info(__all__)