feat: add support for ParameterIndexingProxy

Author: AayushSabharwal

Project.toml

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

src/integrator_interface.jl

@@ -445,6 +445,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
 SymbolicIndexingInterface.parameter_values(prob::AbstractSciMLProblem) = prob.p
 

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

@@ -119,6 +119,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
@@ -303,3 +312,8 @@ prob = ODEProblem(sys, [], (0, 1.0))
 integrator = init(prob, Tsit5())
 @test integrator[x] isa Vector{Float64}
 @test integrator[@nonamespace sys.x] isa Vector{Float64}
+@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]

test/downstream/problem_interface.jl

@@ -33,15 +33,15 @@ oprob = ODEProblem(sys, u0, tspan, p, jac = true)
 getσ1 = getp(sys, σ)
 getσ2 = getp(sys, sys.σ)
 getσ3 = getp(sys, :σ)
-@test getσ1(oprob) == getσ2(oprob) == getσ3(oprob) == 28.0
+@test getσ1(oprob) == getσ2(oprob) == getσ3(oprob) == oprob.ps[σ] == oprob.ps[sys.σ] == oprob.ps[:σ] == 28.0
 getρ1 = getp(sys, ρ)
 getρ2 = getp(sys, sys.ρ)
 getρ3 = getp(sys, :ρ)
-@test getρ1(oprob) == getρ2(oprob) == getρ3(oprob) == 10.0
+@test getρ1(oprob) == getρ2(oprob) == getρ3(oprob) == oprob.ps[ρ] == oprob.ps[sys.ρ] == oprob.ps[:ρ] == 10.0
 getβ1 = getp(sys, β)
 getβ2 = getp(sys, sys.β)
 getβ3 = getp(sys, :β)
-@test getβ1(oprob) == getβ2(oprob) == getβ3(oprob) == 8 / 3
+@test getβ1(oprob) == getβ2(oprob) == getβ3(oprob) == oprob.ps[β] == oprob.ps[sys.β] == oprob.ps[:β] == 8 / 3
 
 @test oprob[x] == oprob[sys.x] == oprob[:x] == 1.0
 @test oprob[y] == oprob[sys.y] == oprob[:y] == 0.0
@@ -51,13 +51,20 @@ getβ3 = getp(sys, :β)
 
 setσ = setp(sys, σ)
 setσ(oprob, 10.0)
-@test getσ1(oprob) == getσ2(oprob) == getσ3(oprob) == 10.0
+@test getσ1(oprob) == getσ2(oprob) == getσ3(oprob) == oprob.ps[σ] == oprob.ps[sys.σ] == oprob.ps[:σ] == 10.0
 setρ = setp(sys, sys.ρ)
 setρ(oprob, 20.0)
-@test getρ1(oprob) == getρ2(oprob) == getρ3(oprob) == 20.0
+@test getρ1(oprob) == getρ2(oprob) == getρ3(oprob) == oprob.ps[ρ] == oprob.ps[sys.ρ] == oprob.ps[:ρ] == 20.0
 setβ = setp(sys, :β)
 setβ(oprob, 30.0)
-@test getβ1(oprob) == getβ2(oprob) == getβ3(oprob) == 30.0
+@test getβ1(oprob) == getβ2(oprob) == getβ3(oprob) == oprob.ps[β] == oprob.ps[sys.β] == oprob.ps[:β] == 30.0
+
+oprob.ps[σ] = 11.0
+@test getσ1(oprob) == getσ2(oprob) == getσ3(oprob) == oprob.ps[σ] == oprob.ps[sys.σ] == oprob.ps[:σ] == 11.0
+oprob.ps[sys.ρ] = 21.0
+@test getρ1(oprob) == getρ2(oprob) == getρ3(oprob) == oprob.ps[ρ] == oprob.ps[sys.ρ] == oprob.ps[:ρ] == 21.0
+oprob.ps[:β] = 31.0
+@test getβ1(oprob) == getβ2(oprob) == getβ3(oprob) == oprob.ps[β] == oprob.ps[sys.β] == oprob.ps[:β] == 31.0
 
 oprob[x] = 10.0
 @test oprob[x] == oprob[sys.x] == oprob[:x] == 10.0
