Add optional threaded sub-samples

julia/sr.jl

@@ -1,5 +1,6 @@
 import Optim
 import Printf: @printf
+import Random: shuffle!
 
 const maxdegree = 2
 const actualMaxsize = maxsize + maxdegree
@@ -625,19 +626,49 @@ end
 
 # Mutate the best sampled member of the population
 function iterateSample(pop::Population, T::Float32)::PopMember
-    allstar = bestOfSample(pop)
-    return iterate(allstar, T)
 end
 
 # Pass through the population several times, replacing the oldest
 # with the fittest of a small subsample
 function regEvolCycle(pop::Population, T::Float32)::Population
-    for i=1:round(Integer, pop.n/ns)
-        baby = iterateSample(pop, T)
-        #printTree(baby.tree)
-        oldest = argmin([pop.members[member].birth for member=1:pop.n])
-        pop.members[oldest] = baby
+    # 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.
+    if fast_cycle:
+        shuffle!(pop.members)
+        n_evol_cycles = round(Integer, pop.n/ns)
+        babies = Array{PopMember}(undef, n_evol_cycles)
+
+        # Iterate each ns-member sub-sample
+        @inbounds Threads.@threads for i=1:n_evol_cycles
+            best_score = Inf32
+            best_idx = 1+(i-1)*ns
+            # Calculate best member of the subsample:
+            for sub_i=1+(i-1)*ns:i*ns
+                if pop.members[sub_i].score < best_score
+                    best_score = pop.members[sub_i].score
+                    best_idx = sub_i
+                end
+            end
+            allstar = pop.members[best_idx]
+            babies[i] = iterate(allstar, T)
+        end
+
+        # Replace the n_evol_cycles-oldest members of each population
+        @inbounds for i=1:n_evol_cycles
+            oldest = argmin([pop.members[member].birth for member=1:pop.n])
+            pop.members[oldest] = babies[i]
+        end
+    else
+        for i=1:round(Integer, pop.n/ns)
+            allstar = bestOfSample(pop)
+            baby = iterate(allstar, T)
+            #printTree(baby.tree)
+            oldest = argmin([pop.members[member].birth for member=1:pop.n])
+            pop.members[oldest] = baby
+        end
     end
+
     return pop
 end
 

pysr/sr.py

@@ -73,6 +73,7 @@ def pysr(X=None, y=None, weights=None,
             test='simple1',
             verbosity=1e9,
             maxsize=20,
+            fast_cycle=False,
             maxdepth=None,
             threads=None, #deprecated
             julia_optimization=3,
@@ -129,6 +130,11 @@ def pysr(X=None, y=None, weights=None,
     :param equation_file: str, Where to save the files (.csv separated by |)
     :param test: str, What test to run, if X,y not passed.
     :param maxsize: int, Max size of an equation.
+    :param maxdepth: int, Max depth of an equation. You can use both maxsize and maxdepth.
+        maxdepth is by default set to = maxsize, which means that it is redundant.
+    :param fast_cycle: bool, (experimental) - batch over population subsamples. This
+        is a slightly different algorithm than regularized evolution, but does cycles
+        15% faster. May be algorithmically less efficient.
     :param julia_optimization: int, Optimization level (0, 1, 2, 3)
     :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
         (as strings).
@@ -204,6 +210,7 @@ const parsimony = {parsimony:f}f0
 const alpha = {alpha:f}f0
 const maxsize = {maxsize:d}
 const maxdepth = {maxdepth:d}
+const fast_cycle = {'true' if fast_cycle else 'false'}
 const migration = {'true' if migration else 'false'}
 const hofMigration = {'true' if hofMigration else 'false'}
 const fractionReplacedHof = {fractionReplacedHof}f0