Test for pruning models based on p_max

opengridcc · Dec 7, 2017 · 16d9b86 · 16d9b86
1 parent f7a403d
commit 16d9b86
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/opengrid/tests/test_regression.py b/opengrid/tests/test_regression.py
@@ -87,6 +87,17 @@ def test_alternative_metrics(self):
         self.assertEqual(best_rsquared, best_akaike)
         self.assertEqual(best_rsquared, best_bic)
 
+    def test_prune(self):
+        "Create overfitted model and prune it"
+        df = datasets.get('gas_2016_hour')
+        df_month = df.resample('MS').sum()
+        mvlr = og.MultiVarLinReg(df_month, '313b')
+        self.assertTrue("Q('d5a7')" in mvlr.fit.model.exog_names)
+        pruned = mvlr._prune(mvlr.fit, 0.05)
+        self.assertTrue("Q('d5a7')" in pruned.model.exog_names)
+        pruned = mvlr._prune(mvlr.fit, 0.0001)
+        self.assertFalse("Q('d5a7')" in pruned.model.exog_names)
+
 
 
 if __name__ == '__main__':