Add feature for operator-level size constraints

docs/options.md

@@ -14,7 +14,7 @@ may find useful include:
 - `maxsize`, `maxdepth`
 - `batching`, `batchSize`
 - `variable_names` (or pandas input)
-- Limiting pow complexity
+- Constraining operator complexity
 - LaTeX, SymPy, and callable equation output
 
 These are described below
@@ -129,13 +129,19 @@ alphabetical characters and `_` are used in these names.
 
 ## Limiting pow complexity
 
-One can limit the complexity of power laws
-with the `limitPowComplexity` argument.
-This will prevent the exponent part of the expression
-having complexity greater than one. This prevents uninterpretable
-expressions such as `x^(y+z^(5+y))`, which sometimes
-occur during training. But it still allows for, e.g., `5^y` or
-`(x+y)^5`: anything where the exponent is not overly complex.
+One can limit the complexity of specific operators with the `constraints` parameter.
+There is a "maxsize" parameter to PySR, but there is also an operator-level
+"constraints" parameter. One supplies a dict, like so:
+
+```python
+constraints={'pow': (-1, 1), 'mult': (3, 3), 'cos': 5}
+```
+
+What this says is that: a power law x^y can have an expression of arbitrary (-1) complexity in the x, but only complexity 1 (e.g., a constant or variable) in the y. So (x0 + 3)^5.5 is allowed, but 5.5^(x0 + 3) is not.
+I find this helps a lot for getting more interpretable equations.
+The other terms say that each multiplication can only have sub-expressions
+of up to complexity 3 (e.g., 5.0 + x2) in each side, and cosine can only operate on
+expressions of complexity 5 (e.g., 5.0 + x2*exp(x3)).
 
 ## LaTeX, SymPy, callables
 

julia/sr.jl

@@ -646,24 +646,46 @@ mutable struct PopMember
 
 end
 
-# Check if any power operator is to the power of a complex expression
-function deepPow(tree::Node)::Integer
+# Check if any binary operator are overly complex
+function flagBinOperatorComplexity(tree::Node, op::Int)::Bool
     if tree.degree == 0
-        return 0
+        return false
     elseif tree.degree == 1
-        return 0 + deepPow(tree.l)
+        return flagBinOperatorComplexity(tree.l, op)
     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)
+        if tree.op == op
+            overly_complex = (
+                    ((bin_constraints[op][1] > -1) &&
+                     (countNodes(tree.l) > bin_constraints[op][1]))
+                      ||
+                    ((bin_constraints[op][2] > -1) &&
+                     (countNodes(tree.r) > bin_constraints[op][2]))
+                )
+            if overly_complex
+                return true
+            end
+        end
+        return (flagBinOperatorComplexity(tree.l, op) || flagBinOperatorComplexity(tree.r, op))
+    end
+end
+
+# Check if any unary operators are overly complex
+function flagUnaOperatorComplexity(tree::Node, op::Int)::Bool
+    if tree.degree == 0
+        return false
+    elseif tree.degree == 1
+        if tree.op == op
+            overly_complex = (
+                      (una_constraints[op] > -1) &&
+                      (countNodes(tree.l) > una_constraints[op])
+                )
+            if overly_complex
+                return true
             end
-        else
-            return 0 + deepPow(tree.l) + deepPow(tree.r)
         end
+        return flagUnaOperatorComplexity(tree.l, op)
+    else
+        return (flagUnaOperatorComplexity(tree.l, op) || flagUnaOperatorComplexity(tree.r, op))
     end
 end
 
@@ -671,61 +693,104 @@ end
 #  exp(-delta/T) defines probability of accepting a change
 function iterate(member::PopMember, T::Float32, curmaxsize::Integer)::PopMember
     prev = member.tree
-    tree = copyNode(prev)
+    tree = prev
     #TODO - reconsider this
     if batching
-        beforeLoss = scoreFuncBatch(member.tree)
+        beforeLoss = scoreFuncBatch(prev)
     else
         beforeLoss = member.score
     end
 
     mutationChoice = rand()
-    weightAdjustmentMutateConstant = min(8, countConstants(tree))/8.0
-    cur_weights = copy(mutationWeights) .* 1.0
     #More constants => more likely to do constant mutation
+    weightAdjustmentMutateConstant = min(8, countConstants(prev))/8.0
+    cur_weights = copy(mutationWeights) .* 1.0
     cur_weights[1] *= weightAdjustmentMutateConstant
-    n = countNodes(tree)
-    depth = countDepth(tree)
+    n = countNodes(prev)
+    depth = countDepth(prev)
 
     # If equation too big, don't add new operators
     if n >= curmaxsize || depth >= maxdepth
         cur_weights[3] = 0.0
         cur_weights[4] = 0.0
     end
-
     cur_weights /= sum(cur_weights)
     cweights = cumsum(cur_weights)
 
