Merge branch 'SciML:master' into iss3707

Author: fchen121

.github/workflows/Tests.yml

@@ -41,7 +41,7 @@ jobs:
           - Downstream
           - RegressionI
           - FMI
-    uses: "SciML/.github/.github/workflows/tests.yml@v1"
+    uses: "SciML/.github/.github/workflows/tests.yml@master"
     with:
       julia-version: "${{ matrix.version }}"
       group: "${{ matrix.group }}"

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Yingbo Ma <[email protected]>", "Chris Rackauckas <[email protected]> and contributors"]
-version = "10.10.0"
+version = "10.13.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -151,7 +151,7 @@ RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.0.0"
-SciMLBase = "2.100.0"
+SciMLBase = "2.104.0"
 SciMLPublic = "1.0.0"
 SciMLStructures = "1.7"
 Serialization = "1"

src/ModelingToolkit.jl

@@ -176,6 +176,7 @@ include("problems/docs.jl")
 include("systems/codegen.jl")
 include("systems/problem_utils.jl")
 include("linearization.jl")
+include("systems/solver_nlprob.jl")
 
 include("problems/compatibility.jl")
 include("problems/odeproblem.jl")

src/linearization.jl

@@ -852,7 +852,7 @@ function reorder_unknowns(sys::NamedTuple, old, new)
     issorted(perm) && return sys # shortcut return, no reordering
     P = zeros(Int, nx, nx)
     for i in 1:nx # Build permutation matrix
-        P[i, perm[i]] = 1
+        P[perm[i], i] = 1
     end
     similarity_transform(sys, P; unitary = true)
 end

src/problems/odeproblem.jl

@@ -3,7 +3,7 @@
         t = nothing, eval_expression = false, eval_module = @__MODULE__, sparse = false,
         steady_state = false, checkbounds = false, sparsity = false, analytic = nothing,
         simplify = false, cse = true, initialization_data = nothing, expression = Val{false},
-        check_compatibility = true, kwargs...) where {iip, spec}
+        check_compatibility = true, nlstep = false, kwargs...) where {iip, spec}
     check_complete(sys, ODEFunction)
     check_compatibility && check_compatible_system(ODEFunction, sys)
 
@@ -41,6 +41,12 @@
     M = calculate_massmatrix(sys)
     _M = concrete_massmatrix(M; sparse, u0)
 
+    if nlstep
+        ode_nlstep = generate_ODENLStepData(sys, u0, p, M)
+    else
+        ode_nlstep = nothing
+    end
+
     observedfun = ObservedFunctionCache(
         sys; expression, steady_state, eval_expression, eval_module, checkbounds, cse)
 
@@ -57,7 +63,8 @@
         observed = observedfun,
         sparsity = sparsity ? _W_sparsity : nothing,
         analytic = analytic,
-        initialization_data)
+        initialization_data,
+        nlstep_data = ode_nlstep)
 
     maybe_codegen_scimlfn(expression, ODEFunction{iip, spec}, args; kwargs...)
 end

src/systems/abstractsystem.jl

@@ -2757,7 +2757,9 @@ function process_parameter_equations(sys::AbstractSystem)
             is_parameter(sys, sym) ||
                 symbolic_type(sym) == ArraySymbolic() &&
                 is_sized_array_symbolic(sym) &&
-                all(Base.Fix1(is_parameter, sys), collect(sym))
+                all(Base.Fix1(is_parameter, sys), collect(sym)) ||
+                iscall(sym) &&
+                operation(sym) === getindex && is_parameter(sys, arguments(sym)[1])
         end
             # Everything in `varsbuf` is a parameter, so this is a cheap `is_parameter`
             # check.

src/systems/solver_nlprob.jl

