Merge pull request #228 from MilesCranmer/optimize=3

Make Julia startup options configurable; set optimize=3

docs/param_groupings.yml

@@ -82,6 +82,7 @@
   - delete_tempfiles
   - julia_project
   - update
+  - julia_kwargs
 - Exporting the Results:
   - equation_file
   - output_jax_format

pysr/julia_helpers.py

@@ -10,6 +10,8 @@ from .version import __version__, __symbolic_regression_jl_version__
 
 juliainfo = None
 julia_initialized = False
+julia_kwargs_at_initialization = None
+julia_activated_env = None
 
 
 def _load_juliainfo():
@@ -143,13 +145,18 @@ def _check_for_conflicting_libraries():  # pragma: no cover
         )
 
 
-def init_julia(julia_project=None, quiet=False):
+def init_julia(julia_project=None, quiet=False, julia_kwargs=None):
     """Initialize julia binary, turning off compiled modules if needed."""
     global julia_initialized
+    global julia_kwargs_at_initialization
+    global julia_activated_env
 
     if not julia_initialized:
         _check_for_conflicting_libraries()
 
+    if julia_kwargs is None:
+        julia_kwargs = {"optimize": 3}
+
     from julia.core import JuliaInfo, UnsupportedPythonError
 
     _julia_version_assertion()
@@ -167,21 +174,37 @@ def init_julia(julia_project=None, quiet=False):
     if not info.is_pycall_built():
         raise ImportError(_import_error())
 
-    Main = None
-    try:
-        from julia import Main as _Main
+    from julia.core import Julia
 
-        Main = _Main
+    try:
+        Julia(**julia_kwargs)
     except UnsupportedPythonError:
         # Static python binary, so we turn off pre-compiled modules.
-        from julia.core import Julia
+        julia_kwargs = {**julia_kwargs, "compiled_modules": False}
+        Julia(**julia_kwargs)
 
-        jl = Julia(compiled_modules=False)
-        from julia import Main as _Main
+    from julia import Main as _Main
 
-        Main = _Main
+    Main = _Main
 
-    if julia_initialized:
+    if julia_activated_env is None:
+        julia_activated_env = processed_julia_project
+
+    if julia_initialized and julia_kwargs_at_initialization is not None:
+        # Check if the kwargs are the same as the previous initialization
+        init_set = set(julia_kwargs_at_initialization.items())
+        new_set = set(julia_kwargs.items())
+        set_diff = new_set - init_set
+        # Remove the `compiled_modules` key, since it is not a user-specified kwarg:
+        set_diff = {k: v for k, v in set_diff if k != "compiled_modules"}
+        if len(set_diff) > 0:
+            warnings.warn(
+                "Julia has already started. The new Julia options "
+                + str(set_diff)
+                + " will be ignored."
+            )
+
+    if julia_initialized and julia_activated_env != processed_julia_project:
         Main.eval("using Pkg")
 
         io_arg = _get_io_arg(quiet)
@@ -193,6 +216,11 @@ def init_julia(julia_project=None, quiet=False):
             f"{io_arg})"
         )
 
+        julia_activated_env = processed_julia_project
+
+    if not julia_initialized:
+        julia_kwargs_at_initialization = julia_kwargs
+
     julia_initialized = True
     return Main
 
@@ -234,7 +262,7 @@ def _backend_version_assertion(Main):
         if backend_version != expected_backend_version:  # pragma: no cover
             warnings.warn(
                 f"PySR backend (SymbolicRegression.jl) version {backend_version} "
-                "does not match expected version {expected_backend_version}. "
+                f"does not match expected version {expected_backend_version}. "
                 "Things may break. "
                 "Please update your PySR installation with "
                 "`python -c 'import pysr; pysr.install()'`."
@@ -257,6 +285,7 @@ def _update_julia_project(Main, is_shared, io_arg):
     try:
         if is_shared:
             _add_sr_to_julia_project(Main, io_arg)
+        Main.eval("using Pkg")
         Main.eval(f"Pkg.resolve({io_arg})")
     except (JuliaError, RuntimeError) as e:
         raise ImportError(_import_error()) from e

pysr/sr.py

@@ -581,10 +581,15 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         inputting to PySR. Can help PySR fit noisy data.
         Default is `False`.
     select_k_features : int
-         whether to run feature selection in Python using random forests,
-         before passing to the symbolic regression code. None means no
-         feature selection; an int means select that many features.
-         Default is `None`.
+        Whether to run feature selection in Python using random forests,
+        before passing to the symbolic regression code. None means no
+        feature selection; an int means select that many features.
+        Default is `None`.
+    julia_kwargs : dict
+        Keyword arguments to pass to `julia.core.Julia(...)` to initialize
+        the Julia runtime. The default, when `None`, is to set `threads` equal
+        to `procs`, and `optimize` to 3.
+        Default is `None`.
     **kwargs : dict
         Supports deprecated keyword arguments. Other arguments will
         result in an error.
@@ -733,6 +738,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         extra_jax_mappings=None,
         denoise=False,
         select_k_features=None,
+        julia_kwargs=None,
         **kwargs,
     ):
 
