Add feature to set arbitrary variable names

README.md

@@ -304,14 +304,14 @@ pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
 - [x] Use @fastmath
 - [x] Try @spawn over each sub-population. Do random sort, compute mutation for each, then replace 10% oldest.
 - [x] Control max depth, rather than max number of nodes?
+- [x] Allow user to pass names for variables - use these when printing
 - [ ] Sort these todo lists by priority
 
 ## Feature ideas
 
 - [ ] Sympy printing
-- [ ] Allow user to pass names for variables - use these when printing
 - [ ] Hierarchical model, so can re-use functional forms. Output of one equation goes into second equation?
-- [ ] Call function to read from csv after running
+- [ ] Call function to read from csv after running, so dont need to run again
 - [ ] Add function to plot equations
 - [ ] Refresh screen rather than dumping to stdout?
 - [ ] Add ability to save state from python

julia/sr.jl

@@ -121,7 +121,11 @@ function stringTree(tree::Node)::String
         if tree.constant
             return string(tree.val)
         else
-            return "x$(tree.val - 1)"
+            if useVarMap
+                return varMap[tree.val]
+            else
+                return "x$(tree.val - 1)"
+            end
         end
     elseif tree.degree == 1
         return "$(unaops[tree.op])($(stringTree(tree.l)))"

pysr/sr.py

@@ -75,6 +75,7 @@ def pysr(X=None, y=None, weights=None,
             maxsize=20,
             fast_cycle=False,
             maxdepth=None,
+            variable_names=[],
             threads=None, #deprecated
             julia_optimization=3,
         ):
@@ -135,6 +136,8 @@ def pysr(X=None, y=None, weights=None,
     :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 variable_names: list, a list of names for the variables, other
+        than "x0", "x1", etc.
     :param julia_optimization: int, Optimization level (0, 1, 2, 3)
     :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
         (as strings).
@@ -153,6 +156,8 @@ def pysr(X=None, y=None, weights=None,
     if weights is not None:
         assert len(weights.shape) == 1
         assert X.shape[0] == weights.shape[0]
+    if len(variable_names) != 0:
+        assert len(variable_names) == X.shape[1]
 
     if populations is None:
         populations = procs
@@ -222,6 +227,7 @@ const nrestarts = {nrestarts:d}
 const perturbationFactor = {perturbationFactor:f}f0
 const annealing = {"true" if annealing else "false"}
 const weighted = {"true" if weights is not None else "false"}
+const useVarMap = {"false" if len(variable_names) == 0 else "true"}
 const mutationWeights = [
     {weightMutateConstant:f},
     {weightMutateOperator:f},
@@ -248,6 +254,10 @@ const y = convert(Array{Float32, 1}, """f"{y_str})"
         def_datasets += """
 const weights = convert(Array{Float32, 1}, """f"{weight_str})"
 
+    if len(variable_names) != 0:
+        def_hyperparams += f"""
+const varMap = {'["' + '", "'.join(variable_names) + '"]'}"""
+
     with open(f'/tmp/.hyperparams_{rand_string}.jl', 'w') as f:
         print(def_hyperparams, file=f)
 

setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="pysr", # Replace with your own username
-    version="0.3.18",
+    version="0.3.19",
     author="Miles Cranmer",
     author_email="miles.cranmer@gmail.com",
     description="Simple and efficient symbolic regression",