Add option to control max depth instead of size

README.md

@@ -320,6 +320,7 @@ pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
 - [ ] Create flexible way of providing "simplification recipes." I.e., plus(plus(T, C), C) => plus(T, +(C, C)). The user could pass these.
 - [ ] Consider allowing multi-threading turned off, for faster testing (cache issue on travis). Or could simply fix the caching issue there.
 - [ ] Consider returning only the equation of interest; rather than all equations.
+- [x] Control max depth, rather than max number of nodes?
 
 ## Algorithmic performance ideas:
 
@@ -332,15 +333,18 @@ pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
 - [ ] Calculate feature importances based on features we've already seen, then weight those features up in all random generations.
 - [ ] Calculate feature importances of future mutations, by looking at correlation between residual of model, and the features.
     - Store feature importances of future, and periodically update it.
+- [ ] Punish depth rather than size, as depth really hurts during optimization.
 
 
 ## Code performance ideas:
 
+- [ ] **Try @spawn over each sub-population. Do random sort, compute mutation for each, then replace 10% oldest.**
+- [ ] **Try defining a binary tree as an array, rather than a linked list. See https://stackoverflow.com/a/6384714/2689923**
 - [ ] Add true multi-node processing, with MPI, or just file sharing. Multiple populations per core.
     - Ongoing in cluster branch
-- [ ] Try @spawn over each sub-population. Do random sort, compute mutation for each, then replace 10% oldest.
 - [ ] Performance: try inling things?
-- [ ] Try defining a binary tree as an array, rather than a linked list. See https://stackoverflow.com/a/6384714/2689923
+- [ ] Try storing things like number nodes in a tree; then can iterate instead of counting
+
 ```julia
 mutable struct Tree
     degree::Array{Integer, 1}
@@ -350,8 +354,14 @@ mutable struct Tree
     Tree(s::Integer) = new(zeros(Integer, s), zeros(Float32, s), zeros(Bool, s), zeros(Integer, s))
 end
 ```
-    - Then, we could even work with trees on the GPU, since they are all pre-allocated arrays.
-    - A population could be a Tree, but with degree 2 on all the degrees. So a slice of population arrays forms a tree.
+
+- Then, we could even work with trees on the GPU, since they are all pre-allocated arrays.
+- A population could be a Tree, but with degree 2 on all the degrees. So a slice of population arrays forms a tree.
+- How many operations can we do via matrix ops? Mutate node=>easy.
+- Can probably batch and do many operations at once across a population.
+    - Or, across all populations! Mutate operator: index 2D array and set it to random vector? But the indexing might hurt.
+- The big advantage: can evaluate all new mutated trees at once; as massive matrix operation.
+    - Can control depth, rather than maxsize. Then just pretend all trees are full and same depth. Then we really don't need to care about depth.
 
 - [ ] Can we cache calculations, or does the compiler do that? E.g., I should only have to run exp(x0) once; after that it should be read from memory.
     - Done on caching branch. Currently am finding that this is quiet slow (presumably because memory allocation is the main issue).

benchmarks/benchmark.sh

@@ -6,7 +6,9 @@ import numpy as np
 from pysr import pysr
 X=np.random.randn(100, 2)*5
 y=2*np.sin((X[:, 0]+X[:, 1]))*np.exp(X[:, 1]/3)
-if version[1] >= 3 and version[2] >= 2:
+if version[1] >= 3 and version[2] >= 16:
+    eq = pysr(X, y, binary_operators=["plus", "mult", "div", "pow"], unary_operators=["sin"], niterations=20, procs=4, parsimony=1e-10, npop=1000, ncyclesperiteration=1000, maxdepth=6)
+elif version[1] >= 3 and version[2] >= 2:
     eq = pysr(X, y, binary_operators=["plus", "mult", "div", "pow"], unary_operators=["sin"], niterations=20, procs=4, parsimony=1e-10, npop=1000, ncyclesperiteration=1000)
 else:
     eq = pysr(X, y, binary_operators=["plus", "mult", "div", "pow"], unary_operators=["sin"], niterations=20, threads=4, parsimony=1e-10, npop=1000, ncyclesperiteration=1000)

julia/sr.jl

@@ -103,6 +103,17 @@ function countNodes(tree::Node)::Integer
     end
 end
 
+# Count the max depth of a tree
+function countDepth(tree::Node)::Integer
+    if tree.degree == 0
+        return 1
+    elseif tree.degree == 1
+        return 1 + countDepth(tree.l)
+    else
+        return 1 + max(countDepth(tree.l), countDepth(tree.r))
+    end
+end
+
 # Convert an equation to a string
 function stringTree(tree::Node)::String
     if tree.degree == 0
@@ -535,14 +546,15 @@ function iterate(member::PopMember, T::Float32)::PopMember
     cur_weights /= sum(cur_weights)
     cweights = cumsum(cur_weights)
     n = countNodes(tree)
+    depth = countDepth(tree)
 
     if mutationChoice < cweights[1]
         tree = mutateConstant(tree, T)
     elseif mutationChoice < cweights[2]
         tree = mutateOperator(tree)
-    elseif mutationChoice < cweights[3] && n < maxsize
+    elseif mutationChoice < cweights[3] && n < maxsize && depth < maxdepth
         tree = appendRandomOp(tree)
-    elseif mutationChoice < cweights[4] && n < maxsize
+    elseif mutationChoice < cweights[4] && n < maxsize && depth < maxdepth
         tree = insertRandomOp(tree)
     elseif mutationChoice < cweights[5]
         tree = deleteRandomOp(tree)
@@ -551,7 +563,7 @@ function iterate(member::PopMember, T::Float32)::PopMember
         tree = combineOperators(tree) # See if repeated constants at outer levels
         return PopMember(tree, beforeLoss)
     elseif mutationChoice < cweights[7]
-        tree = genRandomTree(5) # Sometimes we simplify tree
+        tree = genRandomTree(5) # Sometimes we generate a new tree completely tree
     else
         return PopMember(tree, beforeLoss)
     end

pysr/sr.py

@@ -73,6 +73,7 @@ def pysr(X=None, y=None, weights=None,
             test='simple1',
             verbosity=1e9,
             maxsize=20,
+            maxdepth=None,
             threads=None, #deprecated
             julia_optimization=3,
         ):
@@ -135,6 +136,8 @@ def pysr(X=None, y=None, weights=None,
     """
     if threads is not None:
         raise ValueError("The threads kwarg is deprecated. Use procs.")
+    if maxdepth is None:
+        maxdepth = maxsize
 
     # Check for potential errors before they happen
     assert len(unary_operators) + len(binary_operators) > 0
@@ -200,6 +203,7 @@ const ns=10;
 const parsimony = {parsimony:f}f0
 const alpha = {alpha:f}f0
 const maxsize = {maxsize:d}
+const maxdepth = {maxdepth:d}
 const migration = {'true' if migration else 'false'}
 const hofMigration = {'true' if hofMigration else 'false'}
 const fractionReplacedHof = {fractionReplacedHof}f0