wip: working printing for assignment operator

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -12,14 +12,14 @@ using DispatchDoctor: @stable, @unstable
     include("NodePreallocation.jl")
     include("Strings.jl")
     include("Evaluate.jl")
-    include("SpecialOperators.jl")
     include("EvaluateDerivative.jl")
     include("ChainRules.jl")
     include("EvaluationHelpers.jl")
     include("Simplify.jl")
     include("OperatorEnumConstruction.jl")
     include("Expression.jl")
     include("ExpressionAlgebra.jl")
+    include("SpecialOperators.jl")
     include("Random.jl")
     include("Parse.jl")
     include("ParametricExpression.jl")

src/Evaluate.jl

@@ -10,6 +10,14 @@ import ..NodeUtilsModule: is_constant
 import ..ExtensionInterfaceModule: bumper_eval_tree_array, _is_loopvectorization_loaded
 import ..ValueInterfaceModule: is_valid, is_valid_array
 
+# Overloaded by SpecialOperators.jl:
+function any_special_operators(_)
+    return false
+end
+function special_operator end
+function deg2_eval_special end
+function deg1_eval_special end
+
 const OPERATOR_LIMIT_BEFORE_SLOWDOWN = 15
 
 macro return_on_nonfinite_val(eval_options, val, X)
@@ -268,7 +276,7 @@ function _eval_tree_array(
     # we can just return the constant result.
     if tree.degree == 0
         return deg0_eval(tree, cX, eval_options)
-    elseif is_constant(tree)
+    elseif !any_special_operators(operators) && is_constant(tree)
         # Speed hack for constant trees.
         const_result = dispatch_constant_tree(tree, operators)::ResultOk{T}
         !const_result.ok &&
@@ -330,6 +338,7 @@ end
     eval_options::EvalOptions,
 ) where {T}
     nbin = get_nbin(operators)
+    special_operators = any_special_operators(operators)
     long_compilation_time = nbin > OPERATOR_LIMIT_BEFORE_SLOWDOWN
     if long_compilation_time
         return quote
@@ -352,15 +361,15 @@ end
             i -> let op = operators.binops[i]
                 if special_operator(op)
                     deg2_eval_special(tree, cX, op, eval_options)
-                elseif tree.l.degree == 0 && tree.r.degree == 0
+                elseif !$(special_operators) && tree.l.degree == 0 && tree.r.degree == 0
                     deg2_l0_r0_eval(tree, cX, op, eval_options)
-                elseif tree.r.degree == 0
+                elseif !$(special_operators) && tree.r.degree == 0
                     result_l = _eval_tree_array(tree.l, cX, operators, eval_options)
                     !result_l.ok && return result_l
                     @return_on_nonfinite_array(eval_options, result_l.x)
                     # op(x, y), where y is a constant or variable but x is not.
                     deg2_r0_eval(tree, result_l.x, cX, op, eval_options)
-                elseif tree.l.degree == 0
+                elseif !$(special_operators) && tree.l.degree == 0
                     result_r = _eval_tree_array(tree.r, cX, operators, eval_options)
                     !result_r.ok && return result_r
                     @return_on_nonfinite_array(eval_options, result_r.x)
@@ -393,7 +402,8 @@ end
     if long_compilation_time
         return quote
             op = operators.unaops[op_idx]
-            special_operator(op) && return deg1_eval_special(tree, cX, op, eval_options)
+            special_operator(op) &&
+                return deg1_eval_special(tree, cX, op, eval_options, operators)
             result = _eval_tree_array(tree.l, cX, operators, eval_options)
             !result.ok && return result
             @return_on_nonfinite_array(eval_options, result.x)
@@ -408,8 +418,8 @@ end
             i -> i == op_idx,
             i -> let op = operators.unaops[i]
                 if special_operator(op)
-                    deg1_eval_special(tree, cX, op, eval_options)
-                elseif !special_operators &&
+                    deg1_eval_special(tree, cX, op, eval_options, operators)
+                elseif !$(special_operators) &&
                     tree.l.degree == 2 &&
                     tree.l.l.degree == 0 &&
                     tree.l.r.degree == 0
@@ -418,7 +428,7 @@ end
                     dispatch_deg1_l2_ll0_lr0_eval(
                         tree, cX, op, l_op_idx, operators.binops, eval_options
                     )
