FEATURE: add subsampling

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tsururu"
-version = "1.1.0"
+version = "1.1.1"
 description = "Python tool for time series forecasting"
 readme = "README.md"
 requires-python = ">=3.9,<3.14"

tsururu/models/torch_based/layers/convolution.py

@@ -4,9 +4,10 @@
 
 torch = OptionalImport("torch")
 nn = OptionalImport("torch.nn")
+Module = OptionalImport("torch.nn.Module")
 
 
-class Inception_Block_V1(nn.Module):
+class Inception_Block_V1(Module):
     """Inception Block Version 1.
 
     Args:
@@ -43,7 +44,7 @@ def _initialize_weights(self):
                 if m.bias is not None:
                     nn.init.constant_(m.bias, 0)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the Inception block.
 
         Args:

tsururu/models/torch_based/layers/embedding.py

@@ -7,9 +7,10 @@
 
 torch = OptionalImport("torch")
 nn = OptionalImport("torch.nn")
+Module = OptionalImport("torch.nn.Module")
 
 
-class TokenEmbedding(nn.Module):
+class TokenEmbedding(Module):
     """Token embedding layer using 1D convolution.
 
     Args:
@@ -33,7 +34,7 @@ def __init__(self, c_in: int, d_model: int):
             if isinstance(m, nn.Conv1d):
                 nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="leaky_relu")
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the token embedding.
 
         Args:
@@ -47,7 +48,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
-class PositionalEmbedding(nn.Module):
+class PositionalEmbedding(Module):
     """Positional encoding using sine and cosine functions.
 
     Args:
@@ -71,7 +72,7 @@ def __init__(self, d_model: int, max_len: int = 5000):
         pe = pe.unsqueeze(0)
         self.register_buffer("pe", pe)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the positional embedding.
 
         Args:
@@ -84,7 +85,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.pe[:, : x.size(1)]
 
 
-class FixedEmbedding(nn.Module):
+class FixedEmbedding(Module):
     """Fixed embedding layer using precomputed sine and cosine values.
 
     Args:
@@ -108,7 +109,7 @@ def __init__(self, c_in: int, d_model: int):
         self.emb = nn.Embedding(c_in, d_model)
         self.emb.weight = nn.Parameter(w, requires_grad=False)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the fixed embedding.
 
         Args:
@@ -121,7 +122,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.emb(x).detach()
 
 
-class TemporalEmbedding(nn.Module):
+class TemporalEmbedding(Module):
     """Temporal embedding layer for time-related features.
 
     Args:
@@ -148,7 +149,7 @@ def __init__(self, d_model: int, embed_type: str = "fixed", freq: str = "h"):
         self.day_embed = Embed(day_size, d_model)
         self.month_embed = Embed(month_size, d_model)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the temporal embedding.
 
         Args:
@@ -168,7 +169,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return hour_x + weekday_x + day_x + month_x + minute_x
 
 
-class TimeFeatureEmbedding(nn.Module):
+class TimeFeatureEmbedding(Module):
     """Time feature embedding layer using linear transformation.
 
     Args:
@@ -181,7 +182,7 @@ def __init__(self, d_model: int, d_inp: int = 1):
         super(TimeFeatureEmbedding, self).__init__()
         self.embed = nn.Linear(d_inp, d_model, bias=False)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the time feature embedding.
 
         Args:
@@ -194,7 +195,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.embed(x)
 
 
-class Embedding(nn.Module):
+class Embedding(Module):
     """Data embedding layer combining token, positional, and temporal embeddings.
 
     Args:
@@ -242,7 +243,7 @@ def __init__(
 
         self.dropout = nn.Dropout(p=dropout)
 
-    def forward(self, x: torch.Tensor, x_mark: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor", x_mark: Optional["torch.Tensor"] = None) -> "torch.Tensor":
         """Forward pass of the data embedding.
 
         Args:

tsururu/models/torch_based/layers/patch_tst.py

@@ -12,7 +12,6 @@
 torch = OptionalImport("torch")
 nn = OptionalImport("torch.nn")
 F = OptionalImport("torch.nn.functional")
-Tensor = OptionalImport("torch.Tensor")
 Module = OptionalImport("torch.nn.Module")
 
 
