Merge pull request #564 from MilesCranmer/create-pull-request/patch

.deepsource.toml

@@ -1,16 +0,0 @@
-version = 1
-
-test_patterns = ["test/*.py"]
-
-exclude_patterns = ["Project.toml"]
-
-[[analyzers]]
-name = "python"
-enabled = true
-
-  [analyzers.meta]
-  runtime_version = "3.x.x"
-
-[[transformers]]
-name = "black"
-enabled = true

README.md

@@ -287,7 +287,7 @@ model = PySRRegressor(
     # ^ Higher precision calculations.
     warm_start=True,
     # ^ Start from where left off.
-    turbo=True,
+    bumper=True,
     # ^ Faster evaluation (experimental)
     julia_project=None,
     # ^ Can set to the path of a folder containing the

docs/tuning.md

@@ -20,7 +20,7 @@ I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node o
 8. I typically don't use `maxdepth`, but if I do, I set it strictly, while also leaving a bit of room for exploration. e.g., if you want a final equation limited to a depth of `5`, you might set this to `6` or `7`, so that it has a bit of room to explore.
 9.  Set `parsimony` equal to about the minimum loss you would expect, divided by 5-10. e.g., if you expect the final equation to have a loss of `0.001`, you might set `parsimony=0.0001`.
 10. Set `weight_optimize` to some larger value, maybe `0.001`. This is very important if `ncycles_per_iteration` is large, so that optimization happens more frequently.
-11. Set `turbo` to `True`. This may or not work, if there's an error just turn it off (some operators are not SIMD-capable). If it does work, it should give you a nice 20% speedup.
+11. Set `bumper` to `True`. This turns on bump allocation but is experimental. It should give you a nice 20% speedup.
 12. For final runs, after I have tuned everything, I typically set `niterations` to some very large value, and just let it run for a week until my job finishes (genetic algorithms tend not to converge, they can look like they settle down, but then find a new family of expression, and explore a new space). If I am satisfied with the current equations (which are visible either in the terminal or in the saved csv file), I quit the job early.
 
 Since I am running in IPython, I can just hit `q` and then `<enter>` to stop the job, tweak the hyperparameters, and then start the search again.

pysr/_cli/main.py

@@ -1,3 +1,4 @@
+import fnmatch
 import sys
 import unittest
 import warnings
@@ -52,7 +53,14 @@ TEST_OPTIONS = {"main", "jax", "torch", "cli", "dev", "startup"}
 
 @pysr.command("test")
 @click.argument("tests", nargs=1)
-def _tests(tests):
+@click.option(
+    "-k",
+    "expressions",
+    multiple=True,
+    type=str,
+    help="Filter expressions to select specific tests.",
+)
+def _tests(tests, expressions):
     """Run parts of the PySR test suite.
 
     Choose from main, jax, torch, cli, dev, and startup. You can give multiple tests, separated by commas.
@@ -78,11 +86,16 @@ def _tests(tests):
     loader = unittest.TestLoader()
     suite = unittest.TestSuite()
     for test_case in test_cases:
-        suite.addTests(loader.loadTestsFromTestCase(test_case))
+        loaded_tests = loader.loadTestsFromTestCase(test_case)
+        for test in loaded_tests:
+            if len(expressions) == 0 or any(
+                fnmatch.fnmatch(test.id(), "*" + expression + "*")
+                for expression in expressions
+            ):
+                suite.addTest(test)
+
     runner = unittest.TextTestRunner()
     results = runner.run(suite)
-    # Normally unittest would run this, but here we have
-    # to do it manually to get the exit code.
 
     if not results.wasSuccessful():
         sys.exit(1)

pysr/julia_extensions.py

@@ -0,0 +1,32 @@
+"""This file installs and loads extensions for SymbolicRegression."""
+
+from .julia_import import jl
+
+
+def load_required_packages(
+    *, turbo=False, bumper=False, enable_autodiff=False, cluster_manager=None
+):
+    if turbo:
+        load_package("LoopVectorization", "bdcacae8-1622-11e9-2a5c-532679323890")
+    if bumper:
+        load_package("Bumper", "8ce10254-0962-460f-a3d8-1f77fea1446e")
+    if enable_autodiff:
+        load_package("Zygote", "e88e6eb3-aa80-5325-afca-941959d7151f")
+    if cluster_manager is not None:
+        load_package("ClusterManagers", "34f1f09b-3a8b-5176-ab39-66d58a4d544e")
+
+
+def load_package(package_name, uuid):
+    jl.seval(
+        f"""
+        try
+            using {package_name}
+        catch e
+            isa(e, ArgumentError) || throw(e)
+            using Pkg: Pkg
+            Pkg.add(name="{package_name}", uuid="{uuid}")
+            using {package_name}
+        end
+    """
+    )
+    return None

pysr/juliapkg.json

@@ -3,19 +3,11 @@
     "packages": {
         "SymbolicRegression": {
             "uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
-            "version": "=0.23.1"
-        },
-        "ClusterManagers": {
-            "uuid": "34f1f09b-3a8b-5176-ab39-66d58a4d544e",
-            "version": "0.4"
+            "version": "=0.24.0"
         },
         "Serialization": {
             "uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",
             "version": "1"
-        },
-        "Zygote": {
-            "uuid": "e88e6eb3-aa80-5325-afca-941959d7151f",
-            "version": "0.6"
         }
     }
 }

pysr/param_groupings.yml

@@ -74,6 +74,7 @@
   - precision
   - fast_cycle
   - turbo
+  - bumper
   - enable_autodiff
 - Determinism:
   - random_state

pysr/sr.py

@@ -1,4 +1,5 @@
 """Define the PySRRegressor scikit-learn interface."""
+
 import copy
 import os
 import pickle as pkl
@@ -32,6 +33,7 @@ from .export_numpy import sympy2numpy
 from .export_sympy import assert_valid_sympy_symbol, create_sympy_symbols, pysr2sympy
 from .export_torch import sympy2torch
 from .feature_selection import run_feature_selection
+from .julia_extensions import load_required_packages
 from .julia_helpers import (
     PythonCall,
     _escape_filename,
@@ -482,6 +484,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         search evaluation. Certain operators may not be supported.
         Does not support 16-bit precision floats.
         Default is `False`.
+    bumper: bool
+        (Experimental) Whether to use Bumper.jl to speed up the search
+        evaluation. Does not support 16-bit precision floats.
+        Default is `False`.
     precision : int
         What precision to use for the data. By default this is `32`
         (float32), but you can select `64` or `16` as well, giving
@@ -697,7 +703,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         weight_do_nothing: float = 0.21,
         weight_mutate_constant: float = 0.048,
         weight_mutate_operator: float = 0.47,
-        weight_swap_operands: float = 0.0,
+        weight_swap_operands: float = 0.1,
         weight_randomize: float = 0.00023,
         weight_simplify: float = 0.0020,
         weight_optimize: float = 0.0,
@@ -725,6 +731,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         batch_size: int = 50,
         fast_cycle: bool = False,
         turbo: bool = False,
+        bumper: bool = False,
         precision: int = 32,
         enable_autodiff: bool = False,
         random_state=None,
@@ -820,6 +827,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         self.batch_size = batch_size
         self.fast_cycle = fast_cycle
         self.turbo = turbo
+        self.bumper = bumper
         self.precision = precision
         self.enable_autodiff = enable_autodiff
         self.random_state = random_state
@@ -1263,9 +1271,9 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
                 f"PySR currently only supports the following optimizer algorithms: {VALID_OPTIMIZER_ALGORITHMS}"
             )
 
+        progress = self.progress
         # 'Mutable' parameter validation
-        buffer_available = "buffer" in sys.stdout.__dir__()
-        # Params and their default values, if None is given:
+        #  (Params and their default values, if None is given:)
         default_param_mapping = {
             "binary_operators": "+ * - /".split(" "),
             "unary_operators": [],
@@ -1274,7 +1282,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             "multithreading": self.procs != 0 and self.cluster_manager is None,
             "batch_size": 1,
             "update_verbosity": int(self.verbosity),
-            "progress": buffer_available,
+            "progress": progress,
         }
         packed_modified_params = {}
         for parameter, default_value in default_param_mapping.items():
@@ -1293,7 +1301,11 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
                         "`batch_size` has been increased to equal one."
                     )
                     parameter_value = 1
-                elif parameter == "progress" and not buffer_available:
+                elif (
+                    parameter == "progress"
+                    and parameter_value
+                    and "buffer" not in sys.stdout.__dir__()
+                ):
                     warnings.warn(
                         "Note: it looks like you are running in Jupyter. "
                         "The progress bar will be turned off."
@@ -1605,6 +1617,13 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             else "nothing"
         )
 
+        load_required_packages(
+            turbo=self.turbo,
+            bumper=self.bumper,
+            enable_autodiff=self.enable_autodiff,
+            cluster_manager=cluster_manager,
+        )
+
         mutation_weights = SymbolicRegression.MutationWeights(
             mutate_constant=self.weight_mutate_constant,
             mutate_operator=self.weight_mutate_operator,
@@ -1646,15 +1665,16 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             maxdepth=maxdepth,
             fast_cycle=self.fast_cycle,
             turbo=self.turbo,
+            bumper=self.bumper,
             enable_autodiff=self.enable_autodiff,
             migration=self.migration,
             hof_migration=self.hof_migration,
             fraction_replaced_hof=self.fraction_replaced_hof,
             should_simplify=self.should_simplify,
             should_optimize_constants=self.should_optimize_constants,
-            warmup_maxsize_by=0.0
-            if self.warmup_maxsize_by is None
-            else self.warmup_maxsize_by,
+            warmup_maxsize_by=(
+                0.0 if self.warmup_maxsize_by is None else self.warmup_maxsize_by
+            ),
             use_frequency=self.use_frequency,
             use_frequency_in_tournament=self.use_frequency_in_tournament,
             adaptive_parsimony_scaling=self.adaptive_parsimony_scaling,
@@ -1736,9 +1756,11 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             ),
             y_variable_names=jl_y_variable_names,
             X_units=jl_array(self.X_units_),
-            y_units=jl_array(self.y_units_)
-            if isinstance(self.y_units_, list)
-            else self.y_units_,
+            y_units=(
+                jl_array(self.y_units_)
+                if isinstance(self.y_units_, list)
+                else self.y_units_
+            ),
             options=options,
             numprocs=cprocs,
             parallelism=parallelism,

pysr/test/test.py

@@ -58,16 +58,20 @@ class TestPipeline(unittest.TestCase):
         model.fit(self.X, y, variable_names=["c1", "c2", "c3", "c4", "c5"])
         self.assertIn("c1", model.equations_.iloc[-1]["equation"])
 
-    def test_linear_relation_weighted(self):
+    def test_linear_relation_weighted_bumper(self):
         y = self.X[:, 0]
         weights = np.ones_like(y)
         model = PySRRegressor(
             **self.default_test_kwargs,
             early_stop_condition="stop_if(loss, complexity) = loss < 1e-4 && complexity == 1",
+            bumper=True,
         )
         model.fit(self.X, y, weights=weights)
         print(model.equations_)
         self.assertLessEqual(model.get_best()["loss"], 1e-4)
+        self.assertEqual(
+            jl.seval("((::Val{x}) where x) -> x")(model.julia_options_.bumper), True
+        )
 
     def test_multiprocessing_turbo_custom_objective(self):
         rstate = np.random.RandomState(0)
@@ -97,7 +101,9 @@ class TestPipeline(unittest.TestCase):
         self.assertGreaterEqual(best_loss, 0.0)
 
         # Test options stored:
-        self.assertEqual(model.julia_options_.turbo, True)
+        self.assertEqual(
+            jl.seval("((::Val{x}) where x) -> x")(model.julia_options_.turbo), True
+        )
 
     def test_multiline_seval(self):
         # The user should be able to run multiple things in a single seval call:
@@ -128,7 +134,9 @@ class TestPipeline(unittest.TestCase):
         self.assertTrue(jl.typeof(test_state[1]).parameters[1] == jl.Float64)
 
         # Test options stored:
-        self.assertEqual(model.julia_options_.turbo, False)
+        self.assertEqual(
+            jl.seval("((::Val{x}) where x) -> x")(model.julia_options_.turbo), False
+        )
 
     def test_multioutput_custom_operator_quiet_custom_complexity(self):
         y = self.X[:, [0, 1]] ** 2
@@ -163,10 +171,6 @@ class TestPipeline(unittest.TestCase):
         self.assertLessEqual(mse1, 1e-4)
         self.assertLessEqual(mse2, 1e-4)
 
-        bad_y = model.predict(self.X, index=[0, 0])
-        bad_mse = np.average((bad_y - y) ** 2)
-        self.assertGreater(bad_mse, 1e-4)
-
     def test_multioutput_weighted_with_callable_temp_equation(self):
         X = self.X.copy()
         y = X[:, [0, 1]] ** 2
@@ -1028,9 +1032,8 @@ class TestDimensionalConstraints(unittest.TestCase):
         for i in range(2):
             self.assertGreater(model.get_best()[i]["complexity"], 2)
             self.assertLess(model.get_best()[i]["loss"], 1e-6)
-            self.assertGreater(
-                model.equations_[i].query("complexity <= 2").loss.min(), 1e-6
-            )
+            simple_eqs = model.equations_[i].query("complexity <= 2")
+            self.assertTrue(len(simple_eqs) == 0 or simple_eqs.loss.min() > 1e-6)
 
     def test_unit_checks(self):
         """This just checks the number of units passed"""
@@ -1107,8 +1110,10 @@ class TestDimensionalConstraints(unittest.TestCase):
         self.assertNotIn("x1", best["equation"])
         self.assertIn("x2", best["equation"])
         self.assertEqual(best["complexity"], 3)
-        self.assertEqual(model.equations_.iloc[0].complexity, 1)
-        self.assertGreater(model.equations_.iloc[0].loss, 1e-6)
+        self.assertTrue(
+            model.equations_.iloc[0].complexity > 1
+            or model.equations_.iloc[0].loss > 1e-6
+        )
 
         # With pkl file:
         pkl_file = str(temp_dir / "equation_file.pkl")
@@ -1127,8 +1132,8 @@ class TestDimensionalConstraints(unittest.TestCase):
 
         # Try warm start, but with no units provided (should
         # be a different dataset, and thus different result):
-        model.fit(X, y)
         model.early_stop_condition = "(l, c) -> l < 1e-6 && c == 1"
+        model.fit(X, y)
         self.assertEqual(model.equations_.iloc[0].complexity, 1)
         self.assertLess(model.equations_.iloc[0].loss, 1e-6)
 

pysr/test/test_cli.py

@@ -61,7 +61,8 @@ def get_runtests():
                   tests, separated by commas.
 
                 Options:
-                  --help  Show this message and exit.
+                  -k TEXT  Filter expressions to select specific tests.
+                  --help   Show this message and exit.
                 """
             )
             result = self.cli_runner.invoke(pysr, ["test", "--help"])

