Add way of limiting the complexity of power operators

julia/sr.jl

@@ -599,6 +599,27 @@ mutable struct PopMember
 
 end
 
+# Check if any power operator is to the power of a complex expression
+function deepPow(tree::Node)::Integer
+    if tree.degree == 0
+        return 0
+    elseif tree.degree == 1
+        return 0 + deepPow(tree.l)
+    else
+        if binops[tree.op] == pow
+            complexity_in_power = countNodes(tree.r)
+            is_deep_pow = (complexity_in_power > 1)
+            if is_deep_pow
+                return 1 + deepPow(tree.l)
+            else
+                return 0 + deepPow(tree.l)
+            end
+        else
+            return 0 + deepPow(tree.l) + deepPow(tree.r)
+        end
+    end
+end
+
 # Go through one simulated annealing mutation cycle
 #  exp(-delta/T) defines probability of accepting a change
 function iterate(member::PopMember, T::Float32, curmaxsize::Integer)::PopMember
@@ -652,6 +673,13 @@ function iterate(member::PopMember, T::Float32, curmaxsize::Integer)::PopMember
         return PopMember(tree, beforeLoss)
     end
 
+
+    # Check for illegal functions
+    if limitPowComplexity && (deepPow(tree) > 0)
+        return PopMember(copyNode(prev), beforeLoss)
+    end
+
+
     if batching
         afterLoss = scoreFuncBatch(tree)
     else

pysr/sr.py

@@ -87,6 +87,7 @@ def pysr(X=None, y=None, weights=None,
             batchSize=50,
             select_k_features=None,
             warmupMaxsize=0,
+            limitPowComplexity=False,
             threads=None, #deprecated
             julia_optimization=3,
         ):
@@ -163,6 +164,9 @@ def pysr(X=None, y=None, weights=None,
         a small number up to the maxsize (if greater than 0).
         If greater than 0, says how many cycles before the maxsize
         is increased.
+    :param limitPowComplexity: bool, whether to prevent pow from having
+        complex right arguments. I.e., 3.0^(x+y) becomes impossible,
+        but 3.0^x is possible.
     :param julia_optimization: int, Optimization level (0, 1, 2, 3)
     :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
         (as strings).
@@ -275,6 +279,7 @@ const mutationWeights = [
     {weightDoNothing:f}
 ]
 const warmupMaxsize = {warmupMaxsize:d}
+const limitPowComplexity = {"true" if limitPowComplexity else "false"}
 """
 
     op_runner = ""