feat: use DI to calculate jacobians in LinearizationFunction

Author: AayushSabharwal

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.
   - `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,89 @@ 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]
+
+    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_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
+    hp_jac = PreparedJacobian{true}(hp_fun, similar(prob.u0, size(outputs)), autodiff, inputvals, DI.Constant(prob.u0), DI.Constant(p), DI.Constant(t0))
+    
     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 +178,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 +208,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 +277,16 @@ 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],