Merge branch 'master' into myb/dd

Author: YingboMa

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 = "9.39.1"
+version = "9.41.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -24,6 +24,7 @@ ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 Expronicon = "6b7a57c9-7cc1-4fdf-b7f5-e857abae3636"
 FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -76,10 +77,11 @@ ChainRulesCore = "1"
 Combinatorics = "1"
 Compat = "3.42, 4"
 ConstructionBase = "1"
+DataInterpolations = "6.4"
 DataStructures = "0.17, 0.18"
 DeepDiffs = "1"
 DiffEqBase = "6.103.0"
-DiffEqCallbacks = "2.16, 3"
+DiffEqCallbacks = "2.16, 3, 4"
 DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"
 Distributed = "1"
@@ -91,6 +93,7 @@ ExprTools = "0.1.10"
 Expronicon = "0.8"
 FindFirstFunctions = "1"
 ForwardDiff = "0.10.3"
+FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
 InteractiveUtils = "1"
@@ -118,8 +121,8 @@ SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicIndexingInterface = "0.3.29"
-SymbolicUtils = "3.2"
-Symbolics = "6.3"
+SymbolicUtils = "3.7"
+Symbolics = "6.12"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
@@ -129,6 +132,7 @@ julia = "1.9"
 AmplNLWriter = "7c4d4715-977e-5154-bfe0-e096adeac482"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -154,4 +158,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET"]

docs/src/basics/FAQ.md

@@ -244,3 +244,62 @@ D = Differential(x)
 @variables y(x)
 @named sys = ODESystem([D(y) ~ x], x)
 ```
+
+## Ordering of tunable parameters
+
+Tunable parameters are floating point parameters, not used in callbacks and not marked with `tunable = false` in their metadata. These are expected to be used with AD
+and optimization libraries. As such, they are stored together in one `Vector{T}`. To obtain the ordering of tunable parameters in this buffer, use:
+
+```@docs
+tunable_parameters
+```
+
+If you have an array in which a particular dimension is in the order of tunable parameters (e.g. the jacobian with respect to tunables) then that dimension of the
+array can be reordered into the required permutation using the symbolic variables:
+
+```@docs
+reorder_dimension_by_tunables!
+reorder_dimension_by_tunables
+```
+
+For example:
+
+```@example reorder
+using ModelingToolkit
+
+@parameters p q[1:3] r[1:2, 1:2]
+
+@named sys = ODESystem(Equation[], ModelingToolkit.t_nounits, [], [p, q, r])
+sys = complete(sys)
+```
+
+The canonicalized tunables portion of `MTKParameters` will be in the order of tunables:
+
+```@example reorder
+using SciMLStructures: canonicalize, Tunable
+
+ps = MTKParameters(sys, [p => 1.0, q => [2.0, 3.0, 4.0], r => [5.0 6.0; 7.0 8.0]])
+arr = canonicalize(Tunable(), ps)[1]
+```
+
+We can reorder this to contain the value for `p`, then all values for `q`, then for `r` using:
+
+```@example reorder
+reorder_dimension_by_tunables(sys, arr, [p, q, r])
+```
+
+This also works with interleaved subarrays of symbolics:
+
+```@example reorder
+reorder_dimension_by_tunables(sys, arr, [q[1], r[1, :], q[2], r[2, :], q[3], p])
+```
+
+And arbitrary dimensions of higher dimensional arrays:
+
+```@example reorder
+highdimarr = stack([i * arr for i in 1:5]; dims = 1)
+```
+
+```@example reorder
+reorder_dimension_by_tunables(sys, highdimarr, [q[1:2], r[1, :], q[3], r[2, :], p]; dim = 2)
+```

docs/src/basics/InputOutput.md

@@ -55,7 +55,6 @@ We can inspect the state realization chosen by MTK
 x_sym
 ```
 
