Merge pull request #3006 from AayushSabharwal/as/fix-linearization

fix: fix missing `MTKParameters` handling in `linearize`

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -2285,7 +2285,7 @@ function linearization_function(sys::AbstractSystem, inputs,
 
             function (u, p, t)
                 p_setter!(oldps, p_getter(u, p..., t))
-                newu = u_getter(u, p, t)
+                newu = u_getter(u, p..., t)
                 return newu, oldps
             end
         end
@@ -2303,6 +2303,13 @@ function linearization_function(sys::AbstractSystem, inputs,
     initfn = NonlinearFunction(initsys; eval_expression, eval_module)
     initprobmap = build_explicit_observed_function(
         initsys, unknowns(sys); eval_expression, eval_module)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        initprobmap = let inner = initprobmap
+            fn(u, p::MTKParameters) = inner(u, p...)
+            fn(u, p) = inner(u, p)
+            fn
+        end
+    end
     ps = parameters(sys)
     h = build_explicit_observed_function(sys, outputs; eval_expression, eval_module)
     lin_fun = let diff_idxs = diff_idxs,
@@ -2342,7 +2349,7 @@ function linearization_function(sys::AbstractSystem, inputs,
                         initu0, initp = get_initprob_u_p(u, p, t)
                         initprob = NonlinearLeastSquaresProblem(initfn, initu0, initp)
                         nlsol = solve(initprob, initialization_solver_alg)
-                        u = initprobmap(nlsol)
+                        u = initprobmap(state_values(nlsol), parameter_values(nlsol))
                     end
                 end
                 uf = SciMLBase.UJacobianWrapper(fun, t, p)

test/downstream/linearize.jl

@@ -306,3 +306,13 @@ linfun, _ = linearization_function(sys, [u], []; op = Dict(x => 2.0))
 matrices = linfun([1.0], Dict(p => 3.0), 1.0)
 # this would be 1 if the parameter value isn't respected
 @test matrices.f_u[] == 3.0
+
+@testset "Issue #2941" begin
+    @variables x(t) y(t)
+    @parameters p
+    eqs = [0 ~ x * log(y) - p]
+    @named sys = ODESystem(eqs, t; defaults = [p => 1.0])
+    sys = complete(sys)
+    @test_nowarn linearize(
+        sys, [x], []; op = Dict(x => 1.0), allow_input_derivatives = true)
+end