Merge pull request #3556 from AayushSabharwal/as/jac-header-fix

fix: fix `assert_jac_length_header`

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -303,7 +303,8 @@ export structural_simplify, expand_connections, linearize, linearization_functio
        LinearizationProblem
 export solve
 
-export calculate_jacobian, generate_jacobian, generate_function, generate_custom_function, generate_W
+export calculate_jacobian, generate_jacobian, generate_function, generate_custom_function,
+       generate_W
 export calculate_control_jacobian, generate_control_jacobian
 export calculate_tgrad, generate_tgrad
 export calculate_gradient, generate_gradient

src/systems/diffeqs/abstractodesystem.jl

@@ -141,23 +141,25 @@ end
 
 function assert_jac_length_header(sys)
     W = W_sparsity(sys)
-    identity, expr -> Func([expr.args...], [], LiteralExpr(quote
-        @assert $(findnz)($(expr.args[1]))[1:2] == $(findnz)($W)[1:2]
-        $(expr.body)
-    end))
+    identity,
+    function add_header(expr)
+        Func(expr.args, [], expr.body,
+            [:(@assert $(SymbolicUtils.Code.toexpr(term(findnz, expr.args[1])))[1:2] ==
+                       $(findnz(W)[1:2]))])
+    end
 end
 
-function generate_W(sys::AbstractODESystem, γ = 1., dvs = unknowns(sys),
-        ps = parameters(sys; initial_parameters = true); 
+function generate_W(sys::AbstractODESystem, γ = 1.0, dvs = unknowns(sys),
+        ps = parameters(sys; initial_parameters = true);
         simplify = false, sparse = false, kwargs...)
     @variables ˍ₋gamma
     M = calculate_massmatrix(sys; simplify)
     sparse && (M = SparseArrays.sparse(M))
     J = calculate_jacobian(sys; simplify, sparse, dvs)
-    W = ˍ₋gamma*M + J
+    W = ˍ₋gamma * M + J
 
     p = reorder_parameters(sys, ps)
-    return build_function_wrapper(sys, W, 
+    return build_function_wrapper(sys, W,
         dvs,
         p...,
         ˍ₋gamma,
@@ -302,11 +304,12 @@ function jacobian_dae_sparsity(sys::AbstractODESystem)
     J1 + J2
 end
 
-function W_sparsity(sys::AbstractODESystem) 
+function W_sparsity(sys::AbstractODESystem)
     jac_sparsity = jacobian_sparsity(sys)
     (n, n) = size(jac_sparsity)
     M = calculate_massmatrix(sys)
-    M_sparsity = M isa UniformScaling ? sparse(I(n)) : SparseMatrixCSC{Bool, Int64}((!iszero).(M))
+    M_sparsity = M isa UniformScaling ? sparse(I(n)) :
+                 SparseMatrixCSC{Bool, Int64}((!iszero).(M))
     jac_sparsity .| M_sparsity
 end
 

test/jacobiansparsity.jl

@@ -74,7 +74,7 @@ f = DiffEqBase.ODEFunction(sys, u0 = nothing, sparse = true, jac = false)
 # test when u0 is not Float64
 u0 = similar(init_brusselator_2d(xyd_brusselator), Float32)
 prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
-                                     u0, (0.0, 11.5), p)
+    u0, (0.0, 11.5), p)
 sys = complete(modelingtoolkitize(prob_ode_brusselator_2d))
 
 prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = false)
@@ -94,7 +94,8 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     @mtkbuild pend = ODESystem(eqs, t)
 
     u0 = [x => 1, y => 0]
-    prob = ODEProblem(pend, u0, (0, 11.5), [g => 1], guesses = [λ => 1], sparse = true, jac = true)
+    prob = ODEProblem(
+        pend, u0, (0, 11.5), [g => 1], guesses = [λ => 1], sparse = true, jac = true)
     jac, jac! = generate_jacobian(pend; expression = Val{false}, sparse = true)
     jac_prototype = ModelingToolkit.jacobian_sparsity(pend)
     W_prototype = ModelingToolkit.W_sparsity(pend)
@@ -109,8 +110,28 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     @test_throws AssertionError jac!(similar(jac_prototype, Float64), u, p, t)
 
     W, W! = generate_W(pend; expression = Val{false}, sparse = true)
-    γ = .1
+    γ = 0.1
     M = sparse(calculate_massmatrix(pend))
     @test_throws AssertionError W!(similar(jac_prototype, Float64), u, p, γ, t)
-    @test W!(similar(W_prototype, Float64), u, p, γ, t) == 0.1 * M + jac!(similar(W_prototype, Float64), u, p, t)
+    @test W!(similar(W_prototype, Float64), u, p, γ, t) ==
+          0.1 * M + jac!(similar(W_prototype, Float64), u, p, t)
+end
+
+@testset "Issue#3556: Numerical accuracy" begin
+    t = ModelingToolkit.t_nounits
+    D = ModelingToolkit.D_nounits
+    @parameters g
+    @variables x(t) y(t) [state_priority = 10] λ(t)
+    eqs = [D(D(x)) ~ λ * x
+           D(D(y)) ~ λ * y - g
+           x^2 + y^2 ~ 1]
+    @mtkbuild pend = ODESystem(eqs, t)
+    prob = ODEProblem(pend, [x => 0.0, D(x) => 1.0], (0.0, 1.0), [g => 1.0];
+        guesses = [y => 1.0, λ => 1.0], jac = true, sparse = true)
+    J = deepcopy(prob.f.jac_prototype)
+    prob.f.jac(J, prob.u0, prob.p, 1.0)
+    # this currently works but may not continue to do so
+    # see https://github.com/SciML/ModelingToolkit.jl/pull/3556#issuecomment-2792664039
+    @test J == prob.f.jac(prob.u0, prob.p, 1.0)
+    @test J ≈ prob.f.jac(prob.u0, prob.p, 1.0)
 end