test: integrate Supposition testing

Author: MilesCranmer

test/Project.toml

@@ -17,6 +17,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Supposition = "5a0628fe-1738-4658-9b6d-0b7605a9755b"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/supposition_utils.jl

@@ -0,0 +1,66 @@
+# supposition_utils.jl
+#
+# Helper that builds a Supposition generator returning fully-random
+# DynamicExpressions.Expression objects whose node type is Node{T,D}.
+# D is inferred from `operators`.
+
+module SuppositionUtils
+
+using Supposition: Data
+using DynamicExpressions: Node, Expression, OperatorEnum
+using DynamicExpressions.OperatorEnumConstructionModule: empty_all_globals!
+empty_all_globals!()
+
+function make_expression_generator(
+    ::Type{T};
+    num_features::Int=5,
+    operators::OperatorEnum=OperatorEnum(((abs, cos), (+, -, *, /))),
+    max_layers::Int=3,
+) where {T}
+    D = length(operators.ops)
+
+    val_gen = Data.Floats{T}(; nans=false, infs=false)
+    val_node_gen = map(v -> Node{T,D}(; val=v), val_gen)
+
+    feature_gen = Data.SampledFrom(1:num_features)
+    feature_node_gen = map(i -> Node{T,D}(; feature=i), feature_gen)
+
+    leaf_gen = val_node_gen | feature_node_gen
+
+    wrapper_funcs = ntuple(
+        degree -> let op_list = operators[degree]
+            op_gen = Data.SampledFrom(1:length(op_list))
+
+            child -> map(
+                (op_idx, args...) -> Node{T,D}(; op=op_idx, children=args),
+                op_gen,
+                ntuple(_ -> child, degree)...,
+            )
+        end,
+        Val(D),
+    )
+    expr_wrap(child) = foldl(|, (w(child) for w in wrapper_funcs))
+    tree_gen = Data.Recursive(leaf_gen, expr_wrap; max_layers)
+    return map(
+        t -> Expression(t; operators, variable_names=["x$i" for i in 1:num_features]),
+        tree_gen,
+    )
+end
+
+# inside module SuppositionUtils
+function make_input_matrix_generator(
+    ::Type{T}=Float64; n_features::Int=5, min_batch::Int=1, max_batch::Int=16
+) where {T}
+    elem_gen = Data.Floats{T}(; nans=false, infs=false)
+    batch_gen = Data.Integers(min_batch, max_batch)
+
+    Data.bind(batch_gen) do bs
+        vec_len = n_features * bs
+        vec_gen = Data.Vectors(elem_gen; min_size=vec_len, max_size=vec_len)
+        Data.map(v -> reshape(v, n_features, bs), vec_gen)
+    end
+end
+
+end
+
+using .SuppositionUtils: make_input_matrix_generator, make_expression_generator

test/test_supposition_consistency.jl

@@ -0,0 +1,44 @@
+@testitem "Supposition round-trip consistency" begin
+    using Test
+    using Random
+
+    using Supposition
+    using Supposition: @check, Data
+    using DynamicExpressions
+    using DynamicExpressions:
+        string_tree, parse_expression, eval_tree_array, Node, get_operators, get_tree
+
+    # bring the generator into scope
+    include("supposition_utils.jl")
+
+    n_features = 5
+    max_layers = 20
+    T = Float64
+    operators = OperatorEnum(((abs, cos, exp), (+, -, *, /), (fma, clamp, +, max)))
+
+    expr_gen = make_expression_generator(
+        T; num_features=n_features, max_layers=max_layers, operators=operators
+    )
+
+    @check function roundtrip_string(ex=expr_gen)
+        tree_str = string_tree(ex)
+        ex_parsed = parse_expression(
+            Meta.parse(tree_str);
+            operators=get_operators(ex),
+            variable_names=["x$i" for i in 1:n_features],
+            node_type=Node{Float64,3},
+        )
+        return ex == ex_parsed
+    end
+
+    input_gen = make_input_matrix_generator(T; n_features)
+    @check max_examples = 1024 function eval_against_string(ex=expr_gen, X=input_gen)
+        expression_result, ok = eval_tree_array(ex, X)
+        !ok && return true  # If the expression is not valid, we can't test it
+        tree_str = string_tree(ex)
+        f_sym = gensym("f")
+        f = eval(Meta.parse("(x1, x2, x3, x4, x5) -> ($tree_str)"))
+        true_result = Float64[Base.invokelatest(f, x...) for x in eachcol(X)]
+        return expression_result ≈ true_result
+    end
+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_supposition_consistency.jl")