Merge pull request #2531 from MasonProtter/patch-1

Add `getindex`/`setindex!` methods for `MTKParameters` with `ParameterIndex`

Author: ChrisRackauckas

src/systems/parameter_buffer.jl

@@ -363,6 +363,33 @@ function Base.setindex!(p::MTKParameters, val, i)
     end
 end
 
+function Base.getindex(p::MTKParameters, pind::ParameterIndex)
+    (; portion, idx) = pind
+    i, j, k... = idx
+    if isempty(k)
+        indexer = (v) -> v[i][j]
+    else
+        indexer = (v) -> v[i][j][k...]
+    end
+    if portion isa SciMLStructures.Tunable
+        indexer(p.tunable)
+    elseif portion isa SciMLStructures.Discrete
+        indexer(p.discrete)
+    elseif portion isa SciMLStructures.Constants
+        indexer(p.constant)
+    elseif portion === DEPENDENT_PORTION
+        indexer(p.dependent)
+    elseif portion === NONNUMERIC_PORTION
+        indexer(p.nonnumeric)
+    else
+        error("Unhandled portion ", portion)
+    end
+end
+
+function Base.setindex!(p::MTKParameters, val, pind::ParameterIndex)
+    SymbolicIndexingInterface.set_parameter!(p, val, pind)
+end
+
 function Base.iterate(buf::MTKParameters, state = 1)
     total_len = 0
     total_len += _num_subarrays(buf.tunable)

test/split_parameters.jl

@@ -2,6 +2,8 @@ using ModelingToolkit, Test
 using ModelingToolkitStandardLibrary.Blocks
 using OrdinaryDiffEq
 using ModelingToolkit: t_nounits as t, D_nounits as D
+using ModelingToolkit: MTKParameters, ParameterIndex, DEPENDENT_PORTION, NONNUMERIC_PORTION
+using SciMLStructures: Tunable, Discrete, Constants
 
 x = [1, 2.0, false, [1, 2, 3], Parameter(1.0)]
 
@@ -189,3 +191,25 @@ connections = [[state_feedback.input.u[i] ~ model_outputs[i] for i in 1:4]
                connect(add.output, :u, model.torque.tau)]
 @named closed_loop = ODESystem(connections, t, systems = [model, state_feedback, add, d])
 S = get_sensitivity(closed_loop, :u)
+
+@testset "Indexing MTKParameters with ParameterIndex" begin
+    ps = MTKParameters(([1.0, 2.0], [3, 4]),
+        ([true, false], [[1 2; 3 4]]),
+        ([5, 6],),
+        ([7.0, 8.0],),
+        (["hi", "bye"], [:lie, :die]),
+        nothing,
+        nothing)
+    @test ps[ParameterIndex(Tunable(), (1, 2))] === 2.0
+    @test ps[ParameterIndex(Tunable(), (2, 2))] === 4
+    @test ps[ParameterIndex(Discrete(), (2, 1, 2, 2))] === 4
+    @test ps[ParameterIndex(Discrete(), (2, 1))] == [1 2; 3 4]
+    @test ps[ParameterIndex(Constants(), (1, 1))] === 5
+    @test ps[ParameterIndex(DEPENDENT_PORTION, (1, 1))] === 7.0
+    @test ps[ParameterIndex(NONNUMERIC_PORTION, (2, 2))] === :die
+
+    ps[ParameterIndex(Tunable(), (1, 2))] = 3.0
+    ps[ParameterIndex(Discrete(), (2, 1, 2, 2))] = 5
+    @test ps[ParameterIndex(Tunable(), (1, 2))] === 3.0
+    @test ps[ParameterIndex(Discrete(), (2, 1, 2, 2))] === 5
+end