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

feat: add support for ParameterIndexingProxy

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]

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
@@ -75,7 +82,7 @@ get_obs = getu(oprob, sys.x + sys.z + t + σ)
 @test gety(oprob) == 10.0
 @test get_arr(oprob) == [10.0, 10.0]
 @test get_tuple(oprob) == (10.0, 1.0)
-@test get_obs(oprob) == 39.0
+@test get_obs(oprob) == 22.0
 
 setx! = setu(oprob, x)
 sety! = setu(oprob, :y)
@@ -86,12 +93,8 @@ setx!(oprob, 11.0)
 @test getx(oprob) == 11.0
 sety!(oprob, 12.0)
 @test gety(oprob) == 12.0
-set_arr!(oprob, 10.0)
-@test get_arr(oprob) == [10.0, 10.0]
 set_arr!(oprob, [11.0, 12.0])
 @test get_arr(oprob) == [11.0, 12.0]
-set_tuple!(oprob, 13.0)
-@test get_tuple(oprob) == (13.0, 13.0)
 set_tuple!(oprob, [10.0, 10.0])
 @test get_tuple(oprob) == (10.0, 10.0)
 
@@ -103,22 +106,28 @@ 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
@@ -131,12 +140,14 @@ getx = getu(sprob, x)
 gety = getu(sprob, :y)
 get_arr = getu(sprob, [x, y])
 get_tuple = getu(sprob, (y, z))
-get_obs = getu(sprob, sys.x + sys.z + t + σ)
+# observed doesn't work the same for SDEs
+# uncomment get_obs test below when fixed
+@test_broken get_obs = getu(sprob, sys.x + sys.z + t + σ)
 @test getx(sprob) == 10.0
 @test gety(sprob) == 10.0
 @test get_arr(sprob) == [10.0, 10.0]
 @test get_tuple(sprob) == (10.0, 1.0)
-@test get_obs(sprob) == 39.0
+# @test get_obs(sprob) == 39.0
 
 setx! = setu(sprob, x)
 sety! = setu(sprob, :y)
@@ -147,11 +158,22 @@ setx!(sprob, 11.0)
 @test getx(sprob) == 11.0
 sety!(sprob, 12.0)
 @test gety(sprob) == 12.0
-set_arr!(sprob, 10.0)
-@test get_arr(sprob) == [10.0, 10.0]
 set_arr!(sprob, [11.0, 12.0])
 @test get_arr(sprob) == [11.0, 12.0]
-set_tuple!(sprob, 13.0)
-@test get_tuple(sprob) == (13.0, 13.0)
 set_tuple!(sprob, [10.0, 10.0])
 @test get_tuple(sprob) == (10.0, 10.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]