Expose new models for tuning, add get_class_full_path(), and test visual funcs

pypots/cli/tuning.py

@@ -44,6 +44,9 @@
     BRITS,
     GRUD,
     Transformer,
+    TiDE,
+    Reformer,
+    RevIN_SCINet,
 )
 from ..optim import Adam
 from ..utils.logging import logger
@@ -58,34 +61,37 @@
     )
 
 NN_MODELS = {
-    # imputation models
-    "pypots.imputation.SAITS": SAITS,
-    "pypots.imputation.iTransformer": iTransformer,
-    "pypots.imputation.Transformer": Transformer,
-    "pypots.imputation.FreTS": FreTS,
-    "pypots.imputation.Koopa": Koopa,
+    # imputation models, sorted by the first letter of the model name
+    "pypots.imputation.Autoformer": Autoformer,
+    "pypots.imputation.BRITS": BRITS,
+    "pypots.imputation.CSDI": CSDI,
     "pypots.imputation.Crossformer": Crossformer,
-    "pypots.imputation.PatchTST": PatchTST,
+    "pypots.imputation.DLinear": DLinear,
     "pypots.imputation.ETSformer": ETSformer,
+    "pypots.imputation.FreTS": FreTS,
+    "pypots.imputation.FiLM": FiLM,
+    "pypots.imputation.GPVAE": GPVAE,
+    "pypots.imputation.GRUD": GRUD,
+    "pypots.imputation.Informer": Informer,
+    "pypots.imputation.iTransformer": iTransformer,
+    "pypots.imputation.Koopa": Koopa,
     "pypots.imputation.MICN": MICN,
-    "pypots.imputation.DLinear": DLinear,
-    "pypots.imputation.SCINet": SCINet,
+    "pypots.imputation.MRNN": MRNN,
     "pypots.imputation.NonstationaryTransformer": NonstationaryTransformer,
-    "pypots.imputation.FiLM": FiLM,
+    "pypots.imputation.PatchTST": PatchTST,
     "pypots.imputation.Pyraformer": Pyraformer,
-    "pypots.imputation.Autoformer": Autoformer,
-    "pypots.imputation.Informer": Informer,
+    "pypots.imputation.Reformer": Reformer,
+    "pypots.imputation.RevIN_SCINet": RevIN_SCINet,
+    "pypots.imputation.SAITS": SAITS,
+    "pypots.imputation.SCINet": SCINet,
     "pypots.imputation.StemGNN": StemGNN,
     "pypots.imputation.TimesNet": TimesNet,
-    "pypots.imputation.CSDI": CSDI,
+    "pypots.imputation.TiDE": TiDE,
+    "pypots.imputation.Transformer": Transformer,
     "pypots.imputation.USGAN": USGAN,
-    "pypots.imputation.GPVAE": GPVAE,
-    "pypots.imputation.BRITS": BRITS,
-    "pypots.imputation.MRNN": MRNN,
-    "pypots.imputation.GRUD": GRUD,
     # classification models
-    "pypots.classification.GRUD": GRUD_classification,
     "pypots.classification.BRITS": BRITS_classification,
+    "pypots.classification.GRUD": GRUD_classification,
     "pypots.classification.Raindrop": Raindrop,
     # clustering models
     "pypots.clustering.CRLI": CRLI,

pypots/utils/file.py

@@ -45,3 +45,24 @@ def create_dir_if_not_exist(path: str, is_dir: bool = True) -> None:
     if not os.path.exists(path):
         os.makedirs(path, exist_ok=True)
         logger.info(f"Successfully created the given path {path}")
+
+
+def get_class_full_path(cls) -> str:
+    """Get the full path of the given class.
+
+    Parameters
+    ----------
+    cls:
+        The class to get the full path.
+
+    Returns
+    -------
+    path :
+        The full path of the given class.
+
+    """
+    module = cls.__module__
+    path = cls.__qualname__
+    if module is not None and module != "__builtin__":
+        path = module + "." + path
+    return path

tests/imputation/saits.py

@@ -16,6 +16,7 @@
 from pypots.optim import Adam
 from pypots.utils.logging import logger
 from pypots.utils.metrics import calc_mse
+from pypots.utils.visual.data import plot_data, plot_missingness
 from tests.global_test_config import (
     DATA,
     EPOCHS,
@@ -79,6 +80,12 @@ def test_1_impute(self):
         )
         logger.info(f"SAITS test_MSE: {test_MSE}")
 
+        # plot the missingness and imputed data
+        plot_missingness(
+            ~np.isnan(TEST_SET["X"]), 0, imputation_results["imputation"].shape[1]
+        )
+        plot_data(TEST_SET["X"], TEST_SET["X_ori"], imputation_results["imputation"])
+
     @pytest.mark.xdist_group(name="imputation-saits")
     def test_2_parameters(self):
         assert hasattr(self.saits, "model") and self.saits.model is not None