Merge pull request #192 from bshall/solve-triangular

Add simple forward and backward substitution to solve triangular systems

Author: c42f

src/solve.jl

@@ -1,9 +1,12 @@
-@inline (\)(a::StaticMatrix{<:Any, <:Any, T}, b::StaticVector{<:Any, T}) where {T} = solve(Size(a), Size(b), a, b)
+@inline (\)(a::StaticMatrix, b::StaticVector) = solve(Size(a), Size(b), a, b)
+@inline (\)(a::Union{UpperTriangular{<:Any, <:StaticMatrix}, LowerTriangular{<:Any, <:StaticMatrix}}, b::StaticVecOrMat) = solve(Size(a.data), Size(b), a, b)
 
 # TODO: Ineffecient but requires some infrastructure (e.g. LU or QR) to make efficient so we fall back on inv for now
 @inline solve(::Size, ::Size, a, b) = inv(a) * b
 
-@inline solve(::Size{(1,1)}, ::Size{(1,)}, a, b) = similar_type(b, typeof(b[1] \ a[1]))(b[1] \ a[1])
+@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(b[1] \ a[1]))(b[1] \ a[1])
+end
 
 @inline function solve(::Size{(2,2)}, ::Size{(2,)}, a::StaticMatrix{<:Any, <:Any, Ta}, b::StaticVector{<:Any, Tb}) where {Ta, Tb}
     d = det(a)
@@ -26,3 +29,67 @@ end
             (a[1,2]*a[3,1] - a[1,1]*a[3,2])*b[2] +
             (a[1,1]*a[2,2] - a[1,2]*a[2,1])*b[3]) / d )
 end
+
+@generated function solve(::Size{sa}, ::Size{sb}, a::UpperTriangular{Ta, <:StaticMatrix{<:Any, <:Any, Ta}}, b::StaticVector{<:Any, Tb}) where {sa, sb, Ta, Tb}
+    if sa[1] != sb[1]
+        throw(DimensionMismatch("right hand side b needs first dimension of size $(sa[1]), has size $(sb[1])"))
+    end
+
+    x = [Symbol("x$k") for k = 1:sb[1]]
+    expr = [:($(x[i]) = $(reduce((ex1, ex2) -> :(-($ex1,$ex2)), [j == i ? :(b[$j]) : :(a[$i, $j]*$(x[j])) for j = i:sa[1]]))/a[$i, $i]) for i = sb[1]:-1:1]
+
+    quote
+        @_inline_meta
+        T = typeof((zero(Ta)*zero(Tb) + zero(Ta)*zero(Tb))/one(Ta))
+        @inbounds $(Expr(:block, expr...))
+        @inbounds return similar_type(b, T)(tuple($(x...)))
+    end
+end
+
+@generated function solve(::Size{sa}, ::Size{sb}, a::UpperTriangular{Ta, <:StaticMatrix{<:Any, <:Any, Ta}}, b::StaticMatrix{<:Any, <:Any, Tb}) where {sa, sb, Ta, Tb}
+    if sa[1] != sb[1]
+        throw(DimensionMismatch("right hand side b needs first dimension of size $(sa[1]), has size $(sb[1])"))
+    end
+
+    x = [Symbol("x$k1$k2") for k1 = 1:sb[1], k2 = 1:sb[2]]
+    expr = [:($(x[k1, k2]) = $(reduce((ex1, ex2) -> :(-($ex1,$ex2)), [j == k1 ? :(b[$j, $k2]) : :(a[$k1, $j]*$(x[j, k2])) for j = k1:sa[1]]))/a[$k1, $k1]) for k1 = sb[1]:-1:1, k2 = 1:sb[2]]
+
+    quote
+        @_inline_meta
+        T = typeof((zero(Ta)*zero(Tb) + zero(Ta)*zero(Tb))/one(Ta))
+        @inbounds $(Expr(:block, expr...))
+        @inbounds return similar_type(b, T)(tuple($(x...)))
+    end
+end
+
+@generated function solve(::Size{sa}, ::Size{sb}, a::LowerTriangular{Ta, <:StaticMatrix{<:Any, <:Any, Ta}}, b::StaticVector{<:Any, Tb}) where {sa, sb, Ta, Tb}
+    if sa[1] != sb[1]
+        throw(DimensionMismatch("right hand side b needs first dimension of size $(sa[1]), has size $(sb[1])"))
+    end
+
+    x = [Symbol("x$k") for k = 1:sb[1]]
+    expr = [:($(x[i]) = $(reduce((ex1, ex2) -> :(-($ex1,$ex2)), [j == i ? :(b[$j]) : :(a[$i, $j]*$(x[j])) for j = i:-1:1]))/a[$i, $i]) for i = 1:sb[1]]
+
+    quote
+        @_inline_meta
+        T = typeof((zero(Ta)*zero(Tb) + zero(Ta)*zero(Tb))/one(Ta))
+        @inbounds $(Expr(:block, expr...))
+        @inbounds return similar_type(b, T)(tuple($(x...)))
+    end
+end
+
+@generated function solve(::Size{sa}, ::Size{sb}, a::LowerTriangular{Ta, <:StaticMatrix{<:Any, <:Any, Ta}}, b::StaticMatrix{<:Any, <:Any, Tb}) where {sa, sb, Ta, Tb}
+    if sa[1] != sb[1]
+        throw(DimensionMismatch("right hand side b needs first dimension of size $(sa[1]), has size $(sb[1])"))
+    end
+
+    x = [Symbol("x$k1$k2") for k1 = 1:sb[1], k2 = 1:sb[2]]
+    expr = [:($(x[k1, k2]) = $(reduce((ex1, ex2) -> :(-($ex1,$ex2)), [j == k1 ? :(b[$j, $k2]) : :(a[$k1, $j]*$(x[j, k2])) for j = k1:-1:1]))/a[$k1, $k1]) for k1 = 1:sb[1], k2 = 1:sb[2]]
+
+    quote
+        @_inline_meta
+        T = typeof((zero(Ta)*zero(Tb) + zero(Ta)*zero(Tb))/one(Ta))
+        @inbounds $(Expr(:block, expr...))
+        @inbounds return similar_type(b, T)(tuple($(x...)))
+    end
+end

