Allow for more populations than processes

julia/sr.jl

@@ -748,28 +748,27 @@ function fullRun(niterations::Integer;
                )
     # 1. Start a population on every process
     allPops = Future[]
-    bestSubPops = [Population(1) for j=1:nprocs]
+    bestSubPops = [Population(1) for j=1:npopulations]
     hallOfFame = HallOfFame()
 
-    for i=1:nprocs
-        npop=300
-        future = @spawnat :any Population(npop, 3)
+    for i=1:npopulations
+        future = @spawn Population(npop, 3)
         push!(allPops, future)
     end
 
     # # 2. Start the cycle on every process:
-    for i=1:nprocs
-        allPops[i] = @spawnat :any run(fetch(allPops[i]), ncyclesperiteration, verbosity=verbosity)
+    @sync for i=1:npopulations
+        @async allPops[i] = @spawn run(fetch(allPops[i]), ncyclesperiteration, verbosity=verbosity)
     end
     println("Started!")
-    cycles_complete = nprocs * niterations
+    cycles_complete = npopulations * niterations
 
     last_print_time = time()
     num_equations = 0.0
     print_every_n_seconds = 1
 
     while cycles_complete > 0
-        for i=1:nprocs
+        for i=1:npopulations
             if isready(allPops[i])
                 cur_pop = fetch(allPops[i])
                 bestSubPops[i] = bestSubPop(cur_pop, topn=topn)
@@ -828,7 +827,7 @@ function fullRun(niterations::Integer;
                     end
                 end
 
-                allPops[i] = @spawnat :any let
+                @async allPops[i] = @spawn let
                     tmp_pop = run(cur_pop, ncyclesperiteration, verbosity=verbosity)
                     for j=1:tmp_pop.n
                         if rand() < 0.1

pysr/sr.py

@@ -7,6 +7,7 @@ import pandas as pd
 
 def pysr(X=None, y=None, weights=None,
             procs=4,
+            populations=None,
             niterations=100,
             ncyclesperiteration=300,
             binary_operators=["plus", "mult"],
@@ -49,6 +50,7 @@ def pysr(X=None, y=None, weights=None,
     :param weights: np.ndarray, 1D array. Each row is how to weight the
         mean-square-error loss on weights.
     :param procs: int, Number of processes (=number of populations running).
+    :param populations: int, Number of populations running; by default=procs.
     :param niterations: int, Number of iterations of the algorithm to run. The best
         equations are printed, and migrate between populations, at the
         end of each.
@@ -108,6 +110,8 @@ def pysr(X=None, y=None, weights=None,
         assert len(weights.shape) == 1
         assert X.shape[0] == weights.shape[0]
 
+    if populations is None:
+        populations = procs
 
     rand_string = f'{"".join([str(np.random.rand())[2] for i in range(20)])}'
 
@@ -162,6 +166,7 @@ const fractionReplacedHof = {fractionReplacedHof}f0
 const shouldOptimizeConstants = {'true' if shouldOptimizeConstants else 'false'}
 const hofFile = "{equation_file}"
 const nprocs = {procs:d}
+const npopulations = {populations:d}
 const nrestarts = {nrestarts:d}
 const perturbationFactor = {perturbationFactor:f}f0
 const annealing = {"true" if annealing else "false"}