Merge branch 'master' into ox/nomacrotools

Author: YingboMa

.github/workflows/ci.yml

@@ -34,7 +34,7 @@ jobs:
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
-      - uses: actions/cache@v3
+      - uses: actions/cache@v4
         env:
           cache-name: cache-artifacts
         with:

Project.toml

@@ -19,6 +19,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 DynamicQuantities = "06fc5a27-2a28-4c7c-a15d-362465fb6821"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -73,8 +74,9 @@ Distributed = "1"
 Distributions = "0.23, 0.24, 0.25"
 DocStringExtensions = "0.7, 0.8, 0.9"
 DomainSets = "0.6"
-DynamicQuantities = "0.8, 0.9, 0.10"
+DynamicQuantities = "^0.11.2"
 ExprTools = "0.1.10"
+FindFirstFunctions = "1"
 ForwardDiff = "0.10.3"
 FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
@@ -111,6 +113,7 @@ julia = "1.9"
 [extras]
 AmplNLWriter = "7c4d4715-977e-5154-bfe0-e096adeac482"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"
 ControlSystemsMTK = "687d7614-c7e5-45fc-bfc3-9ee385575c88"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -134,4 +137,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsMTK", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "ControlSystemsMTK", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg"]

src/systems/alias_elimination.jl

@@ -153,7 +153,7 @@ the `constraint`.
         mask,
         constraint)
     eadj = M.row_cols
-    for i in range
+    @inbounds for i in range
         vertices = eadj[i]
         if constraint(length(vertices))
             for (j, v) in enumerate(vertices)
@@ -170,7 +170,7 @@ end
         range,
         mask,
         constraint)
-    for i in range
+    @inbounds for i in range
         row = @view M[i, :]
         if constraint(count(!iszero, row))
             for (v, val) in enumerate(row)
@@ -382,13 +382,6 @@ end
 
 swap!(v, i, j) = v[i], v[j] = v[j], v[i]
 
-function getcoeff(vars, coeffs, var)
-    for (vj, v) in enumerate(vars)
-        v == var && return coeffs[vj]
-    end
-    return 0
-end
-
 """
 $(SIGNATURES)
 

src/systems/clock_inference.jl

@@ -150,6 +150,7 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
     param_to_idx = Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
     offset = length(appended_parameters)
     affect_funs = []
+    init_funs = []
     svs = []
     clocks = TimeDomain[]
     for (i, (sys, input)) in enumerate(zip(syss, inputs))
@@ -202,6 +203,18 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
             push!(save_vec.args, :(p[$(input_offset + i)]))
         end
         empty_disc = isempty(disc_range)
+
+        disc_init = :(function (p, t)
+            d2c_obs = $disc_to_cont_obs
+            d2c_view = view(p, $disc_to_cont_idxs)
+            disc_state = view(p, $disc_range)
+            copyto!(d2c_view, d2c_obs(disc_state, p, t))
+        end)
+
+        # @show disc_to_cont_idxs
+        # @show cont_to_disc_idxs
+        # @show disc_range
+
         affect! = :(function (integrator, saved_values)
             @unpack u, p, t = integrator
             c2d_obs = $cont_to_disc_obs
@@ -212,27 +225,42 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
             d2c_view = view(p, $disc_to_cont_idxs)
             disc_state = view(p, $disc_range)
             disc = $disc
-            # Write continuous into to discrete: handles `Sample`
-            copyto!(c2d_view, c2d_obs(integrator.u, p, t))
-            # Write discrete into to continuous
-            # get old discrete states
-            copyto!(d2c_view, d2c_obs(disc_state, p, t))
+
             push!(saved_values.t, t)
             push!(saved_values.saveval, $save_vec)
-            # update discrete states
+
+            # Write continuous into to discrete: handles `Sample`
+            # Write discrete into to continuous
+            # Update discrete states
+
+            # At a tick, c2d must come first
+            # state update comes in the middle
+            # d2c comes last
+            # @show t
+            # @show "incoming", p
+            copyto!(c2d_view, c2d_obs(integrator.u, p, t))
+            # @show "after c2d", p
             $empty_disc || disc(disc_state, disc_state, p, t)
+            # @show "after state update", p
+            copyto!(d2c_view, d2c_obs(disc_state, p, t))
+            # @show "after d2c", p
         end)
         sv = SavedValues(Float64, Vector{Float64})
         push!(affect_funs, affect!)
+        push!(init_funs, disc_init)
         push!(svs, sv)
     end
     if eval_expression
         affects = map(affect_funs) do a
             drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
         end
+        inits = map(init_funs) do a
+            drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
+        end
     else
         affects = map(a -> toexpr(LiteralExpr(a)), affect_funs)
+        inits = map(a -> toexpr(LiteralExpr(a)), init_funs)
     end
     defaults = Dict{Any, Any}(v => 0.0 for v in Iterators.flatten(inputs))
-    return affects, clocks, svs, appended_parameters, defaults
+    return affects, inits, clocks, svs, appended_parameters, defaults
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -945,8 +945,9 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         has_difference = has_difference,
         check_length, kwargs...)
     cbs = process_events(sys; callback, has_difference, kwargs...)
+    inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
@@ -976,7 +977,13 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     if svs !== nothing
         kwargs1 = merge(kwargs1, (disc_saved_values = svs,))
     end
-    ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
+    prob = ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
+    if !isempty(inits)
+        for init in inits
+            init(prob.p, tspan[1])
+        end
+    end
+    prob
 end
 get_callback(prob::ODEProblem) = prob.kwargs[:callback]
 
@@ -1045,8 +1052,9 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     h = h_oop
     u0 = h(p, tspan[1])
     cbs = process_events(sys; callback, has_difference, kwargs...)
+    inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
@@ -1075,7 +1083,13 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     if svs !== nothing
         kwargs1 = merge(kwargs1, (disc_saved_values = svs,))
     end
-    DDEProblem{iip}(f, u0, h, tspan, p; kwargs1..., kwargs...)
+    prob = DDEProblem{iip}(f, u0, h, tspan, p; kwargs1..., kwargs...)
+    if !isempty(inits)
+        for init in inits
+            init(prob.p, tspan[1])
+        end
+    end
+    prob
 end
 
 function DiffEqBase.SDDEProblem(sys::AbstractODESystem, args...; kwargs...)
@@ -1099,8 +1113,9 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     h(p, t) = h_oop(p, t)
     u0 = h(p, tspan[1])
     cbs = process_events(sys; callback, has_difference, kwargs...)
+    inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
@@ -1140,8 +1155,15 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     else
         noise_rate_prototype = zeros(eltype(u0), size(noiseeqs))
     end
-    SDDEProblem{iip}(f, f.g, u0, h, tspan, p; noise_rate_prototype =
+    prob = SDDEProblem{iip}(f, f.g, u0, h, tspan, p;
+        noise_rate_prototype =
         noise_rate_prototype, kwargs1..., kwargs...)
+    if !isempty(inits)
+        for init in inits
+            init(prob.p, tspan[1])
+        end
+    end
+    prob
 end
 
 """

