tirex-mirror 2025.12.2__py3-none-any.whl → 2025.12.16__py3-none-any.whl

tirex/models/base/base_classifier.py

@@ -0,0 +1,36 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from .base_tirex import BaseTirexEmbeddingModel
+
+
+class BaseTirexClassifier(BaseTirexEmbeddingModel):
+    """Abstract base class for TiRex classification models."""
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class labels for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted class labels with shape (batch_size,).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict(emb)).long()
+
+    @torch.inference_mode()
+    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class probabilities for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict_proba(emb))

tirex/models/base/base_regressor.py

@@ -0,0 +1,23 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from .base_tirex import BaseTirexEmbeddingModel
+
+
+class BaseTirexRegressor(BaseTirexEmbeddingModel):
+    """Abstract base class for TiRex regression models."""
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict values for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted values with shape (batch_size,).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict(emb)).float()

tirex/models/{classification/heads/base_classifier.py → base/base_tirex.py}

@@ -3,15 +3,18 @@
 
 from abc import ABC, abstractmethod
 
+import numpy as np
 import torch
 
+from tirex.util import train_val_split
+
 from ..embedding import TiRexEmbedding
 
 
-class BaseTirexClassifier(ABC):
-    """Abstract base class for TiRex classification models.
+class BaseTirexEmbeddingModel(ABC):
+    """Abstract base class for TiRex models.
 
-    This base class provides common functionality for all TiRex classifiers,
+    This base class provides common functionality for all TiRex classifier and regression models,
     including embedding model initialization and a consistent interface.
 
     """
@@ -19,7 +22,10 @@ class BaseTirexClassifier(ABC):
     def __init__(
         self, data_augmentation: bool = False, device: str | None = None, compile: bool = False, batch_size: int = 512
     ) -> None:
-        """Initializes a TiRex classification model.
+        """Initializes a base TiRex model.
+
+        This base class initializes the embedding model and common configuration
+        used by both classification and regression models.
 
         Args:
             data_augmentation : bool
@@ -49,45 +55,55 @@ class BaseTirexClassifier(ABC):
 
     @abstractmethod
     def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
+        """Abstract method for model training"""
         pass
 
-    @torch.inference_mode()
-    def predict(self, x: torch.Tensor) -> torch.Tensor:
-        """Predict class labels for input time series data.
+    def _compute_embeddings(self, x: torch.Tensor) -> np.ndarray:
+        """Compute embeddings for input time series data.
 
         Args:
             x: Input time series data as torch.Tensor with shape
                 (batch_size, num_variates, seq_len).
-        Returns:
-            torch.Tensor: Predicted class labels with shape (batch_size,).
-        """
-        self.emb_model.eval()
-        x = x.to(self.device)
-        embeddings = self.emb_model(x).cpu().numpy()
-        return torch.from_numpy(self.head.predict(embeddings)).long()
-
-    @torch.inference_mode()
-    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
-        """Predict class probabilities for input time series data.
 
-        Args:
-            x: Input time series data as torch.Tensor with shape
-                (batch_size, num_variates, seq_len).
         Returns:
-            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+            np.ndarray: Embeddings with shape (batch_size, embedding_dim).
         """
         self.emb_model.eval()
         x = x.to(self.device)
-        embeddings = self.emb_model(x).cpu().numpy()
-        return torch.from_numpy(self.head.predict_proba(embeddings))
+        return self.emb_model(x).cpu().numpy()
+
+    def _create_train_val_datasets(
+        self,
+        train_data: tuple[torch.Tensor, torch.Tensor],
+        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
+        val_split_ratio: float = 0.2,
+        stratify: bool = False,
+        seed: int | None = None,
+    ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+        if val_data is None:
+            train_data, val_data = train_val_split(
+                train_data=train_data, val_split_ratio=val_split_ratio, stratify=stratify, seed=seed
+            )
+        return train_data, val_data
 
     @abstractmethod
     def save_model(self, path: str) -> None:
-        """Saving model abstract method"""
+        """Save model to file.
+
+        Args:
+            path: File path where the model should be saved.
+        """
         pass
 
     @classmethod
     @abstractmethod
     def load_model(cls, path: str):
-        """Loading model abstract method"""
+        """Load model from file.
+
+        Args:
+            path: File path to the saved model checkpoint.
+
+        Returns:
+            Instance of the model class with loaded weights and configuration.
+        """
         pass

tirex/models/classification/__init__.py

@@ -1,8 +1,8 @@
 # Copyright (c) NXAI GmbH.
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
-from .heads.gbm_classifier import TirexGBMClassifier
-from .heads.linear_classifier import TirexClassifierTorch
-from .heads.rf_classifier import TirexRFClassifier
+from .gbm_classifier import TirexGBMClassifier
+from .linear_classifier import TirexLinearClassifier
+from .rf_classifier import TirexRFClassifier
 
-__all__ = ["TirexClassifierTorch", "TirexRFClassifier", "TirexGBMClassifier"]
+__all__ = ["TirexLinearClassifier", "TirexRFClassifier", "TirexGBMClassifier"]

tirex/models/classification/{heads/gbm_classifier.py → gbm_classifier.py}

@@ -5,8 +5,7 @@ import joblib
 import torch
 from lightgbm import LGBMClassifier, early_stopping
 
-from ..utils import train_val_split
-from .base_classifier import BaseTirexClassifier
+from ..base.base_classifier import BaseTirexClassifier
 
 
 class TirexGBMClassifier(BaseTirexClassifier):
@@ -118,11 +117,11 @@ class TirexGBMClassifier(BaseTirexClassifier):
             seed=self.random_state,
         )
 
-        self.emb_model.eval()
         X_train = X_train.to(self.device)
         X_val = X_val.to(self.device)
