Merge pull request #81 from SymbolicML/parametric-expressions

Create `ParametricExpression` type

Author: MilesCranmer

Project.toml

@@ -7,12 +7,14 @@ version = "0.18.0-alpha.1"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
+Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 
 [weakdeps]
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
@@ -33,13 +35,15 @@ Bumper = "0.6"
 ChainRulesCore = "1"
 Compat = "3.37, 4"
 DispatchDoctor = "0.4"
+Interfaces = "0.3"
 LoopVectorization = "0.12"
 MacroTools = "0.4, 0.5"
 Optim = "0.19, 1"
 PackageExtensionCompat = "1"
 PrecompileTools = "1"
 Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5, 2"
+TestItems = "0.1"
 Zygote = "0.6"
 julia = "1.6"
 

README.md

@@ -4,7 +4,7 @@
 
 *Ridiculously fast dynamic expressions.*
 
-[![](https://img.shields.io/badge/docs-dev-blue.svg)](https://symbolicml.org/DynamicExpressions.jl/dev) [![CI](https://github.com/SymbolicML/DynamicExpressions.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/SymbolicML/DynamicExpressions.jl/actions/workflows/CI.yml) [![Coverage Status](https://coveralls.io/repos/github/SymbolicML/DynamicExpressions.jl/badge.svg?branch=master)](https://coveralls.io/github/SymbolicML/DynamicExpressions.jl?branch=master) [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
+[![](https://img.shields.io/badge/docs-dev-blue.svg)](https://symbolicml.org/DynamicExpressions.jl/dev) [![CI](https://github.com/SymbolicML/DynamicExpressions.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/SymbolicML/DynamicExpressions.jl/actions/workflows/CI.yml) [![Coverage Status](https://coveralls.io/repos/github/SymbolicML/DynamicExpressions.jl/badge.svg?branch=master)](https://coveralls.io/github/SymbolicML/DynamicExpressions.jl?branch=master) [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl) [![DispatchDoctor](https://img.shields.io/badge/%F0%9F%A9%BA_tested_with-DispatchDoctor.jl-blue?labelColor=white)](https://github.com/MilesCranmer/DispatchDoctor.jl)
 
 DynamicExpressions.jl is the backbone of [SymbolicRegression.jl](https://github.com/MilesCranmer/SymbolicRegression.jl) and
 [PySR](https://github.com/MilesCranmer/PySR).

benchmark/benchmarks.jl

@@ -113,9 +113,18 @@ end
     PACKAGE_VERSION < v"0.14.0" && return :(copy_node(t; preserve_sharing=preserve_sharing))
     return :(copy_node(t))  # Assume type used to infer sharing
 end
-@generated function get_set_constants!(tree)
-    !(@isdefined set_constants!) && return :(set_constants(tree, get_constants(tree)))
-    return :(set_constants!(tree, get_constants(tree)))
+@generated function get_set_constants!(tree::N) where {T,N<:AbstractExpressionNode{T}}
+    if !(@isdefined set_constants!)
+        return :(set_constants(tree, get_constants(tree)))
+    elseif hasmethod(set_constants!, Tuple{N, Vector{T}})
+        return :(set_constants!(tree, get_constants(tree)))
+    else
+        return quote
+            let (x, refs) = get_constants(tree)
+                set_constants!(tree, x, refs)
+            end
+        end
+    end
 end
 #! format: on
 

docs/src/types.md

@@ -42,7 +42,7 @@ Note that you are free to use the `Node` constructors directly.
 This is a more robust approach, and should be used when creating libraries
 which use `DynamicExpressions.jl`.
 
-## Equations
+## Nodes
 
 Equations are specified as binary trees with the `Node` type, defined
 as follows:
@@ -75,7 +75,7 @@ You can create a copy of a node with `copy_node`:
 copy_node
 ```
 
-## Graph-Like Equations
+## Graph Nodes
 
 You can describe an equation as a *graph* rather than a tree
 by using the `GraphNode` type:
@@ -144,3 +144,55 @@ which is more generic but does not have all of the same methods:
 ```@docs
 AbstractNode{T}
 ```
+
+## Expressions
+
+A higher-level user-facing type is the `Expression`:
+
+```@docs
+Expression
+```
+
+This is a subtype of `AbstractExpression`.
+
+```@docs
+AbstractExpression
+```
+
+which can be used for defining custom types, such as the `ParametricExpression`:
+
+```@docs
+ParametricExpression
+ParametricNode
+```
+
+## Interfaces
+
+The interfaces for `AbstractExpression` and `AbstractExpressionNode` are
+tested using Interfaces.jl. You can see the interfaces with:
+
+```@docs
+DynamicExpressions.ExpressionInterface
+DynamicExpressions.NodeInterface
+```
+
+You can declare a new type as implementing these with, e.g.,
+
+```julia
+using DynamicExpressions: ExpressionInterface, all_ei_methods_except
+using Interface: @implements, Arguments, Interface
+
+# Add all optional methods:
+valid_optional_methods = all_ei_methods_except(())
+
+@implements ExpressionInterface{valid_optional_methods} MyCustomExpression [Arguments()]
+```
+
+You can then test the interface is implemented correctly using, for example,
+
+```julia
+@test Interface.test(ExpressionInterface, MyCustomExpression, [ex::MyCustomExpression])
+```
+
+Note that this may not flag all potential issues, so be sure to still read the details about
+what methods can be implemented and customized.

ext/DynamicExpressionsOptimExt.jl

@@ -1,6 +1,12 @@
 module DynamicExpressionsOptimExt
 
-using DynamicExpressions: AbstractExpressionNode, filter_map, eval_tree_array
+using DynamicExpressions:
+    AbstractExpression,
+    AbstractExpressionNode,
+    filter_map,
+    eval_tree_array,
+    get_constants,
+    set_constants!
 using Compat: @inline
 
 import Optim: Optim, OptimizationResults, NLSolversBase
@@ -12,7 +18,7 @@ import Optim: Optim, OptimizationResults, NLSolversBase
 Optimization results for an expression, which wraps the base optimization results
 on a vector of constants.
 """
-struct ExpressionOptimizationResults{R<:OptimizationResults,N<:AbstractExpressionNode} <: OptimizationResults
+struct ExpressionOptimizationResults{R<:OptimizationResults,N<:Union{AbstractExpressionNode,AbstractExpression}} <: OptimizationResults
     _results::R # The raw results from Optim.
     tree::N # The final expression tree
 end
@@ -33,36 +39,34 @@ end
 
 """Wrap function or objective with insertion of values of the constant nodes."""
 function wrap_func(
-    f::F, tree::N, constant_refs::AbstractArray
-) where {F<:Function,T,N<:AbstractExpressionNode{T}}
+    f::F, tree::N, refs
+) where {F<:Function,T,N<:Union{AbstractExpressionNode{T},AbstractExpression{T}}}
     function wrapped_f(args::Vararg{Any,M}) where {M}
-        first_args = args[1:(end - 1)]
-        x = last(args)
-        @inbounds for i in eachindex(constant_refs, x)
-            constant_refs[i][].val = x[i]
-        end
+        first_args = args[begin:(end - 1)]
+        x = args[end]
+        set_constants!(tree, x, refs)
         return @inline(f(first_args..., tree))
     end
     return wrapped_f
 end
 function wrap_func(
-    ::Nothing, tree::N, constant_refs::AbstractArray
-) where {N<:AbstractExpressionNode}
+    ::Nothing, tree::N, refs
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
     return nothing
 end
 function wrap_func(
-    f::NLSolversBase.InplaceObjective, tree::N, constant_refs::AbstractArray
-) where {N<:AbstractExpressionNode}
+    f::NLSolversBase.InplaceObjective, tree::N, refs
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
     # Some objectives, like `Optim.only_fg!(fg!)`, are not functions but instead
     # `InplaceObjective`. These contain multiple functions, each of which needs to be
     # wrapped. Some functions are `nothing`; those can be left as-is.
     @assert fieldnames(NLSolversBase.InplaceObjective) == (:df, :fdf, :fgh, :hv, :fghv)
     return NLSolversBase.InplaceObjective(
-        wrap_func(f.df, tree, constant_refs),
-        wrap_func(f.fdf, tree, constant_refs),
-        wrap_func(f.fgh, tree, constant_refs),
-        wrap_func(f.hv, tree, constant_refs),
-        wrap_func(f.fghv, tree, constant_refs),
+        wrap_func(f.df, tree, refs),
+        wrap_func(f.fdf, tree, refs),
+        wrap_func(f.fgh, tree, refs),
+        wrap_func(f.hv, tree, refs),
+        wrap_func(f.fghv, tree, refs),
     )
 end
 
@@ -74,24 +78,29 @@ Returns an `ExpressionOptimizationResults` object, which wraps the base
 optimization results on a vector of constants. You may use `res.minimizer`
 to view the optimized expression tree.
 """
-function Optim.optimize(f::F, tree::AbstractExpressionNode, args...; kwargs...) where {F}
+function Optim.optimize(
+    f::F, tree::Union{AbstractExpressionNode,AbstractExpression}, args...; kwargs...
+) where {F}
     return Optim.optimize(f, nothing, tree, args...; kwargs...)
 end
 function Optim.optimize(
-    f::F, g!::G, tree::AbstractExpressionNode, args...; kwargs...
+    f::F, g!::G, tree::Union{AbstractExpressionNode,AbstractExpression}, args...; kwargs...
 ) where {F,G}
     return Optim.optimize(f, g!, nothing, tree, args...; kwargs...)
 end
 function Optim.optimize(
-    f::F, g!::G, h!::H, tree::AbstractExpressionNode{T}, args...; make_copy=true, kwargs...
+    f::F,
+    g!::G,
+    h!::H,
+    tree::Union{AbstractExpressionNode{T},AbstractExpression{T}},
+    args...;
+    make_copy=true,
+    kwargs...,
 ) where {F,G,H,T}
     if make_copy
         tree = copy(tree)
     end
-    constant_refs = filter_map(
-        t -> t.degree == 0 && t.constant, t -> Ref(t), tree, Ref{typeof(tree)}
-    )
-    x0 = T[copy(t[].val) for t in constant_refs]
+    x0, refs = get_constants(tree)
     if !isnothing(h!)
         throw(
             ArgumentError(
@@ -101,20 +110,14 @@ function Optim.optimize(
         )
     end
     base_res = if isnothing(g!)
-        Optim.optimize(wrap_func(f, tree, constant_refs), x0, args...; kwargs...)
+        Optim.optimize(wrap_func(f, tree, refs), x0, args...; kwargs...)
     else
         Optim.optimize(
-            wrap_func(f, tree, constant_refs),
-            wrap_func(g!, tree, constant_refs),
-            x0,
-            args...;
-            kwargs...,
+            wrap_func(f, tree, refs), wrap_func(g!, tree, refs), x0, args...; kwargs...
         )
     end
     minimizer = Optim.minimizer(base_res)
-    @inbounds for i in eachindex(constant_refs, minimizer)
-        constant_refs[i][].val = minimizer[i]
-    end
+    set_constants!(tree, minimizer, refs)
     return ExpressionOptimizationResults(base_res, tree)
 end
 

ext/DynamicExpressionsSymbolicUtilsExt.jl

@@ -6,6 +6,7 @@ import DynamicExpressions.NodeModule:
 import DynamicExpressions.OperatorEnumModule: AbstractOperatorEnum
 import DynamicExpressions.UtilsModule: isgood, isbad, deprecate_varmap
 import DynamicExpressions.ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node
+import DynamicExpressions: AbstractExpression, get_tree, get_operators
 
 const SYMBOLIC_UTILS_TYPES = Union{<:Number,SymbolicUtils.Symbolic{<:Number}}
 const SUPPORTED_OPS = (cos, sin, exp, cot, tan, csc, sec, +, -, *, /)
@@ -109,7 +110,7 @@ end
 
 function Base.convert(
     ::typeof(SymbolicUtils.Symbolic),
-    tree::AbstractExpressionNode,
+    tree::Union{AbstractExpression,AbstractExpressionNode},
     operators::AbstractOperatorEnum;
     variable_names::Union{Array{String,1},Nothing}=nothing,
     index_functions::Bool=false,
@@ -206,6 +207,11 @@ function node_to_symbolic(
     )
     return substitute(expr, subs)
 end
+function node_to_symbolic(
+    tree::AbstractExpression, operators::Union{AbstractOperatorEnum,Nothing}=nothing; kws...
+)
+    return node_to_symbolic(get_tree(tree), get_operators(tree, operators); kws...)
+end
 
 function symbolic_to_node(
     eqn::SymbolicUtils.Symbolic,

src/DynamicExpressions.jl

@@ -15,9 +15,10 @@ using DispatchDoctor: @stable, @unstable
     include("EvaluationHelpers.jl")
     include("Simplify.jl")
     include("OperatorEnumConstruction.jl")
-    include("Random.jl")
     include("Expression.jl")
+    include("Random.jl")
     include("Parse.jl")
+    include("ParametricExpression.jl")
 end
 
 import PackageExtensionCompat: @require_extensions
@@ -65,11 +66,20 @@ import .NodeModule:
 @reexport import .EvaluationHelpersModule
 @reexport import .ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node
 @reexport import .RandomModule: NodeSampler
-@reexport import .ExpressionModule: AbstractExpression, Expression
-# Not for export; just for overloading
-import .ExpressionModule: get_tree, get_operators, get_variable_names, Metadata
+@reexport import .ExpressionModule: AbstractExpression, Expression, with_tree
+import .ExpressionModule:
+    get_tree, get_operators, get_variable_names, Metadata, default_node_type, node_type
 @reexport import .ParseModule: @parse_expression, parse_expression
 import .ParseModule: parse_leaf
+@reexport import .ParametricExpressionModule: ParametricExpression, ParametricNode
+
+@stable default_mode = "disable" begin
+    include("Interfaces.jl")
+    include("PatchMethods.jl")
+end
+
+import .InterfacesModule:
+    ExpressionInterface, NodeInterface, all_ei_methods_except, all_ni_methods_except
 
 function __init__()
     @require_extensions