save du and k for DAEs when dense==true

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

Hide file tree

Showing 2 changed files with 32 additions and 10 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.alg isa DAEAlgorithm
+                        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
@@ -134,6 +142,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.alg isa DAEAlgorithm
+                    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 +362,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 +499,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 +511,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
@@ -413,10 +413,17 @@ function DiffEqBase.__init(prob::Union{DiffEqBase.AbstractODEProblem,
     differential_vars = prob isa DAEProblem ? prob.differential_vars : get_differential_vars(f, u)
 
     id = InterpolationData(f, timeseries, ts, ks, alg_choice, dense, cache, differential_vars, false)
-    sol = DiffEqBase.build_solution(prob, _alg, ts, timeseries,
-        dense = dense, k = ks, interp = id,
-        alg_choice = alg_choice,
-        calculate_error = false, stats = stats)
+    if _alg isa DAEAlgorithm
+        sol = DiffEqBase.build_solution(prob, _alg, ts, timeseries, Vector{typeof(du)}(undef,0),
+            dense = dense, k = ks, interp = id,
+            alg_choice = alg_choice,
+            calculate_error = false, stats = stats)
+    else
+        sol = DiffEqBase.build_solution(prob, _alg, ts, timeseries,
+            dense = dense, k = ks, interp = id,
+            alg_choice = alg_choice,
+            calculate_error = false, stats = stats)
+    end
 
     if recompile_flag == true
         FType = typeof(f)