-        embeddings_train = self.emb_model(X_train).cpu().numpy()
-        embeddings_val = self.emb_model(X_val).cpu().numpy()
+
+        embeddings_train = self._compute_embeddings(X_train)
+        embeddings_val = self._compute_embeddings(X_val)
 
         y_train = y_train.detach().cpu().numpy() if isinstance(y_train, torch.Tensor) else y_train
         y_val = y_val.detach().cpu().numpy() if isinstance(y_val, torch.Tensor) else y_val
@@ -187,24 +186,3 @@ class TirexGBMClassifier(BaseTirexClassifier):
         model.random_state = getattr(model.head, "random_state", None)
 
         return model
-
-    def _create_train_val_datasets(
-        self,
-        train_data: tuple[torch.Tensor, torch.Tensor],
-        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
-        val_split_ratio: float = 0.2,
-        stratify: bool = True,
-        seed: int | None = None,
-    ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
-        X_train, y_train = train_data
-
-        if val_data is None:
-            train_data, val_data = train_val_split(
-                train_data=train_data, val_split_ratio=val_split_ratio, stratify=stratify, seed=seed
-            )
-            X_train, y_train = train_data
-            X_val, y_val = val_data
-        else:
-            X_val, y_val = val_data
-
-        return (X_train, y_train), (X_val, y_val)

tirex/models/classification/{heads/linear_classifier.py → linear_classifier.py}

@@ -5,11 +5,11 @@ from dataclasses import asdict
 
 import torch
 
-from ..trainer import TrainConfig, Trainer
-from .base_classifier import BaseTirexClassifier
+from ..base.base_classifier import BaseTirexClassifier
+from ..trainer import TrainConfig, Trainer, TrainingMetrics
 
 
-class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
+class TirexLinearClassifier(BaseTirexClassifier, torch.nn.Module):
     """
     A PyTorch classifier that combines time series embeddings with a linear classification head.
 
@@ -19,10 +19,10 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
 
     Example:
         >>> import torch
-        >>> from tirex.models.classification import TirexClassifierTorch
+        >>> from tirex.models.classification import TirexLinearClassifier
         >>>
         >>> # Create model with TiRex embeddings
-        >>> model = TirexClassifierTorch(
+        >>> model = TirexLinearClassifier(
         ...     data_augmentation=True,
         ...     max_epochs=2,
         ...     lr=1e-4,
@@ -72,7 +72,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
             compile: bool
                 Whether to compile the frozen embedding model. Default: False
             max_epochs : int
-                Maximum number of training epochs. Default: 50
+                Maximum number of training epochs. Default: 10
             lr : float
                 Learning rate for the optimizer. Default: 1e-4
             weight_decay : float
@@ -115,6 +115,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
             lr=lr,
             weight_decay=weight_decay,
             class_weights=class_weights,
+            task_type="classification",
             batch_size=batch_size,
             val_split_ratio=val_split_ratio,
             stratify=stratify,
@@ -132,9 +133,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
 
     @torch.inference_mode()
     def _identify_head_dims(self, x: torch.Tensor, y: torch.Tensor) -> None:
-        self.emb_model.eval()
-        sample_emb = self.emb_model(x[:1])
-        self.emb_dim = sample_emb.shape[-1]
+        self.emb_dim = self._compute_embeddings(x[:1]).shape[-1]
         self.num_classes = len(torch.unique(y))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -155,7 +154,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
 
     def fit(
         self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
-    ) -> dict[str, float]:
+    ) -> TrainingMetrics:
         """Train the classification head on the provided data.
 
         This method initializes the classification head based on the data dimensions,
@@ -231,7 +230,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
         )
 
     @classmethod
-    def load_model(cls, path: str) -> "TirexClassifierTorch":
+    def load_model(cls, path: str) -> "TirexLinearClassifier":
         """Load a saved model from file.
 
         This reconstructs the model architecture and loads the trained weights from
@@ -240,7 +239,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
         Args:
             path: File path to the saved model checkpoint.
         Returns:
-            TirexClassifierTorch: The loaded model with trained weights, ready for inference.
+            TirexLinearClassifier: The loaded model with trained weights, ready for inference.
         """
         checkpoint = torch.load(path)
 

tirex/models/classification/{heads/rf_classifier.py → rf_classifier.py}

@@ -5,7 +5,7 @@ import joblib
 import torch
 from sklearn.ensemble import RandomForestClassifier
 
-from .base_classifier import BaseTirexClassifier
+from ..base.base_classifier import BaseTirexClassifier
 
 
 class TirexRFClassifier(BaseTirexClassifier):
@@ -84,9 +84,9 @@ class TirexRFClassifier(BaseTirexClassifier):
         if isinstance(y_train, torch.Tensor):
             y_train = y_train.detach().cpu().numpy()
 
-        self.emb_model.eval()
         X_train = X_train.to(self.device)
-        embeddings = self.emb_model(X_train).cpu().numpy()
+        embeddings = self._compute_embeddings(X_train)
+
         self.head.fit(embeddings, y_train)
 
     def save_model(self, path: str) -> None:

tirex/models/{classification/embedding.py → embedding.py}

@@ -6,8 +6,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from tirex import load_model
-
-from .utils import nanmax, nanmin, nanstd
+from tirex.util import nanmax, nanmin, nanstd
 
 
 class TiRexEmbedding(nn.Module):
@@ -20,7 +19,7 @@ class TiRexEmbedding(nn.Module):
         self.batch_size = batch_size
 
         if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
         self.device = device
         self._compile = compile
 

