Force tests to import locally

pysr/test/__init__.py

@@ -0,0 +1,4 @@
+from .test import runtests
+from .test_env import runtests as runtests_env
+from .test_jax import runtests as runtests_jax
+from .test_torch import runtests as runtests_torch

pysr/test/test.py

@@ -4,14 +4,6 @@ import inspect
 import unittest
 import numpy as np
 from sklearn import model_selection
-from pysr import PySRRegressor
-from pysr.sr import (
-    run_feature_selection,
-    _handle_feature_selection,
-    _csv_filename_to_pkl_filename,
-    idx_model_selection,
-)
-from pysr.export_latex import to_latex
 from sklearn.utils.estimator_checks import check_estimator
 import sympy
 import pandas as pd
@@ -20,6 +12,15 @@ import pickle as pkl
 import tempfile
 from pathlib import Path
 
+from .. import PySRRegressor
+from ..sr import (
+    run_feature_selection,
+    _handle_feature_selection,
+    _csv_filename_to_pkl_filename,
+    idx_model_selection,
+)
+from ..export_latex import to_latex
+
 DEFAULT_PARAMS = inspect.signature(PySRRegressor.__init__).parameters
 DEFAULT_NITERATIONS = DEFAULT_PARAMS["niterations"].default
 DEFAULT_POPULATIONS = DEFAULT_PARAMS["populations"].default
@@ -856,3 +857,21 @@ class TestLaTeXTable(unittest.TestCase):
             + self.create_true_latex(middle_part, include_score=True)
         )
         self.assertEqual(latex_table_str, true_latex_table_str)
+
+
+def runtests():
+    """Run all tests in test.py."""
+    suite = unittest.TestSuite()
+    loader = unittest.TestLoader()
+    test_cases = [
+        TestPipeline,
+        TestBest,
+        TestFeatureSelection,
+        TestMiscellaneous,
+        TestLaTeXTable,
+    ]
+    for test_case in test_cases:
+        tests = loader.loadTestsFromTestCase(test_case)
+        suite.addTests(tests)
+    runner = unittest.TextTestRunner()
+    return runner.run(suite)

pysr/test/test_env.py

@@ -2,9 +2,10 @@
 
 import unittest
 import os
-from pysr import julia_helpers
 from tempfile import TemporaryDirectory
 
+from .. import julia_helpers
+
 
 class TestJuliaProject(unittest.TestCase):
     """Various tests for working with Julia projects."""
@@ -46,3 +47,12 @@ class TestJuliaProject(unittest.TestCase):
                 del os.environ["JULIA_DEPOT_PATH"]
             else:
                 os.environ["JULIA_DEPOT_PATH"] = old_env
+
+
+def runtests():
+    """Run all tests in test_env.py."""
+    loader = unittest.TestLoader()
+    suite = unittest.TestSuite()
+    suite.addTests(loader.loadTestsFromTestCase(TestJuliaProject))
+    runner = unittest.TextTestRunner()
+    return runner.run(suite)

pysr/test/test_jax.py

@@ -1,18 +1,20 @@
 import unittest
 import numpy as np
-from pysr import sympy2jax, PySRRegressor
 import pandas as pd
-from jax import numpy as jnp
-from jax import random
 import sympy
 from functools import partial
 
+from .. import sympy2jax, PySRRegressor
+
 
 class TestJAX(unittest.TestCase):
     def setUp(self):
         np.random.seed(0)
 
     def test_sympy2jax(self):
+        from jax import numpy as jnp
+        from jax import random
+
         x, y, z = sympy.symbols("x y z")
         cosx = 1.0 * sympy.cos(x) + y
         key = random.PRNGKey(0)
@@ -22,6 +24,8 @@ class TestJAX(unittest.TestCase):
         self.assertTrue(jnp.all(jnp.isclose(f(X, params), true)).item())
 
     def test_pipeline_pandas(self):
+        from jax import numpy as jnp
+
         X = pd.DataFrame(np.random.randn(100, 10))
         y = np.ones(X.shape[0])
         model = PySRRegressor(
@@ -53,6 +57,8 @@ class TestJAX(unittest.TestCase):
         )
 
     def test_pipeline(self):
