refactor: make cached_interpolation work with Duals

Author: SebastianM-C

Project.toml

@@ -9,6 +9,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [weakdeps]
@@ -27,6 +28,7 @@ IfElse = "0.1"
 LinearAlgebra = "1.10"
 ModelingToolkit = "9.28"
 OrdinaryDiffEq = "6.87"
+PreallocationTools = "0.4.23"
 SafeTestsets = "0.1"
 Symbolics = "5.35.1, 6"
 Test = "1"

src/Blocks/sources.jl

@@ -1,5 +1,6 @@
 using DiffEqBase
 import ChainRulesCore
+using PreallocationTools
 
 # Define and register smooth functions
 # These are "smooth" aka differentiable and avoid Gibbs effect
@@ -739,22 +740,28 @@ function Symbolics.derivative(::typeof(apply_interpolation), args::NTuple{2, Any
 end
 
 function cached_interpolation(interpolation_type, u, x, args)
-    prev_u = Ref(collect(copy(u)))
-    prev_x = Ref(collect(copy(x)))
-    # MTKParameters use views, so we want to ensure that the type is the same
-    interp = Ref(interpolation_type(prev_u[], prev_x[], args...))
+    prev_u = DiffCache(u)
+    # Interpolation points can be a range, but we want to be able
+    # to update the cache if needed (and setindex! is not defined on ranges)
+    # with a view from MTKParameters, so we collect to get a vector
+    prev_x = DiffCache(collect(x))
+    interp = GeneralLazyBufferCache() do (u,x)
+        interpolation_type(get_tmp(prev_u, u), get_tmp(prev_x, x), args...)
+    end
 
     let prev_u = prev_u,
         prev_x = prev_x,
         interp = interp,
         interpolation_type = interpolation_type
+
         function build_interpolation(u, x, args)
-            if (u, x) ≠ (prev_u[], prev_x[])
-                prev_u[] = collect(u)
-                prev_x[] = collect(x)
-                interp[] = interpolation_type(prev_u[], prev_x[], args...)
+            if (u, x) ≠ (get_tmp(prev_u, u), get_tmp(prev_x, x))
+                get_tmp(prev_u, u) .= u
+                get_tmp(prev_x, x) .= x
+                interp.bufs[(u,x)] = interpolation_type(
+                    get_tmp(prev_u, u), get_tmp(prev_x, x), args...)
             else
-                interp[]
+                interp[(u,x)]
             end
         end
     end