Merge pull request #790 from jClugstor/strip_sol_fix

Fix interpolation for stripped solutions
SciML · Sep 23, 2024 · 1e6720c · 1e6720c
2 parents 4a0c5e8 + d777184
commit 1e6720c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/src/solutions/ode_solutions.jl b/src/solutions/ode_solutions.jl
@@ -619,14 +619,20 @@ function Base.showerror(io::IO, e::LazyInterpolationException)
         " uses lazy interpolation, which is incompatible with `strip_solution`.")
 end
 
-function strip_solution(sol::ODESolution)
+function strip_solution(sol::ODESolution; strip_alg = false)
     if has_lazy_interpolation(sol.alg)
         throw(LazyInterpolationException(nameof(typeof(sol.alg))))
     end
 
     interp = strip_interpolation(sol.interp)
 
     @reset sol.interp = interp
-    @reset sol.prob = nothing
-    return @set sol.alg = nothing
+
+    @reset sol.prob = (; p = nothing)
+
+    if strip_alg
+        @reset sol.alg = nothing
+    end
+
+    return sol
 end
diff --git a/test/downstream/ode_stripping.jl b/test/downstream/ode_stripping.jl
@@ -12,12 +12,12 @@ prob = ODEProblem(lorenz!, u0, tspan)
 # implicit solver so we can test cache stripping worked
 sol = solve(prob, Rosenbrock23())
 
-@test isnothing(SciMLBase.strip_solution(sol).prob)
+stripped_sol = SciMLBase.strip_solution(sol)
 
-@test isnothing(SciMLBase.strip_solution(sol).alg)
+@test isnothing(SciMLBase.strip_solution(sol, strip_alg = true).alg)
 
-@test isnothing(SciMLBase.strip_solution(sol).interp.f)
+@test isnothing(stripped_sol.interp.f)
 
-@test isnothing(SciMLBase.strip_solution(sol).interp.cache.jac_config)
+@test isnothing(stripped_sol.interp.cache.jac_config)
 
-@test isnothing(SciMLBase.strip_solution(sol).interp.cache.grad_config)
+@test isnothing(stripped_sol.interp.cache.grad_config)