Update parameter_buffer.jl

Author: MasonProtter

src/systems/parameter_buffer.jl

@@ -362,37 +362,50 @@ function Base.setindex!(p::MTKParameters, val, i)
 end
 
 function Base.getindex(p::MTKParameters, pind::ParameterIndex)
-    (;portion, idx) = pind, (i,j) = idx
+    (;portion, idx) = pind
+    if length(idx) > 2
+        i, j, k... = idx
+        indexer = (v) -> v[i][j][k...] 
+    else
+        i, j = idx
+        indexer = (v) -> v[i][j]
+    end
     if portion isa SciMLStructures.Tunable
-        p.tunable[i][j]
+        indexer(p.tunable)
     elseif portion isa SciMLStructures.Discrete
-        p.discrete[i][j]
+        indexer(p.discrete)
     elseif portion isa SciMLStructures.Constants
-        p.constant[i][j]
+        indexer(p.constant)
     elseif portion === DEPENDENT_PORTION
-        p.dependent[i][j]
+        indexer(p.dependent)
     elseif portion === NONNUMERIC_PORTION
-        p.nonnumeric[i][j]
+        indexer(p.nonnumeric)
     else
-        error("Unhandled portion $portion")
+        error("Unhandled portion ", portion)
     end
 end
 
 function Base.setindex!(p::MTKParameters, val, pind::ParameterIndex)
     (;portion, idx) = pind
-    (i,j) = idx
+    if length(idx) > 2
+        i, j, k... = idx
+        setindexer = (v) -> v[i][j][k...] = val
+    else
+        i, j = idx
+        setindexer = (v) -> v[i][j] = val
+    end
     if portion isa SciMLStructures.Tunable
-        p.tunable[i][j] = val
+        setindexer(p.tunable)
     elseif portion isa SciMLStructures.Discrete
-        p.discrete[i][j] = val
+        setindexer(p.discrete)
     elseif portion isa SciMLStructures.Constants
-        p.constant[i][j] = val
+        setindexer(p.constant)
     elseif portion === DEPENDENT_PORTION
-        p.dependent[i][j] = val
+        setindexer(p.dependent)
     elseif portion === NONNUMERIC_PORTION
-        p.nonnumeric[i][j] = val
+        setindexer(p.nonnumeric)
     else
-        error("Unhandled portion $portion")
+        error("Unhandled portion", portion)
     end
 end