Allow trans/adj in triangular ldiv! (#383)

* Allow trans/adj in triangular ldiv!

* Update matmul.jl

* Update Project.toml

* Use dualadjoint

* Increase codecov

* banded and dense matrix tests

Author: dlfivefifty

Project.toml

@@ -17,7 +17,7 @@ BandedMatricesSparseArraysExt = "SparseArrays"
 
 [compat]
 Aqua = "0.6"
-ArrayLayouts = "1"
+ArrayLayouts = "1.1"
 Documenter = "0.27"
 FillArrays = "1.3"
 InfiniteArrays = "0.12"

src/BandedMatrices.jl

@@ -32,7 +32,7 @@ import ArrayLayouts: MemoryLayout, transposelayout, triangulardata,
                     sublayout, sub_materialize, _fill_lmul!, _copy_oftype,
                     reflector!, reflectorApply!, _copyto!, checkdimensions,
                     _qr!, _qr, _lu!, _lu, _factorize, AbstractTridiagonalLayout, TridiagonalLayout,
-                    BidiagonalLayout, bidiagonaluplo, diagonaldata, supdiagonaldata, subdiagonaldata
+                    BidiagonalLayout, bidiagonaluplo, diagonaldata, supdiagonaldata, subdiagonaldata, copymutable_oftype_layout, dualadjoint
 
 import FillArrays: AbstractFill, getindex_value, _broadcasted_zeros, unique_value, OneElement
 

src/banded/BandedMatrix.jl

@@ -148,15 +148,16 @@ convert(::Type{DefaultBandedMatrix}, M::AbstractMatrix{T}) where T = convert(Def
 
 copy(B::BandedMatrix) = _BandedMatrix(copy(B.data), B.raxis, B.l, B.u)
 
-if isdefined(LinearAlgebra, :copymutable_oftype)
-    LinearAlgebra.copymutable_oftype(B::BandedMatrix, ::Type{S}) where S =
-        _BandedMatrix(LinearAlgebra.copymutable_oftype(B.data, S), B.raxis, B.l, B.u)
 
-    LinearAlgebra.copymutable_oftype(B::Adjoint{<:Any,<:BandedMatrix}, ::Type{S}) where S =
-        LinearAlgebra.copymutable_oftype(parent(B), S)'
-    LinearAlgebra.copymutable_oftype(B::Transpose{<:Any,<:BandedMatrix}, ::Type{S}) where S =
-        transpose(LinearAlgebra.copymutable_oftype(parent(B), S))
-end
+copymutable_oftype_layout(::BandedColumns, B, ::Type{S}) where S =
+    _BandedMatrix(LinearAlgebra.copymutable_oftype(bandeddata(B), S), axes(B,1), bandwidths(B)...)
+
+copymutable_oftype_layout(::BandedRows, B, ::Type{S}) where S =
+    dualadjoint(B)(LinearAlgebra.copymutable_oftype(parent(B), S))
+
+copymutable_oftype_layout(::AbstractBandedLayout, B, ::Type{S}) where S =
+    copyto!(BandedMatrix{S}(undef, axes(B), bandwidths(B)), B)
+
 
 promote_rule(::Type{BandedMatrix{T1, C1}}, ::Type{BandedMatrix{T2, C2}}) where {T1,C1, T2,C2} =
     BandedMatrix{promote_type(T1,T2), promote_type(C1, C2)}

src/generic/matmul.jl

@@ -239,7 +239,7 @@ end
 
 ### BandedMatrix * dense matrix
 
-function materialize!(M::MulAdd{BandedColumns{DenseColumnMajor}, <:AbstractColumnMajor, <:AbstractColumnMajor,
+function materialize!(M::MulAdd{<:BandedColumns, <:AbstractColumnMajor, <:AbstractColumnMajor,
         T, <:AbstractMatrix, <:AbstractMatrix, <:AbstractMatrix}) where {T}
     α, β, A, B, C = M.α, M.β, M.A, M.B, M.C
 
@@ -258,3 +258,22 @@ function materialize!(M::MulAdd{BandedColumns{DenseColumnMajor}, <:AbstractColum
 
     return C
 end
+
+function materialize!(M::MulAdd{<:AbstractColumnMajor, <:BandedColumns, <:AbstractColumnMajor,
+    T, <:AbstractMatrix, <:AbstractMatrix, <:AbstractMatrix}) where {T}
+    α, β, A, B, C = M.α, M.β, M.A, M.B, M.C
+    mA, nA = size(A)
+    mB, nB = size(B)
+    mC, nC = size(C)
+    (nA == mB && mC == mA && nC == nB) || throw(DimensionMismatch("A has size ($mA, $nA), B has size ($mB, $nB), C has size ($mC, $nC)"))
+
+    if iszero(α)
+        lmul!(β, C)
+    else
+        for (rowC, rowA) in zip(eachrow(C), eachrow(A))
+            mul!(rowC, transpose(B), rowA, α, β)
+        end
+    end
+
+    return C
+end

src/tribanded.jl

@@ -79,17 +79,27 @@ end
 # Ldiv
 for UNIT in ('N', 'U')
     @eval begin
-        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'U',$UNIT,<:BandedColumnMajor},
-                                        <:AbstractStridedLayout})
+        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'U',$UNIT,<:BandedColumnMajor}, <:AbstractStridedLayout})
             A,x = M.A,M.B
             tbsv!('U', 'N', $UNIT, size(A,1), bandwidth(A,2), bandeddata(A), x)
         end
 
