Added multiprocessing experiment

pysr/Problems.py

@@ -89,20 +89,51 @@ class FeynmanProblem(Problem):
         return ret
 
 
-def run_on_problem(problem, verbosity=0):
+def run_on_problem(problem, verbosity=0, multiprocessing=True):
     """
     Takes in a problem and returns a tuple: (equations, best predicted equation, actual equation)
     """
     from time import time
     starting = time()
-    equations = pysr(problem.X, problem.y, variable_names=problem.variable_names, verbosity=verbosity)
+    equations = pysr(problem.X, problem.y, variable_names=problem.variable_names, verbosity=verbosity,)
     timing = time()-starting
-    others = {"equations": equations, "time": timing}
-    return best(equations), problem.form, others
+    others = {"time": timing, "problem": problem}
+    if not multiprocessing:
+        others['equations'] = equations
+    return str(best(equations)), problem.form, others
 
 
+def do_feynman_experiments_parallel(first=100, verbosity=0, dp=500, output_file_path="experiments/FeynmanExperiment.csv", data_dir="datasets/FeynmanEquations.csv"):
+    import multiprocessing as mp
+    from tqdm import tqdm
+    problems = FeynmanProblem.mk_problems(first=first, gen=True, dp=dp, data_dir=data_dir)
+    ids = []
+    predictions = []
+    true_equations = []
+    time_takens = []
+    pool = mp.Pool()
+    results = []
+    with tqdm(total=len(problems)) as pbar:
+        for i, res in enumerate(pool.imap(run_on_problem, problems)):
+            results.append(res)
+            pbar.update()
+    for res in results:
+        prediction, true_equation, others = res
+        problem = others['problem']
+        ids.append(problem.eq_id)
+        predictions.append(prediction)
+        true_equations.append(true_equation)
+        time_takens.append(others['time'])
+    with open(output_file_path, 'a') as f:
+        writer = csv.writer(f, delimiter=',')
+        writer.writerow(['ID', 'Predicted', 'True', 'Time'])
+        for i in range(len(ids)):
+            writer.writerow([ids[i], predictions[i], true_equations[i], time_takens[i]])
+    return
+
 def do_feynman_experiments(first=100, verbosity=0, dp=500, output_file_path="experiments/FeynmanExperiment.csv", data_dir="datasets/FeynmanEquations.csv"):
     from tqdm import tqdm
+
     problems = FeynmanProblem.mk_problems(first=first, gen=True, dp=dp, data_dir=data_dir)
     indx = range(len(problems))
     ids = []
@@ -116,7 +147,7 @@ def do_feynman_experiments(first=100, verbosity=0, dp=500, output_file_path="exp
         true_equations.append(true_equation)
         time_takens.append(others['time'])
     with open(output_file_path, 'a') as f:
-        writer = csv.writer(outcsv, delimiter=',')
+        writer = csv.writer(f, delimiter=',')
         writer.writerow(['ID', 'Predicted', 'True', 'Time'])
         for i in range(len(ids)):
             writer.writerow([ids[i], predictions[i], true_equations[i], time_takens[i]])
@@ -124,4 +155,4 @@ def do_feynman_experiments(first=100, verbosity=0, dp=500, output_file_path="exp
 
 
 if __name__ == "__main__":
-    do_feynman_experiments(first=4)
+    do_feynman_experiments_parallel(first=10)