Add several helper functions; improve the integration demo

Author: tansongchen

Project.toml

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

examples/Project.toml

@@ -0,0 +1,8 @@
+[deps]
+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"

examples/integration.jl

@@ -1,100 +1,111 @@
-using TaylorDiff: TaylorDiff, TaylorScalar, TaylorArray, make_seed, value, partials, flatten
-using TaylorSeries
-using TaylorIntegration: jetcoeffs!
-using BenchmarkTools
+using TaylorDiff: TaylorDiff, TaylorScalar, make_seed, flatten, get_coefficient,
+                  set_coefficient, append_coefficient
+using TaylorSeries, TaylorIntegration
+using ODEProblemLibrary, OrdinaryDiffEq, BenchmarkTools, Symbolics
 
-"""
-No magic, just a type-stable way to generate a new object since TaylorScalar is immutable
-"""
-function setindex(x::TaylorScalar{T, P}, index, d) where {T, P}
-    v = flatten(x)
-    ntuple(i -> i == index + 1 ? d : v[i], Val(P + 1)) |> TaylorScalar
-end
-
-function setindex(x::TaylorArray{T, N, A, P}, index, d) where {T, N, A, P}
-    v = flatten(x)
-    ntuple(i -> i == index + 1 ? d : v[i], Val(P + 1)) |> TaylorArray
-end
-
-"""
-Computes the taylor integration of order P
+# There are two ways to compute the Taylor coefficients of a ODE solution
+# 1. Using naive repeated differentiation
+# 2. First simplify the expansion using Symbolics and then evaluate the expression
 
-- `f`: ODE function
-- `t`: constructed by TaylorScalar{P}(t0, one(t0))
-- `x0`: initial value
-- `p`: parameters
 """
-function my_jetcoeffs(f, t::TaylorScalar{T, P}, x0, p) where {T, P}
-    x = x0 isa AbstractArray ? TaylorArray{P}(x0) : TaylorScalar{P}(x0)
-    for index in 1:P # computes x.partials[index]
-        fx = f(x, p, t)
-        d = index == 1 ? value(fx) : partials(fx)[index - 1] / index
-        x = setindex(x, index, d)
-    end
-    x
-end
-
-"""
-Computes the taylor integration of order P
+# The first method
 
-- `f!`: ODE function, in non-allocating form
-- `t`: constructed by TaylorScalar{P}(t0, one(t0))
-- `x0`: initial value
-- `p`: parameters
+For ODE u' = f(u, p, t) and initial condition (u0, t0), computes Taylor expansion of the solution `u` up to order `P` using repeated differentiation.
 """
-function my_jetcoeffs!(f!, t::TaylorScalar{T, P}, x0, p) where {T, P}
-    x = x0 isa AbstractArray ? TaylorArray{P}(x0) : TaylorScalar{P}(x0)
-    out = similar(x)
-    for index in 1:P # computes x.partials[index]
-        f!(out, x, p, t)
-        d = index == 1 ? value(out) : partials(out)[index - 1] / index
-        x = setindex(x, index, d)
+function jetcoeffs(f::ODEFunction{iip}, u0, p, t0, ::Val{P}) where {P, iip}
+    t = TaylorScalar{P}(t0, one(t0))
+    u = make_seed(u0, zero(u0), Val(P))
+    fu = copy(u)
+    for index in 1:P
+        if iip
+            f(fu, u, p, t)
+        else
+            fu = f(u, p, t)
+        end
+        d = get_coefficient(fu, index - 1) / index
+        u = set_coefficient(u, index, d)
     end
-    x
+    u
 end
 
 function scalar_test()
-    f(x, p, t) = x * x
-    x0 = 0.1
-    t0 = 0.0
     P = 6
-
+    prob = prob_ode_linear
+    t0 = prob.tspan[1]
     # TaylorIntegration test
-    t = t0 + Taylor1(typeof(t0), P)
-    x = Taylor1(x0, P)
-    @btime jetcoeffs!($f, $t, $x, nothing)
+    ts = t0 + Taylor1(typeof(t0), P)
+    u_ts = Taylor1(prob.u0, P)
+    @btime TaylorIntegration.jetcoeffs!($prob.f, $ts, $u_ts, $prob.p)
 
     # TaylorDiff test
