Add option to slowly increase maxsize

TODO.md

@@ -56,6 +56,7 @@
 - [x] Create backup csv file so always something to copy from for `PySR`. Also use random hall of fame file by default. Call function to read from csv after running, so dont need to run again. Dump scores alongside MSE to .csv (and return with Pandas).
 - [x] Better cleanup of zombie processes after <ctl-c>
 - [x] Consider printing output sorted by score, not by complexity.
+- [x] Increase max complexity slowly over time up to the actual max.
 - [ ] Sort these todo lists by priority
 
 ## Feature ideas

docs/options.md

@@ -107,6 +107,10 @@ constants, variables). `maxdepth` is by default not used, but can be set
 to control the maximum depth of an equation. These will make processing
 faster, as longer equations take longer to test.
 
+One can warm up the maxsize from a small number to encourage
+PySR to start simple, by using the `warmupMaxsize` argument.
+This specifies that maxsize increases every `warmupMaxsize`.
+
 
 ## Batching
 One can turn on mini-batching, with the `batching` flag,

julia/sr.jl

@@ -597,7 +597,7 @@ end
 
 # Go through one simulated annealing mutation cycle
 #  exp(-delta/T) defines probability of accepting a change
-function iterate(member::PopMember, T::Float32)::PopMember
+function iterate(member::PopMember, T::Float32, curmaxsize::Integer)::PopMember
     prev = member.tree
     tree = copyNode(prev)
     #TODO - reconsider this
@@ -610,23 +610,31 @@ function iterate(member::PopMember, T::Float32)::PopMember
     mutationChoice = rand()
     weightAdjustmentMutateConstant = min(8, countConstants(tree))/8.0
     cur_weights = copy(mutationWeights) .* 1.0
+    #More constants => more likely to do constant mutation
     cur_weights[1] *= weightAdjustmentMutateConstant
-    cur_weights /= sum(cur_weights)
-    cweights = cumsum(cur_weights)
     n = countNodes(tree)
     depth = countDepth(tree)
 
+    # If equation too big, don't add new operators
+    if n >= curmaxsize || depth >= maxdepth
+        cur_weights[3] = 0.0
+        cur_weights[4] = 0.0
+    end
+
+    cur_weights /= sum(cur_weights)
+    cweights = cumsum(cur_weights)
+
     if mutationChoice < cweights[1]
         tree = mutateConstant(tree, T)
     elseif mutationChoice < cweights[2]
         tree = mutateOperator(tree)
-    elseif mutationChoice < cweights[3] && n < maxsize && depth < maxdepth
+    elseif mutationChoice < cweights[3]
         if rand() < 0.5
             tree = appendRandomOp(tree)
         else
             tree = prependRandomOp(tree)
         end
-    elseif mutationChoice < cweights[4] && n < maxsize && depth < maxdepth
+    elseif mutationChoice < cweights[4]
         tree = insertRandomOp(tree)
     elseif mutationChoice < cweights[5]
         tree = deleteRandomOp(tree)
@@ -711,7 +719,7 @@ end
 
 # Pass through the population several times, replacing the oldest
 # with the fittest of a small subsample
-function regEvolCycle(pop::Population, T::Float32)::Population
+function regEvolCycle(pop::Population, T::Float32, curmaxsize::Integer)::Population
     # Batch over each subsample. Can give 15% improvement in speed; probably moreso for large pops.
     # but is ultimately a different algorithm than regularized evolution, and might not be
     # as good.
@@ -732,7 +740,7 @@ function regEvolCycle(pop::Population, T::Float32)::Population
                 end
             end
             allstar = pop.members[best_idx]
-            babies[i] = iterate(allstar, T)
+            babies[i] = iterate(allstar, T, curmaxsize)
         end
 
         # Replace the n_evol_cycles-oldest members of each population
@@ -743,7 +751,7 @@ function regEvolCycle(pop::Population, T::Float32)::Population
     else
         for i=1:round(Integer, pop.n/ns)
             allstar = bestOfSample(pop)
