Repository: WenjieDu/PyPOTS (language: python)

Issue:
Add SegRNN testing cases
<!-- 🚨Please create your PR to merge code from your branch to `dev` branch here, rather than `main`. -->

# What does this PR do?

<!--
Congrats! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution 😉.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, I will review your PR shortly. I may suggest changes to make the code even better 🤝.
-->

<!-- Remove if not applicable -->

1. fixing #508;


## Before submitting

<!-- You can remove checks that are not relevant to this PR. -->

- [x] Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have written necessary tests and already run them locally.



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
index 75372c89..8aadbf6f 100644
--- a/.github/workflows/stale.yml
+++ b/.github/workflows/stale.yml
@@ -9,11 +9,12 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/stale@v8
+      - uses: actions/stale@v9
         with:
           stale-issue-message: 'This issue had no activity for **14 days**. It will be closed in **1 week** unless there is some new activity. Is this issue already resolved?'
           stale-issue-label: 'stale'
-          exempt-issue-labels: bug, enhancement, new feature, new model, good first issue, keep
+          exempt-issue-labels: bug, enhancement, new feature, new algo, good first issue, keep
+          exempt-pr-labels: bug, enhancement, new feature, new algo, keep
           stale-pr-message: 'This pull request had no activity for **14 days**. It will be closed in **1 week** unless there is some new activity.'
           stale-pr-label: 'stale'
           days-before-stale: 14
diff --git a/README.md b/README.md
index fcfda433..597c02cd 100644
--- a/README.md
+++ b/README.md
@@ -136,7 +136,7 @@ The paper references and links are all listed at the bottom of this file.
 | Neural Net    | MICN🧑‍🔧[^27]                                                                                                                   |    ✅     |          |          |          |          | `2023 - ICLR`      |
 | Neural Net    | DLinear🧑‍🔧[^17]                                                                                                                |    ✅     |          |          |          |          | `2023 - AAAI`      |
 | Neural Net    | TiDE🧑‍🔧[^28]                                                                                                                   |    ✅     |          |          |          |          | `2023 - TMLR`      |
-| Neural Net    | CSAI[^42]                                                                                                                        |    ✅     |          |          |          |          | `2023 - arXiv`     |
+| Neural Net    | CSAI[^42]                                                                                                                        |    ✅     |          |    ✅     |          |          | `2023 - arXiv`     |
 | Neural Net    | SegRNN🧑‍🔧[^43]                                                                                                                 |    ✅     |          |          |          |          | `2023 - arXiv`     |
 | Neural Net    | SCINet🧑‍🔧[^30]                                                                                                                 |    ✅     |          |          |          |          | `2022 - NeurIPS`   |
 | Neural Net    | Nonstationary Tr.🧑‍🔧[^25]                                                                                                      |    ✅     |          |          |          |          | `2022 - NeurIPS`   |
diff --git a/README_zh.md b/README_zh.md
index d85af8ba..dc64581c 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -121,7 +121,7 @@ PyPOTS当前支持多变量POTS数据的插补, 预测, 分类, 聚类以及异
 | Neural Net    | MICN🧑‍🔧[^27]                                                                                                                   |   ✅    |        |        |        |          | `2023 - ICLR`      |
 | Neural Net    | DLinear🧑‍🔧[^17]                                                                                                                |   ✅    |        |        |        |          | `2023 - AAAI`      |
 | Neural Net    | TiDE🧑‍🔧[^28]                                                                                                                   |   ✅    |        |        |        |          | `2023 - TMLR`      |
-| Neural Net    | CSAI[^42]                                                                                                                        |    ✅     |          |          |          |          | `2023 - arXiv`     |
+| Neural Net    | CSAI[^42]                                                                                                                        |    ✅     |          |    ✅     |          |          | `2023 - arXiv`     |
 | Neural Net    | SegRNN🧑‍🔧[^43]                                                                                                                 |    ✅     |          |          |          |          | `2023 - arXiv`     |
 | Neural Net    | SCINet🧑‍🔧[^30]                                                                                                                 |   ✅    |        |        |        |          | `2022 - NeurIPS`   |
 | Neural Net    | Nonstationary Tr.🧑‍🔧[^25]                                                                                                      |   ✅    |        |        |        |          | `2022 - NeurIPS`   |
diff --git a/docs/index.rst b/docs/index.rst
index fd3b9219..76aee376 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -165,7 +165,7 @@ The paper references are all listed at the bottom of this readme file.
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | TiDE🧑‍🔧 :cite:`das2023tide`                             |  ✅  |      |      |      |      | ``2023 - TMLR``       |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
-| Neural Net     | CSAI :cite:`qian2023csai`                                 |  ✅  |      |      |      |      | ``2023 - arXiv``      |
+| Neural Net     | CSAI :cite:`qian2023csai`                                 |  ✅  |      |  ✅  |      |      | ``2023 - arXiv``      |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | SegRNN🧑‍🔧 :cite:`lin2023segrnn`                         |  ✅  |      |      |      |      | ``2023 - arXiv``      |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
diff --git a/docs/pypots.classification.rst b/docs/pypots.classification.rst
index 48b887c9..91acec09 100644
--- a/docs/pypots.classification.rst
+++ b/docs/pypots.classification.rst
@@ -10,6 +10,15 @@ pypots.classification.brits
    :show-inheritance:
    :inherited-members:
 
+pypots.classification.csai
+----------------------------------
+
+.. automodule:: pypots.classification.csai
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
 pypots.classification.grud
 ---------------------------------
 


Test patch:
diff --git a/tests/imputation/segrnn.py b/tests/imputation/segrnn.py
new file mode 100644
index 00000000..17b877ad
--- /dev/null
+++ b/tests/imputation/segrnn.py
@@ -0,0 +1,117 @@
+"""
+Test cases for SegRNN imputation model.
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+
+import os.path
+import unittest
+
+import numpy as np
+import pytest
+
+from pypots.imputation import SegRNN
+from pypots.optim import Adam
+from pypots.utils.logging import logger
+from pypots.utils.metrics import calc_mse
+from tests.global_test_config import (
+    DATA,
+    EPOCHS,
+    DEVICE,
+    TRAIN_SET,
+    VAL_SET,
+    TEST_SET,
+    GENERAL_H5_TRAIN_SET_PATH,
+    GENERAL_H5_VAL_SET_PATH,
+    GENERAL_H5_TEST_SET_PATH,
+    RESULT_SAVING_DIR_FOR_IMPUTATION,
+    check_tb_and_model_checkpoints_existence,
+)
+
+
+class TestSegRNN(unittest.TestCase):
+    logger.info("Running tests for an imputation model SegRNN...")
+
+    # set the log and model saving path
+    saving_path = os.path.join(RESULT_SAVING_DIR_FOR_IMPUTATION, "SegRNN")
+    model_save_name = "saved_segrnn_model.pypots"
+
+    # initialize an Adam optimizer
+    optimizer = Adam(lr=0.001, weight_decay=1e-5)
+
+    # initialize a SegRNN model
+    segrnn = SegRNN(
+        DATA["n_steps"],
+        DATA["n_features"],
+        seg_len=4,
+        dropout=0,
+        epochs=EPOCHS,
+        saving_path=saving_path,
+        optimizer=optimizer,
+        device=DEVICE,
+    )
+
+    @pytest.mark.xdist_group(name="imputation-segrnn")
+    def test_0_fit(self):
+        self.segrnn.fit(TRAIN_SET, VAL_SET)
+
+    @pytest.mark.xdist_group(name="imputation-segrnn")
+    def test_1_impute(self):
+        imputation_results = self.segrnn.predict(TEST_SET)
+        assert not np.isnan(
+            imputation_results["imputation"]
+        ).any(), "Output still has missing values after running impute()."
+
+        test_MSE = calc_mse(
+            imputation_results["imputation"],
+            DATA["test_X_ori"],
+            DATA["test_X_indicating_mask"],
+        )
+        logger.info(f"SegRNN test_MSE: {test_MSE}")
+
+    @pytest.mark.xdist_group(name="imputation-segrnn")
+    def test_2_parameters(self):
+        assert hasattr(self.segrnn, "model") and self.segrnn.model is not None
+
+        assert hasattr(self.segrnn, "optimizer") and self.segrnn.optimizer is not None
+
+        assert hasattr(self.segrnn, "best_loss")
+        self.assertNotEqual(self.segrnn.best_loss, float("inf"))
+
+        assert hasattr(self.segrnn, "best_model_dict") and self.segrnn.best_model_dict is not None
+
+    @pytest.mark.xdist_group(name="imputation-segrnn")
+    def test_3_saving_path(self):
+        # whether the root saving dir exists, which should be created by save_log_into_tb_file
+        assert os.path.exists(self.saving_path), f"file {self.saving_path} does not exist"
+
+        # check if the tensorboard file and model checkpoints exist
+        check_tb_and_model_checkpoints_existence(self.segrnn)
+
+        # save the trained model into file, and check if the path exists
+        saved_model_path = os.path.join(self.saving_path, self.model_save_name)
+        self.segrnn.save(saved_model_path)
+
+        # test loading the saved model, not necessary, but need to test
+        self.segrnn.load(saved_model_path)
+
+    @pytest.mark.xdist_group(name="imputation-segrnn")
+    def test_4_lazy_loading(self):
+        self.segrnn.fit(GENERAL_H5_TRAIN_SET_PATH, GENERAL_H5_VAL_SET_PATH)
+        imputation_results = self.segrnn.predict(GENERAL_H5_TEST_SET_PATH)
+        assert not np.isnan(
+            imputation_results["imputation"]
+        ).any(), "Output still has missing values after running impute()."
+
+        test_MSE = calc_mse(
+            imputation_results["imputation"],
+            DATA["test_X_ori"],
+            DATA["test_X_indicating_mask"],
+        )
+        logger.info(f"Lazy-loading SegRNN test_MSE: {test_MSE}")
+
+
+if __name__ == "__main__":
+    unittest.main()