@@ -0,0 +1,76 @@
+function generate_ODENLStepData(sys::System, u0, p, mm = calculate_massmatrix(sys))
+    nlsys, outer_tmp, inner_tmp = inner_nlsystem(sys, mm)
+    state = ProblemState(; u = u0, p)
+    op = Dict()
+    op[ODE_GAMMA[1]] = one(eltype(u0))
+    op[ODE_GAMMA[2]] = one(eltype(u0))
+    op[ODE_GAMMA[3]] = one(eltype(u0))
+    op[ODE_C] = zero(eltype(u0))
+    op[outer_tmp] = zeros(eltype(u0), size(outer_tmp))
+    op[inner_tmp] = zeros(eltype(u0), size(inner_tmp))
+    for v in [unknowns(nlsys); parameters(nlsys)]
+        haskey(op, v) && continue
+        op[v] = getsym(sys, v)(state)
+    end
+    nlprob = NonlinearProblem(nlsys, op; build_initializeprob = false)
+
+    subsetidxs = [findfirst(isequal(y),unknowns(sys)) for y in unknowns(nlsys)]
+    set_gamma_c = setsym(nlsys, (ODE_GAMMA..., ODE_C))
+    set_outer_tmp = setsym(nlsys, outer_tmp)
+    set_inner_tmp = setsym(nlsys, inner_tmp)
+    nlprobmap = generate_nlprobmap(sys, nlsys)
+
+    return SciMLBase.ODENLStepData(nlprob, subsetidxs, set_gamma_c, set_outer_tmp, set_inner_tmp, nlprobmap)
+end
+
+const ODE_GAMMA = @parameters γ₁ₘₜₖ, γ₂ₘₜₖ, γ₃ₘₜₖ 
+const ODE_C = only(@parameters cₘₜₖ)
+
+function get_outer_tmp(n::Int)
+    only(@parameters outer_tmpₘₜₖ[1:n])
+end
+
+function get_inner_tmp(n::Int)
+    only(@parameters inner_tmpₘₜₖ[1:n])
+end
+
+function inner_nlsystem(sys::System, mm)
+    dvs = unknowns(sys)
+    eqs = full_equations(sys)
+    t = get_iv(sys)
+    N = length(dvs)
+    @assert length(eqs) == N
+    @assert mm == I || size(mm) == (N, N)
+    rhss = [eq.rhs for eq in eqs]
+    gamma1, gamma2, gamma3 = ODE_GAMMA
+    c = ODE_C
+    outer_tmp = get_outer_tmp(N)
+    inner_tmp = get_inner_tmp(N)
+
+    subrules = Dict([v => gamma2*v + inner_tmp[i] for (i, v) in enumerate(dvs)])
+    subrules[t] = c
+    new_rhss = map(Base.Fix2(fast_substitute, subrules), rhss)
+    new_rhss = collect(outer_tmp) .+ gamma1 .* new_rhss .- gamma3 * mm * dvs
+    new_eqs = [0 ~ rhs for rhs in new_rhss]
+
+    new_dvs = unknowns(sys)
+    new_ps = [parameters(sys); [gamma1, gamma2, gamma3, c, inner_tmp, outer_tmp]]
+    nlsys = mtkcompile(System(new_eqs, new_dvs, new_ps; name = :nlsys); split = is_split(sys))
+    return nlsys, outer_tmp, inner_tmp
+end
+
+struct NLStep_probmap{F}
+    f::F
+end
+
+function (nlp::NLStep_probmap)(buffer, nlsol)
+    nlp.f(buffer, state_values(nlsol), parameter_values(nlsol))
+end
+
+function (nlp::NLStep_probmap)(nlsol)
+    nlp.f(state_values(nlsol), parameter_values(nlsol))
+end
+
+function generate_nlprobmap(sys::System, nlsys::System)
+    return NLStep_probmap(build_explicit_observed_function(nlsys, unknowns(sys)))
+end

test/linearize.jl

@@ -1,6 +1,26 @@
 using ModelingToolkit, ADTypes, Test
 using CommonSolve: solve
 
+# Test reorder_unknowns
+# sys = ssrand(1,1,4);
+mats = let
+    A = [-1.617708540405859 0.14199864151523162 1.8120551022076838 -1.246419696614408;
+         0.6704209450894298 -2.4251566699889575 0.6782529705706082 -1.3731519847672025;
+         -0.09336677360807291 -0.11211714788917712 -3.6877851408229523 -0.7073967284605489;
+         -1.1743200892334098 1.1808779444006103 1.5721685015907167 -0.10858833182921268]
+    B = [-0.3286766047686936; -1.8473436385672866; -2.4092567234250954; -0.06371974677173559;;]
+    C = [-0.7144567541084362 0.18898849455229796 0.023473101245754475 1.0369097263843963]
+    D = [0.6397583934617636;;]
+    (; A, B, C, D)
+end
+@variables x1 x2 x3 x4
+new = [x4, x1, x3, x2]
+old = [x1, x2, x3, x4]
+lsys = ModelingToolkit.reorder_unknowns(mats, old, new)
+P = [0 1 0 0; 0 0 0 1; 0 0 1 0; 1 0 0 0]
+@test isequal(P*new, old)
+@test lsys.A == ModelingToolkit.similarity_transform(mats, P).A
+
 # r is an input, and y is an output.
 @independent_variables t
 @variables x(t)=0 y(t)=0 u(t)=0 r(t)=0

test/parameter_dependencies.jl

@@ -389,3 +389,12 @@ end
         [x(k - 1) ~ x(k) + y(k) + p2, p2 ~ 2p1], t)
     @test is_parameter(sys, p1)
 end
+
+@testset "Scalarized array as RHS of parameter dependency" begin
+    @parameters p[1:2] p1 p2
+    @variables x(t)
+    @named sys = System([D(x) ~ x, p1 ~ p[1], p2 ~ p[2]], t)
+    @test any(isequal(p), ModelingToolkit.get_ps(sys))
+    sys = mtkcompile(sys)
+    @test length(ModelingToolkit.parameter_dependencies(sys)) == 2
+end