Merge pull request SciML#584 from AayushSabharwal/as/parameter-indexing-proxy

feat: add support for ParameterIndexingProxy

Author: ChrisRackauckas

Project.toml

@@ -81,7 +81,7 @@ SciMLOperators = "0.3.7"
 StaticArrays = "1.7"
 StaticArraysCore = "1.4"
 Statistics = "1.9"
-SymbolicIndexingInterface = "0.3.2"
+SymbolicIndexingInterface = "0.3.3"
 Tables = "1.11"
 TruncatedStacktraces = "1.4"
 Zygote = "0.6.67"

src/integrator_interface.jl

@@ -451,6 +451,8 @@ function Base.getproperty(A::DEIntegrator, sym::Symbol)
     if sym === :destats && hasfield(typeof(A), :stats)
         @warn "destats has been deprecated for stats"
         getfield(A, :stats)
+    elseif sym === :ps
+        return ParameterIndexingProxy(A)
     else
         return getfield(A, sym)
     end

src/problems/basic_problems.jl

@@ -502,6 +502,8 @@ function Base.getproperty(prob::IntegralProblem, name::Symbol)
         domain = getfield(prob, :domain)
         lb, ub = domain
         return ub
+    elseif name === :ps
+        return ParameterIndexingProxy(prob)
     end
     return Base.getfield(prob, name)
 end

src/problems/problem_interface.jl

@@ -1,3 +1,10 @@
+Base.@propagate_inbounds function Base.getproperty(prob::AbstractSciMLProblem, sym::Symbol)
+    if sym === :ps
+        return ParameterIndexingProxy(prob)
+    end
+    return getfield(prob, sym)
+end
+
 SymbolicIndexingInterface.symbolic_container(prob::AbstractSciMLProblem) = prob.f
 function SymbolicIndexingInterface.is_observed(A::AbstractSciMLProblem, sym)
     return !is_variable(A, sym) && !is_parameter(A, sym) && !is_independent_variable(A, sym) && symbolic_type(sym) == ScalarSymbolic()
@@ -36,7 +43,7 @@ Base.@propagate_inbounds function Base.getindex(prob::AbstractSciMLProblem, sym)
             error("Invalid indexing of problem: $sym is not a state, parameter, or independent variable")
         end
     elseif symbolic_type(sym) == ArraySymbolic()
-        return map(s -> prob[s], sym)
+        return map(s -> prob[s], collect(sym))
     else
         sym isa AbstractArray || error("Invalid indexing of problem")
         return map(s -> prob[s], sym)

src/solutions/dae_solutions.jl

