Add simple forward and backward substitution to solve triangular systems

Author: bshall

src/solve.jl

@@ -1,9 +1,13 @@
 @inline (\)(a::StaticMatrix{<:Any, <:Any, T}, b::StaticVector{<:Any, T}) where {T} = solve(Size(a), Size(b), a, b)
+@inline (\)(a::Union{UpperTriangular{T, S}, LowerTriangular{T, S}} where {S<:StaticMatrix{<:Any, <:Any, T}}, b::StaticVector{<:Any, T}) where {T} = solve(Size(a.data), Size(b), a, b)
+@inline (\)(a::Union{UpperTriangular{T, S}, LowerTriangular{T, S}} where {S<:StaticMatrix{<:Any, <:Any, T}}, b::StaticMatrix{<:Any, <:Any, T}) where {T} = 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 +30,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, Sa} where {Sa<: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, Sa} where {Sa<: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, Sa} where {Sa<: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, Sa} where {Sa<: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

@@ -25,3 +25,20 @@
         @test m(A)\v(b) ≈ A\b
     end =#
 end
+
+@testset "Solving triangular system" begin
+    for n in (1,2,3,4),
+          (t1, uplo1) 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})),
+                elty in (Float32, Float64, Int)
+
+        eval(quote
+            A = $(t1)($elty == Int ? rand(1:7, $n, $n) : convert(Matrix{$elty}, randn($n, $n)) |> t -> chol(t't) |> t -> $(uplo1 == :U) ? t : ctranspose(t))
+            b = convert(Matrix{$elty}, A*ones($n, 2))
+            SA = $t1($m(A.data))
+            @test SA \ $v(b[:, 1]) ≈ A \ b[:, 1]
+            @test SA \ $u(b) ≈ A \ b
+        end)
+    end
+end