File size: 1,333 Bytes
bed9614
05cf610
 
 
 
 
bed9614
 
48b723d
8cfda07
ddb4d52
3663ed6
 
bed9614
022de02
ddb4d52
48b723d
ddb4d52
 
b078c6a
e394c0e
3663ed6
 
ddb4d52
022de02
8cfda07
48b723d
 
 
 
 
3663ed6
 
48b723d
022de02
05cf610
 
 
 
 
 
 
 
 
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
import numpy as np
from pysr import pysr, sympy2jax
from jax import numpy as jnp
from jax import random
from jax import grad
import sympy
X = np.random.randn(100, 5)

print("Test 1 - defaults; simple linear relation")
y = X[:, 0]
equations = pysr(X, y,
                 niterations=10,
                 user_input=False)
print(equations)
assert equations.iloc[-1]['MSE'] < 1e-4

print("Test 2 - test custom operator")
y = X[:, 0]**2
equations = pysr(X, y,
                 unary_operators=["sq(x) = x^2"], binary_operators=["plus"],
                 extra_sympy_mappings={'square': lambda x: x**2},
                 niterations=10,
                 user_input=False)
print(equations)
assert equations.iloc[-1]['MSE'] < 1e-4

X = np.random.randn(100, 1)
y = X[:, 0] + 3.0
print("Test 3 - empty operator list, and single dimension input")
equations = pysr(X, y,
                 unary_operators=[], binary_operators=["plus"],
                 niterations=10,
                 user_input=False)
print(equations)
assert equations.iloc[-1]['MSE'] < 1e-4

print("Test 4 - text JAX export")
x, y, z = sympy.symbols('x y z')
cosx = 1.0 * sympy.cos(x) + y
key = random.PRNGKey(0)
X = random.normal(key, (1000, 2))
true = 1.0 * jnp.cos(X[:, 0]) + X[:, 1]
f, params = sympy2jax(cosx, [x])
assert jnp.all(jnp.isclose(f(X, params), true)).item()