fix: backport SymbolicUtils v4 compat to release-v1

Restore index_functions round-trip for custom operators under SymbolicUtils v4.

Author: MilesCranmerBot

Project.toml

@@ -37,7 +37,7 @@ MacroTools = "0.4, 0.5"
 Optim = "0.19, 1"
 PrecompileTools = "1"
 Reexport = "1"
-SymbolicUtils = "0.19, ^1.0.5, 2, 3"
+SymbolicUtils = "4"
 Zygote = "0.7"
 julia = "1.10"
 

ext/DynamicExpressionsSymbolicUtilsExt.jl

@@ -8,19 +8,14 @@ using DynamicExpressions.OperatorEnumModule: AbstractOperatorEnum
 using DynamicExpressions.UtilsModule: deprecate_varmap
 
 using SymbolicUtils
+using SymbolicUtils: BasicSymbolic, SymReal, iscall, issym, isconst, unwrap_const
 
 import DynamicExpressions.ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node
 import DynamicExpressions.ValueInterfaceModule: is_valid
 
-const SYMBOLIC_UTILS_TYPES = Union{<:Number,SymbolicUtils.Symbolic{<:Number}}
+const SYMBOLIC_UTILS_TYPES = Union{<:Number,BasicSymbolic}
 const SUPPORTED_OPS = (cos, sin, exp, cot, tan, csc, sec, +, -, *, /)
 
-@static if isdefined(SymbolicUtils, :iscall)
-    iscall(x) = SymbolicUtils.iscall(x)
-else
-    iscall(x) = SymbolicUtils.istree(x)
-end
-
 macro return_on_false(flag, retval)
     :(
         if !$(esc(flag))
@@ -29,7 +24,7 @@ macro return_on_false(flag, retval)
     )
 end
 
-function is_valid(x::SymbolicUtils.Symbolic)
+function is_valid(x::BasicSymbolic)
     return if iscall(x)
         all(is_valid.([SymbolicUtils.operation(x); SymbolicUtils.arguments(x)]))
     else