+        from jax import numpy as jnp
+
         X = np.random.randn(100, 10)
         y = np.ones(X.shape[0])
         model = PySRRegressor(progress=False, max_evals=10000, output_jax_format=True)
@@ -112,3 +118,12 @@ class TestJAX(unittest.TestCase):
 
         np.testing.assert_almost_equal(y.values, np_output, decimal=3)
         np.testing.assert_almost_equal(y.values, jax_output, decimal=3)
+
+
+def runtests():
+    """Run all tests in test_jax.py."""
+    loader = unittest.TestLoader()
+    suite = unittest.TestSuite()
+    suite.addTests(loader.loadTestsFromTestCase(TestJAX))
+    runner = unittest.TextTestRunner()
+    return runner.run(suite)

pysr/test/test_torch.py

@@ -1,22 +1,26 @@
 import unittest
 import numpy as np
 import pandas as pd
-from pysr import sympy2torch, PySRRegressor
+import platform
+import sympy
+from .. import sympy2torch, PySRRegressor
 
 # Need to initialize Julia before importing torch...
-import platform
 
-if platform.system() == "Darwin":
-    # Import PyJulia, then Torch
-    from pysr.julia_helpers import init_julia
 
-    Main = init_julia()
-    import torch
-else:
-    # Import Torch, then PyJulia
-    # https://github.com/pytorch/pytorch/issues/78829
-    import torch
-import sympy
+def _import_torch():
+    if platform.system() == "Darwin":
+        # Import PyJulia, then Torch
+        from ..julia_helpers import init_julia
+
+        init_julia()
+
+        import torch
+    else:
+        # Import Torch, then PyJulia
+        # https://github.com/pytorch/pytorch/issues/78829
+        import torch
+    return torch
 
 
 class TestTorch(unittest.TestCase):
@@ -24,6 +28,7 @@ class TestTorch(unittest.TestCase):
         np.random.seed(0)
 
     def test_sympy2torch(self):
+        torch = _import_torch()
         x, y, z = sympy.symbols("x y z")
         cosx = 1.0 * sympy.cos(x) + y
 
@@ -35,6 +40,7 @@ class TestTorch(unittest.TestCase):
         )
 
     def test_pipeline_pandas(self):
+        torch = _import_torch()
         X = pd.DataFrame(np.random.randn(100, 10))
         y = np.ones(X.shape[0])
         model = PySRRegressor(
@@ -69,6 +75,7 @@ class TestTorch(unittest.TestCase):
         )
 
     def test_pipeline(self):
+        torch = _import_torch()
         X = np.random.randn(100, 10)
         y = np.ones(X.shape[0])
         model = PySRRegressor(
@@ -103,6 +110,7 @@ class TestTorch(unittest.TestCase):
         )
 
     def test_mod_mapping(self):
+        torch = _import_torch()
         x, y, z = sympy.symbols("x y z")
         expression = x**2 + sympy.atanh(sympy.Mod(y + 1, 2) - 1) * 3.2 * z
 
@@ -120,6 +128,7 @@ class TestTorch(unittest.TestCase):
         )
 
     def test_custom_operator(self):
+        torch = _import_torch()
         X = np.random.randn(100, 3)
         y = np.ones(X.shape[0])
         model = PySRRegressor(
@@ -160,6 +169,7 @@ class TestTorch(unittest.TestCase):
         )
 
     def test_feature_selection_custom_operators(self):
+        torch = _import_torch()
         rstate = np.random.RandomState(0)
         X = pd.DataFrame({f"k{i}": rstate.randn(2000) for i in range(10, 21)})
         cos_approx = lambda x: 1 - (x**2) / 2 + (x**4) / 24 + (x**6) / 720
@@ -188,3 +198,12 @@ class TestTorch(unittest.TestCase):
 
         np.testing.assert_almost_equal(y.values, np_output, decimal=3)
         np.testing.assert_almost_equal(y.values, torch_output, decimal=3)
+
+
+def runtests():
+    """Run all tests in test_torch.py."""
+    loader = unittest.TestLoader()
+    suite = unittest.TestSuite()
+    suite.addTests(loader.loadTestsFromTestCase(TestTorch))
+    runner = unittest.TextTestRunner()
+    return runner.run(suite)