-    td = TaylorScalar{P}(t0, one(t0))
-    @btime my_jetcoeffs($f, $td, $x0, nothing)
-
-    result = my_jetcoeffs(f, td, x0, nothing)
-    @assert x.coeffs ≈ collect(flatten(result))
+    @btime jetcoeffs($prob.f, $prob.u0, $prob.p, $t0, Val($P))
+    u_td = jetcoeffs(prob.f, prob.u0, prob.p, t0, Val(P))
+    @assert u_ts.coeffs ≈ collect(flatten(u_td))
 end
 
 function array_test()
-    function lorenz(du, u, p, t)
-        du[1] = 10.0(u[2] - u[1])
-        du[2] = u[1] * (28.0 - u[3]) - u[2]
-        du[3] = u[1] * u[2] - (8 / 3) * u[3]
-        return nothing
-    end
-    u0 = [1.0; 0.0; 0.0]
-    t0 = 0.0
     P = 6
-
+    prob = prob_ode_lotkavolterra
+    t0 = prob.tspan[1]
     # TaylorIntegration test
-    t = t0 + Taylor1(typeof(t0), P)
-    u = [Taylor1(x, P) for x in u0]
-    du = similar(u)
-    uaux = similar(u)
-    @btime jetcoeffs!($lorenz, $t, $u, $du, $uaux, nothing)
+    ts = t0 + Taylor1(typeof(t0), P)
+    u_ts = [Taylor1(x, P) for x in prob.u0]
+    du = similar(u_ts)
+    uaux = similar(u_ts)
+    @btime TaylorIntegration.jetcoeffs!($prob.f, $ts, $u_ts, $du, $uaux, $prob.p)
 
     # TaylorDiff test
-    td = TaylorScalar{P}(t0, one(t0))
-    @btime my_jetcoeffs!($lorenz, $td, $u0, nothing)
-    result = my_jetcoeffs!(lorenz, td, u0, nothing)
-    for i in eachindex(u)
-        @assert u[i].coeffs ≈ collect(flatten(result[i]))
+    @btime jetcoeffs($prob.f, $prob.u0, $prob.p, $t0, Val($P))
+    u_td = jetcoeffs(prob.f, prob.u0, prob.p, t0, Val(P))
+    for i in eachindex(u_ts)
+        @assert u_ts[i].coeffs ≈ collect(flatten(u_td[i]))
+    end
+end
+
+"""
+# The second method
+
+For ODE u' = f(u, p, t) and initial condition (u0, t0), symbolically computes Taylor expansion of the solution `u` up to order `P`, and then builds a function to evaluate the expression.
+"""
+function build_jetcoeffs(f::ODEFunction{iip}, p, ::Val{P}, length = nothing) where {P, iip}
+    @variables t0::Real
+    u0 = isnothing(length) ? Symbolics.variable(:u0) : Symbolics.variables(:u0, 1:length)
+    if iip
+        @assert length isa Integer
+        f0 = similar(u0)
+        f(f0, u0, p, t0)
+    else
+        f0 = f(u0, p, t0)
     end
+    u = TaylorDiff.make_seed(u0, f0, Val(1))
+    for index in 2:P
+        t = TaylorScalar{index - 1}(t0, one(t0))
+        if iip
+            fu = similar(u)
+            f(fu, u, p, t)
+        else
+            fu = f(u, p, t)
+        end
+        d = get_coefficient(fu, index - 1) / index
+        u = append_coefficient(u, d)
+    end
+    build_function(u, u0, t0; expression = Val(false), cse = true)
+end
+
+function simplify_scalar_test()
+    P = 6
+    prob = prob_ode_linear
+    t0 = prob.tspan[1]
+    @btime jetcoeffs($prob.f, $prob.u0, $prob.p, $t0, Val($P))
+    fast_jetcoeffs = build_jetcoeffs(prob.f, prob.p, Val(P))
+    @btime $fast_jetcoeffs($prob.u0, $t0)
+end
+
+function simplify_array_test()
+    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(P), length(prob.u0))
+    @btime $fast_oop($prob.u0, $t0)
 end

examples/multi.jl

@@ -0,0 +1,11 @@
+using TaylorDiff
+
+f(x, y) = exp(x) * exp(y)
+x0, y0 = 0.0, 0.0
+x_dx = TaylorScalar{2}(x0, one(x0))
+x_dxdy = TaylorScalar{2}(x_dx, zero(x_dx))
+y_dx = TaylorScalar{2}(y0)
+y_dxdy = TaylorScalar{2}(y_dx, one(y_dx))
+
+result = f(x_dxdy, y_dxdy)
+map(TaylorDiff.flatten, TaylorDiff.flatten(result))

