BREAKING: Change expression types to DynamicExpressions.Expression (from DynamicExpressions.Node)

.github/workflows/CI.yml

@@ -26,8 +26,9 @@ jobs:
       fail-fast: false
       matrix:
         test:
-          - "unit"
-          - "integration"
+          - "part1"
+          - "part2"
+          - "part3"
         julia-version:
           - "1.6"
           - "1.8"
@@ -37,22 +38,31 @@ jobs:
         include:
           - os: windows-latest
             julia-version: "1"
-            test: "unit"
+            test: "part1"
           - os: windows-latest
             julia-version: "1"
-            test: "integration"
+            test: "part2"
+          - os: windows-latest
+            julia-version: "1"
+            test: "part3"
+          - os: macOS-latest
+            julia-version: "1"
+            test: "part1"
           - os: macOS-latest
             julia-version: "1"
-            test: "unit"
+            test: "part2"
           - os: macOS-latest
             julia-version: "1"
-            test: "integration"
+            test: "part3"
+          - os: ubuntu-latest
+            julia-version: "~1.11.0-0"
+            test: "part1"
           - os: ubuntu-latest
             julia-version: "~1.11.0-0"
-            test: "unit"
+            test: "part2"
           - os: ubuntu-latest
             julia-version: "~1.11.0-0"
-            test: "integration"
+            test: "part3"
 
     steps:
       - uses: actions/checkout@v4

.gitignore

