Merge pull request #44 from SymbolicML/dont-overwrite

Avoid method invalidation in helper functions

Author: MilesCranmer

src/Equation.jl

@@ -190,6 +190,7 @@ const OP_NAMES = Dict(
     "safe_pow" => "^",
 )
 
+get_op_name(op::String) = op
 @generated function get_op_name(op::F) where {F}
     try
         # Bit faster to just cache the name of the operator:
@@ -266,7 +267,7 @@ Convert an equation to a string.
 """
 function string_tree(
     tree::Node{T},
-    operators::AbstractOperatorEnum;
+    operators::Union{AbstractOperatorEnum,Nothing}=nothing;
     bracketed::Bool=false,
     f_variable::F1=string_variable,
     f_constant::F2=string_constant,
@@ -284,7 +285,11 @@ function string_tree(
     elseif tree.degree == 1
         return string_op(
             Val(1),
-            operators.unaops[tree.op],
+            if operators === nothing
+                "unary_operator[" * string(tree.op) * "]"
+            else
+                operators.unaops[tree.op]
+            end,
             tree,
             operators;
             bracketed,
@@ -295,7 +300,11 @@ function string_tree(
     else
         return string_op(
             Val(2),
-            operators.binops[tree.op],
+            if operators === nothing
+                "binary_operator[" * string(tree.op) * "]"
+            else
+                operators.binops[tree.op]
+            end,
             tree,
             operators;
             bracketed,

src/EvaluationHelpers.jl

@@ -54,13 +54,6 @@ function (tree::Node)(X, operators::GenericOperatorEnum; kws...)
     !did_finish && return nothing
     return out
 end
-function (tree::Node)(X; kws...)
-    ## This will be overwritten by OperatorEnumConstructionModule, and turned
-    ## into a depwarn.
-    return error(
-        "The `tree(X; kws...)` syntax is deprecated. Use `tree(X, operators; kws...)` instead.",
-    )
-end
 
 # Gradients:
 function _grad_evaluator(tree::Node, X, operators::OperatorEnum; variable=true, kws...)
@@ -73,13 +66,6 @@ end
 function _grad_evaluator(tree::Node, X, operators::GenericOperatorEnum; kws...)
     return error("Gradients are not implemented for `GenericOperatorEnum`.")
 end
-function _grad_evaluator(tree::Node, X; kws...)
-    ## This will be overwritten by OperatorEnumConstructionModule, and turned
-    ## into a depwarn
-    return error(
-        "The `tree'(X; kws...)` syntax is deprecated. Use `tree'(X, operators; kws...)` instead.",
-    )
-end
 
 """
     (tree::Node{T})'(X::AbstractMatrix{T}, operators::OperatorEnum; turbo::Bool=false, variable::Bool=true)

src/OperatorEnumConstruction.jl

@@ -6,82 +6,135 @@ import ..EvaluateEquationModule: eval_tree_array
 import ..EvaluateEquationDerivativeModule: eval_grad_tree_array, _zygote_gradient
 import ..EvaluationHelpersModule: _grad_evaluator
 