-            baby = iterate(allstar, T)
+            baby = iterate(allstar, T, curmaxsize)
             #printTree(baby.tree)
             oldest = argmin([pop.members[member].birth for member=1:pop.n])
             pop.members[oldest] = baby
@@ -757,16 +765,17 @@ end
 # printing the fittest equation every 10% through
 function run(
         pop::Population,
-        ncycles::Integer;
+        ncycles::Integer,
+        curmaxsize::Integer;
         verbosity::Integer=0
-        )::Population
+       )::Population
 
     allT = LinRange(1.0f0, 0.0f0, ncycles)
     for iT in 1:size(allT)[1]
         if annealing
-            pop = regEvolCycle(pop, allT[iT])
+            pop = regEvolCycle(pop, allT[iT], curmaxsize)
         else
-            pop = regEvolCycle(pop, 1.0f0)
+            pop = regEvolCycle(pop, 1.0f0, curmaxsize)
         end
 
         if verbosity > 0 && (iT % verbosity == 0)
@@ -909,6 +918,10 @@ function fullRun(niterations::Integer;
     channels = [RemoteChannel(1) for j=1:npopulations]
     bestSubPops = [Population(1) for j=1:npopulations]
     hallOfFame = HallOfFame()
+    curmaxsize = 3
+    if warmupMaxsize == 0
+        curmaxsize = maxsize
+    end
 
     for i=1:npopulations
         future = @spawnat :any Population(npop, 3)
@@ -917,10 +930,11 @@ function fullRun(niterations::Integer;
 
     # # 2. Start the cycle on every process:
     @sync for i=1:npopulations
-        @async allPops[i] = @spawnat :any run(fetch(allPops[i]), ncyclesperiteration, verbosity=verbosity)
+        @async allPops[i] = @spawnat :any run(fetch(allPops[i]), ncyclesperiteration, curmaxsize, verbosity=verbosity)
     end
     println("Started!")
     cycles_complete = npopulations * niterations
+    curmaxsize += 1
 
     last_print_time = time()
     num_equations = 0.0
@@ -1006,7 +1020,7 @@ function fullRun(niterations::Integer;
 
                 @async begin
                     allPops[i] = @spawnat :any let
-                        tmp_pop = run(cur_pop, ncyclesperiteration, verbosity=verbosity)
+                        tmp_pop = run(cur_pop, ncyclesperiteration, curmaxsize, verbosity=verbosity)
                         @inbounds @simd for j=1:tmp_pop.n
                             if rand() < 0.1
                                 tmp_pop.members[j].tree = simplifyTree(tmp_pop.members[j].tree)
@@ -1027,6 +1041,13 @@ function fullRun(niterations::Integer;
                 end
 
                 cycles_complete -= 1
+                cycles_elapsed = npopulations * niterations - cycles_complete
+                if warmupMaxsize != 0 && cycles_elapsed % warmupMaxsize == 0
+                    curmaxsize += 1
+                    if curmaxsize > maxsize
+                        curmaxsize = maxsize
+                    end
+                end
                 num_equations += ncyclesperiteration * npop / 10.0
             end
         end

pysr/sr.py

@@ -85,6 +85,7 @@ def pysr(X=None, y=None, weights=None,
             batching=False,
             batchSize=50,
             select_k_features=None,
+            warmupMaxsize=0,
             threads=None, #deprecated
             julia_optimization=3,
         ):
@@ -157,6 +158,10 @@ def pysr(X=None, y=None, weights=None,
         Python using random forests, before passing to the symbolic regression
         code. None means no feature selection; an int means select that many
         features.
+    :param warmupMaxsize: int, whether to slowly increase max size from
+        a small number up to the maxsize (if greater than 0).
+        If greater than 0, says how many cycles before the maxsize
+        is increased.
     :param julia_optimization: int, Optimization level (0, 1, 2, 3)
     :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
         (as strings).
@@ -268,6 +273,7 @@ const mutationWeights = [
     {weightRandomize:f},
     {weightDoNothing:f}
 ]
+const warmupMaxsize = {warmupMaxsize:d}
     """
 
     if X.shape[1] == 1: