File size: 1,452 Bytes
2f38c9c
bed9614
41e5fd5
05cf610
bed9614
2f38c9c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10ff16a
2f38c9c
 
 
 
 
 
 
 
 
ddb4d52
2f38c9c
 
 
 
 
 
8cfda07
2f38c9c
 
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
import unittest
import numpy as np
from pysr import pysr
import sympy

class TestPipeline(unittest.TestCase):
    def setUp(self):
        self.default_test_kwargs = dict(
            niterations=10,
            populations=4,
            user_input=False,
            annealing=True,
            useFrequency=False,
        )
        np.random.seed(0)
        self.X = np.random.randn(100, 5)
    
    def test_linear_relation(self):
        y = self.X[:, 0]
        equations = pysr(self.X, y, **self.default_test_kwargs)
        print(equations)
        self.assertLessEqual(equations.iloc[-1]['MSE'], 1e-4)

    def test_multioutput_custom_operator(self):
        y = self.X[:, [0, 1]]**2
        equations = pysr(self.X, y,
                         unary_operators=["sq(x) = x^2"], binary_operators=["plus"],
                         extra_sympy_mappings={'square': lambda x: x**2},
                         **self.default_test_kwargs)
        print(equations)
        self.assertLessEqual(equations[0].iloc[-1]['MSE'], 1e-4)
        self.assertLessEqual(equations[1].iloc[-1]['MSE'], 1e-4)

    def test_empty_operators_single_input(self):
        X = np.random.randn(100, 1)
        y = X[:, 0] + 3.0
        equations = pysr(X, y,
                         unary_operators=[], binary_operators=["plus"],
                         **self.default_test_kwargs)

        print(equations)
        self.assertLessEqual(equations.iloc[-1]['MSE'], 1e-4)