-function create_evaluation_helpers!(operators::OperatorEnum)
-    @eval begin
-        Base.print(io::IO, tree::Node) = print(io, string_tree(tree, $operators))
-        Base.show(io::IO, tree::Node) = print(io, string_tree(tree, $operators))
-        function (tree::Node)(X; kws...)
-            Base.depwarn(
-                "The `tree(X; kws...)` syntax is deprecated. Use `tree(X, operators; kws...)` instead.",
-                :Node,
-            )
-            return tree(X, $operators; kws...)
-        end
-        # Gradients:
-        function _grad_evaluator(tree::Node, X; kws...)
-            Base.depwarn(
-                "The `tree'(X; kws...)` syntax is deprecated. Use `tree'(X, operators; kws...)` instead.",
-                :Node,
-            )
-            return _grad_evaluator(tree, X, $operators; kws...)
-        end
+"""Used to set a default value for `operators` for ease of use."""
+@enum AvailableOperatorTypes begin
+    IsNothing
+    IsOperatorEnum
+    IsGenericOperatorEnum
+end
+
+# These constants are purely for convenience. Internal code
+# should make use of `Node`, `string_tree`, `eval_tree_array`,
+# and `eval_grad_tree_array` directly.
+
+const LATEST_OPERATORS = Ref{Union{Nothing,AbstractOperatorEnum}}(nothing)
+const LATEST_OPERATORS_TYPE = Ref{AvailableOperatorTypes}(IsNothing)
+const LATEST_UNARY_OPERATOR_MAPPING = Dict{Function,Int}()
+const LATEST_BINARY_OPERATOR_MAPPING = Dict{Function,Int}()
+const ALREADY_DEFINED_UNARY_OPERATORS = (;
+    operator_enum=Dict{Function,Bool}(), generic_operator_enum=Dict{Function,Bool}()
+)
+const ALREADY_DEFINED_BINARY_OPERATORS = (;
+    operator_enum=Dict{Function,Bool}(), generic_operator_enum=Dict{Function,Bool}()
+)
+
+function Base.show(io::IO, tree::Node)
+    latest_operators_type = LATEST_OPERATORS_TYPE.x
+    if latest_operators_type == IsNothing
+        return print(io, string_tree(tree))
+    elseif latest_operators_type == IsOperatorEnum
+        latest_operators = LATEST_OPERATORS.x::OperatorEnum
+        return print(io, string_tree(tree, latest_operators))
+    else
+        latest_operators = LATEST_OPERATORS.x::GenericOperatorEnum
+        return print(io, string_tree(tree, latest_operators))
     end
 end
+function (tree::Node)(X; kws...)
+    Base.depwarn(
+        "The `tree(X; kws...)` syntax is deprecated. Use `tree(X, operators; kws...)` instead.",
+        :Node,
+    )
+    latest_operators_type = LATEST_OPERATORS_TYPE.x
+    if latest_operators_type == IsNothing
+        error("Please use the `tree(X, operators; kws...)` syntax instead.")
+    elseif latest_operators_type == IsOperatorEnum
+        latest_operators = LATEST_OPERATORS.x::OperatorEnum
+        return tree(X, latest_operators; kws...)
+    else
+        latest_operators = LATEST_OPERATORS.x::GenericOperatorEnum
+        return tree(X, latest_operators; kws...)
+    end
+end
+
+function _grad_evaluator(tree::Node, X; kws...)
+    Base.depwarn(
+        "The `tree'(X; kws...)` syntax is deprecated. Use `tree'(X, operators; kws...)` instead.",
+        :Node,
+    )
+    latest_operators_type = LATEST_OPERATORS_TYPE.x
+    # return _grad_evaluator(tree, X, $operators; kws...)
+    if latest_operators_type == IsNothing
+        error("Please use the `tree'(X, operators; kws...)` syntax instead.")
+    elseif latest_operators_type == IsOperatorEnum
+        latest_operators = LATEST_OPERATORS.x::OperatorEnum
+        return _grad_evaluator(tree, X, latest_operators; kws...)
+    else
+        error("Gradients are not implemented for `GenericOperatorEnum`.")
+    end
+end
+
+function create_evaluation_helpers!(operators::OperatorEnum)
+    LATEST_OPERATORS.x = operators
+    return LATEST_OPERATORS_TYPE.x = IsOperatorEnum
+end
 
 function create_evaluation_helpers!(operators::GenericOperatorEnum)
