Merge pull request #484 from MilesCranmer/autoconvert_composable

feat: add Number input support to ComposableExpression

Author: MilesCranmer

src/ComposableExpression.jl

@@ -187,17 +187,42 @@ function (ex::AbstractComposableExpression)(
         return x .* nan
     end
 end
+# Method for all-Number arguments (scalars)
+function (ex::AbstractComposableExpression)(x::Number, _xs::Vararg{Number,N}) where {N}
+    xs = (x, _xs...)
+
+    vectors = ntuple(i -> ValidVector([float(xs[i])], true), length(xs))
+    return only(_get_value(ex(vectors...)))
+end
+
 function (ex::AbstractComposableExpression)(
-    x::ValidVector, _xs::Vararg{ValidVector,N}
+    x::Union{ValidVector,Number}, _xs::Vararg{Union{ValidVector,Number},N}
 ) where {N}
     xs = (x, _xs...)
-    valid = all(_is_valid, xs)
-    if !valid
-        return ValidVector(_get_value(first(xs)), false)
-    else
-        X = Matrix(stack(map(_get_value, xs))')
+    sample_vector =
+        let first_valid_vector_idx = findfirst(arg -> arg isa ValidVector, xs)::Int
+            xs[first_valid_vector_idx]::ValidVector
+        end
+
+    # Convert Numbers to ValidVectors based on first ValidVector's size
+    valid_args = ntuple(length(xs)) do i
+        arg = xs[i]
+        if arg isa ValidVector
+            arg
+        else
+            # Convert Number to ValidVector with repeated values
+            filled_array = similar(sample_vector.x)
+            fill!(filled_array, arg)
+            ValidVector(filled_array, true)
+        end
+    end
+
+    if all(_is_valid, valid_args)
+        X = stack(map(_get_value, valid_args); dims=1)
         eval_options = get_eval_options(ex)
         return ValidVector(eval_tree_array(ex, X; eval_options))
+    else
+        return ValidVector(_get_value(first(valid_args)), false)
     end
 end
 function (ex::AbstractComposableExpression{T})() where {T}
@@ -252,6 +277,55 @@ _is_valid(x) = true
 _get_value(x::ValidVector) = x.x
 _get_value(x) = x
 
+struct ValidVectorMixError <: Exception end
+struct ValidVectorAccessError <: Exception end
+
+function Base.showerror(io::IO, ::ValidVectorMixError)
+    return print(
+        io,
+        """
+ValidVectorMixError: Cannot mix ValidVector with regular Vector.
+
+ValidVector handles validity checks, auto-vectorization, and batching in template expressions.
+The .valid field tracks whether any upstream computation failed (false = failed, true = valid).
+
+Wrap your vectors in ValidVector:
+
+    ```julia
+    valid_ar1 = ValidVector(ar1, all(isfinite, ar1))
+    valid_ar1 + valid_ar2
+    ```
+
+Alternatively, you can access the vector from a ValidVector with `my_validvector.x`,
+but you must be sure to propagate the `.valid` field. For example:
+
+    ```julia
+    out = ar1 .+ valid_ar2.x
+    ValidVector(out, all(isfinite, out) && valid_ar2.valid)
+    ```
+
+""",
+    )
+end
+
+function Base.showerror(io::IO, ::ValidVectorAccessError)
+    return print(
+        io,
+        """
+ValidVectorAccessError: ValidVector doesn't support direct array operations.
+
+Use .x for data and .valid for validity:
+
+    ```julia
+    valid_ar.x[1]        # indexing
+    length(valid_ar.x)   # length
+    valid_ar.valid       # check validity (false = any upstream computation failed)
+    ```
+
+ValidVector handles validity/batching automatically in template expressions.""",
+    )
+end
+
 #! format: off
 # First, binary operators:
 for op in (
@@ -264,6 +338,9 @@ for op in (
         Base.$(op)(x::ValidVector, y::ValidVector) = apply_operator(Base.$(op), x, y)
         Base.$(op)(x::ValidVector, y::Number) = apply_operator(Base.$(op), x, y)
         Base.$(op)(x::Number, y::ValidVector) = apply_operator(Base.$(op), x, y)
+
+        Base.$(op)(::ValidVector, ::AbstractVector) = throw(ValidVectorMixError())
+        Base.$(op)(::AbstractVector, ::ValidVector) = throw(ValidVectorMixError())
     end
 end
 function Base.literal_pow(::typeof(^), x::ValidVector, ::Val{p}) where {p}
@@ -286,6 +363,12 @@ for op in (
 end
 #! format: on
 
+Base.length(::ValidVector) = throw(ValidVectorAccessError())
+Base.push!(::ValidVector, ::Any) = throw(ValidVectorAccessError())
+for op in (:getindex, :size, :append!, :setindex!)
+    @eval Base.$(op)(::ValidVector, ::Any...) = throw(ValidVectorAccessError())
+end
+
 # TODO: Support for 3-ary operators
 
 end

src/SymbolicRegression.jl

@@ -341,8 +341,9 @@ using .SearchUtilsModule:
 using .LoggingModule: AbstractSRLogger, SRLogger, get_logger
 using .TemplateExpressionModule:
     TemplateExpression, TemplateStructure, TemplateExpressionSpec, ParamVector, has_params
-using .TemplateExpressionModule: ValidVector
-using .ComposableExpressionModule: ComposableExpression
+using .TemplateExpressionModule: ValidVector, TemplateReturnError
+using .ComposableExpressionModule:
+    ComposableExpression, ValidVectorMixError, ValidVectorAccessError
 using .ExpressionBuilderModule: embed_metadata, strip_metadata
 using .ParametricExpressionModule: ParametricExpressionSpec
 using .TemplateExpressionMacroModule: @template_spec

src/TemplateExpression.jl

@@ -631,6 +631,40 @@ function HOF.make_prefix(::TemplateExpression, ::AbstractOptions, ::Dataset)
     return ""
 end
 
+struct TemplateReturnError <: Exception end
+
+function Base.showerror(io::IO, ::TemplateReturnError)
+    return print(
+        io,
+        """
+TemplateReturnError: Template expression returned a regular Vector, but ValidVector is required.
+
+Template expressions must return ValidVector for proper handling:
+
+    ```julia
+    return ValidVector(my_data, computation_is_valid)
+    ```
+
+The .valid field is used to track whether any upstream computation failed.
+It's important to handle this correctly.
+
+Example of manually propagating validity:
+
+    ```julia
+    _f_result = f(x1, x2)  # Returns ValidVector
+    _g_result = g(x3)      # Returns ValidVector
+
+    # Combine results manually and propagate validity
+    combined_data = _f_result.x .+ _g_result.x
+    return ValidVector(combined_data, _f_result.valid && _g_result.valid)
+    ```
+
+Note that normally we could simply write `_f_result + _g_result`,
+and this would automatically handle the validity and vectorization.
+""",
+    )
+end
+
 @stable(
     default_mode = "disable",
     default_union_limit = 2,
@@ -657,6 +691,10 @@ end
                 extra_args...,
                 map(x -> ValidVector(copy(x), true), eachrow(cX)),
             )
+            # Validate that template expression returned a ValidVector
+            if !(result isa ValidVector)
+                throw(TemplateReturnError())
+            end
             return result.x, result.valid
         end
         function (ex::TemplateExpression)(

test/test_composable_expression.jl

@@ -129,3 +129,56 @@ end
     x2_val = ValidVector([1.0, 2.0], false)
     @test ex(x1_val, x2_val).valid == false
 end
+
+@testitem "ValidVector helpful error messages" tags = [:part2] begin
+    using SymbolicRegression
+    using SymbolicRegression: ValidVector, ValidVectorMixError, ValidVectorAccessError
+
+    vv = ValidVector([1.0, 2.0], true)
+    v = [3.0, 4.0]
+
+    # Helper function to get error message
+    get_error_msg(err) =
+        let io = IOBuffer()
+            Base.showerror(io, err)
+            String(take!(io))
+        end
+
+    # Test vector arithmetic errors encourage ValidVector wrapping
+    err_mix = @test_throws ValidVectorMixError vv + v
+    @test_throws ValidVectorMixError v * vv  # Test other direction too
+
+    mix_msg = get_error_msg(err_mix.value)
+    @test contains(
+        mix_msg,
+        "ValidVector handles validity checks, auto-vectorization, and batching in template expressions",
+    )
+
+    # Test array access errors mention .x and .valid
+    err_access = @test_throws ValidVectorAccessError vv[1]
+    @test_throws ValidVectorAccessError length(vv)
+    @test_throws ValidVectorAccessError push!(vv, 5.0)
+
+    access_msg = get_error_msg(err_access.value)
+    @test contains(access_msg, "valid_ar.x[1]")
+    @test contains(access_msg, "valid_ar.valid")
+    @test contains(access_msg, "length(valid_ar.x)")
+    @test contains(access_msg, "doesn't support direct array operations")
+    @test contains(access_msg, "ValidVector handles validity/batching automatically")
+end
+
+@testitem "Test Number inputs" tags = [:part2] begin
+    using SymbolicRegression: ComposableExpression, Node, ValidVector
+    using DynamicExpressions: OperatorEnum
+
+    operators = OperatorEnum(; binary_operators=(+, *))
+    x1 = ComposableExpression(Node{Float64}(; feature=1); operators)
+    x2 = ComposableExpression(Node{Float64}(; feature=2); operators)
+    ex = x1 + x2
+
+    @test ex(2.0, 3.0) ≈ 5.0
+    @test isnan(ex(NaN, 3.0))
+    @test ex(ValidVector([1.0], true), 2.0).x ≈ [3.0]
+    @test ex(ValidVector([1.0, 1.0], true), 2.0).x ≈ [3.0, 3.0]
+    @test ex(ValidVector([1.0, 1.0], false), 2.0).valid == false
+end

test/test_template_expression.jl

@@ -657,3 +657,38 @@ end
         parse_guesses(PopMember{Float64,Float64}, [bad_guess], [dataset], options)
     )
 end
+
+@testitem "Template expression return validation" tags = [:part2] begin
+    using SymbolicRegression:
+        TemplateReturnError,
+        ValidVector,
+        ComposableExpression,
+        TemplateStructure,
+        TemplateExpression
+    using DynamicExpressions: OperatorEnum, Node
+
+    operators = OperatorEnum(; binary_operators=(+, *, /, -), unary_operators=(sin, cos))
+    x1 = ComposableExpression(Node{Float64}(; feature=1); operators, variable_names=nothing)
+
+    # Test that returning a regular vector from template expression throws TemplateReturnError
+    bad_structure = TemplateStructure{(:f,)}(
+        ((; f), (x,)) -> [1.0, 2.0];  # Returns regular Vector instead of ValidVector
+        num_features=(; f=1),
+    )
+    bad_expr = TemplateExpression(
+        (; f=x1); structure=bad_structure, operators, variable_names=nothing
+    )
+    X = [1.0 2.0]'
+
+    function get_error_msg(err)
+        io = IOBuffer()
+        Base.showerror(io, err)
+        return String(take!(io))
+    end
+
+    err = @test_throws TemplateReturnError bad_expr(X)
+    msg = get_error_msg(err.value)
+    @test contains(msg, "Template expression returned a regular Vector")
+    @test contains(msg, "ValidVector is required")
+    @test contains(msg, "ValidVector(my_data, computation_is_valid)")
+end