From 35ed20a7882645be072390d861d1318d2adc3b65 Mon Sep 17 00:00:00 2001
From: ioannisPApapadopoulos <john.papadopoulos0@gmail.com>
Date: Thu, 10 Aug 2023 18:29:30 +0100
Subject: [PATCH] more fixes for lazyarrays + some methods for ADI

---
 src/RadialPiecewisePolynomials.jl |  2 +-
 src/annuluselement.jl             |  6 ++--
 src/finitecontinuous.jl           | 57 +++++++++++++++++++++++++++++++
 src/finitezernike.jl              | 23 +++++++++++--
 4 files changed, 83 insertions(+), 5 deletions(-)

diff --git a/src/RadialPiecewisePolynomials.jl b/src/RadialPiecewisePolynomials.jl
index 6a809c8..f64c630 100644
--- a/src/RadialPiecewisePolynomials.jl
+++ b/src/RadialPiecewisePolynomials.jl
@@ -26,7 +26,7 @@ export SVector, Zeros, Ones, Vcat, Derivative, pad, paddeddata, Hcat, RadialCoor
         FiniteContinuousZernikeMode, FiniteContinuousZernike, inf_error,
         FiniteZernikeBasis, ZernikeBasisMode, FiniteZernikeBasisMode,
         ArrowheadMatrix,
-        modaltravcellwise, listtravcellwise
+        list_2_modaltrav, modaltrav_2_list, adi_2_modaltrav, adi_2_list
 
 include("arrowhead.jl")
 include("diskelement.jl")
diff --git a/src/annuluselement.jl b/src/annuluselement.jl
index c160403..6da1502 100644
--- a/src/annuluselement.jl
+++ b/src/annuluselement.jl
@@ -317,11 +317,13 @@ end
 
 function bubble2ann(C::ContinuousZernikeAnnulusElementMode, c::AbstractVector{T}) where T
     L₁₁, L₀₁, L₁₀ = C.L₁₁, C.L₀₁, C.L₁₀
-    R̃ = [L₁₀[:,1] L₀₁[:,1] L₁₁]
     if c isa LazyArray
         c = paddeddata(c)
     end
-    R̃[1:length(c), 1:length(c)] * c # coefficients for ZernikeAnnulus(ρ,0,0)
+    N = length(c)
+    R̃ = Hcat(view(L₁₀,1:N,1), view(L₀₁,1:N,1), view(L₁₁,1:N, 1:N-2))
+
+    R̃ * c # coefficients for ZernikeAnnulus(ρ,0,0)
 end
 
 function plotvalues(u::ApplyQuasiVector{T,typeof(*),<:Tuple{ContinuousZernikeAnnulusElementMode, AbstractVector}}) where T
diff --git a/src/finitecontinuous.jl b/src/finitecontinuous.jl
index 628ef5d..b12dafb 100644
--- a/src/finitecontinuous.jl
+++ b/src/finitecontinuous.jl
@@ -260,4 +260,61 @@ end
 function inf_error(F::FiniteContinuousZernike{T}, θs::AbstractVector, rs::AbstractVector, vals::AbstractVector, u::Function) where T
     K = lastindex(F.points)-1
     _inf_error(K, θs, rs, vals, u)