@@ -827,6 +833,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         # Pre-modelling transformation
         self.denoise = denoise
         self.select_k_features = select_k_features
+        self.julia_kwargs = julia_kwargs
 
         # Once all valid parameters have been assigned handle the
         # deprecated kwargs
@@ -1259,6 +1266,17 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             + len(packed_modified_params["unary_operators"])
             > 0
         )
+
+        julia_kwargs = {}
+        if self.julia_kwargs is not None:
+            for key, value in self.julia_kwargs.items():
+                julia_kwargs[key] = value
+        if "optimize" not in julia_kwargs:
+            julia_kwargs["optimize"] = 3
+        if "threads" not in julia_kwargs and packed_modified_params["multithreading"]:
+            julia_kwargs["threads"] = self.procs
+        packed_modified_params["julia_kwargs"] = julia_kwargs
+
         return packed_modified_params
 
     def _validate_and_set_fit_params(self, X, y, Xresampled, weights, variable_names):
@@ -1469,31 +1487,21 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         batch_size = mutated_params["batch_size"]
         update_verbosity = mutated_params["update_verbosity"]
         progress = mutated_params["progress"]
+        julia_kwargs = mutated_params["julia_kwargs"]
 
         # Start julia backend processes
-        if Main is None:
-            if multithreading:
-                os.environ["JULIA_NUM_THREADS"] = str(self.procs)
-
-            Main = init_julia(self.julia_project)
+        Main = init_julia(self.julia_project, julia_kwargs=julia_kwargs)
 
         if cluster_manager is not None:
             cluster_manager = _load_cluster_manager(cluster_manager)
 
-        if not already_ran:
-            julia_project, is_shared = _process_julia_project(self.julia_project)
-            Main.eval("using Pkg")
+        if self.update:
+            _, is_shared = _process_julia_project(self.julia_project)
             io = "devnull" if update_verbosity == 0 else "stderr"
             io_arg = (
                 f"io={io}" if is_julia_version_greater_eq(version=(1, 6, 0)) else ""
             )
-
-            Main.eval(
-                f'Pkg.activate("{_escape_filename(julia_project)}", shared = Bool({int(is_shared)}), {io_arg})'
-            )
-
-            if self.update:
-                _update_julia_project(Main, is_shared, io_arg)
+            _update_julia_project(Main, is_shared, io_arg)
 
         SymbolicRegression = _load_backend(Main)
 

pysr/test/test.py

@@ -12,6 +12,7 @@ import pickle as pkl
 import tempfile
 from pathlib import Path
 
+from .. import julia_helpers
 from .. import PySRRegressor
 from ..sr import (
     run_feature_selection,
@@ -566,6 +567,23 @@ class TestMiscellaneous(unittest.TestCase):
         with self.assertRaises(ValueError):
             model.fit(X, y)
 
+    def test_changed_options_warning(self):
+        """Check that a warning is given if Julia options are changed."""
+        if julia_helpers.julia_kwargs_at_initialization is None:
+            julia_helpers.init_julia(julia_kwargs={"threads": 2, "optimize": 3})
+
+        cur_init = julia_helpers.julia_kwargs_at_initialization
+
+        threads_to_change = cur_init["threads"] + 1
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            with self.assertRaises(Exception) as context:
+                julia_helpers.init_julia(
+                    julia_kwargs={"threads": threads_to_change, "optimize": 3}
+                )
+            self.assertIn("Julia has already started", str(context.exception))
+            self.assertIn("threads", str(context.exception))
+
     def test_extra_sympy_mappings_undefined(self):
         """extra_sympy_mappings=None errors for custom operators"""
         model = PySRRegressor(unary_operators=["square2(x) = x^2"])

pysr/version.py

@@ -1,2 +1,2 @@
-__version__ = "0.11.10"
+__version__ = "0.11.11"
 __symbolic_regression_jl_version__ = "0.14.4"