tirex-mirror 2025.11.26__tar.gz → 2025.11.28__tar.gz

{tirex_mirror-2025.11.26 → tirex_mirror-2025.11.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.11.26
+Version: 2025.11.28
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -84,6 +84,8 @@ Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev>=0.6.0; extra == "test"
 Requires-Dist: pytest; extra == "test"
+Provides-Extra: classification
+Requires-Dist: scikit-learn; extra == "classification"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -95,6 +97,7 @@ Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
 Requires-Dist: pytest; extra == "all"
 Requires-Dist: fev>=0.6.0; extra == "all"
+Requires-Dist: scikit-learn; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

{tirex_mirror-2025.11.26 → tirex_mirror-2025.11.28}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tirex-mirror"
-version = "2025.11.26"
+version = "2025.11.28"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -29,7 +29,8 @@ plotting = [ "matplotlib",]
 gluonts = [ "gluonts", "pandas",]
 hfdataset = [ "datasets",]
 test = [ "fev>=0.6.0", "pytest",]
-all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0",]
+classification = [ "scikit-learn",]
+all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0", "scikit-learn",]
 
 [tool.docformatter]
 diff = false

tirex_mirror-2025.11.28/src/tirex/models/classification/__init__.py

@@ -0,0 +1,8 @@
+# classification/__init__.py
+from .linear_classifier import TirexClassifierTorch
+from .rf_classifier import TirexRFClassifier
+
+__all__ = [
+    "TirexClassifierTorch",
+    "TirexRFClassifier",
+]

tirex_mirror-2025.11.28/src/tirex/models/classification/embedding.py

@@ -0,0 +1,125 @@
+import inspect
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from tirex import load_model
+
+from .utils import nanmax, nanmin, nanstd
+
+
+class TiRexEmbedding(nn.Module):
+    def __init__(self, device: str | None = None, data_augmentation: bool = False, batch_size: int = 512) -> None:
+        super().__init__()
+        self.data_augmentation = data_augmentation
+        self.number_of_patches = 8
+        self.batch_size = batch_size
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+        self.model = load_model(path="NX-AI/TiRex", device=self.device)
+
+    def _gen_emb_batched(self, data: torch.Tensor) -> torch.Tensor:
+        batches = list(torch.split(data, self.batch_size))
+        embedding_list = []
+        for batch in batches:
+            embedding = self.model._embed_context(batch)
+            embedding_list.append(embedding.cpu())
+        return torch.cat(embedding_list, dim=0)
+
+    def _calculate_n_patches(self, data: torch.Tensor) -> int:
+        _, _, n_steps = data.shape
+        n_patches = -(-n_steps // self.model.config.input_patch_size)
+        return n_patches
+
+    def forward(self, data: torch.Tensor) -> torch.Tensor:
+        n_patches = self._calculate_n_patches(data)
+
+        embedding = torch.stack(
+            [self._gen_emb_batched(var_slice) for var_slice in torch.unbind(data, dim=1)], dim=1
+        )  # Stack in case of multivar
+        embedding = self.process_embedding(embedding, n_patches)
+
+        if self.data_augmentation:
+            # Difference Embedding
+            diff_data = torch.diff(data, dim=-1, prepend=data[..., :0])
+            n_patches = self._calculate_n_patches(diff_data)
+
+            diff_embedding = torch.stack(
+                [self._gen_emb_batched(var_slice) for var_slice in torch.unbind(diff_data, dim=1)], dim=1
+            )
+            diff_embedding = self.process_embedding(diff_embedding, n_patches)
+            embedding = torch.cat((diff_embedding, embedding), dim=-1)
+
+            # Stats Embedding
+            stat_features = self._generate_stats_features(data)
+            normalized_stats = self._normalize_stats(stat_features)
+            normalized_stats = normalized_stats.to(embedding.device)
+
+            # Concat all together
+            embedding = torch.cat((embedding, normalized_stats), dim=-1)
+
+        return embedding
+
+    def process_embedding(self, embedding: torch.Tensor, n_patches: int) -> torch.Tensor:
+        # embedding shape: (bs, var_dim, n_patches, n_layer, emb_dim)
+        embedding = embedding[:, :, -n_patches:, :, :]
+        embedding = torch.mean(embedding, dim=2)  # sequence
+        embedding = torch.nn.functional.normalize(embedding, p=2, dim=-1)
+        embedding = torch.transpose(embedding, 1, -2).flatten(start_dim=-2)  # var
+        embedding = torch.transpose(embedding, 1, -2).flatten(start_dim=-2)  # layer
+        embedding = F.layer_norm(embedding, (embedding.shape[-1],))
+        return embedding
+
+    def _normalize_stats(self, stat_features: torch.Tensor) -> torch.Tensor:
+        dataset_mean = torch.nanmean(stat_features, dim=0, keepdim=True)
+        dataset_std = nanstd(stat_features, dim=0, keepdim=True)
+        stat_features = (stat_features - dataset_mean) / (dataset_std + 1e-8)
+        stat_features = torch.nan_to_num(stat_features, nan=0.0)
+
+        stat_features = (stat_features - stat_features.nanmean(dim=-1, keepdim=True)) / (
+            stat_features.std(dim=-1, keepdim=True) + 1e-8
+        )
+        return stat_features
+
+    def _generate_stats_features(self, data: torch.Tensor) -> torch.Tensor:
+        bs, variates, n_steps = data.shape
+
+        patch_size = max(1, n_steps // self.number_of_patches)
+        n_full_patches = n_steps // patch_size
+        n_remain = n_steps % patch_size
+
+        # [batch, variates, n_patches, patch_size]
+        patches = data[..., : n_full_patches * patch_size].unfold(-1, patch_size, patch_size)
+
+        # Stats for full patches
+        patch_means = torch.nanmean(patches, dim=-1)
+        patch_stds = nanstd(patches, dim=-1)
+        patch_maxes = nanmax(patches, dim=-1)
+        patch_mins = nanmin(patches, dim=-1)
+
+        stats = [patch_means, patch_stds, patch_maxes, patch_mins]
+
+        # Handle last smaller patch if needed
+        if n_remain > 0:
+            self._handle_remaining_patch(data, stats, n_full_patches * patch_size)
+
+        stats = torch.stack(stats, dim=-1)  # [batch, variates, n_patches(+1), 4]
+        return stats.flatten(start_dim=1)  # [batch, variates * n_patches * 4]
+
+    def _handle_remaining_patch(self, data: torch.Tensor, stats: list[torch.Tensor], full_patch_length: int) -> None:
+        last_patch = data[..., full_patch_length:]
+
+        mean_last = last_patch.mean(dim=-1, keepdim=True)
+        std_last = last_patch.std(dim=-1, keepdim=True)
+        max_last = last_patch.max(dim=-1, keepdim=True)
+        min_last = last_patch.min(dim=-1, keepdim=True)
+
+        stats[0] = torch.cat([stats[0], mean_last], dim=-1)
+        stats[1] = torch.cat([stats[1], std_last], dim=-1)
+        stats[2] = torch.cat([stats[2], max_last], dim=-1)
+        stats[3] = torch.cat([stats[3], min_last], dim=-1)

tirex_mirror-2025.11.28/src/tirex/models/classification/linear_classifier.py

@@ -0,0 +1,274 @@
+from dataclasses import asdict
+
+import torch
+
+from .embedding import TiRexEmbedding
+from .trainer import TrainConfig, Trainer
+
+
+class TirexClassifierTorch(torch.nn.Module):
+    """
+    A PyTorch classifier that combines time series embeddings with a linear classification 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 classification. The embedding backbone
+    is frozen during training, and only the classification head is trained.
+
+    Example:
+        >>> import torch
+        >>> from tirex.models.classification import TirexClassifierTorch
+        >>>
+        >>> # Create model with TIREX embeddings
+        >>> model = TirexClassifierTorch(
+        ...     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.randint(0, 3, (100,))  # 3 classes
+        >>>
+        >>> # 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)
+        >>> probabilities = model.predict_proba(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        # Training parameters
+        max_epochs: int = 50,
+        lr: float = 1e-4,
+        weight_decay: float = 0.01,
+        batch_size: int = 512,
+        val_split_ratio: float = 0.2,
+        stratify: bool = True,
+        patience: int = 7,
+        delta: float = 0.001,
+        log_every_n_steps: int = 5,
+        seed: int | None = None,
+        class_weights: torch.Tensor | None = None,
+        # Head parameters
+        dropout: float | None = None,
+    ) -> None:
+        """Initializes Embedding Based Linear Classification model.
+
+        Args:
+            data_augmentation : bool | None
+                Whether to use data_augmentation for embeddings (stats and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the model on. If None, uses CUDA if available, else CPU. Default: None
+            max_epochs : int
+                Maximum number of training epochs. Default: 50
+            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
+            stratify : bool
+                Whether to stratify the train/validation split by class labels. Default: True
+            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
+            class_weights: torch.Tensor | None
+                Weight classes according to given values, has to be a Tensor of size number of categories. Default: None
+            dropout : float | None
+                Dropout probability for the classification head. If None, no dropout is used. Default: None
+        """
+
+        super().__init__()
+
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+        # Create embedding model
+        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
+        self.data_augmentation = data_augmentation
+
+        # Head parameters
+        self.dropout = dropout
+        self.head = None
+        self.emb_dim = None
+        self.num_classes = 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=class_weights,
+            batch_size=batch_size,
+            val_split_ratio=val_split_ratio,
+            stratify=stratify,
+            patience=patience,
+            delta=delta,
+            seed=seed,
+        )
+        self.trainer = Trainer(self, train_config=train_config)
+
+    def _init_classifier(self, emb_dim: int, num_classes: int, dropout: float | None) -> torch.nn.Module:
+        if dropout:
+            return torch.nn.Sequential(torch.nn.Dropout(p=dropout), torch.nn.Linear(emb_dim, num_classes))
+        else:
+            return torch.nn.Linear(emb_dim, num_classes)
+
+    @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.num_classes = len(torch.unique(y))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the embedding model and classification head.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Logits for each class with shape (batch_size, num_classes).
+        Raises:
+            RuntimeError: If the classification 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
+    ) -> dict[str, float]:
+        """Train the classification head on the provided data.
+
+        This method initializes the classification 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 class labels.
+            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_classifier(self.emb_dim, self.num_classes, 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 class labels 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 class labels with shape (batch_size,).
+        """
+        self.eval()
+        x = x.to(self.device)
+        logits = self.forward(x)
+        return torch.argmax(logits, dim=1)
+
+    @torch.inference_mode()
+    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class probabilities for input time series data.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        self.eval()
+        x = x.to(self.device)
+        logits = self.forward(x)
+        return torch.softmax(logits, dim=1)
+
+    def save_model(self, path: str) -> None:
+        """Save the trained classification head.
+
+        This function saves the trained classification 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,
+                "emb_dim": self.emb_dim,
+                "num_classes": self.num_classes,
+                "dropout": self.dropout,
+                "train_config": train_config_dict,
+            },
+            path,
+        )
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexClassifierTorch":
+        """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:
+            TirexClassifierTorch: 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"],
+            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),
+            stratify=train_config_dict.get("stratify", True),
+            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),
+            class_weights=train_config_dict.get("class_weights", None),
+        )
+
+        # Initialize head with dimensions
+        model.emb_dim = checkpoint["emb_dim"]
+        model.num_classes = checkpoint["num_classes"]
+        model.head = model._init_classifier(model.emb_dim, model.num_classes, model.dropout)
+
+        # Load the trained weights
+        model.head.load_state_dict(checkpoint["head_state_dict"])
+        model.to(model.device)
+
+        return model

tirex_mirror-2025.11.28/src/tirex/models/classification/rf_classifier.py

@@ -0,0 +1,155 @@
+import joblib
+import numpy as np
+import torch
+from sklearn.ensemble import RandomForestClassifier
+
+from .embedding import TiRexEmbedding
+
+
+class TirexRFClassifier:
+    """
+    A Random Forest classifier that uses time series embeddings as features.
+
+    This classifier combines a pre-trained embedding model for feature extraction with a scikit-learn
+    Random Forest classifier. The embedding model generates fixed-size feature vectors from variable-length
+    time series, which are then used to train the Random Forest.
+
+    Example:
+        >>> import numpy as np
+        >>> from tirex.models.classification import TirexRFClassifier
+        >>>
+        >>> # Create model with custom Random Forest parameters
+        >>> model = TirexRFClassifier(
+        ...     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.randint(0, 3, (100,))  # 3 classes
+        >>>
+        >>> # Train the model
+        >>> model.fit((X_train, y_train))
+        >>>
+        >>> # Make predictions
+        >>> X_test = torch.randn(20, 1, 128)
+        >>> predictions = model.predict(X_test)
+        >>> probabilities = model.predict_proba(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        batch_size: int = 512,
+        # Random Forest parameters
+        **rf_kwargs,
+    ) -> None:
+        """Initializes Embedding Based Random Forest Classification model.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (stats 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
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+            **rf_kwargs
+                Additional keyword arguments to pass to sklearn's RandomForestClassifier.
+                Common options include n_estimators, max_depth, min_samples_split, random_state, etc.
+        """
+
+        # Set device
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
+        self.data_augmentation = data_augmentation
+
+        self.head = RandomForestClassifier(**rf_kwargs)
+
+    @torch.inference_mode()
+    def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
+        """Train the Random Forest classifier 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 numpy array
+                of class labels.
+        """
+        X_train, y_train = train_data
+
+        if isinstance(y_train, torch.Tensor):
+            y_train = y_train.detach().cpu().numpy()
+
+        embeddings = self.emb_model(X_train).cpu().numpy()
+        self.head.fit(embeddings, y_train)
+
+    @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 or np.ndarray with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted class labels with shape (batch_size,).
+        """
+
+        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 or np.ndarray with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        embeddings = self.emb_model(x).cpu().numpy()
+        return torch.from_numpy(self.head.predict_proba(embeddings))
+
+    def save_model(self, path: str) -> None:
+        """This method saves the trained Random Forest classifier 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,
+            "head": self.head,
+        }
+        joblib.dump(payload, path)
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexRFClassifier":
+        """Load a saved model from file.
+
+        This reconstructs the model with the embedding configuration and loads
+        the trained Random Forest classifier from a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexRFClassifier: The loaded model with trained Random Forest, ready for inference.
+        """
+        checkpoint = joblib.load(path)
+
+        # Create new instance with saved configuration
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+        )
+
+        # Load the trained Random Forest head
+        model.head = checkpoint["head"]
+
+        return model