add ported code for wide matrix solve with qr; add more tests; use pivoted QR by default in solve

Author: mateuszbaran

src/solve.jl

@@ -10,10 +10,66 @@
         R₁ = UpperTriangular(@view R[SOneTo(Sa[2]), SOneTo(Sa[2])])
         return R₁ \ y
     else
-        return R' * ((R * R') \ y)
+        return _wide_qr_solve(q, b)
     end
 end
 
+# based on https://github.com/JuliaLang/LinearAlgebra.jl/blob/16f64e78769d788376df0f36447affdb7b1b3df6/src/qr.jl#L652C1-L697C4
+function _wide_qr_solve(A::QR{T}, B::StaticMatrix{mB,nB,T}) where {mB,nB,T}
+    m, n = size(A)
+    minmn = min(m, n)
+    Bbuffer = similar(B)
+    copyto!(Bbuffer, B)
+    lmul!(adjoint(A.Q), view(Bbuffer, 1:m, :))
+    Rbuffer = similar(A.R)
+    copyto!(Rbuffer, A.R)
+
+    @inbounds begin
+        if n > m # minimum norm solution
+            τ = zeros(T,m)
+            for k = m:-1:1 # Trapezoid to triangular by elementary operation
+                x = view(Rbuffer, k, [k; m + 1:n])
+                τk = LinearAlgebra.reflector!(x)
+                τ[k] = conj(τk)
+                for i = 1:k - 1
+                    vRi = Rbuffer[i,k]
+                    for j = m + 1:n
+                        vRi += Rbuffer[i,j]*x[j - m + 1]'
+                    end
+                    vRi *= τk
+                    Rbuffer[i,k] -= vRi
+                    for j = m + 1:n
+                        Rbuffer[i,j] -= vRi*x[j - m + 1]
+                    end
+                end
+            end
+        end
+        ldiv!(UpperTriangular(view(Rbuffer, :, SOneTo(minmn))), view(Bbuffer, SOneTo(minmn), :))
+        if n > m # Apply elementary transformation to solution
+            Bbuffer[m + 1:mB,1:nB] .= zero(T)
+            for j = 1:nB
+                for k = 1:m
+                    vBj = Bbuffer[k,j]'
+                    for i = m + 1:n
+                        vBj += Bbuffer[i,j]'*Rbuffer[k,i]'
+                    end
+                    vBj *= τ[k]
+                    Bbuffer[k,j] -= vBj'
+                    for i = m + 1:n
+                        Bbuffer[i,j] -= Rbuffer[k,i]'*vBj'
+                    end
+                end
+            end
+        end
+    end
+    return similar_type(B)(Bbuffer)
+end
+function _wide_qr_solve(q::QR, b::StaticVecOrMat)
+    Q, R = q.Q, q.R
+    y = Q' * b
+    return R' * ((R * R') \ y)
+end
+
 @inline function _solve(::Size{(1,1)}, ::Size{(1,)}, a::StaticMatrix{<:Any, <:Any, Ta}, b::StaticVector{<:Any, Tb}) where {Ta, Tb}
     @inbounds return similar_type(b, typeof(a[1] \ b[1]))(a[1] \ b[1])
 end
@@ -80,7 +136,7 @@ end
         else
             quote
                 @_inline_meta
-                q = qr(a)
+                q = qr(a, ColumnNorm())
                 q \ b
             end
         end

test/solve.jl

@@ -33,6 +33,48 @@ using StaticArrays, Test, LinearAlgebra
             @test_throws MethodError m \ Array(v) # TODO: requires adjoint(::QR) method
         end
     end
+    @testset "More static tests" begin
+        # 1) 3×5 real, two RHS
+        A1 = @SMatrix [1.0  2.0  3.0  4.0  5.0;
+            0.0  1.0  0.0  1.0  0.0;
+            -1.0  0.0  2.0 -2.0  1.0]
+        B1 = @SMatrix [1.0 0.0;
+            0.0 1.0;
+            1.0 1.0]
+
+        # 2) 4×6 real
+        A2 = @SMatrix [ 2.0  -1.0  0.0  4.0  1.0  3.0;
+            -3.0   2.0  5.0 -1.0  0.0  2.0;
+            1.0   0.0  1.0  0.0  2.0 -2.0;
+            0.0   3.0 -1.0  1.0  1.0  0.0]
+        b2_1 = @SVector [1.0, 4.0, -2.0, 0.5]
+        b2_2 = @SMatrix [1.0 1.0
+            4.0 6.0
+            -2.0 2.0
+            0.5 1.5]
+
+        # 3) 3×4 complex
+        A3 = @SMatrix [1+2im   0+1im  2-1im  3+0im;
+            0+0im   2+0im  1+1im  0-2im;
+            3-1im  -1+0im  0+2im  1+0im]
+        b3_1 = @SVector [1+0im, 2-1im, -1+3im]
+        b3_2 = @SMatrix [
+            1+0im -9+0im
+            2-1im 2-4im
+            -1+3im 2+3im]
+
+        # 4) 3×6 rank-deficient (cols 3 = 1+2, col 4 = col 1, col 5 = col 2, col 6 = zeros)
+        A4 = @SMatrix [1.0 2.0 3.0 1.0 2.0 0.0;
+            0.0 1.0 1.0 0.0 1.0 0.0;
+            1.0 3.0 4.0 1.0 3.0 0.0]
+        b4_1 = @SVector [1.0, 0.0, 1.0]
+        b4_2 = @SMatrix [1.0 0.0
+                    0.0 1.0
+                    1.0 0.0]
+        for (A, B) in [(A1, B1), (A2, b2_1), (A2, b2_2), (A3, b3_1), (A3, b3_2), (A4, b4_1), (A4, b4_2)]
+            @test A \ B ≈ Array(A) \ Array(B)
+        end
+    end
 end