Merge pull request #3429 from SciML/fb/linextras

output extra information from linearization

Author: AayushSabharwal

src/linearization.jl

@@ -275,7 +275,7 @@ end
 """
     $(TYPEDSIGNATURES)
 
-Linearize the wrapped system at the point given by `(u, p, t)`.
+Linearize the wrapped system at the point given by `(unknowns, p, t)`.
 """
 function (linfun::LinearizationFunction)(u, p, t)
     if eltype(p) <: Pair
@@ -301,7 +301,7 @@ function (linfun::LinearizationFunction)(u, p, t)
             linfun.prob, integ, fun, linfun.initializealg, Val(true);
             linfun.initialize_kwargs...)
         if !success
-            error("Initialization algorithm $(linfun.initializealg) failed with `u = $u` and `p = $p`.")
+            error("Initialization algorithm $(linfun.initializealg) failed with `unknowns = $u` and `p = $p`.")
         end
         fg_xz = linfun.uf_jac(u, DI.Constant(p), DI.Constant(t))
         h_xz = linfun.h_jac(u, DI.Constant(p), DI.Constant(t))
@@ -323,7 +323,10 @@ function (linfun::LinearizationFunction)(u, p, t)
         g_u = fg_u[linfun.alge_idxs, :],
         h_x = h_xz[:, linfun.diff_idxs],
         h_z = h_xz[:, linfun.alge_idxs],
-        h_u = h_u)
+        h_u = h_u,
+        x = u,
+        p,
+        t)
 end
 
 """
@@ -436,7 +439,7 @@ function CommonSolve.solve(prob::LinearizationProblem; allow_input_derivatives =
     p = parameter_values(prob)
     t = current_time(prob)
     linres = prob.f(u0, p, t)
-    f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u = linres
+    f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u, x, p, t = linres
 
     nx, nu = size(f_u)
     nz = size(f_z, 2)
@@ -473,7 +476,7 @@ function CommonSolve.solve(prob::LinearizationProblem; allow_input_derivatives =
         end
     end
 
-    (; A, B, C, D)
+    (; A, B, C, D), (; x, p, t)
 end
 
 """
@@ -618,8 +621,8 @@ function markio!(state, orig_inputs, inputs, outputs, disturbances; check = true
 end
 
 """
-    (; A, B, C, D), simplified_sys = linearize(sys, inputs, outputs;    t=0.0, op = Dict(), allow_input_derivatives = false, zero_dummy_der=false, kwargs...)
-    (; A, B, C, D)                 = linearize(simplified_sys, lin_fun; t=0.0, op = Dict(), allow_input_derivatives = false, zero_dummy_der=false)
+    (; A, B, C, D), simplified_sys, extras = linearize(sys, inputs, outputs;    t=0.0, op = Dict(), allow_input_derivatives = false, zero_dummy_der=false, kwargs...)
+    (; A, B, C, D), extras                 = linearize(simplified_sys, lin_fun; t=0.0, op = Dict(), allow_input_derivatives = false, zero_dummy_der=false)
 
 Linearize `sys` between `inputs` and `outputs`, both vectors of variables. Return a NamedTuple with the matrices of a linear statespace representation
 on the form
@@ -641,6 +644,8 @@ If `allow_input_derivatives = false`, an error will be thrown if input derivativ
 
 `zero_dummy_der` can be set to automatically set the operating point to zero for all dummy derivatives.
 
+The return value `extras` is a NamedTuple `(; x, p, t)` containing the result of the initialization problem that was solved to determine the operating point.
+
 See also [`linearization_function`](@ref) which provides a lower-level interface, [`linearize_symbolic`](@ref) and [`ModelingToolkit.reorder_unknowns`](@ref).
 
 See extended help for an example.
@@ -750,7 +755,8 @@ function linearize(sys, inputs, outputs; op = Dict(), t = 0.0,
         zero_dummy_der,
         op,
         kwargs...)
-    linearize(ssys, lin_fun; op, t, allow_input_derivatives), ssys
+    mats, extras = linearize(ssys, lin_fun; op, t, allow_input_derivatives)
+    mats, ssys, extras
 end
 
 """

src/systems/analysis_points.jl

@@ -908,7 +908,8 @@ for f in [:get_sensitivity, :get_comp_sensitivity, :get_looptransfer]
             sys, ap, args...; loop_openings = [], system_modifier = identity, kwargs...)
         lin_fun, ssys = $(utility_fun)(
             sys, ap, args...; loop_openings, system_modifier, kwargs...)
