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

ml_tools/ML_sequence_evaluation.py

@@ -0,0 +1,219 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import (
+    mean_squared_error,
+    mean_absolute_error,
+    r2_score, 
+    median_absolute_error,
+)
+from pathlib import Path
+from typing import Union, Optional
+
+from .path_manager import make_fullpath
+from ._logger import _LOGGER
+from ._script_info import _script_info
+from .ML_configuration import SequenceValueMetricsFormat, SequenceSequenceMetricsFormat, _BaseSequenceValueFormat, _BaseSequenceSequenceFormat
+
+
+__all__ = [
+    "sequence_to_value_metrics", 
+    "sequence_to_sequence_metrics"
+]
+
+DPI_value = 250
+
+
+def sequence_to_value_metrics(
+    y_true: np.ndarray, 
+    y_pred: np.ndarray, 
+    save_dir: Union[str, Path],
+    config: Optional[SequenceValueMetricsFormat] = None
+):
+    """
+    Saves regression metrics and plots for sequence-to-value (many-to-one) tasks.
+
+    Args:
+        y_true (np.ndarray): Ground truth values (1D array).
+        y_pred (np.ndarray): Predicted values (1D array).
+        save_dir (str | Path): Directory to save plots and report.
+        config (object): Formatting configuration object.
+    """
+    
+    # --- Ensure 1D input ---
+    if y_true.ndim > 1: y_true = y_true.flatten()
+    if y_pred.ndim > 1: y_pred = y_pred.flatten()
+
+    # --- Parse Config or use defaults ---
+    if config is None:
+        # Create a default config if one wasn't provided
+        format_config = _BaseSequenceValueFormat()
+    else:
+        format_config = config
+        
+    # --- Set Matplotlib font size ---
+    original_rc_params = plt.rcParams.copy()
+    plt.rcParams.update({'font.size': format_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)
+    medae = median_absolute_error(y_true, y_pred)
+    
+    report_lines = [
+        "--- Sequence-to-Value Regression Report ---",
+        f"  Root Mean Squared Error (RMSE): {rmse:.4f}",
+        f"  Mean Absolute Error (MAE):      {mae:.4f}",
+        f"  Median Absolute Error (MedAE):  {medae:.4f}",
+        f"  Coefficient of Determination (R²): {r2:.4f}"
+    ]
+    report_string = "\n".join(report_lines)
+
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+    # Save text report
+    report_path = save_dir_path / "sequence_to_value_report.txt"
+    report_path.write_text(report_string)
+    _LOGGER.info(f"📝 Seq-to-Value report saved as '{report_path.name}'")
+
+    # --- Save residual plot ---
+    residuals = y_true - y_pred
+    fig_res, ax_res = plt.subplots(figsize=(8, 6), dpi=DPI_value)
+    ax_res.scatter(y_pred, residuals, 
+                   alpha=format_config.scatter_alpha, 
+                   color=format_config.scatter_color)
+    ax_res.axhline(0, color=format_config.residual_line_color, linestyle='--')
+    ax_res.set_xlabel("Predicted Values")
+    ax_res.set_ylabel("Residuals")
+    ax_res.set_title("Sequence-to-Value Residual Plot")
+    ax_res.grid(True)
+    plt.tight_layout()
+    res_path = save_dir_path / "sequence_to_value_residual_plot.svg"
+    plt.savefig(res_path)
+    _LOGGER.info(f"📈 Seq-to-Value 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=DPI_value)
+    ax_tvp.scatter(y_true, y_pred, 
+                   alpha=format_config.scatter_alpha, 
+                   color=format_config.scatter_color)
+    ax_tvp.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 
+                linestyle='--', 
+                lw=2,
+                color=format_config.ideal_line_color)
+    ax_tvp.set_xlabel('True Values')
+    ax_tvp.set_ylabel('Predictions')
+    ax_tvp.set_title('Sequence-to-Value: True vs. Predicted')
+    ax_tvp.grid(True)
+    plt.tight_layout()
+    tvp_path = save_dir_path / "sequence_to_value_true_vs_predicted_plot.svg"
+    plt.savefig(tvp_path)
+    _LOGGER.info(f"📉 Seq-to-Value True vs. Predicted plot saved as '{tvp_path.name}'")
+    plt.close(fig_tvp)
+    
+    # --- Restore RC params ---
+    plt.rcParams.update(original_rc_params)
+
+
+def sequence_to_sequence_metrics(
+    y_true: np.ndarray, 
+    y_pred: np.ndarray, 
+    save_dir: Union[str, Path],
+    config: Optional[SequenceSequenceMetricsFormat] = None
+):
+    """
+    Saves per-step regression metrics for sequence-to-sequence (many-to-many) tasks.
+
+    Args:
+        y_true (np.ndarray): Ground truth sequences (n_samples, sequence_length).
+        y_pred (np.ndarray): Predicted sequences (n_samples, sequence_length).
+        save_dir (str | Path): Directory to save plots and report.
+        config (object): Formatting configuration object.
+    """
+
+    if y_true.ndim != 2 or y_pred.ndim != 2:
+        _LOGGER.error(f"Input arrays must be 2D (n_samples, sequence_length). Got y_true: {y_true.shape}, y_pred: {y_pred.shape}")
+        raise ValueError("Invalid input dimensions for sequence-to-sequence metrics.")
+
+    if y_true.shape != y_pred.shape:
+        _LOGGER.error(f"Input shapes do not match. Got y_true: {y_true.shape}, y_pred: {y_pred.shape}")
+        raise ValueError("Mismatched input shapes.")
+
+    # --- Parse Config or use defaults ---
+    if config is None:
+        format_config = _BaseSequenceSequenceFormat()
+    else:
+        format_config = config
+        
+    # --- Set Matplotlib font size ---
+    original_rc_params = plt.rcParams.copy()
+    plt.rcParams.update({'font.size': format_config.font_size})
+
+    sequence_length = y_true.shape[1]
+    steps = list(range(1, sequence_length + 1))
+    per_step_rmse = []
+    per_step_mae = []
+
+    # --- Calculate metrics for each time step ---
+    for i in range(sequence_length):
+        y_true_step = y_true[:, i]
+        y_pred_step = y_pred[:, i]
+        
+        rmse = np.sqrt(mean_squared_error(y_true_step, y_pred_step))
+        mae = mean_absolute_error(y_true_step, y_pred_step)
+        
+        per_step_rmse.append(rmse)
+        per_step_mae.append(mae)
+
+    # --- Create and save DataFrame ---
+    report_df = pd.DataFrame({
+        "step": steps,
+        "rmse": per_step_rmse,
+        "mae": per_step_mae
+    })
+
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+    report_path = save_dir_path / "sequence_to_sequence_report.csv"
+    report_df.to_csv(report_path, index=False)
+    _LOGGER.info(f"📝 Seq-to-Seq per-step report saved as '{report_path.name}'")
+
+    # --- Create and save plot ---
+    fig, ax1 = plt.subplots(figsize=format_config.plot_figsize, dpi=DPI_value)
+
+    # Plot RMSE
+    color_rmse = format_config.rmse_color
+    ax1.set_xlabel('Prediction Step')
+    ax1.set_ylabel('RMSE', color=color_rmse)
+    ax1.plot(steps, per_step_rmse, format_config.rmse_marker, color=color_rmse, label='RMSE')
+    ax1.tick_params(axis='y', labelcolor=color_rmse)
+    ax1.grid(True, linestyle=format_config.grid_style)
+
+    # Create a second y-axis for MAE
+    ax2 = ax1.twinx()
+    color_mae = format_config.mae_color
+    ax2.set_ylabel('MAE', color=color_mae)
+    ax2.plot(steps, per_step_mae, format_config.mae_marker, color=color_mae, label='MAE')
+    ax2.tick_params(axis='y', labelcolor=color_mae)
+
+    fig.suptitle('Sequence-to-Sequence Metrics (Per-Step)')
+    
+    # Add a single legend
+    lines, labels = ax1.get_legend_handles_labels()
+    lines2, labels2 = ax2.get_legend_handles_labels()
+    ax2.legend(lines + lines2, labels + labels2, loc='best')
+    
+    fig.tight_layout(rect=(0, 0.03, 1, 0.95)) # Adjust for suptitle
+    
+    plot_path = save_dir_path / "sequence_to_sequence_metrics_plot.svg"
+    plt.savefig(plot_path)
+    _LOGGER.info(f"📈 Seq-to-Seq per-step metrics plot saved as '{plot_path.name}'")
+    plt.close(fig)
+    
+    # --- Restore RC params ---
+    plt.rcParams.update(original_rc_params)
+
+
+def info():
+    _script_info(__all__)

ml_tools/ML_sequence_inference.py

@@ -0,0 +1,391 @@
+import torch
+from torch import nn
+import numpy as np
+from pathlib import Path
+from typing import Union, Literal, Dict, Any, Optional
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+from .ML_scaler import DragonScaler
+from ._script_info import _script_info
+from ._logger import _LOGGER
+from .path_manager import make_fullpath, sanitize_filename
+from ._keys import PyTorchInferenceKeys, MLTaskKeys, PyTorchCheckpointKeys
+from .ML_inference import _BaseInferenceHandler
+
+
+__all__ = [
+    "DragonSequenceInferenceHandler"
+]
+
+
+class DragonSequenceInferenceHandler(_BaseInferenceHandler):
+    """
+    Handles loading a PyTorch sequence model's state and performing inference
+    for univariate sequence tasks.
+    
+    This handler automatically scales inputs and de-scales outputs.
+    """
+    def __init__(self,
+                 model: nn.Module,
+                 state_dict: Union[str, Path],
+                 prediction_mode: Literal["sequence-to-sequence", "sequence-to-value"],
+                 scaler: Union[DragonScaler, str, Path],
+                 device: str = 'cpu'):
+        """
+        Initializes the handler for sequence tasks.
+
+        Args:
+            model (nn.Module): An instantiated PyTorch model architecture.
+            state_dict (str | Path): Path to the saved .pth model state_dict file.
+            prediction_mode (str): The type of sequence task.
+            device (str): The device to run inference on ('cpu', 'cuda', 'mps').
+            scaler (DragonScaler | str | Path): A DragonScaler instance or 
+                the file path to a saved DragonScaler state. This is required
+                to correctly scale inputs and de-scale predictions.
+        """
+        # Call the parent constructor to handle model loading and device
+        super().__init__(model, state_dict, device, scaler)
+        
+        self.sequence_length: Optional[int] = None
+        self.initial_sequence: Optional[np.ndarray] = None
+
+        if prediction_mode not in [MLTaskKeys.SEQUENCE_SEQUENCE, MLTaskKeys.SEQUENCE_VALUE]:
+            _LOGGER.error(f"'prediction_mode' not recognized: '{prediction_mode}'.")
+            raise ValueError()
+        self.prediction_mode = prediction_mode
+        
+        if self.scaler is None:
+            _LOGGER.error("A 'scaler' is required for DragonSequenceInferenceHandler to scale inputs and de-scale predictions.")
+            raise ValueError()
+        
+        # Load sequence length from the loaded dict (populated by _BaseInferenceHandler)
+        if PyTorchCheckpointKeys.SEQUENCE_LENGTH in self._loaded_data_dict:
+            try:
+                self.sequence_length = int(self._loaded_data_dict[PyTorchCheckpointKeys.SEQUENCE_LENGTH])
+                _LOGGER.info(f"'{PyTorchCheckpointKeys.SEQUENCE_LENGTH}' found and set to {self.sequence_length}")
+            except Exception as e_int:
+                _LOGGER.warning(f"State Dictionary has the key '{PyTorchCheckpointKeys.SEQUENCE_LENGTH}' but an error occurred when retrieving it:\n{e_int}")
+        else:
+            _LOGGER.warning(f"'{PyTorchCheckpointKeys.SEQUENCE_LENGTH}' not found in model file. Forecasting validation will be skipped.")
+            
+        # Load initial sequence
+        if PyTorchCheckpointKeys.INITIAL_SEQUENCE in self._loaded_data_dict:
+            try:
+                self.initial_sequence = self._loaded_data_dict[PyTorchCheckpointKeys.INITIAL_SEQUENCE]
+                _LOGGER.info(f"Default 'initial_sequence' for forecasting loaded from model file.")
+                # Optional: Validate shape
+                if self.sequence_length and len(self.initial_sequence) != self.sequence_length: # type: ignore
+                    _LOGGER.warning(f"Loaded 'initial_sequence' length ({len(self.initial_sequence)}) mismatches 'sequence_length' ({self.sequence_length}).") # type: ignore
+            except Exception as e_seq:
+                _LOGGER.warning(f"State Dictionary has the key '{PyTorchCheckpointKeys.INITIAL_SEQUENCE}' but an error occurred when retrieving it:\n{e_seq}")
+        else:
+            _LOGGER.info("No default 'initial_sequence' found in model file. Must be provided for forecasting.")
+
+    def _preprocess_input(self, features: torch.Tensor) -> torch.Tensor:
+        """
+        Converts input sequence to a torch.Tensor, applies scaling, and moves it to the correct device.
+
+        Overrides _BaseInferenceHandler._preprocess_input.
+
+        Args:
+            features (torch.Tensor): Input tensor of shape (batch_size, sequence_length).
+
+        Returns:
+            torch.Tensor: Scaled tensor on the correct device.
+        """
+        if self.scaler is None:
+            # This check is redundant due to __init__ check, but good for safety.
+            _LOGGER.error("Scaler is not available for preprocessing.")
+            raise RuntimeError()
+            
+        features_tensor = features.float()
+        
+        # Scale the sequence values
+        # (batch, seq_len) -> (batch * seq_len, 1)
+        batch_size, seq_len = features_tensor.shape
+        features_flat = features_tensor.reshape(-1, 1)
+        
+        scaled_flat = self.scaler.transform(features_flat)
+        
+        # (batch * seq_len, 1) -> (batch, seq_len)
+        scaled_features = scaled_flat.reshape(batch_size, seq_len)
+
+        return scaled_features.to(self.device)
+
+    def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Core batch prediction method for sequences.
+        Runs a batch of sequences through the model, de-scales the output,
+        and returns the predictions.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 2D array/tensor of input sequences, shape (batch_size, sequence_length).
+
+        Returns:
+            A dictionary containing the de-scaled prediction tensors.
+        """
+        if features.ndim != 2:
+            _LOGGER.error("Input for batch prediction must be a 2D array or tensor (batch_size, sequence_length).")
+            raise ValueError()
+        
+        if isinstance(features, np.ndarray):
+            features_tensor = torch.from_numpy(features).float()
+        else:
+            features_tensor = features.float()
+
+        # _preprocess_input scales the data and moves it to the correct device
+        input_tensor = self._preprocess_input(features_tensor) 
+
+        with torch.no_grad():
+            scaled_output = self.model(input_tensor)
+
+        # De-scale the output using the scaler
+        if self.scaler is None: # Should be impossible due to __init__
+             raise RuntimeError("Scaler not found for de-scaling.")
+
+        if self.prediction_mode == MLTaskKeys.SEQUENCE_VALUE:
+            # scaled_output is (batch)
+            # Reshape to (batch, 1) for scaler
+            scaled_output_reshaped = scaled_output.reshape(-1, 1)
+            descaled_output = self.scaler.inverse_transform(scaled_output_reshaped)
+            descaled_output = descaled_output.squeeze(-1) # (batch)
+        
+        elif self.prediction_mode == MLTaskKeys.SEQUENCE_SEQUENCE:
+            # scaled_output is (batch, seq_len)
+            batch_size, seq_len = scaled_output.shape
+            scaled_flat = scaled_output.reshape(-1, 1)
+            descaled_flat = self.scaler.inverse_transform(scaled_flat)
+            descaled_output = descaled_flat.reshape(batch_size, seq_len)
+        
+        else:
+             # Should not happen
+            _LOGGER.error(f"Invalid prediction mode: {self.prediction_mode}")
+            raise RuntimeError()
+
+        return {PyTorchInferenceKeys.PREDICTIONS: descaled_output}
+
+    def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Core single-sample prediction method for sequences.
+        Runs a single sequence through the model, de-scales the output,
+        and returns the prediction.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 1D array/tensor of 
+                input features, shape (sequence_length).
+
+        Returns:
+            A dictionary containing the de-scaled prediction tensor.
+        """
+        if features.ndim == 1:
+            features = features.reshape(1, -1) # Reshape (seq_len) to (1, seq_len)
+        
+        if features.shape[0] != 1 or features.ndim != 2:
+            _LOGGER.error("The 'predict()' method is for a single sequence (1D tensor). Use 'predict_batch()' for multiple sequences (2D tensor).")
+            raise ValueError()
+
+        batch_results = self.predict_batch(features)
+
+        # Extract the first (and only) result from the batch output
+        # For seq-to-value, result is shape ()
+        # For seq-to-seq, result is shape (seq_len)
+        single_results = {key: value[0] for key, value in batch_results.items()}
+        return single_results
+    
+    # --- NumPy Convenience Wrappers (on CPU) ---
+
+    def predict_batch_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, np.ndarray]:
+        """
+        Convenience wrapper for predict_batch that returns NumPy arrays.
+        
+        Args:
+            features (np.ndarray | torch.Tensor): A 2D array/tensor of 
+                input sequences, shape (batch_size, sequence_length).
+
+        Returns:
+            A dictionary containing the de-scaled prediction as a NumPy array.
+        """
+        tensor_results = self.predict_batch(features)
+        numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
+        return numpy_results
+
+    def predict_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
+        """
+        Convenience wrapper for predict that returns NumPy arrays or scalars.
+        
+        Args:
+            features (np.ndarray | torch.Tensor): A 1D array/tensor of 
+                input features, shape (sequence_length).
+                
+        Returns:
+            A dictionary containing the de-scaled prediction.
+            - For 'sequence-to-value', the value is a Python scalar.
+            - For 'sequence-to-sequence', the value is a 1D NumPy array.
+        """
+        tensor_results = self.predict(features)
+        
+        if self.prediction_mode == MLTaskKeys.SEQUENCE_VALUE:
+             # Prediction is a 0-dim tensor, .item() gets the scalar
+             return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
+        else: # sequence-to-sequence
+             # Prediction is a 1D tensor
+             return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].cpu().numpy()}
+        
+    def forecast(self, 
+                 n_steps: int,
+                 initial_sequence: Optional[Union[np.ndarray, torch.Tensor]]=None) -> np.ndarray:
+        """
+        Autoregressively forecasts 'n_steps' into the future.
+
+        This method works for both 'sequence-to-value' and 
+        'sequence-to-sequence' models.
+        
+        If 'initial_sequence' is not provided, this method will use the
+        default sequence that was saved with the model (if available).
+
+        Args:
+            initial_sequence (np.ndarray | torch.Tensor): The sequence
+                to start forecasting from. If None, uses the loaded default.
+                This should be a 1D array of *un-scaled* data.
+            n_steps (int): The number of future time steps to predict.
+
+        Returns:
+            np.ndarray: A 1D array containing the 'n_steps' forecasted values.
+        """
+        # --- Validation ---
+        if initial_sequence is None:
+            if self.initial_sequence is None:
+                _LOGGER.error("No 'initial_sequence' provided and no default sequence was loaded. Cannot forecast.")
+                raise ValueError()
+            _LOGGER.info("Using default 'initial_sequence' loaded from model file for forecast.")
+            initial_sequence_tensor = torch.from_numpy(self.initial_sequence).float()
+        elif isinstance(initial_sequence, np.ndarray):
+            initial_sequence_tensor = torch.from_numpy(initial_sequence).float()
+        else:
+            initial_sequence_tensor = initial_sequence.float()
+
+        if initial_sequence_tensor.ndim != 1:
+             _LOGGER.error(f"initial_sequence must be a 1D array. Got {initial_sequence_tensor.ndim} dimensions.")
+             raise ValueError()
+        
+        if self.sequence_length is not None:
+            if len(initial_sequence_tensor) != self.sequence_length:
+                _LOGGER.error(f"Input sequence length ({len(initial_sequence_tensor)}) does not match model's required sequence_length ({self.sequence_length}).")
+                raise ValueError()
+        else:
+            _LOGGER.warning("Model's 'sequence_length' is unknown. Cannot validate input sequence length. Assuming it is correct.")
+        
+        # --- Pre-processing ---
+        # 1. Scale the entire initial sequence
+        # We need to use the scaler: (seq_len) -> (seq_len, 1)
+        if self.scaler is None: # Should be impossible due to __init__
+            raise RuntimeError("Scaler not found for forecasting.")
+            
+        scaled_sequence_flat = self.scaler.transform(initial_sequence_tensor.reshape(-1, 1))
+        # (seq_len, 1) -> (seq_len)
+        current_scaled_sequence = scaled_sequence_flat.squeeze(-1).to(self.device)
+        
+        descaled_predictions = []
+
+        # --- Autoregressive Loop ---
+        self.model.eval() # Ensure model is in eval mode
+        with torch.no_grad():
+            for _ in range(n_steps):
+                # (seq_len) -> (1, seq_len)
+                input_tensor = current_scaled_sequence.reshape(1, -1)
+                
+                # Run the model
+                # input_tensor is (1, seq_len)
+                model_output = self.model(input_tensor).squeeze() # remove batch dim
+                
+                # Extract the single new prediction
+                if self.prediction_mode == MLTaskKeys.SEQUENCE_VALUE:
+                    # Output is shape (), a single scalar tensor
+                    scaled_prediction = model_output
+                else: # MLTaskKeys.SEQUENCE_SEQUENCE
+                    # Output is shape (seq_len), we need the last value
+                    scaled_prediction = model_output[-1]
+                
+                # De-scale the prediction for storage
+                # scaler input (1, 1)
+                descaled_prediction = self.scaler.inverse_transform(scaled_prediction.reshape(1, 1)).item()
+                descaled_predictions.append(descaled_prediction)
+                
+                # Create the new input sequence for the next loop
+                # "autoregression": roll the window by dropping the first value and appending the new scaled prediction.
+                # .unsqueeze(0) is needed to make the 0-dim tensor 1-dim for cat
+                current_scaled_sequence = torch.cat((current_scaled_sequence[1:], scaled_prediction.unsqueeze(0)))
+                
+        return np.array(descaled_predictions)
+    
+    def plot_forecast(self,  
+                      n_steps: int, 
+                      save_dir: Union[str, Path], 
+                      filename: str = "forecast_plot.svg",
+                      initial_sequence: Optional[Union[np.ndarray, torch.Tensor]]=None):
+        """
+        Runs a forecast and saves a plot of the results.
+
+        Args:
+            n_steps (int): The number of future time steps to predict.
+            save_dir (str | Path): Directory to save the plot.
+            filename (str, optional): Name for the saved plot file.
+            initial_sequence (np.ndarray | torch.Tensor | None): The sequence
+                to start forecasting from. If None, uses the loaded default.
+        """
+        # --- 1. Get Forecast Data ---
+        predictions = self.forecast(n_steps=n_steps, 
+                                    initial_sequence=initial_sequence)
+        
+        # --- 2. Determine which initial sequence was used for plotting ---
+        if initial_sequence is None:
+            plot_initial_sequence = self.initial_sequence
+            if plot_initial_sequence is None: # Should be caught by forecast() but good to check
+                 _LOGGER.error("Cannot plot: No 'initial_sequence' provided and no default found.")
+                 return
+        elif isinstance(initial_sequence, torch.Tensor):
+            plot_initial_sequence = initial_sequence.cpu().numpy()
+        else: # Is numpy array
+            plot_initial_sequence = initial_sequence
+            
+        # --- 3. Create X-axis indices ---
+        # The x-axis will be integer time steps
+        seq_len = len(plot_initial_sequence)
+        history_x = np.arange(0, seq_len)
+        forecast_x = np.arange(seq_len, seq_len + n_steps)
+
+        # --- 4. Plot ---
+        sns.set_theme(style="darkgrid")
+        plt.figure(figsize=(12, 6))
+
+        # Plot the historical data
+        plt.plot(history_x, plot_initial_sequence, label="Historical Data")
+        
+        # Plot the forecasted data
+        plt.plot(forecast_x, predictions, label="Forecasted Data", linestyle="--")
+        
+        # Add a vertical line to mark the start of the forecast
+        plt.axvline(x=history_x[-1], color='red', linestyle=':', label='Forecast Start')
+
+        plt.title(f"{n_steps}-Step Forecast")
+        plt.xlabel("Time Step")
+        plt.ylabel("Value")
+        plt.legend()
+        plt.tight_layout()
+
+        # --- 5. Save Plot ---
+        dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+        full_path = dir_path / sanitize_filename(filename)
+        
+        try:
+            plt.savefig(full_path)
+            _LOGGER.info(f"📈 Forecast plot saved to '{full_path.name}'.")
+        except Exception as e:
+            _LOGGER.error(f"Failed to save plot:\n{e}")
+        finally:
+            plt.close()
+    
+
+def info():
+    _script_info(__all__)