update scaler; support subset forecasting (#147)

basicts/__init__.py

@@ -1,6 +1,6 @@
 from .launcher import launch_training, launch_evaluation
 from .runners import BaseRunner
 
-__version__ = '0.4.1'
+__version__ = '0.4.2'
 
 __all__ = ['__version__', 'launch_training', 'launch_evaluation', 'BaseRunner']

basicts/runners/runner_zoo/simple_tsf_runner.py

@@ -17,16 +17,17 @@ def __init__(self, cfg: Dict):
         super().__init__(cfg)
         self.forward_features = cfg['MODEL'].get('FORWARD_FEATURES', None)
         self.target_features = cfg['MODEL'].get('TARGET_FEATURES', None)
+        self.target_time_series = cfg['MODEL'].get('TARGET_TIME_SERIES', None)
 
     def select_input_features(self, data: torch.Tensor) -> torch.Tensor:
         """
         Selects input features based on the forward features specified in the configuration.
 
         Args:
-            data (torch.Tensor): Input history data with shape [B, L, N, C].
+            data (torch.Tensor): Input history data with shape [B, L, N, C1].
 
         Returns:
-            torch.Tensor: Data with selected features.
+            torch.Tensor: Data with selected features with shape [B, L, N, C2].
         """
 
         if self.forward_features is not None:
@@ -38,15 +39,29 @@ def select_target_features(self, data: torch.Tensor) -> torch.Tensor:
         Selects target features based on the target features specified in the configuration.
 
         Args:
-            data (torch.Tensor): Model prediction data with arbitrary shape.
+            data (torch.Tensor): Model prediction data with shape [B, L, N, C1].
 
         Returns:
-            torch.Tensor: Data with selected target features and shape [B, L, N, C].
+            torch.Tensor: Data with selected target features and shape [B, L, N, C2].
         """
 
         data = data[:, :, :, self.target_features]
         return data
 
+    def select_target_time_series(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Select target time series based on the target time series specified in the configuration.
+
+        Args:
+            data (torch.Tensor): Model prediction data with shape [B, L, N1, C].
+
+        Returns:
+            torch.Tensor: Data with selected target time series and shape [B, L, N2, C].
+        """
+
+        data = data[:, :, self.target_time_series, :]
+        return data
+
     def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bool = True, **kwargs) -> Dict:
         """
         Performs the forward pass for training, validation, and testing. 
@@ -93,8 +108,12 @@ def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bo
         if 'target' not in model_return:
             model_return['target'] = self.select_target_features(future_data)
 
+        if self.target_time_series is not None:
+            model_return['target'] = self.select_target_time_series(model_return['target'])
+            model_return['prediction'] = self.select_target_time_series(model_return['prediction'])
+
         # Ensure the output shape is correct
-        assert list(model_return['prediction'].shape)[:3] == [batch_size, length, num_nodes], \
+        assert list(model_return['prediction'].shape)[:3] == [batch_size, length, num_nodes if self.target_time_series is None else len(self.target_time_series)], \
             "The shape of the output is incorrect. Ensure it matches [B, L, N, C]."
 
         return model_return

basicts/scaler/min_max_scaler.py

@@ -52,11 +52,12 @@ def __init__(self, dataset_name: str, train_ratio: float, norm_each_channel: boo
 
         # compute minimum and maximum values for normalization
         if norm_each_channel:
-            self.min = np.min(train_data, axis=0, keepdims=True)
-            self.max = np.max(train_data, axis=0, keepdims=True)
+            self.min = torch.tensor(np.min(train_data, axis=0, keepdims=True))
+            self.max = torch.tensor(np.max(train_data, axis=0, keepdims=True))
         else:
             self.min = np.min(train_data)
             self.max = np.max(train_data)
+        self.min, self.max = torch.tensor(self.min), torch.tensor(self.max)
 
     def transform(self, input_data: torch.Tensor) -> torch.Tensor:
         """
@@ -71,8 +72,10 @@ def transform(self, input_data: torch.Tensor) -> torch.Tensor:
         Returns:
             torch.Tensor: The normalized data with the same shape as the input.
         """
-
-        input_data[..., self.target_channel] = (input_data[..., self.target_channel] - self.min) / (self.max - self.min)
+
+        _min = self.min.to(input_data.device)
+        _max = self.max.to(input_data.device)
+        input_data[..., self.target_channel] = (input_data[..., self.target_channel] - _min) / (_max - _min)
         return input_data
 
     def inverse_transform(self, input_data: torch.Tensor) -> torch.Tensor:
@@ -90,5 +93,7 @@ def inverse_transform(self, input_data: torch.Tensor) -> torch.Tensor:
             torch.Tensor: The data transformed back to its original scale.
         """
 
-        input_data[..., self.target_channel] = input_data[..., self.target_channel] * (self.max - self.min) + self.min
+        _min = self.min.to(input_data.device)
+        _max = self.max.to(input_data.device)
+        input_data[..., self.target_channel] = input_data[..., self.target_channel] * (_max - _min) + _min
         return input_data

basicts/scaler/z_score_scaler.py

@@ -61,6 +61,7 @@ def __init__(self, dataset_name: str, train_ratio: float, norm_each_channel: boo
             self.std = np.std(train_data)
             if self.std == 0:
                 self.std = 1.0  # prevent division by zero by setting std to 1 where it's 0
+        self.mean, self.std = torch.tensor(self.mean), torch.tensor(self.std)
 
     def transform(self, input_data: torch.Tensor) -> torch.Tensor:
         """
@@ -76,7 +77,9 @@ def transform(self, input_data: torch.Tensor) -> torch.Tensor:
             torch.Tensor: The normalized data with the same shape as the input.
         """
 
-        input_data[..., self.target_channel] = (input_data[..., self.target_channel] - self.mean) / self.std
+        mean = self.mean.to(input_data.device)
+        std = self.std.to(input_data.device)
+        input_data[..., self.target_channel] = (input_data[..., self.target_channel] - mean) / std
         return input_data
 
     def inverse_transform(self, input_data: torch.Tensor) -> torch.Tensor:
@@ -93,10 +96,9 @@ def inverse_transform(self, input_data: torch.Tensor) -> torch.Tensor:
             torch.Tensor: The data transformed back to its original scale.
         """
 
-        if isinstance(self.mean, np.ndarray):
-            self.mean = torch.tensor(self.mean, device=input_data.device)
-            self.std = torch.tensor(self.std, device=input_data.device)
+        mean = self.mean.to(input_data.device)
+        std = self.std.to(input_data.device)
         # Clone the input data to prevent in-place modification (which is not allowed in PyTorch)
         input_data = input_data.clone()
-        input_data[..., self.target_channel] = input_data[..., self.target_channel] * self.std + self.mean
+        input_data[..., self.target_channel] = input_data[..., self.target_channel] * std + mean
         return input_data