Merge pull request #221 from MilesCranmer/test-refactor

.github/workflows/CI.yml

@@ -36,21 +36,15 @@ jobs:
         uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
-      - name: "Cache dependencies"
-        uses: actions/cache@v1 # Thanks FromFile.jl
-        env:
-          cache-name: cache-artifacts
+      - name: "Cache Julia"
+        uses: julia-actions/cache@v1
         with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-build-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-build-${{ env.cache-name }}-
-            ${{ runner.os }}-build-
-            ${{ runner.os }}-
+          cache-name: ${{ matrix.os }}-test-${{ matrix.julia-version }}-${{ matrix.python-version }}
       - name: "Set up Python"
         uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
+          cache: pip
       - name: "Install PySR"
         run: |
             python -m pip install --upgrade pip
@@ -60,17 +54,17 @@ jobs:
       - name: "Install Coverage tool"
         run: pip install coverage coveralls
       - name: "Run tests"
-        run: coverage run --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
+        run: coverage run --source=pysr --omit='*/test/*,*/feynman_problems.py' -m pysr.test main
       - name: "Install JAX"
         run: pip install jax jaxlib # (optional import)
       - name: "Run JAX tests"
-        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_jax
+        run: coverage run --append --source=pysr --omit='*/test/*,*/feynman_problems.py' -m pysr.test jax
       - name: "Install Torch"
         run: pip install torch # (optional import)
       - name: "Run Torch tests"
-        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_torch
+        run: coverage run --append --source=pysr --omit='*/test/*,*/feynman_problems.py' -m pysr.test torch
       - name: "Run custom env tests"
-        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_env
+        run: coverage run --append --source=pysr --omit='*/test/*,*/feynman_problems.py' -m pysr.test env
       - name: Coveralls
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -91,21 +85,13 @@ jobs:
     
     steps:
       - uses: actions/checkout@v3
-      - name: "Set up Julia"
-        uses: julia-actions/setup-julia@v1
-        with:
-          version: ${{ matrix.julia-version }}
-      - name: "Cache dependencies"
-        uses: actions/cache@v1 # Thanks FromFile.jl
+      - name: "Cache conda"
+        uses: actions/cache@v2
         env:
-          cache-name: cache-artifacts
+          CACHE_NUMBER: 0
         with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-build-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-build-${{ env.cache-name }}-
