Add warm_start test

pysr/_cli/main.py

@@ -2,7 +2,13 @@ import warnings
 
 import click
 
-from ..test import get_runtests_cli, runtests, runtests_jax, runtests_torch
+from ..test import (
+    get_runtests_cli,
+    runtests,
+    runtests_jax,
+    runtests_torch,
+    runtests_warm_start,
+)
 
 
 @click.group("pysr")
@@ -42,7 +48,7 @@ def _install(julia_project, quiet, precompile):
     )
 
 
-TEST_OPTIONS = {"main", "jax", "torch", "cli"}
+TEST_OPTIONS = {"main", "jax", "torch", "cli", "warm_start"}
 
 
 @pysr.command("test", help="Run PySR test suite.")
@@ -50,7 +56,7 @@ TEST_OPTIONS = {"main", "jax", "torch", "cli"}
 def _tests(tests):
     """Run part of the PySR test suite.
 
-    Choose from main, jax, torch, and cli.
+    Choose from main, jax, torch, cli, and warm_start.
     """
     if len(tests) == 0:
         raise click.UsageError(
@@ -71,5 +77,7 @@ def _tests(tests):
                 elif test == "cli":
                     runtests_cli = get_runtests_cli()
                     runtests_cli()
+                elif test == "warm_start":
+                    runtests_warm_start()
             else:
                 warnings.warn(f"Invalid test {test}. Skipping.")

pysr/test/__init__.py

@@ -2,10 +2,12 @@ from .test import runtests
 from .test_cli import get_runtests as get_runtests_cli
 from .test_jax import runtests as runtests_jax
 from .test_torch import runtests as runtests_torch
+from .test_warm_start import runtests as runtests_warm_start
 
 __all__ = [
     "runtests",
     "runtests_jax",
     "runtests_torch",
     "get_runtests_cli",
+    "runtests_warm_start",
 ]

pysr/test/params.py

@@ -0,0 +1,8 @@
+import inspect
+
+from .. import PySRRegressor
+
+DEFAULT_PARAMS = inspect.signature(PySRRegressor.__init__).parameters
+DEFAULT_NITERATIONS = DEFAULT_PARAMS["niterations"].default
+DEFAULT_POPULATIONS = DEFAULT_PARAMS["populations"].default
+DEFAULT_NCYCLES = DEFAULT_PARAMS["ncyclesperiteration"].default

pysr/test/test.py

@@ -1,4 +1,3 @@
-import inspect
 import os
 import pickle as pkl
 import tempfile
@@ -12,16 +11,17 @@ import pandas as pd
 import sympy
 from sklearn.utils.estimator_checks import check_estimator
 
-from .. import PySRRegressor, julia_helpers
+from .. import PySRRegressor
 from ..export_latex import sympy2latex
 from ..feature_selection import _handle_feature_selection, run_feature_selection
 from ..sr import _check_assertions, _process_constraints, idx_model_selection
 from ..utils import _csv_filename_to_pkl_filename
-
-DEFAULT_PARAMS = inspect.signature(PySRRegressor.__init__).parameters
-DEFAULT_NITERATIONS = DEFAULT_PARAMS["niterations"].default
-DEFAULT_POPULATIONS = DEFAULT_PARAMS["populations"].default
-DEFAULT_NCYCLES = DEFAULT_PARAMS["ncyclesperiteration"].default
+from .params import (
+    DEFAULT_NCYCLES,
+    DEFAULT_NITERATIONS,
+    DEFAULT_PARAMS,
+    DEFAULT_POPULATIONS,
+)
 
 
 class TestPipeline(unittest.TestCase):

pysr/test/test_warm_start.py

@@ -0,0 +1,104 @@
+import subprocess
+import tempfile
+import textwrap
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+from .. import PySRRegressor
+from .params import (
+    DEFAULT_NCYCLES,
+    DEFAULT_NITERATIONS,
+    DEFAULT_PARAMS,
+    DEFAULT_POPULATIONS,
+)
+
+
+class TestWarmStart(unittest.TestCase):
+    def setUp(self):
+        # Using inspect,
+        # get default niterations from PySRRegressor, and double them:
+        self.default_test_kwargs = dict(
+            progress=False,
+            model_selection="accuracy",
+            niterations=DEFAULT_NITERATIONS * 2,
+            populations=DEFAULT_POPULATIONS * 2,
+            temp_equation_file=True,
+        )
+        self.rstate = np.random.RandomState(0)
+        self.X = self.rstate.randn(100, 5)
+
+    def test_warm_start_from_file(self):
+        """Test that we can warm start in another process."""
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            model = PySRRegressor(
+                **self.default_test_kwargs,
+                unary_operators=["cos"],
+            )
+            model.warm_start = True
+            model.temp_equation_file = False
+            model.equation_file = Path(tmpdirname) / "equations.csv"
+            model.deterministic = True
+            model.multithreading = False
+            model.random_state = 0
+            model.procs = 0
+            model.early_stop_condition = 1e-10
+
+            rstate = np.random.RandomState(0)
+            X = rstate.randn(100, 2)
+            y = np.cos(X[:, 0]) ** 2
+            model.fit(X, y)
+
+            best_loss = model.equations_.iloc[-1]["loss"]
+
+            # Save X and y to a file:
+            X_file = Path(tmpdirname) / "X.npy"
+            y_file = Path(tmpdirname) / "y.npy"
+            np.save(X_file, X)
+            np.save(y_file, y)
+            # Now, create a new process and warm start from the file:
+            result = subprocess.run(
+                [
+                    "python",
+                    "-c",
+                    textwrap.dedent(
+                        f"""
+                        from pysr import PySRRegressor
+                        import numpy as np
+
+                        X = np.load("{X_file}")
+                        y = np.load("{y_file}")
+
+                        print("Loading model from file")
+                        model = PySRRegressor.from_file("{model.equation_file}")
+
+                        assert model.julia_state_ is not None
+
+                        model.warm_start = True
+                        model.niterations = 0
+                        model.max_evals = 0
+                        model.ncyclesperiteration = 0
+
+                        model.fit(X, y)
+
+                        best_loss = model.equations_.iloc[-1]["loss"]
+
+                        assert best_loss <= {best_loss}
+                    """
+                    ),
+                ],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            self.assertEqual(result.returncode, 0)
+            self.assertIn("Loading model from file", result.stdout.decode())
+            self.assertIn("Started!", result.stderr.decode())
+
+
+def runtests():
+    suite = unittest.TestSuite()
+    loader = unittest.TestLoader()
+    suite.addTests(loader.loadTestsFromTestCase(TestWarmStart))
+    runner = unittest.TextTestRunner()
+    return runner.run(suite)