-    @eval begin
-        Base.print(io::IO, tree::Node) = print(io, string_tree(tree, $operators))
-        Base.show(io::IO, tree::Node) = print(io, string_tree(tree, $operators))
-
-        function (tree::Node)(X; kws...)
-            Base.depwarn(
-                "The `tree(X; kws...)` syntax is deprecated. Use `tree(X, operators; kws...)` instead.",
-                :Node,
-            )
-            return tree(X, $operators; kws...)
-        end
-        function _grad_evaluator(tree::Node, X; kws...)
-            return error("Gradients are not implemented for `GenericOperatorEnum`.")
-        end
+    LATEST_OPERATORS.x = operators
+    return LATEST_OPERATORS_TYPE.x = IsGenericOperatorEnum
+end
+function lookup_op(@nospecialize(f), ::Val{degree}) where {degree}
+    mapping = degree == 1 ? LATEST_UNARY_OPERATOR_MAPPING : LATEST_BINARY_OPERATOR_MAPPING
+    if !haskey(mapping, f)
+        error(
+            "Convenience constructor for `Node` using operator `$(f)` is out-of-date. " *
+            "Please create an `OperatorEnum` (or `GenericOperatorEnum`) with " *
+            "`define_helper_functions=true` and pass `$(f)`.",
+        )
     end
+    return mapping[f]
 end
 
-function _extend_unary_operator(f::Symbol, op, type_requirements)
+function _extend_unary_operator(f::Symbol, type_requirements)
     quote
         quote
             function $($f)(l::Node{T})::Node{T} where {T<:$($type_requirements)}
                 return if (l.degree == 0 && l.constant)
                     Node(T; val=$($f)(l.val::T))
                 else
-                    Node($($op), l)
+                    latest_op_idx = $($lookup_op)($($f), Val(1))
+                    Node(latest_op_idx, l)
                 end
             end
         end
     end
 end
 
-function _extend_binary_operator(f::Symbol, op, type_requirements, build_converters)
+function _extend_binary_operator(f::Symbol, type_requirements, build_converters)
     quote
         quote
             function $($f)(l::Node{T}, r::Node{T}) where {T<:$($type_requirements)}
                 if (l.degree == 0 && l.constant && r.degree == 0 && r.constant)
                     Node(T; val=$($f)(l.val::T, r.val::T))
                 else
-                    Node($($op), l, r)
+                    latest_op_idx = $($lookup_op)($($f), Val(2))
+                    Node(latest_op_idx, l, r)
                 end
             end
             function $($f)(l::Node{T}, r::T) where {T<:$($type_requirements)}
                 if l.degree == 0 && l.constant
                     Node(T; val=$($f)(l.val::T, r))
                 else
-                    Node($($op), l, Node(T; val=r))
+                    latest_op_idx = $($lookup_op)($($f), Val(2))
+                    Node(latest_op_idx, l, Node(T; val=r))
                 end
             end
             function $($f)(l::T, r::Node{T}) where {T<:$($type_requirements)}
                 if r.degree == 0 && r.constant
                     Node(T; val=$($f)(l, r.val::T))
                 else
-                    Node($($op), Node(T; val=l), r)
+                    latest_op_idx = $($lookup_op)($($f), Val(2))
+                    Node(latest_op_idx, Node(T; val=l), r)
                 end
             end
             if $($build_converters)
