Merge pull request #2 from DhananjayAshok/refactoring

Refactoring pf sr.py and integration with Miles's changes last night

pysr/sr.py

@@ -191,15 +191,9 @@ def pysr(X=None, y=None, weights=None,
         (as strings).
 
     """
-    if threads is not None:
-        raise ValueError("The threads kwarg is deprecated. Use procs.")
-    if limitPowComplexity:
-        raise ValueError("The limitPowComplexity kwarg is deprecated. Use constraints.")
-    if maxdepth is None:
-        maxdepth = maxsize
-    if equation_file is None:
-        date_time = datetime.now().strftime("%Y-%m-%d_%H%M%S.%f")[:-3]
-        equation_file = 'hall_of_fame_' + date_time + '.csv'
+    raise_depreciation_errors(limitPowComplexity, threads)
+    auxiliary_filename, X_filename, dataset_filename, hyperparam_filename, julia_auxiliary_filenames, operator_filename \
+        ,pkg_filename, runfile_filename, tmpdir, weights_filename, y_filename = set_paths(tempdir)
 
     if isinstance(X, pd.DataFrame):
         variable_names = list(X.columns)
@@ -210,128 +204,144 @@ def pysr(X=None, y=None, weights=None,
     if len(X.shape) == 1:
         X = X[:, None]
 
-    # Check for potential errors before they happen
-    assert len(unary_operators) + len(binary_operators) > 0
-    assert len(X.shape) == 2
-    assert len(y.shape) == 1
-    assert X.shape[0] == y.shape[0]
-    if weights is not None:
-        assert len(weights.shape) == 1
-        assert X.shape[0] == weights.shape[0]
-    if use_custom_variable_names:
-        assert len(variable_names) == X.shape[1]
+    check_assertions(X, binary_operators, unary_operators, use_custom_variable_names, variable_names, weights, y)
 
-    if select_k_features is not None:
-        selection = run_feature_selection(X, y, select_k_features)
-        print(f"Using features {selection}")
-        X = X[:, selection]
-
-        if use_custom_variable_names:
-            variable_names = [variable_names[selection[i]] for i in range(len(selection))]
+    X, variable_names = handle_feature_selection(X, select_k_features, use_custom_variable_names, variable_names, y)
 
+    if maxdepth is None:
+        maxdepth = maxsize
+    if equation_file is None:
+        date_time = datetime.now().strftime("%Y-%m-%d_%H%M%S.%f")[:-3]
+        equation_file = 'hall_of_fame_' + date_time + '.csv'
     if populations is None:
         populations = procs
-
-    if isinstance(binary_operators, str): binary_operators = [binary_operators]
-    if isinstance(unary_operators, str): unary_operators = [unary_operators]
-
+    if isinstance(binary_operators, str):
+        binary_operators = [binary_operators]
+    if isinstance(unary_operators, str):
+        unary_operators = [unary_operators]
     if X is None:
-        if test == 'simple1':
-            eval_str = "np.sign(X[:, 2])*np.abs(X[:, 2])**2.5 + 5*np.cos(X[:, 3]) - 5"
-        elif test == 'simple2':
-            eval_str = "np.sign(X[:, 2])*np.abs(X[:, 2])**3.5 + 1/(np.abs(X[:, 0])+1)"
-        elif test == 'simple3':
-            eval_str = "np.exp(X[:, 0]/2) + 12.0 + np.log(np.abs(X[:, 0])*10 + 1)"
-        elif test == 'simple4':
-            eval_str = "1.0 + 3*X[:, 0]**2 - 0.5*X[:, 0]**3 + 0.1*X[:, 0]**4"
-        elif test == 'simple5':
-            eval_str = "(np.exp(X[:, 3]) + 3)/(np.abs(X[:, 1]) + np.cos(X[:, 0]) + 1.1)"
-
-        X = np.random.randn(100, 5)*3
-        y = eval(eval_str)
-        print("Running on", eval_str)
-
-    # System-independent paths
-    pkg_directory = Path(__file__).parents[1] / 'julia'
-    pkg_filename = pkg_directory / "sr.jl"
-    operator_filename = pkg_directory / "Operators.jl"
-    julia_auxiliaries = [
-        "Equation.jl", "ProgramConstants.jl",
-        "LossFunctions.jl", "Utils.jl", "EvaluateEquation.jl",
-        "MutationFunctions.jl", "SimplifyEquation.jl", "PopMember.jl",
-        "HallOfFame.jl", "CheckConstraints.jl", "Mutate.jl",
-        "Population.jl", "RegularizedEvolution.jl", "SingleIteration.jl",
-        "ConstantOptimization.jl"
-    ]
-    julia_auxiliary_filenames = [
-        pkg_directory / fname
-        for fname in julia_auxiliaries
-    ]
-
-    tmpdir = Path(tempfile.mkdtemp(dir=tempdir))
-    hyperparam_filename = tmpdir / f'hyperparams.jl'
-    dataset_filename = tmpdir / f'dataset.jl'
-    auxiliary_filename = tmpdir / f'auxiliary.jl'
-    runfile_filename = tmpdir / f'runfile.jl'
-    X_filename = tmpdir / "X.csv"
-    y_filename = tmpdir / "y.csv"
-    weights_filename = tmpdir / "weights.csv"
+        X, y = using_test_input(X, test, y)
 
     def_hyperparams = ""
 
     # Add pre-defined functions to Julia
-    for op_list in [binary_operators, unary_operators]:
-        for i in range(len(op_list)):
-            op = op_list[i]
-            is_user_defined_operator = '(' in op
+    def_hyperparams = predefined_function_addition(binary_operators, def_hyperparams, unary_operators)
 
-            if is_user_defined_operator:
-                def_hyperparams += op + "\n"
-                # Cut off from the first non-alphanumeric char:
-                first_non_char = [
-                        j for j in range(len(op))
-                        if not (op[j].isalpha() or op[j].isdigit())][0]
-                function_name = op[:first_non_char]
-                op_list[i] = function_name
 
     #arbitrary complexity by default
-    for op in unary_operators:
-        if op not in constraints:
-            constraints[op] = -1
-    for op in binary_operators:
-        if op not in constraints:
-            constraints[op] = (-1, -1)
-        if op in ['plus', 'sub']:
-            if constraints[op][0] != constraints[op][1]:
-                raise NotImplementedError("You need equal constraints on both sides for - and *, due to simplification strategies.")
-        elif op == 'mult':
-            # Make sure the complex expression is in the left side.
-            if constraints[op][0] == -1:
-                continue
-            elif constraints[op][1] == -1 or constraints[op][0] < constraints[op][1]:
-                constraints[op][0], constraints[op][1] = constraints[op][1], constraints[op][0]
+    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 = ["""
 
-    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_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_auxiliary = make_auxiliary_julia_str(julia_auxiliary_filenames)
+
+    def_datasets = make_datasets_julia_str(X, X_filename, weights, weights_filename, y, y_filename)
 
+    create_julia_files(auxiliary_filename, dataset_filename, def_auxiliary, def_datasets, def_hyperparams, fractionReplaced, hyperparam_filename,
+                       ncyclesperiteration, niterations, npop, pkg_filename, runfile_filename, topn, verbosity)
+
+    final_pysr_process(julia_optimization, procs, runfile_filename, timeout)
+
+    set_globals(X, equation_file, extra_sympy_mappings, variable_names)
+
+    if delete_tempfiles:
+        shutil.rmtree(tmpdir)
+
+    return get_hof()
+
+
+def make_auxiliary_julia_str(julia_auxiliary_filenames):
+    def_auxiliary = '\n'.join([
+        f"""include("{_escape_filename(aux_fname)}")""" for aux_fname in julia_auxiliary_filenames
+    ])
+    return def_auxiliary
+
+
+def set_globals(X, equation_file, extra_sympy_mappings, variable_names):
+    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
+
+
+def final_pysr_process(julia_optimization, procs, runfile_filename, timeout):
+    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
+    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()
+
+
+def create_julia_files(auxiliary_filename, dataset_filename, def_auxiliary, 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(auxiliary_filename, 'w') as f:
+        print(def_auxiliary, 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(auxiliary_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)
+
+
+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) + ']'}
@@ -370,7 +380,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 += """
@@ -381,14 +390,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)
@@ -398,93 +406,155 @@ 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
+    if use_custom_variable_names:
+        def_hyperparams += f"""
+    const varMap = {'["' + '", "'.join(variable_names) + '"]'}"""
+    return def_hyperparams
 
-    def_auxiliary = '\n'.join([
-    f"""include("{_escape_filename(aux_fname)}")""" for aux_fname in julia_auxiliary_filenames
-    ])
-
-    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')"""
+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
 
-    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) + '"]'}"""
+def handle_constraints(binary_operators, constraints, unary_operators):
+    for op in unary_operators:
+        if op not in constraints:
+            constraints[op] = -1
+    for op in binary_operators:
+        if op not in constraints:
+            constraints[op] = (-1, -1)
+        if op in ['plus', 'sub']:
+            if constraints[op][0] != constraints[op][1]:
+                raise NotImplementedError(
+                    "You need equal constraints on both sides for - and *, due to simplification strategies.")
+        elif op == 'mult':
+            # Make sure the complex expression is in the left side.
+            if constraints[op][0] == -1:
+                continue
+            elif constraints[op][1] == -1 or constraints[op][0] < constraints[op][1]:
+                constraints[op][0], constraints[op][1] = constraints[op][1], constraints[op][0]
 
-    with open(hyperparam_filename, 'w') as f:
-        print(def_hyperparams, file=f)
 
-    with open(dataset_filename, 'w') as f:
-        print(def_datasets, file=f)
+def predefined_function_addition(binary_operators, def_hyperparams, unary_operators):
+    for op_list in [binary_operators, unary_operators]:
+        for i in range(len(op_list)):
+            op = op_list[i]
+            is_user_defined_operator = '(' in op
 
-    with open(auxiliary_filename, 'w') as f:
-        print(def_auxiliary, file=f)
+            if is_user_defined_operator:
+                def_hyperparams += op + "\n"
+                # Cut off from the first non-alphanumeric char:
+                first_non_char = [
+                    j for j in range(len(op))
+                    if not (op[j].isalpha() or op[j].isdigit())][0]
+                function_name = op[:first_non_char]
+                op_list[i] = function_name
+    return def_hyperparams
+
+
+def using_test_input(X, test, y):
+    if test == 'simple1':
+        eval_str = "np.sign(X[:, 2])*np.abs(X[:, 2])**2.5 + 5*np.cos(X[:, 3]) - 5"
+    elif test == 'simple2':
+        eval_str = "np.sign(X[:, 2])*np.abs(X[:, 2])**3.5 + 1/(np.abs(X[:, 0])+1)"
+    elif test == 'simple3':
+        eval_str = "np.exp(X[:, 0]/2) + 12.0 + np.log(np.abs(X[:, 0])*10 + 1)"
+    elif test == 'simple4':
+        eval_str = "1.0 + 3*X[:, 0]**2 - 0.5*X[:, 0]**3 + 0.1*X[:, 0]**4"
+    elif test == 'simple5':
+        eval_str = "(np.exp(X[:, 3]) + 3)/(np.abs(X[:, 1]) + np.cos(X[:, 0]) + 1.1)"
+    X = np.random.randn(100, 5) * 3
+    y = eval(eval_str)
+    print("Running on", eval_str)
+    return X, y
+
+
+def handle_feature_selection(X, select_k_features, use_custom_variable_names, variable_names, y):
+    if select_k_features is not None:
+        selection = run_feature_selection(X, y, select_k_features)
+        print(f"Using features {selection}")
+        X = X[:, selection]
 
-    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(auxiliary_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)
+        if use_custom_variable_names:
+            variable_names = [variable_names[selection[i]] for i in range(len(selection))]
+    return X, variable_names
 
 
-    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
+def set_paths(tempdir):
+    # System-independent paths
+    pkg_directory = Path(__file__).parents[1] / 'julia'
+    pkg_filename = pkg_directory / "sr.jl"
+    operator_filename = pkg_directory / "Operators.jl"
+    julia_auxiliaries = [
+        "Equation.jl", "ProgramConstants.jl",
+        "LossFunctions.jl", "Utils.jl", "EvaluateEquation.jl",
+        "MutationFunctions.jl", "SimplifyEquation.jl", "PopMember.jl",
+        "HallOfFame.jl", "CheckConstraints.jl", "Mutate.jl",
+        "Population.jl", "RegularizedEvolution.jl", "SingleIteration.jl",
+        "ConstantOptimization.jl"
+    ]
+    julia_auxiliary_filenames = [
+        pkg_directory / fname
+        for fname in julia_auxiliaries
+    ]
 
-    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
+    tmpdir = Path(tempfile.mkdtemp(dir=tempdir))
+    hyperparam_filename = tmpdir / f'hyperparams.jl'
+    dataset_filename = tmpdir / f'dataset.jl'
+    auxiliary_filename = tmpdir / f'auxiliary.jl'
+    runfile_filename = tmpdir / f'runfile.jl'
+    X_filename = tmpdir / "X.csv"
+    y_filename = tmpdir / "y.csv"
+    weights_filename = tmpdir / "weights.csv"
+    return auxiliary_filename, X_filename, dataset_filename, hyperparam_filename, julia_auxiliary_filenames, \
+           operator_filename, pkg_filename, runfile_filename, tmpdir, weights_filename, y_filename
 
-    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()
+def check_assertions(X, binary_operators, unary_operators, use_custom_variable_names, variable_names, weights, y):
+    # Check for potential errors before they happen
+    assert len(unary_operators) + len(binary_operators) > 0
+    assert len(X.shape) == 2
+    assert len(y.shape) == 1
+    assert X.shape[0] == y.shape[0]
+    if weights is not None:
+        assert len(weights.shape) == 1
+        assert X.shape[0] == weights.shape[0]
+    if use_custom_variable_names:
+        assert len(variable_names) == X.shape[1]
 
-    if delete_tempfiles:
-        shutil.rmtree(tmpdir)
 
-    return get_hof()
+def raise_depreciation_errors(limitPowComplexity, threads):
+    if threads is not None:
+        raise ValueError("The threads kwarg is deprecated. Use procs.")
+    if limitPowComplexity:
+        raise ValueError("The limitPowComplexity kwarg is deprecated. Use constraints.")
 
 
 def run_feature_selection(X, y, select_k_features):