feat: create WhileOperator

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -77,7 +77,7 @@ import .StringsModule: get_op_name
 @reexport import .EvaluateModule:
     eval_tree_array, differentiable_eval_tree_array, EvalOptions
 import .EvaluateModule: ArrayBuffer
-@reexport import .SpecialOperatorsModule: AssignOperator
+@reexport import .SpecialOperatorsModule: AssignOperator, WhileOperator
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
 @reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!

src/Evaluate.jl

@@ -343,7 +343,8 @@ end
     if long_compilation_time
         return quote
             op = operators.binops[op_idx]
-            special_operator(op) && return deg2_eval_special(tree, cX, op, eval_options)
+            special_operator(op) &&
+                return deg2_eval_special(tree, cX, op, eval_options, operators)
             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)
@@ -360,7 +361,7 @@ end
             i -> i == op_idx,
             i -> let op = operators.binops[i]
                 if special_operator(op)
-                    deg2_eval_special(tree, cX, op, eval_options)
+                    deg2_eval_special(tree, cX, op, eval_options, operators)
                 elseif tree.l.degree == 0 && tree.r.degree == 0
                     deg2_l0_r0_eval(tree, cX, op, eval_options)
                 elseif tree.r.degree == 0

src/SpecialOperators.jl

@@ -1,7 +1,8 @@
 module SpecialOperatorsModule
 
 using ..OperatorEnumModule: OperatorEnum
-using ..EvaluateModule: _eval_tree_array, @return_on_nonfinite_array, deg2_eval
+using ..EvaluateModule:
+    _eval_tree_array, @return_on_nonfinite_array, deg2_eval, ResultOk, get_filled_array
 using ..ExpressionModule: AbstractExpression
 using ..ExpressionAlgebraModule: @declare_expression_operator
 
@@ -24,29 +25,6 @@ end
 @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
-#     for _ in 1:(op.max_iters)
-#         let cond_result = _eval_tree_array(cond, cX, operators, eval_options)
-#             !cond_result.ok && return cond_result
-#             @return_on_nonfinite_array(eval_options, cond_result.x)
-#         end
-#         let body_result = _eval_tree_array(body, cX, operators, eval_options)
-#             !body_result.ok && return body_result
-#             @return_on_nonfinite_array(eval_options, body_result.x)
-#             # TODO: Need to somehow mask instances 
-#         end
-#     end
-
-#     return get_filled_array(eval_options.buffer, zero(eltype(cX)), cX, axes(cX, 2))
-# end
-# TODO: Need to void any instance of buffer when using while loop.
-
 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