-                elseif !special_operators && tree.l.degree == 1 && tree.l.l.degree == 0
+                elseif !$(special_operators) && tree.l.degree == 1 && tree.l.l.degree == 0
                     # op(op2(x)), where x is a constant or variable.
                     l_op_idx = tree.l.op
                     dispatch_deg1_l1_ll0_eval(
@@ -941,10 +951,4 @@ end
     end
 end
 
-# Overloaded by SpecialOperators.jl:
-function any_special_operators end
-function special_operator end
-function deg2_eval_special end
-function deg1_eval_special end
-
 end

src/SpecialOperators.jl

@@ -2,28 +2,32 @@ module SpecialOperatorsModule
 
 using ..OperatorEnumModule: OperatorEnum
 using ..EvaluateModule: _eval_tree_array, @return_on_nonfinite_array, deg2_eval
+using ..ExpressionModule: AbstractExpression
+using ..ExpressionAlgebraModule: @declare_expression_operator
 
 import ..EvaluateModule:
     special_operator, deg2_eval_special, deg1_eval_special, any_special_operators
+import ..StringsModule: get_op_name
 
-function any_special_operators(::Type{OperatorEnum{B,U}}) where {B,U}
+function any_special_operators(::Union{O,Type{O}}) where {B,U,O<:OperatorEnum{B,U}}
     return any(special_operator, B.types) || any(special_operator, U.types)
 end
 
 # Use this to customize evaluation behavior for operators:
 @inline special_operator(::Type) = false
 @inline special_operator(f) = special_operator(typeof(f))
 
-# Base.@kwdef struct WhileOperator <: Function
-#     max_iters::Int = 100
-# end
 Base.@kwdef struct AssignOperator <: Function
     target_register::Int
 end
-
-# @inline special_operator(::Type{WhileOperator}) = true
+@declare_expression_operator((op::AssignOperator), 1)
 @inline special_operator(::Type{AssignOperator}) = true
+get_op_name(o::AssignOperator) = "[{FEATURE_" * string(o.target_register) * "} =]"
 
+# Base.@kwdef struct WhileOperator <: Function
+#     max_iters::Int = 100
+# end
+# @inline special_operator(::Type{WhileOperator}) = true
 # function deg2_eval_special(tree, cX, op::WhileOperator, eval_options)
 #     cond = tree.l
 #     body = tree.r
@@ -43,7 +47,7 @@ end
 # end
 # TODO: Need to void any instance of buffer when using while loop.
 
-function deg1_eval_special(tree, cX, op::AssignOperator, eval_options)
+function deg1_eval_special(tree, cX, op::AssignOperator, eval_options, operators)
     result = _eval_tree_array(tree.l, cX, operators, eval_options)
     !result.ok && return result
     @return_on_nonfinite_array(eval_options, result.x)
@@ -54,4 +58,4 @@ function deg1_eval_special(tree, cX, op::AssignOperator, eval_options)
     return result
 end
 
-end 
+end

src/Strings.jl

@@ -56,6 +56,18 @@ end
     end
 end
 
+const FEATURE_PLACEHOLDER_FIRST_HALF_LENGTH = length("{FEATURE_")
+function replace_feature_placeholders(s::String, f_variable::Function, variable_names)
+    return replace(
+        s,
+        r"\{FEATURE_(\d+)\}" =>
+            m -> f_variable(
+                parse(Int, m[(begin + FEATURE_PLACEHOLDER_FIRST_HALF_LENGTH):(end - 1)]),
+                variable_names,
+            ),
+    )
+end
+
 # Can overload these for custom behavior:
 needs_brackets(val::Real) = false
 needs_brackets(val::AbstractArray) = false
@@ -179,7 +191,9 @@ function string_tree(
             c
         end,
     )
-    return String(strip_brackets(raw_output))
+    string_output = String(strip_brackets(raw_output))
+    string_output = replace_feature_placeholders(string_output, f_variable, variable_names)
+    return string_output
 end
 
 # Print an equation

test/test_special_operators.jl

@@ -0,0 +1,43 @@
+using TestItems: @testitem
+
+@testitem "AssignOperator basic functionality" begin
+    using DynamicExpressions
+    using DynamicExpressions.SpecialOperatorsModule: AssignOperator
+    using DynamicExpressions.EvaluateModule: eval_tree_array
+    using Test
+    using Random
+
+    # Define operators and variable names
+    assign_op2 = AssignOperator(; target_register=2)
+    operators = OperatorEnum(;
+        binary_operators=[+, -, *, /], unary_operators=[sin, cos, assign_op2]
+    )
+    variable_names = ["x1", "x2", "x3", "x4", "x5"]
+
+    # Test data
+    X = zeros(Float64, 2, 3)
+    X[1, :] .= [1.0, 2.0, 3.0]
+    X[2, :] .= [0.5, 1.5, 2.5]
+
+    # 1. Basic register assignment - assign constant to register 2,
+    #    and then add the return to `x2` (which should now be 3.0!)
+    x1 = Expression(Node(; feature=1); operators, variable_names)
+    x2 = Expression(Node(; feature=2); operators, variable_names)
+    assign_expr = assign_op2(0.0 * x1 + 3.0) + x2
+
+    @test string_tree(assign_expr) == "[x2 =]((0.0 * x1) + 3.0) + x2"
+
+    # We should see that x2 will become 3.0 _before_ adding
+    result, completed = eval_tree_array(assign_expr, X)
+    @test completed == true
+    @test all(==(6.0), result)
+
+    # We should also see that X is not changed by this
+    @test X[2, :] == [0.5, 1.5, 2.5]
+
+    # But, with the reverse order, we get the x2 _before_ it was reassigned
+    assign_expr_reverse = x2 + assign_op2(0.0 * x1 + 3.0)
+    result, completed = eval_tree_array(assign_expr_reverse, X)
+    @test completed == true
+    @test result == [3.5, 4.5, 5.5]
+end

test/unittest.jl

@@ -133,3 +133,4 @@ include("test_expression_math.jl")
 include("test_structured_expression.jl")
 include("test_readonlynode.jl")
 include("test_zygote_gradient_wrapper.jl")
+include("test_special_operators.jl")