test/solve.jl

@@ -31,3 +31,28 @@
     @test_throws DimensionMismatch m1\v
     @test_throws DimensionMismatch m1\m2
 end
+
+@testset "Solving triangular system" begin
+    for n in (1, 2, 3, 4),
+        (t, uplo) in ((UpperTriangular, :U),
+                      (LowerTriangular, :L)),
+            (m, v, u) in ((SMatrix{n,n}, SVector{n}, SMatrix{n,2}),
+                          (MMatrix{n,n}, MVector{n}, SMatrix{n,2})),
+                eltya in (Float32, Float64, BigFloat, Complex64, Complex128, Complex{BigFloat}, Int),
+                    eltyb in (Float32, Float64, BigFloat, Complex64, Complex128, Complex{BigFloat})
+
+        A = t(eltya == Int ? rand(1:7, n, n) : convert(Matrix{eltya}, (eltya <: Complex ? complex.(randn(n,n), randn(n,n)) : randn(n,n)) |> z -> chol(z'z) |> z -> uplo == :U ? z : ctranspose(z)))
+        b = convert(Matrix{eltyb}, eltya <: Complex ? real(A)*ones(n,2) : A*ones(n,2))
+        SA = t(m(A.data))
+        Sx = SA \ v(b[:, 1])
+        x = A \ b[:, 1]
+        @test Sx isa StaticVector # test not falling back to Base
+        @test Sx ≈ x
+        @test eltype(Sx) == eltype(x)
+        SX = SA \ u(b)
+        X = A \ b
+        @test SX isa StaticMatrix # test not falling back to Base
+        @test SX ≈ X
+        @test eltype(SX) == eltype(X)
+    end
+end