Merge pull request #3477 from AayushSabharwal/as/linearize-ad

feat: support alternative AD backends in linearization

Author: ChrisRackauckas

Project.toml

@@ -4,6 +4,7 @@ authors = ["Yingbo Ma <[email protected]>", "Chris Rackauckas <accounts@chrisr
 version = "9.68.1"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -16,6 +17,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -77,6 +79,7 @@ MTKInfiniteOptExt = "InfiniteOpt"
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
+ADTypes = "1.14.0"
 AbstractTrees = "0.3, 0.4"
 ArrayInterface = "6, 7"
 BifurcationKit = "0.4"
@@ -96,6 +99,7 @@ DiffEqBase = "6.165.1"
 DiffEqCallbacks = "2.16, 3, 4"
 DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"
+DifferentiationInterface = "0.6.47"
 Distributed = "1"
 Distributions = "0.23, 0.24, 0.25"
 DocStringExtensions = "0.7, 0.8, 0.9"
@@ -156,8 +160,8 @@ julia = "1.9"
 [extras]
 AmplNLWriter = "7c4d4715-977e-5154-bfe0-e096adeac482"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-BoundaryValueDiffEqMIRK = "1a22d4ce-7765-49ea-b6f2-13c8438986a6"
 BoundaryValueDiffEqAscher = "7227322d-7511-4e07-9247-ad6ff830280e"
+BoundaryValueDiffEqMIRK = "1a22d4ce-7765-49ea-b6f2-13c8438986a6"
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"

src/ModelingToolkit.jl

@@ -94,6 +94,9 @@ RuntimeGeneratedFunctions.init(@__MODULE__)
 import DynamicQuantities, Unitful
 const DQ = DynamicQuantities
 
+import DifferentiationInterface as DI
+using ADTypes: AutoForwardDiff
+
 export @derivatives
 
 for fun in [:toexpr]

src/linearization.jl

@@ -25,6 +25,7 @@ The `simplified_sys` has undergone [`structural_simplify`](@ref) and had any occ
   - `simplify`: Apply simplification in tearing.
   - `initialize`: If true, a check is performed to ensure that the operating point is consistent (satisfies algebraic equations). If the op is not consistent, initialization is performed.
   - `initialization_solver_alg`: A NonlinearSolve algorithm to use for solving for a feasible set of state and algebraic variables that satisfies the specified operating point.
+  - `autodiff`: An `ADType` supported by DifferentiationInterface.jl to use for calculating the necessary jacobians. Defaults to using `AutoForwardDiff()`
   - `kwargs`: Are passed on to `find_solvables!`
 
 See also [`linearize`](@ref) which provides a higher-level interface.
@@ -39,6 +40,7 @@ function linearization_function(sys::AbstractSystem, inputs,
         p = DiffEqBase.NullParameters(),
         zero_dummy_der = false,
         initialization_solver_alg = TrustRegion(),
+        autodiff = AutoForwardDiff(),
         eval_expression = false, eval_module = @__MODULE__,
         warn_initialize_determined = true,
         guesses = Dict(),
@@ -82,13 +84,104 @@ function linearization_function(sys::AbstractSystem, inputs,
     initialization_kwargs = (;
         abstol = initialization_abstol, reltol = initialization_reltol,
         nlsolve_alg = initialization_solver_alg)
+
+    p = parameter_values(prob)
+    t0 = current_time(prob)
+    inputvals = [p[idx] for idx in input_idxs]
+
+    hp_fun = let fun = h, setter = setp_oop(sys, input_idxs)
+        function hpf(du, input, u, p, t)
+            p = setter(p, input)
+            fun(du, u, p, t)
+            return du
+        end
+    end
+    if u0 === nothing
+        uf_jac = h_jac = pf_jac = nothing
+        T = p isa MTKParameters ? eltype(p.tunable) : eltype(p)
+        hp_jac = PreparedJacobian{true}(
+            hp_fun, zeros(T, size(outputs)), autodiff, inputvals,
+            DI.Constant(prob.u0), DI.Constant(p), DI.Constant(t0))
+    else
+        uf_fun = let fun = prob.f
+            function uff(du, u, p, t)
+                SciMLBase.UJacobianWrapper(fun, t, p)(du, u)
+            end
+        end
+        uf_jac = PreparedJacobian{true}(
+            uf_fun, similar(prob.u0), autodiff, prob.u0, DI.Constant(p), DI.Constant(t0))
+        # observed function is a `GeneratedFunctionWrapper` with iip component
+        h_jac = PreparedJacobian{true}(h, similar(prob.u0, size(outputs)), autodiff,
+            prob.u0, DI.Constant(p), DI.Constant(t0))
+        pf_fun = let fun = prob.f, setter = setp_oop(sys, input_idxs)
+            function pff(du, input, u, p, t)
+                p = setter(p, input)
+                SciMLBase.ParamJacobianWrapper(fun, t, u)(du, p)
+            end
+        end
+        pf_jac = PreparedJacobian{true}(pf_fun, similar(prob.u0), autodiff, inputvals,
+            DI.Constant(prob.u0), DI.Constant(p), DI.Constant(t0))
+        hp_jac = PreparedJacobian{true}(
+            hp_fun, similar(prob.u0, size(outputs)), autodiff, inputvals,
+            DI.Constant(prob.u0), DI.Constant(p), DI.Constant(t0))
+    end
+
     lin_fun = LinearizationFunction(
         diff_idxs, alge_idxs, input_idxs, length(unknowns(sys)),
-        prob, h, u0 === nothing ? nothing : similar(u0),
-        ForwardDiff.Chunk(input_idxs), initializealg, initialization_kwargs)
+        prob, h, u0 === nothing ? nothing : similar(u0), uf_jac, h_jac, pf_jac,
+        hp_jac, initializealg, initialization_kwargs)
     return lin_fun, sys
 end
 
+"""
+    $(TYPEDEF)
+
+Callable struct which stores a function and its prepared `DI.jacobian`. Calling with the
+appropriate arguments for DI returns the jacobian.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
+struct PreparedJacobian{iip, P, F, B, A}
+    """
+    The preparation object.
+    """
+    prep::P
+    """
+    The function whose jacobian is calculated.
+    """
+    f::F
+    """
+    Buffer for in-place functions.
+    """
+    buf::B
+    """
+    ADType to use for differentiation.
+    """
+    autodiff::A
+end
+
+function PreparedJacobian{true}(f, buf, autodiff, args...)
+    prep = DI.prepare_jacobian(f, buf, autodiff, args...)
+    return PreparedJacobian{true, typeof(prep), typeof(f), typeof(buf), typeof(autodiff)}(
+        prep, f, buf, autodiff)
+end
+
+function PreparedJacobian{false}(f, autodiff, args...)
+    prep = DI.prepare_jacobian(f, autodiff, args...)
+    return PreparedJacobian{true, typeof(prep), typeof(f), Nothing, typeof(autodiff)}(
+        prep, f, nothing)
+end
+
+function (pj::PreparedJacobian{true})(args...)
+    DI.jacobian(pj.f, pj.buf, pj.prep, pj.autodiff, args...)
+end
+
+function (pj::PreparedJacobian{false})(args...)
+    DI.jacobian(pj.f, pj.prep, pj.autodiff, args...)
+end
+
 """
     $(TYPEDEF)
 
@@ -100,7 +193,7 @@ $(TYPEDFIELDS)
 """
 struct LinearizationFunction{
     DI <: AbstractVector{Int}, AI <: AbstractVector{Int}, II, P <: ODEProblem,
-    H, C, Ch, IA <: SciMLBase.DAEInitializationAlgorithm, IK}
+    H, C, J1, J2, J3, J4, IA <: SciMLBase.DAEInitializationAlgorithm, IK}
     """
     The indexes of differential equations in the linearized system.
     """
@@ -130,11 +223,22 @@ struct LinearizationFunction{
     Any required cache buffers.
     """
     caches::C
-    # TODO: Use DI?
     """
-    A `ForwardDiff.Chunk` for taking the jacobian with respect to the inputs.
+    `PreparedJacobian` for calculating jacobian of `prob.f` w.r.t. `u`
+    """
+    uf_jac::J1
+    """
+    `PreparedJacobian` for calculating jacobian of `h` w.r.t. `u`
     """
-    chunk::Ch
+    h_jac::J2
+    """
+    `PreparedJacobian` for calculating jacobian of `prob.f` w.r.t. `p`
+    """
+    pf_jac::J3
+    """
+    `PreparedJacobian` for calculating jacobian of `h` w.r.t. `p`
+    """
+    hp_jac::J4
     """
     The initialization algorithm to use.
     """
@@ -188,25 +292,18 @@ function (linfun::LinearizationFunction)(u, p, t)
         if !success
             error("Initialization algorithm $(linfun.initializealg) failed with `u = $u` and `p = $p`.")
         end
-        uf = SciMLBase.UJacobianWrapper(fun, t, p)
-        fg_xz = ForwardDiff.jacobian(uf, u)
-        h_xz = ForwardDiff.jacobian(
-            let p = p, t = t, h = linfun.h
-                xz -> h(xz, p, t)
-            end, u)
-        pf = SciMLBase.ParamJacobianWrapper(fun, t, u)
-        fg_u = jacobian_wrt_vars(pf, p, linfun.input_idxs, linfun.chunk)
+        fg_xz = linfun.uf_jac(u, DI.Constant(p), DI.Constant(t))
+        h_xz = linfun.h_jac(u, DI.Constant(p), DI.Constant(t))
+        fg_u = linfun.pf_jac([p[idx] for idx in linfun.input_idxs],
+            DI.Constant(u), DI.Constant(p), DI.Constant(t))
     else
         linfun.num_states == 0 ||
             error("Number of unknown variables (0) does not match the number of input unknowns ($(length(u)))")
         fg_xz = zeros(0, 0)
         h_xz = fg_u = zeros(0, length(linfun.input_idxs))
     end
-    hp = let u = u, t = t, h = linfun.h
-        _hp(p) = h(u, p, t)
-        _hp
-    end
-    h_u = jacobian_wrt_vars(hp, p, linfun.input_idxs, linfun.chunk)
+    h_u = linfun.hp_jac([p[idx] for idx in linfun.input_idxs],
+        DI.Constant(u), DI.Constant(p), DI.Constant(t))
     (f_x = fg_xz[linfun.diff_idxs, linfun.diff_idxs],
         f_z = fg_xz[linfun.diff_idxs, linfun.alge_idxs],
         g_x = fg_xz[linfun.alge_idxs, linfun.diff_idxs],

src/systems/parameter_buffer.jl

@@ -734,35 +734,6 @@ function Base.:(==)(a::MTKParameters, b::MTKParameters)
            end)
 end
 
-# to support linearize/linearization_function
-function jacobian_wrt_vars(pf::F, p::MTKParameters, input_idxs, chunk::C) where {F, C}
-    tunable, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
-    T = eltype(tunable)
-    tag = ForwardDiff.Tag(pf, T)
-    dualtype = ForwardDiff.Dual{typeof(tag), T, ForwardDiff.chunksize(chunk)}
-    p_big = SciMLStructures.replace(SciMLStructures.Tunable(), p, dualtype.(tunable))
-    p_closure = let pf = pf,
-        input_idxs = input_idxs,
-        p_big = p_big
-
-        function (p_small_inner)
-            for (i, val) in zip(input_idxs, p_small_inner)
-                set_parameter!(p_big, val, i)
-            end
-            return if pf isa SciMLBase.ParamJacobianWrapper
-                buffer = Array{dualtype}(undef, size(pf.u))
-                pf(buffer, p_big)
-                buffer
-            else
-                pf(p_big)
-            end
-        end
-    end
-    p_small = parameter_values.((p,), input_idxs)
-    cfg = ForwardDiff.JacobianConfig(p_closure, p_small, chunk, tag)
-    ForwardDiff.jacobian(p_closure, p_small, cfg, Val(false))
-end
-
 const MISSING_PARAMETERS_MESSAGE = """
                                 Some parameters are missing from the variable map.
                                 Please provide a value or default for the following variables:

src/utils.jl

@@ -963,27 +963,6 @@ function Base.iterate(it::StatefulBFS, queue = (eltype(it)[(0, it.t)]))
     return (lv, t), queue
 end
 
-function jacobian_wrt_vars(pf::F, p, input_idxs, chunk::C) where {F, C}
-    E = eltype(p)
-    tag = ForwardDiff.Tag(pf, E)
-    T = typeof(tag)
-    dualtype = ForwardDiff.Dual{T, E, ForwardDiff.chunksize(chunk)}
-    p_big = similar(p, dualtype)
-    copyto!(p_big, p)
-    p_closure = let pf = pf,
-        input_idxs = input_idxs,
-        p_big = p_big
-
-        function (p_small_inner)
-            p_big[input_idxs] .= p_small_inner
-            pf(p_big)
-        end
-    end
-    p_small = p[input_idxs]
-    cfg = ForwardDiff.JacobianConfig(p_closure, p_small, chunk, tag)
-    ForwardDiff.jacobian(p_closure, p_small, cfg, Val(false))
-end
-
 function fold_constants(ex)
     if iscall(ex)
         maketerm(typeof(ex), operation(ex), map(fold_constants, arguments(ex)),

test/downstream/linearize.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, Test
+using ModelingToolkit, ADTypes, Test
 using CommonSolve: solve
 
 # r is an input, and y is an output.
@@ -17,11 +17,12 @@ eqs = [u ~ kp * (r - y)
 lsys, ssys = linearize(sys, [r], [y])
 lprob = LinearizationProblem(sys, [r], [y])
 lsys2 = solve(lprob)
+lsys3, _ = linearize(sys, [r], [y]; autodiff = AutoFiniteDiff())
 
-@test lsys.A[] == lsys2.A[] == -2
-@test lsys.B[] == lsys2.B[] == 1
-@test lsys.C[] == lsys2.C[] == 1
-@test lsys.D[] == lsys2.D[] == 0
+@test lsys.A[] == lsys2.A[] == lsys3.A[] == -2
+@test lsys.B[] == lsys2.B[] == lsys3.B[] == 1
+@test lsys.C[] == lsys2.C[] == lsys3.C[] == 1
+@test lsys.D[] == lsys2.D[] == lsys3.D[] == 0
 
 lsys, ssys = linearize(sys, [r], [r])
 
@@ -89,11 +90,13 @@ connections = [f.y ~ c.r # filtered reference to controller reference
 lsys0, ssys = linearize(cl, [f.u], [p.x])
 desired_order = [f.x, p.x]
 lsys = ModelingToolkit.reorder_unknowns(lsys0, unknowns(ssys), desired_order)
+lsys1, ssys = linearize(cl, [f.u], [p.x]; autodiff = AutoFiniteDiff())
+lsys2 = ModelingToolkit.reorder_unknowns(lsys1, unknowns(ssys), desired_order)
 
-@test lsys.A == [-2 0; 1 -2]
-@test lsys.B == reshape([1, 0], 2, 1)
-@test lsys.C == [0 1]
-@test lsys.D[] == 0
+@test lsys.A == lsys2.A == [-2 0; 1 -2]
+@test lsys.B == lsys2.B == reshape([1, 0], 2, 1)
+@test lsys.C == lsys2.C == [0 1]
+@test lsys.D[] == lsys2.D[] == 0
 
 ## Symbolic linearization
 lsyss, _ = ModelingToolkit.linearize_symbolic(cl, [f.u], [p.x])