Update the docs for TimeMixer

docs/examples.rst

@@ -29,15 +29,13 @@ You can also find a simple and quick-start tutorial notebook on Google Colab
 
     # Data preprocessing. Tedious, but PyPOTS can help. 🤓
     data = load_specific_dataset('physionet_2012')  # PyPOTS will automatically download and extract it.
-    X = data['X']
-    num_samples = len(X['RecordID'].unique())
-    X = X.drop(['RecordID', 'Time'], axis = 1)
-    X = StandardScaler().fit_transform(X.to_numpy())
-    X = X.reshape(num_samples, 48, -1)
+    X = data['train_X']
+    num_samples = len(X)
+    X = StandardScaler().fit_transform(X.reshape(-1, X.shape[-1])).reshape(X.shape)
     X_ori = X  # keep X_ori for validation
     X = mcar(X, 0.1)  # randomly hold out 10% observed values as ground truth
     dataset = {"X": X}  # X for model input
-    print(X.shape)  # (11988, 48, 37), 11988 samples, 48 time steps, 37 features
+    print(X.shape)  # (7671, 48, 37), 7671 samples, 48 time steps, 37 features
 
     # initialize the model
     saits = SAITS(
@@ -55,7 +53,7 @@ You can also find a simple and quick-start tutorial notebook on Google Colab
         model_saving_strategy="best", # only save the model with the best validation performance
     )
 
-    # train the model. Here I use the whole dataset as the training set, because ground truth is not visible to the model.
+    # train the model. Here I consider the train dataset only, and evaluate on it, because ground truth is not visible to the model.
     saits.fit(dataset)
     # impute the originally-missing values and artificially-missing values
     imputation = saits.impute(dataset)
@@ -64,6 +62,6 @@ You can also find a simple and quick-start tutorial notebook on Google Colab
     mae = calc_mae(imputation, np.nan_to_num(X_ori), indicating_mask)  # calculate mean absolute error on the ground truth (artificially-missing values)
 
     # the best model has been already saved, but you can still manually save it with function save_model() as below
-    saits.save_model(saving_dir="examples/saits",file_name="manually_saved_saits_model")
+    saits.save(saving_path="examples/saits/manually_saved_saits_model")
     # you can load the saved model into a new initialized model
-    saits.load_model("examples/saits/manually_saved_saits_model")
+    saits.load("examples/saits/manually_saved_saits_model.pypots")

docs/index.rst

@@ -133,10 +133,12 @@ The paper references are all listed at the bottom of this readme file.
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Type           | Algorithm                                                 | IMPU | FORE | CLAS | CLUS | ANOD | Year - Venue          |
 +================+===========================================================+======+======+======+======+======+=======================+
-| Neural Net     | ImputeFormer      :cite:`nie2024imputeformer`             |  ✅  |      |      |      |      | ``2024 - KDD``        |
+| Neural Net     | TimeMixer      :cite:`wang2024timemixer`                  |  ✅  |      |      |      |      | ``2024 - ICLR``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | iTransformer🧑‍🔧 :cite:`liu2024itransformer`             |  ✅  |      |      |      |      | ``2024 - ICLR``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
+| Neural Net     | ImputeFormer      :cite:`nie2024imputeformer`             |  ✅  |      |      |      |      | ``2024 - KDD``        |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | SAITS :cite:`du2023SAITS`                                 |  ✅  |      |      |      |      | ``2023 - ESWA``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | FreTS🧑‍🔧 :cite:`yi2023frets`                            |  ✅  |      |      |      |      | ``2023 - NeurIPS``    |
@@ -333,7 +335,7 @@ By committing your code, you'll
    `pypots/imputation/template <https://github.com/WenjieDu/PyPOTS/tree/main/pypots/imputation/template>`_) to quickly start;
 2. become one of `PyPOTS contributors <https://github.com/WenjieDu/PyPOTS/graphs/contributors>`_ and
    be listed as a volunteer developer `on the PyPOTS website <https://pypots.com/about/#volunteer-developers>`_;
-3. get mentioned in our `release notes <https://github.com/WenjieDu/PyPOTS/releases>`_;
+3. get mentioned in PyPOTS `release notes <https://github.com/WenjieDu/PyPOTS/releases>`_;
 
 You can also contribute to PyPOTS by simply staring🌟 this repo to help more people notice it.
 Your star is your recognition to PyPOTS, and it matters!

docs/pypots.imputation.rst

@@ -19,6 +19,15 @@ pypots.imputation.transformer
    :show-inheritance:
    :inherited-members:
 
+pypots.imputation.timemixer
+------------------------------------
+
+.. automodule:: pypots.imputation.timemixer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
 pypots.imputation.imputeformer
 ------------------------------------
 

pypots/data/load_specific_datasets.py

@@ -35,7 +35,7 @@ def list_supported_datasets() -> list:
 
 def load_specific_dataset(dataset_name: str, use_cache: bool = True) -> dict:
     """Load specific datasets supported by PyPOTS.
-    Different from tsdb.load_dataset(), which only produces merely raw data,
+    Different from tsdb.load(), which only produces merely raw data,
     load_specific_dataset here does some preprocessing operations,
     like truncating time series to generate samples with the same length.
 
@@ -45,7 +45,7 @@ def load_specific_dataset(dataset_name: str, use_cache: bool = True) -> dict:
         The name of the dataset to be loaded, which should be supported, i.e. in SUPPORTED_DATASETS.
 
     use_cache :
-        Whether to use cache. This is an argument of tsdb.load_dataset().
+        Whether to use cache. This is an argument of tsdb.load().
 
     Returns
     -------

tests/imputation/timemixer.py

@@ -47,8 +47,8 @@ class TestTimeMixer(unittest.TestCase):
         DATA["n_features"],
         n_layers=2,
         top_k=5,
-        d_model=512,
-        d_ffn=512,
+        d_model=32,
+        d_ffn=32,
         dropout=0.1,
         epochs=EPOCHS,
         saving_path=saving_path,