5
5
from causal_testing .estimation .genetic_programming_regression_fitter import GP
6
6
7
7
8
+ def root (x ):
9
+ return x ** 0.5
10
+
11
+
8
12
class TestGP (unittest .TestCase ):
9
13
def test_init_invalid_fun_name (self ):
10
14
with self .assertRaises (ValueError ):
@@ -17,7 +21,7 @@ def test_simplify_string(self):
17
21
outcome = None ,
18
22
max_order = 1 ,
19
23
)
20
- self .assertEqual (str (gp .simplify ("power_1(x1)" )), "x1" )
24
+ self .assertEquals (str (gp .simplify ("power_1(x1)" )), "x1" )
21
25
22
26
def test_fitness_string (self ):
23
27
gp = GP (
@@ -26,13 +30,14 @@ def test_fitness_string(self):
26
30
outcome = "outcome" ,
27
31
max_order = 0 ,
28
32
)
29
- self .assertEqual (gp .fitness ("add(x1, 1)" ), (0 ,))
33
+ self .assertEquals (gp .fitness ("add(x1, 1)" ), (0 ,))
30
34
31
- def test_fitness_invalid (self ):
35
+ def test_fitness_inf (self ):
32
36
gp = GP (
33
- df = pd .DataFrame ({"x1" : [1 , None ], "outcome" : [2 , None ]}),
37
+ df = pd .DataFrame ({"x1" : [1 , 2 , 3 ], "outcome" : [2 , 3 , 4 ]}),
34
38
features = ["x1" ],
35
39
outcome = "outcome" ,
36
40
max_order = 0 ,
41
+ extra_operators = [(root , 1 )],
37
42
)
38
- self .assertEqual (gp .fitness ("add(x1, x1 )" ), (float ("inf" ),))
43
+ self .assertEquals (gp .fitness ("root(-1 )" ), (float ("inf" ),))
0 commit comments