fix default handling

SciML · Aug 18, 2024 · dd4ac16 · dd4ac16
1 parent 9fbeeff
commit dd4ac16
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 21 deletions.
diff --git a/lib/OrdinaryDiffEqCore/src/caches/basic_caches.jl b/lib/OrdinaryDiffEqCore/src/caches/basic_caches.jl
@@ -15,7 +15,7 @@ end
 
 get_fsalfirstlast(cache::CompositeCache) = get_fsalfirstlast(cache.caches[1])
 
-mutable struct DefaultCache{T1, T2, T3, T4, T5, T6, A, F} <: OrdinaryDiffEqCache
+mutable struct DefaultCache{T1, T2, T3, T4, T5, T6, A, F, uType} <: OrdinaryDiffEqCache
     args::A
     choice_function::F
     current::Int
@@ -25,31 +25,15 @@ mutable struct DefaultCache{T1, T2, T3, T4, T5, T6, A, F} <: OrdinaryDiffEqCache
     cache4::T4
     cache5::T5
     cache6::T6
+    u::uType
     function DefaultCache{T1, T2, T3, T4, T5, T6, F}(
             args, choice_function, current) where {T1, T2, T3, T4, T5, T6, F}
         new{T1, T2, T3, T4, T5, T6, typeof(args), F}(args, choice_function, current)
     end
 end
 
-function get_fsalfirstlast(cache::DefaultCache{T1, T2, T3, T4, T5, T6}) where {T1, T2, T3, T4, T5, T6} 
-    if isbitstype(T1)
-        get_fsalfirstlast(cache.cache1)
-    end
-    if isbitstype(T2)
-        get_fsalfirstlast(cache.cache2)
-    end
-    if isbitstype(T3)
-        get_fsalfirstlast(cache.cache3)
-    end
-    if isbitstype(T4)
-        get_fsalfirstlast(cache.cache4)
-    end
-    if isbitstype(T5)
-        get_fsalfirstlast(cache.cache5)
-    end
-    if isbitstype(T6)
-        get_fsalfirstlast(cache.cache6)
-    end
+function get_fsalfirstlast(cache::DefaultCache) 
+    (u,u)
 end
 
 function alg_cache(alg::CompositeAlgorithm, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -77,7 +61,7 @@ function alg_cache(alg::CompositeAlgorithm{CS, Tuple{A1, A2, A3, A4, A5, A6}}, u
     T5 = Base.promote_op(alg_cache, A5, argT...)
     T6 = Base.promote_op(alg_cache, A6, argT...)
     cache = DefaultCache{T1, T2, T3, T4, T5, T6, typeof(alg.choice_function)}(
-        args, alg.choice_function, 1)
+        args, alg.choice_function, 1, u)
     algs = alg.algs
     # If the type is a bitstype we need to initialize it correctly here since isdefined will always return true.
     if isbitstype(T1)

diff --git a/lib/OrdinaryDiffEqCore/src/perform_step/composite_perform_step.jl b/lib/OrdinaryDiffEqCore/src/perform_step/composite_perform_step.jl
@@ -32,24 +32,42 @@ function initialize!(integrator, cache::DefaultCache)
     init_ith_default_cache(cache, algs, cache.current)
     if cache.current == 1
         initialize!(integrator, cache.cache1)
+        fsalfirst, fsallast = get_fsalfirstlast(cache.cache1)
+        integrator.fsalfirst = fsalfirst
+        integrator.fsallast = fsallast
     elseif cache.current == 2
         initialize!(integrator, cache.cache2)
+        fsalfirst, fsallast = get_fsalfirstlast(cache.cache2)
+        integrator.fsalfirst = fsalfirst
+        integrator.fsallast = fsallast
         # the controller was initialized by default for algs[1]
         reset_alg_dependent_opts!(integrator.opts.controller, algs[1], algs[2])
     elseif cache.current == 3
         initialize!(integrator, cache.cache3)
+        fsalfirst, fsallast = get_fsalfirstlast(cache.cache3)
+        integrator.fsalfirst = fsalfirst
+        integrator.fsallast = fsallast
         # the controller was initialized by default for algs[1]
         reset_alg_dependent_opts!(integrator.opts.controller, algs[1], algs[3])
     elseif cache.current == 4
         initialize!(integrator, cache.cache4)
+        fsalfirst, fsallast = get_fsalfirstlast(cache.cache4)
+        integrator.fsalfirst = fsalfirst
+        integrator.fsallast = fsallast
         # the controller was initialized by default for algs[1]
         reset_alg_dependent_opts!(integrator.opts.controller, algs[1], algs[4])
     elseif cache.current == 5
         initialize!(integrator, cache.cache5)
+        fsalfirst, fsallast = get_fsalfirstlast(cache.cache5)
+        integrator.fsalfirst = fsalfirst
+        integrator.fsallast = fsallast
         # the controller was initialized by default for algs[1]
         reset_alg_dependent_opts!(integrator.opts.controller, algs[1], algs[5])
     elseif cache.current == 6
         initialize!(integrator, cache.cache6)
+        fsalfirst, fsallast = get_fsalfirstlast(cache.cache6)
+        integrator.fsalfirst = fsalfirst
+        integrator.fsallast = fsallast
         # the controller was initialized by default for algs[1]
         reset_alg_dependent_opts!(integrator.opts.controller, algs[1], algs[6])
     end