Move testing to unittest format

test/test.py

@@ -1,38 +1,42 @@
+import unittest
 import numpy as np
 from pysr import pysr
 import sympy
-X = np.random.randn(100, 5)
 
-default_test_kwargs = dict(
-    niterations=10,
-    populations=4,
-    user_input=False,
-    annealing=True,
-    useFrequency=False,
-)
+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)
 
-print("Test 1 - defaults; simple linear relation")
-y = X[:, 0]
-equations = pysr(X, y, **default_test_kwargs)
-print(equations)
-assert 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)
 
-print("Test 2 - test custom operator, and multiple outputs")
-y = X[:, [0, 1]]**2
-equations = pysr(X, y,
-                 unary_operators=["sq(x) = x^2"], binary_operators=["plus"],
-                 extra_sympy_mappings={'square': lambda x: x**2},
-                 **default_test_kwargs)
-print(equations)
-assert equations[0].iloc[-1]['MSE'] < 1e-4
-assert 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)
 
-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"],
-                 **default_test_kwargs)
-
-print(equations)
-assert equations.iloc[-1]['MSE'] < 1e-4
+        print(equations)
+        self.assertLessEqual(equations.iloc[-1]['MSE'], 1e-4)

test/test_jax.py

@@ -1,3 +1,4 @@
+import unittest
 import numpy as np
 from pysr import pysr, sympy2jax
 from jax import numpy as jnp
@@ -5,11 +6,12 @@ from jax import random
 from jax import grad
 import sympy
 
-print("Test JAX 1 - test 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, y, z])
-assert jnp.all(jnp.isclose(f(X, params), true)).item()
+class TestJAX(unittest.TestCase):
+    def test_sympy2jax(self):
+        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, y, z])
+        self.assertTrue(jnp.all(jnp.isclose(f(X, params), true)).item())