From e568de76df5aa7eac4852cf222b6add33b10793c Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Fri, 5 Jul 2024 09:53:43 +0200
Subject: [PATCH] Support inverting the integral of ConstantInterpolation

---
 src/integral_inverses.jl       | 52 ++++++++++++++++++++++++++++------
 test/integral_inverse_tests.jl | 16 +++++++++++
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/src/integral_inverses.jl b/src/integral_inverses.jl
index daf5d598..a6219b35 100644
--- a/src/integral_inverses.jl
+++ b/src/integral_inverses.jl
@@ -1,9 +1,12 @@
 abstract type AbstractIntegralInverseInterpolation{T} <: AbstractInterpolation{T} end
 
-struct LinearInterpolationIntInv{tType, IType, itpType, T} <:
+_integral(A::AbstractIntegralInverseInterpolation, idx, t) = throw(IntegralNotFoundError())
+invert_integral(A::AbstractInterpolation) = throw(IntegralInverseNotFoundError())
+
+struct LinearInterpolationIntInv{uType, tType, itpType, T} <:
        AbstractIntegralInverseInterpolation{T}
-    u::tType
-    t::IType
+    u::uType
+    t::tType
     extrapolate::Bool
     idx_prev::Base.RefValue{Int}
     itp::itpType
@@ -17,8 +20,6 @@ function invertible_integral(A::LinearInterpolation{<:AbstractVector{<:Number}})
     return all(A.u .> 0)
 end
 
-invert_integral(A::AbstractInterpolation) = throw(IntegralInverseNotFoundError())
-
 function invert_integral(A::LinearInterpolation{<:AbstractVector{<:Number}})
     !invertible_integral(A) && throw(IntegralNotInvertibleError())
     return LinearInterpolationIntInv(A.t, A.I, A)
@@ -26,12 +27,47 @@ end
 
 function _interpolate(
         A::LinearInterpolationIntInv{<:AbstractVector{<:Number}}, t::Number, iguess)
-    (; itp) = A
     idx = get_idx(A.t, t, iguess)
     Δt = t - A.t[idx]
     # TODO: Get from parameter cache: itp.p.slope[idx]
-    slope = (itp.u[idx + 1] - itp.u[idx]) / (itp.t[idx + 1] - itp.t[idx])
-    x = itp.u[idx]
+    slope = (A.itp.u[idx + 1] - A.itp.u[idx]) / (A.itp.t[idx + 1] - A.itp.t[idx])
+    x = A.itp.u[idx]
     u = A.u[idx] + 2Δt / (x + sqrt(x^2 + 2slope * Δt))
     u, idx
 end
+
+struct ConstantInterpolationIntInv{uType, tType, itpType, T} <:
+       AbstractIntegralInverseInterpolation{T}
+    u::uType
+    t::tType
+    extrapolate::Bool
+    idx_prev::Base.RefValue{Int}
+    itp::itpType
+    function ConstantInterpolationIntInv(u, t, A)
+        new{typeof(u), typeof(t), typeof(A), eltype(u)}(
+            u, t, A.extrapolate, Ref(1), A
+        )
+    end
+end
+
+function invertible_integral(A::ConstantInterpolation{<:AbstractVector{<:Number}})
+    return all(A.u .> 0)
+end
+
+function invert_integral(A::ConstantInterpolation{<:AbstractVector{<:Number}})
+    !invertible_integral(A) && throw(IntegralNotInvertibleError())
+    return ConstantInterpolationIntInv(A.t, A.I, A)
+end
+
+function _interpolate(
+        A::ConstantInterpolationIntInv{<:AbstractVector{<:Number}}, t::Number, iguess)
+    idx = get_idx(A.t, t, iguess; ub_shift = 0)
+    if A.itp.dir === :left
+        # :left means that value to the left is used for interpolation
+        idx_ = get_idx(A.t, t, idx; lb = 1, ub_shift = 0)
+    else
+        # :right means that value to the right is used for interpolation
+        idx_ = get_idx(A.t, t, idx; side = :first, lb = 1, ub_shift = 0)
+    end
+    A.u[idx] + (t - A.t[idx]) / A.itp.u[idx_], idx
+end
diff --git a/test/integral_inverse_tests.jl b/test/integral_inverse_tests.jl
index b8a835f7..229dca6f 100644
--- a/test/integral_inverse_tests.jl
+++ b/test/integral_inverse_tests.jl
@@ -21,3 +21,19 @@ end
     A = LinearInterpolation(u, t)
     @test_throws DataInterpolations.IntegralNotInvertibleError invert_integral(A)
 end
+
+@testset "Constant Interpolation" begin
+    t = collect(1:5)
+    u = [1.0, 1.0, 2.0, 4.0, 3.0]
+    test_integral_inverses(ConstantInterpolation; args = [u, t])
+    test_integral_inverses(ConstantInterpolation; args = [u, t], kwargs = [:dir => :right])
+
+    u = [1.0, -1.0, 2.0, 4.0, 3.0]
+    A = ConstantInterpolation(u, t)
+    @test_throws DataInterpolations.IntegralNotInvertibleError invert_integral(A)
+end
+
+t = collect(1:5)
+u = [1.0, 1.0, 2.0, 4.0, 3.0]
+A = QuadraticInterpolation(u, t)
+@test_throws DataInterpolations.IntegralInverseNotFoundError invert_integral(A)