From 56890c96294bdc2073c6e3af8b48b4b4c30e1bc8 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@juliahub.com>
Date: Sat, 3 Feb 2024 20:40:29 +0530
Subject: [PATCH] refactor: make getp generic of the parameter container

---
 src/parameter_indexing.jl | 16 ++++++++++------
 src/state_indexing.jl     |  2 +-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/parameter_indexing.jl b/src/parameter_indexing.jl
index e28720b..c6f5ba8 100644
--- a/src/parameter_indexing.jl
+++ b/src/parameter_indexing.jl
@@ -1,13 +1,17 @@
 """
     parameter_values(p)
+    parameter_values(p, i)
 
-Return an indexable collection containing the value of each parameter in `p`.
+Return an indexable collection containing the value of each parameter in `p`. The two-
+argument version of this function returns the parameter value at index `i`.
 
-If this function is called with an `AbstractArray`, it will return the same array.
+If this function is called with an `AbstractArray`, it will return the same array. The
+two-argument version of this function will return `getindex(parameter_values(p), i)`
 """
 function parameter_values end
 
 parameter_values(arr::AbstractArray) = arr
+parameter_values(prob, i) = parameter_values(prob)[i]
 
 """
     set_parameter!(sys, val, idx)
@@ -28,7 +32,7 @@ end
 
 Return a function that takes an array representing the parameter vector or an integrator
 or solution of `sys`, and returns the value of the parameter `p`. Note that `p` can be a
-direct numerical index or a symbolic value, or an array/tuple of the aforementioned.
+direct index or a symbolic value, or an array/tuple of the aforementioned.
 
 Requires that the integrator or solution implement [`parameter_values`](@ref). This function
 typically does not need to be implemented, and has a default implementation relying on
@@ -42,14 +46,14 @@ end
 
 function _getp(sys, ::NotSymbolic, ::NotSymbolic, p)
     return function getter(sol)
-        return parameter_values(sol)[p]
+        return parameter_values(sol, p)
     end
 end
 
 function _getp(sys, ::ScalarSymbolic, ::SymbolicTypeTrait, p)
     idx = parameter_index(sys, p)
     return function getter(sol)
-        return parameter_values(sol)[idx]
+        return parameter_values(sol, idx)
     end
 end
 
@@ -76,7 +80,7 @@ end
 
 Return a function that takes an array representing the parameter vector or an integrator
 or problem of `sys`, and a value, and sets the parameter `p` to that value. Note that `p`
-can be a direct numerical index or a symbolic value.
+can be a direct index or a symbolic value.
 
 Requires that the integrator implement [`parameter_values`](@ref) and the returned
 collection be a mutable reference to the parameter vector in the integrator. In
diff --git a/src/state_indexing.jl b/src/state_indexing.jl
index f2cd7f2..5426770 100644
--- a/src/state_indexing.jl
+++ b/src/state_indexing.jl
@@ -266,7 +266,7 @@ end
 
 Return a function that takes an array representing the state vector or an integrator or
 problem of `sys`, and a value, and sets the the state `sym` to that value. Note that `sym`
-can be a direct numerical index, a symbolic state, or an array/tuple of the aforementioned.
+can be a direct index, a symbolic state, or an array/tuple of the aforementioned.
 
 Requires that the integrator implement [`state_values`](@ref) and the
 returned collection be a mutable reference to the state vector in the integrator/problem. Alternatively, if this is not possible or additional actions need to