@@ -46,40 +41,53 @@ function parse_tree_to_eqs(
     if tree.degree == 0
         # Return constant if needed
         tree.constant && return subs_bad(tree.val)
-        return SymbolicUtils.Sym{LiteralReal}(Symbol("x$(tree.feature)"))
+        return SymbolicUtils.Sym{SymReal}(Symbol("x$(tree.feature)"); type=Number)
     end
     # Collect the next children
     # TODO: Type instability!
     children = tree.degree == 2 ? (tree.l, tree.r) : (tree.l,)
     # Get the operation
     op = tree.degree == 2 ? operators.binops[tree.op] : operators.unaops[tree.op]
-    # Create an N tuple of Numbers for each argument
-    dtypes = map(x -> Number, 1:(tree.degree))
-    #
-    if !(op ∈ SUPPORTED_OPS) && index_functions
-        op = SymbolicUtils.Sym{(SymbolicUtils.FnType){Tuple{dtypes...},Number}}(Symbol(op))
+
+    # For custom operators, SymbolicUtils can't represent the Julia function directly.
+    # When `index_functions=true`, represent the operator by its name as an uninterpreted
+    # SymbolicUtils function symbol so we can round-trip back to a DynamicExpressions node.
+    if !(op ∈ SUPPORTED_OPS)
+        if index_functions
+            dtypes = ntuple(_ -> Number, tree.degree)
+            op = SymbolicUtils.Sym{SymReal}(
+                Symbol(op); type=SymbolicUtils.FnType{Tuple{dtypes...},Number,Nothing}
+            )
+        else
+            error(
+                "Custom operator '$op' is not supported with SymbolicUtils unless " *
+                "index_functions=true. Supported operators without indexing: $SUPPORTED_OPS",
+            )
+        end
     end
 
-    return subs_bad(
-        op(map(x -> parse_tree_to_eqs(x, operators, index_functions), children)...)
-    )
+    # Convert children to symbolic form
+    sym_children = map(x -> parse_tree_to_eqs(x, operators, index_functions), children)
+
+    return subs_bad(op(sym_children...))
 end
 
-# For operators which are indexed, we need to convert them back
-# using the string:
-function convert_to_function(
-    x::SymbolicUtils.Sym{SymbolicUtils.FnType{T,Number}}, operators::AbstractOperatorEnum
-) where {T<:Tuple}
-    degree = length(T.types)
-    if degree == 1
-        ind = findoperation(x.name, operators.unaops)
-        return operators.unaops[ind]
-    elseif degree == 2
-        ind = findoperation(x.name, operators.binops)
-        return operators.binops[ind]
-    else
-        throw(AssertionError("Function $(String(x.name)) has degree > 2 !"))
+function convert_to_function(x::BasicSymbolic, operators::AbstractOperatorEnum)
+    if issym(x) && SymbolicUtils.symtype(x) <: SymbolicUtils.FnType
+        signature, _ = SymbolicUtils.fntype_X_Y(SymbolicUtils.symtype(x))
+        degree = length(signature.parameters)
+        name = nameof(x)
+        if degree == 1
+            ind = findoperation(name, operators.unaops)
+            return operators.unaops[ind]
+        elseif degree == 2
+            ind = findoperation(name, operators.binops)
+            return operators.binops[ind]
+        else
+            throw(AssertionError("Function $(String(name)) has degree > 2 !"))
+        end
     end
+    return x
 end
 
 # For normal operators, simply return the function itself:
@@ -90,7 +98,7 @@ function split_eq(
     op,
     args,
     operators::AbstractOperatorEnum,
-    ::Type{N}=Node;
+    (::Type{N})=Node;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
     # Deprecated:
     varMap=nothing,
@@ -119,7 +127,7 @@ function findoperation(op, ops)
 end
 
 function Base.convert(
-    ::typeof(SymbolicUtils.Symbolic),
+    ::typeof(BasicSymbolic),
     tree::Union{AbstractExpression,AbstractExpressionNode},
     operators::Union{AbstractOperatorEnum,Nothing}=nothing;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
@@ -141,22 +149,32 @@ end
 
 function Base.convert(
     ::Type{N},
-    expr::SymbolicUtils.Symbolic,
+    expr::BasicSymbolic,
     operators::AbstractOperatorEnum;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
 ) where {N<:AbstractExpressionNode}
     variable_names = deprecate_varmap(variable_names, nothing, :convert)
-    if !iscall(expr)
+    # Handle constants (v4 wraps numbers in Const variant)
+    if isconst(expr)
+        return constructorof(N)(; val=DEFAULT_NODE_TYPE(unwrap_const(expr)))
+    end
+    # Handle symbols (variables)
+    if issym(expr)
+        exprname = nameof(expr)
         if variable_names === nothing
-            s = String(expr.name)
+            s = String(exprname)
             # Verify it is of the format "x{num}":
             @assert(
                 occursin(r"^x\d+$", s),
                 "Variable name $s is not of the format x{num}. Please provide the `variable_names` explicitly."
             )
             return constructorof(N)(s)
         end
-        return constructorof(N)(String(expr.name), variable_names)
+        return constructorof(N)(String(exprname), variable_names)
+    end
+    # Handle function calls
+    if !iscall(expr)
+        error("Unknown symbolic expression type: $(typeof(expr))")
     end
 
     # First, we remove integer powers:
@@ -185,7 +203,7 @@ _node_type(::Type{E}) where {E<:AbstractExpression} = default_node_type(E)
 
 function Base.convert(
     ::Type{E},
-    x::Union{SymbolicUtils.Symbolic,Number},
+    x::Union{BasicSymbolic,Number},
     operators::AbstractOperatorEnum;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
     kws...,
@@ -209,10 +227,9 @@ will generate a symbolic equation in SymbolicUtils.jl format.
 - `operators::AbstractOperatorEnum`: OperatorEnum, which contains the operators used in the equation.
 - `variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing`: What variable names to use for
     each feature. Default is [x1, x2, x3, ...].
-- `index_functions::Bool=false`: Whether to generate special names for the
-    operators, which then allows one to convert back to a `AbstractExpressionNode` format
-    using `symbolic_to_node`.
-    (CURRENTLY UNAVAILABLE - See https://github.com/MilesCranmer/SymbolicRegression.jl/pull/84).
+- `index_functions::Bool=false`: Whether to represent custom operators by name as
+    uninterpreted SymbolicUtils function symbols. This allows round-tripping back to a
+    `AbstractExpressionNode` using `symbolic_to_node`.
 """
 function node_to_symbolic(
     tree::AbstractExpressionNode,
@@ -231,8 +248,8 @@ function node_to_symbolic(
     # Create a substitution tuple
     subs = Dict(
         [
-            SymbolicUtils.Sym{LiteralReal}(Symbol("x$(i)")) =>
-                SymbolicUtils.Sym{LiteralReal}(Symbol(variable_names[i])) for
+            SymbolicUtils.Sym{SymReal}(Symbol("x$(i)"); type=Number) =>
+                SymbolicUtils.Sym{SymReal}(Symbol(variable_names[i]); type=Number) for
             i in 1:length(variable_names)
         ]...,
     )
@@ -253,9 +270,9 @@ function node_to_symbolic(
 end
 
 function symbolic_to_node(
-    eqn::SymbolicUtils.Symbolic,
+    eqn::BasicSymbolic,
     operators::AbstractOperatorEnum,
-    ::Type{N}=Node;
+    (::Type{N})=Node;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
     # Deprecated:
     varMap=nothing,
@@ -268,7 +285,7 @@ function multiply_powers(eqn::Number)::Tuple{SYMBOLIC_UTILS_TYPES,Bool}
     return eqn, true
 end
 
-function multiply_powers(eqn::SymbolicUtils.Symbolic)::Tuple{SYMBOLIC_UTILS_TYPES,Bool}
+function multiply_powers(eqn::BasicSymbolic)::Tuple{SYMBOLIC_UTILS_TYPES,Bool}
     if !iscall(eqn)
         return eqn, true
     end
@@ -277,7 +294,7 @@ function multiply_powers(eqn::SymbolicUtils.Symbolic)::Tuple{SYMBOLIC_UTILS_TYPE
 end
 
 function multiply_powers(
-    eqn::SymbolicUtils.Symbolic, op::F
+    eqn::BasicSymbolic, op::F
 )::Tuple{SYMBOLIC_UTILS_TYPES,Bool} where {F}
     args = SymbolicUtils.arguments(eqn)
     nargs = length(args)
@@ -291,15 +308,23 @@ function multiply_powers(
         @return_on_false complete eqn
         @return_on_false is_valid(l) eqn
         n = args[2]
-        if typeof(n) <: Integer
-            if n == 1
+        # In SymbolicUtils v4, integer constants are wrapped in Const
+        n_val = if isconst(n)
+            unwrap_const(n)
+        elseif typeof(n) <: Integer
+            n
+        else
+            nothing
+        end
+        if n_val !== nothing && typeof(n_val) <: Integer
+            if n_val == 1
                 return l, true
-            elseif n == -1
+            elseif n_val == -1
                 return 1.0 / l, true
-            elseif n > 1
-                return reduce(*, [l for i in 1:n]), true
-            elseif n < -1
-                return reduce(/, vcat([1], [l for i in 1:abs(n)])), true
+            elseif n_val > 1
+                return reduce(*, [l for i in 1:n_val]), true
+            elseif n_val < -1
+                return reduce(/, vcat([1], [l for i in 1:abs(n_val)])), true
             else
                 return 1.0, true
             end