docs: ✏️ update the loss function of STWave

baselines/STWave/METR-LA.py

@@ -12,10 +12,10 @@
 from easydict import EasyDict
 from basicts.runners import SimpleTimeSeriesForecastingRunner
 from basicts.data import TimeSeriesForecastingDataset
-from basicts.losses import stwave_masked_mae
 from basicts.utils import load_adj
 
 from .arch import STWave
+from .loss import stwave_masked_mae
 
 def laplacian(W):
     """Return the Laplacian of the weight matrix."""

baselines/STWave/PEMS-BAY.py

@@ -12,10 +12,10 @@
 from easydict import EasyDict
 from basicts.runners import SimpleTimeSeriesForecastingRunner
 from basicts.data import TimeSeriesForecastingDataset
-from basicts.losses import stwave_masked_mae
 from basicts.utils import load_adj
 
 from .arch import STWave
+from .loss import stwave_masked_mae
 
 def laplacian(W):
     """Return the Laplacian of the weight matrix."""

baselines/STWave/PEMS03.py

@@ -12,10 +12,10 @@
 from easydict import EasyDict
 from basicts.runners import SimpleTimeSeriesForecastingRunner
 from basicts.data import TimeSeriesForecastingDataset
-from basicts.losses import stwave_masked_mae
 from basicts.utils import load_adj
 
 from .arch import STWave
+from .loss import stwave_masked_mae
 
 def laplacian(W):
     """Return the Laplacian of the weight matrix."""

baselines/STWave/PEMS04.py

@@ -12,10 +12,10 @@
 from easydict import EasyDict
 from basicts.runners import SimpleTimeSeriesForecastingRunner
 from basicts.data import TimeSeriesForecastingDataset
-from basicts.losses import stwave_masked_mae
 from basicts.utils import load_adj
 
 from .arch import STWave
+from .loss import stwave_masked_mae
 
 def laplacian(W):
     """Return the Laplacian of the weight matrix."""

baselines/STWave/PEMS07.py

@@ -12,10 +12,10 @@
 from easydict import EasyDict
 from basicts.runners import SimpleTimeSeriesForecastingRunner
 from basicts.data import TimeSeriesForecastingDataset
-from basicts.losses import stwave_masked_mae
 from basicts.utils import load_adj
 
 from .arch import STWave
+from .loss import stwave_masked_mae
 
 def laplacian(W):
     """Return the Laplacian of the weight matrix."""

baselines/STWave/PEMS08.py

@@ -12,10 +12,10 @@
 from easydict import EasyDict
 from basicts.runners import SimpleTimeSeriesForecastingRunner
 from basicts.data import TimeSeriesForecastingDataset
-from basicts.losses import stwave_masked_mae
 from basicts.utils import load_adj
 
 from .arch import STWave
+from .loss import stwave_masked_mae
 
 def laplacian(W):
     """Return the Laplacian of the weight matrix."""

baselines/STWave/loss.py

@@ -0,0 +1,21 @@
+import torch
+import numpy as np
+
+from basicts.losses import masked_mae
+
+
+def stwave_masked_mae(preds: list, labels: torch.Tensor, null_val: float = np.nan) -> torch.Tensor:
+    """Masked mean absolute error.
+
+    Args:
+        preds (torch.Tensor): predicted values
+        labels (torch.Tensor): labels
+        null_val (float, optional): null value. Defaults to np.nan.
+
+    Returns:
+        torch.Tensor: masked mean absolute error
+    """
+    lloss = masked_mae(preds[...,1:2], preds[...,2:])
+    loss = masked_mae(preds[...,:1], labels)
+
+    return loss + lloss

basicts/losses/__init__.py

@@ -1,3 +1,3 @@
-from .losses import l1_loss, l2_loss, masked_mae, masked_mape, masked_rmse, masked_mse, stwave_masked_mae
+from .losses import l1_loss, l2_loss, masked_mae, masked_mape, masked_rmse, masked_mse
 
-__all__ = ["l1_loss", "l2_loss", "masked_mae", "masked_mape", "masked_rmse", "masked_mse", "stwave_masked_mae"]
+__all__ = ["l1_loss", "l2_loss", "masked_mae", "masked_mape", "masked_rmse", "masked_mse"]

basicts/losses/losses.py

@@ -44,22 +44,6 @@ def masked_mae(preds: torch.Tensor, labels: torch.Tensor, null_val: float = np.n
     loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
     return torch.mean(loss)
 
-def stwave_masked_mae(preds: list, labels: torch.Tensor, null_val: float = np.nan) -> torch.Tensor:
-    """Masked mean absolute error.
-
-    Args:
-        preds (torch.Tensor): predicted values
-        labels (torch.Tensor): labels
-        null_val (float, optional): null value. Defaults to np.nan.
-
-    Returns:
-        torch.Tensor: masked mean absolute error
-    """
-    lloss = masked_mae(preds[...,1:2], preds[...,2:])
-    loss = masked_mae(preds[...,:1], labels)
-
-    return loss + lloss
-
 
 def masked_mse(preds: torch.Tensor, labels: torch.Tensor, null_val: float = np.nan) -> torch.Tensor:
     """Masked mean squared error.