Merge pull request #208 from MilesCranmer/update-backend

Speed up evaluation with `turbo` parameter

Author: MilesCranmer

pysr/sr.py

@@ -476,6 +476,11 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         algorithm than regularized evolution, but does cycles 15%
         faster. May be algorithmically less efficient.
         Default is `False`.
+    turbo: bool
+        (Experimental) Whether to use LoopVectorization.jl to speed up the
+        search evaluation. Certain operators may not be supported.
+        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
@@ -692,6 +697,7 @@ def __init__(
         batching=False,
         batch_size=50,
         fast_cycle=False,
+        turbo=False,
         precision=32,
         random_state=None,
         deterministic=False,
@@ -779,6 +785,7 @@ def __init__(
         self.batching = batching
         self.batch_size = batch_size
         self.fast_cycle = fast_cycle
+        self.turbo = turbo
         self.precision = precision
         self.random_state = random_state
         self.deterministic = deterministic
@@ -1518,25 +1525,22 @@ def _run(self, X, y, mutated_params, weights, seed):
             str(self.early_stop_condition) if self.early_stop_condition else None
         )
 
-        mutation_weights = np.array(
-            [
-                self.weight_mutate_constant,
-                self.weight_mutate_operator,
-                self.weight_add_node,
-                self.weight_insert_node,
-                self.weight_delete_node,
-                self.weight_simplify,
-                self.weight_randomize,
-                self.weight_do_nothing,
-            ],
-            dtype=float,
+        mutation_weights = SymbolicRegression.MutationWeights(
+            mutate_constant=self.weight_mutate_constant,
+            mutate_operator=self.weight_mutate_operator,
+            add_node=self.weight_add_node,
+            insert_node=self.weight_insert_node,
+            delete_node=self.weight_delete_node,
+            simplify=self.weight_simplify,
+            randomize=self.weight_randomize,
+            do_nothing=self.weight_do_nothing,
         )
 
         # Call to Julia backend.
         # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/OptionsStruct.jl
         options = SymbolicRegression.Options(
-            binary_operators=Main.eval(str(tuple(binary_operators)).replace("'", "")),
-            unary_operators=Main.eval(str(tuple(unary_operators)).replace("'", "")),
+            binary_operators=Main.eval(str(binary_operators).replace("'", "")),
+            unary_operators=Main.eval(str(unary_operators).replace("'", "")),
             bin_constraints=bin_constraints,
             una_constraints=una_constraints,
             complexity_of_operators=complexity_of_operators,
@@ -1545,45 +1549,47 @@ def _run(self, X, y, mutated_params, weights, seed):
             nested_constraints=nested_constraints,
             loss=custom_loss,
             maxsize=int(self.maxsize),
-            hofFile=_escape_filename(self.equation_file_),
+            output_file=_escape_filename(self.equation_file_),
             npopulations=int(self.populations),
             batching=self.batching,
-            batchSize=int(min([batch_size, len(X)]) if self.batching else len(X)),
-            mutationWeights=mutation_weights,
-            probPickFirst=self.tournament_selection_p,
-            ns=self.tournament_selection_n,
+            batch_size=int(min([batch_size, len(X)]) if self.batching else len(X)),
+            mutation_weights=mutation_weights,
+            tournament_selection_p=self.tournament_selection_p,
+            tournament_selection_n=self.tournament_selection_n,
             # These have the same name:
             parsimony=self.parsimony,
             alpha=self.alpha,
             maxdepth=maxdepth,
             fast_cycle=self.fast_cycle,
+            turbo=self.turbo,
             migration=self.migration,
-            hofMigration=self.hof_migration,
-            fractionReplacedHof=self.fraction_replaced_hof,
-            shouldOptimizeConstants=self.should_optimize_constants,
-            warmupMaxsizeBy=self.warmup_maxsize_by,
-            useFrequency=self.use_frequency,
-            useFrequencyInTournament=self.use_frequency_in_tournament,
+            hof_migration=self.hof_migration,
+            fraction_replaced_hof=self.fraction_replaced_hof,
+            should_optimize_constants=self.should_optimize_constants,
+            warmup_maxsize_by=self.warmup_maxsize_by,
+            use_frequency=self.use_frequency,
+            use_frequency_in_tournament=self.use_frequency_in_tournament,
             npop=self.population_size,
-            ncyclesperiteration=self.ncyclesperiteration,
-            fractionReplaced=self.fraction_replaced,
+            ncycles_per_iteration=self.ncyclesperiteration,
+            fraction_replaced=self.fraction_replaced,
             topn=self.topn,
             verbosity=self.verbosity,
             optimizer_algorithm=self.optimizer_algorithm,
             optimizer_nrestarts=self.optimizer_nrestarts,
-            optimize_probability=self.optimize_probability,
+            optimizer_probability=self.optimize_probability,
             optimizer_iterations=self.optimizer_iterations,
-            perturbationFactor=self.perturbation_factor,
+            perturbation_factor=self.perturbation_factor,
             annealing=self.annealing,
-            stateReturn=True,  # Required for state saving.
+            return_state=True,  # Required for state saving.
             progress=progress,
             timeout_in_seconds=self.timeout_in_seconds,
-            crossoverProbability=self.crossover_probability,
+            crossover_probability=self.crossover_probability,
             skip_mutation_failures=self.skip_mutation_failures,
             max_evals=self.max_evals,
-            earlyStopCondition=early_stop_condition,
+            early_stop_condition=early_stop_condition,
             seed=seed,
             deterministic=self.deterministic,
+            define_helper_functions=False,
         )
 
         # Convert data to desired precision
@@ -1603,7 +1609,16 @@ def _run(self, X, y, mutated_params, weights, seed):
         else:
             Main.weights = None
 
-        cprocs = 0 if multithreading else self.procs
+        if self.procs == 0 and not multithreading:
+            parallelism = "serial"
+        elif multithreading:
+            parallelism = "multithreading"
+        else:
+            parallelism = "multiprocessing"
+
+        cprocs = (
+            None if parallelism in ["serial", "multithreading"] else int(self.procs)
+        )
 
         # Call to Julia backend.
         # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl
@@ -1614,8 +1629,8 @@ def _run(self, X, y, mutated_params, weights, seed):
             niterations=int(self.niterations),
             varMap=self.feature_names_in_.tolist(),
             options=options,
-            numprocs=int(cprocs),
-            multithreading=bool(multithreading),
+            numprocs=cprocs,
+            parallelism=parallelism,
             saved_state=self.raw_julia_state_,
             addprocs_function=cluster_manager,
         )

pysr/version.py

@@ -1,2 +1,2 @@
-__version__ = "0.11.5"
-__symbolic_regression_jl_version__ = "0.12.6"
+__version__ = "0.11.6"
+__symbolic_regression_jl_version__ = "0.14.0"

test/test.py

@@ -70,12 +70,13 @@ def test_linear_relation_weighted(self):
         print(model.equations_)
         self.assertLessEqual(model.get_best()["loss"], 1e-4)
 
-    def test_multiprocessing(self):
+    def test_multiprocessing_turbo(self):
         y = self.X[:, 0]
         model = PySRRegressor(
             **self.default_test_kwargs,
             procs=2,
             multithreading=False,
+            turbo=True,
             early_stop_condition="stop_if(loss, complexity) = loss < 1e-4 && complexity == 1",
         )
         model.fit(self.X, y)
@@ -108,6 +109,8 @@ def test_multioutput_custom_operator_quiet_custom_complexity(self):
             verbosity=0,
             **self.default_test_kwargs,
             procs=0,
+            # Test custom operators with turbo:
+            turbo=True,
             # Test custom operators with constraints:
             nested_constraints={"square_op": {"square_op": 3}},
             constraints={"square_op": 10},