fix: properly eliminate identically zero variables in alias_elimination!

Project.toml

@@ -86,9 +86,9 @@ Libdl = "1"
 LinearAlgebra = "1"
 LinearSolve = "3.66"
 Logging = "1"
-ModelingToolkitBase = "1.50"
+ModelingToolkitBase = "1.51"
 ModelingToolkitStandardLibrary = "2.20"
-ModelingToolkitTearing = "1.17.2"
+ModelingToolkitTearing = "1.18"
 Moshi = "0.3.6"
 NonlinearSolve = "4.3"
 OffsetArrays = "1.3.0"
@@ -117,7 +117,7 @@ Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "2.10.0"
 SparseArrays = "1"
-StateSelection = "1.10.1"
+StateSelection = "1.10.5"
 StaticArrays = "1.9.14"
 StochasticDiffEq = "6.82.0, 7"
 SymbolicIndexingInterface = "0.3.46"

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -89,7 +89,7 @@ import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
 
 import DiffEqBase: @add_kwonly
 export independent_variables, unknowns, observables, parameters, bound_parameters,
-    continuous_events, discrete_events
+    continuous_events, discrete_events, analytically_integrated
 @reexport using Symbolics
 @reexport using UnPack
 RuntimeGeneratedFunctions.init(@__MODULE__)

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -706,15 +706,19 @@ function complete(
         if has_continuous_events(sys) && is_time_dependent(sys)
             cevts = SymbolicContinuousCallback[]
             for ev in get_continuous_events(sys)
-                ev = complete(ev; iv = get_iv(sys)::SymbolicT, extra_eqs = cb_alg_eqs)
+                ev = complete(
+                    ev; iv = get_iv(sys)::SymbolicT, parent_sys = sys
+                )
                 push!(cevts, ev)
             end
             @set! sys.continuous_events = cevts
         end
         if has_discrete_events(sys) && is_time_dependent(sys)
             devts = SymbolicDiscreteCallback[]
             for ev in get_discrete_events(sys)
-                ev = complete(ev; iv = get_iv(sys)::SymbolicT, extra_eqs = cb_alg_eqs)
+                ev = complete(
+                    ev; iv = get_iv(sys)::SymbolicT, parent_sys = sys
+                )
                 push!(devts, ev)
             end
             @set! sys.discrete_events = devts
@@ -961,6 +965,7 @@ const SYS_PROPS = [
     :isscheduled
     :costs
     :consolidate
+    :analytically_integrated
 ]
 
 for prop in SYS_PROPS
@@ -1593,6 +1598,17 @@ function unknowns(sys::AbstractSystem)
     return result
 end
 
+"""
+    $TYPEDSIGNATURES
+
+Return the list of analytically integrated variables in `sys`. Requires that `sys` is
+flattened, since analytically integrated variables cannot be specified for unflattened systems.
+"""
+function analytically_integrated(sys::AbstractSystem)
+    @assert isempty(get_systems(sys))
+    return get_analytically_integrated(sys)
+end
+
 """
     unknowns_toplevel(sys::AbstractSystem)
 
@@ -2589,17 +2605,26 @@ function Base.show(
     end
 
     # Print variables
-    for varfunc in [unknowns, parameters]
+    varsets = Any[unknowns]
+    if has_iv(sys) && get_iv(sys) isa SymbolicT && isempty(get_systems(sys)) &&
+            has_analytically_integrated(sys) && !isempty(analytically_integrated(sys))
+        push!(varsets, analytically_integrated)
+    end
+    push!(varsets, parameters)
+    for varfunc in varsets
         vars = varfunc(sys)
+        if varfunc === analytically_integrated
+            vars = keys(vars)
+        end
         nvars = length(vars)
         nvars == 0 && continue # skip
         header = titlecase(String(nameof(varfunc))) # e.g. "Unknowns"
+        header = replace(header, '_' => ' ')
         printstyled(io, "\n$header ($nvars):"; bold)
         hint && print(io, " see $(nameof(varfunc))($name)")
         nrows = min(nvars, limit ? rows : nvars)
         defs = has_bindings(sys) ? bindings(sys) : nothing
-        for i in 1:nrows
-            s = vars[i]
+        for s in Iterators.take(vars, nrows)
             print(io, "\n  ", s)
             if !isnothing(defs)
                 val = get(defs, s, nothing)

lib/ModelingToolkitBase/src/systems/callbacks.jl

@@ -133,13 +133,15 @@ end
 
 function AffectSystem(
         affect::Vector{Equation}; discrete_parameters = SymbolicT[],
-        iv = nothing, extra_eqs = Equation[], kwargs...
+        iv = nothing, parent_sys, kwargs...
     )
     isempty(affect) && return nothing
     if isnothing(iv)
         iv = t_nounits
         @warn "No independent variable specified. Defaulting to t_nounits."
     end
+    _unhack_sys = reverse_all_default_reversible_transformations(parent_sys)
+    extra_eqs = Equation[alg_equations(_unhack_sys); observed(_unhack_sys)]
     affect = [affect; extra_eqs]
 
     discrete_parameters = SymbolicAffect(affect; discrete_parameters).discrete_parameters
@@ -149,17 +151,47 @@ function AffectSystem(
             error("Non-time dependent parameter $p passed in as a discrete. Must be declared as @parameters $p(t).")
     end
 
+    parent_dvs = Set{SymbolicT}()
+    for var in [unknowns(parent_sys); observables(parent_sys)]
+        push!(parent_dvs, var)
+        push!(parent_dvs, split_indexed_var(var)[1])
+    end
+    for var in values(analytically_integrated(parent_sys))
+        push!(discrete_parameters, var)
+    end
+    parent_ps = Set{SymbolicT}(parameters(parent_sys))
+    for p in parameters(parent_sys)
+        p in parent_dvs && continue
+        push!(parent_ps, split_indexed_var(p)[1])
+    end
+    iv = get_iv(parent_sys)::SymbolicT
+
     dvs = OrderedSet{SymbolicT}()
     params = OrderedSet{SymbolicT}()
     _varsbuf = Set{SymbolicT}()
     for eq in affect
-        collect_vars!(dvs, params, eq, iv, Pre)
         empty!(_varsbuf)
-        SU.search_variables!(_varsbuf, eq; is_atomic = OperatorIsAtomic{Pre}())
-        filter!(x -> iscall(x) && operation(x) === Pre(), _varsbuf)
-        union!(params, _varsbuf)
-        diffvs = collect_applied_operators(eq, Differential)
-        union!(dvs, diffvs)
+        SU.search_variables!(_varsbuf, eq; is_atomic = OperatorIsAtomic{Union{Pre, Differential}}())
+        for var in _varsbuf
+            isequal(var, iv) && continue
+            Moshi.Match.@match var begin
+                BSImpl.Term(; f) && if f isa Pre end => begin
+                    push!(params, var)
+                    continue
+                end
+                BSImpl.Term(; f) && if f isa Differential end => begin
+                    push!(dvs, var)
+                    continue
+                end
+                _ => nothing
+            end
+
+            if var in parent_dvs
+                push!(dvs, var)
+                continue
+            end
+            push!(params, var)
+        end
     end
     pre_params = filter(haspre, params)
     sys_params = SymbolicT[]
@@ -1382,7 +1414,7 @@ Base.@nospecializeinfer function compile_equational_affect(
         @nospecialize(op = nothing), kwargs...
     )
     if aff isa AbstractVector
-        aff = make_affect(aff; iv = get_iv(sys))
+        aff = make_affect(aff; iv = get_iv(sys), parent_sys = sys)
     end
     if op === nothing
         op = default_operating_point(aff)

lib/ModelingToolkitBase/src/systems/diffeqs/basic_transformations.jl

@@ -1311,7 +1311,9 @@ function discover_maybe_zeros(sys::System)
     defs = copy(parent(bindings(sys)))
     left_merge!(defs, initial_conditions(sys))
     filter!(Base.Fix2(!==, COMMON_MISSING) ∘ last, defs)
-    subber = Symbolics.FixpointSubstituter{true}(AADSubWrapper(defs))
+    subber = Symbolics.FixpointSubstituter{true}(
+        AADSubWrapper(defs); maxiters = clamp(length(defs), 10, 1000), warn_maxiters = false
+    )
     mbzs = copy(maybe_zeros(sys))
     for k in [parameters(sys); unknowns(sys)]
         is_variable_numeric(k) || continue

lib/ModelingToolkitBase/src/systems/parameter_buffer.jl

@@ -108,7 +108,10 @@ function MTKParameters(
     # specified array values are never substituted into expressions (issue #4607).
     wrapped_op = AADSubWrapper(op)
     ir = get_irstructure(sys)
-    evaluate_varmap!(ir, wrapped_op, all_ps; limit = substitution_limit)
+    all_ps_no_init = filter(!isinitial, all_ps)
+    evaluate_varmap!(ir, wrapped_op, all_ps_no_init; limit = substitution_limit)
+    all_ps_init = setdiff(all_ps, all_ps_no_init)
+    evaluate_varmap!(ir, wrapped_op, all_ps_init; limit = substitution_limit, allow_symbolic = true)
 
     tunable_buffer = Vector{ic.tunable_buffer_size.type}(
         undef, ic.tunable_buffer_size.length
@@ -171,7 +174,11 @@ function MTKParameters(
 
     for sym in all_ps
         haskey(diffcache_params, sym) && continue
-        val = fixpoint_sub(sym, wrapped_op; maxiters = max(div(substitution_limit, 2), 2), fold = Val(true))
+        _sym = Moshi.Match.@match sym begin
+            BSImpl.Term(; f, args) && if f isa Pre end => args[1]
+            _ => sym
+        end 
+        val = get(wrapped_op, _sym, sym)
         ctype = symtype(sym)
         if !SU.isconst(val) && isinitial(sym)
             # The value of an `Initial` parameter that cannot be fully evaluated

lib/ModelingToolkitBase/src/systems/system.jl

@@ -340,6 +340,19 @@ struct System <: IntermediateDeprecationSystem
     The `Schedule` containing additional information about the simplified system.
     """
     schedule::Union{Schedule, Nothing}
+    """
+    Keys are variables that `mtkcompile` managed to integrate analytically. The analytically integrated
+    expression is present as an observed equation. Despite being observed, these variables
+    require initial conditions, typically to solve for auxililar parameters. For example,
+    `D(x) ~ 0` may be analytically solved as `x ~ x0`, in which case the new parameter `x0`
+    requires a value. Initialization solves for this value given an initial condition for `x`.
+    Currently, setting this field is only valid for flattened systems.
+
+    The values correspond to the introduced parameter for the constant of integration. This
+    must be solvable (have a binding of `missing`) and is allowed to be modified by callbacks
+    as a proxy for updating the analytically integrated variable.
+    """
+    analytically_integrated::Dict{SymbolicT, SymbolicT}
 
     function System(
             tag, eqs, noise_eqs, jumps, constraints, costs, consolidate, unknowns, ps,
@@ -354,8 +367,8 @@ struct System <: IntermediateDeprecationSystem
             preface = nothing, parent = nothing, initializesystem = nothing,
             is_initializesystem = false, is_discrete = false, state_priorities = AtomicMapT{Int}(),
             irreducibles = AtomicSetT(), maybe_zeros = AtomicSetT(),
-            irstructure_tlv = __new_irstructure_tlv(), isscheduled = false, schedule = nothing;
-            checks::Union{Bool, Int} = true
+            irstructure_tlv = __new_irstructure_tlv(), isscheduled = false, schedule = nothing,
+            analytically_integrated = Dict{SymbolicT, SymbolicT}(); checks::Union{Bool, Int} = true
         )
         if is_initializesystem && iv !== nothing
             throw(
@@ -418,7 +431,7 @@ struct System <: IntermediateDeprecationSystem
             complete, index_cache, parameter_bindings_graph, ignored_connections,
             preface, parent, initializesystem, is_initializesystem, is_discrete,
             state_priorities, irreducibles, maybe_zeros, irstructure_tlv,
-            isscheduled, schedule
+            isscheduled, schedule, analytically_integrated
         )
     end
 end

lib/ModelingToolkitBase/test/basic_transformations.jl

@@ -112,7 +112,7 @@ end
     M3 = change_independent_variable(M2, b, extraeqs)
 
     if @isdefined(ModelingToolkit)
-        M1 = mtkcompile(M1)
+        M1 = mtkcompile(M1; eliminate_mm_zeros = false)
         M2 = mtkcompile(M2; allow_symbolic = true)
         M3 = mtkcompile(M3; allow_symbolic = true)
         @test length(unknowns(M3)) == length(unknowns(M2)) == length(unknowns(M1)) - 1

lib/ModelingToolkitBase/test/initializationsystem.jl

@@ -656,7 +656,7 @@ end
         sys1,
         System([D(y) ~ 0], t; initialization_eqs = [y ~ 2], name = :sys2)
     ) |> mtkcompile