@@ -116,37 +169,62 @@ function _extend_binary_operator(f::Symbol, op, type_requirements, build_convert
 end
 
 function _extend_operators(operators, skip_user_operators, __module__::Module)
-    binary_ex = _extend_binary_operator(:f, :op, :type_requirements, :build_converters)
-    unary_ex = _extend_unary_operator(:f, :op, :type_requirements)
+    binary_ex = _extend_binary_operator(:f, :type_requirements, :build_converters)
+    unary_ex = _extend_unary_operator(:f, :type_requirements)
     return quote
         local type_requirements
         local build_converters
+        local binary_exists
+        local unary_exists
         if isa($operators, OperatorEnum)
             type_requirements = Number
             build_converters = true
+            binary_exists = $(ALREADY_DEFINED_BINARY_OPERATORS).operator_enum
+            unary_exists = $(ALREADY_DEFINED_UNARY_OPERATORS).operator_enum
         else
             type_requirements = Any
             build_converters = false
+            binary_exists = $(ALREADY_DEFINED_BINARY_OPERATORS).generic_operator_enum
+            unary_exists = $(ALREADY_DEFINED_UNARY_OPERATORS).generic_operator_enum
         end
-        for (op, f) in enumerate(map(Symbol, $(operators).binops))
+        # Trigger errors if operators are not yet defined:
+        empty!($(LATEST_BINARY_OPERATOR_MAPPING))
+        empty!($(LATEST_UNARY_OPERATOR_MAPPING))
+        for (op, func) in enumerate($(operators).binops)
+            local f = Symbol(func)
+            local skip = false
             if isdefined(Base, f)
                 f = :(Base.$(f))
             elseif $(skip_user_operators)
-                continue
+                skip = true
             else
                 f = :($($__module__).$(f))
             end
-            eval($binary_ex)
+            $(LATEST_BINARY_OPERATOR_MAPPING)[func] = op
+            skip && continue
+            # Avoid redefining methods:
+            if !haskey(unary_exists, func)
+                eval($binary_ex)
+                unary_exists[func] = true
+            end
         end
-        for (op, f) in enumerate(map(Symbol, $(operators).unaops))
+        for (op, func) in enumerate($(operators).unaops)
+            local f = Symbol(func)
+            local skip = false
             if isdefined(Base, f)
                 f = :(Base.$(f))
             elseif $(skip_user_operators)
-                continue
+                skip = true
             else
                 f = :($($__module__).$(f))
             end
-            eval($unary_ex)
+            $(LATEST_UNARY_OPERATOR_MAPPING)[func] = op
+            skip && continue
+            # Avoid redefining methods:
+            if !haskey(binary_exists, func)
+                eval($unary_ex)
+                binary_exists[func] = true
+            end
         end
     end
 end
@@ -162,14 +240,16 @@ apply this macro to the operator enum in the same module you have the operators
 defined.
 """
 macro extend_operators(operators)
-    ex = _extend_operators(esc(operators), false, __module__)
+    ex = _extend_operators(operators, false, __module__)
     expected_type = AbstractOperatorEnum
-    quote
-        if !isa($(esc(operators)), $expected_type)
-            error("You must pass an operator enum to `@extend_operators`.")
-        end
-        $ex
-    end
+    return esc(
+        quote
+            if !isa($(operators), $expected_type)
+                error("You must pass an operator enum to `@extend_operators`.")
+            end
+            $ex
+        end,
+    )
 end
 
 """
@@ -179,14 +259,16 @@ Similar to `@extend_operators`, but only extends operators already
 defined in `Base`.
 """
 macro extend_operators_base(operators)
-    ex = _extend_operators(esc(operators), true, __module__)
+    ex = _extend_operators(operators, true, __module__)
     expected_type = AbstractOperatorEnum
-    quote
-        if !isa($(esc(operators)), $expected_type)
-            error("You must pass an operator enum to `@extend_operators_base`.")
-        end
-        $ex
-    end
+    return esc(
+        quote
+            if !isa($(operators), $expected_type)
+                error("You must pass an operator enum to `@extend_operators_base`.")
+            end
+            $ex
+        end,
+    )
 end
 
 """

test/test_derivatives.jl

@@ -4,11 +4,11 @@ using DynamicExpressions: eval_diff_tree_array, eval_grad_tree_array
 using Random
 using Zygote
 using LinearAlgebra
+include("test_params.jl")
 
 seed = 0
 # SIMD doesn't like abs(x) ^ y for some reason.
 pow_abs2(x, y) = exp(y * log(abs(x)))
-custom_cos(x) = cos(x)^2
 
 equation1(x1, x2, x3) = x1 + x2 + x3 + 3.2
 equation2(x1, x2, x3) = pow_abs2(x1, x2) + x3 + custom_cos(1.0 + x3) + 3.0 / x1