Merge pull request #974 from vyudu/rode

feat: add discretes to RODESolution

Author: ChrisRackauckas

src/remake.jl

@@ -191,7 +191,7 @@ function remake(
                 props = @delete props._func_cache
                 props = @insert props._func_cache = forig._func_cache
             end
-            
+
             args = (args..., f2)
         end
     end

src/scimlfunctions.jl

@@ -3067,7 +3067,7 @@ function ImplicitDiscreteFunction{iip, specialize}(f;
                    f.observed :
                    DEFAULT_OBSERVED,
         resid_prototype = __has_resid_prototype(f) ?
-                          f.resid_prototype : 
+                          f.resid_prototype :
                           nothing,
         sys = __has_sys(f) ? f.sys : nothing,
         initialization_data = __has_initialization_data(f) ? f.initialization_data :
@@ -3083,11 +3083,12 @@ function ImplicitDiscreteFunction{iip, specialize}(f;
             analytic,
             observed,
             sys,
-            resid_prototype, 
+            resid_prototype,
             initialization_data)
     else
         ImplicitDiscreteFunction{
-                                 iip, specialize, typeof(_f), typeof(analytic), typeof(observed), typeof(sys), typeof(resid_prototype),
+            iip, specialize, typeof(_f), typeof(analytic),
+            typeof(observed), typeof(sys), typeof(resid_prototype),
             typeof(initialization_data)}(
             _f, analytic, observed, sys, resid_prototype, initialization_data)
     end
@@ -3097,7 +3098,9 @@ function ImplicitDiscreteFunction{iip}(f; kwargs...) where {iip}
     ImplicitDiscreteFunction{iip, FullSpecialize}(f; kwargs...)
 end
 ImplicitDiscreteFunction{iip}(f::ImplicitDiscreteFunction; kwargs...) where {iip} = f
-function ImplicitDiscreteFunction(f; resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing, kwargs...)
+function ImplicitDiscreteFunction(
+        f; resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing,
+        kwargs...)
     ImplicitDiscreteFunction{isinplace(f, 5), FullSpecialize}(f; resid_prototype, kwargs...)
 end
 ImplicitDiscreteFunction(f::ImplicitDiscreteFunction; kwargs...) = f
@@ -3107,11 +3110,13 @@ function unwrapped_f(f::ImplicitDiscreteFunction, newf = unwrapped_f(f.f))
 
     if specialize === NoSpecialize
         ImplicitDiscreteFunction{isinplace(f, 5), specialize, Any, Any, Any,
-            Any, Any, Any}(newf, f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
+            Any, Any, Any}(
+            newf, f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
     else
         ImplicitDiscreteFunction{isinplace(f, 5), specialize, typeof(newf),
             typeof(f.analytic),
-            typeof(f.observed), typeof(f.sys), typeof(resid_prototype), typeof(f.initialization_data)}(newf,
+            typeof(f.observed), typeof(f.sys), typeof(resid_prototype), typeof(f.initialization_data)}(
+            newf,
             f.analytic, f.observed, f.sys, f.resid_prototype, f.initialization_data)
     end
 end

src/solutions/rode_solutions.jl

@@ -32,14 +32,15 @@ https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/
   exited due to an error. For more details, see
   [the return code documentation](https://docs.sciml.ai/SciMLBase/stable/interfaces/Solutions/#retcodes).
 """
-struct RODESolution{T, N, uType, uType2, DType, tType, randType, P, A, IType, S,
+struct RODESolution{T, N, uType, uType2, DType, tType, randType, discType, P, A, IType, S,
     AC <: Union{Nothing, Vector{Int}}, V} <:
        AbstractRODESolution{T, N, uType}
     u::uType
     u_analytic::uType2
     errors::DType
     t::tType
     W::randType
+    discretes::discType
     prob::P
     alg::A
     interp::IType
@@ -63,9 +64,10 @@ function ConstructionBase.setproperties(sol::RODESolution, patch::NamedTuple)
     patch = merge(getproperties(sol), patch)
     return RODESolution{
         T, N, typeof(patch.u), typeof(patch.u_analytic), typeof(patch.errors),
-        typeof(patch.t), typeof(patch.W), typeof(patch.prob), typeof(patch.alg), typeof(patch.interp),
+        typeof(patch.t), typeof(patch.W), typeof(patch.discretes),
+        typeof(patch.prob), typeof(patch.alg), typeof(patch.interp),
         typeof(patch.stats), typeof(patch.alg_choice), typeof(patch.saved_subsystem)}(
-        patch.u, patch.u_analytic, patch.errors, patch.t, patch.W,
+        patch.u, patch.u_analytic, patch.errors, patch.t, patch.W, patch.discretes,
         patch.prob, patch.alg, patch.interp, patch.dense, patch.tslocation, patch.stats,
         patch.alg_choice, patch.retcode, patch.seed, patch.saved_subsystem)
 end
@@ -120,16 +122,28 @@ function build_solution(prob::Union{AbstractRODEProblem, AbstractSDDEProblem},
         Base.depwarn(msg, :build_solution)
     end
 
+    ps = parameter_values(prob)
+    if has_sys(prob.f)
+        sswf = if saved_subsystem === nothing
+            prob.f.sys
+        else
+            SavedSubsystemWithFallback(saved_subsystem, prob.f.sys)
+        end
+        discretes = create_parameter_timeseries_collection(sswf, ps, prob.tspan)
+    else
+        discretes = nothing
+    end
     if has_analytic(f)
         u_analytic = Vector{typeof(prob.u0)}()
         errors = Dict{Symbol, real(eltype(prob.u0))}()
         sol = RODESolution{T, N, typeof(u), typeof(u_analytic), typeof(errors), typeof(t),
-            typeof(W),
+            typeof(W), typeof(discretes),
             typeof(prob), typeof(alg), typeof(interp), typeof(stats),
             typeof(alg_choice), typeof(saved_subsystem)}(u,
             u_analytic,
             errors,
             t, W,
+            discretes,
             prob,
             alg,
             interp,
@@ -149,15 +163,37 @@ function build_solution(prob::Union{AbstractRODEProblem, AbstractSDDEProblem},
         return sol
     else
         return RODESolution{T, N, typeof(u), Nothing, Nothing, typeof(t),
-            typeof(W), typeof(prob), typeof(alg), typeof(interp),
+            typeof(W), typeof(discretes), typeof(prob), typeof(alg), typeof(interp),
             typeof(stats), typeof(alg_choice), typeof(saved_subsystem)}(
-            u, nothing, nothing, t, W,
+            u, nothing, nothing, t, W, discretes,
             prob, alg, interp,
             dense, 0, stats,
             alg_choice, retcode, seed, saved_subsystem)
     end
 end
 
+function save_discretes!(sol::AbstractRODESolution, t, vals, timeseries_idx)
+    RecursiveArrayTools.has_discretes(sol) || return
+    disc = RecursiveArrayTools.get_discretes(sol)
+    _save_discretes_internal!(disc[timeseries_idx], t, vals)
+end
+
+function get_interpolated_discretes(sol::AbstractRODESolution, t, deriv, continuity)
+    is_parameter_timeseries(sol) == Timeseries() || return nothing
+
+    discs::ParameterTimeseriesCollection = RecursiveArrayTools.get_discretes(sol)
+    interp_discs = map(discs) do partition
+        hold_discrete(partition.u, partition.t, t)
+    end
+    return ParameterTimeseriesCollection(interp_discs, parameter_values(discs))
+end
+
+function SymbolicIndexingInterface.is_parameter_timeseries(::Type{S}) where {
+        T1, T2, T3, T4, T5, T6, T7,
+        S <: RODESolution{T1, T2, T3, T4, T5, T6, T7, <:ParameterTimeseriesCollection}}
+    Timeseries()
+end
+
 function calculate_solution_errors!(sol::AbstractRODESolution; fill_uanalytic = true,
         timeseries_errors = true, dense_errors = true)
     if sol.prob.f isa Tuple

test/downstream/comprehensive_indexing.jl

@@ -977,3 +977,19 @@ end
     sol = solve(prob, ImplicitEM())
     @test sol[sym] ≈ sol(sol.t .- sol.ps[delay]; idxs = original)
 end
+
+@testset "RODESolutions save discretes" begin
+    @parameters k(t)
+    @variables A(t)
+    function affect2!(integ, u, p, ctx)
+        integ.ps[p.k] += 1.0
+    end
+    db = 1.0 => (affect2!, [], [k], [k], nothing)
+
+    @named ssys = SDESystem(D(A) ~ k * A, [0.0], t, [A], [k], discrete_events = db)
+    ssys = complete(ssys)
+    prob = SDEProblem(ssys, [A => 1.0], (0.0, 4.0), [k => 1.0])
+    sol = solve(prob, RI5())
+    @test sol[k] isa AbstractVector
+    @test sol[k] == [1.0, 2.0, 3.0, 4.0]
+end