@@ -47,6 +47,8 @@ Base.@propagate_inbounds function Base.getproperty(x::AbstractDAESolution, s::Sy
     if s === :destats
         Base.depwarn("`sol.destats` is deprecated. Use `sol.stats` instead.", "sol.destats")
         return getfield(x, :stats)
+    elseif s === :ps
+        return ParameterIndexingProxy(x)
     end
     return getfield(x, s)
 end

src/solutions/ode_solutions.jl

@@ -121,6 +121,8 @@ Base.@propagate_inbounds function Base.getproperty(x::AbstractODESolution, s::Sy
     if s === :destats
         Base.depwarn("`sol.destats` is deprecated. Use `sol.stats` instead.", "sol.destats")
         return getfield(x, :stats)
+    elseif s === :ps
+        return ParameterIndexingProxy(x)
     end
     return getfield(x, s)
 end

src/solutions/optimization_solutions.jl

@@ -149,6 +149,8 @@ Base.@propagate_inbounds function Base.getproperty(x::AbstractOptimizationSoluti
         Base.depwarn("`sol.prob` is deprecated. Use getters like `get_p` or `get_syms` on `sol` instead.",
             "sol.prob")
         return getfield(x, :cache)
+    elseif s === :ps
+        return ParameterIndexingProxy(x)
     end
     return getfield(x, s)
 end

src/solutions/rode_solutions.jl

@@ -55,6 +55,8 @@ Base.@propagate_inbounds function Base.getproperty(x::AbstractRODESolution, s::S
     if s === :destats
         Base.depwarn("`sol.destats` is deprecated. Use `sol.stats` instead.", "sol.destats")
         return getfield(x, :stats)
+    elseif s === :ps
+        return ParameterIndexingProxy(x)
     end
     return getfield(x, s)
 end

src/solutions/solution_interface.jl

@@ -29,6 +29,15 @@ end
 
 # SymbolicIndexingInterface.jl
 const AbstractSolution = Union{AbstractTimeseriesSolution,AbstractNoTimeSolution}
+
+Base.@propagate_inbounds function Base.getproperty(A::AbstractSolution, sym::Symbol)
+    if sym === :ps
+        return ParameterIndexingProxy(A)
+    else
+        return getfield(A, sym)
+    end
+end
+
 SymbolicIndexingInterface.symbolic_container(A::AbstractSolution) = A.prob.f
 SymbolicIndexingInterface.parameter_values(A::AbstractSolution) = A.prob.p
 

test/downstream/integrator_indexing.jl

@@ -22,9 +22,13 @@ integrator = init(oprob, Rodas4())
 @test_throws Exception integrator[population_model.a]
 @test_throws Exception integrator[:a]
 @test getp(oprob, a)(integrator) == getp(oprob, population_model.a)(integrator) == getp(oprob, :a)(integrator) == 2.0
+@test integrator.ps[a] == integrator.ps[population_model.a] == integrator.ps[:a] == 2.0
 @test getp(oprob, b)(integrator) == getp(oprob, population_model.b)(integrator) == getp(oprob, :b)(integrator) == 1.0
+@test integrator.ps[b] == integrator.ps[population_model.b] == integrator.ps[:b] == 1.0
 @test getp(oprob, c)(integrator) == getp(oprob, population_model.c)(integrator) == getp(oprob, :c)(integrator) == 1.0
+@test integrator.ps[d] == integrator.ps[population_model.d] == integrator.ps[:d] == 1.0
 @test getp(oprob, d)(integrator) == getp(oprob, population_model.d)(integrator) == getp(oprob, :d)(integrator) == 1.0
+@test integrator.ps[d] == integrator.ps[population_model.d] == integrator.ps[:d] == 1.0
 
 @test integrator[s1] == integrator[population_model.s1] == integrator[:s1] == 2.0
 @test integrator[s2] == integrator[population_model.s2] == integrator[:s2] == 1.0
@@ -42,10 +46,15 @@ step!(integrator, 100.0, true)
 
 setp(oprob, a)(integrator, 10.0)
 @test getp(integrator, a)(integrator) == getp(integrator, population_model.a)(integrator) == getp(integrator, :a)(integrator) == 10.0
+@test integrator.ps[a] == integrator.ps[population_model.a] == integrator.ps[:a] == 10.0
 setp(population_model, population_model.b)(integrator, 20.0)
 @test getp(integrator, b)(integrator) == getp(integrator, population_model.b)(integrator) == getp(integrator, :b)(integrator) == 20.0
+@test integrator.ps[b] == integrator.ps[population_model.b] == integrator.ps[:b] == 20.0
 setp(integrator, c)(integrator, 30.0)
 @test getp(integrator, c)(integrator) == getp(integrator, population_model.c)(integrator) == getp(integrator, :c)(integrator) == 30.0
+@test integrator.ps[c] == integrator.ps[population_model.c] == integrator.ps[:c] == 30.0
+integrator.ps[d] = 40.0
+@test integrator.ps[d] == integrator.ps[population_model.d] == integrator.ps[:d] == 40.0
 
 integrator[s1] = 10.0
 @test integrator[s1] == integrator[population_model.s1] == integrator[:s1] == 10.0
@@ -300,7 +309,7 @@ eqs = [collect(D.(x) .~ x)
     D(y) ~ norm(x) * y - x[1]]
 @named sys = ODESystem(eqs, t, [sts...;], [ps...;])
 prob = ODEProblem(sys, [], (0, 1.0))
-integrator = init(prob, Tsit5())
+integrator = init(prob, Tsit5(), save_everystep = false)
 @test integrator[x] isa Vector{Float64}
 @test integrator[@nonamespace sys.x] isa Vector{Float64}
 
@@ -324,9 +333,12 @@ setx!(integrator, [4.0, 5.0, 6.0])
 @test getx(integrator) == [4.0, 5.0, 6.0]
 sety!(integrator, 3.0)
 @test gety(integrator) == 3.0
-set_arr!(integrator, [1.0, 2.0])
-@test get_arr(integrator) == [[1.0, 1.0, 1.0], 2.0]
 set_arr!(integrator, [[1.0, 2.0, 3.0], 1.0])
 @test get_arr(integrator) == [[1.0, 2.0, 3.0], 1.0]
 set_tuple!(integrator, ([2.0, 4.0, 6.0], 2.0))
 @test get_tuple(integrator) == ([2.0, 4.0, 6.0], 2.0)
+@test getp(sys, p)(integrator) == integrator.ps[p] == [1, 2, 3]
+setp(sys, p)(integrator, [4, 5, 6])
+@test getp(sys, p)(integrator) == integrator.ps[p] == [4, 5, 6]
+integrator.ps[p] = [7, 8, 9]
+@test getp(sys, p)(integrator) == integrator.ps[p] == [7, 8, 9]