Merge pull request #368 from martinholters/mh/lu

Non-allocating LU decomposition

Author: andyferris

src/lu.jl

@@ -1,29 +1,126 @@
 # LU decomposition
 function lu(A::StaticMatrix, pivot::Union{Type{Val{false}},Type{Val{true}}}=Val{true})
-    L,U,p = _lu(Size(A), A, pivot)
+    L,U,p = _lu(A, pivot)
     (L,U,p)
 end
 
 # For the square version, return explicit lower and upper triangular matrices.
 # We would do this for the rectangular case too, but Base doesn't support that.
 function lu(A::StaticMatrix{N,N}, pivot::Union{Type{Val{false}},Type{Val{true}}}=Val{true}) where {N}
-    L,U,p = _lu(Size(A), A, pivot)
+    L,U,p = _lu(A, pivot)
     (LowerTriangular(L), UpperTriangular(U), p)
 end
 
+@generated function _lu(A::StaticMatrix{M,N,T}, pivot) where {M,N,T}
+    if M*N ≤ 14*14
+        :(__lu(A, pivot))
+    else
+        quote
+            # call through to Base to avoid excessive time spent on type inference for large matrices
+            f = lufact(Matrix(A), pivot)
+            # Trick to get the output eltype - can't rely on the result of f[:L] as
+            # it's not type inferrable.
+            T2 = arithmetic_closure(T)
+            L = similar_type(A, T2, Size($M, $(min(M,N))))(f[:L])
+            U = similar_type(A, T2, Size($(min(M,N)), $N))(f[:U])
+            p = similar_type(A, Int, Size($M))(f[:p])
+            (L,U,p)
+        end
+    end
+end
 
-@inline function _lu(::Size{S}, A::StaticMatrix, pivot) where {S}
-    # For now, just call through to Base.
-    # TODO: statically sized LU without allocations!
-    f = lufact(Matrix(A), pivot)
-    T = eltype(A)
-    # Trick to get the output eltype - can't rely on the result of f[:L] as
-    # it's not type inferrable.
-    T2 = arithmetic_closure(T)
-    L = similar_type(A, T2, Size(Size(A)[1], diagsize(A)))(f[:L])
-    U = similar_type(A, T2, Size(diagsize(A), Size(A)[2]))(f[:U])
-    p = similar_type(A, Int, Size(Size(A)[1]))(f[:p])
-    (L,U,p)
+__lu(A::StaticMatrix{0,0,T}, ::Type{Val{Pivot}}) where {T,Pivot} =
+    (SMatrix{0,0,typeof(one(T))}(), A, SVector{0,Int}())
+
+__lu(A::StaticMatrix{0,1,T}, ::Type{Val{Pivot}}) where {T,Pivot} =
+    (SMatrix{0,0,typeof(one(T))}(), A, SVector{0,Int}())
+
+__lu(A::StaticMatrix{0,N,T}, ::Type{Val{Pivot}}) where {T,N,Pivot} =
+    (SMatrix{0,0,typeof(one(T))}(), A, SVector{0,Int}())
+
+__lu(A::StaticMatrix{1,0,T}, ::Type{Val{Pivot}}) where {T,Pivot} =
+    (SMatrix{1,0,typeof(one(T))}(), SMatrix{0,0,T}(), SVector{1,Int}(1))
+
+__lu(A::StaticMatrix{M,0,T}, ::Type{Val{Pivot}}) where {T,M,Pivot} =
+    (SMatrix{M,0,typeof(one(T))}(), SMatrix{0,0,T}(), SVector{M,Int}(1:M))
+
+__lu(A::StaticMatrix{1,1,T}, ::Type{Val{Pivot}}) where {T,Pivot} =
+    (SMatrix{1,1}(one(T)), A, SVector(1))
+
+__lu(A::StaticMatrix{1,N,T}, ::Type{Val{Pivot}}) where {N,T,Pivot} =
+    (SMatrix{1,1,T}(one(T)), A, SVector{1,Int}(1))
+
+function __lu(A::StaticMatrix{M,1}, ::Type{Val{Pivot}}) where {M,Pivot}
+    @inbounds begin
+        kp = 1
+        if Pivot
+            amax = abs(A[1,1])
+            for i = 2:M
+                absi = abs(A[i,1])
+                if absi > amax
+                    kp = i
+                    amax = absi
+                end
+            end
+        end
+        ps = tailindices(Val{M})
+        if kp != 1
+            ps = setindex(ps, 1, kp-1)
+        end
+        U = SMatrix{1,1}(A[kp,1])
+        # Scale first column
+        Akkinv = inv(A[kp,1])
+        Ls = A[ps,1] * Akkinv
+        if !isfinite(Akkinv)
+            Ls = zeros(Ls)
+        end
+        L = [SVector{1}(one(eltype(Ls))); Ls]
+        p = [SVector{1,Int}(kp); ps]
+    end
+    return (SMatrix{M,1}(L), U, p)
+end
+
+function __lu(A::StaticMatrix{M,N,T}, ::Type{Val{Pivot}}) where {M,N,T,Pivot}
+    @inbounds begin
+        kp = 1
+        if Pivot
+            amax = abs(A[1,1])
+            for i = 2:M
+                absi = abs(A[i,1])
+                if absi > amax
+                    kp = i
+                    amax = absi
+                end
+            end
+        end
+        ps = tailindices(Val{M})
+        if kp != 1
+            ps = setindex(ps, 1, kp-1)
+        end
+        Ufirst = SMatrix{1,N}(A[kp,:])
+        # Scale first column
+        Akkinv = inv(A[kp,1])
+        Ls = A[ps,1] * Akkinv
+        if !isfinite(Akkinv)
+            Ls = zeros(Ls)
+        end
+
+        # Update the rest
+        Arest = A[ps,tailindices(Val{N})] - Ls*Ufirst[:,tailindices(Val{N})]
+        Lrest, Urest, prest = __lu(Arest, Val{Pivot})
+        p = [SVector{1,Int}(kp); ps[prest]]
+        L = [[SVector{1}(one(eltype(Ls))); Ls[prest]] [zeros(SMatrix{1}(Lrest[1,:])); Lrest]]
+        U = [Ufirst; [zeros(Urest[:,1]) Urest]]
+    end
+    return (L, U, p)
+end
+
+# Create SVector(2,3,...,M)
+# Note that
+#     tailindices(::Type{Val{M}}) where {M} = SVector(Base.tail(ntuple(identity, Val{M})))
+# works, too, but is only inferrable for M ≤ 14 (at least up to Julia 0.7.0-DEV.4021)
+@generated function tailindices(::Type{Val{M}}) where {M}
+    :(SVector{$(M-1),Int}($(tuple(2:M...))))
 end
 
 # Base.lufact() interface is fairly inherently type unstable.  Punt on

