Merge pull request #200 from oxinabox/ox/do_fail_with_grace

Fail with grace with the ILS after initialization failure

Author: ChrisRackauckas

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,7 +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
-            finish!(ils)
+            finish!(ils, check_error(integ))
             caches = nothing
         end,
         # Don't add tstops to the left and right.

test/independentlylinearizedtests.jl

@@ -99,20 +99,20 @@ display(@benchmark sample(ils, many_ts))
     # Muck with the `ilsc` here to explore some failure options
     # Error: timepoints longer than time_matrix
     ilsc.t_offset += 1
-    @test_throws ArgumentError finish!(ils)
+    @test_throws ArgumentError finish!(ils, ReturnCode.Default)
     ilsc.t_offset -= 1
 
     # Error: time matrix row two has N elements, but `us` has N+1
     ilsc.u_offsets[2] += 1
-    @test_throws ArgumentError finish!(ils)
+    @test_throws ArgumentError finish!(ils, ReturnCode.Default)
     ilsc.u_offsets[2] -= 1
 
     # Error: one of our time masks doesn't start with `1`
     ilsc.time_masks[1][1, 1] = 0
-    @test_throws ArgumentError finish!(ils)
+    @test_throws ArgumentError finish!(ils, ReturnCode.Default)
     ilsc.time_masks[1][1, 1] = 1
 
-    finish!(ils)
+    finish!(ils, ReturnCode.Default)
     @test sample(ils, ils.ts) == repeat(1:num_timepoints, 1, num_us)
     @test sample(ils, ils.ts, 1) == 2*repeat(1:num_timepoints, 1, num_us)
 end

test/saving_tests.jl

@@ -220,6 +220,22 @@ if VERSION >= v"1.9" # stack
         end
     end
 
+
+    @testset "fail gracefully" begin
+        f_error2(du, u, p, t) = du .= u ./ t .- 1
+        u0 = [1.0];
+        du0 = [1.0];
+        prob = DAEProblem(f_error2, u0, du0, (0.0, 1.0); differential_vars = [true])
+        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
+        @test isempty(ils.ts)
+        @test isempty(ils.us)
+        @test isempty(ils.time_mask)
+    end
+    
+
     # We do not support 2d states yet.
     #test_linearization(prob_ode_2Dlinear, Tsit5())
 end