Merge pull request #63 from JuliaComputing/as/fix-diffcache-reference

fix: improve reference calculation for `DiffCache`

Author: AayushSabharwal

lib/ModelingToolkitTearing/Project.toml

@@ -26,6 +26,7 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 BipartiteGraphs = "0.1.3"
 CommonSolve = "0.2"
 DocStringExtensions = "0.7, 0.8, 0.9"
+ForwardDiff = "1.3"
 Graphs = "1"
 LinearAlgebra = "1"
 ModelingToolkit = "11"
@@ -44,8 +45,9 @@ UUIDs = "1"
 julia = "1.10"
 
 [extras]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "ModelingToolkit"]
+test = ["Test", "ModelingToolkit", "ForwardDiff"]

lib/ModelingToolkitTearing/src/reassemble.jl

@@ -606,6 +606,10 @@ function get_linear_scc_linsol(state::TearingState, alg_eqs::Vector{Int},
     else
         reference = fullvars[state_idx]
     end
+    reference = Symbolics.STerm(
+        promote, Symbolics.SArgsT((reference, MTKBase.get_iv(sys)::SymbolicT));
+        type = Vector{Real}, shape = [1:2]
+    )[1]
     sys, A_cache = MTKBase.add_diffcache(sys, length(A))
     A_allocator = A_cache(reference)
     A = SU.Code.with_allocator(A_allocator, SU.Const{VartypeT}(A))

lib/ModelingToolkitTearing/test/runtests.jl

@@ -8,6 +8,7 @@ using Graphs
 import StateSelection
 using ModelingToolkit: t_nounits as t, D_nounits as D
 import SymbolicUtils as SU
+using ForwardDiff
 
 @testset "`InferredDiscrete` validation" begin
     k = ShiftIndex()
@@ -110,3 +111,30 @@ end
     @test isempty(unknowns(sys))
     @test length(observed(sys)) == 2
 end
+
+@testset "Duals through inline linear SCC DiffCaches" begin
+    @variables x(t) y(t) z(t) w(t) q(t)
+    reassemble_alg = MTKTearing.DefaultReassembleAlgorithm(; inline_linear_sccs = true)
+    @mtkcompile sys = System(
+        [
+            D(x) ~ 2t + 1,
+            t * y + x * z + w ~ 4,
+            4y + 3z + 2w ~ 7,
+            2x * y + 3t * z + w ~ 10,
+            D(q) ~ 2w,
+        ],
+        t
+    ) reassemble_alg = reassemble_alg
+
+    prob = ODEProblem(sys, [x => 1, q => 1], (0.0, 1.0); guesses = [x => 1, y => 1, z => 1, w => 1, q => 1])
+
+    buffer = similar(prob.u0)
+    @test_nowarn ForwardDiff.jacobian(buffer, prob.u0) do du, u
+        @test eltype(u) <: ForwardDiff.Dual
+        prob.f.f.f_iip(du, u, prob.p, 1.0)
+    end
+    @test_nowarn ForwardDiff.derivative(buffer, 1.2) do du, t
+        @test t isa ForwardDiff.Dual
+        prob.f.f.f_iip(du, prob.u0, prob.p, t)
+    end
+end