@@ -74,26 +81,47 @@ noiseeqs = [0.1 * x,
 sprob = SDEProblem(noise_sys, u0, (0.0, 100.0), p)
 u0
 
-@test getσ1(sprob) == getσ2(sprob) == getσ3(sprob) == 28.0
-@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == 10.0
-@test getβ1(sprob) == getβ2(sprob) == getβ3(sprob) == 8 / 3
+@test getσ1(sprob) == getσ2(sprob) == getσ3(sprob) == sprob.ps[σ] == sprob.ps[sys.σ] == sprob.ps[:σ] == 28.0
+@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == sprob.ps[ρ] == sprob.ps[sys.ρ] == sprob.ps[:ρ] == 10.0
+@test getβ1(sprob) == getβ2(sprob) == getβ3(sprob) == sprob.ps[β] == sprob.ps[sys.β] == sprob.ps[:β] == 8 / 3
 
 @test sprob[x] == sprob[noise_sys.x] == sprob[:x] == 1.0
 @test sprob[y] == sprob[noise_sys.y] == sprob[:y] == 0.0
 @test sprob[z] == sprob[noise_sys.z] == sprob[:z] == 0.0
 
 setσ(sprob, 10.0)
-@test getσ1(sprob) == getσ2(sprob) == getσ3(sprob) == 10.0
+@test getσ1(sprob) == getσ2(sprob) == getσ3(sprob) == sprob.ps[σ] == sprob.ps[sys.σ] == sprob.ps[:σ] == 10.0
 setρ(sprob, 20.0)
-@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == 20.0
+@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == sprob.ps[ρ] == sprob.ps[sys.ρ] == sprob.ps[:ρ] == 20.0
 setp(noise_sys, noise_sys.ρ)(sprob, 25.0)
-@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == 25.0
+@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == sprob.ps[ρ] == sprob.ps[sys.ρ] == sprob.ps[:ρ] == 25.0
 setβ(sprob, 30.0)
-@test getβ1(sprob) == getβ2(sprob) == getβ3(sprob) == 30.0
+@test getβ1(sprob) == getβ2(sprob) == getβ3(sprob) == sprob.ps[β] == sprob.ps[sys.β] == sprob.ps[:β] == 30.0
+sprob.ps[σ] = 11.0
+@test getσ1(sprob) == getσ2(sprob) == getσ3(sprob) == sprob.ps[σ] == sprob.ps[sys.σ] == sprob.ps[:σ] == 11.0
+sprob.ps[sys.ρ] = 21.0
+@test getρ1(sprob) == getρ2(sprob) == getρ3(sprob) == sprob.ps[ρ] == sprob.ps[sys.ρ] == sprob.ps[:ρ] == 21.0
+sprob.ps[:β] = 31.0
+@test getβ1(sprob) == getβ2(sprob) == getβ3(sprob) == sprob.ps[β] == sprob.ps[sys.β] == sprob.ps[:β] == 31.0
 
 sprob[x] = 10.0
 @test sprob[x] == sprob[noise_sys.x] == sprob[:x] == 10.0
 sprob[noise_sys.y] = 10.0
 @test sprob[y] == sprob[noise_sys.y] == sprob[:y] == 10.0
 sprob[:z] = 1.0
 @test sprob[z] == sprob[noise_sys.z] == sprob[:z] == 1.0
+
+using LinearAlgebra
+@variables t
+sts = @variables x(t)[1:3]=[1, 2, 3.0] y(t)=1.0
+ps = @parameters p[1:3] = [1, 2, 3]
+D = Differential(t)
+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))
+@test getp(sys, p)(prob) == prob.ps[p] == [1, 2, 3]
+setp(sys, p)(prob, [4, 5, 6])
+@test getp(sys, p)(prob) == prob.ps[p] == [4, 5, 6]
+prob.ps[p] = [7, 8, 9]
+@test getp(sys, p)(prob) == prob.ps[p] == [7, 8, 9]

test/downstream/symbol_indexing.jl

@@ -72,8 +72,10 @@ sol = solve(prob, Rodas4())
 @test sol[α, 3] isa Float64
 @test length(sol[α, 5:10]) == 6
 @test getp(prob, γ)(sol) isa Real
-@test getp(prob, γ)(sol) == getp(prob, :γ)(sol) == 2.0
+@test sol.ps[γ] isa Real
+@test getp(prob, γ)(sol) == getp(prob, :γ)(sol) == sol.ps[γ] == sol.ps[:γ] == 2.0
 @test getp(prob, (lorenz1.σ, lorenz1.ρ))(sol) isa Tuple
+@test sol.ps[(lorenz1.σ, lorenz1.ρ)] isa Tuple
 
 @test sol[[lorenz1.x, lorenz2.x]] isa Vector{Vector{Float64}}
 @test length(sol[[lorenz1.x, lorenz2.x]]) == length(sol)
@@ -196,6 +198,7 @@ prob = ODEProblem(sys, [], (0, 1.0))
 sol = solve(prob, Tsit5())
 @test sol[x] isa Vector{<:Vector}
 @test sol[@nonamespace sys.x] isa Vector{<:Vector}
+@test sol.ps[p] == [1, 2, 3]
 
 # accessing parameters
 @variables t x(t)
@@ -225,4 +228,6 @@ sol = solve(prob, Tsit5())
     @test sol[y]≈1 atol=1e-3
     @test getp(sys, a)(sol) ≈ 1
     @test getp(sys, b)(sol) ≈ 100
+    @test sol.ps[a] ≈ 1
+    @test sol.ps[b] ≈ 100
 end