Better Support for getindex for Ldiv (#35)

* Better Support for getindex for Ldiv

* lmul! for LayoutArray

* lmul! for Symmetric/Hermitian

* Test scalar lmul!/rmul!/ldiv!/rdiv!

* Update runtests.jl

* Update runtests.jl

* 3-arg dot, getindex for LayoutArrays

* x'A should behave like a dual-vector

* Update runtests.jl

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.4.2"
+version = "0.4.3"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/ArrayLayouts.jl

@@ -128,6 +128,10 @@ macro layoutgetindex(Typ)
     esc(quote
         ArrayLayouts.@_layoutgetindex $Typ
         ArrayLayouts.@_layoutgetindex LinearAlgebra.AbstractTriangular{<:Any,<:$Typ}
+        ArrayLayouts.@_layoutgetindex LinearAlgebra.Symmetric{<:Any,<:$Typ}
+        ArrayLayouts.@_layoutgetindex LinearAlgebra.Hermitian{<:Any,<:$Typ}
+        ArrayLayouts.@_layoutgetindex LinearAlgebra.Adjoint{<:Any,<:$Typ}
+        ArrayLayouts.@_layoutgetindex LinearAlgebra.Transpose{<:Any,<:$Typ}
     end)
 end
 
@@ -144,6 +148,15 @@ end
 
 @layoutmatrix LayoutMatrix
 
+for Typ in (:LayoutArray, :(Transpose{<:Any,<:LayoutMatrix}), :(Adjoint{<:Any,<:LayoutMatrix}), :(Symmetric{<:Any,<:LayoutMatrix}), :(Hermitian{<:Any,<:LayoutMatrix}))
+    @eval begin
+        LinearAlgebra.lmul!(α::Number, A::$Typ) = lmul!(α, A)
+        LinearAlgebra.rmul!(A::$Typ, α::Number) = rmul!(A, α)
+        LinearAlgebra.ldiv!(α::Number, A::$Typ) = ldiv!(α, A)
+        LinearAlgebra.rdiv!(A::$Typ, α::Number) = rdiv!(A, α)
+    end
+end
+
 getindex(A::LayoutVector, kr::AbstractVector) = layout_getindex(A, kr)
 
 _copyto!(_, _, dest::AbstractArray{T,N}, src::AbstractArray{V,N}) where {T,V,N} =

src/ldiv.jl

@@ -45,8 +45,12 @@ end
 @inline eltype(M::Rdiv) = promote_type(eltype(M.A), Base.promote_op(inv, eltype(M.B)))
 
 # Lazy getindex
-getindex(L::Ldiv{<:Any,<:Any,<:AbstractMatrix,<:AbstractVector}, k::Integer) = copyto!(similar(L), L)[k]
-getindex(L::Ldiv{<:Any,<:Any,<:AbstractMatrix,<:AbstractMatrix}, k::Integer,j::Integer) = Ldiv(L.A, L.B[:,j])[k]
+
+getindex(L::Ldiv, k...) = _getindex(indextype(L), L, k)
+concretize(L::AbstractArray) = convert(Array,L)
+concretize(L::Ldiv) = ldiv(concretize(L.A), concretize(L.B))
+_getindex(::Type{Tuple{I}}, L::Ldiv, (k,)::Tuple{I}) where I = concretize(L)[k]
+_getindex(::Type{Tuple{I,J}}, L::Ldiv, (k,j)::Tuple{I,J}) where {I,J} = Ldiv(L.A, L.B[:,j])[k]
 
 check_ldiv_axes(A, B) =
     axes(A,1) == axes(B,1) || throw(DimensionMismatch("First axis of A, $(axes(A,1)), and first axis of B, $(axes(B,1)) must match"))
@@ -105,10 +109,25 @@ const BlasMatLdivMat{styleA, styleB, T<:BlasFloat} = MatLdivMat{styleA, styleB,
 const MatRdivMat{styleA, styleB, T, V} = Rdiv{styleA, styleB, <:AbstractMatrix{T}, <:AbstractMatrix{V}}
 const BlasMatRdivMat{styleA, styleB, T<:BlasFloat} = MatRdivMat{styleA, styleB, T, T}
 
-# function materialize!(L::BlasMatLdivVec{<:AbstractColumnMajor,<:AbstractColumnMajor})
-
-# end
+materialize!(M::Ldiv{ScalarLayout}) = Base.invoke(LinearAlgebra.ldiv!, Tuple{Number,AbstractArray}, M.A, M.B)
+materialize!(M::Rdiv{<:Any,ScalarLayout}) = Base.invoke(LinearAlgebra.rdiv!, Tuple{AbstractArray,Number}, M.A, M.B)
 
+function materialize!(M::Ldiv{ScalarLayout,<:SymmetricLayout})
+    ldiv!(M.A, symmetricdata(M.B))
+    M.B
+end
+function materialize!(M::Ldiv{ScalarLayout,<:HermitianLayout})
+    ldiv!(M.A, hermitiandata(M.B))
+    M.B
+end
+function materialize!(M::Rdiv{<:SymmetricLayout,ScalarLayout})
+    rdiv!(symmetricdata(M.A), M.B)
+    M.A
+end
+function materialize!(M::Rdiv{<:HermitianLayout,ScalarLayout})
+    rdiv!(hermitiandata(M.A), M.B)
+    M.A
+end
 
 macro _layoutldiv(Typ)
     ret = quote

src/lmul.jl

@@ -73,7 +73,22 @@ materialize!(M::Rmul) = LinearAlgebra.rmul!(M.A,M.B)
 materialize!(M::Lmul{ScalarLayout}) = Base.invoke(LinearAlgebra.lmul!, Tuple{Number,AbstractArray}, M.A, M.B)
 materialize!(M::Rmul{<:Any,ScalarLayout}) = Base.invoke(LinearAlgebra.rmul!, Tuple{AbstractArray,Number}, M.A, M.B)
 
-
+function materialize!(M::Lmul{ScalarLayout,<:SymmetricLayout})
+    lmul!(M.A, symmetricdata(M.B))
+    M.B
+end
+function materialize!(M::Lmul{ScalarLayout,<:HermitianLayout})
+    lmul!(M.A, hermitiandata(M.B))
+    M.B
+end
+function materialize!(M::Rmul{<:SymmetricLayout,ScalarLayout})
+    rmul!(symmetricdata(M.A), M.B)
+    M.A
+end
+function materialize!(M::Rmul{<:HermitianLayout,ScalarLayout})
+    rmul!(hermitiandata(M.A), M.B)
+    M.A
+end
 
 macro _layoutlmul(Typ)
     esc(quote

src/mul.jl

@@ -46,10 +46,10 @@ function _getindex(::Type{Tuple{AA,BB}}, M::Mul, (k, j)::Tuple{AA,BB}) where {AA
 end
 
 # linear indexing
-_getindex(::Type{NTuple{2,Int}}, M::Mul, k::Tuple{Int}) = M[Base._ind2sub(axes(M), k...)...]
+_getindex(::Type{NTuple{2,Int}}, M, k::Tuple{Int}) = M[Base._ind2sub(axes(M), k...)...]
 
-_getindex(::Type{Tuple{Int}}, M::Mul, (k,)::Tuple{CartesianIndex{1}}) = M[convert(Int, k)]
-_getindex(::Type{NTuple{2,Int}}, M::Mul, (kj,)::Tuple{CartesianIndex{2}}) = M[kj[1], kj[2]]
+_getindex(::Type{Tuple{Int}}, M, (k,)::Tuple{CartesianIndex{1}}) = M[convert(Int, k)]
+_getindex(::Type{NTuple{2,Int}}, M, (kj,)::Tuple{CartesianIndex{2}}) = M[kj[1], kj[2]]
 
 """
    indextype(A)
@@ -245,3 +245,7 @@ dot(a, b) = materialize(Dot(a, b))
 @inline LinearAlgebra.dot(a::LayoutArray, b::SubArray{<:Any,N,<:LayoutArray}) where N = dot(a,b)
 @inline LinearAlgebra.dot(a::SubArray{<:Any,N,<:LayoutArray}, b::SubArray{<:Any,N,<:LayoutArray}) where N = dot(a,b)
 
+# Temporary until layout 3-arg dot is added. 
+# We go to generic fallback as layout-arrays are structured
+dot(x, A, y) = dot(x, mul(A, y))
+LinearAlgebra.dot(x::AbstractVector, A::LayoutMatrix, y::AbstractVector) = dot(x, A, y)

src/muladd.jl

@@ -417,4 +417,13 @@ fillzeros(::Type{T}, ax) where T = Zeros{T}(ax)
 # Fill 
 ###
 
-copy(M::MulAdd{<:AbstractFillLayout,<:AbstractFillLayout,<:AbstractFillLayout}) = M.α*M.A*M.B + M.β*M.C
+copy(M::MulAdd{<:AbstractFillLayout,<:AbstractFillLayout,<:AbstractFillLayout}) = M.α*M.A*M.B + M.β*M.C
+
+###
+# DualLayout
+###
+
+function similar(M::MulAdd{<:DualLayout,<:Any,ZerosLayout}, ::Type{T}, (x,y)) where T
+    @assert x ≡ Base.OneTo(1)
+    similar(M.A', T, y)'
+end

test/runtests.jl

@@ -126,11 +126,36 @@ MemoryLayout(::Type{MyVector}) = DenseColumnMajor()
             @test A\MyVector(x) ≈ A\x
             @test A\MyMatrix(X) ≈ A\X
         end
+
+        @testset "dot" begin
+            A = MyMatrix(randn(5,5))
+            b = randn(5)
+            @test dot(b, A, b) ≈ b'*(A*b) ≈ b'A*b
+        end
     end
 
     @testset "l/rmul!" begin
-        b = randn(5)
-        @test ArrayLayouts.lmul!(2, MyVector(copy(b))) == ArrayLayouts.rmul!(MyVector(copy(b)), 2) == 2b
+        b = MyVector(randn(5))
+        A = MyMatrix(randn(5,5))
+        @test lmul!(2, deepcopy(b)) == rmul!(deepcopy(b), 2) == 2b
+        @test lmul!(2, deepcopy(A)) == rmul!(deepcopy(A), 2) == 2A
+        @test lmul!(2, deepcopy(A)') == rmul!(deepcopy(A)', 2) == 2A'
+        @test lmul!(2, transpose(deepcopy(A))) == rmul!(transpose(deepcopy(A)), 2) == 2transpose(A)
+        @test lmul!(2, Symmetric(deepcopy(A))) == rmul!(Symmetric(deepcopy(A)), 2) == 2Symmetric(A)
+        @test lmul!(2, Hermitian(deepcopy(A))) == rmul!(Hermitian(deepcopy(A)), 2) == 2Hermitian(A)
+
+        C = randn(ComplexF64,5,5)
+        @test ArrayLayouts.lmul!(2, Hermitian(copy(C))) == ArrayLayouts.rmul!(Hermitian(copy(C)), 2) == 2Hermitian(C)
+
+        if VERSION ≥ v"1.5"
+            @test ldiv!(2, deepcopy(b)) == rdiv!(deepcopy(b), 2) == 2\b
+            @test ldiv!(2, deepcopy(A)) == rdiv!(deepcopy(A), 2) == 2\A
+            @test ldiv!(2, deepcopy(A)') == rdiv!(deepcopy(A)', 2) == 2\A'
+            @test ldiv!(2, transpose(deepcopy(A))) == rdiv!(transpose(deepcopy(A)), 2) == 2\transpose(A)
+            @test ldiv!(2, Symmetric(deepcopy(A))) == rdiv!(Symmetric(deepcopy(A)), 2) == 2\Symmetric(A)
+            @test ldiv!(2, Hermitian(deepcopy(A))) == rdiv!(Hermitian(deepcopy(A)), 2) == 2\Hermitian(A)
+            @test ArrayLayouts.ldiv!(2, Hermitian(copy(C))) == ArrayLayouts.rdiv!(Hermitian(copy(C)), 2) == 2\Hermitian(C)
+        end
     end
 end