@@ -58,4 +36,47 @@ function deg1_eval_special(tree, cX, op::AssignOperator, eval_options, operators
     return result
 end
 
+Base.@kwdef struct WhileOperator <: Function
+    max_iters::Int = 100
+end
+
+@declare_expression_operator((op::WhileOperator), 2)
+@inline special_operator(::Type{WhileOperator}) = true
+get_op_name(o::WhileOperator) = "while"
+
+# TODO: Need to void any instance of buffer when using while loop.
+function deg2_eval_special(tree, cX, op::WhileOperator, eval_options, operators)
+    cond = tree.l
+    body = tree.r
+    mask = trues(size(cX, 2))
+    X = @view cX[:, mask]
+    # Initialize the result array for all columns
+    result_array = get_filled_array(eval_options.buffer, zero(eltype(cX)), cX, axes(cX, 2))
+    body_result = ResultOk(result_array, true)
+
+    for _ in 1:(op.max_iters)
+        cond_result = _eval_tree_array(cond, X, operators, eval_options)
+        !cond_result.ok && return cond_result
+        @return_on_nonfinite_array(eval_options, cond_result.x)
+
+        new_mask = cond_result.x .> 0.0
+        any(new_mask) || return body_result
+
+        # Track which columns are still active
+        mask[mask] .= new_mask
+        X = @view cX[:, mask]
+
+        # Evaluate just for active columns
+        iter_result = _eval_tree_array(body, X, operators, eval_options)
+        !iter_result.ok && return iter_result
+
+        # Update the corresponding elements in the result array
+        body_result.x[mask] .= iter_result.x
+        @return_on_nonfinite_array(eval_options, body_result.x)
+    end
+
+    # We passed max_iters, so this result is invalid
+    return ResultOk(body_result.x, false)
+end
+
 end

test/test_special_operators.jl

@@ -100,3 +100,74 @@ end
     not_simplified = simplify_tree!(simple_expr_special)
     @test string_tree(not_simplified) == "2.0 + 2.0"
 end
+
+@testitem "WhileOperator basic functionality" begin
+    using DynamicExpressions
+    using Test
+
+    # Define operators
+    while_op = WhileOperator(; max_iters=100)
+    assign_x2 = AssignOperator(; target_register=2)
+    operators = OperatorEnum(;
+        binary_operators=[+, -, *, /, while_op],  # While is binary operator
+        unary_operators=[assign_x2],
+    )
+    variable_names = ["x1", "x2", "x3"]
+
+    # Test data - x2 starts at 1.0 for all samples
+    X = zeros(Float64, 2, 3)
+    X[2, :] .= 1.0  # x2 initial value
+
+    # Build expression: while (3.0 - x2 > 0) do x2 = x2 + 1.0
+    x2 = Expression(Node(; feature=2); operators, variable_names)
+    expr = while_op(3.0 - x2, assign_x2(x2 + 1.0))
+
+    @test string_tree(expr) == "while(3.0 - x2, [x2 =](x2 + 1.0))"
+
+    result, completed = eval_tree_array(expr, X)
+    @test completed == true
+    @test all(result .≈ 3.0)  # After 2 iterations, x2 becomes 3.0
+    @test X[2, :] == [1.0, 1.0, 1.0]  # Original data unchanged
+end
+
+@testitem "Fibonacci sequence with WhileOperator" begin
+    using DynamicExpressions
+    using Test
+
+    # Define operators
+    while_op = WhileOperator(; max_iters=100)
+    assign_ops = [AssignOperator(; target_register=i) for i in 1:5]
+    operators = OperatorEnum(;
+        binary_operators=[+, -, *, /, while_op], unary_operators=assign_ops
+    )
+    variable_names = ["x1", "x2", "x3", "x4", "x5"]
+
+    # Test data - x2=5 (counter), x3=0 (F(0)), x4=1 (F(1))
+    X = zeros(Float64, 5, 4)
+    # Set different Fibonacci sequence positions to calculate
+    X[2, :] = [3.0, 5.0, 7.0, 10.0]  # Calculate F(3), F(5), F(7), F(10)
+
+    # Initialize all rows with F(0)=0, F(1)=1
+    X[3, :] .= 0.0  # x3 = 0.0 (F(0))
+    X[4, :] .= 1.0  # x4 = 1.0 (F(1))
+
+    xs = [Expression(Node(; feature=i); operators, variable_names) for i in 1:5]
+
+    # Build expression: 
+    condition = xs[2]  # WhileOperator implicitly checks if > 0
+    body =
+        assign_ops[5](xs[3]) +
+        assign_ops[3](xs[4]) +
+        assign_ops[4](xs[5] + xs[4]) +
+        assign_ops[2](xs[2] - 1.0)
+    expr = (while_op(condition, body) * 0.0) + xs[3]
+
+    @test string_tree(expr) ==
+        "(while(x2, (([x5 =](x3) + [x3 =](x4)) + [x4 =](x5 + x4)) + [x2 =](x2 - 1.0)) * 0.0) + x3"
+
+    result, completed = eval_tree_array(expr, X)
+    @test completed == true
+
+    # Test each Fibonacci number is correctly calculated
+    @test result ≈ [2.0, 5.0, 13.0, 55.0]  # F(3)=2, F(5)=5, F(7)=13, F(10)=55
+end