Merge pull request #1724 from SciML/myb/partial_jac

Use efficient partial Jacobian evaluation in `linearization_function`

Author: YingboMa

src/systems/abstractsystem.jl

@@ -1019,28 +1019,35 @@ function linearization_function(sys::AbstractSystem, inputs,
                                 kwargs...)
     sys, diff_idxs, alge_idxs, input_idxs = io_preprocessing(sys, inputs, outputs; simplify,
                                                              kwargs...)
-    sts = states(sys)
-    fun = ODEFunction(sys)
-    lin_fun = let fun = fun,
-        h = ModelingToolkit.build_explicit_observed_function(sys, outputs)
+    lin_fun = let diff_idxs = diff_idxs,
+        alge_idxs = alge_idxs,
+        input_idxs = input_idxs,
+        sts = states(sys),
+        fun = ODEFunction(sys),
+        h = ModelingToolkit.build_explicit_observed_function(sys, outputs),
+        chunk = ForwardDiff.Chunk(input_idxs)
 
         function (u, p, t)
             if u !== nothing # Handle systems without states
                 length(sts) == length(u) ||
                     error("Number of state variables ($(length(sts))) does not match the number of input states ($(length(u)))")
                 uf = SciMLBase.UJacobianWrapper(fun, t, p)
                 fg_xz = ForwardDiff.jacobian(uf, u)
-                h_xz = ForwardDiff.jacobian(xz -> h(xz, p, t), u)
+                h_xz = ForwardDiff.jacobian(let p = p, t = t
+                                                xz -> h(xz, p, t)
+                                            end, u)
                 pf = SciMLBase.ParamJacobianWrapper(fun, t, u)
-                # TODO: this is very inefficient, p contains all parameters of the system
-                fg_u = ForwardDiff.jacobian(pf, p)[:, input_idxs]
+                fg_u = jacobian_wrt_vars(pf, p, input_idxs, chunk)
             else
                 length(sts) == 0 ||
                     error("Number of state variables (0) does not match the number of input states ($(length(u)))")
                 fg_xz = zeros(0, 0)
                 h_xz = fg_u = zeros(0, length(inputs))
             end
-            h_u = ForwardDiff.jacobian(p -> h(u, p, t), p)[:, input_idxs]
+            hp = let u = u, t = t
+                p -> h(u, p, t)
+            end
+            h_u = jacobian_wrt_vars(hp, p, input_idxs, chunk)
             (f_x = fg_xz[diff_idxs, diff_idxs],
              f_z = fg_xz[diff_idxs, alge_idxs],
              g_x = fg_xz[alge_idxs, diff_idxs],

src/utils.jl

@@ -710,3 +710,24 @@ function Base.iterate(it::StatefulBFS, queue = (eltype(it)[(0, it.t)]))
     end
     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