fix: fix generate_W

Author: vyudu

src/ModelingToolkit.jl

@@ -303,7 +303,7 @@ export structural_simplify, expand_connections, linearize, linearization_functio
        LinearizationProblem
 export solve
 
-export calculate_jacobian, generate_jacobian, generate_function, generate_custom_function
+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/structural_transformation/symbolics_tearing.jl

@@ -378,6 +378,7 @@ function generate_derivative_variables!(
     # For variable x, make dummy derivative x_t if the
     # derivative is in the system
     for v in 1:length(var_to_diff)
+        @show graph
         dv = var_to_diff[v]
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
@@ -403,6 +404,10 @@ function generate_derivative_variables!(
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)
         # Add `D(x) - x_t ~ 0` to the graph
         dummy_eq = add_dd_equation!(structure, neweqs, 0 ~ dx - x_t, dv, v_t)
+        # Update graph to say, all the equations featuring D(x) also feature x_t
+        for e in 𝑑neighbors(graph, dv)
+            add_edge!(graph, e, v_t)
+        end
 
         # Update matching
         push!(var_eq_matching, unassigned)

src/systems/diffeqs/abstractodesystem.jl

@@ -143,15 +143,20 @@ function generate_W(sys::AbstractODESystem, γ = 1., dvs = unknowns(sys),
         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
+    @show W
 
     p = reorder_parameters(sys, ps)
+    @show length(p)
     return build_function_wrapper(sys, W, 
         dvs,
         p...,
         ˍ₋gamma,
         get_iv(sys);
+        wrap_code = sparse ? assert_jac_length_header(sys) : (identity, identity),
+        p_end = 1 + length(p),
         kwargs...)
 end
 

test/jacobiansparsity.jl

@@ -86,7 +86,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
 @testset "W matrix sparsity" begin
     t = ModelingToolkit.t_nounits
     @parameters g
-    @variables x(t) y(t) [state_priority = 10] λ(t)
+    @variables x(t) y(t) λ(t)
     eqs = [D(D(x)) ~ λ * x
            D(D(y)) ~ λ * y - g
            x^2 + y^2 ~ 1]
@@ -105,4 +105,10 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     p = prob.p
     t = 0.0
     @test_throws AssertionError jac!(similar(jac_prototype, Float64), u, p, t)
+
+    W, W! = generate_W(pend; expression = Val{false}, sparse = true)
+    γ = .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)
 end