feat: expose backsolve options

docs/src/_api_template.md

@@ -87,6 +87,14 @@ PARAMSKEY
         heading_level: 3
         show_root_full_path: false
 
+## Backsolve Options
+
+::: pysr.BacksolveOptions
+    options:
+        show_root_heading: true
+        heading_level: 3
+        show_root_full_path: false
+
 ## Logger Specifications
 
 ::: pysr.TensorBoardLoggerSpec

pysr/__init__.py

@@ -23,6 +23,7 @@
 from importlib.metadata import PackageNotFoundError, version
 
 from . import sklearn_monkeypatch
+from .backsolve_options import BacksolveOptions
 from .deprecated import best, best_callable, best_row, best_tex, install, pysr
 from .export_jax import sympy2jax
 from .export_torch import sympy2torch
@@ -51,6 +52,7 @@
     "install",
     "load_all_packages",
     "PySRRegressor",
+    "BacksolveOptions",
     "AbstractExpressionSpec",
     "ExpressionSpec",
     "TemplateExpressionSpec",

pysr/backsolve_options.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .julia_import import AnyValue, SymbolicRegression
+
+
+@dataclass
+class BacksolveOptions:
+    """Options for the experimental backsolve mutation.
+
+    Parameters
+    ----------
+    max_library_size : int
+        Maximum number of candidate library terms. Default is `500`.
+    lambda_ : float
+        STLSQ sparsity threshold. Default is `0.01`.
+    max_iter : int
+        Maximum number of STLSQ iterations. Default is `10`.
+    """
+
+    max_library_size: int = 500
+    lambda_: float = 0.01
+    max_iter: int = 10
+
+    def julia_options(self) -> AnyValue:
+        """Create the corresponding `SymbolicRegression.BacksolveOptions`."""
+        return SymbolicRegression.BacksolveOptions(
+            max_library_size=int(self.max_library_size),
+            max_iter=int(self.max_iter),
+            **{"lambda": float(self.lambda_)},
+        )

pysr/juliapkg.json

@@ -4,7 +4,7 @@
         "SymbolicRegression": {
             "uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
             "url": "https://github.com/MilesCranmer/SymbolicRegression.jl",
-            "rev": "v2.0.0-alpha.9"
+            "rev": "v2.0.0-alpha.11"
         },
         "Serialization": {
             "uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",

pysr/param_groupings.yml

@@ -45,12 +45,15 @@
     - weight_randomize
     - weight_simplify
     - weight_optimize
+    - weight_backsolve
     - crossover_probability
     - annealing
     - alpha
     - perturbation_factor
     - probability_negate_constant
     - skip_mutation_failures
+  - Backsolve:
+    - backsolve_options
   - Tournament Selection:
     - tournament_selection_n
     - tournament_selection_p

pysr/sr.py

@@ -26,6 +26,7 @@
 from sklearn.utils.validation import _check_feature_names_in  # type: ignore
 from sklearn.utils.validation import check_is_fitted
 
+from .backsolve_options import BacksolveOptions
 from .denoising import denoise, multi_denoise
 from .deprecated import DEPRECATED_KWARGS
 from .export_latex import (
@@ -629,13 +630,19 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         every iteration. Using it as a mutation is useful if you want to use
         a large `ncycles_periteration`, and may not optimize very often.
         Default is `0.0`.
+    weight_backsolve: float
+        Relative likelihood for the experimental backsolve mutation.
+        Default is `0.0`.
     crossover_probability : float
         Absolute probability of crossover-type genetic operation, instead of a mutation.
         Default is `0.0259`.
     skip_mutation_failures : bool
         Whether to skip mutation and crossover failures, rather than
         simply re-sampling the current member.
         Default is `True`.
+    backsolve_options : BacksolveOptions | None
+        Options for the experimental backsolve mutation. Default is `None`,
+        which uses the Julia backend defaults.
     migration : bool
         Whether to migrate.  Default is `True`.
     hof_migration : bool
@@ -996,8 +1003,10 @@ def __init__(
         weight_randomize: float = 0.000502,
         weight_simplify: float = 0.00209,
         weight_optimize: float = 0.0,
+        weight_backsolve: float = 0.0,
         crossover_probability: float = 0.0259,
         skip_mutation_failures: bool = True,
+        backsolve_options: BacksolveOptions | None = None,
         migration: bool = True,
         hof_migration: bool = True,
         topn: int = 12,
@@ -1113,8 +1122,10 @@ def __init__(
         self.weight_randomize = weight_randomize
         self.weight_simplify = weight_simplify
         self.weight_optimize = weight_optimize
+        self.weight_backsolve = weight_backsolve
         self.crossover_probability = crossover_probability
         self.skip_mutation_failures = skip_mutation_failures
+        self.backsolve_options = backsolve_options
         # -- Migration parameters
         self.migration = migration
         self.hof_migration = hof_migration
@@ -2190,6 +2201,7 @@ def _run(
             randomize=self.weight_randomize,
             do_nothing=self.weight_do_nothing,
             optimize=self.weight_optimize,
+            backsolve=self.weight_backsolve,
         )
 
         # Convert operators dict to Julia format and create OperatorEnum
@@ -2269,6 +2281,11 @@ def _run(
                 else len(X)
             ),
             mutation_weights=mutation_weights,
+            backsolve=(
+                self.backsolve_options.julia_options()
+                if self.backsolve_options is not None
+                else None
+            ),
             tournament_selection_p=self.tournament_selection_p,
             tournament_selection_n=self.tournament_selection_n,
             # These have the same name:

pysr/test/test_main.py

@@ -24,6 +24,7 @@
     estimator_checks_generator = functools.partial(check_estimator, generate_only=True)
 
 from pysr import (
+    BacksolveOptions,
     ParametricExpressionSpec,
     PySRRegressor,
     TemplateExpressionSpec,
@@ -160,6 +161,32 @@ def test_multiline_seval(self):
         """)
         self.assertEqual(num, 1.5)
 
+    def test_backsolve_options(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            population_size=5,
+            tournament_selection_n=2,
+            ncycles_per_iteration=1,
+            binary_operators=["+", "*"],
+            weight_backsolve=0.5,
+            backsolve_options=BacksolveOptions(
+                max_library_size=37,
+                lambda_=0.2,
+                max_iter=3,
+            ),
+            progress=False,
+            temp_equation_file=True,
+        )
+        model.fit(self.X[:, :1], self.X[:, 0])
+
+        self.assertEqual(model.julia_options_.mutation_weights.backsolve, 0.5)
+        self.assertEqual(model.julia_options_.backsolve.max_library_size, 37)
+        self.assertEqual(
+            jl.getproperty(model.julia_options_.backsolve, jl.Symbol("lambda")), 0.2
+        )
+        self.assertEqual(model.julia_options_.backsolve.max_iter, 3)
+
     def test_high_precision_search_custom_loss(self):
         y = 1.23456789 * self.X[:, 0]
         model = PySRRegressor(