Remove toexpr overload

Author: tansongchen

Project.toml

@@ -1,7 +1,7 @@
 name = "TaylorDiff"
 uuid = "b36ab563-344f-407b-a36a-4f200bebf99c"
 authors = ["Songchen Tan <[email protected]>"]
-version = "0.3.2"
+version = "0.3.3"
 
 [deps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"

examples/Project.toml

@@ -1,8 +1,14 @@
 [deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+LayeredLayouts = "f4a74d36-062a-4d48-97cd-1356bad1de4e"
+NetworkLayout = "46757867-2c16-5918-afeb-47bfcb05e46a"
 ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
 TaylorIntegration = "92b13dbe-c966-51a2-8445-caca9f8a7d42"
 TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
+TikzGraphs = "b4f28e30-c73f-5eaf-a395-8a9db949a742"

examples/integration.jl

@@ -90,7 +90,12 @@ function build_jetcoeffs(f::ODEFunction{iip}, p, ::Val{P}, length = nothing) whe
         d = get_coefficient(fu, index - 1) / index
         u = append_coefficient(u, d)
     end
-    build_function(u, u0, t0; expression = Val(false), cse = true)
+    u_term = make_term.(u)
+    build_function(u_term, u0, t0; expression = Val(false), cse = true)
+end
+
+function make_term(a)
+    term(TaylorScalar, Symbolics.unwrap(a.value), map(Symbolics.unwrap, a.partials))
 end
 
 function simplify_scalar_test()
@@ -111,6 +116,13 @@ function simplify_array_test()
     @btime $fast_oop($prob.u0, $t0)
 end
 
+P = 6
+prob = prob_ode_lotkavolterra
+t0 = prob.tspan[1]
+@btime jetcoeffs($prob.f, $prob.u0, $prob.p, $t0, Val($P))
+fast_oop, fast_iip = build_jetcoeffs(prob.f, prob.p, Val(10), length(prob.u0));
+@btime $fast_oop($prob.u0, $t0)
+
 @generated function evaluate_polynomial(t::TaylorScalar{T, P}, z) where {T, P}
     ex = :(v[$(P + 1)])
     for i in P:-1:1

examples/plot_expression.jl

@@ -0,0 +1,69 @@
+using Symbolics, SymbolicUtils
+using Graphs, GraphMakie, CairoMakie, LayeredLayouts
+using Printf
+
+function generate_dag(rexs...)
+    exs = map(Symbolics.unwrap, rexs)
+    dag = SimpleDiGraph()
+    vertex_map = IdDict{Any, Int}()
+    label_map = Dict{Int, String}()
+    unicode_replace = Dict(
+        "-" => "−",
+        "*" => "×",
+        "/" => "÷"
+    )
+    function dfs(node)
+        if haskey(vertex_map, node)
+            return vertex_map[node]
+        end
+        add_vertex!(dag)
+        v = nv(dag)
+        vertex_map[node] = v
+        if iscall(node)
+            op = operation(node)
+            args = map(dfs, arguments(node))
+            for arg in args
+                add_edge!(dag, v, arg)
+            end
+            rop = repr(op)
+            label_map[v] = get(unicode_replace, rop, rop)
+        elseif node isa Number
+            label_map[v] = @sprintf "%.1g" node
+        else
+            label_map[v] = repr(node)
+        end
+        return v
+    end
+    for ex in exs
+        dfs(ex)
+    end
+    return dag, label_map
+end
+
+function plot_dag(dag, label_map)
+    xs, ys, paths = solve_positions(Zarate(), dag)
+    for (key, value) in paths
+        paths[key] = (value[2], -value[1])
+    end
+    lay = Point.(zip(ys, -xs))
+    wp = [Point2f.(zip(paths[e]...)) for e in edges(dag)]
+    fig, ax, p = graphplot(
+        dag; layout = lay, ilabels = [label_map[v] for v in 1:nv(dag)],
+        waypoints = wp, node_color = :white)
+    hidedecorations!(ax)
+    hidespines!(ax)
+    fig
+end
+
+@variables a b c d
+x = (a + b) * (c + d)
+y = (a - b) * (c - d)
+z = (a + b) * (c - d)
+w = (a - b) * (c + d)
+p = x + y
+q = z - w
+rex = (x + y) / (z - w)
+
+dag, label_map = generate_dag(p, q)
+fig = plot_dag(dag, label_map)
+save("expression_dag.png", fig)

examples/scaling.jl

@@ -0,0 +1,45 @@
+using Symbolics, SymbolicUtils
+
+@variables x
+p = Symbolics.variables(:p, 0:20)
+
+struct A{T, P} <: Real
+    a::T
+    b::NTuple{P, T}
+end
+
+function make_nested_expressions(order)
+    exprs = [x]
+    for i in 1:order
+        term = p[i + 1]
+        for j in i:-1:1
+            term += p[j] * exprs[i - j + 1]
+        end
+        push!(exprs, term)
+    end
+    exprs
+end
+
+for order in 1:17
+    final_expr = make_nested_expressions(order)[end]
+    print("Order: $order")
+    build_function(final_expr, x, p; expression = Val(false), cse = true)
+    @time build_function(final_expr, x, p; expression = Val(false), cse = true)
+end
+
+for order in 1:17
+    exprs = make_nested_expressions(order)
+    tuple_expr = Symbolics.Code.MakeTuple(Symbolics.unwrap.(exprs))
+    print("Order: $order")
+    @time f = build_function(tuple_expr, x, p; expression = Val(false), cse = true)
+    @show f(1.0, 1.0:5.0)
+end
+
+# But it still doesn't scale with the unwrap trick. What else can I do?
+for order in 1:17
+    exprs = make_nested_expressions(order)
+    struct_expr = term(A, Symbolics.unwrap(exprs[1]), (Symbolics.unwrap.(exprs[2:end])...,))
+    print("Order: $order")
+    @time f = build_function(struct_expr, x, p; expression = Val(false), cse = true)
+    @show f(1.0, 1.0:5.0)
+end

examples/test.jl

@@ -57,7 +57,8 @@ function rrule(::typeof(Base.getindex), a::TaylorArray, i::Int...)
 end
 
 function rrule(::typeof(*), A::AbstractMatrix{S},
-        t::AbstractVector{TaylorScalar{T, N}}) where {N, S, T}
+        t::AbstractVector{TaylorScalar{T, N}}) where {
+        N, S <: Union{Real, Complex}, T <: Union{Real, Complex}}
     project_A = ProjectTo(A)
     function gemv_pullback(x̄)
         x̂ = reinterpret(reshape, T, x̄)
@@ -68,7 +69,8 @@ function rrule(::typeof(*), A::AbstractMatrix{S},
 end
 
 function rrule(::typeof(*), A::AbstractMatrix{S},
-        B::AbstractMatrix{TaylorScalar{T, N}}) where {N, S, T}
+        B::AbstractMatrix{TaylorScalar{T, N}}) where {
+        N, S <: Union{Real, Complex}, T <: Union{Real, Complex}}
     project_A = ProjectTo(A)
     project_B = ProjectTo(B)
     function gemm_pullback(x̄)

src/chainrules.jl

@@ -119,7 +119,8 @@ sincos(t::TaylorScalar) = (sin(t), cos(t))
 @inline *(a::TaylorScalar, b::Number) = TaylorScalar(value(a) * b, partials(a) .* b)
 @inline /(a::TaylorScalar, b::Number) = TaylorScalar(value(a) / b, partials(a) ./ b)
 
-const AMBIGUOUS_TYPES = (AbstractFloat, Irrational, Integer, Rational, Real, Complex, RoundingMode)
+const AMBIGUOUS_TYPES = (
+    AbstractFloat, Irrational, Integer, Rational, Real, Complex, RoundingMode)
 
 for op in [:>, :<, :(==), :(>=), :(<=)]
     for R in AMBIGUOUS_TYPES

src/primitive.jl

@@ -3,17 +3,11 @@
 
 using ChainRules
 using ChainRulesCore
-using Symbolics: Symbolics, @variables, @rule, @register_symbolic, unwrap, isdiv
+using Symbolics: Symbolics, @variables, @rule, unwrap, isdiv
 using SymbolicUtils.Code: toexpr
 using MacroTools
 using MacroTools: prewalk, postwalk
 
-@register_symbolic TaylorScalar(x, y)
-function Symbolics.Code.toexpr(t::TaylorScalar, st)
-    :($TaylorScalar($(Symbolics.Code.toexpr(t.value, st)),
-        $(Symbolics.Code.toexpr(Symbolics.Code.MakeTuple(t.partials), st))))
-end
-
 """
 Pick a strategy for raising the derivative of a function.
 If the derivative is like 1 over something, raise with the division rule;