split tri_mul_invupper in 3

Author: dlfivefifty

src/banded/tridiagonalconjugation.jl

@@ -1,22 +1,20 @@
-"""
-upper_mul_tri_triview(U, X) == Tridiagonal(U*X) where U is Upper triangular BandedMatrix and X is Tridiagonal
-"""
+
+# upper_mul_tri_triview(U, X) == Tridiagonal(U*X) where U is Upper triangular BandedMatrix and X is Tridiagonal
 function upper_mul_tri_triview(U::BandedMatrix, X::Tridiagonal)
     T = promote_type(eltype(U), eltype(X))
     n = size(U,1)
     UX = Tridiagonal(Vector{T}(undef, n-1), Vector{T}(undef, n), Vector{T}(undef, n-1))
-    
+
     upper_mul_tri_triview!(UX, U, X)
 end
 
 function upper_mul_tri_triview!(UX::Tridiagonal, U::BandedMatrix, X::Tridiagonal)
     n = size(UX,1)
 
-    
+
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     UXdl, UXd, UXdu = UX.dl, UX.d, UX.du
-    Udat = U.data
-    
+
     l,u = bandwidths(U)
 
     @assert size(U) == (n,n)
@@ -29,7 +27,7 @@ function upper_mul_tri_triview!(UX::Tridiagonal, U::BandedMatrix, X::Tridiagonal
     finalize_upper_mul_tri_triview!(UX, U, X, n-1, bₖ, aₖ, cₖ, cₖ₋₁)
 end
 
-
+# populate first row of UX with UX
 function initiate_upper_mul_tri_triview!(UX, U, X)
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     UXdl, UXd, UXdu = UX.dl, UX.d, UX.du
@@ -65,6 +63,7 @@ function main_upper_mul_tri_triview!(UX, U, X, kr, bₖ=X.du[kr[1]-1], aₖ=X.d[
     UX, bₖ, aₖ, cₖ, cₖ₋₁
 end
 
+# populate rows k and k+1 of UX, assuming we are at the bottom-right
 function finalize_upper_mul_tri_triview!(UX, U, X, k, bₖ, aₖ, cₖ, cₖ₋₁)
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     UXdl, UXd, UXdu = UX.dl, UX.d, UX.du
@@ -92,43 +91,74 @@ end
 
 tri_mul_invupper_triview(X::Tridiagonal, R::BandedMatrix) = tri_mul_invupper_triview!(similar(X, promote_type(eltype(X), eltype(R))), X, R)
 
+
 function tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::BandedMatrix)
     n = size(X,1)
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     Ydl, Yd, Ydu = Y.dl, Y.d, Y.du
-    Rdat = R.data
-    
+
     l,u = bandwidths(R)
-    
+
     @assert size(R) == (n,n)
     @assert l == 0 && u ≥ 2
     # Tridiagonal bands can be resized
     @assert length(Xdl)+1 == length(Xd) == length(Xdu)+1 == length(Ydl)+1 == length(Yd) == length(Ydu)+1 == n
-    
-    
+
+    UX, Rₖₖ, Rₖₖ₊₁ = initiate_tri_mul_invupper_triview!(Y, X, R)
+    UX, Rₖₖ, Rₖₖ₊₁ = main_tri_mul_invupper_triview!(Y, X, R, 2:n-1, Rₖₖ, Rₖₖ₊₁)
+    finalize_tri_mul_invupper_triview!(Y, X, R, n, Rₖₖ, Rₖₖ₊₁)
+end
+
+# populate first row of X/R
+function initiate_tri_mul_invupper_triview!(Y, X, R)
+    Xdl, Xd, Xdu = X.dl, X.d, X.du
+    Ydl, Yd, Ydu = Y.dl, Y.d, Y.du
+    Rdat = R.data
+
+    l,u = bandwidths(R)
+
     k = 1
     aₖ,bₖ = Xd[k], Xdu[k]
     Rₖₖ,Rₖₖ₊₁ = Rdat[u+1,k], Rdat[u,k+1] # R[1,1], R[1,2]
     Yd[k] = aₖ/Rₖₖ
     Ydu[k] = bₖ - aₖ * Rₖₖ₊₁/Rₖₖ
 
-    @inbounds for k = 2:n-1
+    Y, Rₖₖ, Rₖₖ₊₁
+end
+
+
+# populate rows kr of X/R
+function main_tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::BandedMatrix, kr, Rₖₖ=R[first(kr),first(kr)], Rₖₖ₊₁=R[first(kr),first(kr)+1])
+    Xdl, Xd, Xdu = X.dl, X.d, X.du
+    Ydl, Yd, Ydu = Y.dl, Y.d, Y.du
+    Rdat = R.data
+    l,u = bandwidths(R)
+
+    @inbounds for k = kr
         cₖ₋₁,aₖ,bₖ = Xdl[k-1], Xd[k], Xdu[k]
         Ydl[k-1] = cₖ₋₁/Rₖₖ
         Yd[k] = aₖ-cₖ₋₁*Rₖₖ₊₁/Rₖₖ
         Ydu[k] = cₖ₋₁/Rₖₖ
-        Rₖₖ,Rₖₖ₊₁,Rₖ₋₁ₖ₊₁,Rₖ₋₁ₖ = Rdat[u+1,k], Rdat[u,k+1],Rdat[u-1,k+1],Rₖₖ₊₁ # R[2,2], R[2,3], R[1,3]
+        Rₖₖ,Rₖₖ₊₁,Rₖ₋₁ₖ₊₁,Rₖ₋₁ₖ = Rdat[u+1,k], Rdat[u,k+1],Rdat[u-1,k+1],Rₖₖ₊₁ # R[k,k], R[k,k+1], R[k-1,k]
         Yd[k] /= Rₖₖ
         Ydu[k-1] /= Rₖₖ
         Ydu[k] *= Rₖ₋₁ₖ*Rₖₖ₊₁/Rₖₖ - Rₖ₋₁ₖ₊₁
         Ydu[k] += bₖ - aₖ * Rₖₖ₊₁ / Rₖₖ
     end
+    Y, Rₖₖ, Rₖₖ₊₁
+end
+
 
-    k = n
+# populate row k of X/R, assuming we are at the bottom-right
+function finalize_tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::BandedMatrix, k, Rₖₖ=R[k-1,k-1], Rₖₖ₊₁=R[k-1,k])
+    Xdl, Xd, Xdu = X.dl, X.d, X.du
+    Ydl, Yd, Ydu = Y.dl, Y.d, Y.du
+    Rdat = R.data  
+    l,u = bandwidths(R)
     cₖ₋₁,aₖ = Xdl[k-1], Xd[k]
     Ydl[k-1] = cₖ₋₁/Rₖₖ
     Yd[k] = aₖ-cₖ₋₁*Rₖₖ₊₁/Rₖₖ
-    Rₖₖ = Rdat[u+1,k] # R[2,2], R[2,3], R[1,3]
+    Rₖₖ = Rdat[u+1,k] # R[k,k]
     Yd[k] /= Rₖₖ
     Ydu[k-1] /= Rₖₖ
 
@@ -173,7 +203,7 @@ function resizedata!(data::TridiagonalConjugationData, n)
         resize!(data.UX.dl, 2n)
         resize!(data.UX.d, 2n + 1)
         resize!(data.UX.du, 2n)
-    
+
         resize!(data.Y.dl, 2n)
         resize!(data.Y.d, 2n + 1)
         resize!(data.Y.du, 2n)