[Dev] AD

[ocots]

See control-toolbox/CTBase.jl#25.

[ocots]

Pour choisir le backend, faire comme dans CTParser : https://github.com/control-toolbox/CTParser.jl/blob/314399f5d0bd806c00c5e98d2d2a8c2f01926c2a/src/onepass.jl#L42.

[jbcaillau]

FTR:

• https://discourse.julialang.org/t/forwarddiff-adaptive-ode-solvers-timestep-issue-leads-to-incorrect-derivatives/125870
• Enzyme.jl allows to define one's own rules https://enzyme.mit.edu/julia/stable/#Importing-ChainRules

See also: #42

[jbcaillau]

@joseph-gergaud @ocots

Lightweight ODE Integration with Internal Numerical Differentiation in Julia

Problem

We want to solve ODEs in Julia with:

• Adaptive stepsize (explicit Runge-Kutta, typically Tsit5).
• Derivatives (Jacobians, Hessians, higher) via ForwardDiff.jl dual numbers passed through u0 or p.
• Fast precompilation — avoid the heavy dependency footprint of DifferentialEquations.jl or even OrdinaryDiffEq.jl.
• Grid invariance (IND = Internal Numerical Differentiation): the adaptive grid must be determined only by the real (primal) part of the solution, so that derivatives of arbitrary order are computed on a consistent, frozen grid.

Solution

1. Use OrdinaryDiffEqTsit5 (the sub-package, not the monolith)

The splitting of OrdinaryDiffEq.jl into independent sub-packages is done. For Tsit5:

using Pkg
Pkg.add("OrdinaryDiffEqTsit5")

First-call time drops from ~2.5 s (OrdinaryDiffEq) to ~0.5 s (OrdinaryDiffEqTsit5).

2. Define a deepvalue-based internal norm

The adaptive controller uses internalnorm to compute local error estimates and scaling.
By projecting away all dual layers, the step-size selection sees only the real trajectory.

using ForwardDiff: Dual, value
using DiffEqBase: ODE_DEFAULT_NORM

# Strip all dual layers recursively
deepvalue(x::Dual) = deepvalue(value(x))
deepvalue(x::Real) = x

# Internal norm: always based on the innermost real part
real_norm(u::AbstractArray, t) = ODE_DEFAULT_NORM(deepvalue.(u), t)
real_norm(u::Dual, t)          = abs(deepvalue(u))
real_norm(u::Real, t)          = abs(u)

3. Pass internalnorm to solve

using OrdinaryDiffEqTsit5

sol = solve(prob, Tsit5(); reltol=1e-8, abstol=1e-8, internalnorm=real_norm)

This single keyword argument guarantees that the adaptive grid is identical to what a
purely real solve would produce, regardless of derivative order.

4. Compute Jacobians, Hessians, etc.

using ForwardDiff

# Jacobian w.r.t. parameters
function endpoint(p)
    prob = ODEProblem(f!, u0, tspan, p)
    sol  = solve(prob, Tsit5(); reltol=1e-8, abstol=1e-8, internalnorm=real_norm)
    return sol[end]
end

J = ForwardDiff.jacobian(endpoint, p)

# Jacobian w.r.t. initial conditions (= state transition matrix)
function flow(x0)
    prob = ODEProblem(f!, x0, tspan, p)
    sol  = solve(prob, Tsit5(); reltol=1e-8, abstol=1e-8, internalnorm=real_norm)
    return sol[end]
end

Φ = ForwardDiff.jacobian(flow, u0)

# Second-order derivatives (Hessian of a scalar loss, or Jacobian of Jacobian)
H = ForwardDiff.jacobian(y -> ForwardDiff.jacobian(flow, y), u0)

At every nesting level, deepvalue recurses through Dual{Tag1, Dual{Tag2, ...}} down
to Real, so the grid remains invariant.

Why it works

Component	What happens
Error estimate	‖ŷ₅ − ŷ₄‖ computed via real_norm → sees only Float64
Step scaling	atol + rtol * ‖u‖ computed via real_norm → sees only Float64
Accept/reject	Identical decision to a real-only solve
Chunking	Each ForwardDiff chunk triggers a separate solve; all chunks see the same real_norm, hence the same grid
Nesting	deepvalue bottoms out at Real regardless of dual depth

