Merge pull request #3069 from AayushSabharwal/as/tunable-order

feat: add `reorder_dimension_by_tunables!`, ensure `tunable_parameters` is correctly ordered

Author: ChrisRackauckas

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)
+```

src/ModelingToolkit.jl

@@ -278,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/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/index_cache.jl

@@ -505,3 +505,55 @@ end
 fntype_to_function_type(::Type{FnType{A, R, T}}) where {A, R, T} = T
 fntype_to_function_type(::Type{FnType{A, R, Nothing}}) where {A, R} = FunctionWrapper{R, A}
 fntype_to_function_type(::Type{FnType{A, R}}) where {A, R} = FunctionWrapper{R, A}
+
+"""
+    reorder_dimension_by_tunables!(dest::AbstractArray, sys::AbstractSystem, arr::AbstractArray, syms; dim = 1)
+
+Assuming the order of values in dimension `dim` of `arr` correspond to the order of tunable
+parameters in the system, reorder them according to the order described in `syms`. `syms` must
+be a permutation of `tunable_parameters(sys)`. The result is written to `dest`. The `size` of `dest` and
+`arr` must be equal. Return `dest`.
+
+See also: [`MTKParameters`](@ref), [`tunable_parameters`](@ref), [`reorder_dimension_by_tunables`](@ref).
+"""
+function reorder_dimension_by_tunables!(
+        dest::AbstractArray, sys::AbstractSystem, arr::AbstractArray, syms; dim = 1)
+    if !iscomplete(sys)
+        throw(ArgumentError("A completed system is required. Call `complete` or `structural_simplify` on the system."))
+    end
+    if !has_index_cache(sys) || (ic = get_index_cache(sys)) === nothing
+        throw(ArgumentError("The system does not have an index cache. Call `complete` or `structural_simplify` on the system with `split = true`."))
+    end
+    if size(dest) != size(arr)
+        throw(ArgumentError("Source and destination arrays must have the same size. Got source array with size $(size(arr)) and destination with size $(size(dest))."))
+    end
+
+    dsti = 1
+    for sym in syms
+        idx = parameter_index(ic, sym)
+        if !(idx.portion isa SciMLStructures.Tunable)
+            throw(ArgumentError("`syms` must be a permutation of `tunable_parameters(sys)`. Found $sym which is not a tunable parameter."))
+        end
+
+        dstidx = ntuple(
+            i -> i == dim ? (dsti:(dsti + length(sym) - 1)) : (:), Val(ndims(arr)))
+        destv = @view dest[dstidx...]
+        dsti += length(sym)
+        arridx = ntuple(i -> i == dim ? (idx.idx) : (:), Val(ndims(arr)))
+        srcv = @view arr[arridx...]
+        copyto!(destv, srcv)
+    end
+    return dest
+end
+
+"""
+    reorder_dimension_by_tunables(sys::AbstractSystem, arr::AbstractArray, syms; dim = 1)
+
+Out-of-place version of [`reorder_dimension_by_tunables!`](@ref).
+"""
+function reorder_dimension_by_tunables(
+        sys::AbstractSystem, arr::AbstractArray, syms; dim = 1)
+    buffer = similar(arr)
+    reorder_dimension_by_tunables!(buffer, sys, arr, syms; dim)
+    return buffer
+end

src/variables.jl

@@ -381,7 +381,12 @@ Create a tunable parameter by
 @parameters u [tunable=true]
 ```
 
