-
Notifications
You must be signed in to change notification settings - Fork 250
/
unit_tests.py
49 lines (36 loc) · 1.54 KB
/
unit_tests.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import unittest
from boruta import BorutaPy
from sklearn.ensemble import RandomForestClassifier
import numpy as np
class BorutaTestCases(unittest.TestCase):
def test_get_tree_num(self):
rfc = RandomForestClassifier(max_depth=10)
bt = BorutaPy(rfc)
self.assertEqual(bt._get_tree_num(10), 44, "Tree Est. Math Fail")
self.assertEqual(bt._get_tree_num(100), 141, "Tree Est. Math Fail")
def test_if_boruta_extracts_relevant_features(self):
np.random.seed(42)
y = np.random.binomial(1, 0.5, 1000)
X = np.zeros((1000, 10))
z = y - np.random.binomial(1, 0.1, 1000) + np.random.binomial(1, 0.1, 1000)
z[z == -1] = 0
z[z == 2] = 1
# 5 relevant features
X[:, 0] = z
X[:, 1] = y * np.abs(np.random.normal(0, 1, 1000)) + np.random.normal(0, 0.1, 1000)
X[:, 2] = y + np.random.normal(0, 1, 1000)
X[:, 3] = y ** 2 + np.random.normal(0, 1, 1000)
X[:, 4] = np.sqrt(y) + np.random.binomial(2, 0.1, 1000)
# 5 irrelevant features
X[:, 5] = np.random.normal(0, 1, 1000)
X[:, 6] = np.random.poisson(1, 1000)
X[:, 7] = np.random.binomial(1, 0.3, 1000)
X[:, 8] = np.random.normal(0, 1, 1000)
X[:, 9] = np.random.poisson(1, 1000)
rfc = RandomForestClassifier()
bt = BorutaPy(rfc)
bt.fit(X, y)
# make sure that only all the relevant features are returned
self.assertListEqual(list(range(5)), list(np.where(bt.support_)[0]))
if __name__ == '__main__':
unittest.main()