test/lu.jl

@@ -1,21 +1,41 @@
 using StaticArrays, Base.Test
 
-@testset "LU decomposition" begin
-    # Square case
-    m22 = @SMatrix [1 2; 3 4]
-    @test @inferred(lu(m22)) isa Tuple{LowerTriangular{Float64,SMatrix{2,2,Float64,4}}, UpperTriangular{Float64,SMatrix{2,2,Float64,4}}, SVector{2,Int}}
-    @test lu(m22)[1]::LowerTriangular{<:Any,<:StaticMatrix} ≊ lu(Matrix(m22))[1]
-    @test lu(m22)[2]::UpperTriangular{<:Any,<:StaticMatrix} ≊ lu(Matrix(m22))[2]
-    @test lu(m22)[3]::StaticVector ≊ lu(Matrix(m22))[3]
+@testset "LU decomposition (pivot=$pivot)" for pivot in (true, false)
+    @testset "$m×$n" for m in [0:4..., 15, 50], n in [0:4..., 15, 50]
+        a = SMatrix{m,n,Int}(1:(m*n))
+        l, u, p = @inferred(lu(a, Val{pivot}))
 
-    # Rectangular case
-    m23 = @SMatrix Float64[3 9 4; 6 6 2]
-    @test @inferred(lu(m23)) isa Tuple{SMatrix{2,2,Float64,4}, SMatrix{2,3,Float64,6}, SVector{2,Int}}
-    @test lu(m23)[1] ≊ lu(Matrix(m23))[1]
-    @test lu(m23)[2] ≊ lu(Matrix(m23))[2]
-    @test lu(m23)[3] ≊ lu(Matrix(m23))[3]
+        # expected types
+        @test p isa SVector{m,Int}
+        if m==n
+            @test l isa LowerTriangular{<:Any,<:SMatrix{m,n}}
+            @test u isa UpperTriangular{<:Any,<:SMatrix{m,n}}
+        else
+            @test l isa SMatrix{m,min(m,n)}
+            @test u isa SMatrix{min(m,n),n}
+        end
 
-    @test lu(m23')[1] ≊ lu(Matrix(m23'))[1]
-    @test lu(m23')[2] ≊ lu(Matrix(m23'))[2]
-    @test lu(m23')[3] ≊ lu(Matrix(m23'))[3]
+        if pivot
+            # p is a permutation
+            @test sort(p) == collect(1:m)
+        else
+            @test p == collect(1:m)
+        end
+
+        # l is unit lower triangular
+        for i=1:m, j=(i+1):size(l,2)
+            @test iszero(l[i,j])
+        end
+        for i=1:size(l,2)
+            @test l[i,i] == 1
+        end
+
+        # u is upper triangular
+        for i=1:size(u,1), j=1:i-1
+            @test iszero(u[i,j])
+        end
+
+        # decomposition is correct
+        @test l*u ≈ a[p,:]
+    end
 end