Extend modeling of linear incidences with unknown coefficients

Project.toml

@@ -47,7 +47,7 @@ ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 [sources]
 Compiler = {rev = "master", url = "https://github.com/JuliaLang/BaseCompiler.jl.git"}
 Cthulhu = {rev = "master", url = "https://github.com/JuliaDebug/Cthulhu.jl.git"}
-DifferentiationInterface = {rev = "main", subdir = "DifferentiationInterface", url = "https://github.com/Keno/DifferentiationInterface.jl#main"}
+DifferentiationInterface = {rev = "main", subdir = "DifferentiationInterface", url = "https://github.com/Keno/DifferentiationInterface.jl"}
 Diffractor = {rev = "main", url = "https://github.com/JuliaDiff/Diffractor.jl.git"}
 SimpleNonlinearSolve = {rev = "master", subdir = "lib/SimpleNonlinearSolve", url = "https://github.com/SciML/NonlinearSolve.jl.git"}
 StateSelection = {rev = "main", url = "https://github.com/JuliaComputing/StateSelection.jl.git"}

benchmark/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Compiler = "807dbc54-b67e-4c79-8afb-eafe4df6f2e1"
 DAECompiler = "32805668-c3d0-42c2-aafd-0d0a9857a104"
+Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 XSteam = "95ff35a0-be81-11e9-2ca3-5b4e338e8476"

src/analysis/ipoincidence.jl

@@ -45,49 +45,95 @@ end
 apply_linear_incidence(𝕃, ret::Type, caller::CallerMappingState, mapping::CalleeMapping) = ret
 apply_linear_incidence(𝕃, ret::Const, caller::CallerMappingState, mapping::CalleeMapping) = ret
 function apply_linear_incidence(𝕃, ret::Incidence, caller::Union{CallerMappingState, Nothing}, mapping::CalleeMapping)
-    coeffs = mapping.var_coeffs
-
-    const_val = ret.typ
-    new_row = _zero_row()
-
-    for (v_offset, coeff) in zip(rowvals(ret.row), nonzeros(ret.row))
-        v = v_offset - 1
+    # Substitute variables returned by the callee with the incidence defined by the caller.
+    # The composition will be additive in the constant terms, and multiplicative for linear coefficients.
+    caller_variables = mapping.var_coeffs
+
+    typ = ret.typ
+    row = _zero_row()
+
+    used_caller_variables = Int[]
+    for i in rowvals(ret.row)
+        i == 1 && continue # skip time
+        v = i - 1
+        if !isassigned(caller_variables, v)
+            compute_missing_coeff!(caller_variables, caller::CallerMappingState, v)
+        end
+        substitution = caller_variables[i - 1]
+        isa(substitution, Incidence) || continue
+        for j in rowvals(substitution.row)
+            push!(used_caller_variables, j)
+        end
+    end
+    did_use_time = in(1, used_caller_variables)
 
+    for (i, coeff) in zip(rowvals(ret.row), nonzeros(ret.row))
         # Time dependence persists as itself
-        if v == 0
-            new_row[v_offset] += coeff
+        if i == 1
+            row[i] = coeff
             continue
         end
 
-        if !isassigned(coeffs, v)
-            @assert caller !== nothing
-            compute_missing_coeff!(coeffs, caller, v)
-        end
-
-        replacement = coeffs[v]
-        if isa(replacement, Incidence)
-            new_row .+= replacement.row .* coeff
-        else
-            if isa(replacement, Const)
-                if isa(const_val, Const)
-                    new_const_val = const_val.val + replacement.val * coeff
-                    if isa(new_const_val, Float64)
-                        const_val = Const(new_const_val)
-                    else
-                        const_val = widenconst(const_val)
-                    end
+        v = i - 1
+        substitution = caller_variables[v]
+        if isa(substitution, Incidence)
+            # Distribute the coefficient to all terms.
+            # Because the coefficient is expressed in the reference of the callee,
+            # state dependence must be processed carefully.
+            typ = compose_additive_term(typ, substitution.typ, coeff)
+            for (j, substitute) in zip(rowvals(substitution.row), nonzeros(substitution.row))
+                row[j] === nonlinear && continue # no more information to be gained
+                if substitute === nonlinear || coeff === nonlinear
+                    row[j] = nonlinear
+                elseif isa(coeff, Float64)
+                    row[j] += coeff * substitute
                 else
-                    const_val = widenconst(const_val)
+                    time_dependent = coeff.time_dependent
+                    state_dependent = false
+                    if isa(substitute, Linearity)
+                        time_dependent |= substitute.time_dependent
+                        state_dependent |= substitute.state_dependent
+                    end
+                    if coeff.state_dependent
+                        if coeff.time_dependent && did_use_time
+                            # The term is at least bilinear in another state, and this state
+                            # from the callee may alias time from the caller, so we must mark
+                            # time as nonlinear.
+                            row[1] = nonlinear
+                        end
+                        if count(==(j), used_caller_variables) > 1
+                            # The term is at least bilinear in another state, but we don't
+                            # know which state, so we must fall back to nonlinear.
+                            row[j] = nonlinear
+                            continue
+                        end
+                        # We'll only be state-dependent if variables from the callee
+                        # map to at least one other variable than `j`.
+                        state_dependent |= length(used_caller_variables) - did_use_time > 1
+                        # If another state may contain time, we may be time-dependent too.
+                        time_dependent |= did_use_time
+                    end
+                    j == 1 && (time_dependent = false)
+                    row[j] += Linearity(; nonlinear = false, state_dependent, time_dependent)
                 end
-            else
-                # The replacement has some unknown type - we need to widen
-                # all the way here.
-                return widenconst(const_val)
             end
+        elseif isa(substitution, Const)
+            typ = compose_additive_term(typ, substitution, coeff)
+        else
+            return widenconst(typ) # unknown lattice element, we should widen
         end
     end
 
-    return Incidence(const_val, new_row)
+    return Incidence(typ, row)
+end
+
+function compose_additive_term(@nospecialize(a), @nospecialize(b), coeff)
+    isa(a, Const) || return widenconst(a)
+    isa(b, Const) || return widenconst(a)
+    isa(coeff, Linearity) && return b.val == 0 ? a : widenconst(a)
+    val = a.val + b.val * coeff
+    isa(val, Float64) || return widenconst(a)
+    return Const(val)
 end
 
 function apply_linear_incidence(𝕃, ret::Eq, caller::CallerMappingState, mapping::CalleeMapping)