src/systems/model_parsing.jl

@@ -219,7 +219,7 @@ function generate_var!(dict, a, b, varclass;
     vd isa Vector && (vd = first(vd))
     vd[a] = Dict{Symbol, Any}()
     var = if indices === nothing
-        Symbolics.variable(a, T = SymbolicUtils.FnType{Tuple{Real}, Real})(iv)
+        Symbolics.variable(a, T = SymbolicUtils.FnType{Tuple{Any}, Real})(iv)
     else
         vd[a][:size] = Tuple(lastindex.(indices))
         first(@variables $a(iv)[indices...])

src/systems/sparsematrixclil.jl

@@ -129,6 +129,8 @@ end
 # build something that works for us here and worry about it later.
 nonzerosmap(a::CLILVector) = NonZeros(a)
 
+using FindFirstFunctions: findfirstequal
+
 function bareiss_update_virtual_colswap_mtk!(zero!, M::SparseMatrixCLIL, k, swapto, pivot,
         last_pivot; pivot_equal_optimization = true)
     # for ei in nzrows(>= k)
@@ -168,12 +170,11 @@ function bareiss_update_virtual_colswap_mtk!(zero!, M::SparseMatrixCLIL, k, swap
     # conservative, we leave it at this, as this captures the most important
     # case for MTK (where most pivots are `1` or `-1`).
     pivot_equal = pivot_equal_optimization && abs(pivot) == abs(last_pivot)
-
-    for ei in (k + 1):size(M, 1)
+    @inbounds for ei in (k + 1):size(M, 1)
         # eliminate `v`
         coeff = 0
         ivars = eadj[ei]
-        vj = findfirst(isequal(vpivot), ivars)
+        vj = findfirstequal(vpivot, ivars)
         if vj !== nothing
             coeff = old_cadj[ei][vj]
             deleteat!(old_cadj[ei], vj)
@@ -189,24 +190,118 @@ function bareiss_update_virtual_colswap_mtk!(zero!, M::SparseMatrixCLIL, k, swap
         ivars = eadj[ei]
         icoeffs = old_cadj[ei]
 
-        tmp_incidence = similar(eadj[ei], 0)
-        tmp_coeffs = similar(old_cadj[ei], 0)
-        # TODO: We know both ivars and kvars are sorted, we could just write
-        # a quick iterator here that does this without allocation/faster.
-        vars = sort(union(ivars, kvars))
-
-        for v in vars
-            v == vpivot && continue
-            ck = getcoeff(kvars, kcoeffs, v)
-            ci = getcoeff(ivars, icoeffs, v)
-            p1 = Base.Checked.checked_mul(pivot, ci)
-            p2 = Base.Checked.checked_mul(coeff, ck)
-            ci = exactdiv(Base.Checked.checked_sub(p1, p2), last_pivot)
-            if !iszero(ci)
-                push!(tmp_incidence, v)
-                push!(tmp_coeffs, ci)
+        numkvars = length(kvars)
+        numivars = length(ivars)
+        tmp_incidence = similar(eadj[ei], numkvars + numivars)
+        tmp_coeffs = similar(old_cadj[ei], numkvars + numivars)
+        tmp_len = 0
+        kvind = ivind = 0
+        if _debug_mode
+            # in debug mode, we at least check to confirm we're iterating over
+            # `v`s in the correct order
+            vars = sort(union(ivars, kvars))
+            vi = 0
+        end
+        if numivars > 0 && numkvars > 0
+            kvv = kvars[kvind += 1]
+            ivv = ivars[ivind += 1]
+            dobreak = false
+            while true
+                if kvv == ivv
+                    v = kvv
+                    ck = kcoeffs[kvind]
+                    ci = icoeffs[ivind]
+                    kvind += 1
+                    ivind += 1
+                    if kvind > numkvars
+                        dobreak = true
+                    else
+                        kvv = kvars[kvind]
+                    end
+                    if ivind > numivars
+                        dobreak = true
+                    else
+                        ivv = ivars[ivind]
+                    end
+                    p1 = Base.Checked.checked_mul(pivot, ci)
+                    p2 = Base.Checked.checked_mul(coeff, ck)
+                    ci = exactdiv(Base.Checked.checked_sub(p1, p2), last_pivot)
+                elseif kvv < ivv
+                    v = kvv
+                    ck = kcoeffs[kvind]
+                    kvind += 1
+                    if kvind > numkvars
+                        dobreak = true
+                    else
+                        kvv = kvars[kvind]
+                    end
+                    p2 = Base.Checked.checked_mul(coeff, ck)
+                    ci = exactdiv(Base.Checked.checked_neg(p2), last_pivot)
+                else # kvv > ivv
+                    v = ivv
+                    ci = icoeffs[ivind]
+                    ivind += 1
+                    if ivind > numivars
+                        dobreak = true
+                    else
+                        ivv = ivars[ivind]
+                    end
+                    ci = exactdiv(Base.Checked.checked_mul(pivot, ci), last_pivot)
+                end
+                if _debug_mode
+                    @assert v == vars[vi += 1]
+                end
+                if v != vpivot && !iszero(ci)
+                    tmp_incidence[tmp_len += 1] = v
+                    tmp_coeffs[tmp_len] = ci
+                end
+                dobreak && break
+            end
+        elseif numkvars > 0
+            ivind = 1
+            kvv = kvars[kvind += 1]
+        elseif numivars > 0
+            kvind = 1
+            ivv = ivars[ivind += 1]
+        end
+        if kvind <= numkvars
+            v = kvv
+            while true
+                if _debug_mode
+                    @assert v == vars[vi += 1]
+                end
+                if v != vpivot
+                    ck = kcoeffs[kvind]
+                    p2 = Base.Checked.checked_mul(coeff, ck)
+                    ci = exactdiv(Base.Checked.checked_neg(p2), last_pivot)
+                    if !iszero(ci)
+                        tmp_incidence[tmp_len += 1] = v
+                        tmp_coeffs[tmp_len] = ci
+                    end
+                end
+                (kvind == numkvars) && break
+                v = kvars[kvind += 1]
+            end
+        elseif ivind <= numivars
+            v = ivv
+            while true
+                if _debug_mode
+                    @assert v == vars[vi += 1]
+                end
+                if v != vpivot
+                    p1 = Base.Checked.checked_mul(pivot, icoeffs[ivind])
+                    ci = exactdiv(p1, last_pivot)
+                    if !iszero(ci)
+                        tmp_incidence[tmp_len += 1] = v
+                        tmp_coeffs[tmp_len] = ci
+                    end
+                end
+                (ivind == numivars) && break
+                v = ivars[ivind += 1]
             end
         end
+        resize!(tmp_incidence, tmp_len)
+        resize!(tmp_coeffs, tmp_len)
         eadj[ei] = tmp_incidence
         old_cadj[ei] = tmp_coeffs
     end