-as expected, `x` is chosen as the state variable.
 as expected, `x` is chosen as the state variable.
 
 ```@example inputoutput

docs/src/basics/MTKLanguage.md

@@ -63,13 +63,14 @@ end
     @structural_parameters begin
         f = sin
         N = 2
+        M = 3
     end
     begin
         v_var = 1.0
     end
     @variables begin
         v(t) = v_var
-        v_array(t)[1:2, 1:3]
+        v_array(t)[1:N, 1:M]
         v_for_defaults(t)
     end
     @extend ModelB(; p1)
@@ -310,10 +311,10 @@ end
   - `:defaults`: Dictionary of variables and default values specified in the `@defaults`.
   - `:extend`: The list of extended unknowns, name given to the base system, and name of the base system.
   - `:structural_parameters`: Dictionary of structural parameters mapped to their metadata.
-  - `:parameters`: Dictionary of symbolic parameters mapped to their metadata. For
-    parameter arrays, length is added to the metadata as `:size`.
-  - `:variables`: Dictionary of symbolic variables mapped to their metadata. For
-    variable arrays, length is added to the metadata as `:size`.
+  - `:parameters`: Dictionary of symbolic parameters mapped to their metadata. Metadata of
+    the parameter arrays is, for now, omitted.
+  - `:variables`: Dictionary of symbolic variables mapped to their metadata. Metadata of
+    the variable arrays is, for now, omitted.
   - `:kwargs`: Dictionary of keyword arguments mapped to their metadata.
   - `:independent_variable`: Independent variable, which is added while generating the Model.
   - `:equations`: List of equations (represented as strings).
@@ -324,10 +325,10 @@ For example, the structure of `ModelC` is:
 julia> ModelC.structure
 Dict{Symbol, Any} with 10 entries:
   :components            => Any[Union{Expr, Symbol}[:model_a, :ModelA], Union{Expr, Symbol}[:model_array_a, :ModelA, :(1:N)], Union{Expr, Symbol}[:model_array_b, :ModelA, :(1:N)]]
-  :variables             => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var, :type=>Real), :v_array=>Dict(:type=>Real, :size=>(2, 3)), :v_for_defaults=>Dict(:type=>Real))
+  :variables             => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var, :type=>Real), :v_for_defaults=>Dict(:type=>Real))
   :icon                  => URI("https://github.com/SciML/SciMLDocs/blob/main/docs/src/assets/logo.png")
-  :kwargs                => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :N=>Dict(:value=>2), :v=>Dict{Symbol, Any}(:value=>:v_var, :type=>Real), :v_array=>Dict{Symbol, Union{Nothing, UnionAll}}(:value=>nothing, :type=>AbstractArray{Real}), :v_for_defaults=>Dict{Symbol, Union{Nothing, DataType}}(:value=>nothing, :type=>Real), :p1=>Dict(:value=>nothing))
-  :structural_parameters => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :N=>Dict(:value=>2))
+  :kwargs                => Dict{Symbol, Dict}(:f => Dict(:value => :sin), :N => Dict(:value => 2), :M => Dict(:value => 3), :v => Dict{Symbol, Any}(:value => :v_var, :type => Real), :v_for_defaults => Dict{Symbol, Union{Nothing, DataType}}(:value => nothing, :type => Real), :p1 => Dict(:value => nothing)),
+  :structural_parameters => Dict{Symbol, Dict}(:f => Dict(:value => :sin), :N => Dict(:value => 2), :M => Dict(:value => 3))
   :independent_variable  => t
   :constants             => Dict{Symbol, Dict}(:c=>Dict{Symbol, Any}(:value=>1, :type=>Int64, :description=>"Example constant."))
   :extend                => Any[[:p2, :p1], Symbol("#mtkmodel__anonymous__ModelB"), :ModelB]

ext/MTKChainRulesCoreExt.jl

@@ -2,14 +2,99 @@ module MTKChainRulesCoreExt
 
 import ModelingToolkit as MTK
 import ChainRulesCore
-import ChainRulesCore: NoTangent
+import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
 
 function ChainRulesCore.rrule(::Type{MTK.MTKParameters}, tunables, args...)
     function mtp_pullback(dt)
+        dt = unthunk(dt)
         (NoTangent(), dt.tunable[1:length(tunables)],
             ntuple(_ -> NoTangent(), length(args))...)
     end
     MTK.MTKParameters(tunables, args...), mtp_pullback
 end
 
+function subset_idxs(idxs, portion, template)
+    ntuple(Val(length(template))) do subi
+        [Base.tail(idx.idx) for idx in idxs if idx.portion == portion && idx.idx[1] == subi]
+    end
+end
+
+selected_tangents(::NoTangent, _) = ()
+selected_tangents(::ZeroTangent, _) = ZeroTangent()
+function selected_tangents(
+        tangents::AbstractArray{T}, idxs::Vector{Tuple{Int}}) where {T <: Number}
+    selected_tangents(tangents, map(only, idxs))
+end
+function selected_tangents(tangents::AbstractArray{T}, idxs...) where {T <: Number}
+    newtangents = copy(tangents)
+    view(newtangents, idxs...) .= zero(T)
+    newtangents
+end
+function selected_tangents(
+        tangents::AbstractVector{T}, idxs) where {S <: Number, T <: AbstractArray{S}}
+    newtangents = copy(tangents)
+    for i in idxs
+        j, k... = i
+        if k == ()
+            newtangents[j] = zero(newtangents[j])
+        else
+            newtangents[j] = selected_tangents(newtangents[j], k...)
+        end
+    end
+    newtangents
+end
+function selected_tangents(tangents::AbstractVector{T}, idxs) where {T <: AbstractArray}
+    newtangents = similar(tangents, Union{T, NoTangent})
+    copyto!(newtangents, tangents)
+    for i in idxs
+        j, k... = i
+        if k == ()
+            newtangents[j] = NoTangent()
+        else
+            newtangents[j] = selected_tangents(newtangents[j], k...)
+        end
+    end
+    newtangents
+end
+function selected_tangents(
+        tangents::Union{Tangent{<:Tuple}, Tangent{T, <:Tuple}}, idxs) where {T}
+    ntuple(Val(length(tangents))) do i
+        selected_tangents(tangents[i], idxs[i])
+    end
+end
+
+function ChainRulesCore.rrule(
+        ::typeof(MTK.remake_buffer), indp, oldbuf::MTK.MTKParameters, idxs, vals)
+    if idxs isa AbstractSet
+        idxs = collect(idxs)
+    end
+    idxs = map(idxs) do i
+        i isa MTK.ParameterIndex ? i : MTK.parameter_index(indp, i)
+    end
+    newbuf = MTK.remake_buffer(indp, oldbuf, idxs, vals)
+    tunable_idxs = reduce(
+        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Tunable))
+    disc_idxs = subset_idxs(idxs, MTK.SciMLStructures.Discrete(), oldbuf.discrete)
+    const_idxs = subset_idxs(idxs, MTK.SciMLStructures.Constants(), oldbuf.constant)
+    nn_idxs = subset_idxs(idxs, MTK.NONNUMERIC_PORTION, oldbuf.nonnumeric)
+
+    pullback = let idxs = idxs
+        function remake_buffer_pullback(buf′)
+            buf′ = unthunk(buf′)
+            f′ = NoTangent()
+            indp′ = NoTangent()
+
+            tunable = selected_tangents(buf′.tunable, tunable_idxs)
+            discrete = selected_tangents(buf′.discrete, disc_idxs)
+            constant = selected_tangents(buf′.constant, const_idxs)
+            nonnumeric = selected_tangents(buf′.nonnumeric, nn_idxs)
+            oldbuf′ = Tangent{typeof(oldbuf)}(; tunable, discrete, constant, nonnumeric)
+            idxs′ = NoTangent()
+            vals′ = map(i -> MTK._ducktyped_parameter_values(buf′, i), idxs)
+            return f′, indp′, oldbuf′, idxs′, vals′
+        end
+    end
+    newbuf, pullback
+end
+
 end

src/ModelingToolkit.jl

@@ -38,6 +38,7 @@ using Base: RefValue
 using Combinatorics
 import Distributions
 import FunctionWrappersWrappers
+import FunctionWrappers: FunctionWrapper
 using URIs: URI
 using SciMLStructures
 using Compat
@@ -63,7 +64,8 @@ using Symbolics: _parse_vars, value, @derivatives, get_variables,
                  VariableSource, getname, variable, Connection, connect,
                  NAMESPACE_SEPARATOR, set_scalar_metadata, setdefaultval,
                  initial_state, transition, activeState, entry, hasnode,
-                 ticksInState, timeInState, fixpoint_sub, fast_substitute
+                 ticksInState, timeInState, fixpoint_sub, fast_substitute,
+                 CallWithMetadata, CallWithParent
 const NAMESPACE_SEPARATOR_SYMBOL = Symbol(NAMESPACE_SEPARATOR)
 import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
                   jacobian_sparsity, isaffine, islinear, _iszero, _isone,
@@ -276,6 +278,6 @@ export debug_system
 export Sample, Hold, Shift, ShiftIndex, sampletime, SampleTime
 export Clock, SolverStepClock, TimeDomain
 
-export MTKParameters
+export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables
 
 end # module

src/parameters.jl

@@ -26,6 +26,16 @@ function isparameter(x)
     end
 end
 
+function iscalledparameter(x)
+    x = unwrap(x)
+    return isparameter(getmetadata(x, CallWithParent, nothing))
+end
+
+function getcalledparameter(x)
+    x = unwrap(x)
+    return getmetadata(x, CallWithParent)
+end
+
 """
     toparam(s)
 

src/systems/abstractsystem.jl

@@ -882,6 +882,44 @@ namespace its subsystems or variables, i.e. `isequal(complete(sys).v.i, v.i)`.
 function complete(sys::AbstractSystem; split = true)
     if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)
+        all_ps = parameters(sys)
+        if !isempty(all_ps)
+            # reorder parameters by portions
+            ps_split = reorder_parameters(sys, all_ps)
+            # if there are no tunables, vcat them
+            if isempty(get_index_cache(sys).tunable_idx)
+                ordered_ps = reduce(vcat, ps_split)
+            else
+                # if there are tunables, they will all be in `ps_split[1]`
+                # and the arrays will have been scalarized
+                ordered_ps = eltype(all_ps)[]
+                i = 1
+                # go through all the tunables
+                while i <= length(ps_split[1])
+                    sym = ps_split[1][i]
+                    # if the sym is not a scalarized array symbolic OR it was already scalarized,
+                    # just push it as-is
+                    if !iscall(sym) || operation(sym) != getindex ||
+                       any(isequal(sym), all_ps)
+                        push!(ordered_ps, sym)
+                        i += 1
+                        continue
+                    end
+                    # the next `length(sym)` symbols should be scalarized versions of the same
+                    # array symbolic
+                    if !allequal(first(arguments(x))
+                    for x in view(ps_split[1], i:(i + length(sym) - 1)))
+                        error("This should not be possible. Please open an issue in ModelingToolkit.jl with an MWE and stacktrace.")
+                    end
+                    arrsym = first(arguments(sym))
+                    push!(ordered_ps, arrsym)
+                    i += length(arrsym)
+                end
+                ordered_ps = vcat(
+                    ordered_ps, reduce(vcat, ps_split[2:end]; init = eltype(ordered_ps)[]))
+            end
+            @set! sys.ps = ordered_ps
+        end
     end
     if isdefined(sys, :initializesystem) && get_initializesystem(sys) !== nothing
         @set! sys.initializesystem = complete(get_initializesystem(sys); split)

src/systems/diffeqs/abstractodesystem.jl

@@ -817,11 +817,12 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
          ModelingToolkit.get_tearing_state(sys) !== nothing) ||
         !isempty(initialization_equations(sys))) && t !== nothing
         if eltype(u0map) <: Number
-            u0map = unknowns(sys) .=> u0map
+            u0map = unknowns(sys) .=> vec(u0map)
         end
-        if isempty(u0map)
+        if u0map === nothing || isempty(u0map)
             u0map = Dict()
         end
+
         initializeprob = ModelingToolkit.InitializationProblem(
             sys, t, u0map, parammap; guesses, warn_initialize_determined,
             initialization_eqs, eval_expression, eval_module, fully_determined, check_units)

src/systems/discrete_system/discrete_system.jl

@@ -247,6 +247,13 @@ function process_DiscreteProblem(constructor, sys::DiscreteSystem, u0map, paramm
     dvs = unknowns(sys)
     ps = parameters(sys)
 
+    if eltype(u0map) <: Number
+        u0map = unknowns(sys) .=> vec(u0map)
+    end
+    if u0map === nothing || isempty(u0map)
+        u0map = Dict()
+    end
+
     trueu0map = Dict()
     for (k, v) in u0map
         k = unwrap(k)