Clean up API; make const optimization random

julia/sr.jl

@@ -684,18 +684,20 @@ function optimizeConstants(member::PopMember)::PopMember
     x0 = getConstants(member.tree)
     f(x::Array{Float32,1})::Float32 = optFunc(x, member.tree)
     if size(x0)[1] == 1
-        result = Optim.optimize(f, x0, Optim.Newton(), Optim.Options(iterations=20))
+        algorithm = Optim.Newton
     else
-        result = Optim.optimize(f, x0, Optim.NelderMead(), Optim.Options(iterations=100))
+        algorithm = Optim.NelderMead
+    end
 
-        # Try other initial conditions:
-        for i=1:5
-            tmpresult = Optim.optimize(f, x0 .* (1f0 .+ 5f-1*randn(Float32, size(x0)[1])), Optim.NelderMead(), Optim.Options(iterations=100))
-            if tmpresult.minimum < result.minimum
-                result = tmpresult
-            end
+    result = Optim.optimize(f, x0, algorithm(), Optim.Options(iterations=100))
+    # Try other initial conditions:
+    for i=1:nrestarts
+        tmpresult = Optim.optimize(f, x0 .* (1f0 .+ 5f-1*randn(Float32, size(x0)[1])), algorithm(), Optim.Options(iterations=100))
+        if tmpresult.minimum < result.minimum
+            result = tmpresult
         end
     end
+
     if Optim.converged(result)
         setConstants(member.tree, result.minimizer)
         member.score = convert(Float32, result.minimum)
@@ -736,14 +738,16 @@ function fullRun(niterations::Integer;
         # Spawn threads to run indepdent evolutions, then gather them
         @inbounds Threads.@threads for i=1:nthreads
             allPops[i] = run(allPops[i], ncyclesperiteration, verbosity=verbosity)
-            bestSubPops[i] = bestSubPop(allPops[i], topn=topn)
-            for j=1:bestSubPops[i].n
-                bestSubPops[i].members[j].tree = simplifyTree(bestSubPops[i].members[j].tree)
-                bestSubPops[i].members[j].tree = combineOperators(bestSubPops[i].members[j].tree)
-                if shouldOptimizeConstants
-                    bestSubPops[i].members[j] = optimizeConstants(bestSubPops[i].members[j])
+            for j=1:allPops[i].n
+                if rand() < 0.1
+                    allPops[i].members[j].tree = simplifyTree(allPops[i].members[j].tree)
+                    allPops[i].members[j].tree = combineOperators(allPops[i].members[j].tree)
+                    if shouldOptimizeConstants
+                        allPops[i].members[j] = optimizeConstants(allPops[i].members[j])
+                    end
                 end
             end
+            bestSubPops[i] = bestSubPop(allPops[i], topn=topn)
         end
 
         # Get best 10 models from each evolution. Copy because we re-assign later.

pysr/sr.py

@@ -5,48 +5,31 @@ import pathlib
 import numpy as np
 import pandas as pd
 
-# Dumped from hyperparam optimization
-default_alpha =                  1.0#0.1
-default_fractionReplaced =      0.10
-default_fractionReplacedHof =   0.10
-default_npop =                  1000
-default_weightAddNode =            1
-default_weightInsertNode =         3
-default_weightDeleteNode =         3
-default_weightMutateConstant =    10
-default_weightMutateOperator =     1
-default_weightRandomize =          1
-default_weightSimplify =        0.01
-default_weightDoNothing =          1
-default_result =                   1
-default_topn =                    10
-default_parsimony =              1e-4
-default_perturbationFactor =     1.0
-
 def pysr(X=None, y=None, threads=4,
             niterations=100,
             ncyclesperiteration=300,
             binary_operators=["plus", "mult"],
             unary_operators=["cos", "exp", "sin"],
-            alpha=default_alpha,
+            alpha=0.1,
             annealing=True,
-            fractionReplaced=default_fractionReplaced,
-            fractionReplacedHof=default_fractionReplacedHof,
-            npop=int(default_npop),
-            parsimony=default_parsimony,
+            fractionReplaced=0.10,
+            fractionReplacedHof=0.10,
+            npop=1000,
+            parsimony=1e-4,
             migration=True,
             hofMigration=True,
             shouldOptimizeConstants=True,
-            topn=int(default_topn),
-            weightAddNode=default_weightAddNode,
-            weightInsertNode=default_weightInsertNode,
-            weightDeleteNode=default_weightDeleteNode,
-            weightDoNothing=default_weightDoNothing,
-            weightMutateConstant=default_weightMutateConstant,
-            weightMutateOperator=default_weightMutateOperator,
-            weightRandomize=default_weightRandomize,
-            weightSimplify=default_weightSimplify,
-            perturbationFactor=default_perturbationFactor,
+            topn=10,
+            weightAddNode=1,
+            weightInsertNode=3,
+            weightDeleteNode=3,
+            weightDoNothing=1,
+            weightMutateConstant=10,
+            weightMutateOperator=1,
+            weightRandomize=1,
+            weightSimplify=0.01,
+            perturbationFactor=1.0,
+            nrestarts=3,
             timeout=None,
             equation_file='hall_of_fame.csv',
             test='simple1',
@@ -84,6 +67,7 @@ def pysr(X=None, y=None, threads=4,
     :param shouldOptimizeConstants: bool, Whether to numerically optimize
         constants (Nelder-Mead/Newton) at the end of each iteration.
     :param topn: int, How many top individuals migrate from each population.
+    :param nrestarts: int, Number of times to restart the constant optimizer
     :param weightAddNode: float, Relative likelihood for mutation to add a node
     :param weightInsertNode: float, Relative likelihood for mutation to insert a node
     :param weightDeleteNode: float, Relative likelihood for mutation to delete a node
@@ -141,6 +125,7 @@ const fractionReplacedHof = {fractionReplacedHof}f0
 const shouldOptimizeConstants = {'true' if shouldOptimizeConstants else 'false'}
 const hofFile = "{equation_file}"
 const nthreads = {threads:d}
+const nrestarts = {nrestarts:d}
 const perturbationFactor = {perturbationFactor:f}f0
 const annealing = {"true" if annealing else "false"}
 const mutationWeights = [