-        ModelingToolkit.linearize(ssys, lin_fun), ssys
+        mats, extras = ModelingToolkit.linearize(ssys, lin_fun)
+        mats, ssys, extras
     end
 end
 

test/downstream/inversemodel.jl

@@ -132,28 +132,28 @@ sol = solve(prob, Rodas5P())
 
 # we need to provide `op` so the initialization system knows what to hold constant
 # the values don't matter
-Sf, simplified_sys = Blocks.get_sensitivity_function(model, :y; op); # This should work without providing an operating opint containing a dummy derivative
+Sf, simplified_sys = get_sensitivity_function(model, :y; op); # This should work without providing an operating opint containing a dummy derivative
 x = state_values(Sf)
 p = parameter_values(Sf)
 # If this somehow passes, mention it on
 # https://github.com/SciML/ModelingToolkit.jl/issues/2786
 matrices1 = Sf(x, p, 0)
-matrices2, _ = Blocks.get_sensitivity(model, :y; op); # Test that we get the same result when calling the higher-level API
+matrices2, _ = get_sensitivity(model, :y; op); # Test that we get the same result when calling the higher-level API
 @test matrices1.f_x ≈ matrices2.A[1:6, 1:6]
 nsys = get_named_sensitivity(model, :y; op) # Test that we get the same result when calling an even higher-level API
 @test matrices2.A ≈ nsys.A
 
 # Test the same thing for comp sensitivities
 
 # This should work without providing an operating opint containing a dummy derivative
-Sf, simplified_sys = Blocks.get_comp_sensitivity_function(model, :y; op);
+Sf, simplified_sys = get_comp_sensitivity_function(model, :y; op);
 x = state_values(Sf)
 p = parameter_values(Sf)
 # If this somehow passes, mention it on
 # https://github.com/SciML/ModelingToolkit.jl/issues/2786
 matrices1 = Sf(x, p, 0)
 # Test that we get the same result when calling the higher-level API
-matrices2, _ = Blocks.get_comp_sensitivity(model, :y; op)
+matrices2, _ = get_comp_sensitivity(model, :y; op)
 @test matrices1.f_x ≈ matrices2.A[1:6, 1:6]
 # Test that we get the same result when calling an even higher-level API
 nsys = get_named_comp_sensitivity(model, :y; op)
@@ -173,15 +173,16 @@ nsys = get_named_comp_sensitivity(model, :y; op)
     output = :y
     # we need to provide `op` so the initialization system knows which
     # values to hold constant
-    lin_fun, ssys = Blocks.get_sensitivity_function(model, output; op = op1)
-    matrices1 = linearize(ssys, lin_fun, op = op1)
-    matrices2 = linearize(ssys, lin_fun, op = op2)
+    lin_fun, ssys = get_sensitivity_function(model, output; op = op1)
+    matrices1, extras1 = linearize(ssys, lin_fun, op = op1)
+    matrices2, extras2 = linearize(ssys, lin_fun, op = op2)
+    @test extras1.x != extras2.x
     S1f = ss(matrices1...)
     S2f = ss(matrices2...)
     @test S1f != S2f
 
-    matrices1, ssys = Blocks.get_sensitivity(model, output; op = op1)
-    matrices2, ssys = Blocks.get_sensitivity(model, output; op = op2)
+    matrices1, ssys = get_sensitivity(model, output; op = op1)
+    matrices2, ssys = get_sensitivity(model, output; op = op2)
     S1 = ss(matrices1...)
     S2 = ss(matrices2...)
     @test S1 != S2

test/linearize.jl

@@ -14,15 +14,16 @@ eqs = [u ~ kp * (r - y)
 
 @named sys = System(eqs, t)
 
-lsys, ssys = linearize(sys, [r], [y])
+lsys, ssys, extras = linearize(sys, [r], [y])
 lprob = LinearizationProblem(sys, [r], [y])
-lsys2 = solve(lprob)
+lsys2, extras2 = solve(lprob)
 lsys3, _ = linearize(sys, [r], [y]; autodiff = AutoFiniteDiff())
 
 @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
+@test extras == extras2
 
 lsys, ssys = linearize(sys, [r], [r])