-            ${{ runner.os }}-build-
-            ${{ runner.os }}-
+          path: ~/conda_pkgs_dir
+          key: ${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-${{ hashFiles('environment.yml') }}
       - name: "Set up Conda"
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -113,30 +99,23 @@ jobs:
           miniforge-version: latest
           auto-activate-base: true
           python-version: ${{ matrix.python-version }}
-          activate-environment: test
+          activate-environment: pysr-test
           environment-file: environment.yml
+      - name: "Cache Julia"
+        uses: julia-actions/cache@v1
+        with:
+          cache-name: ${{ matrix.os }}-conda-${{ matrix.julia-version }}-${{ matrix.python-version }}
       - name: "Install PySR"
         run: |
             python3 -m pip install .
             python3 -c 'import pysr; pysr.install()'
-      - name: "Install Coverage tool"
-        run: python3 -m pip install coverage coveralls
       - name: "Run tests"
-        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
-      - name: "Run custom env tests"
-        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_env
-      - name: Coveralls
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          COVERALLS_FLAG_NAME: conda-test-${{ matrix.test-name }}
-          COVERALLS_PARALLEL: true
-        run: coveralls --service=github
+        run: cd /tmp && python -m pysr.test main
 
   coveralls:
     name: Indicate completion to coveralls.io
     needs: 
       - test
-      - conda_test
     runs-on: ubuntu-latest
     defaults:
       run:

.github/workflows/CI_Windows.yml

@@ -36,21 +36,15 @@ jobs:
         uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
-      - name: "Cache dependencies"
-        uses: actions/cache@v1 # Thanks FromFile.jl
-        env:
-          cache-name: cache-artifacts
+      - name: "Cache Julia"
+        uses: julia-actions/cache@v1
         with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-build-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-build-${{ env.cache-name }}-
-            ${{ runner.os }}-build-
-            ${{ runner.os }}-
+          cache-name: ${{ matrix.os }}-test-${{ matrix.julia-version }}-${{ matrix.python-version }}
       - name: "Set up Python"
         uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
+          cache: pip
       - name: "Install PySR"
         run: |
             python -m pip install --upgrade pip
@@ -58,10 +52,10 @@ jobs:
             python setup.py install
             python -c 'import pysr; pysr.install()'
       - name: "Run tests"
-        run: python -m unittest test.test
+        run: python -m pysr.test main
       - name: "Install Torch"
         run: pip install torch # (optional import)
       - name: "Run Torch tests"
-        run: python -m unittest test.test_torch
+        run: python -m pysr.test torch
       - name: "Run custom env tests"
-        run: python -m unittest test.test_env
+        run: python -m pysr.test env

.github/workflows/CI_conda_forge.yml

@@ -24,7 +24,6 @@ jobs:
         use-mamba: [true, false]
     
     steps:
-      - uses: actions/checkout@v3
       - name: "Set up Conda"
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -40,7 +39,4 @@ jobs:
         run: conda activate pysr-test && conda install pysr
         if: ${{ !matrix.use-mamba }}
       - name: "Run tests"
-        run: |
-          conda activate pysr-test
-          cd $(python -c 'import pysr; from pathlib import Path; print(Path(pysr.__file__).parent.parent)')
-          python -m unittest test.test
+        run: python -m pysr.test main

.github/workflows/CI_docker.yml

@@ -38,4 +38,4 @@ jobs:
       - name: Build docker
         run: docker build -t pysr --build-arg ARCH=${{ matrix.arch }} --build-arg VERSION=${{ matrix.julia-version }} --build-arg PYVERSION=${{ matrix.python-version }} .
       - name: Test docker
-        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m unittest test.test && python3 -m unittest test.test_env'
+        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m pysr.test main && python3 -m pysr.test env'

.github/workflows/CI_docker_large_nightly.yml

@@ -31,4 +31,4 @@ jobs:
       - name: Build docker
         run: docker build -t pysr --build-arg ARCH=${{ matrix.arch }} --build-arg VERSION=${{ matrix.julia-version }} --build-arg PYVERSION=${{ matrix.python-version }} .
       - name: Test docker
-        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m unittest test.test && python3 -m unittest test.test_env'
+        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m pysr.test main && python3 -m pysr.test env'

.github/workflows/CI_large_nightly.yml

@@ -41,7 +41,7 @@ jobs:
             python setup.py install
             python -c 'import pysr; pysr.install()'
       - name: "Run tests"
-        run: python -m unittest test.test
+        run: python -m pysr.test main
       - name: "Run new env test"
-        run: python -m unittest test.test_env
+        run: python -m pysr.test env
         if: ${{ !(matrix.os == 'windows-latest' && matrix.python-version == '3.7') }}

.github/workflows/CI_mac.yml

@@ -36,21 +36,15 @@ jobs:
         uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
-      - name: "Cache dependencies"
-        uses: actions/cache@v3
-        env:
-          cache-name: cache-artifacts
+      - name: "Cache Julia"
+        uses: julia-actions/cache@v1
         with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-build-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-build-${{ env.cache-name }}-
-            ${{ runner.os }}-build-
-            ${{ runner.os }}-
+          cache-name: ${{ matrix.os }}-test-${{ matrix.julia-version }}-${{ matrix.python-version }}
       - name: "Set up Python"
         uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
+          cache: pip
       - name: "Install PySR"
         run: |
             python -m pip install --upgrade pip
@@ -58,14 +52,14 @@ jobs:
             python setup.py install
             python -c 'import pysr; pysr.install()'
       - name: "Run tests"
-        run: python -m unittest test.test
+        run: python -m pysr.test main
       - name: "Install JAX"
         run: pip install jax jaxlib # (optional import)
       - name: "Run JAX tests"
-        run: python -m unittest test.test_jax
+        run: python -m pysr.test jax
       - name: "Install Torch"
         run: pip install torch # (optional import)
       - name: "Run Torch tests"
-        run: python -m unittest test.test_torch
+        run: python -m pysr.test torch
       - name: "Run custom env tests"
-        run: python -m unittest test.test_env
+        run: python -m pysr.test env

.github/workflows/docs.yml

@@ -25,6 +25,7 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: 3.9
+          cache: pip
       - name: "Install packages for docs building"
         run: pip install mkdocs-material mkdocs-autorefs 'mkdocstrings[python]' docstring_parser
       - name: "Install PySR"

.gitignore

@@ -19,3 +19,6 @@ Manifest.toml
 workflow
 docs/index.md
 site
+**/.DS_Store
+**/*.code-workspace
+**/*.tar.gz

Dockerfile

@@ -50,7 +50,4 @@ RUN pip3 install .
 # Install Julia pre-requisites:
 RUN python3 -c 'import pysr; pysr.install()'
 
-# Add tests
-ADD ./test/ /pysr/test/
-
 CMD ["bash"]

environment.yml

@@ -1,7 +1,6 @@
 name: test
 channels:
   - conda-forge
-  - defaults
 dependencies:
   - sympy
   - pandas

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/__main__.py

@@ -0,0 +1,41 @@
+"""CLI for running PySR's test suite."""
+import argparse
+import os
+
+from . import *
+
+if __name__ == "__main__":
+    # Get args:
+    parser = argparse.ArgumentParser()
+    parser.usage = "python -m pysr.test [tests...]"
+    parser.add_argument(
+        "test",
+        nargs="*",
+        help="Test to run. One or more of 'main', 'env', 'jax', 'torch'.",
+    )
+
+    # Parse args:
+    args = parser.parse_args()
+    tests = args.test
+
+    if len(tests) == 0:
+        # Raise help message:
+        parser.print_help()
+        raise SystemExit(1)
+
+    # Run tests:
+    for test in tests:
+        if test in {"main", "env", "jax", "torch"}:
+            cur_dir = os.path.dirname(os.path.abspath(__file__))
+            print(f"Running test from {cur_dir}")
+            if test == "main":
+                runtests()
+            elif test == "env":
+                runtests_env()
+            elif test == "jax":
+                runtests_jax()
+            elif test == "torch":
+                runtests_torch()
+        else:
+            parser.print_help()
+            raise SystemExit(1)

{test → 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)

{test → 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)

{test → 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)

{test → 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)