+end
+
+function modaltrav_2_list(F::FiniteContinuousZernike{T}, u::AbstractArray{Vector{T}}) where T
+    N, points = F.N, F.points
+    K = length(points) - 1
+    Ns, _, _ = _getMs_ms_js(N)
+
+    cs = []
+    u = ModalTrav.(u)
+    for i in 1:2N-1
+        v = zeros(T, Ns[i]-1, K)
+        for k in 1:K
+            v[:, k] = u[k].matrix[1:Ns[i]-1,i]
+        end
+        append!(cs, [pad(vec(v'), blockedrange(Fill(K, Ns[i]-1)))])
+    end
+    return cs
+end
+
+function adi_2_modaltrav(F::FiniteContinuousZernike{T}, wQ::Weighted{<:Any, <:Jacobi}, Us::AbstractArray, z::AbstractArray{T}) where T
+    N, points = F.N, F.points
+    K = length(points) - 1
+    Ns, _, _ = _getMs_ms_js(N)
+
+    Y =  [zeros(T, sum(1:N), length(z)) for k in 1:K]
+    for zi in 1:lastindex(z)
+        X = [zeros(T, Ns[1], 2N-1) for k in 1:K]
+        for k in 1:K
+            for n in 1:2N-1
+                us = Us[n][k:K:end, zi]
+                X[k][1:lastindex(us), n] = us
+            end
+            Y[k][:,zi] = ModalTrav(X[k])
+        end
+    end
+    return Y
+end
+
+function adi_2_list(F::FiniteContinuousZernike{T}, wQ::Weighted{<:Any, <:Jacobi}, Us::AbstractArray, z::AbstractArray{T}) where T
+    N, points = F.N, F.points
+    K = length(points) - 1
+    Ns, _, _ = _getMs_ms_js(N)
+
+    Y = []
+
+    for zi in 1:lastindex(z)
+        cs = []
+        for n in 1:2N-1
+            v = zeros(T, Ns[n]-1, K)
+            for k in 1:K
+                v[:, k] = Us[n][k:K:end, zi]
+            end
+            append!(cs, [pad(vec(v'), blockedrange(Fill(K, Ns[n]-1)))])
+        end
+        append!(Y, [cs])
+    end
+    return Y
 end
\ No newline at end of file
diff --git a/src/finitezernike.jl b/src/finitezernike.jl
index 1908689..5c7c585 100644
--- a/src/finitezernike.jl
+++ b/src/finitezernike.jl
@@ -346,7 +346,7 @@ function inf_error(Z::FiniteZernikeBasis{T}, θs::AbstractVector, rs::AbstractVe
 end
 
 ## Coefficient storage conversion
-function modaltravcellwise(Z::FiniteZernikeBasis{T}, u::AbstractArray) where T
+function list_2_modaltrav(Z::FiniteZernikeBasis{T}, u::AbstractArray) where T
     N, points = Z.N, Z.points
     K = length(points) - 1
     U = [zeros(T, N ÷ K, 2N-1) for i in 1:K]
@@ -359,7 +359,7 @@ function modaltravcellwise(Z::FiniteZernikeBasis{T}, u::AbstractArray) where T
     return ModalTrav.(U)
 end
 
-function listtravcellwise(Z::FiniteZernikeBasis{T}, u::AbstractArray{Vector{T}}) where T
+function modaltrav_2_list(Z::FiniteZernikeBasis{T}, u::AbstractArray{Vector{T}}) where T
     N, points = Z.N, Z.points
     K = length(points) - 1
     Ns, _, _ = _getMs_ms_js(N)
@@ -374,4 +374,23 @@ function listtravcellwise(Z::FiniteZernikeBasis{T}, u::AbstractArray{Vector{T}})
         append!(cs, [pad(vec(v'), blockedrange(Fill(K, Ns[i])))])
     end
     return cs
+end
+
+function adi_2_modaltrav(Z::FiniteZernikeBasis{T}, P::Legendre{T}, Us::AbstractArray, z::AbstractArray{T}) where T
+    N, points = Z.N, Z.points
+    K = length(points) - 1
+    Ns, _, _ = _getMs_ms_js(N)
+
+    Y =  [zeros(T, sum(1:N), length(z)) for k in 1:K]
+    for i in 1:lastindex(z)
+        X = [zeros(T, Ns[1], 2N-1) for k in 1:K]
+        for k in 1:K
+            for n in 1:2N-1
+                us = Us[n][k:K:end, i]
+                X[k][1:lastindex(us), n] = us
+            end
+            Y[k][:,i] = ModalTrav(X[k])
+        end
+    end
+    return Y
 end
\ No newline at end of file