-    if mutationChoice < cweights[1]
-        tree = mutateConstant(tree, T)
-    elseif mutationChoice < cweights[2]
-        tree = mutateOperator(tree)
-    elseif mutationChoice < cweights[3]
-        if rand() < 0.5
-            tree = appendRandomOp(tree)
-        else
-            tree = prependRandomOp(tree)
+    successful_mutation = false
+    #TODO: Currently we dont take this \/ into account
+    is_success_always_possible = true
+    attempts = 0
+    max_attempts = 10
+    
+    #############################################
+    # Mutations
+    #############################################
+    while (!successful_mutation) && attempts < max_attempts
+        tree = copyNode(prev)
+        successful_mutation = true
+        if mutationChoice < cweights[1]
+            tree = mutateConstant(tree, T)
+
+            is_success_always_possible = true
+            # Mutating a constant shouldn't invalidate an already-valid function
+
+        elseif mutationChoice < cweights[2]
+            tree = mutateOperator(tree)
+
+            is_success_always_possible = true
+            # Can always mutate to the same operator
+
+        elseif mutationChoice < cweights[3]
+            if rand() < 0.5
+                tree = appendRandomOp(tree)
+            else
+                tree = prependRandomOp(tree)
+            end
+            is_success_always_possible = false
+            # Can potentially have a situation without success
+        elseif mutationChoice < cweights[4]
+            tree = insertRandomOp(tree)
+            is_success_always_possible = false
+        elseif mutationChoice < cweights[5]
+            tree = deleteRandomOp(tree)
+            is_success_always_possible = true
+        elseif mutationChoice < cweights[6]
+            tree = simplifyTree(tree) # Sometimes we simplify tree
+            tree = combineOperators(tree) # See if repeated constants at outer levels
+            return PopMember(tree, beforeLoss)
+
+            is_success_always_possible = true
+            # Simplification shouldn't hurt complexity; unless some non-symmetric constraint
+            # to commutative operator...
+
+        elseif mutationChoice < cweights[7]
+            tree = genRandomTree(5) # Sometimes we generate a new tree completely tree
+
+            is_success_always_possible = true
+        else # no mutation applied
+            return PopMember(tree, beforeLoss)
         end
-    elseif mutationChoice < cweights[4]
-        tree = insertRandomOp(tree)
-    elseif mutationChoice < cweights[5]
-        tree = deleteRandomOp(tree)
-    elseif mutationChoice < cweights[6]
-        tree = simplifyTree(tree) # Sometimes we simplify tree
-        tree = combineOperators(tree) # See if repeated constants at outer levels
-        return PopMember(tree, beforeLoss)
-    elseif mutationChoice < cweights[7]
-        tree = genRandomTree(5) # Sometimes we generate a new tree completely tree
-    else
-        return PopMember(tree, beforeLoss)
-    end
 
+        # Check for illegal equations
+        for i=1:nbin
+            if successful_mutation && flagBinOperatorComplexity(tree, i)
+                successful_mutation = false
+            end
+        end
+        for i=1:nuna
+            if successful_mutation && flagUnaOperatorComplexity(tree, i)
+                successful_mutation = false
+            end
+        end
 
-    # Check for illegal functions
-    if limitPowComplexity && (deepPow(tree) > 0)
-        return PopMember(copyNode(prev), beforeLoss)
+        attempts += 1
     end
+    #############################################
 
+    if !successful_mutation
+        return PopMember(copyNode(prev), beforeLoss)
+    end
 
     if batching
         afterLoss = scoreFuncBatch(tree)

pysr/sr.py

@@ -89,7 +89,8 @@ def pysr(X=None, y=None, weights=None,
             batchSize=50,
             select_k_features=None,
             warmupMaxsize=0,
-            limitPowComplexity=False,
+            constraints={},
+            limitPowComplexity=False, #deprecated
             threads=None, #deprecated
             julia_optimization=3,
         ):
@@ -166,9 +167,11 @@ 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 constraints: dict of int (unary) or 2-tuples (binary),
+        this enforces maxsize constraints on the individual
+        arguments of operators. E.g., `'pow': (-1, 1)`
+        says that power laws can have any complexity left argument, but only
+        1 complexity exponent. Use this to force more interpretable solutions.
     :param julia_optimization: int, Optimization level (0, 1, 2, 3)
     :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
         (as strings).
@@ -176,6 +179,8 @@ def pysr(X=None, y=None, weights=None,
     """
     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
 
@@ -207,6 +212,17 @@ def pysr(X=None, y=None, weights=None,
     if populations is None:
         populations = procs
 
+    #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 ['mult', 'plus', 'sub']:
+            if constraints[op][0] != constraints[op][1]:
+                raise NotImplementedError("You need equal constraints on both sides for +, -, and *, due to simplification strategies.")
+
     rand_string = f'{"".join([str(np.random.rand())[2] for i in range(20)])}'
 
     if isinstance(binary_operators, str): binary_operators = [binary_operators]
@@ -247,7 +263,30 @@ def pysr(X=None, y=None, weights=None,
                 function_name = op[:first_non_char]
                 op_list[i] = function_name
 
+    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 += "]"
+
+
     def_hyperparams += f"""include("{pkg_directory}/operators.jl")
+{constraints_str}
 const binops = {'[' + ', '.join(binary_operators) + ']'}
 const unaops = {'[' + ', '.join(unary_operators) + ']'}
 const ns=10;