Refactored till file creation

pysr/sr.py

@@ -195,7 +195,6 @@ def pysr(X=None, y=None, weights=None,
     X_filename, dataset_filename, hyperparam_filename, operator_filename, pkg_filename, runfile_filename, tmpdir, \
     weights_filename, y_filename = set_paths(tempdir)
 
-
     if isinstance(X, pd.DataFrame):
         variable_names = list(X.columns)
         X = np.array(X)
@@ -231,27 +230,99 @@ def pysr(X=None, y=None, weights=None,
     #arbitrary complexity by default
     handle_constraints(binary_operators, constraints, unary_operators)
 
-    constraints_str = "const una_constraints = ["
-    first = True
-    for op in unary_operators:
-        val = constraints[op]
-        if not first:
-            constraints_str += ", "
-        constraints_str += f"{val:d}"
-        first = False
+    constraints_str = make_constraints_str(binary_operators, constraints, unary_operators)
 
-    constraints_str += """]
-const bin_constraints = ["""
+    def_hyperparams = make_hyperparams_julia_str(X, alpha, annealing, batchSize, batching, binary_operators,
+                                                 constraints_str, def_hyperparams, equation_file, fast_cycle,
+                                                 fractionReplacedHof, hofMigration, limitPowComplexity, maxdepth,
+                                                 maxsize, migration, nrestarts, operator_filename, parsimony,
+                                                 perturbationFactor, populations, procs, shouldOptimizeConstants,
+                                                 unary_operators, useFrequency, use_custom_variable_names, variable_names,
+                                                 warmupMaxsize, weightAddNode, weightDeleteNode, weightDoNothing,
+                                                 weightInsertNode, weightMutateConstant, weightMutateOperator,
+                                                 weightRandomize, weightSimplify, weights)
+
+    def_datasets = make_datasets_julia_str(X, X_filename, weights, weights_filename, y, y_filename)
+
+    create_julia_files(dataset_filename, def_datasets, def_hyperparams, fractionReplaced, hyperparam_filename,
+                       ncyclesperiteration, niterations, npop, pkg_filename, runfile_filename, topn, verbosity)
+
+    command = [
+        f'julia', f'-O{julia_optimization:d}',
+        f'-p', f'{procs}',
+        str(runfile_filename),
+        ]
+    if timeout is not None:
+        command = [f'timeout', f'{timeout}'] + command
+
+    global global_n_features
+    global global_equation_file
+    global global_variable_names
+    global global_extra_sympy_mappings
+
+    global_n_features = X.shape[1]
+    global_equation_file = equation_file
+    global_variable_names = variable_names
+    global_extra_sympy_mappings = extra_sympy_mappings
+
+    print("Running on", ' '.join(command))
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=1, shell=True)
+    try:
+        while True:
+            line = process.stdout.readline()
+            if not line: break
+            print(line.decode('utf-8').replace('\n', ''))
+
+        process.stdout.close()
+        process.wait()
+    except KeyboardInterrupt:
+        print("Killing process... will return when done.")
+        process.kill()
+
+    if delete_tempfiles:
+        shutil.rmtree(tmpdir)
+
+    return get_hof()
+
+
+def create_julia_files(dataset_filename, def_datasets, def_hyperparams, fractionReplaced, hyperparam_filename,
+                       ncyclesperiteration, niterations, npop, pkg_filename, runfile_filename, topn, verbosity):
+    with open(hyperparam_filename, 'w') as f:
+        print(def_hyperparams, file=f)
+    with open(dataset_filename, 'w') as f:
+        print(def_datasets, file=f)
+    with open(runfile_filename, 'w') as f:
+        print(f'@everywhere include("{_escape_filename(hyperparam_filename)}")', file=f)
+        print(f'@everywhere include("{_escape_filename(dataset_filename)}")', file=f)
+        print(f'@everywhere include("{_escape_filename(pkg_filename)}")', file=f)
+        print(
+            f'fullRun({niterations:d}, npop={npop:d}, ncyclesperiteration={ncyclesperiteration:d}, fractionReplaced={fractionReplaced:f}f0, verbosity=round(Int32, {verbosity:f}), topn={topn:d})',
+            file=f)
+        print(f'rmprocs(nprocs)', file=f)
 
-    first = True
-    for op in binary_operators:
-        tup = constraints[op]
-        if not first:
-            constraints_str += ", "
-        constraints_str += f"({tup[0]:d}, {tup[1]:d})"
-        first = False
-    constraints_str += "]"
 
+def make_datasets_julia_str(X, X_filename, weights, weights_filename, y, y_filename):
+    def_datasets = """using DelimitedFiles"""
+    np.savetxt(X_filename, X, delimiter=',')
+    np.savetxt(y_filename, y, delimiter=',')
+    if weights is not None:
+        np.savetxt(weights_filename, weights, delimiter=',')
+    def_datasets += f"""
+const X = readdlm("{_escape_filename(X_filename)}", ',', Float32, '\\n')
+const y = readdlm("{_escape_filename(y_filename)}", ',', Float32, '\\n')"""
+    if weights is not None:
+        def_datasets += f"""
+const weights = readdlm("{_escape_filename(weights_filename)}", ',', Float32, '\\n')"""
+    return def_datasets
+
+
+def make_hyperparams_julia_str(X, alpha, annealing, batchSize, batching, binary_operators, constraints_str,
+                               def_hyperparams, equation_file, fast_cycle, fractionReplacedHof, hofMigration,
+                               limitPowComplexity, maxdepth, maxsize, migration, nrestarts, operator_filename,
+                               parsimony, perturbationFactor, populations, procs, shouldOptimizeConstants,
+                               unary_operators, useFrequency, use_custom_variable_names, variable_names, warmupMaxsize, weightAddNode,
+                               weightDeleteNode, weightDoNothing, weightInsertNode, weightMutateConstant,
+                               weightMutateOperator, weightRandomize, weightSimplify, weights):
     def_hyperparams += f"""include("{_escape_filename(operator_filename)}")
 {constraints_str}
 const binops = {'[' + ', '.join(binary_operators) + ']'}
@@ -290,7 +361,6 @@ const warmupMaxsize = {warmupMaxsize:d}
 const limitPowComplexity = {"true" if limitPowComplexity else "false"}
 const useFrequency = {"true" if useFrequency else "false"}
 """
-
     op_runner = ""
     if len(binary_operators) > 0:
         op_runner += """
@@ -301,14 +371,13 @@ const useFrequency = {"true" if useFrequency else "false"}
         end"""
         for i in range(1, len(binary_operators)):
             op_runner += f"""
-    elseif i === {i+1}
+    elseif i === {i + 1}
         @inbounds @simd for j=1:clen
             x[j] = {binary_operators[i]}(x[j], y[j])
         end"""
         op_runner += """
     end
 end"""
-
     if len(unary_operators) > 0:
         op_runner += """
 @inline function UNAOP!(x::Array{Float32, 1}, i::Int, clen::Int)
@@ -318,85 +387,40 @@ end"""
         end"""
         for i in range(1, len(unary_operators)):
             op_runner += f"""
-    elseif i === {i+1}
+    elseif i === {i + 1}
         @inbounds @simd for j=1:clen
             x[j] = {unary_operators[i]}(x[j])
         end"""
         op_runner += """
     end
 end"""
-
     def_hyperparams += op_runner
-
-    def_datasets = """using DelimitedFiles"""
-
-    np.savetxt(X_filename, X, delimiter=',')
-    np.savetxt(y_filename, y, delimiter=',')
-    if weights is not None:
-        np.savetxt(weights_filename, weights, delimiter=',')
-
-    def_datasets += f"""
-const X = readdlm("{_escape_filename(X_filename)}", ',', Float32, '\\n')
-const y = readdlm("{_escape_filename(y_filename)}", ',', Float32, '\\n')"""
-
-    if weights is not None:
-        def_datasets += f"""
-const weights = readdlm("{_escape_filename(weights_filename)}", ',', Float32, '\\n')"""
-
     if use_custom_variable_names:
         def_hyperparams += f"""
-const varMap = {'["' + '", "'.join(variable_names) + '"]'}"""
-
-    with open(hyperparam_filename, 'w') as f:
-        print(def_hyperparams, file=f)
-
-    with open(dataset_filename, 'w') as f:
-        print(def_datasets, file=f)
-
-    with open(runfile_filename, 'w') as f:
-        print(f'@everywhere include("{_escape_filename(hyperparam_filename)}")', file=f)
-        print(f'@everywhere include("{_escape_filename(dataset_filename)}")', file=f)
-        print(f'@everywhere include("{_escape_filename(pkg_filename)}")', file=f)
-        print(f'fullRun({niterations:d}, npop={npop:d}, ncyclesperiteration={ncyclesperiteration:d}, fractionReplaced={fractionReplaced:f}f0, verbosity=round(Int32, {verbosity:f}), topn={topn:d})', file=f)
-        print(f'rmprocs(nprocs)', file=f)
-
-
-    command = [
-        f'julia', f'-O{julia_optimization:d}',
-        f'-p', f'{procs}',
-        str(runfile_filename),
-        ]
-    if timeout is not None:
-        command = [f'timeout', f'{timeout}'] + command
-
-    global global_n_features
-    global global_equation_file
-    global global_variable_names
-    global global_extra_sympy_mappings
-
-    global_n_features = X.shape[1]
-    global_equation_file = equation_file
-    global_variable_names = variable_names
-    global_extra_sympy_mappings = extra_sympy_mappings
-
-    print("Running on", ' '.join(command))
-    process = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=1, shell=True)
-    try:
-        while True:
-            line = process.stdout.readline()
-            if not line: break
-            print(line.decode('utf-8').replace('\n', ''))
-
-        process.stdout.close()
-        process.wait()
-    except KeyboardInterrupt:
-        print("Killing process... will return when done.")
-        process.kill()
+    const varMap = {'["' + '", "'.join(variable_names) + '"]'}"""
+    return def_hyperparams
 
-    if delete_tempfiles:
-        shutil.rmtree(tmpdir)
 
-    return get_hof()
+def make_constraints_str(binary_operators, constraints, unary_operators):
+    constraints_str = "const una_constraints = ["
+    first = True
+    for op in unary_operators:
+        val = constraints[op]
+        if not first:
+            constraints_str += ", "
+        constraints_str += f"{val:d}"
+        first = False
+    constraints_str += """]
+const bin_constraints = ["""
+    first = True
+    for op in binary_operators:
+        tup = constraints[op]
+        if not first:
+            constraints_str += ", "
+        constraints_str += f"({tup[0]:d}, {tup[1]:d})"
+        first = False
+    constraints_str += "]"
+    return constraints_str
 
 
 def handle_constraints(binary_operators, constraints, unary_operators):