diff --git a/README.md b/README.md
index 58ff8bfd..ab8b860c 100644
--- a/README.md
+++ b/README.md
@@ -124,7 +124,7 @@ The paper references and links are all listed at the bottom of this file.
| **Type** | **Algo** | **IMPU** | **FORE** | **CLAS** | **CLUS** | **ANOD** | **Year - Venue** |
|:--------------|:---------------------------------------------------------------------------------------------------------------------------------|:--------:|:--------:|:--------:|:--------:|:--------:|:-------------------|
| LLM | 
Time-Series.AI [^36] | ✅ | ✅ | ✅ | ✅ | ✅ | `Later in 2024` |
-| Neural Net | TEFN[^39] | ✅ | | | | | `2024 - arXiv` |
+| Neural Net | TEFN🧑🔧[^39] | ✅ | | | | | `2024 - arXiv` |
| Neural Net | TimeMixer[^37] | ✅ | | | | | `2024 - ICLR` |
| Neural Net | iTransformer🧑🔧[^24] | ✅ | | | | | `2024 - ICLR` |
| Neural Net | ModernTCN[^38] | ✅ | | | | | `2024 - ICLR` |
@@ -161,7 +161,7 @@ The paper references and links are all listed at the bottom of this file.
| Neural Net | GRU-D[^4] | ✅ | | ✅ | | | `2018 - Sci. Rep.` |
| Neural Net | TCN🧑🔧[^35] | ✅ | | | | | `2018 - arXiv` |
| Neural Net | Transformer🧑🔧[^2] | ✅ | | | | | `2017 - NeurIPS` |
-| Naive | Lerp | ✅ | | | | | |
+| Naive | Lerp[^40] | ✅ | | | | | |
| Naive | LOCF/NOCB | ✅ | | | | | |
| Naive | Mean | ✅ | | | | | |
| Naive | Median | ✅ | | | | | |
@@ -508,4 +508,4 @@ Time-Series.AI
[^39]: Zhan, T., He, Y., Li, Z., & Deng, Y. (
2024). [Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting](https://arxiv.org/abs/2405.06419).
*arXiv 2024*
-
+[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
\ No newline at end of file
diff --git a/README_zh.md b/README_zh.md
index d43b5e83..1812ae1e 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -105,7 +105,7 @@ PyPOTS当前支持多变量POTS数据的插补，预测，分类，聚类以及
| **类型** | **算法** | **插补** | **预测** | **分类** | **聚类** | **异常检测** | **年份 - 刊物** |
|:--------------|:---------------------------------------------------------------------------------------------------------------------------------|:------:|:------:|:------:|:------:|:--------:|:-------------------|
| LLM | Time-Series.AI [^36] | ✅ | ✅ | ✅ | ✅ | ✅ | `Later in 2024` |
-| Neural Net | TEFN[^39] | ✅ | | | | | `2024 - arXiv` |
+| Neural Net | TEFN🧑🔧[^39] | ✅ | | | | | `2024 - arXiv` |
| Neural Net | TimeMixer[^37] | ✅ | | | | | `2024 - ICLR` |
| Neural Net | iTransformer🧑🔧[^24] | ✅ | | | | | `2024 - ICLR` |
| Neural Net | ModernTCN[^38] | ✅ | | | | | `2024 - ICLR` |
@@ -142,7 +142,7 @@ PyPOTS当前支持多变量POTS数据的插补，预测，分类，聚类以及
| Neural Net | GRU-D[^4] | ✅ | | ✅ | | | `2018 - Sci. Rep.` |
| Neural Net | TCN🧑🔧[^35] | ✅ | | | | | `2018 - arXiv` |
| Neural Net | Transformer🧑🔧[^2] | ✅ | | | | | `2017 - NeurIPS` |
-| Naive | Lerp | ✅ | | | | | |
+| Naive | Lerp[^40] | ✅ | | | | | |
| Naive | LOCF/NOCB | ✅ | | | | | |
| Naive | Mean | ✅ | | | | | |
| Naive | Median | ✅ | | | | | |
@@ -468,3 +468,4 @@ Time-Series.AI
[^39]: Zhan, T., He, Y., Li, Z., & Deng, Y. (
2024). [Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting](https://arxiv.org/abs/2405.06419).
*arXiv 2024*
+[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
\ No newline at end of file
diff --git a/docs/index.rst b/docs/index.rst
index 113fbb65..c64c9196 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -133,6 +133,8 @@ The paper references are all listed at the bottom of this readme file.
+----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
| Type | Algorithm | IMPU | FORE | CLAS | CLUS | ANOD | Year - Venue |
+================+===========================================================+======+======+======+======+======+=======================+
+| Neural Net | TEFN🧑🔧 :cite:`zhan2024tefn` | ✅ | | | | | ``2024 - arXiv`` |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
| Neural Net | TimeMixer :cite:`wang2024timemixer` | ✅ | | | | | ``2024 - ICLR`` |
+----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
| Neural Net | iTransformer🧑🔧 :cite:`liu2024itransformer` | ✅ | | | | | ``2024 - ICLR`` |
diff --git a/pypots/imputation/__init__.py b/pypots/imputation/__init__.py
index 4ae883f8..19a7e2c6 100644
--- a/pypots/imputation/__init__.py
+++ b/pypots/imputation/__init__.py
@@ -84,5 +84,5 @@
"Mean",
"Median",
"Lerp",
- "TEFN"
+ "TEFN",
]
diff --git a/pypots/imputation/tefn/core.py b/pypots/imputation/tefn/core.py
index c6dc6d8b..f71927a6 100644
--- a/pypots/imputation/tefn/core.py
+++ b/pypots/imputation/tefn/core.py
@@ -14,11 +14,11 @@
class _TEFN(nn.Module):
def __init__(
- self,
- n_steps,
- n_features,
- n_fod,
- apply_nonstationary_norm,
+ self,
+ n_steps,
+ n_features,
+ n_fod,
+ apply_nonstationary_norm,
):
super().__init__()
diff --git a/pypots/imputation/tefn/data.py b/pypots/imputation/tefn/data.py
index 3239a6e2..e8aa19a2 100644
--- a/pypots/imputation/tefn/data.py
+++ b/pypots/imputation/tefn/data.py
@@ -14,11 +14,11 @@ class DatasetForTEFN(DatasetForSAITS):
"""Actually TEFN uses the same data strategy as SAITS, needs MIT for training."""
def __init__(
- self,
- data: Union[dict, str],
- return_X_ori: bool,
- return_y: bool,
- file_type: str = "hdf5",
- rate: float = 0.2,
+ self,
+ data: Union[dict, str],
+ return_X_ori: bool,
+ return_y: bool,
+ file_type: str = "hdf5",
+ rate: float = 0.2,
):
super().__init__(data, return_X_ori, return_y, file_type, rate)
diff --git a/pypots/imputation/tefn/model.py b/pypots/imputation/tefn/model.py
index c7912164..02024f28 100644
--- a/pypots/imputation/tefn/model.py
+++ b/pypots/imputation/tefn/model.py
@@ -23,83 +23,83 @@
class TEFN(BaseNNImputer):
"""The PyTorch implementation of the TEFN model.
- TEFN is originally proposed by Zhan et al. in :cite:`zhan2024tefn`.
-
- Parameters
- ----------
- n_steps :
- The number of time steps in the time-series data sample.
-
- n_features :
- The number of features in the time-series data sample.
-
- n_fod :
- The number of frame of discernment in the TEFN model.
-
- apply_nonstationary_norm :
- Whether to apply non-stationary normalization to the input data for TimesNet.
- Please refer to :cite:`liu2022nonstationary` for details about non-stationary normalization,
- which is not the idea of the original TimesNet paper. Hence, we make it optional
- and default not to use here.
-
- batch_size :
- The batch size for training and evaluating the model.
-
- epochs :
- The number of epochs for training the model.
-
- patience :
- The patience for the early-stopping mechanism. Given a positive integer, the training process will be
- stopped when the model does not perform better after that number of epochs.
- Leaving it default as None will disable the early-stopping.
-
- optimizer :
- The optimizer for model training.
- If not given, will use a default Adam optimizer.
-
- num_workers :
- The number of subprocesses to use for data loading.
- `0` means data loading will be in the main process, i.e. there won't be subprocesses.
-
- device :
- The device for the model to run on. It can be a string, a :class:`torch.device` object, or a list of them.
- If not given, will try to use CUDA devices first (will use the default CUDA device if there are multiple),
- then CPUs, considering CUDA and CPU are so far the main devices for people to train ML models.
- If given a list of devices, e.g. ['cuda:0', 'cuda:1'], or [torch.device('cuda:0'), torch.device('cuda:1')] ,
- the model will be parallely trained on the multiple devices (so far only support parallel training on CUDA
- devices). Other devices like Google TPU and Apple Silicon accelerator MPS may be added in the future.
-
- saving_path :
- The path for automatically saving model checkpoints and tensorboard files (i.e. loss values recorded during
- training into a tensorboard file). Will not save if not given.
-
- model_saving_strategy :
- The strategy to save model checkpoints. It has to be one of [None, "best", "better", "all"].
- No model will be saved when it is set as None.
- The "best" strategy will only automatically save the best model after the training finished.
- The "better" strategy will automatically save the model during training whenever the model performs
- better than in previous epochs.
- The "all" strategy will save every model after each epoch training.
-
- verbose :
- Whether to print out the training logs during the training process.
- """
+ TEFN is originally proposed by Zhan et al. in :cite:`zhan2024tefn`.
+
+ Parameters
+ ----------
+ n_steps :
+ The number of time steps in the time-series data sample.
+
+ n_features :
+ The number of features in the time-series data sample.
+
+ n_fod :
+ The number of frame of discernment in the TEFN model.
+
+ apply_nonstationary_norm :
+ Whether to apply non-stationary normalization to the input data for TimesNet.
+ Please refer to :cite:`liu2022nonstationary` for details about non-stationary normalization,
+ which is not the idea of the original TimesNet paper. Hence, we make it optional
+ and default not to use here.
+
+ batch_size :
+ The batch size for training and evaluating the model.
+
+ epochs :
+ The number of epochs for training the model.
+
+ patience :
+ The patience for the early-stopping mechanism. Given a positive integer, the training process will be
+ stopped when the model does not perform better after that number of epochs.
+ Leaving it default as None will disable the early-stopping.
+
+ optimizer :
+ The optimizer for model training.
+ If not given, will use a default Adam optimizer.
+
+ num_workers :
+ The number of subprocesses to use for data loading.
+ `0` means data loading will be in the main process, i.e. there won't be subprocesses.
+
+ device :
+ The device for the model to run on. It can be a string, a :class:`torch.device` object, or a list of them.
+ If not given, will try to use CUDA devices first (will use the default CUDA device if there are multiple),
+ then CPUs, considering CUDA and CPU are so far the main devices for people to train ML models.
+ If given a list of devices, e.g. ['cuda:0', 'cuda:1'], or [torch.device('cuda:0'), torch.device('cuda:1')] ,
+ the model will be parallely trained on the multiple devices (so far only support parallel training on CUDA
+ devices). Other devices like Google TPU and Apple Silicon accelerator MPS may be added in the future.
+
+ saving_path :
+ The path for automatically saving model checkpoints and tensorboard files (i.e. loss values recorded during
+ training into a tensorboard file). Will not save if not given.
+
+ model_saving_strategy :
+ The strategy to save model checkpoints. It has to be one of [None, "best", "better", "all"].
+ No model will be saved when it is set as None.
+ The "best" strategy will only automatically save the best model after the training finished.
+ The "better" strategy will automatically save the model during training whenever the model performs
+ better than in previous epochs.
+ The "all" strategy will save every model after each epoch training.
+
+ verbose :
+ Whether to print out the training logs during the training process.
+ """
def __init__(
- self,
- n_steps: int,
- n_features: int,
- n_fod: int = 2,
- apply_nonstationary_norm: bool = True,
- batch_size: int = 32,
- epochs: int = 100,
- patience: int = None,
- optimizer: Optional[Optimizer] = Adam(),
- num_workers: int = 0,
- device: Optional[Union[str, torch.device, list]] = None,
- saving_path: str = None,
- model_saving_strategy: Optional[str] = "best",
- verbose: bool = True,
+ self,
+ n_steps: int,
+ n_features: int,
+ n_fod: int = 2,
+ apply_nonstationary_norm: bool = True,
+ batch_size: int = 32,
+ epochs: int = 100,
+ patience: int = None,
+ optimizer: Optional[Optimizer] = Adam(),
+ num_workers: int = 0,
+ device: Optional[Union[str, torch.device, list]] = None,
+ saving_path: str = None,
+ model_saving_strategy: Optional[str] = "best",
+ verbose: bool = True,
):
super().__init__(
batch_size,
@@ -119,12 +119,7 @@ def __init__(
self.n_fod = n_fod
# set up the model
- self.model = _TEFN(
- n_steps,
- n_features,
- n_fod,
- self.apply_nonstationary_norm
- )
+ self.model = _TEFN(n_steps, n_features, n_fod, self.apply_nonstationary_norm)
self._send_model_to_given_device()
self._print_model_size()
@@ -164,10 +159,10 @@ def _assemble_input_for_testing(self, data: list) -> dict:
return inputs
def fit(
- self,
- train_set: Union[dict, str],
- val_set: Optional[Union[dict, str]] = None,
- file_type: str = "hdf5",
+ self,
+ train_set: Union[dict, str],
+ val_set: Optional[Union[dict, str]] = None,
+ file_type: str = "hdf5",
) -> None:
# Step 1: wrap the input data with classes Dataset and DataLoader
training_set = DatasetForTEFN(
@@ -202,9 +197,9 @@ def fit(
self._auto_save_model_if_necessary(confirm_saving=True)
def predict(
- self,
- test_set: Union[dict, str],
- file_type: str = "hdf5",
+ self,
+ test_set: Union[dict, str],
+ file_type: str = "hdf5",
) -> dict:
"""Make predictions for the input data with the trained model.
@@ -260,9 +255,9 @@ def predict(
return result_dict
def impute(
- self,
- test_set: Union[dict, str],
- file_type: str = "hdf5",
+ self,
+ test_set: Union[dict, str],
+ file_type: str = "hdf5",
) -> np.ndarray:
"""Impute missing values in the given data with the trained model.
diff --git a/pypots/nn/modules/tefn/backbone.py b/pypots/nn/modules/tefn/backbone.py
index 78708a62..22dd7c88 100644
--- a/pypots/nn/modules/tefn/backbone.py
+++ b/pypots/nn/modules/tefn/backbone.py
@@ -1,22 +1,18 @@
"""
"""
+
+# Created by Tianxiang Zhan
+# License: BSD-3-Clause
+
import torch
import torch.nn as nn
from .layers import EvidenceMachineKernel
-# Created by Tianxiang Zhan >
-# License: BSD-3-Clause
-
-
class BackboneTEFN(nn.Module):
-
- def __init__(self,
- n_steps,
- n_features,
- n_fod):
+ def __init__(self, n_steps, n_features, n_fod):
super().__init__()
self.n_steps = n_steps
@@ -28,5 +24,5 @@ def __init__(self,
def forward(self, X) -> torch.Tensor:
X = self.T_model(X.permute(0, 2, 1)).permute(0, 2, 1, 3) + self.C_model(X)
- X = torch.einsum('btcf->btc', X)
+ X = torch.einsum("btcf->btc", X)
return X
diff --git a/pypots/nn/modules/tefn/layers.py b/pypots/nn/modules/tefn/layers.py
index f0b8fa0f..5f2e5934 100644
--- a/pypots/nn/modules/tefn/layers.py
+++ b/pypots/nn/modules/tefn/layers.py
@@ -2,22 +2,21 @@
"""
+# Created by Tianxiang Zhan
+# License: BSD-3-Clause
+
import torch
import torch.nn as nn
-# Created by Tianxiang Zhan >
-# License: BSD-3-Clause
-
-
class EvidenceMachineKernel(nn.Module):
def __init__(self, C, F):
- super(EvidenceMachineKernel, self).__init__()
+ super().__init__()
self.C = C
- self.F = 2 ** F
+ self.F = 2**F
self.C_weight = nn.Parameter(torch.randn(self.C, self.F))
self.C_bias = nn.Parameter(torch.randn(self.C, self.F))
def forward(self, x):
- x = torch.einsum('btc,cf->btcf', x, self.C_weight) + self.C_bias
+ x = torch.einsum("btc,cf->btcf", x, self.C_weight) + self.C_bias
return x