Merge pull request #935 from AayushSabharwal/as/fix-disc-interp

feat: allow interpolating discretes past the final time point

Author: ChrisRackauckas

src/remake.jl

@@ -1083,6 +1083,10 @@ calling `SymbolicIndexingInterface.symbolic_container`, provided for dispatch. R
 the updated `newu0` and `newp`.
 """
 function late_binding_update_u0_p(prob, root_indp, u0, p, t0, newu0, newp)
+    if hasmethod(symbolic_container, Tuple{typeof(root_indp)}) &&
+       (sc = symbolic_container(root_indp)) !== root_indp
+        return late_binding_update_u0_p(prob, sc, u0, p, t0, newu0, newp)
+    end
     return newu0, newp
 end
 
@@ -1094,10 +1098,6 @@ Calls `late_binding_update_u0_p(prob, root_indp, u0, p, t0, newu0, newp)` after
 """
 function late_binding_update_u0_p(prob, u0, p, t0, newu0, newp)
     root_indp = prob
-    while hasmethod(symbolic_container, Tuple{typeof(root_indp)}) &&
-        (sc = symbolic_container(root_indp)) !== root_indp
-        root_indp = sc
-    end
     return late_binding_update_u0_p(prob, root_indp, u0, p, t0, newu0, newp)
 end
 

src/solutions/ode_solutions.jl

@@ -287,8 +287,7 @@ function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs,
         ps = parameter_values(discs)
         for ts_idx in eachindex(discs)
             partition = discs[ts_idx]
-            interp_val = ConstantInterpolation(partition.t, partition.u)(
-                t, nothing, deriv, nothing, continuity)
+            interp_val = _hold_discrete(partition.u, partition.t, t)
             ps = with_updated_parameter_timeseries_values(sol, ps, ts_idx => interp_val)
         end
     end
@@ -312,8 +311,7 @@ function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs::AbstractVect
         ps = parameter_values(discs)
         for ts_idx in eachindex(discs)
             partition = discs[ts_idx]
-            interp_val = ConstantInterpolation(partition.t, partition.u)(
-                t, nothing, deriv, nothing, continuity)
+            interp_val = _hold_discrete(partition.u, partition.t, t)
             ps = with_updated_parameter_timeseries_values(sol, ps, ts_idx => interp_val)
         end
     end

test/downstream/Project.toml

@@ -34,7 +34,7 @@ DelayDiffEq = "5"
 DiffEqCallbacks = "3, 4"
 ForwardDiff = "0.10"
 JumpProcesses = "9.10"
-ModelingToolkit = "9.56"
+ModelingToolkit = "9.64.1"
 ModelingToolkitStandardLibrary = "2.7"
 NonlinearSolve = "2, 3, 4"
 Optimization = "4"

test/downstream/solution_interface.jl

@@ -215,7 +215,7 @@ end
             @test sol[x] == xvals
             @test is_parameter(sol, p)
             @test parameter_index(sol, p) == parameter_index(sys, p)
-            @test isequal(only(parameter_symbols(sol)), p)
+            @test any(isequal(p), parameter_symbols(sol))
             @test is_independent_variable(sol, t)
 
             tmp = copy(prob.u0)
@@ -341,3 +341,12 @@ end
         @test _ss isa SciMLBase.SavedSubsystem
     end
 end
+
+@testset "Interpolation after final discrete save" begin
+    @variables x(t) y(t)
+    @parameters start
+    @mtkbuild sys=ODESystem([D(x) ~ y, y ~ ifelse(t < start, 1.0, 2.0)], t) additional_passes=[ModelingToolkit.IfLifting]
+    prob = ODEProblem(sys, [x => 0.0], (0.0, 1.0), [start => 0.5])
+    sol = solve(prob)
+    @test sol(0.6, idxs = y) ≈ 2.0
+end