-        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'L',$UNIT,<:BandedColumnMajor},
-                                                        <:AbstractStridedLayout})
+        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'L',$UNIT,<:BandedColumnMajor}, <:AbstractStridedLayout})
             A,x = M.A,M.B
             tbsv!('L', 'N', $UNIT, size(A,1), bandwidth(A,1), bandeddata(A), x)
         end
+
+        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'U',$UNIT,<:BandedRowMajor}, <:AbstractStridedLayout})
+            U,x = M.A,M.B
+            A = triangulardata(U)
+            tbsv!('L', 'T', $UNIT, size(A,1), bandwidth(A,2), bandeddata(LowerTriangular(A')), x)
+        end
+
+        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'L',$UNIT,<:BandedRowMajor}, <:AbstractStridedLayout})
+            L,x = M.A,M.B
+            A = triangulardata(L)
+            tbsv!('U', 'T', $UNIT, size(A,1), bandwidth(A,1), bandeddata(UpperTriangular(A')), x)
+        end
     end
     # for UPLO in ('U', 'L')
     #     @eval begin

test/test_banded.jl

@@ -508,13 +508,15 @@ Base.similar(::MyMatrix, ::Type{T}, m::Int, n::Int) where T = MyMatrix{T}(undef,
 
     if isdefined(LinearAlgebra, :copymutable_oftype)
         @testset "copymutable_oftype" begin
-            B = _BandedMatrix((2:3)', 4, -2, 2)
+            B = _BandedMatrix((2:5)', 4, -2, 2)
             @test LinearAlgebra.copymutable_oftype(B, Float64) == B
             @test LinearAlgebra.copymutable_oftype(B, Float64) isa BandedMatrix{Float64}
             @test LinearAlgebra.copymutable_oftype(B', Float64) == B'
             @test LinearAlgebra.copymutable_oftype(B', Float64) isa Adjoint{Float64,<:BandedMatrix{Float64}}
             @test LinearAlgebra.copymutable_oftype(transpose(B), Float64) == transpose(B)
             @test LinearAlgebra.copymutable_oftype(transpose(B), Float64) isa Transpose{Float64,<:BandedMatrix{Float64}}
+
+            @test LinearAlgebra.copymutable_oftype(UpperTriangular(B), Float64) == B
         end
     end
 end

test/test_interface.jl

@@ -205,6 +205,12 @@ LinearAlgebra.fill!(A::PseudoBandedMatrix, v) = fill!(A.data,v)
             bandwidths(convert(BandedMatrix{Float64}, A)) ==
             bandwidths(convert(BandedMatrix{Float64,Matrix{Float64}},A)) ==
             bandwidths(A)
+
+    @testset "copymutable_oftype" begin
+        A = PseudoBandedMatrix(rand(5, 4), 1, 2)
+        @test ArrayLayouts.copymutable_oftype_layout(MemoryLayout(A), A, BigFloat) == A
+        @test ArrayLayouts.copymutable_oftype_layout(MemoryLayout(A), A, BigFloat) isa BandedMatrix
+    end
 end
 
 @testset "Bi/Tri/Diagonal" begin

test/test_linalg.jl

@@ -426,5 +426,12 @@ import BandedMatrices: BandedColumns, _BandedMatrix
             end
         end
     end
+
+    @testset "mul with α = 0" begin
+        A = randn(5,5)
+        B = brand(5,5,1,2)
+        @test muladd!(0.0, A, B, 2.0, copy(A)) == 2A
+        @test muladd!(0.0, B, A, 2.0, copy(A)) == 2A
+    end
 end
 

test/test_tribanded.jl

@@ -73,7 +73,31 @@ import BandedMatrices: BandedColumns, BandedRows
         A = brand(5,5,2,1)
         b = randn(5)
         U = UpperTriangular(A')
+        L = LowerTriangular(A')
+        Ũ = LowerTriangular(A)'
+        L̃ = UpperTriangular(A)'
         @test MemoryLayout(U) isa TriangularLayout{'U','N',BandedRows{DenseColumnMajor}}
-        @test ArrayLayouts.lmul!(U,copy(b)) ≈ U*b
+        @test U*b ≈ ArrayLayouts.lmul!(U,copy(b)) ≈ Matrix(U)*b
+        @test L*b ≈ ArrayLayouts.lmul!(L,copy(b)) ≈ Matrix(L)*b
+        @test ldiv!(U,copy(b)) ≈ ArrayLayouts.ldiv!(U,copy(b)) ≈ U\b ≈ Matrix(U)\b
+        @test ldiv!(L,copy(b)) ≈ ArrayLayouts.ldiv!(L,copy(b)) ≈ L\b ≈ Matrix(L)\b
+        @test ldiv!(Ũ,copy(b)) ≈ ArrayLayouts.ldiv!(Ũ,copy(b)) ≈ Ũ\b ≈ Matrix(Ũ)\b
+        @test ldiv!(L̃,copy(b)) ≈ ArrayLayouts.ldiv!(L̃,copy(b)) ≈ L̃\b ≈ Matrix(L̃)\b
+
+        B = randn(5,5)
+
+        @test U*B ≈ ArrayLayouts.lmul!(U,copy(B)) ≈ Matrix(U)*B
+        @test L*B ≈ ArrayLayouts.lmul!(L,copy(B)) ≈ Matrix(L)*B
+        @test B*U ≈ ArrayLayouts.rmul!(copy(B),U) ≈ B*Matrix(U)
+        @test B*L ≈ ArrayLayouts.rmul!(copy(B),L) ≈ B*Matrix(L)
+
+        @test U \ B ≈ Matrix(U) \ B
+        @test B / U ≈ B / Matrix(U)
+        @test L \ B ≈ Matrix(L) \ B
+        @test B / L ≈ B / Matrix(L)
+        @test Ũ \ B ≈ Matrix(Ũ) \ B
+        @test B / Ũ ≈ B / Matrix(Ũ)
+        @test L̃ \ B ≈ Matrix(L̃) \ B
+        @test B / L̃ ≈ B / Matrix(L̃)
     end
 end