Merge pull request #2480 from AayushSabharwal/as/nonlinear-jacobian

fix: fix jacobian generation for NonlinearSystem

Author: ChrisRackauckas

docs/src/basics/Composition.md

@@ -204,7 +204,7 @@ N = S + I + R
 
 @named seqn = ODESystem([D(S) ~ -β * S * I / N], t)
 @named ieqn = ODESystem([D(I) ~ β * S * I / N - γ * I], t)
-@named reqn = ODESystem([D(R) ~ γ * I],t )
+@named reqn = ODESystem([D(R) ~ γ * I], t)
 
 sir = compose(
     ODESystem(

docs/src/examples/higher_order.md

@@ -22,7 +22,7 @@ eqs = [D(D(x)) ~ σ * (y - x),
     D(y) ~ x * (ρ - z) - y,
     D(z) ~ x * y - β * z]
 
-@named sys = ODESystem(eqs,t)
+@named sys = ODESystem(eqs, t)
 ```
 
 Note that we could've used an alternative syntax for 2nd order, i.e.

docs/src/tutorials/bifurcation_diagram_computation.md

@@ -91,7 +91,7 @@ It is also possible to use `ODESystem`s (rather than `NonlinearSystem`s) as inpu
 using BifurcationKit, ModelingToolkit, Plots
 using ModelingToolkit: t_nounits as t, D_nounits as D
 
-@variables  x(t) y(t)
+@variables x(t) y(t)
 @parameters μ
 eqs = [D(x) ~ μ * x - y - x * (x^2 + y^2),
     D(y) ~ x + μ * y - y * (x^2 + y^2)]

src/systems/diffeqs/abstractodesystem.jl

@@ -807,7 +807,8 @@ function get_u0(sys, u0map, parammap = nothing; symbolic_u0 = false)
     end
     defs = mergedefaults(defs, u0map, dvs)
     for (k, v) in defs
-        if Symbolics.isarraysymbolic(k)
+        if Symbolics.isarraysymbolic(k) &&
+           Symbolics.shape(unwrap(k)) !== Symbolics.Unknown()
             ks = scalarize(k)
             length(ks) == length(v) || error("$k has default value $v with unmatched size")
             for (kk, vv) in zip(ks, v)

src/systems/index_cache.jl

@@ -49,13 +49,13 @@ function IndexCache(sys::AbstractSystem)
         end
     end
 
-    disc_buffers = Dict{DataType, Set{BasicSymbolic}}()
-    tunable_buffers = Dict{DataType, Set{BasicSymbolic}}()
-    constant_buffers = Dict{DataType, Set{BasicSymbolic}}()
-    dependent_buffers = Dict{DataType, Set{BasicSymbolic}}()
-    nonnumeric_buffers = Dict{DataType, Set{BasicSymbolic}}()
+    disc_buffers = Dict{Any, Set{BasicSymbolic}}()
+    tunable_buffers = Dict{Any, Set{BasicSymbolic}}()
+    constant_buffers = Dict{Any, Set{BasicSymbolic}}()
+    dependent_buffers = Dict{Any, Set{BasicSymbolic}}()
+    nonnumeric_buffers = Dict{Any, Set{BasicSymbolic}}()
 
-    function insert_by_type!(buffers::Dict{DataType, Set{BasicSymbolic}}, sym)
+    function insert_by_type!(buffers::Dict{Any, Set{BasicSymbolic}}, sym)
         sym = unwrap(sym)
         ctype = concrete_symtype(sym)
         buf = get!(buffers, ctype, Set{BasicSymbolic}())
@@ -101,7 +101,7 @@ function IndexCache(sys::AbstractSystem)
             if ctype <: Real || ctype <: AbstractArray{<:Real}
                 if is_discrete_domain(p)
                     disc_buffers
-                elseif istunable(p, true) && size(p) !== Symbolics.Unknown()
+                elseif istunable(p, true) && Symbolics.shape(p) !== Symbolics.Unknown()
                     tunable_buffers
                 else
                     constant_buffers
@@ -113,7 +113,7 @@ function IndexCache(sys::AbstractSystem)
         )
     end
 
-    function get_buffer_sizes_and_idxs(buffers::Dict{DataType, Set{BasicSymbolic}})
+    function get_buffer_sizes_and_idxs(buffers::Dict{Any, Set{BasicSymbolic}})
         idxs = IndexMap()
         buffer_sizes = BufferTemplate[]
         for (i, (T, buf)) in enumerate(buffers)

src/systems/nonlinear/nonlinearsystem.jl

@@ -175,9 +175,10 @@ function generate_jacobian(
         sys::NonlinearSystem, vs = unknowns(sys), ps = full_parameters(sys);
         sparse = false, simplify = false, kwargs...)
     jac = calculate_jacobian(sys, sparse = sparse, simplify = simplify)
-    pre = get_preprocess_constants(jac)
+    pre, sol_states = get_substitutions_and_solved_unknowns(sys)
     p = reorder_parameters(sys, ps)
-    return build_function(jac, vs, p...; postprocess_fbody = pre, kwargs...)
+    return build_function(
+        jac, vs, p...; postprocess_fbody = pre, states = sol_states, kwargs...)
 end
 
 function calculate_hessian(sys::NonlinearSystem; sparse = false, simplify = false)

test/nonlinearsystem.jl

@@ -87,6 +87,9 @@ prob = NonlinearProblem(ns, ones(3), [σ => 1.0, ρ => 1.0, β => 1.0])
 sol = solve(prob, NewtonRaphson())
 @test sol.u[1] ≈ sol.u[2]
 
+prob = NonlinearProblem(ns, ones(3), [σ => 1.0, ρ => 1.0, β => 1.0], jac = true)
+@test_nowarn solve(prob, NewtonRaphson())
+
 @test_throws ArgumentError NonlinearProblem(ns, ones(4), [σ => 1.0, ρ => 1.0, β => 1.0])
 
 @variables u F s a