Make periodic extrapolation smooth

SciML · Dec 3, 2024 · 8395f22 · 8395f22
1 parent 77a5875
commit 8395f22
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 14 deletions.
diff --git a/src/interpolation_methods.jl b/src/interpolation_methods.jl
@@ -192,13 +192,6 @@ function _interpolate(A::SmoothedConstantInterpolation{<:AbstractVector}, t::Num
     idx = get_idx(A, t, iguess)
     d_lower, d_upper, c_lower, c_upper = get_parameters(A, idx)
 
-    # Fix extrapolation behavior as constant for now
-    if t <= first(A.t)
-        return first(A.u)
-    elseif t >= last(A.t)
-        return A.u[end - 1]
-    end
-
     out = A.u[idx]
 
     if (t - A.t[idx]) < d_lower

diff --git a/src/interpolation_utils.jl b/src/interpolation_utils.jl
@@ -215,9 +215,9 @@ function get_parameters(A::SmoothedConstantInterpolation, idx)
         c_upper = A.p.c[idx + 1]
         d_lower, d_upper, c_lower, c_upper
     else
-        d_lower, c_lower = smoothed_linear_interpolation_parameters(A.u, A.t, A.d_max, idx)
-        d_upper, c_upper = smoothed_linear_interpolation_parameters(
-            A.u, A.t, A.d_max, idx + 1)
+        d_lower, c_lower = smoothed_constant_interpolation_parameters(A.u, A.t, A.d_max, idx, A.extrapolation_left, A.extrapolation_right)
+        d_upper, c_upper = smoothed_constant_interpolation_parameters(
+            A.u, A.t, A.d_max, idx + 1, A.extrapolation_left, A.extrapolation_right)
         d_lower, d_upper, c_lower, c_upper
     end
 end

diff --git a/src/parameter_caches.jl b/src/parameter_caches.jl
@@ -39,7 +39,7 @@ end
 
 function SmoothedConstantParameterCache(u, t, cache_parameters, d_max)
     if cache_parameters
-        parameters = smoothed_linear_interpolation_parameters.(
+        parameters = smoothed_constant_interpolation_parameters.(
             Ref(u), Ref(t), d_max, eachindex(t))
         d, c = collect.(eachrow(stack(collect.(parameters))))
         SmoothedConstantParameterCache(d, c)
@@ -48,10 +48,18 @@ function SmoothedConstantParameterCache(u, t, cache_parameters, d_max)
     end
 end
 
-function smoothed_linear_interpolation_parameters(u, t, d_max, idx)
-    # TODO: Add support for making periodic extrapolation smooth
+function smoothed_constant_interpolation_parameters(u, t, d_max, idx, extrapolation_left, extrapolation_right)
     if isone(idx) || (idx == length(t))
-        zero(one(eltype(t))) / 2, zero(one(eltype(u)) / 2)
+        # If extrapolation is periodic, make the transition differentiable
+        if extrapolation_left == extrapolation_right == ExtrapolationType.Periodic
+            if isone(idx)
+                min(t[end] - t[end - 1], t[2] - t[1], 2d_max) / 2, (u[1] - u[end - 1]) / 2
+            else
+                min(t[end] - t[end - 1], t[2] - t[1], 2d_max) / 2, (u[1] - u[end - 1]) / 2
+            end
+        else
+            zero(one(eltype(t)) / 2), zero(one(eltype(u)) / 2)
+        end
     else
         min(t[idx] - t[idx - 1], t[idx + 1] - t[idx], 2d_max) / 2, (u[idx] - u[idx - 1]) / 2
     end