Handle failure in finish!

Author: oxinabox

src/independentlylinearizedutils.jl

@@ -199,7 +199,7 @@ num_us(ils::IndependentlyLinearizedSolution) = length(ils.us)
 Base.size(ils::IndependentlyLinearizedSolution) = size(ils.time_mask)
 Base.length(ils::IndependentlyLinearizedSolution) = length(ils.ts)
 
-function finish!(ils::IndependentlyLinearizedSolution)
+function finish!(ils::IndependentlyLinearizedSolution{T, S}, return_code) where {T,S}
     function trim_chunk(chunks::Vector, offset)
         chunks = [chunk for chunk in chunks]
         if eltype(chunks) <: AbstractVector
@@ -216,18 +216,25 @@ function finish!(ils::IndependentlyLinearizedSolution)
     end
 
     ilsc = ils.ilsc::IndependentlyLinearizedSolutionChunks
-    ts = vcat(trim_chunk(ilsc.t_chunks, ilsc.t_offset)...)
-    time_mask = hcat(trim_chunk(ilsc.time_masks, ilsc.t_offset)...)
-    us = [hcat(trim_chunk(ilsc.u_chunks[u_idx], ilsc.u_offsets[u_idx])...)
-          for u_idx in 1:length(ilsc.u_chunks)]
+    if return_code == ReturnCode.InitialFailure
+        # then no (consistent) data to put in, so just put in empty values
+        ts = Vector{T}()
+        us = Vector{Matrix{S}}()
+        time_mask = BitMatrix(undef, 0, 0)
+    else
+        ts = vcat(trim_chunk(ilsc.t_chunks, ilsc.t_offset)...)
+        time_mask = hcat(trim_chunk(ilsc.time_masks, ilsc.t_offset)...)
+        us = [hcat(trim_chunk(ilsc.u_chunks[u_idx], ilsc.u_offsets[u_idx])...)
+            for u_idx in 1:length(ilsc.u_chunks)]
+    end
 
     # Sanity-check lengths
     if length(ts) != size(time_mask, 2)
         throw(ArgumentError("`length(ts)` must equal `size(time_mask, 2)`!"))
     end
 
     # All time masks must start and end with `1`:
-    if !all(@view time_mask[:, 1]) || !all(@view time_mask[:, end])
+    if !isempty(time_mask) && (!all(@view time_mask[:, 1]) || !all(@view time_mask[:, end]))
         throw(ArgumentError("Time mask must start and end with 1s!"))
     end
 

src/saving.jl

@@ -433,10 +433,7 @@ function LinearizingSavingCallback(ils::IndependentlyLinearizedSolution{T,S};
         end,
         # We need to finalize the ils and free our caches
         finalize = (c, u, t, integ) -> begin
-            # Don't run the `finish!` if ils is in an inconsistent state
-            if check_error(integ) != ReturnCode.InitialFailure
-                finish!(ils)
-            end
+            finish!(ils, check_error(integ))
             caches = nothing
         end,
         # Don't add tstops to the left and right.

test/saving_tests.jl

@@ -226,7 +226,7 @@ if VERSION >= v"1.9" # stack
         u0 = [1.0];
         du0 = [1.0];
         prob = DAEProblem(f_error2, u0, du0, (0.0, 1.0); differential_vars = [true])
-        ils = IndependentlyLinearizedSolution(unstable_prob, 0)
+        ils = IndependentlyLinearizedSolution(prob, 0)
         lsc = LinearizingSavingCallback(ils)
         sol = solve(prob, DFBDF(); callback = lsc)  # this would if we were not failing with grace
         @test sol.retcode == ReturnCode.InitialFailure