AbstractVector -> AbstractArray and minor fixes

Author: mohamed82008

src/multilevel.jl

@@ -116,7 +116,7 @@ struct V <: Cycle
 end
 
 """
-    solve(ml::MultiLevel, b::AbstractVector, cycle, kwargs...)
+    solve(ml::MultiLevel, b::AbstractArray, cycle, kwargs...)
 
 Execute multigrid cycling.
 
@@ -134,14 +134,14 @@ Keyword Arguments
 * log::Bool - return vector of residuals along with solution
 
 """
-function solve(ml::MultiLevel, b::AbstractVector, args...; kwargs...)
+function solve(ml::MultiLevel, b::AbstractArray, args...; kwargs...)
     n = length(ml) == 1 ? size(ml.final_A, 1) : size(ml.levels[1].A, 1) 
     V = promote_type(eltype(ml.workspace), eltype(b))
-    bs = blocksize(ml.workspace)
-    x = bs === 1 ? zeros(V, n) : zeros(V, n, bs)
+    ndims(b) == blocksize(ml.workspace) && throw("Dimension of MultiLevel workspace and dimension of vector b must be equal.")
+    x = zeros(V, size(b))
     return solve!(x, ml, b, args...; kwargs...)
 end
-function solve!(x, ml::MultiLevel, b::AbstractVector{T},
+function solve!(x, ml::MultiLevel, b::AbstractArray{T},
                                     cycle::Cycle = V();
                                     maxiter::Int = 100,
                                     tol::Float64 = 1e-5,
@@ -169,7 +169,7 @@ function solve!(x, ml::MultiLevel, b::AbstractVector{T},
         end
         if calculate_residual
             mul!(res, A, x)
-            res .= b .- res
+            reshape(res, size(b)) .= b .- reshape(res, size(b))
             normres = norm(res)
             log && push!(residuals, normres)
         end
@@ -186,7 +186,7 @@ function __solve!(x, ml, v::V, b, lvl)
 
     res = ml.workspace.res_vecs[lvl]
     mul!(res, A, x)
-    res .= b .- res
+    reshape(res, size(b)) .= b .- reshape(res, size(b))
 
     coarse_b = ml.workspace.coarse_bs[lvl]
     mul!(coarse_b, ml.levels[lvl].R, res)

src/smoother.jl

@@ -10,7 +10,7 @@ struct GaussSeidel{S} <: Smoother
     sweep::S
     iter::Int
 end
-GaussSeidel(iter = 1) = GaussSeidel(SymmetricSweep(), iter)
+GaussSeidel(; iter = 1) = GaussSeidel(SymmetricSweep(), iter)
 GaussSeidel(f::ForwardSweep) = GaussSeidel(f, 1)
 GaussSeidel(b::BackwardSweep) = GaussSeidel(b, 1)
 GaussSeidel(s::SymmetricSweep) = GaussSeidel(s, 1)
@@ -47,35 +47,38 @@ end
 struct Jacobi{T,TX} <: Smoother
     ω::T
     temp::TX
+    iter::Int
 end
-Jacobi(ω, x::TX) where {T, TX<:AbstractVector{T}} = Jacobi{T,TX}(ω, similar(x))
-Jacobi(x::TX, ω = 0.5) where {T, TX<:AbstractVector{T}} = Jacobi{T,TX}(ω, similar(x))
+Jacobi(ω, x::TX; iter=1) where {T, TX<:AbstractArray{T}} = Jacobi{T,TX}(ω, similar(x), iter)
+Jacobi(x::TX, ω=0.5; iter=1) where {T, TX<:AbstractArray{T}} = Jacobi{T,TX}(ω, similar(x), iter)
 
 function (jacobi::Jacobi)(A, x, b)
 
     ω = jacobi.ω
     one = Base.one(eltype(A))
     temp = jacobi.temp
 
-    @inbounds for col = 1:size(x, 2)
-        for i = 1:size(A, 1)
-            temp[i, col] = x[i, col]
-        end
+    for i in 1:jacobi.iter
+        @inbounds for col = 1:size(x, 2)
+            for i = 1:size(A, 1)
+                temp[i, col] = x[i, col]
+            end
 
-        for i = 1:size(A, 1)
-            rsum = zero(eltype(A))
-            diag = zero(eltype(A))
+            for i = 1:size(A, 1)
+                rsum = zero(eltype(A))
+                diag = zero(eltype(A))
 
-            for j in nzrange(A, i)
-                row = A.rowval[j]
-                val = A.nzval[j]
+                for j in nzrange(A, i)
+                    row = A.rowval[j]
+                    val = A.nzval[j]
 
-                diag = ifelse(row == i, val, diag)
-                rsum += ifelse(row == i, 0, val * temp[row, col])
-            end
+                    diag = ifelse(row == i, val, diag)
+                    rsum += ifelse(row == i, 0, val * temp[row, col])
+                end
 
-            xcand = (one - ω) * temp[i, col] + ω * ((b[i, col] - rsum) / diag)
-            x[i, col] = ifelse(diag == 0, x[i, col], xcand)
+                xcand = (one - ω) * temp[i, col] + ω * ((b[i, col] - rsum) / diag)
+                x[i, col] = ifelse(diag == 0, x[i, col], xcand)
+            end
         end
     end
 end
@@ -87,8 +90,8 @@ struct ParallelJacobi{T,TX} <: Smoother
     ω::T
     temp::TX
 end
-ParallelJacobi(ω, x::TX) where {T, TX<:AbstractVector{T}} = ParallelJacobi{T,TX}(ω, similar(x))
-ParallelJacobi(x::TX, ω = 0.5) where {T, TX<:AbstractVector{T}} = ParallelJacobi{T,TX}(ω, similar(x))
+ParallelJacobi(ω, x::TX) where {T, TX<:AbstractArray{T}} = ParallelJacobi{T,TX}(ω, similar(x))
+ParallelJacobi(x::TX, ω = 0.5) where {T, TX<:AbstractArray{T}} = ParallelJacobi{T,TX}(ω, similar(x))
 
 struct ParallelTemp{TX, TI}
     temp::TX