pysr/test/test_nb.ipynb

@@ -11,18 +11,6 @@
      "text": [
       "Detected Jupyter notebook. Loading juliacall extension. Set `PYSR_AUTOLOAD_EXTENSIONS=no` to disable.\n"
      ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Precompiling SymbolicRegression\n",
-      "\u001b[32m  ✓ \u001b[39mSymbolicRegression\n",
-      "  1 dependency successfully precompiled in 26 seconds. 106 already precompiled.\n",
-      "Precompiling SymbolicRegressionJSON3Ext\n",
-      "\u001b[32m  ✓ \u001b[39m\u001b[90mSymbolicRegression → SymbolicRegressionJSON3Ext\u001b[39m\n",
-      "  1 dependency successfully precompiled in 2 seconds. 110 already precompiled.\n"
-     ]
     }
    ],
    "source": [
@@ -143,14 +131,6 @@
    "execution_count": 7,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/mcranmer/PermaDocuments/SymbolicRegressionMonorepo/.venv/lib/python3.12/site-packages/pysr/sr.py:1297: UserWarning: Note: it looks like you are running in Jupyter. The progress bar will be turned off.\n",
-      "  warnings.warn(\n"
-     ]
-    },
     {
      "data": {
       "text/plain": [
@@ -166,13 +146,6 @@
     "model.fit(X, y)\n",
     "type(model.equations_)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -191,7 +164,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.1"
+   "version": "3.11.2"
   }
  },
  "nbformat": 4,