Merge pull request #70 from JuliaComputing/as/strict-inline-scc

fix: use strict form of `linear_expansion` in inline linear SCC pass

Author: AayushSabharwal

lib/ModelingToolkitTearing/src/reassemble.jl

@@ -565,15 +565,24 @@ function get_linear_scc_linsol(state::TearingState, alg_eqs::Vector{Int},
         # Substitute solved differential equations
         resid = Symbolics.fixpoint_sub(
             resid, total_sub, MTKBase.Shift; maxiters = max(2length(total_sub), 10))
-        # Standard `linear_expansion`-based process
-        for (varidx, var) in enumerate(vars)
-            a, resid, islinear = Symbolics.linear_expansion(resid, var)
+        b[eqidx] = resid
+    end
+
+    for (varidx, var) in enumerate(vars)
+        lex = Symbolics.LinearExpander(var; strict = true)
+        for (eqidx, resid) in enumerate(b)
+            p, q, islinear = lex(resid)
             islinear || return nothing
-            A[eqidx, varidx] = a
+            A[eqidx, varidx] = p
+            b[eqidx] = q
         end
-        # `-` is important! `b` is on the other side of the equality.
-        b[eqidx] = -resid
     end
+
+    # `-` is important! `b` is on the other side of the equality.
+    for i in eachindex(b)
+        b[i] = -b[i]
+    end
+
     if N <= analytical_linear_scc_limit && _check_allow_symbolic_parameter(
             state, A, allow_symbolic, allow_parameter
         )

lib/ModelingToolkitTearing/test/runtests.jl

@@ -155,3 +155,17 @@ end
     # variables is removed from `ts.structure.graph`. This would cause incorrect values in `linear_subsys_adjmat!`
     @test coeffs == [-1]
 end
+
+@testset "Inline linear SCC uses strict form of `linear_expansion`" begin
+    @variables x(t) y(t) z(t) w(t) p(t)
+    # Without strict form, the `ifelse` is considered linear in `z`
+    @test_nowarn @mtkcompile sys = System(
+        [
+            D(x) ~ 2x + t
+            y + 2z + 3w ~ 4
+            2y + 4ifelse(z < 0, 2z, 3z) + 7w ~ 12
+            3y - 2z + 5w ~ 13
+            D(p) ~ 2y + 3w + z
+        ], t
+    ) reassemble_alg = MTKTearing.DefaultReassembleAlgorithm(; inline_linear_sccs = true)
+end