Edge case

If u0 is zero in the real part but nonzero only in the dual seed (e.g., differentiating
w.r.t. initial conditions at a fixed point), the initial step-size estimate may degenerate.
In that case, provide a hint:

sol = solve(prob, Tsit5(); dt=1e-3, reltol=1e-8, abstol=1e-8, internalnorm=real_norm)

Dependencies (minimal)

OrdinaryDiffEqTsit5   # Tsit5 solver only
ForwardDiff            # dual number AD

No need for DifferentialEquations.jl, OrdinaryDiffEq.jl, SciMLSensitivity.jl,
or any adjoint machinery. This is pure forward-mode IND.

[jbcaillau]

@ocots for the extensions:

• reason as for optimisation solvers: the modellers are imported by OC.jl (CTFlows.jl), not the solvers
• for the solvers, Tsit5 can be the default, and the user can either import the lightweight package OrdinaryDiffEqTsit5.jl (or another ODE solver), or OrdinaryDiffEq / DifferentialEquations.jl (bigger and slower)

[joseph-gergaud]

Here my results
PS : I am in holidays, come back on 2nd may

First numerical results

Test example

The numerical results are obtained on the following example $\lambda = (1,2)$

$$(IVP)\left\{\begin{array}{l} \dot{x}_1 = \lambda_1x_1\\\ \dot{x}_2 = \lambda_2x_2\\\ \dot{x}_3 = (\lambda_1-\lambda_2)x_3\\\ x_1(0) = \lambda_2\\\ x_2(0) = 1\\\ x_3(0) = 1. \end{array} \right.$$

Numerical results for fixed steps

include("../../test/test_ForwardDiff.jl")
df_sol = DataFrame(adaptive=Bool[], VAR_IND=String[], internalnorm=String[], norm_∞_error=Real[], norm_∞_diff=Real[], time_steps=Vector[])

test_FD!(df_sol,fun_lin, tspan, x0, λ, sol_∂xO_flow,false)

println(df_sol)

Row	adaptive	VAR_IND	internalnorm	norm_∞_error	norm_∞_diff	time_steps
	Bool	String	String	Real	Real	Vector
1	false	myode43	default	NaN	NaN	[0.05, 0.1]
2	false	IND	default	1.13316e-5	NaN	[0.05, 0.1]
3	false	VAR2	default	1.13316e-5	0.0	[0.05, 0.1]
4	false	VAR1	default	1.13316e-5	0.0	[0.05, 0.1]
5	false	RK4	default	NaN	NaN	[0.05, 0.1]
6	false	IND	default	1.13316e-5	NaN	[0.05, 0.1]
7	false	VAR2	default	1.13316e-5	0.0	[0.05, 0.1]
8	false	VAR1	default	1.13316e-5	0.0	[0.05, 0.1]
9	false	Tsit5	default	NaN	NaN
10	false	IND	default	3.10631e-9	NaN	[0.05, 0.1]
11	false	VAR2	default	3.10631e-9	0.0	[0.05, 0.1]
12	false	VAR1	default	3.10631e-9		0.0
13	false	RadauIIA5	default	NaN	NaN	[0.05, 0.1]
14	false	IND	default	2.08895e-8	NaN	[0.05, 0.1]
15	false	VAR2	default	2.08895e-8	0.0	[0.05, 0.1]
16	false	VAR1	default	2.08895e-8	0.0	[0.05, 0.1]

• VAR_IND
  • name of the numerical integration
  • IND : IND of Bock
  • VAR2 : variational equation with steps control on the initial variables
  • VAR1 : variational equation with steps control on all the variables
• internalnorm : internalnorm for tests on errors
• norm_∞_error : error with the analytique solution
• norm_∞_diff : error with the solution of the previous line

Numerical results for variable steps

with my_norm

df_sol = DataFrame(adaptive=Bool[], VAR_IND=String[], internalnorm=String[], norm_∞_error=Real[], norm_∞_diff=Real[], time_steps=Vector[])


# with my_norm the diagram switches 
sse(x::Number) = x^2
sse(x::ForwardDiff.Dual) = sse(ForwardDiff.value(x)) #+ sum(sse, ForwardDiff.partials(x))
totallength(x::Number) = 1
function totallength(x::ForwardDiff.Dual)
  totallength(ForwardDiff.value(x)) #+ sum(totallength, ForwardDiff.partials(x))
end
totallength(x::AbstractArray) = sum(totallength, x)
function my_norm(u, t)
  return sqrt(sum(x -> sse(x), u) / totallength(u))
end

test_FD!(df_sol,fun_lin, tspan, x0, λ, sol_∂xO_flow,true,internalnorm=my_norm)
println(df_sol)

Row	adaptive	VAR_IND	internalnorm	norm_∞_error	norm_∞_diff	time_steps
	Bool	String	String	Real	Real	Vector
1	true	myode43	my_norm	NaN	NaN	[0.044956, 0.269736]
2	true	IND	my_norm	0.00984532	NaN	[0.044956, 0.269736]
3	true	VAR2	my_norm	0.00984532	0.0	[0.044956, 0.269736]
4	true	VAR1	my_norm	0.013199	0.00335368	[0.0490278, 0.294167]
5	true	RK4	my_norm	NaN	NaN	[0.044956, 0.102922]
6	true	IND	my_norm	0.000368325	NaN	[0.044956, 0.102922]
7	true	VAR2	my_norm	0.000368325	0.0	[0.044956, 0.102922]
8	true	VAR1	my_norm	0.000513877	0.000145553	[0.0490278, 0.112227]
9	true	Tsit5	my_norm	NaN	NaN	[0.0836045, 0.241822]
10	true	IND	my_norm	1.58866e-5	NaN	[0.0836045, 0.241822]
11	true	VAR2	my_norm	1.58866e-5	0.0	[0.0836045, 0.241822]
12	true	VAR1	my_norm	1.45191e-5	1.36745e-6	[0.0896094, 0.25958]
13	true	RadauIIA5	my_norm	NaN	NaN	[0.0836045, 0.421748]
14	true	IND	my_norm	0.000278511	NaN	[0.0836045, 0.421748]
15	true	VAR2	my_norm	0.000388466	0.000109955	[0.0836045, 0.441859]
16	true	VAR1	my_norm	0.000462982	7.4516e-5	[0.0896094, 0.45849]

Conclusion : no differences between VAR2 and IND in this case

[jbcaillau]

@joseph-gergaud @ocots FTR https://link.springer.com/chapter/10.1007/978-3-642-68220-9_8

Bock, H.G. (1981). Numerical Treatment of Inverse Problems in Chemical Reaction Kinetics. In: Ebert, K.H., Deuflhard, P., Jäger, W. (eds) Modelling of Chemical Reaction Systems. Springer Series in Chemical Physics, vol 18. Springer, Berlin, Heidelberg. https://doi.org/10.1007/978-3-642-68220-9_8

[jbcaillau]

cc @abavoil

[ocots]

For infos:

• From https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/#solver_options:

internalnorm: The norm function internalnorm(u,t) which error estimates are calculated. Required are two dispatches: one dispatch for the state variable and the other on the elements of the state variable (scalar norm). Defaults are package-dependent.

Two dispatches are required.

• From https://docs.sciml.ai/DiffEqDocs/stable/extras/timestepping/. There is the computation of the error but also different "controller" to compute the new stepsize are presented.

[ocots]

Norm in DiffEqGPU: https://github.com/SciML/DiffEqGPU.jl/blob/dbce79054a2f6438c4c55d606db1417929a88e57/src/utils.jl#L1-L6

diffeqgpunorm(u::AbstractArray, t) = sqrt.(sum(abs2, u) ./ length(u))
diffeqgpunorm(u::Union{AbstractFloat, Complex}, t) = abs(u)
function diffeqgpunorm(u::AbstractArray{<:ForwardDiff.Dual}, t)
    return sqrt.(sum(abs2 ∘ ForwardDiff.value, u) ./ length(u))
end
diffeqgpunorm(u::ForwardDiff.Dual, t) = abs(ForwardDiff.value(u))

ODE Init : https://github.com/SciML/OrdinaryDiffEq.jl/blob/ae48b8b56261756029d35677c389f75ac3b08a72/lib/OrdinaryDiffEqCore/src/solve.jl#L47-L119

There are two norms:

        internalnorm = ODE_DEFAULT_NORM,
        internalopnorm = opnorm,

Default norm: DiffEqBase.ODE_DEFAULT_NORM:

• https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/DiffEqBase/src/common_defaults.jl#L53-L122

ODE_DEFAULT_NORM(u::Union{AbstractFloat, Complex}, t) = @fastmath abs(u)

function ODE_DEFAULT_NORM(f::F, u::Union{AbstractFloat, Complex}, t) where {F}
    return @fastmath abs(f(u))
end

function ODE_DEFAULT_NORM(u::Array{T}, t) where {T <: Union{AbstractFloat, Complex}}
    x = zero(T)
    @inbounds @fastmath for ui in u
        x += abs2(ui)
    end
    return Base.FastMath.sqrt_fast(real(x) / max(length(u), 1))
end

function ODE_DEFAULT_NORM(
        f::F,
        u::Array{T},
        t
    ) where {F, T <: Union{AbstractFloat, Complex}}
    x = zero(T)
    @inbounds @fastmath for ui in u
        x += abs2(f(ui))
    end
    return Base.FastMath.sqrt_fast(real(x) / max(length(u), 1))
end

function ODE_DEFAULT_NORM(
        f::F,
        u::Iterators.Zip{Z},
        t
    ) where {F, Z <: Tuple{Vararg{Array{<:Union{AbstractFloat, Complex}}}}}
    T = eltype(first(u.is))
    x = zero(T)
    @inbounds @fastmath for ui in u
        x += abs2(f(ui))
    end
    return Base.FastMath.sqrt_fast(real(x) / max(length(u), 1))
end

function ODE_DEFAULT_NORM(
        u::StaticArraysCore.StaticArray{<:Tuple, T},
        t
    ) where {T <: Union{AbstractFloat, Complex}}
    return Base.FastMath.sqrt_fast(real(sum(abs2, u)) / max(length(u), 1))
end

function ODE_DEFAULT_NORM(
        f::F, u::StaticArraysCore.StaticArray{<:Tuple, T},
        t
    ) where {F, T <: Union{AbstractFloat, Complex}}
    return Base.FastMath.sqrt_fast(real(sum(abs2 ∘ f, u)) / max(length(u), 1))
end

function ODE_DEFAULT_NORM(
        u::Union{
            AbstractArray,
            RecursiveArrayTools.AbstractVectorOfArray,
            RecursiveArrayTools.AbstractRaggedVectorOfArray,
        },
        t
    )
    return Base.FastMath.sqrt_fast(UNITLESS_ABS2(u) / max(recursive_length(u), 1))
end

function ODE_DEFAULT_NORM(f::F, u::AbstractArray, t) where {F}
    return Base.FastMath.sqrt_fast(UNITLESS_ABS2(f, u) / max(recursive_length(u), 1))
end

ODE_DEFAULT_NORM(u, t) = norm(u)
ODE_DEFAULT_NORM(f::F, u, t) where {F} = norm(f.(u))

• https://github.com/SciML/OrdinaryDiffEq.jl/blob/ae48b8b56261756029d35677c389f75ac3b08a72/lib/DiffEqBase/ext/DiffEqBaseCUDAExt.jl#L1-L11:

module DiffEqBaseCUDAExt

using DiffEqBase, CUDA

function DiffEqBase.ODE_DEFAULT_NORM(
        u::CuArray{T}, t
    ) where {T <: Union{AbstractFloat, Complex}}
    return sqrt(sum(DiffEqBase.sse, u; init = DiffEqBase.sse(zero(T))) / DiffEqBase.totallength(u))
end

end

• https://github.com/SciML/OrdinaryDiffEq.jl/blob/ae48b8b56261756029d35677c389f75ac3b08a72/lib/DiffEqBase/ext/DiffEqBaseForwardDiffExt.jl#L182-L195:

@inline ODE_DEFAULT_NORM(u::ForwardDiff.Dual, ::Any) = sqrt(sse(u))
@inline function ODE_DEFAULT_NORM(
        u::AbstractArray{<:ForwardDiff.Dual{Tag, T}},
        t::Any
    ) where {Tag, T}
    return sqrt(DiffEqBase.__sum(sse, u; init = sse(zero(T))) / DiffEqBase.totallength(u))
end
@inline ODE_DEFAULT_NORM(u::ForwardDiff.Dual, ::ForwardDiff.Dual) = sqrt(sse(u))
@inline function ODE_DEFAULT_NORM(
        u::AbstractArray{<:ForwardDiff.Dual{Tag, T}},
        ::ForwardDiff.Dual
    ) where {Tag, T}
    return sqrt(DiffEqBase.__sum(sse, u; init = sse(zero(T))) / DiffEqBase.totallength(u))
end

[ocots]

And in the extension DiffEqBaseForwardDiffExt:

@inline ODE_DEFAULT_NORM(u::ForwardDiff.Dual, ::Any) = sqrt(sse(u))
@inline function ODE_DEFAULT_NORM(
        u::AbstractArray{<:ForwardDiff.Dual{Tag, T}},
        t::Any
    ) where {Tag, T}
    return sqrt(DiffEqBase.__sum(sse, u; init = sse(zero(T))) / DiffEqBase.totallength(u))
end
@inline ODE_DEFAULT_NORM(u::ForwardDiff.Dual, ::ForwardDiff.Dual) = sqrt(sse(u))
@inline function ODE_DEFAULT_NORM(
        u::AbstractArray{<:ForwardDiff.Dual{Tag, T}},
        ::ForwardDiff.Dual
    ) where {Tag, T}
    return sqrt(DiffEqBase.__sum(sse, u; init = sse(zero(T))) / DiffEqBase.totallength(u))
end

[ocots]

@joseph-gergaud @jbcaillau I have started to implement IND but just for ForwardDiff. See https://github.com/control-toolbox/CTFlows.jl/blob/develop/ext/CTFlowsForwardDiffExt.jl

CTFlows.jl/ext/CTFlowsForwardDiffExt.jl

Line 59 in 887c4ea

Common.deepvalue(x::ForwardDiff.Dual) = Common.deepvalue(ForwardDiff.value(x))

CTFlows.jl/ext/CTFlowsForwardDiffExt.jl

Line 78 in 887c4ea

Common.real_norm(u::ForwardDiff.Dual, t) = Common.real_norm(Common.deepvalue(u), t)

CTFlows.jl/ext/CTFlowsSciMLExt.jl

Lines 266 to 275 in 887c4ea

	Strategies.OptionDefinition(;
	name = :internalnorm,
	type = Function,
	default = Common.real_norm,
	description = "Internal norm for adaptive step-size control. " *
	"Defaults to `real_norm`, which extracts the primal (Float64) " *
	"part of ForwardDiff dual numbers to ensure grid invariance (IND) " *
	"when ForwardDiff is loaded. Set to `DiffEqBase.ODE_DEFAULT_NORM` to use the SciML default.",
	aliases = (:internal_norm, :norm),
	),

[ocots]

I need the same for other AD backends.

[jbcaillau]

@joseph-gergaud as soon as sth is testable, would be nice that you check we get the expected behaviour regarding your previous tests @ocots green light?

@joseph-gergaud @jbcaillau I have started to implement IND but just for ForwardDiff. See https://github.com/control-toolbox/CTFlows.jl/blob/develop/ext/CTFlowsForwardDiffExt.jl

[ocots]

I can make a beta version of CTFlows. We have the possibility to build a flow from a vector field, a Hamiltonian vector field, a ODEFunction and a ODEProblem.

We don't have yet from a Hamiltonian and from an optimal control problem (with control law).

GPU is not supported natively but should work selecting the right options and algorithms.