inv, size, and \ for QR objects

Fixes JuliaArrays#1192

Author: gustafsson

src/qr.jl

@@ -1,3 +1,6 @@
+import Base: \, size, inv
+import LinearAlgebra: ldiv!, checksquare
+
 # define our own struct since LinearAlgebra.QR are restricted to Matrix
 struct QR{Q,R,P}
     Q::Q
@@ -245,3 +248,83 @@ end
 #    end
 #end
 
+
+size(F::QR) = size(F.Q)
+
+is_identity_perm(p::AbstractVector{T}) where {T<:Integer} = all(i->i==p[i], first(axes(p)))
+
+ldiv!(x::StaticArray, F::QR, y::AbstractVecOrMat) = (x .= F \ y) # Compatibility. Note that \ already allocates 0 bytes.
+
+
+function \(F::QR, y::AbstractVecOrMat)
+    checksquare(F.R)
+    v = F.Q' * y
+
+    x = backsub(F.R, v)
+
+    invpivot(x, F.p)
+end
+
+
+@inline function invpivot(x::AbstractVecOrMat, p)
+    if is_identity_perm(p)
+        x
+    else
+        extra = ntuple(_ -> Colon(), ndims(x) - 1)
+        getindex(x, invperm(p), extra...)
+    end
+end
+
+
+# Simple back substitution for an upper–triangular system R*x = y.
+function backsub(R::StaticMatrix{r,c,T}, y::AbstractVector{T}) where {r,c,T}
+    x = MVector{c,T}(undef)
+    backsub!(x, R, y)
+    SVector(x)
+end
+
+function backsub(R::StaticMatrix{r,c,T}, y::AbstractMatrix{T}) where {r,c,T}
+    x = MMatrix{c,size(y,2),T}(undef)
+    for i in 1:size(y,2)
+        @views backsub!(x[:,i], R, y[:,i])
+    end
+    SMatrix(x)
+end
+
+@inline function backsub!(x::StaticVector{c}, R::StaticMatrix{r,c,T}, y::AbstractVector{T}) where {r,c,T}
+    Base.@boundscheck Base.checkbounds(x, Base.OneTo(r))
+    Base.@boundscheck Base.checkbounds(y, Base.OneTo(r))
+
+    @inbounds for i in r:-1:1
+        s = zero(T)
+        for j in i+1:c
+            s += R[i, j] * x[j]
+        end
+        x[i] = (y[i] - s) / R[i, i]
+    end
+    if r < c
+        @inbounds for i in r+1:c
+            x[i] = 0
+        end
+    end
+    x
+end
+
+
+function inv(F::QR)
+    checksquare(F.Q)
+    n = checksquare(F.R)
+
+    T = eltype(F.R)
+    
+    # Compute inverse of R via back substitution on each column of the identity.
+    R_inv_cols = ntuple(j -> begin
+        # Build the j-th unit vector.
+        e_j = SVector{n, T}(ntuple(i -> i == j ? one(T) : zero(T), n))
+        backsub(F.R, e_j)
+    end, n)
+
+    R⁻¹ = hcat(R_inv_cols...)
+    A⁻¹ = R⁻¹ * F.Q'
+    invpivot(A⁻¹, F.p)
+end

test/qr.jl

@@ -69,10 +69,94 @@ Random.seed!(42)
     end
 end
 
+
 @testset "QR method ambiguity" begin
     # Issue #931; just test that methods do not throw an ambiguity error when called
     A = @SMatrix [1.0 2.0 3.0; 4.0 5.0 6.0]
     @test isa(qr(A),              StaticArrays.QR)
     @test isa(qr(A, Val(true)),   StaticArrays.QR)
     @test isa(qr(A, Val(false)),  StaticArrays.QR)
-end
+end
+
+
+@testset "#1192 The following functions are available for the QR objects: inv, size, and \\." begin
+    @testset "pivot=$pivot" for pivot in [Val(true), Val(false)] #, ColumnNorm()]
+        y = @SVector rand(5)
+        Y = @SMatrix rand(5,5)
+        A = @SMatrix rand(5,5)
+        A_over = @SMatrix rand(5,6)
+        A_under = @SMatrix rand(5,4)
+
+        F = qr(A, pivot)
+        F_over = qr(A_over, pivot)
+        F_under = qr(A_under, pivot)
+
+        @testset "size" begin
+            @test size(A) == (5,5)
+            @test size(A_over) == (5,6)
+            @test size(A_under) == (5,4)
+        end
+
+        @testset "square inversion" begin
+            A_inv = inv(F)
+            @test inv(F) * A ≈ I(5)
+            @test inv(F) ≈ inv(qr(Matrix(A)))
+            @test_throws DimensionMismatch inv(F_under)
+            @test_throws DimensionMismatch inv(F_over)
+        end
+
+        @testset "solve linear system" begin
+            x = Matrix(A) \ Vector(y)
+            @test x ≈ A \ y ≈ F \ y ≈ F \ Vector(y)
+
+            x_under = Matrix(A_under) \ Vector(y)
+            @test x_under == A_under \ y
+            @test x_under ≈ F_under \ y
+            @test F_under \ y == F_under \ Vector(y)
+
+            x_over = Matrix(A_over) \ Vector(y)
+            @test x_over ≈ A_over \ y
+            @test A_over * x_over ≈ y
+
+            @test_throws DimensionMismatch F_over \ y
+            @test_throws DimensionMismatch qr(Matrix(A_over)) \ y
+        end
+        
+        @testset "solve several linear systems" begin
+            @test F \ Y ≈ A \ Y
+            @test F_under \ Y ≈ A_under \ Y
+        end
+
+        @testset "ldiv!" begin
+            x = @MVector zeros(5)
+            ldiv!(x, F, y)
+            @test x ≈ A \ y
+
+            X = @MMatrix zeros(5,5)
+            Y = @SMatrix rand(5,5)
+            ldiv!(X, F, Y)
+            @test X ≈ A \ Y
+        end
+        
+        @testset "invperm" begin
+            x = @SVector [10,15,3,7]
+            p = @SVector [4,2,1,3]
+            @test x == x[p][invperm(p)]
+            @test StaticArrays.is_identity_perm(p[invperm(p)])
+            @test_throws Union{BoundsError,ArgumentError} invperm(x)
+        end
+
+        @testset "10x faster" begin
+            time_to_test = @elapsed (function()
+                y2 = @SVector rand(50)
+                A2 = @SMatrix rand(50,5)
+                F2 = qr(A2, pivot)
+
+                min_time_to_solve = minimum(@elapsed(A2 \ y2) for _ in 1:1_000)
+                min_time_to_solve_qr = minimum(@elapsed(F2 \ y2) for _ in 1:1_000)
+                @test 10min_time_to_solve_qr < min_time_to_solve
+            end)()
+            @test time_to_test < 10
+        end
+    end
+end