Overload Optim.optimize for Node as x0

Author: MilesCranmer

Project.toml

@@ -12,6 +12,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

src/ConstantOptimization.jl

@@ -0,0 +1,67 @@
+module ConstantOptimizationModule
+
+import Compat: @inline
+import Optim: optimize
+import ..EquationModule: Node
+import ..EvaluateEquationModule: eval_tree_array
+
+"""Wrap f with insertion of values of the constant nodes."""
+function get_wrapped_f(
+    f::F, tree::N, constant_nodes::AbstractArray{N}
+) where {F,T,N<:Node{T}}
+    function wrapped_f(x)
+        for (ci, xi) in zip(constant_nodes, x)
+            ci.val::T = xi::T
+        end
+        return @inline(f(tree))
+    end
+    return wrapped_f
+end
+
+"""Wrap g! or h! with insertion of values of the constant nodes."""
+function get_wrapped_gh!(
+    gh!::GH, tree::N, constant_nodes::AbstractArray{N}
+) where {GH,T,N<:Node{T}}
+    function wrapped_gh!(G, x)
+        for (ci, xi) in zip(constant_nodes, x)
+            ci.val::T = xi::T
+        end
+        @inline(gh!(G, tree))
+        return nothing
+    end
+    return wrapped_gh!
+end
+
+function optimize(f::F, g!::G, h!::H, tree::Node{T}, args...; kwargs...) where {F,G,H,T}
+    constant_nodes = filter(t -> t.degree == 0 && t.constant, tree)
+    x0 = [t.val::T for t in constant_nodes]
+    if g! === nothing
+        @assert h! === nothing
+        return optimize(get_wrapped_f(f, tree, constant_nodes), x0, args...; kwargs...)
+    elseif h! === nothing
+        return optimize(
+            get_wrapped_f(f, tree, constant_nodes),
+            get_wrapped_gh!(g!, tree, constant_nodes),
+            x0,
+            args...;
+            kwargs...,
+        )
+    else
+        return optimize(
+            get_wrapped_f(f, tree, constant_nodes),
+            get_wrapped_gh!(g!, tree, constant_nodes),
+            get_wrapped_gh!(h!, tree, constant_nodes),
+            x0,
+            args...;
+            kwargs...,
+        )
+    end
+end
+function optimize(f::F, g!::G, tree::Node, args...; kwargs...) where {F,G}
+    return optimize(f, g!, nothing, tree, args...; kwargs...)
+end
+function optimize(f::F, tree::Node, args...; kwargs...) where {F}
+    return optimize(f, nothing, tree, args...; kwargs...)
+end
+
+end

src/DynamicExpressions.jl

@@ -1,5 +1,7 @@
 module DynamicExpressions
 
+using Requires: @require
+
 include("Utils.jl")
 include("OperatorEnum.jl")
 include("Equation.jl")
@@ -10,8 +12,9 @@ include("EvaluationHelpers.jl")
 include("InterfaceSymbolicUtils.jl")
 include("SimplifyEquation.jl")
 include("OperatorEnumConstruction.jl")
+@require Optim = "429524aa-4258-5aef-a3af-852621145aeb" include("ConstantOptimization.jl")
 
-using Reexport
+using Reexport: @reexport
 @reexport import .EquationModule:
     Node,
     string_tree,
@@ -41,6 +44,8 @@ using Reexport
 @reexport import .InterfaceSymbolicUtilsModule: node_to_symbolic, symbolic_to_node
 @reexport import .SimplifyEquationModule: combine_operators, simplify_tree
 @reexport import .EvaluationHelpersModule
+@require Optim = "429524aa-4258-5aef-a3af-852621145aeb" @reexport import .ConstantOptimizationModule:
+    optimize
 
 include("deprecated.jl")