tirex/models/regression/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from .gbm_regressor import TirexGBMRegressor
+from .linear_regressor import TirexLinearRegressor
+from .rf_regressor import TirexRFRegressor
+
+__all__ = ["TirexLinearRegressor", "TirexRFRegressor", "TirexGBMRegressor"]

tirex/models/regression/gbm_regressor.py

@@ -0,0 +1,181 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import joblib
+import torch
+from lightgbm import LGBMRegressor, early_stopping
+
+from ..base.base_regressor import BaseTirexRegressor
+
+
+class TirexGBMRegressor(BaseTirexRegressor):
+    """
+    A Gradient Boosting regressor that uses time series embeddings as features.
+
+    This regressor combines a pre-trained embedding model for feature extraction with a
+    Gradient Boosting regressor.
+
+    Example:
+        >>> import torch
+        >>> from tirex.models.regression import TirexGBMRegressor
+        >>>
+        >>> # Create model with custom LightGBM parameters
+        >>> model = TirexGBMRegressor(
+        ...     data_augmentation=True,
+        ...     n_estimators=50,
+        ...     random_state=42
+        ... )
+        >>>
+        >>> # Prepare data (can use NumPy arrays or PyTorch tensors)
+        >>> X_train = torch.randn(100, 1, 128)  # 100 samples, 1 number of variates, 128 sequence length
+        >>> y_train = torch.randn(100,)  # target values
+        >>>
+        >>> # Train the model
+        >>> model.fit((X_train, y_train))
+        >>>
+        >>> # Make predictions
+        >>> X_test = torch.randn(20, 1, 128)
+        >>> predictions = model.predict(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        compile: bool = False,
+        batch_size: int = 512,
+        early_stopping_rounds: int | None = 10,
+        min_delta: float = 0.0,
+        val_split_ratio: float = 0.2,
+        # LightGBM kwargs
+        **lgbm_kwargs,
+    ) -> None:
+        """Initializes Embedding Based Gradient Boosting Regression model.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+            early_stopping_rounds: int | None
+                Number of rounds without improvement of all metrics for Early Stopping. Default: 10
+            min_delta: float
+                Minimum improvement in score to keep training. Default 0.0
+            val_split_ratio : float
+                Proportion of training data to use for validation, if validation data are not provided. Default: 0.2
+            **lgbm_kwargs
+                Additional keyword arguments to pass to LightGBM's LGBMRegressor.
+                Common options include n_estimators, max_depth, learning_rate, random_state, etc.
+        """
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
+
+        # Early Stopping callback
+        self.early_stopping_rounds = early_stopping_rounds
+        self.min_delta = min_delta
+
+        # Data split parameters:
+        self.val_split_ratio = val_split_ratio
+
+        # Extract random_state for train_val_split if provided
+        self.random_state = lgbm_kwargs.get("random_state", None)
+
+        self.head = LGBMRegressor(**lgbm_kwargs)
+
+    @torch.inference_mode()
+    def fit(
+        self,
+        train_data: tuple[torch.Tensor, torch.Tensor],
+        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> None:
+        """Train the LightGBM regressor on embedded time series data.
+
+        This method generates embeddings for the training data using the embedding
+        model, then trains the LightGBM regressor on these embeddings.
+
+        Args:
+            train_data: Tuple of (X_train, y_train) where X_train is the input time
+                series data (torch.Tensor) and y_train is a torch.Tensor
+                of target values.
+            val_data: Optional tuple of (X_val, y_val) for validation where X_val is the input time
+                series data (torch.Tensor) and y_val is a torch.Tensor
+                of target values. If None, validation data will be automatically split from train_data (20% split).
+        """
+
+        (X_train, y_train), (X_val, y_val) = self._create_train_val_datasets(
+            train_data=train_data,
+            val_data=val_data,
+            val_split_ratio=self.val_split_ratio,
+            seed=self.random_state,
+        )
+
+        X_train = X_train.to(self.device)
+        X_val = X_val.to(self.device)
+
+        embeddings_train = self._compute_embeddings(X_train)
+        embeddings_val = self._compute_embeddings(X_val)
+
+        y_train = y_train.detach().cpu().numpy() if isinstance(y_train, torch.Tensor) else y_train
+        y_val = y_val.detach().cpu().numpy() if isinstance(y_val, torch.Tensor) else y_val
+
+        self.head.fit(
+            embeddings_train,
+            y_train,
+            eval_set=[(embeddings_val, y_val)],
+            callbacks=[early_stopping(stopping_rounds=self.early_stopping_rounds, min_delta=self.min_delta)]
+            if self.early_stopping_rounds is not None
+            else None,
+        )
+
+    def save_model(self, path: str) -> None:
+        """This method saves the trained LightGBM regressor head (joblib format) and embedding information.
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.joblib').
+        """
+
+        payload = {
+            "data_augmentation": self.data_augmentation,
+            "compile": self._compile,
+            "batch_size": self.batch_size,
+            "early_stopping_rounds": self.early_stopping_rounds,
+            "min_delta": self.min_delta,
+            "val_split_ratio": self.val_split_ratio,
+            "head": self.head,
+        }
+        joblib.dump(payload, path)
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexGBMRegressor":
+        """Load a saved model from file.
+
+        This reconstructs the model with the embedding configuration and loads
+        the trained LightGBM regressor from a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexGBMRegressor: The loaded model with trained Gradient Boosting regressor, ready for inference.
+        """
+        checkpoint = joblib.load(path)
+
+        # Create new instance with saved configuration
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
+            batch_size=checkpoint["batch_size"],
+            early_stopping_rounds=checkpoint["early_stopping_rounds"],
+            min_delta=checkpoint["min_delta"],
+            val_split_ratio=checkpoint["val_split_ratio"],
+        )
+
+        # Load the trained LightGBM head
+        model.head = checkpoint["head"]
+
+        # Extract random_state from the loaded head if available
+        model.random_state = getattr(model.head, "random_state", None)
+
+        return model

tirex/models/regression/linear_regressor.py

@@ -0,0 +1,250 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from dataclasses import asdict
+
+import torch
+
+from ..base.base_regressor import BaseTirexRegressor
+from ..trainer import TrainConfig, Trainer, TrainingMetrics
+
+
+class TirexLinearRegressor(BaseTirexRegressor, torch.nn.Module):
+    """
+    A PyTorch regressor that combines time series embeddings with a linear regression head.
+
+    This model uses a pre-trained TiRex embedding model to generate feature representations from time series
+    data, followed by a linear layer (with optional dropout) for regression. The embedding backbone
+    is frozen during training, and only the regression head is trained.
+
+    Example:
+        >>> import torch
+        >>> from tirex.models.regression import TirexLinearRegressor
+        >>>
+        >>> # Create model with TiRex embeddings
+        >>> model = TirexLinearRegressor(
+        ...     data_augmentation=True,
+        ...     max_epochs=2,
+        ...     lr=1e-4,
+        ...     batch_size=32
+        ... )
+        >>>
+        >>> # Prepare data
+        >>> X_train = torch.randn(100, 1, 128)  # 100 samples, 1 number of variates, 128 sequence length
+        >>> y_train = torch.randn(100, 1)  # target values
+        >>>
+        >>> # Train the model
+        >>> metrics = model.fit((X_train, y_train)) # doctest: +ELLIPSIS
+        Epoch 1, Train Loss: ...
+        >>> # Make predictions
+        >>> X_test = torch.randn(20, 1, 128)
+        >>> predictions = model.predict(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        compile: bool = False,
+        # Training parameters
+        max_epochs: int = 10,
+        lr: float = 1e-4,
+        weight_decay: float = 0.01,
+        batch_size: int = 512,
+        val_split_ratio: float = 0.2,
+        patience: int = 7,
+        delta: float = 0.001,
+        log_every_n_steps: int = 5,
+        seed: int | None = None,
+        # Head parameters
+        dropout: float | None = None,
+    ) -> None:
+        """Initializes Embedding Based Linear Regression model.
+
+        Args:
+            data_augmentation : bool | None
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            max_epochs : int
+                Maximum number of training epochs. Default: 10
+            lr : float
+                Learning rate for the optimizer. Default: 1e-4
+            weight_decay : float
+                Weight decay coefficient. Default: 0.01
+            batch_size : int
+                Batch size for training and embedding calculations. Default: 512
+            val_split_ratio : float
+                Proportion of training data to use for validation, if validation data are not provided. Default: 0.2
+            patience : int
+                Number of epochs to wait for improvement before early stopping. Default: 7
+            delta : float
+                Minimum change in validation loss to qualify as an improvement. Default: 0.001
+            log_every_n_steps : int
+                Frequency of logging during training. Default: 5
+            seed : int | None
+                Random seed for reproducibility. If None, no seed is set. Default: None
+            dropout : float | None
+                Dropout probability for the regression head. If None, no dropout is used. Default: None
+        """
+
+        torch.nn.Module.__init__(self)
+
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
+
+        # Head parameters
+        self.dropout = dropout
+        self.head = None
+        self.emb_dim = None
+        self.output_dim = None
+
+        # Train config
+        train_config = TrainConfig(
+            max_epochs=max_epochs,
+            log_every_n_steps=log_every_n_steps,
+            device=self.device,
+            lr=lr,
+            weight_decay=weight_decay,
+            class_weights=None,
+            task_type="regression",
+            batch_size=batch_size,
+            val_split_ratio=val_split_ratio,
+            patience=patience,
+            delta=delta,
+            seed=seed,
+        )
+        self.trainer = Trainer(self, train_config=train_config)
+
+    def _init_regressor(self, emb_dim: int, output_dim: int, dropout: float | None) -> torch.nn.Module:
+        if dropout:
+            return torch.nn.Sequential(torch.nn.Dropout(p=dropout), torch.nn.Linear(emb_dim, output_dim))
+        else:
+            return torch.nn.Linear(emb_dim, output_dim)
+
+    @torch.inference_mode()
+    def _identify_head_dims(self, x: torch.Tensor, y: torch.Tensor) -> None:
+        self.emb_dim = self._compute_embeddings(x[:1]).shape[-1]
+        self.output_dim = 1
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the embedding model and regression head.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted values with shape (batch_size, 1).
+        Raises:
+            RuntimeError: If the regression head has not been initialized via fit().
+        """
+        if self.head is None:
+            raise RuntimeError("Head not initialized. Call fit() first to automatically build the head.")
+
+        embedding = self.emb_model(x).to(self.device)
+        return self.head(embedding)
+
+    def fit(
+        self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
+    ) -> TrainingMetrics:
+        """Train the regression head on the provided data.
+
+        This method initializes the regression head based on the data dimensions,
+        then trains it on provided data. The embedding model remains frozen.
+
+        Args:
+            train_data: Tuple of (X_train, y_train) where X_train is the input time series
+                data and y_train are the corresponding target values.
+            val_data: Optional tuple of (X_val, y_val) for validation. If None and
+                val_split_ratio > 0, validation data will be split from train_data.
+
+        Returns:
+            dict[str, float]: Dictionary containing final training and validation losses.
+        """
+        X_train, y_train = train_data
+
+        self._identify_head_dims(X_train, y_train)
+        self.head = self._init_regressor(self.emb_dim, self.output_dim, self.dropout)
+        self.head = self.head.to(self.trainer.device)
+
+        return self.trainer.fit(train_data, val_data=val_data)
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict values for input time series data.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted values with shape (batch_size, 1).
+        """
+        self.eval()
+        x = x.to(self.device)
+        return self.forward(x)
+
+    def save_model(self, path: str) -> None:
+        """Save the trained regression head.
+
+        This function saves the trained regression head weights (.pt format), embedding configuration,
+        model dimensions, and device information. The embedding model itself is not
+        saved as it uses a pre-trained backbone that can be reloaded.
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.pt').
+        """
+        train_config_dict = asdict(self.trainer.train_config)
+        torch.save(
+            {
+                "head_state_dict": self.head.state_dict(),  # need to save only head, embedding is frozen
+                "data_augmentation": self.data_augmentation,
+                "compile": self._compile,
+                "emb_dim": self.emb_dim,
+                "output_dim": self.output_dim,
+                "dropout": self.dropout,
+                "train_config": train_config_dict,
+            },
+            path,
+        )
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexLinearRegressor":
+        """Load a saved model from file.
+
+        This reconstructs the model architecture and loads the trained weights from
+        a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexLinearRegressor: The loaded model with trained weights, ready for inference.
+        """
+        checkpoint = torch.load(path)
+
+        # Extract train_config if available, otherwise use defaults
+        train_config_dict = checkpoint.get("train_config", {})
+
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
+            dropout=checkpoint["dropout"],
+            max_epochs=train_config_dict.get("max_epochs", 50),
+            lr=train_config_dict.get("lr", 1e-4),
+            weight_decay=train_config_dict.get("weight_decay", 0.01),
+            batch_size=train_config_dict.get("batch_size", 512),
+            val_split_ratio=train_config_dict.get("val_split_ratio", 0.2),
+            patience=train_config_dict.get("patience", 7),
+            delta=train_config_dict.get("delta", 0.001),
+            log_every_n_steps=train_config_dict.get("log_every_n_steps", 5),
+            seed=train_config_dict.get("seed", None),
+        )
+
+        # Initialize head with dimensions
+        model.emb_dim = checkpoint["emb_dim"]
+        model.output_dim = checkpoint.get("output_dim", checkpoint.get("num_classes", 1))  # Backward compatibility
+        model.head = model._init_regressor(model.emb_dim, model.output_dim, model.dropout)
+
+        # Load the trained weights
+        model.head.load_state_dict(checkpoint["head_state_dict"])
+        model.to(model.device)
+
+        return model

