From f4893f29986c4ea0aa61d10b4cf71f24da9aa36a Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Tue, 2 Jan 2024 13:41:42 +0530 Subject: [PATCH 1/4] fix: JumpProblem SII methods, parameter indexing with tuples --- src/problems/problem_interface.jl | 4 +++- test/downstream/problem_interface.jl | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/problems/problem_interface.jl b/src/problems/problem_interface.jl index 32f912055..5ecbee4c3 100644 --- a/src/problems/problem_interface.jl +++ b/src/problems/problem_interface.jl @@ -1,5 +1,7 @@ SymbolicIndexingInterface.symbolic_container(prob::AbstractSciMLProblem) = prob.f +SymbolicIndexingInterface.symbolic_container(prob::AbstractJumpProblem) = prob.prob SymbolicIndexingInterface.parameter_values(prob::AbstractSciMLProblem) = prob.p +SymbolicIndexingInterface.parameter_values(prob::AbstractJumpProblem) = parameter_values(prob.prob) Base.@propagate_inbounds function Base.getindex(prob::AbstractSciMLProblem, ::SymbolicIndexingInterface.SolvedVariables) return getindex(prob, variable_symbols(prob)) @@ -30,7 +32,7 @@ Base.@propagate_inbounds function Base.getindex(prob::AbstractSciMLProblem, sym) elseif symbolic_type(sym) == ArraySymbolic() return map(s -> prob[s], sym) else - sym isa AbstractArray || error("Invalid indexing of problem") + sym isa Union{<:AbstractArray, <:Tuple} || error("Invalid indexing of problem") return map(s -> prob[s], sym) end end diff --git a/test/downstream/problem_interface.jl b/test/downstream/problem_interface.jl index 67c10c8a0..d7a736c67 100644 --- a/test/downstream/problem_interface.jl +++ b/test/downstream/problem_interface.jl @@ -46,6 +46,8 @@ getβ3 = getp(sys, :β) @test oprob[x] == oprob[sys.x] == oprob[:x] == 1.0 @test oprob[y] == oprob[sys.y] == oprob[:y] == 0.0 @test oprob[z] == oprob[sys.z] == oprob[:z] == 0.0 +@test oprob[[x, y]] == oprob[[sys.x, sys.y]] == oprob[[:x, :y]] == [1.0, 0.0] +@test oprob[(x, y)] == oprob[(sys.x, sys.y)] == oprob[(:x, :y)] == (1.0, 0.0) @test oprob[solvedvariables] == oprob[variable_symbols(sys)] @test oprob[allvariables] == oprob[all_variable_symbols(sys)] From 98f357b4f5f5ffbc93c5b14d6d9cd6277fc9f0b7 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Tue, 2 Jan 2024 14:37:59 +0530 Subject: [PATCH 2/4] fix: u0 and p generation in remake --- src/remake.jl | 129 +++++++++++++++++++------------------------------- 1 file changed, 48 insertions(+), 81 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index a03c11ec5..7ca8116a4 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -19,13 +19,7 @@ for T in [ @eval remaker_of(::$T) = $T end -""" - remake(thing; ) - -Re-construct `thing` with new field values specified by the keyword -arguments. -""" -function remake(thing; kwargs...) +function _remake_internal(thing; kwargs...) T = remaker_of(thing) if :kwargs ∈ fieldnames(typeof(thing)) if :kwargs ∉ keys(kwargs) @@ -38,6 +32,21 @@ function remake(thing; kwargs...) end end +""" + remake(thing; ) + +Re-construct `thing` with new field values specified by the keyword +arguments. +""" +function remake(thing; kwargs...) + _remake_internal(thing; kwargs...) +end + +function remake(prob::DiscreteProblem; u0 = missing, p = missing, kwargs...) + p, u0 = _remake_get_p_u0(prob; p, u0) + _remake_internal(prob; p, u0, kwargs...) +end + function isrecompile(prob::ODEProblem{iip}) where {iip} (prob.f isa ODEFunction) ? !isfunctionwrapper(prob.f.f) : true end @@ -59,25 +68,7 @@ function remake(prob::ODEProblem; f = missing, tspan = prob.tspan end - if p === missing && u0 === missing - p, u0 = prob.p, prob.u0 - else # at least one of them has a value - if p === missing - p = prob.p - end - if u0 === missing - u0 = prob.u0 - end - if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map - hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) || - throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * - " Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order.")) - hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) || - throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * - " Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order.")) - p, u0 = process_p_u0_symbolic(prob, p, u0) - end - end + p, u0 = _remake_get_p_u0(prob; p, u0) iip = isinplace(prob) @@ -132,15 +123,11 @@ function remake(prob::BVProblem; f = missing, bc = missing, u0 = missing, tspan tspan = prob.tspan end - if p === missing && u0 === missing - p, u0 = prob.p, prob.u0 - else # at least one of them has a value - if p === missing - p = prob.p - end - if u0 === missing - u0 = prob.u0 - end + if p === missing + p = prob.p + end + if u0 === missing + u0 = prob.u0 end iip = isinplace(prob) @@ -202,13 +189,7 @@ function remake(prob::SDEProblem; tspan = prob.tspan end - if p === missing - p = prob.p - end - - if u0 === missing - u0 = prob.u0 - end + p, u0 = _remake_get_p_u0(prob; p, u0) if noise === missing noise = prob.noise @@ -263,26 +244,7 @@ function remake(prob::OptimizationProblem; sense = missing, kwargs = missing, _kwargs...) - if p === missing && u0 === missing - p, u0 = prob.p, prob.u0 - else # at least one of them has a value - if p === missing - p = prob.p - end - if u0 === missing - u0 = prob.u0 - end - if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map - hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) || - throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * - " Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order.")) - hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) || - throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * - " Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order.")) - p, u0 = process_p_u0_symbolic(prob, p, u0) - end - end - + p, u0 = _remake_get_p_u0(prob; p, u0) if f === missing f = prob.f end @@ -332,25 +294,7 @@ function remake(prob::NonlinearProblem; problem_type = missing, kwargs = missing, _kwargs...) - if p === missing && u0 === missing - p, u0 = prob.p, prob.u0 - else # at least one of them has a value - if p === missing - p = prob.p - end - if u0 === missing - u0 = prob.u0 - end - if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map - hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) || - throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * - " Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order.")) - hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) || - throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * - " Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order.")) - p, u0 = process_p_u0_symbolic(prob, p, u0) - end - end + p, u0 = _remake_get_p_u0(prob; p, u0) if f === missing f = prob.f @@ -419,3 +363,26 @@ function remake(thing::AbstractEnsembleProblem; kwargs...) en_kwargs = [k for k in kwargs if first(k) ∈ fieldnames(T)] T(remake(thing.prob; setdiff(kwargs, en_kwargs)...); en_kwargs...) end + +function _remake_get_p_u0(prob; p = missing, u0 = missing) + if p === missing && u0 === missing + p, u0 = prob.p, prob.u0 + else # at least one of them has a value + if p === missing + p = prob.p + end + if u0 === missing + u0 = prob.u0 + end + if (eltype(p) <: Pair && !isempty(p)) || (eltype(u0) <: Pair && !isempty(u0)) # one is a non-empty symbolic map + hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :ps) || + throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * + " Please use `remake` with the `p` keyword argument as a vector of values, paying attention to parameter order.")) + hasproperty(prob.f, :sys) && hasfield(typeof(prob.f.sys), :states) || + throw(ArgumentError("This problem does not support symbolic maps with `remake`, i.e. it does not have a symbolic origin." * + " Please use `remake` with the `u0` keyword argument as a vector of values, paying attention to state order.")) + p, u0 = process_p_u0_symbolic(prob, p, u0) + end + end + return p, u0 +end From 9ea5b75249f2bc93101bbeb23cd91dcf2e304da7 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Tue, 2 Jan 2024 14:38:17 +0530 Subject: [PATCH 3/4] fix: interpolation using symbolic idxs for RODESolution --- src/solutions/rode_solutions.jl | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/solutions/rode_solutions.jl b/src/solutions/rode_solutions.jl index b6ec13358..14ca2f200 100644 --- a/src/solutions/rode_solutions.jl +++ b/src/solutions/rode_solutions.jl @@ -63,6 +63,22 @@ TruncatedStacktraces.@truncate_stacktrace RODESolution 1 2 function (sol::RODESolution)(t, ::Type{deriv} = Val{0}; idxs = nothing, continuity = :left) where {deriv} + if idxs !== nothing + if !(idxs isa Union{<:AbstractArray, <:Tuple}) + idxs = [idxs] + end + idxs = map(idxs) do idx + if symbolic_type(idx) === NotSymbolic() + return idx + elseif symbolic_type(idx) === ScalarSymbolic() + return variable_index(sol, idx) + else + return variable_index.((sol,), collect(idx)) + end + end + any(i === nothing for i in idxs) && error("All idxs must be variables") + idxs = reduce(vcat, idxs) + end sol.interp(t, idxs, deriv, sol.prob.p, continuity) end function (sol::RODESolution)(v, t, ::Type{deriv} = Val{0}; idxs = nothing, From 7ca2602079fe0fbfcbf3ba2c6b99baf6d25f0a90 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Tue, 2 Jan 2024 15:49:52 +0530 Subject: [PATCH 4/4] test: add remake tests refactor: add JumpProblem to remake tests --- Project.toml | 7 +- test/remake.jl | 163 +++++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 3 files changed, 171 insertions(+), 2 deletions(-) create mode 100644 test/remake.jl diff --git a/Project.toml b/Project.toml index ceb1fabb2..cecfe3bee 100644 --- a/Project.toml +++ b/Project.toml @@ -62,6 +62,7 @@ EnumX = "1" FillArrays = "1.9" FunctionWrappersWrappers = "0.1.3" IteratorInterfaceExtensions = "^1" +JumpProcesses = "9.10.1" LinearAlgebra = "1.9" Logging = "1.9" Markdown = "1.9" @@ -81,7 +82,7 @@ SciMLOperators = "0.3.7" StaticArrays = "1.7" StaticArraysCore = "1.4" Statistics = "1.9" -SymbolicIndexingInterface = "0.3" +SymbolicIndexingInterface = "0.3.2" Tables = "1.11" TruncatedStacktraces = "1.4" Zygote = "0.6.67" @@ -92,6 +93,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb" +JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b" @@ -102,8 +104,9 @@ RCall = "6f49c342-dc21-5d91-9882-a32aef131414" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" +SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Pkg", "PyCall", "PythonCall", "SafeTestsets", "Test", "StaticArrays", "StochasticDiffEq", "Aqua", "Zygote", "PartialFunctions", "DataFrames", "ModelingToolkit", "OrdinaryDiffEq"] +test = ["Pkg", "PyCall", "PythonCall", "SafeTestsets", "Test", "StaticArrays", "StochasticDiffEq", "Aqua", "Zygote", "PartialFunctions", "DataFrames", "ModelingToolkit", "OrdinaryDiffEq", "JumpProcesses", "SymbolicIndexingInterface"] diff --git a/test/remake.jl b/test/remake.jl new file mode 100644 index 000000000..f47ea84a9 --- /dev/null +++ b/test/remake.jl @@ -0,0 +1,163 @@ +using ModelingToolkit, SymbolicIndexingInterface +using JumpProcesses + +@parameters σ ρ β +@variables t x(t) y(t) z(t) +D = Differential(t) + +eqs = [D(D(x)) ~ σ * (y - x), + D(y) ~ x * (ρ - z) - y, + D(z) ~ x * y - β * z] + +@named sys = ODESystem(eqs) +sys = structural_simplify(sys) +u0 = [D(x) => 2.0, + x => 1.0, + y => 0.0, + z => 0.0] + +p = [σ => 28.0, + ρ => 10.0, + β => 8 / 3] + +tspan = (0.0, 100.0) +oprob = ODEProblem(sys, u0, tspan, p, jac = true) + +@inferred typeof(oprob) remake(oprob; u0 = [x => 2.0], p = [σ => 29.0]) +oprob2 = remake( + oprob; + u0 = [x => 2.0, sys.y => 1.2, :z => 1.0], + p = [σ => 29.0, sys.ρ => 11.0, :β => 3.0] +) +@test oprob2.u0 isa Vector{<:Number} +@test oprob2.p isa Vector{<:Number} +@test oprob2[x] == oprob2[sys.x] == oprob2[:x] == 2.0 +@test oprob2[y] == oprob2[sys.y] == oprob2[:y] == 1.2 +@test oprob2[z] == oprob2[sys.z] == oprob2[:z] == 1.0 +@test getp(sys, σ)(oprob2) == 29.0 +@test getp(sys, sys.ρ)(oprob2) == 11.0 +@test getp(sys, :β)(oprob2) == 3.0 + +oprob3 = remake(oprob; u0 = [x => 3.0], p = [σ => 30.0]) # partial update +@test oprob3[x] == 3.0 +@test getp(sys, σ)(oprob3) == 30.0 + +# SDEProblem. +noiseeqs = [0.1 * x, + 0.1 * y, + 0.1 * z] +@named noise_sys = SDESystem(sys, noiseeqs) +sprob = SDEProblem(noise_sys, u0, (0.0, 100.0), p) + +@inferred typeof(sprob) remake(sprob; u0 = [x => 2.0], p = [σ => 29.0]) +sprob2 = remake( + sprob; + u0 = [x => 2.0, sys.y => 1.2, :z => 1.0], + p = [σ => 29.0, sys.ρ => 11.0, :β => 3.0] +) +@test sprob2.u0 isa Vector{<:Number} +@test sprob2.p isa Vector{<:Number} +@test sprob2[x] == sprob2[sys.x] == sprob2[:x] == 2.0 +@test sprob2[y] == sprob2[sys.y] == sprob2[:y] == 1.2 +@test sprob2[z] == sprob2[sys.z] == sprob2[:z] == 1.0 +@test getp(sys, σ)(sprob2) == 29.0 +@test getp(sys, sys.ρ)(sprob2) == 11.0 +@test getp(sys, :β)(sprob2) == 3.0 + +sprob3 = remake(sprob; u0 = [x => 3.0], p = [σ => 30.0]) # partial update +@test sprob3[x] == 3.0 +@test getp(sys, σ)(sprob3) == 30.0 + +# DiscreteProblem +@named de = DiscreteSystem( + [D(x) ~ σ*(y-x), + D(y) ~ x*(ρ-z)-y, + D(z) ~ x*y - β*z], + t, + [x, y, z], + [σ, ρ, β], +) +dprob = DiscreteProblem(de, u0, tspan, p) + +@inferred typeof(dprob) remake(dprob; u0 = [x => 2.0], p = [σ => 29.0]) +dprob2 = remake( + dprob; + u0 = [x => 2.0, sys.y => 1.2, :z => 1.0], + p = [σ => 29.0, sys.ρ => 11.0, :β => 3.0] +) +@test dprob2.u0 isa Vector{<:Number} +@test dprob2.p isa Vector{<:Number} +@test dprob2[x] == dprob2[sys.x] == dprob2[:x] == 2.0 +@test dprob2[y] == dprob2[sys.y] == dprob2[:y] == 1.2 +@test dprob2[z] == dprob2[sys.z] == dprob2[:z] == 1.0 +@test getp(de, σ)(dprob2) == 29.0 +@test getp(de, sys.ρ)(dprob2) == 11.0 +@test getp(de, :β)(dprob2) == 3.0 + +dprob3 = remake(dprob; u0 = [x => 3.0], p = [σ => 30.0]) # partial update +@test dprob3[x] == 3.0 +@test getp(de, σ)(dprob3) == 30.0 + +# NonlinearProblem +@named ns = NonlinearSystem( + [0 ~ σ*(y-x), + 0 ~ x*(ρ-z)-y, + 0 ~ x*y - β*z], + [x,y,z], + [σ,ρ,β] +) +nlprob = NonlinearProblem(ns, u0, p) + +@inferred typeof(nlprob) remake(nlprob; u0 = [x => 2.0], p = [σ => 29.0]) +nlprob2 = remake( + nlprob; + u0 = [x => 2.0, sys.y => 1.2, :z => 1.0], + p = [σ => 29.0, sys.ρ => 11.0, :β => 3.0] +) +@test nlprob2.u0 isa Vector{<:Number} +@test nlprob2.p isa Vector{<:Number} +@test nlprob2[x] == nlprob2[sys.x] == nlprob2[:x] == 2.0 +@test nlprob2[y] == nlprob2[sys.y] == nlprob2[:y] == 1.2 +@test nlprob2[z] == nlprob2[sys.z] == nlprob2[:z] == 1.0 +@test getp(ns, σ)(nlprob2) == 29.0 +@test getp(ns, sys.ρ)(nlprob2) == 11.0 +@test getp(ns, :β)(nlprob2) == 3.0 + +nlprob3 = remake(nlprob; u0 = [x => 3.0], p = [σ => 30.0]) # partial update +@test nlprob3[x] == 3.0 +@test getp(ns, σ)(nlprob3) == 30.0 + +@parameters β γ +@variables t S(t) I(t) R(t) +rate₁ = β*S*I +affect₁ = [S ~ S - 1, I ~ I + 1] +rate₂ = γ*I +affect₂ = [I ~ I - 1, R ~ R + 1] +j₁ = ConstantRateJump(rate₁,affect₁) +j₂ = ConstantRateJump(rate₂,affect₂) +j₃ = MassActionJump(2*β+γ, [R => 1], [S => 1, R => -1]) +@named js = JumpSystem([j₁,j₂,j₃], t, [S,I,R], [β,γ]) + +u₀map = [S => 999, I => 1, R => 0] +parammap = [β => 0.1 / 1000, γ => 0.01] +tspan = (0.0, 250.0) +jump_dprob = DiscreteProblem(js, u₀map, tspan, parammap) +jprob = JumpProblem(js, jump_dprob, Direct()) + +@test_broken @inferred typeof(jprob) remake(jprob; u0 = [S => 900], p = [β => 0.2e-3]) +jprob2 = remake( + jprob; + u0 = [S => 900, js.I => 2, :R => 0.1], + p = [β => 0.2 / 1000, js.γ => 11.0] +) +@test jprob2.prob.u0 isa Vector{<:Number} +@test jprob2.prob.p isa Vector{<:Number} +@test jprob2[S] == jprob2[js.S] == jprob2[:S] == 900.0 +@test jprob2[I] == jprob2[js.I] == jprob2[:I] == 2.0 +@test jprob2[R] == jprob2[js.R] == jprob2[:R] == 0.1 +@test getp(js, β)(jprob2) == 0.2 / 1000 +@test getp(js, js.γ)(jprob2) == 11.0 + +jprob3 = remake(jprob; u0 = [S => 901], p = [:β => 0.3 / 1000]) # partial update +@test jprob3[S] == 901 +@test getp(js, β)(jprob3) == 0.3 / 1000 diff --git a/test/runtests.jl b/test/runtests.jl index 0c58fb841..5113c9a5b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -62,6 +62,9 @@ end @time @safetestset "Problem building tests" begin include("problem_building_test.jl") end + @time @safetestset "Remake tests" begin + include("remake.jl") + end end if !is_APPVEYOR && GROUP == "Downstream"