feat: add reorder_dimension_by_tunables!

Author: AayushSabharwal

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

@@ -496,3 +496,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

test/index_cache.jl

@@ -63,3 +63,32 @@ ic = ModelingToolkit.get_index_cache(sys)
     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