test: test visual funcs in SAITS cases;

WenjieDu · Jun 23, 2024 · b3c419a · b3c419a
1 parent a59a7a7
commit b3c419a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/imputation/saits.py b/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