save du and k for DAEs when dense==true

SciML · Feb 12, 2024 · b5e2a0d · b5e2a0d
1 parent 18fc6d1
commit b5e2a0d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/src/integrators/integrator_utils.jl b/src/integrators/integrator_utils.jl
@@ -102,6 +102,14 @@ function _savevalues!(integrator, force_save, reduce_size)::Tuple{Bool, Bool}
                             [k[integrator.opts.save_idxs] for k in integrator.k],
                             false)
                     end
+                    if integrator.prob isa DAEProblem
+                        if integrator.opts.save_idxs === nothing
+                            copyat_or_push!(integrator.sol.du, integrator.saveiter, integrator.du)
+                        else
+                            copyat_or_push!(integrator.sol.du, integrator.saveiter,
+                                integrator.du[integrator.opts.save_idxs], false)
+                        end
+                    end
                 end
             end
             if integrator.alg isa OrdinaryDiffEqCompositeAlgorithm
@@ -346,6 +354,7 @@ function handle_callbacks!(integrator)
         discrete_modified, saved_in_cb = DiffEqBase.apply_discrete_callback!(integrator,
                                                                              discrete_callbacks...)
     end
+
     if !saved_in_cb
         savevalues!(integrator)
     end
@@ -482,12 +491,10 @@ function handle_tstop!(integrator)
 end
 
 function reset_fsal!(integrator)
-    # Under these conditions, these algorithms are not FSAL anymore
+    # call this when the conditions for FSAL are no longer true (e.g. u_modified)
     integrator.stats.nf += 1
 
-    if integrator.sol.prob isa DAEProblem
-        DiffEqBase.initialize_dae!(integrator)
-    else
+    if integrator.sol.prob isa ODEProblem
         if integrator.cache isa OrdinaryDiffEqMutableCache ||
            (integrator.cache isa CompositeCache &&
             integrator.cache.caches[1] isa OrdinaryDiffEqMutableCache)
@@ -496,7 +503,7 @@ function reset_fsal!(integrator)
             integrator.fsalfirst = integrator.f(integrator.u, integrator.p, integrator.t)
         end
     end
-    # Do not set false here so it can be checked in the algorithm
+    # Do not set reeval_fsal false here so it can be checked in the algorithm
     # integrator.reeval_fsal = false
 end
 

diff --git a/src/solve.jl b/src/solve.jl
@@ -25,7 +25,7 @@ function DiffEqBase.__init(prob::Union{DiffEqBase.AbstractODEProblem,
     save_end = nothing,
     callback = nothing,
     dense = save_everystep &&
-                !(alg isa Union{DAEAlgorithm, FunctionMap}) &&
+                !(alg isa FunctionMap) &&
                 isempty(saveat),
     calck = (callback !== nothing && callback !== CallbackSet()) ||
                 (dense) || !isempty(saveat), # and no dense output