tirex/models/regression/rf_regressor.py

@@ -0,0 +1,130 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import joblib
+import torch
+from sklearn.ensemble import RandomForestRegressor
+
+from ..base.base_regressor import BaseTirexRegressor
+
+
+class TirexRFRegressor(BaseTirexRegressor):
+    """
+    A Random Forest regressor that uses time series embeddings as features.
+
+    This regressor combines a pre-trained embedding model for feature extraction with a scikit-learn
+    Random Forest regressor. The embedding model generates fixed-size feature vectors from variable-length
+    time series, which are then used to train the Random Forest.
+
+    Example:
+        >>> import torch
+        >>> from tirex.models.regression import TirexRFRegressor
+        >>>
+        >>> # Create model with custom Random Forest parameters
+        >>> model = TirexRFRegressor(
+        ...     data_augmentation=True,
+        ...     n_estimators=50,
+        ...     max_depth=10,
+        ...     random_state=42
+        ... )
+        >>>
+        >>> # Prepare data (can use NumPy arrays or PyTorch tensors)
+        >>> X_train = torch.randn(100, 1, 128)  # 100 samples, 1 number of variates, 128 sequence length
+        >>> y_train = torch.randn(100,)  # target values
+        >>>
+        >>> # Train the model
+        >>> model.fit((X_train, y_train))
+        >>>
+        >>> # Make predictions
+        >>> X_test = torch.randn(20, 1, 128)
+        >>> predictions = model.predict(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        compile: bool = False,
+        batch_size: int = 512,
+        # Random Forest parameters
+        **rf_kwargs,
+    ) -> None:
+        """Initializes Embedding Based Random Forest Regression model.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+            **rf_kwargs
+                Additional keyword arguments to pass to sklearn's RandomForestRegressor.
+                Common options include n_estimators, max_depth, min_samples_split, random_state, etc.
+        """
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
+        self.head = RandomForestRegressor(**rf_kwargs)
+
+    @torch.inference_mode()
+    def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
+        """Train the Random Forest regressor on embedded time series data.
+
+        This method generates embeddings for the training data using the embedding
+        model, then trains the Random Forest on these embeddings.
+
+        Args:
+            train_data: Tuple of (X_train, y_train) where X_train is the input time
+                series data (torch.Tensor) and y_train is a torch.Tensor
+                of target values.
+        """
+        X_train, y_train = train_data
+
+        if isinstance(y_train, torch.Tensor):
+            y_train = y_train.detach().cpu().numpy()
+
+        X_train = X_train.to(self.device)
+        embeddings = self._compute_embeddings(X_train)
+
+        self.head.fit(embeddings, y_train)
+
+    def save_model(self, path: str) -> None:
+        """This method saves the trained Random Forest regressor head and embedding information in joblib format
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.joblib').
+        """
+        payload = {
+            "data_augmentation": self.data_augmentation,
+            "compile": self._compile,
+            "batch_size": self.batch_size,
+            "head": self.head,
+        }
+        joblib.dump(payload, path)
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexRFRegressor":
+        """Load a saved model from file.
+
+        This reconstructs the model with the embedding configuration and loads
+        the trained Random Forest regressor from a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexRFRegressor: The loaded model with trained Random Forest regressor, ready for inference.
+        """
+        checkpoint = joblib.load(path)
+
+        # Create new instance with saved configuration
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
+            batch_size=checkpoint["batch_size"],
+        )
+
+        # Load the trained Random Forest head
+        model.head = checkpoint["head"]
+
+        return model

tirex/models/{classification/trainer.py → trainer.py}

@@ -2,11 +2,18 @@
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
 from dataclasses import dataclass
+from typing import TypedDict
 
 import torch
+from torch.optim import Optimizer
 from torch.utils.data import DataLoader, TensorDataset
 
-from .utils import EarlyStopping, set_seed, train_val_split
+from ..util import EarlyStopping, set_seed, train_val_split
+
+
+class TrainingMetrics(TypedDict):
+    train_loss: float
+    val_loss: float
 
 
 @dataclass
@@ -22,11 +29,7 @@ class TrainConfig:
 
     # Loss parameters
     class_weights: torch.Tensor | None
-
-    # Data loading parameters
-    batch_size: int
-    val_split_ratio: float
-    stratify: bool
+    task_type: str
 
     # Earlystopping parameters
     patience: int
@@ -35,6 +38,11 @@ class TrainConfig:
     # Reproducability
     seed: int | None
 
+    # Data loading parameters
+    batch_size: int
+    val_split_ratio: float
+    stratify: bool = False
+
     def __post_init__(self) -> None:
         if self.max_epochs <= 0:
             raise ValueError(f"max_epochs must be positive, got {self.max_epochs}")
@@ -60,6 +68,12 @@ class TrainConfig:
         if self.delta < 0:
             raise ValueError(f"delta must be non-negative, got {self.delta}")
 
+        if self.task_type not in ["classification", "regression"]:
+            raise ValueError(f"task_type must be 'classification' or 'regression', got {self.task_type}")
+
+        if self.stratify and self.task_type == "regression":
+            raise ValueError("stratify=True is not valid for regression tasks")
+
 
 class Trainer:
     def __init__(
@@ -74,14 +88,19 @@ class Trainer:
         class_weights = (
             self.train_config.class_weights.to(self.device) if self.train_config.class_weights is not None else None
         )
-        self.loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights).to(self.device)
-
-        self.optimizer = None
+        if self.train_config.task_type == "classification":
+            self.loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights).to(self.device)
+        elif self.train_config.task_type == "regression":
+            self.loss_fn = torch.nn.MSELoss().to(self.device)
+        else:
+            raise ValueError(f"Unsupported task_type: {self.train_config.task_type}")
+
+        self.optimizer: Optimizer | None = None
         self.early_stopper = EarlyStopping(patience=self.train_config.patience, delta=self.train_config.delta)
 
     def fit(
         self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
-    ) -> dict[str, float]:
+    ) -> TrainingMetrics:
         if self.train_config.seed is not None:
             set_seed(self.train_config.seed)
 

tirex/util.py

@@ -9,6 +9,7 @@ from typing import Literal, Optional
 
 import numpy as np
 import torch
+from sklearn.model_selection import train_test_split
 
 COLOR_CONTEXT = "#4a90d0"
 COLOR_FORECAST = "#d94e4e"
@@ -845,3 +846,84 @@ def plot_forecast(
     ax.grid()
 
     return ax
+
+
+# ==== Classification and Regression Utilities ====
+
+
+# Remove after Issue will be solved: https://github.com/pytorch/pytorch/issues/61474
+def nanmax(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    min_value = torch.finfo(tensor.dtype).min
+    output = tensor.nan_to_num(min_value).max(dim=dim, keepdim=keepdim)
+    return output.values
+
+
+def nanmin(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    max_value = torch.finfo(tensor.dtype).max
+    output = tensor.nan_to_num(max_value).min(dim=dim, keepdim=keepdim)
+    return output.values
+
+
+def nanvar(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    tensor_mean = tensor.nanmean(dim=dim, keepdim=True)
+    output = (tensor - tensor_mean).square().nanmean(dim=dim, keepdim=keepdim)
+    return output
+
+
+def nanstd(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    output = nanvar(tensor, dim=dim, keepdim=keepdim)
+    output = output.sqrt()
+    return output
+
+
+def train_val_split(
+    train_data: tuple[torch.Tensor, torch.Tensor],
+    val_split_ratio: float,
+    stratify: bool = False,
+    seed: int | None = None,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+    idx_train, idx_val = train_test_split(
+        np.arange(len(train_data[0])),
+        test_size=val_split_ratio,
+        random_state=seed,
+        shuffle=True,
+        stratify=train_data[1] if stratify else None,
+    )
+
+    return (
+        (train_data[0][idx_train], train_data[1][idx_train]),
+        (train_data[0][idx_val], train_data[1][idx_val]),
+    )
+
+
+def set_seed(seed: int) -> None:
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+    np.random.seed(seed)
+
+
+class EarlyStopping:
+    def __init__(
+        self,
+        patience: int = 7,
+        delta: float = 0.0001,
+    ) -> None:
+        self.patience: int = patience
+        self.delta: float = delta
+
+        self.best: float = np.inf
+        self.wait_count: int = 0
+        self.early_stop: bool = False
+
+    def __call__(self, epoch: int, val_loss: float) -> bool:
+        improved = val_loss < (self.best - self.delta)
+        if improved:
+            self.best = val_loss
+            self.wait_count = 0
+        else:
+            self.wait_count += 1
+            if self.wait_count >= self.patience:
+                self.early_stop = True
+                print(f"Early stopping triggered at epoch {epoch}.")
+        return self.early_stop

{tirex_mirror-2025.12.2.dist-info → tirex_mirror-2025.12.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.12.2
+Version: 2025.12.16
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -55,7 +55,7 @@ License: NXAI COMMUNITY LICENSE AGREEMENT
         
 Project-URL: Repository, https://github.com/rozsasarpi/tirex-mirror
 Project-URL: Issues, https://github.com/rozsasarpi/tirex-mirror/issues
-Keywords: TiRex,xLSTM,Time Series,Zero-shot,Deep Learning,Classification,Timeseries-Classification
+Keywords: TiRex,xLSTM,Timeseries,Zero-shot,Deep Learning,Timeseries-Forecasting,Timeseries-Classification,Timeseries-Regression
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.11
@@ -66,6 +66,7 @@ License-File: NOTICE.txt
 Requires-Dist: torch
 Requires-Dist: huggingface-hub
 Requires-Dist: numpy
+Requires-Dist: scikit-learn
 Provides-Extra: cuda
 Requires-Dist: xlstm; extra == "cuda"
 Requires-Dist: ninja; extra == "cuda"
@@ -84,9 +85,11 @@ Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev>=0.6.0; extra == "test"
 Requires-Dist: pytest; extra == "test"
+Requires-Dist: aeon; extra == "test"
 Provides-Extra: classification
-Requires-Dist: scikit-learn; extra == "classification"
 Requires-Dist: lightgbm[scikit-learn]; extra == "classification"
+Provides-Extra: regression
+Requires-Dist: lightgbm[scikit-learn]; extra == "regression"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"

tirex_mirror-2025.12.16.dist-info/RECORD

@@ -0,0 +1,34 @@
+tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
+tirex/base.py,sha256=pGShyI3LwPAwkl0rIYqik5qb8uBAWqbQrVoVjq7UH_8,5096
+tirex/util.py,sha256=D1HaHiAfZAEiOA_q2saOpugn0gozepBcwXwvQHS85cg,34308
+tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/api_adapter/forecast.py,sha256=FnRgdI_4vJ7iFTqyWwxPE_C7MBB3nhLTnVMzTr5pihc,17581
+tirex/api_adapter/gluon.py,sha256=0KfKX7dGSooVIUumGtjK6Va2YgI6Isa_8kz9j0oPXUM,1883
+tirex/api_adapter/hf_data.py,sha256=TRyys2xKIGZS0Yhq2Eb61lWCMg5CWWn1yRlLIN1mU7o,1369
+tirex/api_adapter/standard_adapter.py,sha256=vdlxNs8mTUtPgK_5WMqYqNdMj8W44igqWsAgtggt_xk,2809
+tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
+tirex/models/embedding.py,sha256=quDT_xKaaMmhXvZxz7LV2r4ICou2pzcjam7YnfOzxIs,5457
+tirex/models/patcher.py,sha256=8T4c3PZnOAsEpahhrjtt7S7405WUjN6g3cV33E55PD4,1911
+tirex/models/tirex.py,sha256=wtdjrdE1TuoueIvhlKf-deLH3oKyQuGmd2k15My7SWA,9710
+tirex/models/trainer.py,sha256=lkNSt5aHzjoAQTRVPj6QlFVSP81jWbkMcKl-2FsPNhQ,6680
+tirex/models/base/base_classifier.py,sha256=QU_eiwTe8B3QRT6Xs83utnnvlo6iaTzIv8VCId_S900,1310
+tirex/models/base/base_regressor.py,sha256=RA9FfRgJeAuE_pbaP_4pGIiYWRmMs4zgkb9ajqAYfsg,786
+tirex/models/base/base_tirex.py,sha256=JS82k0qzi19LNl3aW45ty4AQAIZktvZZCuT9lhMlpcg,3646
+tirex/models/classification/__init__.py,sha256=uBoN5BvZV_bHDFFQpOQ8E0CKAyU7BYkmkQndb48aqTk,361
+tirex/models/classification/gbm_classifier.py,sha256=ylcMJKC63ps4h-qoVFB7JaEzyuS8X3Mqt2B-gMYzN4I,7523
+tirex/models/classification/linear_classifier.py,sha256=I7etrhctJgjTRJISIXVs2v_efXkJvTz3G9YOuAIcL5E,11366
+tirex/models/classification/rf_classifier.py,sha256=jo9otjwhMhNkKCmVUdDbkvb2gbCLt5eTYMKclWyeQks,5036
+tirex/models/regression/__init__.py,sha256=GjIcpJXA1n-RzP3JjwLknIu1FgEQX6hh4C6U1fHLoL0,352
+tirex/models/regression/gbm_regressor.py,sha256=8RVk4gmQ3o0N58zDE11CnNZIAWZa1HQrw234-2nHu5o,7162
+tirex/models/regression/linear_regressor.py,sha256=n4Nr6HtTQVQ-T45yMCcOuk9Nj3kGt5V9vyX5q2UkQ_k,10152
+tirex/models/regression/rf_regressor.py,sha256=03c39NCDKmHZTB7yfytOXREqF1Lfdcc0sWV2FJQ0Ouw,4985
+tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
+tirex/models/slstm/cell.py,sha256=Otyil_AjpJbUckkINWGHxlqP14J5epm_J_zdWPzvD2g,7290
+tirex/models/slstm/layer.py,sha256=hrDydQJIAHf5W0A0Rt0hXG4yKXrOSY-HPL0UbigR6Q8,2867
+tirex_mirror-2025.12.16.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.12.16.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.12.16.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.12.16.dist-info/METADATA,sha256=z_2g2gtkAJMFpzCnIxTVA1R0k10-Fnq20VpocikYpBM,11911
+tirex_mirror-2025.12.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.12.16.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.12.16.dist-info/RECORD,,

tirex/models/classification/utils.py

@@ -1,84 +0,0 @@
-# Copyright (c) NXAI GmbH.
-# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
-
-import numpy as np
-import torch
-from sklearn.model_selection import train_test_split
-
-
-# Remove after Issue will be solved: https://github.com/pytorch/pytorch/issues/61474
-def nanmax(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
-    min_value = torch.finfo(tensor.dtype).min
-    output = tensor.nan_to_num(min_value).max(dim=dim, keepdim=keepdim)
-    return output.values
-
-
-def nanmin(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
-    max_value = torch.finfo(tensor.dtype).max
-    output = tensor.nan_to_num(max_value).min(dim=dim, keepdim=keepdim)
-    return output.values
-
-
-def nanvar(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
-    tensor_mean = tensor.nanmean(dim=dim, keepdim=True)
-    output = (tensor - tensor_mean).square().nanmean(dim=dim, keepdim=keepdim)
-    return output
-
-
-def nanstd(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
-    output = nanvar(tensor, dim=dim, keepdim=keepdim)
-    output = output.sqrt()
-    return output
-
-
-def train_val_split(
-    train_data: tuple[torch.Tensor, torch.Tensor],
-    val_split_ratio: float,
-    stratify: bool,
-    seed: int | None,
-) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
-    idx_train, idx_val = train_test_split(
-        np.arange(len(train_data[0])),
-        test_size=val_split_ratio,
-        random_state=seed,
-        shuffle=True,
-        stratify=train_data[1] if stratify else None,
-    )
-
-    return (
-        (train_data[0][idx_train], train_data[1][idx_train]),
-        (train_data[0][idx_val], train_data[1][idx_val]),
-    )
-
-
-def set_seed(seed: int) -> None:
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-    np.random.seed(seed)
-
-
-class EarlyStopping:
-    def __init__(
-        self,
-        patience: int = 7,
-        delta: float = 0.0001,
-    ) -> None:
-        self.patience: int = patience
-        self.delta: float = delta
-
-        self.best: float = np.inf
-        self.wait_count: int = 0
-        self.early_stop: bool = False
-
-    def __call__(self, epoch: int, val_loss: float) -> bool:
-        improved = val_loss < (self.best - self.delta)
-        if improved:
-            self.best = val_loss
-            self.wait_count = 0
-        else:
-            self.wait_count += 1
-            if self.wait_count >= self.patience:
-                self.early_stop = True
-                print(f"Early stopping triggered at epoch {epoch}.")
-        return self.early_stop

tirex_mirror-2025.12.2.dist-info/RECORD

@@ -1,29 +0,0 @@
-tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=pGShyI3LwPAwkl0rIYqik5qb8uBAWqbQrVoVjq7UH_8,5096
-tirex/util.py,sha256=ggfETirJ589Dr5o3QThxTnjhgNnMCk10bNoJghnoeoA,31672
-tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/api_adapter/forecast.py,sha256=FnRgdI_4vJ7iFTqyWwxPE_C7MBB3nhLTnVMzTr5pihc,17581
-tirex/api_adapter/gluon.py,sha256=0KfKX7dGSooVIUumGtjK6Va2YgI6Isa_8kz9j0oPXUM,1883
-tirex/api_adapter/hf_data.py,sha256=TRyys2xKIGZS0Yhq2Eb61lWCMg5CWWn1yRlLIN1mU7o,1369
-tirex/api_adapter/standard_adapter.py,sha256=vdlxNs8mTUtPgK_5WMqYqNdMj8W44igqWsAgtggt_xk,2809
-tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
-tirex/models/patcher.py,sha256=8T4c3PZnOAsEpahhrjtt7S7405WUjN6g3cV33E55PD4,1911
-tirex/models/tirex.py,sha256=wtdjrdE1TuoueIvhlKf-deLH3oKyQuGmd2k15My7SWA,9710
-tirex/models/classification/__init__.py,sha256=0PKH8Cr-VWT4x4SXzGAyj90jMJbxzlIyrBIiV1ZXfm4,377
-tirex/models/classification/embedding.py,sha256=wxIefIpccHe64Cb2JlmL0l0PhqYB_zdBHSh5auPqdlY,5452
-tirex/models/classification/trainer.py,sha256=CbYfdTNdj5iucchxNwgg35EF-MzrdjrMm9auNOLae4U,5895
-tirex/models/classification/utils.py,sha256=JbloQu4KlBHfdlH1NYYg7_VIoH-t2HBo5ovdyAEoCGs,2735
-tirex/models/classification/heads/base_classifier.py,sha256=mBJV7jShBAUQgMpidow8lVfqhufNHvEIwDnidtxA1Hw,3236
-tirex/models/classification/heads/gbm_classifier.py,sha256=V1CzfmJF6Tp2j14F3dAmRZ4wxBsuKoK1jQu4HbjttCc,8365
-tirex/models/classification/heads/linear_classifier.py,sha256=2BI2GqQvBiHjFJ7_dS8rGQdQ-pnpF7-9sZv_YRUL-8w,11351
-tirex/models/classification/heads/rf_classifier.py,sha256=Nsw3mxToa804txYzkvkRQwq34whEb5x-9_6ywskylBk,5063
-tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
-tirex/models/slstm/cell.py,sha256=Otyil_AjpJbUckkINWGHxlqP14J5epm_J_zdWPzvD2g,7290
-tirex/models/slstm/layer.py,sha256=hrDydQJIAHf5W0A0Rt0hXG4yKXrOSY-HPL0UbigR6Q8,2867
-tirex_mirror-2025.12.2.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.12.2.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.12.2.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.12.2.dist-info/METADATA,sha256=XU9yHQPFUYg5qiqMSNoOJRGmSE9dY3GUhGfn8m4LouM,11783
-tirex_mirror-2025.12.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.12.2.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.12.2.dist-info/RECORD,,