lazy loading pytests

chenyangkang · Oct 27, 2024 · 045f552 · 045f552
1 parent 9cf0b70
commit 045f552
Showing 1 changed file with 9 additions and 91 deletions.
diff --git a/tests/test_model_lazy_loading.py b/tests/test_model_lazy_loading.py
@@ -22,36 +22,6 @@
 )
 
 
-def test_STEMClassifier_lazy():
-    model = make_STEMClassifier(lazy_loading=True, lazy_loading_dir='./tmp')
-    model = model.fit(X_train, np.where(y_train > 0, 1, 0))
-
-    pred_mean, pred_std = model.predict(X_test.reset_index(drop=True), return_std=True, verbosity=1, n_jobs=1)
-    assert np.sum(~np.isnan(pred_mean)) > 0
-    assert np.sum(~np.isnan(pred_std)) > 0
-
-    pred = model.predict(X_test)
-    assert len(pred) == len(X_test)
-    assert np.sum(np.isnan(pred)) / len(pred) <= 0.3
-
-    pred_df = pd.DataFrame(
-        {"y_true": y_test.flatten(), "y_pred": np.where(pred.flatten() < 0, 0, pred.flatten())}
-    ).dropna()
-    assert len(pred_df) > 0
-
-    eval = AdaSTEM.eval_STEM_res("classification", pred_df.y_true, pred_df.y_pred)
-    assert eval["AUC"] >= 0.5
-    assert eval["kappa"] >= 0.2
-    # assert eval["Spearman_r"] >= 0.2
-
-    model.calculate_feature_importances()
-    assert model.feature_importances_.shape[0] > 0
-
-    importances_by_points = model.assign_feature_importances_by_points(verbosity=0, n_jobs=1)
-    assert importances_by_points.shape[0] > 0
-    assert importances_by_points.shape[1] == len(x_names) + 3
-
-
 def test_parallel_STEMClassifier_lazy():
     model = make_parallel_STEMClassifier(lazy_loading=True)
     model = model.fit(X_train, np.where(y_train > 0, 1, 0))
@@ -80,7 +50,11 @@ def test_parallel_STEMClassifier_lazy():
     importances_by_points = model.assign_feature_importances_by_points(verbosity=0)
     assert importances_by_points.shape[0] > 0
     assert importances_by_points.shape[1] == len(x_names) + 3
-
+
+    model.save(tar_gz_file='./my_model.tar.gz', remove_temporary_file=True) 
+    model = AdaSTEM.load(tar_gz_file='./my_model.tar.gz', remove_original_file=False)
+    model = AdaSTEM.load(tar_gz_file='./my_model.tar.gz', remove_original_file=True)
+
 
 def test_STEMRegressor_lazy():
     model = make_STEMRegressor(lazy_loading=True)
@@ -111,36 +85,10 @@ def test_STEMRegressor_lazy():
     assert importances_by_points.shape[0] > 0
     assert importances_by_points.shape[1] == len(x_names) + 3
 
-
-def test_AdaSTEMClassifier_lazy():
-    model = make_AdaSTEMClassifier(lazy_loading=True)
-    model = model.fit(X_train, np.where(y_train > 0, 1, 0))
-
-    pred_mean, pred_std = model.predict(X_test.reset_index(drop=True), return_std=True, verbosity=1, n_jobs=1)
-    assert np.sum(~np.isnan(pred_mean)) > 0
-    assert np.sum(~np.isnan(pred_std)) > 0
-
-    pred = model.predict(X_test)
-    assert len(pred) == len(X_test)
-    assert np.sum(np.isnan(pred)) / len(pred) <= 0.3
-
-    pred_df = pd.DataFrame(
-        {"y_true": y_test.flatten(), "y_pred": np.where(pred.flatten() < 0, 0, pred.flatten())}
-    ).dropna()
-    assert len(pred_df) > 0
-
-    eval = AdaSTEM.eval_STEM_res("classification", pred_df.y_true, pred_df.y_pred)
-    assert eval["AUC"] >= 0.5
-    assert eval["kappa"] >= 0.2
-    # assert eval["Spearman_r"] >= 0.2
-
-    model.calculate_feature_importances()
-    assert model.feature_importances_.shape[0] > 0
-
-    importances_by_points = model.assign_feature_importances_by_points(verbosity=0, n_jobs=1)
-    assert importances_by_points.shape[0] > 0
-    assert importances_by_points.shape[1] == len(x_names) + 3
-
+    model.save(tar_gz_file='./my_model1.tar.gz', remove_temporary_file=True) 
+    model = AdaSTEM.load(tar_gz_file='./my_model1.tar.gz', remove_original_file=False)
+    model = AdaSTEM.load(tar_gz_file='./my_model1.tar.gz', remove_original_file=True)
+
 
 def test_AdaSTEMRegressor_lazy():
     model = make_AdaSTEMRegressor(lazy_loading=True)
@@ -172,36 +120,6 @@ def test_AdaSTEMRegressor_lazy():
     assert importances_by_points.shape[1] == len(x_names) + 3
 
 
-def test_SphereAdaClassifier_lazy():
-    model = make_SphereAdaClassifier(lazy_loading=True)
-    model = model.fit(X_train, np.where(y_train > 0, 1, 0))
-
-    pred_mean, pred_std = model.predict(X_test.reset_index(drop=True), return_std=True, verbosity=1, n_jobs=1)
-    assert np.sum(~np.isnan(pred_mean)) > 0
-    assert np.sum(~np.isnan(pred_std)) > 0
-
-    pred = model.predict(X_test)
-    assert len(pred) == len(X_test)
-    assert np.sum(np.isnan(pred)) / len(pred) <= 0.3
-
-    pred_df = pd.DataFrame(
-        {"y_true": y_test.flatten(), "y_pred": np.where(pred.flatten() < 0, 0, pred.flatten())}
-    ).dropna()
-    assert len(pred_df) > 0
-
-    eval = AdaSTEM.eval_STEM_res("classification", pred_df.y_true, pred_df.y_pred)
-    assert eval["AUC"] >= 0.5
-    assert eval["kappa"] >= 0.2
-    # assert eval["Spearman_r"] >= 0.2
-
-    model.calculate_feature_importances()
-    assert model.feature_importances_.shape[0] > 0
-
-    importances_by_points = model.assign_feature_importances_by_points(verbosity=0, n_jobs=1)
-    assert importances_by_points.shape[0] > 0
-    assert importances_by_points.shape[1] == len(x_names) + 3
-
-
 def test_parallel_SphereAdaClassifier_lazy():
     model = make_parallel_SphereAdaClassifier(lazy_loading=True)
     model = model.fit(X_train, np.where(y_train > 0, 1, 0))