@@ -82,7 +81,7 @@ def __init__(
         act: str = "gelu",
         key_padding_mask: Union[bool, str] = "auto",
         padding_var: Optional[int] = None,
-        attn_mask: Optional[Tensor] = None,
+        attn_mask: Optional["torch.Tensor"] = None,
         res_attention: bool = True,
         pre_norm: bool = False,
         store_attn: bool = False,
@@ -166,7 +165,7 @@ def __init__(
                 head_dropout=head_dropout,
             )
 
-    def forward(self, z: Tensor) -> Tensor:
+    def forward(self, z: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the PatchTST backbone.
 
         Args:
@@ -330,7 +329,7 @@ def __init__(
         store_attn: bool = False,
         key_padding_mask: Union[bool, str] = "auto",
         padding_var: Optional[int] = None,
-        attn_mask: Optional[Tensor] = None,
+        attn_mask: Optional["torch.Tensor"] = None,
         res_attention: bool = True,
         pre_norm: bool = False,
         pe: str = "zeros",
@@ -383,7 +382,7 @@ def __init__(
             store_attn=store_attn,
         )
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass of the TSTi encoder.
 
         Args:
@@ -423,7 +422,7 @@ def forward(self, x: Tensor) -> Tensor:
         return x
 
     @staticmethod
-    def _reorganize_tensor(x: Tensor, n_vars: int) -> Tensor:
+    def _reorganize_tensor(x: "torch.Tensor", n_vars: int) -> "torch.Tensor":
         """Reorganizes input tensor to pair time series channels with exogenous features.
 
         Args:
@@ -520,10 +519,10 @@ def __init__(
 
     def forward(
         self,
-        src: Tensor,
-        key_padding_mask: Optional[Tensor] = None,
-        attn_mask: Optional[Tensor] = None,
-    ) -> Tensor:
+        src: "torch.Tensor",
+        key_padding_mask: Optional["torch.Tensor"] = None,
+        attn_mask: Optional["torch.Tensor"] = None,
+    ) -> "torch.Tensor":
         """Forward pass of the TST encoder.
 
         Args:
@@ -639,11 +638,11 @@ def __init__(
 
     def forward(
         self,
-        src: Tensor,
-        prev: Optional[Tensor] = None,
-        key_padding_mask: Optional[Tensor] = None,
-        attn_mask: Optional[Tensor] = None,
-    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        src: "torch.Tensor",
+        prev: Optional["torch.Tensor"] = None,
+        key_padding_mask: Optional["torch.Tensor"] = None,
+        attn_mask: Optional["torch.Tensor"] = None,
+    ) -> Union["torch.Tensor", Tuple["torch.Tensor", "torch.Tensor"]]:
         """Forward pass of the TST encoder layer.
 
         Args:
@@ -747,13 +746,13 @@ def __init__(
 
     def forward(
         self,
-        Q: Tensor,
-        K: Optional[Tensor] = None,
-        V: Optional[Tensor] = None,
-        prev: Optional[Tensor] = None,
-        key_padding_mask: Optional[Tensor] = None,
-        attn_mask: Optional[Tensor] = None,
-    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        Q: "torch.Tensor",
+        K: Optional["torch.Tensor"] = None,
+        V: Optional["torch.Tensor"] = None,
+        prev: Optional["torch.Tensor"] = None,
+        key_padding_mask: Optional["torch.Tensor"] = None,
+        attn_mask: Optional["torch.Tensor"] = None,
+    ) -> Union["torch.Tensor", Tuple["torch.Tensor", "torch.Tensor"]]:
         """Forward pass of the multi-head attention layer.
 
         Args:
@@ -845,13 +844,13 @@ def __init__(
 
     def forward(
         self,
-        q: Tensor,
-        k: Tensor,
-        v: Tensor,
-        prev: Optional[Tensor] = None,
-        key_padding_mask: Optional[Tensor] = None,
-        attn_mask: Optional[Tensor] = None,
-    ) -> Union[Tensor, Tuple[Tensor, Tensor, Tensor]]:
+        q: "torch.Tensor",
+        k: "torch.Tensor",
+        v: "torch.Tensor",
+        prev: Optional["torch.Tensor"] = None,
+        key_padding_mask: Optional["torch.Tensor"] = None,
+        attn_mask: Optional["torch.Tensor"] = None,
+    ) -> Union["torch.Tensor", Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]]:
         """Forward pass of the scaled dot-product attention.
 
         Args:

tsururu/models/torch_based/layers/positional_encoding.py

@@ -78,9 +78,7 @@ def Coord2dPosEncoding(
     return cpe
 
 
-def Coord1dPosEncoding(
-    q_len: int, exponential: bool = False, normalize: bool = True
-) -> "torch.Tensor":
+def Coord1dPosEncoding(q_len: int, exponential: bool = False, normalize: bool = True) -> "torch.Tensor":
     """Generate 1D coordinate positional encoding.
 
     Args:

tsururu/models/torch_based/times_net.py

@@ -14,9 +14,10 @@
 nn = OptionalImport("torch.nn")
 F = OptionalImport("torch.nn.functional")
 rearrange = OptionalImport("einops.rearrange")
+Module = OptionalImport("torch.nn.Module")
 
 
-def FFT_for_Period(x: torch.Tensor, k: int = 2) -> Tuple[np.ndarray, torch.Tensor]:
+def FFT_for_Period(x: "torch.Tensor", k: int = 2) -> Tuple[np.ndarray, "torch.Tensor"]:
     """Compute the FFT for the input tensor and find the top-k periods.
 
     Args:
@@ -41,7 +42,7 @@ def FFT_for_Period(x: torch.Tensor, k: int = 2) -> Tuple[np.ndarray, torch.Tenso
     return period, abs(xf).mean(-1)[:, top_list]
 
 
-class TimesBlock(nn.Module):
+class TimesBlock(Module):
     """TimesBlock module for time series forecasting.
 
     Args:
@@ -69,7 +70,7 @@ def __init__(
             Inception_Block_V1(d_ff, d_model, num_kernels=num_kernels),
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: "torch.Tensor") -> "torch.Tensor":
         """Forward pass for the TimesBlock module.
 
         Args:

tsururu/models/torch_based/utils.py

@@ -33,10 +33,10 @@ def adjust_features_groups(features_groups: Dict[str, int], num_lags: int) -> Di
 
 
 def slice_features(
-    X: torch.Tensor,
+    X: "torch.Tensor",
     feature_list: List[str],
     features_groups_corrected,
-) -> torch.Tensor:
+) -> "torch.Tensor":
     """Slice the input tensor X based on the corrected feature groups.
 
     Args:
@@ -88,11 +88,11 @@ def slice_features(
 
 
 def slice_features_4d(
-    X: torch.Tensor,
+    X: "torch.Tensor",
     features_list: List[str],
     features_groups_corrected,
     num_series,
-) -> torch.Tensor:
+) -> "torch.Tensor":
     """Slice the input tensor X based on the corrected feature groups and reshape it to 4D."""
     groups_order: List[str] = [
         "series",

tsururu/strategies/direct.py

@@ -1,6 +1,8 @@
 from copy import deepcopy
 from typing import Union
 
+import numpy as np
+
 from tsururu.dataset.dataset import TSDataset
 from tsururu.dataset.pipeline import Pipeline
 from tsururu.dataset.slice import IndexSlicer
@@ -56,6 +58,8 @@ def __init__(
     def fit(
         self,
         dataset: TSDataset,
+        subsampling_rate: float = 1.0,
+        subsampling_seed: int = 42,
     ) -> "DirectStrategy":
         """Fits the direct strategy to the given dataset.
 
@@ -92,6 +96,15 @@ def fit(
                 delta=dataset.delta,
             )
 
+            if subsampling_rate < 1.0:
+                all_idx = np.arange(features_idx.shape[0])
+                np.random.seed(subsampling_seed)
+                sampled_idx = np.random.choice(
+                    all_idx, size=int(subsampling_rate * len(all_idx)), replace=False
+                )
+                features_idx = features_idx[sampled_idx]
+                target_idx = target_idx[sampled_idx]
+
             data = self.pipeline.create_data_dict_for_pipeline(dataset, features_idx, target_idx)
             data = self.pipeline.fit_transform(data, self.strategy_name)