Support complex vectors in olver (#10)

* Support complex vectors in olve * add test
JuliaApproximation · Aug 19, 2024 · e83e9e7 · e83e9e7
1 parent 2088ce9
commit e83e9e7
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 10 deletions.
diff --git a/src/olver.jl b/src/olver.jl
@@ -19,7 +19,7 @@ function olver_forward!(d::AbstractVector{T}, r, a, b, c, f, N; atol=eps(T)) whe
     d[1] = f[1]
     for k = 2:N
         d[k],r[k] = olver_forward_next(d, r, a, b, c, f, k)
-        M = max(M*abs(c[k-1]),one(T)) / abs(r[k])
+        M = max(M*abs(c[k-1]),one(M)) / abs(r[k])
     end
 
     for k = N+1:length(f)
@@ -45,7 +45,7 @@ function olver_forward!(d::AbstractVector{T}, r, a, b, c, f; atol=eps(T)) where
     resize!(d, Ñ)
 
     r[1] = b[1]
-    M = zero(T)
+    M = inv(abs(r[1]))
     # d[1] is left unmodified
     d[1] = f[1]
 
@@ -56,7 +56,7 @@ function olver_forward!(d::AbstractVector{T}, r, a, b, c, f; atol=eps(T)) where
             resize!(d, Ñ)
         end
         d[k],r[k] = olver_forward_next(d, r, a, b, c, f, k)
-        M = max(M*abs(c[k-1]),1) / abs(r[k])
+        M = max(M*abs(c[k-1]),one(M)) / abs(r[k])
         ε = M * abs(d[k])
         if ε ≤ atol
             return resize!(d,k-1), resize!(r,k), ε
@@ -100,8 +100,8 @@ returns a vector `u` satisfying the 3-term recurrence relationship
 It will compute at least `N` entries, but possibly more, returning the result
 when the backward error of the first `N` entries between consective truncations is less than `atol`.
 """
-function olver(a, b, c, f::AbstractVector{T}, N; atol=eps(float(T))) where T
-    dest = Vector{float(T)}(undef, N)
+function olver(a::AbstractVector{A}, b::AbstractVector{B}, c::AbstractVector{C}, f::AbstractVector{T}, N; atol=eps(float(real(promote_type(A,B,C,T))))) where {A,B,C,T}
+    dest = Vector{float(promote_type(A,B,C,T))}(undef, N)
     olver!(dest, similar(dest), a, b, c, f, N; atol)
 end
 
@@ -116,7 +116,7 @@ returns a vector `u` satisfying the 3-term recurrence relationship
 It will compute the result
 when the backward error between consective truncations is less than `atol`.
 """
-function olver(a, b, c, f::AbstractVector{T}; atol=eps(float(T))) where T
-    dest = Vector{float(T)}(undef, 1)
+function olver(a::AbstractVector{A}, b::AbstractVector{B}, c::AbstractVector{C}, f::AbstractVector{T}; atol=eps(float(real(promote_type(A,B,C,T))))) where {A,B,C,T}
+    dest = Vector{float(promote_type(A,B,C,T))}(undef, 1)
     olver!(dest, similar(dest), a, b, c, f; atol)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -212,7 +212,7 @@ end
         N = 2000
         z = 30.1
         a,b,c = 2ones(N-1) .- 0.5*cos.(1:N-1), -range(2; step=2, length=N)/z, 2ones(N-1)  .+ sin.(1:N-1)
-        
+
         f = [cos.(-(1:50)); exp.(-(1:N-50))]
         u = olver(a, b, c, f)
         T = Tridiagonal(a, Vector(b), c)
@@ -225,7 +225,7 @@ end
             d,r,ε = RecurrenceRelationships.olver_forward!(d, r, a, b, c, f; atol=0.1)
             n = length(d)
             @test Bidiagonal(ones(n+1), a[1:n] ./ r[1:n], :L) ≈ L[1:n+1,1:n+1]
-            
+
             g = L[1:n+1,1:n+1] \ f[1:n+1]
             @test g[1:n] ≈ d
             er = U[1:n+1,1:n+1] \ ([d; 0] -  g)
@@ -247,7 +247,7 @@ end
                 d = [0.0]; r = [0.0];
                 d,r,ε = RecurrenceRelationships.olver_forward!(d, r, a, b, c, f,k; atol=0.1)
                 n = length(d)
-            
+
                 g = L[1:n+1,1:n+1] \ f[1:n+1]
                 @test g[1:n] ≈ d
                 er = U[1:n+1,1:n+1] \ ([d; 0] -  g)
@@ -261,4 +261,16 @@ end
             @test length(olver(a, b, c, f, N)) == N
         end
     end
+
+    @testset "complex" begin
+        N = 100
+        a,b,c = Fill(1/2,N-1), Zeros{Int}(N), Fill(1/2,N-1)
+        z = 1 + im
+        f = [-π/2; zeros(N-1)]
+        u = olver(a, b .- z, c, f)
+        ζ = z - sqrt(z-1)*sqrt(z+1)
+        @test π*ζ .^ (1:length(u)) ≈ u
+        @test u ≈ olver(a, b .- z, c, f .+ 0im)
+        @test u[1:5] ≈ olver(a, b .- z, c, f .+ 0im, 5)
+    end
 end