examples/test.jl

@@ -0,0 +1,34 @@
+using Symbolics, SymbolicUtils
+struct A1{T} <: Real
+    x::T
+    y::T
+    function A1(value::T, partials::T) where {T}
+        new{T}(value, partials)
+    end
+end
+
+struct A2{T, P} <: Real
+    x::T
+    y::NTuple{P, T}
+    function A2(value::T, partials::NTuple{P, T}) where {T, P}
+        new{T, P}(value, partials)
+    end
+end
+
+@register_symbolic A1(x, y)
+@register_symbolic A2(x, y)
+@variables x::Real
+a1 = A1(x, x^2)
+a2 = A2(x, (x^2, x^3))
+
+fa1 = build_function(a1, x; expression = Val(false))
+fa1(2.0)
+fa2 = build_function(a2, x; expression = Val(false))
+fa2(2.0)
+
+function Symbolics.Code.toexpr(a2::A2, st)
+    :($A2($(Symbolics.Code.toexpr(a2.x, st)),
+        $(Symbolics.Code.toexpr(Symbolics.Code.MakeTuple(a2.y), st))))
+end
+fa2 = build_function(a2, x; expression = Val(false))
+fa2(2.0)

src/TaylorDiff.jl

@@ -14,9 +14,9 @@ can_taylorize(::Type) = false
                         " If the type behaves as a scalar, define TaylorDiff.can_taylorize(::Type{$V}) = true."))
 end
 
-include("utils.jl")
 include("scalar.jl")
 include("array.jl")
+include("utils.jl")
 include("primitive.jl")
 include("codegen.jl")
 include("derivative.jl")

src/primitive.jl

@@ -9,6 +9,8 @@ Taylor = Union{TaylorScalar, TaylorArray}
 
 @inline value(t::Taylor) = t.value
 @inline partials(t::Taylor) = t.partials
+@inline value(ts::AbstractArray{<:Taylor}) = map(value, ts)
+@inline partials(ts::AbstractArray{<:Taylor}) = map(partials, ts)
 @inline @generated extract_derivative(t::Taylor, ::Val{P}) where {P} = :(t.partials[P] *
                                                                          $(factorial(P)))
 @inline extract_derivative(a::AbstractArray{<:TaylorScalar}, p) = map(
@@ -19,6 +21,35 @@ Taylor = Union{TaylorScalar, TaylorArray}
 
 @inline flatten(t::Taylor) = (value(t), partials(t)...)
 
+get_coefficient(t::TaylorScalar, order) = order == 0 ? value(t) : partials(t)[order]
+function get_coefficient(a::AbstractArray{<:TaylorScalar}, order)
+    map(t -> get_coefficient(t, order), a)
+end
+
+function set_coefficient(t::TaylorScalar{T, P}, order, d) where {T, P}
+    order == 0 && return TaylorScalar(d, partials(t))
+    new_partials = ntuple(i -> i == order ? d : partials(t)[i], Val(P))
+    TaylorScalar(value(t), new_partials)
+end
+
+function set_coefficient(a::AbstractArray{<:TaylorScalar}, index, d)
+    for i in eachindex(a)
+        a[i] = set_coefficient(a[i], index, d[i])
+    end
+    a
+end
+
+function append_coefficient(x::TaylorScalar{T, P}, d) where {T, P}
+    TaylorScalar(value(x), ntuple(i -> i == P + 1 ? d : partials(x)[i], Val(P + 1)))
+end
+
+append_coefficient(x::AbstractArray{<:TaylorScalar}, d) = map(append_coefficient, x, d)
+
+Base.hash(t::Taylor, h::UInt) = hash(flatten(t), h)
+function Base.isequal(a::T, b::T) where {T <: Taylor}
+    isequal(value(a), value(b)) && isequal(partials(a), partials(b))
+end
+
 function Base.promote_rule(::Type{TaylorScalar{T, P}},
         ::Type{S}) where {T, S, P}
     TaylorScalar{promote_type(T, S), P}

src/utils.jl

@@ -3,11 +3,17 @@
 
 using ChainRules
 using ChainRulesCore
-using Symbolics: @variables, @rule, unwrap, isdiv
+using Symbolics: Symbolics, @variables, @rule, @register_symbolic, 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;