-    ics2 = unknowns(sys1) .=> 2 # should be equivalent to "ics2 = [x => 2]"
+    ics2 = [x => 2]
     prob2 = ODEProblem(sys2, ics2, (0.0, 1.0); fully_determined = true)
     sol2 = solve(prob2, Tsit5(); abstol = 1.0e-6, reltol = 1.0e-6)
     @test SciMLBase.successful_retcode(sol2)

lib/ModelingToolkitBase/test/odesystem.jl

@@ -1062,9 +1062,9 @@ end
 
 @testset "Non-1-indexed variable array (issue #2670)" begin
     @variables x(t)[0:1] # 0-indexed variable array
-    @named sys = System([x[0] ~ 0.0, D(x[1]) ~ x[0]], t, [x], [])
+    @named sys = System([x[0] ~ 1.0, D(x[1]) ~ x[0]], t, [x], [])
     sys = @test_nowarn mtkcompile(sys)
-    @test full_equations(sys) == [D(x[1]) ~ 0.0]
+    @test full_equations(sys) == [D(x[1]) ~ 1.0]
 end
 
 # Namespacing of array variables

lib/ModelingToolkitBase/test/symbolic_events.jl

@@ -963,9 +963,9 @@ if @isdefined(ModelingToolkit)
             ]
 
             discrete_events = [
-                [30] => [binary_valve_1.S ~ 0.0, binary_valve_2.Δp ~ 0.0]
+                [30] => [binary_valve_1.S ~ 0.0, binary_valve_2.Δp ~ 0.0, binary_valve_2.S ~ 0.0]
                 [60] => [binary_valve_1.S ~ 1.0, binary_valve_2.Δp ~ 1.0]
-                [120] => [binary_valve_1.S ~ 0.0, binary_valve_2.Δp ~ 0.0]
+                [120] => [binary_valve_1.S ~ 0.0, binary_valve_2.S ~ Pre(binary_valve_2.S)]
             ]
 
             return System(equations, t, vars, pars; name, systems, discrete_events)
@@ -981,7 +981,8 @@ if @isdefined(ModelingToolkit)
         # constant after that point anyway. Just make sure it hits the last event and
         # had the correct `u`.
         @test sol.t[end] >= 120.0
-        @test sol[[sys.binary_valve_1.S, sys.binary_valve_2.Δp]][end] == [0.0, 0.0]
+        @test sol[sys.binary_valve_1.S][end] == 0.0
+        @test sol[sys.binary_valve_2.S][end] == 0.0
     end
 end