Add LSQR

Author: timholy

src/krylov_workspaces.jl

@@ -2030,22 +2030,22 @@ The following outer constructors can be used to initialize this workspace:
     workspace = LsqrWorkspace(A, b)
     workspace = LsqrWorkspace(kc::KrylovConstructor)
 """
-mutable struct LsqrWorkspace{T,FC,S} <: KrylovWorkspace{T,FC,S,S}
+mutable struct LsqrWorkspace{T,FC,Sm,Sn} <: KrylovWorkspace{T,FC,Sm,Sn}
   m       :: Int
   n       :: Int
-  x       :: S
-  Nv      :: S
-  Aᴴu     :: S
-  w       :: S
-  Mu      :: S
-  Av      :: S
-  u       :: S
-  v       :: S
+  x       :: Sn
+  Nv      :: Sn
+  Aᴴu     :: Sn
+  w       :: Sn
+  Mu      :: Sm
+  Av      :: Sm
+  u       :: Sm
+  v       :: Sn
   err_vec :: Vector{T}
   stats   :: SimpleStats{T}
 end
 
-function LsqrWorkspace(kc::KrylovConstructor; window::Integer = 5)
+function LsqrWorkspace(kc::KrylovConstructor{Sm,Sn}; window::Integer = 5) where {Sm,Sn}
   S   = typeof(kc.vm)
   FC  = eltype(S)
   T   = real(FC)
@@ -2061,7 +2061,7 @@ function LsqrWorkspace(kc::KrylovConstructor; window::Integer = 5)
   v   = similar(kc.vn_empty)
   err_vec = zeros(T, window)
   stats = SimpleStats(0, false, false, false, 0, T[], T[], T[], 0.0, "unknown")
-  workspace = LsqrWorkspace{T,FC,S}(m, n, x, Nv, Aᴴu, w, Mu, Av, u, v, err_vec, stats)
+  workspace = LsqrWorkspace{T,FC,Sm,Sn}(m, n, x, Nv, Aᴴu, w, Mu, Av, u, v, err_vec, stats)
   return workspace
 end
 
@@ -2079,7 +2079,7 @@ function LsqrWorkspace(m::Integer, n::Integer, S::Type; window::Integer = 5)
   err_vec = zeros(T, window)
   S = isconcretetype(S) ? S : typeof(x)
   stats = SimpleStats(0, false, false, false, 0, T[], T[], T[], 0.0, "unknown")
-  workspace = LsqrWorkspace{T,FC,S}(m, n, x, Nv, Aᴴu, w, Mu, Av, u, v, err_vec, stats)
+  workspace = LsqrWorkspace{T,FC,S,S}(m, n, x, Nv, Aᴴu, w, Mu, Av, u, v, err_vec, stats)
   return workspace
 end
 

src/lsqr.jl

@@ -164,7 +164,7 @@ args_lsqr = (:A, :b)
 kwargs_lsqr = (:M, :N, :ldiv, :sqd, :λ, :radius, :etol, :axtol, :btol, :conlim, :atol, :rtol, :itmax, :timemax, :verbose, :history, :callback, :iostream)
 
 @eval begin
-  function lsqr!(workspace :: LsqrWorkspace{T,FC,S}, $(def_args_lsqr...); $(def_kwargs_lsqr...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, S <: AbstractVector{FC}}
+  function lsqr!(workspace :: LsqrWorkspace{T,FC,Sm,Sn}, $(def_args_lsqr...); $(def_kwargs_lsqr...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, Sm <: AbstractVector{FC}, Sn <: AbstractVector{FC}}
 
     # Timer
     start_time = time_ns()
@@ -185,14 +185,14 @@ kwargs_lsqr = (:M, :N, :ldiv, :sqd, :λ, :radius, :etol, :axtol, :btol, :conlim,
 
     # Check type consistency
     eltype(A) == FC || @warn "eltype(A) ≠ $FC. This could lead to errors or additional allocations in operator-vector products."
-    ktypeof(b) == S || error("ktypeof(b) must be equal to $S")
+    ktypeof(b) == Sm || error("ktypeof(b) must be equal to $Sm")
 
     # Compute the adjoint of A
     Aᴴ = A'
 
     # Set up workspace.
-    allocate_if(!MisI, workspace, :u, S, workspace.Av)  # The length of u is m
-    allocate_if(!NisI, workspace, :v, S, workspace.x)   # The length of v is n
+    allocate_if(!MisI, workspace, :u, Sm, workspace.Av)  # The length of u is m
+    allocate_if(!NisI, workspace, :v, Sn, workspace.x)   # The length of v is n
     x, Nv, Aᴴu, w = workspace.x, workspace.Nv, workspace.Aᴴu, workspace.w
     Mu, Av, err_vec, stats = workspace.Mu, workspace.Av, workspace.err_vec, workspace.stats
     rNorms, ArNorms = stats.residuals, stats.Aresiduals