-See also [`getbounds`](@ref), [`istunable`](@ref)
+For systems created with `split = true` (the default) and `default = true` passed to this function, the order
+of parameters returned is the order in which they are stored in the tunables portion of `MTKParameters`. Note
+that array variables will not be scalarized. To obtain the flattened representation of the tunables portion,
+call `Symbolics.scalarize(tunable_parameters(sys))` and concatenate the resulting arrays.
+
+See also [`getbounds`](@ref), [`istunable`](@ref), [`MTKParameters`](@ref), [`complete`](@ref)
 """
 function tunable_parameters(sys, p = parameters(sys); default = true)
     filter(x -> istunable(x, default), p)

test/discrete_system.jl

@@ -42,7 +42,7 @@ for df in [
     # iip
     du = zeros(3)
     u = collect(1:3)
-    p = MTKParameters(syss, parameters(syss) .=> collect(1:5))
+    p = MTKParameters(syss, [c, nsteps, δt, β, γ] .=> collect(1:5))
     df.f(du, u, p, 0)
     @test du ≈ [0.01831563888873422, 0.9816849729159067, 4.999999388195359]
 

test/index_cache.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, SymbolicIndexingInterface
+using ModelingToolkit, SymbolicIndexingInterface, SciMLStructures
 using ModelingToolkit: t_nounits as t
 
 # Ensure indexes of array symbolics are cached appropriately
@@ -43,3 +43,52 @@ ic = ModelingToolkit.get_index_cache(sys)
 @test isequal(ic.symbol_to_variable[:x], x)
 @test isequal(ic.symbol_to_variable[:y], y)
 @test isequal(ic.symbol_to_variable[:z], z)
+
+@testset "tunable_parameters is ordered" begin
+    @parameters p q[1:3] r[1:2, 1:2] s [tunable = false]
+    @named sys = ODESystem(Equation[], t, [], [p, q, r, s])
+    sys = complete(sys)
+    @test all(splat(isequal), zip(tunable_parameters(sys), parameters(sys)[1:3]))
+
+    offset = 1
+    for par in tunable_parameters(sys)
+        idx = parameter_index(sys, par)
+        @test idx.portion isa SciMLStructures.Tunable
+        if Symbolics.isarraysymbolic(par)
+            @test vec(idx.idx) == offset:(offset + length(par) - 1)
+        else
+            @test idx.idx == offset
+        end
+        offset += length(par)
+    end
+end
+
+@testset "reorder_dimension_by_tunables" begin
+    @parameters p q[1:3] r[1:2, 1:2] s [tunable = false]
+    @named sys = ODESystem(Equation[], t, [], [p, q, r, s])
+    src = ones(8)
+    dst = zeros(8)
+    # system must be complete...
+    @test_throws ArgumentError reorder_dimension_by_tunables!(dst, sys, src, [p, q, r])
+    @test_throws ArgumentError reorder_dimension_by_tunables(sys, src, [p, q, r])
+    sys = complete(sys; split = false)
+    # with split = true...
+    @test_throws ArgumentError reorder_dimension_by_tunables!(dst, sys, src, [p, q, r])
+    @test_throws ArgumentError reorder_dimension_by_tunables(sys, src, [p, q, r])
+    sys = complete(sys)
+    # and the arrays must have matching size
+    @test_throws ArgumentError reorder_dimension_by_tunables!(
+        zeros(2, 4), sys, src, [p, q, r])
+
+    ps = MTKParameters(sys, [p => 1.0, q => 3ones(3), r => 4ones(2, 2), s => 0.0])
+    src = ps.tunable
+    reorder_dimension_by_tunables!(dst, sys, src, [q, r, p])
+    @test dst ≈ vcat(3ones(3), 4ones(4), 1.0)
+    @test reorder_dimension_by_tunables(sys, src, [r, p, q]) ≈ vcat(4ones(4), 1.0, 3ones(3))
+    reorder_dimension_by_tunables!(dst, sys, src, [q[1], r[:, 1], q[2], r[:, 2], q[3], p])
+    @test dst ≈ vcat(3.0, 4ones(2), 3.0, 4ones(2), 3.0, 1.0)
+    src = stack([copy(ps.tunable) for i in 1:5]; dims = 1)
+    dst = zeros(size(src))
+    reorder_dimension_by_tunables!(dst, sys, src, [r, q, p]; dim = 2)
+    @test dst ≈ stack([vcat(4ones(4), 3ones(3), 1.0) for i in 1:5]; dims = 1)
+end

test/labelledarrays.jl

@@ -87,5 +87,6 @@ u0 = @LArray [9998.0, 1.0, 1.0, 1.0] (:S, :I, :R, :C)
 problem = ODEProblem(SIR!, u0, tspan, p)
 sys = complete(modelingtoolkitize(problem))
 
-@test all(isequal.(parameters(sys), getproperty.(@variables(β, η, ω, φ, σ, μ), :val)))
+@test all(any(isequal(x), parameters(sys))
+for x in ModelingToolkit.unwrap.(@variables(β, η, ω, φ, σ, μ)))
 @test all(isequal.(Symbol.(unknowns(sys)), Symbol.(@variables(S(t), I(t), R(t), C(t)))))

test/model_parsing.jl

@@ -534,9 +534,9 @@ end
     @named else_in_sys = InsideTheBlock(flag = 3)
     else_in_sys = complete(else_in_sys)
 
-    @test getname.(parameters(if_in_sys)) == [:if_parameter, :eq]
-    @test getname.(parameters(elseif_in_sys)) == [:elseif_parameter, :eq]
-    @test getname.(parameters(else_in_sys)) == [:else_parameter, :eq]
+    @test sort(getname.(parameters(if_in_sys))) == [:eq, :if_parameter]
+    @test sort(getname.(parameters(elseif_in_sys))) == [:elseif_parameter, :eq]
+    @test sort(getname.(parameters(else_in_sys))) == [:else_parameter, :eq]
 
     @test getdefault(if_in_sys.if_parameter) == 100
     @test getdefault(elseif_in_sys.elseif_parameter) == 101

test/precompile_test.jl

@@ -10,7 +10,7 @@ using ODEPrecompileTest
 
 u = collect(1:3)
 p = ModelingToolkit.MTKParameters(ODEPrecompileTest.f_noeval_good.sys,
-    parameters(ODEPrecompileTest.f_noeval_good.sys) .=> collect(4:6))
+    [:σ, :ρ, :β] .=> collect(4:6))
 
 # These cases do not work, because they get defined in the ModelingToolkit's RGF cache.
 @test parentmodule(typeof(ODEPrecompileTest.f_bad.f.f_iip).parameters[2]) == ModelingToolkit