From 8395f226336b16141136e6848307f54546c7d6d6 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Tue, 3 Dec 2024 16:04:45 +0100
Subject: [PATCH] Make periodic extrapolation smooth

---
 src/interpolation_methods.jl |  7 -------
 src/interpolation_utils.jl   |  6 +++---
 src/parameter_caches.jl      | 16 ++++++++++++----
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/interpolation_methods.jl b/src/interpolation_methods.jl
index 0882e67b..ccbc668f 100644
--- 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
index f6537e4d..bb1f7492 100644
--- 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
index 912f497e..ca4c5773 100644
--- 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