@@ -20,3 +20,4 @@ docs/src/index.md
 *.code-workspace
 .vscode
 **/*.json
+LocalPreferences.toml

Project.toml

@@ -1,11 +1,13 @@
 name = "SymbolicRegression"
 uuid = "8254be44-1295-4e6a-a16d-46603ac705cb"
 authors = ["MilesCranmer <[email protected]>"]
-version = "0.24.5"
+version = "0.25.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
@@ -35,11 +37,13 @@ SymbolicRegressionJSON3Ext = "JSON3"
 SymbolicRegressionSymbolicUtilsExt = "SymbolicUtils"
 
 [compat]
+ADTypes = "~1.4"
 Compat = "^4.2"
 Dates = "1"
+DifferentiationInterface = "0.5"
 DispatchDoctor = "0.4"
 Distributed = "1"
-DynamicExpressions = "0.16"
+DynamicExpressions = "0.18.2"
 DynamicQuantities = "0.10, 0.11, 0.12, 0.13, 0.14"
 JSON3 = "1"
 LineSearches = "7"
@@ -56,7 +60,7 @@ Random = "1"
 Reexport = "1"
 SpecialFunctions = "0.10.1, 1, 2"
 StatsBase = "0.33, 0.34"
-SymbolicUtils = "0.19, ^1.0.5"
+SymbolicUtils = "0.19, ^1.0.5, 2"
 TOML = "1"
 julia = "1.6"
 

examples/parameterized_function.jl

@@ -0,0 +1,30 @@
+using SymbolicRegression
+using Random: MersenneTwister
+using Zygote
+using MLJBase: machine, fit!, predict
+
+rng = MersenneTwister(0)
+X = NamedTuple{(:x1, :x2, :x3, :x4, :x5)}(ntuple(_ -> randn(rng, Float32, 30), Val(5)))
+X = (; X..., classes=rand(rng, 1:2, 30))
+p1 = rand(rng, Float32, 2)
+p2 = rand(rng, Float32, 2)
+
+y = [
+    2 * cos(X.x4[i] + p1[X.classes[i]]) + X.x1[i]^2 - p2[X.classes[i]] for
+    i in eachindex(X.classes)
+]
+
+model = SRRegressor(;
+    niterations=10,
+    binary_operators=[+, *, /, -],
+    unary_operators=[cos, exp],
+    populations=10,
+    expression_type=ParametricExpression,
+    expression_options=(; max_parameters=2),
+    autodiff_backend=:Zygote,
+    parallelism=:multithreading,
+)
+
+mach = machine(model, X, y)
+fit!(mach)
+ypred = predict(mach, X)

ext/SymbolicRegressionSymbolicUtilsExt.jl

@@ -1,8 +1,9 @@
 module SymbolicRegressionSymbolicUtilsExt
 
 using SymbolicUtils: Symbolic
-using SymbolicRegression: AbstractExpressionNode, Node, Options
+using SymbolicRegression: AbstractExpressionNode, AbstractExpression, Node, Options
 using SymbolicRegression.MLJInterfaceModule: AbstractSRRegressor, get_options
+using DynamicExpressions: get_tree, get_operators
 
 import SymbolicRegression: node_to_symbolic, symbolic_to_node
 
@@ -11,10 +12,14 @@ import SymbolicRegression: node_to_symbolic, symbolic_to_node
 
 Convert an expression to SymbolicUtils.jl form.
 """
-function node_to_symbolic(tree::AbstractExpressionNode, options::Options; kws...)
-    return node_to_symbolic(tree, options.operators; kws...)
+function node_to_symbolic(
+    tree::Union{AbstractExpressionNode,AbstractExpression}, options::Options; kws...
+)
+    return node_to_symbolic(get_tree(tree), get_operators(tree, options); kws...)
 end
-function node_to_symbolic(tree::AbstractExpressionNode, m::AbstractSRRegressor; kws...)
+function node_to_symbolic(
+    tree::Union{AbstractExpressionNode,AbstractExpression}, m::AbstractSRRegressor; kws...
+)
     return node_to_symbolic(tree, get_options(m); kws...)
 end
 
@@ -31,24 +36,30 @@ function symbolic_to_node(eqn::Symbolic, m::AbstractSRRegressor; kws...)
 end
 
 function Base.convert(
-    ::Type{Symbolic}, tree::AbstractExpressionNode, options::Options; kws...
+    ::Type{Symbolic},
+    tree::Union{AbstractExpressionNode,AbstractExpression},
+    options::Union{Options,Nothing}=nothing;
+    kws...,
 )
-    return convert(Symbolic, tree, options.operators; kws...)
+    return convert(Symbolic, get_tree(tree), get_operators(tree, options); kws...)
 end
 function Base.convert(
-    ::Type{Symbolic}, tree::AbstractExpressionNode, m::AbstractSRRegressor; kws...
+    ::Type{Symbolic},
+    tree::Union{AbstractExpressionNode,AbstractExpression},
+    m::AbstractSRRegressor;
+    kws...,
 )
     return convert(Symbolic, tree, get_options(m); kws...)
 end
 
 function Base.convert(
     ::Type{N}, x::Union{Number,Symbolic}, options::Options; kws...
-) where {N<:AbstractExpressionNode}
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
     return convert(N, x, options.operators; kws...)
 end
 function Base.convert(
     ::Type{N}, x::Union{Number,Symbolic}, m::AbstractSRRegressor; kws...
-) where {N<:AbstractExpressionNode}
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
     return convert(N, x, get_options(m); kws...)
 end
 

src/CheckConstraints.jl

@@ -1,6 +1,7 @@
 module CheckConstraintsModule
 
-using DynamicExpressions: AbstractExpressionNode, count_depth, tree_mapreduce
+using DynamicExpressions:
+    AbstractExpressionNode, AbstractExpression, get_tree, count_depth, tree_mapreduce
 using ..UtilsModule: vals
 using ..CoreModule: Options
 using ..ComplexityModule: compute_complexity, past_complexity_limit
@@ -70,6 +71,15 @@ function flag_illegal_nests(tree::AbstractExpressionNode, options::Options)::Boo
 end
 
 """Check if user-passed constraints are violated or not"""
+function check_constraints(
+    ex::AbstractExpression,
+    options::Options,
+    maxsize::Int,
+    cursize::Union{Int,Nothing}=nothing,
+)::Bool
+    tree = get_tree(ex)
+    return check_constraints(tree, options, maxsize, cursize)
+end
 function check_constraints(
     tree::AbstractExpressionNode,
     options::Options,
@@ -93,7 +103,8 @@ function check_constraints(
     return true
 end
 
-check_constraints(tree::AbstractExpressionNode, options::Options)::Bool =
-    check_constraints(tree, options, options.maxsize)
+check_constraints(
+    ex::Union{AbstractExpression,AbstractExpressionNode}, options::Options
+)::Bool = check_constraints(ex, options, options.maxsize)
 
 end

src/Complexity.jl

@@ -1,9 +1,12 @@
 module ComplexityModule
 
-using DynamicExpressions: AbstractExpressionNode, count_nodes, tree_mapreduce
+using DynamicExpressions:
+    AbstractExpression, AbstractExpressionNode, get_tree, count_nodes, tree_mapreduce
 using ..CoreModule: Options, ComplexityMapping
 
-function past_complexity_limit(tree::AbstractExpressionNode, options::Options, limit)::Bool
+function past_complexity_limit(
+    tree::Union{AbstractExpression,AbstractExpressionNode}, options::Options, limit
+)::Bool
     return compute_complexity(tree, options) > limit
 end
 
@@ -14,6 +17,11 @@ By default, this is the number of nodes in a tree.
 However, it could use the custom settings in options.complexity_mapping
 if these are defined.
 """
+function compute_complexity(
+    tree::AbstractExpression, options::Options; break_sharing=Val(false)
+)
+    return compute_complexity(get_tree(tree), options; break_sharing)
+end
 function compute_complexity(
     tree::AbstractExpressionNode, options::Options; break_sharing=Val(false)
 )::Int

src/ConstantOptimization.jl

@@ -2,7 +2,8 @@ module ConstantOptimizationModule
 
 using LineSearches: LineSearches
 using Optim: Optim
-using DynamicExpressions: Node, count_constants
+using DifferentiationInterface: value_and_gradient
+using DynamicExpressions: Expression, Node, count_constants, get_constants, set_constants!
 using ..CoreModule: Options, Dataset, DATA_TYPE, LOSS_TYPE
 using ..UtilsModule: get_birth_order
 using ..LossFunctionsModule: eval_loss, loss_to_score, batch_sample
@@ -22,7 +23,7 @@ end
 function dispatch_optimize_constants(
     dataset::Dataset{T,L}, member::P, options::Options, idx
 ) where {T<:DATA_TYPE,L<:LOSS_TYPE,P<:PopMember{T,L}}
-    nconst = count_constants(member.tree)
+    nconst = count_constants_for_optimization(member.tree)
     nconst == 0 && return (member, 0.0)
     if nconst == 1 && !(T <: Complex)
         algorithm = Optim.Newton(; linesearch=LineSearches.BackTracking())
@@ -39,36 +40,43 @@ function dispatch_optimize_constants(
         idx,
     )
 end
+count_constants_for_optimization(ex::Expression) = count_constants(ex)
 
 function _optimize_constants(
     dataset, member::P, options, algorithm, optimizer_options, idx
 )::Tuple{P,Float64} where {T,L,P<:PopMember{T,L}}
     tree = member.tree
     eval_fraction = options.batching ? (options.batch_size / dataset.n) : 1.0
-    f(t) = eval_loss(t, dataset, options; regularization=false, idx=idx)::L
+    f = Evaluator(dataset, options, idx)
+    fg! = GradEvaluator(f)
+    obj = if algorithm isa Optim.Newton || options.autodiff_backend === nothing
+        f
+    else
+        Optim.only_fg!(fg!)
+    end
     baseline = f(tree)
-    result = Optim.optimize(f, tree, algorithm, optimizer_options)
+    x0, refs = get_constants(tree)
+    result = Optim.optimize(obj, tree, algorithm, optimizer_options)
     num_evals = result.f_calls * eval_fraction
     # Try other initial conditions:
     for _ in 1:(options.optimizer_nrestarts)
         tmptree = copy(tree)
-        foreach(tmptree) do node
-            if node.degree == 0 && node.constant
-                node.val = (node.val) * (T(1) + T(1//2) * randn(T))
-            end
-        end
+        eps = randn(T, size(x0)...)
+        xt = @. x0 * (T(1) + T(1//2) * eps)
+        set_constants!(tmptree, xt, refs)
         tmpresult = Optim.optimize(
-            f, tmptree, algorithm, optimizer_options; make_copy=false
+            obj, tmptree, algorithm, optimizer_options; make_copy=false
         )
         num_evals += tmpresult.f_calls * eval_fraction
+        # TODO: Does this need to take into account h_calls?
 
         if tmpresult.minimum < result.minimum
             result = tmpresult
         end
     end
 
     if result.minimum < baseline
-        member.tree = result.minimizer
+        member.tree = result.minimizer::typeof(member.tree)
         member.loss = eval_loss(member.tree, dataset, options; regularization=true, idx=idx)
         member.score = loss_to_score(
             member.loss, dataset.use_baseline, dataset.baseline_loss, member, options
@@ -80,4 +88,21 @@ function _optimize_constants(
     return member, num_evals
 end
 
+struct Evaluator{D<:Dataset,O<:Options,I} <: Function
+    dataset::D
+    options::O
+    idx::I
+end
+(e::Evaluator)(t) = eval_loss(t, e.dataset, e.options; regularization=false, idx=e.idx)
+struct GradEvaluator{F<:Evaluator} <: Function
+    f::F
+end
+function (g::GradEvaluator)(F, G, t)
+    (val, grad) = value_and_gradient(g.f, g.f.options.autodiff_backend, t)
+    if G !== nothing && grad !== nothing && grad.tree !== nothing
+        G .= grad.tree.gradient
+    end
+    return val
+end
+
 end

src/Core.jl

@@ -1,5 +1,7 @@
 module CoreModule
 
+function create_expression end
+
 include("Utils.jl")
 include("ProgramConstants.jl")
 include("Dataset.jl")

src/Dataset.jl

@@ -22,6 +22,8 @@ import ...deprecate_varmap
 
 - `X::AbstractMatrix{T}`: The input features, with shape `(nfeatures, n)`.
 - `y::AbstractVector{T}`: The desired output values, with shape `(n,)`.
+- `index::Int`: The index of the output feature corresponding to this
+    dataset, if any.
 - `n::Int`: The number of samples.
 - `nfeatures::Int`: The number of features.
 - `weighted::Bool`: Whether the dataset is non-uniformly weighted.
@@ -64,6 +66,7 @@ mutable struct Dataset{
 }
     @constfield X::AX
     @constfield y::AY
+    @constfield index::Int
     @constfield n::Int
     @constfield nfeatures::Int
     @constfield weighted::Bool
@@ -99,6 +102,7 @@ function Dataset(
     X::AbstractMatrix{T},
     y::Union{AbstractVector{T},Nothing}=nothing,
     loss_type::Type{L}=Nothing;
+    index::Int=1,
     weights::Union{AbstractVector{T},Nothing}=nothing,
     variable_names::Union{Array{String,1},Nothing}=nothing,
     display_variable_names=variable_names,
@@ -123,6 +127,7 @@ function Dataset(
             X,
             y,
             kws[:loss_type];
+            index,
             weights,
             variable_names,
             display_variable_names,
@@ -206,6 +211,7 @@ function Dataset(
     }(
         X,
         y,
+        index,
         n,
         nfeatures,
         weighted,